[MoE] Add FlashInfer SM90 MXFP4 W4A8 CUTLASS MoE (#34967)
Co-authored-by: luoyuan.luo <luoyuan.luo@antgroup.com>
This commit is contained in:
@@ -337,7 +337,7 @@ automatically.
|
|||||||
|
|
||||||
Two options are available for running DeepSeek-V4 on Hopper:
|
Two options are available for running DeepSeek-V4 on Hopper:
|
||||||
|
|
||||||
- **Original FP4 checkpoints** — apply the W4A16 MoE kernels (Marlin) as the command generator picks for Hopper cells. This path works on both H100 and H200 and is the only option for H100 (no FP8 path). It is TP-only; on H200 the Pro variant fits on a single 8-GPU node, while H100 Pro needs 2 nodes (TP=16).
|
- **Original FP4 checkpoints** — run the MoE experts with W4A16 kernels (Marlin or the FlashInfer SM90 CUTLASS runner) as the command generator picks for Hopper cells. With FlashInfer >= 0.6.18 you can instead select the **W4A8** path — MXFP4 weights with FP8 activations via FlashInfer's Humming kernels — by adding `--flashinfer-mxfp4-moe-precision fp8` to `--moe-runner-backend flashinfer_mxfp4`; the low-latency Hopper cells now generate this form. Both work on H100 and H200; FP4 is the only option for H100 (no FP8 path). It is TP-only; on H200 the Pro variant fits on a single 8-GPU node, while H100 Pro needs 2 nodes (TP=16).
|
||||||
- **Converted FP8 checkpoints** (H100 and H200 only) — pre-repackaged FP8 weights at [`sgl-project/DeepSeek-V4-Flash-FP8`](https://huggingface.co/sgl-project/DeepSeek-V4-Flash-FP8) and [`sgl-project/DeepSeek-V4-Pro-FP8`](https://huggingface.co/sgl-project/DeepSeek-V4-Pro-FP8) unlock DP-attention + DeepEP and richer parallelism (e.g. Pro TP=16 across 2 nodes).
|
- **Converted FP8 checkpoints** (H100 and H200 only) — pre-repackaged FP8 weights at [`sgl-project/DeepSeek-V4-Flash-FP8`](https://huggingface.co/sgl-project/DeepSeek-V4-Flash-FP8) and [`sgl-project/DeepSeek-V4-Pro-FP8`](https://huggingface.co/sgl-project/DeepSeek-V4-Pro-FP8) unlock DP-attention + DeepEP and richer parallelism (e.g. Pro TP=16 across 2 nodes).
|
||||||
|
|
||||||
On these FP8 checkpoints you can additionally enable the all-FP8 **MegaMoE** path on SM90 for higher long-context / large-decode throughput — see the **SM90 (Hopper) FP8 MegaMoE** note in Configuration Tips below.
|
On these FP8 checkpoints you can additionally enable the all-FP8 **MegaMoE** path on SM90 for higher long-context / large-decode throughput — see the **SM90 (Hopper) FP8 MegaMoE** note in Configuration Tips below.
|
||||||
|
|||||||
@@ -1808,9 +1808,9 @@ Please consult the documentation below and [server_args.py](https://github.com/s
|
|||||||
</tr>
|
</tr>
|
||||||
<tr>
|
<tr>
|
||||||
<td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}>`--flashinfer-mxfp4-moe-precision`</td>
|
<td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}>`--flashinfer-mxfp4-moe-precision`</td>
|
||||||
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>Choose the computation precision of flashinfer mxfp4 moe</td>
|
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>Choose the computation precision of flashinfer mxfp4 moe. On SM90, `fp8` selects the Humming-style MXFP4-weight x FP8-activation path introduced by FlashInfer #3738 and requires FlashInfer >= 0.6.18.</td>
|
||||||
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.02)"}}>`default`</td>
|
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.02)"}}>`default`</td>
|
||||||
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}><code>default</code>, <code>bf16</code></td>
|
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}><code>default</code>, <code>bf16</code>, <code>fp8</code></td>
|
||||||
</tr>
|
</tr>
|
||||||
<tr>
|
<tr>
|
||||||
<td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}>`--enable-flashinfer-allreduce-fusion`</td>
|
<td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}>`--enable-flashinfer-allreduce-fusion`</td>
|
||||||
|
|||||||
@@ -1638,7 +1638,7 @@ click [Server Arguments](../../../advanced_features/server_arguments).
|
|||||||
<tr>
|
<tr>
|
||||||
<td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}>`--flashinfer-mxfp4-moe-precision`</td>
|
<td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}>`--flashinfer-mxfp4-moe-precision`</td>
|
||||||
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>`default`</td>
|
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>`default`</td>
|
||||||
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.02)"}}>`default`,<br/> `bf16`</td>
|
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.02)"}}>`default`,<br/> `bf16`,<br/> `fp8`</td>
|
||||||
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>Special for GPU</td>
|
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>Special for GPU</td>
|
||||||
</tr>
|
</tr>
|
||||||
<tr>
|
<tr>
|
||||||
|
|||||||
@@ -1904,13 +1904,17 @@ sgl-eval run mmmu_pro \\
|
|||||||
|
|
||||||
{
|
{
|
||||||
match: { hw: "h200", variant: "flash-official", quant: "fp4", strategy: "low-latency", nodes: "single" },
|
match: { hw: "h200", variant: "flash-official", quant: "fp4", strategy: "low-latency", nodes: "single" },
|
||||||
verified: true,
|
// W4A8 (MXFP4 weights x FP8 activations, FlashInfer Humming kernels);
|
||||||
|
// requires FlashInfer >= 0.6.18. Falls back: drop the precision flag
|
||||||
|
// for the W4A16 path, or use --moe-runner-backend marlin.
|
||||||
|
verificationStatus: "in-progress",
|
||||||
env: [],
|
env: [],
|
||||||
flags: [
|
flags: [
|
||||||
"--trust-remote-code",
|
"--trust-remote-code",
|
||||||
"--model-path {{MODEL_NAME}}",
|
"--model-path {{MODEL_NAME}}",
|
||||||
"--tp 4",
|
"--tp 4",
|
||||||
"--moe-runner-backend marlin",
|
"--moe-runner-backend flashinfer_mxfp4",
|
||||||
|
"--flashinfer-mxfp4-moe-precision fp8",
|
||||||
"--speculative-algorithm DSPARK",
|
"--speculative-algorithm DSPARK",
|
||||||
"--host {{HOST_IP}}",
|
"--host {{HOST_IP}}",
|
||||||
"--port {{PORT}}",
|
"--port {{PORT}}",
|
||||||
@@ -1918,13 +1922,15 @@ sgl-eval run mmmu_pro \\
|
|||||||
},
|
},
|
||||||
{
|
{
|
||||||
match: { hw: "h200", variant: "flash", quant: "fp4", strategy: "low-latency", nodes: "single" },
|
match: { hw: "h200", variant: "flash", quant: "fp4", strategy: "low-latency", nodes: "single" },
|
||||||
verified: true,
|
// W4A8 Humming path -- see the flash-official cell above.
|
||||||
|
verificationStatus: "in-progress",
|
||||||
env: [],
|
env: [],
|
||||||
flags: [
|
flags: [
|
||||||
"--trust-remote-code",
|
"--trust-remote-code",
|
||||||
"--model-path {{MODEL_NAME}}",
|
"--model-path {{MODEL_NAME}}",
|
||||||
"--tp 4",
|
"--tp 4",
|
||||||
"--moe-runner-backend marlin",
|
"--moe-runner-backend flashinfer_mxfp4",
|
||||||
|
"--flashinfer-mxfp4-moe-precision fp8",
|
||||||
"--speculative-algorithm EAGLE",
|
"--speculative-algorithm EAGLE",
|
||||||
"--speculative-num-steps 3",
|
"--speculative-num-steps 3",
|
||||||
"--speculative-eagle-topk 1",
|
"--speculative-eagle-topk 1",
|
||||||
@@ -1993,13 +1999,15 @@ sgl-eval run mmmu_pro \\
|
|||||||
},
|
},
|
||||||
{
|
{
|
||||||
match: { hw: "h200", variant: "pro", quant: "fp4", strategy: "low-latency", nodes: "single" },
|
match: { hw: "h200", variant: "pro", quant: "fp4", strategy: "low-latency", nodes: "single" },
|
||||||
verified: true,
|
// W4A8 Humming path -- see the flash-official cell above.
|
||||||
|
verificationStatus: "in-progress",
|
||||||
env: [],
|
env: [],
|
||||||
flags: [
|
flags: [
|
||||||
"--trust-remote-code",
|
"--trust-remote-code",
|
||||||
"--model-path {{MODEL_NAME}}",
|
"--model-path {{MODEL_NAME}}",
|
||||||
"--tp 8",
|
"--tp 8",
|
||||||
"--moe-runner-backend flashinfer_mxfp4",
|
"--moe-runner-backend flashinfer_mxfp4",
|
||||||
|
"--flashinfer-mxfp4-moe-precision fp8",
|
||||||
"--speculative-algorithm EAGLE",
|
"--speculative-algorithm EAGLE",
|
||||||
"--speculative-num-steps 3",
|
"--speculative-num-steps 3",
|
||||||
"--speculative-eagle-topk 1",
|
"--speculative-eagle-topk 1",
|
||||||
|
|||||||
@@ -401,6 +401,11 @@ class FusedMoE(torch.nn.Module):
|
|||||||
self.quant_config = quant_config
|
self.quant_config = quant_config
|
||||||
self.use_flashinfer_mxfp4_moe = get_moe_runner_backend().is_flashinfer_mxfp4()
|
self.use_flashinfer_mxfp4_moe = get_moe_runner_backend().is_flashinfer_mxfp4()
|
||||||
# TODO maybe we should remove this `if`, since `Mxfp4MoEMethod` does another round-up logic
|
# TODO maybe we should remove this `if`, since `Mxfp4MoEMethod` does another round-up logic
|
||||||
|
# Keep the pre-round value: the mxfp4 SM90 CUTLASS post-load processor
|
||||||
|
# needs to know which trailing K columns are padding, because the
|
||||||
|
# loader never writes them and their scale entries keep the buffer's
|
||||||
|
# _UE8M0_ONE fill.
|
||||||
|
self.hidden_size_unpadded = hidden_size
|
||||||
if (
|
if (
|
||||||
self.quant_config is not None
|
self.quant_config is not None
|
||||||
and self.quant_config.get_name() == "mxfp4"
|
and self.quant_config.get_name() == "mxfp4"
|
||||||
|
|||||||
@@ -65,8 +65,9 @@ class FlashInferCutlassMoeQuantInfo(MoeQuantInfo):
|
|||||||
class FlashInferCutlassMxfp4MoeQuantInfo(MoeQuantInfo):
|
class FlashInferCutlassMxfp4MoeQuantInfo(MoeQuantInfo):
|
||||||
"""Quantization payload for CUTLASS MXFP4 MoE.
|
"""Quantization payload for CUTLASS MXFP4 MoE.
|
||||||
|
|
||||||
SM90 consumes W4A16-interleaved weights and scales. SM120 consumes packed
|
SM90 consumes either W4A16-interleaved weights/scales or the Humming-style
|
||||||
MXFP4 weights and block-interleaved scales with MXFP8 activations.
|
W4A8 layouts. SM120 consumes packed MXFP4 weights and block-interleaved
|
||||||
|
scales with MXFP8 activations.
|
||||||
"""
|
"""
|
||||||
|
|
||||||
# SM90 weights are interleaved; SM120 weights remain checkpoint-packed.
|
# SM90 weights are interleaved; SM120 weights remain checkpoint-packed.
|
||||||
@@ -80,6 +81,13 @@ class FlashInferCutlassMxfp4MoeQuantInfo(MoeQuantInfo):
|
|||||||
# A non-None global scale selects the SM120 MXFP8 activation path.
|
# A non-None global scale selects the SM120 MXFP8 activation path.
|
||||||
mxfp4_weight_global_scale: Optional[torch.Tensor] = None
|
mxfp4_weight_global_scale: Optional[torch.Tensor] = None
|
||||||
|
|
||||||
|
# A complete non-None triplet selects the SM90 Humming W4A8 path. The
|
||||||
|
# residuals are FP32 [num_local_experts] and already include the fixed 2^6
|
||||||
|
# compensation required by FlashInfer's epilogue.
|
||||||
|
w13_humming_residual_scale: Optional[torch.Tensor] = None
|
||||||
|
w2_humming_residual_scale: Optional[torch.Tensor] = None
|
||||||
|
humming_fc2_act_scale: Optional[torch.Tensor] = None
|
||||||
|
|
||||||
# Per-expert bias. GPT-OSS has both; DSv4 leaves both None.
|
# Per-expert bias. GPT-OSS has both; DSv4 leaves both None.
|
||||||
w13_bias: Optional[torch.Tensor] = None # bf16 [E, 2*N]
|
w13_bias: Optional[torch.Tensor] = None # bf16 [E, 2*N]
|
||||||
w2_bias: Optional[torch.Tensor] = None # bf16 [E, K]
|
w2_bias: Optional[torch.Tensor] = None # bf16 [E, K]
|
||||||
@@ -340,6 +348,22 @@ def fused_experts_none_to_flashinfer_mxfp4(
|
|||||||
|
|
||||||
weight_global_scale = quant_info.mxfp4_weight_global_scale
|
weight_global_scale = quant_info.mxfp4_weight_global_scale
|
||||||
use_mxfp8_act_scaling = weight_global_scale is not None
|
use_mxfp8_act_scaling = weight_global_scale is not None
|
||||||
|
w13_humming_residual_scale = quant_info.w13_humming_residual_scale
|
||||||
|
w2_humming_residual_scale = quant_info.w2_humming_residual_scale
|
||||||
|
humming_fc2_act_scale = quant_info.humming_fc2_act_scale
|
||||||
|
humming_scales = (
|
||||||
|
w13_humming_residual_scale,
|
||||||
|
w2_humming_residual_scale,
|
||||||
|
humming_fc2_act_scale,
|
||||||
|
)
|
||||||
|
use_wfp4afp8_humming = any(scale is not None for scale in humming_scales)
|
||||||
|
if use_wfp4afp8_humming and not all(scale is not None for scale in humming_scales):
|
||||||
|
raise ValueError(
|
||||||
|
"SM90 Humming MXFP4 MoE requires both expert residual scales "
|
||||||
|
"and the FC2 activation scale."
|
||||||
|
)
|
||||||
|
if use_wfp4afp8_humming and use_mxfp8_act_scaling:
|
||||||
|
raise ValueError("SM90 Humming and SM120 MXFP8 scaling are mutually exclusive.")
|
||||||
input_sf = None
|
input_sf = None
|
||||||
fc1_expert_weights = quant_info.w13_weight
|
fc1_expert_weights = quant_info.w13_weight
|
||||||
fc2_expert_weights = quant_info.w2_weight
|
fc2_expert_weights = quant_info.w2_weight
|
||||||
@@ -359,6 +383,17 @@ def fused_experts_none_to_flashinfer_mxfp4(
|
|||||||
quant_info.w2_weight_scale.view(torch.int32),
|
quant_info.w2_weight_scale.view(torch.int32),
|
||||||
weight_global_scale,
|
weight_global_scale,
|
||||||
]
|
]
|
||||||
|
elif use_wfp4afp8_humming:
|
||||||
|
assert w13_humming_residual_scale is not None
|
||||||
|
assert w2_humming_residual_scale is not None
|
||||||
|
assert humming_fc2_act_scale is not None
|
||||||
|
quant_scales = [
|
||||||
|
quant_info.w13_weight_scale.view(torch.int32),
|
||||||
|
w13_humming_residual_scale,
|
||||||
|
humming_fc2_act_scale,
|
||||||
|
quant_info.w2_weight_scale.view(torch.int32),
|
||||||
|
w2_humming_residual_scale,
|
||||||
|
]
|
||||||
else:
|
else:
|
||||||
quant_scales = [
|
quant_scales = [
|
||||||
quant_info.w13_weight_scale.view(torch.int32),
|
quant_info.w13_weight_scale.view(torch.int32),
|
||||||
@@ -367,6 +402,10 @@ def fused_experts_none_to_flashinfer_mxfp4(
|
|||||||
|
|
||||||
out_hidden = padded_hidden if do_pad else origin_hidden
|
out_hidden = padded_hidden if do_pad else origin_hidden
|
||||||
output_dtype = torch.bfloat16
|
output_dtype = torch.bfloat16
|
||||||
|
# FlashInfer 0.6.17 intentionally reverted the Humming API. Do not pass the
|
||||||
|
# new keyword at all on the existing W4A16/MXFP8 paths, so those paths keep
|
||||||
|
# working with SGLang's currently pinned release.
|
||||||
|
humming_kwargs = {"use_wfp4afp8_humming": True} if use_wfp4afp8_humming else {}
|
||||||
with use_symmetric_memory(get_tp_group(), disabled=not is_allocation_symmetric()):
|
with use_symmetric_memory(get_tp_group(), disabled=not is_allocation_symmetric()):
|
||||||
out = torch.empty(x.shape[0], out_hidden, dtype=output_dtype, device=x.device)
|
out = torch.empty(x.shape[0], out_hidden, dtype=output_dtype, device=x.device)
|
||||||
|
|
||||||
@@ -398,6 +437,7 @@ def fused_experts_none_to_flashinfer_mxfp4(
|
|||||||
tune_max_num_tokens=next_power_of_2(x.shape[0]),
|
tune_max_num_tokens=next_power_of_2(x.shape[0]),
|
||||||
output=out,
|
output=out,
|
||||||
use_fused_finalize=envs.SGLANG_FLASHINFER_MOE_FUSED_FINALIZE.get(),
|
use_fused_finalize=envs.SGLANG_FLASHINFER_MOE_FUSED_FINALIZE.get(),
|
||||||
|
**humming_kwargs,
|
||||||
)
|
)
|
||||||
|
|
||||||
if do_pad:
|
if do_pad:
|
||||||
|
|||||||
@@ -405,14 +405,19 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
|
|||||||
self.flashinfer_mxfp4_moe_precision = (
|
self.flashinfer_mxfp4_moe_precision = (
|
||||||
get_exec().moe.flashinfer_mxfp4_moe_precision
|
get_exec().moe.flashinfer_mxfp4_moe_precision
|
||||||
)
|
)
|
||||||
|
self._use_sm90_humming = False
|
||||||
# When `flashinfer_mxfp4` is enabled, dispatch to one of three FlashInfer
|
# When `flashinfer_mxfp4` is enabled, dispatch to one of three FlashInfer
|
||||||
# entry points depending on the GPU:
|
# entry points depending on the GPU:
|
||||||
# - SM100 (Blackwell) -> trtllm_fp4_block_scale_moe (existing)
|
# - SM100 (Blackwell) -> trtllm_fp4_block_scale_moe (existing)
|
||||||
# - SM120 (Blackwell) -> cutlass_fused_moe(MXFP8 x MXFP4)
|
# - SM120 (Blackwell) -> cutlass_fused_moe(MXFP8 x MXFP4)
|
||||||
# - SM90 (Hopper) -> cutlass_fused_moe(use_w4_group_scaling=True)
|
# - SM90 (Hopper) -> cutlass_fused_moe(use_w4_group_scaling=True),
|
||||||
# (FlashInfer PR #3084, post-0.6.10)
|
# W4A16 by default (PR #3084) or opt-in
|
||||||
|
# Humming W4A8 (PR #3738/#4431)
|
||||||
self._fi_kernel: Optional[str] = None
|
self._fi_kernel: Optional[str] = None
|
||||||
if self.use_flashinfer:
|
if self.use_flashinfer:
|
||||||
|
# precision=fp8 is an SM90 knob (Humming W4A8). The Blackwell
|
||||||
|
# paths already run MXFP8 activations, so the flag is inert there
|
||||||
|
# rather than an error -- one config can move across hardware.
|
||||||
if get_platform().is_sm100:
|
if get_platform().is_sm100:
|
||||||
self._fi_kernel = "trtllm_sm100"
|
self._fi_kernel = "trtllm_sm100"
|
||||||
elif get_platform().is_sm120:
|
elif get_platform().is_sm120:
|
||||||
@@ -425,6 +430,7 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
|
|||||||
"from FlashInfer PR #3084 (>= 0.6.11). Upgrade flashinfer-python "
|
"from FlashInfer PR #3084 (>= 0.6.11). Upgrade flashinfer-python "
|
||||||
"or pick a different backend (e.g. marlin / triton_kernel)."
|
"or pick a different backend (e.g. marlin / triton_kernel)."
|
||||||
)
|
)
|
||||||
|
self._use_sm90_humming = self.flashinfer_mxfp4_moe_precision == "fp8"
|
||||||
self._fi_kernel = "cutlass_sm90"
|
self._fi_kernel = "cutlass_sm90"
|
||||||
else:
|
else:
|
||||||
raise NotImplementedError(
|
raise NotImplementedError(
|
||||||
@@ -497,6 +503,11 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
|
|||||||
# CUTLASS post-load processor after the load completes.
|
# CUTLASS post-load processor after the load completes.
|
||||||
self._padded_intermediate = round_up(intermediate_size_per_partition, 128)
|
self._padded_intermediate = round_up(intermediate_size_per_partition, 128)
|
||||||
self._padded_hidden = round_up(hidden_size, 128)
|
self._padded_hidden = round_up(hidden_size, 128)
|
||||||
|
# `hidden_size` here may ALREADY be FusedMoE's rounded value (GPT-OSS
|
||||||
|
# 2880 -> 3072). Remember the checkpoint's K so the post-load
|
||||||
|
# processor can exclude the never-written tail from Humming's
|
||||||
|
# per-expert scale range.
|
||||||
|
self._unpadded_hidden = getattr(layer, "hidden_size_unpadded", hidden_size)
|
||||||
# create_weights below uses the *unpadded* sizes so the loader's
|
# create_weights below uses the *unpadded* sizes so the loader's
|
||||||
# naive-copy fast path is correct.
|
# naive-copy fast path is correct.
|
||||||
intermediate_size_per_partition_after_pad = intermediate_size_per_partition
|
intermediate_size_per_partition_after_pad = intermediate_size_per_partition
|
||||||
@@ -1113,7 +1124,22 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
|
|||||||
|
|
||||||
_interleaved = getattr(layer.moe_runner_config, "gate_up_interleaved", True)
|
_interleaved = getattr(layer.moe_runner_config, "gate_up_interleaved", True)
|
||||||
|
|
||||||
def _stack_up_gate_w13(unpadded_w13, last_pad, last_un):
|
# FusedMoE may have rounded hidden up (GPT-OSS 2880 -> 3072) BEFORE
|
||||||
|
# create_weights, so K_un above is that rounded value, not the
|
||||||
|
# checkpoint's K. The loader never writes the trailing columns, so they
|
||||||
|
# keep the scale buffer's _UE8M0_ONE (2^0) fill -- far above a real
|
||||||
|
# per-expert max. Humming derives its residual from each expert's
|
||||||
|
# min/max E8M0 exponent, so letting those columns through would shift
|
||||||
|
# the residual and perturb the REAL weights. Copy only the checkpoint's
|
||||||
|
# columns and let the preserve_expert_range fill cover the rest; the
|
||||||
|
# packed weights there are zero, so the scale is numerically inert.
|
||||||
|
K_real = min(getattr(self, "_unpadded_hidden", None) or K_un, K_un)
|
||||||
|
# ceil: a partial trailing group is still real and must be kept.
|
||||||
|
w13_scale_real = -(-K_real // sf_block_size)
|
||||||
|
|
||||||
|
def _stack_up_gate_w13(
|
||||||
|
unpadded_w13, last_pad, last_un, preserve_expert_range=False, last_real=None
|
||||||
|
):
|
||||||
# unpadded_w13: [E, 2*N_un, last_un]
|
# unpadded_w13: [E, 2*N_un, last_un]
|
||||||
# Returns: [E, 2*N_pad, last_pad] in [up_padded; gate_padded] order.
|
# Returns: [E, 2*N_pad, last_pad] in [up_padded; gate_padded] order.
|
||||||
if _interleaved:
|
if _interleaved:
|
||||||
@@ -1126,10 +1152,22 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
|
|||||||
out = torch.zeros(
|
out = torch.zeros(
|
||||||
E, 2 * N_pad, last_pad, dtype=unpadded_w13.dtype, device=device
|
E, 2 * N_pad, last_pad, dtype=unpadded_w13.dtype, device=device
|
||||||
)
|
)
|
||||||
|
if preserve_expert_range:
|
||||||
|
# Humming derives one residual from each expert's min/max E8M0
|
||||||
|
# exponents. Fill padding with an existing expert value so
|
||||||
|
# padding cannot change that range.
|
||||||
|
out.copy_(unpadded_w13[:, :1, :1])
|
||||||
|
# When protecting the expert range, stop at the checkpoint's K so
|
||||||
|
# the stale tail is covered by the fill instead of copied through.
|
||||||
|
copy_un = (
|
||||||
|
min(last_real, last_un)
|
||||||
|
if (preserve_expert_range and last_real is not None)
|
||||||
|
else last_un
|
||||||
|
)
|
||||||
# First half: up (with row + col padding zeros).
|
# First half: up (with row + col padding zeros).
|
||||||
out[:, :N_un, :last_un] = up_rows
|
out[:, :N_un, :copy_un] = up_rows[:, :, :copy_un]
|
||||||
# Second half: gate.
|
# Second half: gate.
|
||||||
out[:, N_pad : N_pad + N_un, :last_un] = gate_rows
|
out[:, N_pad : N_pad + N_un, :copy_un] = gate_rows[:, :, :copy_un]
|
||||||
return out
|
return out
|
||||||
|
|
||||||
w13_padded = _stack_up_gate_w13(
|
w13_padded = _stack_up_gate_w13(
|
||||||
@@ -1139,6 +1177,8 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
|
|||||||
layer.w13_weight_scale.data,
|
layer.w13_weight_scale.data,
|
||||||
K_pad // sf_block_size,
|
K_pad // sf_block_size,
|
||||||
K_un // sf_block_size,
|
K_un // sf_block_size,
|
||||||
|
preserve_expert_range=self._use_sm90_humming,
|
||||||
|
last_real=w13_scale_real,
|
||||||
)
|
)
|
||||||
# Bias: same de-interleave on dim=-1.
|
# Bias: same de-interleave on dim=-1.
|
||||||
if _interleaved:
|
if _interleaved:
|
||||||
@@ -1151,9 +1191,19 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
|
|||||||
w13_bias_padded[:, :N_un] = w13_bias_up
|
w13_bias_padded[:, :N_un] = w13_bias_up
|
||||||
w13_bias_padded[:, N_pad : N_pad + N_un] = w13_bias_gate
|
w13_bias_padded[:, N_pad : N_pad + N_un] = w13_bias_gate
|
||||||
|
|
||||||
def _pad_w2_3d(unpadded, last_pad, last_un):
|
def _pad_w2_3d(
|
||||||
|
unpadded, last_pad, last_un, preserve_expert_range=False, k_real=None
|
||||||
|
):
|
||||||
out = torch.zeros(E, K_pad, last_pad, dtype=unpadded.dtype, device=device)
|
out = torch.zeros(E, K_pad, last_pad, dtype=unpadded.dtype, device=device)
|
||||||
out[:, :K_un, :last_un] = unpadded[:, :K_un, :]
|
if preserve_expert_range:
|
||||||
|
out.copy_(unpadded[:, :1, :1])
|
||||||
|
# Same stale-tail exclusion as _stack_up_gate_w13, on w2's K rows.
|
||||||
|
k_copy = (
|
||||||
|
min(k_real, K_un)
|
||||||
|
if (preserve_expert_range and k_real is not None)
|
||||||
|
else K_un
|
||||||
|
)
|
||||||
|
out[:, :k_copy, :last_un] = unpadded[:, :k_copy, :]
|
||||||
return out
|
return out
|
||||||
|
|
||||||
# ---- w2 (no halving, just pad to [E, K_pad, N_pad/2]) ----------------
|
# ---- w2 (no halving, just pad to [E, K_pad, N_pad/2]) ----------------
|
||||||
@@ -1164,6 +1214,8 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
|
|||||||
layer.w2_weight_scale.data,
|
layer.w2_weight_scale.data,
|
||||||
N_pad // sf_block_size,
|
N_pad // sf_block_size,
|
||||||
N_un // sf_block_size,
|
N_un // sf_block_size,
|
||||||
|
preserve_expert_range=self._use_sm90_humming,
|
||||||
|
k_real=K_real,
|
||||||
)
|
)
|
||||||
w2_bias_padded = torch.zeros(E, K_pad, dtype=bias_dtype, device=device)
|
w2_bias_padded = torch.zeros(E, K_pad, dtype=bias_dtype, device=device)
|
||||||
w2_bias_padded[:, :K_un] = layer.w2_weight_bias.data
|
w2_bias_padded[:, :K_un] = layer.w2_weight_bias.data
|
||||||
@@ -1187,26 +1239,47 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
|
|||||||
# ---- FlashInfer SM90 byte / scale interleave -----------------------
|
# ---- FlashInfer SM90 byte / scale interleave -----------------------
|
||||||
# The padded buffers above are contiguous by construction (allocated
|
# The padded buffers above are contiguous by construction (allocated
|
||||||
# via torch.zeros + slice assignment), so we feed them straight in.
|
# via torch.zeros + slice assignment), so we feed them straight in.
|
||||||
layer.w13_weight = Parameter(
|
if self._use_sm90_humming:
|
||||||
interleave_moe_weights_for_sm90_mixed_gemm(w13_padded, "fp4"),
|
from flashinfer.fused_moe import (
|
||||||
requires_grad=False,
|
preprocess_moe_weights_for_sm90_mixed_gemm_humming,
|
||||||
)
|
)
|
||||||
layer.w2_weight = Parameter(
|
|
||||||
interleave_moe_weights_for_sm90_mixed_gemm(w2_padded, "fp4"),
|
w13_il, w13_scale_il, w13_residual = (
|
||||||
requires_grad=False,
|
preprocess_moe_weights_for_sm90_mixed_gemm_humming(
|
||||||
)
|
w13_padded, w13_scale_padded
|
||||||
layer.w13_weight_scale = Parameter(
|
)
|
||||||
interleave_moe_scales_for_sm90_mixed_gemm(
|
)
|
||||||
|
w2_il, w2_scale_il, w2_residual = (
|
||||||
|
preprocess_moe_weights_for_sm90_mixed_gemm_humming(
|
||||||
|
w2_padded, w2_scale_padded
|
||||||
|
)
|
||||||
|
)
|
||||||
|
# Humming keeps the FP4->FP8 exponent-bias compensation in the
|
||||||
|
# epilogue. FlashInfer #4431 consumes these in local expert order.
|
||||||
|
layer.w13_humming_residual_scale = Parameter(
|
||||||
|
(w13_residual * 64.0).contiguous(), requires_grad=False
|
||||||
|
)
|
||||||
|
layer.w2_humming_residual_scale = Parameter(
|
||||||
|
(w2_residual * 64.0).contiguous(), requires_grad=False
|
||||||
|
)
|
||||||
|
layer.humming_fc2_act_scale = Parameter(
|
||||||
|
torch.ones((), dtype=torch.float32, device=device),
|
||||||
|
requires_grad=False,
|
||||||
|
)
|
||||||
|
else:
|
||||||
|
w13_il = interleave_moe_weights_for_sm90_mixed_gemm(w13_padded, "fp4")
|
||||||
|
w2_il = interleave_moe_weights_for_sm90_mixed_gemm(w2_padded, "fp4")
|
||||||
|
w13_scale_il = interleave_moe_scales_for_sm90_mixed_gemm(
|
||||||
w13_scale_padded, group_size=sf_block_size
|
w13_scale_padded, group_size=sf_block_size
|
||||||
),
|
)
|
||||||
requires_grad=False,
|
w2_scale_il = interleave_moe_scales_for_sm90_mixed_gemm(
|
||||||
)
|
|
||||||
layer.w2_weight_scale = Parameter(
|
|
||||||
interleave_moe_scales_for_sm90_mixed_gemm(
|
|
||||||
w2_scale_padded, group_size=sf_block_size
|
w2_scale_padded, group_size=sf_block_size
|
||||||
),
|
)
|
||||||
requires_grad=False,
|
|
||||||
)
|
layer.w13_weight = Parameter(w13_il, requires_grad=False)
|
||||||
|
layer.w2_weight = Parameter(w2_il, requires_grad=False)
|
||||||
|
layer.w13_weight_scale = Parameter(w13_scale_il, requires_grad=False)
|
||||||
|
layer.w2_weight_scale = Parameter(w2_scale_il, requires_grad=False)
|
||||||
layer.w13_weight_bias = Parameter(w13_bias_padded, requires_grad=False)
|
layer.w13_weight_bias = Parameter(w13_bias_padded, requires_grad=False)
|
||||||
layer.w2_weight_bias = Parameter(w2_bias_padded, requires_grad=False)
|
layer.w2_weight_bias = Parameter(w2_bias_padded, requires_grad=False)
|
||||||
|
|
||||||
@@ -1349,10 +1422,12 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
|
|||||||
pass
|
pass
|
||||||
|
|
||||||
def _apply_sm90_cutlass(self, layer, dispatch_output):
|
def _apply_sm90_cutlass(self, layer, dispatch_output):
|
||||||
"""SM90 (Hopper) MXFP4 x BF16 MoE via FlashInfer's cutlass mixed-input
|
"""SM90 MXFP4 x BF16/FP8 MoE via FlashInfer's mixed-input kernels.
|
||||||
path (PR #3084). Routed through the unified ``MoeRunner`` -- this
|
|
||||||
helper only builds the quant_info; the actual kernel call lives in
|
Routed through the unified ``MoeRunner``; this helper only builds the
|
||||||
:mod:`sglang.srt.layers.moe.moe_runner.flashinfer_cutlass`."""
|
quant_info. The actual kernel call lives in
|
||||||
|
:mod:`sglang.srt.layers.moe.moe_runner.flashinfer_cutlass`.
|
||||||
|
"""
|
||||||
from sglang.srt.layers.moe.moe_runner.flashinfer_cutlass import (
|
from sglang.srt.layers.moe.moe_runner.flashinfer_cutlass import (
|
||||||
FlashInferCutlassMxfp4MoeQuantInfo,
|
FlashInferCutlassMxfp4MoeQuantInfo,
|
||||||
)
|
)
|
||||||
@@ -1362,6 +1437,11 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
|
|||||||
w2_weight=layer.w2_weight,
|
w2_weight=layer.w2_weight,
|
||||||
w13_weight_scale=layer.w13_weight_scale,
|
w13_weight_scale=layer.w13_weight_scale,
|
||||||
w2_weight_scale=layer.w2_weight_scale,
|
w2_weight_scale=layer.w2_weight_scale,
|
||||||
|
w13_humming_residual_scale=getattr(
|
||||||
|
layer, "w13_humming_residual_scale", None
|
||||||
|
),
|
||||||
|
w2_humming_residual_scale=getattr(layer, "w2_humming_residual_scale", None),
|
||||||
|
humming_fc2_act_scale=getattr(layer, "humming_fc2_act_scale", None),
|
||||||
w13_bias=layer.w13_weight_bias,
|
w13_bias=layer.w13_weight_bias,
|
||||||
w2_bias=layer.w2_weight_bias,
|
w2_bias=layer.w2_weight_bias,
|
||||||
swiglu_alpha=layer.swiglu_alpha,
|
swiglu_alpha=layer.swiglu_alpha,
|
||||||
|
|||||||
@@ -14,7 +14,7 @@ import torch
|
|||||||
from torch.nn import Module
|
from torch.nn import Module
|
||||||
from torch.nn.parameter import Parameter
|
from torch.nn.parameter import Parameter
|
||||||
|
|
||||||
from sglang.srt.runtime_context import get_platform
|
from sglang.srt.runtime_context import get_exec, get_platform
|
||||||
from sglang.srt.utils import is_flashinfer_available, log_info_on_rank0
|
from sglang.srt.utils import is_flashinfer_available, log_info_on_rank0
|
||||||
|
|
||||||
# Suppress TRT-LLM CUTLASS trace logs without overriding user configuration.
|
# Suppress TRT-LLM CUTLASS trace logs without overriding user configuration.
|
||||||
@@ -30,7 +30,7 @@ _GROUP_SIZE = 32
|
|||||||
|
|
||||||
|
|
||||||
class Mxfp4FlashinferCutlassMoEMethod:
|
class Mxfp4FlashinferCutlassMoEMethod:
|
||||||
"""FlashInfer MXFP4 MoE: W4A16 on SM90 and W4A8 on SM120."""
|
"""FlashInfer MXFP4 MoE: W4A16/W4A8 on SM90 and W4A8 on SM120."""
|
||||||
|
|
||||||
fuse_routed_scaling_factor_in_topk = True
|
fuse_routed_scaling_factor_in_topk = True
|
||||||
|
|
||||||
@@ -38,6 +38,11 @@ class Mxfp4FlashinferCutlassMoEMethod:
|
|||||||
if not is_flashinfer_available():
|
if not is_flashinfer_available():
|
||||||
raise RuntimeError("Mxfp4FlashinferCutlassMoEMethod requires FlashInfer.")
|
raise RuntimeError("Mxfp4FlashinferCutlassMoEMethod requires FlashInfer.")
|
||||||
self._use_mxfp8_act_scaling = get_platform().is_sm120
|
self._use_mxfp8_act_scaling = get_platform().is_sm120
|
||||||
|
precision = get_exec().moe.flashinfer_mxfp4_moe_precision
|
||||||
|
# precision=fp8 is an SM90 knob (Humming W4A8); on SM120 the MXFP8
|
||||||
|
# activation path already computes in FP8, so the flag is simply inert
|
||||||
|
# there rather than an error -- one config can move across hardware.
|
||||||
|
self._use_sm90_humming = not self._use_mxfp8_act_scaling and precision == "fp8"
|
||||||
self._fp8 = fp8_method
|
self._fp8 = fp8_method
|
||||||
self.prefix = prefix
|
self.prefix = prefix
|
||||||
self._swiglu_limit_tensor: torch.Tensor | None = None
|
self._swiglu_limit_tensor: torch.Tensor | None = None
|
||||||
@@ -124,9 +129,12 @@ class Mxfp4FlashinferCutlassMoEMethod:
|
|||||||
return
|
return
|
||||||
|
|
||||||
arch = "SM120" if self._use_mxfp8_act_scaling else "SM90"
|
arch = "SM120" if self._use_mxfp8_act_scaling else "SM90"
|
||||||
|
precision = (
|
||||||
|
"W4A8" if self._use_sm90_humming or self._use_mxfp8_act_scaling else "W4A16"
|
||||||
|
)
|
||||||
log_info_on_rank0(
|
log_info_on_rank0(
|
||||||
logger,
|
logger,
|
||||||
f"Preparing DSv4 MXFP4 experts for FlashInfer {arch} CUTLASS "
|
f"Preparing DSv4 MXFP4 experts for FlashInfer {arch} CUTLASS {precision} "
|
||||||
f"(layer: {self.prefix})...",
|
f"(layer: {self.prefix})...",
|
||||||
)
|
)
|
||||||
|
|
||||||
@@ -152,23 +160,51 @@ class Mxfp4FlashinferCutlassMoEMethod:
|
|||||||
for scale_u8 in (w13_scale_u8, w2_scale_u8):
|
for scale_u8 in (w13_scale_u8, w2_scale_u8):
|
||||||
scale_u8.copy_(block_scale_interleave(scale_u8).reshape_as(scale_u8))
|
scale_u8.copy_(block_scale_interleave(scale_u8).reshape_as(scale_u8))
|
||||||
else:
|
else:
|
||||||
from flashinfer.fused_moe import (
|
if self._use_sm90_humming:
|
||||||
interleave_moe_scales_for_sm90_mixed_gemm,
|
from flashinfer.fused_moe import (
|
||||||
interleave_moe_weights_for_sm90_mixed_gemm,
|
preprocess_moe_weights_for_sm90_mixed_gemm_humming,
|
||||||
)
|
)
|
||||||
|
|
||||||
w13_il = interleave_moe_weights_for_sm90_mixed_gemm(
|
w13_il, w13_s_il, w13_residual = (
|
||||||
layer.w13_weight.data.view(torch.uint8).contiguous(), "fp4"
|
preprocess_moe_weights_for_sm90_mixed_gemm_humming(
|
||||||
)
|
layer.w13_weight.data.view(torch.uint8).contiguous(),
|
||||||
w2_il = interleave_moe_weights_for_sm90_mixed_gemm(
|
w13_scale_u8,
|
||||||
layer.w2_weight.data.view(torch.uint8).contiguous(), "fp4"
|
)
|
||||||
)
|
)
|
||||||
w13_s_il = interleave_moe_scales_for_sm90_mixed_gemm(
|
w2_il, w2_s_il, w2_residual = (
|
||||||
w13_scale_u8, group_size=_GROUP_SIZE
|
preprocess_moe_weights_for_sm90_mixed_gemm_humming(
|
||||||
)
|
layer.w2_weight.data.view(torch.uint8).contiguous(),
|
||||||
w2_s_il = interleave_moe_scales_for_sm90_mixed_gemm(
|
w2_scale_u8,
|
||||||
w2_scale_u8, group_size=_GROUP_SIZE
|
)
|
||||||
)
|
)
|
||||||
|
layer.w13_humming_residual_scale = Parameter(
|
||||||
|
(w13_residual * 64.0).contiguous(), requires_grad=False
|
||||||
|
)
|
||||||
|
layer.w2_humming_residual_scale = Parameter(
|
||||||
|
(w2_residual * 64.0).contiguous(), requires_grad=False
|
||||||
|
)
|
||||||
|
layer.humming_fc2_act_scale = Parameter(
|
||||||
|
torch.ones((), dtype=torch.float32, device=w13_scale_u8.device),
|
||||||
|
requires_grad=False,
|
||||||
|
)
|
||||||
|
else:
|
||||||
|
from flashinfer.fused_moe import (
|
||||||
|
interleave_moe_scales_for_sm90_mixed_gemm,
|
||||||
|
interleave_moe_weights_for_sm90_mixed_gemm,
|
||||||
|
)
|
||||||
|
|
||||||
|
w13_il = interleave_moe_weights_for_sm90_mixed_gemm(
|
||||||
|
layer.w13_weight.data.view(torch.uint8).contiguous(), "fp4"
|
||||||
|
)
|
||||||
|
w2_il = interleave_moe_weights_for_sm90_mixed_gemm(
|
||||||
|
layer.w2_weight.data.view(torch.uint8).contiguous(), "fp4"
|
||||||
|
)
|
||||||
|
w13_s_il = interleave_moe_scales_for_sm90_mixed_gemm(
|
||||||
|
w13_scale_u8, group_size=_GROUP_SIZE
|
||||||
|
)
|
||||||
|
w2_s_il = interleave_moe_scales_for_sm90_mixed_gemm(
|
||||||
|
w2_scale_u8, group_size=_GROUP_SIZE
|
||||||
|
)
|
||||||
layer.w13_weight = Parameter(w13_il, requires_grad=False)
|
layer.w13_weight = Parameter(w13_il, requires_grad=False)
|
||||||
layer.w2_weight = Parameter(w2_il, requires_grad=False)
|
layer.w2_weight = Parameter(w2_il, requires_grad=False)
|
||||||
layer.w13_weight_scale_inv = Parameter(w13_s_il, requires_grad=False)
|
layer.w13_weight_scale_inv = Parameter(w13_s_il, requires_grad=False)
|
||||||
@@ -177,7 +213,11 @@ class Mxfp4FlashinferCutlassMoEMethod:
|
|||||||
layer._dsv4_mxfp4_backend = (
|
layer._dsv4_mxfp4_backend = (
|
||||||
"flashinfer_cutlass_sm120"
|
"flashinfer_cutlass_sm120"
|
||||||
if self._use_mxfp8_act_scaling
|
if self._use_mxfp8_act_scaling
|
||||||
else "flashinfer_cutlass_sm90"
|
else (
|
||||||
|
"flashinfer_cutlass_sm90_fp8"
|
||||||
|
if self._use_sm90_humming
|
||||||
|
else "flashinfer_cutlass_sm90"
|
||||||
|
)
|
||||||
)
|
)
|
||||||
# SM90 creates full-size interleaved copies; release old layouts per layer.
|
# SM90 creates full-size interleaved copies; release old layouts per layer.
|
||||||
if not self._use_mxfp8_act_scaling:
|
if not self._use_mxfp8_act_scaling:
|
||||||
@@ -198,6 +238,11 @@ class Mxfp4FlashinferCutlassMoEMethod:
|
|||||||
w13_weight_scale=layer.w13_weight_scale_inv,
|
w13_weight_scale=layer.w13_weight_scale_inv,
|
||||||
w2_weight_scale=layer.w2_weight_scale_inv,
|
w2_weight_scale=layer.w2_weight_scale_inv,
|
||||||
mxfp4_weight_global_scale=self._mxfp4_weight_global_scale_tensor,
|
mxfp4_weight_global_scale=self._mxfp4_weight_global_scale_tensor,
|
||||||
|
w13_humming_residual_scale=getattr(
|
||||||
|
layer, "w13_humming_residual_scale", None
|
||||||
|
),
|
||||||
|
w2_humming_residual_scale=getattr(layer, "w2_humming_residual_scale", None),
|
||||||
|
humming_fc2_act_scale=getattr(layer, "humming_fc2_act_scale", None),
|
||||||
w13_bias=None,
|
w13_bias=None,
|
||||||
w2_bias=None,
|
w2_bias=None,
|
||||||
swiglu_alpha=None,
|
swiglu_alpha=None,
|
||||||
|
|||||||
@@ -54,6 +54,9 @@ class Mxfp4FlashinferTrtllmMoEMethod:
|
|||||||
def __init__(self, fp8_method, prefix: str):
|
def __init__(self, fp8_method, prefix: str):
|
||||||
self._fp8 = fp8_method
|
self._fp8 = fp8_method
|
||||||
self.prefix = prefix
|
self.prefix = prefix
|
||||||
|
# precision=fp8 is an SM90 knob (Humming W4A8); this SM100 trtllm path
|
||||||
|
# already runs MXFP8 activations, so the flag is inert here rather than
|
||||||
|
# an error -- one config can move across hardware.
|
||||||
self.flashinfer_mxfp4_moe_precision = (
|
self.flashinfer_mxfp4_moe_precision = (
|
||||||
get_exec().moe.flashinfer_mxfp4_moe_precision
|
get_exec().moe.flashinfer_mxfp4_moe_precision
|
||||||
)
|
)
|
||||||
|
|||||||
@@ -2422,8 +2422,10 @@ class ServerArgs:
|
|||||||
NS("exec.moe"),
|
NS("exec.moe"),
|
||||||
] = "auto"
|
] = "auto"
|
||||||
flashinfer_mxfp4_moe_precision: A[
|
flashinfer_mxfp4_moe_precision: A[
|
||||||
Literal["default", "bf16"],
|
Literal["default", "bf16", "fp8"],
|
||||||
"Choose the computation precision of flashinfer mxfp4 moe",
|
"Choose the computation precision of flashinfer mxfp4 moe. "
|
||||||
|
"On SM90, `fp8` selects the Humming-style MXFP4-weight x FP8-activation "
|
||||||
|
"path introduced by FlashInfer #3738 and requires FlashInfer >= 0.6.18.",
|
||||||
NS("exec.moe"),
|
NS("exec.moe"),
|
||||||
] = "default"
|
] = "default"
|
||||||
deepep_mode: A[
|
deepep_mode: A[
|
||||||
|
|||||||
@@ -1,10 +1,10 @@
|
|||||||
"""Benchmark MXFP4 MoE kernels on H100/H200: SGLang Marlin vs FlashInfer cutlass.
|
"""Benchmark MXFP4 MoE kernels on H100/H200: Marlin vs FlashInfer CUTLASS.
|
||||||
|
|
||||||
Compares per-call latency of:
|
Compares per-call latency of:
|
||||||
|
|
||||||
* Marlin path : ``fused_marlin_moe(...)`` after Marlin weight repack
|
* Marlin path : ``fused_marlin_moe(...)`` after Marlin weight repack
|
||||||
* FlashInfer : ``cutlass_fused_moe(use_w4_group_scaling=True, ...)``
|
* FlashInfer W4A16 : PR #3084's SM90 mixed-input path
|
||||||
(PR #3084's SM90 mixed-input path)
|
* FlashInfer W4A8 : PR #3738/#4431's corrected Humming path, when available
|
||||||
|
|
||||||
Both run on the same random MXFP4 weights/scales (semantics differ slightly --
|
Both run on the same random MXFP4 weights/scales (semantics differ slightly --
|
||||||
Marlin uses a scalar swiglu clamp + no bias, FlashInfer fuses per-expert
|
Marlin uses a scalar swiglu clamp + no bias, FlashInfer fuses per-expert
|
||||||
@@ -15,7 +15,7 @@ Run on H100/H200:
|
|||||||
|
|
||||||
cd /sgl-workspace/sglang_dev3 && \\
|
cd /sgl-workspace/sglang_dev3 && \\
|
||||||
PYTHONPATH=python:/sgl-workspace/flashinfer FLASHINFER_DISABLE_VERSION_CHECK=1 \\
|
PYTHONPATH=python:/sgl-workspace/flashinfer FLASHINFER_DISABLE_VERSION_CHECK=1 \\
|
||||||
python python/sglang/test/bench_mxfp4_sm90_kernels.py
|
python test/manual/layers/moe/bench_mxfp4_sm90_kernels.py
|
||||||
"""
|
"""
|
||||||
|
|
||||||
from __future__ import annotations
|
from __future__ import annotations
|
||||||
@@ -25,15 +25,29 @@ from dataclasses import dataclass
|
|||||||
from typing import Callable, List, Tuple
|
from typing import Callable, List, Tuple
|
||||||
|
|
||||||
import torch
|
import torch
|
||||||
from flashinfer.autotuner import autotune
|
|
||||||
|
|
||||||
# ---- FlashInfer ----
|
# ---- FlashInfer ----
|
||||||
|
from flashinfer import __version__ as flashinfer_version
|
||||||
|
from flashinfer.autotuner import autotune
|
||||||
from flashinfer.fused_moe import (
|
from flashinfer.fused_moe import (
|
||||||
cutlass_fused_moe,
|
cutlass_fused_moe,
|
||||||
interleave_moe_scales_for_sm90_mixed_gemm,
|
interleave_moe_scales_for_sm90_mixed_gemm,
|
||||||
interleave_moe_weights_for_sm90_mixed_gemm,
|
interleave_moe_weights_for_sm90_mixed_gemm,
|
||||||
)
|
)
|
||||||
from flashinfer.fused_moe.core import ActivationType
|
from flashinfer.fused_moe.core import ActivationType
|
||||||
|
from packaging.version import Version
|
||||||
|
|
||||||
|
try:
|
||||||
|
from flashinfer.fused_moe import (
|
||||||
|
preprocess_moe_weights_for_sm90_mixed_gemm_humming,
|
||||||
|
)
|
||||||
|
except ImportError:
|
||||||
|
preprocess_moe_weights_for_sm90_mixed_gemm_humming = None
|
||||||
|
|
||||||
|
_fi_release = Version(flashinfer_version).release
|
||||||
|
_fi_release = _fi_release + (0,) * (3 - len(_fi_release))
|
||||||
|
if _fi_release[:3] < (0, 6, 18):
|
||||||
|
preprocess_moe_weights_for_sm90_mixed_gemm_humming = None
|
||||||
|
|
||||||
# ---- SGLang Marlin ----
|
# ---- SGLang Marlin ----
|
||||||
from sglang.kernels.ops.quantization.gptq_marlin_repack import gptq_marlin_repack
|
from sglang.kernels.ops.quantization.gptq_marlin_repack import gptq_marlin_repack
|
||||||
@@ -162,14 +176,50 @@ def build_flashinfer_inputs(shape: Shape, w13, w2, w13_s, w2_s, w13_b, w2_b):
|
|||||||
}
|
}
|
||||||
|
|
||||||
|
|
||||||
|
def build_flashinfer_humming_inputs(shape: Shape, w13, w2, w13_s, w2_s, w13_b, w2_b):
|
||||||
|
if preprocess_moe_weights_for_sm90_mixed_gemm_humming is None:
|
||||||
|
raise RuntimeError("FlashInfer does not provide the corrected Humming API.")
|
||||||
|
w13_il, w13_s_il, w13_residual = preprocess_moe_weights_for_sm90_mixed_gemm_humming(
|
||||||
|
w13, w13_s
|
||||||
|
)
|
||||||
|
w2_il, w2_s_il, w2_residual = preprocess_moe_weights_for_sm90_mixed_gemm_humming(
|
||||||
|
w2, w2_s
|
||||||
|
)
|
||||||
|
e = shape.num_experts
|
||||||
|
return {
|
||||||
|
"w13": w13_il,
|
||||||
|
"w2": w2_il,
|
||||||
|
"quant_scales": [
|
||||||
|
w13_s_il.view(torch.int32),
|
||||||
|
(w13_residual * 64.0).contiguous(),
|
||||||
|
torch.ones((), dtype=torch.float32, device="cuda"),
|
||||||
|
w2_s_il.view(torch.int32),
|
||||||
|
(w2_residual * 64.0).contiguous(),
|
||||||
|
],
|
||||||
|
"w13_b": w13_b,
|
||||||
|
"w2_b": w2_b,
|
||||||
|
"swiglu_alpha": torch.full((e,), 1.702, dtype=torch.float32, device="cuda"),
|
||||||
|
"swiglu_beta": torch.full((e,), 1.0, dtype=torch.float32, device="cuda"),
|
||||||
|
"swiglu_limit": torch.full((e,), 7.0, dtype=torch.float32, device="cuda"),
|
||||||
|
}
|
||||||
|
|
||||||
|
|
||||||
def make_flashinfer_runner(
|
def make_flashinfer_runner(
|
||||||
shape: Shape, prep, x, topk_w, topk_i, autotuned: bool, with_bias: bool = True
|
shape: Shape,
|
||||||
|
prep,
|
||||||
|
x,
|
||||||
|
topk_w,
|
||||||
|
topk_i,
|
||||||
|
autotuned: bool,
|
||||||
|
with_bias: bool = True,
|
||||||
|
use_humming: bool = False,
|
||||||
):
|
):
|
||||||
out = torch.empty(shape.tokens, shape.hidden, dtype=torch.bfloat16, device="cuda")
|
out = torch.empty(shape.tokens, shape.hidden, dtype=torch.bfloat16, device="cuda")
|
||||||
fc1_b = prep["w13_b"] if with_bias else None
|
fc1_b = prep["w13_b"] if with_bias else None
|
||||||
fc2_b = prep["w2_b"] if with_bias else None
|
fc2_b = prep["w2_b"] if with_bias else None
|
||||||
|
|
||||||
def _call():
|
def _call():
|
||||||
|
humming_kwargs = {"use_wfp4afp8_humming": True} if use_humming else {}
|
||||||
cutlass_fused_moe(
|
cutlass_fused_moe(
|
||||||
input=x,
|
input=x,
|
||||||
token_selected_experts=topk_i,
|
token_selected_experts=topk_i,
|
||||||
@@ -186,6 +236,7 @@ def make_flashinfer_runner(
|
|||||||
use_w4_group_scaling=True,
|
use_w4_group_scaling=True,
|
||||||
activation_type=ActivationType.Swiglu,
|
activation_type=ActivationType.Swiglu,
|
||||||
output=out,
|
output=out,
|
||||||
|
**humming_kwargs,
|
||||||
)
|
)
|
||||||
|
|
||||||
if autotuned:
|
if autotuned:
|
||||||
@@ -328,6 +379,29 @@ def run_one_shape(shape: Shape, run_marlin: bool):
|
|||||||
)
|
)
|
||||||
fi_med = fi_at_med # alias for downstream speedup print
|
fi_med = fi_at_med # alias for downstream speedup print
|
||||||
|
|
||||||
|
if preprocess_moe_weights_for_sm90_mixed_gemm_humming is not None:
|
||||||
|
humming_prep = build_flashinfer_humming_inputs(
|
||||||
|
shape, w13, w2, w13_s, w2_s, w13_b, w2_b
|
||||||
|
)
|
||||||
|
humming_call = make_flashinfer_runner(
|
||||||
|
shape,
|
||||||
|
humming_prep,
|
||||||
|
x,
|
||||||
|
topk_w,
|
||||||
|
topk_i,
|
||||||
|
autotuned=True,
|
||||||
|
with_bias=True,
|
||||||
|
use_humming=True,
|
||||||
|
)
|
||||||
|
humming_med, humming_min = time_call(humming_call)
|
||||||
|
print(
|
||||||
|
f" FlashInfer Humming W4A8: median={humming_med:.3f} ms "
|
||||||
|
f"min={humming_min:.3f} ms"
|
||||||
|
)
|
||||||
|
print(f" speedup (FI W4A16 / FI W4A8): {fi_med / humming_med:.2f}x")
|
||||||
|
else:
|
||||||
|
print(" FlashInfer Humming W4A8: SKIPPED (requires >= 0.6.18)")
|
||||||
|
|
||||||
# Marlin
|
# Marlin
|
||||||
if run_marlin:
|
if run_marlin:
|
||||||
try:
|
try:
|
||||||
|
|||||||
@@ -84,6 +84,7 @@ def test_cutlass_adapter_import_does_not_require_flashinfer(monkeypatch):
|
|||||||
|
|
||||||
def test_dsv4_sm120_load_contract(monkeypatch, request):
|
def test_dsv4_sm120_load_contract(monkeypatch, request):
|
||||||
import sglang.srt.layers.quantization.mxfp4_flashinfer_cutlass_moe as adapter_module
|
import sglang.srt.layers.quantization.mxfp4_flashinfer_cutlass_moe as adapter_module
|
||||||
|
from sglang.srt.runtime_context import get_context
|
||||||
|
|
||||||
platform = override_platform(is_sm120=True)
|
platform = override_platform(is_sm120=True)
|
||||||
platform.install()
|
platform.install()
|
||||||
@@ -95,7 +96,8 @@ def test_dsv4_sm120_load_contract(monkeypatch, request):
|
|||||||
def create_weights(self, *args, **kwargs):
|
def create_weights(self, *args, **kwargs):
|
||||||
captured.update(kwargs)
|
captured.update(kwargs)
|
||||||
|
|
||||||
method = adapter_module.Mxfp4FlashinferCutlassMoEMethod(_Fp8Method(), "test")
|
with get_context().override_server_args(flashinfer_mxfp4_moe_precision="default"):
|
||||||
|
method = adapter_module.Mxfp4FlashinferCutlassMoEMethod(_Fp8Method(), "test")
|
||||||
method.create_weights(
|
method.create_weights(
|
||||||
SimpleNamespace(),
|
SimpleNamespace(),
|
||||||
num_experts=4,
|
num_experts=4,
|
||||||
@@ -126,6 +128,7 @@ def test_dsv4_sm120_matches_direct_flashinfer(monkeypatch):
|
|||||||
from sglang.srt.layers.quantization.mxfp4_flashinfer_cutlass_moe import (
|
from sglang.srt.layers.quantization.mxfp4_flashinfer_cutlass_moe import (
|
||||||
Mxfp4FlashinferCutlassMoEMethod,
|
Mxfp4FlashinferCutlassMoEMethod,
|
||||||
)
|
)
|
||||||
|
from sglang.srt.runtime_context import get_context
|
||||||
|
|
||||||
monkeypatch.setattr(
|
monkeypatch.setattr(
|
||||||
runner_module, "use_symmetric_memory", lambda *args, **kwargs: nullcontext()
|
runner_module, "use_symmetric_memory", lambda *args, **kwargs: nullcontext()
|
||||||
@@ -155,9 +158,10 @@ def test_dsv4_sm120_matches_direct_flashinfer(monkeypatch):
|
|||||||
moe_ep_rank=0,
|
moe_ep_rank=0,
|
||||||
)
|
)
|
||||||
|
|
||||||
method = Mxfp4FlashinferCutlassMoEMethod(
|
with get_context().override_server_args(flashinfer_mxfp4_moe_precision="default"):
|
||||||
SimpleNamespace(process_weights_after_loading=lambda layer: None), "test"
|
method = Mxfp4FlashinferCutlassMoEMethod(
|
||||||
)
|
SimpleNamespace(process_weights_after_loading=lambda layer: None), "test"
|
||||||
|
)
|
||||||
config = MoeRunnerConfig(
|
config = MoeRunnerConfig(
|
||||||
num_experts=num_experts,
|
num_experts=num_experts,
|
||||||
num_local_experts=num_experts,
|
num_local_experts=num_experts,
|
||||||
|
|||||||
@@ -2,9 +2,9 @@
|
|||||||
|
|
||||||
Builds a single-layer GPT-OSS-style MoE with random MXFP4 weights, drives the
|
Builds a single-layer GPT-OSS-style MoE with random MXFP4 weights, drives the
|
||||||
SGLang plumbing (``_process_weights_for_sm90_cutlass`` + ``_apply_sm90_cutlass``)
|
SGLang plumbing (``_process_weights_for_sm90_cutlass`` + ``_apply_sm90_cutlass``)
|
||||||
and compares against a direct FlashInfer ``cutlass_fused_moe`` call with the
|
and compares against direct FlashInfer ``cutlass_fused_moe`` calls. It covers
|
||||||
same inputs. Both paths invoke the same SM90 kernel from FlashInfer PR #3084,
|
both PR #3084's W4A16 path and PR #3738/#4431's corrected Humming W4A8 path;
|
||||||
so outputs must be bit-exact.
|
outputs must be bit-exact within each path.
|
||||||
|
|
||||||
Run on H100/H200:
|
Run on H100/H200:
|
||||||
|
|
||||||
@@ -24,6 +24,16 @@ register_cuda_ci(est_time=120, stage="base-b", runner_config="1-gpu-large")
|
|||||||
|
|
||||||
flashinfer_fused_moe = pytest.importorskip("flashinfer.fused_moe")
|
flashinfer_fused_moe = pytest.importorskip("flashinfer.fused_moe")
|
||||||
|
|
||||||
|
HAS_CORRECTED_HUMMING_API = hasattr(
|
||||||
|
flashinfer_fused_moe,
|
||||||
|
"preprocess_moe_weights_for_sm90_mixed_gemm_humming",
|
||||||
|
)
|
||||||
|
preprocess_humming = getattr(
|
||||||
|
flashinfer_fused_moe,
|
||||||
|
"preprocess_moe_weights_for_sm90_mixed_gemm_humming",
|
||||||
|
None,
|
||||||
|
)
|
||||||
|
|
||||||
if not hasattr(flashinfer_fused_moe, "interleave_moe_weights_for_sm90_mixed_gemm"):
|
if not hasattr(flashinfer_fused_moe, "interleave_moe_weights_for_sm90_mixed_gemm"):
|
||||||
pytest.skip(
|
pytest.skip(
|
||||||
"FlashInfer build does not include PR #3084 SM90 mixed-input helpers",
|
"FlashInfer build does not include PR #3084 SM90 mixed-input helpers",
|
||||||
@@ -154,11 +164,12 @@ def _round_up(x, base):
|
|||||||
return ((x + base - 1) // base) * base
|
return ((x + base - 1) // base) * base
|
||||||
|
|
||||||
|
|
||||||
def _build_method(num_experts, hidden, inter):
|
def _build_method(num_experts, hidden, inter, *, use_humming=False):
|
||||||
from sglang.srt.layers.quantization.mxfp4 import Mxfp4MoEMethod
|
from sglang.srt.layers.quantization.mxfp4 import Mxfp4MoEMethod
|
||||||
|
|
||||||
method = Mxfp4MoEMethod.__new__(Mxfp4MoEMethod)
|
method = Mxfp4MoEMethod.__new__(Mxfp4MoEMethod)
|
||||||
method._fi_kernel = "cutlass_sm90"
|
method._fi_kernel = "cutlass_sm90"
|
||||||
|
method._use_sm90_humming = use_humming
|
||||||
method.num_experts = num_experts
|
method.num_experts = num_experts
|
||||||
# The new SM90 cutlass path tracks padded sizes in dedicated attrs;
|
# The new SM90 cutlass path tracks padded sizes in dedicated attrs;
|
||||||
# ``hidden_size`` / ``intermediate_size_per_partition`` keep the unpadded
|
# ``hidden_size`` / ``intermediate_size_per_partition`` keep the unpadded
|
||||||
@@ -348,6 +359,11 @@ def test_apply_sm90_cutlass_matches_flashinfer_direct(
|
|||||||
)
|
)
|
||||||
monkeypatch.setattr(fi_cutlass_mod, "is_allocation_symmetric", lambda: False)
|
monkeypatch.setattr(fi_cutlass_mod, "is_allocation_symmetric", lambda: False)
|
||||||
monkeypatch.setattr(fi_cutlass_mod, "get_tp_group", lambda: None)
|
monkeypatch.setattr(fi_cutlass_mod, "get_tp_group", lambda: None)
|
||||||
|
monkeypatch.setattr(
|
||||||
|
fi_cutlass_mod.envs.SGLANG_FLASHINFER_MOE_FUSED_FINALIZE,
|
||||||
|
"get",
|
||||||
|
lambda: False,
|
||||||
|
)
|
||||||
|
|
||||||
w13, w2, w13_s, w2_s, w13_b, w2_b = _make_random_mxfp4(num_experts, hidden, inter)
|
w13, w2, w13_s, w2_s, w13_b, w2_b = _make_random_mxfp4(num_experts, hidden, inter)
|
||||||
x = torch.randn(tokens, hidden, dtype=torch.bfloat16, device="cuda") * 0.1
|
x = torch.randn(tokens, hidden, dtype=torch.bfloat16, device="cuda") * 0.1
|
||||||
@@ -392,6 +408,7 @@ def test_apply_sm90_cutlass_matches_flashinfer_direct(
|
|||||||
swiglu_limit=layer.swiglu_limit,
|
swiglu_limit=layer.swiglu_limit,
|
||||||
use_w4_group_scaling=True,
|
use_w4_group_scaling=True,
|
||||||
activation_type=ActivationType.Swiglu,
|
activation_type=ActivationType.Swiglu,
|
||||||
|
use_fused_finalize=False,
|
||||||
output=out_ref_padded,
|
output=out_ref_padded,
|
||||||
)
|
)
|
||||||
out_ref = (
|
out_ref = (
|
||||||
@@ -404,6 +421,258 @@ def test_apply_sm90_cutlass_matches_flashinfer_direct(
|
|||||||
)
|
)
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.skipif(
|
||||||
|
not HAS_CORRECTED_HUMMING_API,
|
||||||
|
reason="requires corrected per-expert Humming API from FlashInfer >= 0.6.18",
|
||||||
|
)
|
||||||
|
def test_process_weights_humming_matches_flashinfer_direct():
|
||||||
|
"""The SM90 fp8 option must use #3738's preprocessing and retain #4431's
|
||||||
|
per-local-expert residual contract."""
|
||||||
|
num_experts, hidden, inter = 4, 256, 256
|
||||||
|
w13, w2, w13_s, w2_s, w13_b, w2_b = _make_random_mxfp4(num_experts, hidden, inter)
|
||||||
|
|
||||||
|
# Make each expert's residual distinct so an accidental scalar/broadcast
|
||||||
|
# contract cannot pass this check.
|
||||||
|
expert_offsets = torch.arange(num_experts, dtype=torch.uint8, device="cuda").view(
|
||||||
|
-1, 1, 1
|
||||||
|
)
|
||||||
|
w13_s = w13_s + expert_offsets
|
||||||
|
w2_s = w2_s + expert_offsets + 1
|
||||||
|
|
||||||
|
layer = _build_mock_layer(
|
||||||
|
num_experts, hidden, inter, w13, w2, w13_s, w2_s, w13_b, w2_b
|
||||||
|
)
|
||||||
|
method = _build_method(num_experts, hidden, inter, use_humming=True)
|
||||||
|
method._process_weights_for_sm90_cutlass(layer)
|
||||||
|
|
||||||
|
# GPT-OSS loads pair-wise [gate, up]; FlashInfer consumes halved [up; gate].
|
||||||
|
ref_w13 = torch.cat((w13[:, 1::2], w13[:, 0::2]), dim=1).contiguous()
|
||||||
|
ref_w13_s = torch.cat((w13_s[:, 1::2], w13_s[:, 0::2]), dim=1).contiguous()
|
||||||
|
expected_w13, expected_w13_s, expected_w13_residual = preprocess_humming(
|
||||||
|
ref_w13, ref_w13_s
|
||||||
|
)
|
||||||
|
expected_w2, expected_w2_s, expected_w2_residual = preprocess_humming(w2, w2_s)
|
||||||
|
|
||||||
|
assert torch.equal(layer.w13_weight, expected_w13)
|
||||||
|
assert torch.equal(layer.w2_weight, expected_w2)
|
||||||
|
assert torch.equal(layer.w13_weight_scale, expected_w13_s)
|
||||||
|
assert torch.equal(layer.w2_weight_scale, expected_w2_s)
|
||||||
|
assert torch.equal(layer.w13_humming_residual_scale, expected_w13_residual * 64.0)
|
||||||
|
assert torch.equal(layer.w2_humming_residual_scale, expected_w2_residual * 64.0)
|
||||||
|
assert layer.w13_humming_residual_scale.shape == (num_experts,)
|
||||||
|
assert layer.w2_humming_residual_scale.shape == (num_experts,)
|
||||||
|
assert layer.humming_fc2_act_scale.shape == ()
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.skipif(
|
||||||
|
not HAS_CORRECTED_HUMMING_API,
|
||||||
|
reason="requires corrected per-expert Humming API from FlashInfer >= 0.6.18",
|
||||||
|
)
|
||||||
|
def test_humming_padding_preserves_per_expert_residual():
|
||||||
|
"""Synthetic alignment padding must not change an expert's E8M0 range."""
|
||||||
|
num_experts, hidden, inter = 4, 192, 192
|
||||||
|
w13, w2, w13_s, w2_s, w13_b, w2_b = _make_random_mxfp4(num_experts, hidden, inter)
|
||||||
|
ref_w13 = torch.cat((w13[:, 1::2], w13[:, 0::2]), dim=1).contiguous()
|
||||||
|
ref_w13_s = torch.cat((w13_s[:, 1::2], w13_s[:, 0::2]), dim=1).contiguous()
|
||||||
|
_, _, expected_w13_residual = preprocess_humming(
|
||||||
|
ref_w13, ref_w13_s, interleave=False
|
||||||
|
)
|
||||||
|
_, _, expected_w2_residual = preprocess_humming(w2, w2_s, interleave=False)
|
||||||
|
|
||||||
|
layer = _build_mock_layer(
|
||||||
|
num_experts, hidden, inter, w13, w2, w13_s, w2_s, w13_b, w2_b
|
||||||
|
)
|
||||||
|
method = _build_method(num_experts, hidden, inter, use_humming=True)
|
||||||
|
method._process_weights_for_sm90_cutlass(layer)
|
||||||
|
|
||||||
|
assert torch.equal(layer.w13_humming_residual_scale, expected_w13_residual * 64.0)
|
||||||
|
assert torch.equal(layer.w2_humming_residual_scale, expected_w2_residual * 64.0)
|
||||||
|
|
||||||
|
|
||||||
|
def _build_prerounded_case(E, hidden_real, hidden_rounded, inter, tail_fill, seed=0):
|
||||||
|
"""Weights as ``create_weights`` leaves them when FusedMoE pre-rounds hidden.
|
||||||
|
|
||||||
|
Buffers are allocated at ``hidden_rounded``; the loader only ever writes the
|
||||||
|
first ``hidden_real`` columns, so the tail keeps whatever the buffer was
|
||||||
|
filled with (``_UE8M0_ONE`` in production). Real scales sit well BELOW 2^0
|
||||||
|
so a leaked tail moves the per-expert max.
|
||||||
|
"""
|
||||||
|
g = torch.Generator(device="cuda").manual_seed(seed)
|
||||||
|
kr_bytes = hidden_real // 2
|
||||||
|
kr_grp = hidden_real // GROUP_SIZE
|
||||||
|
|
||||||
|
w13 = torch.zeros(
|
||||||
|
(E, 2 * inter, hidden_rounded // 2), dtype=torch.uint8, device="cuda"
|
||||||
|
)
|
||||||
|
w13[:, :, :kr_bytes] = torch.randint(
|
||||||
|
0, 256, (E, 2 * inter, kr_bytes), dtype=torch.uint8, device="cuda", generator=g
|
||||||
|
)
|
||||||
|
w2 = torch.zeros((E, hidden_rounded, inter // 2), dtype=torch.uint8, device="cuda")
|
||||||
|
w2[:, :hidden_real, :] = torch.randint(
|
||||||
|
0,
|
||||||
|
256,
|
||||||
|
(E, hidden_real, inter // 2),
|
||||||
|
dtype=torch.uint8,
|
||||||
|
device="cuda",
|
||||||
|
generator=g,
|
||||||
|
)
|
||||||
|
|
||||||
|
w13_s = torch.full(
|
||||||
|
(E, 2 * inter, hidden_rounded // GROUP_SIZE),
|
||||||
|
tail_fill,
|
||||||
|
dtype=torch.uint8,
|
||||||
|
device="cuda",
|
||||||
|
)
|
||||||
|
w13_s[:, :, :kr_grp] = torch.randint(
|
||||||
|
100, 110, (E, 2 * inter, kr_grp), dtype=torch.uint8, device="cuda", generator=g
|
||||||
|
)
|
||||||
|
w2_s = torch.full(
|
||||||
|
(E, hidden_rounded, inter // GROUP_SIZE),
|
||||||
|
tail_fill,
|
||||||
|
dtype=torch.uint8,
|
||||||
|
device="cuda",
|
||||||
|
)
|
||||||
|
w2_s[:, :hidden_real, :] = torch.randint(
|
||||||
|
100,
|
||||||
|
110,
|
||||||
|
(E, hidden_real, inter // GROUP_SIZE),
|
||||||
|
dtype=torch.uint8,
|
||||||
|
device="cuda",
|
||||||
|
generator=g,
|
||||||
|
)
|
||||||
|
|
||||||
|
w13_b = torch.zeros((E, 2 * inter), dtype=torch.bfloat16, device="cuda")
|
||||||
|
w2_b = torch.zeros((E, hidden_rounded), dtype=torch.bfloat16, device="cuda")
|
||||||
|
return w13, w2, w13_s, w2_s, w13_b, w2_b
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.skipif(
|
||||||
|
not HAS_CORRECTED_HUMMING_API,
|
||||||
|
reason="requires corrected per-expert Humming API from FlashInfer >= 0.6.18",
|
||||||
|
)
|
||||||
|
def test_humming_range_ignores_prerounded_hidden_tail():
|
||||||
|
"""FusedMoE rounds GPT-OSS hidden 2880 -> 3072 BEFORE ``create_weights``, so
|
||||||
|
the trailing scale columns keep the ``_UE8M0_ONE`` buffer fill. Those bytes
|
||||||
|
are 2^0 -- above any real per-expert max -- and must not reach Humming's
|
||||||
|
min/max, or the residual shifts and perturbs the real weights.
|
||||||
|
|
||||||
|
Invariant: the residual must not depend on what the never-written tail holds.
|
||||||
|
"""
|
||||||
|
from sglang.srt.layers.quantization.mxfp4 import _UE8M0_ONE
|
||||||
|
|
||||||
|
E, hidden_real, hidden_rounded, inter = 4, 2880, 3072, 256
|
||||||
|
|
||||||
|
residuals = []
|
||||||
|
for tail_fill in (_UE8M0_ONE, 105): # 105 sits inside the real 100..110 band
|
||||||
|
w13, w2, w13_s, w2_s, w13_b, w2_b = _build_prerounded_case(
|
||||||
|
E, hidden_real, hidden_rounded, inter, tail_fill
|
||||||
|
)
|
||||||
|
layer = _build_mock_layer(
|
||||||
|
E, hidden_rounded, inter, w13, w2, w13_s, w2_s, w13_b, w2_b
|
||||||
|
)
|
||||||
|
method = _build_method(E, hidden_rounded, inter, use_humming=True)
|
||||||
|
# What create_weights records from layer.hidden_size_unpadded.
|
||||||
|
method._unpadded_hidden = hidden_real
|
||||||
|
method._process_weights_for_sm90_cutlass(layer)
|
||||||
|
residuals.append(
|
||||||
|
(
|
||||||
|
layer.w13_humming_residual_scale.clone(),
|
||||||
|
layer.w2_humming_residual_scale.clone(),
|
||||||
|
)
|
||||||
|
)
|
||||||
|
|
||||||
|
assert torch.equal(residuals[0][0], residuals[1][0]), (
|
||||||
|
"w13 Humming residual changed with the never-written hidden tail; "
|
||||||
|
"the _UE8M0_ONE fill leaked into the per-expert E8M0 range"
|
||||||
|
)
|
||||||
|
assert torch.equal(
|
||||||
|
residuals[0][1], residuals[1][1]
|
||||||
|
), "w2 Humming residual changed with the never-written hidden tail"
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.skipif(
|
||||||
|
not HAS_CORRECTED_HUMMING_API,
|
||||||
|
reason="requires corrected per-expert Humming API from FlashInfer >= 0.6.18",
|
||||||
|
)
|
||||||
|
@pytest.mark.parametrize(
|
||||||
|
"tokens,hidden,inter,ep_size,ep_rank",
|
||||||
|
[(8, 256, 256, 1, 0), (8, 192, 192, 1, 0), (8, 256, 256, 2, 1)],
|
||||||
|
)
|
||||||
|
def test_apply_sm90_humming_matches_flashinfer_direct(
|
||||||
|
tokens, hidden, inter, ep_size, ep_rank, monkeypatch
|
||||||
|
):
|
||||||
|
"""SGLang must forward the five Humming scales and enable the new kernel."""
|
||||||
|
import sglang.srt.layers.moe.moe_runner.flashinfer_cutlass as fi_cutlass_mod
|
||||||
|
|
||||||
|
monkeypatch.setattr(
|
||||||
|
fi_cutlass_mod, "use_symmetric_memory", lambda *a, **kw: nullcontext()
|
||||||
|
)
|
||||||
|
monkeypatch.setattr(fi_cutlass_mod, "is_allocation_symmetric", lambda: False)
|
||||||
|
monkeypatch.setattr(fi_cutlass_mod, "get_tp_group", lambda: None)
|
||||||
|
monkeypatch.setattr(
|
||||||
|
fi_cutlass_mod.envs.SGLANG_FLASHINFER_MOE_FUSED_FINALIZE,
|
||||||
|
"get",
|
||||||
|
lambda: False,
|
||||||
|
)
|
||||||
|
|
||||||
|
num_experts, top_k = 4, 2
|
||||||
|
w13, w2, w13_s, w2_s, w13_b, w2_b = _make_random_mxfp4(num_experts, hidden, inter)
|
||||||
|
x = torch.randn(tokens, hidden, dtype=torch.bfloat16, device="cuda") * 0.1
|
||||||
|
topk_w, topk_i = _make_topk(tokens, num_experts, top_k)
|
||||||
|
topk_i = topk_i + ep_rank * num_experts
|
||||||
|
|
||||||
|
layer = _build_mock_layer(
|
||||||
|
num_experts, hidden, inter, w13, w2, w13_s, w2_s, w13_b, w2_b
|
||||||
|
)
|
||||||
|
layer.moe_ep_size = ep_size
|
||||||
|
layer.moe_ep_rank = ep_rank
|
||||||
|
method = _build_method(num_experts, hidden, inter, use_humming=True)
|
||||||
|
method._process_weights_for_sm90_cutlass(layer)
|
||||||
|
out_sglang = method._apply_sm90_cutlass(
|
||||||
|
layer, _MockDispatchOutput(x.clone(), topk_w, topk_i)
|
||||||
|
).hidden_states
|
||||||
|
|
||||||
|
padded_hidden = method._padded_hidden
|
||||||
|
x_ref = (
|
||||||
|
torch.nn.functional.pad(x, (0, padded_hidden - hidden))
|
||||||
|
if padded_hidden != hidden
|
||||||
|
else x
|
||||||
|
)
|
||||||
|
out_ref_padded = torch.empty(
|
||||||
|
tokens, padded_hidden, dtype=torch.bfloat16, device="cuda"
|
||||||
|
)
|
||||||
|
cutlass_fused_moe(
|
||||||
|
input=x_ref,
|
||||||
|
token_selected_experts=topk_i,
|
||||||
|
token_final_scales=topk_w,
|
||||||
|
fc1_expert_weights=layer.w13_weight,
|
||||||
|
fc2_expert_weights=layer.w2_weight,
|
||||||
|
output_dtype=torch.bfloat16,
|
||||||
|
quant_scales=[
|
||||||
|
layer.w13_weight_scale.view(torch.int32),
|
||||||
|
layer.w13_humming_residual_scale,
|
||||||
|
layer.humming_fc2_act_scale,
|
||||||
|
layer.w2_weight_scale.view(torch.int32),
|
||||||
|
layer.w2_humming_residual_scale,
|
||||||
|
],
|
||||||
|
fc1_expert_biases=layer.w13_weight_bias,
|
||||||
|
fc2_expert_biases=layer.w2_weight_bias,
|
||||||
|
swiglu_alpha=layer.swiglu_alpha,
|
||||||
|
swiglu_beta=layer.swiglu_beta,
|
||||||
|
swiglu_limit=layer.swiglu_limit,
|
||||||
|
ep_size=ep_size,
|
||||||
|
ep_rank=ep_rank,
|
||||||
|
use_w4_group_scaling=True,
|
||||||
|
use_wfp4afp8_humming=True,
|
||||||
|
activation_type=ActivationType.Swiglu,
|
||||||
|
tune_max_num_tokens=tokens,
|
||||||
|
use_fused_finalize=False,
|
||||||
|
output=out_ref_padded,
|
||||||
|
)
|
||||||
|
out_ref = out_ref_padded[:, :hidden].contiguous()
|
||||||
|
assert torch.equal(out_sglang, out_ref)
|
||||||
|
|
||||||
|
|
||||||
# =============================================================================
|
# =============================================================================
|
||||||
# DeepSeek-V4 path: Mxfp4FlashinferCutlassMoEMethod (sibling of Marlin /
|
# DeepSeek-V4 path: Mxfp4FlashinferCutlassMoEMethod (sibling of Marlin /
|
||||||
# trtllm-gen). Wired into fp8.py's get_quant_method when SM90 +
|
# trtllm-gen). Wired into fp8.py's get_quant_method when SM90 +
|
||||||
@@ -494,10 +763,13 @@ def test_dsv4_apply_matches_flashinfer_direct(
|
|||||||
|
|
||||||
# ---- SGLang DSv4 path ----
|
# ---- SGLang DSv4 path ----
|
||||||
# plain SiLU * up — all three SwiGLU scalars None (no clamp configured).
|
# plain SiLU * up — all three SwiGLU scalars None (no clamp configured).
|
||||||
method = ds_mod.Mxfp4FlashinferCutlassMoEMethod(
|
from sglang.srt.runtime_context import get_context
|
||||||
SimpleNamespace(process_weights_after_loading=lambda layer: None),
|
|
||||||
"test",
|
with get_context().override_server_args(flashinfer_mxfp4_moe_precision="default"):
|
||||||
)
|
method = ds_mod.Mxfp4FlashinferCutlassMoEMethod(
|
||||||
|
SimpleNamespace(process_weights_after_loading=lambda layer: None),
|
||||||
|
"test",
|
||||||
|
)
|
||||||
# Wire the unified MoeRunner -> flashinfer_mxfp4 fused func that
|
# Wire the unified MoeRunner -> flashinfer_mxfp4 fused func that
|
||||||
# ``apply`` now dispatches through.
|
# ``apply`` now dispatches through.
|
||||||
method.runner = _build_flashinfer_mxfp4_runner(num_experts, hidden, inter)
|
method.runner = _build_flashinfer_mxfp4_runner(num_experts, hidden, inter)
|
||||||
@@ -559,6 +831,54 @@ def test_dsv4_apply_matches_flashinfer_direct(
|
|||||||
)
|
)
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.skipif(
|
||||||
|
not HAS_CORRECTED_HUMMING_API,
|
||||||
|
reason="requires corrected per-expert Humming API from FlashInfer >= 0.6.18",
|
||||||
|
)
|
||||||
|
def test_dsv4_process_weights_humming_matches_flashinfer_direct():
|
||||||
|
"""DSv4's native [up; gate] layout must use the same #3738 transform."""
|
||||||
|
from types import SimpleNamespace
|
||||||
|
|
||||||
|
import sglang.srt.layers.quantization.mxfp4_flashinfer_cutlass_moe as ds_mod
|
||||||
|
from sglang.srt.runtime_context import get_context
|
||||||
|
|
||||||
|
num_experts, hidden, inter = 4, 256, 256
|
||||||
|
w13, w2, w13_s, w2_s = _make_random_dsv4_mxfp4(num_experts, hidden, inter)
|
||||||
|
w1, w3 = w13.chunk(2, dim=1)
|
||||||
|
w1_s, w3_s = w13_s.chunk(2, dim=1)
|
||||||
|
w31 = torch.cat((w3, w1), dim=1).contiguous()
|
||||||
|
w31_s = torch.cat((w3_s.view(torch.uint8), w1_s.view(torch.uint8)), dim=1).view(
|
||||||
|
torch.float8_e8m0fnu
|
||||||
|
)
|
||||||
|
|
||||||
|
with get_context().override_server_args(flashinfer_mxfp4_moe_precision="fp8"):
|
||||||
|
method = ds_mod.Mxfp4FlashinferCutlassMoEMethod(
|
||||||
|
SimpleNamespace(process_weights_after_loading=lambda layer: None),
|
||||||
|
"test",
|
||||||
|
)
|
||||||
|
|
||||||
|
layer = _MockLayer()
|
||||||
|
layer.w13_weight = torch.nn.Parameter(w31.clone(), requires_grad=False)
|
||||||
|
layer.w2_weight = torch.nn.Parameter(w2.clone(), requires_grad=False)
|
||||||
|
layer.w13_weight_scale_inv = torch.nn.Parameter(w31_s.clone(), requires_grad=False)
|
||||||
|
layer.w2_weight_scale_inv = torch.nn.Parameter(w2_s.clone(), requires_grad=False)
|
||||||
|
layer.num_local_experts = num_experts
|
||||||
|
method.process_weights_after_loading(layer)
|
||||||
|
|
||||||
|
ref_w13, ref_w13_s, ref_w13_residual = preprocess_humming(
|
||||||
|
w31.view(torch.uint8), w31_s.view(torch.uint8)
|
||||||
|
)
|
||||||
|
ref_w2, ref_w2_s, ref_w2_residual = preprocess_humming(
|
||||||
|
w2.view(torch.uint8), w2_s.view(torch.uint8)
|
||||||
|
)
|
||||||
|
assert torch.equal(layer.w13_weight, ref_w13)
|
||||||
|
assert torch.equal(layer.w2_weight, ref_w2)
|
||||||
|
assert torch.equal(layer.w13_weight_scale_inv, ref_w13_s)
|
||||||
|
assert torch.equal(layer.w2_weight_scale_inv, ref_w2_s)
|
||||||
|
assert torch.equal(layer.w13_humming_residual_scale, ref_w13_residual * 64.0)
|
||||||
|
assert torch.equal(layer.w2_humming_residual_scale, ref_w2_residual * 64.0)
|
||||||
|
|
||||||
|
|
||||||
class _MockDispatchOutput:
|
class _MockDispatchOutput:
|
||||||
"""Stand-in for StandardDispatchOutput. ``topk_output`` is a real
|
"""Stand-in for StandardDispatchOutput. ``topk_output`` is a real
|
||||||
``StandardTopKOutput`` so ``TopKOutputChecker.format_is_standard``
|
``StandardTopKOutput`` so ``TopKOutputChecker.format_is_standard``
|
||||||
|
|||||||
Reference in New Issue
Block a user