minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

fix legacy deepep path for flashinfer_cutedsl (#22925)

Browse Source
This commit is contained in:
Lee Nau
2026-04-20 11:49:33 -07:00
committed by GitHub
parent 4698f4cd10
commit b4bb036b73
5 changed files with 664 additions and 193 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/layers/moe/token_dispatcher/deepep.py
+18 -5
View File
@@ -620,9 +620,25 @@ class _DeepEPDispatcherImplLowLatency(_DeepEPDispatcherImplBase):
input_global_scale = self.quant_config.get("input_global_scale", None) input_global_scale = self.quant_config.get("input_global_scale", None)
if input_global_scale is not None: if input_global_scale is not None:
use_nvfp4 = True 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 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() buffer = self._get_buffer()
_deepep_precompile_tp_barrier() _deepep_precompile_tp_barrier()
packed_recv_hidden, self.packed_recv_count, self.handle, event, hook = ( 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, async_finish=not self.return_recv_hook,
return_recv_hook=self.return_recv_hook, return_recv_hook=self.return_recv_hook,
round_scale=deep_gemm_wrapper.ENABLE_JIT_DEEPGEMM **fp8_deepgemm_scale_opts,
and deep_gemm_wrapper.DEEPGEMM_BLACKWELL,
use_ue8m0=deep_gemm_wrapper.ENABLE_JIT_DEEPGEMM
and deep_gemm_wrapper.DEEPGEMM_BLACKWELL,
) )
) )
return packed_recv_hidden, self.packed_recv_count, event, hook return packed_recv_hidden, self.packed_recv_count, event, hook
python/sglang/srt/layers/moe/utils.py
+8
View File
@@ -263,6 +263,14 @@ def is_deepep_class_backend() -> bool:
return b.is_deepep() or b.is_mooncake() or b.is_mori() 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: def get_tbo_token_distribution_threshold() -> float:
global TBO_TOKEN_DISTRIBUTION_THRESHOLD global TBO_TOKEN_DISTRIBUTION_THRESHOLD
if TBO_TOKEN_DISTRIBUTION_THRESHOLD is None: if TBO_TOKEN_DISTRIBUTION_THRESHOLD is None:
python/sglang/srt/layers/quantization/modelopt_quant.py
+70 -8
View File
@@ -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.cutlass_moe_params import CutlassMoEParams, CutlassMoEType
from sglang.srt.layers.moe.moe_runner.triton import TritonMoeQuantInfo 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.parameter import ModelWeightParameter, PerTensorScaleParameter
from sglang.srt.layers.quantization.base_config import ( from sglang.srt.layers.quantization.base_config import (
FusedMoEMethodBase, FusedMoEMethodBase,
@@ -1546,6 +1549,29 @@ class ModelOptNvFp4FusedMoEMethod(FusedMoEMethodBase):
return get_moe_runner_backend().is_flashinfer_cutedsl() 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( def create_weights(
self, self,
layer: torch.nn.Module, layer: torch.nn.Module,
@@ -1812,9 +1838,11 @@ class ModelOptNvFp4FusedMoEMethod(FusedMoEMethodBase):
else: else:
# CUTLASS processing - handle w13 and w2 separately # CUTLASS processing - handle w13 and w2 separately
if self.enable_flashinfer_cutedsl_moe and layer.moe_runner_config.is_gated: if self._is_cutedsl_v2_standard and layer.moe_runner_config.is_gated:
# For the CuteDSL FP4 path, interleave the two logical W13 halves # CuteDSL v2 only: interleave the two logical W13 halves in
# in 64-row chunks before swizzling the block-scales. # 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 ( from sglang.srt.layers.moe.moe_runner.flashinfer_cutedsl import (
interleave_w13_halves, interleave_w13_halves,
) )
@@ -1876,8 +1904,10 @@ class ModelOptNvFp4FusedMoEMethod(FusedMoEMethodBase):
layer, "w2_blockscale_swizzled", w2_blockscale_swizzled layer, "w2_blockscale_swizzled", w2_blockscale_swizzled
) )
if self.enable_flashinfer_cutedsl_moe: if self._is_cutedsl_v2_standard:
# CuteDSL expects MMA layout for weight scales. Convert from swizzled bytes. # 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 flashinfer.cute_dsl.utils import convert_sf_to_mma_layout
from sglang.srt.layers.moe.moe_runner.flashinfer_cutedsl import ( from sglang.srt.layers.moe.moe_runner.flashinfer_cutedsl import (
@@ -1940,9 +1970,10 @@ class ModelOptNvFp4FusedMoEMethod(FusedMoEMethodBase):
@property @property
def load_up_proj_weight_first(self) -> bool: 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 ( 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( def create_moe_runner(
@@ -1959,6 +1990,11 @@ class ModelOptNvFp4FusedMoEMethod(FusedMoEMethodBase):
if moe_runner_backend.is_flashinfer_cutedsl(): if moe_runner_backend.is_flashinfer_cutedsl():
import sglang.srt.layers.moe.moe_runner.flashinfer_cutedsl # noqa: F401 – triggers @register_fused_func 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(): if not moe_runner_backend.is_flashinfer_cutlass():
self.runner = MoeRunner(moe_runner_backend, moe_runner_config) self.runner = MoeRunner(moe_runner_backend, moe_runner_config)
@@ -2010,6 +2046,9 @@ class ModelOptNvFp4FusedMoEMethod(FusedMoEMethodBase):
return self.runner.run(dispatch_output, quant_info) 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: if self.enable_flashinfer_cutedsl_moe:
from sglang.srt.layers.moe.moe_runner.flashinfer_cutedsl import ( from sglang.srt.layers.moe.moe_runner.flashinfer_cutedsl import (
CuteDslFp4MoeQuantInfo, CuteDslFp4MoeQuantInfo,
@@ -2124,6 +2163,14 @@ class ModelOptNvFp4FusedMoEMethod(FusedMoEMethodBase):
masked_m: torch.Tensor, masked_m: torch.Tensor,
moe_runner_config: MoeRunnerConfig, moe_runner_config: MoeRunnerConfig,
) -> torch.Tensor: ) -> 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 ( assert (
moe_runner_config.activation == "silu" moe_runner_config.activation == "silu"
), "Only SiLU activation is supported." ), "Only SiLU activation is supported."
@@ -2137,6 +2184,21 @@ class ModelOptNvFp4FusedMoEMethod(FusedMoEMethodBase):
flashinfer_cutedsl_moe_masked, 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( down_gemm_overlap_args: Optional[DownGemmOverlapArgs] = getattr(
layer, "down_gemm_overlap_args", None layer, "down_gemm_overlap_args", None
) )
python/sglang/srt/model_executor/model_runner.py
+11
View File
@@ -2193,6 +2193,17 @@ class ModelRunner(ModelRunnerKVCacheMixin):
if self.server_args.disable_flashinfer_autotune: if self.server_args.disable_flashinfer_autotune:
return False 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 backend_str = self.server_args.moe_runner_backend
# TODO smor- support other cases for flashinfer autotune, such as, mamba backend # TODO smor- support other cases for flashinfer autotune, such as, mamba backend
test/registered/moe/test_cutedsl_moe.py
+557 -180
View File
@@ -473,188 +473,25 @@ def torch_moe_nvfp4(a, w1, w2, topk, topk_weight, topk_ids):
).sum(dim=1) ).sum(dim=1)
class TestFlashinferCutedslMoe(unittest.TestCase): class TestCuteDslV2(unittest.TestCase):
@unittest.skipIf(SKIP_TEST, SKIP_REASON) """Correctness tests for the CuteDSL v2 (standard) path.
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),
]
for bs, hidden_dim, inter_dim, topk in test_cases: The v2 path uses CuteDslMoEWrapper with:
with self.subTest( - W13 in [Up, Gate] order (load_up_proj_weight_first = True)
bs=bs, hidden_dim=hidden_dim, inter_dim=inter_dim, topk=topk - W13 interleaved in 64-row chunks (interleave_w13_halves)
): - MMA-layout blockscales (convert_sf_to_mma_layout)
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)
hidden_states_expanded = ( This is the path used with --moe-runner-backend flashinfer_cutedsl and
hidden_states.view(bs, -1, hidden_dim) --moe-a2a-backend none or flashinfer (i.e. NOT deepep).
.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
)
@unittest.skipIf(SKIP_TEST, SKIP_REASON) @unittest.skipIf(SKIP_TEST, SKIP_REASON)
@unittest.skipIf( @unittest.skipIf(
CuteDslMoEWrapper is None or convert_sf_to_mma_layout is None, CuteDslMoEWrapper is None or convert_sf_to_mma_layout is None,
"CuteDslMoEWrapper / convert_sf_to_mma_layout not available", "CuteDslMoEWrapper / convert_sf_to_mma_layout not available",
) )
def test_cutedsl_moe_wrapper_run(self): def test_v2_wrapper_correctness(self):
"""Call CuteDslMoEWrapper.run() with MMA-layout tensors and verify against reference.""" """CuteDslMoEWrapper.run() with MMA-layout tensors vs PyTorch reference."""
test_cases = [ test_cases = [
# (num_tokens, hidden_size, intermediate_size, num_experts, top_k) # (num_tokens, hidden_size, intermediate_size, num_experts, top_k)
# Minimum dimensions match FlashInfer's test_wrapper_accuracy: # 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 is None or convert_sf_to_mma_layout is None,
"CuteDslMoEWrapper / convert_sf_to_mma_layout not available", "CuteDslMoEWrapper / convert_sf_to_mma_layout not available",
) )
def test_cutedsl_cuda_graph_parity(self): def test_v2_cuda_graph_parity(self):
"""Verify non-graph and cuda_graph wrappers produce identical results. """Verify non-graph and cuda_graph v2 wrappers produce identical results.
Also checks both match the pure-PyTorch reference, and that a second Also checks both match the pure-PyTorch reference, and that a second
cuda_graph pass reuses buffers deterministically (subsumes the former 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 is None or convert_sf_to_mma_layout is None,
"CuteDslMoEWrapper / convert_sf_to_mma_layout not available", "CuteDslMoEWrapper / convert_sf_to_mma_layout not available",
) )
def test_cutedsl_ep_sharded_allreduce(self): def test_v2_ep_sharded_allreduce(self):
"""Verify EP-sharded execution: partial outputs from EP ranks sum to full result. """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 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, 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. real runtime contract where each rank holds only its own expert partition.
The partial outputs are summed (simulating tensor_model_parallel_all_reduce) The partial outputs are summed (simulating tensor_model_parallel_all_reduce)
and compared against a single wrapper processing all experts. 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__": if __name__ == "__main__":
unittest.main() unittest.main()