@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 552a9dbe59

[NVFP4][Emulation] Fuse NVFP4 weight dequantization with compute in triton kernel for w13/w2 MOE MLP linears (#44667)

Signed-off-by: Felix Marty <Felix.Marty@amd.com>

fxmarty-amd committed Jun 26, 2026 at 04:33 UTC 552a9dbe59bf2e6a35654440c64c5e52bed90586
3 files changed +906 -68
tests/kernels/quantization/test_nvfp4_emulation.py
+463
@@ -1,10 +1,25 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +from typing import cast
4 +
5 import huggingface_hub
6 import pytest
7 import torch
8 from safetensors import safe_open
9
10 +import vllm.model_executor.layers.fused_moe.modular_kernel as mk
11 +from vllm.model_executor.layers.fused_moe.activation import MoEActivation
12 +from vllm.model_executor.layers.fused_moe.config import (
13 + FusedMoEConfig,
14 + FusedMoEParallelConfig,
15 + FusedMoEQuantConfig,
16 + RoutingMethodType,
17 + nvfp4_moe_quant_config,
18 +)
19 +from vllm.model_executor.layers.fused_moe.experts.nvfp4_emulation_moe import (
20 + Nvfp4QuantizationEmulationTritonExperts,
21 +)
22 +from vllm.model_executor.layers.fused_moe.experts.triton_moe import TritonExperts
23 from vllm.model_executor.layers.quantization.utils import (
24 nvfp4_emulation_utils,
25 )
@@ -12,9 +27,167 @@ from vllm.model_executor.layers.quantization.utils.nvfp4_emulation_utils import
27 dequantize_to_dtype,
28 ref_nvfp4_quant_dequant,
29 )
30 +from vllm.model_executor.layers.quantization.utils.quant_utils import (
31 + QuantKey,
32 + kNvfp4Dynamic,
33 + kNvfp4Static,
34 +)
35 from vllm.platforms import current_platform
36 from vllm.triton_utils import triton
37
38 +if current_platform.is_rocm():
39 + from vllm.platforms.rocm import on_gfx950
40 +else:
41 +
42 + def on_gfx950() -> bool:
43 + return False
44 +
45 +
46 +class Nvfp4QuantizationEmulationTritonExpertsReference(TritonExperts):
47 + """
48 + Extension of TritonExperts to support emulated NVFP4 MoE experts.
49 +
50 + It may be used for NVFP4 models when the device does not have
51 + native support for this dtype.
52 + """
53 +
54 + def __init__(
55 + self,
56 + moe_config: FusedMoEConfig,
57 + quant_config: FusedMoEQuantConfig,
58 + ):
59 + super().__init__(moe_config, quant_config)
60 +
61 + # `TritonExperts.apply` expects pre-dequantized weights,
62 + # which we handle in `apply` below.
63 + self.w1_scale_val = self.quant_config.w1_scale
64 + self.w2_scale_val = self.quant_config.w2_scale
65 +
66 + self.quant_config._w1.scale = None
67 + self.quant_config._w2.scale = None
68 +
69 + self.quantization_emulation = True
70 +
71 + @property
72 + def quant_dtype(self) -> torch.dtype | str | None:
73 + return "nvfp4"
74 +
75 + @property
76 + def a1_scale(self) -> torch.Tensor | None:
77 + return self.quant_config.a1_gscale
78 +
79 + @property
80 + def expects_unquantized_inputs(self) -> bool:
81 + return True
82 +
83 + @staticmethod
84 + def _supports_quant_scheme(
85 + weight_key: QuantKey | None,
86 + activation_key: QuantKey | None,
87 + ) -> bool:
88 + return (weight_key, activation_key) == (kNvfp4Static, kNvfp4Dynamic)
89 +
90 + def apply(
91 + self,
92 + output: torch.Tensor,
93 + hidden_states: torch.Tensor,
94 + w1: torch.Tensor,
95 + w2: torch.Tensor,
96 + topk_weights: torch.Tensor,
97 + topk_ids: torch.Tensor,
98 + activation: MoEActivation,
99 + global_num_experts: int,
100 + expert_map: torch.Tensor | None,
101 + a1q_scale: torch.Tensor | None,
102 + a2_scale: torch.Tensor | None,
103 + workspace13: torch.Tensor,
104 + workspace2: torch.Tensor,
105 + expert_tokens_meta: mk.ExpertTokensMetadata | None,
106 + apply_router_weight_on_input: bool,
107 + ):
108 + assert w1.dtype == torch.uint8
109 + assert w2.dtype == torch.uint8
110 +
111 + # Dequantize w1 from packed NVFP4 to fp16/bf16
112 + w13_global_scale = self.quant_config.g1_alphas
113 +
114 + w1_dequant = dequantize_to_dtype(
115 + tensor_fp4=w1,
116 + tensor_sf=self.w1_scale_val,
117 + global_scale=w13_global_scale,
118 + dtype=hidden_states.dtype,
119 + block_size=16,
120 + swizzle=False,
121 + )
122 +
123 + # Dequantize w2 from packed NVFP4 to fp16/bf16
124 + w2_global_scale = self.quant_config.g2_alphas
125 +
126 + w2_dequant = dequantize_to_dtype(
127 + tensor_fp4=w2,
128 + tensor_sf=self.w2_scale_val,
129 + global_scale=w2_global_scale,
130 + dtype=hidden_states.dtype,
131 + block_size=16,
132 + swizzle=False,
133 + )
134 +
135 + super().apply(
136 + output=output,
137 + hidden_states=hidden_states,
138 + w1=w1_dequant,
139 + w2=w2_dequant,
140 + topk_weights=topk_weights,
141 + topk_ids=topk_ids,
142 + activation=activation,
143 + global_num_experts=global_num_experts,
144 + expert_map=expert_map,
145 + a1q_scale=None,
146 + a2_scale=self.quant_config.a2_gscale,
147 + workspace13=workspace13,
148 + workspace2=workspace2,
149 + expert_tokens_meta=expert_tokens_meta,
150 + apply_router_weight_on_input=apply_router_weight_on_input,
151 + )
152 +
153 +
154 +@pytest.mark.parametrize(
155 + ("config_kwargs", "expected_reason"),
156 + [
157 + ({"has_bias": True}, "kernel does not support bias"),
158 + ({"is_lora_enabled": True}, "kernel does not support LoRA"),
159 + ],
160 +)
161 +def test_nvfp4_emulation_support_check_rejects_bias_and_lora(
162 + config_kwargs: dict[str, bool],
163 + expected_reason: str,
164 +) -> None:
165 + moe_config = FusedMoEConfig(
166 + num_experts=2,
167 + experts_per_token=1,
168 + hidden_dim=16,
169 + intermediate_size=16,
170 + num_local_experts=2,
171 + num_logical_experts=2,
172 + moe_parallel_config=FusedMoEParallelConfig.make_no_parallel(),
173 + activation=MoEActivation.SILU,
174 + in_dtype=torch.bfloat16,
175 + device="cuda",
176 + routing_method=RoutingMethodType.TopK,
177 + **config_kwargs,
178 + )
179 +
180 + supported, reason = Nvfp4QuantizationEmulationTritonExperts.is_supported_config(
181 + Nvfp4QuantizationEmulationTritonExperts,
182 + moe_config,
183 + kNvfp4Static,
184 + kNvfp4Dynamic,
185 + mk.FusedMoEActivationFormat.Standard,
186 + )
187 +
188 + assert not supported
189 + assert reason == expected_reason
190 +
191
192 @pytest.mark.skipif(
193 not current_platform.is_cuda_alike(),
@@ -306,3 +479,293 @@ def test_triton_nvfp4_quant_dequant(
479 f"min={ref_min:.3f}ms, max={ref_max:.3f}ms"
480 )
481 print(f" speedup: {speedup:.2f}x")
482 +
483 +
484 +MOE_MODEL_CONFIGS = {
485 + "nvidia/Qwen3-30B-A3B-NVFP4": {
486 + "shards": ["model-00001-of-00004.safetensors"],
487 + "expert_prefix": "model.layers.9.mlp.experts.",
488 + # Position of the expert index in the dot-split key.
489 + "expert_idx_pos": 5,
490 + },
491 + "nvidia/Kimi-K2.6-NVFP4": {
492 + "shards": [
493 + "model-00001-of-00060.safetensors",
494 + "model-00002-of-00060.safetensors",
495 + ],
496 + "expert_prefix": "language_model.model.layers.1.mlp.experts.",
497 + "expert_idx_pos": 6,
498 + },
499 +}
500 +
501 +
502 +def _load_nvfp4_moe_weights(
503 + model_id: str,
504 + tensor_parallel_size: int,
505 + max_experts: int | None = None,
506 +):
507 + """Load and stack NVFP4 MoE weights from checkpoint shards.
508 +
509 + Returns (w1, w1_scale, w1_gscale, w2, w2_scale, w2_gscale,
510 + a1_gscale, a2_gscale, num_experts, hidden_dim,
511 + intermediate_size).
512 +
513 + When max_experts is set, only the first max_experts experts are loaded.
514 +
515 + When tensor_parallel_size > 1, the N dimension of w1 and the K
516 + dimension of w2 are narrowed to the first TP shard (simulating
517 + column-parallel on w1 / row-parallel on w2).
518 + """
519 + cfg = MOE_MODEL_CONFIGS[model_id]
520 + shards = cast(list[str], cfg["shards"])
521 + checkpoint_path = huggingface_hub.snapshot_download(
522 + model_id,
523 + allow_patterns=shards,
524 + )
525 + expert_prefix = cfg["expert_prefix"]
526 + idx_pos = cast(int, cfg["expert_idx_pos"])
527 +
528 + # Collect all tensors across shards into a flat dict — an expert's
529 + # tensors may be split across multiple shard files.
530 + all_tensors: dict[str, torch.Tensor] = {}
531 + for shard_name in shards:
532 + shard_path = f"{checkpoint_path}/{shard_name}"
533 + with safe_open(shard_path, framework="pt", device="cpu") as f:
534 + for key in f.keys(): # noqa: SIM118
535 + if key.startswith(expert_prefix):
536 + all_tensors[key] = f.get_tensor(key)
537 +
538 + expert_indices = sorted(
539 + {
540 + int(key.split(".")[idx_pos])
541 + for key in all_tensors
542 + if key.endswith(".gate_proj.weight")
543 + }
544 + )
545 + if max_experts is not None:
546 + expert_indices = expert_indices[:max_experts]
547 + num_experts = len(expert_indices)
548 +
549 + gate_weights, up_weights, down_weights = [], [], []
550 + gate_scales, up_scales, down_scales = [], [], []
551 + gate_gscales, up_gscales, down_gscales = [], [], []
552 + a1_scales, a2_scales = [], []
553 +
554 + for idx in expert_indices:
555 + prefix = f"{expert_prefix}{idx}"
556 + gate_weights.append(all_tensors[f"{prefix}.gate_proj.weight"])
557 + gate_scales.append(all_tensors[f"{prefix}.gate_proj.weight_scale"])
558 + gate_gscales.append(all_tensors[f"{prefix}.gate_proj.weight_scale_2"])
559 + up_weights.append(all_tensors[f"{prefix}.up_proj.weight"])
560 + up_scales.append(all_tensors[f"{prefix}.up_proj.weight_scale"])
561 + up_gscales.append(all_tensors[f"{prefix}.up_proj.weight_scale_2"])
562 + down_weights.append(all_tensors[f"{prefix}.down_proj.weight"])
563 + down_scales.append(all_tensors[f"{prefix}.down_proj.weight_scale"])
564 + down_gscales.append(all_tensors[f"{prefix}.down_proj.weight_scale_2"])
565 + a1_scales.append(all_tensors[f"{prefix}.gate_proj.input_scale"])
566 + a2_scales.append(all_tensors[f"{prefix}.down_proj.input_scale"])
567 +
568 + # Stack into MoE format.
569 + # w1 = [E, 2*intermediate, hidden//2] (gate + up concatenated)
570 + w1 = torch.stack(
571 + [torch.cat([g, u], dim=0) for g, u in zip(gate_weights, up_weights)]
572 + ).cuda()
573 + w1_scale = torch.stack(
574 + [torch.cat([g, u], dim=0) for g, u in zip(gate_scales, up_scales)]
575 + ).cuda()
576 + w1_gscale = torch.stack(gate_gscales).cuda()
577 +
578 + # w2 = [E, hidden, intermediate//2]
579 + w2 = torch.stack(down_weights).cuda()
580 + w2_scale = torch.stack(down_scales).cuda()
581 + w2_gscale = torch.stack(down_gscales).cuda()
582 +
583 + a13_scale_raw = torch.stack(a1_scales).cuda()
584 + a2_scale_raw = torch.stack(a2_scales).cuda()
585 +
586 + # Apply EMULATION transforms (matches oracle/nvfp4.py).
587 + nvfp4_emulation_utils.kE2M1ToFloat_handle.val = (
588 + nvfp4_emulation_utils.kE2M1ToFloat_handle.val.cuda()
589 + )
590 + a1_gscale = 1.0 / a13_scale_raw.max().to(torch.float32)
591 + a2_gscale = 1.0 / a2_scale_raw.max().to(torch.float32)
592 +
593 + # ── Simulate TP sharding ──
594 + # w1 (gate_up): column-parallel → shard the N dimension (dim 1).
595 + # w2 (down): row-parallel → shard the K dimension (dim 2,
596 + # which is the packed K//2 dim).
597 + # Scales follow the same sharding on the corresponding dimension.
598 + tp = tensor_parallel_size
599 + if tp > 1:
600 + n1 = w1.size(1) // tp
601 + w1 = w1[:, :n1, :].contiguous()
602 + w1_scale = w1_scale[:, :n1, :].contiguous()
603 +
604 + k2_packed = w2.size(2) // tp
605 + k2_scale = w2_scale.size(2) // tp
606 + w2 = w2[:, :, :k2_packed].contiguous()
607 + w2_scale = w2_scale[:, :, :k2_scale].contiguous()
608 +
609 + hidden_dim = w1.size(2) * 2
610 + intermediate_size = w1.size(1) // 2
611 +
612 + return (
613 + w1,
614 + w1_scale,
615 + w1_gscale,
616 + w2,
617 + w2_scale,
618 + w2_gscale,
619 + a1_gscale,
620 + a2_gscale,
621 + num_experts,
622 + hidden_dim,
623 + intermediate_size,
624 + )
625 +
626 +
627 +@pytest.mark.skipif(
628 + not current_platform.is_cuda_alike(),
629 + reason="Triton NVFP4 kernel requires CUDA.",
630 +)
631 +@pytest.mark.parametrize("num_tokens", [1, 2, 4, 1024])
632 +@pytest.mark.parametrize("top_k", [4])
633 +@pytest.mark.parametrize("model_id", list(MOE_MODEL_CONFIGS.keys()))
634 +@pytest.mark.parametrize(
635 + "tensor_parallel_size",
636 + [pytest.param(val, id=f"tensor_parallel_size:{val}") for val in [1, 2, 4, 8]],
637 +)
638 +def test_nvfp4_moe_correctness(
639 + num_tokens: int,
640 + top_k: int,
641 + model_id: str,
642 + tensor_parallel_size: int,
643 +) -> None:
644 + """Compare Nvfp4QuantizationEmulationTritonExperts (fused weight dequant + compute)
645 + against the unfused reference Nvfp4QuantizationEmulationTritonExpertsReference.
646 +
647 + Both must produce bit-identical results.
648 + """
649 + num_test_experts = max(8, top_k)
650 + (
651 + w1,
652 + w1_scale,
653 + w1_gscale,
654 + w2,
655 + w2_scale,
656 + w2_gscale,
657 + a1_gscale,
658 + a2_gscale,
659 + num_experts,
660 + hidden_dim,
661 + intermediate_size,
662 + ) = _load_nvfp4_moe_weights(
663 + model_id,
664 + tensor_parallel_size,
665 + max_experts=num_test_experts,
666 + )
667 +
668 + moe_config = FusedMoEConfig(
669 + num_experts=num_experts,
670 + experts_per_token=top_k,
671 + hidden_dim=hidden_dim,
672 + intermediate_size=intermediate_size,
673 + num_local_experts=num_experts,
674 + num_logical_experts=num_experts,
675 + moe_parallel_config=FusedMoEParallelConfig.make_no_parallel(),
676 + activation=MoEActivation.SILU,
677 + in_dtype=torch.bfloat16,
678 + device="cuda",
679 + routing_method=RoutingMethodType.TopK,
680 + max_num_tokens=512,
681 + )
682 +
683 + def _make_quant_config():
684 + return nvfp4_moe_quant_config(
685 + g1_alphas=w1_gscale.clone(),
686 + g2_alphas=w2_gscale.clone(),
687 + a1_gscale=a1_gscale.clone(),
688 + a2_gscale=a2_gscale.clone(),
689 + w1_scale=w1_scale.clone(),
690 + w2_scale=w2_scale.clone(),
691 + )
692 +
693 + ref_experts = Nvfp4QuantizationEmulationTritonExpertsReference(
694 + moe_config=moe_config,
695 + quant_config=_make_quant_config(),
696 + )
697 + fused_experts = Nvfp4QuantizationEmulationTritonExperts(
698 + moe_config=moe_config,
699 + quant_config=_make_quant_config(),
700 + )
701 +
702 + torch.manual_seed(42)
703 + hidden_states = torch.randn(
704 + num_tokens, hidden_dim, dtype=torch.bfloat16, device="cuda"
705 + )
706 +
707 + topk_weights = torch.randn(
708 + num_tokens, top_k, dtype=torch.float32, device="cuda"
709 + ).softmax(dim=-1)
710 + topk_ids = torch.stack(
711 + [torch.randperm(num_experts, device="cuda")[:top_k] for _ in range(num_tokens)]
712 + ).to(torch.int32)
713 +
714 + N = w1.size(1) # 2 * intermediate
715 + K = hidden_dim
716 +
717 + ws13_size = num_tokens * top_k * max(intermediate_size, K)
718 + ws2_size = num_tokens * top_k * max(N, K)
719 +
720 + workspace13_ref = torch.zeros(ws13_size, dtype=torch.bfloat16, device="cuda")
721 + workspace2_ref = torch.zeros(ws2_size, dtype=torch.bfloat16, device="cuda")
722 + output_ref = torch.zeros(num_tokens, K, dtype=torch.bfloat16, device="cuda")
723 +
724 + workspace13_fused = torch.zeros_like(workspace13_ref)
725 + workspace2_fused = torch.zeros_like(workspace2_ref)
726 + output_fused = torch.zeros_like(output_ref)
727 +
728 + apply_kwargs = dict(
729 + hidden_states=hidden_states,
730 + w1=w1,
731 + w2=w2,
732 + topk_weights=topk_weights,
733 + topk_ids=topk_ids,
734 + activation=MoEActivation.SILU,
735 + global_num_experts=num_experts,
736 + expert_map=None,
737 + a1q_scale=None,
738 + a2_scale=None,
739 + expert_tokens_meta=None,
740 + apply_router_weight_on_input=False,
741 + )
742 +
743 + # Unfused reference.
744 + ref_experts.apply(
745 + output=output_ref,
746 + workspace13=workspace13_ref,
747 + workspace2=workspace2_ref,
748 + **apply_kwargs,
749 + )
750 +
751 + # Fused implementation.
752 + fused_experts.apply(
753 + output=output_fused,
754 + workspace13=workspace13_fused,
755 + workspace2=workspace2_fused,
756 + **apply_kwargs,
757 + )
758 +
759 + # Not strict equality on H100, MI325, MI300 (< 0.1% elements).
760 + # The fused on-the-fly dequant path can lower to a slightly
761 + # different Triton/MMA tiling than the pre-dequantized
762 + # reference; experiments with reference-like tiling/masking
763 + # reduced some diffs were not kept because they regress
764 + # the fused kernel speed.
765 + # Strict equality validated on MI355.
766 + torch.testing.assert_close(
767 + output_fused,
768 + output_ref,
769 + atol=0.0 if on_gfx950() else 0.02,
770 + rtol=0,
771 + )
vllm/model_executor/layers/fused_moe/experts/nvfp4_emulation_moe.py
+424 -37
@@ -11,6 +11,8 @@ Weights are dequantized on the fly during each forward, we fall back to calling
11 is applied on `a13`, `a2`.
12 """
13
14 +from typing import Any
15 +
16 import torch
17
18 import vllm.model_executor.layers.fused_moe.modular_kernel as mk
@@ -21,18 +23,316 @@ from vllm.model_executor.layers.fused_moe.config import (
23 FusedMoEQuantConfig,
24 )
25 from vllm.model_executor.layers.fused_moe.experts.triton_moe import TritonExperts
26 +from vllm.model_executor.layers.fused_moe.fused_moe import (
27 + try_get_optimal_moe_config,
28 + write_zeros_to_output,
29 +)
30 +from vllm.model_executor.layers.fused_moe.moe_align_block_size import (
31 + moe_align_block_size,
32 +)
33 +from vllm.model_executor.layers.fused_moe.utils import (
34 + _resize_cache,
35 + moe_kernel_quantize_input,
36 +)
37 from vllm.model_executor.layers.quantization.utils.nvfp4_emulation_utils import (
25 - dequantize_to_dtype,
38 + _e2m1_inline,
39 )
40 from vllm.model_executor.layers.quantization.utils.quant_utils import (
41 QuantKey,
42 kNvfp4Dynamic,
43 kNvfp4Static,
44 )
45 +from vllm.triton_utils import tl, triton
46
47 logger = init_logger(__name__)
48
49
50 +@triton.jit
51 +def fused_moe_nvfp4_emulation_kernel(
52 + a_ptr,
53 + b_ptr,
54 + c_ptr,
55 + b_scale_ptr,
56 + w_global_scale_ptr,
57 + topk_weights_ptr,
58 + sorted_token_ids_ptr,
59 + expert_ids_ptr,
60 + num_tokens_post_padded_ptr,
61 + N: tl.constexpr,
62 + K: tl.constexpr,
63 + EM,
64 + num_valid_tokens,
65 + # Strides — A [M, K]
66 + stride_am,
67 + stride_ak,
68 + # Strides — B [E, N, K//2], passed as (expert, K-packed, N)
69 + stride_be,
70 + stride_bk,
71 + stride_bn,
72 + # Strides — C [M, topk, N]
73 + stride_cm,
74 + stride_cn,
75 + # Strides — B_scale [E, N, K//BLOCK], passed as (expert, K-scale, N)
76 + stride_bse,
77 + stride_bsk,
78 + stride_bsn,
79 + block_k_diviable: tl.constexpr,
80 + BLOCK_SIZE_M: tl.constexpr,
81 + BLOCK_SIZE_N: tl.constexpr,
82 + BLOCK_SIZE_K: tl.constexpr,
83 + GROUP_SIZE_M: tl.constexpr,
84 + MUL_ROUTED_WEIGHT: tl.constexpr,
85 + top_k: tl.constexpr,
86 + compute_type: tl.constexpr,
87 + group_size: tl.constexpr,
88 +):
89 + """
90 + Fused MoE kernel for emulated NVFP4 weight-only dequantization + GEMM.
91 +
92 + Activations A are BF16 (already QDQ'd externally).
93 + Weights B are packed uint8 NVFP4 [E, N, K//2] — two FP4 values per byte
94 + along the K dimension.
95 + B_scale holds per-block FP8-E4M3 scales [E, N, K // group_size].
96 + w_global_scale is a per-expert scalar global scale.
97 +
98 + The dequantization formula per element is:
99 + w_float = e2m1_decode(nibble) * (block_scale_fp8 * global_scale)
100 +
101 + Weight loading optimization: each packed byte is loaded exactly once as
102 + a [BLOCK_SIZE_N, BLOCK_SIZE_K // 2] tile (N-major), both nibbles are
103 + extracted, decoded and scaled, then tl.interleave produces the
104 + [BLOCK_SIZE_N, BLOCK_SIZE_K] dequantized tile which is transposed to
105 + [BLOCK_SIZE_K, BLOCK_SIZE_N] for tl.dot.
106 + """
107 + BLOCK_SIZE_K_PACKED: tl.constexpr = BLOCK_SIZE_K // 2
108 +
109 + # Map program ids to the block of C it should compute.
110 + pid = tl.program_id(axis=0)
111 + num_pid_m = tl.cdiv(EM, BLOCK_SIZE_M)
112 + num_pid_n = tl.cdiv(N, BLOCK_SIZE_N)
113 + num_pid_in_group = GROUP_SIZE_M * num_pid_n
114 + group_id = pid // num_pid_in_group
115 + first_pid_m = group_id * GROUP_SIZE_M
116 + group_size_m = min(num_pid_m - first_pid_m, GROUP_SIZE_M)
117 + pid_m = first_pid_m + ((pid % num_pid_in_group) % group_size_m)
118 + pid_n = (pid % num_pid_in_group) // group_size_m
119 +
120 + # Token / expert setup
121 + num_tokens_post_padded = tl.load(num_tokens_post_padded_ptr)
122 + if pid_m * BLOCK_SIZE_M >= num_tokens_post_padded:
123 + return
124 +
125 + offs_token_id = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M).to(tl.int64)
126 + offs_token = tl.load(sorted_token_ids_ptr + offs_token_id).to(tl.int64)
127 + token_mask = offs_token < num_valid_tokens
128 +
129 + off_experts = tl.load(expert_ids_ptr + pid_m).to(tl.int64)
130 + if off_experts == -1:
131 + write_zeros_to_output(
132 + c_ptr,
133 + stride_cm,
134 + stride_cn,
135 + pid_n,
136 + N,
137 + offs_token,
138 + token_mask,
139 + BLOCK_SIZE_M,
140 + BLOCK_SIZE_N,
141 + compute_type,
142 + )
143 + return
144 +
145 + # Pointer setup
146 + offs_bn = (pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N).to(tl.int64)) % N
147 + offs_k = tl.arange(0, BLOCK_SIZE_K)
148 + offs_k_packed = tl.arange(0, BLOCK_SIZE_K_PACKED)
149 +
150 + # A pointers: [BLOCK_SIZE_M, BLOCK_SIZE_K]
151 + a_ptrs = a_ptr + (
152 + offs_token[:, None] // top_k * stride_am + offs_k[None, :] * stride_ak
153 + )
154 +
155 + # B pointers: [BLOCK_SIZE_N, BLOCK_SIZE_K_PACKED] — N-major so that
156 + # tl.interleave (which operates on the last dim) produces a
157 + # [BLOCK_SIZE_N, BLOCK_SIZE_K] tile that we transpose for tl.dot.
158 + # Each unique byte is loaded exactly once.
159 + b_ptrs = (
160 + b_ptr
161 + + off_experts * stride_be
162 + + offs_bn[:, None] * stride_bn
163 + + offs_k_packed[None, :] * stride_bk
164 + )
165 +
166 + # B_scale pointers: [BLOCK_SIZE_N, BLOCK_SIZE_K_PACKED] — same
167 + # N-major layout. Each packed byte index covers 2 K elements that
168 + # always fall within the same group (group_size=16, so each group
169 + # spans 8 packed bytes). We can therefore index the scale using
170 + # offs_k_packed directly.
171 + # Note: group_size_packed = group_size // 2 maps packed indices to
172 + # scale indices the same way unpacked indices map via group_size.
173 + group_size_packed: tl.constexpr = group_size // 2
174 +
175 + # Load per-expert global scale (scalar).
176 + w_global_scale = tl.load(w_global_scale_ptr + off_experts).to(tl.float32)
177 +
178 + # K-loop with FP32 accumulation
179 + accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32)
180 +
181 + for k in range(0, tl.cdiv(K, BLOCK_SIZE_K)):
182 + # Load A tile [BLOCK_SIZE_M, BLOCK_SIZE_K].
183 + if block_k_diviable:
184 + a = tl.load(
185 + a_ptrs,
186 + mask=token_mask[:, None],
187 + other=0.0,
188 + )
189 + else:
190 + a = tl.load(
191 + a_ptrs,
192 + mask=token_mask[:, None] & (offs_k[None, :] < K - k * BLOCK_SIZE_K),
193 + other=0.0,
194 + )
195 +
196 + # Load packed weight tile [BLOCK_SIZE_N, BLOCK_SIZE_K_PACKED].
197 + if block_k_diviable:
198 + raw_bytes = tl.load(b_ptrs)
199 + else:
200 + kp_mask = offs_k_packed[None, :] < (K // 2) - k * BLOCK_SIZE_K_PACKED
201 + raw_bytes = tl.load(b_ptrs, mask=kp_mask, other=0)
202 +
203 + # Extract both nibbles from each byte (each [N, K_packed]).
204 + low_nibble = raw_bytes & 0x0F
205 + high_nibble = (raw_bytes >> 4) & 0x0F
206 +
207 + low_decoded = _e2m1_inline(low_nibble)
208 + high_decoded = _e2m1_inline(high_nibble)
209 +
210 + # Load and apply per-block FP8 scales.
211 + # Scale shape: [BLOCK_SIZE_N, BLOCK_SIZE_K_PACKED], one scale per
212 + # group_size_packed packed elements.
213 + b_scale_ptrs = (
214 + b_scale_ptr
215 + + off_experts * stride_bse
216 + + offs_bn[:, None] * stride_bsn
217 + + ((offs_k_packed[None, :] + BLOCK_SIZE_K_PACKED * k) // group_size_packed)
218 + * stride_bsk
219 + )
220 + if block_k_diviable:
221 + b_scale_raw = tl.load(b_scale_ptrs)
222 + else:
223 + b_scale_raw = tl.load(b_scale_ptrs, mask=kp_mask, other=0.0)
224 +
225 + b_scale = tl.cast(b_scale_raw, tl.float8e4nv, bitcast=True).to(tl.float32)
226 + b_scale = b_scale * w_global_scale
227 +
228 + # Scale both halves with the same per-block scale (the two
229 + # elements packed in one byte always belong to the same group).
230 + low_scaled = low_decoded * b_scale
231 + high_scaled = high_decoded * b_scale
232 +
233 + # Interleave along last dim: [N, K_packed] x2 -> [N, K],
234 + # then transpose to [K, N] for tl.dot.
235 + b = tl.trans(tl.interleave(low_scaled, high_scaled)).to(compute_type)
236 +
237 + accumulator = tl.dot(a, b, acc=accumulator)
238 +
239 + # Advance pointers along K.
240 + a_ptrs += BLOCK_SIZE_K * stride_ak
241 + b_ptrs += BLOCK_SIZE_K_PACKED * stride_bk
242 +
243 + # Router weight multiplication (in float32 for stability)
244 + if MUL_ROUTED_WEIGHT:
245 + moe_weight = tl.load(topk_weights_ptr + offs_token, mask=token_mask, other=0)
246 + accumulator = accumulator * moe_weight[:, None]
247 +
248 + accumulator = accumulator.to(compute_type)
249 +
250 + # Write output
251 + offs_cn = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)
252 + c_ptrs = c_ptr + stride_cm * offs_token[:, None] + stride_cn * offs_cn[None, :]
253 + c_mask = token_mask[:, None] & (offs_cn[None, :] < N)
254 + tl.store(c_ptrs, accumulator, mask=c_mask)
255 +
256 +
257 +def invoke_fused_moe_nvfp4_emulation_kernel(
258 + A: torch.Tensor,
259 + B: torch.Tensor,
260 + C: torch.Tensor,
261 + B_scale: torch.Tensor,
262 + act_global_scale: torch.Tensor,
263 + w_global_scale: torch.Tensor,
264 + topk_weights: torch.Tensor | None,
265 + sorted_token_ids: torch.Tensor,
266 + expert_ids: torch.Tensor,
267 + num_tokens_post_padded: torch.Tensor,
268 + mul_routed_weight: bool,
269 + top_k: int,
270 + config: dict[str, Any],
271 + compute_type: tl.dtype,
272 +):
273 + """Launch the fused NVFP4 emulation MoE kernel.
274 +
275 + B has shape [E, N, K_packed] where K_packed = K // 2 (two FP4 per byte).
276 + B_scale has shape [E, N, K // group_size] in FP8-E4M3 (stored as uint8).
277 + w_global_scale has shape [E] (per-expert scalar).
278 + """
279 + assert B_scale is not None and B_scale.ndim == 3
280 +
281 + N = B.size(1)
282 + K = A.size(1)
283 +
284 + M = A.size(0)
285 + num_tokens = M * top_k
286 +
287 + EM = sorted_token_ids.size(0)
288 + if A.size(0) < config["BLOCK_SIZE_M"]:
289 + EM = min(
290 + sorted_token_ids.size(0),
291 + A.size(0) * top_k * config["BLOCK_SIZE_M"],
292 + )
293 +
294 + grid = lambda META: (
295 + triton.cdiv(EM, META["BLOCK_SIZE_M"]) * triton.cdiv(N, META["BLOCK_SIZE_N"]),
296 + )
297 +
298 + fused_moe_nvfp4_emulation_kernel[grid](
299 + A,
300 + B,
301 + C,
302 + B_scale,
303 + w_global_scale,
304 + topk_weights,
305 + sorted_token_ids,
306 + expert_ids,
307 + num_tokens_post_padded,
308 + N,
309 + K,
310 + EM,
311 + num_tokens,
312 + A.stride(0),
313 + A.stride(1),
314 + # B is [E, N, K//2]: swap N and K strides so kernel indexes [K, N].
315 + B.stride(0),
316 + B.stride(2),
317 + B.stride(1),
318 + C.stride(1),
319 + C.stride(2),
320 + # B_scale is [E, N, K//group]: swap N and K strides likewise.
321 + B_scale.stride(0),
322 + B_scale.stride(2),
323 + B_scale.stride(1),
324 + block_k_diviable=K % config["BLOCK_SIZE_K"] == 0,
325 + MUL_ROUTED_WEIGHT=mul_routed_weight,
326 + top_k=top_k,
327 + compute_type=compute_type,
328 + group_size=16,
329 + BLOCK_SIZE_M=config["BLOCK_SIZE_M"],
330 + BLOCK_SIZE_N=config["BLOCK_SIZE_N"],
331 + BLOCK_SIZE_K=config["BLOCK_SIZE_K"],
332 + GROUP_SIZE_M=config["GROUP_SIZE_M"],
333 + )
334 +
335 +
336 class Nvfp4QuantizationEmulationTritonExperts(TritonExperts):
337 """
338 Extension of TritonExperts to support emulated NVFP4 MoE experts.
@@ -72,6 +372,31 @@ class Nvfp4QuantizationEmulationTritonExperts(TritonExperts):
372 def expects_unquantized_inputs(self) -> bool:
373 return True
374
375 + @staticmethod
376 + def supports_lora() -> bool:
377 + return False
378 +
379 + @staticmethod
380 + def is_supported_config(
381 + cls: type[mk.FusedMoEExperts],
382 + moe_config: FusedMoEConfig,
383 + weight_key: QuantKey | None,
384 + activation_key: QuantKey | None,
385 + activation_format: mk.FusedMoEActivationFormat,
386 + ) -> tuple[bool, str | None]:
387 + if moe_config.is_lora_enabled:
388 + return False, "kernel does not support LoRA"
389 + if moe_config.has_bias:
390 + return False, "kernel does not support bias"
391 +
392 + return TritonExperts.is_supported_config(
393 + cls,
394 + moe_config,
395 + weight_key,
396 + activation_key,
397 + activation_format,
398 + )
399 +
400 @staticmethod
401 def _supports_quant_scheme(
402 weight_key: QuantKey | None,
@@ -109,47 +434,109 @@ class Nvfp4QuantizationEmulationTritonExperts(TritonExperts):
434 # w2 shape: [num_experts, hidden_size, intermediate_size//2]
435 assert w1.dtype == torch.uint8
436 assert w2.dtype == torch.uint8
437 + assert hidden_states.is_contiguous()
438 + assert hidden_states.dim() == 2
439
113 - # Dequantize w1 from packed NVFP4 to fp16/bf16
114 - w13_global_scale = self.quant_config.g1_alphas
440 + K = hidden_states.size(-1)
441 + assert w1.size(2) * 2 == K, f"Hidden size mismatch: {K} != {w1.size(2) * 2}"
442
116 - w1_dequant = dequantize_to_dtype(
117 - tensor_fp4=w1,
118 - tensor_sf=self.w1_scale_val,
119 - global_scale=w13_global_scale,
120 - dtype=hidden_states.dtype,
121 - block_size=16,
122 - swizzle=False,
443 + E, num_tokens, N, _, top_k_num = self.moe_problem_size(
444 + hidden_states, w1, w2, topk_ids
445 )
446
125 - # Dequantize w2 from packed NVFP4 to fp16/bf16
126 - w2_global_scale = self.quant_config.g2_alphas
447 + if global_num_experts == -1:
448 + global_num_experts = E
449
128 - w2_dequant = dequantize_to_dtype(
129 - tensor_fp4=w2,
130 - tensor_sf=self.w2_scale_val,
131 - global_scale=w2_global_scale,
132 - dtype=hidden_states.dtype,
133 - block_size=16,
134 - swizzle=False,
450 + # TODO: There is actually no support for tuning of the underlying triton
451 + # hyperparameters in benchmarks/kernels/benchmark_moe.py, to be added.
452 + config = try_get_optimal_moe_config(
453 + w1.size(),
454 + w2.size(),
455 + top_k_num,
456 + self.quant_config.config_name(hidden_states.dtype),
457 + num_tokens,
458 + block_shape=None,
459 )
460
137 - # Activation quantization/dequantization is deferred to
138 - # `moe_kernel_quantize_input` in TritonExperts.apply.
139 - super().apply(
140 - output=output,
141 - hidden_states=hidden_states,
142 - w1=w1_dequant,
143 - w2=w2_dequant,
144 - topk_weights=topk_weights,
145 - topk_ids=topk_ids,
146 - activation=activation,
147 - global_num_experts=global_num_experts,
148 - expert_map=expert_map,
149 - a1q_scale=None,
150 - a2_scale=self.quant_config.a2_gscale,
151 - workspace13=workspace13,
152 - workspace2=workspace2,
153 - expert_tokens_meta=expert_tokens_meta,
154 - apply_router_weight_on_input=apply_router_weight_on_input,
461 + if hidden_states.dtype == torch.bfloat16:
462 + compute_type = tl.bfloat16
463 + elif hidden_states.dtype == torch.float16:
464 + compute_type = tl.float16
465 + elif hidden_states.dtype == torch.float32:
466 + compute_type = tl.float32
467 + else:
468 + raise ValueError(f"Unsupported compute_type: {hidden_states.dtype}")
469 +
470 + intermediate_cache1 = _resize_cache(workspace2, (num_tokens, top_k_num, N))
471 + activation_out_dim = self.adjust_N_for_activation(N, activation)
472 + intermediate_cache2 = _resize_cache(
473 + workspace13, (num_tokens * top_k_num, activation_out_dim)
474 )
475 + intermediate_cache3 = _resize_cache(workspace2, (num_tokens, top_k_num, K))
476 +
477 + sorted_token_ids, expert_ids, num_tokens_post_padded = moe_align_block_size(
478 + topk_ids,
479 + config["BLOCK_SIZE_M"],
480 + global_num_experts,
481 + expert_map,
482 + )
483 +
484 + # Activation NVFP4 QDQ.
485 + hidden_states_qdq, _ = moe_kernel_quantize_input(
486 + A=hidden_states,
487 + A_scale=self.quant_config.a1_gscale,
488 + quant_dtype="nvfp4",
489 + per_act_token_quant=False,
490 + quantization_emulation=True,
491 + )
492 +
493 + # w13: fused weight dequant + GEMM.
494 + invoke_fused_moe_nvfp4_emulation_kernel(
495 + hidden_states_qdq,
496 + w1,
497 + intermediate_cache1,
498 + self.w1_scale_val,
499 + self.quant_config.a1_gscale,
500 + self.quant_config.g1_alphas,
501 + None, # topk_weights — applied after w2
502 + sorted_token_ids,
503 + expert_ids,
504 + num_tokens_post_padded,
505 + False, # mul_routed_weight
506 + top_k_num,
507 + config,
508 + compute_type=compute_type,
509 + )
510 +
511 + self.activation(
512 + activation, intermediate_cache2, intermediate_cache1.view(-1, N)
513 + )
514 +
515 + # Activation NVFP4 QDQ.
516 + intermediate_cache2_qdq, _ = moe_kernel_quantize_input(
517 + A=intermediate_cache2,
518 + A_scale=self.quant_config.a2_gscale,
519 + quant_dtype="nvfp4",
520 + per_act_token_quant=False,
521 + quantization_emulation=True,
522 + )
523 +
524 + # w2: fused weight dequant + GEMM.
525 + invoke_fused_moe_nvfp4_emulation_kernel(
526 + intermediate_cache2_qdq,
527 + w2,
528 + intermediate_cache3,
529 + self.w2_scale_val,
530 + self.quant_config.a2_gscale,
531 + self.quant_config.g2_alphas,
532 + topk_weights,
533 + sorted_token_ids,
534 + expert_ids,
535 + num_tokens_post_padded,
536 + not apply_router_weight_on_input,
537 + 1,
538 + config,
539 + compute_type=compute_type,
540 + )
541 +
542 + self.moe_sum(intermediate_cache3, output)
vllm/model_executor/layers/quantization/utils/nvfp4_emulation_utils.py
+19 -31
@@ -23,28 +23,24 @@ kE2M1ToFloat_handle = SimpleNamespace(
23
24
25 @triton.jit
26 -def _e2m1_inline(magnitude):
27 - """Inline E2M1 lookup using binary tree - 3 levels instead of 7 sequential.
26 +def _e2m1_inline(nibble):
27 + """Decode an NVFP4 nibble (4 bits: 1 sign + 3 magnitude) to float32.
28
29 - Maps 3-bit magnitude to float: [0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0]
30 - Uses bit decomposition for fewer comparisons.
29 + Uses direct IEEE 754 bit construction.
30 + For magnitudes 2-7 the FP32 bit pattern is 0x3F000000 + (mag << 22),
31 + which is a single shift + add + bitcast. Magnitudes 0 (zero) and 1
32 + (E2M1 subnormal = 0.5) are patched with two tl.where ops.
33 """
32 - # Bit 2 (MSB): separates 0-3 from 4-7
33 - # Bit 1: separates within groups
34 - # Bit 0 (LSB): separates within pairs
35 - b2 = (magnitude >> 2) & 1 # 0 for mag 0-3, 1 for mag 4-7
36 - b1 = (magnitude >> 1) & 1 # middle bit
37 - b0 = magnitude & 1 # LSB
38 -
39 - # For mag 0-3: [0.0, 0.5, 1.0, 1.5]
40 - low_group = tl.where(
41 - b1 == 1, tl.where(b0 == 1, 1.5, 1.0), tl.where(b0 == 1, 0.5, 0.0)
42 - )
43 - # For mag 4-7: [2.0, 3.0, 4.0, 6.0]
44 - high_group = tl.where(
45 - b1 == 1, tl.where(b0 == 1, 6.0, 4.0), tl.where(b0 == 1, 3.0, 2.0)
46 - )
47 - return tl.where(b2 == 1, high_group, low_group)
34 + magnitude = nibble & 0x07
35 + sign = (nibble >> 3) & 1
36 +
37 + fp32_bits = 0x3F000000 + (magnitude.to(tl.int32) << 22)
38 + val = fp32_bits.to(tl.float32, bitcast=True)
39 +
40 + val = tl.where(magnitude == 0, 0.0, val)
41 + val = tl.where(magnitude == 1, 0.5, val)
42 +
43 + return tl.where(sign == 1, -val, val)
44
45
46 @triton.jit
@@ -65,7 +61,7 @@ def _dequantize_nvfp4_kernel(
61 """
62 BLOCK_PACKED: tl.constexpr = BLOCK_SIZE // 2
63
68 - row_idx = tl.program_id(0)
64 + row_idx = tl.program_id(0).to(tl.int64)
65 tile_idx = tl.program_id(1)
66
67 if has_batch_global_scale:
@@ -105,16 +101,8 @@ def _dequantize_nvfp4_kernel(
101 low_nibble = raw_bytes & 0x0F
102 high_nibble = (raw_bytes >> 4) & 0x0F
103
108 - # Binary tree E2M1 decode
109 - low_mag = low_nibble & 0x07
110 - low_val = _e2m1_inline(low_mag)
111 - low_sign = (low_nibble >> 3) & 1
112 - low_result = tl.where(low_sign == 1, -low_val, low_val) * scale_values
113 -
114 - high_mag = high_nibble & 0x07
115 - high_val = _e2m1_inline(high_mag)
116 - high_sign = (high_nibble >> 3) & 1
117 - high_result = tl.where(high_sign == 1, -high_val, high_val) * scale_values
104 + low_result = _e2m1_inline(low_nibble) * scale_values
105 + high_result = _e2m1_inline(high_nibble) * scale_values
106
107 # Interleave for coalesced contiguous store
108 result = tl.interleave(low_result, high_result)