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, 03:45 UTC
[W8A8 Block Linear Refactor][1/N] Keep all quantization types into `QuantFP8` class. (#33047)
Signed-off-by: maral <maralbahari.98@gmail.com> Co-authored-by: Cyrus Leung <tlleungac@connect.ust.hk>
Maral committed
Feb 1, 2026 at 17:28 UTC
b5f8c3092d1e1466b2b9c516fb39e5b2c15e774b
3 files changed
+51
-33
vllm/_aiter_ops.py
+1
-1
@@ -1570,7 +1570,7 @@ class rocm_aiter_ops:
1570
def group_fp8_quant(
1571
input_2d: torch.Tensor,
1572
group_size: int = 128,
1573
- ) -> tuple[torch.Tensor, ...]:
1573
+ ) -> tuple[torch.Tensor, torch.Tensor]:
1574
assert group_size == 128, "Group size must be 128"
1575
return torch.ops.vllm.rocm_aiter_group_fp8_quant(input_2d, group_size)
1576
vllm/model_executor/layers/quantization/input_quant_fp8.py
+40
-12
@@ -14,6 +14,11 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import (
14
prep_scale_for_group_broadcast,
15
)
16
from vllm.platforms import current_platform
17
+from vllm.utils.deep_gemm import (
18
+ DeepGemmQuantScaleFMT,
19
+ is_deep_gemm_e8m0_used,
20
+ is_deep_gemm_supported,
21
+)
22
23
_FP8_DTYPE = current_platform.fp8_dtype()
24
_FP8_MIN, _FP8_MAX = get_fp8_min_max()
@@ -59,7 +64,8 @@ class QuantFP8(CustomOp):
64
self.num_token_padding = num_token_padding
65
self.column_major_scales = column_major_scales
66
self.tma_aligned_scales = tma_aligned_scales
62
- self.use_ue8m0 = use_ue8m0
67
+ self.use_ue8m0 = is_deep_gemm_e8m0_used() if use_ue8m0 is None else use_ue8m0
68
+ self.use_deep_gemm_supported = is_deep_gemm_supported()
69
70
self.use_aiter = rocm_aiter_ops.is_linear_fp8_enabled()
71
@@ -79,10 +85,23 @@ class QuantFP8(CustomOp):
85
x: torch.Tensor,
86
scale: torch.Tensor | None = None,
87
scale_ub: torch.Tensor | None = None,
88
+ **kwargs,
89
) -> tuple[torch.Tensor, torch.Tensor]:
90
+ from vllm.model_executor.layers.quantization.utils import fp8_utils
91
+
92
+ if (
93
+ self.is_group_quant
94
+ and self.use_deep_gemm_supported
95
+ and (DeepGemmQuantScaleFMT.from_oracle() == DeepGemmQuantScaleFMT.UE8M0)
96
+ ):
97
+ return fp8_utils.per_token_group_quant_fp8_packed_for_deepgemm(
98
+ x,
99
+ group_size=self.group_size,
100
+ use_ue8m0=True,
101
+ )
102
+
103
if self.is_group_quant and not self.static:
104
assert scale is None, "Dynamic group quantization does not use scale"
85
- from vllm.model_executor.layers.quantization.utils import fp8_utils
105
106
return fp8_utils.per_token_group_quant_fp8(
107
x,
@@ -116,25 +135,34 @@ class QuantFP8(CustomOp):
135
x: torch.Tensor,
136
scale: torch.Tensor | None = None,
137
scale_ub: torch.Tensor | None = None,
138
+ **kwargs,
139
) -> tuple[torch.Tensor, torch.Tensor]:
120
- use_aiter_quant = (
121
- not self.is_group_quant
122
- and self.use_aiter
123
- and scale_ub is None
124
- and x.is_contiguous()
125
- )
140
+ use_triton = kwargs.get("use_triton", False)
141
+ if self.is_group_quant and use_triton:
142
+ assert scale is None, "Dynamic group quantization does not use scale"
143
+
144
+ return torch.ops.vllm.triton_per_token_group_quant_fp8(x, self.group_size)
145
+
146
+ use_aiter_quant = self.use_aiter and scale_ub is None and x.is_contiguous()
147
use_aiter_per_tensor_quant = (
127
- use_aiter_quant and self.group_shape == GroupShape.PER_TENSOR
128
- )
129
- use_aiter_per_token_quant = (
130
- use_aiter_quant and self.group_shape == GroupShape.PER_TOKEN
148
+ use_aiter_quant and self.group_shape.is_per_tensor()
149
)
150
+ use_aiter_per_token_quant = use_aiter_quant and self.group_shape.is_per_token()
151
152
+ use_aiter_per_group_quant = use_aiter_quant and self.group_shape.is_per_group()
153
+
154
+ if use_aiter_per_group_quant:
155
+ return rocm_aiter_ops.group_fp8_quant(x, self.group_size)
156
if use_aiter_per_tensor_quant:
157
return rocm_aiter_ops.per_tensor_quant(x, _FP8_DTYPE, scale)
158
if use_aiter_per_token_quant:
159
return rocm_aiter_ops.per_token_quant(x, _FP8_DTYPE, scale)
160
161
+ # Fallback to native implementation for group quantization.
162
+ if self.is_group_quant:
163
+ assert scale is None, "Dynamic group quantization does not use scale"
164
+ return self._quantize_group_native(x)
165
+
166
# Fallback to CUDA implementation
167
return self.forward_cuda(x, scale, scale_ub)
168
vllm/model_executor/layers/quantization/utils/fp8_utils.py
+10
-20
@@ -33,7 +33,6 @@ from vllm.model_executor.utils import replace_parameter, set_weight_attrs
33
from vllm.platforms import current_platform
34
from vllm.triton_utils import tl, triton
35
from vllm.utils.deep_gemm import (
36
- DeepGemmQuantScaleFMT,
36
fp8_gemm_nt,
37
get_tma_aligned_size,
38
is_deep_gemm_e8m0_used,
@@ -426,15 +425,8 @@ class W8A8BlockFp8LinearOp:
425
weight: torch.Tensor,
426
weight_scale: torch.Tensor,
427
) -> torch.Tensor:
429
- if DeepGemmQuantScaleFMT.from_oracle() == DeepGemmQuantScaleFMT.UE8M0:
430
- q_input, input_scale = per_token_group_quant_fp8_packed_for_deepgemm(
431
- input_2d,
432
- group_size=self.act_quant_group_shape.col,
433
- use_ue8m0=True,
434
- )
435
- else:
436
- assert self.deepgemm_input_quant_op is not None
437
- q_input, input_scale = self.deepgemm_input_quant_op(input_2d)
428
+ assert self.deepgemm_input_quant_op is not None
429
+ q_input, input_scale = self.deepgemm_input_quant_op(input_2d)
430
output = torch.empty(
431
(q_input.shape[0], weight.shape[0]),
432
dtype=torch.bfloat16,
@@ -497,15 +489,8 @@ class W8A8BlockFp8LinearOp:
489
490
if input_scale is not None:
491
q_input = input_2d
500
- elif use_triton:
501
- q_input, input_scale = torch.ops.vllm.triton_per_token_group_quant_fp8(
502
- input_2d,
503
- self.act_quant_group_shape.col,
504
- )
492
else:
506
- q_input, input_scale = rocm_aiter_ops.group_fp8_quant(
507
- input_2d, self.act_quant_group_shape.col
508
- )
493
+ q_input, input_scale = self.input_quant_op(input_2d, use_triton=use_triton)
494
495
return gemm_a8w8_blockscale_op(
496
q_input,
@@ -572,7 +557,7 @@ class W8A8BlockFp8LinearOp:
557
],
558
torch.Tensor,
559
],
575
- QuantFP8 | None,
560
+ QuantFP8,
561
]:
562
if use_cutlass:
563
return self._run_cutlass, (
@@ -584,7 +569,12 @@ class W8A8BlockFp8LinearOp:
569
)
570
)
571
if use_aiter_and_is_supported:
587
- return self._run_aiter, None
572
+ return self._run_aiter, QuantFP8(
573
+ False,
574
+ self.act_quant_group_shape,
575
+ column_major_scales=False,
576
+ use_ue8m0=False,
577
+ )
578
return self._run_triton, (
579
QuantFP8(
580
False,