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 15, 19:45 UTC
[FP8] Refactor apply_fp8_linear and apply_fp8_linear_generic into an object (#14390)
Signed-off-by: luka <luka@neuralmagic.com>
Luka Govedič committed
Mar 7, 2025 at 00:20 UTC
e1744502c21f38d913c0905b10471e96fe755104
11 files changed
+268
-242
tests/compile/test_fusion.py
+7
-13
@@ -13,7 +13,7 @@ from vllm.compilation.noop_elimination import NoOpEliminationPass
13
from vllm.config import CompilationConfig, CompilationLevel, VllmConfig
14
from vllm.model_executor.layers.layernorm import RMSNorm
15
from vllm.model_executor.layers.quantization.utils.w8a8_utils import (
16
- CUTLASS_FP8_SUPPORTED, apply_fp8_linear, maybe_create_device_identity)
16
+ CUTLASS_FP8_SUPPORTED, Fp8LinearOp, maybe_create_device_identity)
17
18
from .backend import TestBackend
19
@@ -34,26 +34,20 @@ class TestModel(torch.nn.Module):
34
torch.rand(hidden_size, hidden_size).to(dtype=FP8_DTYPE).t()
35
for _ in range(2)
36
]
37
+ self.fp8_linear = Fp8LinearOp(
38
+ cutlass_fp8_supported=cutlass_fp8_enabled,
39
+ use_per_token_if_dynamic=True)
40
41
def forward(self, x):
42
resid = torch.sqrt(x)
43
y = self.norm[0](x)
44
42
- x2 = apply_fp8_linear(y,
43
- self.w[0],
44
- self.wscale[0],
45
- self.scale[0],
46
- use_per_token_if_dynamic=True,
47
- cutlass_fp8_supported=self.cutlass_fp8_enabled)
45
+ x2 = self.fp8_linear.apply(y, self.w[0], self.wscale[0], self.scale[0])
46
# make sure resid is used for replacement to work
47
y2, resid = self.norm[1](x2, resid)
48
51
- x3 = apply_fp8_linear(y2,
52
- self.w[1],
53
- self.wscale[1],
54
- self.scale[1],
55
- use_per_token_if_dynamic=True,
56
- cutlass_fp8_supported=self.cutlass_fp8_enabled)
49
+ x3 = self.fp8_linear.apply(y2, self.w[1], self.wscale[1],
50
+ self.scale[1])
51
y3, resid = self.norm[2](x3, resid) # use resid here
52
return y3
53
vllm/attention/backends/mla/common.py
+4
-3
@@ -226,7 +226,7 @@ from vllm.model_executor.layers.quantization.compressed_tensors.schemes import (
226
CompressedTensorsW8A8Fp8)
227
from vllm.model_executor.layers.quantization.fp8 import Fp8LinearMethod
228
from vllm.model_executor.layers.quantization.utils.fp8_utils import (
229
- apply_fp8_linear_generic, current_platform_fp8_dtype, is_fp8)
229
+ Fp8LinearGenericOp, current_platform_fp8_dtype, is_fp8)
230
from vllm.model_executor.layers.quantization.utils.quant_utils import (
231
scaled_quantize)
232
from vllm.model_executor.layers.rotary_embedding import (
@@ -1057,6 +1057,7 @@ class MLACommonImpl(MLAAttentionImpl[T], Generic[T]):
1057
self.kv_b_proj = kv_b_proj
1058
self.o_proj = o_proj
1059
self.triton_fa_func = triton_attention
1060
+ self.fp8_linear_generic = Fp8LinearGenericOp()
1061
1062
# Handle the differences between the flash_attn_varlen from flash_attn
1063
# and the one from vllm_flash_attn. The former is used on RoCM and the
@@ -1071,7 +1072,7 @@ class MLACommonImpl(MLAAttentionImpl[T], Generic[T]):
1072
def _v_up_proj_and_o_proj(self, x):
1073
if envs.VLLM_MLA_PERFORM_MATRIX_ABSORPTION:
1074
if is_fp8(self.W_UV_O):
1074
- output_parallel = apply_fp8_linear_generic(
1075
+ output_parallel = self.fp8_linear_generic.apply(
1076
x.flatten(start_dim=1), self.W_UV_O, self.W_UV_O_scales,
1077
self.reqaunt_input_group_shape,
1078
self.reqaunt_weight_group_shape)
@@ -1091,7 +1092,7 @@ class MLACommonImpl(MLAAttentionImpl[T], Generic[T]):
1092
def _q_proj_and_k_up_proj(self, x):
1093
if envs.VLLM_MLA_PERFORM_MATRIX_ABSORPTION:
1094
if is_fp8(self.W_Q_UK):
1094
- return apply_fp8_linear_generic(
1095
+ return self.fp8_linear_generic.apply(
1096
x, self.W_Q_UK, self.W_Q_UK_scales,
1097
self.reqaunt_input_group_shape,
1098
self.reqaunt_weight_group_shape).view(
vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_w8a8_fp8.py
+8
-11
@@ -9,8 +9,8 @@ from torch.nn import Parameter
9
from vllm.model_executor.layers.quantization.compressed_tensors.schemes import (
10
CompressedTensorsScheme)
11
from vllm.model_executor.layers.quantization.utils.w8a8_utils import (
12
- apply_fp8_linear, cutlass_fp8_supported, maybe_create_device_identity,
13
- normalize_e4m3fn_to_e4m3fnuz, requantize_with_max_scale)
12
+ Fp8LinearOp, maybe_create_device_identity, normalize_e4m3fn_to_e4m3fnuz,
13
+ requantize_with_max_scale)
14
from vllm.model_executor.parameter import (ChannelQuantScaleParameter,
15
ModelWeightParameter,
16
PerTensorScaleParameter)
@@ -24,7 +24,7 @@ class CompressedTensorsW8A8Fp8(CompressedTensorsScheme):
24
def __init__(self, strategy: str, is_static_input_scheme: bool):
25
self.strategy = strategy
26
self.is_static_input_scheme = is_static_input_scheme
27
- self.cutlass_fp8_supported = cutlass_fp8_supported()
27
+ self.fp8_linear = Fp8LinearOp(use_per_token_if_dynamic=True)
28
29
@classmethod
30
def get_min_capability(cls) -> int:
@@ -140,11 +140,8 @@ class CompressedTensorsW8A8Fp8(CompressedTensorsScheme):
140
x: torch.Tensor,
141
bias: Optional[torch.Tensor] = None) -> torch.Tensor:
142
143
- return apply_fp8_linear(
144
- input=x,
145
- weight=layer.weight,
146
- weight_scale=layer.weight_scale,
147
- input_scale=layer.input_scale,
148
- bias=bias,
149
- cutlass_fp8_supported=self.cutlass_fp8_supported,
150
- use_per_token_if_dynamic=True)
143
+ return self.fp8_linear.apply(input=x,
144
+ weight=layer.weight,
145
+ weight_scale=layer.weight_scale,
146
+ input_scale=layer.input_scale,
147
+ bias=bias)
vllm/model_executor/layers/quantization/fbgemm_fp8.py
+9
-13
@@ -11,14 +11,12 @@ from vllm.model_executor.layers.linear import (LinearBase, LinearMethodBase,
11
UnquantizedLinearMethod)
12
from vllm.model_executor.layers.quantization.base_config import (
13
QuantizationConfig, QuantizeMethodBase)
14
-from vllm.model_executor.layers.quantization.fp8 import cutlass_fp8_supported
14
from vllm.model_executor.layers.quantization.utils.marlin_utils_fp8 import (
15
apply_fp8_marlin_linear, prepare_fp8_layer_for_marlin)
16
from vllm.model_executor.layers.quantization.utils.quant_utils import (
17
is_layer_skipped)
18
from vllm.model_executor.layers.quantization.utils.w8a8_utils import (
20
- apply_fp8_linear, maybe_create_device_identity,
21
- normalize_e4m3fn_to_e4m3fnuz)
19
+ Fp8LinearOp, maybe_create_device_identity, normalize_e4m3fn_to_e4m3fnuz)
20
from vllm.model_executor.parameter import (ChannelQuantScaleParameter,
21
ModelWeightParameter)
22
from vllm.platforms import current_platform
@@ -37,6 +35,7 @@ class FBGEMMFp8Config(QuantizationConfig):
35
# For GPUs that lack FP8 hardware support, we can leverage the Marlin
36
# kernel for fast weight-only FP8 quantization
37
self.use_marlin = not current_platform.has_device_capability(89)
38
+ self.fp8_linear = Fp8LinearOp()
39
40
@classmethod
41
def get_name(cls) -> str:
@@ -73,7 +72,7 @@ class FBGEMMFp8LinearMethod(LinearMethodBase):
72
73
def __init__(self, quant_config: FBGEMMFp8Config):
74
self.quant_config = quant_config
76
- self.cutlass_fp8_supported = cutlass_fp8_supported()
75
+ self.fp8_linear = Fp8LinearOp(use_per_token_if_dynamic=True)
76
77
def create_weights(
78
self,
@@ -159,12 +158,9 @@ class FBGEMMFp8LinearMethod(LinearMethodBase):
158
size_k=layer.input_size_per_partition,
159
bias=bias)
160
162
- return apply_fp8_linear(
163
- input=x,
164
- weight=layer.weight,
165
- weight_scale=layer.weight_scale,
166
- input_scale=None,
167
- input_scale_ub=layer.input_scale_ub,
168
- bias=bias,
169
- cutlass_fp8_supported=self.cutlass_fp8_supported,
170
- use_per_token_if_dynamic=True)
161
+ return self.fp8_linear.apply(input=x,
162
+ weight=layer.weight,
163
+ weight_scale=layer.weight_scale,
164
+ input_scale=None,
165
+ input_scale_ub=layer.input_scale_ub,
166
+ bias=bias)
vllm/model_executor/layers/quantization/fp8.py
+10
-11
@@ -23,7 +23,7 @@ from vllm.model_executor.layers.quantization.utils.marlin_utils_fp8 import (
23
from vllm.model_executor.layers.quantization.utils.quant_utils import (
24
is_layer_skipped)
25
from vllm.model_executor.layers.quantization.utils.w8a8_utils import (
26
- all_close_1d, apply_fp8_linear, convert_to_channelwise,
26
+ Fp8LinearOp, all_close_1d, convert_to_channelwise,
27
cutlass_block_fp8_supported, cutlass_fp8_supported,
28
maybe_create_device_identity, normalize_e4m3fn_to_e4m3fnuz,
29
per_tensor_dequantize, requantize_with_max_scale)
@@ -137,7 +137,6 @@ class Fp8LinearMethod(LinearMethodBase):
137
138
def __init__(self, quant_config: Fp8Config):
139
self.quant_config = quant_config
140
- self.cutlass_fp8_supported = cutlass_fp8_supported()
140
self.cutlass_block_fp8_supported = cutlass_block_fp8_supported()
141
142
# For GPUs that lack FP8 hardware support, we can leverage the Marlin
@@ -153,6 +152,10 @@ class Fp8LinearMethod(LinearMethodBase):
152
# Marlin doesn't support block-wise fp8
153
self.use_marlin = False
154
155
+ self.fp8_linear = Fp8LinearOp(
156
+ # Default to using per_token quantization if cutlass is supported
157
+ use_per_token_if_dynamic=cutlass_fp8_supported())
158
+
159
def create_weights(
160
self,
161
layer: torch.nn.Module,
@@ -381,15 +384,11 @@ class Fp8LinearMethod(LinearMethodBase):
384
cutlass_block_fp8_supported=self.cutlass_block_fp8_supported,
385
)
386
384
- return apply_fp8_linear(
385
- input=x,
386
- weight=layer.weight,
387
- weight_scale=layer.weight_scale,
388
- input_scale=layer.input_scale,
389
- bias=bias,
390
- cutlass_fp8_supported=self.cutlass_fp8_supported,
391
- # Default to using per_token quantization if cutlass is supported
392
- use_per_token_if_dynamic=self.cutlass_fp8_supported)
387
+ return self.fp8_linear.apply(input=x,
388
+ weight=layer.weight,
389
+ weight_scale=layer.weight_scale,
390
+ input_scale=layer.input_scale,
391
+ bias=bias)
392
393
394
class Fp8MoEMethod(FusedMoEMethodBase):
vllm/model_executor/layers/quantization/modelopt.py
+7
-9
@@ -12,7 +12,7 @@ from vllm.model_executor.layers.quantization.base_config import (
12
QuantizationConfig, QuantizeMethodBase)
13
from vllm.model_executor.layers.quantization.kv_cache import BaseKVCacheMethod
14
from vllm.model_executor.layers.quantization.utils.w8a8_utils import (
15
- apply_fp8_linear, cutlass_fp8_supported, requantize_with_max_scale)
15
+ Fp8LinearOp, requantize_with_max_scale)
16
from vllm.model_executor.parameter import (ModelWeightParameter,
17
PerTensorScaleParameter)
18
@@ -95,7 +95,7 @@ class ModelOptFp8LinearMethod(LinearMethodBase):
95
96
def __init__(self, quant_config: ModelOptFp8Config):
97
self.quant_config = quant_config
98
- self.cutlass_fp8_supported = cutlass_fp8_supported()
98
+ self.fp8_linear = Fp8LinearOp()
99
100
def create_weights(
101
self,
@@ -157,10 +157,8 @@ class ModelOptFp8LinearMethod(LinearMethodBase):
157
x: torch.Tensor,
158
bias: Optional[torch.Tensor] = None,
159
) -> torch.Tensor:
160
- return apply_fp8_linear(
161
- input=x,
162
- weight=layer.weight,
163
- weight_scale=layer.weight_scale,
164
- input_scale=layer.input_scale,
165
- bias=bias,
166
- cutlass_fp8_supported=self.cutlass_fp8_supported)
160
+ return self.fp8_linear.apply(input=x,
161
+ weight=layer.weight,
162
+ weight_scale=layer.weight_scale,
163
+ input_scale=layer.input_scale,
164
+ bias=bias)
vllm/model_executor/layers/quantization/ptpc_fp8.py
+9
-9
@@ -17,7 +17,7 @@ from vllm.model_executor.layers.quantization.fp8 import (Fp8Config,
17
from vllm.model_executor.layers.quantization.utils.quant_utils import (
18
is_layer_skipped)
19
from vllm.model_executor.layers.quantization.utils.w8a8_utils import (
20
- apply_fp8_linear)
20
+ Fp8LinearOp)
21
from vllm.platforms import current_platform
22
23
ACTIVATION_SCHEMES = ["static", "dynamic"]
@@ -93,6 +93,8 @@ class PTPCFp8LinearMethod(Fp8LinearMethod):
93
super().__init__(quant_config=quant_config)
94
# Force weight quantization
95
self.quant_config.is_checkpoint_fp8_serialized = False
96
+ self.fp8_linear = Fp8LinearOp(cutlass_fp8_supported=False,
97
+ use_per_token_if_dynamic=True)
98
99
def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
100
layer.weight = torch.nn.Parameter(layer.weight.data,
@@ -115,11 +117,9 @@ class PTPCFp8LinearMethod(Fp8LinearMethod):
117
x: torch.Tensor,
118
bias: Optional[torch.Tensor] = None) -> torch.Tensor:
119
118
- return apply_fp8_linear(input=x,
119
- weight=layer.weight,
120
- weight_scale=layer.weight_scale,
121
- input_scale=None,
122
- input_scale_ub=None,
123
- bias=bias,
124
- cutlass_fp8_supported=False,
125
- use_per_token_if_dynamic=True)
120
+ return self.fp8_linear.apply(input=x,
121
+ weight=layer.weight,
122
+ weight_scale=layer.weight_scale,
123
+ input_scale=None,
124
+ input_scale_ub=None,
125
+ bias=bias)
vllm/model_executor/layers/quantization/quark/schemes/quark_w8a8_fp8.py
+7
-11
@@ -7,8 +7,7 @@ from torch.nn import Parameter
7
8
from vllm.model_executor.layers.quantization.quark.schemes import QuarkScheme
9
from vllm.model_executor.layers.quantization.utils.w8a8_utils import (
10
- apply_fp8_linear, cutlass_fp8_supported, normalize_e4m3fn_to_e4m3fnuz,
11
- requantize_with_max_scale)
10
+ Fp8LinearOp, normalize_e4m3fn_to_e4m3fnuz, requantize_with_max_scale)
11
from vllm.model_executor.parameter import (ChannelQuantScaleParameter,
12
ModelWeightParameter,
13
PerTensorScaleParameter)
@@ -22,7 +21,7 @@ class QuarkW8A8Fp8(QuarkScheme):
21
def __init__(self, qscheme: str, is_static_input_scheme: Optional[bool]):
22
self.qscheme = qscheme
23
self.is_static_input_scheme = is_static_input_scheme
25
- self.cutlass_fp8_supported = cutlass_fp8_supported()
24
+ self.fp8_linear = Fp8LinearOp(use_per_token_if_dynamic=True)
25
26
@classmethod
27
def get_min_capability(cls) -> int:
@@ -132,11 +131,8 @@ class QuarkW8A8Fp8(QuarkScheme):
131
x: torch.Tensor,
132
bias: Optional[torch.Tensor] = None) -> torch.Tensor:
133
135
- return apply_fp8_linear(
136
- input=x,
137
- weight=layer.weight,
138
- weight_scale=layer.weight_scale,
139
- input_scale=layer.input_scale,
140
- bias=bias,
141
- cutlass_fp8_supported=self.cutlass_fp8_supported,
142
- use_per_token_if_dynamic=True)
134
+ return self.fp8_linear.apply(input=x,
135
+ weight=layer.weight,
136
+ weight_scale=layer.weight_scale,
137
+ input_scale=layer.input_scale,
138
+ bias=bias)
vllm/model_executor/layers/quantization/utils/fp8_utils.py
+54
-38
@@ -15,7 +15,8 @@ from vllm.logger import init_logger
15
from vllm.model_executor.layers.quantization.utils.quant_utils import (
16
_normalize_quant_group_shape, scaled_dequantize)
17
from vllm.model_executor.layers.quantization.utils.w8a8_utils import (
18
- CUTLASS_BLOCK_FP8_SUPPORTED, CUTLASS_FP8_SUPPORTED, apply_fp8_linear)
18
+ CUTLASS_BLOCK_FP8_SUPPORTED, Fp8LinearOp, cutlass_block_fp8_supported,
19
+ cutlass_fp8_supported)
20
from vllm.platforms import current_platform
21
from vllm.utils import direct_register_custom_op
22
@@ -32,6 +33,8 @@ def is_fp8(x: Union[torch.dtype, torch.Tensor]) -> bool:
33
return x == torch.float8_e4m3fn or x == torch.float8_e4m3fnuz
34
35
36
+# TODO fix ROCm->Triton custom path:
37
+# https://github.com/vllm-project/vllm/issues/14397
38
def apply_w8a8_block_fp8_linear(
39
input: torch.Tensor,
40
weight: torch.Tensor,
@@ -49,6 +52,7 @@ def apply_w8a8_block_fp8_linear(
52
shape_supported_by_cutlass = (weight.shape[0] % 128 == 0
53
and weight.shape[1] % 128 == 0)
54
if current_platform.is_rocm():
55
+ # TODO this is never used, as cutlass_block_fp8_supported is False
56
scale_a_shape = ((input_2d.shape[-1] // block_size[1], ) +
57
input_2d.shape[:-1])[::-1]
58
scale_b_shape = (weight_scale.view(-1, 1)
@@ -104,43 +108,55 @@ direct_register_custom_op(
108
# Unify the interface between `apply_w8a8_block_fp8_linear` and
109
# `apply_fp8_linear`
110
# NOTE(lucas): this is quite messy, we should think through this more formally
107
-def apply_fp8_linear_generic(
108
- input: torch.Tensor,
109
- weight: torch.Tensor,
110
- weight_scale: torch.Tensor,
111
- input_group_shape: Tuple[int, int],
112
- weight_group_shape: Tuple[int, int],
113
- input_scale: Optional[torch.Tensor] = None, # static scale if one
114
- cutlass_fp8_supported: bool = CUTLASS_FP8_SUPPORTED,
115
- cutlass_block_fp8_supported: bool = CUTLASS_BLOCK_FP8_SUPPORTED,
116
-) -> torch.Tensor:
117
- # View input as 2D matrix for fp8 methods
118
- input = input.view(-1, input.shape[-1])
119
-
120
- weight_group_shape = _normalize_quant_group_shape(\
121
- weight, weight_group_shape)
122
- input_group_shape = _normalize_quant_group_shape(input, input_group_shape)
123
-
124
- def is_dim_blocked(dim, shape, group_shape):
125
- return group_shape < shape[dim] and group_shape > 1
126
-
127
- if is_dim_blocked(0, weight.shape, weight_group_shape[0])\
128
- and is_dim_blocked(1, weight.shape, weight_group_shape[1]) and\
129
- input_group_shape == (1, weight_group_shape[1]):
130
- return apply_w8a8_block_fp8_linear(
131
- input,
132
- weight,
133
- list(weight_group_shape),
134
- weight_scale,
135
- cutlass_block_fp8_supported=cutlass_block_fp8_supported)
136
- else:
137
- # Despite having linear in the it doesn't conform to
138
- # `torch.nn.functional.linear` which is defined as `input @ weight.T`
139
- # so we explicitly transpose the weight matrix here
140
- return apply_fp8_linear(input, weight.T, weight_scale.T,
141
- cutlass_fp8_supported=cutlass_fp8_supported,
142
- use_per_token_if_dynamic=\
143
- (input_group_shape == (1, input.shape[1])))
111
+# TODO(luka): unify this better
112
+# https://github.com/vllm-project/vllm/issues/14397
113
+class Fp8LinearGenericOp:
114
+
115
+ def __init__(
116
+ self,
117
+ cutlass_fp8_supported: bool = cutlass_fp8_supported(),
118
+ cutlass_block_fp8_supported: bool = cutlass_block_fp8_supported(),
119
+ ):
120
+ self.cutlass_block_fp8_supported = cutlass_block_fp8_supported
121
+ self.fp8_linear = Fp8LinearOp(
122
+ cutlass_fp8_supported=cutlass_fp8_supported)
123
+
124
+ def apply(
125
+ self,
126
+ input: torch.Tensor,
127
+ weight: torch.Tensor,
128
+ weight_scale: torch.Tensor,
129
+ input_group_shape: Tuple[int, int],
130
+ weight_group_shape: Tuple[int, int],
131
+ input_scale: Optional[torch.Tensor] = None, # static scale if one
132
+ ) -> torch.Tensor:
133
+ # View input as 2D matrix for fp8 methods
134
+ input = input.view(-1, input.shape[-1])
135
+
136
+ weight_group_shape = _normalize_quant_group_shape( \
137
+ weight, weight_group_shape)
138
+ input_group_shape = _normalize_quant_group_shape(
139
+ input, input_group_shape)
140
+
141
+ def is_dim_blocked(dim, shape, group_shape):
142
+ return group_shape < shape[dim] and group_shape > 1
143
+
144
+ if is_dim_blocked(0, weight.shape, weight_group_shape[0])\
145
+ and is_dim_blocked(1, weight.shape, weight_group_shape[1]) and\
146
+ input_group_shape == (1, weight_group_shape[1]):
147
+ return apply_w8a8_block_fp8_linear(
148
+ input,
149
+ weight,
150
+ list(weight_group_shape),
151
+ weight_scale,
152
+ cutlass_block_fp8_supported=self.cutlass_block_fp8_supported)
153
+ else:
154
+ # Despite having linear in the name it doesn't conform to
155
+ # `torch.nn.functional.linear` which is defined as
156
+ # `input @ weight.T` so we explicitly transpose the weight matrix
157
+ return self.fp8_linear.apply(input, weight.T, weight_scale.T,
158
+ use_per_token_if_dynamic=\
159
+ (input_group_shape == (1, input.shape[1])))
160
161
162
def input_to_float8(
vllm/model_executor/layers/quantization/utils/w8a8_utils.py
+149
-121
@@ -121,134 +121,162 @@ def maybe_create_device_identity():
121
TORCH_DEVICE_IDENTITY = torch.ones(1, dtype=torch.float32)
122
123
124
-def apply_fp8_linear(
125
- input: torch.Tensor,
126
- weight: torch.Tensor,
127
- weight_scale: torch.Tensor,
128
- input_scale: Optional[torch.Tensor] = None,
129
- input_scale_ub: Optional[torch.Tensor] = None,
130
- bias: Optional[torch.Tensor] = None,
131
- cutlass_fp8_supported: bool = CUTLASS_FP8_SUPPORTED,
132
- use_per_token_if_dynamic: bool = False,
133
-) -> torch.Tensor:
134
- # ops.scaled_fp8_quant supports both dynamic and static quant.
135
- # If dynamic, layer.input_scale is None and x_scale computed from x.
136
- # If static, layer.input_scale is scalar and x_scale is input_scale.
137
-
138
- # View input as 2D matrix for fp8 methods
139
- input_2d = input.view(-1, input.shape[-1])
140
- output_shape = [*input.shape[:-1], weight.shape[1]]
141
-
142
- # cutlass_scaled_mm supports per tensor/channel W and per tensor/token A
143
- if cutlass_fp8_supported:
144
- qinput, x_scale = ops.scaled_fp8_quant(
145
- input_2d,
146
- input_scale,
147
- scale_ub=input_scale_ub,
148
- use_per_token_if_dynamic=use_per_token_if_dynamic)
149
-
150
- # Fused GEMM_DQ
151
- output = ops.cutlass_scaled_mm(qinput,
152
- weight,
153
- out_dtype=input.dtype,
154
- scale_a=x_scale,
155
- scale_b=weight_scale,
156
- bias=bias)
157
- return output.view(*output_shape)
158
-
159
- # torch.scaled_mm supports per tensor weights + activations only
160
- # so fallback to naive if per channel or per token
161
- else:
124
+# TODO(luka): follow similar pattern for marlin and block-fp8-linear
125
+# https://github.com/vllm-project/vllm/issues/14397
126
+class Fp8LinearOp:
127
+ """
128
+ This class executes a FP8 linear layer using cutlass if supported and
129
+ torch.scaled_mm otherwise.
130
+ It needs to be a class instead of a method so that config can be read
131
+ in the __init__ method, as reading config is not allowed inside forward.
132
+ """
133
+
134
+ def __init__(self,
135
+ cutlass_fp8_supported: bool = cutlass_fp8_supported(),
136
+ use_per_token_if_dynamic: bool = False,
137
+ pad_output: Optional[bool] = None):
138
+ self.cutlass_fp8_supported = cutlass_fp8_supported
139
+ self.use_per_token_if_dynamic = use_per_token_if_dynamic
140
+
141
# Note: we pad the input because torch._scaled_mm is more performant
142
# for matrices with batch dimension > 16.
143
# This could change in the future.
144
# We also don't pad when using torch.compile,
145
# as it breaks with dynamic shapes.
167
- config = get_current_vllm_config().compilation_config
168
- do_pad = config.level < CompilationLevel.PIECEWISE
169
- qinput, x_scale = ops.scaled_fp8_quant(
170
- input_2d,
171
- input_scale,
172
- num_token_padding=17 if do_pad else None,
173
- use_per_token_if_dynamic=use_per_token_if_dynamic)
174
-
175
- per_tensor_weights = (weight_scale.numel() == 1)
176
- per_tensor_activations = (x_scale.numel() == 1)
177
-
178
- if per_tensor_weights and per_tensor_activations:
179
- # Fused GEMM_DQ
180
- output = torch._scaled_mm(qinput,
181
- weight,
182
- out_dtype=input.dtype,
183
- scale_a=x_scale,
184
- scale_b=weight_scale,
185
- bias=bias)
186
- # A fix for discrepancy in scaled_mm which returns tuple
187
- # for torch < 2.5 and a single value in torch >= 2.5
188
- if type(output) is tuple and len(output) == 2:
189
- output = output[0]
190
-
191
- return torch.narrow(output, 0, 0,
192
- input_2d.shape[0]).view(*output_shape)
193
-
194
- elif (use_per_token_if_dynamic and not per_tensor_weights
195
- and not per_tensor_activations and USE_ROWWISE_TORCH_SCALED_MM):
196
- # For now validated on ROCm platform
197
- # fp8 rowwise scaling in torch._scaled_mm is introduced in
198
- # https://github.com/pytorch/pytorch/pull/144432 using
199
- # hipBLASLt and ROCm 6.3, which only exists in torch 2.7 and above.
200
- # For CUDA platform please validate if the
201
- # torch._scaled_mm support rowwise scaled GEMM
202
- # Fused GEMM_DQ Rowwise GEMM
203
- output = torch._scaled_mm(qinput,
204
- weight,
205
- out_dtype=input.dtype,
206
- scale_a=x_scale,
207
- scale_b=weight_scale.t(),
208
- bias=bias)
209
-
210
- output = torch.narrow(output, 0, 0, input_2d.shape[0])
211
- output = output.view(*output_shape)
212
- return output
146
+ if pad_output is None:
147
+ config = get_current_vllm_config().compilation_config
148
+ pad_output = config.level < CompilationLevel.PIECEWISE
149
+ self.output_padding = 17 if pad_output else None
150
+
151
+ def apply(
152
+ self,
153
+ input: torch.Tensor,
154
+ weight: torch.Tensor,
155
+ weight_scale: torch.Tensor,
156
+ input_scale: Optional[torch.Tensor] = None,
157
+ input_scale_ub: Optional[torch.Tensor] = None,
158
+ bias: Optional[torch.Tensor] = None,
159
+ # TODO(luka) remove this parameter in favor of __init__
160
+ use_per_token_if_dynamic: Optional[bool] = None
161
+ ) -> torch.Tensor:
162
+ # ops.scaled_fp8_quant supports both dynamic and static quant.
163
+ # If dynamic, layer.input_scale is None and x_scale computed from x.
164
+ # If static, layer.input_scale is scalar and x_scale is input_scale.
165
+
166
+ # View input as 2D matrix for fp8 methods
167
+ input_2d = input.view(-1, input.shape[-1])
168
+ output_shape = [*input.shape[:-1], weight.shape[1]]
169
+
170
+ # TODO(luka) this is here because currently MLA only decides this
171
+ # during the forward method instead of in __init__.
172
+ if use_per_token_if_dynamic is None:
173
+ use_per_token_if_dynamic = self.use_per_token_if_dynamic
174
+
175
+ # cutlass_scaled_mm supports per tensor/channel W and per tensor/token A
176
+ if self.cutlass_fp8_supported:
177
+ qinput, x_scale = ops.scaled_fp8_quant(
178
+ input_2d,
179
+ input_scale,
180
+ scale_ub=input_scale_ub,
181
+ use_per_token_if_dynamic=use_per_token_if_dynamic)
182
183
+ # Fused GEMM_DQ
184
+ output = ops.cutlass_scaled_mm(qinput,
185
+ weight,
186
+ out_dtype=input.dtype,
187
+ scale_a=x_scale,
188
+ scale_b=weight_scale,
189
+ bias=bias)
190
+ return output.view(*output_shape)
191
+
192
+ # torch.scaled_mm supports per tensor weights + activations only
193
+ # so fallback to naive if per channel or per token
194
else:
215
- # Fallback for channelwise case, where we use unfused DQ
216
- # due to limitations with scaled_mm
217
-
218
- # Symmetric quantized GEMM by definition computes the following:
219
- # C = (s_x * X) (s_w * W) + bias
220
- # This is equivalent to dequantizing the weights and activations
221
- # before applying a GEMM.
222
- #
223
- # In order to compute quantized operands, a quantized kernel
224
- # will rewrite the above like so:
225
- # C = s_w * s_x * (X * W) + bias
226
- #
227
- # For the scaled_mm fallback case, we break this down, since it
228
- # does not support s_w being a vector.
229
-
230
- # GEMM
231
- # This computes C = (X * W).
232
- # Output in fp32 to allow subsequent ops to happen in-place
233
- output = torch._scaled_mm(qinput,
234
- weight,
235
- scale_a=TORCH_DEVICE_IDENTITY,
236
- scale_b=TORCH_DEVICE_IDENTITY,
237
- out_dtype=torch.float32)
238
- # A fix for discrepancy in scaled_mm which returns tuple
239
- # for torch < 2.5 and a single value in torch >= 2.5
240
- if type(output) is tuple and len(output) == 2:
241
- output = output[0]
242
- # Unpad (undo num_token_padding)
243
- output = torch.narrow(output, 0, 0, input_2d.shape[0])
244
- x_scale = torch.narrow(x_scale, 0, 0, input_2d.shape[0])
245
-
246
- # DQ
247
- # C = sw * sx * (X * W) + bias
248
- output = output * x_scale * weight_scale.t()
249
- if bias is not None:
250
- output = output + bias
251
- return output.to(dtype=input.dtype).view(*output_shape)
195
+ # Maybe apply padding to output, see comment in __init__
196
+ qinput, x_scale = ops.scaled_fp8_quant(
197
+ input_2d,
198
+ input_scale,
199
+ num_token_padding=self.output_padding,
200
+ use_per_token_if_dynamic=use_per_token_if_dynamic)
201
+
202
+ per_tensor_weights = (weight_scale.numel() == 1)
203
+ per_tensor_activations = (x_scale.numel() == 1)
204
+
205
+ if per_tensor_weights and per_tensor_activations:
206
+ # Fused GEMM_DQ
207
+ output = torch._scaled_mm(qinput,
208
+ weight,
209
+ out_dtype=input.dtype,
210
+ scale_a=x_scale,
211
+ scale_b=weight_scale,
212
+ bias=bias)
213
+ # A fix for discrepancy in scaled_mm which returns tuple
214
+ # for torch < 2.5 and a single value in torch >= 2.5
215
+ if type(output) is tuple and len(output) == 2:
216
+ output = output[0]
217
+
218
+ return torch.narrow(output, 0, 0,
219
+ input_2d.shape[0]).view(*output_shape)
220
+
221
+ elif (use_per_token_if_dynamic and not per_tensor_weights
222
+ and not per_tensor_activations
223
+ and USE_ROWWISE_TORCH_SCALED_MM):
224
+ # For now validated on ROCm platform
225
+ # fp8 rowwise scaling in torch._scaled_mm is introduced in
226
+ # https://github.com/pytorch/pytorch/pull/144432 using hipBLASLt
227
+ # and ROCm 6.3, which only exists in torch 2.7 and above.
228
+ # For CUDA platform please validate if the
229
+ # torch._scaled_mm support rowwise scaled GEMM
230
+ # Fused GEMM_DQ Rowwise GEMM
231
+ output = torch._scaled_mm(qinput,
232
+ weight,
233
+ out_dtype=input.dtype,
234
+ scale_a=x_scale,
235
+ scale_b=weight_scale.t(),
236
+ bias=bias)
237
+
238
+ output = torch.narrow(output, 0, 0, input_2d.shape[0])
239
+ output = output.view(*output_shape)
240
+ return output
241
+
242
+ else:
243
+ # Fallback for channelwise case, where we use unfused DQ
244
+ # due to limitations with scaled_mm
245
+
246
+ # Symmetric quantized GEMM by definition computes the following:
247
+ # C = (s_x * X) (s_w * W) + bias
248
+ # This is equivalent to dequantizing the weights and activations
249
+ # before applying a GEMM.
250
+ #
251
+ # In order to compute quantized operands, a quantized kernel
252
+ # will rewrite the above like so:
253
+ # C = s_w * s_x * (X * W) + bias
254
+ #
255
+ # For the scaled_mm fallback case, we break this down, since it
256
+ # does not support s_w being a vector.
257
+
258
+ # GEMM
259
+ # This computes C = (X * W).
260
+ # Output in fp32 to allow subsequent ops to happen in-place
261
+ output = torch._scaled_mm(qinput,
262
+ weight,
263
+ scale_a=TORCH_DEVICE_IDENTITY,
264
+ scale_b=TORCH_DEVICE_IDENTITY,
265
+ out_dtype=torch.float32)
266
+ # A fix for discrepancy in scaled_mm which returns tuple
267
+ # for torch < 2.5 and a single value in torch >= 2.5
268
+ if type(output) is tuple and len(output) == 2:
269
+ output = output[0]
270
+ # Unpad (undo num_token_padding)
271
+ output = torch.narrow(output, 0, 0, input_2d.shape[0])
272
+ x_scale = torch.narrow(x_scale, 0, 0, input_2d.shape[0])
273
+
274
+ # DQ
275
+ # C = sw * sx * (X * W) + bias
276
+ output = output * x_scale * weight_scale.t()
277
+ if bias is not None:
278
+ output = output + bias
279
+ return output.to(dtype=input.dtype).view(*output_shape)
280
281
282
def normalize_e4m3fn_to_e4m3fnuz(
vllm/v1/attention/backends/mla/common.py
+4
-3
@@ -219,7 +219,7 @@ from vllm.model_executor.layers.quantization.compressed_tensors.schemes import (
219
CompressedTensorsW8A8Fp8)
220
from vllm.model_executor.layers.quantization.fp8 import Fp8LinearMethod
221
from vllm.model_executor.layers.quantization.utils.fp8_utils import (
222
- apply_fp8_linear_generic, current_platform_fp8_dtype, is_fp8)
222
+ Fp8LinearGenericOp, current_platform_fp8_dtype, is_fp8)
223
from vllm.model_executor.layers.quantization.utils.quant_utils import (
224
scaled_quantize)
225
from vllm.model_executor.layers.rotary_embedding import RotaryEmbedding
@@ -640,6 +640,7 @@ class MLACommonImpl(MLAAttentionImpl[M], Generic[M]):
640
self.kv_b_proj = kv_b_proj
641
self.o_proj = o_proj
642
self.vllm_flash_attn_version = get_flash_attn_version()
643
+ self.fp8_linear_generic = Fp8LinearGenericOp()
644
645
# Handle the differences between the flash_attn_varlen from flash_attn
646
# and the one from vllm_flash_attn. The former is used on RoCM and the
@@ -653,7 +654,7 @@ class MLACommonImpl(MLAAttentionImpl[M], Generic[M]):
654
def _v_up_proj_and_o_proj(self, x):
655
if envs.VLLM_MLA_PERFORM_MATRIX_ABSORPTION:
656
if is_fp8(self.W_UV_O):
656
- output_parallel = apply_fp8_linear_generic(
657
+ output_parallel = self.fp8_linear_generic.apply(
658
x.flatten(start_dim=1), self.W_UV_O, self.W_UV_O_scales,
659
self.reqaunt_input_group_shape,
660
self.reqaunt_weight_group_shape)
@@ -673,7 +674,7 @@ class MLACommonImpl(MLAAttentionImpl[M], Generic[M]):
674
def _q_proj_and_k_up_proj(self, x):
675
if envs.VLLM_MLA_PERFORM_MATRIX_ABSORPTION:
676
if is_fp8(self.W_Q_UK):
676
- return apply_fp8_linear_generic(
677
+ return self.fp8_linear_generic.apply(
678
x, self.W_Q_UK, self.W_Q_UK_scales,
679
self.reqaunt_input_group_shape,
680
self.reqaunt_weight_group_shape).view(