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 + )