A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 14, 05:30 UTC
[Bugfix] Fix gpt-oss w4a8 DP/EP on B200 (#26729)
Signed-off-by: Varun Sundar Rabindranath <vsundarr@redhat.com> Co-authored-by: Varun Sundar Rabindranath <vsundarr@redhat.com> Co-authored-by: Michael Goin <mgoin64@gmail.com>
Varun Sundar Rabindranath committed
Oct 21, 2025 at 01:51 UTC
5ff5d94e77851f8ca11592bbb9aa414e65f4c353
5 files changed
+82
-2
tests/quantization/test_blackwell_moe.py
+20
@@ -170,3 +170,23 @@ def test_gptoss_mxfp4mxfp8_moe_flashinfer_cutlass(monkeypatch: pytest.MonkeyPatc
170
def test_gptoss_mxfp4mxfp8_moe_flashinfer_trtllm(monkeypatch: pytest.MonkeyPatch):
171
monkeypatch.setenv("VLLM_USE_FLASHINFER_MOE_MXFP4_MXFP8", "1")
172
can_initialize("openai/gpt-oss-20b", hf_overrides=HF_OVERRIDE_TEXT)
173
+
174
+
175
+def test_gptoss_dp2_mxfp4mxfp8_moe_flashinfer_trtllm(monkeypatch: pytest.MonkeyPatch):
176
+ monkeypatch.setenv("VLLM_USE_FLASHINFER_MOE_MXFP4_MXFP8", "1")
177
+ monkeypatch.setenv("VLLM_ALL2ALL_BACKEND", "deepep_high_throughput")
178
+ can_initialize(
179
+ "openai/gpt-oss-20b",
180
+ extra_args=["--data-parallel-size", "2", "--enable-expert-parallel"],
181
+ hf_overrides=HF_OVERRIDE_TEXT,
182
+ )
183
+
184
+
185
+def test_gptoss_dp2_mxfp4bf16_moe_flashinfer_trtllm(monkeypatch: pytest.MonkeyPatch):
186
+ monkeypatch.setenv("VLLM_USE_FLASHINFER_MOE_MXFP4_BF16", "1")
187
+ monkeypatch.setenv("VLLM_ALL2ALL_BACKEND", "deepep_high_throughput")
188
+ can_initialize(
189
+ "openai/gpt-oss-20b",
190
+ extra_args=["--data-parallel-size", "2", "--enable-expert-parallel"],
191
+ hf_overrides=HF_OVERRIDE_TEXT,
192
+ )
vllm/model_executor/layers/fused_moe/config.py
+20
@@ -517,6 +517,26 @@ def mxfp4_w4a16_moe_quant_config(
517
)
518
519
520
+def mxfp4_mxfp8_moe_quant_config(
521
+ w1_scale: Union[torch.Tensor, "PrecisionConfig"],
522
+ w2_scale: Union[torch.Tensor, "PrecisionConfig"],
523
+ a1_scale: torch.Tensor | None = None,
524
+ a2_scale: torch.Tensor | None = None,
525
+ w1_bias: torch.Tensor | None = None,
526
+ w2_bias: torch.Tensor | None = None,
527
+ block_shape: list[int] | None = None,
528
+) -> FusedMoEQuantConfig:
529
+ """
530
+ Construct a quant config for mxfp4 activations and mxfp4 weights.
531
+ """
532
+ return FusedMoEQuantConfig(
533
+ _a1=FusedMoEQuantDesc("mxfp8"),
534
+ _a2=FusedMoEQuantDesc("mxfp8"),
535
+ _w1=FusedMoEQuantDesc("mxfp4", None, w1_scale, None, None, w1_bias),
536
+ _w2=FusedMoEQuantDesc("mxfp4", None, w2_scale, None, None, w2_bias),
537
+ )
538
+
539
+
540
def ocp_mx_moe_quant_config(
541
quant_dtype: str,
542
w1_scale: Union[torch.Tensor, "PrecisionConfig"],
vllm/model_executor/layers/quantization/mxfp4.py
+18
@@ -18,6 +18,7 @@ from vllm.model_executor.layers.fused_moe import (
18
from vllm.model_executor.layers.fused_moe import modular_kernel as mk
19
from vllm.model_executor.layers.fused_moe.config import (
20
FusedMoEQuantConfig,
21
+ mxfp4_mxfp8_moe_quant_config,
22
mxfp4_w4a16_moe_quant_config,
23
ocp_mx_moe_quant_config,
24
)
@@ -747,6 +748,23 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
748
w1_scale=w1_scale,
749
w2_scale=w2_scale,
750
)
751
+ elif self.mxfp4_backend in [
752
+ Mxfp4Backend.SM100_FI_MXFP4_MXFP8_TRTLLM,
753
+ Mxfp4Backend.SM100_FI_MXFP4_MXFP8_CUTLASS,
754
+ ]:
755
+ return mxfp4_mxfp8_moe_quant_config(
756
+ w1_bias=layer.w13_bias,
757
+ w2_bias=layer.w2_bias,
758
+ w1_scale=layer.w13_weight_scale,
759
+ w2_scale=layer.w2_weight_scale,
760
+ )
761
+ elif self.mxfp4_backend in [Mxfp4Backend.SM100_FI_MXFP4_BF16]:
762
+ return mxfp4_w4a16_moe_quant_config(
763
+ w1_bias=layer.w13_bias,
764
+ w2_bias=layer.w2_bias,
765
+ w1_scale=layer.w13_weight_scale,
766
+ w2_scale=layer.w2_weight_scale,
767
+ )
768
else:
769
w1_scale = layer.w13_weight_scale
770
w2_scale = layer.w2_weight_scale
vllm/model_executor/layers/quantization/utils/mxfp8_utils.py
+4
-1
@@ -18,4 +18,7 @@ def mxfp8_e4m3_quantize(x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
18
"`pip install flashinfer`"
19
) from err
20
21
- return mxfp8_e4m3_quantize(x, is_sf_swizzled_layout=False)
21
+ x_q, x_scales = mxfp8_e4m3_quantize(x, is_sf_swizzled_layout=False)
22
+ if x_scales.ndim == 1:
23
+ x_scales = x_scales.view(x.size(0), -1)
24
+ return x_q, x_scales
vllm/model_executor/warmup/kernel_warmup.py
+20
-1
@@ -11,6 +11,7 @@ from typing import TYPE_CHECKING
11
import torch
12
13
import vllm.envs as envs
14
+from vllm.config import VllmConfig
15
from vllm.logger import init_logger
16
from vllm.model_executor.warmup.deep_gemm_warmup import deep_gemm_warmup
17
from vllm.platforms import current_platform
@@ -24,6 +25,20 @@ if TYPE_CHECKING:
25
logger = init_logger(__name__)
26
27
28
+def flashinfer_autotune_supported(vllm_config: VllmConfig) -> bool:
29
+ """
30
+ Record known issues with vllm + flashinfer autotune here. Return True if
31
+ and only if flashinfer autotune will run through without issues.
32
+ """
33
+ return not (
34
+ vllm_config.parallel_config.data_parallel_size > 1
35
+ and (
36
+ envs.VLLM_USE_FLASHINFER_MOE_MXFP4_BF16
37
+ or envs.VLLM_USE_FLASHINFER_MOE_MXFP4_MXFP8
38
+ )
39
+ )
40
+
41
+
42
def kernel_warmup(worker: "Worker"):
43
# Deep GEMM warmup
44
do_deep_gemm_warmup = (
@@ -37,7 +52,11 @@ def kernel_warmup(worker: "Worker"):
52
deep_gemm_warmup(model, max_tokens)
53
54
# FlashInfer autotune for Hopper (SM 9.0) and Blackwell (SM 10.0) GPUs
40
- if has_flashinfer() and current_platform.has_device_capability(90):
55
+ if (
56
+ has_flashinfer()
57
+ and current_platform.has_device_capability(90)
58
+ and flashinfer_autotune_supported(worker.vllm_config)
59
+ ):
60
flashinfer_autotune(worker.model_runner)
61
62
# FlashInfer attention warmup