diff --git a/tests/kernels/quantization/test_nvfp4_emulation.py b/tests/kernels/quantization/test_nvfp4_emulation.py index 71072d9e9ff..f5652af6e92 100644 --- a/tests/kernels/quantization/test_nvfp4_emulation.py +++ b/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, + ) diff --git a/vllm/model_executor/layers/fused_moe/experts/nvfp4_emulation_moe.py b/vllm/model_executor/layers/fused_moe/experts/nvfp4_emulation_moe.py index d7ed53612e0..f93c67a97dd 100644 --- a/vllm/model_executor/layers/fused_moe/experts/nvfp4_emulation_moe.py +++ b/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) diff --git a/vllm/model_executor/layers/quantization/utils/nvfp4_emulation_utils.py b/vllm/model_executor/layers/quantization/utils/nvfp4_emulation_utils.py index 39c78a9062b..ad6b272371e 100644 --- a/vllm/model_executor/layers/quantization/utils/nvfp4_emulation_utils.py +++ b/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)