From 78a78f5e0025b77ae2b6e6863cb2513dac8d4919 Mon Sep 17 00:00:00 2001 From: Peuqui Date: Sat, 12 Sep 2026 11:50:22 +0200 Subject: [PATCH] [Feature][SM75] Run ModelOpt NVFP4 and FP8 linears on Turing through the SM70 QPN kernels The compiled SM70 ops carry the QPN2 (NVFP4) and QPN8 (FP8) decode kernels and the QPN8 dense prefill, none of which depend on the TurboMind GEMM registry, but every gate admitted exact SM70 only and the registry itself is Arch<700, 750>. A Turing worker therefore took Marlin for NVFP4 and, for the mixed ModelOpt checkpoints, did not load at all (minimum capability 89; with the gate lowered the Marlin FP8 path fails on orig_dtype). Turing now prepares NVFP4 linears as the QPN2 prepack only and runs the QPN2 kernels for M <= 32 and a Triton dequantization into a transient fp16 buffer plus cuBLAS above that, split inside one registered op so the compiled decode graph keeps the kernel path; FP8 per-tensor linears become QPN8 codes with channel scales and run the existing QPN8 dispatch. Gates read the worker's own device. Volta keeps its TurboMind path unchanged. Co-authored-by: Claude Fable 5.1 Signed-off-by: Peuqui --- .../quantization/test_nvfp4_qpn2_dequant.py | 138 ++++++++ .../test_sm70_modelopt_mixed_nvfp4.py | 10 +- .../test_sm70_turbomind_turing_gates.py | 95 ++++++ .../layers/quantization/modelopt.py | 49 ++- .../layers/quantization/sm70_turbomind.py | 308 +++++++++++++++++- .../quantization/utils/nvfp4_qpn2_dequant.py | 288 ++++++++++++++++ 6 files changed, 875 insertions(+), 13 deletions(-) create mode 100644 tests/kernels/quantization/test_nvfp4_qpn2_dequant.py create mode 100644 tests/quantization/test_sm70_turbomind_turing_gates.py create mode 100644 vllm/model_executor/layers/quantization/utils/nvfp4_qpn2_dequant.py diff --git a/tests/kernels/quantization/test_nvfp4_qpn2_dequant.py b/tests/kernels/quantization/test_nvfp4_qpn2_dequant.py new file mode 100644 index 0000000000..86e6a7e225 --- /dev/null +++ b/tests/kernels/quantization/test_nvfp4_qpn2_dequant.py @@ -0,0 +1,138 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""The QPN2 dequantization kernel reproduces the checkpoint weight exactly. + +Round trip: random NVFP4 codes and e4m3 block scales -> QPN2 prepack +(``nvfp4_qpn2_prepare_sm70``) -> Triton dequantization -> compare with the +direct dequantization of the checkpoint tensors and with the pure-torch +inverse of the prepack. NVFP4 values times an e4m3 scale times the global +scale are exact in fp32 and round once to fp16, so the comparison is exact. +""" + +import pytest +import torch + +pytestmark = pytest.mark.skipif( + not torch.cuda.is_available(), reason="needs a CUDA device" +) + +E2M1 = torch.tensor([0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0]) + + +def _checkpoint_dequant( + packed: torch.Tensor, scales: torch.Tensor, global_scale: float +) -> torch.Tensor: + n, k_half = packed.shape + low = packed & 0xF + high = packed >> 4 + nib = torch.stack([low, high], dim=-1).view(n, k_half * 2) + values = E2M1.to(packed.device)[(nib & 7).long()] * torch.where( + nib & 8 > 0, -1.0, 1.0 + ) + scale = scales.view(torch.float8_e4m3fn).to(torch.float32) + values = values.view(n, -1, 16) * scale.unsqueeze(-1) * global_scale + return values.view(n, k_half * 2).to(torch.float16) + + +@pytest.mark.parametrize( + ("n", "k"), + [(32, 64), (64, 256), (3584, 5120), (8704, 5120), (1536, 5120)], +) +def test_qpn2_dequant_matches_checkpoint(n: int, k: int): + from vllm import _sm70_ops as sm70_ops + from vllm.model_executor.layers.quantization.utils.nvfp4_qpn2_dequant import ( + nvfp4_qpn2_dequant, + nvfp4_qpn2_dequant_reference, + ) + + if not hasattr(torch.ops._C, "nvfp4_qpn2_prepare_sm70"): + pytest.skip("build without the SM70 QPN2 extension") + generator = torch.Generator(device="cuda").manual_seed(n * 31 + k) + packed = torch.randint( + 0, 256, (n, k // 2), dtype=torch.uint8, device="cuda", generator=generator + ) + # e4m3 scales without the NaN code and without the sign bit, like + # ModelOpt block scales. + scales = torch.randint( + 0, 0x7F, (n, k // 16), dtype=torch.uint8, device="cuda", generator=generator + ) + global_scale = 0.0123 + expected = _checkpoint_dequant(packed, scales, global_scale) + + codes, qpn2_scales = sm70_ops.nvfp4_qpn2_prepare_sm70( + packed, scales.view(torch.float8_e4m3fn) + ) + reference = nvfp4_qpn2_dequant_reference(codes, qpn2_scales, global_scale, n, k) + kernel = nvfp4_qpn2_dequant(codes, qpn2_scales, global_scale, n, k) + + assert torch.equal(reference, expected) + assert torch.equal(kernel, expected) + + +def test_qpn2_dense_linear_matches_fp16_matmul(): + from vllm import _sm70_ops as sm70_ops + from vllm.model_executor.layers.quantization.utils.nvfp4_qpn2_dequant import ( + nvfp4_qpn2_dense_linear, + ) + + if not hasattr(torch.ops._C, "nvfp4_qpn2_prepare_sm70"): + pytest.skip("build without the SM70 QPN2 extension") + n, k, m = 3584, 5120, 64 + generator = torch.Generator(device="cuda").manual_seed(7) + packed = torch.randint( + 0, 256, (n, k // 2), dtype=torch.uint8, device="cuda", generator=generator + ) + scales = torch.randint( + 0, 0x7F, (n, k // 16), dtype=torch.uint8, device="cuda", generator=generator + ) + global_scale = 0.0123 + x = torch.randn(m, k, dtype=torch.float16, device="cuda", generator=generator) + codes, qpn2_scales = sm70_ops.nvfp4_qpn2_prepare_sm70( + packed, scales.view(torch.float8_e4m3fn) + ) + + expected = torch.nn.functional.linear( + x, _checkpoint_dequant(packed, scales, global_scale) + ) + out = nvfp4_qpn2_dense_linear(x, codes, qpn2_scales, global_scale, n, k) + assert out.shape == (m, n) + assert torch.equal(out, expected) + + +@pytest.mark.parametrize("m", [1, 4, 32, 33, 64]) +def test_qpn2_dispatch_linear_both_sides_of_the_threshold(m: int): + from vllm import _sm70_ops as sm70_ops + from vllm.model_executor.layers.quantization import sm70_turbomind as sm70_tm + from vllm.model_executor.layers.quantization.utils.nvfp4_qpn2_dequant import ( + nvfp4_qpn2_dispatch_linear, + ) + + if not hasattr(torch.ops._C, "nvfp4_qpn2_prepare_sm70"): + pytest.skip("build without the SM70 QPN2 extension") + n, k = 3584, 5120 + generator = torch.Generator(device="cuda").manual_seed(11) + packed = torch.randint( + 0, 256, (n, k // 2), dtype=torch.uint8, device="cuda", generator=generator + ) + # Block scales between 0.25 and 1.5 (e4m3 0x28..0x3c): a K=5120 dot + # product of such rows stays inside fp16, which the GEMM output is. + scales = torch.randint( + 0x28, 0x3D, (n, k // 16), dtype=torch.uint8, device="cuda", generator=generator + ) + global_scale = 0.01 + x = torch.randn(m, k, dtype=torch.float16, device="cuda", generator=generator) + codes, qpn2_scales = sm70_ops.nvfp4_qpn2_prepare_sm70( + packed, scales.view(torch.float8_e4m3fn) + ) + split_k, chains = sm70_tm.qpn2_launch_config(k, n) + + expected = torch.nn.functional.linear( + x, _checkpoint_dequant(packed, scales, global_scale) + ) + out = nvfp4_qpn2_dispatch_linear( + x, codes, qpn2_scales, global_scale, n, k, split_k, chains + ) + assert out.shape == (m, n) + # The QPN2 kernel accumulates in a different order than cuBLAS; both are + # fp32 accumulations of exact products, so they agree to fp16 rounding. + torch.testing.assert_close(out, expected, rtol=2e-3, atol=2e-2) diff --git a/tests/quantization/test_sm70_modelopt_mixed_nvfp4.py b/tests/quantization/test_sm70_modelopt_mixed_nvfp4.py index 19f351d570..6d661a7c59 100644 --- a/tests/quantization/test_sm70_modelopt_mixed_nvfp4.py +++ b/tests/quantization/test_sm70_modelopt_mixed_nvfp4.py @@ -109,20 +109,22 @@ def _qwen4_moe_contract(**overrides): return SimpleNamespace(**values) -def test_mixed_min_capability_requires_exact_sm70_and_both_turbomind_routes(): +def test_mixed_min_capability_requires_pre_ampere_and_both_turbomind_routes(): + # Volta takes the TurboMind routes, Turing the QPN routes; both admit the + # mixed checkpoint when the FP8 and NVFP4 switches are on. with ( - patch.object(sm70_tm, "is_exact_sm70_cuda_platform", return_value=True), + patch.object(sm70_tm, "is_pre_ampere_cuda_platform", return_value=True), patch.object(sm70_tm, "use_turbomind", side_effect=[True, True]), ): assert ModelOptMixedPrecisionConfig.get_min_capability() == 70 with ( - patch.object(sm70_tm, "is_exact_sm70_cuda_platform", return_value=True), + patch.object(sm70_tm, "is_pre_ampere_cuda_platform", return_value=True), patch.object(sm70_tm, "use_turbomind", side_effect=[True, False]), ): assert ModelOptMixedPrecisionConfig.get_min_capability() == 89 - with patch.object(sm70_tm, "is_exact_sm70_cuda_platform", return_value=False): + with patch.object(sm70_tm, "is_pre_ampere_cuda_platform", return_value=False): assert ModelOptMixedPrecisionConfig.get_min_capability() == 89 diff --git a/tests/quantization/test_sm70_turbomind_turing_gates.py b/tests/quantization/test_sm70_turbomind_turing_gates.py new file mode 100644 index 0000000000..d91ae71b3e --- /dev/null +++ b/tests/quantization/test_sm70_turbomind_turing_gates.py @@ -0,0 +1,95 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Turing takes the QPN2 route of the SM70 linear path; Volta is unchanged. + +CPU-only: device capabilities are mocked, no CUDA context is created. +""" + +from unittest.mock import patch + +import pytest +import torch + +from vllm.model_executor.layers.quantization import sm70_turbomind as sm70_tm + +MODULE = "vllm.model_executor.layers.quantization.sm70_turbomind" + + +@pytest.fixture +def cuda_tensor(): + tensor = torch.empty(0) + with patch.object( + torch.Tensor, "is_cuda", new_callable=lambda: property(lambda _: True) + ): + yield tensor + + +@pytest.mark.parametrize( + ("capability", "turbomind", "turing_qpn2"), + [ + ((7, 0), True, False), + ((7, 5), False, True), + ((8, 0), False, False), + ((8, 9), False, False), + ], +) +def test_prepare_routes_follow_the_tensor_device( + cuda_tensor, capability, turbomind, turing_qpn2 +): + with ( + patch(f"{MODULE}.torch.cuda.get_device_capability", return_value=capability), + patch(f"{MODULE}.use_turbomind", return_value=True), + ): + assert sm70_tm.should_prepare_turbomind(cuda_tensor, True) is turbomind + assert sm70_tm.should_prepare_turing_qpn2(cuda_tensor, True) is turing_qpn2 + + +def test_turing_route_honours_the_backend_switch(cuda_tensor): + with ( + patch(f"{MODULE}.torch.cuda.get_device_capability", return_value=(7, 5)), + patch(f"{MODULE}.use_turbomind", return_value=False), + ): + assert sm70_tm.should_prepare_turing_qpn2(cuda_tensor, True) is False + + +def test_turing_route_needs_a_cuda_tensor(): + with patch(f"{MODULE}.torch.cuda.get_device_capability", return_value=(7, 5)): + assert sm70_tm.should_prepare_turing_qpn2(torch.empty(0), True) is False + + +@pytest.mark.parametrize( + ("k", "n", "expected"), + [ + (5120, 8704, (8, 2)), # table entry + (5120, 3584, (16, 2)), # table entry + (4096, 4096, (8, 1)), # heuristic: 128 tiles x 8 = 1024 warps in flight + (2048, 1024, (32, 2)), # heuristic: needs the largest split to fill the GPU + (256, 64, (16, 2)), # heuristic fallback: 16 groups, split 16 + ], +) +def test_qpn2_launch_config(k, n, expected): + assert sm70_tm.qpn2_launch_config(k, n) == expected + + +def test_qpn2_launch_config_rejects_unsplittable_k(): + with pytest.raises(RuntimeError): + sm70_tm.qpn2_launch_config(16 * 3, 64) + + +def test_pad_qpn2_output_rows_pads_to_32(): + weight = torch.arange(40 * 8, dtype=torch.uint8).view(40, 8) + scales = torch.arange(40 * 2, dtype=torch.uint8).view(40, 2) + padded_weight, padded_scales, physical_n = sm70_tm.pad_qpn2_output_rows( + weight, scales + ) + assert physical_n == 64 + assert padded_weight.shape == (64, 8) and padded_scales.shape == (64, 2) + assert torch.equal(padded_weight[:40], weight) + assert torch.equal(padded_scales[:40], scales) + assert int(padded_weight[40:].abs().sum()) == 0 + + same_weight, same_scales, physical_n = sm70_tm.pad_qpn2_output_rows( + weight[:32], scales[:32] + ) + assert physical_n == 32 + assert same_weight is weight[:32] or torch.equal(same_weight, weight[:32]) diff --git a/vllm/model_executor/layers/quantization/modelopt.py b/vllm/model_executor/layers/quantization/modelopt.py index 97c59d988e..06cfc8f38f 100644 --- a/vllm/model_executor/layers/quantization/modelopt.py +++ b/vllm/model_executor/layers/quantization/modelopt.py @@ -461,6 +461,12 @@ def __init__(self, quant_config: ModelOptFp8Config) -> None: sm70_tm.is_exact_sm70_cuda_platform() and sm70_tm.use_turbomind(envs.VLLM_SM70_FP8_TURBOMIND) ) + # Turing takes the QPN8 kernels with the QPN8 dense prefill; the + # TurboMind GEMMs are registered for exact SM70 only. + self.use_sm75_fp8_qpn8 = ( + sm70_tm.is_turing_cuda_platform() + and sm70_tm.use_turbomind(envs.VLLM_SM70_FP8_TURBOMIND) + ) def create_weights( self, @@ -563,6 +569,26 @@ def process_weights_after_loading(self, layer: torch.nn.Module) -> None: logger.info_once("SM70 ModelOpt FP8 TurboMind W8A16 dense path enabled.") return + if self.use_sm75_fp8_qpn8: + if self.input_dtype != torch.float16: + raise RuntimeError( + "ModelOpt FP8 QPN8 on Turing requires FP16 activations, " + f"got {self.input_dtype}." + ) + sm70_tm.prepare_fp8_qpn8_dense_linear(layer, weight, max_w_scale) + replace_parameter( + layer, + "weight", + torch.empty(0, dtype=weight.dtype, device=weight.device), + ) + layer.weight_scale = Parameter(max_w_scale, requires_grad=False) + layer.input_scale = None + logger.info_once( + "SM75 ModelOpt FP8 QPN8 W8A16 dense path enabled " + "(QPN8 decode kernels, dense fp16 prefill)." + ) + return + layer.weight = Parameter(weight.t(), requires_grad=False) layer.weight_scale = Parameter(max_w_scale, requires_grad=False) layer.input_scale = Parameter(layer.input_scale.max(), requires_grad=False) @@ -599,6 +625,8 @@ def apply( if bias is not None: out.add_(bias) return out.reshape(*x.shape[:-1], layer.output_size_per_partition) + if sm70_tm.has_prepared_fp8_qpn8_linear(layer): + return sm70_tm.apply_prepared_fp8_qpn8_linear(layer, x, bias) return self.fp8_linear.apply_weights(layer, x, bias) @@ -1099,15 +1127,22 @@ def _try_prepare_sm70_modelopt_nvfp4(layer: torch.nn.Module) -> bool: ``amax / (6 * 448)`` (the Marlin multiplier). TurboMind combine is ``block * global``. Do not infer convention from scale magnitude. """ - if not sm70_tm.should_prepare_turbomind( + if sm70_tm.should_prepare_turbomind(layer.weight, envs.VLLM_SM70_NVFP4_TURBOMIND): + logger.info_once( + "SM70 ModelOpt NVFP4 TurboMind dense path enabled " + "(weight-only; activations remain half)." + ) + sm70_tm.prepare_nvfp4_linear(layer) + elif sm70_tm.should_prepare_turing_qpn2( layer.weight, envs.VLLM_SM70_NVFP4_TURBOMIND ): + logger.info_once( + "SM75 ModelOpt NVFP4 QPN2 path enabled (QPN2 decode kernels, dense " + "fp16 prefill; weight-only, activations remain half)." + ) + sm70_tm.prepare_nvfp4_qpn2_dense_linear(layer) + else: return False - logger.info_once( - "SM70 ModelOpt NVFP4 TurboMind dense path enabled " - "(weight-only; activations remain half)." - ) - sm70_tm.prepare_nvfp4_linear(layer) layer.weight = Parameter( torch.empty(0, dtype=torch.uint8, device=layer.weight.device), requires_grad=False, @@ -2402,7 +2437,7 @@ def get_supported_act_dtypes(self) -> list[torch.dtype]: @classmethod def get_min_capability(cls) -> int: if ( - sm70_tm.is_exact_sm70_cuda_platform() + sm70_tm.is_pre_ampere_cuda_platform() and sm70_tm.use_turbomind(envs.VLLM_SM70_FP8_TURBOMIND) and sm70_tm.use_turbomind(envs.VLLM_SM70_NVFP4_TURBOMIND) ): diff --git a/vllm/model_executor/layers/quantization/sm70_turbomind.py b/vllm/model_executor/layers/quantization/sm70_turbomind.py index 889a99873c..36c73448f6 100644 --- a/vllm/model_executor/layers/quantization/sm70_turbomind.py +++ b/vllm/model_executor/layers/quantization/sm70_turbomind.py @@ -30,12 +30,15 @@ class SM70TurboMindLinearState: k_ld: int q_ld: int output_size: int - op_kind: Literal["uint4", "mxfp4", "nvfp4", "nvfp4_qpn4"] + op_kind: Literal["uint4", "mxfp4", "nvfp4", "nvfp4_qpn4", "nvfp4_qpn2_dense"] gated_silu: bool = False dense_weight_ptr: int = 0 global_scale: float = 0.0 use_scale_code: bool = False padded_output_size: int = 0 + # QPN2 launch configuration (split-K, independent accumulator chains). + split_k: int = 0 + accumulator_chains: int = 0 # States retain only data_ptr(), so this cache owns the bounded allocation. @@ -45,6 +48,12 @@ class SM70TurboMindLinearState: def clear_sm70_turbomind_workspaces() -> None: """Release process-global NVFP4 QPN4 dense workspaces.""" _nvfp4_qpn4_dense_workspaces.clear() + from vllm.model_executor.layers.quantization.utils.nvfp4_qpn2_dequant import ( + clear_nvfp4_qpn2_dense_workspaces, + ) + + clear_nvfp4_qpn2_dense_workspaces() + _fp8_qpn8_dense_workspaces.clear() def quant_backend() -> SM70QuantBackend: @@ -103,6 +112,54 @@ def should_prepare_turbomind_or_marlin( return is_exact_sm70_cuda(tensor, use_turbomind(default_enabled) or forces_marlin()) +def is_turing_cuda(tensor: torch.Tensor, enabled: bool) -> bool: + if not enabled or not tensor.is_cuda: + return False + return torch.cuda.get_device_capability(tensor.device) == (7, 5) + + +def is_pre_ampere_cuda_platform() -> bool: + """Return true for Volta and Turing workers, judged on the worker's device. + + Quant-method selection runs before a layer owns a CUDA tensor. The + capability is read from the device this process computes on: on a node + that mixes card generations, device 0 of the visibility list answers for + another card. + """ + if not current_platform.is_cuda(): + return False + device_id = ( + torch.accelerator.current_device_index() if torch.cuda.is_initialized() else 0 + ) + return current_platform.has_device_capability(70, device_id=device_id) and ( + not current_platform.has_device_capability(80, device_id=device_id) + ) + + +def is_turing_cuda_platform() -> bool: + """Return true for Turing workers, judged on the worker's device.""" + if not current_platform.is_cuda(): + return False + device_id = ( + torch.accelerator.current_device_index() if torch.cuda.is_initialized() else 0 + ) + return current_platform.is_device_capability((7, 5), device_id=device_id) + + +def should_prepare_turing_qpn2( + tensor: torch.Tensor, + default_enabled: bool, +) -> bool: + """Turing takes the QPN2 decode kernels with a dense fp16 prefill. + + The TurboMind GEMMs are registered for exact SM70 only (``Sm70`` is + ``Arch<700, 750>``), so Turing cannot take the Volta path; the QPN2 + kernels and the dequantization do not depend on the TurboMind registry. + The same switches as on Volta apply. + """ + return is_turing_cuda(tensor, use_turbomind(default_enabled)) + + def _get_u4_slices(x: torch.Tensor, dtype: torch.dtype) -> list[torch.Tensor]: if x.dtype == torch.int32: count = 8 @@ -162,12 +219,14 @@ def _store_state( meta: torch.Tensor | None, group_size: int, output_size: int, - op_kind: Literal["uint4", "mxfp4", "nvfp4", "nvfp4_qpn4"], + op_kind: Literal["uint4", "mxfp4", "nvfp4", "nvfp4_qpn4", "nvfp4_qpn2_dense"], gated_silu: bool = False, dense_weight_ptr: int = 0, global_scale: float = 0.0, use_scale_code: bool = False, padded_output_size: int = 0, + split_k: int = 0, + accumulator_chains: int = 0, ) -> None: state = SM70TurboMindLinearState( weight=weight, @@ -182,6 +241,8 @@ def _store_state( global_scale=global_scale, use_scale_code=use_scale_code, padded_output_size=padded_output_size, + split_k=split_k, + accumulator_chains=accumulator_chains, ) setattr(layer, STATE_ATTR, state) @@ -347,6 +408,227 @@ def prepare_nvfp4_linear( ) +# Mirror of ``kQpn2DispatchMaxRows`` in nvfp4_qpn2_sm70.cu: rows up to this +# take the QPN2 decode kernels, larger M takes the dense prefill. +QPN2_DISPATCH_MAX_ROWS = 32 +QPN2_GROUP_SIZE = 16 + +# Launch configurations (split-K, accumulator chains) measured on the +# Qwen3.8 TP2 shapes; other shapes take the heuristic below. +_QPN2_LAUNCH_TABLE: dict[tuple[int, int], tuple[int, int]] = { + (1536, 5120): (16, 2), + (4352, 5120): (16, 2), + (5120, 8704): (8, 2), + (5120, 4096): (16, 2), + (5120, 2048): (32, 2), + (5120, 62080): (8, 1), + (5120, 3584): (16, 2), +} + + +def qpn2_launch_config(k: int, n: int) -> tuple[int, int]: + """Split-K and accumulator chains for a QPN2 GEMM of ``[n, k]``.""" + groups = k // QPN2_GROUP_SIZE + config = _QPN2_LAUNCH_TABLE.get((k, n)) + if config is not None and groups % config[0] == 0: + return config + # Smallest split that puts about 640 warps in flight (80 SMs x 8). + for split_k in (8, 16, 32): + if groups % split_k == 0 and (n // 32) * split_k >= 640: + return split_k, 2 if split_k >= 16 else 1 + for split_k in (16, 8): + if groups % split_k == 0: + return split_k, 2 if split_k >= 16 else 1 + raise RuntimeError(f"no QPN2 launch configuration for K={k}, N={n}") + + +def pad_qpn2_output_rows( + weight: torch.Tensor, scales: torch.Tensor +) -> tuple[torch.Tensor, torch.Tensor, int]: + """Pad checkpoint-native output rows to QPN2's 32-column contract.""" + logical_n = weight.shape[0] + physical_n = (logical_n + 31) // 32 * 32 + if physical_n == logical_n: + return weight, scales, physical_n + padded_weight = weight.new_zeros((physical_n, weight.shape[1])) + padded_scales = scales.new_zeros((physical_n, scales.shape[1])) + padded_weight[:logical_n].copy_(weight) + padded_scales[:logical_n].copy_(scales) + return padded_weight, padded_scales, physical_n + + +def prepare_nvfp4_qpn2_dense_linear(layer: torch.nn.Module) -> None: + """Prepare the QPN2 prepack as the only resident layout of an NVFP4 linear. + + Decode (M <= ``QPN2_DISPATCH_MAX_ROWS``) runs the QPN2 kernels on it; + larger M dequantizes into a transient fp16 buffer and runs cuBLAS. No + TurboMind weight is built, so this does not need the TurboMind registry. + """ + if not hasattr(torch.ops._C, "nvfp4_qpn2_prepare_sm70"): + raise RuntimeError( + "The pre-Ampere NVFP4 QPN2 path requires a build with CUDA arch " + "7.0 and the SM70 TurboMind NVFP4 extension." + ) + from vllm import _sm70_ops as sm70_ops + + # Registers the dispatch op now, at weight loading, so it exists before + # the first forward is traced. + from vllm.model_executor.layers.quantization.utils import ( # noqa: F401 + nvfp4_qpn2_dequant, + ) + + weight, scales, padded_output_size = pad_qpn2_output_rows( + layer.weight.data, layer.weight_scale.data + ) + codes, qpn2_scales = sm70_ops.nvfp4_qpn2_prepare_sm70(weight, scales) + output_size = int(layer.weight.shape[0]) + input_size = int(layer.weight.shape[1]) * 2 + split_k, accumulator_chains = qpn2_launch_config(input_size, padded_output_size) + _store_state( + layer, + codes, + qpn2_scales, + None, + QPN2_GROUP_SIZE, + output_size, + "nvfp4_qpn2_dense", + global_scale=float(layer.weight_global_scale.item()), + padded_output_size=padded_output_size, + split_k=split_k, + accumulator_chains=accumulator_chains, + ) + + +# FP8 weights on Turing: QPN8 decode kernels plus the QPN8 dense prefill +# (dequantization into a transient fp16 [K, N] workspace and cuBLAS), both +# provided by fp8_qpn8_sm70.cu without the TurboMind registry. +_fp8_qpn8_dense_workspaces: dict[tuple[int, int, torch.device], torch.Tensor] = {} +FP8_QPN8_STATE_ATTR = "_sm70_fp8_qpn8_state" + + +class FP8QPN8LinearState: + def __init__( + self, + codes: torch.Tensor, + group_scales: torch.Tensor, + output_size: int, + split_k: int, + accumulator_chains: int, + prefetch_codes: bool, + dense_weight_ptr: int, + ) -> None: + self.codes = codes + self.group_scales = group_scales + self.output_size = output_size + self.split_k = split_k + self.accumulator_chains = accumulator_chains + self.prefetch_codes = prefetch_codes + self.dense_weight_ptr = dense_weight_ptr + + +def fp8_qpn8_launch_config(k: int) -> tuple[int, int, bool]: + """Split-K, accumulator chains and code prefetch for a QPN8 GEMM. + + The dispatcher admits split-K up to 16; the measured Qwen3.8 shapes use + 16 (K = 5120) and 12 (K = 1536) with two accumulator chains. + """ + groups = k // QPN2_GROUP_SIZE + for split_k in (16, 12, 8): + if groups % split_k == 0: + return split_k, 2, False + raise RuntimeError(f"no QPN8 launch configuration for K={k}") + + +def get_fp8_qpn8_dense_workspace(k: int, n: int, device: torch.device) -> torch.Tensor: + key = (k, n, device) + workspace = _fp8_qpn8_dense_workspaces.get(key) + if workspace is None: + workspace = torch.empty(k, n, dtype=torch.float16, device=device) + _fp8_qpn8_dense_workspaces[key] = workspace + return workspace + + +def prepare_fp8_qpn8_dense_linear( + layer: torch.nn.Module, weight: torch.Tensor, weight_scale: torch.Tensor +) -> None: + """Prepare a per-tensor FP8 linear as QPN8 codes with channel scales. + + ``weight`` is the checkpoint-native fp8-e4m3fn ``[N, K]`` tensor, + ``weight_scale`` its single scale. The QPN8 prepack takes channel scales, + so the scale is broadcast over the rows; the dequantization is then the + same product the reference path computes. + """ + if not hasattr(torch.ops._C, "fp8_qpn8_prepare_sm70"): + raise RuntimeError( + "The Turing FP8 QPN8 path requires a build with CUDA arch 7.0 and " + "the SM70 TurboMind FP8 extension." + ) + from vllm import _sm70_ops as sm70_ops + + n, k = (int(dim) for dim in weight.shape) + if n % 32 != 0 or k % 16 != 0: + raise RuntimeError(f"QPN8 needs N % 32 == 0 and K % 16 == 0, got N={n}, K={k}.") + channel_scales = ( + weight_scale.to(torch.float32).reshape(1, 1).expand(n, 1).contiguous() + ) + codes, group_scales = sm70_ops.fp8_qpn8_prepare_sm70( + weight.contiguous(), channel_scales + ) + split_k, accumulator_chains, prefetch_codes = fp8_qpn8_launch_config(k) + workspace = get_fp8_qpn8_dense_workspace(k, n, weight.device) + setattr( + layer, + FP8_QPN8_STATE_ATTR, + FP8QPN8LinearState( + codes, + group_scales, + n, + split_k, + accumulator_chains, + prefetch_codes, + workspace.data_ptr(), + ), + ) + + +def has_prepared_fp8_qpn8_linear(layer: torch.nn.Module) -> bool: + return getattr(layer, FP8_QPN8_STATE_ATTR, None) is not None + + +def apply_prepared_fp8_qpn8_linear( + layer: torch.nn.Module, + x: torch.Tensor, + bias: torch.Tensor | None, +) -> torch.Tensor: + state = getattr(layer, FP8_QPN8_STATE_ATTR) + if x.dtype != torch.float16: + raise RuntimeError( + f"The Turing FP8 QPN8 path requires float16 activations, got {x.dtype}." + ) + reshaped_x = x.reshape(-1, x.shape[-1]) + if reshaped_x.stride(-1) != 1: + reshaped_x = reshaped_x.contiguous() + out = torch.empty( + (reshaped_x.shape[0], state.output_size), dtype=x.dtype, device=x.device + ) + from vllm import _sm70_ops as sm70_ops + + sm70_ops.fp8_qpn8_dispatch_sm70_out( + out, + state.dense_weight_ptr, + reshaped_x, + state.codes, + state.group_scales, + state.split_k, + state.accumulator_chains, + state.prefetch_codes, + False, + ) + if bias is not None: + out.add_(bias) + return out.reshape(x.shape[:-1] + (state.output_size,)) + + def get_nvfp4_qpn4_dense_workspace(weight: torch.Tensor) -> torch.Tensor | None: device_index = weight.device.index if device_index is None: @@ -474,6 +756,28 @@ def apply_prepared_linear( state.use_scale_code, False, ) + elif state.op_kind == "nvfp4_qpn2_dense": + if reshaped_x.dtype != torch.float16: + raise RuntimeError( + "The pre-Ampere NVFP4 QPN2 path requires float16 activations, " + f"got {reshaped_x.dtype}." + ) + if reshaped_x.stride(-1) != 1: + reshaped_x = reshaped_x.contiguous() + from vllm.model_executor.layers.quantization.utils import ( + nvfp4_qpn2_dequant, + ) + + out = nvfp4_qpn2_dequant.nvfp4_qpn2_dispatch_linear( + reshaped_x, + state.weight, + state.scales, + state.global_scale, + kernel_output_size, + reshaped_x.shape[1], + state.split_k, + state.accumulator_chains, + ) else: raise AssertionError(f"unknown SM70 TurboMind op kind: {state.op_kind}") if kernel_output_size != state.output_size: diff --git a/vllm/model_executor/layers/quantization/utils/nvfp4_qpn2_dequant.py b/vllm/model_executor/layers/quantization/utils/nvfp4_qpn2_dequant.py new file mode 100644 index 0000000000..ba09829a2e --- /dev/null +++ b/vllm/model_executor/layers/quantization/utils/nvfp4_qpn2_dequant.py @@ -0,0 +1,288 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# The dequantization of the QPN2 prepack is derived from dnv2003/v100-skinny +# (MIT), where it serves the same purpose on Turing. +"""Dense fp16 GEMM on the QPN2-packed NVFP4 layout for large M. + +The QPN2 prepack (``nvfp4_qpn2_prepare_sm70``) is the only resident weight +layout of the pre-Ampere-but-not-Volta NVFP4 linear path. Decode runs the +QPN2 kernels on it directly; prefill dequantizes one layer at a time into a +transient fp16 ``[N, K]`` buffer and runs ``torch.matmul`` on the fp16 tensor +cores. Marlin's FP4 GEMM reaches only about 27 TFLOPS on Turing, cuBLAS fp16 +does better, and keeping a second weight layout for prefill would double the +weight memory. + +Layout (see ``nvfp4_qpn2_prepack_codes_kernel``): codes are +``[tiles = N/32][groups = K/16][lane = 32][8 bytes]``; a lane owns row +``n = tile * 32 + col(lane)`` with +``col = ((lane >> 2) & 3) * 8 + (lane & 3) + ((lane & 16) > 0) * 4``. Its 8 +bytes hold the 16 nibbles of the group's 16 k in the order +``(0, 2, 4, 6, 1, 3, 5, 7, 8, 10, 12, 14, 9, 11, 13, 15)``: byte ``b`` holds +``korder[2b]`` in its low nibble and ``korder[2b + 1]`` in its high nibble. +Scales are one fp8-e4m3fn byte per ``(tile, group, lane)``. +""" + +import torch + +from vllm.triton_utils import tl, triton +from vllm.utils.torch_utils import direct_register_custom_op + +# e2m1 magnitudes indexed by the low three code bits; bit 3 is the sign. +_E2M1_MAGNITUDES = (0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0) +_KORDER = (0, 2, 4, 6, 1, 3, 5, 7, 8, 10, 12, 14, 9, 11, 13, 15) + +# Bounded transient workspace: one fp16 [N, K] buffer per distinct shape and +# device, reused by every layer of that shape. +_dense_workspaces: dict[tuple[int, int, torch.device], torch.Tensor] = {} + + +def clear_nvfp4_qpn2_dense_workspaces() -> None: + _dense_workspaces.clear() + + +def _lane_to_col() -> torch.Tensor: + lane = torch.arange(32) + return ((lane >> 2) & 3) * 8 + (lane & 3) + ((lane & 16) > 0).long() * 4 + + +def nvfp4_qpn2_dequant_reference( + codes: torch.Tensor, + scales: torch.Tensor, + global_scale: float, + n: int, + k: int, +) -> torch.Tensor: + """Pure-torch inverse of the QPN2 prepack (slow; the spec for the kernel).""" + tiles, groups = n // 32, k // 16 + device = codes.device + qc = codes.view(tiles, groups, 32, 8) + qs = scales.view(torch.uint8).view(tiles, groups, 32) + nib = torch.stack([qc & 0xF, qc >> 4], dim=-1).view(tiles, groups, 32, 16) + korder = torch.tensor(_KORDER, device=device) + inverse = torch.empty(16, dtype=torch.long, device=device) + inverse[korder] = torch.arange(16, device=device) + nib = nib[..., inverse] + magnitudes = torch.tensor(_E2M1_MAGNITUDES, device=device, dtype=torch.float32) + values = magnitudes[(nib & 7).long()] * torch.where(nib & 8 > 0, -1.0, 1.0) + scale = qs.view(torch.float8_e4m3fn).to(torch.float32) + values = values * scale.unsqueeze(-1) * global_scale + col = _lane_to_col().to(device) + out = torch.empty(n, k, dtype=torch.float32, device=device) + rows = torch.arange(tiles, device=device).view(tiles, 1) * 32 + col.view(1, 32) + out[rows.view(-1)] = values.permute(0, 2, 1, 3).reshape(tiles * 32, k) + return out.to(torch.float16) + + +@triton.jit +def _e2m1_value(code): + """e2m1 nibble -> float: magnitudes 0, .5, 1, 1.5, 2, 3, 4, 6; bit 3 is the sign.""" + mag = code & 7 + m_f = mag.to(tl.float32) + val = tl.where( + mag < 4, m_f * 0.5, tl.where(mag < 6, m_f - 2.0, tl.where(mag == 6, 4.0, 6.0)) + ) + return tl.where((code & 8) > 0, -val, val) + + +@triton.jit +def _e4m3_value(b): + """fp8-e4m3fn byte -> float in integer arithmetic (no fp8 hardware here). + + Bias 7; exponent 0 is subnormal (mantissa / 8 * 2^-6). The NaN code 0x7f + does not occur in NVFP4 block scales. + """ + sign = tl.where((b & 0x80) > 0, -1.0, 1.0) + exp = ((b >> 3) & 0xF).to(tl.float32) + mant = (b & 7).to(tl.float32) / 8.0 + normal = (1.0 + mant) * tl.exp2(exp - 7.0) + subnormal = mant * tl.exp2(-6.0) + return sign * tl.where(exp == 0, subnormal, normal) + + +@triton.jit +def _nvfp4_qpn2_dequant_kernel( + codes32_ptr, + scales_ptr, + out_ptr, + global_scale, + groups, + K, + GROUPS_PER_BLOCK: tl.constexpr, +): + # One program: one tile (32 rows) x GROUPS_PER_BLOCK groups. A lane's + # 8-byte payload is read as two 32-bit words (64-bit shifts are slow on + # Volta and Turing); nibble j sits at bit 4 * (j & 7) of word j >> 3 and + # holds k offset korder[j]. Scales are read once per (lane, group). + tile = tl.program_id(0) + gblock = tl.program_id(1) + lane = tl.arange(0, 32) + col = ((lane >> 2) & 3) * 8 + (lane & 3) + ((lane & 16) > 0).to(tl.int32) * 4 + row = tile * 32 + col + g = gblock * GROUPS_PER_BLOCK + tl.arange(0, GROUPS_PER_BLOCK) + valid = g < groups + lane_base = (tile * groups + g) * 32 + idx = lane_base[None, :] + lane[:, None] + w0 = tl.load(codes32_ptr + idx * 2, mask=valid[None, :], other=0) + w1 = tl.load(codes32_ptr + idx * 2 + 1, mask=valid[None, :], other=0) + sc = _e4m3_value(tl.load(scales_ptr + idx, mask=valid[None, :], other=0)) + sc = sc * global_scale + # Produce the 8 k of each word in natural k order so the stores are + # contiguous 16-byte runs: k offset p (0..7) lives in nibble + # j = (p & 1) * 4 + (p >> 1), the inverse of korder 0, 2, 4, 6, 1, 3, 5, 7. + p = tl.arange(0, 8) + shift = (4 * ((p & 1) * 4 + (p >> 1)))[None, None, :] + v0 = _e2m1_value((w0[:, :, None] >> shift) & 0xF) * sc[:, :, None] + v1 = _e2m1_value((w1[:, :, None] >> shift) & 0xF) * sc[:, :, None] + base_idx = row[:, None, None] * K + g[None, :, None] * 16 + p[None, None, :] + m3 = valid[None, :, None] + tl.store(out_ptr + base_idx, v0.to(tl.float16), mask=m3) + tl.store(out_ptr + base_idx + 8, v1.to(tl.float16), mask=m3) + + +def nvfp4_qpn2_dequant( + codes: torch.Tensor, + scales: torch.Tensor, + global_scale: float, + n: int, + k: int, + out: torch.Tensor | None = None, +) -> torch.Tensor: + """Dense fp16 ``[n, k]`` weight from the QPN2 prepack; ``out`` may be reused.""" + groups = k // 16 + tiles = n // 32 + if out is None: + out = torch.empty(n, k, dtype=torch.float16, device=codes.device) + # 64 groups per program with 8 warps was the fastest of the swept + # configurations on Volta (34816 x 5120: 1.83 ms vs 2.19 ms for 32 / 4). + groups_per_block = 64 + grid = (tiles, triton.cdiv(groups, groups_per_block)) + codes32 = codes.view(torch.int32) + # The scales are fp8-e4m3fn bytes; Triton on these devices reads them as + # uint8 and converts in integer arithmetic. + _nvfp4_qpn2_dequant_kernel[grid]( + codes32, + scales.view(torch.uint8), + out, + global_scale, + groups, + k, + GROUPS_PER_BLOCK=groups_per_block, + num_warps=8, + ) + return out + + +def _nvfp4_qpn2_dense_linear( + x: torch.Tensor, + codes: torch.Tensor, + scales: torch.Tensor, + global_scale: float, + n: int, + k: int, +) -> torch.Tensor: + key = (n, k, x.device) + workspace = _dense_workspaces.get(key) + if workspace is None: + workspace = torch.empty(n, k, dtype=torch.float16, device=x.device) + _dense_workspaces[key] = workspace + weight = nvfp4_qpn2_dequant(codes, scales, global_scale, n, k, out=workspace) + return torch.nn.functional.linear(x, weight) + + +def _nvfp4_qpn2_dense_linear_fake( + x: torch.Tensor, + codes: torch.Tensor, + scales: torch.Tensor, + global_scale: float, + n: int, + k: int, +) -> torch.Tensor: + return x.new_empty((x.shape[0], n)) + + +direct_register_custom_op( + op_name="nvfp4_qpn2_dense_linear", + op_func=_nvfp4_qpn2_dense_linear, + mutates_args=[], + fake_impl=_nvfp4_qpn2_dense_linear_fake, +) + + +def nvfp4_qpn2_dense_linear( + x: torch.Tensor, + codes: torch.Tensor, + scales: torch.Tensor, + global_scale: float, + n: int, + k: int, +) -> torch.Tensor: + """``x @ W.T`` for fp16 ``x`` of ``[M, k]`` and the QPN2 prepack of ``[n, k]``.""" + return torch.ops.vllm.nvfp4_qpn2_dense_linear(x, codes, scales, global_scale, n, k) + + +# Mirror of ``kQpn2DispatchMaxRows`` in nvfp4_qpn2_sm70.cu. +QPN2_DISPATCH_MAX_ROWS = 32 + + +def _nvfp4_qpn2_dispatch_linear( + x: torch.Tensor, + codes: torch.Tensor, + scales: torch.Tensor, + global_scale: float, + n: int, + k: int, + split_k: int, + accumulator_chains: int, +) -> torch.Tensor: + # The split on M happens here at run time, inside one opaque op: a Python + # branch in the model's forward would be traced once by torch.compile at + # the warm-up M and keep the dense path in the decode graph. + if x.shape[0] <= QPN2_DISPATCH_MAX_ROWS: + from vllm import _sm70_ops as sm70_ops + + out = torch.empty((x.shape[0], n), dtype=x.dtype, device=x.device) + sm70_ops.nvfp4_qpn2_gemm_sm70_out( + out, x, codes, scales, global_scale, split_k, accumulator_chains + ) + return out + return _nvfp4_qpn2_dense_linear(x, codes, scales, global_scale, n, k) + + +def _nvfp4_qpn2_dispatch_linear_fake( + x: torch.Tensor, + codes: torch.Tensor, + scales: torch.Tensor, + global_scale: float, + n: int, + k: int, + split_k: int, + accumulator_chains: int, +) -> torch.Tensor: + return x.new_empty((x.shape[0], n)) + + +direct_register_custom_op( + op_name="nvfp4_qpn2_dispatch_linear", + op_func=_nvfp4_qpn2_dispatch_linear, + mutates_args=[], + fake_impl=_nvfp4_qpn2_dispatch_linear_fake, +) + + +def nvfp4_qpn2_dispatch_linear( + x: torch.Tensor, + codes: torch.Tensor, + scales: torch.Tensor, + global_scale: float, + n: int, + k: int, + split_k: int, + accumulator_chains: int, +) -> torch.Tensor: + """QPN2 kernels for M <= 32, dequantization plus cuBLAS above. + + The split is decided at run time inside the op. + """ + return torch.ops.vllm.nvfp4_qpn2_dispatch_linear( + x, codes, scales, global_scale, n, k, split_k, accumulator_chains + )