iHore/vLLM
8
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

[Fusion] Support 3D activation quantization fusion and enable it for Pixtral (#59612)

Browse Source 
Signed-off-by: Jakub Zakrzewski <jzakrzewski@nvidia.com>
Co-authored-by: Codex <noreply@openai.com>
Co-authored-by: Misha Goin <mgoin64@gmail.com>
This commit is contained in:
Jakub Zakrzewski
2026-10-09 15:12:03 +08:00
committed by GitHub
co-authored by Codex Misha Goin
parent dd30786fef
commit 19c0bed584
6 changed files with 101 additions and 29 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/fusion/test_quant_activation_contract.py
+2
View File
@@ -21,6 +21,7 @@ from vllm.model_executor.kernels.linear.nvfp4.base import (
NvFp4LinearLayerConfig,
)
from vllm.model_executor.kernels.linear.nvfp4.flashinfer import (
FlashInferCuteDslNvFp4LinearKernel,
FlashInferCutlassNvFp4LinearKernel,
FlashInferTrtllmNvFp4LinearKernel,
)
@@ -59,6 +60,7 @@ from vllm.platforms import current_platform
SUPPORTING = {
CutlassFP8ScaledMMLinearKernel,
FlashInferFP8ScaledMMLinearKernel,
FlashInferCuteDslNvFp4LinearKernel,
FlashInferCutlassNvFp4LinearKernel,
PerTensorTorchFP8ScaledMMLinearKernel,
AiterHipbMMPerTokenFp8ScaledMMLinearKernel,
tests/kernels/quantization/test_flashinfer_nvfp4_scaled_mm.py
+52
View File
@@ -13,11 +13,15 @@ from nvfp4_utils import (
from vllm import _custom_ops as ops
from vllm.model_executor.kernels.linear.nvfp4 import NvFp4LinearLayerConfig
from vllm.model_executor.kernels.linear.nvfp4.flashinfer import (
FlashInferCuteDslNvFp4LinearKernel,
FlashInferCuteDslNvFp4W4A16LinearKernel,
FlashInferCutlassNvFp4LinearKernel,
)
from vllm.model_executor.layers.fusion.quant_activation import QuantizedActivation
from vllm.model_executor.layers.quantization.utils.nvfp4_emulation_utils import (
dequantize_to_dtype,
)
from vllm.model_executor.layers.quantization.utils.quant_utils import kNvfp4Dynamic
from vllm.platforms import current_platform
from vllm.utils.flashinfer import (
flashinfer_scaled_fp4_mm,
@@ -47,6 +51,54 @@ SEEDS = [42]
CUDA_DEVICES = ["cuda:0"]
@pytest.mark.parametrize("token_shape", [(17,), (1, 17), (2, 17)])
@pytest.mark.parametrize(
"kernel_cls",
[FlashInferCutlassNvFp4LinearKernel, FlashInferCuteDslNvFp4LinearKernel],
)
@torch.inference_mode()
def test_flashinfer_nvfp4_preserves_quantized_input_shape(
token_shape: tuple[int, ...],
kernel_cls: type[FlashInferCutlassNvFp4LinearKernel]
| type[FlashInferCuteDslNvFp4LinearKernel],
) -> None:
"""Pre-quantized inputs must match the tensor path and retain token dimensions."""
supported, reason = kernel_cls.is_supported()
if not supported:
pytest.skip(reason)
set_random_seed(42)
hidden_size, output_size = 128, 128
x = torch.randn(*token_shape, hidden_size, dtype=torch.bfloat16, device="cuda")
weight = torch.randn(output_size, hidden_size, dtype=x.dtype, device=x.device)
input_scale = (FLOAT8_E4M3_MAX * FLOAT4_E2M1_MAX / x.abs().max()).float()
weight_scale = (FLOAT8_E4M3_MAX * FLOAT4_E2M1_MAX / weight.abs().max()).float()
x_fp4, x_blockscale = ops.scaled_fp4_quant(x.reshape(-1, hidden_size), input_scale)
weight_fp4, weight_blockscale = ops.scaled_fp4_quant(weight, weight_scale)
quantized_input = QuantizedActivation(
data=x_fp4.reshape(*token_shape, hidden_size // 2),
scale=x_blockscale,
orig_dtype=x.dtype,
orig_shape=x.shape,
quant_key=kNvfp4Dynamic,
)
layer = torch.nn.Module()
layer.input_size_per_partition = hidden_size
layer.output_size_per_partition = output_size
layer.weight = weight_fp4
layer.weight_scale = weight_blockscale
layer.input_global_scale_inv = input_scale
layer.alpha = (input_scale * weight_scale).reciprocal()
kernel = kernel_cls(NvFp4LinearLayerConfig())
expected = kernel.apply_weights(layer, x.reshape(-1, hidden_size))
actual = kernel.apply_weights(layer, quantized_input)
assert actual.shape == (*token_shape, output_size)
assert actual.dtype == x.dtype
torch.testing.assert_close(actual.reshape(expected.shape), expected)
def get_ref_results(
a_fp4,
b_fp4,
tests/kernels/test_fused_quant_activation.py
+19 -13
View File
@@ -115,13 +115,13 @@ class MockLinearRequiresUnquantized(MockLinearFp8Static):
requires_unquantized_input = True
@pytest.mark.parametrize("num_tokens", [1, 16, 128])
@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)
@torch.inference_mode()
def test_maybe_fused_act_quant_fp8_static(
default_vllm_config,
num_tokens: int,
token_shape: tuple[int, ...],
hidden_size: int,
dtype: torch.dtype,
) -> None:
@@ -133,26 +133,29 @@ def test_maybe_fused_act_quant_fp8_static(
scale = torch.tensor([0.5], device=device, dtype=torch.float32)
linear = MockLinearFp8Static(scale)
x = torch.randn(num_tokens, hidden_size * 2, dtype=dtype, device=device)
x = torch.randn(*token_shape, hidden_size * 2, dtype=dtype, device=device)
result = maybe_fused_act_quant(act_fn, x, linear)
assert isinstance(result, QuantizedActivation)
assert result.quant_key == kFp8StaticTensorSym
assert result.data.dtype == current_platform.fp8_dtype()
assert result.orig_dtype == dtype
assert result.orig_shape == (num_tokens, hidden_size)
assert result.orig_shape == (*token_shape, hidden_size)
assert result.data.shape == result.orig_shape
ref_out = ref_impl(act_fn, x, scale)
ref_out = ref_impl(act_fn, x.reshape(-1, hidden_size * 2), scale).reshape(
result.orig_shape
)
torch.testing.assert_close(result.data.to(torch.float32), ref_out.to(torch.float32))
@pytest.mark.parametrize("num_tokens", [1, 16, 128])
@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)
@torch.inference_mode()
def test_maybe_fused_act_quant_fp8_dynamic_block(
default_vllm_config,
num_tokens: int,
token_shape: tuple[int, ...],
hidden_size: int,
dtype: torch.dtype,
) -> None:
@@ -166,31 +169,34 @@ def test_maybe_fused_act_quant_fp8_dynamic_block(
linear = MockLinearFp8Dynamic128()
scale = 1 / hidden_size
x = torch.randn(num_tokens, hidden_size * 2, dtype=dtype, device=device) * scale
x = torch.randn(*token_shape, hidden_size * 2, dtype=dtype, device=device) * scale
result = maybe_fused_act_quant(act_fn, x, linear)
assert isinstance(result, QuantizedActivation)
assert result.quant_key == kFp8Dynamic128Sym
assert result.data.dtype == current_platform.fp8_dtype()
assert result.orig_dtype == dtype
assert result.orig_shape == (num_tokens, hidden_size)
assert result.orig_shape == (*token_shape, hidden_size)
assert result.data.shape == result.orig_shape
num_groups = hidden_size // group_size
assert result.scale.shape == (num_tokens, num_groups)
assert result.scale.shape == (*token_shape, num_groups)
gate, up = x.split(hidden_size, dim=-1)
silu_out = F.silu(gate) * up
ref_out, ref_scales = per_token_group_quant_fp8(
silu_out, group_size=group_size, use_ue8m0=False
silu_out.reshape(-1, hidden_size), group_size=group_size, use_ue8m0=False
)
ref_out = ref_out.reshape(result.orig_shape)
ref_scales = ref_scales.reshape(*token_shape, num_groups)
torch.testing.assert_close(result.scale, ref_scales, rtol=1e-5, atol=1e-5)
ref_deq = ref_out.to(torch.float32) * ref_scales.repeat_interleave(
group_size, dim=1
group_size, dim=-1
)
result_deq = result.data.to(torch.float32) * result.scale.repeat_interleave(
group_size, dim=1
group_size, dim=-1
)
torch.testing.assert_close(ref_deq, result_deq, atol=5e-2, rtol=5e-2)
vllm/model_executor/kernels/linear/nvfp4/flashinfer.py
+21 -10
View File
@@ -112,6 +112,9 @@ class FlashInferCuteDslNvFp4W4A16LinearKernel(NvFp4LinearKernel):
class FlashInferCuteDslNvFp4LinearKernel(NvFp4LinearKernel):
"""NVFP4 GEMM via FlashInfer's cutedsl backend."""
def input_quant_key(self) -> QuantKey | None:
return kNvfp4Dynamic
@classmethod
def is_supported(
cls, compute_capability: int | None = None
@@ -137,19 +140,26 @@ class FlashInferCuteDslNvFp4LinearKernel(NvFp4LinearKernel):
def apply_weights(
self,
layer: torch.nn.Module,
x: torch.Tensor,
x: torch.Tensor | QuantizedActivation,
bias: torch.Tensor | None = None,
) -> torch.Tensor:
output_size = layer.output_size_per_partition
output_dtype = x.dtype
output_shape = [*x.shape[:-1], output_size]
x_fp4, x_blockscale = scaled_fp4_quant(
x,
layer.input_global_scale_inv,
is_sf_swizzled_layout=True,
backend="flashinfer-cutedsl",
)
qa = as_quantized_activation(x, self.input_quant_key())
if qa is not None:
x_fp4, x_blockscale = qa.data, qa.scale
x_fp4 = x_fp4.reshape(-1, x_fp4.shape[-1])
output_dtype = qa.orig_dtype
output_shape = [*qa.orig_shape[:-1], output_size]
else:
assert isinstance(x, torch.Tensor)
output_dtype = x.dtype
output_shape = [*x.shape[:-1], output_size]
x_fp4, x_blockscale = scaled_fp4_quant(
x,
layer.input_global_scale_inv,
is_sf_swizzled_layout=True,
backend="flashinfer-cutedsl",
)
x_fp4 = pad_nvfp4_activation_for_cutlass(
x_fp4, nvfp4_weight_padding_bytes(layer)
@@ -219,6 +229,7 @@ class FlashInferCutlassNvFp4LinearKernel(NvFp4LinearKernel):
qa = as_quantized_activation(x, self.input_quant_key())
if qa is not None:
x_fp4, x_blockscale = qa.data, qa.scale
x_fp4 = x_fp4.reshape(-1, x_fp4.shape[-1])
x_fp4 = pad_nvfp4_activation_for_cutlass(x_fp4, weights_padding_bytes)
output_dtype = qa.orig_dtype
output_shape = [*qa.orig_shape[:-1], output_size]
vllm/model_executor/layers/fusion/fused_act_quant.py
+4 -4
View File
@@ -63,10 +63,10 @@ def _silu_and_mul_fp8_dynamic_block(
x: torch.Tensor, linear: LinearBase, group_size: int, quant_key: QuantKey
) -> QuantizedActivation:
"""SiluAndMul + FP8 dynamic per-block quantization."""
assert x.ndim == 2, f"Input must be 2D [batch, hidden*2], got {x.shape}"
assert x.ndim in (2, 3), f"Input must be 2D or 3D, got {x.shape}"
d = x.shape[-1] // 2
out_shape = x.shape[:-1] + (d,)
x = x.reshape(-1, x.shape[-1])
num_tokens = x.shape[0]
num_groups = d // group_size
@@ -102,10 +102,10 @@ def _silu_and_mul_nvfp4_dynamic(
x: torch.Tensor, linear: LinearBase
) -> QuantizedActivation:
"""SiluAndMul + NVFP4 dynamic quantization."""
assert x.ndim == 2, f"Input must be 2D [batch, hidden*2], got {x.shape}"
assert x.ndim in (2, 3), f"Input must be 2D or 3D, got {x.shape}"
d = x.shape[-1] // 2
out_shape = x.shape[:-1] + (d,)
x = x.reshape(-1, x.shape[-1])
num_tokens = x.shape[0]
# NVFP4 packs 2 values into 1 byte
vllm/model_executor/models/pixtral.py
+3 -2
View File
@@ -32,6 +32,7 @@ from vllm.inputs import MultiModalDataDict
from vllm.model_executor.layers.activation import SiluAndMul, get_act_and_mul_fn
from vllm.model_executor.layers.attention import MMEncoderAttention
from vllm.model_executor.layers.conv import Conv2dLayer
from vllm.model_executor.layers.fusion.fused_act_quant import maybe_fused_act_quant
from vllm.model_executor.layers.layernorm import RMSNorm
from vllm.model_executor.layers.linear import (
MergedColumnParallelLinear,
@@ -758,7 +759,7 @@ class FeedForward(nn.Module):
def forward(self, x):
x, _ = self.gate_up_proj(x)
x = self.act_fn(x)
x = maybe_fused_act_quant(self.act_fn, x, self.down_proj)
x, _ = self.down_proj(x)
return x
@@ -1259,7 +1260,7 @@ class PixtralHFMLP(nn.Module):
def forward(self, x: torch.Tensor) -> torch.Tensor:
gate_up, _ = self.gate_up_proj(x)
x = self.act_and_mul(gate_up)
x = maybe_fused_act_quant(self.act_and_mul, gate_up, self.down_proj)
x, _ = self.down_proj(x)
return x