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)
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
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()
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:
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.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
)
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:
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
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)
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()