From 78cb38ed5ec4ca5b220ed2765b19b367dc704f79 Mon Sep 17 00:00:00 2001 From: Ziang Li Date: Tue, 19 May 2026 01:04:48 -0700 Subject: [PATCH] [FlashInfer v0.6.11] [RL] Support FlashInfer per-token NVFP4 MoE (#22918) --- docs/references/environment_variables.md | 1 + .../docs/references/environment_variables.mdx | 5 ++++ python/sglang/srt/environ.py | 2 ++ .../moe/moe_runner/flashinfer_trtllm.py | 26 ++++++++++++++++--- .../srt/layers/quantization/modelopt_quant.py | 7 +++++ .../test_flashinfer_trtllm_gen_moe_backend.py | 24 +++-------------- 6 files changed, 42 insertions(+), 23 deletions(-) diff --git a/docs/references/environment_variables.md b/docs/references/environment_variables.md index 0171bdc05..99f65c7c5 100644 --- a/docs/references/environment_variables.md +++ b/docs/references/environment_variables.md @@ -127,6 +127,7 @@ SGLang supports various environment variables that can be used to configure its | `SGLANG_FORCE_FP8_MARLIN` | Force using FP8 MARLIN kernels even if other FP8 kernels are available | `false` | | `SGLANG_NVFP4_CKPT_FP8_GEMM_IN_ATTN` | Quantize q_b_proj from BF16 to FP8 when launching DeepSeek NVFP4 checkpoint | `false` | | `SGLANG_MOE_NVFP4_DISPATCH` | Use nvfp4 for moe dispatch (on flashinfer_cutlass or flashinfer_cutedsl moe runner backend) | `"false"` | +| `SGLANG_FLASHINFER_NVFP4_PER_TOKEN_ACTIVATION` | Enable FlashInfer TRTLLM per-token NVFP4 activation scaling; ignores checkpoint activation FP32 scale by treating it as `1` | `false` | | `SGLANG_NVFP4_CKPT_FP8_NEXTN_MOE` | Quantize moe of nextn layer from BF16 to FP8 when launching DeepSeek NVFP4 checkpoint | `false` | | `SGLANG_QUANT_ALLOW_DOWNCASTING` | Allow weight dtype downcasting during loading (e.g., fp32 → fp16). By default, SGLang rejects this kind of downcasting when using quantization. | `false` | | `SGLANG_FP8_IGNORED_LAYERS` | A comma-separated list of layer names to ignore during FP8 quantization. For example: `model.layers.0,model.layers.1.,qkv_proj`. | `""` | diff --git a/docs_new/docs/references/environment_variables.mdx b/docs_new/docs/references/environment_variables.mdx index f412c4718..af52bd385 100644 --- a/docs_new/docs/references/environment_variables.mdx +++ b/docs_new/docs/references/environment_variables.mdx @@ -549,6 +549,11 @@ SGLang supports various environment variables that can be used to configure its Use nvfp4 for moe dispatch (on flashinfer_cutlass or flashinfer_cutedsl moe runner backend) "false" + + SGLANG_FLASHINFER_NVFP4_PER_TOKEN_ACTIVATION + Enable FlashInfer TRTLLM per-token NVFP4 activation scaling; ignores checkpoint activation FP32 scale by treating it as 1 + false + SGLANG_NVFP4_CKPT_FP8_NEXTN_MOE Quantize moe of nextn layer from BF16 to FP8 when launching DeepSeek NVFP4 checkpoint diff --git a/python/sglang/srt/environ.py b/python/sglang/srt/environ.py index 9f791751a..d0a4ac8b1 100644 --- a/python/sglang/srt/environ.py +++ b/python/sglang/srt/environ.py @@ -369,6 +369,8 @@ class Envs: SGLANG_FLASHINFER_USE_PAGED = EnvBool(False) # Default to the pick from flashinfer SGLANG_FLASHINFER_WORKSPACE_SIZE = EnvInt(384 * 1024 * 1024) + # Enable per-token NVFP4 activation scaling path for FlashInfer TRT-LLM MoE. + SGLANG_FLASHINFER_NVFP4_PER_TOKEN_ACTIVATION = EnvBool(False) # Skip-softmax threshold scale factor for TRT-LLM attention (prefill and decode separately). # None = standard attention. See https://arxiv.org/abs/2512.12087 SGLANG_SKIP_SOFTMAX_PREFILL_THRESHOLD_SCALE_FACTOR = EnvFloat(None) diff --git a/python/sglang/srt/layers/moe/moe_runner/flashinfer_trtllm.py b/python/sglang/srt/layers/moe/moe_runner/flashinfer_trtllm.py index 61dddd6ac..e3336636b 100644 --- a/python/sglang/srt/layers/moe/moe_runner/flashinfer_trtllm.py +++ b/python/sglang/srt/layers/moe/moe_runner/flashinfer_trtllm.py @@ -14,6 +14,7 @@ from sglang.srt.distributed.device_communicators.pynccl_allocator import ( is_tensor_in_symmetric_mempool, use_symmetric_memory, ) +from sglang.srt.environ import envs from sglang.srt.layers.dp_attention import is_allocation_symmetric from sglang.srt.layers.moe.flashinfer_trtllm_moe import ( trtllm_fp8_block_scale_moe_wrapper, @@ -872,9 +873,26 @@ def fused_experts_none_to_flashinfer_trtllm_fp4( topk_output = dispatch_output.topk_output # Quantize hidden states to FP4 - hs_fp4, hs_scale_linear = quantize_hidden_states_fp4( - hidden_states, quant_info.w13_input_scale_quant - ) + if envs.SGLANG_FLASHINFER_NVFP4_PER_TOKEN_ACTIVATION.get(): + from flashinfer import SfLayout, nvfp4_quantize + + hs_fp4_bytes, hs_sf_bytes, per_token_scale = nvfp4_quantize( + hidden_states, + 1.0 / (448.0 * 6.0), + sfLayout=SfLayout.layout_linear, + per_token_activation=True, + ) + + seq_len, hidden_size = hidden_states.shape + hs_fp4 = hs_fp4_bytes.reshape(seq_len, hidden_size // 2) + hs_scale_linear = hs_sf_bytes.view(torch.float8_e4m3fn).reshape( + seq_len, hidden_size // 16 + ) + else: + per_token_scale = None + hs_fp4, hs_scale_linear = quantize_hidden_states_fp4( + hidden_states, quant_info.w13_input_scale_quant + ) hs_scale = hs_scale_linear.view(torch.float8_e4m3fn).reshape( *hs_scale_linear.shape[:-1], -1 ) @@ -927,6 +945,7 @@ def fused_experts_none_to_flashinfer_trtllm_fp4( output1_scale_scalar=quant_info.g1_scale_c, output1_scale_gate_scalar=quant_info.g1_alphas, output2_scale_scalar=quant_info.g2_alphas, + per_token_scale=per_token_scale, num_experts=quant_info.global_num_experts, top_k=topk_output.topk_ids.shape[1], n_group=0, @@ -974,6 +993,7 @@ def fused_experts_none_to_flashinfer_trtllm_fp4( output1_scale_scalar=quant_info.g1_scale_c, output1_scale_gate_scalar=quant_info.g1_alphas, output2_scale_scalar=quant_info.g2_alphas, + per_token_scale=per_token_scale, num_experts=quant_info.global_num_experts, top_k=topk_config.top_k, n_group=topk_config.num_expert_group, diff --git a/python/sglang/srt/layers/quantization/modelopt_quant.py b/python/sglang/srt/layers/quantization/modelopt_quant.py index c0d9b300f..fd7cfb4ff 100755 --- a/python/sglang/srt/layers/quantization/modelopt_quant.py +++ b/python/sglang/srt/layers/quantization/modelopt_quant.py @@ -1758,6 +1758,13 @@ class ModelOptNvFp4FusedMoEMethod(FusedMoEMethodBase): w13_input_scale = layer.w13_input_scale.max(dim=-1).values.to(torch.float32) w2_input_scale = layer.w2_input_scale + if ( + self.enable_flashinfer_trtllm_moe + and envs.SGLANG_FLASHINFER_NVFP4_PER_TOKEN_ACTIVATION.get() + ): + w13_input_scale = torch.ones_like(w13_input_scale, dtype=torch.float32) + w2_input_scale = torch.ones_like(w2_input_scale, dtype=torch.float32) + # Create shared parameters copy_or_rebind_param( layer, diff --git a/test/registered/backends/test_flashinfer_trtllm_gen_moe_backend.py b/test/registered/backends/test_flashinfer_trtllm_gen_moe_backend.py index aff581054..e483fa562 100644 --- a/test/registered/backends/test_flashinfer_trtllm_gen_moe_backend.py +++ b/test/registered/backends/test_flashinfer_trtllm_gen_moe_backend.py @@ -157,6 +157,7 @@ class FlashinferTrtllmGenMoeBackendMXFP8Base: class FlashinferTrtllmGenMoeBackendNVFP4Base: backend = None + extra_env = {} @classmethod def setUpClass(cls): @@ -166,7 +167,7 @@ class FlashinferTrtllmGenMoeBackendNVFP4Base: cls.model, cls.base_url, timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, - env={**os.environ, "SGLANG_ENABLE_JIT_DEEPGEMM": "False"}, + env={**os.environ, **cls.extra_env, "SGLANG_ENABLE_JIT_DEEPGEMM": "False"}, other_args=[ "--moe-runner-backend", cls.backend, @@ -204,30 +205,12 @@ class TestFlashinferTrtllmGenMoeBackendFP8( backend = "flashinfer_trtllm" -class TestFlashinferTrtllmGenMoeBackendMXFP8( - FlashinferTrtllmGenMoeBackendMXFP8Base, CustomTestCase -): - backend = "flashinfer_trtllm" - - -class TestFlashinferTrtllmGenMoeBackendBF16( - FlashinferTrtllmGenMoeBackendBF16Base, CustomTestCase -): - backend = "flashinfer_trtllm" - - class TestFlashinferTrtllmGenMoeBackendNVFP4( FlashinferTrtllmGenMoeBackendNVFP4Base, CustomTestCase ): backend = "flashinfer_trtllm" -class TestFlashinferTrtllmGenMoeBackendFP8Routed( - FlashinferTrtllmGenMoeBackendFP8Base, CustomTestCase -): - backend = "flashinfer_trtllm_routed" - - class TestFlashinferTrtllmGenMoeBackendMXFP8Routed( FlashinferTrtllmGenMoeBackendMXFP8Base, CustomTestCase ): @@ -240,9 +223,10 @@ class TestFlashinferTrtllmGenMoeBackendBF16Routed( backend = "flashinfer_trtllm_routed" -class TestFlashinferTrtllmGenMoeBackendNVFP4Routed( +class TestFlashinferTrtllmGenMoeBackendPerTokenNVFP4Routed( FlashinferTrtllmGenMoeBackendNVFP4Base, CustomTestCase ): + extra_env = {"SGLANG_FLASHINFER_NVFP4_PER_TOKEN_ACTIVATION": "1"} backend = "flashinfer_trtllm_routed"