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

[FlashInfer v0.6.16] Support FlashInfer CuTe DSL NVFP4 MoE quantization (#28354)

Browse Source
This commit is contained in:
Ziang Li
2026-08-13 17:33:46 -07:00
committed by GitHub
parent c4271c3fe1
commit 9d34c2809f
14 changed files with 405 additions and 133 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/arg_groups/overrides.py
+2 -1
View File
@@ -2396,11 +2396,12 @@ def _moe_runner_backend_quant_constraints(view: Any) -> dict:
elif moe_runner_backend not in [
"flashinfer_trtllm",
"flashinfer_trtllm_routed",
"flashinfer_cutedsl",
]:
raise ValueError(
"--quantization nvfp4_online supports only "
"--moe-runner-backend flashinfer_trtllm or "
"flashinfer_trtllm_routed."
"flashinfer_trtllm_routed, or flashinfer_cutedsl."
)
# Ascend runs MXFP8 MoE on the Ascend runner; every backend selected below is
# CUDA/ROCm-only. Forcing one here would not merely pick the wrong runner:
python/sglang/srt/configs/model_config.py
+1 -1
View File
@@ -1466,7 +1466,7 @@ class ModelConfig:
# so eligible MoE experts are requantized online.
"modelopt_fp4": ["modelopt", "fp8"],
"modelopt_mixed": ["modelopt"],
"nvfp4_online": ["fp8"],
"nvfp4_online": ["fp8", "modelopt_fp8"],
"petit_nvfp4": ["modelopt"],
"w8a8_int8": ["compressed-tensors", "compressed_tensors"],
"w8a8_fp8": ["compressed-tensors", "compressed_tensors"],
python/sglang/srt/environ.py
+4 -1
View File
@@ -866,11 +866,14 @@ class Envs:
# Per-rank dispatch capacity of the FlashInfer MoE A2A dispatcher. Unset
# means each call site keeps its own default.
SGLANG_FLASHINFER_NUM_MAX_DISPATCH_TOKENS_PER_RANK = EnvInt(None)
# Enable NVFP4 per-token activation scaling path for FlashInfer TRT-LLM MoE.
# Enable per-token FP32 activation scaling for serialized ModelOpt FP4 with
# FlashInfer TRT-LLM or CuTe DSL v2 MoE.
SGLANG_FLASHINFER_NVFP4_PER_TOKEN_ACTIVATION = EnvBool(False)
# Launch the TRT-LLM MoE grouped GEMMs with PDL only at or below this
# token count.
SGLANG_TRTLLM_MOE_PDL_MAX_TOKENS = EnvInt(8192)
# Use FlashInfer's fused atomic CUTLASS/CuTe DSL MoE finalize.
SGLANG_FLASHINFER_MOE_FUSED_FINALIZE = EnvBool(True)
# Master switch for the experimental TRT-LLM LoRA fast path; when OFF (default) every
# fine-grained opt switch reads False, keeping non-experimental paths byte-identical.
SGLANG_EXPERIMENTAL_LORA_OPTI = EnvBool(False)
python/sglang/srt/layers/moe/moe_runner/flashinfer_cutedsl.py
+125 -15
View File
@@ -6,6 +6,7 @@ from typing import TYPE_CHECKING, Any, Optional
import torch
from sglang.srt.environ import envs
from sglang.srt.layers.moe.moe_runner.base import (
MoeQuantInfo,
MoeRunnerConfig,
@@ -237,6 +238,64 @@ def _cutedsl_wrapper_activation_type(activation: str, activation_type_cls: Any)
)
def _make_per_token_global_scale(input_tensor: torch.Tensor) -> torch.Tensor:
from flashinfer.quantization.nvfp4_quantization_utils import (
current_nvfp4_4over6_config,
make_nvfp4_global_scale,
)
return make_nvfp4_global_scale(
input_tensor,
per_token_activation=True,
nvfp4_4over6_config=current_nvfp4_4over6_config(),
)
def refresh_cutedsl_standard_scales_for_weight_update(
layer: torch.nn.Module,
) -> None:
w1_alpha, fc2_input_scale, w2_alpha, used_input_scale = (
resolve_cutedsl_standard_scales(layer)
)
if layer.quant_config.use_per_token_activation:
used_input_scale = _make_per_token_global_scale(used_input_scale)
new_scales = (w1_alpha, fc2_input_scale, w2_alpha)
current_scales = layer._cutedsl_scales
current_input_scale = layer._cutedsl_input_scale
# Decode CUDA graphs capture these tensor addresses, so reloads must update
# their values without replacing the tensors.
if (
not isinstance(current_scales, tuple)
or len(current_scales) != len(new_scales)
or not isinstance(current_input_scale, torch.Tensor)
):
raise RuntimeError(
"CuTe DSL scale metadata changed during weight reload; "
"CUDA graph recapture is required."
)
scale_pairs = (
*zip(current_scales, new_scales),
(current_input_scale, used_input_scale),
)
for current, new in scale_pairs:
if (
not isinstance(current, torch.Tensor)
or current.shape != new.shape
or current.dtype != new.dtype
or current.device != new.device
):
raise RuntimeError(
"CuTe DSL scale metadata changed during weight reload; "
"CUDA graph recapture is required."
)
with torch.no_grad():
for current, new in scale_pairs:
current.copy_(new)
def ensure_cutedsl_wrapper(layer: torch.nn.Module) -> None:
"""Lazily create CuteDslMoEWrapper and resolve scales on first forward.
@@ -246,9 +305,10 @@ def ensure_cutedsl_wrapper(layer: torch.nn.Module) -> None:
typically runs during the autotune dummy forward under inference_mode().
We wrap the creation in inference_mode(False) so that those pre-allocated
buffers are normal tensors -- inference tensors cannot be inplace-updated
during later CUDA graph capture, which runs outside inference_mode.
during later CUDA graph capture, which runs outside inference_mode. The
resolved scale tensors share this scope because reload updates them in place.
"""
if getattr(layer, "_cutedsl_wrapper", None) is not None:
if layer._cutedsl_wrapper is not None:
return
try:
@@ -296,14 +356,17 @@ def ensure_cutedsl_wrapper(layer: torch.nn.Module) -> None:
local_expert_offset=layer.moe_ep_rank * layer.num_local_experts,
output_dtype=layer.moe_runner_config.params_dtype,
device=str(layer.w13_weight.device),
use_fused_finalize=envs.SGLANG_FLASHINFER_MOE_FUSED_FINALIZE.get(),
activation_type=_cutedsl_wrapper_activation_type(
layer.moe_runner_config.activation, ActivationType
),
)
w1_alpha, fc2_input_scale, w2_alpha, used_input_scale = (
resolve_cutedsl_standard_scales(layer)
)
w1_alpha, fc2_input_scale, w2_alpha, used_input_scale = (
resolve_cutedsl_standard_scales(layer)
)
if layer.quant_config.use_per_token_activation:
used_input_scale = _make_per_token_global_scale(used_input_scale)
layer._cutedsl_scales = (w1_alpha, fc2_input_scale, w2_alpha)
layer._cutedsl_input_scale = used_input_scale
@@ -356,6 +419,9 @@ class CuteDslFp4MoeQuantInfo(MoeQuantInfo):
# v1 only: True when DeepEP pre-quantizes activations to NVFP4.
use_nvfp4_dispatch: bool = False
# v2 only: quantize hidden states with per-token dynamic activation scales.
use_per_token_activation: bool = False
# v1 only: SBO down-GEMM overlap args.
down_gemm_overlap_args: Optional[DownGemmOverlapArgs] = None
@@ -385,11 +451,29 @@ def fused_experts_none_to_flashinfer_cutedsl_fp4(
if topk_ids.dtype != torch.int32:
topk_ids = topk_ids.to(torch.int32)
x_fp4, x_sf = fp4_quantize(
hidden_states,
quant_info.a1_scale,
sf_vec_size=_FP4_SF_VEC_SIZE,
is_sf_swizzled_layout=False,
if quant_info.use_per_token_activation:
from flashinfer import SfLayout, nvfp4_quantize
x_fp4, x_sf, per_token_scale = nvfp4_quantize(
hidden_states,
quant_info.a1_scale,
sfLayout=SfLayout.layout_linear,
per_token_activation=True,
backend="cute-dsl",
)
else:
x_fp4, x_sf = fp4_quantize(
hidden_states,
quant_info.a1_scale,
sf_vec_size=_FP4_SF_VEC_SIZE,
is_sf_swizzled_layout=False,
)
per_token_scale = None
seq_len, hidden_size = hidden_states.shape
x_fp4 = x_fp4.reshape(seq_len, hidden_size // 2)
x_sf = x_sf.view(torch.float8_e4m3fn).reshape(
seq_len, hidden_size // _FP4_SF_VEC_SIZE
)
output = quant_info.wrapper.run(
@@ -404,6 +488,7 @@ def fused_experts_none_to_flashinfer_cutedsl_fp4(
w2_weight=quant_info.w2_weight,
w2_weight_sf=quant_info.w2_weight_sf,
w2_alpha=quant_info.w2_alpha,
per_token_scale=per_token_scale,
)
return StandardCombineInput(hidden_states=output)
@@ -444,14 +529,38 @@ def fused_experts_flashinfer_to_flashinfer_cutedsl_fp4(
topk_ids = topk_ids.to(torch.int32)
if x_sf is not None:
if quant_info.use_per_token_activation:
raise ValueError(
"flashinfer_cutedsl per-token activation requires BF16 dispatch "
"so the runner can forward per_token_scale to FlashInfer."
)
# NVFP4 dispatch, inputs are already quantized.
x_fp4 = hidden_states
per_token_scale = None
else:
x_fp4, x_sf = fp4_quantize(
hidden_states,
quant_info.a1_scale,
sf_vec_size=_FP4_SF_VEC_SIZE,
is_sf_swizzled_layout=False,
if quant_info.use_per_token_activation:
from flashinfer import SfLayout, nvfp4_quantize
x_fp4, x_sf, per_token_scale = nvfp4_quantize(
hidden_states,
quant_info.a1_scale,
sfLayout=SfLayout.layout_linear,
per_token_activation=True,
backend="cute-dsl",
)
else:
x_fp4, x_sf = fp4_quantize(
hidden_states,
quant_info.a1_scale,
sf_vec_size=_FP4_SF_VEC_SIZE,
is_sf_swizzled_layout=False,
)
per_token_scale = None
seq_len, hidden_size = hidden_states.shape
x_fp4 = x_fp4.reshape(seq_len, hidden_size // 2)
x_sf = x_sf.view(torch.float8_e4m3fn).reshape(
seq_len, hidden_size // _FP4_SF_VEC_SIZE
)
output = quant_info.wrapper.run(
@@ -466,6 +575,7 @@ def fused_experts_flashinfer_to_flashinfer_cutedsl_fp4(
w2_weight=quant_info.w2_weight,
w2_weight_sf=quant_info.w2_weight_sf,
w2_alpha=quant_info.w2_alpha,
per_token_scale=per_token_scale,
)
# Note: output contains routed expert results; shared_expert is handled separately
python/sglang/srt/layers/moe/moe_runner/flashinfer_cutlass.py
+3
View File
@@ -18,6 +18,7 @@ from sglang.srt.distributed import get_tp_group
from sglang.srt.distributed.device_communicators.pynccl_allocator import (
use_symmetric_memory,
)
from sglang.srt.environ import envs
from sglang.srt.layers.dp_attention import is_allocation_symmetric
from sglang.srt.layers.moe.moe_runner.base import (
MoeQuantInfo,
@@ -235,6 +236,7 @@ def _run_flashinfer_cutlass(
tune_max_num_tokens=next_power_of_2(x.shape[0]),
activation_type=_activation_type(runner_config),
enable_alltoall=enable_alltoall,
use_fused_finalize=envs.SGLANG_FLASHINFER_MOE_FUSED_FINALIZE.get(),
)[0]
if quant_info.quant_type in ("bf16", "fp8"):
@@ -387,6 +389,7 @@ def fused_experts_none_to_flashinfer_mxfp4(
activation_type=ActivationType.Swiglu,
tune_max_num_tokens=next_power_of_2(x.shape[0]),
output=out,
use_fused_finalize=envs.SGLANG_FLASHINFER_MOE_FUSED_FINALIZE.get(),
)
if do_pad:
python/sglang/srt/layers/quantization/base_config.py
+3
View File
@@ -185,6 +185,9 @@ class QuantizationConfig(ABC):
if hf_quant_config is None:
return None
if user_quant == "nvfp4_online":
return None
# Check if this is a ModelOpt config
quant_algo = hf_quant_config.get("quant_algo", "").upper()
python/sglang/srt/layers/quantization/kv_cache.py
+2 -2
View File
@@ -56,8 +56,8 @@ class BaseKVCacheMethod(QuantizeMethodBase):
if is_fp8_fnuz():
k_scale *= 2
v_scale *= 2
elif layer.k_scale < 0.0 and layer.v_scale < 0.0:
# If no scales were loaded (both scales are invalid negative
elif layer.k_scale <= 0.0 and layer.v_scale <= 0.0:
# If no scales were loaded (both scales are invalid non-positive
# values), use the default value of 1.0
k_scale = 1.0
v_scale = 1.0
python/sglang/srt/layers/quantization/modelopt_quant.py
+25 -13
View File
@@ -1330,7 +1330,8 @@ class ModelOptFp4Config(ModelOptQuantConfig):
and checkpoint-provided scales.
- Serialized + per-token FP32 activation scales: set
`SGLANG_FLASHINFER_NVFP4_PER_TOKEN_ACTIVATION=1`; use
`flashinfer_trtllm` or `flashinfer_trtllm_routed`.
`flashinfer_trtllm`, `flashinfer_trtllm_routed`, or `flashinfer_cutedsl`
v2 with no A2A or FlashInfer A2A.
- BF16/FP16/FP8 MoE + per-tensor FP32 activation scales: quantize expert
weights on load, keep dense weights in source precision or FP8, and use
1.0 when the checkpoint has no NVFP4 activation scale.
@@ -2498,13 +2499,14 @@ class ModelOptNvFp4FusedMoEMethod(FusedMoEMethodBase):
)
# TODO: for flashinfer always do MOE_NVFP4_DISPATCH
use_dispatch_fp4 = not self.quant_config.use_per_token_activation and (
MOE_NVFP4_DISPATCH or should_use_flashinfer_cutlass_moe_fp4_allgather()
)
layer.dispatcher.set_quant_config(
{
"input_global_scale": (
layer.w13_input_scale_quant
if MOE_NVFP4_DISPATCH
or should_use_flashinfer_cutlass_moe_fp4_allgather()
else None
layer.w13_input_scale_quant if use_dispatch_fp4 else None
)
}
)
@@ -2565,17 +2567,19 @@ class ModelOptNvFp4FusedMoEMethod(FusedMoEMethodBase):
interleave_w13_halves,
)
layer.w13_weight = Parameter(
copy_or_rebind_param(
layer,
"w13_weight",
interleave_w13_halves(
layer.w13_weight.view(torch.uint8), group_size=64, dim=1
).contiguous(),
requires_grad=False,
)
layer.w13_weight_scale = Parameter(
copy_or_rebind_param(
layer,
"w13_weight_scale",
interleave_w13_halves(
layer.w13_weight_scale, group_size=64, dim=1
).contiguous(),
requires_grad=False,
)
# Process w13 weights
@@ -2636,6 +2640,7 @@ class ModelOptNvFp4FusedMoEMethod(FusedMoEMethodBase):
from sglang.srt.layers.moe.moe_runner.flashinfer_cutedsl import (
_FP4_SF_VEC_SIZE,
refresh_cutedsl_standard_scales_for_weight_update,
)
sf_vec_size = _FP4_SF_VEC_SIZE
@@ -2644,7 +2649,9 @@ class ModelOptNvFp4FusedMoEMethod(FusedMoEMethodBase):
w13_k = layer.w13_weight.shape[2] * 2
w2_m = layer.w2_weight.shape[1]
w2_k = layer.w2_weight.shape[2] * 2
layer.w13_blockscale_mma = Parameter(
copy_or_rebind_param(
layer,
"w13_blockscale_mma",
convert_sf_to_mma_layout(
layer.w13_blockscale_swizzled.contiguous()
.view(torch.uint8)
@@ -2654,9 +2661,10 @@ class ModelOptNvFp4FusedMoEMethod(FusedMoEMethodBase):
num_groups=num_local_experts,
sf_vec_size=sf_vec_size,
),
requires_grad=False,
)
layer.w2_blockscale_mma = Parameter(
copy_or_rebind_param(
layer,
"w2_blockscale_mma",
convert_sf_to_mma_layout(
layer.w2_blockscale_swizzled.contiguous()
.view(torch.uint8)
@@ -2666,8 +2674,9 @@ class ModelOptNvFp4FusedMoEMethod(FusedMoEMethodBase):
num_groups=num_local_experts,
sf_vec_size=sf_vec_size,
),
requires_grad=False,
)
if layer._cutedsl_wrapper is not None:
refresh_cutedsl_standard_scales_for_weight_update(layer)
@property
def load_up_proj_weight_first(self) -> bool:
@@ -2696,6 +2705,8 @@ 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
layer._cutedsl_wrapper = None
if moe_runner_backend.is_flashinfer_cutlass():
import sglang.srt.layers.moe.moe_runner.flashinfer_cutlass # noqa: F401
@@ -2840,6 +2851,7 @@ class ModelOptNvFp4FusedMoEMethod(FusedMoEMethodBase):
a1_scale=layer._cutedsl_input_scale,
a2_scale=fc2_input_scale,
wrapper=layer._cutedsl_wrapper,
use_per_token_activation=self.quant_config.use_per_token_activation,
)
return self.runner.run(dispatch_output, quant_info)
python/sglang/srt/layers/quantization/nvfp4_online.py
+17 -12
View File
@@ -69,7 +69,9 @@ class NvFp4OnlineConfig(ModelOptQuantConfig):
use_mxfp8: bool = False,
) -> None:
source_ignored_layers = self._normalize_ignored_layers(exclude_modules)
fp4_ignored_layers = list(source_ignored_layers)
fp4_ignored_layers = (
[] if self._use_per_token_activation else list(source_ignored_layers)
)
if ignored_layers_str := envs.SGLANG_FP4_IGNORED_LAYERS.get():
fp4_ignored_layers.extend(
layer.strip()
@@ -114,8 +116,10 @@ class NvFp4OnlineConfig(ModelOptQuantConfig):
quant_method = str(config.get("quant_method", "")).lower()
use_mxfp8 = "mxfp8" in quant_method
is_checkpoint_fp8_serialized = "fp8" in quant_method or use_mxfp8
ignored_layers = config.get("ignored_layers") or config.get(
"modules_to_not_convert"
ignored_layers = (
config.get("ignore")
or config.get("ignored_layers")
or config.get("modules_to_not_convert")
)
if isinstance(ignored_layers, str):
ignored_layers = [ignored_layers]
@@ -142,14 +146,15 @@ class NvFp4OnlineConfig(ModelOptQuantConfig):
return Fp8LinearMethod(self)
return UnquantizedLinearMethod()
if isinstance(layer, FusedMoE):
if is_layer_skipped(
source_layer_ignored = is_layer_skipped(
prefix, self.exclude_modules, self.packed_modules_mapping
) or self.is_layer_excluded(prefix):
) or self.is_layer_excluded(prefix)
if source_layer_ignored and not self.use_per_token_activation:
return None
if is_layer_skipped(
prefix, self.fp4_ignored_layers, self.packed_modules_mapping
):
if self.is_checkpoint_fp8_serialized:
if self.is_checkpoint_fp8_serialized and not source_layer_ignored:
return Fp8MoEMethod(self)
return None
return ModelOptNvFp4OnlineFusedMoEMethod(self, prefix)
@@ -201,15 +206,15 @@ class ModelOptNvFp4OnlineFusedMoEMethod(ModelOptNvFp4FusedMoEMethod):
if layer_match is not None
else layer_prefix
)
if (
quant_config.use_per_token_activation
and not self.enable_flashinfer_trtllm_moe
if quant_config.use_per_token_activation and not (
self.enable_flashinfer_trtllm_moe or self._is_cutedsl_v2_standard
):
raise ValueError(
"--quantization nvfp4_online requires online per-token FP32 "
"activation scales and supports only flashinfer_trtllm or "
"flashinfer_trtllm_routed. Use --quantization modelopt_fp4 "
"for per-tensor FP32 activation scales."
"activation scales and supports flashinfer_trtllm, "
"flashinfer_trtllm_routed, or flashinfer_cutedsl with no A2A "
"or FlashInfer A2A. Use --quantization modelopt_fp4 for "
"per-tensor FP32 activation scales."
)
def prepare_weight_loader(self, layer, weight_loader):
python/sglang/srt/server_args.py
+12 -2
View File
@@ -6670,9 +6670,9 @@ class ServerArgs:
if view.moe_runner_backend == "flashinfer_cutedsl":
# modelopt_mixed with non-NVFP4 MoE layers is rejected at load time.
assert (
view.quantization in ["modelopt_fp4", "modelopt_mixed"]
view.quantization in ["modelopt_fp4", "modelopt_mixed", "nvfp4_online"]
or self.get_model_config().nvfp4_moe_meta is not None
), f"Invalid quantization '{view.quantization}'. \nFlashInfer CuteDSL MOE currently supports only: 'modelopt_fp4', 'modelopt_mixed' (with NVFP4 MoE layers), or hybrid NVFP4 models."
), f"Invalid quantization '{view.quantization}'. \nFlashInfer CuteDSL MOE currently supports only: 'modelopt_fp4', 'modelopt_mixed' (with NVFP4 MoE layers), 'nvfp4_online', or hybrid NVFP4 models."
assert view.ep_size in [
1,
self.tp_size,
@@ -6685,6 +6685,14 @@ class ServerArgs:
f"flashinfer_cutedsl supports moe_a2a_backend='none', 'deepep', or 'flashinfer', "
f"got '{view.moe_a2a_backend}'."
)
if view.moe_a2a_backend == "deepep" and (
view.quantization == "nvfp4_online"
or envs.SGLANG_FLASHINFER_NVFP4_PER_TOKEN_ACTIVATION.get()
):
raise ValueError(
"flashinfer_cutedsl per-token NVFP4 activation requires "
"moe_a2a_backend='none' or 'flashinfer'."
)
if view.moe_runner_backend in ["flashinfer_trtllm", "experimental_sgl_trtllm"]:
assert view.quantization in [
@@ -7983,6 +7991,8 @@ class ServerArgs:
"1" if self.enable_deterministic_inference else "0"
)
self._handle_custom_all_reduce_v2_multinode()
if self.enable_deterministic_inference:
envs.SGLANG_FLASHINFER_MOE_FUSED_FINALIZE.set("0")
if self.debug_cuda_graph:
if not (is_cuda() or is_hip()):
logger.warning(