add SUPPORT_MOE_MARLIN_W16A16 to use moe marlin on bw

vllm/platforms/rocm.py

@@ -16,8 +16,12 @@ from vllm.utils import cuda_device_count_stateless
 
 from .interface import DeviceCapability, Platform, PlatformEnum, _Backend
 
-from vllm.utils import SUPPORT_TC
+from vllm.utils import SUPPORT_TC, SUPPORT_MOE_MARLIN_W16A16
 
+if SUPPORT_MOE_MARLIN_W16A16:
+    os.environ['VLLM_USE_MARLIN_W16A16_MOE'] = '1'
+    os.environ['MOE_NN'] = '0'
+    
 if not SUPPORT_TC:
     os.environ['VLLM_USE_V1'] = '0'
     os.environ['VLLM_USE_FLASH_ATTN_PA'] = '0'

vllm/utils/__init__.py

@@ -87,6 +87,7 @@ MULTIMODAL_MODEL_MAX_NUM_BATCHED_TOKENS = 5120
 
 GPU_ARCH = torch.cuda.get_device_properties("cuda").gcnArchName
 SUPPORT_TC = any(arch in GPU_ARCH for arch in ["gfx928", "gfx936", "gfx938"])
+SUPPORT_MOE_MARLIN_W16A16 = any(arch in GPU_ARCH for arch in ["gfx936"])
 
 def _generate_random_int8(
     tensor: torch.Tensor,