diff --git a/python/sglang/srt/layers/moe/token_dispatcher/deepep.py b/python/sglang/srt/layers/moe/token_dispatcher/deepep.py index b3fb47d3f..17d5bf1a8 100644 --- a/python/sglang/srt/layers/moe/token_dispatcher/deepep.py +++ b/python/sglang/srt/layers/moe/token_dispatcher/deepep.py @@ -620,9 +620,25 @@ class _DeepEPDispatcherImplLowLatency(_DeepEPDispatcherImplBase): input_global_scale = self.quant_config.get("input_global_scale", None) if input_global_scale is not None: use_nvfp4 = True - else: + elif not get_moe_runner_backend().is_flashinfer_cutedsl(): + # flashinfer_cutedsl expects BF16 dispatch when NVFP4 dispatch is + # off; its kernel quantizes to NVFP4 internally. use_fp8 = True + # round_scale / use_ue8m0 are FP8-DeepGEMM specific; they cause DeepEP + # to return int32-packed UE8M0 scales that don't feed the flashinfer + # cutedsl kernel. + fp8_deepgemm_scale_opts = ( + dict( + round_scale=deep_gemm_wrapper.ENABLE_JIT_DEEPGEMM + and deep_gemm_wrapper.DEEPGEMM_BLACKWELL, + use_ue8m0=deep_gemm_wrapper.ENABLE_JIT_DEEPGEMM + and deep_gemm_wrapper.DEEPGEMM_BLACKWELL, + ) + if use_fp8 + else dict() + ) + buffer = self._get_buffer() _deepep_precompile_tp_barrier() packed_recv_hidden, self.packed_recv_count, self.handle, event, hook = ( @@ -640,10 +656,7 @@ class _DeepEPDispatcherImplLowLatency(_DeepEPDispatcherImplBase): ), async_finish=not self.return_recv_hook, return_recv_hook=self.return_recv_hook, - round_scale=deep_gemm_wrapper.ENABLE_JIT_DEEPGEMM - and deep_gemm_wrapper.DEEPGEMM_BLACKWELL, - use_ue8m0=deep_gemm_wrapper.ENABLE_JIT_DEEPGEMM - and deep_gemm_wrapper.DEEPGEMM_BLACKWELL, + **fp8_deepgemm_scale_opts, ) ) return packed_recv_hidden, self.packed_recv_count, event, hook diff --git a/python/sglang/srt/layers/moe/utils.py b/python/sglang/srt/layers/moe/utils.py index 9908d4005..61ba72177 100644 --- a/python/sglang/srt/layers/moe/utils.py +++ b/python/sglang/srt/layers/moe/utils.py @@ -263,6 +263,14 @@ def is_deepep_class_backend() -> bool: return b.is_deepep() or b.is_mooncake() or b.is_mori() +def is_flashinfer_cutedsl_v1_path() -> bool: + """CuteDSL v1 + DeepEP low-latency path (no MoeRunner, no autotune).""" + return ( + get_moe_runner_backend().is_flashinfer_cutedsl() + and get_moe_a2a_backend().is_deepep() + ) + + def get_tbo_token_distribution_threshold() -> float: global TBO_TOKEN_DISTRIBUTION_THRESHOLD if TBO_TOKEN_DISTRIBUTION_THRESHOLD is None: diff --git a/python/sglang/srt/layers/quantization/modelopt_quant.py b/python/sglang/srt/layers/quantization/modelopt_quant.py index c24185bdf..42753f31c 100755 --- a/python/sglang/srt/layers/quantization/modelopt_quant.py +++ b/python/sglang/srt/layers/quantization/modelopt_quant.py @@ -24,7 +24,10 @@ from sglang.srt.layers.moe import ( ) from sglang.srt.layers.moe.cutlass_moe_params import CutlassMoEParams, CutlassMoEType from sglang.srt.layers.moe.moe_runner.triton import TritonMoeQuantInfo -from sglang.srt.layers.moe.utils import should_use_flashinfer_cutlass_moe_fp4_allgather +from sglang.srt.layers.moe.utils import ( + is_flashinfer_cutedsl_v1_path, + should_use_flashinfer_cutlass_moe_fp4_allgather, +) from sglang.srt.layers.parameter import ModelWeightParameter, PerTensorScaleParameter from sglang.srt.layers.quantization.base_config import ( FusedMoEMethodBase, @@ -1546,6 +1549,29 @@ class ModelOptNvFp4FusedMoEMethod(FusedMoEMethodBase): return get_moe_runner_backend().is_flashinfer_cutedsl() + # ----- CuteDSL v1 vs v2 path helpers ----- + # + # "v1": cutedsl + deepep low-latency. + # - Bypasses MoeRunner entirely; calls apply_without_routing_weights -> + # flashinfer_cutedsl_moe_masked (grouped_gemm_nt_masked). + # - Expects W13 in default [Gate, Up] order, NOT interleaved. + # - Uses swizzled blockscales directly (w13_blockscale_swizzled). + # + # "v2" (standard): cutedsl + none/flashinfer a2a. + # - Uses MoeRunner with @register_fused_func CuteDslMoEWrapper kernels. + # - Expects W13 in [Up, Gate] order, interleaved in 64-row chunks. + # - Uses MMA-layout blockscales (w13_blockscale_mma). + + @property + def _is_cutedsl_v1_deepep(self) -> bool: + """CuteDSL v1 + DeepEP low-latency path (no MoeRunner).""" + return is_flashinfer_cutedsl_v1_path() + + @property + def _is_cutedsl_v2_standard(self) -> bool: + """New CuteDSL standard path (a2a=none or flashinfer, uses MoeRunner).""" + return self.enable_flashinfer_cutedsl_moe and not self._is_cutedsl_v1_deepep + def create_weights( self, layer: torch.nn.Module, @@ -1812,9 +1838,11 @@ class ModelOptNvFp4FusedMoEMethod(FusedMoEMethodBase): else: # CUTLASS processing - handle w13 and w2 separately - if self.enable_flashinfer_cutedsl_moe and layer.moe_runner_config.is_gated: - # For the CuteDSL FP4 path, interleave the two logical W13 halves - # in 64-row chunks before swizzling the block-scales. + if self._is_cutedsl_v2_standard and layer.moe_runner_config.is_gated: + # CuteDSL v2 only: interleave the two logical W13 halves in + # 64-row chunks for the fused SwiGLU GEMM1 layout expected by + # CuteDslMoEWrapper. The v1 (deepep) path uses + # grouped_gemm_nt_masked which expects plain contiguous halves. from sglang.srt.layers.moe.moe_runner.flashinfer_cutedsl import ( interleave_w13_halves, ) @@ -1876,8 +1904,10 @@ class ModelOptNvFp4FusedMoEMethod(FusedMoEMethodBase): layer, "w2_blockscale_swizzled", w2_blockscale_swizzled ) - if self.enable_flashinfer_cutedsl_moe: - # CuteDSL expects MMA layout for weight scales. Convert from swizzled bytes. + if self._is_cutedsl_v2_standard: + # CuteDSL v2 only: convert blockscales to MMA layout for + # CuteDslMoEWrapper. The v1 (deepep) path uses the + # swizzled blockscales directly via flashinfer_cutedsl_moe_masked. from flashinfer.cute_dsl.utils import convert_sf_to_mma_layout from sglang.srt.layers.moe.moe_runner.flashinfer_cutedsl import ( @@ -1940,9 +1970,10 @@ class ModelOptNvFp4FusedMoEMethod(FusedMoEMethodBase): @property def load_up_proj_weight_first(self) -> bool: - # Load W13 as [Up, Gate] for FlashInfer CUTLASS/CuteDSL kernels. + # Load W13 as [Up, Gate] for FlashInfer CUTLASS and CuteDSL v2 kernels. + # The CuteDSL v1 (deepep) path uses [Gate, Up] -- do NOT flip. return self.moe_runner_config.is_gated and ( - self.enable_flashinfer_cutlass_moe or self.enable_flashinfer_cutedsl_moe + self.enable_flashinfer_cutlass_moe or self._is_cutedsl_v2_standard ) def create_moe_runner( @@ -1959,6 +1990,11 @@ class ModelOptNvFp4FusedMoEMethod(FusedMoEMethodBase): if moe_runner_backend.is_flashinfer_cutedsl(): import sglang.srt.layers.moe.moe_runner.flashinfer_cutedsl # noqa: F401 – triggers @register_fused_func + # CuteDSL v1 (deepep) uses the apply_without_routing_weights + # path (flashinfer_cutedsl_moe_masked) and does not need a MoeRunner. + if self._is_cutedsl_v1_deepep: + return + if not moe_runner_backend.is_flashinfer_cutlass(): self.runner = MoeRunner(moe_runner_backend, moe_runner_config) @@ -2010,6 +2046,9 @@ class ModelOptNvFp4FusedMoEMethod(FusedMoEMethodBase): return self.runner.run(dispatch_output, quant_info) + # CuteDSL v2 standard path (a2a=none/flashinfer). + # The v1 (deepep) path never reaches apply(); it goes through + # apply_without_routing_weights instead. if self.enable_flashinfer_cutedsl_moe: from sglang.srt.layers.moe.moe_runner.flashinfer_cutedsl import ( CuteDslFp4MoeQuantInfo, @@ -2124,6 +2163,14 @@ class ModelOptNvFp4FusedMoEMethod(FusedMoEMethodBase): masked_m: torch.Tensor, moe_runner_config: MoeRunnerConfig, ) -> torch.Tensor: + """CuteDSL v1 (deepep low-latency) path. + + Called by the DeepEP dispatcher instead of apply(). Uses + flashinfer_cutedsl_moe_masked (grouped_gemm_nt_masked) directly, + bypassing MoeRunner. Weights must be in default [Gate, Up] order + and NOT interleaved -- see _is_cutedsl_v1_deepep guards in + process_weights_after_loading and load_up_proj_weight_first. + """ assert ( moe_runner_config.activation == "silu" ), "Only SiLU activation is supported." @@ -2137,6 +2184,21 @@ class ModelOptNvFp4FusedMoEMethod(FusedMoEMethodBase): flashinfer_cutedsl_moe_masked, ) + # flashinfer_cutedsl_moe_masked reinterprets scales as float8_e4m3fn. + # Same-dtype .view is a no-op; only wider dtypes (e.g. int32-packed + # UE8M0) need stride(-1)==1. + if ( + MOE_NVFP4_DISPATCH + and x[1] is not None + and x[1].element_size() != 1 + and x[1].stride(-1) != 1 + ): + raise AssertionError( + f"NVFP4 dispatch scale has stride(-1)={x[1].stride(-1)}, " + f"dtype={x[1].dtype}; .view(float8_e4m3fn) requires stride(-1)==1. " + "Try SGLANG_MOE_NVFP4_DISPATCH=0 or check DeepEP version." + ) + down_gemm_overlap_args: Optional[DownGemmOverlapArgs] = getattr( layer, "down_gemm_overlap_args", None ) diff --git a/python/sglang/srt/model_executor/model_runner.py b/python/sglang/srt/model_executor/model_runner.py index c46321ed3..878d1f756 100644 --- a/python/sglang/srt/model_executor/model_runner.py +++ b/python/sglang/srt/model_executor/model_runner.py @@ -2193,6 +2193,17 @@ class ModelRunner(ModelRunnerKVCacheMixin): if self.server_args.disable_flashinfer_autotune: return False + # CuteDSL v1 (cutedsl runner + deepep a2a) bypasses MoeRunner and must not + # be autotuned -- its _dummy_run would dispatch more tokens per rank than + # SGLANG_DEEPEP_NUM_MAX_DISPATCH_TOKENS_PER_RANK, tripping a DeepEP assert. + # Read server_args directly to avoid depending on initialize_moe_config() + # having already populated the MoE backend globals. + if ( + self.server_args.moe_runner_backend == "flashinfer_cutedsl" + and self.server_args.moe_a2a_backend == "deepep" + ): + return False + backend_str = self.server_args.moe_runner_backend # TODO smor- support other cases for flashinfer autotune, such as, mamba backend diff --git a/test/registered/moe/test_cutedsl_moe.py b/test/registered/moe/test_cutedsl_moe.py index 302aaba4b..25d91bc45 100644 --- a/test/registered/moe/test_cutedsl_moe.py +++ b/test/registered/moe/test_cutedsl_moe.py @@ -473,188 +473,25 @@ def torch_moe_nvfp4(a, w1, w2, topk, topk_weight, topk_ids): ).sum(dim=1) -class TestFlashinferCutedslMoe(unittest.TestCase): - @unittest.skipIf(SKIP_TEST, SKIP_REASON) - def test_flashinfer_cutedsl_moe_masked(self): - # Test parameters - test_cases = [ - (2, 128, 256, 1), - (2, 128, 256, 2), - (2, 128, 256, 4), - (16, 128, 512, 1), - (16, 128, 512, 2), - (16, 128, 512, 4), - ] +class TestCuteDslV2(unittest.TestCase): + """Correctness tests for the CuteDSL v2 (standard) path. - for bs, hidden_dim, inter_dim, topk in test_cases: - with self.subTest( - bs=bs, hidden_dim=hidden_dim, inter_dim=inter_dim, topk=topk - ): - with torch.inference_mode(): - torch.manual_seed(42) - device = "cuda" - dtype = torch.bfloat16 - num_experts = 8 - hidden_states = ( - torch.randn(bs, hidden_dim, dtype=torch.bfloat16, device=device) - / 5.0 - ) - w1 = ( - torch.randn( - num_experts, - 2 * inter_dim, - hidden_dim, - dtype=torch.bfloat16, - device=device, - ) - / 10.0 - ) - w2 = ( - torch.randn( - num_experts, - hidden_dim, - inter_dim, - dtype=torch.bfloat16, - device=device, - ) - / 10.0 - ) - router_logits = torch.randn(bs, num_experts, dtype=torch.float32) + The v2 path uses CuteDslMoEWrapper with: + - W13 in [Up, Gate] order (load_up_proj_weight_first = True) + - W13 interleaved in 64-row chunks (interleave_w13_halves) + - MMA-layout blockscales (convert_sf_to_mma_layout) - hidden_states_expanded = ( - hidden_states.view(bs, -1, hidden_dim) - .repeat(1, topk, 1) - .reshape(-1, hidden_dim) - ) - hidden_states_3d, masked_m, topk_idx, routing_weights = ( - prepare_inputs( - hidden_states_expanded, router_logits, num_experts, topk - ) - ) - - w1_amax = w1.abs().amax(dim=(1, 2)).to(torch.float32).to(w1.device) - w2_amax = w2.abs().amax(dim=(1, 2)).to(torch.float32).to(w2.device) - input_global_scale = torch.ones( - (num_experts,), dtype=torch.float32, device=hidden_states.device - ) - - w1_global_scale = FLOAT8_E4M3_MAX * FLOAT4_E2M1_MAX / w1_amax - w2_global_scale = FLOAT8_E4M3_MAX * FLOAT4_E2M1_MAX / w2_amax - a2_global_scale = torch.ones( - (num_experts,), dtype=torch.float32, device=hidden_states.device - ) # assume intermediate scale is 1.0 - - w1_fp4, w1_blockscale = scaled_fp4_grouped_quantize( - w1, - torch.ones(num_experts, dtype=torch.int32, device=w1.device) - * 2 - * inter_dim, - w1_global_scale, - ) - w2_fp4, w2_blockscale = scaled_fp4_grouped_quantize( - w2, - torch.ones(num_experts, dtype=torch.int32, device=w2.device) - * hidden_dim, - w2_global_scale, - ) - - w1_alpha = 1.0 / (input_global_scale * w1_global_scale) - w2_alpha = 1.0 / (a2_global_scale * w2_global_scale) - - out = flashinfer_cutedsl_moe_masked( - (hidden_states_3d.to(hidden_states.device), None), - input_global_scale, - w1_fp4.permute(2, 0, 1), - w1_blockscale, - w1_alpha, - w2_fp4.permute(2, 0, 1), - a2_global_scale, - w2_blockscale, - w2_alpha, - masked_m.to(hidden_states.device), - ) - - # reference - a_fp4, a_scale_interleaved = fp4_quantize( - hidden_states, input_global_scale - ) - a_in_dtype = dequantize_nvfp4_to_dtype( - a_fp4, - a_scale_interleaved, - input_global_scale, - dtype=hidden_states.dtype, - device=hidden_states.device, - block_size=16, - ) - w1_d = torch.empty( - (num_experts, 2 * inter_dim, hidden_dim), - device=w1.device, - dtype=w1.dtype, - ) - w2_d = torch.empty( - (num_experts, hidden_dim, inter_dim), - device=w2.device, - dtype=w2.dtype, - ) - - for idx in range(0, num_experts): - w1_fp4_sliced, w1_blockscale_sliced = fp4_quantize( - w1[idx], w1_global_scale[idx] - ) - w2_fp4_sliced, w2_blockscale_sliced = fp4_quantize( - w2[idx], w2_global_scale[idx] - ) - w1_d[idx] = dequantize_nvfp4_to_dtype( - w1_fp4_sliced, - w1_blockscale_sliced, - w1_global_scale[idx], - dtype=w1.dtype, - device=w1.device, - block_size=16, - ) - w2_d[idx] = dequantize_nvfp4_to_dtype( - w2_fp4_sliced, - w2_blockscale_sliced, - w2_global_scale[idx], - dtype=w2.dtype, - device=w2.device, - block_size=16, - ) - - ref_output = torch_moe_nvfp4( - a_in_dtype, - w1_d, - w2_d, - topk, - routing_weights.to(a_in_dtype.device), - topk_idx.to(a_in_dtype.device), - ) - out_weighted = torch.zeros_like( - ref_output, device=out.device, dtype=out.dtype - ) - - positions = torch.nonzero(masked_m[topk_idx], as_tuple=False) - rows, cols = positions[:, 0], positions[:, 1] - experts = topk_idx[rows, cols] - for i in range(num_experts): - mask = experts == i - if mask.any(): - idx = torch.nonzero(mask, as_tuple=False).squeeze(-1) - r, c = rows[idx], cols[idx] - out_weighted[r] += out[i, : len(r), :] * routing_weights[ - r, c - ].to(out.device).unsqueeze(-1) - torch.testing.assert_close( - out_weighted.cpu(), ref_output.cpu(), atol=5e-2, rtol=5e-2 - ) + This is the path used with --moe-runner-backend flashinfer_cutedsl and + --moe-a2a-backend none or flashinfer (i.e. NOT deepep). + """ @unittest.skipIf(SKIP_TEST, SKIP_REASON) @unittest.skipIf( CuteDslMoEWrapper is None or convert_sf_to_mma_layout is None, "CuteDslMoEWrapper / convert_sf_to_mma_layout not available", ) - def test_cutedsl_moe_wrapper_run(self): - """Call CuteDslMoEWrapper.run() with MMA-layout tensors and verify against reference.""" + def test_v2_wrapper_correctness(self): + """CuteDslMoEWrapper.run() with MMA-layout tensors vs PyTorch reference.""" test_cases = [ # (num_tokens, hidden_size, intermediate_size, num_experts, top_k) # Minimum dimensions match FlashInfer's test_wrapper_accuracy: @@ -739,8 +576,8 @@ class TestFlashinferCutedslMoe(unittest.TestCase): CuteDslMoEWrapper is None or convert_sf_to_mma_layout is None, "CuteDslMoEWrapper / convert_sf_to_mma_layout not available", ) - def test_cutedsl_cuda_graph_parity(self): - """Verify non-graph and cuda_graph wrappers produce identical results. + def test_v2_cuda_graph_parity(self): + """Verify non-graph and cuda_graph v2 wrappers produce identical results. Also checks both match the pure-PyTorch reference, and that a second cuda_graph pass reuses buffers deterministically (subsumes the former @@ -841,13 +678,13 @@ class TestFlashinferCutedslMoe(unittest.TestCase): CuteDslMoEWrapper is None or convert_sf_to_mma_layout is None, "CuteDslMoEWrapper / convert_sf_to_mma_layout not available", ) - def test_cutedsl_ep_sharded_allreduce(self): - """Verify EP-sharded execution: partial outputs from EP ranks sum to full result. + def test_v2_ep_sharded_allreduce(self): + """Verify EP-sharded v2 execution: partial outputs from EP ranks sum to full result. Simulates the EP=TP all-reduce pattern used by the CuteDSL moe_runner when - ep_size > 1 and moe_a2a_backend=none. Each "rank" runs a wrapper with + ep_size > 1 and moe_a2a_backend=none. Each "rank" runs a v2 wrapper with num_local_experts < num_experts and a corresponding local_expert_offset, - receiving only the local slice of weights/scales/alphas — matching the + receiving only the local slice of weights/scales/alphas -- matching the real runtime contract where each rank holds only its own expert partition. The partial outputs are summed (simulating tensor_model_parallel_all_reduce) and compared against a single wrapper processing all experts. @@ -940,5 +777,545 @@ class TestFlashinferCutedslMoe(unittest.TestCase): ) +class TestCuteDslV1(unittest.TestCase): + """Correctness tests for the CuteDSL v1 (deepep) path. + + The v1 path (apply_without_routing_weights -> flashinfer_cutedsl_moe_masked) + is used when --moe-runner-backend flashinfer_cutedsl and --moe-a2a-backend + deepep are combined. It expects: + - W13 in default [Gate, Up] order (load_up_proj_weight_first = False) + - W13 NOT interleaved (no interleave_w13_halves) + - Swizzled blockscales (w13_blockscale_swizzled, not MMA layout) + + A regression that accidentally applies v2 transforms (interleave, + [Up,Gate] flip, MMA blockscales) to v1 weights would cause these tests + to fail with numerical mismatch against the PyTorch reference. + + The companion v2 (standard) path correctness is covered by TestCuteDslV2. + """ + + @unittest.skipIf(SKIP_TEST, SKIP_REASON) + def test_v1_masked_kernel_bf16_input(self): + """V1 masked kernel with BF16 activations (kernel quantizes internally). + + Weights are in v1 layout: [Gate, Up] order, non-interleaved, swizzled + blockscales. This mirrors the production path when DeepEP dispatch + does NOT pre-quantize activations (MOE_NVFP4_DISPATCH is off). + """ + test_cases = [ + # (bs, hidden_dim, inter_dim, topk) + (2, 128, 256, 1), + (2, 128, 256, 2), + (2, 128, 256, 4), + (16, 128, 512, 1), + (16, 128, 512, 2), + (16, 128, 512, 4), + ] + + for bs, hidden_dim, inter_dim, topk in test_cases: + with self.subTest( + bs=bs, hidden_dim=hidden_dim, inter_dim=inter_dim, topk=topk + ): + with torch.inference_mode(): + torch.manual_seed(42) + device = "cuda" + num_experts = 8 + hidden_states = ( + torch.randn(bs, hidden_dim, dtype=torch.bfloat16, device=device) + / 5.0 + ) + w1 = ( + torch.randn( + num_experts, + 2 * inter_dim, + hidden_dim, + dtype=torch.bfloat16, + device=device, + ) + / 10.0 + ) + w2 = ( + torch.randn( + num_experts, + hidden_dim, + inter_dim, + dtype=torch.bfloat16, + device=device, + ) + / 10.0 + ) + router_logits = torch.randn(bs, num_experts, dtype=torch.float32) + + hidden_states_expanded = ( + hidden_states.view(bs, -1, hidden_dim) + .repeat(1, topk, 1) + .reshape(-1, hidden_dim) + ) + hidden_states_3d, masked_m, topk_idx, routing_weights = ( + prepare_inputs( + hidden_states_expanded, router_logits, num_experts, topk + ) + ) + + w1_amax = w1.abs().amax(dim=(1, 2)).to(torch.float32).to(w1.device) + w2_amax = w2.abs().amax(dim=(1, 2)).to(torch.float32).to(w2.device) + input_global_scale = torch.ones( + (num_experts,), dtype=torch.float32, device=hidden_states.device + ) + + w1_global_scale = FLOAT8_E4M3_MAX * FLOAT4_E2M1_MAX / w1_amax + w2_global_scale = FLOAT8_E4M3_MAX * FLOAT4_E2M1_MAX / w2_amax + a2_global_scale = torch.ones( + (num_experts,), dtype=torch.float32, device=hidden_states.device + ) + + w1_fp4, w1_blockscale = scaled_fp4_grouped_quantize( + w1, + torch.ones(num_experts, dtype=torch.int32, device=w1.device) + * 2 + * inter_dim, + w1_global_scale, + ) + w2_fp4, w2_blockscale = scaled_fp4_grouped_quantize( + w2, + torch.ones(num_experts, dtype=torch.int32, device=w2.device) + * hidden_dim, + w2_global_scale, + ) + + w1_alpha = 1.0 / (input_global_scale * w1_global_scale) + w2_alpha = 1.0 / (a2_global_scale * w2_global_scale) + + out = flashinfer_cutedsl_moe_masked( + (hidden_states_3d.to(hidden_states.device), None), + input_global_scale, + w1_fp4.permute(2, 0, 1), + w1_blockscale, + w1_alpha, + w2_fp4.permute(2, 0, 1), + a2_global_scale, + w2_blockscale, + w2_alpha, + masked_m.to(hidden_states.device), + ) + + a_fp4, a_scale_interleaved = fp4_quantize( + hidden_states, input_global_scale + ) + a_in_dtype = dequantize_nvfp4_to_dtype( + a_fp4, + a_scale_interleaved, + input_global_scale, + dtype=hidden_states.dtype, + device=hidden_states.device, + block_size=16, + ) + w1_d = torch.empty( + (num_experts, 2 * inter_dim, hidden_dim), + device=w1.device, + dtype=w1.dtype, + ) + w2_d = torch.empty( + (num_experts, hidden_dim, inter_dim), + device=w2.device, + dtype=w2.dtype, + ) + + for idx in range(0, num_experts): + w1_fp4_sliced, w1_blockscale_sliced = fp4_quantize( + w1[idx], w1_global_scale[idx] + ) + w2_fp4_sliced, w2_blockscale_sliced = fp4_quantize( + w2[idx], w2_global_scale[idx] + ) + w1_d[idx] = dequantize_nvfp4_to_dtype( + w1_fp4_sliced, + w1_blockscale_sliced, + w1_global_scale[idx], + dtype=w1.dtype, + device=w1.device, + block_size=16, + ) + w2_d[idx] = dequantize_nvfp4_to_dtype( + w2_fp4_sliced, + w2_blockscale_sliced, + w2_global_scale[idx], + dtype=w2.dtype, + device=w2.device, + block_size=16, + ) + + ref_output = torch_moe_nvfp4( + a_in_dtype, + w1_d, + w2_d, + topk, + routing_weights.to(a_in_dtype.device), + topk_idx.to(a_in_dtype.device), + ) + out_weighted = torch.zeros_like( + ref_output, device=out.device, dtype=out.dtype + ) + + positions = torch.nonzero(masked_m[topk_idx], as_tuple=False) + rows, cols = positions[:, 0], positions[:, 1] + experts = topk_idx[rows, cols] + for i in range(num_experts): + mask = experts == i + if mask.any(): + idx = torch.nonzero(mask, as_tuple=False).squeeze(-1) + r, c = rows[idx], cols[idx] + out_weighted[r] += out[i, : len(r), :] * routing_weights[ + r, c + ].to(out.device).unsqueeze(-1) + torch.testing.assert_close( + out_weighted.cpu(), ref_output.cpu(), atol=5e-2, rtol=5e-2 + ) + + @unittest.skipIf(SKIP_TEST, SKIP_REASON) + def test_v1_masked_kernel_rejects_v2_w13_layout(self): + """Applying the v2 W13 transform must break the v1 masked path.""" + with torch.inference_mode(): + torch.manual_seed(42) + device = "cuda" + num_experts, bs, hidden_dim, inter_dim, topk = 8, 16, 128, 512, 2 + + hidden_states = ( + torch.randn(bs, hidden_dim, dtype=torch.bfloat16, device=device) / 5.0 + ) + w1 = ( + torch.randn( + num_experts, + 2 * inter_dim, + hidden_dim, + dtype=torch.bfloat16, + device=device, + ) + / 10.0 + ) + w2 = ( + torch.randn( + num_experts, + hidden_dim, + inter_dim, + dtype=torch.bfloat16, + device=device, + ) + / 10.0 + ) + router_logits = torch.randn(bs, num_experts, dtype=torch.float32) + + hidden_expanded = ( + hidden_states.view(bs, -1, hidden_dim) + .repeat(1, topk, 1) + .reshape(-1, hidden_dim) + ) + hidden_3d, masked_m, topk_idx, routing_weights = prepare_inputs( + hidden_expanded, router_logits, num_experts, topk + ) + + input_global_scale = torch.ones( + (num_experts,), dtype=torch.float32, device=device + ) + w1_amax = w1.abs().amax(dim=(1, 2)).to(torch.float32) + w2_amax = w2.abs().amax(dim=(1, 2)).to(torch.float32) + w1_global_scale = FLOAT8_E4M3_MAX * FLOAT4_E2M1_MAX / w1_amax + w2_global_scale = FLOAT8_E4M3_MAX * FLOAT4_E2M1_MAX / w2_amax + a2_global_scale = torch.ones( + (num_experts,), dtype=torch.float32, device=device + ) + + expert_sizes_w1 = ( + torch.ones(num_experts, dtype=torch.int32, device=device) + * 2 + * inter_dim + ) + expert_sizes_w2 = ( + torch.ones(num_experts, dtype=torch.int32, device=device) * hidden_dim + ) + w1_fp4, w1_blockscale = scaled_fp4_grouped_quantize( + w1, expert_sizes_w1, w1_global_scale + ) + w2_fp4, w2_blockscale = scaled_fp4_grouped_quantize( + w2, expert_sizes_w2, w2_global_scale + ) + + # The v2 standard path flips W13 to [Up, Gate] order and interleaves + # 64-row chunks for CuteDslMoEWrapper. The v1 masked kernel must not + # receive that transformed layout. + w1_v2 = torch.cat((w1[:, inter_dim:, :], w1[:, :inter_dim, :]), dim=1) + w1_v2 = _interleave_w13_halves(w1_v2, group_size=64, dim=1).contiguous() + w1_fp4_v2, w1_blockscale_v2 = scaled_fp4_grouped_quantize( + w1_v2, expert_sizes_w1, w1_global_scale + ) + + w1_alpha = 1.0 / (input_global_scale * w1_global_scale) + w2_alpha = 1.0 / (a2_global_scale * w2_global_scale) + + out_v1 = flashinfer_cutedsl_moe_masked( + (hidden_3d.to(device), None), + input_global_scale, + w1_fp4.permute(2, 0, 1), + w1_blockscale, + w1_alpha, + w2_fp4.permute(2, 0, 1), + a2_global_scale, + w2_blockscale, + w2_alpha, + masked_m.to(device), + ) + out_v2_layout = flashinfer_cutedsl_moe_masked( + (hidden_3d.to(device), None), + input_global_scale, + w1_fp4_v2.permute(2, 0, 1), + w1_blockscale_v2, + w1_alpha, + w2_fp4.permute(2, 0, 1), + a2_global_scale, + w2_blockscale, + w2_alpha, + masked_m.to(device), + ) + + a_fp4, a_scale_interleaved = fp4_quantize(hidden_states, input_global_scale) + a_in_dtype = dequantize_nvfp4_to_dtype( + a_fp4, + a_scale_interleaved, + input_global_scale, + dtype=hidden_states.dtype, + device=device, + block_size=16, + ) + w1_d = torch.empty( + (num_experts, 2 * inter_dim, hidden_dim), + device=device, + dtype=w1.dtype, + ) + w2_d = torch.empty( + (num_experts, hidden_dim, inter_dim), device=device, dtype=w2.dtype + ) + + for idx in range(num_experts): + w1_fp4_sliced, w1_blockscale_sliced = fp4_quantize( + w1[idx], w1_global_scale[idx] + ) + w2_fp4_sliced, w2_blockscale_sliced = fp4_quantize( + w2[idx], w2_global_scale[idx] + ) + w1_d[idx] = dequantize_nvfp4_to_dtype( + w1_fp4_sliced, + w1_blockscale_sliced, + w1_global_scale[idx], + dtype=w1.dtype, + device=device, + block_size=16, + ) + w2_d[idx] = dequantize_nvfp4_to_dtype( + w2_fp4_sliced, + w2_blockscale_sliced, + w2_global_scale[idx], + dtype=w2.dtype, + device=device, + block_size=16, + ) + + ref_output = torch_moe_nvfp4( + a_in_dtype, + w1_d, + w2_d, + topk, + routing_weights.to(device), + topk_idx.to(device), + ) + + positions = torch.nonzero(masked_m[topk_idx], as_tuple=False) + rows, cols = positions[:, 0], positions[:, 1] + experts = topk_idx[rows, cols] + + def combine_weighted_output(out: torch.Tensor) -> torch.Tensor: + out_weighted = torch.zeros_like( + ref_output, device=device, dtype=out.dtype + ) + for i in range(num_experts): + mask = experts == i + if mask.any(): + idx = torch.nonzero(mask, as_tuple=False).squeeze(-1) + r, c = rows[idx], cols[idx] + out_weighted[r] += out[i, : len(r), :] * routing_weights[ + r, c + ].to(device).unsqueeze(-1) + return out_weighted + + out_v1_weighted = combine_weighted_output(out_v1) + out_v2_layout_weighted = combine_weighted_output(out_v2_layout) + + torch.testing.assert_close( + out_v1_weighted.cpu(), ref_output.cpu(), atol=5e-2, rtol=5e-2 + ) + with self.assertRaises(AssertionError): + torch.testing.assert_close( + out_v2_layout_weighted.cpu(), + ref_output.cpu(), + atol=5e-2, + rtol=5e-2, + ) + + @unittest.skipIf(SKIP_TEST, SKIP_REASON) + def test_v1_masked_kernel_fp4_input(self): + """V1 masked kernel with pre-quantized FP4 activations. + + In production with MOE_NVFP4_DISPATCH, the DeepEP dispatcher quantizes + activations during dispatch. The v1 kernel receives + hidden_states=(fp4_data, blockscale) instead of (bf16_data, None) and + skips its internal scaled_fp4_grouped_quantize call. + """ + with torch.inference_mode(): + torch.manual_seed(42) + device = "cuda" + num_experts, bs, hidden_dim, inter_dim, topk = 8, 16, 128, 512, 2 + + hidden_states = ( + torch.randn(bs, hidden_dim, dtype=torch.bfloat16, device=device) / 5.0 + ) + w1 = ( + torch.randn( + num_experts, + 2 * inter_dim, + hidden_dim, + dtype=torch.bfloat16, + device=device, + ) + / 10.0 + ) + w2 = ( + torch.randn( + num_experts, + hidden_dim, + inter_dim, + dtype=torch.bfloat16, + device=device, + ) + / 10.0 + ) + router_logits = torch.randn(bs, num_experts, dtype=torch.float32) + + hidden_expanded = ( + hidden_states.view(bs, -1, hidden_dim) + .repeat(1, topk, 1) + .reshape(-1, hidden_dim) + ) + hidden_3d, masked_m, topk_idx, routing_weights = prepare_inputs( + hidden_expanded, router_logits, num_experts, topk + ) + + input_gs = torch.ones(num_experts, dtype=torch.float32, device=device) + w1_amax = w1.abs().amax(dim=(1, 2)).to(torch.float32) + w2_amax = w2.abs().amax(dim=(1, 2)).to(torch.float32) + w1_gs = FLOAT8_E4M3_MAX * FLOAT4_E2M1_MAX / w1_amax + w2_gs = FLOAT8_E4M3_MAX * FLOAT4_E2M1_MAX / w2_amax + a2_gs = torch.ones(num_experts, dtype=torch.float32, device=device) + + expert_sizes_w1 = ( + torch.ones(num_experts, dtype=torch.int32, device=device) + * 2 + * inter_dim + ) + expert_sizes_w2 = ( + torch.ones(num_experts, dtype=torch.int32, device=device) * hidden_dim + ) + w1_fp4, w1_bs = scaled_fp4_grouped_quantize(w1, expert_sizes_w1, w1_gs) + w2_fp4, w2_bs = scaled_fp4_grouped_quantize(w2, expert_sizes_w2, w2_gs) + w1_alpha = 1.0 / (input_gs * w1_gs) + w2_alpha = 1.0 / (a2_gs * w2_gs) + + # Pre-quantize activations -- simulates what DeepEP dispatch does + # when MOE_NVFP4_DISPATCH is enabled. The kernel expects + # (m, k//2, num_experts) layout from scaled_fp4_grouped_quantize. + a_q, a_q_sf = scaled_fp4_grouped_quantize( + hidden_3d.to(device), + masked_m.to(device), + input_gs, + ) + + out = flashinfer_cutedsl_moe_masked( + (a_q, a_q_sf), + input_gs, + w1_fp4.permute(2, 0, 1), + w1_bs, + w1_alpha, + w2_fp4.permute(2, 0, 1), + a2_gs, + w2_bs, + w2_alpha, + masked_m.to(device), + ) + + # PyTorch reference (same as the bf16 input test) + a_fp4, a_scale = fp4_quantize(hidden_states, input_gs) + a_deq = dequantize_nvfp4_to_dtype( + a_fp4, + a_scale, + input_gs, + dtype=torch.bfloat16, + device=device, + block_size=16, + ) + w1_d = torch.empty( + (num_experts, 2 * inter_dim, hidden_dim), + device=device, + dtype=w1.dtype, + ) + w2_d = torch.empty( + (num_experts, hidden_dim, inter_dim), device=device, dtype=w2.dtype + ) + for idx in range(num_experts): + w1_fp4_sliced, w1_blockscale_sliced = fp4_quantize(w1[idx], w1_gs[idx]) + w2_fp4_sliced, w2_blockscale_sliced = fp4_quantize(w2[idx], w2_gs[idx]) + w1_d[idx] = dequantize_nvfp4_to_dtype( + w1_fp4_sliced, + w1_blockscale_sliced, + w1_gs[idx], + dtype=w1.dtype, + device=device, + block_size=16, + ) + w2_d[idx] = dequantize_nvfp4_to_dtype( + w2_fp4_sliced, + w2_blockscale_sliced, + w2_gs[idx], + dtype=w2.dtype, + device=device, + block_size=16, + ) + ref = torch_moe_nvfp4( + a_deq, + w1_d, + w2_d, + topk, + routing_weights.to(device), + topk_idx.to(device), + ) + + out_weighted = torch.zeros_like(ref, device=device) + positions = torch.nonzero(masked_m[topk_idx], as_tuple=False) + rows, cols = positions[:, 0], positions[:, 1] + experts = topk_idx[rows, cols] + for i in range(num_experts): + mask = experts == i + if mask.any(): + idx = torch.nonzero(mask, as_tuple=False).squeeze(-1) + r, c = rows[idx], cols[idx] + out_weighted[r] += out[i, : len(r), :] * routing_weights[r, c].to( + device + ).unsqueeze(-1) + + torch.testing.assert_close( + out_weighted.cpu(), + ref.cpu(), + atol=5e-2, + rtol=5e-2, + ) + + if __name__ == "__main__": unittest.main()