A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 5, 09:45 UTC
[W8A8 Block Linear Refactor][2/N] Remove W8A8Fp8BlockLinearOp and adopt Fp8 block linear kernel selections. (#33892)
Signed-off-by: maral <maralbahari.98@gmail.com> Signed-off-by: Maral <maralbahari.98@gmail.com>
Maral committed
Apr 9, 2026 at 08:50 UTC
2e9034c998e231634a75ab74c5812e0dba2cf3a1
35 files changed
+1710
-904
benchmarks/kernels/benchmark_block_fp8_gemm.py
+12
-7
@@ -9,11 +9,12 @@ os.environ["VLLM_USE_DEEP_GEMM"] = "0"
9
import torch
10
11
from vllm.benchmarks.lib.utils import default_vllm_config
12
-from vllm.model_executor.layers.quantization.utils.fp8_utils import (
13
- W8A8BlockFp8LinearOp,
12
+from vllm.model_executor.kernels.linear import (
13
+ init_fp8_linear_kernel,
14
)
15
from vllm.model_executor.layers.quantization.utils.quant_utils import (
16
GroupShape,
17
+ create_fp8_quant_key,
18
)
19
from vllm.model_executor.layers.quantization.utils.w8a8_utils import (
20
CUTLASS_BLOCK_FP8_SUPPORTED,
@@ -70,11 +71,15 @@ def build_w8a8_block_fp8_runner(M, N, K, block_size, device, use_cutlass):
71
weight_group_shape = GroupShape(block_n, block_k)
72
act_quant_group_shape = GroupShape(1, block_k) # Per-token, per-group quantization
73
73
- linear_op = W8A8BlockFp8LinearOp(
74
- weight_group_shape=weight_group_shape,
75
- act_quant_group_shape=act_quant_group_shape,
76
- cutlass_block_fp8_supported=use_cutlass,
77
- use_aiter_and_is_supported=False,
74
+ linear_op = init_fp8_linear_kernel(
75
+ weight_quant_key=create_fp8_quant_key(
76
+ static=True, group_shape=weight_group_shape
77
+ ),
78
+ activation_quant_key=create_fp8_quant_key(
79
+ static=False, group_shape=act_quant_group_shape
80
+ ),
81
+ out_dtype=torch.get_default_dtype(),
82
+ module_name="build_w8a8_block_fp8_runner",
83
)
84
85
def run():
tests/compile/passes/distributed/test_fusion_all_reduce.py
+11
-4
@@ -39,7 +39,9 @@ from vllm.utils.torch_utils import set_random_seed
39
40
41
class TestAllReduceRMSNormModel(torch.nn.Module):
42
- def __init__(self, hidden_size=16, token_num=16, eps=1e-6):
42
+ def __init__(
43
+ self, hidden_size=16, token_num=16, eps=1e-6, dtype: torch.dtype = torch.float16
44
+ ):
45
super().__init__()
46
self.hidden_size = hidden_size
47
self.eps = eps
@@ -78,7 +80,9 @@ class TestAllReduceRMSNormModel(torch.nn.Module):
80
class TestAllReduceRMSNormStaticQuantFP8Model(torch.nn.Module):
81
quant_key = kFp8StaticTensorSym
82
81
- def __init__(self, hidden_size=16, token_num=16, eps=1e-6):
83
+ def __init__(
84
+ self, hidden_size=16, token_num=16, eps=1e-6, dtype: torch.dtype = torch.float16
85
+ ):
86
super().__init__()
87
self.hidden_size = hidden_size
88
self.eps = eps
@@ -88,6 +92,7 @@ class TestAllReduceRMSNormStaticQuantFP8Model(torch.nn.Module):
92
weight_shape=(hidden_size, hidden_size),
93
activation_quant_key=self.quant_key,
94
weight_quant_key=self.quant_key,
95
+ input_dtype=dtype,
96
)
97
for i in range(3)
98
]
@@ -127,7 +132,9 @@ class TestAllReduceRMSNormStaticQuantFP8Model(torch.nn.Module):
132
133
134
class TestAllReduceFusedAddRMSNormStaticQuantFP4Model(torch.nn.Module):
130
- def __init__(self, hidden_size=16, token_num=16, eps=1e-6):
135
+ def __init__(
136
+ self, hidden_size=16, token_num=16, eps=1e-6, dtype: torch.dtype = torch.float16
137
+ ):
138
super().__init__()
139
self.hidden_size = hidden_size
140
self.eps = eps
@@ -314,7 +321,7 @@ def all_reduce_fusion_pass_on_test_model(
321
)
322
323
token_num = batch_size * seq_len
317
- model = test_model_cls(hidden_size, token_num)
324
+ model = test_model_cls(hidden_size, token_num, dtype=dtype)
325
326
hidden_states = torch.randn((token_num, hidden_size), requires_grad=False)
327
tests/compile/passes/distributed/test_sequence_parallelism.py
+1
@@ -109,6 +109,7 @@ class TestAllReduceRMSNormStaticQuantFP8Model(torch.nn.Module):
109
weight_shape=(hidden_size, hidden_size),
110
activation_quant_key=self.quant_key,
111
weight_quant_key=self.quant_key,
112
+ input_dtype=self.vllm_config.model_config.dtype,
113
)
114
for i in range(3)
115
]
tests/compile/passes/test_functionalization.py
+3
@@ -23,6 +23,7 @@ from vllm.config import (
23
ModelConfig,
24
PassConfig,
25
VllmConfig,
26
+ get_current_vllm_config,
27
set_current_vllm_config,
28
)
29
from vllm.model_executor.layers.activation import SiluAndMul
@@ -49,6 +50,7 @@ class TestSiluMul(torch.nn.Module):
50
weight_shape=(hidden_size, hidden_size),
51
activation_quant_key=self.quant_key,
52
weight_quant_key=self.quant_key,
53
+ input_dtype=get_current_vllm_config().model_config.dtype,
54
)
55
56
def forward(self, x):
@@ -92,6 +94,7 @@ class TestFusedAddRMSNorm(torch.nn.Module):
94
weight_shape=(hidden_size, intermediate_size),
95
activation_quant_key=self.quant_key,
96
weight_quant_key=self.quant_key,
97
+ input_dtype=get_current_vllm_config().model_config.dtype,
98
)
99
100
def forward(self, hidden_states, residual):
tests/compile/passes/test_fusion.py
+49
-51
@@ -9,7 +9,7 @@ import vllm.config
9
import vllm.ir.ops
10
import vllm.plugins
11
from tests.compile.backend import TestBackend
12
-from tests.utils import TestBlockFP8Layer, TestFP8Layer
12
+from tests.utils import TestFP8Layer
13
from vllm._aiter_ops import IS_AITER_FOUND, rocm_aiter_ops
14
from vllm.compilation.passes.fusion.matcher_utils import QUANT_OPS
15
from vllm.compilation.passes.fusion.rms_quant_fusion import (
@@ -28,19 +28,23 @@ from vllm.config import (
28
VllmConfig,
29
)
30
from vllm.model_executor.kernels.linear import (
31
+ AiterFp8BlockScaledMMKernel,
32
ChannelWiseTorchFP8ScaledMMLinearKernel,
33
+ CutlassFp8BlockScaledMMKernel,
34
CutlassFP8ScaledMMLinearKernel,
35
+ DeepGemmFp8BlockScaledMMKernel,
36
+ FlashInferFp8DeepGEMMDynamicBlockScaledKernel,
37
FlashInferFP8ScaledMMLinearKernel,
34
- FP8ScaledMMLinearKernel,
38
PerTensorTorchFP8ScaledMMLinearKernel,
39
ROCmFP8ScaledMMLinearKernel,
40
RowWiseTorchFP8ScaledMMLinearKernel,
41
+ TritonFp8BlockScaledMMKernel,
42
+ _KernelT,
43
)
44
from vllm.model_executor.layers.layernorm import RMSNorm
45
from vllm.model_executor.layers.quantization.utils.quant_utils import (
46
GroupShape,
42
- QuantKey,
43
- ScaleDesc,
47
+ create_fp8_quant_key,
48
)
49
from vllm.model_executor.layers.quantization.utils.w8a8_utils import (
50
cutlass_block_fp8_supported,
@@ -66,9 +70,12 @@ CUDA_KERNEL_GROUPSHAPE_COMBINATIONS = [
70
(PerTensorTorchFP8ScaledMMLinearKernel, GroupShape.PER_TENSOR),
71
# ChannelWiseTorchFP8ScaledMMLinearKernel only supports per-token
72
(ChannelWiseTorchFP8ScaledMMLinearKernel, GroupShape.PER_TOKEN),
69
- # Blockwise group shapes (no kernel abstraction)
70
- (None, GroupShape(1, 128)),
71
- (None, GroupShape(1, 64)),
73
+ # Blockwise group shapes
74
+ (FlashInferFp8DeepGEMMDynamicBlockScaledKernel, GroupShape(1, 128)),
75
+ (CutlassFp8BlockScaledMMKernel, GroupShape(1, 128)),
76
+ (DeepGemmFp8BlockScaledMMKernel, GroupShape(1, 128)),
77
+ (TritonFp8BlockScaledMMKernel, GroupShape(1, 128)),
78
+ (TritonFp8BlockScaledMMKernel, GroupShape(1, 64)),
79
]
80
81
# ROCm kernels
@@ -80,8 +87,8 @@ ROCM_KERNEL_GROUPSHAPE_COMBINATIONS = [
87
# ChannelWiseTorchFP8ScaledMMLinearKernel only supports per-token
88
(ChannelWiseTorchFP8ScaledMMLinearKernel, GroupShape.PER_TOKEN),
89
# Blockwise group shapes (no kernel abstraction)
83
- (None, GroupShape(1, 128)),
84
- (None, GroupShape(1, 64)),
90
+ (TritonFp8BlockScaledMMKernel, GroupShape(1, 128)),
91
+ (TritonFp8BlockScaledMMKernel, GroupShape(1, 64)),
92
]
93
94
KERNEL_GROUPSHAPE_COMBINATIONS = (
@@ -100,8 +107,8 @@ AITER_KERNEL_GROUPSHAPE_COMBINATIONS = [
107
# Per-token with ChannelWiseTorchFP8ScaledMMLinearKernel
108
(ChannelWiseTorchFP8ScaledMMLinearKernel, GroupShape.PER_TOKEN, True),
109
(ChannelWiseTorchFP8ScaledMMLinearKernel, GroupShape.PER_TOKEN, False),
103
- # Blockwise (no kernel abstraction)
104
- (None, GroupShape(1, 128), True),
110
+ # Blockwise
111
+ (AiterFp8BlockScaledMMKernel, GroupShape(1, 128), True),
112
]
113
114
@@ -110,8 +117,9 @@ class TestModel(torch.nn.Module):
117
self,
118
hidden_size: int,
119
eps: float,
113
- force_kernel: FP8ScaledMMLinearKernel | None,
120
+ force_kernel: type[_KernelT] | None,
121
group_shape: GroupShape,
122
+ dtype: torch.dtype,
123
use_aiter_fusion: bool = False,
124
use_aiter_quant: bool = False,
125
*args,
@@ -129,54 +137,42 @@ class TestModel(torch.nn.Module):
137
is_blockwise = group_shape.is_per_group()
138
139
if is_blockwise:
132
- act_quant_scale_desc = ScaleDesc(torch.float32, False, group_shape)
133
- self.activation_quant_key = QuantKey(
134
- dtype=FP8_DTYPE, scale=act_quant_scale_desc, symmetric=True
140
+ block_size = group_shape.col
141
+ self.activation_quant_key = create_fp8_quant_key(
142
+ static=False, group_shape=group_shape
143
)
136
- self.fp8_linear_layers = [
137
- TestBlockFP8Layer(
138
- weight_shape=(hidden_size, hidden_size),
139
- group_shape=group_shape,
140
- cutlass_block_fp8_supported=cutlass_block_fp8_supported(),
141
- use_aiter_and_is_supported=use_aiter_quant,
142
- transpose_weights=use_aiter_fusion,
143
- )
144
- for _ in range(3)
145
- ]
146
-
147
- self.enable_quant_fp8_custom_op = (
148
- False
149
- if use_aiter_quant
150
- else self.fp8_linear_layers[0].linear_op.input_quant_op.enabled()
144
+ self.weight_quant_key = create_fp8_quant_key(
145
+ static=True, group_shape=GroupShape(block_size, block_size)
146
)
147
148
else:
149
is_static = group_shape == GroupShape.PER_TENSOR
155
- act_quant_scale_desc = ScaleDesc(torch.float32, is_static, group_shape)
156
- w_quant_scale_desc = ScaleDesc(torch.float32, True, group_shape)
157
- self.activation_quant_key = QuantKey(
158
- dtype=FP8_DTYPE, scale=act_quant_scale_desc, symmetric=True
150
+ self.activation_quant_key = create_fp8_quant_key(
151
+ is_static, group_shape=group_shape
152
)
160
- self.weight_quant_key = QuantKey(
161
- dtype=FP8_DTYPE, scale=w_quant_scale_desc, symmetric=True
153
+ self.weight_quant_key = create_fp8_quant_key(
154
+ static=True, group_shape=group_shape
155
)
163
- self.fp8_linear_layers = [
164
- TestFP8Layer(
165
- weight_shape=(hidden_size, hidden_size),
166
- activation_quant_key=self.activation_quant_key,
167
- weight_quant_key=self.weight_quant_key,
168
- force_kernel=force_kernel,
169
- )
170
- for _ in range(3)
171
- ]
156
173
- # Enable aiter quantization if requested
174
- for layer in self.fp8_linear_layers:
175
- layer.kernel.quant_fp8.use_aiter = use_aiter_quant
157
+ self.fp8_linear_layers = [
158
+ TestFP8Layer(
159
+ weight_shape=(hidden_size, hidden_size),
160
+ activation_quant_key=self.activation_quant_key,
161
+ weight_quant_key=self.weight_quant_key,
162
+ force_kernel=force_kernel,
163
+ transpose_weights=use_aiter_fusion,
164
+ input_dtype=dtype,
165
+ )
166
+ for _ in range(3)
167
+ ]
168
+
169
+ # Enable aiter quantization if requested
170
+ for layer in self.fp8_linear_layers:
171
+ layer.kernel.quant_fp8.use_aiter = use_aiter_quant
172
177
- self.enable_quant_fp8_custom_op = self.fp8_linear_layers[
178
- 0
179
- ].is_quant_fp8_enabled()
173
+ self.enable_quant_fp8_custom_op = self.fp8_linear_layers[
174
+ 0
175
+ ].is_quant_fp8_enabled()
176
177
def forward(self, x):
178
# avoid having graph input be an arg to a pattern directly
@@ -354,6 +350,7 @@ def test_fusion_rmsnorm_quant(
350
eps=eps,
351
force_kernel=force_kernel,
352
group_shape=group_shape,
353
+ dtype=dtype,
354
use_aiter_fusion=False,
355
use_aiter_quant=False,
356
)
@@ -426,6 +423,7 @@ def test_aiter_fusion_rmsnorm_quant(
423
eps=eps,
424
force_kernel=force_kernel,
425
group_shape=group_shape,
426
+ dtype=dtype,
427
use_aiter_fusion=True, # Always use aiter fusion ops in aiter test
428
use_aiter_quant=use_aiter_quant_op, # Toggle aiter quantization
429
)
tests/compile/passes/test_fusion_attn.py
+2
@@ -66,6 +66,7 @@ class AttentionQuantPatternModel(torch.nn.Module):
66
self.kv_cache_dtype = kv_cache_dtype
67
self.device = device
68
self.vllm_config = vllm_config
69
+ self.dtype = vllm_config.model_config.dtype
70
71
self.attn = Attention(
72
num_heads=self.num_qo_heads,
@@ -155,6 +156,7 @@ class TestAttentionFp8StaticQuantPatternModel(AttentionQuantPatternModel):
156
activation_quant_key=self.quant_key,
157
weight_quant_key=self.quant_key,
158
device=self.device,
159
+ input_dtype=self.dtype,
160
)
161
162
w = kwargs.get("w")
tests/compile/passes/test_mla_attn_quant_fusion.py
+2
@@ -74,6 +74,7 @@ class MLAAttentionQuantPatternModel(torch.nn.Module):
74
self.kv_cache_dtype = kv_cache_dtype
75
self.device = device
76
self.vllm_config = vllm_config
77
+ self.dtype = vllm_config.model_config.dtype
78
79
# Create kv_b_proj (ColumnParallelLinear) on device.
80
# Reuse weights from prior model instance when available, because
@@ -190,6 +191,7 @@ class TestMLAAttentionFp8StaticQuantPatternModel(MLAAttentionQuantPatternModel):
191
activation_quant_key=self.quant_key,
192
weight_quant_key=self.quant_key,
193
device=self.device,
194
+ input_dtype=self.dtype,
195
)
196
197
w = kwargs.get("w")
tests/compile/passes/test_silu_mul_quant_fusion.py
+23
-10
@@ -36,9 +36,9 @@ from vllm.model_executor.kernels.linear import (
36
)
37
from vllm.model_executor.layers.activation import SiluAndMul
38
from vllm.model_executor.layers.quantization.input_quant_fp8 import QuantFP8
39
-from vllm.model_executor.layers.quantization.utils.fp8_utils import W8A8BlockFp8LinearOp
39
from vllm.model_executor.layers.quantization.utils.quant_utils import (
40
GroupShape,
41
+ create_fp8_quant_key,
42
kFp8Dynamic128Sym,
43
kFp8StaticTensorSym,
44
kNvfp4Dynamic,
@@ -58,7 +58,11 @@ class TestSiluMulFp8QuantModel(torch.nn.Module):
58
quant_key = kFp8StaticTensorSym
59
60
def __init__(
61
- self, hidden_size: int, force_kernel: FP8ScaledMMLinearKernel, **kwargs
61
+ self,
62
+ hidden_size: int,
63
+ force_kernel: FP8ScaledMMLinearKernel,
64
+ dtype: torch.dtype,
65
+ **kwargs,
66
):
67
super().__init__()
68
self.silu_and_mul = SiluAndMul()
@@ -68,6 +72,7 @@ class TestSiluMulFp8QuantModel(torch.nn.Module):
72
activation_quant_key=self.quant_key,
73
weight_quant_key=self.quant_key,
74
force_kernel=force_kernel,
75
+ input_dtype=dtype,
76
)
77
78
self.enable_silu_mul_custom_op = self.silu_and_mul.enabled()
@@ -137,14 +142,20 @@ class TestSiluMulNvfp4QuantModel(torch.nn.Module):
142
143
144
class TestSiluMulGroupFp8QuantModel(torch.nn.Module):
140
- def __init__(self, hidden_size: int, **kwargs):
145
+ act_quant_key = kFp8Dynamic128Sym
146
+
147
+ def __init__(self, hidden_size: int, dtype: torch.dtype, **kwargs):
148
super().__init__()
149
self.silu_and_mul = SiluAndMul()
143
- self.w8a8_block_fp8_linear = W8A8BlockFp8LinearOp(
144
- weight_group_shape=GroupShape(128, 128),
145
- act_quant_group_shape=GroupShape(1, 128),
146
- cutlass_block_fp8_supported=False,
147
- use_aiter_and_is_supported=True,
150
+ self.weight_quant_key = create_fp8_quant_key(
151
+ static=True, group_shape=GroupShape(hidden_size, hidden_size)
152
+ )
153
+
154
+ self.w8a8_block_fp8_linear = TestFP8Layer(
155
+ weight_shape=(hidden_size, hidden_size),
156
+ weight_quant_key=self.weight_quant_key,
157
+ activation_quant_key=self.act_quant_key,
158
+ input_dtype=dtype,
159
)
160
self.w = torch.rand(hidden_size, hidden_size).to(dtype=FP8_DTYPE).t()
161
@@ -157,7 +168,7 @@ class TestSiluMulGroupFp8QuantModel(torch.nn.Module):
168
169
def forward(self, x):
170
y = self.silu_and_mul(x)
160
- x2 = self.w8a8_block_fp8_linear.apply(y, self.w, self.wscale)
171
+ x2 = self.w8a8_block_fp8_linear(y, self.w, self.wscale)
172
return x2
173
174
def ops_in_model_before(self):
@@ -324,7 +335,9 @@ def test_fusion_silu_and_mul_quant(
335
336
passes = [NoOpEliminationPass(config), *fusion_passes, PostCleanupPass(config)]
337
backend = TestBackend(*passes)
327
- model = model_class(hidden_size=hidden_size, force_kernel=force_kernel, x=x)
338
+ model = model_class(
339
+ hidden_size=hidden_size, force_kernel=force_kernel, x=x, dtype=dtype
340
+ )
341
342
# First dimension dynamic
343
torch._dynamo.mark_dynamic(x, 0)
tests/conftest.py
+3
-2
@@ -246,8 +246,9 @@ def default_vllm_config():
246
"""
247
from vllm.config import VllmConfig, set_current_vllm_config
248
249
- with set_current_vllm_config(VllmConfig()):
250
- yield
249
+ config = VllmConfig()
250
+ with set_current_vllm_config(config):
251
+ yield config
252
253
254
@pytest.fixture()
tests/kernels/quantization/test_block_fp8.py
+2
-2
@@ -12,8 +12,8 @@ from tests.kernels.quant_utils import (
12
native_w8a8_block_matmul,
13
)
14
from vllm.config import VllmConfig
15
+from vllm.model_executor.kernels.linear.scaled_mm.cutlass import cutlass_scaled_mm
16
from vllm.model_executor.layers.quantization.utils.fp8_utils import (
16
- cutlass_scaled_mm,
17
per_token_group_quant_fp8,
18
w8a8_triton_block_scaled_mm,
19
)
@@ -202,7 +202,7 @@ def test_w8a8_block_fp8_deep_gemm_matmul(M, N, K, block_size, out_dtype, seed):
202
203
# only aligned sizes are supported by deepgemm
204
if not should_use_deepgemm_for_fp8_linear(
205
- output_dtype=out_dtype, weight=B_fp32, supports_deep_gemm=True
205
+ output_dtype=out_dtype, weight_shape=B_fp32.shape, supports_deep_gemm=True
206
):
207
pytest.skip(f"Skipping test; invalid size {M}, {N}, {K}")
208
tests/quantization/test_compressed_tensors.py
+5
-5
@@ -16,6 +16,9 @@ from compressed_tensors.quantization import (
16
)
17
18
from tests.models.utils import check_logprobs_close
19
+from vllm.model_executor.kernels.linear import (
20
+ Fp8BlockScaledMMLinearKernel,
21
+)
22
from vllm.model_executor.layers.fused_moe import UnquantizedFusedMoEMethod
23
from vllm.model_executor.layers.quantization.compressed_tensors.compressed_tensors import ( # noqa: E501
24
CompressedTensorsConfig,
@@ -29,7 +32,6 @@ from vllm.model_executor.layers.quantization.compressed_tensors.compressed_tenso
32
CompressedTensorsWNA16,
33
)
34
from vllm.model_executor.layers.quantization.input_quant_fp8 import QuantFP8
32
-from vllm.model_executor.layers.quantization.utils.fp8_utils import W8A8BlockFp8LinearOp
35
from vllm.model_executor.layers.quantization.utils.nvfp4_utils import (
36
cutlass_fp4_supported,
37
)
@@ -473,16 +475,14 @@ def test_compressed_tensors_fp8_block_enabled(vllm_runner):
475
qkv_proj = layer.self_attn.qkv_proj
476
assert isinstance(qkv_proj.quant_method, CompressedTensorsLinearMethod)
477
assert isinstance(qkv_proj.scheme, CompressedTensorsW8A8Fp8)
476
- assert isinstance(
477
- qkv_proj.scheme.w8a8_block_fp8_linear, W8A8BlockFp8LinearOp
478
- )
478
+ assert isinstance(qkv_proj.scheme.fp8_linear, Fp8BlockScaledMMLinearKernel)
479
480
assert qkv_proj.weight.dtype is fp8_dtype
481
assert qkv_proj.weight_scale.dtype is torch.float32
482
assert len(qkv_proj.weight.shape) == 2
483
assert len(qkv_proj.weight_scale.shape) == 2
484
485
- input_quant_op = qkv_proj.scheme.w8a8_block_fp8_linear.input_quant_op
485
+ input_quant_op = qkv_proj.scheme.fp8_linear.quant_fp8
486
assert isinstance(input_quant_op, QuantFP8)
487
assert input_quant_op._forward_method in (
488
input_quant_op.forward_cuda,
tests/quantization/test_fp8.py
+3
@@ -13,6 +13,7 @@ import torch
13
14
from tests.quantization.utils import is_quant_method_supported
15
from vllm import _custom_ops as ops
16
+from vllm.config.model import ModelConfig
17
from vllm.model_executor.layers.fused_moe import FusedMoE
18
from vllm.model_executor.layers.quantization.fp8 import (
19
Fp8Config,
@@ -406,6 +407,8 @@ def test_fp8_reloading(
407
"If this is your use case, consider using a restore function like #26327"
408
)
409
410
+ # Set model config as model_config.dtype is required in Fp8LinearMethod.
411
+ default_vllm_config.model_config = ModelConfig()
412
with torch.device("cuda:0"):
413
config = Fp8Config(
414
is_checkpoint_fp8_serialized=is_checkpoint_fp8_serialized,
tests/quantization/test_modelopt.py
+10
-3
@@ -12,6 +12,7 @@ import pytest
12
import torch
13
14
from tests.quantization.utils import is_quant_method_supported
15
+from vllm.config.model import ModelConfig
16
17
18
@pytest.fixture(scope="function", autouse=True)
@@ -46,7 +47,7 @@ def _snapshot_download_or_skip(model_id: str) -> str:
47
not is_quant_method_supported("modelopt"),
48
reason="ModelOpt FP8 is not supported on this GPU type.",
49
)
49
-def test_modelopt_fp8_checkpoint_setup(vllm_runner):
50
+def test_modelopt_fp8_checkpoint_setup(default_vllm_config, vllm_runner):
51
"""Test ModelOpt FP8 checkpoint loading and structure validation."""
52
# TODO: provide a small publicly available test checkpoint
53
model_path = (
@@ -61,6 +62,8 @@ def test_modelopt_fp8_checkpoint_setup(vllm_runner):
62
"This test requires a local ModelOpt FP8 checkpoint."
63
)
64
65
+ # Set model config as model_config.dtype is required in ModelOptFp8LinearMethod.
66
+ default_vllm_config.model_config = ModelConfig()
67
with vllm_runner(model_path, quantization="modelopt", enforce_eager=True) as llm:
68
69
def check_model(model):
@@ -120,11 +123,13 @@ def test_modelopt_fp8_checkpoint_setup(vllm_runner):
123
not is_quant_method_supported("modelopt"),
124
reason="ModelOpt FP8 is not supported on this GPU type.",
125
)
123
-def test_modelopt_fp8_pc_pt_checkpoint_setup(vllm_runner):
126
+def test_modelopt_fp8_pc_pt_checkpoint_setup(default_vllm_config, vllm_runner):
127
"""Test ModelOpt FP8_PER_CHANNEL_PER_TOKEN checkpoint setup."""
128
model_id = "CedricHwang/qwen2.5-0.5b-modelopt-fp8-pc-pt"
129
model_path = _snapshot_download_or_skip(model_id)
130
131
+ # Set model config as model_config.dtype is required in ModelOptFp8LinearMethod.
132
+ default_vllm_config.model_config = ModelConfig()
133
with vllm_runner(model_path, quantization="modelopt", enforce_eager=True) as llm:
134
135
def check_model(model):
@@ -181,11 +186,13 @@ def test_modelopt_fp8_pc_pt_checkpoint_setup(vllm_runner):
186
not is_quant_method_supported("modelopt"),
187
reason="ModelOpt FP8 is not supported on this GPU type.",
188
)
184
-def test_modelopt_fp8_pb_wo_checkpoint_setup(vllm_runner):
189
+def test_modelopt_fp8_pb_wo_checkpoint_setup(default_vllm_config, vllm_runner):
190
"""Test ModelOpt FP8_PB_WO checkpoint setup."""
191
model_id = "CedricHwang/qwen2.5-0.5b-modelopt-fp8-pb-wo"
192
model_path = _snapshot_download_or_skip(model_id)
193
194
+ # Set model config as model_config.dtype is required in ModelOptFp8LinearMethod.
195
+ default_vllm_config.model_config = ModelConfig()
196
with vllm_runner(model_path, quantization="modelopt", enforce_eager=True) as llm:
197
198
def check_model(model):
tests/utils.py
+37
-76
@@ -43,12 +43,10 @@ from vllm.distributed import (
43
from vllm.engine.arg_utils import AsyncEngineArgs
44
from vllm.entrypoints.cli.serve import ServeSubcommand
45
from vllm.model_executor.kernels.linear import (
46
- FP8ScaledMMLinearKernel,
46
+ _KernelT,
47
init_fp8_linear_kernel,
48
)
49
-from vllm.model_executor.layers.quantization.utils.fp8_utils import W8A8BlockFp8LinearOp
49
from vllm.model_executor.layers.quantization.utils.quant_utils import (
51
- GroupShape,
50
QuantKey,
51
)
52
from vllm.model_executor.model_loader import get_model_loader
@@ -1811,31 +1809,52 @@ class TestFP8Layer(torch.nn.Module):
1809
weight_shape: tuple[int, int],
1810
activation_quant_key: QuantKey,
1811
weight_quant_key: QuantKey,
1812
+ input_dtype: torch.dtype,
1813
out_dtype: torch.dtype | None = None,
1814
+ transpose_weights: bool = False,
1815
device: torch.device | None = None,
1816
- force_kernel: FP8ScaledMMLinearKernel | None = None,
1816
+ force_kernel: type[_KernelT] | None = None,
1817
):
1818
super().__init__()
1819
- per_tensor_weights = weight_quant_key.scale.group_shape.is_per_tensor()
1820
- is_static_activation_scale = activation_quant_key.scale.static
1821
- weight_scale_shape = (1,) if per_tensor_weights else (weight_shape[0], 1)
1822
-
1823
- self.weight_scale = torch.rand(
1824
- weight_scale_shape, dtype=torch.float32, device=device
1825
- )
1826
- self.input_scale = (
1827
- torch.rand(1, dtype=torch.float32, device=device)
1828
- if is_static_activation_scale
1829
- else None
1830
- )
1831
- self.weight = torch.rand(weight_shape, device=device).to(dtype=FP8_DTYPE).t()
1832
- self.input_scale_ub = None
1819
+ act_scale_desc = activation_quant_key.scale
1820
+ weight_scale_desc = weight_quant_key.scale
1821
+ is_block_wise = act_scale_desc.group_shape.is_per_group()
1822
+ if is_block_wise:
1823
+ block_size = weight_scale_desc.group_shape.col
1824
+ weight_scale_shape = weight_shape[0] // block_size
1825
+ self.weight_scale_inv = torch.rand(
1826
+ (weight_scale_shape, weight_scale_shape), dtype=torch.float32
1827
+ )
1828
+ self.weight = torch.rand(weight_shape).to(dtype=FP8_DTYPE)
1829
+ self.input_scale = None
1830
+ self.weight_scale = None
1831
+ if transpose_weights:
1832
+ self.weight = self.weight.t()
1833
+ else:
1834
+ per_tensor_weights = weight_scale_desc.group_shape.is_per_tensor()
1835
+ is_static_activation_scale = act_scale_desc.static
1836
+ weight_scale_shape = (1,) if per_tensor_weights else (weight_shape[0], 1)
1837
+ self.weight_scale_inv = None
1838
+ self.weight_scale = torch.rand(
1839
+ weight_scale_shape, dtype=torch.float32, device=device
1840
+ )
1841
+ self.input_scale = (
1842
+ torch.rand(1, dtype=torch.float32, device=device)
1843
+ if is_static_activation_scale
1844
+ else None
1845
+ )
1846
+ self.weight = (
1847
+ torch.rand(weight_shape, device=device).to(dtype=FP8_DTYPE).t()
1848
+ )
1849
+ self.input_scale_ub = None
1850
1851
out_dtype = torch.get_default_dtype() if out_dtype is None else out_dtype
1852
1853
self.kernel = init_fp8_linear_kernel(
1854
activation_quant_key=activation_quant_key,
1855
weight_quant_key=weight_quant_key,
1856
+ weight_shape=weight_shape,
1857
+ input_dtype=input_dtype,
1858
out_dtype=out_dtype,
1859
force_kernel=force_kernel,
1860
)
@@ -1847,61 +1866,3 @@ class TestFP8Layer(torch.nn.Module):
1866
self, y: torch.Tensor, bias: torch.Tensor | None = None
1867
) -> torch.Tensor:
1868
return self.kernel.apply_weights(self, y, bias)
1850
-
1851
-
1852
-# TODO: Drop TestBlockFP8Layer in favour of a unified TestFP8Layer
1853
-# after refactoring W8A8BlockFp8LinearOp.
1854
-# https://github.com/vllm-project/vllm/issues/31818
1855
-class TestBlockFP8Layer:
1856
- """
1857
- Test helper for blockwise FP8 linear operations. Creates random weights
1858
- and scales for W8A8BlockFp8LinearOp.
1859
-
1860
- This is a workaround until W8A8BlockFp8LinearOp implements the kernel
1861
- abstraction (ScaledMMLinearKernel) for blockwise quantization.
1862
-
1863
- Args:
1864
- weight_shape: Shape of the weight tensor (out_features, in_features).
1865
- group_shape: Blockwise quantization group shape.
1866
- cutlass_block_fp8_supported: Whether CUTLASS blockwise FP8 is available.
1867
- use_aiter_and_is_supported: Whether to use aiter quantization ops.
1868
- transpose_weights: Whether to transpose weights after creation.
1869
- """
1870
-
1871
- def __init__(
1872
- self,
1873
- weight_shape: tuple[int, int],
1874
- group_shape: GroupShape,
1875
- cutlass_block_fp8_supported: bool = False,
1876
- use_aiter_and_is_supported: bool = False,
1877
- transpose_weights: bool = False,
1878
- ):
1879
- weight_scale_shape = weight_shape[0] // group_shape[1]
1880
- self.weight_scale = torch.rand(
1881
- (weight_scale_shape, weight_scale_shape), dtype=torch.float32
1882
- )
1883
- self.weight = torch.rand(weight_shape).to(dtype=FP8_DTYPE)
1884
- self.input_scale = None
1885
- if transpose_weights:
1886
- self.weight = self.weight.t()
1887
-
1888
- self.linear_op = W8A8BlockFp8LinearOp(
1889
- weight_group_shape=GroupShape(group_shape[1], group_shape[1]),
1890
- act_quant_group_shape=group_shape,
1891
- cutlass_block_fp8_supported=cutlass_block_fp8_supported,
1892
- use_aiter_and_is_supported=use_aiter_and_is_supported,
1893
- )
1894
-
1895
- def __call__(
1896
- self, y: torch.Tensor, bias: torch.Tensor | None = None
1897
- ) -> torch.Tensor:
1898
- return self.linear_op.apply(
1899
- input=y,
1900
- weight=self.weight,
1901
- weight_scale=self.weight_scale,
1902
- input_scale=self.input_scale,
1903
- bias=bias,
1904
- )
1905
-
1906
- def is_quant_fp8_enabled(self) -> bool:
1907
- return self.linear_op.input_quant_op.enabled()
vllm/compilation/backends.py
+1
-1
@@ -1002,11 +1002,11 @@ class VllmBackend:
1002
)
1003
hash_content = []
1004
for filepath in forward_code_files:
1005
- hash_content.append(filepath)
1005
if filepath == "<string>":
1006
# This means the function was dynamically generated, with
1007
# e.g. exec(). We can't actually check these.
1008
continue
1009
+ hash_content.append(filepath)
1010
try:
1011
with open(filepath) as f:
1012
hash_content.append(f.read())
vllm/model_executor/kernels/linear/__init__.py
+80
-18
@@ -19,6 +19,10 @@ import torch
19
20
import vllm.envs as envs
21
from vllm.logger import init_logger
22
+from vllm.model_executor.kernels.linear.base import (
23
+ MMLinearKernel,
24
+ MMLinearLayerConfig,
25
+)
26
from vllm.model_executor.kernels.linear.mixed_precision import (
27
MPLinearKernel,
28
MPLinearLayerConfig,
@@ -52,24 +56,30 @@ from vllm.model_executor.kernels.linear.mixed_precision.xpu import (
56
XPUwNa16LinearKernel,
57
)
58
from vllm.model_executor.kernels.linear.scaled_mm import (
59
+ Fp8BlockScaledMMLinearKernel,
60
FP8ScaledMMLinearKernel,
61
FP8ScaledMMLinearLayerConfig,
62
Int8ScaledMMLinearKernel,
63
Int8ScaledMMLinearLayerConfig,
64
ScaledMMLinearKernel,
60
- ScaledMMLinearLayerConfig,
65
)
66
from vllm.model_executor.kernels.linear.scaled_mm.aiter import (
67
+ AiterFp8BlockScaledMMKernel,
68
AiterInt8ScaledMMLinearKernel,
69
)
70
from vllm.model_executor.kernels.linear.scaled_mm.cpu import (
71
CPUInt8ScaledMMLinearKernel,
72
)
73
from vllm.model_executor.kernels.linear.scaled_mm.cutlass import (
74
+ CutlassFp8BlockScaledMMKernel,
75
CutlassFP8ScaledMMLinearKernel,
76
CutlassInt8ScaledMMLinearKernel,
77
)
78
+from vllm.model_executor.kernels.linear.scaled_mm.deep_gemm import (
79
+ DeepGemmFp8BlockScaledMMKernel,
80
+)
81
from vllm.model_executor.kernels.linear.scaled_mm.flashinfer import (
82
+ FlashInferFp8DeepGEMMDynamicBlockScaledKernel,
83
FlashInferFP8ScaledMMLinearKernel,
84
)
85
from vllm.model_executor.kernels.linear.scaled_mm.marlin import (
@@ -84,6 +94,7 @@ from vllm.model_executor.kernels.linear.scaled_mm.rocm import (
94
ROCmFP8ScaledMMLinearKernel,
95
)
96
from vllm.model_executor.kernels.linear.scaled_mm.triton import (
97
+ TritonFp8BlockScaledMMKernel,
98
TritonInt8ScaledMMLinearKernel,
99
)
100
from vllm.model_executor.kernels.linear.scaled_mm.xpu import (
@@ -128,6 +139,23 @@ _POSSIBLE_FP8_KERNELS: dict[PlatformEnum, list[type[FP8ScaledMMLinearKernel]]] =
139
],
140
}
141
142
+
143
+# in priority/performance order (when available)
144
+_POSSIBLE_FP8_BLOCK_KERNELS: dict[
145
+ PlatformEnum, list[type[Fp8BlockScaledMMLinearKernel]]
146
+] = {
147
+ PlatformEnum.CUDA: [
148
+ FlashInferFp8DeepGEMMDynamicBlockScaledKernel,
149
+ DeepGemmFp8BlockScaledMMKernel,
150
+ CutlassFp8BlockScaledMMKernel,
151
+ TritonFp8BlockScaledMMKernel,
152
+ ],
153
+ PlatformEnum.ROCM: [
154
+ AiterFp8BlockScaledMMKernel,
155
+ TritonFp8BlockScaledMMKernel,
156
+ ],
157
+}
158
+
159
# in priority/performance order (when available)
160
_POSSIBLE_KERNELS: dict[PlatformEnum, list[type[MPLinearKernel]]] = {
161
PlatformEnum.CUDA: [
@@ -152,8 +180,10 @@ _POSSIBLE_KERNELS: dict[PlatformEnum, list[type[MPLinearKernel]]] = {
180
],
181
}
182
155
-_KernelT = TypeVar("_KernelT", bound=ScaledMMLinearKernel)
156
-_KernelConfigT = TypeVar("_KernelConfigT", bound=ScaledMMLinearLayerConfig)
183
+# TODO make all kernels inherit from MMLinearKernel
184
+# then bound _KernelT only to MMLinearKernel
185
+_KernelT = TypeVar("_KernelT", bound=ScaledMMLinearKernel | MMLinearKernel)
186
+_KernelConfigT = TypeVar("_KernelConfigT", bound=MMLinearLayerConfig)
187
188
189
def is_supported_and_can_implement_kernel(
@@ -243,32 +273,61 @@ def choose_scaled_mm_linear_kernel(
273
def init_fp8_linear_kernel(
274
activation_quant_key: QuantKey,
275
weight_quant_key: QuantKey,
276
+ weight_shape: tuple[int, int],
277
+ input_dtype: torch.dtype,
278
out_dtype: torch.dtype,
247
- force_kernel: type[FP8ScaledMMLinearKernel] | None = None,
279
+ force_kernel: type[_KernelT] | None = None,
280
module_name: str | None = None,
249
-) -> FP8ScaledMMLinearKernel:
281
+) -> FP8ScaledMMLinearKernel | Fp8BlockScaledMMLinearKernel:
282
scaled_mm_linear_kernel_config = FP8ScaledMMLinearLayerConfig(
283
weight_quant_key=weight_quant_key,
284
activation_quant_key=activation_quant_key,
285
+ weight_shape=weight_shape,
286
+ input_dtype=input_dtype,
287
out_dtype=out_dtype,
288
)
289
256
- kernel_type = choose_scaled_mm_linear_kernel(
257
- scaled_mm_linear_kernel_config, _POSSIBLE_FP8_KERNELS, force_kernel=force_kernel
258
- )
290
+ if activation_quant_key.scale.group_shape.is_per_group():
291
+ kernel_type = choose_scaled_mm_linear_kernel(
292
+ config=scaled_mm_linear_kernel_config,
293
+ possible_kernels=_POSSIBLE_FP8_BLOCK_KERNELS, # type: ignore[misc]
294
+ force_kernel=force_kernel,
295
+ )
296
+ if module_name:
297
+ logger.info_once(
298
+ "Selected %s for %s",
299
+ kernel_type.__name__,
300
+ module_name,
301
+ scope="global",
302
+ )
303
260
- if module_name:
261
- logger.info_once(
262
- "Selected %s for %s",
263
- kernel_type.__name__,
264
- module_name,
265
- scope="global",
304
+ return kernel_type(
305
+ scaled_mm_linear_kernel_config,
306
)
307
268
- return kernel_type(
269
- scaled_mm_linear_kernel_config,
270
- layer_param_names=["weight", "weight_scale", "input_scale", "input_scale_ub"],
271
- )
308
+ else:
309
+ kernel_type = choose_scaled_mm_linear_kernel(
310
+ config=scaled_mm_linear_kernel_config,
311
+ possible_kernels=_POSSIBLE_FP8_KERNELS, # type: ignore[misc]
312
+ force_kernel=force_kernel,
313
+ )
314
+ if module_name:
315
+ logger.info_once(
316
+ "Selected %s for %s",
317
+ kernel_type.__name__,
318
+ module_name,
319
+ scope="global",
320
+ )
321
+
322
+ return kernel_type(
323
+ scaled_mm_linear_kernel_config,
324
+ layer_param_names=[
325
+ "weight",
326
+ "weight_scale",
327
+ "input_scale",
328
+ "input_scale_ub",
329
+ ],
330
+ )
331
332
333
def init_int8_linear_kernel(
@@ -433,4 +492,7 @@ __all__ = [
492
"MarlinLinearKernel",
493
"XPUW4A8IntLinearKernel",
494
"XPUwNa16LinearKernel",
495
+ "_KernelT",
496
+ "DeepGemmFp8BlockScaledMMKernel",
497
+ "FlashInferFp8DeepGEMMDynamicBlockScaledKernel",
498
]
vllm/model_executor/kernels/linear/base.py
new
+324
@@ -0,0 +1,324 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+from abc import ABC, abstractmethod
5
+from dataclasses import dataclass
6
+from typing import Any, ClassVar, Generic, TypeVar
7
+
8
+import torch
9
+from typing_extensions import Self
10
+
11
+
12
+@dataclass
13
+class MMLinearLayerConfig: ...
14
+
15
+
16
+@dataclass
17
+class Params:
18
+ """Base class for quantized layer parameters.
19
+
20
+ This class provides a typed interface for accessing quantized weights and scales
21
+ from layer modules. It serves as a parameter container that can be extracted from
22
+ layers and passed to kernel implementations.
23
+
24
+ Attributes:
25
+ weight: The quantized weight tensor
26
+ weight_scale: weight scaling factors
27
+ input_scale: Optional input scaling factors
28
+
29
+ Class Variables:
30
+ WEIGHT: Attribute name for weight tensor on the layer module
31
+ WEIGHT_SCALE: Attribute name for weight scale tensor on the layer module
32
+ INPUT_SCALE: Attribute name for input scale tensor on the layer module
33
+
34
+ Important:
35
+ The string values of WEIGHT, WEIGHT_SCALE, and INPUT_SCALE class variables
36
+ MUST match the attribute names used in the corresponding quantization method's
37
+ create_weights() implementation.
38
+ For example, if FP8LinearMethod.create_weights()
39
+ sets layer.weight and layer.weight_scale,
40
+ then WEIGHT="weight" and
41
+ WEIGHT_SCALE="weight_scale" must be used here.
42
+
43
+ Usage:
44
+ ```python
45
+ # Extract parameters from a quantized layer
46
+ params = Params.from_layer(layer)
47
+
48
+ # Access typed parameters
49
+ output = func(input, params.weight, params.weight_scale)
50
+ ```
51
+ """
52
+
53
+ weight: torch.Tensor
54
+ weight_scale: torch.Tensor
55
+ input_scale: torch.Tensor | None
56
+
57
+ # Attribute names on the layer
58
+ WEIGHT: ClassVar[str] = "weight"
59
+ WEIGHT_SCALE: ClassVar[str] = "weight_scale"
60
+ INPUT_SCALE: ClassVar[str] = "input_scale"
61
+
62
+ @classmethod
63
+ def from_layer(cls, layer: torch.nn.Module) -> Self:
64
+ return cls(
65
+ weight=getattr(layer, cls.WEIGHT),
66
+ weight_scale=getattr(layer, cls.WEIGHT_SCALE),
67
+ input_scale=getattr(layer, cls.INPUT_SCALE, None),
68
+ )
69
+
70
+
71
+@dataclass
72
+class FP8Params(Params):
73
+ """FP8 layer parameters with typed fields"""
74
+
75
+ input_scale_ub: torch.Tensor | None
76
+
77
+ INPUT_SCALE_UB: ClassVar[str] = "input_scale_ub"
78
+
79
+ @classmethod
80
+ def from_layer(cls, layer: torch.nn.Module) -> "FP8Params":
81
+ """Extract parameters from layer"""
82
+ return cls(
83
+ weight=getattr(layer, cls.WEIGHT),
84
+ weight_scale=getattr(layer, cls.WEIGHT_SCALE),
85
+ input_scale=getattr(layer, cls.INPUT_SCALE, None),
86
+ input_scale_ub=getattr(layer, cls.INPUT_SCALE_UB, None),
87
+ )
88
+
89
+
90
+@dataclass
91
+class Int8Params(Params):
92
+ """Int8 layer parameters with typed fields"""
93
+
94
+ input_zero_point: torch.Tensor | None
95
+ azp_adj: torch.Tensor | None
96
+
97
+ INPUT_ZERO_POINT: ClassVar[str] = "input_zero_point"
98
+ AZP_ADJ: ClassVar[str] = "azp_adj"
99
+
100
+ @classmethod
101
+ def from_layer(cls, layer: torch.nn.Module) -> "Int8Params":
102
+ """Extract parameters from layer"""
103
+ return cls(
104
+ weight=getattr(layer, cls.WEIGHT),
105
+ weight_scale=getattr(layer, cls.WEIGHT_SCALE),
106
+ input_scale=getattr(layer, cls.INPUT_SCALE, None),
107
+ input_zero_point=getattr(layer, cls.INPUT_ZERO_POINT, None),
108
+ azp_adj=getattr(layer, cls.AZP_ADJ, None),
109
+ )
110
+
111
+
112
+_ParamsT = TypeVar("_ParamsT", bound=Params)
113
+_ConfigT = TypeVar("_ConfigT", bound=MMLinearLayerConfig)
114
+
115
+
116
+class MMLinearKernel(ABC, Generic[_ConfigT, _ParamsT]):
117
+ """Abstract base class for quantized matrix multiplication kernels.
118
+
119
+ This class provides the interface for implementing custom quantized linear layer
120
+ kernels in vLLM. Subclasses should implement specific quantization strategies
121
+ (e.g., FP8, INT8) and their corresponding compute kernels.
122
+
123
+ Generic Type Parameters:
124
+ _ConfigT: Configuration type for the kernel (subclass of MMLinearLayerConfig).
125
+ Contains kernel-specific settings like quantization keys, dtypes, etc.
126
+ _ParamsT: Parameter type for the kernel (subclass of Params).
127
+ Defines the quantized weights and scales needed by the kernel.
128
+
129
+ Typical Usage:
130
+ 1. Define a config dataclass inheriting from MMLinearLayerConfig
131
+ 2. Define a params dataclass inheriting from Params (or FP8Params/Int8Params)
132
+ 3. Subclass MMLinearKernel with your config and params types
133
+ 4. Implement all abstract methods
134
+ 5. Register the kernel with the quantization method
135
+
136
+ Example:
137
+ ```python
138
+ @dataclass
139
+ class MyKernelConfig(MMLinearLayerConfig):
140
+ static: bool
141
+ output_dtype: torch.dtype
142
+
143
+
144
+ @dataclass
145
+ class MyKernelParams(FP8Params):
146
+ custom_scale: torch.Tensor
147
+ CUSTOM_SCALE: ClassVar[str] = "custom_scale"
148
+
149
+
150
+ class MyKernel(MMLinearKernel[MyKernelConfig, MyKernelParams]):
151
+ @classmethod
152
+ def is_supported(cls, compute_capability=None):
153
+ if compute_capability and compute_capability < 90:
154
+ return False, "Requires compute capability >= 9.0"
155
+ return True, None
156
+
157
+ @classmethod
158
+ def can_implement(cls, config):
159
+ if not config.static:
160
+ return False, "Only static quantization supported"
161
+ return True, None
162
+
163
+ def process_weights_after_loading(self, layer):
164
+ # Preprocess weights for the kernel
165
+ params = self._get_layer_params(layer)
166
+ processed = preprocess_weights(params.weight)
167
+ replace_parameter(layer, params.WEIGHT, processed)
168
+
169
+ def _get_layer_params(self, layer, **kwargs):
170
+ return MyKernelParams.from_layer(layer)
171
+
172
+ def apply_weights(self, layer, x, bias=None, **kwargs):
173
+ params = self._get_layer_params(layer)
174
+ # Call your custom kernel
175
+ output = my_custom_kernel(x, params.weight, params.weight_scale)
176
+ if bias is not None:
177
+ output += bias
178
+ return output
179
+ ```
180
+
181
+ Lifecycle:
182
+ 1. Kernel selection: is_supported() and can_implement() check compatibility
183
+ 2. Initialization: __init__() creates kernel instance with config
184
+ 3. Weight loading: process_weights_after_loading() preprocesses weights
185
+ 4. Inference: apply_weights() executes the quantized matmul
186
+ """
187
+
188
+ @classmethod
189
+ @abstractmethod
190
+ def is_supported(
191
+ cls, compute_capability: int | None = None
192
+ ) -> tuple[bool, str | None]:
193
+ """Check if this kernel is supported on the current hardware.
194
+
195
+ This method checks hardware-level compatibility (e.g., GPU architecture,
196
+ compute capability, available instructions). It's called during kernel
197
+ selection to filter out kernels that cannot run on the current device.
198
+
199
+ Args:
200
+ compute_capability: GPU compute capability (e.g., 80 for A100, 90 for H100).
201
+ If None, should check the current device.
202
+
203
+ Returns:
204
+ A tuple of (is_supported, reason):
205
+ - is_supported: True if the kernel can run on this hardware
206
+ - reason: If not supported, a string explaining why; otherwise None
207
+ """
208
+ raise NotImplementedError
209
+
210
+ @classmethod
211
+ @abstractmethod
212
+ def can_implement(cls, config: _ConfigT) -> tuple[bool, str | None]:
213
+ """Check if this kernel can implement the given configuration.
214
+
215
+ This method checks configuration-level compatibility (e.g., quantization
216
+ scheme, group sizes, static vs dynamic quantization). It's called after
217
+ is_supported() to determine if this kernel can handle the specific
218
+ quantization configuration.
219
+
220
+ Args:
221
+ config: The kernel configuration to check
222
+
223
+ Returns:
224
+ A tuple of (can_implement, reason):
225
+ - can_implement: True if this kernel supports the config
226
+ - reason: If not supported, a string explaining why; otherwise None
227
+ ```
228
+ """
229
+ raise NotImplementedError
230
+
231
+ def __init__(self, config: _ConfigT) -> None:
232
+ """Initialize the kernel with the given configuration.
233
+
234
+ Args:
235
+ config: Kernel-specific configuration containing settings like
236
+ quantization keys, output dtypes, etc.
237
+ """
238
+ self.config = config
239
+
240
+ @abstractmethod
241
+ def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
242
+ """Process and transform weights after loading from checkpoint.
243
+
244
+ This method is called once after weights are loaded but before inference.
245
+ Use it to preprocess weights into the format required by your kernel
246
+ (e.g., reordering, padding, format conversion).
247
+
248
+ Modifications should be done in-place using replace_parameter() to ensure
249
+ the layer's parameters are properly updated.
250
+
251
+ Args:
252
+ layer: The layer module containing the weights to process
253
+
254
+ Example:
255
+ ```python
256
+ def process_weights_after_loading(self, layer):
257
+ params = self._get_layer_params(layer)
258
+ # Reorder weights for better memory access
259
+ weight_reordered = reorder_weights(params.weight)
260
+ replace_parameter(layer, params.WEIGHT, weight_reordered)
261
+ ```
262
+ """
263
+ raise NotImplementedError
264
+
265
+ # return a covariant type in the subclass
266
+ @abstractmethod
267
+ def _get_layer_params(self, layer: torch.nn.Module, **kwargs: Any) -> _ParamsT:
268
+ """Extract typed parameters from the layer module.
269
+
270
+ This internal method retrieves the quantized weights and scales from
271
+ the layer as a typed parameter object. Subclasses should typically
272
+ delegate to ParamsClass.from_layer().
273
+
274
+ Args:
275
+ layer: The layer module containing the parameters
276
+ **kwargs: Additional arguments
277
+
278
+ Returns:
279
+ A typed parameter object containing weights, scales, and other
280
+ quantization parameters
281
+
282
+ Example:
283
+ ```python
284
+ def _get_layer_params(self, layer, **kwargs):
285
+ return MyKernelParams.from_layer(layer)
286
+ ```
287
+ """
288
+ raise NotImplementedError
289
+
290
+ def get_output_padding(self) -> int | None:
291
+ """Get the number of output tokens to pad for this kernel.
292
+
293
+ Some kernels require input padding for optimal performance.
294
+ Override this method to specify padding requirements.
295
+
296
+ Returns:
297
+ Number of tokens to pad, or None for no padding (default)
298
+ """
299
+ return None
300
+
301
+ @abstractmethod
302
+ def apply_weights(
303
+ self,
304
+ layer: torch.nn.Module,
305
+ x: torch.Tensor,
306
+ bias: torch.Tensor | None = None,
307
+ **kwargs: Any,
308
+ ) -> torch.Tensor:
309
+ """Apply the quantized weights to the input tensor.
310
+
311
+ This is the main inference method that performs the quantized matrix
312
+ multiplication. It should handle input quantization (if needed), call
313
+ the underlying kernel, and apply bias.
314
+
315
+ Args:
316
+ layer: The layer module containing the quantized weights
317
+ x: Input tensor of shape [..., in_features]
318
+ bias: Optional bias tensor of shape [out_features]
319
+ **kwargs: Additional kernel-specific arguments
320
+
321
+ Returns:
322
+ Output tensor of shape [..., out_features]
323
+ """
324
+ raise NotImplementedError
vllm/model_executor/kernels/linear/scaled_mm/BlockScaledMMLinearKernel.py
new
+209
@@ -0,0 +1,209 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+from abc import ABC, abstractmethod
5
+from dataclasses import dataclass
6
+from typing import ClassVar
7
+
8
+import torch
9
+from typing_extensions import Self
10
+
11
+from vllm.model_executor.layers.quantization.input_quant_fp8 import QuantFP8
12
+from vllm.model_executor.layers.quantization.utils.fp8_utils import (
13
+ process_fp8_weight_block_strategy,
14
+)
15
+from vllm.model_executor.utils import replace_parameter
16
+
17
+from ..base import (
18
+ FP8Params,
19
+ MMLinearKernel,
20
+)
21
+from .ScaledMMLinearKernel import FP8ScaledMMLinearLayerConfig
22
+
23
+
24
+@dataclass
25
+class FP8BlockParams(FP8Params):
26
+ weight_scale_inv: torch.Tensor | None
27
+ weight_scale: torch.Tensor | None
28
+
29
+ WEIGHT_SCALE_INV: ClassVar[str] = "weight_scale_inv"
30
+
31
+ @classmethod
32
+ def from_layer(cls, layer: torch.nn.Module) -> Self:
33
+ return cls(
34
+ weight=getattr(layer, cls.WEIGHT),
35
+ weight_scale_inv=getattr(layer, cls.WEIGHT_SCALE_INV, None),
36
+ weight_scale=getattr(layer, cls.WEIGHT_SCALE, None),
37
+ input_scale=getattr(layer, cls.INPUT_SCALE, None),
38
+ input_scale_ub=getattr(layer, cls.INPUT_SCALE_UB, None),
39
+ )
40
+
41
+
42
+class Fp8BlockScaledMMLinearKernel(
43
+ MMLinearKernel[FP8ScaledMMLinearLayerConfig, FP8BlockParams], ABC
44
+):
45
+ # Set to False in subclasses that accept BF16 input directly (e.g. FlashInfer)
46
+ # and therefore do not need the input quantization step in apply_weights.
47
+ apply_input_quant: ClassVar[bool] = True
48
+
49
+ def __init__(self, config: FP8ScaledMMLinearLayerConfig) -> None:
50
+ super().__init__(config)
51
+ act_scale_descriptor = config.activation_quant_key.scale
52
+ self.weight_group_shape = config.weight_quant_key.scale.group_shape
53
+ self.quant_fp8 = QuantFP8(
54
+ static=act_scale_descriptor.static,
55
+ group_shape=act_scale_descriptor.group_shape,
56
+ num_token_padding=self.get_output_padding(),
57
+ use_ue8m0=False,
58
+ )
59
+ self.use_triton = False
60
+
61
+ @classmethod
62
+ def can_implement(cls, config: FP8ScaledMMLinearLayerConfig):
63
+ act_quant_key = config.activation_quant_key
64
+ if act_quant_key.scale.static:
65
+ return (
66
+ False,
67
+ "Only dynamic per token group activation quantization is supported.",
68
+ )
69
+
70
+ return True, None
71
+
72
+ def _get_layer_params(self, layer: torch.nn.Module, **kwargs) -> FP8BlockParams:
73
+ return FP8BlockParams.from_layer(layer)
74
+
75
+ def process_weights_after_loading(self, layer: torch.nn.Module):
76
+ params = self._get_layer_params(layer)
77
+ # Fp8LinearMethod registered weight scale
78
+ # buffer as weight_scale_inv unlike compressed tensors.
79
+ weight_scale = (
80
+ params.weight_scale
81
+ if params.weight_scale_inv is None
82
+ else params.weight_scale_inv
83
+ )
84
+ scale_attr_name = (
85
+ params.WEIGHT_SCALE
86
+ if params.weight_scale_inv is None
87
+ else params.WEIGHT_SCALE_INV
88
+ )
89
+ new_weight, new_weight_scale = process_fp8_weight_block_strategy(
90
+ params.weight,
91
+ weight_scale,
92
+ )
93
+
94
+ replace_parameter(layer, params.WEIGHT, new_weight.data)
95
+ replace_parameter(layer, scale_attr_name, new_weight_scale.data)
96
+
97
+ def apply_weights(
98
+ self,
99
+ layer: torch.nn.Module,
100
+ x: torch.Tensor,
101
+ bias: torch.Tensor | None = None,
102
+ **kwargs,
103
+ ) -> torch.Tensor:
104
+ out_dtype = self.config.out_dtype
105
+ params = self._get_layer_params(layer)
106
+ weight = params.weight
107
+ weight_scale = (
108
+ params.weight_scale
109
+ if params.weight_scale_inv is None
110
+ else params.weight_scale_inv
111
+ )
112
+ input_scale = params.input_scale
113
+ scale_up = params.input_scale_ub
114
+
115
+ # View input as 2D matrix for fp8 methods
116
+ input_2d = x.view(-1, x.shape[-1])
117
+ output_shape = [*x.shape[:-1], weight.shape[0]]
118
+
119
+ if self.apply_input_quant:
120
+ q_input, input_scale = self.quant_fp8(
121
+ input_2d, input_scale, scale_up, use_triton=self.use_triton
122
+ )
123
+ else:
124
+ q_input = input_2d
125
+ # Provide a concrete placeholder so apply_block_scaled_mm args are
126
+ # always Tensors. Subclasses with apply_input_quant=False must not
127
+ # use As in apply_block_scaled_mm.
128
+ input_scale = (
129
+ input_scale if input_scale is not None else input_2d.new_ones(1)
130
+ )
131
+
132
+ output = self.apply_block_scaled_mm(
133
+ A=q_input,
134
+ B=weight,
135
+ As=input_scale,
136
+ Bs=weight_scale,
137
+ )
138
+
139
+ if bias is not None:
140
+ output = output + bias
141
+ return output.to(dtype=out_dtype).view(*output_shape)
142
+
143
+ @abstractmethod
144
+ def apply_block_scaled_mm(
145
+ self,
146
+ A: torch.Tensor,
147
+ B: torch.Tensor,
148
+ As: torch.Tensor,
149
+ Bs: torch.Tensor,
150
+ ) -> torch.Tensor:
151
+ raise NotImplementedError
152
+
153
+
154
+class Fp8BlockScaledDynamicMMLinearKernel(Fp8BlockScaledMMLinearKernel, ABC):
155
+ """Dynamic FP8 block-scaled kernel that dispatches at runtime.
156
+
157
+ Extends Fp8BlockScaledMMLinearKernel to inherit apply_weights and overrides
158
+ apply_block_scaled_mm to dispatch between two sub-kernels using torch.cond.
159
+
160
+ Subclasses must define:
161
+ base_type: The primary kernel class.
162
+ fallback_type: The fallback kernel class.
163
+ """
164
+
165
+ base_type: ClassVar[type[Fp8BlockScaledMMLinearKernel]]
166
+ fallback_type: ClassVar[type[Fp8BlockScaledMMLinearKernel]]
167
+
168
+ def __init__(self, config: "FP8ScaledMMLinearLayerConfig") -> None:
169
+ super().__init__(config)
170
+ self.base = self.base_type(config)
171
+ self.fallback = self.fallback_type(config)
172
+
173
+ @classmethod
174
+ def is_supported(
175
+ cls, compute_capability: int | None = None
176
+ ) -> tuple[bool, str | None]:
177
+ is_base_supported, reason_1 = cls.base_type.is_supported(compute_capability)
178
+ is_fallback_supported, reason_2 = cls.fallback_type.is_supported(
179
+ compute_capability
180
+ )
181
+ if is_base_supported and is_fallback_supported:
182
+ return True, None
183
+ if not is_base_supported and not is_fallback_supported:
184
+ return (
185
+ False,
186
+ f"base is not supported due to {reason_1}; "
187
+ f"fallback is not supported due to {reason_2}",
188
+ )
189
+ if not is_base_supported:
190
+ return False, f"base is not supported due to {reason_1}"
191
+ return False, f"fallback is not supported due to {reason_2}"
192
+
193
+ @classmethod
194
+ def can_implement(
195
+ cls, config: "FP8ScaledMMLinearLayerConfig"
196
+ ) -> tuple[bool, str | None]:
197
+ can_implement_base, reason_1 = cls.base_type.can_implement(config)
198
+ can_implement_fallback, reason_2 = cls.fallback_type.can_implement(config)
199
+ if can_implement_base and can_implement_fallback:
200
+ return True, None
201
+ if not can_implement_base and not can_implement_fallback:
202
+ return (
203
+ False,
204
+ f"base cannot implement due to {reason_1}; "
205
+ f"fallback cannot implement due to {reason_2}",
206
+ )
207
+ if not can_implement_base:
208
+ return False, f"base cannot implement due to {reason_1}"
209
+ return False, f"fallback cannot implement due to {reason_2}"
vllm/model_executor/kernels/linear/scaled_mm/ScaledMMLinearKernel.py
+7
-8
@@ -14,14 +14,11 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import (
14
)
15
from vllm.platforms import current_platform
16
17
-
18
-@dataclass
19
-class ScaledMMLinearLayerConfig:
20
- pass
17
+from ..base import MMLinearLayerConfig
18
19
20
@dataclass
24
-class Int8ScaledMMLinearLayerConfig(ScaledMMLinearLayerConfig):
21
+class Int8ScaledMMLinearLayerConfig(MMLinearLayerConfig):
22
# TODO: Change to QuantKey like FP8ScaledMMLinearLayerConfig
23
is_static_input_scheme: bool
24
is_channelwise: bool
@@ -29,10 +26,12 @@ class Int8ScaledMMLinearLayerConfig(ScaledMMLinearLayerConfig):
26
27
28
@dataclass
32
-class FP8ScaledMMLinearLayerConfig(ScaledMMLinearLayerConfig):
29
+class FP8ScaledMMLinearLayerConfig(MMLinearLayerConfig):
30
weight_quant_key: QuantKey
31
activation_quant_key: QuantKey
35
- out_dtype: torch.dtype | None
32
+ weight_shape: tuple[int, int]
33
+ input_dtype: torch.dtype
34
+ out_dtype: torch.dtype
35
36
37
_FP8ParamsT = tuple[
@@ -50,7 +49,7 @@ _Int8ParamsT = tuple[
49
]
50
51
_ParamsT = TypeVar("_ParamsT", _Int8ParamsT, _FP8ParamsT)
53
-_ConfigT = TypeVar("_ConfigT", bound=ScaledMMLinearLayerConfig)
52
+_ConfigT = TypeVar("_ConfigT", bound=MMLinearLayerConfig)
53
54
55
class ScaledMMLinearKernel(Generic[_ConfigT, _ParamsT], ABC):
vllm/model_executor/kernels/linear/scaled_mm/__init__.py
+4
-1
@@ -4,6 +4,9 @@
4
from vllm.model_executor.kernels.linear.scaled_mm.aiter import (
5
AiterInt8ScaledMMLinearKernel,
6
)
7
+from vllm.model_executor.kernels.linear.scaled_mm.BlockScaledMMLinearKernel import (
8
+ Fp8BlockScaledMMLinearKernel,
9
+)
10
from vllm.model_executor.kernels.linear.scaled_mm.cpu import (
11
CPUInt8ScaledMMLinearKernel,
12
)
@@ -31,7 +34,6 @@ from vllm.model_executor.kernels.linear.scaled_mm.ScaledMMLinearKernel import (
34
Int8ScaledMMLinearKernel,
35
Int8ScaledMMLinearLayerConfig,
36
ScaledMMLinearKernel,
34
- ScaledMMLinearLayerConfig,
37
)
38
from vllm.model_executor.kernels.linear.scaled_mm.triton import (
39
TritonInt8ScaledMMLinearKernel,
@@ -55,4 +57,5 @@ __all__ = [
57
"RowWiseTorchFP8ScaledMMLinearKernel",
58
"ROCmFP8ScaledMMLinearKernel",
59
"TritonInt8ScaledMMLinearKernel",
60
+ "Fp8BlockScaledMMLinearKernel",
61
]
vllm/model_executor/kernels/linear/scaled_mm/aiter.py
+58
@@ -6,8 +6,15 @@ import torch
6
7
from vllm import _custom_ops as ops
8
from vllm._aiter_ops import rocm_aiter_ops
9
+from vllm.model_executor.layers.quantization.utils.quant_utils import (
10
+ GroupShape,
11
+)
12
from vllm.platforms import current_platform
13
14
+from .BlockScaledMMLinearKernel import (
15
+ Fp8BlockScaledMMLinearKernel,
16
+ FP8ScaledMMLinearLayerConfig,
17
+)
18
from .cutlass import CutlassInt8ScaledMMLinearKernel
19
from .ScaledMMLinearKernel import Int8ScaledMMLinearLayerConfig
20
@@ -107,3 +114,54 @@ class AiterInt8ScaledMMLinearKernel(CutlassInt8ScaledMMLinearKernel):
114
# b to be [N, K]
115
# CutlassInt8ScaledMMLinearKernel prepare weight `w_q` in [K, N] format
116
return rocm_aiter_ops.gemm_a8w8(x_q, w_q.t(), x_s, w_s, bias, out_dtype)
117
+
118
+
119
+class AiterFp8BlockScaledMMKernel(Fp8BlockScaledMMLinearKernel):
120
+ def __init__(self, config: FP8ScaledMMLinearLayerConfig):
121
+ super().__init__(config)
122
+ n, k = config.weight_shape
123
+
124
+ self.use_triton = (
125
+ not current_platform.is_fp8_fnuz()
126
+ and rocm_aiter_ops.is_triton_gemm_w8a8_tuned(n, k)
127
+ )
128
+
129
+ @classmethod
130
+ def is_supported(cls, compute_capability=None):
131
+ return (
132
+ rocm_aiter_ops.is_linear_enabled(),
133
+ "Only supported on ROCm platform \
134
+ with aiter package installed.",
135
+ )
136
+
137
+ @classmethod
138
+ def can_implement(cls, config: FP8ScaledMMLinearLayerConfig):
139
+ can_implement_base, reason = super().can_implement(config)
140
+ if not can_implement_base:
141
+ return can_implement_base, reason
142
+
143
+ act_quant_desc = config.activation_quant_key.scale
144
+ if act_quant_desc.group_shape != GroupShape(1, 128):
145
+ return (
146
+ False,
147
+ "Supports only dynamic per token group activation "
148
+ "quantization with group_shape=(1,128).",
149
+ )
150
+ return True, None
151
+
152
+ def apply_block_scaled_mm(
153
+ self,
154
+ A: torch.Tensor,
155
+ B: torch.Tensor,
156
+ As: torch.Tensor,
157
+ Bs: torch.Tensor,
158
+ ) -> torch.Tensor:
159
+ out_dtype = self.config.out_dtype
160
+ if self.use_triton:
161
+ gemm_a8w8_blockscale_op = rocm_aiter_ops.triton_gemm_a8w8_blockscale
162
+ else:
163
+ gemm_a8w8_blockscale_op = rocm_aiter_ops.gemm_a8w8_blockscale
164
+
165
+ return gemm_a8w8_blockscale_op(
166
+ A, B, As, Bs, list(self.weight_group_shape), output_dtype=out_dtype
167
+ )
vllm/model_executor/kernels/linear/scaled_mm/cutlass.py
+147
@@ -5,12 +5,19 @@
5
import torch
6
7
from vllm import _custom_ops as ops
8
+from vllm.model_executor.layers.quantization.input_quant_fp8 import QuantFP8
9
from vllm.model_executor.layers.quantization.utils import replace_parameter
10
+from vllm.model_executor.layers.quantization.utils.quant_utils import (
11
+ GroupShape,
12
+)
13
from vllm.model_executor.layers.quantization.utils.w8a8_utils import (
14
+ CUTLASS_BLOCK_FP8_SUPPORTED,
15
convert_to_channelwise,
16
)
17
from vllm.platforms import current_platform
18
+from vllm.utils.torch_utils import direct_register_custom_op
19
20
+from .BlockScaledMMLinearKernel import Fp8BlockScaledMMLinearKernel
21
from .ScaledMMLinearKernel import (
22
FP8ScaledMMLinearKernel,
23
FP8ScaledMMLinearLayerConfig,
@@ -171,3 +178,143 @@ class CutlassFP8ScaledMMLinearKernel(FP8ScaledMMLinearKernel):
178
A, B, out_dtype=out_dtype, scale_a=As, scale_b=Bs, bias=bias
179
)
180
return output.view(*output_shape)
181
+
182
+
183
+class CutlassFp8BlockScaledMMKernel(Fp8BlockScaledMMLinearKernel):
184
+ def __init__(self, config: FP8ScaledMMLinearLayerConfig) -> None:
185
+ super().__init__(config)
186
+ act_scale_descriptor = config.activation_quant_key.scale
187
+ self.weight_group_shape = config.weight_quant_key.scale.group_shape
188
+ self.quant_fp8 = QuantFP8(
189
+ static=act_scale_descriptor.static,
190
+ group_shape=act_scale_descriptor.group_shape,
191
+ num_token_padding=self.get_output_padding(),
192
+ use_ue8m0=False,
193
+ column_major_scales=True,
194
+ )
195
+ self.is_hopper = current_platform.is_device_capability(90)
196
+
197
+ @classmethod
198
+ def is_supported(cls, compute_capability=None):
199
+ if not CUTLASS_BLOCK_FP8_SUPPORTED:
200
+ return (
201
+ False,
202
+ "The device compute capability of"
203
+ f"{compute_capability} is not supported.",
204
+ )
205
+ return True, None
206
+
207
+ @classmethod
208
+ def can_implement(cls, config: FP8ScaledMMLinearLayerConfig):
209
+ can_implement_base, reason = super().can_implement(config)
210
+ if not can_implement_base:
211
+ return can_implement_base, reason
212
+
213
+ act_quant_desc = config.activation_quant_key.scale
214
+ if act_quant_desc.group_shape != GroupShape(1, 128):
215
+ return (
216
+ False,
217
+ "Supports only dynamic per token group activation "
218
+ "quantization with group_shape=(1,128).",
219
+ )
220
+ return True, None
221
+
222
+ def apply_block_scaled_mm(
223
+ self,
224
+ A: torch.Tensor,
225
+ B: torch.Tensor,
226
+ As: torch.Tensor,
227
+ Bs: torch.Tensor,
228
+ ) -> torch.Tensor:
229
+ out_dtype = self.config.out_dtype
230
+ if self.is_hopper:
231
+ return torch.ops.vllm.padded_cutlass(
232
+ A,
233
+ B,
234
+ As,
235
+ Bs,
236
+ list(self.weight_group_shape),
237
+ out_dtype,
238
+ )
239
+ else:
240
+ return ops.cutlass_scaled_mm(
241
+ A,
242
+ B.T,
243
+ out_dtype=out_dtype,
244
+ scale_a=As,
245
+ scale_b=Bs.T,
246
+ )
247
+
248
+
249
+def cutlass_scaled_mm(
250
+ A: torch.Tensor,
251
+ B: torch.Tensor,
252
+ As: torch.Tensor,
253
+ Bs: torch.Tensor,
254
+ block_size: list[int],
255
+ output_dtype: torch.dtype = torch.float16,
256
+) -> torch.Tensor:
257
+ return ops.cutlass_scaled_mm(
258
+ A,
259
+ B.T,
260
+ out_dtype=output_dtype,
261
+ scale_a=As,
262
+ scale_b=Bs.T,
263
+ )
264
+
265
+
266
+def _padded_cutlass(
267
+ qx: torch.Tensor,
268
+ weight: torch.Tensor,
269
+ x_scale: torch.Tensor,
270
+ weight_scale: torch.Tensor,
271
+ block_size: list[int],
272
+ output_dtype: torch.dtype,
273
+) -> torch.Tensor:
274
+ pad_multiple = 4
275
+ dim = qx.shape[0]
276
+ padded = (
277
+ dim if dim % pad_multiple == 0 else dim + pad_multiple - (dim % pad_multiple)
278
+ )
279
+
280
+ has_pad = padded > dim
281
+
282
+ if has_pad:
283
+ padded_shape = [padded, *qx.shape[1:]]
284
+ padded_qx = torch.zeros(padded_shape, device=qx.device, dtype=qx.dtype)
285
+ padded_qx[0 : qx.shape[0], ...].copy_(qx)
286
+
287
+ padded_x_scale_shape = [*x_scale.shape[1:], padded]
288
+ padded_x_scale = torch.ones(
289
+ padded_x_scale_shape, device=x_scale.device, dtype=x_scale.dtype
290
+ ).permute(-1, -2)
291
+ padded_x_scale[0 : x_scale.shape[0], ...].copy_(x_scale)
292
+
293
+ output = cutlass_scaled_mm(
294
+ padded_qx, weight, padded_x_scale, weight_scale, block_size, output_dtype
295
+ )
296
+ return output[0 : qx.shape[0], ...]
297
+ else:
298
+ return cutlass_scaled_mm(
299
+ qx, weight, x_scale, weight_scale, block_size, output_dtype
300
+ )
301
+
302
+
303
+def _padded_cutlass_fake(
304
+ qx: torch.Tensor,
305
+ weight: torch.Tensor,
306
+ x_scale: torch.Tensor,
307
+ weight_scale: torch.Tensor,
308
+ block_size: list[int],
309
+ output_dtype: torch.dtype,
310
+) -> torch.Tensor:
311
+ return torch.empty(
312
+ (qx.size(0), weight.size(0)), dtype=output_dtype, device=qx.device
313
+ )
314
+
315
+
316
+direct_register_custom_op(
317
+ "padded_cutlass",
318
+ _padded_cutlass,
319
+ fake_impl=_padded_cutlass_fake,
320
+)
vllm/model_executor/kernels/linear/scaled_mm/deep_gemm.py
new
+156
@@ -0,0 +1,156 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+import torch
5
+
6
+import vllm.envs as envs
7
+from vllm.config import get_current_vllm_config
8
+from vllm.model_executor.layers.quantization.input_quant_fp8 import QuantFP8
9
+from vllm.model_executor.layers.quantization.utils.fp8_utils import (
10
+ deepgemm_post_process_fp8_weight_block,
11
+)
12
+from vllm.model_executor.layers.quantization.utils.quant_utils import (
13
+ GroupShape,
14
+)
15
+from vllm.model_executor.utils import replace_parameter
16
+from vllm.platforms import current_platform
17
+from vllm.utils.deep_gemm import (
18
+ fp8_gemm_nt,
19
+ is_deep_gemm_e8m0_used,
20
+ is_deep_gemm_supported,
21
+ should_auto_disable_deep_gemm,
22
+ should_use_deepgemm_for_fp8_linear,
23
+)
24
+from vllm.utils.torch_utils import direct_register_custom_op
25
+
26
+from .BlockScaledMMLinearKernel import (
27
+ Fp8BlockScaledMMLinearKernel,
28
+ FP8ScaledMMLinearLayerConfig,
29
+)
30
+
31
+
32
+class DeepGemmFp8BlockScaledMMKernel(Fp8BlockScaledMMLinearKernel):
33
+ def __init__(self, config: FP8ScaledMMLinearLayerConfig):
34
+ super().__init__(config)
35
+ self.use_deep_gemm_e8m0 = is_deep_gemm_e8m0_used()
36
+ act_scale_descriptor = config.activation_quant_key.scale
37
+ self.is_deep_gemm_supported = is_deep_gemm_supported()
38
+ self.quant_fp8 = QuantFP8(
39
+ static=False,
40
+ group_shape=act_scale_descriptor.group_shape,
41
+ use_ue8m0=self.use_deep_gemm_e8m0,
42
+ tma_aligned_scales=envs.VLLM_USE_DEEP_GEMM_TMA_ALIGNED_SCALES,
43
+ column_major_scales=True,
44
+ )
45
+
46
+ @classmethod
47
+ def is_supported(cls, compute_capability=None):
48
+ if not current_platform.is_cuda():
49
+ return False, "DeepGEMM is only supported on cuda platform"
50
+ if not is_deep_gemm_supported():
51
+ return False, "Currently, only Hopper and Blackwell GPUs are supported."
52
+ return True, None
53
+
54
+ @classmethod
55
+ def can_implement(cls, config):
56
+ can_implement_base, reason = super().can_implement(config)
57
+ if not can_implement_base:
58
+ return can_implement_base, reason
59
+ if config.out_dtype != torch.bfloat16:
60
+ return (False, "Supports only output dtype of bfloat16")
61
+
62
+ act_quant_desc = config.activation_quant_key.scale
63
+ if act_quant_desc.group_shape != GroupShape(1, 128):
64
+ return (
65
+ False,
66
+ "Supports only dynamic per token group activation "
67
+ "quantization with group_shape=(1,128).",
68
+ )
69
+ model_config = get_current_vllm_config().model_config
70
+
71
+ if model_config is None:
72
+ return False, "Model configuration is required."
73
+
74
+ model_type = getattr(model_config.hf_text_config, "model_type", None)
75
+ if should_auto_disable_deep_gemm(model_type):
76
+ return False, f"Should not use deepgemm for model {model_type}"
77
+
78
+ if not should_use_deepgemm_for_fp8_linear(
79
+ config.out_dtype, config.weight_shape
80
+ ):
81
+ return False, "The provided metadata is not supported."
82
+ return True, None
83
+
84
+ def process_weights_after_loading(self, layer):
85
+ super().process_weights_after_loading(layer)
86
+ params = self._get_layer_params(layer)
87
+ assert layer.weight_block_size is not None
88
+
89
+ if self.is_deep_gemm_supported:
90
+ weight_scale_invs = params.weight_scale_inv
91
+ scale_attr = (
92
+ params.WEIGHT_SCALE_INV
93
+ if weight_scale_invs is not None
94
+ else params.WEIGHT_SCALE
95
+ )
96
+ dg_weight, dg_weight_scale = deepgemm_post_process_fp8_weight_block(
97
+ wq=params.weight,
98
+ ws=weight_scale_invs
99
+ if weight_scale_invs is not None
100
+ else params.weight_scale,
101
+ quant_block_shape=tuple(layer.weight_block_size),
102
+ use_e8m0=self.use_deep_gemm_e8m0,
103
+ )
104
+ replace_parameter(layer, params.WEIGHT, dg_weight)
105
+ replace_parameter(layer, scale_attr, dg_weight_scale)
106
+
107
+ def apply_block_scaled_mm(
108
+ self,
109
+ A: torch.Tensor,
110
+ B: torch.Tensor,
111
+ As: torch.Tensor,
112
+ Bs: torch.Tensor,
113
+ ) -> torch.Tensor:
114
+ out_dtype = self.config.out_dtype
115
+ output = torch.empty(
116
+ (A.shape[0], B.shape[0]),
117
+ dtype=out_dtype,
118
+ device=A.device,
119
+ )
120
+ torch.ops.vllm.fp8_gemm_nt_op(A, As, B, Bs, output, self.use_deep_gemm_e8m0)
121
+ return output
122
+
123
+
124
+def _fp8_gemm_nt_op(
125
+ q_input: torch.Tensor,
126
+ input_scale: torch.Tensor,
127
+ weight: torch.Tensor,
128
+ weight_scale: torch.Tensor,
129
+ output: torch.Tensor,
130
+ use_deep_gemm_e8m0: bool,
131
+) -> None:
132
+ fp8_gemm_nt(
133
+ (q_input, input_scale),
134
+ (weight, weight_scale),
135
+ output,
136
+ is_deep_gemm_e8m0_used=use_deep_gemm_e8m0,
137
+ )
138
+
139
+
140
+def _fp8_gemm_nt_op_fake(
141
+ q_input: torch.Tensor,
142
+ input_scale: torch.Tensor,
143
+ weight: torch.Tensor,
144
+ weight_scale: torch.Tensor,
145
+ output: torch.Tensor,
146
+ use_deep_gemm_e8m0: bool,
147
+) -> None:
148
+ return None
149
+
150
+
151
+direct_register_custom_op(
152
+ "fp8_gemm_nt_op",
153
+ _fp8_gemm_nt_op,
154
+ mutates_args=["output"],
155
+ fake_impl=_fp8_gemm_nt_op_fake,
156
+)
vllm/model_executor/kernels/linear/scaled_mm/flashinfer.py
+275
-1
@@ -2,11 +2,32 @@
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
5
+from typing import ClassVar
6
+
7
import torch
8
9
+import vllm.envs as envs
10
+from vllm.model_executor.layers.quantization.utils.fp8_utils import (
11
+ per_token_group_quant_fp8,
12
+)
13
+from vllm.model_executor.layers.quantization.utils.quant_utils import (
14
+ GroupShape,
15
+)
16
from vllm.platforms import current_platform
8
-from vllm.utils.flashinfer import flashinfer_scaled_fp8_mm, has_flashinfer
17
+from vllm.utils.flashinfer import (
18
+ flashinfer_fp8_blockscale_gemm,
19
+ flashinfer_scaled_fp8_mm,
20
+ has_flashinfer,
21
+ is_flashinfer_fp8_blockscale_gemm_supported,
22
+ should_use_flashinfer_for_blockscale_fp8_gemm,
23
+)
24
+from vllm.utils.torch_utils import direct_register_custom_op
25
26
+from .BlockScaledMMLinearKernel import (
27
+ Fp8BlockScaledDynamicMMLinearKernel,
28
+ Fp8BlockScaledMMLinearKernel,
29
+)
30
+from .deep_gemm import DeepGemmFp8BlockScaledMMKernel, fp8_gemm_nt
31
from .ScaledMMLinearKernel import (
32
FP8ScaledMMLinearKernel,
33
FP8ScaledMMLinearLayerConfig,
@@ -55,3 +76,256 @@ class FlashInferFP8ScaledMMLinearKernel(FP8ScaledMMLinearKernel):
76
return flashinfer_scaled_fp8_mm(
77
A, B, out_dtype=out_dtype, scale_a=As, scale_b=Bs, bias=bias
78
)
79
+
80
+
81
+class FlashInferFp8BlockScaledMMKernel(Fp8BlockScaledMMLinearKernel):
82
+ # FlashInfer accepts BF16 input and handles FP8 conversion internally.
83
+ apply_input_quant: ClassVar[bool] = False
84
+
85
+ def __init__(self, config: FP8ScaledMMLinearLayerConfig) -> None:
86
+ super().__init__(config)
87
+
88
+ @classmethod
89
+ def can_implement(cls, config: FP8ScaledMMLinearLayerConfig):
90
+ can_implement_base, reason = super().can_implement(config)
91
+ if not can_implement_base:
92
+ return can_implement_base, reason
93
+
94
+ act_quant_desc = config.activation_quant_key.scale
95
+ if act_quant_desc.group_shape != GroupShape(1, 128):
96
+ return (
97
+ False,
98
+ "Supports only dynamic per token group activation "
99
+ "quantization with group_shape=(1,128).",
100
+ )
101
+
102
+ if not should_use_flashinfer_for_blockscale_fp8_gemm(
103
+ is_flashinfer_fp8_blockscale_gemm_supported(),
104
+ config.out_dtype,
105
+ config.input_dtype,
106
+ config.weight_quant_key.dtype,
107
+ config.weight_shape,
108
+ ):
109
+ return (
110
+ False,
111
+ "The provided metadata is not supported.",
112
+ )
113
+
114
+ return True, None
115
+
116
+ @classmethod
117
+ def is_supported(cls, compute_capability=None):
118
+ if not current_platform.is_cuda():
119
+ return False, "only cuda devices are supported."
120
+
121
+ if not is_flashinfer_fp8_blockscale_gemm_supported():
122
+ return False, "FlashInfer block-scale FP8 GEMM is not available."
123
+
124
+ return True, None
125
+
126
+ def apply_block_scaled_mm(
127
+ self,
128
+ A: torch.Tensor,
129
+ B: torch.Tensor,
130
+ As: torch.Tensor,
131
+ Bs: torch.Tensor,
132
+ ) -> torch.Tensor:
133
+ # A is BF16 — FlashInfer handles FP8 conversion internally.
134
+ # As is a placeholder (apply_input_quant=False) and is not used here.
135
+ return torch.ops.vllm.flashinfer_fp8_blockscale_gemm(
136
+ A, # BF16 input
137
+ B, # FP8 weight
138
+ Bs, # Weight scales
139
+ )
140
+
141
+
142
+class FlashInferFp8DeepGEMMDynamicBlockScaledKernel(
143
+ Fp8BlockScaledDynamicMMLinearKernel
144
+):
145
+ """
146
+ Conditional FlashInfer / DeepGEMM FP8 block-scaled GEMM.
147
+
148
+ Dispatches between two kernels based on input batch size:
149
+ - Small batches (M < 32): FlashInfer's swapAB trick for better utilisation.
150
+ - Large batches (M >= 32): DeepGEMM for peak throughput.
151
+
152
+ apply_input_quant is False because FlashInfer accepts BF16 input and
153
+ handles FP8 conversion internally. The DeepGEMM branch therefore
154
+ quantises BF16→FP8 inside apply_mm via a closure before dispatching to
155
+ the DeepGEMM kernel — keeping both branches compatible with the single
156
+ BF16 tensor operand list passed by torch.cond.
157
+ """
158
+
159
+ base_type: ClassVar[type[FlashInferFp8BlockScaledMMKernel]] = (
160
+ FlashInferFp8BlockScaledMMKernel
161
+ )
162
+ fallback_type: ClassVar[type[DeepGemmFp8BlockScaledMMKernel]] = (
163
+ DeepGemmFp8BlockScaledMMKernel
164
+ )
165
+ apply_input_quant: ClassVar[bool] = False
166
+
167
+ def __init__(self, config: FP8ScaledMMLinearLayerConfig):
168
+ super().__init__(config)
169
+ self.base: FlashInferFp8BlockScaledMMKernel
170
+ self.fallback: DeepGemmFp8BlockScaledMMKernel
171
+
172
+ def process_weights_after_loading(self, layer: torch.nn.Module):
173
+ # DeepGEMM need post-processing; both kernels share the same
174
+ # parameter tensor layout so processing once is sufficient.
175
+ self.fallback.process_weights_after_loading(layer)
176
+
177
+ def apply_block_scaled_mm(
178
+ self,
179
+ A: torch.Tensor,
180
+ B: torch.Tensor,
181
+ As: torch.Tensor,
182
+ Bs: torch.Tensor,
183
+ ) -> torch.Tensor:
184
+ group_size = self.weight_group_shape.col
185
+ use_deep_gemm_e8m0 = self.fallback.use_deep_gemm_e8m0
186
+
187
+ return torch.ops.vllm.dynamic_flashinfer_deepgemm_blockscale_gemm(
188
+ A, B, Bs, group_size, use_deep_gemm_e8m0
189
+ )
190
+
191
+
192
+def _flashinfer_fp8_blockscale_gemm_impl(
193
+ input: torch.Tensor,
194
+ weight: torch.Tensor,
195
+ weight_scale: torch.Tensor,
196
+) -> torch.Tensor:
197
+ return flashinfer_fp8_blockscale_gemm(
198
+ input=input,
199
+ weight=weight,
200
+ weight_scale=weight_scale,
201
+ out_dtype=torch.bfloat16,
202
+ )
203
+
204
+
205
+def _flashinfer_fp8_blockscale_gemm_fake(
206
+ input: torch.Tensor,
207
+ weight: torch.Tensor,
208
+ weight_scale: torch.Tensor,
209
+) -> torch.Tensor:
210
+ """
211
+ Required fake/meta implementation for torch.compile graph tracing.
212
+ """
213
+ return torch.empty(
214
+ input.shape[0], weight.shape[0], dtype=torch.bfloat16, device=input.device
215
+ )
216
+
217
+
218
+direct_register_custom_op(
219
+ "flashinfer_fp8_blockscale_gemm",
220
+ _flashinfer_fp8_blockscale_gemm_impl,
221
+ fake_impl=_flashinfer_fp8_blockscale_gemm_fake,
222
+)
223
+
224
+
225
+def _dynamic_flashinfer_deepgemm_blockscale_gemm_impl(
226
+ input: torch.Tensor,
227
+ weight: torch.Tensor,
228
+ weight_scale: torch.Tensor,
229
+ group_size: int,
230
+ use_deep_gemm_e8m0: bool,
231
+) -> torch.Tensor:
232
+ """
233
+ Conditional FlashInfer FP8 blockscale GEMM with batch-size-dependent selection.
234
+
235
+ This function switches between two optimized kernels based on the input batch size:
236
+ - For small batches (M < 32): Uses FlashInfer's DeepGEMM swapAB optimization.
237
+ - For larger batches (M >= 32): Uses the official DeepGEMM kernel.
238
+
239
+ The conditional logic must use torch.cond() instead of a simple if-else statement
240
+ to maintain compatibility with torch.compile graph compilation.
241
+
242
+ This batch-size-dependent selection is essential for maintaining model accuracy.
243
+ Benchmarks on GSM8K show a significant accuracy gap (88% vs 95%) for DeepSeek-V3.1
244
+ when using FlashInfer's DeepGEMM on M>=32. The M < 32 strategy fixes the accuracy
245
+ drop.
246
+
247
+ Args:
248
+ input: Input tensor of shape (batch_size, input_dim) in FP8 format
249
+ weight: Weight tensor of shape (output_dim, input_dim) in FP8 format
250
+ weight_scale: Scale factors for weight quantization (per-group)
251
+ group_size: Quantization group size for the weight tensor
252
+ use_deep_gemm_e8m0: Whether to use the E8M0 format in DeepGEMM quantization
253
+
254
+ Returns:
255
+ Output tensor of shape (batch_size, output_dim) in bfloat16 format
256
+ """
257
+
258
+ def run_flashinfer_deepgemm_swapAB(
259
+ input: torch.Tensor,
260
+ weight: torch.Tensor,
261
+ weight_scale: torch.Tensor,
262
+ ) -> torch.Tensor:
263
+ return flashinfer_fp8_blockscale_gemm(
264
+ input=input,
265
+ weight=weight,
266
+ weight_scale=weight_scale,
267
+ out_dtype=torch.bfloat16,
268
+ )
269
+
270
+ def run_deepgemm(
271
+ input: torch.Tensor,
272
+ weight: torch.Tensor,
273
+ weight_scale: torch.Tensor,
274
+ ) -> torch.Tensor:
275
+ q_input, input_scale = per_token_group_quant_fp8(
276
+ input,
277
+ group_size=group_size,
278
+ column_major_scales=True,
279
+ use_ue8m0=use_deep_gemm_e8m0,
280
+ )
281
+ output = torch.empty(
282
+ (q_input.shape[0], weight.shape[0]),
283
+ dtype=torch.bfloat16,
284
+ device=q_input.device,
285
+ )
286
+ fp8_gemm_nt(
287
+ (q_input, input_scale),
288
+ (weight, weight_scale),
289
+ output,
290
+ is_deep_gemm_e8m0_used=use_deep_gemm_e8m0,
291
+ )
292
+ return output
293
+
294
+ if envs.VLLM_BATCH_INVARIANT:
295
+ return run_deepgemm(input, weight, weight_scale)
296
+
297
+ condition = input.shape[0] < 32
298
+
299
+ # PyTorch's torch.compile cannot handle input-dependent control flow in standard
300
+ # Python conditionals. torch.cond() explicitly registers both code paths in the
301
+ # computation graph, allowing torch.compile to capture both branches.
302
+ # without torch.cond, the M < 32 condition won't be able to be captured by torch
303
+ # compile
304
+ return torch.cond(
305
+ condition,
306
+ run_flashinfer_deepgemm_swapAB,
307
+ run_deepgemm,
308
+ (input, weight, weight_scale),
309
+ )
310
+
311
+
312
+def _dynamic_flashinfer_deepgemm_blockscale_gemm_fake(
313
+ input: torch.Tensor,
314
+ weight: torch.Tensor,
315
+ weight_scale: torch.Tensor,
316
+ group_size: int,
317
+ use_deep_gemm_e8m0: bool,
318
+) -> torch.Tensor:
319
+ """
320
+ Required fake/meta implementation for torch.compile graph tracing.
321
+ """
322
+ return torch.empty(
323
+ input.shape[0], weight.shape[0], dtype=torch.bfloat16, device=input.device
324
+ )
325
+
326
+
327
+direct_register_custom_op(
328
+ "dynamic_flashinfer_deepgemm_blockscale_gemm",
329
+ _dynamic_flashinfer_deepgemm_blockscale_gemm_impl,
330
+ fake_impl=_dynamic_flashinfer_deepgemm_blockscale_gemm_fake,
331
+)
vllm/model_executor/kernels/linear/scaled_mm/triton.py
+68
@@ -13,7 +13,11 @@ from vllm.model_executor.layers.quantization.utils.w8a8_utils import (
13
convert_to_channelwise,
14
)
15
from vllm.platforms import current_platform
16
+from vllm.utils.torch_utils import direct_register_custom_op
17
18
+from .BlockScaledMMLinearKernel import (
19
+ Fp8BlockScaledMMLinearKernel,
20
+)
21
from .cutlass import CutlassInt8ScaledMMLinearKernel
22
from .ScaledMMLinearKernel import (
23
Int8ScaledMMLinearLayerConfig,
@@ -150,3 +154,67 @@ class TritonInt8ScaledMMLinearKernel(CutlassInt8ScaledMMLinearKernel):
154
out -= (x_s * w_s_row * azp_adj).to(x.dtype)
155
156
return out
157
+
158
+
159
+class TritonFp8BlockScaledMMKernel(Fp8BlockScaledMMLinearKernel):
160
+ @classmethod
161
+ def is_supported(cls, compute_capability=None):
162
+ if not current_platform.is_cuda_alike():
163
+ return False, "only cuda like devices are supported."
164
+ return True, None
165
+
166
+ def apply_block_scaled_mm(
167
+ self,
168
+ A: torch.Tensor,
169
+ B: torch.Tensor,
170
+ As: torch.Tensor,
171
+ Bs: torch.Tensor,
172
+ ) -> torch.Tensor:
173
+ return torch.ops.vllm.w8a8_triton_block_scaled_mm_func(
174
+ A,
175
+ B,
176
+ As,
177
+ Bs,
178
+ list(self.weight_group_shape),
179
+ self.config.out_dtype,
180
+ )
181
+
182
+
183
+# TODO we should be able to change the type of block_size to GroupShape
184
+# after we resolve GroupShape compilation issue
185
+# https://github.com/vllm-project/vllm/issues/25270
186
+def _w8a8_triton_block_scaled_mm_func(
187
+ qx: torch.Tensor,
188
+ weight: torch.Tensor,
189
+ x_scale: torch.Tensor,
190
+ weight_scale: torch.Tensor,
191
+ block_size: list[int],
192
+ output_dtype: torch.dtype,
193
+) -> torch.Tensor:
194
+ from vllm.model_executor.layers.quantization.utils.fp8_utils import (
195
+ w8a8_triton_block_scaled_mm,
196
+ )
197
+
198
+ return w8a8_triton_block_scaled_mm(
199
+ qx, weight, x_scale, weight_scale, block_size, output_dtype
200
+ )
201
+
202
+
203
+def _w8a8_triton_block_scaled_mm_fake(
204
+ qx: torch.Tensor,
205
+ weight: torch.Tensor,
206
+ x_scale: torch.Tensor,
207
+ weight_scale: torch.Tensor,
208
+ block_size: list[int],
209
+ output_dtype: torch.dtype,
210
+) -> torch.Tensor:
211
+ return torch.empty(
212
+ (qx.size(0), weight.size(0)), dtype=output_dtype, device=qx.device
213
+ )
214
+
215
+
216
+direct_register_custom_op(
217
+ "w8a8_triton_block_scaled_mm_func",
218
+ _w8a8_triton_block_scaled_mm_func,
219
+ fake_impl=_w8a8_triton_block_scaled_mm_fake,
220
+)
vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_w8a8_fp8.py
+28
-31
@@ -8,6 +8,7 @@ from compressed_tensors.quantization import QuantizationArgs, QuantizationStrate
8
from torch.nn import Parameter
9
10
from vllm._aiter_ops import rocm_aiter_ops
11
+from vllm.config import get_current_vllm_config
12
from vllm.logger import init_logger
13
from vllm.model_executor.kernels.linear import (
14
init_fp8_linear_kernel,
@@ -16,18 +17,16 @@ from vllm.model_executor.layers.quantization.compressed_tensors.schemes import (
17
CompressedTensorsScheme,
18
)
19
from vllm.model_executor.layers.quantization.utils.fp8_utils import (
19
- W8A8BlockFp8LinearOp,
20
create_fp8_input_scale,
21
create_fp8_scale_parameter,
22
create_fp8_weight_parameter,
23
- maybe_post_process_fp8_weight_block,
24
- process_fp8_weight_block_strategy,
23
process_fp8_weight_channel_strategy,
24
process_fp8_weight_tensor_strategy,
25
validate_fp8_block_shape,
26
)
27
from vllm.model_executor.layers.quantization.utils.quant_utils import (
28
GroupShape,
29
+ create_fp8_quant_key,
30
kFp8DynamicTokenSym,
31
kFp8StaticTensorSym,
32
kFp8StaticTokenSym,
@@ -67,6 +66,7 @@ class CompressedTensorsW8A8Fp8(CompressedTensorsScheme):
66
self.weight_quant = weight_quant
67
self.strategy = weight_quant.strategy
68
self.out_dtype = torch.get_default_dtype()
69
+ self.input_dtype = get_current_vllm_config().model_config.dtype
70
self.is_static_input_scheme = is_static_input_scheme
71
self.weight_block_size = self.weight_quant.block_structure
72
@@ -75,21 +75,18 @@ class CompressedTensorsW8A8Fp8(CompressedTensorsScheme):
75
self.use_aiter_and_is_supported = rocm_aiter_ops.is_linear_fp8_enabled()
76
assert not self.is_static_input_scheme
77
self.act_q_group_shape = GroupShape(1, self.weight_block_size[0])
78
- self.w8a8_block_fp8_linear = W8A8BlockFp8LinearOp(
79
- weight_group_shape=GroupShape(*self.weight_block_size),
80
- act_quant_group_shape=self.act_q_group_shape,
81
- cutlass_block_fp8_supported=self.cutlass_block_fp8_supported,
82
- use_aiter_and_is_supported=self.use_aiter_and_is_supported,
78
+
79
+ self.weight_quant_key = create_fp8_quant_key(
80
+ static=True, group_shape=GroupShape(*self.weight_block_size)
81
)
84
- else:
85
- activation_quant_key = activation_quant_key_mapping[is_static_input_scheme]
86
- weight_quant_key = weight_quant_key_mapping[self.strategy]
87
- self.fp8_linear = init_fp8_linear_kernel(
88
- activation_quant_key=activation_quant_key,
89
- weight_quant_key=weight_quant_key,
90
- out_dtype=self.out_dtype,
91
- module_name=self.__class__.__name__,
82
+ self.activation_quant_key = create_fp8_quant_key(
83
+ static=False, group_shape=self.act_q_group_shape
84
)
85
+ else:
86
+ self.activation_quant_key = activation_quant_key_mapping[
87
+ is_static_input_scheme
88
+ ]
89
+ self.weight_quant_key = weight_quant_key_mapping[self.strategy]
90
91
@classmethod
92
def get_min_capability(cls) -> int:
@@ -146,6 +143,15 @@ class CompressedTensorsW8A8Fp8(CompressedTensorsScheme):
143
input_scale = create_fp8_input_scale(output_partition_sizes, weight_loader)
144
layer.register_parameter("input_scale", input_scale)
145
146
+ self.fp8_linear = init_fp8_linear_kernel(
147
+ activation_quant_key=self.activation_quant_key,
148
+ weight_quant_key=self.weight_quant_key,
149
+ weight_shape=layer.weight.shape,
150
+ input_dtype=self.input_dtype,
151
+ out_dtype=self.out_dtype,
152
+ module_name=self.__class__.__name__,
153
+ )
154
+
155
def process_weights_after_loading(self, layer) -> None:
156
if self.strategy == QuantizationStrategy.TENSOR:
157
weight, weight_scale, input_scale = process_fp8_weight_tensor_strategy(
@@ -163,10 +169,12 @@ class CompressedTensorsW8A8Fp8(CompressedTensorsScheme):
169
170
elif self.strategy == QuantizationStrategy.BLOCK:
171
assert self.is_static_input_scheme is False
166
- weight, weight_scale = process_fp8_weight_block_strategy(
167
- layer.weight, layer.weight_scale
168
- )
169
- input_scale = None
172
+ self.fp8_linear.process_weights_after_loading(layer)
173
+
174
+ layer.input_scale = None
175
+ # fp8_linear.process_weights_after_loading applies the post process
176
+ # and reassigns the weight and weight_scale buffers to layer attributes.
177
+ return
178
179
else:
180
raise ValueError(
@@ -185,8 +193,6 @@ class CompressedTensorsW8A8Fp8(CompressedTensorsScheme):
193
layer.input_scale = Parameter(layer.input_scale.max(), requires_grad=False)
194
else:
195
layer.input_scale = None
188
- if self.strategy == QuantizationStrategy.BLOCK:
189
- maybe_post_process_fp8_weight_block(layer)
196
197
if hasattr(self, "fp8_linear"):
198
self.fp8_linear.process_weights_after_loading(layer)
@@ -197,13 +203,4 @@ class CompressedTensorsW8A8Fp8(CompressedTensorsScheme):
203
x: torch.Tensor,
204
bias: torch.Tensor | None = None,
205
) -> torch.Tensor:
200
- if self.weight_block_size is not None:
201
- return self.w8a8_block_fp8_linear.apply(
202
- input=x,
203
- weight=layer.weight,
204
- weight_scale=layer.weight_scale,
205
- input_scale=layer.input_scale,
206
- bias=bias,
207
- )
208
-
206
return self.fp8_linear.apply_weights(layer, x, bias)
vllm/model_executor/layers/quantization/fbgemm_fp8.py
+11
-6
@@ -7,6 +7,7 @@ import torch
7
from torch.nn import Module
8
from torch.nn.parameter import Parameter
9
10
+from vllm.config import get_current_vllm_config
11
from vllm.logger import init_logger
12
from vllm.model_executor.kernels.linear import (
13
init_fp8_linear_kernel,
@@ -93,12 +94,7 @@ class FBGEMMFp8LinearMethod(LinearMethodBase):
94
def __init__(self, quant_config: FBGEMMFp8Config):
95
self.quant_config = quant_config
96
self.out_dtype = torch.get_default_dtype()
96
- self.fp8_linear = init_fp8_linear_kernel(
97
- activation_quant_key=kFp8DynamicTokenSym,
98
- weight_quant_key=kFp8StaticTokenSym,
99
- out_dtype=torch.get_default_dtype(),
100
- module_name=self.__class__.__name__,
101
- )
97
+ self.input_dtype = get_current_vllm_config().model_config.dtype
98
99
def create_weights(
100
self,
@@ -149,6 +145,15 @@ class FBGEMMFp8LinearMethod(LinearMethodBase):
145
)
146
layer.input_scale_ub = input_scale_ub
147
148
+ self.fp8_linear = init_fp8_linear_kernel(
149
+ activation_quant_key=kFp8DynamicTokenSym,
150
+ weight_quant_key=kFp8StaticTokenSym,
151
+ weight_shape=layer.weight.shape,
152
+ input_dtype=self.input_dtype,
153
+ out_dtype=self.out_dtype,
154
+ module_name=self.__class__.__name__,
155
+ )
156
+
157
def process_weights_after_loading(self, layer: Module) -> None:
158
# required by torch.compile
159
layer.weight_scale = Parameter(layer.weight_scale.data, requires_grad=False)
vllm/model_executor/layers/quantization/fp8.py
+54
-59
@@ -10,7 +10,7 @@ from torch.utils._python_dispatch import TorchDispatchMode
10
import vllm.envs as envs
11
import vllm.model_executor.layers.fused_moe.modular_kernel as mk
12
from vllm import _custom_ops as ops
13
-from vllm._aiter_ops import rocm_aiter_ops
13
+from vllm.config import get_current_vllm_config
14
from vllm.distributed import get_tensor_model_parallel_world_size
15
from vllm.logger import init_logger
16
from vllm.model_executor.kernels.linear import (
@@ -45,13 +45,10 @@ from vllm.model_executor.layers.quantization.base_config import (
45
)
46
from vllm.model_executor.layers.quantization.kv_cache import BaseKVCacheMethod
47
from vllm.model_executor.layers.quantization.utils.fp8_utils import (
48
- W8A8BlockFp8LinearOp,
48
create_fp8_input_scale,
49
create_fp8_scale_parameter,
50
create_fp8_weight_parameter,
52
- maybe_post_process_fp8_weight_block,
51
process_fp8_input_tensor_strategy_moe,
54
- process_fp8_weight_block_strategy,
52
process_fp8_weight_tensor_strategy,
53
process_fp8_weight_tensor_strategy_moe,
54
validate_fp8_block_shape,
@@ -61,6 +58,7 @@ from vllm.model_executor.layers.quantization.utils.marlin_utils import (
58
)
59
from vllm.model_executor.layers.quantization.utils.quant_utils import (
60
GroupShape,
61
+ create_fp8_quant_key,
62
is_layer_skipped,
63
kFp8Dynamic128Sym,
64
kFp8DynamicTensorSym,
@@ -273,12 +271,13 @@ class Fp8LinearMethod(LinearMethodBase):
271
self.quant_config = quant_config
272
self.cutlass_block_fp8_supported = cutlass_block_fp8_supported()
273
self.out_dtype = torch.get_default_dtype()
274
+ self.input_dtype = get_current_vllm_config().model_config.dtype
275
276
# For GPUs that lack FP8 hardware support, we can leverage the Marlin
277
# kernel for fast weight-only FP8 quantization
278
self.marlin_input_dtype = None
279
+ self.use_marlin = False
280
281
- self.use_aiter_and_is_supported = rocm_aiter_ops.is_linear_fp8_enabled()
281
if self.quant_config.use_deep_gemm is not None:
282
self.use_deep_gemm = self.quant_config.use_deep_gemm
283
else:
@@ -288,37 +287,26 @@ class Fp8LinearMethod(LinearMethodBase):
287
self.block_quant = self.weight_block_size is not None
288
self.act_q_static = self.quant_config.activation_scheme == "static"
289
291
- # Use per-token quantization for better perf if dynamic and cutlass
292
- if self.act_q_static:
293
- activation_quant_key = kFp8StaticTensorSym
294
- elif cutlass_fp8_supported():
295
- activation_quant_key = kFp8DynamicTokenSym
296
- else:
297
- activation_quant_key = kFp8DynamicTensorSym
298
-
290
if self.block_quant:
300
- weight_quant_key = kFp8Static128BlockSym
301
- else:
302
- weight_quant_key = kFp8StaticTensorSym
303
-
304
- self.fp8_linear = init_fp8_linear_kernel(
305
- activation_quant_key=activation_quant_key,
306
- weight_quant_key=weight_quant_key,
307
- out_dtype=torch.get_default_dtype(),
308
- module_name=self.__class__.__name__,
309
- )
310
- self.use_marlin = isinstance(self.fp8_linear, MarlinFP8ScaledMMLinearKernel)
311
-
312
- if self.block_quant and not self.use_marlin:
291
assert not self.act_q_static
292
assert self.weight_block_size is not None
315
- self.w8a8_block_fp8_linear = W8A8BlockFp8LinearOp(
316
- weight_group_shape=GroupShape(*self.weight_block_size),
317
- act_quant_group_shape=GroupShape(1, self.weight_block_size[0]),
318
- cutlass_block_fp8_supported=self.cutlass_block_fp8_supported,
319
- use_aiter_and_is_supported=self.use_aiter_and_is_supported,
320
- use_deep_gemm=self.use_deep_gemm,
293
+
294
+ self.activation_quant_key = create_fp8_quant_key(
295
+ static=self.act_q_static,
296
+ group_shape=GroupShape(1, self.weight_block_size[0]),
297
+ )
298
+ self.weight_quant_key = create_fp8_quant_key(
299
+ static=True, group_shape=GroupShape(*self.weight_block_size)
300
)
301
+ else:
302
+ self.weight_quant_key = kFp8StaticTensorSym
303
+ # Use per-token quantization for better perf if dynamic and cutlass
304
+ if self.act_q_static:
305
+ self.activation_quant_key = kFp8StaticTensorSym
306
+ elif cutlass_fp8_supported():
307
+ self.activation_quant_key = kFp8DynamicTokenSym
308
+ else:
309
+ self.activation_quant_key = kFp8DynamicTensorSym
310
311
def create_weights(
312
self,
@@ -384,6 +372,17 @@ class Fp8LinearMethod(LinearMethodBase):
372
set_weight_attrs(scale, {"scale_type": "input_scale"})
373
layer.register_parameter("input_scale", scale)
374
375
+ self.fp8_linear = init_fp8_linear_kernel(
376
+ activation_quant_key=self.activation_quant_key,
377
+ weight_quant_key=self.weight_quant_key,
378
+ weight_shape=layer.weight.shape,
379
+ input_dtype=self.input_dtype,
380
+ out_dtype=self.out_dtype,
381
+ module_name=self.__class__.__name__,
382
+ )
383
+
384
+ self.use_marlin = isinstance(self.fp8_linear, MarlinFP8ScaledMMLinearKernel)
385
+
386
def process_weights_after_loading(self, layer: Module) -> None:
387
if self.use_marlin:
388
# Only Marlin kernels support `marlin_input_dtype`; guard to avoid
@@ -398,13 +397,7 @@ class Fp8LinearMethod(LinearMethodBase):
397
if self.block_quant:
398
assert not self.act_q_static
399
401
- weight, weight_scale_inv = process_fp8_weight_block_strategy(
402
- layer.weight, layer.weight_scale_inv
403
- )
404
-
405
- # Update layer with new values
406
- replace_parameter(layer, "weight", weight.data)
407
- replace_parameter(layer, "weight_scale_inv", weight_scale_inv.data)
400
+ self.fp8_linear.process_weights_after_loading(layer)
401
402
# If checkpoint not serialized fp8, quantize the weights.
403
else:
@@ -435,9 +428,6 @@ class Fp8LinearMethod(LinearMethodBase):
428
else:
429
layer.input_scale = None
430
438
- if self.block_quant and self.use_deep_gemm:
439
- maybe_post_process_fp8_weight_block(layer)
440
-
431
def apply(
432
self,
433
layer: torch.nn.Module,
@@ -449,12 +439,10 @@ class Fp8LinearMethod(LinearMethodBase):
439
if envs.VLLM_BATCH_INVARIANT:
440
if self.block_quant:
441
assert self.weight_block_size is not None
452
- return self.w8a8_block_fp8_linear.apply(
453
- input=x,
454
- weight=layer.weight,
455
- weight_scale=layer.weight_scale_inv,
456
- input_scale=layer.input_scale,
457
- bias=bias,
442
+ return self.fp8_linear.apply_weights(
443
+ layer,
444
+ x,
445
+ bias,
446
)
447
else:
448
# per-tensor/channel: dequant to BF16 and run GEMM
@@ -483,17 +471,6 @@ class Fp8LinearMethod(LinearMethodBase):
471
if self.use_marlin:
472
return self.fp8_linear.apply_weights(layer, x, bias)
473
486
- if self.block_quant:
487
- assert self.weight_block_size is not None
488
-
489
- return self.w8a8_block_fp8_linear.apply(
490
- input=x,
491
- weight=layer.weight,
492
- weight_scale=layer.weight_scale_inv,
493
- input_scale=layer.input_scale,
494
- bias=bias,
495
- )
496
-
474
return self.fp8_linear.apply_weights(layer, x, bias)
475
476
@@ -538,6 +515,24 @@ class Fp8OnlineLinearMethod(Fp8LinearMethod):
515
516
initialize_online_processing(layer)
517
518
+ # TODO: remove this check once the following RFC is resolved.
519
+ # https://github.com/vllm-project/vllm/issues/33314
520
+ # This check is required because Mxfp8OnlineLinearMethod inherits from
521
+ # Fp8OnlineLinearMethod but only calls super().create_weights(), so we must
522
+ # skip the fp8_linear kernel creation.
523
+ if hasattr(self, "mxfp8_linear"):
524
+ return
525
+
526
+ self.fp8_linear = init_fp8_linear_kernel(
527
+ activation_quant_key=self.activation_quant_key,
528
+ weight_quant_key=self.weight_quant_key,
529
+ weight_shape=layer.weight.shape,
530
+ input_dtype=self.input_dtype,
531
+ out_dtype=self.out_dtype,
532
+ module_name=self.__class__.__name__,
533
+ )
534
+ self.use_marlin = isinstance(self.fp8_linear, MarlinFP8ScaledMMLinearKernel)
535
+
536
def process_weights_after_loading(self, layer: Module) -> None:
537
if getattr(layer, "_already_called_process_weights_after_loading", False):
538
return
vllm/model_executor/layers/quantization/modelopt.py
+44
-27
@@ -8,6 +8,7 @@ import torch
8
from torch.nn.parameter import Parameter
9
10
import vllm.model_executor.layers.fused_moe.modular_kernel as mk
11
+from vllm.config import get_current_vllm_config
12
from vllm.logger import init_logger
13
from vllm.model_executor.kernels.linear import init_fp8_linear_kernel
14
from vllm.model_executor.layers.attention import Attention, MLAAttention
@@ -56,7 +57,6 @@ from vllm.model_executor.layers.quantization.utils.flashinfer_utils import (
57
swap_w13_to_w31,
58
)
59
from vllm.model_executor.layers.quantization.utils.fp8_utils import (
59
- W8A8BlockFp8LinearOp,
60
process_fp8_input_tensor_strategy_moe,
61
process_fp8_weight_tensor_strategy_moe,
62
)
@@ -78,6 +78,7 @@ from vllm.model_executor.layers.quantization.utils.nvfp4_utils import (
78
)
79
from vllm.model_executor.layers.quantization.utils.quant_utils import (
80
GroupShape,
81
+ create_fp8_quant_key,
82
is_layer_skipped,
83
kFp8DynamicTokenSym,
84
kFp8StaticTensorSym,
@@ -86,7 +87,6 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import (
87
kNvfp4Static,
88
)
89
from vllm.model_executor.layers.quantization.utils.w8a8_utils import (
89
- cutlass_block_fp8_supported,
90
requantize_with_max_scale,
91
)
92
from vllm.model_executor.parameter import (
@@ -450,12 +450,8 @@ class ModelOptFp8LinearMethod(LinearMethodBase):
450
451
def __init__(self, quant_config: ModelOptFp8Config) -> None:
452
self.quant_config = quant_config
453
- self.fp8_linear = init_fp8_linear_kernel(
454
- activation_quant_key=kFp8StaticTensorSym,
455
- weight_quant_key=kFp8StaticTensorSym,
456
- out_dtype=torch.get_default_dtype(),
457
- module_name=self.__class__.__name__,
458
- )
453
+ self.out_dtype = torch.get_default_dtype()
454
+ self.input_dtype = get_current_vllm_config().model_config.dtype
455
456
def create_weights(
457
self,
@@ -505,6 +501,15 @@ class ModelOptFp8LinearMethod(LinearMethodBase):
501
scale[:] = torch.finfo(torch.float32).min
502
layer.register_parameter("input_scale", scale)
503
504
+ self.fp8_linear = init_fp8_linear_kernel(
505
+ activation_quant_key=kFp8StaticTensorSym,
506
+ weight_quant_key=kFp8StaticTensorSym,
507
+ weight_shape=layer.weight.shape,
508
+ input_dtype=self.input_dtype,
509
+ out_dtype=self.out_dtype,
510
+ module_name=self.__class__.__name__,
511
+ )
512
+
513
def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
514
weight = layer.weight
515
max_w_scale = layer.weight_scale.max()
@@ -536,12 +541,8 @@ class ModelOptFp8PcPtLinearMethod(LinearMethodBase):
541
542
def __init__(self, quant_config: ModelOptFp8Config) -> None:
543
self.quant_config = quant_config
539
- self.fp8_linear = init_fp8_linear_kernel(
540
- activation_quant_key=kFp8DynamicTokenSym,
541
- weight_quant_key=kFp8StaticTokenSym,
542
- out_dtype=torch.get_default_dtype(),
543
- module_name=self.__class__.__name__,
544
- )
544
+ self.out_dtype = torch.get_default_dtype()
545
+ self.input_dtype = get_current_vllm_config().model_config.dtype
546
547
def create_weights(
548
self,
@@ -587,6 +588,15 @@ class ModelOptFp8PcPtLinearMethod(LinearMethodBase):
588
weight_scale[:] = torch.finfo(torch.float32).min
589
layer.register_parameter("weight_scale", weight_scale)
590
591
+ self.fp8_linear = init_fp8_linear_kernel(
592
+ activation_quant_key=kFp8DynamicTokenSym,
593
+ weight_quant_key=kFp8StaticTokenSym,
594
+ weight_shape=layer.weight.shape,
595
+ input_dtype=self.input_dtype,
596
+ out_dtype=self.out_dtype,
597
+ module_name=self.__class__.__name__,
598
+ )
599
+
600
def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
601
layer.weight = Parameter(layer.weight.t(), requires_grad=False)
602
layer.weight_scale = Parameter(layer.weight_scale.data, requires_grad=False)
@@ -616,12 +626,16 @@ class ModelOptFp8PbWoLinearMethod(LinearMethodBase):
626
self.quant_config = quant_config
627
block_n, block_k = self._WEIGHT_BLOCK_SIZE
628
self.weight_block_size = list(self._WEIGHT_BLOCK_SIZE)
619
- self.w8a8_block_fp8_linear = W8A8BlockFp8LinearOp(
620
- weight_group_shape=GroupShape(block_n, block_k),
621
- act_quant_group_shape=GroupShape(1, block_k),
622
- cutlass_block_fp8_supported=cutlass_block_fp8_supported(),
623
- use_aiter_and_is_supported=False,
629
+
630
+ self.activation_quant_key = create_fp8_quant_key(
631
+ static=False, group_shape=GroupShape(1, block_k)
632
)
633
+ self.weight_quant_key = create_fp8_quant_key(
634
+ static=True, group_shape=GroupShape(block_n, block_k)
635
+ )
636
+
637
+ self.out_dtype = torch.get_default_dtype()
638
+ self.input_dtype = get_current_vllm_config().model_config.dtype
639
640
def create_weights(
641
self,
@@ -688,8 +702,17 @@ class ModelOptFp8PbWoLinearMethod(LinearMethodBase):
702
weight_scale[:] = torch.finfo(torch.float32).min
703
layer.register_parameter("weight_scale", weight_scale)
704
705
+ self.w8a8_block_fp8_linear = init_fp8_linear_kernel(
706
+ activation_quant_key=self.activation_quant_key,
707
+ weight_quant_key=self.weight_quant_key,
708
+ weight_shape=layer.weight.shape,
709
+ input_dtype=self.input_dtype,
710
+ out_dtype=self.out_dtype,
711
+ module_name=self.__class__.__name__,
712
+ )
713
+
714
def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
692
- # Keep weight in [out, in] layout for W8A8BlockFp8LinearOp.
715
+ # Keep weight in [out, in] layout for Fp8BlockScaledMMLinearKernel.
716
layer.weight = Parameter(layer.weight.data, requires_grad=False)
717
718
scale = layer.weight_scale
@@ -713,13 +736,7 @@ class ModelOptFp8PbWoLinearMethod(LinearMethodBase):
736
x: torch.Tensor,
737
bias: torch.Tensor | None = None,
738
) -> torch.Tensor:
716
- return self.w8a8_block_fp8_linear.apply(
717
- input=x,
718
- weight=layer.weight,
719
- weight_scale=layer.weight_scale,
720
- input_scale=None,
721
- bias=bias,
722
- )
739
+ return self.w8a8_block_fp8_linear.apply_weights(layer, x, bias)
740
741
742
class ModelOptFp8MoEMethod(FusedMoEMethodBase):
vllm/model_executor/layers/quantization/online/fp8.py
+57
-37
@@ -17,7 +17,7 @@ if TYPE_CHECKING:
17
18
import vllm.envs as envs
19
from vllm import _custom_ops as ops
20
-from vllm._aiter_ops import rocm_aiter_ops
20
+from vllm.config import get_current_vllm_config
21
from vllm.model_executor.kernels.linear import init_fp8_linear_kernel
22
from vllm.model_executor.layers.fused_moe import (
23
FusedMoEMethodBase,
@@ -28,13 +28,9 @@ from vllm.model_executor.layers.fused_moe.oracle.fp8 import (
28
from vllm.model_executor.layers.linear import (
29
LinearMethodBase,
30
)
31
-from vllm.model_executor.layers.quantization.utils.fp8_utils import (
32
- W8A8BlockFp8LinearOp,
33
- maybe_post_process_fp8_weight_block,
34
- process_fp8_weight_block_strategy,
35
-)
31
from vllm.model_executor.layers.quantization.utils.quant_utils import (
32
GroupShape,
33
+ create_fp8_quant_key,
34
kFp8Dynamic128Sym,
35
kFp8DynamicTensorSym,
36
kFp8DynamicTokenSym,
@@ -42,7 +38,6 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import (
38
kFp8StaticTensorSym,
39
)
40
from vllm.model_executor.layers.quantization.utils.w8a8_utils import (
45
- cutlass_block_fp8_supported,
41
cutlass_fp8_supported,
42
)
43
from vllm.model_executor.model_loader.reload.layerwise import (
@@ -51,7 +46,7 @@ from vllm.model_executor.model_loader.reload.layerwise import (
46
from vllm.model_executor.parameter import ModelWeightParameter
47
from vllm.model_executor.utils import replace_parameter, set_weight_attrs
48
from vllm.platforms import current_platform
54
-from vllm.utils.deep_gemm import is_deep_gemm_supported, per_block_cast_to_fp8
49
+from vllm.utils.deep_gemm import per_block_cast_to_fp8
50
51
# ---------------------------------------------------------------------------
52
# Online FP8 Linear Methods
@@ -64,6 +59,10 @@ class _Fp8OnlineLinearBase(LinearMethodBase):
59
60
uses_meta_device: bool = True
61
62
+ def __init__(self):
63
+ self.out_dtype = torch.get_default_dtype()
64
+ self.input_dtype = get_current_vllm_config().model_config.dtype
65
+
66
def create_weights(
67
self,
68
layer: torch.nn.Module,
@@ -103,18 +102,41 @@ class Fp8PerTensorOnlineLinearMethod(_Fp8OnlineLinearBase):
102
Loads fp16/bf16 weights and quantizes them per-tensor during loading."""
103
104
def __init__(self):
106
- self.out_dtype = torch.get_default_dtype()
105
+ super().__init__()
106
107
+ self.weight_quant_key = kFp8StaticTensorSym
108
# Use per-token quantization for better perf if dynamic and cutlass
109
if cutlass_fp8_supported():
110
- activation_quant_key = kFp8DynamicTokenSym
110
+ self.activation_quant_key = kFp8DynamicTokenSym
111
else:
112
- activation_quant_key = kFp8DynamicTensorSym
112
+ self.activation_quant_key = kFp8DynamicTensorSym
113
+
114
+ def create_weights(
115
+ self,
116
+ layer: torch.nn.Module,
117
+ input_size_per_partition: int,
118
+ output_partition_sizes: list[int],
119
+ input_size: int,
120
+ output_size: int,
121
+ params_dtype: torch.dtype,
122
+ **extra_weight_attrs,
123
+ ):
124
+ super().create_weights(
125
+ layer,
126
+ input_size_per_partition,
127
+ output_partition_sizes,
128
+ input_size,
129
+ output_size,
130
+ params_dtype,
131
+ **extra_weight_attrs,
132
+ )
133
134
self.fp8_linear = init_fp8_linear_kernel(
115
- activation_quant_key=activation_quant_key,
116
- weight_quant_key=kFp8StaticTensorSym,
117
- out_dtype=torch.get_default_dtype(),
135
+ activation_quant_key=self.activation_quant_key,
136
+ weight_quant_key=self.weight_quant_key,
137
+ weight_shape=layer.weight.shape,
138
+ input_dtype=self.input_dtype,
139
+ out_dtype=self.out_dtype,
140
module_name=self.__class__.__name__,
141
)
142
@@ -166,19 +188,14 @@ class Fp8PerBlockOnlineLinearMethod(_Fp8OnlineLinearBase):
188
Loads fp16/bf16 weights and quantizes them per-block during loading."""
189
190
def __init__(self):
169
- self.out_dtype = torch.get_default_dtype()
191
+ super().__init__()
192
self.weight_block_size = [128, 128]
171
-
172
- self.use_deep_gemm = is_deep_gemm_supported()
173
- self.use_aiter_and_is_supported = rocm_aiter_ops.is_linear_fp8_enabled()
174
- self.cutlass_block_fp8_supported = cutlass_block_fp8_supported()
175
-
176
- self.w8a8_block_fp8_linear = W8A8BlockFp8LinearOp(
177
- weight_group_shape=GroupShape(*self.weight_block_size),
178
- act_quant_group_shape=GroupShape(1, self.weight_block_size[0]),
179
- cutlass_block_fp8_supported=self.cutlass_block_fp8_supported,
180
- use_aiter_and_is_supported=self.use_aiter_and_is_supported,
181
- use_deep_gemm=self.use_deep_gemm,
193
+ self.activation_quant_key = create_fp8_quant_key(
194
+ static=False,
195
+ group_shape=GroupShape(1, self.weight_block_size[0]),
196
+ )
197
+ self.weight_quant_key = create_fp8_quant_key(
198
+ static=True, group_shape=GroupShape(*self.weight_block_size)
199
)
200
201
def create_weights(
@@ -202,6 +219,15 @@ class Fp8PerBlockOnlineLinearMethod(_Fp8OnlineLinearBase):
219
)
220
layer.weight_block_size = self.weight_block_size
221
222
+ self.fp8_linear = init_fp8_linear_kernel(
223
+ activation_quant_key=self.activation_quant_key,
224
+ weight_quant_key=self.weight_quant_key,
225
+ weight_shape=layer.weight.shape,
226
+ input_dtype=self.input_dtype,
227
+ out_dtype=self.out_dtype,
228
+ module_name=self.__class__.__name__,
229
+ )
230
+
231
def process_weights_after_loading(self, layer: Module) -> None:
232
if getattr(layer, "_already_called_process_weights_after_loading", False):
233
return
@@ -213,14 +239,10 @@ class Fp8PerBlockOnlineLinearMethod(_Fp8OnlineLinearBase):
239
layer.weight, block_size=block_size, use_ue8m0=False
240
)
241
216
- qweight, weight_scale_inv = process_fp8_weight_block_strategy(
217
- qweight, weight_scale_inv
218
- )
219
-
242
replace_parameter(layer, "weight", qweight.data)
243
replace_parameter(layer, "weight_scale_inv", weight_scale_inv.data)
244
223
- maybe_post_process_fp8_weight_block(layer)
245
+ self.fp8_linear.process_weights_after_loading(layer)
246
247
# Prevent duplicate processing (e.g., during weight reload)
248
layer._already_called_process_weights_after_loading = True
@@ -234,12 +256,10 @@ class Fp8PerBlockOnlineLinearMethod(_Fp8OnlineLinearBase):
256
assert self.weight_block_size is not None
257
258
# Note: batch invariance already handled in the function below
237
- return self.w8a8_block_fp8_linear.apply(
238
- input=x,
239
- weight=layer.weight,
240
- weight_scale=layer.weight_scale_inv,
241
- input_scale=layer.input_scale,
242
- bias=bias,
259
+ return self.fp8_linear.apply_weights(
260
+ layer,
261
+ x,
262
+ bias,
263
)
264
265
vllm/model_executor/layers/quantization/quark/schemes/quark_w8a8_fp8.py
+5
-1
@@ -7,6 +7,7 @@ from typing import Any, cast
7
import torch
8
from torch.nn import Parameter
9
10
+from vllm.config import get_current_vllm_config
11
from vllm.logger import init_logger
12
from vllm.model_executor.kernels.linear import (
13
init_fp8_linear_kernel,
@@ -57,6 +58,7 @@ class QuarkW8A8Fp8(QuarkScheme):
58
kFp8StaticTokenSym if per_token_weight else kFp8StaticTensorSym
59
)
60
self.out_dtype = torch.get_default_dtype()
61
+ self.input_dtype = get_current_vllm_config().model_config.dtype
62
63
@classmethod
64
def get_min_capability(cls) -> int:
@@ -175,7 +177,9 @@ class QuarkW8A8Fp8(QuarkScheme):
177
self.fp8_linear = init_fp8_linear_kernel(
178
activation_quant_key=self.activation_quant_key,
179
weight_quant_key=self.weight_quant_key,
178
- out_dtype=torch.get_default_dtype(),
180
+ weight_shape=layer.weight.shape,
181
+ input_dtype=self.input_dtype,
182
+ out_dtype=self.out_dtype,
183
module_name=self.__class__.__name__,
184
)
185
vllm/model_executor/layers/quantization/utils/fp8_utils.py
+1
-544
@@ -12,15 +12,11 @@ import torch
12
13
import vllm.envs as envs
14
from vllm import _custom_ops as ops
15
-from vllm._aiter_ops import rocm_aiter_ops
15
from vllm.logger import init_logger
17
-from vllm.model_executor.layers.quantization.input_quant_fp8 import QuantFP8
16
from vllm.model_executor.layers.quantization.utils.quant_utils import (
19
- GroupShape,
17
get_fp8_min_max,
18
)
19
from vllm.model_executor.layers.quantization.utils.w8a8_utils import (
23
- CUTLASS_BLOCK_FP8_SUPPORTED,
20
all_close_1d,
21
per_tensor_dequantize,
22
)
@@ -29,22 +25,14 @@ from vllm.model_executor.parameter import (
25
ChannelQuantScaleParameter,
26
PerTensorScaleParameter,
27
)
32
-from vllm.model_executor.utils import replace_parameter, set_weight_attrs
28
+from vllm.model_executor.utils import set_weight_attrs
29
from vllm.platforms import current_platform
30
from vllm.triton_utils import tl, triton
31
from vllm.utils.deep_gemm import (
36
- fp8_gemm_nt,
32
get_tma_aligned_size,
33
is_deep_gemm_e8m0_used,
39
- is_deep_gemm_supported,
40
- should_use_deepgemm_for_fp8_linear,
34
transform_sf_into_required_layout,
35
)
43
-from vllm.utils.flashinfer import (
44
- flashinfer_fp8_blockscale_gemm,
45
- is_flashinfer_fp8_blockscale_gemm_supported,
46
- should_use_flashinfer_for_blockscale_fp8_gemm,
47
-)
36
from vllm.utils.torch_utils import direct_register_custom_op
37
38
logger = init_logger(__name__)
@@ -56,153 +44,6 @@ def is_fp8(x: torch.dtype | torch.Tensor) -> bool:
44
return x == torch.float8_e4m3fn or x == torch.float8_e4m3fnuz
45
46
59
-# We need to pass in the is_hopper flag as argument because the function
60
-# current_platform.is_device_capability() is not supported by Torch compiler.
61
-def cutlass_scaled_mm(
62
- A: torch.Tensor,
63
- B: torch.Tensor,
64
- As: torch.Tensor,
65
- Bs: torch.Tensor,
66
- block_size: list[int],
67
- output_dtype: torch.dtype = torch.float16,
68
-) -> torch.Tensor:
69
- return ops.cutlass_scaled_mm(
70
- A,
71
- B.T,
72
- out_dtype=output_dtype,
73
- scale_a=As,
74
- scale_b=Bs.T,
75
- )
76
-
77
-
78
-# TODO we should be able to change the type of block_size to GroupShape
79
-# after we resolve GroupShape compilation issue
80
-# https://github.com/vllm-project/vllm/issues/25270
81
-def _w8a8_triton_block_scaled_mm_func(
82
- qx: torch.Tensor,
83
- weight: torch.Tensor,
84
- x_scale: torch.Tensor,
85
- weight_scale: torch.Tensor,
86
- block_size: list[int],
87
- output_dtype: torch.dtype,
88
-) -> torch.Tensor:
89
- return w8a8_triton_block_scaled_mm(
90
- qx, weight, x_scale, weight_scale, block_size, output_dtype
91
- )
92
-
93
-
94
-def _w8a8_triton_block_scaled_mm_fake(
95
- qx: torch.Tensor,
96
- weight: torch.Tensor,
97
- x_scale: torch.Tensor,
98
- weight_scale: torch.Tensor,
99
- block_size: list[int],
100
- output_dtype: torch.dtype,
101
-) -> torch.Tensor:
102
- return torch.empty(
103
- (qx.size(0), weight.size(0)), dtype=output_dtype, device=qx.device
104
- )
105
-
106
-
107
-direct_register_custom_op(
108
- "w8a8_triton_block_scaled_mm_func",
109
- _w8a8_triton_block_scaled_mm_func,
110
- fake_impl=_w8a8_triton_block_scaled_mm_fake,
111
-)
112
-
113
-
114
-def _padded_cutlass(
115
- qx: torch.Tensor,
116
- weight: torch.Tensor,
117
- x_scale: torch.Tensor,
118
- weight_scale: torch.Tensor,
119
- block_size: list[int],
120
- output_dtype: torch.dtype,
121
-) -> torch.Tensor:
122
- pad_multiple = 4
123
- dim = qx.shape[0]
124
- padded = (
125
- dim if dim % pad_multiple == 0 else dim + pad_multiple - (dim % pad_multiple)
126
- )
127
-
128
- has_pad = padded > dim
129
-
130
- if has_pad:
131
- padded_shape = [padded, *qx.shape[1:]]
132
- padded_qx = torch.zeros(padded_shape, device=qx.device, dtype=qx.dtype)
133
- padded_qx[0 : qx.shape[0], ...].copy_(qx)
134
-
135
- padded_x_scale_shape = [*x_scale.shape[1:], padded]
136
- padded_x_scale = torch.ones(
137
- padded_x_scale_shape, device=x_scale.device, dtype=x_scale.dtype
138
- ).permute(-1, -2)
139
- padded_x_scale[0 : x_scale.shape[0], ...].copy_(x_scale)
140
-
141
- output = cutlass_scaled_mm(
142
- padded_qx, weight, padded_x_scale, weight_scale, block_size, output_dtype
143
- )
144
- return output[0 : qx.shape[0], ...]
145
- else:
146
- return cutlass_scaled_mm(
147
- qx, weight, x_scale, weight_scale, block_size, output_dtype
148
- )
149
-
150
-
151
-def _padded_cutlass_fake(
152
- qx: torch.Tensor,
153
- weight: torch.Tensor,
154
- x_scale: torch.Tensor,
155
- weight_scale: torch.Tensor,
156
- block_size: list[int],
157
- output_dtype: torch.dtype,
158
-) -> torch.Tensor:
159
- return torch.empty(
160
- (qx.size(0), weight.size(0)), dtype=output_dtype, device=qx.device
161
- )
162
-
163
-
164
-direct_register_custom_op(
165
- "padded_cutlass",
166
- _padded_cutlass,
167
- fake_impl=_padded_cutlass_fake,
168
-)
169
-
170
-
171
-def _fp8_gemm_nt_op(
172
- q_input: torch.Tensor,
173
- input_scale: torch.Tensor,
174
- weight: torch.Tensor,
175
- weight_scale: torch.Tensor,
176
- output: torch.Tensor,
177
- use_deep_gemm_e8m0: bool,
178
-) -> None:
179
- fp8_gemm_nt(
180
- (q_input, input_scale),
181
- (weight, weight_scale),
182
- output,
183
- is_deep_gemm_e8m0_used=use_deep_gemm_e8m0,
184
- )
185
-
186
-
187
-def _fp8_gemm_nt_op_fake(
188
- q_input: torch.Tensor,
189
- input_scale: torch.Tensor,
190
- weight: torch.Tensor,
191
- weight_scale: torch.Tensor,
192
- output: torch.Tensor,
193
- use_deep_gemm_e8m0: bool,
194
-) -> None:
195
- return None
196
-
197
-
198
-direct_register_custom_op(
199
- "fp8_gemm_nt_op",
200
- _fp8_gemm_nt_op,
201
- mutates_args=["output"],
202
- fake_impl=_fp8_gemm_nt_op_fake,
203
-)
204
-
205
-
47
def _triton_per_token_group_quant_fp8_impl(
48
x: torch.Tensor,
49
group_size: int,
@@ -236,362 +77,6 @@ direct_register_custom_op(
77
)
78
79
239
-def _flashinfer_fp8_blockscale_gemm_impl(
240
- input: torch.Tensor,
241
- weight: torch.Tensor,
242
- weight_scale: torch.Tensor,
243
- group_size: int,
244
- use_deep_gemm_e8m0: bool,
245
-) -> torch.Tensor:
246
- """
247
- Conditional FlashInfer FP8 blockscale GEMM with batch-size-dependent selection.
248
-
249
- This function switches between two optimized kernels based on the input batch size:
250
- - For small batches (M < 32): Uses FlashInfer's DeepGEMM swapAB optimization.
251
- - For larger batches (M >= 32): Uses the official DeepGEMM kernel.
252
-
253
- The conditional logic must use torch.cond() instead of a simple if-else statement
254
- to maintain compatibility with torch.compile graph compilation.
255
-
256
- This batch-size-dependent selection is essential for maintaining model accuracy.
257
- Benchmarks on GSM8K show a significant accuracy gap (88% vs 95%) for DeepSeek-V3.1
258
- when using FlashInfer's DeepGEMM on M>=32. The M < 32 strategy fixes the accuracy
259
- drop.
260
-
261
- Args:
262
- input: Input tensor of shape (batch_size, input_dim) in FP8 format
263
- weight: Weight tensor of shape (output_dim, input_dim) in FP8 format
264
- weight_scale: Scale factors for weight quantization (per-group)
265
- group_size: Quantization group size for the weight tensor
266
- use_deep_gemm_e8m0: Whether to use the E8M0 format in DeepGEMM quantization
267
-
268
- Returns:
269
- Output tensor of shape (batch_size, output_dim) in bfloat16 format
270
- """
271
-
272
- def run_flashinfer_deepgemm_swapAB(
273
- input: torch.Tensor,
274
- weight: torch.Tensor,
275
- weight_scale: torch.Tensor,
276
- ) -> torch.Tensor:
277
- return flashinfer_fp8_blockscale_gemm(
278
- input=input,
279
- weight=weight,
280
- weight_scale=weight_scale,
281
- out_dtype=torch.bfloat16,
282
- )
283
-
284
- def run_deepgemm(
285
- input: torch.Tensor,
286
- weight: torch.Tensor,
287
- weight_scale: torch.Tensor,
288
- ) -> torch.Tensor:
289
- q_input, input_scale = per_token_group_quant_fp8(
290
- input,
291
- group_size=group_size,
292
- column_major_scales=True,
293
- use_ue8m0=use_deep_gemm_e8m0,
294
- )
295
- output = torch.empty(
296
- (q_input.shape[0], weight.shape[0]),
297
- dtype=torch.bfloat16,
298
- device=q_input.device,
299
- )
300
- fp8_gemm_nt(
301
- (q_input, input_scale),
302
- (weight, weight_scale),
303
- output,
304
- is_deep_gemm_e8m0_used=use_deep_gemm_e8m0,
305
- )
306
- return output
307
-
308
- if envs.VLLM_BATCH_INVARIANT:
309
- return run_deepgemm(input, weight, weight_scale)
310
-
311
- condition = input.shape[0] < 32
312
-
313
- # PyTorch's torch.compile cannot handle input-dependent control flow in standard
314
- # Python conditionals. torch.cond() explicitly registers both code paths in the
315
- # computation graph, allowing torch.compile to capture both branches.
316
- # without torch.cond, the M < 32 condition won't be able to be captured by torch
317
- # compile
318
- return torch.cond(
319
- condition,
320
- run_flashinfer_deepgemm_swapAB,
321
- run_deepgemm,
322
- (input, weight, weight_scale),
323
- )
324
-
325
-
326
-def _flashinfer_fp8_blockscale_gemm_fake(
327
- input: torch.Tensor,
328
- weight: torch.Tensor,
329
- weight_scale: torch.Tensor,
330
- group_size: int,
331
- use_deep_gemm_e8m0: bool,
332
-) -> torch.Tensor:
333
- """
334
- Required fake/meta implementation for torch.compile graph tracing.
335
- """
336
- return torch.empty(
337
- input.shape[0], weight.shape[0], dtype=torch.bfloat16, device=input.device
338
- )
339
-
340
-
341
-direct_register_custom_op(
342
- "flashinfer_fp8_blockscale_gemm",
343
- _flashinfer_fp8_blockscale_gemm_impl,
344
- fake_impl=_flashinfer_fp8_blockscale_gemm_fake,
345
-)
346
-
347
-
348
-# TODO fix ROCm->Triton custom path:
349
-# https://github.com/vllm-project/vllm/issues/14397
350
-class W8A8BlockFp8LinearOp:
351
- """
352
- This class executes a Blocked FP8 linear layer using cutlass if supported
353
- and torch.scaled_mm otherwise.
354
- """
355
-
356
- def __init__(
357
- self,
358
- weight_group_shape: GroupShape,
359
- act_quant_group_shape: GroupShape,
360
- cutlass_block_fp8_supported: bool = CUTLASS_BLOCK_FP8_SUPPORTED,
361
- use_aiter_and_is_supported: bool = False,
362
- use_deep_gemm: bool | None = None,
363
- ):
364
- self.weight_group_shape = weight_group_shape
365
- self.act_quant_group_shape = act_quant_group_shape
366
- if use_deep_gemm is not None:
367
- self.is_deep_gemm_supported = use_deep_gemm
368
- else:
369
- self.is_deep_gemm_supported = is_deep_gemm_supported()
370
- self.is_hopper = current_platform.is_device_capability(90)
371
- self.use_deep_gemm_e8m0 = is_deep_gemm_e8m0_used()
372
- self.is_flashinfer_supported = is_flashinfer_fp8_blockscale_gemm_supported()
373
-
374
- # Get the correct blockscale mul and input quant operations.
375
- # We can't use _dispatch_w8a8_blockscale_op to figure out if we want
376
- # to use deepgemm because we don't know the shape of weights (and
377
- # whether deepgemm supports it) at the init time.
378
- self.w8a8_blockscale_op, self.input_quant_op = (
379
- self._dispatch_w8a8_blockscale_op(
380
- cutlass_block_fp8_supported, use_aiter_and_is_supported
381
- )
382
- )
383
- self.deepgemm_input_quant_op = (
384
- QuantFP8(
385
- False,
386
- self.act_quant_group_shape,
387
- column_major_scales=True,
388
- tma_aligned_scales=envs.VLLM_USE_DEEP_GEMM_TMA_ALIGNED_SCALES,
389
- use_ue8m0=self.use_deep_gemm_e8m0,
390
- )
391
- if self.is_deep_gemm_supported
392
- else None
393
- )
394
-
395
- def apply(
396
- self,
397
- input: torch.Tensor,
398
- weight: torch.Tensor,
399
- weight_scale: torch.Tensor,
400
- input_scale: torch.Tensor | None = None,
401
- bias: torch.Tensor | None = None,
402
- ) -> torch.Tensor:
403
- assert input_scale is None
404
- # View input as 2D matrix for fp8 methods
405
- input_2d = input.view(-1, input.shape[-1])
406
- output_shape = [*input.shape[:-1], weight.shape[0]]
407
- output_dtype = input.dtype
408
-
409
- if should_use_flashinfer_for_blockscale_fp8_gemm(
410
- self.is_flashinfer_supported, output_dtype, input_2d, weight
411
- ) and should_use_deepgemm_for_fp8_linear(
412
- output_dtype, weight, self.is_deep_gemm_supported
413
- ):
414
- output = self._run_flashinfer(input_2d, weight, weight_scale)
415
-
416
- elif should_use_deepgemm_for_fp8_linear(
417
- output_dtype, weight, self.is_deep_gemm_supported
418
- ):
419
- output = self._run_deepgemm(input_2d, weight, weight_scale)
420
- else:
421
- output = self.w8a8_blockscale_op(
422
- input_2d, weight, weight_scale, input_scale
423
- )
424
-
425
- if bias is not None:
426
- output = output + bias
427
- return output.to(dtype=input.dtype).view(*output_shape)
428
-
429
- def _run_deepgemm(
430
- self,
431
- input_2d: torch.Tensor,
432
- weight: torch.Tensor,
433
- weight_scale: torch.Tensor,
434
- ) -> torch.Tensor:
435
- assert self.deepgemm_input_quant_op is not None
436
- q_input, input_scale = self.deepgemm_input_quant_op(input_2d)
437
- output = torch.empty(
438
- (q_input.shape[0], weight.shape[0]),
439
- dtype=torch.bfloat16,
440
- device=q_input.device,
441
- )
442
- torch.ops.vllm.fp8_gemm_nt_op(
443
- q_input, input_scale, weight, weight_scale, output, self.use_deep_gemm_e8m0
444
- )
445
- return output
446
-
447
- def _run_cutlass(
448
- self,
449
- input_2d: torch.Tensor,
450
- weight: torch.Tensor,
451
- weight_scale: torch.Tensor,
452
- input_scale: torch.Tensor | None = None,
453
- ) -> torch.Tensor:
454
- assert input_scale is None
455
- assert self.input_quant_op is not None
456
- q_input, input_scale = self.input_quant_op(input_2d)
457
- if self.is_hopper:
458
- return torch.ops.vllm.padded_cutlass(
459
- q_input,
460
- weight,
461
- input_scale,
462
- weight_scale,
463
- list(self.weight_group_shape),
464
- input_2d.dtype,
465
- )
466
- else:
467
- return cutlass_scaled_mm(
468
- q_input,
469
- weight,
470
- input_scale,
471
- weight_scale,
472
- list(self.weight_group_shape),
473
- input_2d.dtype,
474
- )
475
-
476
- def _run_aiter(
477
- self,
478
- input_2d: torch.Tensor,
479
- weight: torch.Tensor,
480
- weight_scale: torch.Tensor,
481
- input_scale: torch.Tensor | None = None,
482
- ) -> torch.Tensor:
483
- assert self.act_quant_group_shape == GroupShape(1, 128)
484
-
485
- n, k = weight.shape
486
-
487
- use_triton = (
488
- not current_platform.is_fp8_fnuz()
489
- and rocm_aiter_ops.is_triton_gemm_w8a8_tuned(n, k)
490
- )
491
-
492
- if use_triton:
493
- gemm_a8w8_blockscale_op = rocm_aiter_ops.triton_gemm_a8w8_blockscale
494
- else:
495
- gemm_a8w8_blockscale_op = rocm_aiter_ops.gemm_a8w8_blockscale
496
-
497
- if input_scale is not None:
498
- q_input = input_2d
499
- else:
500
- q_input, input_scale = self.input_quant_op(input_2d, use_triton=use_triton)
501
-
502
- return gemm_a8w8_blockscale_op(
503
- q_input,
504
- weight,
505
- input_scale,
506
- weight_scale,
507
- list(self.weight_group_shape),
508
- output_dtype=input_2d.dtype,
509
- )
510
-
511
- def _run_triton(
512
- self,
513
- input_2d: torch.Tensor,
514
- weight: torch.Tensor,
515
- weight_scale: torch.Tensor,
516
- input_scale: torch.Tensor | None = None,
517
- ) -> torch.Tensor:
518
- assert input_scale is None
519
- assert self.input_quant_op is not None
520
- q_input, input_scale = self.input_quant_op(input_2d)
521
- return torch.ops.vllm.w8a8_triton_block_scaled_mm_func(
522
- q_input,
523
- weight,
524
- input_scale,
525
- weight_scale,
526
- list(self.weight_group_shape),
527
- input_2d.dtype,
528
- )
529
-
530
- def _run_flashinfer(
531
- self,
532
- input_2d: torch.Tensor,
533
- weight: torch.Tensor,
534
- weight_scale: torch.Tensor,
535
- ) -> torch.Tensor:
536
- """
537
- Run FlashInfer FP8 block-scale GEMM.
538
-
539
- This backend uses TensorRT-LLM's FP8 block-scale GEMM kernels
540
- and supports FP8+FP8 (W8A8 full quantization) on SM90+ (Hopper).
541
- """
542
- # Now call FlashInfer with BF16 input + FP8 weight, input will be
543
- # quantized with FlashInfer kernel (W8A8)
544
- output = torch.ops.vllm.flashinfer_fp8_blockscale_gemm(
545
- input=input_2d, # BF16 input
546
- weight=weight, # FP8 weight
547
- weight_scale=weight_scale, # Weight scales
548
- group_size=self.act_quant_group_shape.col,
549
- use_deep_gemm_e8m0=self.use_deep_gemm_e8m0,
550
- )
551
- return output
552
-
553
- def _dispatch_w8a8_blockscale_op(
554
- self,
555
- use_cutlass: bool,
556
- use_aiter_and_is_supported: bool,
557
- ) -> tuple[
558
- Callable[
559
- [
560
- torch.Tensor,
561
- torch.Tensor,
562
- torch.Tensor,
563
- torch.Tensor | None,
564
- ],
565
- torch.Tensor,
566
- ],
567
- QuantFP8,
568
- ]:
569
- if use_cutlass:
570
- return self._run_cutlass, (
571
- QuantFP8(
572
- False,
573
- self.act_quant_group_shape,
574
- column_major_scales=True,
575
- use_ue8m0=False,
576
- )
577
- )
578
- if use_aiter_and_is_supported:
579
- return self._run_aiter, QuantFP8(
580
- False,
581
- self.act_quant_group_shape,
582
- column_major_scales=False,
583
- use_ue8m0=False,
584
- )
585
- return self._run_triton, (
586
- QuantFP8(
587
- False,
588
- self.act_quant_group_shape,
589
- column_major_scales=False,
590
- use_ue8m0=False,
591
- )
592
- )
593
-
594
-
80
def input_to_float8(
81
x: torch.Tensor, dtype: torch.dtype | None = None
82
) -> tuple[torch.Tensor, torch.Tensor]:
@@ -1612,34 +1097,6 @@ def process_fp8_weight_block_strategy(
1097
return weight, weight_scale
1098
1099
1615
-def maybe_post_process_fp8_weight_block(layer: torch.nn.Module):
1616
- assert layer.weight_block_size is not None
1617
-
1618
- from vllm.utils.deep_gemm import (
1619
- is_deep_gemm_e8m0_used,
1620
- should_use_deepgemm_for_fp8_linear,
1621
- )
1622
-
1623
- # On Blackwell or Hopper, if E8M0 for DeepGemm is used, we need to
1624
- # requantize the weight and input to the specific scale
1625
- # at the same time.
1626
- should_use_deepgemm = should_use_deepgemm_for_fp8_linear(
1627
- layer.orig_dtype, layer.weight
1628
- )
1629
- if should_use_deepgemm:
1630
- scale_attr = (
1631
- "weight_scale_inv" if hasattr(layer, "weight_scale_inv") else "weight_scale"
1632
- )
1633
- dg_weight, dg_weight_scale = deepgemm_post_process_fp8_weight_block(
1634
- wq=layer.weight.data,
1635
- ws=getattr(layer, scale_attr).data,
1636
- quant_block_shape=tuple(layer.weight_block_size),
1637
- use_e8m0=is_deep_gemm_e8m0_used(),
1638
- )
1639
- replace_parameter(layer, "weight", dg_weight)
1640
- replace_parameter(layer, scale_attr, dg_weight_scale)
1641
-
1642
-
1100
def process_fp8_weight_tensor_strategy_moe(
1101
weight: torch.Tensor,
1102
weight_scales: torch.Tensor,
vllm/model_executor/layers/quantization/utils/quant_utils.py
+10
@@ -171,6 +171,16 @@ kMxfp4StaticGroupScale = ScaleDesc(MXFP_SCALE_DTYPE, True, GroupShape(1, 32))
171
kMxfp4Static = QuantKey(FP4_DTYPE, scale=kMxfp4StaticGroupScale, symmetric=True)
172
173
174
+def create_fp8_quant_key(
175
+ static: bool,
176
+ group_shape: GroupShape,
177
+ symmetric: bool = True,
178
+ scale_dtype: torch.dtype = torch.float32,
179
+) -> QuantKey:
180
+ scale_desc = ScaleDesc(scale_dtype, static, group_shape)
181
+ return QuantKey(FP8_DTYPE, scale_desc, symmetric=symmetric)
182
+
183
+
184
# Normalize the group_shape to the full extent for any dims that are -1
185
def _normalize_quant_group_shape(x: torch.Tensor, group_shape: GroupShape):
186
# -1 means full extent
vllm/utils/deep_gemm.py
+3
-3
@@ -413,7 +413,7 @@ def calc_diff(x: torch.Tensor, y: torch.Tensor):
413
414
def should_use_deepgemm_for_fp8_linear(
415
output_dtype: torch.dtype,
416
- weight: torch.Tensor,
416
+ weight_shape: tuple[int, int],
417
supports_deep_gemm: bool | None = None,
418
):
419
if supports_deep_gemm is None:
@@ -428,8 +428,8 @@ def should_use_deepgemm_for_fp8_linear(
428
return (
429
supports_deep_gemm
430
and output_dtype == torch.bfloat16
431
- and weight.shape[0] % N_MULTIPLE == 0
432
- and weight.shape[1] % K_MULTIPLE == 0
431
+ and weight_shape[0] % N_MULTIPLE == 0
432
+ and weight_shape[1] % K_MULTIPLE == 0
433
)
434
435
vllm/utils/flashinfer.py
+5
-7
@@ -748,8 +748,9 @@ def is_flashinfer_fp8_blockscale_gemm_supported() -> bool:
748
def should_use_flashinfer_for_blockscale_fp8_gemm(
749
is_flashinfer_supported: bool,
750
output_dtype: torch.dtype,
751
- input: torch.Tensor,
752
- weight: torch.Tensor,
751
+ input_dtype: torch.dtype,
752
+ weight_dtype: torch.dtype,
753
+ weight_shape: tuple[int, int],
754
):
755
if not is_flashinfer_supported:
756
return False
@@ -760,15 +761,12 @@ def should_use_flashinfer_for_blockscale_fp8_gemm(
761
N_MULTIPLE = 64
762
K_MULTIPLE = 128
763
763
- weight_dtype = weight.dtype
764
- input_dtype = input.dtype
765
-
764
should_use_flashinfer = (
765
output_dtype == torch.bfloat16
766
and input_dtype == torch.bfloat16
767
and weight_dtype == torch.float8_e4m3fn
770
- and weight.shape[0] % N_MULTIPLE == 0
771
- and weight.shape[1] % K_MULTIPLE == 0
768
+ and weight_shape[0] % N_MULTIPLE == 0
769
+ and weight_shape[1] % K_MULTIPLE == 0
770
)
771
772
return should_use_flashinfer