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 9, 23:00 UTC
[Refactor] Make FP8 Linear Ops use kernel abstraction (#27814)
Signed-off-by: vllmellm <vllm.ellm@embeddedllm.com>
vllmellm committed
Jan 20, 2026 at 14:48 UTC
148117ea2e689cd43df4be6892671a17cdae5833
30 files changed
+1458
-1029
.buildkite/lm-eval-harness/configs/models-small-rocm.txt
new
+5
@@ -0,0 +1,5 @@
1
+Qwen2.5-1.5B-Instruct.yaml
2
+Meta-Llama-3.2-1B-Instruct-INT8-compressed-tensors.yaml
3
+Meta-Llama-3-8B-Instruct-nonuniform-compressed-tensors.yaml
4
+Qwen2.5-VL-3B-Instruct-FP8-dynamic.yaml
5
+Qwen1.5-MoE-W4A16-compressed-tensors.yaml
tests/compile/distributed/test_fusion_all_reduce.py
+17
-27
@@ -26,15 +26,14 @@ from vllm.distributed.parallel_state import (
26
initialize_model_parallel,
27
)
28
from vllm.model_executor.layers.layernorm import RMSNorm
29
-from vllm.model_executor.layers.quantization.utils.w8a8_utils import (
30
- Fp8LinearOp,
31
- GroupShape,
29
+from vllm.model_executor.layers.quantization.utils.quant_utils import (
30
+ kFp8StaticTensorSym,
31
)
32
from vllm.platforms import current_platform
33
from vllm.utils.system_utils import update_environment_variables
34
from vllm.utils.torch_utils import set_random_seed
35
37
-from ...utils import has_module_attribute, multi_gpu_test
36
+from ...utils import TestFP8Layer, has_module_attribute, multi_gpu_test
37
from ..backend import TestBackend
38
39
@@ -76,49 +75,40 @@ class TestAllReduceRMSNormModel(torch.nn.Module):
75
76
77
class TestAllReduceRMSNormStaticQuantFP8Model(torch.nn.Module):
78
+ quant_key = kFp8StaticTensorSym
79
+
80
def __init__(self, hidden_size=16, token_num=16, eps=1e-6):
81
super().__init__()
82
self.hidden_size = hidden_size
83
self.eps = eps
84
self.norm = [RMSNorm(hidden_size, eps) for i in range(4)]
84
- self.wscale = [torch.rand(1, dtype=torch.float32) for _ in range(3)]
85
- self.w = [
86
- torch.rand(hidden_size, hidden_size)
87
- .to(dtype=current_platform.fp8_dtype())
88
- .t()
89
- for _ in range(3)
85
+ self.fp8_linear_layers = [
86
+ TestFP8Layer(
87
+ weight_shape=(hidden_size, hidden_size),
88
+ activation_quant_key=self.quant_key,
89
+ weight_quant_key=self.quant_key,
90
+ )
91
+ for i in range(3)
92
]
93
92
- self.fp8_linear = Fp8LinearOp(
93
- act_quant_static=True,
94
- act_quant_group_shape=GroupShape.PER_TENSOR,
95
- )
96
-
97
- self.scale = [torch.rand(1, dtype=torch.float32) for _ in range(3)]
98
-
94
def forward(self, hidden_states):
95
# avoid having graph input be an arg to a pattern directly
96
z = torch.relu(hidden_states)
97
x = resid = tensor_model_parallel_all_reduce(z)
98
y = self.norm[0](x)
99
105
- z2 = self.fp8_linear.apply(
106
- y, self.w[0], self.wscale[0], input_scale=self.scale[0]
107
- )
100
+ z2 = self.fp8_linear_layers[0](y)
101
102
x2 = tensor_model_parallel_all_reduce(z2)
103
y2, resid = self.norm[1](x2, resid)
104
112
- z3 = self.fp8_linear.apply(
113
- y2, self.w[1], self.wscale[1], input_scale=self.scale[1]
114
- )
105
+ z3 = self.fp8_linear_layers[1](y2)
106
107
x3 = tensor_model_parallel_all_reduce(z3)
108
y3, resid = self.norm[2](x3, resid) # use resid here
109
119
- z4 = self.fp8_linear.apply(
120
- y3, self.w[2], self.wscale[2], input_scale=self.scale[2]
121
- )
110
+ z4 = self.fp8_linear_layers[2](y3)
111
+
112
x4 = tensor_model_parallel_all_reduce(z4)
113
y4, resid = self.norm[3](x4, resid) # use resid here
114
return y4
@@ -130,7 +120,7 @@ class TestAllReduceRMSNormStaticQuantFP8Model(torch.nn.Module):
120
return [
121
torch.ops.vllm.all_reduce.default,
122
torch.ops._C.static_scaled_fp8_quant.default
133
- if self.fp8_linear.quant_fp8.enabled()
123
+ if self.fp8_linear_layers[0].is_quant_fp8_enabled()
124
else torch.ops.aten.reciprocal.default,
125
]
126
tests/compile/distributed/test_sequence_parallelism.py
+17
-26
@@ -27,13 +27,14 @@ from vllm.distributed.parallel_state import (
27
initialize_model_parallel,
28
)
29
from vllm.model_executor.layers.layernorm import RMSNorm
30
-from vllm.model_executor.layers.quantization.utils.quant_utils import GroupShape
31
-from vllm.model_executor.layers.quantization.utils.w8a8_utils import Fp8LinearOp
30
+from vllm.model_executor.layers.quantization.utils.quant_utils import (
31
+ kFp8StaticTensorSym,
32
+)
33
from vllm.platforms import current_platform
34
from vllm.utils.system_utils import update_environment_variables
35
from vllm.utils.torch_utils import set_random_seed
36
36
-from ...utils import multi_gpu_test
37
+from ...utils import TestFP8Layer, multi_gpu_test
38
from ..backend import TestBackend
39
40
FP8_DTYPE = current_platform.fp8_dtype()
@@ -94,50 +95,40 @@ class TestAllReduceRMSNormModel(torch.nn.Module):
95
96
97
class TestAllReduceRMSNormStaticQuantFP8Model(torch.nn.Module):
98
+ quant_key = kFp8StaticTensorSym
99
+
100
def __init__(self, hidden_size=16, eps=1e-6):
101
super().__init__()
102
self.vllm_config = get_current_vllm_config()
103
self.hidden_size = hidden_size
104
self.eps = eps
105
self.norm = [RMSNorm(hidden_size, eps) for i in range(4)]
103
- self.wscale = [torch.rand(1, dtype=torch.float32) for _ in range(3)]
104
- self.w = [
105
- torch.rand(hidden_size, hidden_size)
106
- .to(dtype=current_platform.fp8_dtype())
107
- .t()
108
- for _ in range(3)
106
+ self.fp8_linear_layers = [
107
+ TestFP8Layer(
108
+ weight_shape=(hidden_size, hidden_size),
109
+ activation_quant_key=self.quant_key,
110
+ weight_quant_key=self.quant_key,
111
+ )
112
+ for i in range(3)
113
]
114
111
- self.fp8_linear = Fp8LinearOp(
112
- act_quant_static=True,
113
- act_quant_group_shape=GroupShape.PER_TENSOR,
114
- )
115
-
116
- self.scale = [torch.rand(1, dtype=torch.float32) for _ in range(3)]
117
-
115
def forward(self, hidden_states):
116
# avoid having graph input be an arg to a pattern directly
117
z = torch.relu(hidden_states)
118
x = resid = tensor_model_parallel_all_reduce(z)
119
y = self.norm[0](x)
120
124
- z2 = self.fp8_linear.apply(
125
- y, self.w[0], self.wscale[0], input_scale=self.scale[0]
126
- )
121
+ z2 = self.fp8_linear_layers[0](y)
122
123
x2 = tensor_model_parallel_all_reduce(z2)
124
y2, resid = self.norm[1](x2, resid)
125
131
- z3 = self.fp8_linear.apply(
132
- y2, self.w[1], self.wscale[1], input_scale=self.scale[1]
133
- )
126
+ z3 = self.fp8_linear_layers[1](y2)
127
128
x3 = tensor_model_parallel_all_reduce(z3)
129
y3, resid = self.norm[2](x3, resid) # use resid here
130
138
- z4 = self.fp8_linear.apply(
139
- y3, self.w[2], self.wscale[2], input_scale=self.scale[2]
140
- )
131
+ z4 = self.fp8_linear_layers[2](y3)
132
x4 = tensor_model_parallel_all_reduce(z4)
133
y4, resid = self.norm[3](x4, resid) # use resid here
134
return y4
@@ -160,7 +151,7 @@ class TestAllReduceRMSNormStaticQuantFP8Model(torch.nn.Module):
151
return [
152
torch.ops._C.fused_add_rms_norm.default,
153
]
163
- elif self.fp8_linear.quant_fp8.enabled():
154
+ elif any(layer.is_quant_fp8_enabled() for layer in self.fp8_linear_layers):
155
return [
156
torch.ops._C.static_scaled_fp8_quant.default,
157
]
tests/compile/test_functionalization.py
+19
-22
@@ -20,11 +20,13 @@ from vllm.config import (
20
)
21
from vllm.model_executor.layers.activation import SiluAndMul
22
from vllm.model_executor.layers.layernorm import RMSNorm
23
-from vllm.model_executor.layers.quantization.utils.quant_utils import GroupShape
24
-from vllm.model_executor.layers.quantization.utils.w8a8_utils import Fp8LinearOp
23
+from vllm.model_executor.layers.quantization.utils.quant_utils import (
24
+ kFp8StaticTensorSym,
25
+)
26
from vllm.model_executor.layers.rotary_embedding import get_rope
27
from vllm.platforms import current_platform
28
29
+from ..utils import TestFP8Layer
30
from .backend import TestBackend
31
32
TEST_FP8 = current_platform.supports_fp8()
@@ -32,24 +34,22 @@ FP8_DTYPE = current_platform.fp8_dtype()
34
35
36
class TestSiluMul(torch.nn.Module):
37
+ quant_key = kFp8StaticTensorSym
38
+
39
def __init__(self, hidden_size: int = 128):
40
super().__init__()
41
self.silu_and_mul = SiluAndMul()
38
- self.wscale = torch.rand(1, dtype=torch.float32)
39
- self.scale = torch.rand(1, dtype=torch.float32)
40
-
42
if TEST_FP8:
42
- self.w = torch.rand(hidden_size, hidden_size).to(dtype=FP8_DTYPE).t()
43
- self.fp8_linear = Fp8LinearOp(
44
- act_quant_static=True,
45
- act_quant_group_shape=GroupShape.PER_TENSOR,
43
+ self.fp8_linear = TestFP8Layer(
44
+ weight_shape=(hidden_size, hidden_size),
45
+ activation_quant_key=self.quant_key,
46
+ weight_quant_key=self.quant_key,
47
)
48
49
def forward(self, x):
50
y = self.silu_and_mul(x)
51
if TEST_FP8:
51
- x2 = self.fp8_linear.apply(y, self.w, self.wscale, input_scale=self.wscale)
52
- return x2
52
+ return self.fp8_linear(y)
53
else:
54
return y
55
@@ -67,6 +67,8 @@ class TestSiluMul(torch.nn.Module):
67
68
69
class TestFusedAddRMSNorm(torch.nn.Module):
70
+ quant_key = kFp8StaticTensorSym
71
+
72
def __init__(self, hidden_size=16, intermediate_size=32):
73
super().__init__()
74
self.hidden_size = hidden_size
@@ -81,11 +83,11 @@ class TestFusedAddRMSNorm(torch.nn.Module):
83
torch.nn.init.normal_(self.gate_proj, std=0.02)
84
85
if TEST_FP8:
84
- self.fp8_linear = Fp8LinearOp(act_quant_static=True)
85
-
86
- self.scale = torch.rand(1, dtype=torch.float32)
87
- self.w = torch.rand(hidden_size, intermediate_size).to(dtype=FP8_DTYPE).t()
88
- self.wscale = torch.rand(1, dtype=torch.float32)
86
+ self.fp8_linear = TestFP8Layer(
87
+ weight_shape=(hidden_size, intermediate_size),
88
+ activation_quant_key=self.quant_key,
89
+ weight_quant_key=self.quant_key,
90
+ )
91
92
def forward(self, hidden_states, residual):
93
# Reshape input
@@ -100,12 +102,7 @@ class TestFusedAddRMSNorm(torch.nn.Module):
102
103
if TEST_FP8:
104
# scaled_mm with static input quantization
103
- fp8_linear_result = self.fp8_linear.apply(
104
- norm_output,
105
- self.w,
106
- self.wscale,
107
- input_scale=self.scale.to(norm_output.device),
108
- )
105
+ fp8_linear_result = self.fp8_linear(norm_output)
106
107
return fp8_linear_result, residual_output
108
tests/compile/test_fusion.py
+182
-146
@@ -5,6 +5,7 @@
5
import pytest
6
import torch
7
8
+import vllm.config
9
import vllm.plugins
10
from vllm._aiter_ops import IS_AITER_FOUND, rocm_aiter_ops
11
from vllm.compilation.fusion import FUSED_OPS, FusedRMSQuantKey, RMSNormQuantFusionPass
@@ -20,8 +21,22 @@ from vllm.config import (
21
VllmConfig,
22
)
23
from vllm.model_executor.layers.layernorm import RMSNorm
23
-from vllm.model_executor.layers.quantization.utils.fp8_utils import (
24
- W8A8BlockFp8LinearOp,
24
+from vllm.model_executor.layers.quantization.kernels.scaled_mm.cutlass import (
25
+ CutlassFP8ScaledMMLinearKernel,
26
+)
27
+from vllm.model_executor.layers.quantization.kernels.scaled_mm.flashinfer import (
28
+ FlashInferFP8ScaledMMLinearKernel,
29
+)
30
+from vllm.model_executor.layers.quantization.kernels.scaled_mm.pytorch import (
31
+ ChannelWiseTorchFP8ScaledMMLinearKernel,
32
+ PerTensorTorchFP8ScaledMMLinearKernel,
33
+ RowWiseTorchFP8ScaledMMLinearKernel,
34
+)
35
+from vllm.model_executor.layers.quantization.kernels.scaled_mm.rocm import (
36
+ ROCmFP8ScaledMMLinearKernel,
37
+)
38
+from vllm.model_executor.layers.quantization.kernels.scaled_mm.ScaledMMLinearKernel import ( # noqa: E501
39
+ FP8ScaledMMLinearKernel,
40
)
41
from vllm.model_executor.layers.quantization.utils.quant_utils import (
42
GroupShape,
@@ -29,15 +44,14 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import (
44
ScaleDesc,
45
)
46
from vllm.model_executor.layers.quantization.utils.w8a8_utils import (
32
- Fp8LinearOp,
47
cutlass_block_fp8_supported,
34
- cutlass_fp8_supported,
35
- maybe_create_device_identity,
48
)
49
from vllm.platforms import current_platform
38
-from vllm.utils.deep_gemm import is_deep_gemm_supported
50
+from vllm.utils.deep_gemm import (
51
+ is_deep_gemm_supported,
52
+)
53
40
-from ..utils import override_cutlass_fp8_supported
54
+from ..utils import TestBlockFP8Layer, TestFP8Layer
55
from .backend import TestBackend
56
57
FP8_DTYPE = current_platform.fp8_dtype()
@@ -45,157 +59,195 @@ FP8_DTYPE = current_platform.fp8_dtype()
59
RMS_OP = torch.ops._C.rms_norm.default
60
RMS_ADD_OP = torch.ops._C.fused_add_rms_norm.default
61
62
+# Kernel and group_shape combinations: (kernel, group_shape)
63
+# CUDA kernels
64
+CUDA_KERNEL_GROUPSHAPE_COMBINATIONS = [
65
+ # FlashInferFP8ScaledMMLinearKernel supports both per-tensor only
66
+ (FlashInferFP8ScaledMMLinearKernel, GroupShape.PER_TENSOR),
67
+ # CutlassFP8ScaledMMLinearKernel supports both per-tensor and per-token
68
+ (CutlassFP8ScaledMMLinearKernel, GroupShape.PER_TOKEN),
69
+ (CutlassFP8ScaledMMLinearKernel, GroupShape.PER_TENSOR),
70
+ # PerTensorTorchFP8ScaledMMLinearKernel only supports per-tensor
71
+ (PerTensorTorchFP8ScaledMMLinearKernel, GroupShape.PER_TENSOR),
72
+ # ChannelWiseTorchFP8ScaledMMLinearKernel only supports per-token
73
+ (ChannelWiseTorchFP8ScaledMMLinearKernel, GroupShape.PER_TOKEN),
74
+ # Blockwise group shapes (no kernel abstraction)
75
+ (None, GroupShape(1, 128)),
76
+ (None, GroupShape(1, 64)),
77
+]
78
+
79
+# ROCm kernels
80
+ROCM_KERNEL_GROUPSHAPE_COMBINATIONS = [
81
+ # ROCmFP8ScaledMMLinearKernel supports per-tensor only
82
+ (ROCmFP8ScaledMMLinearKernel, GroupShape.PER_TENSOR),
83
+ # RowWiseTorchFP8ScaledMMLinearKernel only supports per-token
84
+ (RowWiseTorchFP8ScaledMMLinearKernel, GroupShape.PER_TOKEN),
85
+ # ChannelWiseTorchFP8ScaledMMLinearKernel only supports per-token
86
+ (ChannelWiseTorchFP8ScaledMMLinearKernel, GroupShape.PER_TOKEN),
87
+ # Blockwise group shapes (no kernel abstraction)
88
+ (None, GroupShape(1, 128)),
89
+ (None, GroupShape(1, 64)),
90
+]
91
+
92
+KERNEL_GROUPSHAPE_COMBINATIONS = (
93
+ CUDA_KERNEL_GROUPSHAPE_COMBINATIONS
94
+ if current_platform.is_cuda()
95
+ else ROCM_KERNEL_GROUPSHAPE_COMBINATIONS
96
+)
97
+
98
+# For Aiter tests we toggle use_aiter_quant_op
99
+AITER_KERNEL_GROUPSHAPE_COMBINATIONS = [
100
+ # Per-token with ROCmFP8ScaledMMLinearKernel
101
+ (ROCmFP8ScaledMMLinearKernel, GroupShape.PER_TENSOR, False),
102
+ # Per-token with RowWiseTorchFP8ScaledMMLinearKernel
103
+ (RowWiseTorchFP8ScaledMMLinearKernel, GroupShape.PER_TOKEN, True),
104
+ (RowWiseTorchFP8ScaledMMLinearKernel, GroupShape.PER_TOKEN, False),
105
+ # Per-token with ChannelWiseTorchFP8ScaledMMLinearKernel
106
+ (ChannelWiseTorchFP8ScaledMMLinearKernel, GroupShape.PER_TOKEN, True),
107
+ (ChannelWiseTorchFP8ScaledMMLinearKernel, GroupShape.PER_TOKEN, False),
108
+ # Blockwise (no kernel abstraction)
109
+ (None, GroupShape(1, 128), True),
110
+]
111
+
112
113
class TestModel(torch.nn.Module):
114
def __init__(
115
self,
116
hidden_size: int,
117
eps: float,
118
+ force_kernel: FP8ScaledMMLinearKernel | None,
119
group_shape: GroupShape,
55
- use_aiter: bool = False,
56
- cuda_force_torch: bool = False,
57
- use_aiter_quant_op: bool = True,
120
+ use_aiter_fusion: bool = False,
121
+ use_aiter_quant: bool = False,
122
*args,
123
**kwargs,
124
):
125
super().__init__(*args, **kwargs)
62
- self.use_aiter = use_aiter
63
- self.use_aiter_quant_op = use_aiter_quant_op
64
- self.cuda_force_torch = cuda_force_torch
126
+ self.fp8_linear_layers: list[torch.nn.Module]
127
self.group_shape = group_shape
66
- self.enable_quant_fp8_custom_op = None # Will be set later if applicable
67
-
128
+ self.use_aiter_quant_op = use_aiter_quant
129
+ self.use_aiter_fusion = use_aiter_fusion
130
self.norm = [RMSNorm(hidden_size, eps) for _ in range(4)]
131
+ self.enable_rms_norm_custom_op = self.norm[0].enabled()
132
70
- # Setup quantization scale descriptor
71
- static = group_shape == GroupShape.PER_TENSOR and not use_aiter
72
- quant_scale = ScaleDesc(torch.float32, static, group_shape)
73
- self.quant_key = QuantKey(dtype=FP8_DTYPE, scale=quant_scale, symmetric=True)
133
+ # Determine if blockwise based on group_shape
134
+ is_blockwise = group_shape.is_per_group()
135
75
- # Setup scales
76
- if static:
77
- self.scale = [torch.rand(1, dtype=torch.float32) for _ in range(3)]
78
- else:
79
- self.scale = [None for _ in range(3)]
136
+ if is_blockwise:
137
+ act_quant_scale_desc = ScaleDesc(torch.float32, False, group_shape)
138
+ self.activation_quant_key = QuantKey(
139
+ dtype=FP8_DTYPE, scale=act_quant_scale_desc, symmetric=True
140
+ )
141
+ self.fp8_linear_layers = [
142
+ TestBlockFP8Layer(
143
+ weight_shape=(hidden_size, hidden_size),
144
+ group_shape=group_shape,
145
+ cutlass_block_fp8_supported=cutlass_block_fp8_supported(),
146
+ use_aiter_and_is_supported=use_aiter_quant,
147
+ transpose_weights=use_aiter_fusion,
148
+ )
149
+ for _ in range(3)
150
+ ]
151
81
- # Setup weights
82
- self.w = [
83
- torch.rand(hidden_size, hidden_size).to(dtype=FP8_DTYPE) for _ in range(3)
84
- ]
85
- if not group_shape.is_per_group() or use_aiter:
86
- self.w = [self.w[0].t() for _ in range(3)]
87
-
88
- # Setup weight scales
89
- if group_shape.is_per_group():
90
- scale_size = (
91
- (hidden_size + 128 - 1) // 128
92
- if use_aiter
93
- else hidden_size // group_shape[1]
152
+ self.enable_quant_fp8_custom_op = (
153
+ False
154
+ if use_aiter_quant
155
+ else self.fp8_linear_layers[0].linear_op.input_quant_op.enabled()
156
)
95
- wscale_shape: tuple[int, ...] = (scale_size, scale_size)
157
+
158
else:
97
- wscale_shape = (1,)
98
- self.wscale = [torch.rand(wscale_shape, dtype=torch.float32) for _ in range(3)]
99
-
100
- # Setup FP8 linear operation
101
- is_per_group = group_shape.is_per_group()
102
- if is_per_group and use_aiter:
103
- self.fp8_linear = W8A8BlockFp8LinearOp(
104
- weight_group_shape=GroupShape(128, 128),
105
- act_quant_group_shape=group_shape,
106
- use_aiter_and_is_supported=use_aiter_quant_op,
107
- )
108
- # AITER blockwise doesn't use enable_quant_fp8_custom_op
109
- elif is_per_group:
110
- self.fp8_linear = W8A8BlockFp8LinearOp(
111
- weight_group_shape=GroupShape(group_shape[1], group_shape[1]),
112
- act_quant_group_shape=group_shape,
113
- cutlass_block_fp8_supported=cutlass_block_fp8_supported(),
114
- use_aiter_and_is_supported=False,
159
+ is_static = group_shape == GroupShape.PER_TENSOR
160
+ act_quant_scale_desc = ScaleDesc(torch.float32, is_static, group_shape)
161
+ w_quant_scale_desc = ScaleDesc(torch.float32, True, group_shape)
162
+ self.activation_quant_key = QuantKey(
163
+ dtype=FP8_DTYPE, scale=act_quant_scale_desc, symmetric=True
164
)
116
- self.enable_quant_fp8_custom_op = self.fp8_linear.input_quant_op.enabled()
117
- elif use_aiter:
118
- self.fp8_linear = Fp8LinearOp(
119
- act_quant_static=False,
120
- act_quant_group_shape=group_shape,
165
+ self.weight_quant_key = QuantKey(
166
+ dtype=FP8_DTYPE, scale=w_quant_scale_desc, symmetric=True
167
)
122
- self.fp8_linear.quant_fp8.use_aiter = use_aiter_quant_op
123
- self.enable_quant_fp8_custom_op = self.fp8_linear.quant_fp8.enabled()
124
- else:
125
- with override_cutlass_fp8_supported(not cuda_force_torch):
126
- self.fp8_linear = Fp8LinearOp(
127
- act_quant_static=static,
128
- act_quant_group_shape=group_shape,
168
+ self.fp8_linear_layers = [
169
+ TestFP8Layer(
170
+ weight_shape=(hidden_size, hidden_size),
171
+ activation_quant_key=self.activation_quant_key,
172
+ weight_quant_key=self.weight_quant_key,
173
+ force_kernel=force_kernel,
174
)
130
- self.enable_quant_fp8_custom_op = self.fp8_linear.quant_fp8.enabled()
175
+ for _ in range(3)
176
+ ]
177
132
- self.enable_rms_norm_custom_op = self.norm[0].enabled()
178
+ # Enable aiter quantization if requested
179
+ for layer in self.fp8_linear_layers:
180
+ layer.kernel.quant_fp8.use_aiter = use_aiter_quant
181
+
182
+ self.enable_quant_fp8_custom_op = self.fp8_linear_layers[
183
+ 0
184
+ ].is_quant_fp8_enabled()
185
186
def forward(self, x):
187
# avoid having graph input be an arg to a pattern directly
188
x = resid = torch.relu(x)
189
y = self.norm[0](x)
190
139
- x2 = self.fp8_linear.apply(
140
- y, self.w[0], self.wscale[0], input_scale=self.scale[0]
141
- )
191
+ x2 = self.fp8_linear_layers[0](y)
192
# make sure resid is used for replacement to work
193
y2, resid = self.norm[1](x2, resid)
194
145
- x3 = self.fp8_linear.apply(
146
- y2, self.w[1], self.wscale[1], input_scale=self.scale[1]
147
- )
195
+ x3 = self.fp8_linear_layers[1](y2)
196
197
y3, resid = self.norm[2](x3, resid) # use resid here
198
151
- x4 = self.fp8_linear.apply(
152
- y3, self.w[2], self.wscale[2], input_scale=self.scale[2]
153
- )
199
+ x4 = self.fp8_linear_layers[2](y3)
200
201
y4, resid = self.norm[3](x4, resid) # use resid here
202
return y4
203
204
def ops_in_model_before(self):
159
- if (
160
- self.use_aiter
161
- and self.group_shape.is_per_group()
162
- and current_platform.is_fp8_fnuz()
163
- ):
164
- return [rocm_aiter_ops.get_group_quant_op()]
165
- if self.use_aiter and self.group_shape.is_per_group():
166
- return [torch.ops.vllm.triton_per_token_group_quant_fp8.default]
167
- if self.use_aiter and self.use_aiter_quant_op:
168
- return [rocm_aiter_ops.get_per_token_quant_op()]
169
- if self.use_aiter:
170
- return [QUANT_OPS[self.quant_key]]
171
- if self.enable_quant_fp8_custom_op:
172
- return [QUANT_OPS[self.quant_key]]
173
- return [torch.ops.aten.reciprocal]
205
+ if self.group_shape.is_per_group():
206
+ # Blockwise path
207
+ if self.use_aiter_fusion and self.use_aiter_quant_op:
208
+ return [rocm_aiter_ops.get_group_quant_op()]
209
+ if self.use_aiter_fusion:
210
+ return [torch.ops.vllm.triton_per_token_group_quant_fp8.default]
211
+ else:
212
+ if self.use_aiter_quant_op:
213
+ return [rocm_aiter_ops.get_per_token_quant_op()]
214
+
215
+ # Common path
216
+ return (
217
+ [QUANT_OPS[self.activation_quant_key]]
218
+ if self.enable_quant_fp8_custom_op
219
+ else [torch.ops.aten.reciprocal]
220
+ )
221
222
def ops_in_model_after(self):
176
- if self.use_aiter and self.group_shape.is_per_group():
177
- from vllm.compilation.rocm_aiter_fusion import (
178
- AiterFusedAddRMSFp8GroupQuantPattern,
179
- AiterRMSFp8GroupQuantPattern,
180
- )
223
+ if self.use_aiter_fusion:
224
+ if self.group_shape.is_per_group():
225
+ # Blockwise aiter fusion
226
+ from vllm.compilation.rocm_aiter_fusion import (
227
+ AiterFusedAddRMSFp8GroupQuantPattern,
228
+ AiterRMSFp8GroupQuantPattern,
229
+ )
230
182
- return [
183
- AiterFusedAddRMSFp8GroupQuantPattern.FUSED_OP,
184
- AiterRMSFp8GroupQuantPattern.FUSED_OP,
185
- ]
186
- if self.use_aiter:
187
- from vllm.compilation.rocm_aiter_fusion import (
188
- AiterFusedAddRMSNormDynamicQuantPattern,
189
- AiterRMSNormDynamicQuantPattern,
190
- )
231
+ return [
232
+ AiterFusedAddRMSFp8GroupQuantPattern.FUSED_OP,
233
+ AiterRMSFp8GroupQuantPattern.FUSED_OP,
234
+ ]
235
+ else:
236
+ # Per-token aiter fusion
237
+ from vllm.compilation.rocm_aiter_fusion import (
238
+ AiterFusedAddRMSNormDynamicQuantPattern,
239
+ AiterRMSNormDynamicQuantPattern,
240
+ )
241
192
- return [
193
- AiterFusedAddRMSNormDynamicQuantPattern.FUSED_OP,
194
- AiterRMSNormDynamicQuantPattern.FUSED_OP,
195
- ]
242
+ return [
243
+ AiterFusedAddRMSNormDynamicQuantPattern.FUSED_OP,
244
+ AiterRMSNormDynamicQuantPattern.FUSED_OP,
245
+ ]
246
+
247
+ # Regular fusion
248
return [
197
- FUSED_OPS[FusedRMSQuantKey(self.quant_key, True)],
198
- FUSED_OPS[FusedRMSQuantKey(self.quant_key, False)],
249
+ FUSED_OPS[FusedRMSQuantKey(self.activation_quant_key, True)],
250
+ FUSED_OPS[FusedRMSQuantKey(self.activation_quant_key, False)],
251
]
252
253
def ops_in_model_before_partial(self):
@@ -206,14 +258,6 @@ class TestModel(torch.nn.Module):
258
)
259
260
209
-GROUP_SHAPES = [
210
- GroupShape.PER_TOKEN,
211
- GroupShape.PER_TENSOR,
212
- GroupShape(1, 128),
213
- GroupShape(1, 64),
214
-]
215
-
216
-
261
def _run_fusion_test(
262
model,
263
fusion_pass,
@@ -259,14 +303,9 @@ def _run_fusion_test(
303
@pytest.mark.parametrize("hidden_size", [256])
304
@pytest.mark.parametrize("num_tokens", [257])
305
@pytest.mark.parametrize("eps", [1e-5, 1e-6])
262
-@pytest.mark.parametrize("group_shape", GROUP_SHAPES)
306
+@pytest.mark.parametrize("kernel_groupshape", KERNEL_GROUPSHAPE_COMBINATIONS)
307
@pytest.mark.parametrize("enable_rms_norm_custom_op", [True, False])
308
@pytest.mark.parametrize("enable_quant_fp8_custom_op", [True, False])
265
-# cuda_force_torch used to test torch code path on platforms that
266
-# cutlass_fp8_supported() == True.
267
-@pytest.mark.parametrize(
268
- "cuda_force_torch", [True, False] if cutlass_fp8_supported() else [True]
269
-)
309
@pytest.mark.skipif(
310
not current_platform.is_cuda_alike(), reason="Only test on CUDA and ROCm"
311
)
@@ -275,11 +314,12 @@ def test_fusion_rmsnorm_quant(
314
hidden_size,
315
num_tokens,
316
eps,
278
- group_shape,
317
+ kernel_groupshape,
318
enable_rms_norm_custom_op,
319
enable_quant_fp8_custom_op,
281
- cuda_force_torch,
320
):
321
+ force_kernel, group_shape = kernel_groupshape
322
+
323
if not enable_quant_fp8_custom_op and group_shape.is_per_group():
324
pytest.skip("Unsupported unwrapped quant fp8 op for blockwise quantization")
325
@@ -310,15 +350,16 @@ def test_fusion_rmsnorm_quant(
350
torch.set_default_device("cuda")
351
torch.set_default_dtype(dtype)
352
torch.manual_seed(1)
313
- maybe_create_device_identity()
353
354
fusion_pass = RMSNormQuantFusionPass(vllm_config)
355
+
356
model = TestModel(
357
hidden_size=hidden_size,
358
eps=eps,
359
+ force_kernel=force_kernel,
360
group_shape=group_shape,
320
- use_aiter=False,
321
- cuda_force_torch=cuda_force_torch,
361
+ use_aiter_fusion=False,
362
+ use_aiter_quant=False,
363
)
364
365
backend, _ = _run_fusion_test(
@@ -339,19 +380,12 @@ def test_fusion_rmsnorm_quant(
380
assert n_add_nodes(backend.graph_post_pass) == 2
381
382
342
-GROUP_SHAPE_QUANT_OPS_MATCHS = [
343
- (GroupShape.PER_TOKEN, True),
344
- (GroupShape.PER_TOKEN, False),
345
- (GroupShape(1, 128), True),
346
-]
347
-
348
-
383
@pytest.mark.parametrize("dtype", [torch.bfloat16])
384
@pytest.mark.parametrize("hidden_size", [256])
385
@pytest.mark.parametrize("num_tokens", [257])
386
@pytest.mark.parametrize("eps", [1e-5, 1e-6])
387
@pytest.mark.parametrize(
354
- "group_shape, use_aiter_quant_op", GROUP_SHAPE_QUANT_OPS_MATCHS
388
+ "kernel_groupshape_quant", AITER_KERNEL_GROUPSHAPE_COMBINATIONS
389
)
390
@pytest.mark.skipif(
391
(not current_platform.is_rocm() or not IS_AITER_FOUND),
@@ -362,10 +396,10 @@ def test_aiter_fusion_rmsnorm_quant(
396
hidden_size: int,
397
num_tokens: int,
398
eps: float,
365
- group_shape: GroupShape,
366
- use_aiter_quant_op: bool,
399
+ kernel_groupshape_quant: tuple,
400
monkeypatch: pytest.MonkeyPatch,
401
):
402
+ force_kernel, group_shape, use_aiter_quant_op = kernel_groupshape_quant
403
vllm_config = VllmConfig(
404
model_config=ModelConfig(dtype=dtype),
405
compilation_config=CompilationConfig(
@@ -379,20 +413,22 @@ def test_aiter_fusion_rmsnorm_quant(
413
from vllm.compilation.rocm_aiter_fusion import RocmAiterRMSNormFusionPass
414
415
m.setenv("VLLM_ROCM_USE_AITER", "1")
416
+
417
rocm_aiter_ops.refresh_env_variables()
418
419
torch.set_default_device("cuda")
420
torch.set_default_dtype(dtype)
421
torch.manual_seed(1)
387
- maybe_create_device_identity()
422
423
fusion_pass = RocmAiterRMSNormFusionPass(vllm_config)
424
+
425
model = TestModel(
426
hidden_size=hidden_size,
427
eps=eps,
428
+ force_kernel=force_kernel,
429
group_shape=group_shape,
394
- use_aiter=True,
395
- use_aiter_quant_op=use_aiter_quant_op,
430
+ use_aiter_fusion=True, # Always use aiter fusion ops in aiter test
431
+ use_aiter_quant=use_aiter_quant_op, # Toggle aiter quantization
432
)
433
434
_run_fusion_test(
tests/compile/test_fusion_attn.py
+20
-21
@@ -45,7 +45,6 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import (
45
kFp8StaticTensorSym,
46
kNvfp4Quant,
47
)
48
-from vllm.model_executor.layers.quantization.utils.w8a8_utils import Fp8LinearOp
48
from vllm.platforms import current_platform
49
from vllm.utils.flashinfer import has_flashinfer
50
from vllm.utils.torch_utils import is_torch_equal_or_newer
@@ -53,6 +52,8 @@ from vllm.v1.attention.backend import AttentionMetadata
52
from vllm.v1.attention.backends.registry import AttentionBackendEnum
53
from vllm.v1.kv_cache_interface import AttentionSpec
54
55
+from ..utils import TestFP8Layer
56
+
57
FP8_DTYPE = current_platform.fp8_dtype()
58
FP4_DTYPE = torch.uint8
59
@@ -185,32 +186,30 @@ class TestAttentionFp8StaticQuantPatternModel(AttentionQuantPatternModel):
186
def __init__(self, *args, **kwargs):
187
super().__init__(*args, **kwargs)
188
188
- self.fp8_linear = Fp8LinearOp(
189
- act_quant_static=self.quant_key.scale.static,
190
- act_quant_group_shape=self.quant_key.scale.group_shape,
191
- )
192
-
189
hidden_size = self.num_qo_heads * self.head_size
194
- self.w = kwargs.get(
195
- "w",
196
- {
197
- "weight": torch.randn(hidden_size, hidden_size)
198
- .to(dtype=FP8_DTYPE, device=self.device)
199
- .t(),
200
- "wscale": torch.tensor([1.0], dtype=torch.float32, device=self.device),
201
- "scale": torch.tensor([1.0], dtype=torch.float32, device=self.device),
202
- },
190
+ self.fp8_linear = TestFP8Layer(
191
+ weight_shape=(hidden_size, hidden_size),
192
+ activation_quant_key=self.quant_key,
193
+ weight_quant_key=self.quant_key,
194
+ device=self.device,
195
)
196
197
+ w = kwargs.get("w")
198
+ if w is not None:
199
+ self.fp8_linear.weight = w["weight"]
200
+ self.fp8_linear.weight_scale = w["wscale"]
201
+ self.fp8_linear.input_scale = w["scale"]
202
+
203
+ self.w = {
204
+ "weight": self.fp8_linear.weight,
205
+ "wscale": self.fp8_linear.weight_scale,
206
+ "scale": self.fp8_linear.input_scale,
207
+ }
208
+
209
def forward(self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor):
210
"""Forward pass that creates the pattern to be fused."""
211
attn_output = self.attn(q, k, v)
208
- return self.fp8_linear.apply(
209
- input=attn_output,
210
- weight=self.w["weight"],
211
- weight_scale=self.w["wscale"],
212
- input_scale=self.w["scale"],
213
- )
212
+ return self.fp8_linear(attn_output)
213
214
215
class TestAttentionNvfp4QuantPatternModel(AttentionQuantPatternModel):
tests/compile/test_silu_mul_quant_fusion.py
+44
-27
@@ -25,19 +25,30 @@ from vllm.config import (
25
set_current_vllm_config,
26
)
27
from vllm.model_executor.layers.activation import SiluAndMul
28
+from vllm.model_executor.layers.quantization.kernels.scaled_mm.cutlass import (
29
+ CutlassFP8ScaledMMLinearKernel,
30
+)
31
+from vllm.model_executor.layers.quantization.kernels.scaled_mm.flashinfer import (
32
+ FlashInferFP8ScaledMMLinearKernel,
33
+)
34
+from vllm.model_executor.layers.quantization.kernels.scaled_mm.pytorch import (
35
+ PerTensorTorchFP8ScaledMMLinearKernel,
36
+)
37
+from vllm.model_executor.layers.quantization.kernels.scaled_mm.rocm import (
38
+ ROCmFP8ScaledMMLinearKernel,
39
+)
40
+from vllm.model_executor.layers.quantization.kernels.scaled_mm.ScaledMMLinearKernel import ( # noqa: E501
41
+ FP8ScaledMMLinearKernel,
42
+)
43
from vllm.model_executor.layers.quantization.utils.fp8_utils import W8A8BlockFp8LinearOp
44
from vllm.model_executor.layers.quantization.utils.quant_utils import (
45
GroupShape,
46
kFp8StaticTensorSym,
47
kNvfp4Quant,
48
)
34
-from vllm.model_executor.layers.quantization.utils.w8a8_utils import (
35
- Fp8LinearOp,
36
- maybe_create_device_identity,
37
-)
49
from vllm.platforms import current_platform
50
40
-from ..utils import override_cutlass_fp8_supported
51
+from ..utils import TestFP8Layer
52
from .backend import TestBackend
53
54
FP8_DTYPE = current_platform.fp8_dtype()
@@ -49,25 +60,27 @@ def is_nvfp4_supported():
60
61
62
class TestSiluMulFp8QuantModel(torch.nn.Module):
52
- def __init__(self, hidden_size: int, cuda_force_torch: bool, **kwargs):
63
+ quant_key = kFp8StaticTensorSym
64
+
65
+ def __init__(
66
+ self, hidden_size: int, force_kernel: FP8ScaledMMLinearKernel, **kwargs
67
+ ):
68
super().__init__()
69
self.silu_and_mul = SiluAndMul()
55
- self.wscale = torch.rand(1, dtype=torch.float32)
56
- self.scale = torch.rand(1, dtype=torch.float32)
70
58
- self.w = torch.rand(hidden_size, hidden_size).to(dtype=FP8_DTYPE).t()
71
+ self.fp8_linear = TestFP8Layer(
72
+ weight_shape=(hidden_size, hidden_size),
73
+ activation_quant_key=self.quant_key,
74
+ weight_quant_key=self.quant_key,
75
+ force_kernel=force_kernel,
76
+ )
77
60
- with override_cutlass_fp8_supported(not cuda_force_torch):
61
- self.fp8_linear = Fp8LinearOp(
62
- act_quant_static=True,
63
- act_quant_group_shape=GroupShape.PER_TENSOR,
64
- )
78
self.enable_silu_mul_custom_op = self.silu_and_mul.enabled()
66
- self.enable_quant_fp8_custom_op = self.fp8_linear.quant_fp8.enabled()
79
+ self.enable_quant_fp8_custom_op = self.fp8_linear.is_quant_fp8_enabled()
80
81
def forward(self, x):
82
y = self.silu_and_mul(x)
70
- x2 = self.fp8_linear.apply(y, self.w, self.wscale, input_scale=self.wscale)
83
+ x2 = self.fp8_linear(y)
84
return x2
85
86
def ops_in_model_before(self):
@@ -161,20 +174,27 @@ class TestSiluMulGroupFp8QuantModel(torch.nn.Module):
174
return [torch.ops.vllm.rocm_aiter_act_mul_and_fp8_group_quant]
175
176
177
+ROCM_KERNELS = [ROCmFP8ScaledMMLinearKernel, PerTensorTorchFP8ScaledMMLinearKernel]
178
+CUDA_KERNELS = [
179
+ FlashInferFP8ScaledMMLinearKernel,
180
+ CutlassFP8ScaledMMLinearKernel,
181
+ PerTensorTorchFP8ScaledMMLinearKernel,
182
+]
183
+TEST_KERNELS = ROCM_KERNELS if current_platform.is_rocm() else CUDA_KERNELS
184
+
185
+
186
@pytest.mark.parametrize("num_tokens", [32, 64])
187
@pytest.mark.parametrize("hidden_size", [128, 256])
188
@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16])
189
@pytest.mark.parametrize("enable_silu_mul_custom_op", [True, False])
190
@pytest.mark.parametrize(
169
- "model_class, enable_quant_fp8_custom_op, cuda_force_torch",
170
- list(itertools.product([TestSiluMulFp8QuantModel], [True, False], [True, False]))
191
+ "model_class, enable_quant_fp8_custom_op, force_kernel",
192
+ list(itertools.product([TestSiluMulFp8QuantModel], [True, False], TEST_KERNELS))
193
+ [
172
- (TestSiluMulNvfp4QuantModel, False, False),
173
- (TestSiluMulGroupFp8QuantModel, False, False),
194
+ (TestSiluMulNvfp4QuantModel, False, None),
195
+ (TestSiluMulGroupFp8QuantModel, False, None),
196
],
197
)
176
-# cuda_force_torch used to test torch code path on platforms that
177
-# cutlass_fp8_supported() == True.
198
@pytest.mark.skipif(
199
envs.VLLM_TARGET_DEVICE not in ["cuda", "rocm"], reason="Only test on CUDA and ROCm"
200
)
@@ -189,7 +209,7 @@ def test_fusion_silu_and_mul_quant(
209
],
210
enable_silu_mul_custom_op: bool,
211
enable_quant_fp8_custom_op: bool,
192
- cuda_force_torch: bool,
212
+ force_kernel: FP8ScaledMMLinearKernel | None,
213
):
214
if model_class is TestSiluMulNvfp4QuantModel and not is_nvfp4_supported():
215
pytest.skip("NVFP4 is not supported on this GPU.")
@@ -198,7 +218,6 @@ def test_fusion_silu_and_mul_quant(
218
219
torch.set_default_device("cuda")
220
torch.set_default_dtype(dtype)
201
- maybe_create_device_identity()
221
222
x = torch.rand(num_tokens, hidden_size * 2)
223
@@ -227,9 +246,7 @@ def test_fusion_silu_and_mul_quant(
246
247
passes = [NoOpEliminationPass(config), *fusion_passes, PostCleanupPass(config)]
248
backend = TestBackend(*passes)
230
- model = model_class(
231
- hidden_size=hidden_size, cuda_force_torch=cuda_force_torch, x=x
232
- )
249
+ model = model_class(hidden_size=hidden_size, force_kernel=force_kernel, x=x)
250
251
# First dimension dynamic
252
torch._dynamo.mark_dynamic(x, 0)
tests/kernels/quantization/test_scaled_mm_kernel_selection.py
+24
-22
@@ -11,13 +11,13 @@ from abc import ABC
11
import pytest
12
13
from vllm.model_executor.layers.quantization.kernels.scaled_mm import (
14
- ScaledMMLinearLayerConfig,
14
+ Int8ScaledMMLinearLayerConfig,
15
)
16
from vllm.model_executor.layers.quantization.kernels.scaled_mm.aiter import (
17
- AiterScaledMMLinearKernel,
17
+ AiterInt8ScaledMMLinearKernel,
18
)
19
from vllm.model_executor.layers.quantization.kernels.scaled_mm.cpu import (
20
- CPUScaledMMLinearKernel,
20
+ CPUInt8ScaledMMLinearKernel,
21
)
22
from vllm.model_executor.layers.quantization.kernels.scaled_mm.ScaledMMLinearKernel import ( # noqa: E501
23
ScaledMMLinearKernel,
@@ -33,36 +33,38 @@ def test_is_supported_is_abstract():
33
34
35
def test_cpu_kernel_implements_is_supported():
36
- """Test that CPUScaledMMLinearKernel implements is_supported() method."""
37
- assert hasattr(CPUScaledMMLinearKernel, "is_supported"), (
38
- "CPUScaledMMLinearKernel missing is_supported() method"
36
+ """Test that CPUInt8ScaledMMLinearKernel implements is_supported() method."""
37
+ assert hasattr(CPUInt8ScaledMMLinearKernel, "is_supported"), (
38
+ "CPUInt8ScaledMMLinearKernel missing is_supported() method"
39
)
40
# Verify it's a classmethod by checking if it can be called with the class
41
# and by checking the method type
42
- assert inspect.ismethod(CPUScaledMMLinearKernel.is_supported) or inspect.isfunction(
43
- CPUScaledMMLinearKernel.is_supported
44
- ), "CPUScaledMMLinearKernel.is_supported() should be a classmethod"
42
+ assert inspect.ismethod(
43
+ CPUInt8ScaledMMLinearKernel.is_supported
44
+ ) or inspect.isfunction(CPUInt8ScaledMMLinearKernel.is_supported), (
45
+ "CPUInt8ScaledMMLinearKernel.is_supported() should be a classmethod"
46
+ )
47
# Verify it can be called as a classmethod
46
- result, reason = CPUScaledMMLinearKernel.is_supported()
48
+ result, reason = CPUInt8ScaledMMLinearKernel.is_supported()
49
assert isinstance(result, bool), "is_supported() should return a bool"
50
assert reason is None or isinstance(reason, str), "reason should be str or None"
51
52
53
def test_aiter_kernel_implements_is_supported():
52
- """Test that AiterScaledMMLinearKernel implements is_supported() method."""
53
- assert hasattr(AiterScaledMMLinearKernel, "is_supported"), (
54
- "AiterScaledMMLinearKernel missing is_supported() method"
54
+ """Test that AiterInt8ScaledMMLinearKernel implements is_supported() method."""
55
+ assert hasattr(AiterInt8ScaledMMLinearKernel, "is_supported"), (
56
+ "AiterInt8ScaledMMLinearKernel missing is_supported() method"
57
)
58
# Verify it's a classmethod by checking if it can be called with the class
59
# and by checking the method type
60
assert inspect.ismethod(
59
- AiterScaledMMLinearKernel.is_supported
60
- ) or inspect.isfunction(AiterScaledMMLinearKernel.is_supported), (
61
- "AiterScaledMMLinearKernel.is_supported() should be a classmethod"
61
+ AiterInt8ScaledMMLinearKernel.is_supported
62
+ ) or inspect.isfunction(AiterInt8ScaledMMLinearKernel.is_supported), (
63
+ "AiterInt8ScaledMMLinearKernel.is_supported() should be a classmethod"
64
)
65
# Verify it can be called as a classmethod
66
# (will return False on CPU, which is expected)
65
- result, reason = AiterScaledMMLinearKernel.is_supported()
67
+ result, reason = AiterInt8ScaledMMLinearKernel.is_supported()
68
assert isinstance(result, bool), "is_supported() should return a bool"
69
assert reason is None or isinstance(reason, str), "reason should be str or None"
70
# On CPU, it should return False with a reason about requiring ROCm
@@ -70,14 +72,14 @@ def test_aiter_kernel_implements_is_supported():
72
73
74
def test_cpu_kernel_accepts_all_configs():
73
- """Test that CPUScaledMMLinearKernel accepts all config combinations."""
75
+ """Test that CPUInt8ScaledMMLinearKernel accepts all config combinations."""
76
configs = [
75
- ScaledMMLinearLayerConfig(
77
+ Int8ScaledMMLinearLayerConfig(
78
is_channelwise=False,
79
is_static_input_scheme=True,
80
input_symmetric=True,
81
),
80
- ScaledMMLinearLayerConfig(
82
+ Int8ScaledMMLinearLayerConfig(
83
is_channelwise=True,
84
is_static_input_scheme=False,
85
input_symmetric=False,
@@ -85,7 +87,7 @@ def test_cpu_kernel_accepts_all_configs():
87
]
88
89
for config in configs:
88
- can_impl, reason = CPUScaledMMLinearKernel.can_implement(config)
90
+ can_impl, reason = CPUInt8ScaledMMLinearKernel.can_implement(config)
91
assert can_impl, (
90
- f"CPUScaledMMLinearKernel should accept config {config}: {reason}"
92
+ f"CPUInt8ScaledMMLinearKernel should accept config {config}: {reason}"
93
)
tests/quantization/test_compressed_tensors.py
+1
-1
@@ -41,7 +41,7 @@ ROCM_AITER_SUPPORTED_INT8_MODEL = [
41
"nm-testing/tinyllama-oneshot-w8a8-channel-dynamic-token-v2",
42
]
43
44
-# TritonScaledMMLinearKernel only supports symmetric quantization.
44
+# TritonInt8ScaledMMLinearKernel only supports symmetric quantization.
45
ROCM_TRITON_SCALED_MM_SUPPORTED_INT8_MODEL = [
46
"nm-testing/tinyllama-oneshot-w8w8-test-static-shape-change",
47
"nm-testing/tinyllama-oneshot-w8-channel-a8-tensor",
tests/utils.py
+127
@@ -42,6 +42,17 @@ from vllm.distributed import (
42
)
43
from vllm.engine.arg_utils import AsyncEngineArgs
44
from vllm.entrypoints.cli.serve import ServeSubcommand
45
+from vllm.model_executor.layers.quantization.kernels.scaled_mm import (
46
+ init_fp8_linear_kernel,
47
+)
48
+from vllm.model_executor.layers.quantization.kernels.scaled_mm.ScaledMMLinearKernel import ( # noqa: E501
49
+ FP8ScaledMMLinearKernel,
50
+)
51
+from vllm.model_executor.layers.quantization.utils.fp8_utils import W8A8BlockFp8LinearOp
52
+from vllm.model_executor.layers.quantization.utils.quant_utils import (
53
+ GroupShape,
54
+ QuantKey,
55
+)
56
from vllm.model_executor.model_loader import get_model_loader
57
from vllm.platforms import current_platform
58
from vllm.tokenizers import get_tokenizer
@@ -50,6 +61,8 @@ from vllm.utils.mem_constants import GB_bytes
61
from vllm.utils.network_utils import get_open_port
62
from vllm.utils.torch_utils import cuda_device_count_stateless
63
64
+FP8_DTYPE = current_platform.fp8_dtype()
65
+
66
if current_platform.is_rocm():
67
from amdsmi import (
68
amdsmi_get_gpu_vram_usage,
@@ -1332,3 +1345,117 @@ def flat_product(*iterables: Iterable[Any]):
1345
for element in itertools.product(*iterables):
1346
normalized = (e if isinstance(e, tuple) else (e,) for e in element)
1347
yield tuple(itertools.chain(*normalized))
1348
+
1349
+
1350
+class TestFP8Layer(torch.nn.Module):
1351
+ """
1352
+ Test helper for FP8 linear operations. Creates random weights and scales
1353
+ based on quantization configuration.
1354
+
1355
+ Args:
1356
+ weight_shape: Shape of the weight tensor (out_features, in_features).
1357
+ activation_quant_key: Activation quantization configuration.
1358
+ weight_quant_key: Weight quantization configuration.
1359
+ out_dtype: Output dtype. Defaults to current default dtype.
1360
+ force_kernel: Optional kernel to force use of specific implementation.
1361
+ """
1362
+
1363
+ def __init__(
1364
+ self,
1365
+ weight_shape: tuple[int, int],
1366
+ activation_quant_key: QuantKey,
1367
+ weight_quant_key: QuantKey,
1368
+ out_dtype: torch.dtype | None = None,
1369
+ device: torch.device | None = None,
1370
+ force_kernel: FP8ScaledMMLinearKernel | None = None,
1371
+ ):
1372
+ super().__init__()
1373
+ per_tensor_weights = weight_quant_key.scale.group_shape.is_per_tensor()
1374
+ is_static_activation_scale = activation_quant_key.scale.static
1375
+ weight_scale_shape = (1,) if per_tensor_weights else (weight_shape[0], 1)
1376
+
1377
+ self.weight_scale = torch.rand(
1378
+ weight_scale_shape, dtype=torch.float32, device=device
1379
+ )
1380
+ self.input_scale = (
1381
+ torch.rand(1, dtype=torch.float32, device=device)
1382
+ if is_static_activation_scale
1383
+ else None
1384
+ )
1385
+ self.weight = torch.rand(weight_shape, device=device).to(dtype=FP8_DTYPE).t()
1386
+ self.input_scale_ub = None
1387
+
1388
+ out_dtype = torch.get_default_dtype() if out_dtype is None else out_dtype
1389
+
1390
+ self.kernel = init_fp8_linear_kernel(
1391
+ activation_quant_key=activation_quant_key,
1392
+ weight_quant_key=weight_quant_key,
1393
+ out_dtype=out_dtype,
1394
+ force_kernel=force_kernel,
1395
+ )
1396
+
1397
+ def is_quant_fp8_enabled(self) -> bool:
1398
+ return self.kernel.quant_fp8.enabled()
1399
+
1400
+ def forward(
1401
+ self, y: torch.Tensor, bias: torch.Tensor | None = None
1402
+ ) -> torch.Tensor:
1403
+ return self.kernel.apply_weights(self, y, bias)
1404
+
1405
+
1406
+# TODO: Drop TestBlockFP8Layer in favour of a unified TestFP8Layer
1407
+# after refactoring W8A8BlockFp8LinearOp.
1408
+# https://github.com/vllm-project/vllm/issues/31818
1409
+class TestBlockFP8Layer:
1410
+ """
1411
+ Test helper for blockwise FP8 linear operations. Creates random weights
1412
+ and scales for W8A8BlockFp8LinearOp.
1413
+
1414
+ This is a workaround until W8A8BlockFp8LinearOp implements the kernel
1415
+ abstraction (ScaledMMLinearKernel) for blockwise quantization.
1416
+
1417
+ Args:
1418
+ weight_shape: Shape of the weight tensor (out_features, in_features).
1419
+ group_shape: Blockwise quantization group shape.
1420
+ cutlass_block_fp8_supported: Whether CUTLASS blockwise FP8 is available.
1421
+ use_aiter_and_is_supported: Whether to use aiter quantization ops.
1422
+ transpose_weights: Whether to transpose weights after creation.
1423
+ """
1424
+
1425
+ def __init__(
1426
+ self,
1427
+ weight_shape: tuple[int, int],
1428
+ group_shape: GroupShape,
1429
+ cutlass_block_fp8_supported: bool = False,
1430
+ use_aiter_and_is_supported: bool = False,
1431
+ transpose_weights: bool = False,
1432
+ ):
1433
+ weight_scale_shape = weight_shape[0] // group_shape[1]
1434
+ self.weight_scale = torch.rand(
1435
+ (weight_scale_shape, weight_scale_shape), dtype=torch.float32
1436
+ )
1437
+ self.weight = torch.rand(weight_shape).to(dtype=FP8_DTYPE)
1438
+ self.input_scale = None
1439
+ if transpose_weights:
1440
+ self.weight = self.weight.t()
1441
+
1442
+ self.linear_op = W8A8BlockFp8LinearOp(
1443
+ weight_group_shape=GroupShape(group_shape[1], group_shape[1]),
1444
+ act_quant_group_shape=group_shape,
1445
+ cutlass_block_fp8_supported=cutlass_block_fp8_supported,
1446
+ use_aiter_and_is_supported=use_aiter_and_is_supported,
1447
+ )
1448
+
1449
+ def __call__(
1450
+ self, y: torch.Tensor, bias: torch.Tensor | None = None
1451
+ ) -> torch.Tensor:
1452
+ return self.linear_op.apply(
1453
+ input=y,
1454
+ weight=self.weight,
1455
+ weight_scale=self.weight_scale,
1456
+ input_scale=self.input_scale,
1457
+ bias=bias,
1458
+ )
1459
+
1460
+ def is_quant_fp8_enabled(self) -> bool:
1461
+ return self.linear_op.input_quant_op.enabled()
vllm/_aiter_ops.py
+1
-1
@@ -372,7 +372,7 @@ def _rocm_aiter_gemm_a8w8_impl(
372
# gemm_a8w8_CK(a, b, scale_a, scale_b, bias) expects
373
# a to be [M, K]
374
# b to be [N, K]
375
- # CutlassScaledMMLinearKernel prepare weight `w_q` in [K, N] format
375
+ # CutlassInt8ScaledMMLinearKernel prepare weight `w_q` in [K, N] format
376
return gemm_a8w8_CK(A, B, As, Bs, bias, output_dtype)
377
378
vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_w8a8_fp8.py
+33
-30
@@ -8,9 +8,13 @@ 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.logger import init_logger
12
from vllm.model_executor.layers.quantization.compressed_tensors.schemes import (
13
CompressedTensorsScheme,
14
)
15
+from vllm.model_executor.layers.quantization.kernels.scaled_mm import (
16
+ init_fp8_linear_kernel,
17
+)
18
from vllm.model_executor.layers.quantization.utils.fp8_utils import (
19
W8A8BlockFp8LinearOp,
20
create_fp8_input_scale,
@@ -22,11 +26,14 @@ from vllm.model_executor.layers.quantization.utils.fp8_utils import (
26
process_fp8_weight_tensor_strategy,
27
validate_fp8_block_shape,
28
)
25
-from vllm.model_executor.layers.quantization.utils.quant_utils import GroupShape
29
+from vllm.model_executor.layers.quantization.utils.quant_utils import (
30
+ GroupShape,
31
+ kFp8DynamicTokenSym,
32
+ kFp8StaticTensorSym,
33
+ kFp8StaticTokenSym,
34
+)
35
from vllm.model_executor.layers.quantization.utils.w8a8_utils import (
27
- Fp8LinearOp,
36
cutlass_block_fp8_supported,
29
- maybe_create_device_identity,
37
)
38
from vllm.model_executor.parameter import (
39
BlockQuantScaleParameter,
@@ -42,6 +49,18 @@ strategy_to_parameter_type = {
49
QuantizationStrategy.TENSOR: PerTensorScaleParameter,
50
}
51
52
+STATIC_QUANT = True
53
+DYNAMIC_QUANT = False
54
+activation_quant_key_mapping = {
55
+ STATIC_QUANT: kFp8StaticTensorSym,
56
+ DYNAMIC_QUANT: kFp8DynamicTokenSym,
57
+}
58
+weight_quant_key_mapping = {
59
+ QuantizationStrategy.CHANNEL: kFp8StaticTokenSym,
60
+ QuantizationStrategy.TENSOR: kFp8StaticTensorSym,
61
+}
62
+logger = init_logger(__name__)
63
+
64
65
class CompressedTensorsW8A8Fp8(CompressedTensorsScheme):
66
def __init__(self, weight_quant: QuantizationArgs, is_static_input_scheme: bool):
@@ -49,22 +68,13 @@ class CompressedTensorsW8A8Fp8(CompressedTensorsScheme):
68
self.strategy = weight_quant.strategy
69
self.out_dtype = torch.get_default_dtype()
70
self.is_static_input_scheme = is_static_input_scheme
52
-
71
self.weight_block_size = self.weight_quant.block_structure
54
- if self.weight_block_size is not None:
55
- self.act_q_group_shape = GroupShape(1, self.weight_block_size[0])
56
- else:
57
- self.act_q_group_shape = (
58
- GroupShape.PER_TENSOR
59
- if is_static_input_scheme
60
- else GroupShape.PER_TOKEN
61
- )
62
-
63
- self.cutlass_block_fp8_supported = cutlass_block_fp8_supported()
64
- self.use_aiter_and_is_supported = rocm_aiter_ops.is_linear_fp8_enabled()
72
73
if self.weight_block_size is not None:
74
+ self.cutlass_block_fp8_supported = cutlass_block_fp8_supported()
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,
@@ -72,9 +82,13 @@ class CompressedTensorsW8A8Fp8(CompressedTensorsScheme):
82
use_aiter_and_is_supported=self.use_aiter_and_is_supported,
83
)
84
else:
75
- self.fp8_linear = Fp8LinearOp(
76
- act_quant_static=self.is_static_input_scheme,
77
- act_quant_group_shape=self.act_q_group_shape,
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__,
92
)
93
94
@classmethod
@@ -93,8 +107,6 @@ class CompressedTensorsW8A8Fp8(CompressedTensorsScheme):
107
weight_loader: Callable,
108
**kwargs,
109
):
96
- maybe_create_device_identity()
97
-
110
output_size_per_partition = sum(output_partition_sizes)
111
layer.logical_widths = output_partition_sizes
112
layer.weight_block_size = None
@@ -143,7 +155,6 @@ class CompressedTensorsW8A8Fp8(CompressedTensorsScheme):
155
getattr(layer, "input_scale", None),
156
)
157
weight = weight.t()
146
-
158
elif self.strategy == QuantizationStrategy.CHANNEL:
159
weight, weight_scale, input_scale = process_fp8_weight_channel_strategy(
160
layer.weight, layer.weight_scale, getattr(layer, "input_scale", None)
@@ -174,7 +185,6 @@ class CompressedTensorsW8A8Fp8(CompressedTensorsScheme):
185
layer.input_scale = Parameter(layer.input_scale.max(), requires_grad=False)
186
else:
187
layer.input_scale = None
177
-
188
if self.strategy == QuantizationStrategy.BLOCK:
189
maybe_post_process_fp8_weight_block(layer)
190
@@ -193,11 +203,4 @@ class CompressedTensorsW8A8Fp8(CompressedTensorsScheme):
203
bias=bias,
204
)
205
196
- return self.fp8_linear.apply(
197
- input=x,
198
- weight=layer.weight,
199
- weight_scale=layer.weight_scale,
200
- out_dtype=self.out_dtype,
201
- input_scale=layer.input_scale,
202
- bias=bias,
203
- )
206
+ return self.fp8_linear.apply_weights(layer, x, bias)
vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_w8a8_int8.py
+10
-23
@@ -11,8 +11,7 @@ from vllm.model_executor.layers.quantization.compressed_tensors.schemes import (
11
CompressedTensorsScheme,
12
)
13
from vllm.model_executor.layers.quantization.kernels.scaled_mm import (
14
- ScaledMMLinearLayerConfig,
15
- choose_scaled_mm_linear_kernel,
14
+ init_int8_linear_kernel,
15
)
16
from vllm.model_executor.parameter import (
17
BasevLLMParameter,
@@ -25,8 +24,6 @@ logger = init_logger(__name__)
24
25
26
class CompressedTensorsW8A8Int8(CompressedTensorsScheme):
28
- _kernel_backends_being_used: set[str] = set()
29
-
27
def __init__(
28
self, strategy: str, is_static_input_scheme: bool, input_symmetric: bool
29
):
@@ -50,18 +47,13 @@ class CompressedTensorsW8A8Int8(CompressedTensorsScheme):
47
):
48
layer.logical_widths = output_partition_sizes
49
53
- scaled_mm_linear_kernel_config = ScaledMMLinearLayerConfig(
50
+ self.kernel = init_int8_linear_kernel(
51
is_channelwise=(self.strategy == QuantizationStrategy.CHANNEL),
52
is_static_input_scheme=self.is_static_input_scheme,
53
input_symmetric=self.input_symmetric,
54
+ module_name=self.__class__.__name__,
55
)
56
59
- kernel_type = choose_scaled_mm_linear_kernel(scaled_mm_linear_kernel_config)
60
-
61
- if kernel_type.__name__ not in self._kernel_backends_being_used:
62
- logger.info("Using %s for CompressedTensorsW8A8Int8", kernel_type.__name__)
63
- self._kernel_backends_being_used.add(kernel_type.__name__)
64
-
57
# WEIGHT
58
weight = ModelWeightParameter(
59
data=torch.empty(
@@ -90,12 +82,12 @@ class CompressedTensorsW8A8Int8(CompressedTensorsScheme):
82
layer.register_parameter("weight_scale", weight_scale)
83
84
# INPUT SCALE
85
+ input_zero_point = None
86
+ input_scale = None
87
if self.is_static_input_scheme:
88
input_scale = BasevLLMParameter(
89
data=torch.empty(1, dtype=torch.float32), weight_loader=weight_loader
90
)
97
- layer.register_parameter("input_scale", input_scale)
98
-
91
if not self.input_symmetric:
92
# Note: compressed-tensors stores the zp using the same dtype
93
# as the weights
@@ -103,16 +95,11 @@ class CompressedTensorsW8A8Int8(CompressedTensorsScheme):
95
input_zero_point = BasevLLMParameter(
96
data=torch.empty(1, dtype=torch.int8), weight_loader=weight_loader
97
)
106
- layer.register_parameter("input_zero_point", input_zero_point)
107
-
108
- self.kernel = kernel_type(
109
- c=scaled_mm_linear_kernel_config,
110
- w_q_param_name="weight",
111
- w_s_param_name="weight_scale",
112
- i_s_param_name="input_scale",
113
- i_zp_param_name="input_zero_point",
114
- azp_adj_param_name="azp_adj",
115
- )
98
+
99
+ layer.register_parameter("input_zero_point", input_zero_point)
100
+ layer.register_parameter("input_scale", input_scale)
101
+ if not hasattr(layer, "azp_adj"):
102
+ layer.register_parameter("azp_adj", None)
103
104
# Checkpoints are serialized in compressed-tensors format, which is
105
# different from the format the kernel may want. Handle repacking here.
vllm/model_executor/layers/quantization/fbgemm_fp8.py
+12
-16
@@ -18,17 +18,19 @@ from vllm.model_executor.layers.quantization.base_config import (
18
QuantizationConfig,
19
QuantizeMethodBase,
20
)
21
+from vllm.model_executor.layers.quantization.kernels.scaled_mm import (
22
+ init_fp8_linear_kernel,
23
+)
24
from vllm.model_executor.layers.quantization.utils.marlin_utils_fp8 import (
25
apply_fp8_marlin_linear,
26
prepare_fp8_layer_for_marlin,
27
)
28
from vllm.model_executor.layers.quantization.utils.quant_utils import (
26
- GroupShape,
29
is_layer_skipped,
30
+ kFp8DynamicTokenSym,
31
+ kFp8StaticTokenSym,
32
)
33
from vllm.model_executor.layers.quantization.utils.w8a8_utils import (
30
- Fp8LinearOp,
31
- maybe_create_device_identity,
34
normalize_e4m3fn_to_e4m3fnuz,
35
)
36
from vllm.model_executor.parameter import (
@@ -91,10 +93,13 @@ class FBGEMMFp8Config(QuantizationConfig):
93
class FBGEMMFp8LinearMethod(LinearMethodBase):
94
def __init__(self, quant_config: FBGEMMFp8Config):
95
self.quant_config = quant_config
94
- self.fp8_linear = Fp8LinearOp(
95
- act_quant_static=False, act_quant_group_shape=GroupShape.PER_TOKEN
96
- )
96
self.out_dtype = torch.get_default_dtype()
97
+ self.fp8_linear = init_fp8_linear_kernel(
98
+ activation_quant_key=kFp8DynamicTokenSym,
99
+ weight_quant_key=kFp8StaticTokenSym,
100
+ out_dtype=torch.get_default_dtype(),
101
+ module_name=self.__class__.__name__,
102
+ )
103
104
def create_weights(
105
self,
@@ -106,7 +111,6 @@ class FBGEMMFp8LinearMethod(LinearMethodBase):
111
params_dtype: torch.dtype,
112
**extra_weight_attrs,
113
):
109
- maybe_create_device_identity()
114
weight_loader = extra_weight_attrs.get("weight_loader")
115
del input_size, output_size
116
output_size_per_partition = sum(output_partition_sizes)
@@ -184,12 +188,4 @@ class FBGEMMFp8LinearMethod(LinearMethodBase):
188
bias=bias,
189
)
190
187
- return self.fp8_linear.apply(
188
- input=x,
189
- weight=layer.weight,
190
- weight_scale=layer.weight_scale,
191
- out_dtype=self.out_dtype,
192
- input_scale=None,
193
- input_scale_ub=layer.input_scale_ub,
194
- bias=bias,
195
- )
191
+ return self.fp8_linear.apply_weights(layer, x, bias)
vllm/model_executor/layers/quantization/fp8.py
+21
-26
@@ -48,6 +48,9 @@ from vllm.model_executor.layers.quantization.base_config import (
48
QuantizationConfig,
49
QuantizeMethodBase,
50
)
51
+from vllm.model_executor.layers.quantization.kernels.scaled_mm import (
52
+ init_fp8_linear_kernel,
53
+)
54
from vllm.model_executor.layers.quantization.kv_cache import BaseKVCacheMethod
55
from vllm.model_executor.layers.quantization.utils.flashinfer_utils import (
56
apply_fi_trtllm_fp8_per_tensor_moe,
@@ -76,12 +79,13 @@ from vllm.model_executor.layers.quantization.utils.marlin_utils_fp8 import (
79
from vllm.model_executor.layers.quantization.utils.quant_utils import (
80
GroupShape,
81
is_layer_skipped,
82
+ kFp8DynamicTensorSym,
83
+ kFp8DynamicTokenSym,
84
+ kFp8StaticTensorSym,
85
)
86
from vllm.model_executor.layers.quantization.utils.w8a8_utils import (
81
- Fp8LinearOp,
87
cutlass_block_fp8_supported,
88
cutlass_fp8_supported,
84
- maybe_create_device_identity,
89
normalize_e4m3fn_to_e4m3fnuz,
90
)
91
from vllm.model_executor.parameter import (
@@ -328,28 +332,30 @@ class Fp8LinearMethod(LinearMethodBase):
332
self.weight_block_size = self.quant_config.weight_block_size
333
self.block_quant = self.weight_block_size is not None
334
self.act_q_static = self.quant_config.activation_scheme == "static"
331
- if self.weight_block_size:
332
- self.act_q_group_shape = GroupShape(1, self.weight_block_size[0])
333
- else:
334
- # Use per-token quantization for better perf if dynamic and cutlass
335
- if not self.act_q_static and cutlass_fp8_supported():
336
- self.act_q_group_shape = GroupShape.PER_TOKEN
337
- else:
338
- self.act_q_group_shape = GroupShape.PER_TENSOR
335
336
if self.block_quant:
337
assert not self.act_q_static
338
assert self.weight_block_size is not None
339
self.w8a8_block_fp8_linear = W8A8BlockFp8LinearOp(
340
weight_group_shape=GroupShape(*self.weight_block_size),
345
- act_quant_group_shape=self.act_q_group_shape,
341
+ act_quant_group_shape=GroupShape(1, self.weight_block_size[0]),
342
cutlass_block_fp8_supported=self.cutlass_block_fp8_supported,
343
use_aiter_and_is_supported=self.use_aiter_and_is_supported,
344
)
345
else:
350
- self.fp8_linear = Fp8LinearOp(
351
- act_quant_static=self.act_q_static,
352
- act_quant_group_shape=self.act_q_group_shape,
346
+ # Use per-token quantization for better perf if dynamic and cutlass
347
+ if self.act_q_static:
348
+ activation_quant_key = kFp8StaticTensorSym
349
+ elif cutlass_fp8_supported():
350
+ activation_quant_key = kFp8DynamicTokenSym
351
+ else:
352
+ activation_quant_key = kFp8DynamicTensorSym
353
+
354
+ self.fp8_linear = init_fp8_linear_kernel(
355
+ activation_quant_key=activation_quant_key,
356
+ weight_quant_key=kFp8StaticTensorSym,
357
+ out_dtype=torch.get_default_dtype(),
358
+ module_name=self.__class__.__name__,
359
)
360
361
def create_weights(
@@ -362,8 +368,6 @@ class Fp8LinearMethod(LinearMethodBase):
368
params_dtype: torch.dtype,
369
**extra_weight_attrs,
370
):
365
- maybe_create_device_identity()
366
-
371
output_size_per_partition = sum(output_partition_sizes)
372
weight_loader = extra_weight_attrs.get("weight_loader")
373
layer.logical_widths = output_partition_sizes
@@ -462,8 +466,6 @@ class Fp8LinearMethod(LinearMethodBase):
466
scale = create_fp8_input_scale(output_partition_sizes, weight_loader)
467
set_weight_attrs(scale, {"scale_type": "input_scale"})
468
layer.register_parameter("input_scale", scale)
465
- else:
466
- layer.register_parameter("input_scale", None)
469
470
def process_weights_after_loading(self, layer: Module) -> None:
471
if getattr(layer, "_already_called_process_weights_after_loading", False):
@@ -602,14 +604,7 @@ class Fp8LinearMethod(LinearMethodBase):
604
bias=bias,
605
)
606
605
- return self.fp8_linear.apply(
606
- input=x,
607
- weight=layer.weight,
608
- weight_scale=layer.weight_scale,
609
- out_dtype=self.out_dtype,
610
- input_scale=layer.input_scale,
611
- bias=bias,
612
- )
607
+ return self.fp8_linear.apply_weights(layer, x, bias)
608
609
610
class Fp8MoEMethod(FusedMoEMethodBase):
vllm/model_executor/layers/quantization/kernels/scaled_mm/ScaledMMLinearKernel.py
+144
-33
@@ -2,19 +2,58 @@
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
from abc import ABC, abstractmethod
5
+from collections.abc import Sequence
6
from dataclasses import dataclass
7
+from typing import Generic, TypeVar
8
9
import torch
10
11
+from vllm.model_executor.layers.quantization.input_quant_fp8 import QuantFP8
12
+from vllm.model_executor.layers.quantization.utils.quant_utils import (
13
+ QuantKey,
14
+)
15
+from vllm.platforms import current_platform
16
+
17
18
@dataclass
19
class ScaledMMLinearLayerConfig:
12
- is_channelwise: bool
20
+ pass
21
+
22
+
23
+@dataclass
24
+class Int8ScaledMMLinearLayerConfig(ScaledMMLinearLayerConfig):
25
+ # TODO: Chnage to QuantKey like FP8ScaledMMLinearLayerConfig
26
is_static_input_scheme: bool
27
+ is_channelwise: bool
28
input_symmetric: bool
29
30
17
-class ScaledMMLinearKernel(ABC):
31
+@dataclass
32
+class FP8ScaledMMLinearLayerConfig(ScaledMMLinearLayerConfig):
33
+ weight_quant_key: QuantKey
34
+ activation_quant_key: QuantKey
35
+ out_dtype: torch.dtype | None
36
+
37
+
38
+_FP8ParamsT = tuple[
39
+ torch.Tensor, # weight
40
+ torch.Tensor, # weight_scale
41
+ torch.Tensor | None, # input_scale,
42
+ torch.Tensor | None, # input_scale_ub,
43
+]
44
+_Int8ParamsT = tuple[
45
+ torch.Tensor, # weight
46
+ torch.Tensor, # weight_scale
47
+ torch.Tensor | None, # input_scale,
48
+ torch.Tensor | None, # input_zp
49
+ torch.Tensor | None, # azp_adj
50
+]
51
+
52
+_ParamsT = TypeVar("_ParamsT", _Int8ParamsT, _FP8ParamsT)
53
+_ConfigT = TypeVar("_ConfigT", bound=ScaledMMLinearLayerConfig)
54
+
55
+
56
+class ScaledMMLinearKernel(Generic[_ConfigT, _ParamsT], ABC):
57
@classmethod
58
@abstractmethod
59
def is_supported(
@@ -24,26 +63,14 @@ class ScaledMMLinearKernel(ABC):
63
64
@classmethod
65
@abstractmethod
27
- def can_implement(cls, c: ScaledMMLinearLayerConfig) -> tuple[bool, str | None]:
66
+ def can_implement(cls, c: _ConfigT) -> tuple[bool, str | None]:
67
raise NotImplementedError
68
30
- def __init__(
31
- self,
32
- c: ScaledMMLinearLayerConfig,
33
- w_q_param_name: str,
34
- w_s_param_name: str,
35
- i_s_param_name: str,
36
- i_zp_param_name: str,
37
- azp_adj_param_name: str,
38
- ) -> None:
39
- assert self.can_implement(c)
40
- assert self.is_supported()
69
+ def __init__(self, c: _ConfigT, layer_param_names: Sequence[str]) -> None:
70
+ assert self.can_implement(c)[0]
71
+ assert self.is_supported()[0]
72
self.config = c
42
- self.w_q_name = w_q_param_name
43
- self.w_s_name = w_s_param_name
44
- self.i_s_name = i_s_param_name
45
- self.i_zp_name = i_zp_param_name
46
- self.azp_adj_name = azp_adj_param_name
73
+ self.layer_param_names = layer_param_names
74
75
@abstractmethod
76
def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
@@ -58,19 +85,103 @@ class ScaledMMLinearKernel(ABC):
85
) -> torch.Tensor:
86
raise NotImplementedError
87
61
- def _get_weight_params(
62
- self, layer: torch.nn.Module
63
- ) -> tuple[
64
- torch.Tensor, # weight
65
- torch.Tensor, # weight_scale
66
- torch.Tensor | None, # input_scale,
67
- torch.Tensor | None, # input_zp
68
- torch.Tensor | None, # azp_adj
69
- ]:
88
+ # return a covariant type in the subclass
89
+ @abstractmethod
90
+ def _get_layer_params(self, layer) -> _ParamsT:
91
+ raise NotImplementedError
92
+
93
+
94
+class FP8ScaledMMLinearKernel(
95
+ ScaledMMLinearKernel[FP8ScaledMMLinearLayerConfig, _FP8ParamsT], ABC
96
+):
97
+ def __init__(
98
+ self, c: FP8ScaledMMLinearLayerConfig, layer_param_names: Sequence[str]
99
+ ) -> None:
100
+ act_scale_descriptor = c.activation_quant_key.scale
101
+ self.quant_fp8 = QuantFP8(
102
+ static=act_scale_descriptor.static,
103
+ group_shape=act_scale_descriptor.group_shape,
104
+ num_token_padding=self.get_output_padding(),
105
+ )
106
+ self.fp8_dtype = current_platform.fp8_dtype()
107
+ super().__init__(c, layer_param_names)
108
+
109
+ def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
110
+ pass
111
+
112
+ def _get_layer_params(self, layer) -> _FP8ParamsT:
113
+ w, w_s, x_s, x_s_ub = self.layer_param_names
114
+ return (
115
+ getattr(layer, w),
116
+ getattr(layer, w_s),
117
+ getattr(layer, x_s, None),
118
+ getattr(layer, x_s_ub, None),
119
+ )
120
+
121
+ def apply_weights(
122
+ self,
123
+ layer: torch.nn.Module,
124
+ x: torch.Tensor,
125
+ bias: torch.Tensor | None = None,
126
+ ) -> torch.Tensor:
127
+ fp8_dtype = self.fp8_dtype
128
+ maybe_out_dtype = self.config.out_dtype
129
+ w, w_s, x_s, x_s_ub = self._get_layer_params(layer)
130
+
131
+ # ops.scaled_fp8_quant supports both dynamic and static quant.
132
+ # If dynamic, layer.input_scale is None and x_s computed from x.
133
+ # If static, layer.input_scale is scalar and x_s is input_scale.
134
+ # View input as 2D matrix for fp8 methods
135
+ x_2d = x.view(-1, x.shape[-1])
136
+ output_shape = [*x.shape[:-1], w.shape[1]]
137
+ out_dtype = x.dtype if maybe_out_dtype is None else maybe_out_dtype
138
+
139
+ # If input not quantized
140
+ # TODO(luka) remove this path if not used anymore
141
+ x_2d_q = x_2d
142
+ if x.dtype != fp8_dtype:
143
+ x_2d_q, x_s = self.quant_fp8(
144
+ x_2d,
145
+ x_s,
146
+ x_s_ub,
147
+ )
148
+ return self.apply_scaled_mm(
149
+ A=x_2d_q,
150
+ B=w,
151
+ out_dtype=out_dtype,
152
+ As=x_s,
153
+ Bs=w_s,
154
+ bias=bias,
155
+ output_shape=output_shape,
156
+ )
157
+
158
+ @abstractmethod
159
+ def apply_scaled_mm(
160
+ self,
161
+ *,
162
+ A: torch.Tensor,
163
+ B: torch.Tensor,
164
+ out_dtype: torch.dtype,
165
+ As: torch.Tensor,
166
+ Bs: torch.Tensor,
167
+ bias: torch.Tensor | None,
168
+ output_shape: list,
169
+ ) -> torch.Tensor:
170
+ raise NotImplementedError
171
+
172
+ def get_output_padding(self) -> int | None:
173
+ return None
174
+
175
+
176
+class Int8ScaledMMLinearKernel(
177
+ ScaledMMLinearKernel[Int8ScaledMMLinearLayerConfig, _Int8ParamsT], ABC
178
+):
179
+ def _get_layer_params(self, layer) -> _Int8ParamsT:
180
+ w_q, w_s, i_s, i_zp, azp_adj = self.layer_param_names
181
return (
71
- getattr(layer, self.w_q_name),
72
- getattr(layer, self.w_s_name),
73
- getattr(layer, self.i_s_name),
74
- getattr(layer, self.i_zp_name),
75
- getattr(layer, self.azp_adj_name),
182
+ getattr(layer, w_q),
183
+ getattr(layer, w_s),
184
+ getattr(layer, i_s, None),
185
+ getattr(layer, i_zp, None),
186
+ getattr(layer, azp_adj, None),
187
)
vllm/model_executor/layers/quantization/kernels/scaled_mm/__init__.py
+183
-30
@@ -2,76 +2,229 @@
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
import os
5
+from typing import TypeVar
6
7
+import torch
8
+
9
+from vllm.logger import init_logger
10
from vllm.model_executor.layers.quantization.kernels.scaled_mm.aiter import (
7
- AiterScaledMMLinearKernel,
11
+ AiterInt8ScaledMMLinearKernel,
12
)
13
from vllm.model_executor.layers.quantization.kernels.scaled_mm.cpu import (
10
- CPUScaledMMLinearKernel,
14
+ CPUInt8ScaledMMLinearKernel,
15
)
16
from vllm.model_executor.layers.quantization.kernels.scaled_mm.cutlass import (
13
- CutlassScaledMMLinearKernel,
17
+ CutlassFP8ScaledMMLinearKernel,
18
+ CutlassInt8ScaledMMLinearKernel,
19
+)
20
+from vllm.model_executor.layers.quantization.kernels.scaled_mm.flashinfer import (
21
+ FlashInferFP8ScaledMMLinearKernel,
22
+)
23
+from vllm.model_executor.layers.quantization.kernels.scaled_mm.pytorch import (
24
+ ChannelWiseTorchFP8ScaledMMLinearKernel,
25
+ PerTensorTorchFP8ScaledMMLinearKernel,
26
+ RowWiseTorchFP8ScaledMMLinearKernel,
27
+)
28
+from vllm.model_executor.layers.quantization.kernels.scaled_mm.rocm import (
29
+ ROCmFP8ScaledMMLinearKernel,
30
)
31
from vllm.model_executor.layers.quantization.kernels.scaled_mm.ScaledMMLinearKernel import ( # noqa: E501
32
+ FP8ScaledMMLinearKernel,
33
+ FP8ScaledMMLinearLayerConfig,
34
+ Int8ScaledMMLinearKernel,
35
+ Int8ScaledMMLinearLayerConfig,
36
ScaledMMLinearKernel,
37
ScaledMMLinearLayerConfig,
38
)
39
from vllm.model_executor.layers.quantization.kernels.scaled_mm.triton import (
20
- TritonScaledMMLinearKernel,
40
+ TritonInt8ScaledMMLinearKernel,
41
)
42
+from vllm.model_executor.layers.quantization.utils.quant_utils import QuantKey
43
from vllm.platforms import PlatformEnum, current_platform
44
45
+logger = init_logger(__name__)
46
+
47
+# in priority/performance order (when available)
48
+_POSSIBLE_INT8_KERNELS: dict[PlatformEnum, list[type[Int8ScaledMMLinearKernel]]] = {
49
+ PlatformEnum.CPU: [CPUInt8ScaledMMLinearKernel],
50
+ PlatformEnum.CUDA: [
51
+ CutlassInt8ScaledMMLinearKernel,
52
+ TritonInt8ScaledMMLinearKernel,
53
+ ],
54
+ PlatformEnum.ROCM: [AiterInt8ScaledMMLinearKernel, TritonInt8ScaledMMLinearKernel],
55
+}
56
+
57
# in priority/performance order (when available)
25
-_POSSIBLE_KERNELS: dict[PlatformEnum, list[type[ScaledMMLinearKernel]]] = {
26
- PlatformEnum.CPU: [CPUScaledMMLinearKernel],
27
- PlatformEnum.CUDA: [CutlassScaledMMLinearKernel, TritonScaledMMLinearKernel],
28
- PlatformEnum.ROCM: [AiterScaledMMLinearKernel, TritonScaledMMLinearKernel],
58
+_POSSIBLE_FP8_KERNELS: dict[PlatformEnum, list[type[FP8ScaledMMLinearKernel]]] = {
59
+ PlatformEnum.CUDA: [
60
+ FlashInferFP8ScaledMMLinearKernel,
61
+ CutlassFP8ScaledMMLinearKernel,
62
+ PerTensorTorchFP8ScaledMMLinearKernel,
63
+ ChannelWiseTorchFP8ScaledMMLinearKernel,
64
+ ],
65
+ PlatformEnum.ROCM: [
66
+ ROCmFP8ScaledMMLinearKernel,
67
+ PerTensorTorchFP8ScaledMMLinearKernel,
68
+ RowWiseTorchFP8ScaledMMLinearKernel,
69
+ ChannelWiseTorchFP8ScaledMMLinearKernel,
70
+ ],
71
+ PlatformEnum.CPU: [
72
+ PerTensorTorchFP8ScaledMMLinearKernel,
73
+ ChannelWiseTorchFP8ScaledMMLinearKernel,
74
+ ],
75
}
76
77
+_KernelT = TypeVar("_KernelT", bound=ScaledMMLinearKernel)
78
+_KernelConfigT = TypeVar("_KernelConfigT", bound=ScaledMMLinearLayerConfig)
79
+
80
+
81
+def is_supported_and_can_implement_kernel(
82
+ kernel: type[_KernelT], config: _KernelConfigT, compute_capability: int | None
83
+) -> tuple[bool, str]:
84
+ # TODO: Fetch `VLLM_DISABLED_KERNELS` from vllm.envs instead.
85
+ if kernel.__name__ in os.environ.get("VLLM_DISABLED_KERNELS", "").split(","):
86
+ return False, f" {kernel.__name__} is disabled by environment variable"
87
+
88
+ if compute_capability is None:
89
+ _cc = current_platform.get_device_capability()
90
+ if _cc is not None:
91
+ compute_capability = _cc[0] * 10 + _cc[1]
92
+
93
+ is_supported, failure_reason = kernel.is_supported(compute_capability)
94
+ if not is_supported:
95
+ return False, f"{kernel.__name__} {failure_reason}."
96
+
97
+ can_implement, failure_reason = kernel.can_implement(config)
98
+ if not can_implement:
99
+ return (
100
+ False,
101
+ f"{kernel.__name__} {failure_reason}.",
102
+ )
103
+
104
+ return True, ""
105
+
106
107
def choose_scaled_mm_linear_kernel(
33
- config: ScaledMMLinearLayerConfig, compute_capability: int | None = None
34
-) -> type[ScaledMMLinearKernel]:
108
+ config: _KernelConfigT,
109
+ possible_kernels: dict[PlatformEnum, list[type[_KernelT]]],
110
+ compute_capability: int | None = None,
111
+ force_kernel: type[_KernelT] | None = None,
112
+) -> type[_KernelT]:
113
"""
36
- Choose an ScaledMMLinearKernel that can implement the given config for the
114
+ Choose a _KernelT that can implement the given config for the
115
given compute capability. Attempts to choose the best kernel in terms of
116
performance.
117
118
Args:
41
- config (ScaledMMLinearLayerConfig): Description of the linear layer
119
+ config (_KernelConfigT): Description of the linear layer
120
to be implemented.
121
+ possible_kernels (dict[PlatformEnum, list[_KernelT]]): A
122
+ dictionary of platforms and their list list of possible kernels.
123
compute_capability (Optional[int], optional): The compute capability of
124
the target device, if None uses `current_platform` to get the
125
compute capability. Defaults to None.
126
+ force_kernel (Optional[type[_KernelT]]): An Optional forced kernel to override
127
+ the possible_kernels if it can be implemented. If None, it will only try the
128
+ possible kernels.
129
130
Raises:
131
ValueError: If no kernel can implement the given config.
132
133
Returns:
51
- type[ScaledMMLinearKernel]: Chosen kernel.
134
+ _KernelT: Chosen kernel.
135
"""
136
54
- failure_reasons = []
55
- for kernel in _POSSIBLE_KERNELS[current_platform._enum]:
56
- if kernel.__name__ in os.environ.get("VLLM_DISABLED_KERNELS", "").split(","):
57
- failure_reasons.append(f"{kernel.__name__}: disabled by env var")
58
- continue
137
+ failure_reason_list = []
138
60
- # If the current platform uses compute_capability,
61
- # make sure the kernel supports the compute capability.
62
- is_supported, reason = kernel.is_supported(compute_capability)
63
- if not is_supported:
64
- failure_reasons.append(f"{kernel.__name__}: {reason}")
65
- continue
139
+ if force_kernel is not None:
140
+ can_implement, failure_reason = is_supported_and_can_implement_kernel(
141
+ force_kernel, config, compute_capability
142
+ )
143
+ if can_implement:
144
+ return force_kernel
145
67
- can_implement, reason = kernel.can_implement(config)
68
- if not can_implement:
69
- failure_reasons.append(f"{kernel.__name__}: {reason}")
70
- continue
146
+ logger.info_once(
147
+ "Tried to force %s, but the kernel couldn't be implemented",
148
+ force_kernel.__name__,
149
+ scope="global",
150
+ )
151
72
- return kernel
152
+ for kernel in possible_kernels[current_platform._enum]:
153
+ is_supported_and_can_implement, failure_reason = (
154
+ is_supported_and_can_implement_kernel(kernel, config, compute_capability)
155
+ )
156
+ if is_supported_and_can_implement:
157
+ return kernel
158
+ failure_reason_list.append(failure_reason)
159
160
raise ValueError(
161
"Failed to find a kernel that can implement the "
76
- "ScaledMM linear layer. Reasons: \n" + "\n".join(failure_reasons)
162
+ "ScaledMM linear layer. Reasons: \n" + "\n".join(failure_reason_list)
163
+ )
164
+
165
+
166
+def init_fp8_linear_kernel(
167
+ activation_quant_key: QuantKey,
168
+ weight_quant_key: QuantKey,
169
+ out_dtype: torch.dtype,
170
+ force_kernel: type[FP8ScaledMMLinearKernel] | None = None,
171
+ module_name: str | None = None,
172
+) -> FP8ScaledMMLinearKernel:
173
+ scaled_mm_linear_kernel_config = FP8ScaledMMLinearLayerConfig(
174
+ weight_quant_key=weight_quant_key,
175
+ activation_quant_key=activation_quant_key,
176
+ out_dtype=out_dtype,
177
+ )
178
+
179
+ kernel_type = choose_scaled_mm_linear_kernel(
180
+ scaled_mm_linear_kernel_config, _POSSIBLE_FP8_KERNELS, force_kernel=force_kernel
181
+ )
182
+
183
+ if module_name:
184
+ logger.info_once(
185
+ "Selected %s for %s",
186
+ kernel_type.__name__,
187
+ module_name,
188
+ scope="global",
189
+ )
190
+
191
+ return kernel_type(
192
+ scaled_mm_linear_kernel_config,
193
+ layer_param_names=["weight", "weight_scale", "input_scale", "input_scale_ub"],
194
+ )
195
+
196
+
197
+def init_int8_linear_kernel(
198
+ is_channelwise: bool,
199
+ is_static_input_scheme: bool,
200
+ input_symmetric: bool,
201
+ module_name: str,
202
+) -> Int8ScaledMMLinearKernel:
203
+ config = Int8ScaledMMLinearLayerConfig(
204
+ is_channelwise=is_channelwise,
205
+ is_static_input_scheme=is_static_input_scheme,
206
+ input_symmetric=input_symmetric,
207
+ )
208
+
209
+ kernel_type = choose_scaled_mm_linear_kernel(
210
+ config,
211
+ _POSSIBLE_INT8_KERNELS,
212
+ )
213
+
214
+ logger.info_once(
215
+ "Selected %s for %s",
216
+ kernel_type.__name__,
217
+ module_name,
218
+ scope="global",
219
+ )
220
+
221
+ return kernel_type(
222
+ config,
223
+ layer_param_names=[
224
+ "weight",
225
+ "weight_scale",
226
+ "input_scale",
227
+ "input_zero_point",
228
+ "azp_adj",
229
+ ],
230
)
vllm/model_executor/layers/quantization/kernels/scaled_mm/aiter.py
+17
-36
@@ -8,60 +8,41 @@ from vllm import _custom_ops as ops
8
from vllm._aiter_ops import rocm_aiter_ops
9
from vllm.platforms import current_platform
10
11
-from .cutlass import CutlassScaledMMLinearKernel
12
-from .ScaledMMLinearKernel import ScaledMMLinearLayerConfig
11
+from .cutlass import CutlassInt8ScaledMMLinearKernel
12
+from .ScaledMMLinearKernel import Int8ScaledMMLinearLayerConfig
13
14
15
-class AiterScaledMMLinearKernel(CutlassScaledMMLinearKernel):
15
+class AiterInt8ScaledMMLinearKernel(CutlassInt8ScaledMMLinearKernel):
16
@classmethod
17
def is_supported(
18
cls, compute_capability: int | None = None
19
) -> tuple[bool, str | None]:
20
if not current_platform.is_rocm():
21
- return (
22
- False,
23
- "AiterScaledMMLinearKernel requires `aiter` which is not "
24
- + "currently supported on non-ROCm platform.",
25
- )
26
- if compute_capability is None:
27
- _cc = current_platform.get_device_capability()
28
- if _cc is not None:
29
- compute_capability = _cc.major * 10 + _cc.minor
21
+ return False, "Requires ROCm."
22
+
23
if compute_capability is not None and compute_capability < 90:
31
- return False, f"requires capability 90, got {compute_capability}"
24
+ return False, "requires compute capability 90 and above."
25
26
try:
27
import aiter # noqa: F401 # deliberately attempt to import aiter
28
except Exception:
36
- return (
37
- False,
38
- "AiterScaledMMLinearKernel requires `aiter` which is not "
39
- + "installed on ROCm.",
40
- )
29
+ return False, "requires `aiter` to be installed."
30
31
if not rocm_aiter_ops.is_linear_enabled():
32
return (
33
False,
45
- "AiterScaledMMLinearKernel is disabled. "
46
- + "Enable by setting `VLLM_ROCM_USE_AITER=1` "
34
+ "requires setting `VLLM_ROCM_USE_AITER=1` "
35
+ "and `VLLM_ROCM_USE_AITER_LINEAR=1`. "
36
+ "`VLLM_ROCM_USE_AITER_LINEAR` default is True.",
37
)
50
-
38
return True, None
39
40
@classmethod
54
- def can_implement(cls, c: ScaledMMLinearLayerConfig) -> tuple[bool, str | None]:
41
+ def can_implement(cls, c: Int8ScaledMMLinearLayerConfig) -> tuple[bool, str | None]:
42
if not c.input_symmetric:
56
- return (
57
- False,
58
- "AiterScaledMMLinearKernel only supports symmetric " + "quantization.",
59
- )
43
+ return False, "supports symmetric quantization only."
44
return True, None
45
62
- def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
63
- super().process_weights_after_loading(layer)
64
-
46
def apply_weights(
47
self,
48
layer: torch.nn.Module,
@@ -69,28 +50,28 @@ class AiterScaledMMLinearKernel(CutlassScaledMMLinearKernel):
50
bias: torch.Tensor | None = None,
51
) -> torch.Tensor:
52
"""
72
- `AiterScaledMMLinearKernel` implements a fused version of
53
+ `AiterInt8ScaledMMLinearKernel` implements a fused version of
54
`output = torch.mm((scale_a * a), (scale_b * b)).to(out_dtype)`
55
where scale_a * a and scale_b * b are implemented using numpy-style
56
broadcasting.
57
Currently only support per-tensor-per-tensor GEMM
58
and per-token-per-channel GEMM through AITER
78
- w8a8 scaled gemm. `AiterScaledMMLinearKernel` also does not support
59
+ w8a8 scaled gemm. `AiterInt8ScaledMMLinearKernel` also does not support
60
ATIER block scaled GEMM and mix-precision GEMM.
61
"""
81
- w_q, w_s, i_s, i_zp, azp_adj = self._get_weight_params(layer)
62
+ w_q, w_s, i_s, i_zp, azp_adj = self._get_layer_params(layer)
63
64
# ops.scaled_int8_quant supports both dynamic and static quant:
65
# * dynamic, i_s is None and x_s computed from x.
66
# * static, i_s is scalar and x_s is i_s.
67
symmetric = azp_adj is None
68
assert symmetric, (
88
- "AiterScaledMMLinearKernel only supports symmetric quantization."
69
+ "AiterInt8ScaledMMLinearKernel only supports symmetric quantization."
70
)
71
x_q, x_s, x_zp = ops.scaled_int8_quant(x, i_s, i_zp, symmetric=symmetric)
72
73
assert x_zp is None, (
93
- "AiterScaledMMLinearKernel only supports symmetric quantization."
74
+ "AiterInt8ScaledMMLinearKernel only supports symmetric quantization."
75
)
76
out_dtype = x.dtype
77
@@ -117,12 +98,12 @@ class AiterScaledMMLinearKernel(CutlassScaledMMLinearKernel):
98
), (
99
"Currently only support per-tensor-per-tensor GEMM "
100
+ " and per-token-per-channel GEMM through AITER"
120
- " w8a8 scaled gemm. `AiterScaledMMLinearKernel` "
101
+ " w8a8 scaled gemm. `AiterInt8ScaledMMLinearKernel` "
102
+ "does not support AITER block scaled GEMM."
103
)
104
105
# gemm_a8w8_CK(a, b, scale_a, scale_b, bias) expects
106
# a to be [M, K]
107
# b to be [N, K]
127
- # CutlassScaledMMLinearKernel prepare weight `w_q` in [K, N] format
108
+ # CutlassInt8ScaledMMLinearKernel prepare weight `w_q` in [K, N] format
109
return rocm_aiter_ops.gemm_a8w8(x_q, w_q.t(), x_s, w_s, bias, out_dtype)
vllm/model_executor/layers/quantization/kernels/scaled_mm/cpu.py
+33
-38
@@ -14,24 +14,28 @@ from vllm.model_executor.layers.utils import check_cpu_sgl_kernel
14
from vllm.platforms import current_platform
15
from vllm.platforms.interface import CpuArchEnum
16
17
-from .ScaledMMLinearKernel import ScaledMMLinearKernel, ScaledMMLinearLayerConfig
17
+from .ScaledMMLinearKernel import (
18
+ Int8ScaledMMLinearKernel,
19
+ Int8ScaledMMLinearLayerConfig,
20
+)
21
22
20
-class CPUScaledMMLinearKernel(ScaledMMLinearKernel):
23
+class CPUInt8ScaledMMLinearKernel(Int8ScaledMMLinearKernel):
24
@classmethod
25
def is_supported(
26
cls, compute_capability: int | None = None
27
) -> tuple[bool, str | None]:
28
if not current_platform.is_cpu():
26
- return False, "Requires CPU."
29
+ return False, "requires CPU."
30
return True, None
31
32
@classmethod
30
- def can_implement(cls, c: ScaledMMLinearLayerConfig) -> tuple[bool, str | None]:
33
+ def can_implement(cls, c: Int8ScaledMMLinearLayerConfig) -> tuple[bool, str | None]:
34
return True, None
35
36
def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
34
- weight = getattr(layer, self.w_q_name)
37
+ w_q_name, _, _, _, _ = self.layer_param_names
38
+ weight = getattr(layer, w_q_name)
39
dtype = weight.dtype
40
N, K = weight.size()
41
if (
@@ -49,10 +53,11 @@ class CPUScaledMMLinearKernel(ScaledMMLinearKernel):
53
def process_weights_for_onednn(self, layer: torch.nn.Module) -> None:
54
# WEIGHT
55
# Transpose to [K, N] for convenience
52
- weight = getattr(layer, self.w_q_name)
56
+ w_q_name, w_s_name, i_s_name, i_zp_name, azp_adj_name = self.layer_param_names
57
+ weight = getattr(layer, w_q_name)
58
replace_parameter(
59
layer,
55
- self.w_q_name,
60
+ w_q_name,
61
torch.nn.Parameter(weight.t().data, requires_grad=False),
62
)
63
@@ -61,28 +66,27 @@ class CPUScaledMMLinearKernel(ScaledMMLinearKernel):
66
# If we have a fused module (QKV, MLP) with per tensor scales (thus N
67
# scales being passed to the kernel), convert to the per-channel case.
68
is_fused_module = len(layer.logical_widths) > 1
64
- weight_scale = getattr(layer, self.w_s_name)
69
+ weight_scale = getattr(layer, w_s_name)
70
if is_fused_module and not self.config.is_channelwise:
71
weight_scale = convert_to_channelwise(weight_scale, layer.logical_widths)
72
replace_parameter(
73
layer,
69
- self.w_s_name,
74
+ w_s_name,
75
torch.nn.Parameter(weight_scale.data, requires_grad=False),
76
)
77
78
# INPUT SCALE
79
if self.config.is_static_input_scheme:
75
- input_scale = getattr(layer, self.i_s_name)
80
+ input_scale = getattr(layer, i_s_name)
81
82
if self.config.input_symmetric:
83
replace_parameter(
84
layer,
80
- self.i_s_name,
85
+ i_s_name,
86
torch.nn.Parameter(input_scale.max(), requires_grad=False),
87
)
83
- setattr(layer, self.i_zp_name, None)
88
else:
85
- input_zero_point = getattr(layer, self.i_zp_name)
89
+ input_zero_point = getattr(layer, i_zp_name)
90
91
# reconstruct the ranges
92
int8_traits = torch.iinfo(torch.int8)
@@ -92,20 +96,16 @@ class CPUScaledMMLinearKernel(ScaledMMLinearKernel):
96
97
scale = (range_max - range_min) / (int8_traits.max - int8_traits.min)
98
replace_parameter(
95
- layer, self.i_s_name, torch.nn.Parameter(scale, requires_grad=False)
99
+ layer, i_s_name, torch.nn.Parameter(scale, requires_grad=False)
100
)
101
102
azp = (
103
(int8_traits.min - range_min / scale).round().to(dtype=torch.int32)
104
)
105
replace_parameter(
102
- layer, self.i_zp_name, torch.nn.Parameter(azp, requires_grad=False)
106
+ layer, i_zp_name, torch.nn.Parameter(azp, requires_grad=False)
107
)
108
105
- else:
106
- setattr(layer, self.i_s_name, None)
107
- setattr(layer, self.i_zp_name, None)
108
-
109
# Different from cutlass, oneDNN kernels only need the AZP adjustment
110
# term for dynamic quantization. And s_b should be folded into the
111
# term. Such as:
@@ -113,38 +113,37 @@ class CPUScaledMMLinearKernel(ScaledMMLinearKernel):
113
# s_a * (s_b * AB) - s_a * s_b * zp_a * B + bias =
114
# s_a * GEMM_output - s_a * zp_a * adj + bias
115
if not (self.config.input_symmetric and self.config.is_static_input_scheme):
116
- weight = getattr(layer, self.w_q_name)
117
- weight_scale = getattr(layer, self.w_s_name)
116
+ weight = getattr(layer, w_q_name)
117
+ weight_scale = getattr(layer, w_s_name)
118
azp_adj = weight.sum(dim=0, keepdim=True, dtype=torch.float32)
119
azp_adj = azp_adj * weight_scale.squeeze()
120
setattr(
121
layer,
122
- self.azp_adj_name,
122
+ azp_adj_name,
123
torch.nn.Parameter(azp_adj, requires_grad=False),
124
)
125
- else:
126
- setattr(layer, self.azp_adj_name, None)
125
128
- weight = getattr(layer, self.w_q_name)
126
+ weight = getattr(layer, w_q_name)
127
self.dnnl_handler = ops.create_onednn_scaled_mm(
128
weight,
131
- getattr(layer, self.w_s_name),
129
+ getattr(layer, w_s_name),
130
torch.get_default_dtype(),
133
- getattr(layer, self.i_s_name) is None,
131
+ getattr(layer, i_s_name) is None,
132
not self.config.input_symmetric,
133
32,
134
)
135
# weight is prepacked and maintained by the dnnl_handler,
136
# release the original weight
139
- setattr(layer, self.w_q_name, None)
137
+ setattr(layer, w_q_name, None)
138
del weight
139
140
def process_weights_for_sgl(self, layer: torch.nn.Module) -> None:
141
+ w_q_name, w_s_name, _, _, _ = self.layer_param_names
142
# WEIGHT
144
- weight = getattr(layer, self.w_q_name)
143
+ weight = getattr(layer, w_q_name)
144
packed_weight = torch.ops._C.convert_weight_packed(weight)
145
replace_parameter(
147
- layer, self.w_q_name, torch.nn.Parameter(packed_weight, requires_grad=False)
146
+ layer, w_q_name, torch.nn.Parameter(packed_weight, requires_grad=False)
147
)
148
149
if layer.bias is not None:
@@ -156,19 +155,15 @@ class CPUScaledMMLinearKernel(ScaledMMLinearKernel):
155
# WEIGHT SCALE
156
# CPU SGL kernels only support per-channel.
157
# For per-tensor quant, convert to the per-channel case.
159
- weight_scale = getattr(layer, self.w_s_name)
158
+ weight_scale = getattr(layer, w_s_name)
159
if not self.config.is_channelwise:
160
weight_scale = convert_to_channelwise(weight_scale, layer.logical_widths)
161
replace_parameter(
162
layer,
164
- self.w_s_name,
163
+ w_s_name,
164
torch.nn.Parameter(weight_scale.data, requires_grad=False),
165
)
166
168
- setattr(layer, self.i_s_name, None)
169
- setattr(layer, self.i_zp_name, None)
170
- setattr(layer, self.azp_adj_name, None)
171
-
167
def apply_weights(
168
self,
169
layer: torch.nn.Module,
@@ -187,7 +182,7 @@ class CPUScaledMMLinearKernel(ScaledMMLinearKernel):
182
x: torch.Tensor,
183
bias: torch.Tensor | None = None,
184
) -> torch.Tensor:
190
- w_q, w_s, i_s, i_zp, azp_adj = self._get_weight_params(layer)
185
+ w_q, w_s, i_s, i_zp, azp_adj = self._get_layer_params(layer)
186
187
# ops.scaled_int8_quant supports both dynamic and static quant:
188
# * dynamic, i_s is None and x_s computed from x.
@@ -209,7 +204,7 @@ class CPUScaledMMLinearKernel(ScaledMMLinearKernel):
204
x: torch.Tensor,
205
bias: torch.Tensor | None = None,
206
) -> torch.Tensor:
212
- w_q, w_s, _, _, _ = self._get_weight_params(layer)
207
+ w_q, w_s, _, _, _ = self._get_layer_params(layer)
208
return torch.ops._C.int8_scaled_mm_with_quant(
209
x,
210
w_q,
vllm/model_executor/layers/quantization/kernels/scaled_mm/cutlass.py
+61
-35
@@ -11,35 +11,36 @@ from vllm.model_executor.layers.quantization.utils.w8a8_utils import (
11
)
12
from vllm.platforms import current_platform
13
14
-from .ScaledMMLinearKernel import ScaledMMLinearKernel, ScaledMMLinearLayerConfig
14
+from .ScaledMMLinearKernel import (
15
+ FP8ScaledMMLinearKernel,
16
+ FP8ScaledMMLinearLayerConfig,
17
+ Int8ScaledMMLinearKernel,
18
+ Int8ScaledMMLinearLayerConfig,
19
+)
20
21
17
-class CutlassScaledMMLinearKernel(ScaledMMLinearKernel):
22
+class CutlassInt8ScaledMMLinearKernel(Int8ScaledMMLinearKernel):
23
@classmethod
24
def is_supported(
25
cls, compute_capability: int | None = None
26
) -> tuple[bool, str | None]:
27
if not current_platform.is_cuda():
23
- return False, "Requires CUDA."
24
- if compute_capability is None:
25
- _cc = current_platform.get_device_capability()
26
- if _cc is not None:
27
- compute_capability = _cc.major * 10 + _cc.minor
28
- if compute_capability is not None and compute_capability < 75:
29
- return False, f"requires capability 75, got {compute_capability}"
28
+ return False, "requires CUDA."
29
return True, None
30
31
@classmethod
33
- def can_implement(cls, c: ScaledMMLinearLayerConfig) -> tuple[bool, str | None]:
32
+ def can_implement(cls, c: Int8ScaledMMLinearLayerConfig) -> tuple[bool, str | None]:
33
return True, None
34
35
def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
36
+ w_q_name, w_s_name, i_s_name, i_zp_name, azp_adj_name = self.layer_param_names
37
+ config = self.config
38
# WEIGHT
39
# Cutlass kernels need transposed weight.
39
- weight = getattr(layer, self.w_q_name)
40
+ weight = getattr(layer, w_q_name)
41
replace_parameter(
42
layer,
42
- self.w_q_name,
43
+ w_q_name,
44
torch.nn.Parameter(weight.t().data, requires_grad=False),
45
)
46
@@ -48,28 +49,28 @@ class CutlassScaledMMLinearKernel(ScaledMMLinearKernel):
49
# If we have a fused module (QKV, MLP) with per tensor scales (thus N
50
# scales being passed to the kernel), convert to the per-channel case.
51
is_fused_module = len(layer.logical_widths) > 1
51
- weight_scale = getattr(layer, self.w_s_name)
52
- if is_fused_module and not self.config.is_channelwise:
52
+ weight_scale = getattr(layer, w_s_name)
53
+ if is_fused_module and not config.is_channelwise:
54
weight_scale = convert_to_channelwise(weight_scale, layer.logical_widths)
55
replace_parameter(
56
layer,
56
- self.w_s_name,
57
+ w_s_name,
58
torch.nn.Parameter(weight_scale.data, requires_grad=False),
59
)
60
61
# INPUT SCALE
61
- if self.config.is_static_input_scheme:
62
- input_scale = getattr(layer, self.i_s_name)
62
+ if config.is_static_input_scheme:
63
+ input_scale = getattr(layer, i_s_name)
64
64
- if self.config.input_symmetric:
65
+ if config.input_symmetric:
66
replace_parameter(
67
layer,
67
- self.i_s_name,
68
+ i_s_name,
69
torch.nn.Parameter(input_scale.max(), requires_grad=False),
70
)
70
- setattr(layer, self.i_zp_name, None)
71
+ setattr(layer, i_zp_name, None)
72
else:
72
- input_zero_point = getattr(layer, self.i_zp_name)
73
+ input_zero_point = getattr(layer, i_zp_name)
74
75
# reconstruct the ranges
76
int8_traits = torch.iinfo(torch.int8)
@@ -79,38 +80,32 @@ class CutlassScaledMMLinearKernel(ScaledMMLinearKernel):
80
81
scale = (range_max - range_min) / (int8_traits.max - int8_traits.min)
82
replace_parameter(
82
- layer, self.i_s_name, torch.nn.Parameter(scale, requires_grad=False)
83
+ layer, i_s_name, torch.nn.Parameter(scale, requires_grad=False)
84
)
85
86
# AZP loaded as int8 but used as int32
87
azp = (int8_traits.min - range_min / scale).to(dtype=torch.int32)
88
replace_parameter(
88
- layer, self.i_zp_name, torch.nn.Parameter(azp, requires_grad=False)
89
+ layer, i_zp_name, torch.nn.Parameter(azp, requires_grad=False)
90
)
91
91
- else:
92
- setattr(layer, self.i_s_name, None)
93
- setattr(layer, self.i_zp_name, None)
94
-
92
# azp_adj is the AZP adjustment term, used to account for weights.
93
# It does not depend on scales or azp, so it is the same for
94
# static and dynamic quantization.
95
# For more details, see csrc/quantization/w8a8/cutlass/Epilogues.md
96
# https://github.com/vllm-project/vllm/blob/main/csrc/quantization/w8a8/cutlass/Epilogues.md
100
- if not self.config.input_symmetric:
101
- weight = getattr(layer, self.w_q_name)
97
+ if not config.input_symmetric:
98
+ weight = getattr(layer, w_q_name)
99
azp_adj = weight.sum(dim=0, keepdim=True, dtype=torch.int32)
103
- if self.config.is_static_input_scheme:
100
+ if config.is_static_input_scheme:
101
# cutlass_w8a8 requires azp to be folded into azp_adj
102
# in the per-tensor case
106
- azp_adj = getattr(layer, self.i_zp_name) * azp_adj
103
+ azp_adj = getattr(layer, i_zp_name) * azp_adj
104
setattr(
105
layer,
109
- self.azp_adj_name,
106
+ azp_adj_name,
107
torch.nn.Parameter(azp_adj, requires_grad=False),
108
)
112
- else:
113
- setattr(layer, self.azp_adj_name, None)
109
110
def apply_weights(
111
self,
@@ -118,7 +113,7 @@ class CutlassScaledMMLinearKernel(ScaledMMLinearKernel):
113
x: torch.Tensor,
114
bias: torch.Tensor | None = None,
115
) -> torch.Tensor:
121
- w_q, w_s, i_s, i_zp, azp_adj = self._get_weight_params(layer)
116
+ w_q, w_s, i_s, i_zp, azp_adj = self._get_layer_params(layer)
117
118
# ops.scaled_int8_quant supports both dynamic and static quant:
119
# * dynamic, i_s is None and x_s computed from x.
@@ -145,3 +140,34 @@ class CutlassScaledMMLinearKernel(ScaledMMLinearKernel):
140
return ops.cutlass_scaled_mm(
141
x_q, w_q, scale_a=x_s, scale_b=w_s, out_dtype=x.dtype, bias=bias
142
)
143
+
144
+
145
+class CutlassFP8ScaledMMLinearKernel(FP8ScaledMMLinearKernel):
146
+ @classmethod
147
+ def is_supported(
148
+ cls, compute_capability: int | None = None
149
+ ) -> tuple[bool, str | None]:
150
+ if not current_platform.is_cuda():
151
+ return False, "requires CUDA."
152
+ return True, None
153
+
154
+ @classmethod
155
+ def can_implement(cls, c: FP8ScaledMMLinearLayerConfig) -> tuple[bool, str | None]:
156
+ return True, None
157
+
158
+ def apply_scaled_mm(
159
+ self,
160
+ *,
161
+ A: torch.Tensor,
162
+ B: torch.Tensor,
163
+ out_dtype: torch.dtype,
164
+ As: torch.Tensor,
165
+ Bs: torch.Tensor,
166
+ bias: torch.Tensor | None,
167
+ output_shape: list,
168
+ ) -> torch.Tensor:
169
+ # Fused GEMM_DQ
170
+ output = ops.cutlass_scaled_mm(
171
+ A, B, out_dtype=out_dtype, scale_a=As, scale_b=Bs, bias=bias
172
+ )
173
+ return output.view(*output_shape)
vllm/model_executor/layers/quantization/kernels/scaled_mm/flashinfer.py
new
+57
@@ -0,0 +1,57 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+
5
+import torch
6
+
7
+from vllm.platforms import current_platform
8
+from vllm.utils.flashinfer import flashinfer_scaled_fp8_mm, has_flashinfer
9
+
10
+from .ScaledMMLinearKernel import (
11
+ FP8ScaledMMLinearKernel,
12
+ FP8ScaledMMLinearLayerConfig,
13
+)
14
+
15
+
16
+class FlashInferFP8ScaledMMLinearKernel(FP8ScaledMMLinearKernel):
17
+ @classmethod
18
+ def is_supported(
19
+ cls, compute_capability: int | None = None
20
+ ) -> tuple[bool, str | None]:
21
+ if not current_platform.is_cuda():
22
+ return False, "requires CUDA."
23
+
24
+ if not has_flashinfer():
25
+ return False, "requires FlashInfer to be installed."
26
+
27
+ if compute_capability is not None and compute_capability < 100:
28
+ return False, "requires compute capability 100 and above."
29
+
30
+ return True, None
31
+
32
+ @classmethod
33
+ def can_implement(cls, c: FP8ScaledMMLinearLayerConfig) -> tuple[bool, str | None]:
34
+ per_tensor_activation_scales = (
35
+ c.activation_quant_key.scale.group_shape.is_per_tensor()
36
+ )
37
+ per_tensor_weight_scales = c.weight_quant_key.scale.group_shape.is_per_tensor()
38
+
39
+ if not (per_tensor_activation_scales and per_tensor_weight_scales):
40
+ return False, "requires per tensor activation and weight scales."
41
+
42
+ return True, None
43
+
44
+ def apply_scaled_mm(
45
+ self,
46
+ *,
47
+ A: torch.Tensor,
48
+ B: torch.Tensor,
49
+ out_dtype: torch.dtype,
50
+ As: torch.Tensor,
51
+ Bs: torch.Tensor,
52
+ bias: torch.Tensor | None,
53
+ output_shape: list,
54
+ ) -> torch.Tensor:
55
+ return flashinfer_scaled_fp8_mm(
56
+ A, B, out_dtype=out_dtype, scale_a=As, scale_b=Bs, bias=bias
57
+ )
vllm/model_executor/layers/quantization/kernels/scaled_mm/pytorch.py
new
+221
@@ -0,0 +1,221 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+
5
+import torch
6
+from packaging import version
7
+
8
+from vllm.config import CompilationMode, get_current_vllm_config
9
+from vllm.platforms import current_platform
10
+
11
+from .ScaledMMLinearKernel import (
12
+ FP8ScaledMMLinearKernel,
13
+ FP8ScaledMMLinearLayerConfig,
14
+)
15
+
16
+
17
+class TorchFP8ScaledMMLinearKernel(FP8ScaledMMLinearKernel):
18
+ """
19
+ Base class for FP8 linear kernels using Torch.
20
+ Each subclass represents a kernel variant for
21
+ specific device capabilities and torch versions.
22
+ """
23
+
24
+ @classmethod
25
+ def is_supported(
26
+ cls, compute_capability: int | None = None
27
+ ) -> tuple[bool, str | None]:
28
+ if not (current_platform.is_cuda_alike() or current_platform.is_cpu()):
29
+ return False, "requires ROCm, CUDA or CPU."
30
+
31
+ if compute_capability is not None and compute_capability < 89:
32
+ return False, "requires compute capability 89 and above."
33
+
34
+ return True, None
35
+
36
+ def get_output_padding(self) -> int | None:
37
+ # Note: we pad the input because torch._scaled_mm is more performant
38
+ # for matrices with batch dimension > 16.
39
+ # This could change in the future.
40
+ # We also don't pad when using torch.compile,
41
+ # as it breaks with dynamic shapes.
42
+ #
43
+ # The perf gain is still relevant as of 16/1/2026
44
+ # torch version == 2.9.0. More details in the link below:
45
+ # https://github.com/vllm-project/vllm/issues/32269
46
+ vllm_config = get_current_vllm_config().compilation_config
47
+ pad_output = vllm_config.mode < CompilationMode.VLLM_COMPILE
48
+ return 17 if pad_output else None
49
+
50
+
51
+class PerTensorTorchFP8ScaledMMLinearKernel(TorchFP8ScaledMMLinearKernel):
52
+ @classmethod
53
+ def can_implement(cls, c: FP8ScaledMMLinearLayerConfig) -> tuple[bool, str | None]:
54
+ per_tensor_activation_scales = (
55
+ c.activation_quant_key.scale.group_shape.is_per_tensor()
56
+ )
57
+ per_tensor_weight_scales = c.weight_quant_key.scale.group_shape.is_per_tensor()
58
+
59
+ if not (per_tensor_activation_scales and per_tensor_weight_scales):
60
+ return False, "requires per tensor activation and weight scales."
61
+ return True, None
62
+
63
+ def apply_scaled_mm(
64
+ self,
65
+ *,
66
+ A: torch.Tensor,
67
+ B: torch.Tensor,
68
+ out_dtype: torch.dtype,
69
+ As: torch.Tensor,
70
+ Bs: torch.Tensor,
71
+ bias: torch.Tensor | None,
72
+ output_shape: list,
73
+ ) -> torch.Tensor:
74
+ output = torch._scaled_mm(
75
+ A, B, out_dtype=out_dtype, scale_a=As, scale_b=Bs, bias=bias
76
+ )
77
+ # A fix for discrepancy in scaled_mm which returns tuple
78
+ # for torch < 2.5 and a single value in torch >= 2.5
79
+ if type(output) is tuple and len(output) == 2:
80
+ output = output[0]
81
+
82
+ return torch.narrow(output, 0, 0, output_shape[0]).view(*output_shape)
83
+
84
+
85
+class RowWiseTorchFP8ScaledMMLinearKernel(TorchFP8ScaledMMLinearKernel):
86
+ @classmethod
87
+ def is_supported(
88
+ cls, compute_capability: int | None = None
89
+ ) -> tuple[bool, str | None]:
90
+ if not current_platform.is_rocm():
91
+ return False, "requires ROCm."
92
+
93
+ from vllm.platforms.rocm import on_mi3xx
94
+
95
+ if not on_mi3xx():
96
+ return False, "requires MI3xx."
97
+
98
+ if compute_capability is not None and compute_capability < 94:
99
+ return False, "requires compute capability 94 and above."
100
+
101
+ if not version.parse(torch.__version__) >= version.parse("2.7"):
102
+ return False, "requires pytorch version >=2.7."
103
+
104
+ return True, None
105
+
106
+ @classmethod
107
+ def can_implement(cls, c: FP8ScaledMMLinearLayerConfig) -> tuple[bool, str | None]:
108
+ per_tensor_activation_scales = (
109
+ c.activation_quant_key.scale.group_shape.is_per_tensor()
110
+ )
111
+ per_tensor_weight_scales = c.weight_quant_key.scale.group_shape.is_per_tensor()
112
+
113
+ if c.out_dtype == torch.float16:
114
+ # hipblaslt rowwise _scaled_mm only supports BFloat16
115
+ return False, "supports BFloat16 output data type only."
116
+
117
+ if per_tensor_activation_scales or per_tensor_weight_scales:
118
+ return False, "cannot be used with per tensor activation and weight scales."
119
+
120
+ return True, None
121
+
122
+ def apply_scaled_mm(
123
+ self,
124
+ *,
125
+ A: torch.Tensor,
126
+ B: torch.Tensor,
127
+ out_dtype: torch.dtype,
128
+ As: torch.Tensor,
129
+ Bs: torch.Tensor,
130
+ bias: torch.Tensor | None,
131
+ output_shape: list,
132
+ ) -> torch.Tensor:
133
+ # Note:
134
+ # For now it has only been validated on ROCm platform.
135
+ # fp8 rowwise scaling in torch._scaled_mm is introduced in
136
+ # https://github.com/pytorch/pytorch/pull/144432 using
137
+ # hipBLASLt and ROCm 6.3, which only exists in torch 2.7 and above.
138
+ #
139
+ # For CUDA platform please validate if the torch._scaled_mm supports
140
+ # rowwise scaled GEMM before using it
141
+
142
+ # Fused GEMM_DQ Rowwise GEMM
143
+ output = torch._scaled_mm(
144
+ A,
145
+ B,
146
+ out_dtype=out_dtype,
147
+ scale_a=As,
148
+ scale_b=Bs.t(),
149
+ bias=bias,
150
+ )
151
+
152
+ return torch.narrow(output, 0, 0, output_shape[0]).view(*output_shape)
153
+
154
+
155
+class ChannelWiseTorchFP8ScaledMMLinearKernel(TorchFP8ScaledMMLinearKernel):
156
+ @classmethod
157
+ def can_implement(cls, c: FP8ScaledMMLinearLayerConfig) -> tuple[bool, str | None]:
158
+ per_tensor_activation_scales = (
159
+ c.activation_quant_key.scale.group_shape.is_per_tensor()
160
+ )
161
+ per_tensor_weight_scales = c.weight_quant_key.scale.group_shape.is_per_tensor()
162
+
163
+ if per_tensor_activation_scales and per_tensor_weight_scales:
164
+ return False, "cannot be used with per tensor activation and weight scales."
165
+
166
+ return True, None
167
+
168
+ def apply_scaled_mm(
169
+ self,
170
+ *,
171
+ A: torch.Tensor,
172
+ B: torch.Tensor,
173
+ out_dtype: torch.dtype,
174
+ As: torch.Tensor,
175
+ Bs: torch.Tensor,
176
+ bias: torch.Tensor | None,
177
+ output_shape: list,
178
+ ) -> torch.Tensor:
179
+ # Use unfused DQ due to limitations with scaled_mm
180
+
181
+ # Symmetric quantized GEMM by definition computes the following:
182
+ # C = (s_x * X) (s_w * W) + bias
183
+ # This is equivalent to dequantizing the weights and activations
184
+ # before applying a GEMM.
185
+ #
186
+ # In order to compute quantized operands, a quantized kernel
187
+ # will rewrite the above like so:
188
+ # C = s_w * s_x * (X * W) + bias
189
+ #
190
+ # For the scaled_mm fallback case, we break this down, since it
191
+ # does not support s_w being a vector.
192
+
193
+ # Input scaling factors are no longer optional in _scaled_mm starting
194
+ # from pytorch 2.5. Allocating a dummy tensor to pass as scales
195
+ dummy_tensor = torch.ones(1, dtype=torch.float32, device=A.device)
196
+
197
+ # GEMM
198
+ # This computes C = (X * W).
199
+ # Output in fp32 to allow subsequent ops to happen in-place
200
+ output = torch._scaled_mm(
201
+ A,
202
+ B,
203
+ scale_a=dummy_tensor,
204
+ scale_b=dummy_tensor,
205
+ out_dtype=torch.float32,
206
+ )
207
+ # A fix for discrepancy in scaled_mm which returns tuple
208
+ # for torch < 2.5 and a single value in torch >= 2.5
209
+ if type(output) is tuple and len(output) == 2:
210
+ output = output[0]
211
+
212
+ # Unpad (undo num_token_padding)
213
+ output = torch.narrow(output, 0, 0, output_shape[0])
214
+ x_scale = torch.narrow(As, 0, 0, output_shape[0])
215
+
216
+ # DQ
217
+ # C = sw * sx * (X * W) + bias
218
+ output = output * x_scale * Bs.t()
219
+ if bias is not None:
220
+ output = output + bias
221
+ return output.to(out_dtype).view(*output_shape)
vllm/model_executor/layers/quantization/kernels/scaled_mm/rocm.py
new
+117
@@ -0,0 +1,117 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+
5
+import torch
6
+
7
+import vllm.envs as envs
8
+from vllm import _custom_ops as ops
9
+from vllm.platforms import current_platform
10
+from vllm.utils.platform_utils import get_cu_count
11
+from vllm.utils.torch_utils import direct_register_custom_op
12
+
13
+from .ScaledMMLinearKernel import (
14
+ FP8ScaledMMLinearKernel,
15
+ FP8ScaledMMLinearLayerConfig,
16
+)
17
+
18
+
19
+def rocm_per_tensor_float_w8a8_scaled_mm_impl(
20
+ A: torch.Tensor,
21
+ B: torch.Tensor,
22
+ out_dtype: torch.dtype,
23
+ As: torch.Tensor,
24
+ Bs: torch.Tensor,
25
+ bias: torch.Tensor,
26
+) -> torch.Tensor:
27
+ if (
28
+ A.shape[0] == 1
29
+ and B.shape[1] % 16 == 0
30
+ and ((bias is None) or (bias.dtype == out_dtype))
31
+ ):
32
+ output = ops.wvSplitKQ(
33
+ B.t(),
34
+ A,
35
+ out_dtype,
36
+ As,
37
+ Bs,
38
+ get_cu_count(),
39
+ bias,
40
+ )
41
+ # Fallback
42
+ else:
43
+ output = torch._scaled_mm(
44
+ A,
45
+ B,
46
+ out_dtype=out_dtype,
47
+ scale_a=As,
48
+ scale_b=Bs,
49
+ bias=bias,
50
+ )
51
+ return output
52
+
53
+
54
+def rocm_per_tensor_float_w8a8_scaled_mm_fake(
55
+ A: torch.Tensor,
56
+ B: torch.Tensor,
57
+ out_dtype: torch.dtype,
58
+ As: torch.Tensor,
59
+ Bs: torch.Tensor,
60
+ bias: torch.Tensor,
61
+) -> torch.Tensor:
62
+ return A.new_empty((*A.shape[:-1], B.shape[1]), dtype=out_dtype)
63
+
64
+
65
+if current_platform.is_rocm():
66
+ direct_register_custom_op(
67
+ op_name="rocm_per_tensor_float_w8a8_scaled_mm_impl",
68
+ op_func=rocm_per_tensor_float_w8a8_scaled_mm_impl,
69
+ fake_impl=rocm_per_tensor_float_w8a8_scaled_mm_fake,
70
+ )
71
+
72
+
73
+class ROCmFP8ScaledMMLinearKernel(FP8ScaledMMLinearKernel):
74
+ @classmethod
75
+ def is_supported(
76
+ cls, compute_capability: int | None = None
77
+ ) -> tuple[bool, str | None]:
78
+ if not current_platform.is_rocm():
79
+ return False, "requires ROCm."
80
+
81
+ from vllm.platforms.rocm import on_mi3xx
82
+
83
+ if not on_mi3xx():
84
+ return False, "requires MI3xx."
85
+
86
+ if not envs.VLLM_ROCM_USE_SKINNY_GEMM:
87
+ return False, "requires VLLM_ROCM_USE_SKINNY_GEMM to be enabled."
88
+
89
+ return True, None
90
+
91
+ @classmethod
92
+ def can_implement(cls, c: FP8ScaledMMLinearLayerConfig) -> tuple[bool, str | None]:
93
+ per_tensor_activation_scales = (
94
+ c.activation_quant_key.scale.group_shape.is_per_tensor()
95
+ )
96
+ per_tensor_weight_scales = c.weight_quant_key.scale.group_shape.is_per_tensor()
97
+
98
+ if not (per_tensor_activation_scales and per_tensor_weight_scales):
99
+ return False, "requires per tensor activation and weight scales."
100
+
101
+ return True, None
102
+
103
+ def apply_scaled_mm(
104
+ self,
105
+ *,
106
+ A: torch.Tensor,
107
+ B: torch.Tensor,
108
+ out_dtype: torch.dtype,
109
+ As: torch.Tensor,
110
+ Bs: torch.Tensor,
111
+ bias: torch.Tensor | None,
112
+ output_shape: list,
113
+ ) -> torch.Tensor:
114
+ output = torch.ops.vllm.rocm_per_tensor_float_w8a8_scaled_mm_impl(
115
+ A, B, out_dtype, As, Bs, bias
116
+ )
117
+ return torch.narrow(output, 0, 0, A.shape[0]).view(*output_shape)
vllm/model_executor/layers/quantization/kernels/scaled_mm/triton.py
+23
-18
@@ -14,30 +14,35 @@ from vllm.model_executor.layers.quantization.utils.w8a8_utils import (
14
)
15
from vllm.platforms import current_platform
16
17
-from .ScaledMMLinearKernel import ScaledMMLinearKernel, ScaledMMLinearLayerConfig
17
+from .cutlass import CutlassInt8ScaledMMLinearKernel
18
+from .ScaledMMLinearKernel import (
19
+ Int8ScaledMMLinearLayerConfig,
20
+)
21
22
20
-class TritonScaledMMLinearKernel(ScaledMMLinearKernel):
23
+class TritonInt8ScaledMMLinearKernel(CutlassInt8ScaledMMLinearKernel):
24
@classmethod
25
def is_supported(
26
cls, compute_capability: int | None = None
27
) -> tuple[bool, str | None]:
28
if current_platform.is_cuda_alike():
29
return True, None
27
- return False, "Requires ROCm or CUDA."
30
+ return False, "requires ROCm or CUDA."
31
32
@classmethod
30
- def can_implement(cls, c: ScaledMMLinearLayerConfig) -> tuple[bool, str | None]:
33
+ def can_implement(cls, c: Int8ScaledMMLinearLayerConfig) -> tuple[bool, str | None]:
34
if not c.input_symmetric:
32
- return False, "Only symmetric input is supported."
35
+ return False, "supports symmetric input only."
36
return True, None
37
38
def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
36
- weight = getattr(layer, self.w_q_name)
39
+ w_q, _, i_s, _, _ = self._get_layer_params(layer)
40
+ w_q_name, w_s_name, i_s_name, i_zp_name, azp_adj_name = self.layer_param_names
41
+
42
replace_parameter(
43
layer,
39
- self.w_q_name,
40
- torch.nn.Parameter(weight.t().data, requires_grad=False),
44
+ w_q_name,
45
+ torch.nn.Parameter(w_q.t().data, requires_grad=False),
46
)
47
48
# WEIGHT SCALE
@@ -45,29 +50,29 @@ class TritonScaledMMLinearKernel(ScaledMMLinearKernel):
50
# If we have a fused module (QKV, MLP) with per tensor scales (thus N
51
# scales being passed to the kernel), convert to the per-channel case.
52
is_fused_module = len(layer.logical_widths) > 1
48
- weight_scale = getattr(layer, self.w_s_name)
53
+ weight_scale = getattr(layer, w_s_name)
54
if is_fused_module and not self.config.is_channelwise:
55
weight_scale = convert_to_channelwise(weight_scale, layer.logical_widths)
56
replace_parameter(
57
layer,
53
- self.w_s_name,
58
+ w_s_name,
59
torch.nn.Parameter(weight_scale.data, requires_grad=False),
60
)
61
62
# INPUT SCALE
63
if self.config.is_static_input_scheme:
59
- input_scale = getattr(layer, self.i_s_name)
64
+ assert i_s is not None
65
replace_parameter(
66
layer,
62
- self.i_s_name,
63
- torch.nn.Parameter(input_scale.max(), requires_grad=False),
67
+ i_s_name,
68
+ torch.nn.Parameter(i_s.max(), requires_grad=False),
69
)
65
- setattr(layer, self.i_zp_name, None)
70
+ setattr(layer, i_zp_name, None)
71
else:
67
- setattr(layer, self.i_s_name, None)
68
- setattr(layer, self.i_zp_name, None)
72
+ setattr(layer, i_s_name, None)
73
+ setattr(layer, i_zp_name, None)
74
70
- setattr(layer, self.azp_adj_name, None)
75
+ setattr(layer, azp_adj_name, None)
76
77
def apply_weights(
78
self,
@@ -75,7 +80,7 @@ class TritonScaledMMLinearKernel(ScaledMMLinearKernel):
80
x: torch.Tensor,
81
bias: torch.Tensor | None = None,
82
) -> torch.Tensor:
78
- w_q, w_s, i_s, i_zp, azp_adj = self._get_weight_params(layer)
83
+ w_q, w_s, i_s, i_zp, _ = self._get_layer_params(layer)
84
85
x_q, x_s, x_zp = ops.scaled_int8_quant(
86
x.contiguous(), i_s, i_zp, symmetric=True
vllm/model_executor/layers/quantization/modelopt.py
+18
-19
@@ -49,6 +49,9 @@ from vllm.model_executor.layers.quantization.base_config import (
49
QuantizationConfig,
50
QuantizeMethodBase,
51
)
52
+from vllm.model_executor.layers.quantization.kernels.scaled_mm import (
53
+ init_fp8_linear_kernel,
54
+)
55
from vllm.model_executor.layers.quantization.kv_cache import BaseKVCacheMethod
56
from vllm.model_executor.layers.quantization.utils.flashinfer_fp4_moe import (
57
build_flashinfer_fp4_cutlass_moe_prepare_finalize,
@@ -78,10 +81,12 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import (
81
GroupShape,
82
cutlass_fp4_supported,
83
is_layer_skipped,
84
+ kFp8DynamicTokenSym,
85
+ kFp8StaticTensorSym,
86
+ kFp8StaticTokenSym,
87
swizzle_blockscale,
88
)
89
from vllm.model_executor.layers.quantization.utils.w8a8_utils import (
84
- Fp8LinearOp,
90
cutlass_block_fp8_supported,
91
requantize_with_max_scale,
92
)
@@ -438,8 +443,11 @@ class ModelOptFp8LinearMethod(LinearMethodBase):
443
444
def __init__(self, quant_config: ModelOptFp8Config) -> None:
445
self.quant_config = quant_config
441
- self.fp8_linear = Fp8LinearOp(
442
- act_quant_static=True, act_quant_group_shape=GroupShape.PER_TENSOR
446
+ self.fp8_linear = init_fp8_linear_kernel(
447
+ activation_quant_key=kFp8StaticTensorSym,
448
+ weight_quant_key=kFp8StaticTensorSym,
449
+ out_dtype=torch.get_default_dtype(),
450
+ module_name=self.__class__.__name__,
451
)
452
453
def create_weights(
@@ -507,13 +515,7 @@ class ModelOptFp8LinearMethod(LinearMethodBase):
515
x: torch.Tensor,
516
bias: torch.Tensor | None = None,
517
) -> torch.Tensor:
510
- return self.fp8_linear.apply(
511
- input=x,
512
- weight=layer.weight,
513
- weight_scale=layer.weight_scale,
514
- input_scale=layer.input_scale,
515
- bias=bias,
516
- )
518
+ return self.fp8_linear.apply_weights(layer, x, bias)
519
520
521
class ModelOptFp8PcPtLinearMethod(LinearMethodBase):
@@ -527,8 +529,11 @@ class ModelOptFp8PcPtLinearMethod(LinearMethodBase):
529
530
def __init__(self, quant_config: ModelOptFp8Config) -> None:
531
self.quant_config = quant_config
530
- self.fp8_linear = Fp8LinearOp(
531
- act_quant_static=False, act_quant_group_shape=GroupShape.PER_TOKEN
532
+ self.fp8_linear = init_fp8_linear_kernel(
533
+ activation_quant_key=kFp8DynamicTokenSym,
534
+ weight_quant_key=kFp8StaticTokenSym,
535
+ out_dtype=torch.get_default_dtype(),
536
+ module_name=self.__class__.__name__,
537
)
538
539
def create_weights(
@@ -585,13 +590,7 @@ class ModelOptFp8PcPtLinearMethod(LinearMethodBase):
590
x: torch.Tensor,
591
bias: torch.Tensor | None = None,
592
) -> torch.Tensor:
588
- return self.fp8_linear.apply(
589
- input=x,
590
- weight=layer.weight,
591
- weight_scale=layer.weight_scale,
592
- input_scale=None,
593
- bias=bias,
594
- )
593
+ return self.fp8_linear.apply_weights(layer, x, bias)
594
595
596
class ModelOptFp8PbWoLinearMethod(LinearMethodBase):
vllm/model_executor/layers/quantization/ptpc_fp8.py
+10
-13
@@ -17,11 +17,13 @@ from vllm.model_executor.layers.quantization.fp8 import (
17
Fp8KVCacheMethod,
18
Fp8LinearMethod,
19
)
20
+from vllm.model_executor.layers.quantization.kernels.scaled_mm import (
21
+ init_fp8_linear_kernel,
22
+)
23
from vllm.model_executor.layers.quantization.utils.quant_utils import (
21
- GroupShape,
24
is_layer_skipped,
25
+ kFp8DynamicTokenSym,
26
)
24
-from vllm.model_executor.layers.quantization.utils.w8a8_utils import Fp8LinearOp
27
from vllm.platforms import current_platform
28
29
ACTIVATION_SCHEMES = ["static", "dynamic"]
@@ -97,9 +99,11 @@ class PTPCFp8LinearMethod(Fp8LinearMethod):
99
)
100
super().__init__(quant_config=quant_config)
101
# Force weight quantization
100
- self.quant_config.is_checkpoint_fp8_serialized = False
101
- self.fp8_linear = Fp8LinearOp(
102
- act_quant_static=False, act_quant_group_shape=GroupShape.PER_TOKEN
102
+ self.fp8_linear = init_fp8_linear_kernel(
103
+ activation_quant_key=kFp8DynamicTokenSym,
104
+ weight_quant_key=kFp8DynamicTokenSym,
105
+ out_dtype=torch.get_default_dtype(),
106
+ module_name=self.__class__.__name__,
107
)
108
109
def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
@@ -130,11 +134,4 @@ class PTPCFp8LinearMethod(Fp8LinearMethod):
134
x: torch.Tensor,
135
bias: torch.Tensor | None = None,
136
) -> torch.Tensor:
133
- return self.fp8_linear.apply(
134
- input=x,
135
- weight=layer.weight,
136
- weight_scale=layer.weight_scale,
137
- input_scale=None,
138
- input_scale_ub=None,
139
- bias=bias,
140
- )
137
+ return self.fp8_linear.apply_weights(layer, x, bias)
vllm/model_executor/layers/quantization/quark/schemes/quark_w8a8_fp8.py
+28
-19
@@ -7,10 +7,18 @@ from typing import Any, cast
7
import torch
8
from torch.nn import Parameter
9
10
+from vllm.logger import init_logger
11
+from vllm.model_executor.layers.quantization.kernels.scaled_mm import (
12
+ init_fp8_linear_kernel,
13
+)
14
from vllm.model_executor.layers.quantization.quark.schemes import QuarkScheme
11
-from vllm.model_executor.layers.quantization.utils.quant_utils import GroupShape
15
+from vllm.model_executor.layers.quantization.utils.quant_utils import (
16
+ GroupShape,
17
+ kFp8DynamicTokenSym,
18
+ kFp8StaticTensorSym,
19
+ kFp8StaticTokenSym,
20
+)
21
from vllm.model_executor.layers.quantization.utils.w8a8_utils import (
13
- Fp8LinearOp,
22
normalize_e4m3fn_to_e4m3fnuz,
23
requantize_with_max_scale,
24
)
@@ -23,6 +31,8 @@ from vllm.platforms import current_platform
31
32
__all__ = ["QuarkW8A8Fp8"]
33
34
+logger = init_logger(__name__)
35
+
36
37
class QuarkW8A8Fp8(QuarkScheme):
38
def __init__(
@@ -35,15 +45,16 @@ class QuarkW8A8Fp8(QuarkScheme):
45
self.is_static_input_scheme = not cast(bool, input_config.get("is_dynamic"))
46
self.input_qscheme = cast(str, input_config.get("qscheme"))
47
38
- per_token = (
48
+ per_token_activation = (
49
not self.is_static_input_scheme and self.input_qscheme == "per_channel"
50
)
41
- self.act_quant_group_shape = (
42
- GroupShape.PER_TOKEN if per_token else GroupShape.PER_TENSOR
51
+ per_token_weight = self.weight_qscheme == "per_channel"
52
+
53
+ self.activation_quant_key = (
54
+ kFp8DynamicTokenSym if per_token_activation else kFp8StaticTensorSym
55
)
44
- self.fp8_linear = Fp8LinearOp(
45
- act_quant_static=self.is_static_input_scheme,
46
- act_quant_group_shape=self.act_quant_group_shape,
56
+ self.weight_quant_key = (
57
+ kFp8StaticTokenSym if per_token_weight else kFp8StaticTensorSym
58
)
59
self.out_dtype = torch.get_default_dtype()
60
@@ -94,7 +105,7 @@ class QuarkW8A8Fp8(QuarkScheme):
105
layer.input_scale = Parameter(input_scale, requires_grad=False)
106
else:
107
weight_scale = layer.weight_scale.data
97
- if self.act_quant_group_shape == GroupShape.PER_TOKEN:
108
+ if self.activation_quant_key.scale.group_shape == GroupShape.PER_TOKEN:
109
weight_scale = weight_scale.view(-1, 1)
110
layer.weight = Parameter(weight.t(), requires_grad=False)
111
# required by torch.compile to be torch.nn.Parameter
@@ -106,8 +117,6 @@ class QuarkW8A8Fp8(QuarkScheme):
117
# INPUT SCALE
118
if self.is_static_input_scheme:
119
layer.input_scale = Parameter(layer.input_scale.max(), requires_grad=False)
109
- else:
110
- layer.input_scale = None
120
121
def create_weights(
122
self,
@@ -163,17 +172,17 @@ class QuarkW8A8Fp8(QuarkScheme):
172
input_scale[:] = torch.finfo(torch.float32).min
173
layer.register_parameter("input_scale", input_scale)
174
175
+ self.fp8_linear = init_fp8_linear_kernel(
176
+ activation_quant_key=self.activation_quant_key,
177
+ weight_quant_key=self.weight_quant_key,
178
+ out_dtype=torch.get_default_dtype(),
179
+ module_name=self.__class__.__name__,
180
+ )
181
+
182
def apply_weights(
183
self,
184
layer: torch.nn.Module,
185
x: torch.Tensor,
186
bias: torch.Tensor | None = None,
187
) -> torch.Tensor:
172
- return self.fp8_linear.apply(
173
- input=x,
174
- weight=layer.weight,
175
- weight_scale=layer.weight_scale,
176
- out_dtype=self.out_dtype,
177
- input_scale=layer.input_scale,
178
- bias=bias,
179
- )
188
+ return self.fp8_linear.apply_weights(layer, x, bias)
vllm/model_executor/layers/quantization/quark/schemes/quark_w8a8_int8.py
+10
-22
@@ -7,8 +7,7 @@ import torch
7
8
from vllm.logger import init_logger
9
from vllm.model_executor.layers.quantization.kernels.scaled_mm import (
10
- ScaledMMLinearLayerConfig,
11
- choose_scaled_mm_linear_kernel,
10
+ init_int8_linear_kernel,
11
)
12
from vllm.model_executor.layers.quantization.quark.schemes import QuarkScheme
13
from vllm.model_executor.parameter import (
@@ -22,8 +21,6 @@ logger = init_logger(__name__)
21
22
23
class QuarkW8A8Int8(QuarkScheme):
25
- _kernel_backends_being_used: set[str] = set()
26
-
24
def __init__(
25
self,
26
qscheme: str,
@@ -50,18 +47,13 @@ class QuarkW8A8Int8(QuarkScheme):
47
):
48
layer.logical_widths = output_partition_sizes
49
53
- scaled_mm_linear_kernel_config = ScaledMMLinearLayerConfig(
50
+ self.kernel = init_int8_linear_kernel(
51
is_channelwise=(self.qscheme == "per_channel"),
52
is_static_input_scheme=(self.is_static_input_scheme is True),
53
input_symmetric=(self.input_symmetric is True),
54
+ module_name=self.__class__.__name__,
55
)
56
59
- kernel_type = choose_scaled_mm_linear_kernel(scaled_mm_linear_kernel_config)
60
-
61
- if kernel_type.__name__ not in self._kernel_backends_being_used:
62
- logger.info("Using %s for QuarkW8A8Int8", kernel_type.__name__)
63
- self._kernel_backends_being_used.add(kernel_type.__name__)
64
-
57
# WEIGHT
58
weight = ModelWeightParameter(
59
data=torch.empty(
@@ -102,25 +94,21 @@ class QuarkW8A8Int8(QuarkScheme):
94
layer.register_parameter("weight_zero_point", weight_zero_point)
95
96
# INPUT SCALE
97
+ input_zero_point = None
98
+ input_scale = None
99
if self.is_static_input_scheme:
100
input_scale = BasevLLMParameter(
101
data=torch.empty(1, dtype=torch.float32), weight_loader=weight_loader
102
)
109
- layer.register_parameter("input_scale", input_scale)
103
104
input_zero_point = BasevLLMParameter(
105
data=torch.empty(1, dtype=torch.int8), weight_loader=weight_loader
106
)
114
- layer.register_parameter("input_zero_point", input_zero_point)
115
-
116
- self.kernel = kernel_type(
117
- c=scaled_mm_linear_kernel_config,
118
- w_q_param_name="weight",
119
- w_s_param_name="weight_scale",
120
- i_s_param_name="input_scale",
121
- i_zp_param_name="input_zero_point",
122
- azp_adj_param_name="azp_adj",
123
- )
107
+
108
+ layer.register_parameter("input_scale", input_scale)
109
+ layer.register_parameter("input_zero_point", input_zero_point)
110
+ if not hasattr(layer, "azp_adj"):
111
+ layer.register_parameter("azp_adj", None)
112
113
# Checkpoints are serialized in quark format, which is
114
# different from the format the kernel may want. Handle repacking here.
vllm/model_executor/layers/quantization/utils/quant_utils.py
+3
@@ -123,6 +123,9 @@ kFp8StaticTensorSym = QuantKey(FP8_DTYPE, kStaticTensorScale, symmetric=True)
123
kDynamicTensorScale = ScaleDesc(torch.float32, False, GroupShape.PER_TENSOR)
124
kFp8DynamicTensorSym = QuantKey(FP8_DTYPE, kDynamicTensorScale, symmetric=True)
125
126
+kStaticTokenScale = ScaleDesc(torch.float32, True, GroupShape.PER_TOKEN)
127
+kFp8StaticTokenSym = QuantKey(FP8_DTYPE, kStaticTokenScale, symmetric=True)
128
+
129
kDynamicTokenScale = ScaleDesc(torch.float32, False, GroupShape.PER_TOKEN)
130
kFp8DynamicTokenSym = QuantKey(FP8_DTYPE, kDynamicTokenScale, symmetric=True)
131
vllm/model_executor/layers/quantization/utils/w8a8_utils.py
-378
@@ -1,34 +1,11 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
-from collections.abc import Callable
4
5
import torch
7
-from packaging import version
6
7
from vllm import _custom_ops as ops
10
-from vllm import envs
11
-from vllm.config import CompilationMode, get_current_vllm_config
12
-from vllm.model_executor.layers.quantization.input_quant_fp8 import QuantFP8
13
-from vllm.model_executor.layers.quantization.utils.quant_utils import GroupShape
8
from vllm.platforms import current_platform
15
-from vllm.utils.flashinfer import flashinfer_scaled_fp8_mm, has_flashinfer
16
-from vllm.utils.platform_utils import get_cu_count
17
-from vllm.utils.torch_utils import direct_register_custom_op
18
-
19
-# Input scaling factors are no longer optional in _scaled_mm starting
20
-# from pytorch 2.5. Allocating a dummy tensor to pass as input_scale
21
-TORCH_DEVICE_IDENTITY = None
22
-
23
-# The condition to determine if it is on a platform that supports
24
-# torch._scaled_mm rowwise feature.
25
-# The condition is determined once as the operations
26
-# are time-consuming.
27
-USE_ROWWISE_TORCH_SCALED_MM = (
28
- current_platform.is_rocm()
29
- and version.parse(torch.__version__) >= version.parse("2.7")
30
- and current_platform.has_device_capability(94)
31
-)
9
10
11
def sparse_cutlass_supported() -> bool:
@@ -140,361 +117,6 @@ def requantize_with_max_scale(
117
return max_w_scale, weight
118
119
143
-def maybe_create_device_identity():
144
- # Allocate dummy ones tensor for torch._scaled_mm
145
- global TORCH_DEVICE_IDENTITY
146
- if TORCH_DEVICE_IDENTITY is None:
147
- TORCH_DEVICE_IDENTITY = torch.ones(1, dtype=torch.float32)
148
-
149
-
150
-def cutlass_w8a8_scaled_mm(
151
- *,
152
- qinput: torch.Tensor,
153
- weight: torch.Tensor,
154
- out_dtype: torch.dtype,
155
- scale_a: torch.Tensor,
156
- scale_b: torch.Tensor,
157
- bias: torch.Tensor,
158
- output_shape: list,
159
- **kwargs,
160
-) -> torch.Tensor:
161
- # Fused GEMM_DQ
162
- output = ops.cutlass_scaled_mm(
163
- qinput, weight, out_dtype=out_dtype, scale_a=scale_a, scale_b=scale_b, bias=bias
164
- )
165
- return output.view(*output_shape)
166
-
167
-
168
-def flashinfer_w8a8_scaled_mm(
169
- *,
170
- qinput: torch.Tensor,
171
- weight: torch.Tensor,
172
- out_dtype: torch.dtype,
173
- scale_a: torch.Tensor,
174
- scale_b: torch.Tensor,
175
- bias: torch.Tensor,
176
- output_shape: list,
177
- **kwargs,
178
-) -> torch.Tensor:
179
- return flashinfer_scaled_fp8_mm(
180
- qinput, weight, out_dtype=out_dtype, scale_a=scale_a, scale_b=scale_b, bias=bias
181
- )
182
-
183
-
184
-def rocm_per_tensor_w8a8_scaled_mm_impl(
185
- qinput: torch.Tensor,
186
- weight: torch.Tensor,
187
- out_dtype: torch.dtype,
188
- scale_a: torch.Tensor,
189
- scale_b: torch.Tensor,
190
- bias: torch.Tensor,
191
-) -> torch.Tensor:
192
- from vllm.platforms.rocm import on_mi3xx
193
-
194
- if (
195
- envs.VLLM_ROCM_USE_SKINNY_GEMM
196
- and on_mi3xx()
197
- and qinput.shape[0] == 1
198
- and qinput.shape[1] % 16 == 0
199
- and ((bias is None) or (bias.dtype == out_dtype))
200
- ):
201
- output = ops.wvSplitKQ(
202
- weight.t(),
203
- qinput,
204
- out_dtype,
205
- scale_a,
206
- scale_b,
207
- get_cu_count(),
208
- bias,
209
- )
210
- else:
211
- output = torch._scaled_mm(
212
- qinput,
213
- weight,
214
- out_dtype=out_dtype,
215
- scale_a=scale_a,
216
- scale_b=scale_b,
217
- bias=bias,
218
- )
219
- return output
220
-
221
-
222
-def rocm_per_tensor_w8a8_scaled_mm_fake(
223
- qinput: torch.Tensor,
224
- weight: torch.Tensor,
225
- out_dtype: torch.dtype,
226
- scale_a: torch.Tensor,
227
- scale_b: torch.Tensor,
228
- bias: torch.Tensor,
229
-) -> torch.Tensor:
230
- return qinput.new_empty((*qinput.shape[:-1], weight.shape[1]), dtype=out_dtype)
231
-
232
-
233
-def rocm_per_tensor_w8a8_scaled_mm(
234
- *,
235
- qinput: torch.Tensor,
236
- weight: torch.Tensor,
237
- out_dtype: torch.dtype,
238
- scale_a: torch.Tensor,
239
- scale_b: torch.Tensor,
240
- bias: torch.Tensor,
241
- output_shape: list,
242
-) -> torch.Tensor:
243
- output = torch.ops.vllm.rocm_per_tensor_w8a8_scaled_mm_impl(
244
- qinput, weight, out_dtype, scale_a, scale_b, bias
245
- )
246
- return torch.narrow(output, 0, 0, qinput.shape[0]).view(*output_shape)
247
-
248
-
249
-direct_register_custom_op(
250
- op_name="rocm_per_tensor_w8a8_scaled_mm_impl",
251
- op_func=rocm_per_tensor_w8a8_scaled_mm_impl,
252
- fake_impl=rocm_per_tensor_w8a8_scaled_mm_fake,
253
-)
254
-
255
-
256
-def torch_per_tensor_w8a8_scaled_mm(
257
- *,
258
- qinput: torch.Tensor,
259
- weight: torch.Tensor,
260
- out_dtype: torch.dtype,
261
- scale_a: torch.Tensor,
262
- scale_b: torch.Tensor,
263
- bias: torch.Tensor,
264
- output_shape: list,
265
-) -> torch.Tensor:
266
- output = torch._scaled_mm(
267
- qinput, weight, out_dtype=out_dtype, scale_a=scale_a, scale_b=scale_b, bias=bias
268
- )
269
- # A fix for discrepancy in scaled_mm which returns tuple
270
- # for torch < 2.5 and a single value in torch >= 2.5
271
- if type(output) is tuple and len(output) == 2:
272
- output = output[0]
273
-
274
- return torch.narrow(output, 0, 0, qinput.shape[0]).view(*output_shape)
275
-
276
-
277
-def torch_per_token_w8a8_scaled_mm(
278
- *,
279
- qinput: torch.Tensor,
280
- weight: torch.Tensor,
281
- out_dtype: torch.dtype,
282
- scale_a: torch.Tensor,
283
- scale_b: torch.Tensor,
284
- bias: torch.Tensor,
285
- output_shape: list,
286
- **kwargs,
287
-) -> torch.Tensor:
288
- # Note: Callers of this function should check USE_ROWWISE_TORCH_SCALED_MM
289
- # when using it.
290
- # For now it has only been validated on ROCm platform.
291
- # fp8 rowwise scaling in torch._scaled_mm is introduced in
292
- # https://github.com/pytorch/pytorch/pull/144432 using
293
- # hipBLASLt and ROCm 6.3, which only exists in torch 2.7 and above.
294
- #
295
- # For CUDA platform please validate if the torch._scaled_mm supports
296
- # rowwise scaled GEMM before using it
297
-
298
- # Fused GEMM_DQ Rowwise GEMM
299
- output = torch._scaled_mm(
300
- qinput,
301
- weight,
302
- out_dtype=out_dtype,
303
- scale_a=scale_a,
304
- scale_b=scale_b.t(),
305
- bias=bias,
306
- )
307
-
308
- output = torch.narrow(output, 0, 0, qinput.shape[0])
309
- output = output.view(*output_shape)
310
- return output
311
-
312
-
313
-def torch_channelwise_w8a8_scaled_mm(
314
- *,
315
- qinput: torch.Tensor,
316
- weight: torch.Tensor,
317
- out_dtype: torch.dtype,
318
- scale_a: torch.Tensor,
319
- scale_b: torch.Tensor,
320
- bias: torch.Tensor,
321
- output_shape: list,
322
- **kwargs,
323
-) -> torch.Tensor:
324
- # Use unfused DQ due to limitations with scaled_mm
325
-
326
- # Symmetric quantized GEMM by definition computes the following:
327
- # C = (s_x * X) (s_w * W) + bias
328
- # This is equivalent to dequantizing the weights and activations
329
- # before applying a GEMM.
330
- #
331
- # In order to compute quantized operands, a quantized kernel
332
- # will rewrite the above like so:
333
- # C = s_w * s_x * (X * W) + bias
334
- #
335
- # For the scaled_mm fallback case, we break this down, since it
336
- # does not support s_w being a vector.
337
-
338
- # GEMM
339
- # This computes C = (X * W).
340
- # Output in fp32 to allow subsequent ops to happen in-place
341
- output = torch._scaled_mm(
342
- qinput,
343
- weight,
344
- scale_a=TORCH_DEVICE_IDENTITY,
345
- scale_b=TORCH_DEVICE_IDENTITY,
346
- out_dtype=torch.float32,
347
- )
348
- # A fix for discrepancy in scaled_mm which returns tuple
349
- # for torch < 2.5 and a single value in torch >= 2.5
350
- if type(output) is tuple and len(output) == 2:
351
- output = output[0]
352
- # Unpad (undo num_token_padding)
353
- output = torch.narrow(output, 0, 0, qinput.shape[0])
354
- x_scale = torch.narrow(scale_a, 0, 0, qinput.shape[0])
355
-
356
- # DQ
357
- # C = sw * sx * (X * W) + bias
358
- output = output * x_scale * scale_b.t()
359
- if bias is not None:
360
- output = output + bias
361
- return output.to(out_dtype).view(*output_shape)
362
-
363
-
364
-def dispatch_w8a8_scaled_mm(
365
- preferred_backend: str, per_tensor_weights: bool, per_tensor_activations: bool
366
-) -> Callable[..., torch.Tensor]:
367
- if per_tensor_weights and per_tensor_activations:
368
- if preferred_backend == "rocm":
369
- return rocm_per_tensor_w8a8_scaled_mm
370
- if preferred_backend == "flashinfer":
371
- return flashinfer_w8a8_scaled_mm
372
- if preferred_backend == "cutlass":
373
- return cutlass_w8a8_scaled_mm
374
- return torch_per_tensor_w8a8_scaled_mm
375
-
376
- # cutlass_scaled_mm supports per tensor/channel W and per tensor/token A
377
- if preferred_backend == "cutlass" or preferred_backend == "flashinfer":
378
- return cutlass_w8a8_scaled_mm
379
-
380
- # If torch.scaled_mm supports per-channel (weights) per-token (inputs)
381
- if (
382
- not per_tensor_weights
383
- and not per_tensor_activations
384
- and USE_ROWWISE_TORCH_SCALED_MM
385
- ):
386
- return torch_per_token_w8a8_scaled_mm
387
- # Normally, torch.scaled_mm supports per tensor weights + activations only
388
- # so fallback to naive if per channel or per token
389
- return torch_channelwise_w8a8_scaled_mm
390
-
391
-
392
-# TODO(luka): follow similar pattern for marlin and block-fp8-linear
393
-# https://github.com/vllm-project/vllm/issues/14397
394
-class Fp8LinearOp:
395
- """
396
- This class executes a FP8 linear layer using cutlass if supported and
397
- torch.scaled_mm otherwise.
398
- It needs to be a class instead of a method so that config can be read
399
- in the __init__ method, as reading config is not allowed inside forward.
400
- """
401
-
402
- def __init__(
403
- self,
404
- act_quant_static: bool,
405
- act_quant_group_shape: GroupShape = GroupShape.PER_TENSOR,
406
- pad_output: bool | None = None,
407
- ):
408
- if current_platform.is_rocm():
409
- self.preferred_backend = "rocm"
410
- elif current_platform.is_cuda() and cutlass_fp8_supported():
411
- if has_flashinfer() and current_platform.has_device_capability(100):
412
- self.preferred_backend = "flashinfer"
413
- else:
414
- self.preferred_backend = "cutlass"
415
- else:
416
- self.preferred_backend = "torch"
417
-
418
- # Note: we pad the input because torch._scaled_mm is more performant
419
- # for matrices with batch dimension > 16.
420
- # This could change in the future.
421
- # We also don't pad when using torch.compile,
422
- # as it breaks with dynamic shapes.
423
- if pad_output is None:
424
- config = get_current_vllm_config().compilation_config
425
- pad_output = (
426
- config.mode < CompilationMode.VLLM_COMPILE
427
- and self.preferred_backend == "torch"
428
- )
429
-
430
- self.output_padding = 17 if pad_output else None
431
- self.act_quant_static = act_quant_static
432
- self.act_quant_group_shape = act_quant_group_shape
433
- self.quant_fp8 = QuantFP8(
434
- static=act_quant_static,
435
- group_shape=act_quant_group_shape,
436
- num_token_padding=self.output_padding,
437
- )
438
-
439
- def apply(
440
- self,
441
- input: torch.Tensor,
442
- weight: torch.Tensor,
443
- weight_scale: torch.Tensor,
444
- out_dtype: torch.dtype | None = None,
445
- input_scale: torch.Tensor | None = None,
446
- input_scale_ub: torch.Tensor | None = None,
447
- bias: torch.Tensor | None = None,
448
- ) -> torch.Tensor:
449
- # ops.scaled_fp8_quant supports both dynamic and static quant.
450
- # If dynamic, layer.input_scale is None and x_scale computed from x.
451
- # If static, layer.input_scale is scalar and x_scale is input_scale.
452
-
453
- # View input as 2D matrix for fp8 methods
454
- input_2d = input.view(-1, input.shape[-1])
455
- output_shape = [*input.shape[:-1], weight.shape[1]]
456
-
457
- if out_dtype is None:
458
- out_dtype = input.dtype
459
-
460
- # If input not quantized
461
- # TODO(luka) remove this path if not used anymore
462
- if input.dtype != current_platform.fp8_dtype():
463
- qinput, x_scale = self.quant_fp8(
464
- input_2d,
465
- input_scale,
466
- input_scale_ub,
467
- )
468
- else:
469
- qinput, x_scale = input_2d, input_scale
470
-
471
- # Must have dim() conditions
472
- # In per-token quant scenario, when the number of token is 1,
473
- # the scale will only have 1 elements.
474
- # Without checking the dim(),
475
- # we cannot distingushes between per-tensor and per-token quant.
476
- # Example:
477
- # When the number of token is 1, per-token scale is [[1]]
478
- # When per-tensor scale is [1] or ().
479
- per_tensor_weights = weight_scale.numel() == 1
480
- per_tensor_activations = (x_scale.numel() == 1) and x_scale.dim() < 2
481
-
482
- # TODO(luka) do this dispatch during init (after ScaledMM refactor)
483
- w8a8_scaled_mm_func = dispatch_w8a8_scaled_mm(
484
- self.preferred_backend, per_tensor_weights, per_tensor_activations
485
- )
486
-
487
- return w8a8_scaled_mm_func(
488
- qinput=qinput,
489
- weight=weight,
490
- out_dtype=out_dtype,
491
- scale_a=x_scale,
492
- scale_b=weight_scale,
493
- bias=bias,
494
- output_shape=output_shape,
495
- )
496
-
497
-
120
def normalize_e4m3fn_to_e4m3fnuz(
121
weight: torch.Tensor,
122
weight_scale: torch.Tensor,