[Fusion] Manual ActivationQuantFusionPass initial application (#51415)
Signed-off-by: mgoin <mgoin64@gmail.com> Signed-off-by: ElizaWszola <ewszola@redhat.com> Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Co-authored-by: ElizaWszola <ewszola@redhat.com>
This commit is contained in:
Misha Goin
co-authored by
Claude Opus 4.8
ElizaWszola
Claude Opus 4.8
ElizaWszola
parent
d4d703caf9
commit
b7624069ff
@@ -43,6 +43,45 @@ is_blackwell = lambda: current_platform.is_device_capability_family(100)
|
||||
"""Are we running on Blackwell, a lot of tests depend on it"""
|
||||
|
||||
|
||||
def nvfp4_kernel_exposes_input_quant_key() -> bool:
|
||||
"""Check if the NVFP4 kernel selected on this platform exposes input_quant_key.
|
||||
|
||||
FlashInferCuteDslNvFp4LinearKernel does not expose input_quant_key() due to
|
||||
layout incompatibility between the manual fusion kernel (silu_and_mul_nvfp4_quant)
|
||||
output format and the cutedsl backend's expected input format.
|
||||
|
||||
FlashInferCutlassNvFp4LinearKernel does expose input_quant_key() and supports
|
||||
manual fusion.
|
||||
"""
|
||||
if not current_platform.is_cuda():
|
||||
return False
|
||||
|
||||
try:
|
||||
from vllm.model_executor.kernels.linear.nvfp4.flashinfer import (
|
||||
FlashInferCuteDslNvFp4LinearKernel,
|
||||
)
|
||||
|
||||
is_supported, _ = FlashInferCuteDslNvFp4LinearKernel.is_supported()
|
||||
if is_supported:
|
||||
return False
|
||||
except ImportError:
|
||||
pass
|
||||
|
||||
try:
|
||||
from vllm.model_executor.kernels.linear.nvfp4.flashinfer import (
|
||||
FlashInferCutlassNvFp4LinearKernel,
|
||||
)
|
||||
|
||||
is_supported, _ = FlashInferCutlassNvFp4LinearKernel.is_supported()
|
||||
if is_supported:
|
||||
return True
|
||||
except ImportError:
|
||||
pass
|
||||
|
||||
# Fallback: assume no manual fusion support
|
||||
return False
|
||||
|
||||
|
||||
def custom_ops_combos(*custom_ops: str) -> Iterable[str]:
|
||||
"""Generate all combinations of custom ops for parametrization."""
|
||||
custom_ops_lists = [[f"-{op}", f"+{op}"] for op in custom_ops]
|
||||
|
||||
@@ -14,6 +14,7 @@ from .common import (
|
||||
Matches,
|
||||
custom_ops_combos,
|
||||
is_blackwell,
|
||||
nvfp4_kernel_exposes_input_quant_key,
|
||||
)
|
||||
from .models import (
|
||||
FLASHINFER_ATTN,
|
||||
@@ -77,7 +78,6 @@ def test_tp1_fp8_fusions(
|
||||
inductor_graph_partition: bool,
|
||||
use_deepgemm: bool,
|
||||
run_e2e_fusion_test,
|
||||
monkeypatch,
|
||||
):
|
||||
if use_deepgemm and not current_platform.is_cuda():
|
||||
pytest.skip("DeepGemm only supported on CUDA")
|
||||
@@ -109,7 +109,7 @@ def test_tp1_fp8_fusions(
|
||||
custom_ops=custom_ops.split(","),
|
||||
pass_config=PassConfig(
|
||||
fuse_norm_quant=True,
|
||||
fuse_act_quant=True,
|
||||
fuse_act_quant=False,
|
||||
fuse_attn_quant=True,
|
||||
enable_qk_norm_rope_fusion=True,
|
||||
),
|
||||
@@ -119,7 +119,6 @@ def test_tp1_fp8_fusions(
|
||||
|
||||
matches_check = [
|
||||
"rms_quant_fusion",
|
||||
"act_quant_fusion",
|
||||
"norm_rope_fusion",
|
||||
"attn_quant_fusion",
|
||||
]
|
||||
@@ -168,6 +167,12 @@ def test_tp1_fp4_fusions(
|
||||
inductor_graph_partition: bool,
|
||||
run_e2e_fusion_test,
|
||||
):
|
||||
if nvfp4_kernel_exposes_input_quant_key():
|
||||
pytest.skip(
|
||||
"NVFP4 kernel exposes input_quant_key; manual fusion fires "
|
||||
"instead of compiler pass-based fusion"
|
||||
)
|
||||
|
||||
matches = matches_fn(n_layers)
|
||||
|
||||
# Reduce size of model and skip weight loading time
|
||||
@@ -181,13 +186,13 @@ def test_tp1_fp4_fusions(
|
||||
custom_ops=custom_ops.split(","),
|
||||
pass_config=PassConfig(
|
||||
fuse_norm_quant=True,
|
||||
fuse_act_quant=True,
|
||||
fuse_act_quant=False,
|
||||
fuse_attn_quant=True,
|
||||
enable_qk_norm_rope_fusion=True,
|
||||
),
|
||||
)
|
||||
|
||||
matches_check = ["act_quant_fusion", "attn_quant_fusion", "norm_rope_fusion"]
|
||||
matches_check = ["attn_quant_fusion", "norm_rope_fusion"]
|
||||
|
||||
run_e2e_fusion_test(
|
||||
model_name,
|
||||
|
||||
@@ -14,6 +14,7 @@ from .common import (
|
||||
Matches,
|
||||
custom_ops_combos,
|
||||
is_blackwell,
|
||||
nvfp4_kernel_exposes_input_quant_key,
|
||||
)
|
||||
from .models import (
|
||||
FLASHINFER_ATTN,
|
||||
@@ -68,7 +69,6 @@ def test_tp2_ar_rms_fp8_fusions(
|
||||
custom_ops: str,
|
||||
inductor_graph_partition: bool,
|
||||
run_e2e_fusion_test,
|
||||
monkeypatch,
|
||||
):
|
||||
matches = matches_fn(n_layers)
|
||||
|
||||
@@ -89,7 +89,7 @@ def test_tp2_ar_rms_fp8_fusions(
|
||||
custom_ops=custom_ops.split(","),
|
||||
pass_config=PassConfig(
|
||||
fuse_norm_quant=True,
|
||||
fuse_act_quant=True,
|
||||
fuse_act_quant=False,
|
||||
fuse_attn_quant=True,
|
||||
enable_qk_norm_rope_fusion=True,
|
||||
fuse_allreduce_rms=True,
|
||||
@@ -98,7 +98,6 @@ def test_tp2_ar_rms_fp8_fusions(
|
||||
|
||||
matches_check = [
|
||||
"rms_quant_fusion",
|
||||
"act_quant_fusion",
|
||||
"norm_rope_fusion",
|
||||
"attn_quant_fusion",
|
||||
"ar_rms_fusion",
|
||||
@@ -139,8 +138,13 @@ def test_tp2_ar_rms_fp4_fusions(
|
||||
custom_ops: str,
|
||||
inductor_graph_partition: bool,
|
||||
run_e2e_fusion_test,
|
||||
monkeypatch,
|
||||
):
|
||||
if nvfp4_kernel_exposes_input_quant_key():
|
||||
pytest.skip(
|
||||
"NVFP4 kernel exposes input_quant_key; manual fusion fires "
|
||||
"instead of compiler pass-based fusion"
|
||||
)
|
||||
|
||||
matches = matches_fn(n_layers)
|
||||
|
||||
# Reduce size of model and skip weight loading time
|
||||
@@ -154,14 +158,13 @@ def test_tp2_ar_rms_fp4_fusions(
|
||||
use_inductor_graph_partition=inductor_graph_partition,
|
||||
custom_ops=custom_ops.split(","),
|
||||
pass_config=PassConfig(
|
||||
fuse_act_quant=True,
|
||||
fuse_act_quant=False,
|
||||
fuse_attn_quant=True,
|
||||
fuse_allreduce_rms=True,
|
||||
),
|
||||
)
|
||||
|
||||
matches_check = [
|
||||
"act_quant_fusion",
|
||||
"attn_quant_fusion",
|
||||
"ar_rms_fusion",
|
||||
]
|
||||
|
||||
@@ -14,6 +14,7 @@ from .common import (
|
||||
Matches,
|
||||
custom_ops_combos,
|
||||
is_blackwell,
|
||||
nvfp4_kernel_exposes_input_quant_key,
|
||||
)
|
||||
from .models import (
|
||||
FLASHINFER_ATTN,
|
||||
@@ -61,7 +62,7 @@ def test_tp2_async_tp_fp8_fusions(
|
||||
custom_ops=custom_ops.split(","),
|
||||
pass_config=PassConfig(
|
||||
fuse_norm_quant=True,
|
||||
fuse_act_quant=True,
|
||||
fuse_act_quant=False,
|
||||
fuse_attn_quant=True,
|
||||
enable_qk_norm_rope_fusion=True,
|
||||
enable_sp=True,
|
||||
@@ -74,7 +75,6 @@ def test_tp2_async_tp_fp8_fusions(
|
||||
|
||||
matches_check = [
|
||||
"rms_quant_fusion",
|
||||
"act_quant_fusion",
|
||||
"norm_rope_fusion",
|
||||
"attn_quant_fusion",
|
||||
"sequence_parallel",
|
||||
@@ -114,6 +114,12 @@ def test_tp2_async_tp_nvfp4_fusions(
|
||||
inductor_graph_partition: bool,
|
||||
run_e2e_fusion_test,
|
||||
):
|
||||
if nvfp4_kernel_exposes_input_quant_key():
|
||||
pytest.skip(
|
||||
"NVFP4 kernel exposes input_quant_key; manual fusion fires "
|
||||
"instead of compiler pass-based fusion"
|
||||
)
|
||||
|
||||
# NVFP4 currently wires the all-gather + GEMM path only.
|
||||
matches = matches_fn(n_layers)._replace(async_tp=n_layers * 2)
|
||||
|
||||
@@ -127,7 +133,7 @@ def test_tp2_async_tp_nvfp4_fusions(
|
||||
use_inductor_graph_partition=inductor_graph_partition,
|
||||
custom_ops=custom_ops.split(","),
|
||||
pass_config=PassConfig(
|
||||
fuse_act_quant=True,
|
||||
fuse_act_quant=False,
|
||||
fuse_attn_quant=True,
|
||||
enable_sp=True,
|
||||
fuse_gemm_comms=True,
|
||||
@@ -138,7 +144,6 @@ def test_tp2_async_tp_nvfp4_fusions(
|
||||
)
|
||||
|
||||
matches_check = [
|
||||
"act_quant_fusion",
|
||||
"attn_quant_fusion",
|
||||
"sequence_parallel",
|
||||
"async_tp",
|
||||
@@ -245,7 +250,7 @@ def test_tp2_sp_ar_rms_fp8_fusions(
|
||||
custom_ops=custom_ops.split(","),
|
||||
pass_config=PassConfig(
|
||||
fuse_norm_quant=True,
|
||||
fuse_act_quant=True,
|
||||
fuse_act_quant=False,
|
||||
fuse_attn_quant=True,
|
||||
enable_qk_norm_rope_fusion=True,
|
||||
enable_sp=True,
|
||||
@@ -258,7 +263,6 @@ def test_tp2_sp_ar_rms_fp8_fusions(
|
||||
|
||||
matches_check = [
|
||||
"rms_quant_fusion",
|
||||
"act_quant_fusion",
|
||||
"norm_rope_fusion",
|
||||
"attn_quant_fusion",
|
||||
"ar_rms_fusion",
|
||||
|
||||
@@ -0,0 +1,348 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""
|
||||
Tests for manual fusion via maybe_fused_act_quant.
|
||||
|
||||
Tests all fusion paths in _FUSED_ACT_QUANT:
|
||||
- kFp8StaticTensorSym: all platforms
|
||||
- kFp8Dynamic128Sym: CUDA only
|
||||
- kNvfp4Dynamic: CUDA SM100+ only
|
||||
"""
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
import vllm.envs as envs
|
||||
from tests.utils import TestFP8Layer
|
||||
from vllm.config import (
|
||||
CompilationConfig,
|
||||
VllmConfig,
|
||||
set_current_vllm_config,
|
||||
)
|
||||
from vllm.model_executor.kernels.linear import (
|
||||
CutlassFP8ScaledMMLinearKernel,
|
||||
FlashInferFP8ScaledMMLinearKernel,
|
||||
FP8ScaledMMLinearKernel,
|
||||
PerTensorTorchFP8ScaledMMLinearKernel,
|
||||
ROCmFP8ScaledMMLinearKernel,
|
||||
)
|
||||
from vllm.model_executor.layers.activation import SiluAndMul
|
||||
from vllm.model_executor.layers.fusion.fused_act_quant import (
|
||||
_FUSED_ACT_QUANT,
|
||||
maybe_fused_act_quant,
|
||||
)
|
||||
from vllm.model_executor.layers.fusion.quant_activation import (
|
||||
QuantizedActivation,
|
||||
expose_input_quant_key,
|
||||
)
|
||||
from vllm.model_executor.layers.quantization.utils.quant_utils import (
|
||||
kFp8Dynamic128Sym,
|
||||
kFp8StaticTensorSym,
|
||||
kNvfp4Dynamic,
|
||||
)
|
||||
from vllm.platforms import current_platform
|
||||
|
||||
|
||||
# Mock linear layer for testing fusion paths that don't have real kernel support
|
||||
class MockLinearForFusion(torch.nn.Module):
|
||||
"""Mock linear layer that exposes input_quant_key for fusion testing."""
|
||||
|
||||
def __init__(self, quant_key, input_scale=None, input_global_scale=None):
|
||||
super().__init__()
|
||||
self.input_quant_key = quant_key
|
||||
if input_scale is not None:
|
||||
self.input_scale = input_scale
|
||||
if input_global_scale is not None:
|
||||
self.input_global_scale = input_global_scale
|
||||
|
||||
|
||||
ROCM_KERNELS = [ROCmFP8ScaledMMLinearKernel, PerTensorTorchFP8ScaledMMLinearKernel]
|
||||
CUDA_KERNELS = [
|
||||
FlashInferFP8ScaledMMLinearKernel,
|
||||
CutlassFP8ScaledMMLinearKernel,
|
||||
PerTensorTorchFP8ScaledMMLinearKernel,
|
||||
]
|
||||
TEST_KERNELS = ROCM_KERNELS if current_platform.is_rocm() else CUDA_KERNELS
|
||||
|
||||
|
||||
@pytest.mark.parametrize("num_tokens", [32, 64])
|
||||
@pytest.mark.parametrize("hidden_size", [128, 256])
|
||||
@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16])
|
||||
@pytest.mark.parametrize("force_kernel", TEST_KERNELS)
|
||||
@pytest.mark.skipif(
|
||||
envs.VLLM_TARGET_DEVICE not in ["cuda", "rocm"], reason="Only test on CUDA and ROCm"
|
||||
)
|
||||
def test_manual_fusion_fp8_static_with_linear(
|
||||
num_tokens: int,
|
||||
hidden_size: int,
|
||||
dtype: torch.dtype,
|
||||
force_kernel: FP8ScaledMMLinearKernel,
|
||||
):
|
||||
"""Test manual fusion with real FP8 linear layer (kFp8StaticTensorSym).
|
||||
|
||||
This is an end-to-end test that verifies the full flow:
|
||||
unfused (silu_and_mul -> in-kernel quant) vs fused (silu_and_mul_quant).
|
||||
"""
|
||||
torch.set_default_device("cuda")
|
||||
torch.set_default_dtype(dtype)
|
||||
|
||||
x = torch.rand(num_tokens, hidden_size * 2)
|
||||
|
||||
config = VllmConfig(
|
||||
compilation_config=CompilationConfig(custom_ops=["none"]),
|
||||
)
|
||||
|
||||
with set_current_vllm_config(config):
|
||||
silu_and_mul = SiluAndMul()
|
||||
fp8_linear = TestFP8Layer(
|
||||
weight_shape=(hidden_size, hidden_size),
|
||||
activation_quant_key=kFp8StaticTensorSym,
|
||||
weight_quant_key=kFp8StaticTensorSym,
|
||||
force_kernel=force_kernel,
|
||||
input_dtype=dtype,
|
||||
)
|
||||
|
||||
# Run without fusion: silu_and_mul returns plain tensor
|
||||
y_unfused = maybe_fused_act_quant(silu_and_mul, x, fp8_linear)
|
||||
assert isinstance(y_unfused, torch.Tensor)
|
||||
result_unfused = fp8_linear(y_unfused)
|
||||
|
||||
# Enable fusion
|
||||
expose_input_quant_key(fp8_linear, fp8_linear.kernel)
|
||||
|
||||
if not hasattr(fp8_linear, "input_quant_key"):
|
||||
pytest.skip(
|
||||
f"Kernel {force_kernel.__name__} doesn't support input_quant_key"
|
||||
)
|
||||
|
||||
# Run with fusion: silu_and_mul returns QuantizedActivation
|
||||
y_fused = maybe_fused_act_quant(silu_and_mul, x, fp8_linear)
|
||||
assert isinstance(y_fused, QuantizedActivation)
|
||||
assert y_fused.quant_key == kFp8StaticTensorSym
|
||||
assert y_fused.data.dtype == current_platform.fp8_dtype()
|
||||
assert y_fused.data.shape == (num_tokens, hidden_size)
|
||||
assert y_fused.orig_dtype == dtype
|
||||
assert y_fused.orig_shape == (num_tokens, hidden_size)
|
||||
|
||||
result_fused = fp8_linear(y_fused)
|
||||
|
||||
torch.testing.assert_close(
|
||||
result_fused.to(dtype=dtype),
|
||||
result_unfused.to(dtype=dtype),
|
||||
atol=5e-2,
|
||||
rtol=5e-2,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16])
|
||||
@pytest.mark.skipif(
|
||||
not current_platform.is_cuda_alike(), reason="Dynamic block quant CUDA only"
|
||||
)
|
||||
@pytest.mark.skipif(
|
||||
envs.VLLM_TARGET_DEVICE not in ["cuda", "rocm"], reason="Only test on CUDA and ROCm"
|
||||
)
|
||||
def test_manual_fusion_fp8_dynamic_128(dtype: torch.dtype):
|
||||
"""Test kFp8Dynamic128Sym fusion path (group_size=128).
|
||||
|
||||
Compares fused (silu_and_mul_per_block_quant) vs unfused (silu_and_mul)
|
||||
by dequantizing the fused result and comparing with unfused.
|
||||
"""
|
||||
if (SiluAndMul, kFp8Dynamic128Sym) not in _FUSED_ACT_QUANT:
|
||||
pytest.skip("kFp8Dynamic128Sym fusion not available")
|
||||
|
||||
torch.set_default_device("cuda")
|
||||
torch.set_default_dtype(dtype)
|
||||
|
||||
# hidden_size must be divisible by group_size (128)
|
||||
num_tokens, hidden_size = 32, 256
|
||||
group_size = 128
|
||||
x = torch.rand(num_tokens, hidden_size * 2)
|
||||
|
||||
config = VllmConfig(
|
||||
compilation_config=CompilationConfig(custom_ops=["none"]),
|
||||
)
|
||||
|
||||
with set_current_vllm_config(config):
|
||||
silu_and_mul = SiluAndMul()
|
||||
|
||||
# Unfused path: just apply silu_and_mul
|
||||
mock_linear_no_key = torch.nn.Linear(hidden_size, hidden_size)
|
||||
result_unfused = maybe_fused_act_quant(silu_and_mul, x, mock_linear_no_key)
|
||||
assert isinstance(result_unfused, torch.Tensor)
|
||||
|
||||
# Fused path: apply silu_and_mul + per-block quantization
|
||||
mock_linear_with_key = MockLinearForFusion(kFp8Dynamic128Sym)
|
||||
result_fused = maybe_fused_act_quant(silu_and_mul, x, mock_linear_with_key)
|
||||
|
||||
# Verify fused result structure
|
||||
assert isinstance(result_fused, QuantizedActivation)
|
||||
assert result_fused.quant_key == kFp8Dynamic128Sym
|
||||
assert result_fused.data.dtype == current_platform.fp8_dtype()
|
||||
assert result_fused.data.shape == (num_tokens, hidden_size)
|
||||
assert result_fused.orig_dtype == dtype
|
||||
|
||||
# Check scale shape
|
||||
expected_num_groups = hidden_size // group_size
|
||||
assert result_fused.scale.shape == (num_tokens, expected_num_groups)
|
||||
|
||||
# Dequantize fused result and compare with unfused
|
||||
# Per-block dequant: data * scale (broadcast scale across group)
|
||||
dequant_data = result_fused.data.to(dtype).view(
|
||||
num_tokens, expected_num_groups, group_size
|
||||
)
|
||||
scales_expanded = result_fused.scale.unsqueeze(
|
||||
-1
|
||||
) # (num_tokens, num_groups, 1)
|
||||
dequant_result = (
|
||||
(dequant_data * scales_expanded).view(num_tokens, hidden_size).to(dtype)
|
||||
)
|
||||
|
||||
torch.testing.assert_close(
|
||||
dequant_result,
|
||||
result_unfused,
|
||||
atol=5e-2,
|
||||
rtol=5e-2,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16])
|
||||
@pytest.mark.skipif(not current_platform.is_cuda(), reason="NVFP4 CUDA only")
|
||||
@pytest.mark.skipif(
|
||||
not current_platform.has_device_capability(100), reason="NVFP4 requires SM100+"
|
||||
)
|
||||
@pytest.mark.skipif(
|
||||
envs.VLLM_TARGET_DEVICE not in ["cuda", "rocm"], reason="Only test on CUDA and ROCm"
|
||||
)
|
||||
def test_manual_fusion_nvfp4_dynamic(dtype: torch.dtype):
|
||||
"""Test kNvfp4Dynamic fusion path.
|
||||
|
||||
Compares fused (silu_and_mul_nvfp4_quant) vs unfused (silu_and_mul)
|
||||
by dequantizing the fused result and comparing with unfused.
|
||||
"""
|
||||
if (SiluAndMul, kNvfp4Dynamic) not in _FUSED_ACT_QUANT:
|
||||
pytest.skip("kNvfp4Dynamic fusion not available")
|
||||
|
||||
from tests.kernels.quantization.nvfp4_utils import dequantize_nvfp4_to_dtype
|
||||
|
||||
torch.set_default_device("cuda")
|
||||
torch.set_default_dtype(dtype)
|
||||
|
||||
# NVFP4 requires hidden_size divisible by 16 (block size) and by 2 (packing).
|
||||
# M (num_tokens) must be >= 128 for the 128x4 swizzled scale layout.
|
||||
num_tokens, hidden_size = 128, 256
|
||||
x = torch.rand(num_tokens, hidden_size * 2)
|
||||
input_global_scale = torch.tensor([1.0], dtype=torch.float32, device="cuda")
|
||||
|
||||
config = VllmConfig(
|
||||
compilation_config=CompilationConfig(custom_ops=["none"]),
|
||||
)
|
||||
|
||||
with set_current_vllm_config(config):
|
||||
silu_and_mul = SiluAndMul()
|
||||
|
||||
# Unfused path: just apply silu_and_mul
|
||||
mock_linear_no_key = torch.nn.Linear(hidden_size, hidden_size)
|
||||
result_unfused = maybe_fused_act_quant(silu_and_mul, x, mock_linear_no_key)
|
||||
assert isinstance(result_unfused, torch.Tensor)
|
||||
|
||||
# Fused path: apply silu_and_mul + NVFP4 quantization
|
||||
mock_linear_with_key = MockLinearForFusion(
|
||||
kNvfp4Dynamic, input_global_scale=input_global_scale
|
||||
)
|
||||
result_fused = maybe_fused_act_quant(silu_and_mul, x, mock_linear_with_key)
|
||||
|
||||
# Verify fused result structure
|
||||
assert isinstance(result_fused, QuantizedActivation)
|
||||
assert result_fused.quant_key == kNvfp4Dynamic
|
||||
# NVFP4 packs 2 values into 1 byte
|
||||
assert result_fused.data.dtype == torch.uint8
|
||||
assert result_fused.data.shape == (num_tokens, hidden_size // 2)
|
||||
assert result_fused.orig_dtype == dtype
|
||||
assert result_fused.orig_shape == (num_tokens, hidden_size)
|
||||
|
||||
# Dequantize fused result and compare with unfused
|
||||
dequant_result = dequantize_nvfp4_to_dtype(
|
||||
tensor_fp4=result_fused.data,
|
||||
tensor_sf=result_fused.scale,
|
||||
global_scale=input_global_scale,
|
||||
dtype=dtype,
|
||||
device="cuda",
|
||||
block_size=16,
|
||||
is_sf_128x4_layout=True,
|
||||
)
|
||||
|
||||
torch.testing.assert_close(
|
||||
dequant_result,
|
||||
result_unfused,
|
||||
atol=3e-1,
|
||||
rtol=3e-1,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16])
|
||||
@pytest.mark.skipif(
|
||||
envs.VLLM_TARGET_DEVICE not in ["cuda", "rocm"], reason="Only test on CUDA and ROCm"
|
||||
)
|
||||
def test_manual_fusion_fallback_no_key(dtype: torch.dtype):
|
||||
"""Test that maybe_fused_act_quant falls back when no input_quant_key."""
|
||||
torch.set_default_device("cuda")
|
||||
torch.set_default_dtype(dtype)
|
||||
|
||||
x = torch.rand(32, 256)
|
||||
|
||||
config = VllmConfig(
|
||||
compilation_config=CompilationConfig(custom_ops=["none"]),
|
||||
)
|
||||
|
||||
with set_current_vllm_config(config):
|
||||
silu_and_mul = SiluAndMul()
|
||||
# Linear without input_quant_key attribute
|
||||
mock_linear = torch.nn.Linear(128, 128)
|
||||
|
||||
result = maybe_fused_act_quant(silu_and_mul, x, mock_linear)
|
||||
|
||||
# Should fall back to plain silu_and_mul
|
||||
assert isinstance(result, torch.Tensor)
|
||||
assert result.shape == (32, 128)
|
||||
assert result.dtype == dtype
|
||||
|
||||
|
||||
@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16])
|
||||
@pytest.mark.skipif(
|
||||
envs.VLLM_TARGET_DEVICE not in ["cuda", "rocm"], reason="Only test on CUDA and ROCm"
|
||||
)
|
||||
def test_manual_fusion_fallback_unsupported_key(dtype: torch.dtype):
|
||||
"""Test that maybe_fused_act_quant falls back for unsupported quant keys."""
|
||||
from vllm.model_executor.layers.quantization.utils.quant_utils import (
|
||||
GroupShape,
|
||||
QuantKey,
|
||||
ScaleDesc,
|
||||
)
|
||||
|
||||
torch.set_default_device("cuda")
|
||||
torch.set_default_dtype(dtype)
|
||||
|
||||
x = torch.rand(32, 256)
|
||||
|
||||
config = VllmConfig(
|
||||
compilation_config=CompilationConfig(custom_ops=["none"]),
|
||||
)
|
||||
|
||||
with set_current_vllm_config(config):
|
||||
# Create an unsupported quant key
|
||||
unsupported_key = QuantKey(
|
||||
dtype=torch.int8,
|
||||
scale=ScaleDesc(
|
||||
dtype=torch.float32, static=False, group_shape=GroupShape(1, 1)
|
||||
),
|
||||
)
|
||||
|
||||
silu_and_mul = SiluAndMul()
|
||||
mock_linear = MockLinearForFusion(unsupported_key)
|
||||
|
||||
result = maybe_fused_act_quant(silu_and_mul, x, mock_linear)
|
||||
|
||||
# Should fall back to plain silu_and_mul since key not in _FUSED_ACT_QUANT
|
||||
assert isinstance(result, torch.Tensor)
|
||||
assert result.shape == (32, 128)
|
||||
assert result.dtype == dtype
|
||||
@@ -25,6 +25,9 @@ from vllm.model_executor.kernels.linear.scaled_mm.cutlass import (
|
||||
from vllm.model_executor.kernels.linear.scaled_mm.flashinfer import (
|
||||
FlashInferFP8ScaledMMLinearKernel,
|
||||
)
|
||||
from vllm.model_executor.kernels.linear.scaled_mm.pytorch import (
|
||||
PerTensorTorchFP8ScaledMMLinearKernel,
|
||||
)
|
||||
from vllm.model_executor.kernels.linear.scaled_mm.ScaledMMLinearKernel import (
|
||||
FP8ScaledMMLinearLayerConfig,
|
||||
Int8ScaledMMLinearKernel,
|
||||
@@ -46,6 +49,7 @@ SUPPORTING = {
|
||||
CutlassFP8ScaledMMLinearKernel,
|
||||
FlashInferFP8ScaledMMLinearKernel,
|
||||
FlashInferCutlassNvFp4LinearKernel,
|
||||
PerTensorTorchFP8ScaledMMLinearKernel,
|
||||
}
|
||||
|
||||
|
||||
|
||||
@@ -2,10 +2,20 @@
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
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.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 (
|
||||
per_token_group_quant_fp8,
|
||||
)
|
||||
from vllm.model_executor.layers.quantization.utils.quant_utils import (
|
||||
kFp8Dynamic128Sym,
|
||||
kFp8StaticTensorSym,
|
||||
)
|
||||
from vllm.platforms import current_platform
|
||||
from vllm.utils.torch_utils import set_random_seed
|
||||
|
||||
@@ -69,3 +79,140 @@ def test_silu_and_mul(
|
||||
ref_out.to(dtype=torch.float32), ops_out.to(dtype=torch.float32)
|
||||
)
|
||||
opcheck(torch.ops._C.silu_and_mul_quant, (ops_out, x, scale))
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Tests for maybe_fused_act_quant interface
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class MockLinearFp8Static(torch.nn.Module):
|
||||
"""Mock linear layer advertising kFp8StaticTensorSym."""
|
||||
|
||||
def __init__(self, input_scale: torch.Tensor):
|
||||
super().__init__()
|
||||
self.input_quant_key = kFp8StaticTensorSym
|
||||
self.input_scale = input_scale
|
||||
|
||||
|
||||
class MockLinearFp8Dynamic128(torch.nn.Module):
|
||||
"""Mock linear layer advertising kFp8Dynamic128Sym."""
|
||||
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
self.input_quant_key = kFp8Dynamic128Sym
|
||||
|
||||
|
||||
class MockLinearNoQuant(torch.nn.Module):
|
||||
"""Mock linear layer with no input_quant_key (no fusion)."""
|
||||
|
||||
pass
|
||||
|
||||
|
||||
@pytest.mark.parametrize("num_tokens", [1, 16, 128])
|
||||
@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,
|
||||
hidden_size: int,
|
||||
dtype: torch.dtype,
|
||||
) -> None:
|
||||
"""Test maybe_fused_act_quant with FP8 static per-tensor quantization."""
|
||||
device = "cuda:0"
|
||||
torch.set_default_device(device)
|
||||
|
||||
act_fn = SiluAndMul()
|
||||
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)
|
||||
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)
|
||||
|
||||
ref_out = ref_impl(act_fn, x, scale)
|
||||
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("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,
|
||||
hidden_size: int,
|
||||
dtype: torch.dtype,
|
||||
) -> None:
|
||||
"""Test maybe_fused_act_quant with FP8 dynamic per-block quantization."""
|
||||
group_size = 128 # We only support 128 for now
|
||||
|
||||
device = "cuda:0"
|
||||
torch.set_default_device(device)
|
||||
|
||||
act_fn = SiluAndMul()
|
||||
linear = MockLinearFp8Dynamic128()
|
||||
|
||||
scale = 1 / hidden_size
|
||||
x = torch.randn(num_tokens, 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)
|
||||
|
||||
num_groups = hidden_size // group_size
|
||||
assert result.scale.shape == (num_tokens, 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
|
||||
)
|
||||
|
||||
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
|
||||
)
|
||||
result_deq = result.data.to(torch.float32) * result.scale.repeat_interleave(
|
||||
group_size, dim=1
|
||||
)
|
||||
torch.testing.assert_close(ref_deq, result_deq, atol=5e-2, rtol=5e-2)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("num_tokens", [1, 16, 128])
|
||||
@pytest.mark.parametrize("hidden_size", [128, 512])
|
||||
@pytest.mark.parametrize("dtype", DTYPES)
|
||||
@torch.inference_mode()
|
||||
def test_maybe_fused_act_quant_fallback(
|
||||
default_vllm_config,
|
||||
num_tokens: int,
|
||||
hidden_size: int,
|
||||
dtype: torch.dtype,
|
||||
) -> None:
|
||||
"""Test maybe_fused_act_quant falls back when no input_quant_key."""
|
||||
device = "cuda:0"
|
||||
torch.set_default_device(device)
|
||||
|
||||
act_fn = SiluAndMul()
|
||||
linear = MockLinearNoQuant()
|
||||
x = torch.randn(num_tokens, hidden_size * 2, dtype=dtype, device=device)
|
||||
|
||||
result = maybe_fused_act_quant(act_fn, x, linear)
|
||||
|
||||
assert isinstance(result, torch.Tensor)
|
||||
assert not isinstance(result, QuantizedActivation)
|
||||
assert result.dtype == dtype
|
||||
assert result.shape == (num_tokens, hidden_size)
|
||||
|
||||
ref_out = act_fn(x)
|
||||
torch.testing.assert_close(result, ref_out)
|
||||
|
||||
@@ -7,7 +7,11 @@ import torch
|
||||
|
||||
from vllm.config import CompilationMode, get_current_vllm_config
|
||||
from vllm.model_executor.layers.quantization.input_quant_fp8 import QuantFP8
|
||||
from vllm.model_executor.layers.quantization.utils.quant_utils import GroupShape
|
||||
from vllm.model_executor.layers.quantization.utils.quant_utils import (
|
||||
GroupShape,
|
||||
QuantKey,
|
||||
kFp8StaticTensorSym,
|
||||
)
|
||||
from vllm.platforms import current_platform
|
||||
|
||||
from .BlockScaledMMLinearKernel import Fp8BlockScaledMMLinearKernel
|
||||
@@ -88,6 +92,11 @@ class PerTensorTorchFP8ScaledMMLinearKernel(TorchFP8ScaledMMLinearKernel):
|
||||
return False, "requires per tensor activation and weight scales."
|
||||
return True, None
|
||||
|
||||
def input_quant_key(self) -> QuantKey | None:
|
||||
if self.config.activation_quant_key == kFp8StaticTensorSym:
|
||||
return kFp8StaticTensorSym
|
||||
return None
|
||||
|
||||
def apply_scaled_mm(
|
||||
self,
|
||||
*,
|
||||
|
||||
@@ -0,0 +1,161 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""
|
||||
Producer side of the QuantizedActivation contract for activation layers.
|
||||
|
||||
Given an activation module and the downstream linear it feeds, fuse the
|
||||
activation with that linear's input quantization into a single kernel when the
|
||||
linear advertises a consumable input_quant_key (see quant_activation.py).
|
||||
Falls back to the plain activation when nothing matches, so a model forward can
|
||||
always call maybe_fused_act_quant unconditionally.
|
||||
|
||||
This is the manual-fusion counterpart to ActivationQuantFusionPass: when fusion
|
||||
fires here the silu_and_mul pattern is already consumed, so the compiler pass
|
||||
finds nothing to rewrite and the two never double-fuse.
|
||||
"""
|
||||
|
||||
from collections.abc import Callable
|
||||
|
||||
import torch
|
||||
|
||||
from vllm.model_executor.layers.activation import SiluAndMul
|
||||
from vllm.model_executor.layers.fusion.quant_activation import QuantizedActivation
|
||||
from vllm.model_executor.layers.linear import LinearBase
|
||||
from vllm.model_executor.layers.quantization.utils.quant_utils import (
|
||||
QuantKey,
|
||||
kFp8Dynamic128Sym,
|
||||
kFp8StaticTensorSym,
|
||||
kNvfp4Dynamic,
|
||||
)
|
||||
from vllm.platforms import current_platform
|
||||
|
||||
FP8_DTYPE = current_platform.fp8_dtype()
|
||||
FP4_DTYPE = torch.uint8
|
||||
|
||||
|
||||
def _silu_and_mul_fp8_static(
|
||||
x: torch.Tensor, linear: LinearBase
|
||||
) -> QuantizedActivation:
|
||||
"""SiluAndMul + FP8 static per-tensor quantization."""
|
||||
d = x.shape[-1] // 2
|
||||
out_shape = x.shape[:-1] + (d,)
|
||||
result = torch.empty(out_shape, dtype=FP8_DTYPE, device=x.device)
|
||||
# TODO(mgoin): read the consumer scale via the contract instead of reaching
|
||||
# into the kernel-specific input_scale attribute.
|
||||
scale = linear.input_scale
|
||||
torch.ops._C.silu_and_mul_quant(result, x, scale)
|
||||
return QuantizedActivation(
|
||||
data=result,
|
||||
scale=scale,
|
||||
orig_dtype=x.dtype,
|
||||
orig_shape=out_shape,
|
||||
quant_key=kFp8StaticTensorSym,
|
||||
)
|
||||
|
||||
|
||||
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}"
|
||||
|
||||
d = x.shape[-1] // 2
|
||||
out_shape = x.shape[:-1] + (d,)
|
||||
num_tokens = x.shape[0]
|
||||
num_groups = d // group_size
|
||||
|
||||
result = torch.empty((num_tokens, d), dtype=FP8_DTYPE, device=x.device)
|
||||
scales = torch.empty((num_tokens, num_groups), dtype=torch.float32, device=x.device)
|
||||
|
||||
torch.ops._C.silu_and_mul_per_block_quant(
|
||||
out=result,
|
||||
input=x,
|
||||
scales=scales,
|
||||
group_size=group_size,
|
||||
scale_ub=None,
|
||||
is_scale_transposed=False,
|
||||
)
|
||||
|
||||
return QuantizedActivation(
|
||||
data=result.view(out_shape),
|
||||
scale=scales.view(out_shape[:-1] + (num_groups,)),
|
||||
orig_dtype=x.dtype,
|
||||
orig_shape=out_shape,
|
||||
quant_key=quant_key,
|
||||
)
|
||||
|
||||
|
||||
def _silu_and_mul_fp8_dynamic_128(
|
||||
x: torch.Tensor, linear: LinearBase
|
||||
) -> QuantizedActivation:
|
||||
"""SiluAndMul + FP8 dynamic per-group (group=128) quantization."""
|
||||
return _silu_and_mul_fp8_dynamic_block(x, linear, 128, kFp8Dynamic128Sym)
|
||||
|
||||
|
||||
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}"
|
||||
|
||||
d = x.shape[-1] // 2
|
||||
out_shape = x.shape[:-1] + (d,)
|
||||
num_tokens = x.shape[0]
|
||||
|
||||
# NVFP4 packs 2 values into 1 byte
|
||||
result = torch.empty((num_tokens, d // 2), dtype=FP4_DTYPE, device=x.device)
|
||||
|
||||
# Block scale output shape: swizzled layout for tensor cores
|
||||
# Each group of 16 elements shares one FP8 scale
|
||||
num_k_tiles = (d + 63) // 64
|
||||
block_scale = torch.empty(
|
||||
(num_tokens, num_k_tiles * 4), dtype=FP8_DTYPE, device=x.device
|
||||
)
|
||||
|
||||
input_global_scale = getattr(linear, "input_global_scale", None)
|
||||
assert input_global_scale is not None, (
|
||||
"input_global_scale is required for NVFP4 quantization"
|
||||
)
|
||||
|
||||
torch.ops._C.silu_and_mul_nvfp4_quant(result, block_scale, x, input_global_scale)
|
||||
|
||||
return QuantizedActivation(
|
||||
data=result.view(out_shape[:-1] + (d // 2,)),
|
||||
scale=block_scale,
|
||||
orig_dtype=x.dtype,
|
||||
orig_shape=out_shape,
|
||||
quant_key=kNvfp4Dynamic,
|
||||
)
|
||||
|
||||
|
||||
# (activation module type, consumer input_quant_key) -> fused producer.
|
||||
# Mirrors ActivationQuantFusionPass.FUSED_OPS; add a row to migrate a scheme.
|
||||
_FUSED_ACT_QUANT: dict[tuple[type, QuantKey], Callable] = {
|
||||
(SiluAndMul, kFp8StaticTensorSym): _silu_and_mul_fp8_static,
|
||||
}
|
||||
|
||||
# Add CUDA-specific entries for dynamic block quantization
|
||||
if current_platform.is_cuda_alike():
|
||||
_FUSED_ACT_QUANT[(SiluAndMul, kFp8Dynamic128Sym)] = _silu_and_mul_fp8_dynamic_128
|
||||
|
||||
# Add NVFP4 if supported (requires SM100+)
|
||||
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
|
||||
|
||||
|
||||
def maybe_fused_act_quant(
|
||||
act_fn: torch.nn.Module,
|
||||
x: torch.Tensor,
|
||||
linear: LinearBase,
|
||||
) -> "torch.Tensor | QuantizedActivation":
|
||||
"""Apply act_fn, fusing the downstream linear's input quant when possible.
|
||||
|
||||
Returns a QuantizedActivation when a fused kernel matches
|
||||
(act_fn, linear.input_quant_key), else the plain activated tensor.
|
||||
"""
|
||||
key = getattr(linear, "input_quant_key", None)
|
||||
if key is not None:
|
||||
producer = _FUSED_ACT_QUANT.get((type(act_fn), key))
|
||||
if producer is not None:
|
||||
return producer(x, linear)
|
||||
return act_fn(x)
|
||||
@@ -39,6 +39,7 @@ from vllm.model_executor.layers.attention import (
|
||||
Attention,
|
||||
EncoderOnlyAttention,
|
||||
)
|
||||
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,
|
||||
@@ -114,7 +115,7 @@ class LlamaMLP(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
|
||||
|
||||
|
||||
Reference in New Issue
Block a user