| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | import pytest |
| 4 | import torch |
| 5 | |
| 6 | import vllm.model_executor.layers.fused_moe.modular_kernel as mk |
| 7 | from tests.kernels.moe.utils import make_dummy_moe_config, make_test_weights |
| 8 | from tests.kernels.quantization.nvfp4_utils import ( |
| 9 | FLOAT4_E2M1_MAX, |
| 10 | FLOAT8_E4M3_MAX, |
| 11 | dequantize_nvfp4_to_dtype, |
| 12 | ) |
| 13 | from tests.kernels.utils import torch_moe |
| 14 | from vllm import _custom_ops as ops |
| 15 | from vllm.config import ParallelConfig, VllmConfig, set_current_vllm_config |
| 16 | from vllm.model_executor.layers.fused_moe import fused_topk |
| 17 | from vllm.model_executor.layers.fused_moe.activation import MoEActivation |
| 18 | from vllm.model_executor.layers.fused_moe.all2all_utils import ( |
| 19 | maybe_make_prepare_finalize, |
| 20 | ) |
| 21 | from vllm.model_executor.layers.fused_moe.config import nvfp4_moe_quant_config |
| 22 | from vllm.model_executor.layers.fused_moe.experts.cutlass_moe import ( |
| 23 | CutlassExpertsFp4, |
| 24 | ) |
| 25 | from vllm.model_executor.layers.fused_moe.prepare_finalize import ( |
| 26 | make_moe_prepare_and_finalize_no_dp_ep, |
| 27 | ) |
| 28 | from vllm.platforms import current_platform |
| 29 | from vllm.utils.torch_utils import set_random_seed |
| 30 | |
| 31 | if not current_platform.has_device_capability(100): |
| 32 | pytest.skip( |
| 33 | "Nvfp4 Requires compute capability of 10 or above.", allow_module_level=True |
| 34 | ) |
| 35 | |
| 36 | MNK_FACTORS = [ |
| 37 | (2, 1024, 1024), |
| 38 | (2, 1024, 1536), |
| 39 | (2, 3072, 1024), |
| 40 | (64, 1024, 1024), |
| 41 | (64, 3072, 1024), |
| 42 | (64, 2048, 1536), |
| 43 | (224, 1024, 1024), |
| 44 | (224, 1024, 1536), |
| 45 | ] |
| 46 | |
| 47 | |
| 48 | @pytest.mark.parametrize("m,n,k", MNK_FACTORS) |
| 49 | @pytest.mark.parametrize("e", [40, 64, 256]) |
| 50 | @pytest.mark.parametrize("topk", [1, 6, 8]) |
| 51 | @pytest.mark.parametrize("dtype", [torch.bfloat16]) |
| 52 | @torch.inference_mode() |
| 53 | def test_cutlass_fp4_moe_no_graph( |
| 54 | m: int, n: int, k: int, e: int, topk: int, dtype: torch.dtype, workspace_init |
| 55 | ): |
| 56 | set_random_seed(7) |
| 57 | with set_current_vllm_config( |
| 58 | VllmConfig(parallel_config=ParallelConfig(pipeline_parallel_size=1)) |
| 59 | ): |
| 60 | quant_blocksize = 16 |
| 61 | |
| 62 | a = torch.randn((m, k), device="cuda", dtype=dtype) / 10 |
| 63 | |
| 64 | (_, w1_q, w1_blockscale, w1_gs), (_, w2_q, w2_blockscale, w2_gs) = ( |
| 65 | make_test_weights( |
| 66 | e, |
| 67 | n, |
| 68 | k, |
| 69 | in_dtype=dtype, |
| 70 | quant_dtype="nvfp4", |
| 71 | block_shape=None, # use quant_blocksize? |
| 72 | per_out_ch_quant=False, |
| 73 | ) |
| 74 | ) |
| 75 | |
| 76 | score = torch.randn((m, e), device="cuda", dtype=dtype) |
| 77 | topk_weights, topk_ids, _ = fused_topk(a, score, topk, renormalize=False) |
| 78 | |
| 79 | a1_gs = torch.ones((e,), device="cuda", dtype=torch.float32) |
| 80 | a2_gs = torch.ones((e,), device="cuda", dtype=torch.float32) |
| 81 | |
| 82 | assert w1_gs is not None |
| 83 | assert w2_gs is not None |
| 84 | assert w1_blockscale is not None |
| 85 | assert w2_blockscale is not None |
| 86 | |
| 87 | quant_config = nvfp4_moe_quant_config( |
| 88 | g1_alphas=(1 / w1_gs), |
| 89 | g2_alphas=(1 / w2_gs), |
| 90 | a1_gscale=a1_gs, |
| 91 | a2_gscale=a2_gs, |
| 92 | w1_scale=w1_blockscale, |
| 93 | w2_scale=w2_blockscale, |
| 94 | ) |
| 95 | moe_config = make_dummy_moe_config() |
| 96 | |
| 97 | kernel = mk.FusedMoEKernel( |
| 98 | maybe_make_prepare_finalize( |
| 99 | moe=moe_config, |
| 100 | quant_config=quant_config, |
| 101 | allow_new_interface=True, |
| 102 | use_monolithic=False, |
| 103 | ), |
| 104 | CutlassExpertsFp4( |
| 105 | moe_config=moe_config, |
| 106 | quant_config=quant_config, |
| 107 | ), |
| 108 | ) |
| 109 | |
| 110 | cutlass_output = kernel.apply( |
| 111 | hidden_states=a, |
| 112 | w1=w1_q, |
| 113 | w2=w2_q, |
| 114 | topk_weights=topk_weights, |
| 115 | topk_ids=topk_ids, |
| 116 | global_num_experts=e, |
| 117 | activation=mk.MoEActivation.SILU, |
| 118 | apply_router_weight_on_input=False, |
| 119 | expert_map=None, |
| 120 | ) |
| 121 | |
| 122 | # Reference check: |
| 123 | a_global_scale = ( |
| 124 | (FLOAT8_E4M3_MAX * FLOAT4_E2M1_MAX) / torch.amax(a.flatten(), dim=-1) |
| 125 | ).to(torch.float32) |
| 126 | a_fp4, a_scale_interleaved = ops.scaled_fp4_quant(a, a_global_scale) |
| 127 | |
| 128 | a_in_dtype = dequantize_nvfp4_to_dtype( |
| 129 | a_fp4, |
| 130 | a_scale_interleaved, |
| 131 | a_global_scale, |
| 132 | dtype=a.dtype, |
| 133 | device=a.device, |
| 134 | block_size=quant_blocksize, |
| 135 | ) |
| 136 | |
| 137 | w1_d = torch.empty((e, 2 * n, k), device="cuda", dtype=dtype) |
| 138 | w2_d = torch.empty((e, k, n), device="cuda", dtype=dtype) |
| 139 | |
| 140 | for idx in range(0, e): |
| 141 | w1_d[idx] = dequantize_nvfp4_to_dtype( |
| 142 | w1_q[idx], |
| 143 | w1_blockscale[idx], |
| 144 | w1_gs[idx], |
| 145 | dtype=dtype, |
| 146 | device=w1_q.device, |
| 147 | block_size=quant_blocksize, |
| 148 | ) |
| 149 | w2_d[idx] = dequantize_nvfp4_to_dtype( |
| 150 | w2_q[idx], |
| 151 | w2_blockscale[idx], |
| 152 | w2_gs[idx], |
| 153 | dtype=dtype, |
| 154 | device=w2_q.device, |
| 155 | block_size=quant_blocksize, |
| 156 | ) |
| 157 | |
| 158 | torch_output = torch_moe(a_in_dtype, w1_d, w2_d, score, topk) |
| 159 | |
| 160 | torch.testing.assert_close(torch_output, cutlass_output, atol=1e-1, rtol=1e-1) |
| 161 | |
| 162 | |
| 163 | # step3.5-flash uses swiglustep activation (clipped SwiGLU with limit=7.0) |
| 164 | # for MoE layers 43-44. This tests the non-fused activation fallback path |
| 165 | # in run_cutlass_moe_fp4 (apply_moe_activation + separate fp4 quantization). |
| 166 | # Model dims: e=288, topk=8, n=1280 (moe_intermediate_size), k=4096 (hidden) |
| 167 | SWIGLUSTEP_MNK_FACTORS = [ |
| 168 | (2, 1280, 4096), |
| 169 | (64, 1280, 4096), |
| 170 | (224, 1280, 4096), |
| 171 | ] |
| 172 | |
| 173 | |
| 174 | @pytest.mark.parametrize("m,n,k", SWIGLUSTEP_MNK_FACTORS) |
| 175 | @pytest.mark.parametrize("e", [64, 288]) |
| 176 | @pytest.mark.parametrize("topk", [1, 8]) |
| 177 | @pytest.mark.parametrize("dtype", [torch.bfloat16]) |
| 178 | @torch.inference_mode() |
| 179 | def test_cutlass_fp4_moe_swiglustep( |
| 180 | m: int, n: int, k: int, e: int, topk: int, dtype: torch.dtype, workspace_init |
| 181 | ): |
| 182 | set_random_seed(7) |
| 183 | with set_current_vllm_config( |
| 184 | VllmConfig(parallel_config=ParallelConfig(pipeline_parallel_size=1)) |
| 185 | ): |
| 186 | quant_blocksize = 16 |
| 187 | |
| 188 | a = torch.randn((m, k), device="cuda", dtype=dtype) / 10 |
| 189 | |
| 190 | (_, w1_q, w1_blockscale, w1_gs), (_, w2_q, w2_blockscale, w2_gs) = ( |
| 191 | make_test_weights( |
| 192 | e, |
| 193 | n, |
| 194 | k, |
| 195 | in_dtype=dtype, |
| 196 | quant_dtype="nvfp4", |
| 197 | block_shape=None, |
| 198 | per_out_ch_quant=False, |
| 199 | ) |
| 200 | ) |
| 201 | |
| 202 | score = torch.randn((m, e), device="cuda", dtype=dtype) |
| 203 | topk_weights, topk_ids, _ = fused_topk(a, score, topk, renormalize=False) |
| 204 | |
| 205 | a1_gs = torch.ones((e,), device="cuda", dtype=torch.float32) |
| 206 | a2_gs = torch.ones((e,), device="cuda", dtype=torch.float32) |
| 207 | |
| 208 | assert w1_gs is not None |
| 209 | assert w2_gs is not None |
| 210 | assert w1_blockscale is not None |
| 211 | assert w2_blockscale is not None |
| 212 | |
| 213 | quant_config = nvfp4_moe_quant_config( |
| 214 | g1_alphas=(1 / w1_gs), |
| 215 | g2_alphas=(1 / w2_gs), |
| 216 | a1_gscale=a1_gs, |
| 217 | a2_gscale=a2_gs, |
| 218 | w1_scale=w1_blockscale, |
| 219 | w2_scale=w2_blockscale, |
| 220 | ) |
| 221 | |
| 222 | kernel = mk.FusedMoEKernel( |
| 223 | make_moe_prepare_and_finalize_no_dp_ep(use_monolithic=False), |
| 224 | CutlassExpertsFp4( |
| 225 | moe_config=make_dummy_moe_config(), |
| 226 | quant_config=quant_config, |
| 227 | ), |
| 228 | ) |
| 229 | |
| 230 | cutlass_output = kernel.apply( |
| 231 | hidden_states=a, |
| 232 | w1=w1_q, |
| 233 | w2=w2_q, |
| 234 | topk_weights=topk_weights, |
| 235 | topk_ids=topk_ids, |
| 236 | activation=MoEActivation.SWIGLUSTEP, |
| 237 | global_num_experts=e, |
| 238 | expert_map=None, |
| 239 | apply_router_weight_on_input=False, |
| 240 | ) |
| 241 | |
| 242 | # Reference: dequantize everything and run torch_moe with swiglustep |
| 243 | a_global_scale = ( |
| 244 | (FLOAT8_E4M3_MAX * FLOAT4_E2M1_MAX) / torch.amax(a.flatten(), dim=-1) |
| 245 | ).to(torch.float32) |
| 246 | a_fp4, a_scale_interleaved = ops.scaled_fp4_quant(a, a_global_scale) |
| 247 | |
| 248 | a_in_dtype = dequantize_nvfp4_to_dtype( |
| 249 | a_fp4, |
| 250 | a_scale_interleaved, |
| 251 | a_global_scale, |
| 252 | dtype=a.dtype, |
| 253 | device=a.device, |
| 254 | block_size=quant_blocksize, |
| 255 | ) |
| 256 | |
| 257 | w1_d = torch.empty((e, 2 * n, k), device="cuda", dtype=dtype) |
| 258 | w2_d = torch.empty((e, k, n), device="cuda", dtype=dtype) |
| 259 | |
| 260 | for idx in range(0, e): |
| 261 | w1_d[idx] = dequantize_nvfp4_to_dtype( |
| 262 | w1_q[idx], |
| 263 | w1_blockscale[idx], |
| 264 | w1_gs[idx], |
| 265 | dtype=dtype, |
| 266 | device=w1_q.device, |
| 267 | block_size=quant_blocksize, |
| 268 | ) |
| 269 | w2_d[idx] = dequantize_nvfp4_to_dtype( |
| 270 | w2_q[idx], |
| 271 | w2_blockscale[idx], |
| 272 | w2_gs[idx], |
| 273 | dtype=dtype, |
| 274 | device=w2_q.device, |
| 275 | block_size=quant_blocksize, |
| 276 | ) |
| 277 | |
| 278 | torch_output = torch_moe( |
| 279 | a_in_dtype, |
| 280 | w1_d, |
| 281 | w2_d, |
| 282 | score, |
| 283 | topk, |
| 284 | activation=MoEActivation.SWIGLUSTEP, |
| 285 | ) |
| 286 | |
| 287 | torch.testing.assert_close(torch_output, cutlass_output, atol=1e-1, rtol=1e-1) |
| 288 | |
| 289 | |
| 290 | if __name__ == "__main__": |
| 291 | test_cutlass_fp4_moe_no_graph((2, 1024, 1024), 40, 1, torch.half) |