From 462744b08e03b35c44a5d46c8465cf430949e2e6 Mon Sep 17 00:00:00 2001 From: samnordmann Date: Wed, 30 Sep 2026 08:43:41 +0200 Subject: [PATCH] [Kernel] Fuse native BF16 ReLU2 with NVFP4 quantization Fuse the native BF16 activation and NVFP4 producer without changing the CuTe-DSL GEMM. Preserve rounding, padded scales, and unsupported-input fallbacks. Keep dynamic-row launch selection outside guard-free compiled graphs and restrict automatic fusion to the validated SM10x CuTe path. Add startup warmup and targeted regression coverage. Co-authored-by: Codex Signed-off-by: samnordmann --- .../fusion/test_quant_activation_contract.py | 47 ++ .../test_flashinfer_nvfp4_scaled_mm.py | 40 ++ tests/kernels/test_fused_quant_activation.py | 404 +++++++++++++++++- .../kernels/linear/nvfp4/flashinfer.py | 6 + .../layers/fusion/fused_act_quant.py | 100 ++++- .../layers/fusion/quant_activation.py | 31 +- .../layers/fusion/relu2_nvfp4_quant.py | 269 ++++++++++++ vllm/model_executor/models/nemotron_h.py | 6 +- 8 files changed, 890 insertions(+), 13 deletions(-) create mode 100644 vllm/model_executor/layers/fusion/relu2_nvfp4_quant.py diff --git a/tests/fusion/test_quant_activation_contract.py b/tests/fusion/test_quant_activation_contract.py index da507e5caf50..dd4ab8c56854 100644 --- a/tests/fusion/test_quant_activation_contract.py +++ b/tests/fusion/test_quant_activation_contract.py @@ -44,10 +44,12 @@ Int8ScaledMMLinearKernel, Int8ScaledMMLinearLayerConfig, ) +from vllm.model_executor.layers.activation import ReLUSquaredActivation, SiluAndMul from vllm.model_executor.layers.fusion.quant_activation import ( QuantizedActivation, as_quantized_activation, expose_input_quant_key, + get_fused_act_quant_key, get_input_quant_key, ) from vllm.model_executor.layers.quantization.utils.quant_utils import ( @@ -138,6 +140,51 @@ def test_bridge_marks_supporting_and_skips_others(): assert get_input_quant_key(layer) is None +@pytest.mark.parametrize("act_type", [ReLUSquaredActivation, SiluAndMul]) +def test_activation_policy_reexposure(default_vllm_config, act_type): + """Backend changes must clear stale keys and producer restrictions.""" + layer = torch.nn.Module() + act = act_type(compile_native=False) + for kernel_cls, allowed in ( + (FlashInferCuteDslNvFp4LinearKernel, True), + (FlashInferCutlassNvFp4LinearKernel, act_type is SiluAndMul), + (FlashInferTrtllmNvFp4LinearKernel, False), + (FlashInferCuteDslNvFp4LinearKernel, True), + (CutlassFP8ScaledMMLinearKernel, True), + ): + kernel = _probe(kernel_cls) + expose_input_quant_key(layer, kernel) + assert get_input_quant_key(layer) == kernel.input_quant_key() + assert (get_fused_act_quant_key(layer, act) is not None) is allowed + layer.requires_unquantized_input = True + assert get_input_quant_key(layer) is get_fused_act_quant_key(layer, act) is None + + +def test_nvfp4_policy_preserves_silu_fusion(default_vllm_config, monkeypatch): + """Both existing NVFP4 consumers retain their SiLU producer selection.""" + from vllm.model_executor.layers.fusion import fused_act_quant as fusion + + act = SiluAndMul(compile_native=False) + act._forward_method = act.forward_native + layer = torch.nn.Module() + x = torch.ones((2, 4), device="cpu") + qa = QuantizedActivation( + data=torch.zeros((2, 1), dtype=torch.uint8, device="cpu"), + scale=torch.ones((), device="cpu"), + orig_dtype=x.dtype, + orig_shape=torch.Size((2, 2)), + quant_key=kNvfp4Dynamic, + ) + key = (SiluAndMul, kNvfp4Dynamic) + monkeypatch.setitem(fusion._FUSED_ACT_QUANT, key, lambda x, linear: qa) + for kernel_cls in ( + FlashInferCuteDslNvFp4LinearKernel, + FlashInferCutlassNvFp4LinearKernel, + ): + expose_input_quant_key(layer, _probe(kernel_cls)) + assert fusion.maybe_fused_act_quant(act, x, layer) is qa + + def test_as_quantized_activation_validates_key(): qa = QuantizedActivation( data=torch.zeros(2, 4, dtype=current_platform.fp8_dtype()), diff --git a/tests/kernels/quantization/test_flashinfer_nvfp4_scaled_mm.py b/tests/kernels/quantization/test_flashinfer_nvfp4_scaled_mm.py index 47c25d046d2c..90354beae966 100644 --- a/tests/kernels/quantization/test_flashinfer_nvfp4_scaled_mm.py +++ b/tests/kernels/quantization/test_flashinfer_nvfp4_scaled_mm.py @@ -18,9 +18,13 @@ FlashInferCutlassNvFp4LinearKernel, ) from vllm.model_executor.layers.fusion.quant_activation import QuantizedActivation +from vllm.model_executor.layers.fusion.relu2_nvfp4_quant import relu_squared_nvfp4_quant from vllm.model_executor.layers.quantization.utils.nvfp4_emulation_utils import ( dequantize_to_dtype, ) +from vllm.model_executor.layers.quantization.utils.nvfp4_utils import ( + pad_nvfp4_weight_for_cutlass, +) from vllm.model_executor.layers.quantization.utils.quant_utils import kNvfp4Dynamic from vllm.platforms import current_platform from vllm.utils.flashinfer import ( @@ -99,6 +103,42 @@ def test_flashinfer_nvfp4_preserves_quantized_input_shape( torch.testing.assert_close(actual.reshape(expected.shape), expected) +@pytest.mark.parametrize("shape", [(1, 1344), (17, 5376), (129, 48)]) +@pytest.mark.parametrize("with_bias", [False, True]) +@torch.inference_mode() +def test_cutedsl_relu2_prequantized_bypasses_quantizer(monkeypatch, shape, with_bias): + """The CuTe-DSL consumer preserves output and padding without requantizing.""" + from vllm.model_executor.kernels.linear.nvfp4 import flashinfer as consumer_module + + supported, reason = FlashInferCuteDslNvFp4LinearKernel.is_supported() + if not supported: + pytest.skip(reason) + m, k = shape + x = torch.randn((m, k), device="cuda", dtype=torch.bfloat16) + n = 130 # Exercise output slicing as well as activation K padding. + w = torch.randn((n, k), device="cuda", dtype=torch.bfloat16) + layer = torch.nn.Module() + layer.input_global_scale_inv = torch.tensor(64.0, device="cuda") + weight_scale_inv = torch.tensor(512.0, device="cuda") + weight, layer.weight_scale = ops.scaled_fp4_quant(w, weight_scale_inv) + layer.weight, _ = pad_nvfp4_weight_for_cutlass(weight) + layer.input_size_per_partition = k + layer.output_size_per_partition = n + layer.alpha = 1.0 / (layer.input_global_scale_inv * weight_scale_inv) + kernel = FlashInferCuteDslNvFp4LinearKernel(NvFp4LinearLayerConfig()) + activation = torch.compile(lambda value: torch.square(torch.relu(value))) + bias = torch.randn(n, device="cuda", dtype=x.dtype) if with_bias else None + expected = kernel.apply_weights(layer, activation(x), bias) + qa = relu_squared_nvfp4_quant(x, layer) + + def reject_requant(*args, **kwargs): + raise AssertionError("CuTe-DSL requantized an existing QuantizedActivation") + + monkeypatch.setattr(consumer_module, "scaled_fp4_quant", reject_requant) + actual = kernel.apply_weights(layer, qa, bias) + torch.testing.assert_close(actual, expected, rtol=0, atol=0) + + def get_ref_results( a_fp4, b_fp4, diff --git a/tests/kernels/test_fused_quant_activation.py b/tests/kernels/test_fused_quant_activation.py index 90309dbd1d9b..e04b6b5d2f91 100644 --- a/tests/kernels/test_fused_quant_activation.py +++ b/tests/kernels/test_fused_quant_activation.py @@ -1,12 +1,20 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from itertools import product + import pytest import torch import torch.nn.functional as F import vllm._custom_ops as ops from tests.kernels.utils import opcheck -from vllm.model_executor.layers.activation import SiluAndMul +from vllm.config import ( + CompilationConfig, + CompilationMode, + VllmConfig, + set_current_vllm_config, +) +from vllm.model_executor.layers.activation import ReLUSquaredActivation, SiluAndMul from vllm.model_executor.layers.fusion.fused_act_quant import maybe_fused_act_quant from vllm.model_executor.layers.fusion.quant_activation import QuantizedActivation from vllm.model_executor.layers.quantization.utils.fp8_utils import ( @@ -15,6 +23,7 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import ( kFp8Dynamic128Sym, kFp8StaticTensorSym, + kNvfp4Dynamic, ) from vllm.platforms import current_platform from vllm.utils.torch_utils import set_random_seed @@ -109,12 +118,233 @@ class MockLinearNoQuant(torch.nn.Module): pass +class MockLinearNvfp4(torch.nn.Module): + def __init__(self, scale: torch.Tensor): + super().__init__() + self._input_quant_key = kNvfp4Dynamic + self.input_global_scale_inv = scale + + +@pytest.fixture +def native_vllm_config(): + config = VllmConfig( + compilation_config=CompilationConfig( + mode=CompilationMode.VLLM_COMPILE, custom_ops=["none"] + ) + ) + with set_current_vllm_config(config): + yield config + + +@pytest.mark.skipif( + not ( + current_platform.is_cuda() and current_platform.is_device_capability_family(100) + ), + reason="ReLU2 NVFP4 fusion is validated on SM10x", +) +@pytest.mark.parametrize("shape", [(1, 16), (17, 80), (129, 1344), (256, 5376)]) +@pytest.mark.parametrize("scale_value", [0.0009765625, 1.0, 17.125, 224.0]) +@torch.inference_mode() +def test_relu2_nvfp4_exact_and_changed_graph(native_vllm_config, shape, scale_value): + """Preserve BF16 rounding, packing, padded scales, and replay data dependence.""" + from tests.kernels.quantization.test_nvfp4_quant import recover_swizzled_scales + + x = torch.randn(shape, dtype=torch.bfloat16, device="cuda") + scale = torch.tensor(scale_value, dtype=torch.float32, device=x.device) + linear = MockLinearNvfp4(scale) + act = ReLUSquaredActivation(compile_native=False) + compiled_relu = torch.compile(act.forward_native, fullgraph=True) + expected = ops.scaled_fp4_quant(compiled_relu(x), scale) + before = x.clone() + actual = maybe_fused_act_quant(act, x, linear) + assert isinstance(actual, QuantizedActivation) + assert torch.equal(x, before) + assert torch.equal(actual.data, expected[0]) + expected_scale = recover_swizzled_scales(expected[1], *shape) + actual_scale = recover_swizzled_scales(actual.scale, *shape, include_padding=True) + torch.testing.assert_close( + actual_scale[: shape[0], : shape[1] // 16], expected_scale, rtol=0, atol=0 + ) + actual_scale[: shape[0], : shape[1] // 16] = 0 + assert not torch.count_nonzero(actual_scale) + old_data = actual.data.clone() + old_scale = actual.scale.view(torch.uint8).clone() + for _ in range(3): + maybe_fused_act_quant(act, x, linear) + torch.accelerator.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + replayed = maybe_fused_act_quant(act, x, linear) + x.normal_() + scale.mul_(1.25) + graph.replay() + expected = ops.scaled_fp4_quant(compiled_relu(x), scale) + assert torch.equal(replayed.data, expected[0]) + torch.testing.assert_close( + recover_swizzled_scales(replayed.scale, *shape), + recover_swizzled_scales(expected[1], *shape), + rtol=0, + atol=0, + ) + assert torch.equal(actual.data, old_data) + assert torch.equal(actual.scale.view(torch.uint8), old_scale) + + +@pytest.mark.skipif( + not ( + current_platform.is_cuda() and current_platform.is_device_capability_family(100) + ), + reason="ReLU2 NVFP4 fusion is validated on SM10x", +) +@pytest.mark.parametrize("width", [0, 15, 17, 33]) +@torch.inference_mode() +def test_relu2_nvfp4_unsupported_width_preserves_fallback(native_vllm_config, width): + x = torch.randn((3, width), dtype=torch.bfloat16, device="cuda") + scale = torch.ones((), dtype=torch.float32, device=x.device) + act = ReLUSquaredActivation(compile_native=False) + result = maybe_fused_act_quant(act, x, MockLinearNvfp4(scale)) + assert isinstance(result, torch.Tensor) + torch.testing.assert_close(result, act.forward_native(x), rtol=0, atol=0) + + class MockLinearRequiresUnquantized(MockLinearFp8Static): """Mock consumer that needs the original activation for another branch.""" requires_unquantized_input = True +@pytest.mark.skipif( + not ( + current_platform.is_cuda() and current_platform.is_device_capability_family(100) + ), + reason="ReLU2 NVFP4 fusion is validated on SM10x", +) +@pytest.mark.parametrize("value", [float("nan"), float("inf"), 3.0e20, 100.0]) +@pytest.mark.parametrize("scale_value", [2.0**-30, 1.0, 64.0, 1.0e20]) +@torch.inference_mode() +def test_relu2_nvfp4_nonfinite_and_saturation(native_vllm_config, value, scale_value): + x = torch.tensor([[value] + [1.0] * 15], dtype=torch.bfloat16, device="cuda") + scale = torch.tensor(scale_value, dtype=torch.float32, device=x.device) + act = ReLUSquaredActivation(compile_native=False) + compiled_relu = torch.compile(act.forward_native, fullgraph=True) + expected = ops.scaled_fp4_quant( + compiled_relu(x), scale, backend="flashinfer-cutedsl" + ) + actual = maybe_fused_act_quant(act, x, MockLinearNvfp4(scale)) + assert isinstance(actual, QuantizedActivation) + assert torch.equal(actual.data, expected[0]) + assert torch.equal(actual.scale.view(torch.uint8), expected[1].view(torch.uint8)) + + +@pytest.mark.skipif( + not ( + current_platform.is_cuda() and current_platform.is_device_capability_family(100) + ), + reason="ReLU2 NVFP4 fusion is validated on SM10x", +) +@pytest.mark.parametrize("mode", ["native", "compiled_native", "cuda"]) +@torch.inference_mode() +def test_relu2_nvfp4_preserves_selected_activation(mode): + """CUDA ReLU2 suppresses NaN and must not use the native-semantics producer.""" + config = VllmConfig( + compilation_config=CompilationConfig( + mode=CompilationMode.VLLM_COMPILE, + custom_ops=["all" if mode == "cuda" else "none"], + ) + ) + with set_current_vllm_config(config): + act = ReLUSquaredActivation(compile_native=mode == "compiled_native") + x = torch.tensor( + [[float("nan")] + [1.0] * 15], device="cuda", dtype=torch.bfloat16 + ) + scale = torch.ones((), device=x.device, dtype=torch.float32) + expected = act(x) + dispatch = maybe_fused_act_quant + if mode == "compiled_native": + dispatch = torch.compile(dispatch, fullgraph=True) + actual = dispatch(act, x, MockLinearNvfp4(scale)) + if mode == "cuda": + assert isinstance(actual, torch.Tensor) + assert expected[0, 0].item() == 0 + torch.testing.assert_close(actual, expected, rtol=0, atol=0) + else: + assert isinstance(actual, QuantizedActivation) + expected_data, expected_scale = ops.scaled_fp4_quant(expected, scale) + assert torch.equal(actual.data, expected_data) + assert torch.equal( + actual.scale.view(torch.uint8), expected_scale.view(torch.uint8) + ) + + +@pytest.mark.skipif( + not ( + current_platform.is_cuda() and current_platform.is_device_capability_family(100) + ), + reason="ReLU2 NVFP4 fusion is validated on SM10x", +) +@torch.inference_mode() +@pytest.mark.parametrize( + "use_aot,use_bytecode_hook", [(False, False), (False, True), (True, False)] +) +def test_relu2_nvfp4_fullgraph_changed_rows( + native_vllm_config, monkeypatch, use_aot, use_bytecode_hook +): + """One guard-free vLLM graph must handle unseen M, including K padding.""" + from tests.kernels.quantization.test_nvfp4_quant import recover_swizzled_scales + from vllm.compilation.counter import compilation_counter + from vllm.compilation.wrapper import TorchCompileWithNoGuardsWrapper + from vllm.config import CUDAGraphMode + + monkeypatch.setenv("VLLM_DISABLE_COMPILE_CACHE", "1") + monkeypatch.setenv("VLLM_USE_AOT_COMPILE", str(int(use_aot))) + monkeypatch.setenv("VLLM_USE_BYTECODE_HOOK", str(int(use_bytecode_hook))) + native_vllm_config.compilation_config.cudagraph_mode = CUDAGraphMode.NONE + native_vllm_config.compilation_config.compile_sizes = [] + native_vllm_config.scheduler_config.max_num_batched_tokens = 8192 + native_vllm_config.compilation_config.compile_ranges_endpoints = [8192] + + class Wrapper(TorchCompileWithNoGuardsWrapper): + def __init__(self): + self.act = ReLUSquaredActivation(compile_native=True) + self.linear = MockLinearNvfp4(torch.tensor(17.125, device="cuda")) + super().__init__() + + def forward(self, x): + qa = maybe_fused_act_quant(self.act, x, self.linear) + return qa.data, qa.scale + + torch._dynamo.reset() + dispatch = Wrapper() + with compilation_counter.expect(num_graphs_seen=1, num_inductor_compiles=1): + for m in (128, 17, 129, 1, 257, 33, 4097): + x = torch.randn((m, 80), device="cuda", dtype=torch.bfloat16) + if m == 128: + torch._dynamo.mark_dynamic(x, 0) + actual_data, actual_scale = dispatch(x) + expected_data, expected_scale = ops.scaled_fp4_quant( + dispatch.act.forward_native(x), + dispatch.linear.input_global_scale_inv, + backend="flashinfer-cutedsl", + ) + assert torch.equal(actual_data, expected_data) + scales = recover_swizzled_scales( + actual_scale, *x.shape, include_padding=True + ) + torch.testing.assert_close( + scales[:m, :5], + recover_swizzled_scales(expected_scale, *x.shape), + rtol=0, + atol=0, + ) + scales[:m, :5] = 0 + assert not torch.count_nonzero(scales) + opcheck( + torch.ops.vllm.relu_squared_nvfp4_quant, + (x, dispatch.linear.input_global_scale_inv), + ) + torch._dynamo.reset() + + @pytest.mark.parametrize("token_shape", [(1,), (16,), (128,), (1, 17), (2, 17)]) @pytest.mark.parametrize("hidden_size", [128, 512, 1024]) @pytest.mark.parametrize("dtype", DTYPES) @@ -244,3 +474,175 @@ def test_maybe_fused_act_quant_preserves_required_unquantized_input( assert isinstance(result, torch.Tensor) torch.testing.assert_close(result, act_fn(x)) + + +def _relu2_nvfp4_warmup_mlp(width): + module = torch.nn.Module() + module.act_fn = ReLUSquaredActivation(compile_native=False) + module.down_proj = MockLinearNvfp4( + torch.ones(1, dtype=torch.float32, device="cuda") + ) + module.down_proj.params_dtype = torch.bfloat16 + module.down_proj.input_size_per_partition = width + return module + + +@pytest.mark.skipif( + not ( + current_platform.is_cuda() and current_platform.is_device_capability_family(100) + ), + reason="ReLU2 NVFP4 fusion is validated on SM10x", +) +def test_relu2_nvfp4_warmup_resolves_post_load_consumer(native_vllm_config): + """Resolve post-load scales and a replaced LoRA consumer at warmup time.""" + from vllm.model_executor.layers.fusion.fused_act_quant import ( + register_relu_squared_nvfp4_quant_warmup, + ) + from vllm.model_executor.layers.fusion.relu2_nvfp4_quant import ( + _relu_squared_nvfp4_warmup_inputs, + ) + from vllm.model_executor.warmup.jit_warmup import JitWarmupRegistry + + native_vllm_config.kernel_config.enable_jit_warmup = True + module = _relu2_nvfp4_warmup_mlp(80) + scale = module.down_proj.input_global_scale_inv + del module.down_proj.input_global_scale_inv + registry = JitWarmupRegistry(native_vllm_config) + with registry.activate(): + register_relu_squared_nvfp4_quant_warmup(module) + assert len(registry) == 1 + assert not list(_relu_squared_nvfp4_warmup_inputs(module, 128)) + + module.down_proj.input_global_scale_inv = scale + assert list(_relu_squared_nvfp4_warmup_inputs(module, 128)) + module.down_proj = _relu2_nvfp4_warmup_mlp(80).down_proj + module.down_proj.requires_unquantized_input = True + assert not list(_relu_squared_nvfp4_warmup_inputs(module, 128)) + + +@pytest.mark.skipif( + not ( + current_platform.is_cuda() and current_platform.is_device_capability_family(100) + ), + reason="ReLU2 NVFP4 fusion is validated on SM10x", +) +@pytest.mark.parametrize("width", [16, 48, 80, 1344, 5376]) +def test_relu2_nvfp4_oversized_dispatch_falls_back( + native_vllm_config, monkeypatch, width +): + """Reject unsafe index products with fake inputs, without huge allocations.""" + from torch._subclasses.fake_tensor import FakeTensorMode + + from vllm.model_executor.layers.fusion.fused_act_quant import ( + _FUSED_ACT_QUANT, + _relu_squared_nvfp4_quant_supported, + ) + from vllm.model_executor.layers.fusion.relu2_nvfp4_quant import ( + _relu_squared_nvfp4_max_rows, + ) + + def fail_fused(*args, **kwargs): + pytest.fail("Oversized input reached the int32-indexed producer") + + monkeypatch.setitem( + _FUSED_ACT_QUANT, (ReLUSquaredActivation, kNvfp4Dynamic), fail_fused + ) + act = ReLUSquaredActivation(compile_native=False) + with FakeTensorMode(): + x = torch.empty( + (_relu_squared_nvfp4_max_rows(width) + 1, width), + dtype=torch.bfloat16, + device="cuda", + ) + scale = torch.ones(1, dtype=torch.float32, device="cuda") + linear = MockLinearNvfp4(scale) + assert _relu_squared_nvfp4_quant_supported(act, x[:-1], linear) + actual = maybe_fused_act_quant(act, x, linear) + assert isinstance(actual, torch.Tensor) + assert actual.shape == x.shape + assert actual.dtype == x.dtype + + +@pytest.mark.skipif( + not ( + current_platform.is_cuda() and current_platform.is_device_capability_family(100) + ), + reason="ReLU2 NVFP4 fusion is validated on SM10x", +) +@pytest.mark.parametrize("width", [16, 48, 80, 1344, 5376]) +@torch.inference_mode() +def test_relu2_nvfp4_startup_warmup_covers_unseen_rows( + native_vllm_config, monkeypatch, width +): + """Startup must cover unseen row counts and offset input/scale pointers.""" + from vllm.model_executor.layers.fusion.fused_act_quant import ( + register_relu_squared_nvfp4_quant_warmup, + ) + from vllm.model_executor.warmup.jit_warmup import JitWarmupRegistry + from vllm.triton_utils import triton + + native_vllm_config.kernel_config.enable_jit_warmup = True + native_vllm_config.scheduler_config.max_num_batched_tokens = 32768 + module = _relu2_nvfp4_warmup_mlp(width) + registry = JitWarmupRegistry(native_vllm_config) + with registry.activate(): + register_relu_squared_nvfp4_quant_warmup(module) + registry.warmup() + + def fail_compile(**kwargs): + pytest.fail(f"Producer compiled after startup warmup: {kwargs.get('fn')}") + + monkeypatch.setattr(triton.knobs.runtime, "jit_post_compile_hook", fail_compile) + rows = (0, 1, 2, 7, 17, 33, 65, 127, 128, 129, 255, 256, 257, 1023, 4097, 32768) + for input_offset, scale_offset in product((0, 1), repeat=2): + scale = torch.full((2,), 17.125, dtype=torch.float32, device="cuda") + module.down_proj.input_global_scale_inv = scale[scale_offset : scale_offset + 1] + for m in rows: + x = torch.randn( + m * width + input_offset, dtype=torch.bfloat16, device="cuda" + )[input_offset:].view(m, width) + actual = maybe_fused_act_quant(module.act_fn, x, module.down_proj) + assert isinstance(actual, QuantizedActivation) + if m == 129: + expected = ops.scaled_fp4_quant( + module.act_fn.forward_native(x), + module.down_proj.input_global_scale_inv, + ) + assert torch.equal(actual.data, expected[0]) + assert torch.equal( + actual.scale.view(torch.uint8), expected[1].view(torch.uint8) + ) + if m == 129 and input_offset == scale_offset == 1: + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + maybe_fused_act_quant(module.act_fn, x, module.down_proj) + graph.replay() + + +@pytest.mark.parametrize( + "width", + [16, 48, 80, 1344, 5376, 16 * ((1 << 23) - 1), 16 << 23, 16 * ((1 << 23) + 1)], +) +def test_relu2_nvfp4_row_bound_keeps_index_products_i32(width: int) -> None: + from vllm.model_executor.layers.fusion.relu2_nvfp4_quant import ( + _relu_squared_nvfp4_max_rows, + _relu_squared_nvfp4_quant_launch, + _relu_squared_nvfp4_warmup_rows, + ) + from vllm.utils.math_utils import round_up + + max_rows = _relu_squared_nvfp4_max_rows(width) + rows = {0, 1, 127, 128, 129, max_rows, max_rows + 1, max(0, max_rows - 1)} + assert len(_relu_squared_nvfp4_warmup_rows(width, (1 << 63))) == min(128, max_rows) + for m in rows: + groups = m * (width // 16) + scale_size = round_up(m, 128) * round_up(width // 16, 4) + assert (m <= max_rows) == (groups <= 1 << 30 and scale_size <= 1 << 30) + if 0 < m <= max_rows: + grid, args = _relu_squared_nvfp4_quant_launch(m, width) + assert grid[0] * args["GROUPS_PER_CTA"] <= 1 << 31 + assert grid[0] * args["PAD_PER_CTA"] <= 1 << 31 + assert all( + args[name] <= 1 << 30 + for name in ("M", "NUM_GROUPS", "PADDED_SCALE_SIZE") + ) diff --git a/vllm/model_executor/kernels/linear/nvfp4/flashinfer.py b/vllm/model_executor/kernels/linear/nvfp4/flashinfer.py index 2a69a3be2100..90eb8478e3bc 100644 --- a/vllm/model_executor/kernels/linear/nvfp4/flashinfer.py +++ b/vllm/model_executor/kernels/linear/nvfp4/flashinfer.py @@ -190,6 +190,12 @@ def input_quant_key(self) -> QuantKey | None: convention, pre-quantized blockscales must use the swizzled layout.""" return kNvfp4Dynamic + def input_quant_activation_types(self) -> tuple[type[torch.nn.Module], ...]: + from vllm.model_executor.layers.activation import SiluAndMul + + # Preserve the existing producer; ReLU2 is validated only with CuTe-DSL. + return (SiluAndMul,) + @classmethod def is_supported( cls, compute_capability: int | None = None diff --git a/vllm/model_executor/layers/fusion/fused_act_quant.py b/vllm/model_executor/layers/fusion/fused_act_quant.py index 6261a1bdd019..dba4285e9150 100644 --- a/vllm/model_executor/layers/fusion/fused_act_quant.py +++ b/vllm/model_executor/layers/fusion/fused_act_quant.py @@ -14,17 +14,24 @@ """ from collections.abc import Callable +from types import FunctionType, MethodType import torch +from vllm.config import get_current_vllm_config_or_none from vllm.model_executor.layers.activation import ReLUSquaredActivation, SiluAndMul from vllm.model_executor.layers.fusion.quant_activation import ( QuantizedActivation, - get_input_quant_key, + get_fused_act_quant_key, ) from vllm.model_executor.layers.fusion.relu2_fp8_quant import ( relu_squared_static_fp8_quant, ) +from vllm.model_executor.layers.fusion.relu2_nvfp4_quant import ( + _relu_squared_nvfp4_max_rows, + _relu_squared_nvfp4_quant_kernel, + relu_squared_nvfp4_quant, +) from vllm.model_executor.layers.linear import LinearBase from vllm.model_executor.layers.quantization.utils.quant_utils import ( QuantKey, @@ -170,6 +177,46 @@ def _relu_squared_static_fp8_quant_supported( ) +def _uses_native_relu_squared(act_fn: torch.nn.Module) -> bool: + if not isinstance(act_fn, ReLUSquaredActivation): + return False + forward = act_fn._forward_method + if type(forward) is FunctionType: + try: + # Dynamo handles this direct attribute access when unwrapping a + # compiled bound method; getattr(..., default) loses that handling. + forward = forward._torchdynamo_inline # type: ignore[attr-defined] + except AttributeError: + return False + return ( + type(forward) is MethodType + and forward.__func__ is ReLUSquaredActivation.forward_native + ) + + +def _relu_squared_nvfp4_quant_supported( + act_fn: torch.nn.Module, + x: torch.Tensor, + linear: LinearBase, +) -> bool: + scale = getattr(linear, "input_global_scale_inv", None) + return ( + # CUDA ReLU2 suppresses NaNs; this producer matches native semantics. + _uses_native_relu_squared(act_fn) + and x.is_cuda + and x.dtype == torch.bfloat16 + and x.ndim == 2 + and x.shape[-1] > 0 + and x.shape[-1] % 16 == 0 + and x.shape[0] <= _relu_squared_nvfp4_max_rows(x.shape[-1]) + and x.is_contiguous() + and isinstance(scale, torch.Tensor) + and scale.dtype == torch.float32 + and scale.device == x.device + and scale.numel() == 1 + ) + + # Optional per-entry predicates preserve the unconditional fallback contract # for producers whose supported inputs are narrower than their registry key. _FUSED_ACT_QUANT_SUPPORT: dict[ @@ -192,6 +239,55 @@ def _relu_squared_static_fp8_quant_supported( if current_platform.is_cuda() and hasattr(torch.ops._C, "silu_and_mul_nvfp4_quant"): _FUSED_ACT_QUANT[(SiluAndMul, kNvfp4Dynamic)] = _silu_and_mul_nvfp4_dynamic +if current_platform.is_cuda() and current_platform.is_device_capability_family(100): + _RELU2_NVFP4_KEY = (ReLUSquaredActivation, kNvfp4Dynamic) + _FUSED_ACT_QUANT[_RELU2_NVFP4_KEY] = relu_squared_nvfp4_quant + _FUSED_ACT_QUANT_SUPPORT[_RELU2_NVFP4_KEY] = _relu_squared_nvfp4_quant_supported + + +def _relu_squared_nvfp4_warmup_width(module: torch.nn.Module) -> int | None: + act_fn, linear = module.act_fn, module.down_proj + key = get_fused_act_quant_key(linear, act_fn) + if ( + key is None + or _FUSED_ACT_QUANT.get((type(act_fn), key)) is not relu_squared_nvfp4_quant + ): + return None + k = getattr(linear, "input_size_per_partition", None) + scale = getattr(linear, "input_global_scale_inv", None) + if ( + _uses_native_relu_squared(act_fn) + and getattr(linear, "params_dtype", None) == torch.bfloat16 + and isinstance(k, int) + and k > 0 + and k % 16 == 0 + and isinstance(scale, torch.Tensor) + and scale.is_cuda + and scale.dtype == torch.float32 + and scale.numel() == 1 + ): + return k + return None + + +def register_relu_squared_nvfp4_quant_warmup(module: torch.nn.Module) -> None: + """Defer a Nemotron MLP's producer selection until startup kernel warmup.""" + config = get_current_vllm_config_or_none() + if config is None or not config.kernel_config.enable_jit_warmup: + return + if ( + _FUSED_ACT_QUANT.get((ReLUSquaredActivation, kNvfp4Dynamic)) + is not relu_squared_nvfp4_quant + ): + return + max_tokens = max( + config.scheduler_config.max_num_batched_tokens, + config.compilation_config.max_cudagraph_capture_size or 0, + ) + _relu_squared_nvfp4_quant_kernel.register_warmup( + module=module, max_tokens=max_tokens + ) + def maybe_fused_act_quant( act_fn: torch.nn.Module, @@ -203,7 +299,7 @@ def maybe_fused_act_quant( Returns a QuantizedActivation when a fused kernel matches the activation and the consumer's effective input quantization key, else the plain activation. """ - key = get_input_quant_key(linear) + key = get_fused_act_quant_key(linear, act_fn) if key is not None: registry_key = (type(act_fn), key) producer = _FUSED_ACT_QUANT.get(registry_key) diff --git a/vllm/model_executor/layers/fusion/quant_activation.py b/vllm/model_executor/layers/fusion/quant_activation.py index c7dbe840db1d..a057d9ac1eb5 100644 --- a/vllm/model_executor/layers/fusion/quant_activation.py +++ b/vllm/model_executor/layers/fusion/quant_activation.py @@ -4,9 +4,9 @@ and consumed directly by a linear layer, letting the layer skip its own input quantization. A linear advertises the key its kernel can consume via expose_input_quant_key; the kernel validates and reads the activation via -as_quantized_activation. Producers query the effective capability through -get_input_quant_key, which hides the key when another consumer branch needs the -original activation. +as_quantized_activation. Activation producers query get_fused_act_quant_key; +get_input_quant_key exposes the underlying consumer capability. Both hide the +key when another consumer branch needs the original activation. """ from dataclasses import dataclass, replace @@ -47,28 +47,41 @@ def weak_ref(self) -> "QuantizedActivation": def expose_input_quant_key(layer: torch.nn.Module, kernel) -> None: - """Store the kernel's pre-quantized input key on the layer, if any. + """Store the kernel's input key and optional automatic-fusion policy. This is the bridge from a kernel's input_quant_key() to the layer capability - that fusion call sites read through get_input_quant_key. The raw key is left - unset when the kernel quantizes its own input. + that fusion call sites read through get_fused_act_quant_key. Re-exposure replaces + both fields, including clearing stale capabilities for unsupported kernels. + Kernels without an activation policy retain unrestricted producer selection. TODO(mgoin): Producers also need the consumer's quantization scales (e.g. static input scale, global scale). Expose those here as well so producers do not reach into kernel-specific layer attributes. """ key = kernel.input_quant_key() - if key is not None: - layer._input_quant_key = key + get_types = getattr(kernel, "input_quant_activation_types", None) + types = get_types() if key is not None and get_types is not None else None + layer._input_quant_key = key + layer._input_quant_activation_types = tuple(types) if types is not None else None def get_input_quant_key(layer: torch.nn.Module) -> QuantKey | None: - """Return the pre-quantized input key the complete layer can consume.""" + """Return the consumer key, without selecting an activation producer.""" if getattr(layer, "requires_unquantized_input", False): return None return getattr(layer, "_input_quant_key", None) +def get_fused_act_quant_key( + layer: torch.nn.Module, act_fn: torch.nn.Module +) -> QuantKey | None: + """Return the key only for an allowed automatic activation producer.""" + allowed = getattr(layer, "_input_quant_activation_types", None) + if allowed is not None and type(act_fn) not in allowed: + return None + return get_input_quant_key(layer) + + def as_quantized_activation( x: "torch.Tensor | QuantizedActivation", expected_key: QuantKey | None ) -> "QuantizedActivation | None": diff --git a/vllm/model_executor/layers/fusion/relu2_nvfp4_quant.py b/vllm/model_executor/layers/fusion/relu2_nvfp4_quant.py new file mode 100644 index 000000000000..ccf1074291af --- /dev/null +++ b/vllm/model_executor/layers/fusion/relu2_nvfp4_quant.py @@ -0,0 +1,269 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from __future__ import annotations + +from collections.abc import Iterator +from itertools import product +from typing import TYPE_CHECKING, Any + +import torch + +from vllm.model_executor.warmup.jit_warmup_triton_helper import ( + TritonWarmupTensor, + triton_kernel_dispatcher_with_warmup, +) +from vllm.triton_utils import tl, triton +from vllm.utils.math_utils import round_up +from vllm.utils.torch_utils import direct_register_custom_op + +if TYPE_CHECKING: + from vllm.model_executor.layers.fusion.quant_activation import QuantizedActivation + from vllm.model_executor.layers.linear import LinearBase + + +def _relu_squared_nvfp4_quant_launch( + m: int, k: int +) -> tuple[tuple[int], dict[str, Any]]: + scale_cols = round_up(k // 16, 4) + scale_size = round_up(m, 128) * scale_cols + groups = m * (k // 16) + ctas = max(triton.cdiv(groups, 128), min(triton.cdiv(scale_size, 512), 128)) + return (ctas,), dict( + M=m, + K=k, + NUM_GROUPS=groups, + PADDED_SCALE_COLS=scale_cols, + PADDED_SCALE_SIZE=scale_size, + GROUPS_PER_CTA=128, + PAD_PER_CTA=triton.next_power_of_2(triton.cdiv(scale_size, ctas)), + num_warps=4, + enable_fp_fusion=False, + ) + + +def _relu_squared_nvfp4_max_rows(k: int) -> int: + # Keep groups and padded scale size <= 2^30. The power-of-two padding + # tile then keeps every pid * PAD_PER_CTA offset within signed int32. + groups_per_row = k // 16 + scale_cols = round_up(groups_per_row, 4) + return min((1 << 30) // groups_per_row, (1 << 30) // scale_cols // 128 * 128) + + +def _relu_squared_nvfp4_warmup_rows(k: int, max_tokens: int) -> range: + # All padding-tile buckets occur in the first 128 rows. For M > 128, + # CTAs = ceil(M * (K/16) / 128), so the padding/CTA ratio is between + # its M=128 and M=64 values. The row bound keeps all runtime scalars i32. + return range(1, min(128, max_tokens, _relu_squared_nvfp4_max_rows(k)) + 1) + + +def _relu_squared_nvfp4_warmup_inputs( + module: torch.nn.Module, max_tokens: int +) -> Iterator[dict[str, Any]]: + # Resolve after weight processing and LoRA replacement, not at construction. + from vllm.model_executor.layers.fusion.fused_act_quant import ( + _relu_squared_nvfp4_warmup_width, + ) + + k = _relu_squared_nvfp4_warmup_width(module) + if k is None: + return + for m in _relu_squared_nvfp4_warmup_rows(k, max_tokens): + grid, kwargs = _relu_squared_nvfp4_quant_launch(m, k) + for input_aligned, scale_aligned in product((True, False), repeat=2): + yield dict( + grid=grid, + x_ptr=TritonWarmupTensor(torch.bfloat16, aligned=input_aligned), + global_scale_ptr=TritonWarmupTensor( + torch.float32, aligned=scale_aligned + ), + output_ptr=TritonWarmupTensor(torch.uint8), + scale_ptr=TritonWarmupTensor(torch.float8_e4m3fn), + **kwargs, + ) + + +@triton.jit +def _rcp_ftz(x): + return tl.inline_asm_elementwise( + "rcp.approx.ftz.f32 $0, $1;", + constraints="=f,f", + args=[x], + dtype=tl.float32, + is_pure=True, + pack=1, + ) + + +@triton.jit +def _max_ignore_nan(a, b): + return tl.maximum(a, b, propagate_nan=tl.PropagateNan.NONE) + + +@triton.jit +def _e4m3_satfinite(x): + packed = tl.inline_asm_elementwise( + "{ .reg .b16 v; cvt.rn.satfinite.e4m3x2.f32 v, 0f00000000, $1; " + "cvt.u32.u16 $0, v; }", + constraints="=r,f", + args=[x], + dtype=tl.int32, + is_pure=True, + pack=1, + ) + return packed.to(tl.uint8).to(tl.float8e4nv, bitcast=True) + + +@triton_kernel_dispatcher_with_warmup(warmup_inputs=_relu_squared_nvfp4_warmup_inputs) +@triton.jit(do_not_specialize=["M", "NUM_GROUPS", "PADDED_SCALE_SIZE"]) +def _relu_squared_nvfp4_quant_kernel( + x_ptr, + global_scale_ptr, + output_ptr, + scale_ptr, + M, + K: tl.constexpr, + NUM_GROUPS, + PADDED_SCALE_COLS: tl.constexpr, + PADDED_SCALE_SIZE, + GROUPS_PER_CTA: tl.constexpr, + PAD_PER_CTA: tl.constexpr, +): + pid = tl.program_id(0) + groups = pid * GROUPS_PER_CTA + tl.arange(0, GROUPS_PER_CTA) + row = groups // (K // 16) + col = groups % (K // 16) + if pid * GROUPS_PER_CTA < NUM_GROUPS: + pair = tl.arange(0, 8) + index = row[:, None].to(tl.int64) * K + col[:, None] * 16 + pair[None, :] * 2 + valid = groups[:, None] < NUM_GROUPS + values = tl.load( + x_ptr + + row[:, None].to(tl.int64) * K + + col[:, None] * 16 + + tl.arange(0, 16)[None, :], + valid, + other=0, + ) + lo, hi = tl.split(tl.reshape(values, (GROUPS_PER_CTA, 8, 2))) + lo, hi = lo.to(tl.float32), hi.to(tl.float32) + # ReLU propagates NaN; tl.maximum otherwise defaults to maxnum semantics. + lo = tl.maximum(lo, 0.0, propagate_nan=tl.PropagateNan.ALL) + hi = tl.maximum(hi, 0.0, propagate_nan=tl.PropagateNan.ALL) + # Preserve the materialized BF16 ReLU2 boundary before computing amax. + lo = (lo * lo).to(tl.bfloat16).to(tl.float32) + hi = (hi * hi).to(tl.bfloat16).to(tl.float32) + # Match __hmax2 / __hmax in the CUDA quantizer: ignore individual NaNs. + amax = tl.reduce(_max_ignore_nan(lo, hi), 1, _max_ignore_nan) + global_scale = tl.load(global_scale_ptr) + scale = _e4m3_satfinite(global_scale * (amax * _rcp_ftz(6.0))) + scale_f32 = scale.to(tl.float32) + inverse = tl.where( + scale_f32 != 0.0, + _rcp_ftz(scale_f32 * _rcp_ftz(global_scale)), + 0.0, + ) + lo = lo * inverse[:, None] + hi = hi * inverse[:, None] + packed = tl.inline_asm_elementwise( + "{ .reg .b8 v; cvt.rn.satfinite.e2m1x2.f32 v, $2, $1; cvt.u32.u8 $0, v; }", + constraints="=r,f,f", + args=[lo, hi], + dtype=tl.int32, + is_pure=True, + pack=1, + ).to(tl.uint8) + tl.store(output_ptr + index // 2, packed, valid) + scale_index = ( + ((row.to(tl.int64) // 128) * (PADDED_SCALE_COLS // 4) + col // 4) * 512 + + (row % 32) * 16 + + ((row % 128) // 32) * 4 + + col % 4 + ) + tl.store(scale_ptr + scale_index, scale, groups < NUM_GROUPS) + + # Initialize only padding; these stores never overlap the live scale stores. + # Distribute padding across the same launch, including small-M cases. + offsets = pid * PAD_PER_CTA + tl.arange(0, PAD_PER_CTA) + tile = offsets // 512 + padded_row = ( + (tile // (PADDED_SCALE_COLS // 4)) * 128 + + (offsets % 16 // 4) * 32 + + offsets % 512 // 16 + ) + padded_col = (tile % (PADDED_SCALE_COLS // 4)) * 4 + offsets % 4 + is_padding = (padded_row >= M) | (padded_col >= K // 16) + tl.store(scale_ptr + offsets, 0.0, (offsets < PADDED_SCALE_SIZE) & is_padding) + + +def relu_squared_nvfp4_quant_out( + x: torch.Tensor, + global_scale: torch.Tensor, + output: torch.Tensor, + block_scale: torch.Tensor, +) -> None: + """Write BF16 ReLU2 + NVFP4 into distinct, preallocated output buffers. + + Requires M <= _relu_squared_nvfp4_max_rows(K). + """ + m, k = x.shape + # vLLM can reuse compiled graphs without checking Dynamo's shape guards. + assert m <= _relu_squared_nvfp4_max_rows(k), "NVFP4 row limit exceeded" + if m == 0: + return + grid, kwargs = _relu_squared_nvfp4_quant_launch(m, k) + _relu_squared_nvfp4_quant_kernel[grid]( + x, + global_scale, + output, + block_scale, + **kwargs, + ) + + +def _relu_squared_nvfp4_quant_fake( + x: torch.Tensor, global_scale: torch.Tensor +) -> tuple[torch.Tensor, torch.Tensor]: + m, k = x.shape + output = torch.empty((m, k // 2), dtype=torch.uint8, device=x.device) + block_scale = torch.empty( + (round_up(m, 128), round_up(k // 16, 4)), + dtype=torch.float8_e4m3fn, + device=x.device, + ) + return output, block_scale + + +def _relu_squared_nvfp4_quant( + x: torch.Tensor, global_scale: torch.Tensor +) -> tuple[torch.Tensor, torch.Tensor]: + output, block_scale = _relu_squared_nvfp4_quant_fake(x, global_scale) + relu_squared_nvfp4_quant_out(x, global_scale, output, block_scale) + return output, block_scale + + +# Keep the M-dependent launch geometry outside Dynamo: vLLM drops shape guards. +direct_register_custom_op( + op_name="relu_squared_nvfp4_quant", + op_func=_relu_squared_nvfp4_quant, + fake_impl=_relu_squared_nvfp4_quant_fake, +) + + +def relu_squared_nvfp4_quant( + x: torch.Tensor, linear: LinearBase +) -> QuantizedActivation: + """Fuse BF16 ReLU2 and block16 NVFP4 with 128x4 swizzled scales.""" + from vllm.model_executor.layers.fusion.quant_activation import QuantizedActivation + from vllm.model_executor.layers.quantization.utils.quant_utils import kNvfp4Dynamic + + output, block_scale = torch.ops.vllm.relu_squared_nvfp4_quant( + x, linear.input_global_scale_inv + ) + return QuantizedActivation( + data=output, + scale=block_scale, + orig_dtype=x.dtype, + orig_shape=x.shape, + quant_key=kNvfp4Dynamic, + ) diff --git a/vllm/model_executor/models/nemotron_h.py b/vllm/model_executor/models/nemotron_h.py index e477c65302b6..47a13f84aa58 100644 --- a/vllm/model_executor/models/nemotron_h.py +++ b/vllm/model_executor/models/nemotron_h.py @@ -42,7 +42,10 @@ activation_without_mul, fused_moe_make_expert_params_mapping, ) -from vllm.model_executor.layers.fusion.fused_act_quant import maybe_fused_act_quant +from vllm.model_executor.layers.fusion.fused_act_quant import ( + maybe_fused_act_quant, + register_relu_squared_nvfp4_quant_warmup, +) from vllm.model_executor.layers.layernorm import RMSNorm from vllm.model_executor.layers.linear import ( ColumnParallelLinear, @@ -125,6 +128,7 @@ def __init__( prefix=f"{prefix}.down_proj", ) self.act_fn = get_act_fn(hidden_act) + register_relu_squared_nvfp4_quant_warmup(self) def forward(self, x: torch.Tensor): x, _ = self.up_proj(x)