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

[Bugfix] Fix Mistral Large 3 NVFP4 TRTLLM MoE (#18065)

Browse Source
This commit is contained in:
elvischenv
2026-02-03 20:32:49 +08:00
committed by GitHub
parent a45647bce1
commit 99fab2ce67
2 changed files with 122 additions and 118 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/layers/quantization/compressed_tensors/compressed_tensors_moe.py
+24 -33
View File
@@ -461,45 +461,14 @@ class CompressedTensorsW4A4Nvfp4MoEMethod(CompressedTensorsMoEMethod):
dispatch_output: StandardDispatchOutput,
) -> CombineInput:
from sglang.srt.layers.moe.cutlass_moe import cutlass_moe_fp4
from sglang.srt.layers.moe.token_dispatcher import StandardCombineInput
x = dispatch_output.hidden_states
topk_output = dispatch_output.topk_output
topk_weights, topk_ids = topk_output.topk_weights, topk_output.topk_ids
output = cutlass_moe_fp4(
a=x,
a1_gscale=layer.w13_input_scale_quant,
w1_fp4=layer.w13_weight,
w1_blockscale=layer.w13_weight_scale,
w1_alphas=layer.g1_alphas,
a2_gscale=layer.w2_input_scale_quant,
w2_fp4=layer.w2_weight,
w2_blockscale=layer.w2_weight_scale,
w2_alphas=layer.g2_alphas,
topk_weights=topk_weights,
topk_ids=topk_ids,
params=layer.cutlass_moe_params,
apply_router_weight_on_input=self.moe_runner_config.apply_router_weight_on_input,
).to(x.dtype)
return StandardCombineInput(hidden_states=output)
def apply_with_router_logits(
self,
layer: torch.nn.Module,
dispatch_output: StandardDispatchOutput,
) -> torch.Tensor:
assert self.use_flashinfer_trtllm
x = dispatch_output.hidden_states
topk_output = dispatch_output.topk_output
if self.use_flashinfer_trtllm:
from flashinfer import fp4_quantize, trtllm_fp4_block_scale_moe
from sglang.srt.layers.moe.utils import RoutingMethodType
router_logits = topk_output.router_logits
topk_config = topk_output.topk_config
@@ -544,7 +513,7 @@ class CompressedTensorsW4A4Nvfp4MoEMethod(CompressedTensorsMoEMethod):
num_tokens, hidden_size, dtype=torch.bfloat16, device=hs_fp4.device
)
return trtllm_fp4_block_scale_moe(
output = trtllm_fp4_block_scale_moe(
routing_logits=router_logits,
routing_bias=correction_bias,
hidden_states=hs_fp4,
@@ -578,6 +547,28 @@ class CompressedTensorsW4A4Nvfp4MoEMethod(CompressedTensorsMoEMethod):
tune_max_num_tokens=next_power_of_2(hs_fp4.shape[0]),
output=symm_output,
)[0]
else:
from sglang.srt.layers.moe.cutlass_moe import cutlass_moe_fp4
topk_weights, topk_ids = topk_output.topk_weights, topk_output.topk_ids
output = cutlass_moe_fp4(
a=x,
a1_gscale=layer.w13_input_scale_quant,
w1_fp4=layer.w13_weight,
w1_blockscale=layer.w13_weight_scale,
w1_alphas=layer.g1_alphas,
a2_gscale=layer.w2_input_scale_quant,
w2_fp4=layer.w2_weight,
w2_blockscale=layer.w2_weight_scale,
w2_alphas=layer.g2_alphas,
topk_weights=topk_weights,
topk_ids=topk_ids,
params=layer.cutlass_moe_params,
apply_router_weight_on_input=self.moe_runner_config.apply_router_weight_on_input,
).to(x.dtype)
return StandardCombineInput(hidden_states=output)
class CompressedTensorsW8A8Fp8MoEMethod(CompressedTensorsMoEMethod):
test/registered/8-gpu-models/test_mistral_large3.py
+21 -8
View File
@@ -9,9 +9,10 @@ from sglang.test.test_utils import ModelLaunchSettings, is_blackwell_system
# Runs on both H200 and B200 via nightly-8-gpu-common suite
# Note: trtllm_mla backend may have hardware-specific behavior
register_cuda_ci(est_time=1800, suite="nightly-8-gpu-common", nightly=True)
register_cuda_ci(est_time=3000, suite="nightly-8-gpu-common", nightly=True)
MISTRAL_LARGE3_MODEL_PATH = "mistralai/Mistral-Large-3-675B-Instruct-2512"
MISTRAL_LARGE3_FP8_MODEL_PATH = "mistralai/Mistral-Large-3-675B-Instruct-2512"
MISTRAL_LARGE3_NVFP4_MODEL_PATH = "mistralai/Mistral-Large-3-675B-Instruct-2512-NVFP4"
MISTRAL_LARGE3_EAGLE_MODEL_PATH = "mistralai/Mistral-Large-3-675B-Instruct-2512-Eagle"
@@ -19,9 +20,10 @@ MISTRAL_LARGE3_EAGLE_MODEL_PATH = "mistralai/Mistral-Large-3-675B-Instruct-2512-
class TestMistralLarge3(unittest.TestCase):
"""Unified test class for Mistral-Large-3 performance and accuracy.
Two variants:
- basic: TP=8 + trtllm_mla backend
Three variants:
- basic: FP8 model + TP=8 + trtllm_mla backend
- eagle: basic + EAGLE speculative decoding with draft model
- nvfp4: NVFP4 model + TP=8 + trtllm_mla backend
Each variant runs BOTH:
- Performance test (using NightlyBenchmarkRunner)
@@ -56,22 +58,33 @@ class TestMistralLarge3(unittest.TestCase):
"--speculative-num-draft-tokens=4",
"--kv-cache-dtype=auto",
]
# TODO: add this to base args when FP8 TRTLLM moe is supported
nvfp4_args = [
"--moe-runner-backend=flashinfer_trtllm",
]
variants = [
# Variant: "basic" - TP=8 + trtllm_mla backend
# Variant: "basic" - FP8 model + TP=8 + trtllm_mla backend
ModelLaunchSettings(
MISTRAL_LARGE3_MODEL_PATH,
MISTRAL_LARGE3_FP8_MODEL_PATH,
tp_size=8,
extra_args=base_args,
variant="TP8",
),
# Variant: "eagle" - TP=8 + trtllm_mla + EAGLE with draft model
# Variant: "eagle" - FP8 model + TP=8 + trtllm_mla + EAGLE with draft model
ModelLaunchSettings(
MISTRAL_LARGE3_MODEL_PATH,
MISTRAL_LARGE3_FP8_MODEL_PATH,
tp_size=8,
extra_args=base_args + eagle_args,
variant="TP8+MTP",
),
# Variant: "nvfp4" - NVFP4 model + TP=8 + trtllm_mla backend
ModelLaunchSettings(
MISTRAL_LARGE3_NVFP4_MODEL_PATH,
tp_size=8,
extra_args=base_args + nvfp4_args,
variant="NVFP4",
),
]
run_combined_tests(