update test_moe.py

set USE_FUSED_RMS_QUANT=1 and USE_FUSED_SILU_MUL_QUANT=1

tests/kernels/moe/test_moe.py

@@ -114,9 +114,9 @@ def run_moe_test(
     return baseline_output
 
 
-@pytest.mark.parametrize("m", [1, 33, 64, 222, 32768, 40000])
+@pytest.mark.parametrize("m", [1, 33, 64, 32768, 40000])
 @pytest.mark.parametrize("n", [128, 1024, 2048])
-@pytest.mark.parametrize("k", [128, 511, 1024])
+@pytest.mark.parametrize("k", [128, 512, 1024])
 @pytest.mark.parametrize("e", NUM_EXPERTS)
 @pytest.mark.parametrize("topk", TOP_KS)
 @pytest.mark.parametrize("ep_size", EP_SIZE)
@@ -734,7 +734,7 @@ def test_moe_align_block_size_opcheck():
 
 @pytest.mark.parametrize("m", [1, 33, 64, 222])
 @pytest.mark.parametrize("topk", TOP_KS)
-@pytest.mark.parametrize("k", [128, 512, 1024])
+@pytest.mark.parametrize("k", [128, 511, 1024])
 @pytest.mark.parametrize("dtype",
                          [torch.float32, torch.float16, torch.bfloat16])
 @pytest.mark.skipif(current_platform.is_rocm(), reason="Skip for rocm")

vllm/model_executor/model_loader/utils.py

@@ -260,6 +260,10 @@ def get_model_architecture(
                     os.environ['VLLM_USE_LIGHTOP'] = '1'
                 if not envs.is_set("VLLM_USE_OPT_CAT"):
                     os.environ['VLLM_USE_OPT_CAT'] = '1'
+                if not envs.is_set("USE_FUSED_RMS_QUANT"):
+                    os.environ['USE_FUSED_RMS_QUANT'] = '1'
+                if not envs.is_set("USE_FUSED_SILU_MUL_QUANT"):
+                    os.environ['USE_FUSED_SILU_MUL_QUANT'] = '1'
         # awq相关配置
         try:
             if os.getenv('AWQ_MOE_SZ') == None: