@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 9127170b88

[Quantization] Select per-token NVFP4 MoE backends explicitly (#57176)

Signed-off-by: aoshen02 <aoshen@inferact.ai> Signed-off-by: S1ro1 <matej.sirovatka@gmail.com> Signed-off-by: aoshen02 <aoshen02@users.noreply.github.com> Co-authored-by: aoshen02 <aoshen@inferact.ai> Co-authored-by: OpenAI Codex <noreply@openai.com> Co-authored-by: aoshen02 <aoshen02@users.noreply.github.com> Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>

Matej Sirovatka committed Sep 23, 2026 at 00:30 UTC 9127170b880dc86e032c4acaf649c225fdeff3b1

6 files changed +97 -21

tests/quantization/test_online.py

+71

@@ -36,6 +36,7 @@ from vllm.model_executor.kernels.linear.mxfp8.marlin import (

36 )

37 from vllm.model_executor.layers.attention import Attention

38 from vllm.model_executor.layers.fused_moe import FusedMoEFactory

39 +from vllm.model_executor.layers.fused_moe.oracle import nvfp4 as nvfp4_oracle

40 from vllm.model_executor.layers.linear import (

41 ColumnParallelLinear,

42 LinearBase,

@@ -503,6 +504,76 @@ def test_checkpoint_quantization_rejects_online_shorthand(tmp_path) -> None:

504 ModelConfig(model=str(tmp_path), quantization="fp8_per_channel")

505

506

507 +@pytest.mark.parametrize("per_token_activation", [False, True])

508 +@pytest.mark.parametrize("input_dtype", [torch.bfloat16, torch.float16])

509 +def test_nvfp4_one_sided_sizes_dispatched_activations(

510 + monkeypatch, per_token_activation: bool, input_dtype: torch.dtype

511 +) -> None:

512 + from tests.kernels.moe.utils import make_dummy_moe_config

513 + from vllm.model_executor.layers.fused_moe.all2all_utils import (

514 + flashinfer_one_sided_dispatch_layout,

515 + )

516 + from vllm.model_executor.layers.fused_moe.config import FusedMoEQuantConfig

517 + from vllm.model_executor.layers.fused_moe.experts.trtllm_nvfp4_moe import (

518 + TrtLlmNvFp4ExpertsModular,

519 + )

520 +

521 + config = make_dummy_moe_config(hidden_dim=6144, in_dtype=input_dtype)

522 + config.moe_parallel_config.dp_size = 2

523 + config.moe_parallel_config.ep_size = 2

524 + config.moe_parallel_config.use_ep = True

525 + config.moe_parallel_config.all2all_backend = "flashinfer_nvlink_one_sided"

526 + expert_quant_config = FusedMoEQuantConfig.make("nvfp4")

527 +

528 + def check_dispatch_layout(*, quant_config, input_dtype, **kwargs):

529 + assert quant_config is expert_quant_config

530 + layout = flashinfer_one_sided_dispatch_layout(

531 + config.hidden_dim, quant_config, input_dtype=input_dtype

532 + )

533 + if per_token_activation:

534 + assert input_dtype == config.in_dtype

535 + assert (

536 + layout.x_bytes_per_token == config.hidden_dim * config.in_dtype.itemsize

537 + )

538 + assert layout.x_sf_bytes_per_token == 0

539 + else:

540 + assert input_dtype is None

541 + assert layout.x_bytes_per_token == config.hidden_dim // 2

542 + assert layout.x_sf_bytes_per_token == config.hidden_dim // 16

543 + assert expert_quant_config.quant_dtype == "nvfp4"

544 + raise RuntimeError("dispatch layout verified before allocation")

545 +

546 + monkeypatch.setattr(

547 + nvfp4_oracle, "maybe_make_prepare_finalize", check_dispatch_layout

548 + )

549 + with pytest.raises(

550 + RuntimeError, match="dispatch layout verified before allocation"

551 + ):

552 + nvfp4_oracle.make_nvfp4_moe_kernel(

553 + moe_quant_config=expert_quant_config,

554 + moe_config=config,

555 + experts_cls=TrtLlmNvFp4ExpertsModular,

556 + backend=nvfp4_oracle.NvFp4MoeBackend.FLASHINFER_TRTLLM,

557 + per_token_activation=per_token_activation,

558 + )

559 +

560 +

561 +@pytest.mark.parametrize(

562 + "input_dtype",

563 + [torch.float32, torch.float8_e4m3fn, torch.float8_e5m2, torch.int8],

564 +)

565 +def test_nvfp4_one_sided_rejects_unsupported_input_dtype(input_dtype: torch.dtype):

566 + from vllm.model_executor.layers.fused_moe.all2all_utils import (

567 + flashinfer_one_sided_dispatch_layout,

568 + )

569 + from vllm.model_executor.layers.fused_moe.config import FusedMoEQuantConfig

570 +

571 + with pytest.raises(ValueError, match="unpacked inputs must be float16 or bfloat16"):

572 + flashinfer_one_sided_dispatch_layout(

573 + 6144, FusedMoEQuantConfig.make("nvfp4"), input_dtype=input_dtype

574 + )

575 +

576 +

577 @pytest.mark.skipif(

578 not is_quant_method_supported("fp8"),

579 reason="FP8 is not supported on this GPU type.",

vllm/model_executor/layers/fused_moe/all2all_utils.py

+12 -2

@@ -48,9 +48,18 @@ class FlashInferOneSidedDispatchLayout:

50 def flashinfer_one_sided_dispatch_layout(

51 - hidden_dim: int, quant_config: FusedMoEQuantConfig

51 + hidden_dim: int,

52 + quant_config: FusedMoEQuantConfig,

53 + input_dtype: torch.dtype | None = None,

54 ) -> FlashInferOneSidedDispatchLayout:

55 """Return the one-sided activation payload layout."""

56 + if input_dtype is not None:

57 + if input_dtype not in (torch.float16, torch.bfloat16):

58 + raise ValueError(

59 + "flashinfer_nvlink_one_sided unpacked inputs must be float16 "

60 + f"or bfloat16, got {input_dtype}"

61 + )

62 + return FlashInferOneSidedDispatchLayout(hidden_dim * input_dtype.itemsize, 0)

63 if quant_config.quant_dtype is None:

64 return FlashInferOneSidedDispatchLayout(hidden_dim * 2, 0)

65 if quant_config.quant_dtype == "nvfp4":

@@ -174,6 +183,7 @@ def maybe_make_prepare_finalize(

183 allow_new_interface: bool = False,

184 use_monolithic: bool = False,

185 all2all_manager: Any | None = None,

186 + input_dtype: torch.dtype | None = None,

187 ) -> FusedMoEPrepareAndFinalize | None:

188 if not moe.moe_parallel_config.use_all2all_kernels:

189 if not allow_new_interface:

@@ -367,7 +377,7 @@ def maybe_make_prepare_finalize(

377 get_current_vllm_config().scheduler_config.max_num_batched_tokens

378 )

379 dispatch_layout = flashinfer_one_sided_dispatch_layout(

370 - moe.hidden_dim, quant_config

380 + moe.hidden_dim, quant_config, input_dtype=input_dtype

381 )

382 prepare_finalize = FlashInferNVLinkOneSidedPrepareAndFinalize(

383 max_num_tokens=max_num_tokens,

vllm/model_executor/layers/fused_moe/experts/trtllm_nvfp4_moe.py

+8 -17

@@ -28,6 +28,7 @@ from vllm.model_executor.layers.quantization.utils.flashinfer_utils import (

28 from vllm.model_executor.layers.quantization.utils.quant_utils import (

29 QuantKey,

30 kNvfp4Dynamic,

31 + kNvfp4DynamicToken,

32 kNvfp4Static,

33 )

34 from vllm.model_executor.utils import is_weights_pre_processed

@@ -222,10 +223,10 @@ class TrtLlmNvFp4ExpertsBase:

223 activation_key: QuantKey | None,

224 ) -> bool:

225 """Supports Nvfp4 quantization."""

225 - SUPPORTED_W_A = [

226 - (kNvfp4Static, kNvfp4Dynamic),

227 - ]

228 - return (weight_key, activation_key) in SUPPORTED_W_A

226 + return weight_key == kNvfp4Static and activation_key in (

227 + kNvfp4Dynamic,

228 + kNvfp4DynamicToken,

229 + )

230

231 @staticmethod

232 def _supports_activation(activation: MoEActivation) -> bool:

@@ -316,22 +317,12 @@ class TrtLlmNvFp4ExpertsModular(TrtLlmNvFp4ExpertsBase, mk.FusedMoEExpertsModula

317 expert_tokens_meta: mk.ExpertTokensMetadata | None,

318 activation: MoEActivation,

319 ) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...]]:

319 - if self.per_token_activation:

320 - # Deferred input quant leaves K unpacked here, breaking the

321 - # workspace assumptions below. Per-token NVFP4 is only supported on

322 - # the monolithic (non-EP) path for now.

323 - raise NotImplementedError(

324 - "NVFP4 per-token activation is only supported on the monolithic "

325 - "(non-EP) FlashInfer TRTLLM MoE path."

326 - )

327 -

320 # The workspaces for this implementation are managed by flashinfer.

321 workspace1 = (0,)

322 workspace2 = (0,)

323

332 - # Hidden states are Nvfp4, packed into int8 dtype, so we

333 - # need to multiply K by 2 to get the output shape right.

334 - assert self.hidden_dim == K * 2

324 + # Per-token inputs are unpacked; otherwise each byte holds two FP4 values.

325 + assert self.hidden_dim == (K if self.expects_unquantized_inputs else K * 2)

326 output = (M, self.hidden_dim)

327

328 return (workspace1, workspace2, output)

@@ -349,7 +340,7 @@ class TrtLlmNvFp4ExpertsModular(TrtLlmNvFp4ExpertsBase, mk.FusedMoEExpertsModula

340 topk_ids: torch.Tensor,

341 activation: MoEActivation,

342 global_num_experts: int,

352 - a1q_scale: torch.Tensor,

343 + a1q_scale: torch.Tensor | None,

344 ):

345 import flashinfer

346

vllm/model_executor/layers/fused_moe/oracle/nvfp4.py

@@ -605,6 +605,7 @@ def make_nvfp4_moe_kernel(

605 routing_tables=routing_tables,

606 allow_new_interface=True,

607 use_monolithic=issubclass(experts_cls, mk.FusedMoEExpertsMonolithic),

608 + input_dtype=moe_config.in_dtype if per_token_activation else None,

609 )

610 assert prepare_finalize is not None

611

vllm/model_executor/layers/quantization/online/nvfp4.py

+2 -2

@@ -24,7 +24,7 @@ from vllm.model_executor.layers.quantization.utils.nvfp4_emulation_utils import

24 )

25 from vllm.model_executor.layers.quantization.utils.quant_utils import (

26 amax_for_moe_weight_quant,

27 - kNvfp4Dynamic,

27 + kNvfp4DynamicToken,

28 kNvfp4Static,

29 weight_amax,

30 )

@@ -100,7 +100,7 @@ class Nvfp4OnlineMoEMethod(OnlineMoEMethodBase):

100 self.nvfp4_backend, self.experts_cls = select_nvfp4_moe_backend(

101 config=self.moe,

102 weight_key=kNvfp4Static,

103 - activation_key=kNvfp4Dynamic,

103 + activation_key=kNvfp4DynamicToken,

104 )

105

106 def process_weights_after_loading(self, layer: Module) -> None:

vllm/model_executor/layers/quantization/utils/quant_utils.py

@@ -208,6 +208,9 @@ kNvfp4DynamicGroupScale = ScaleDesc(FP8_DTYPE, False, GroupShape(1, 16))

208 kNvfp4Dynamic = QuantKey(

209 FP4_DTYPE, scale=kNvfp4DynamicGroupScale, scale2=kStaticTensorScale

210 )

211 +kNvfp4DynamicToken = QuantKey(

212 + FP4_DTYPE, scale=kNvfp4DynamicGroupScale, scale2=kDynamicTokenScale

213 +)

214

215 kNvfp4StaticGroupScale = ScaleDesc(FP8_DTYPE, True, GroupShape(1, 16))

216 kNvfp4Static = QuantKey(