[Kernel] Fuse native BF16 ReLU2 with NVFP4 quantization

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()),

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,