[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>

tests/kernels/quantization/test_nvfp4_emulation.py

@@ -1,10 +1,25 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+from typing import cast
+
 import huggingface_hub
 import pytest
 import torch
 from safetensors import safe_open
 
+import vllm.model_executor.layers.fused_moe.modular_kernel as mk
+from vllm.model_executor.layers.fused_moe.activation import MoEActivation
+from vllm.model_executor.layers.fused_moe.config import (
+    FusedMoEConfig,
+    FusedMoEParallelConfig,
+    FusedMoEQuantConfig,
+    RoutingMethodType,
+    nvfp4_moe_quant_config,
+)
+from vllm.model_executor.layers.fused_moe.experts.nvfp4_emulation_moe import (
+    Nvfp4QuantizationEmulationTritonExperts,
+)
+from vllm.model_executor.layers.fused_moe.experts.triton_moe import TritonExperts
 from vllm.model_executor.layers.quantization.utils import (
     nvfp4_emulation_utils,
 )
@@ -12,9 +27,167 @@ from vllm.model_executor.layers.quantization.utils.nvfp4_emulation_utils import
     dequantize_to_dtype,
     ref_nvfp4_quant_dequant,
 )
+from vllm.model_executor.layers.quantization.utils.quant_utils import (
+    QuantKey,
+    kNvfp4Dynamic,
+    kNvfp4Static,
+)
 from vllm.platforms import current_platform
 from vllm.triton_utils import triton
 
+if current_platform.is_rocm():
+    from vllm.platforms.rocm import on_gfx950
+else:
+
+    def on_gfx950() -> bool:
+        return False
+
+
+class Nvfp4QuantizationEmulationTritonExpertsReference(TritonExperts):
+    """
+    Extension of TritonExperts to support emulated NVFP4 MoE experts.
+
+    It may be used for NVFP4 models when the device does not have
+    native support for this dtype.
+    """
+
+    def __init__(
+        self,
+        moe_config: FusedMoEConfig,
+        quant_config: FusedMoEQuantConfig,
+    ):
+        super().__init__(moe_config, quant_config)
+
+        # `TritonExperts.apply` expects pre-dequantized weights,
+        # which we handle in `apply` below.
+        self.w1_scale_val = self.quant_config.w1_scale
+        self.w2_scale_val = self.quant_config.w2_scale
+
+        self.quant_config._w1.scale = None
+        self.quant_config._w2.scale = None
+
+        self.quantization_emulation = True
+
+    @property
+    def quant_dtype(self) -> torch.dtype | str | None:
+        return "nvfp4"
+
+    @property
+    def a1_scale(self) -> torch.Tensor | None:
+        return self.quant_config.a1_gscale
+
+    @property
+    def expects_unquantized_inputs(self) -> bool:
+        return True
+
+    @staticmethod
+    def _supports_quant_scheme(
+        weight_key: QuantKey | None,
+        activation_key: QuantKey | None,
+    ) -> bool:
+        return (weight_key, activation_key) == (kNvfp4Static, kNvfp4Dynamic)
+
+    def apply(
+        self,
+        output: torch.Tensor,
+        hidden_states: torch.Tensor,
+        w1: torch.Tensor,
+        w2: torch.Tensor,
+        topk_weights: torch.Tensor,
+        topk_ids: torch.Tensor,
+        activation: MoEActivation,
+        global_num_experts: int,
+        expert_map: torch.Tensor | None,
+        a1q_scale: torch.Tensor | None,
+        a2_scale: torch.Tensor | None,
+        workspace13: torch.Tensor,
+        workspace2: torch.Tensor,
+        expert_tokens_meta: mk.ExpertTokensMetadata | None,
+        apply_router_weight_on_input: bool,
+    ):
+        assert w1.dtype == torch.uint8
+        assert w2.dtype == torch.uint8
+
+        # Dequantize w1 from packed NVFP4 to fp16/bf16
+        w13_global_scale = self.quant_config.g1_alphas
+
+        w1_dequant = dequantize_to_dtype(
+            tensor_fp4=w1,
+            tensor_sf=self.w1_scale_val,
+            global_scale=w13_global_scale,
+            dtype=hidden_states.dtype,
+            block_size=16,
+            swizzle=False,
+        )
+
+        # Dequantize w2 from packed NVFP4 to fp16/bf16
+        w2_global_scale = self.quant_config.g2_alphas
+
+        w2_dequant = dequantize_to_dtype(
+            tensor_fp4=w2,
+            tensor_sf=self.w2_scale_val,
+            global_scale=w2_global_scale,
+            dtype=hidden_states.dtype,
+            block_size=16,
+            swizzle=False,
+        )
+
+        super().apply(
+            output=output,
+            hidden_states=hidden_states,
+            w1=w1_dequant,
+            w2=w2_dequant,
+            topk_weights=topk_weights,
+            topk_ids=topk_ids,
+            activation=activation,
+            global_num_experts=global_num_experts,
+            expert_map=expert_map,
+            a1q_scale=None,
+            a2_scale=self.quant_config.a2_gscale,
+            workspace13=workspace13,
+            workspace2=workspace2,
+            expert_tokens_meta=expert_tokens_meta,
+            apply_router_weight_on_input=apply_router_weight_on_input,
+        )
+
+
+@pytest.mark.parametrize(
+    ("config_kwargs", "expected_reason"),
+    [
+        ({"has_bias": True}, "kernel does not support bias"),
+        ({"is_lora_enabled": True}, "kernel does not support LoRA"),
+    ],
+)
+def test_nvfp4_emulation_support_check_rejects_bias_and_lora(
+    config_kwargs: dict[str, bool],
+    expected_reason: str,
+) -> None:
+    moe_config = FusedMoEConfig(
+        num_experts=2,
+        experts_per_token=1,
+        hidden_dim=16,
+        intermediate_size=16,
+        num_local_experts=2,
+        num_logical_experts=2,
+        moe_parallel_config=FusedMoEParallelConfig.make_no_parallel(),
+        activation=MoEActivation.SILU,
+        in_dtype=torch.bfloat16,
+        device="cuda",
+        routing_method=RoutingMethodType.TopK,
+        **config_kwargs,
+    )
+
+    supported, reason = Nvfp4QuantizationEmulationTritonExperts.is_supported_config(
+        Nvfp4QuantizationEmulationTritonExperts,
+        moe_config,
+        kNvfp4Static,
+        kNvfp4Dynamic,
+        mk.FusedMoEActivationFormat.Standard,
+    )
+
+    assert not supported
+    assert reason == expected_reason
+
 
 @pytest.mark.skipif(
     not current_platform.is_cuda_alike(),
@@ -306,3 +479,293 @@ def test_triton_nvfp4_quant_dequant(
         f"min={ref_min:.3f}ms, max={ref_max:.3f}ms"
     )
     print(f"    speedup:   {speedup:.2f}x")
+
+
+MOE_MODEL_CONFIGS = {
+    "nvidia/Qwen3-30B-A3B-NVFP4": {
+        "shards": ["model-00001-of-00004.safetensors"],
+        "expert_prefix": "model.layers.9.mlp.experts.",
+        # Position of the expert index in the dot-split key.
+        "expert_idx_pos": 5,
+    },
+    "nvidia/Kimi-K2.6-NVFP4": {
+        "shards": [
+            "model-00001-of-00060.safetensors",
+            "model-00002-of-00060.safetensors",
+        ],
+        "expert_prefix": "language_model.model.layers.1.mlp.experts.",
+        "expert_idx_pos": 6,
+    },
+}
+
+
+def _load_nvfp4_moe_weights(
+    model_id: str,
+    tensor_parallel_size: int,
+    max_experts: int | None = None,
+):
+    """Load and stack NVFP4 MoE weights from checkpoint shards.
+
+    Returns (w1, w1_scale, w1_gscale, w2, w2_scale, w2_gscale,
+             a1_gscale, a2_gscale, num_experts, hidden_dim,
+             intermediate_size).
+
+    When max_experts is set, only the first max_experts experts are loaded.
+
+    When tensor_parallel_size > 1, the N dimension of w1 and the K
+    dimension of w2 are narrowed to the first TP shard (simulating
+    column-parallel on w1 / row-parallel on w2).
+    """
+    cfg = MOE_MODEL_CONFIGS[model_id]
+    shards = cast(list[str], cfg["shards"])
+    checkpoint_path = huggingface_hub.snapshot_download(
+        model_id,
+        allow_patterns=shards,
+    )
+    expert_prefix = cfg["expert_prefix"]
+    idx_pos = cast(int, cfg["expert_idx_pos"])
+
+    # Collect all tensors across shards into a flat dict — an expert's
+    # tensors may be split across multiple shard files.
+    all_tensors: dict[str, torch.Tensor] = {}
+    for shard_name in shards:
+        shard_path = f"{checkpoint_path}/{shard_name}"
+        with safe_open(shard_path, framework="pt", device="cpu") as f:
+            for key in f.keys():  # noqa: SIM118
+                if key.startswith(expert_prefix):
+                    all_tensors[key] = f.get_tensor(key)
+
+    expert_indices = sorted(
+        {
+            int(key.split(".")[idx_pos])
+            for key in all_tensors
+            if key.endswith(".gate_proj.weight")
+        }
+    )
+    if max_experts is not None:
+        expert_indices = expert_indices[:max_experts]
+    num_experts = len(expert_indices)
+
+    gate_weights, up_weights, down_weights = [], [], []
+    gate_scales, up_scales, down_scales = [], [], []
+    gate_gscales, up_gscales, down_gscales = [], [], []
+    a1_scales, a2_scales = [], []
+
+    for idx in expert_indices:
+        prefix = f"{expert_prefix}{idx}"
+        gate_weights.append(all_tensors[f"{prefix}.gate_proj.weight"])
+        gate_scales.append(all_tensors[f"{prefix}.gate_proj.weight_scale"])
+        gate_gscales.append(all_tensors[f"{prefix}.gate_proj.weight_scale_2"])
+        up_weights.append(all_tensors[f"{prefix}.up_proj.weight"])
+        up_scales.append(all_tensors[f"{prefix}.up_proj.weight_scale"])
+        up_gscales.append(all_tensors[f"{prefix}.up_proj.weight_scale_2"])
+        down_weights.append(all_tensors[f"{prefix}.down_proj.weight"])
+        down_scales.append(all_tensors[f"{prefix}.down_proj.weight_scale"])
+        down_gscales.append(all_tensors[f"{prefix}.down_proj.weight_scale_2"])
+        a1_scales.append(all_tensors[f"{prefix}.gate_proj.input_scale"])
+        a2_scales.append(all_tensors[f"{prefix}.down_proj.input_scale"])
+
+    # Stack into MoE format.
+    # w1 = [E, 2*intermediate, hidden//2]  (gate + up concatenated)
+    w1 = torch.stack(
+        [torch.cat([g, u], dim=0) for g, u in zip(gate_weights, up_weights)]
+    ).cuda()
+    w1_scale = torch.stack(
+        [torch.cat([g, u], dim=0) for g, u in zip(gate_scales, up_scales)]
+    ).cuda()
+    w1_gscale = torch.stack(gate_gscales).cuda()
+
+    # w2 = [E, hidden, intermediate//2]
+    w2 = torch.stack(down_weights).cuda()
+    w2_scale = torch.stack(down_scales).cuda()
+    w2_gscale = torch.stack(down_gscales).cuda()
+
+    a13_scale_raw = torch.stack(a1_scales).cuda()
+    a2_scale_raw = torch.stack(a2_scales).cuda()
+
+    # Apply EMULATION transforms (matches oracle/nvfp4.py).
+    nvfp4_emulation_utils.kE2M1ToFloat_handle.val = (
+        nvfp4_emulation_utils.kE2M1ToFloat_handle.val.cuda()
+    )
+    a1_gscale = 1.0 / a13_scale_raw.max().to(torch.float32)
+    a2_gscale = 1.0 / a2_scale_raw.max().to(torch.float32)
+
+    # ── Simulate TP sharding ──
+    # w1 (gate_up): column-parallel → shard the N dimension (dim 1).
+    # w2 (down):    row-parallel   → shard the K dimension (dim 2,
+    #               which is the packed K//2 dim).
+    # Scales follow the same sharding on the corresponding dimension.
+    tp = tensor_parallel_size
+    if tp > 1:
+        n1 = w1.size(1) // tp
+        w1 = w1[:, :n1, :].contiguous()
+        w1_scale = w1_scale[:, :n1, :].contiguous()
+
+        k2_packed = w2.size(2) // tp
+        k2_scale = w2_scale.size(2) // tp
+        w2 = w2[:, :, :k2_packed].contiguous()
+        w2_scale = w2_scale[:, :, :k2_scale].contiguous()
+
+    hidden_dim = w1.size(2) * 2
+    intermediate_size = w1.size(1) // 2
+
+    return (
+        w1,
+        w1_scale,
+        w1_gscale,
+        w2,
+        w2_scale,
+        w2_gscale,
+        a1_gscale,
+        a2_gscale,
+        num_experts,
+        hidden_dim,
+        intermediate_size,
+    )
+
+
+@pytest.mark.skipif(
+    not current_platform.is_cuda_alike(),
+    reason="Triton NVFP4 kernel requires CUDA.",
+)
+@pytest.mark.parametrize("num_tokens", [1, 2, 4, 1024])
+@pytest.mark.parametrize("top_k", [4])
+@pytest.mark.parametrize("model_id", list(MOE_MODEL_CONFIGS.keys()))
+@pytest.mark.parametrize(
+    "tensor_parallel_size",
+    [pytest.param(val, id=f"tensor_parallel_size:{val}") for val in [1, 2, 4, 8]],
+)
+def test_nvfp4_moe_correctness(
+    num_tokens: int,
+    top_k: int,
+    model_id: str,
+    tensor_parallel_size: int,
+) -> None:
+    """Compare Nvfp4QuantizationEmulationTritonExperts (fused weight dequant + compute)
+    against the unfused reference Nvfp4QuantizationEmulationTritonExpertsReference.
+
+    Both must produce bit-identical results.
+    """
+    num_test_experts = max(8, top_k)
+    (
+        w1,
+        w1_scale,
+        w1_gscale,
+        w2,
+        w2_scale,
+        w2_gscale,
+        a1_gscale,
+        a2_gscale,
+        num_experts,
+        hidden_dim,
+        intermediate_size,
+    ) = _load_nvfp4_moe_weights(
+        model_id,
+        tensor_parallel_size,
+        max_experts=num_test_experts,
+    )
+
+    moe_config = FusedMoEConfig(
+        num_experts=num_experts,
+        experts_per_token=top_k,
+        hidden_dim=hidden_dim,
+        intermediate_size=intermediate_size,
+        num_local_experts=num_experts,
+        num_logical_experts=num_experts,
+        moe_parallel_config=FusedMoEParallelConfig.make_no_parallel(),
+        activation=MoEActivation.SILU,
+        in_dtype=torch.bfloat16,
+        device="cuda",
+        routing_method=RoutingMethodType.TopK,
+        max_num_tokens=512,
+    )
+
+    def _make_quant_config():
+        return nvfp4_moe_quant_config(
+            g1_alphas=w1_gscale.clone(),
+            g2_alphas=w2_gscale.clone(),
+            a1_gscale=a1_gscale.clone(),
+            a2_gscale=a2_gscale.clone(),
+            w1_scale=w1_scale.clone(),
+            w2_scale=w2_scale.clone(),
+        )
+
+    ref_experts = Nvfp4QuantizationEmulationTritonExpertsReference(
+        moe_config=moe_config,
+        quant_config=_make_quant_config(),
+    )
+    fused_experts = Nvfp4QuantizationEmulationTritonExperts(
+        moe_config=moe_config,
+        quant_config=_make_quant_config(),
+    )
+
+    torch.manual_seed(42)
+    hidden_states = torch.randn(
+        num_tokens, hidden_dim, dtype=torch.bfloat16, device="cuda"
+    )
+
+    topk_weights = torch.randn(
+        num_tokens, top_k, dtype=torch.float32, device="cuda"
+    ).softmax(dim=-1)
+    topk_ids = torch.stack(
+        [torch.randperm(num_experts, device="cuda")[:top_k] for _ in range(num_tokens)]
+    ).to(torch.int32)
+
+    N = w1.size(1)  # 2 * intermediate
+    K = hidden_dim
+
+    ws13_size = num_tokens * top_k * max(intermediate_size, K)
+    ws2_size = num_tokens * top_k * max(N, K)
+
+    workspace13_ref = torch.zeros(ws13_size, dtype=torch.bfloat16, device="cuda")
+    workspace2_ref = torch.zeros(ws2_size, dtype=torch.bfloat16, device="cuda")
+    output_ref = torch.zeros(num_tokens, K, dtype=torch.bfloat16, device="cuda")
+
+    workspace13_fused = torch.zeros_like(workspace13_ref)
+    workspace2_fused = torch.zeros_like(workspace2_ref)
+    output_fused = torch.zeros_like(output_ref)
+
+    apply_kwargs = dict(
+        hidden_states=hidden_states,
+        w1=w1,
+        w2=w2,
+        topk_weights=topk_weights,
+        topk_ids=topk_ids,
+        activation=MoEActivation.SILU,
+        global_num_experts=num_experts,
+        expert_map=None,
+        a1q_scale=None,
+        a2_scale=None,
+        expert_tokens_meta=None,
+        apply_router_weight_on_input=False,
+    )
+
+    # Unfused reference.
+    ref_experts.apply(
+        output=output_ref,
+        workspace13=workspace13_ref,
+        workspace2=workspace2_ref,
+        **apply_kwargs,
+    )
+
+    # Fused implementation.
+    fused_experts.apply(
+        output=output_fused,
+        workspace13=workspace13_fused,
+        workspace2=workspace2_fused,
+        **apply_kwargs,
+    )
+
+    # Not strict equality on H100, MI325, MI300 (< 0.1% elements).
+    # The fused on-the-fly dequant path can lower to a slightly
+    # different Triton/MMA tiling than the pre-dequantized
+    # reference; experiments with reference-like tiling/masking
+    # reduced some diffs were not kept because they regress
+    # the fused kernel speed.
+    # Strict equality validated on MI355.
+    torch.testing.assert_close(
+        output_fused,
+        output_ref,
+        atol=0.0 if on_gfx950() else 0.02,
+        rtol=0,
+    )

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

@@ -11,6 +11,8 @@ Weights are dequantized on the fly during each forward, we fall back to calling
 is applied on `a13`, `a2`.
 """
 
+from typing import Any
+
 import torch
 
 import vllm.model_executor.layers.fused_moe.modular_kernel as mk
@@ -21,18 +23,316 @@ from vllm.model_executor.layers.fused_moe.config import (
     FusedMoEQuantConfig,
 )
 from vllm.model_executor.layers.fused_moe.experts.triton_moe import TritonExperts
+from vllm.model_executor.layers.fused_moe.fused_moe import (
+    try_get_optimal_moe_config,
+    write_zeros_to_output,
+)
+from vllm.model_executor.layers.fused_moe.moe_align_block_size import (
+    moe_align_block_size,
+)
+from vllm.model_executor.layers.fused_moe.utils import (
+    _resize_cache,
+    moe_kernel_quantize_input,
+)
 from vllm.model_executor.layers.quantization.utils.nvfp4_emulation_utils import (
-    dequantize_to_dtype,
+    _e2m1_inline,
 )
 from vllm.model_executor.layers.quantization.utils.quant_utils import (
     QuantKey,
     kNvfp4Dynamic,
     kNvfp4Static,
 )
+from vllm.triton_utils import tl, triton
 
 logger = init_logger(__name__)
 
 
+@triton.jit
+def fused_moe_nvfp4_emulation_kernel(
+    a_ptr,
+    b_ptr,
+    c_ptr,
+    b_scale_ptr,
+    w_global_scale_ptr,
+    topk_weights_ptr,
+    sorted_token_ids_ptr,
+    expert_ids_ptr,
+    num_tokens_post_padded_ptr,
+    N: tl.constexpr,
+    K: tl.constexpr,
+    EM,
+    num_valid_tokens,
+    # Strides — A [M, K]
+    stride_am,
+    stride_ak,
+    # Strides — B [E, N, K//2], passed as (expert, K-packed, N)
+    stride_be,
+    stride_bk,
+    stride_bn,
+    # Strides — C [M, topk, N]
+    stride_cm,
+    stride_cn,
+    # Strides — B_scale [E, N, K//BLOCK], passed as (expert, K-scale, N)
+    stride_bse,
+    stride_bsk,
+    stride_bsn,
+    block_k_diviable: tl.constexpr,
+    BLOCK_SIZE_M: tl.constexpr,
+    BLOCK_SIZE_N: tl.constexpr,
+    BLOCK_SIZE_K: tl.constexpr,
+    GROUP_SIZE_M: tl.constexpr,
+    MUL_ROUTED_WEIGHT: tl.constexpr,
+    top_k: tl.constexpr,
+    compute_type: tl.constexpr,
+    group_size: tl.constexpr,
+):
+    """
+    Fused MoE kernel for emulated NVFP4 weight-only dequantization + GEMM.
+
+    Activations A are BF16 (already QDQ'd externally).
+    Weights B are packed uint8 NVFP4 [E, N, K//2] — two FP4 values per byte
+    along the K dimension.
+    B_scale holds per-block FP8-E4M3 scales [E, N, K // group_size].
+    w_global_scale is a per-expert scalar global scale.
+
+    The dequantization formula per element is:
+        w_float = e2m1_decode(nibble) * (block_scale_fp8 * global_scale)
+
+    Weight loading optimization: each packed byte is loaded exactly once as
+    a [BLOCK_SIZE_N, BLOCK_SIZE_K // 2] tile (N-major), both nibbles are
+    extracted, decoded and scaled, then tl.interleave produces the
+    [BLOCK_SIZE_N, BLOCK_SIZE_K] dequantized tile which is transposed to
+    [BLOCK_SIZE_K, BLOCK_SIZE_N] for tl.dot.
+    """
+    BLOCK_SIZE_K_PACKED: tl.constexpr = BLOCK_SIZE_K // 2
+
+    # Map program ids to the block of C it should compute.
+    pid = tl.program_id(axis=0)
+    num_pid_m = tl.cdiv(EM, BLOCK_SIZE_M)
+    num_pid_n = tl.cdiv(N, BLOCK_SIZE_N)
+    num_pid_in_group = GROUP_SIZE_M * num_pid_n
+    group_id = pid // num_pid_in_group
+    first_pid_m = group_id * GROUP_SIZE_M
+    group_size_m = min(num_pid_m - first_pid_m, GROUP_SIZE_M)
+    pid_m = first_pid_m + ((pid % num_pid_in_group) % group_size_m)
+    pid_n = (pid % num_pid_in_group) // group_size_m
+
+    # Token / expert setup
+    num_tokens_post_padded = tl.load(num_tokens_post_padded_ptr)
+    if pid_m * BLOCK_SIZE_M >= num_tokens_post_padded:
+        return
+
+    offs_token_id = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M).to(tl.int64)
+    offs_token = tl.load(sorted_token_ids_ptr + offs_token_id).to(tl.int64)
+    token_mask = offs_token < num_valid_tokens
+
+    off_experts = tl.load(expert_ids_ptr + pid_m).to(tl.int64)
+    if off_experts == -1:
+        write_zeros_to_output(
+            c_ptr,
+            stride_cm,
+            stride_cn,
+            pid_n,
+            N,
+            offs_token,
+            token_mask,
+            BLOCK_SIZE_M,
+            BLOCK_SIZE_N,
+            compute_type,
+        )
+        return
+
+    # Pointer setup
+    offs_bn = (pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N).to(tl.int64)) % N
+    offs_k = tl.arange(0, BLOCK_SIZE_K)
+    offs_k_packed = tl.arange(0, BLOCK_SIZE_K_PACKED)
+
+    # A pointers: [BLOCK_SIZE_M, BLOCK_SIZE_K]
+    a_ptrs = a_ptr + (
+        offs_token[:, None] // top_k * stride_am + offs_k[None, :] * stride_ak
+    )
+
+    # B pointers: [BLOCK_SIZE_N, BLOCK_SIZE_K_PACKED] — N-major so that
+    # tl.interleave (which operates on the last dim) produces a
+    # [BLOCK_SIZE_N, BLOCK_SIZE_K] tile that we transpose for tl.dot.
+    # Each unique byte is loaded exactly once.
+    b_ptrs = (
+        b_ptr
+        + off_experts * stride_be
+        + offs_bn[:, None] * stride_bn
+        + offs_k_packed[None, :] * stride_bk
+    )
+
+    # B_scale pointers: [BLOCK_SIZE_N, BLOCK_SIZE_K_PACKED] — same
+    # N-major layout.  Each packed byte index covers 2 K elements that
+    # always fall within the same group (group_size=16, so each group
+    # spans 8 packed bytes).  We can therefore index the scale using
+    # offs_k_packed directly.
+    # Note: group_size_packed = group_size // 2 maps packed indices to
+    # scale indices the same way unpacked indices map via group_size.
+    group_size_packed: tl.constexpr = group_size // 2
+
+    # Load per-expert global scale (scalar).
+    w_global_scale = tl.load(w_global_scale_ptr + off_experts).to(tl.float32)
+
+    # K-loop with FP32 accumulation
+    accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32)
+
+    for k in range(0, tl.cdiv(K, BLOCK_SIZE_K)):
+        # Load A tile [BLOCK_SIZE_M, BLOCK_SIZE_K].
+        if block_k_diviable:
+            a = tl.load(
+                a_ptrs,
+                mask=token_mask[:, None],
+                other=0.0,
+            )
+        else:
+            a = tl.load(
+                a_ptrs,
+                mask=token_mask[:, None] & (offs_k[None, :] < K - k * BLOCK_SIZE_K),
+                other=0.0,
+            )
+
+        # Load packed weight tile [BLOCK_SIZE_N, BLOCK_SIZE_K_PACKED].
+        if block_k_diviable:
+            raw_bytes = tl.load(b_ptrs)
+        else:
+            kp_mask = offs_k_packed[None, :] < (K // 2) - k * BLOCK_SIZE_K_PACKED
+            raw_bytes = tl.load(b_ptrs, mask=kp_mask, other=0)
+
+        # Extract both nibbles from each byte (each [N, K_packed]).
+        low_nibble = raw_bytes & 0x0F
+        high_nibble = (raw_bytes >> 4) & 0x0F
+
+        low_decoded = _e2m1_inline(low_nibble)
+        high_decoded = _e2m1_inline(high_nibble)
+
+        # Load and apply per-block FP8 scales.
+        # Scale shape: [BLOCK_SIZE_N, BLOCK_SIZE_K_PACKED], one scale per
+        # group_size_packed packed elements.
+        b_scale_ptrs = (
+            b_scale_ptr
+            + off_experts * stride_bse
+            + offs_bn[:, None] * stride_bsn
+            + ((offs_k_packed[None, :] + BLOCK_SIZE_K_PACKED * k) // group_size_packed)
+            * stride_bsk
+        )
+        if block_k_diviable:
+            b_scale_raw = tl.load(b_scale_ptrs)
+        else:
+            b_scale_raw = tl.load(b_scale_ptrs, mask=kp_mask, other=0.0)
+
+        b_scale = tl.cast(b_scale_raw, tl.float8e4nv, bitcast=True).to(tl.float32)
+        b_scale = b_scale * w_global_scale
+
+        # Scale both halves with the same per-block scale (the two
+        # elements packed in one byte always belong to the same group).
+        low_scaled = low_decoded * b_scale
+        high_scaled = high_decoded * b_scale
+
+        # Interleave along last dim: [N, K_packed] x2 -> [N, K],
+        # then transpose to [K, N] for tl.dot.
+        b = tl.trans(tl.interleave(low_scaled, high_scaled)).to(compute_type)
+
+        accumulator = tl.dot(a, b, acc=accumulator)
+
+        # Advance pointers along K.
+        a_ptrs += BLOCK_SIZE_K * stride_ak
+        b_ptrs += BLOCK_SIZE_K_PACKED * stride_bk
+
+    # Router weight multiplication (in float32 for stability)
+    if MUL_ROUTED_WEIGHT:
+        moe_weight = tl.load(topk_weights_ptr + offs_token, mask=token_mask, other=0)
+        accumulator = accumulator * moe_weight[:, None]
+
+    accumulator = accumulator.to(compute_type)
+
+    # Write output
+    offs_cn = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)
+    c_ptrs = c_ptr + stride_cm * offs_token[:, None] + stride_cn * offs_cn[None, :]
+    c_mask = token_mask[:, None] & (offs_cn[None, :] < N)
+    tl.store(c_ptrs, accumulator, mask=c_mask)
+
+
+def invoke_fused_moe_nvfp4_emulation_kernel(
+    A: torch.Tensor,
+    B: torch.Tensor,
+    C: torch.Tensor,
+    B_scale: torch.Tensor,
+    act_global_scale: torch.Tensor,
+    w_global_scale: torch.Tensor,
+    topk_weights: torch.Tensor | None,
+    sorted_token_ids: torch.Tensor,
+    expert_ids: torch.Tensor,
+    num_tokens_post_padded: torch.Tensor,
+    mul_routed_weight: bool,
+    top_k: int,
+    config: dict[str, Any],
+    compute_type: tl.dtype,
+):
+    """Launch the fused NVFP4 emulation MoE kernel.
+
+    B has shape [E, N, K_packed] where K_packed = K // 2 (two FP4 per byte).
+    B_scale has shape [E, N, K // group_size] in FP8-E4M3 (stored as uint8).
+    w_global_scale has shape [E] (per-expert scalar).
+    """
+    assert B_scale is not None and B_scale.ndim == 3
+
+    N = B.size(1)
+    K = A.size(1)
+
+    M = A.size(0)
+    num_tokens = M * top_k
+
+    EM = sorted_token_ids.size(0)
+    if A.size(0) < config["BLOCK_SIZE_M"]:
+        EM = min(
+            sorted_token_ids.size(0),
+            A.size(0) * top_k * config["BLOCK_SIZE_M"],
+        )
+
+    grid = lambda META: (
+        triton.cdiv(EM, META["BLOCK_SIZE_M"]) * triton.cdiv(N, META["BLOCK_SIZE_N"]),
+    )
+
+    fused_moe_nvfp4_emulation_kernel[grid](
+        A,
+        B,
+        C,
+        B_scale,
+        w_global_scale,
+        topk_weights,
+        sorted_token_ids,
+        expert_ids,
+        num_tokens_post_padded,
+        N,
+        K,
+        EM,
+        num_tokens,
+        A.stride(0),
+        A.stride(1),
+        # B is [E, N, K//2]: swap N and K strides so kernel indexes [K, N].
+        B.stride(0),
+        B.stride(2),
+        B.stride(1),
+        C.stride(1),
+        C.stride(2),
+        # B_scale is [E, N, K//group]: swap N and K strides likewise.
+        B_scale.stride(0),
+        B_scale.stride(2),
+        B_scale.stride(1),
+        block_k_diviable=K % config["BLOCK_SIZE_K"] == 0,
+        MUL_ROUTED_WEIGHT=mul_routed_weight,
+        top_k=top_k,
+        compute_type=compute_type,
+        group_size=16,
+        BLOCK_SIZE_M=config["BLOCK_SIZE_M"],
+        BLOCK_SIZE_N=config["BLOCK_SIZE_N"],
+        BLOCK_SIZE_K=config["BLOCK_SIZE_K"],
+        GROUP_SIZE_M=config["GROUP_SIZE_M"],
+    )
+
+
 class Nvfp4QuantizationEmulationTritonExperts(TritonExperts):
     """
     Extension of TritonExperts to support emulated NVFP4 MoE experts.
@@ -72,6 +372,31 @@ class Nvfp4QuantizationEmulationTritonExperts(TritonExperts):
     def expects_unquantized_inputs(self) -> bool:
         return True
 
+    @staticmethod
+    def supports_lora() -> bool:
+        return False
+
+    @staticmethod
+    def is_supported_config(
+        cls: type[mk.FusedMoEExperts],
+        moe_config: FusedMoEConfig,
+        weight_key: QuantKey | None,
+        activation_key: QuantKey | None,
+        activation_format: mk.FusedMoEActivationFormat,
+    ) -> tuple[bool, str | None]:
+        if moe_config.is_lora_enabled:
+            return False, "kernel does not support LoRA"
+        if moe_config.has_bias:
+            return False, "kernel does not support bias"
+
+        return TritonExperts.is_supported_config(
+            cls,
+            moe_config,
+            weight_key,
+            activation_key,
+            activation_format,
+        )
+
     @staticmethod
     def _supports_quant_scheme(
         weight_key: QuantKey | None,
@@ -109,47 +434,109 @@ class Nvfp4QuantizationEmulationTritonExperts(TritonExperts):
         # w2 shape: [num_experts, hidden_size, intermediate_size//2]
         assert w1.dtype == torch.uint8
         assert w2.dtype == torch.uint8
+        assert hidden_states.is_contiguous()
+        assert hidden_states.dim() == 2
 
-        # Dequantize w1 from packed NVFP4 to fp16/bf16
-        w13_global_scale = self.quant_config.g1_alphas
+        K = hidden_states.size(-1)
+        assert w1.size(2) * 2 == K, f"Hidden size mismatch: {K} != {w1.size(2) * 2}"
 
-        w1_dequant = dequantize_to_dtype(
-            tensor_fp4=w1,
-            tensor_sf=self.w1_scale_val,
-            global_scale=w13_global_scale,
-            dtype=hidden_states.dtype,
-            block_size=16,
-            swizzle=False,
+        E, num_tokens, N, _, top_k_num = self.moe_problem_size(
+            hidden_states, w1, w2, topk_ids
         )
 
-        # Dequantize w2 from packed NVFP4 to fp16/bf16
-        w2_global_scale = self.quant_config.g2_alphas
+        if global_num_experts == -1:
+            global_num_experts = E
 
-        w2_dequant = dequantize_to_dtype(
-            tensor_fp4=w2,
-            tensor_sf=self.w2_scale_val,
-            global_scale=w2_global_scale,
-            dtype=hidden_states.dtype,
-            block_size=16,
-            swizzle=False,
+        # TODO: There is actually no support for tuning of the underlying triton
+        # hyperparameters in benchmarks/kernels/benchmark_moe.py, to be added.
+        config = try_get_optimal_moe_config(
+            w1.size(),
+            w2.size(),
+            top_k_num,
+            self.quant_config.config_name(hidden_states.dtype),
+            num_tokens,
+            block_shape=None,
         )
 
-        # Activation quantization/dequantization is deferred to
-        # `moe_kernel_quantize_input` in TritonExperts.apply.
-        super().apply(
-            output=output,
-            hidden_states=hidden_states,
-            w1=w1_dequant,
-            w2=w2_dequant,
-            topk_weights=topk_weights,
-            topk_ids=topk_ids,
-            activation=activation,
-            global_num_experts=global_num_experts,
-            expert_map=expert_map,
-            a1q_scale=None,
-            a2_scale=self.quant_config.a2_gscale,
-            workspace13=workspace13,
-            workspace2=workspace2,
-            expert_tokens_meta=expert_tokens_meta,
-            apply_router_weight_on_input=apply_router_weight_on_input,
+        if hidden_states.dtype == torch.bfloat16:
+            compute_type = tl.bfloat16
+        elif hidden_states.dtype == torch.float16:
+            compute_type = tl.float16
+        elif hidden_states.dtype == torch.float32:
+            compute_type = tl.float32
+        else:
+            raise ValueError(f"Unsupported compute_type: {hidden_states.dtype}")
+
+        intermediate_cache1 = _resize_cache(workspace2, (num_tokens, top_k_num, N))
+        activation_out_dim = self.adjust_N_for_activation(N, activation)
+        intermediate_cache2 = _resize_cache(
+            workspace13, (num_tokens * top_k_num, activation_out_dim)
         )
+        intermediate_cache3 = _resize_cache(workspace2, (num_tokens, top_k_num, K))
+
+        sorted_token_ids, expert_ids, num_tokens_post_padded = moe_align_block_size(
+            topk_ids,
+            config["BLOCK_SIZE_M"],
+            global_num_experts,
+            expert_map,
+        )
+
+        # Activation NVFP4 QDQ.
+        hidden_states_qdq, _ = moe_kernel_quantize_input(
+            A=hidden_states,
+            A_scale=self.quant_config.a1_gscale,
+            quant_dtype="nvfp4",
+            per_act_token_quant=False,
+            quantization_emulation=True,
+        )
+
+        # w13: fused weight dequant + GEMM.
+        invoke_fused_moe_nvfp4_emulation_kernel(
+            hidden_states_qdq,
+            w1,
+            intermediate_cache1,
+            self.w1_scale_val,
+            self.quant_config.a1_gscale,
+            self.quant_config.g1_alphas,
+            None,  # topk_weights — applied after w2
+            sorted_token_ids,
+            expert_ids,
+            num_tokens_post_padded,
+            False,  # mul_routed_weight
+            top_k_num,
+            config,
+            compute_type=compute_type,
+        )
+
+        self.activation(
+            activation, intermediate_cache2, intermediate_cache1.view(-1, N)
+        )
+
+        # Activation NVFP4 QDQ.
+        intermediate_cache2_qdq, _ = moe_kernel_quantize_input(
+            A=intermediate_cache2,
+            A_scale=self.quant_config.a2_gscale,
+            quant_dtype="nvfp4",
+            per_act_token_quant=False,
+            quantization_emulation=True,
+        )
+
+        # w2: fused weight dequant + GEMM.
+        invoke_fused_moe_nvfp4_emulation_kernel(
+            intermediate_cache2_qdq,
+            w2,
+            intermediate_cache3,
+            self.w2_scale_val,
+            self.quant_config.a2_gscale,
+            self.quant_config.g2_alphas,
+            topk_weights,
+            sorted_token_ids,
+            expert_ids,
+            num_tokens_post_padded,
+            not apply_router_weight_on_input,
+            1,
+            config,
+            compute_type=compute_type,
+        )
+
+        self.moe_sum(intermediate_cache3, output)

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

@@ -23,28 +23,24 @@ kE2M1ToFloat_handle = SimpleNamespace(
 
 
 @triton.jit
-def _e2m1_inline(magnitude):
-    """Inline E2M1 lookup using binary tree - 3 levels instead of 7 sequential.
+def _e2m1_inline(nibble):
+    """Decode an NVFP4 nibble (4 bits: 1 sign + 3 magnitude) to float32.
 
-    Maps 3-bit magnitude to float: [0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0]
-    Uses bit decomposition for fewer comparisons.
+    Uses direct IEEE 754 bit construction.
+    For magnitudes 2-7 the FP32 bit pattern is 0x3F000000 + (mag << 22),
+    which is a single shift + add + bitcast.  Magnitudes 0 (zero) and 1
+    (E2M1 subnormal = 0.5) are patched with two tl.where ops.
     """
-    # Bit 2 (MSB): separates 0-3 from 4-7
-    # Bit 1: separates within groups
-    # Bit 0 (LSB): separates within pairs
-    b2 = (magnitude >> 2) & 1  # 0 for mag 0-3, 1 for mag 4-7
-    b1 = (magnitude >> 1) & 1  # middle bit
-    b0 = magnitude & 1  # LSB
+    magnitude = nibble & 0x07
+    sign = (nibble >> 3) & 1
 
-    # For mag 0-3: [0.0, 0.5, 1.0, 1.5]
-    low_group = tl.where(
-        b1 == 1, tl.where(b0 == 1, 1.5, 1.0), tl.where(b0 == 1, 0.5, 0.0)
-    )
-    # For mag 4-7: [2.0, 3.0, 4.0, 6.0]
-    high_group = tl.where(
-        b1 == 1, tl.where(b0 == 1, 6.0, 4.0), tl.where(b0 == 1, 3.0, 2.0)
-    )
-    return tl.where(b2 == 1, high_group, low_group)
+    fp32_bits = 0x3F000000 + (magnitude.to(tl.int32) << 22)
+    val = fp32_bits.to(tl.float32, bitcast=True)
+
+    val = tl.where(magnitude == 0, 0.0, val)
+    val = tl.where(magnitude == 1, 0.5, val)
+
+    return tl.where(sign == 1, -val, val)
 
 
 @triton.jit
@@ -65,7 +61,7 @@ def _dequantize_nvfp4_kernel(
     """
     BLOCK_PACKED: tl.constexpr = BLOCK_SIZE // 2
 
-    row_idx = tl.program_id(0)
+    row_idx = tl.program_id(0).to(tl.int64)
     tile_idx = tl.program_id(1)
 
     if has_batch_global_scale:
@@ -105,16 +101,8 @@ def _dequantize_nvfp4_kernel(
     low_nibble = raw_bytes & 0x0F
     high_nibble = (raw_bytes >> 4) & 0x0F
 
-    # Binary tree E2M1 decode
-    low_mag = low_nibble & 0x07
-    low_val = _e2m1_inline(low_mag)
-    low_sign = (low_nibble >> 3) & 1
-    low_result = tl.where(low_sign == 1, -low_val, low_val) * scale_values
-
-    high_mag = high_nibble & 0x07
-    high_val = _e2m1_inline(high_mag)
-    high_sign = (high_nibble >> 3) & 1
-    high_result = tl.where(high_sign == 1, -high_val, high_val) * scale_values
+    low_result = _e2m1_inline(low_nibble) * scale_values
+    high_result = _e2m1_inline(high_nibble) * scale_values
 
     # Interleave for coalesced contiguous store
     result = tl.interleave(low_result, high_result)