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 14, 05:30 UTC
Revert "[Performance] Move apply_w8a8_block_fp8_linear to an op class… (#25607)
Signed-off-by: Tyler Michael Smith <tlrmchlsmth@gmail.com>
Tyler Michael Smith committed
Sep 25, 2025 at 04:05 UTC
1260180c678b1ca0a5863b907d9afadcb72399f1
14 files changed
+205
-346
benchmarks/cutlass_benchmarks/w8a8_benchmarks.py
+2
-2
@@ -17,7 +17,7 @@ from weight_shapes import WEIGHT_SHAPES
17
18
from vllm import _custom_ops as ops
19
from vllm.model_executor.layers.quantization.utils.fp8_utils import (
20
- w8a8_triton_block_scaled_mm,
20
+ w8a8_block_fp8_matmul,
21
)
22
from vllm.utils import FlexibleArgumentParser, cdiv
23
@@ -158,7 +158,7 @@ def bench_fp8(
158
"cutlass_fp8_fp8_fp16_scaled_mm_bias": lambda: ops.cutlass_scaled_mm(
159
a, b, scale_a, scale_b, torch.float16, bias.to(dtype=torch.float16)
160
),
161
- "triton_fp8_fp8_fp16_scaled_mm_blockwise": lambda: w8a8_triton_block_scaled_mm(
161
+ "triton_fp8_fp8_fp16_scaled_mm_blockwise": lambda: w8a8_block_fp8_matmul(
162
a_cont, b.t(), block_scale_a, block_scale_b.t(), (128, 128)
163
),
164
"cutlass_fp8_fp8_fp16_scaled_mm_blockwise": lambda: ops.cutlass_scaled_mm(
benchmarks/kernels/deepgemm/benchmark_fp8_block_dense_gemm.py
+2
-2
@@ -9,7 +9,7 @@ import torch
9
from vllm import _custom_ops as ops
10
from vllm.model_executor.layers.quantization.utils.fp8_utils import (
11
per_token_group_quant_fp8,
12
- w8a8_triton_block_scaled_mm,
12
+ w8a8_block_fp8_matmul,
13
)
14
from vllm.triton_utils import triton
15
from vllm.utils.deep_gemm import (
@@ -63,7 +63,7 @@ def benchmark_shape(m: int,
63
64
# === vLLM Triton Implementation ===
65
def vllm_triton_gemm():
66
- return w8a8_triton_block_scaled_mm(A_vllm,
66
+ return w8a8_block_fp8_matmul(A_vllm,
67
B_vllm,
68
A_scale_vllm,
69
B_scale_vllm,
tests/kernels/quantization/test_block_fp8.py
+2
-3
@@ -11,7 +11,7 @@ from tests.kernels.quant_utils import (native_per_token_group_quant_fp8,
11
native_w8a8_block_matmul)
12
from vllm.config import VllmConfig
13
from vllm.model_executor.layers.quantization.utils.fp8_utils import (
14
- cutlass_scaled_mm, per_token_group_quant_fp8, w8a8_triton_block_scaled_mm)
14
+ cutlass_scaled_mm, per_token_group_quant_fp8, w8a8_block_fp8_matmul)
15
from vllm.platforms import current_platform
16
from vllm.utils import has_deep_gemm
17
from vllm.utils.deep_gemm import (fp8_gemm_nt,
@@ -91,8 +91,7 @@ def test_w8a8_block_fp8_matmul(M, N, K, block_size, out_dtype, seed):
91
92
ref_out = native_w8a8_block_matmul(A_fp8, B_fp8, As, Bs, block_size,
93
out_dtype)
94
- out = w8a8_triton_block_scaled_mm(A_fp8, B_fp8, As, Bs, block_size,
95
- out_dtype)
94
+ out = w8a8_block_fp8_matmul(A_fp8, B_fp8, As, Bs, block_size, out_dtype)
95
96
rel_diff = (torch.mean(
97
torch.abs(out.to(torch.float32) - ref_out.to(torch.float32))) /
tests/kernels/quantization/test_fp8_quant_group.py
+7
-19
@@ -20,11 +20,9 @@ from vllm.platforms import current_platform
20
(8, 513, 64), # Non-divisible (native only)
21
])
22
@pytest.mark.parametrize("seed", [42])
23
-@pytest.mark.parametrize("use_ue8m0", [True, False])
23
@torch.inference_mode()
24
def test_quantfp8_group_functionality(batch_size: int, hidden_dim: int,
26
- group_size: int, seed: int,
27
- use_ue8m0: bool) -> None:
25
+ group_size: int, seed: int) -> None:
26
"""Test QuantFP8 group quantization with various configurations.
27
28
Tests both CUDA and native implementations, column-major scales,
@@ -40,8 +38,7 @@ def test_quantfp8_group_functionality(batch_size: int, hidden_dim: int,
38
group_shape = GroupShape(1, group_size)
39
quant_op = QuantFP8(static=False,
40
group_shape=group_shape,
43
- column_major_scales=False,
44
- use_ue8m0=use_ue8m0)
41
+ column_major_scales=False)
42
43
# 1. Test native implementation (always available)
44
x_quant_native, scales_native = quant_op.forward_native(x.clone())
@@ -51,15 +48,9 @@ def test_quantfp8_group_functionality(batch_size: int, hidden_dim: int,
48
# 2. Test column-major scales configuration
49
quant_op_col = QuantFP8(static=False,
50
group_shape=group_shape,
54
- column_major_scales=True,
55
- use_ue8m0=use_ue8m0)
51
+ column_major_scales=True)
52
_, scales_col = quant_op_col.forward_native(x.clone())
57
- assert scales_col.shape == (batch_size, expected_num_groups)
58
- assert scales_col.stride(0) == 1
59
- assert scales_col.stride(1) == batch_size
60
-
61
- # Test column-major scales consistency
62
- assert torch.allclose(scales_col, scales_native, rtol=1e-9, atol=1e-8)
53
+ assert scales_col.shape == (expected_num_groups, batch_size)
54
55
# 3. Test CUDA implementation (only for divisible dimensions)
56
if is_divisible:
@@ -77,9 +68,8 @@ def test_quantfp8_group_functionality(batch_size: int, hidden_dim: int,
68
69
70
@pytest.mark.parametrize("seed", [42])
80
-@pytest.mark.parametrize("use_ue8m0", [True, False])
71
@torch.inference_mode()
82
-def test_quantfp8_group_multidimensional(seed: int, use_ue8m0: bool) -> None:
72
+def test_quantfp8_group_multidimensional(seed: int) -> None:
73
current_platform.seed_everything(seed)
74
75
group_size = 64
@@ -92,8 +82,7 @@ def test_quantfp8_group_multidimensional(seed: int, use_ue8m0: bool) -> None:
82
group_shape = GroupShape(1, group_size)
83
quant_op = QuantFP8(static=False,
84
group_shape=group_shape,
95
- column_major_scales=False,
96
- use_ue8m0=use_ue8m0)
85
+ column_major_scales=False)
86
87
x_quant, scales = quant_op.forward_native(x_3d.clone())
88
assert x_quant.shape == x_3d.shape
@@ -102,8 +91,7 @@ def test_quantfp8_group_multidimensional(seed: int, use_ue8m0: bool) -> None:
91
# Test column_major_scales with multi-dim
92
quant_op_col = QuantFP8(static=False,
93
group_shape=group_shape,
105
- column_major_scales=True,
106
- use_ue8m0=use_ue8m0)
94
+ column_major_scales=True)
95
_, scales_col = quant_op_col.forward_native(x_3d.clone())
96
assert scales_col.shape == (batch1, hidden_dim // group_size, batch2)
97
tests/model_executor/test_enabled_custom_ops.py
+30
@@ -17,6 +17,8 @@ from vllm.model_executor.layers.fused_moe.rocm_aiter_fused_moe import (
17
from vllm.model_executor.layers.layernorm import (RMSNorm,
18
dispatch_rocm_rmsnorm_func,
19
fused_add_rms_norm, rms_norm)
20
+from vllm.model_executor.layers.quantization.utils.fp8_utils import (
21
+ cutlass_scaled_mm, dispatch_w8a8_blockscale_func, w8a8_block_fp8_matmul)
22
from vllm.platforms import current_platform
23
24
RMS_NORM_SUPPORTED_DTYPES = [torch.float16, torch.bfloat16]
@@ -109,6 +111,34 @@ def test_enabled_ops_invalid(env: str):
111
RMSNorm(1024).enabled()
112
113
114
+@pytest.mark.skipif(
115
+ not current_platform.is_rocm() or not current_platform.is_fp8_fnuz(),
116
+ reason="AITER is a feature exclusive for ROCm and FP8_FNUZ")
117
+@pytest.mark.parametrize("use_cutlass", [True, False])
118
+@pytest.mark.parametrize("use_rocm_aiter", ["0", "1"])
119
+@pytest.mark.parametrize("use_rocm_aiter_gemm_w8a8_blockscale", ["0", "1"])
120
+def test_w8a8_blockscale_dispatch(use_cutlass: bool, use_rocm_aiter: str,
121
+ use_rocm_aiter_gemm_w8a8_blockscale: str,
122
+ monkeypatch):
123
+
124
+ monkeypatch.setenv("VLLM_ROCM_USE_AITER", use_rocm_aiter)
125
+ monkeypatch.setenv("VLLM_ROCM_USE_AITER_LINEAR",
126
+ use_rocm_aiter_gemm_w8a8_blockscale)
127
+
128
+ use_aiter_and_is_supported = (bool(int(use_rocm_aiter)) and bool(
129
+ int(use_rocm_aiter_gemm_w8a8_blockscale)))
130
+ block_scale_func = dispatch_w8a8_blockscale_func(
131
+ use_cutlass, use_aiter_and_is_supported=use_aiter_and_is_supported)
132
+ if use_cutlass:
133
+ assert block_scale_func == cutlass_scaled_mm
134
+ elif current_platform.is_rocm() and int(use_rocm_aiter) and int(
135
+ use_rocm_aiter_gemm_w8a8_blockscale):
136
+ assert block_scale_func == (
137
+ torch.ops.vllm.rocm_aiter_gemm_w8a8_blockscale)
138
+ else:
139
+ assert block_scale_func == w8a8_block_fp8_matmul
140
+
141
+
142
@pytest.mark.parametrize("use_rocm_aiter", ["0", "1"])
143
def test_topk_dispatch(use_rocm_aiter: str, monkeypatch):
144
monkeypatch.setenv("VLLM_ROCM_USE_AITER", use_rocm_aiter)
tests/quantization/test_compressed_tensors.py
-35
@@ -18,9 +18,6 @@ from vllm.model_executor.layers.quantization.compressed_tensors.compressed_tenso
18
CompressedTensorsW4A16Fp4, CompressedTensorsW4A16Sparse24,
19
CompressedTensorsW8A8Fp8, CompressedTensorsW8A8Int8,
20
CompressedTensorsW8A16Fp8, CompressedTensorsWNA16)
21
-from vllm.model_executor.layers.quantization.input_quant_fp8 import QuantFP8
22
-from vllm.model_executor.layers.quantization.utils.fp8_utils import (
23
- W8A8BlockFp8LinearOp)
21
from vllm.model_executor.layers.quantization.utils.quant_utils import (
22
cutlass_fp4_supported)
23
from vllm.model_executor.layers.quantization.utils.w8a8_utils import (
@@ -745,35 +742,3 @@ def test_compressed_tensors_transforms_perplexity(vllm_runner, model, prompt,
742
perplexity = llm.generate_prompt_perplexity([prompt])[0]
743
print(perplexity)
744
assert perplexity <= exp_perplexity
748
-
749
-
750
-def test_compressed_tensors_fp8_block_enabled(vllm_runner):
751
- model_path = "RedHatAI/Qwen3-0.6B-FP8-BLOCK"
752
- with vllm_runner(model_path) as llm:
753
-
754
- fp8_dtype = current_platform.fp8_dtype()
755
-
756
- def check_model(model):
757
- layer = model.model.layers[0]
758
-
759
- qkv_proj = layer.self_attn.qkv_proj
760
- assert isinstance(qkv_proj.quant_method,
761
- CompressedTensorsLinearMethod)
762
- assert isinstance(qkv_proj.scheme, CompressedTensorsW8A8Fp8)
763
- assert isinstance(qkv_proj.scheme.w8a8_block_fp8_linear,
764
- W8A8BlockFp8LinearOp)
765
-
766
- assert qkv_proj.weight.dtype is fp8_dtype
767
- assert qkv_proj.weight_scale.dtype is torch.float32
768
- assert len(qkv_proj.weight.shape) == 2
769
- assert len(qkv_proj.weight_scale.shape) == 2
770
-
771
- input_quant_op = \
772
- qkv_proj.scheme.w8a8_block_fp8_linear.input_quant_op
773
- assert isinstance(input_quant_op, QuantFP8)
774
- assert input_quant_op._forward_method == input_quant_op.forward_cuda
775
-
776
- llm.apply_model(check_model)
777
-
778
- output = llm.generate_greedy("Hello my name is", max_tokens=20)
779
- assert output
vllm/config/__init__.py
-17
@@ -545,23 +545,6 @@ class VllmConfig:
545
# local attention.
546
self.scheduler_config.disable_hybrid_kv_cache_manager = True
547
548
- def has_blocked_weights():
549
- if self.quant_config is not None:
550
- if hasattr(self.quant_config, "weight_block_size"):
551
- return self.quant_config.weight_block_size is not None
552
- elif hasattr(self.quant_config, "has_blocked_weights"):
553
- return self.quant_config.has_blocked_weights()
554
- return False
555
-
556
- # Enable quant_fp8 CUDA ops (TODO disable in follow up)
557
- # On H100 the CUDA kernel is faster than
558
- # native implementation
559
- # https://github.com/vllm-project/vllm/issues/25094
560
- if has_blocked_weights():
561
- custom_ops = self.compilation_config.custom_ops
562
- if "none" not in custom_ops and "-quant_fp8" not in custom_ops:
563
- custom_ops.append("+quant_fp8")
564
-
548
def update_sizes_for_sequence_parallelism(self,
549
possible_sizes: list) -> list:
550
# remove the sizes that not multiple of tp_size when
vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors.py
-8
@@ -644,14 +644,6 @@ class CompressedTensorsConfig(QuantizationConfig):
644
# If no matches, return None
645
return None
646
647
- def has_blocked_weights(self) -> bool:
648
- for scheme in self.target_scheme_map.values():
649
- weight_quant = scheme.get("weights")
650
- if (weight_quant is not None
651
- and weight_quant.strategy == QuantizationStrategy.BLOCK):
652
- return True
653
- return False
654
-
647
@staticmethod
648
def supports_cutlass_24(
649
weight_quant: Optional[QuantizationArgs],
vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_w8a8_fp8.py
+11
-26
@@ -11,7 +11,7 @@ from torch.nn import Parameter
11
from vllm.model_executor.layers.quantization.compressed_tensors.schemes import (
12
CompressedTensorsScheme)
13
from vllm.model_executor.layers.quantization.utils.fp8_utils import (
14
- W8A8BlockFp8LinearOp, check_aiter_fp8_linear_support,
14
+ apply_fp8_block_linear, check_aiter_fp8_linear_support,
15
create_fp8_input_scale, create_fp8_scale_parameter,
16
create_fp8_weight_parameter, maybe_post_process_fp8_weight_block,
17
process_fp8_weight_block_strategy, process_fp8_weight_channel_strategy,
@@ -41,30 +41,16 @@ class CompressedTensorsW8A8Fp8(CompressedTensorsScheme):
41
self.strategy = weight_quant.strategy
42
self.out_dtype = torch.get_default_dtype()
43
self.is_static_input_scheme = is_static_input_scheme
44
+ self.act_q_group_shape = GroupShape.PER_TENSOR \
45
+ if is_static_input_scheme else GroupShape.PER_TOKEN
46
+ self.fp8_linear = Fp8LinearOp(
47
+ act_quant_static=self.is_static_input_scheme,
48
+ act_quant_group_shape=self.act_q_group_shape)
49
50
self.weight_block_size = self.weight_quant.block_structure
46
- if self.weight_block_size is not None:
47
- self.act_q_group_shape = GroupShape(1, self.weight_block_size[0])
48
- else:
49
- self.act_q_group_shape = GroupShape.PER_TENSOR \
50
- if is_static_input_scheme else GroupShape.PER_TOKEN
51
-
51
self.cutlass_block_fp8_supported = cutlass_block_fp8_supported()
52
self.use_aiter_and_is_supported = check_aiter_fp8_linear_support()
53
55
- if self.weight_block_size is not None:
56
- assert not self.is_static_input_scheme
57
- self.w8a8_block_fp8_linear = W8A8BlockFp8LinearOp(
58
- weight_group_shape=GroupShape(*self.weight_block_size),
59
- act_quant_group_shape=self.act_q_group_shape,
60
- cutlass_block_fp8_supported=self.cutlass_block_fp8_supported,
61
- use_aiter_and_is_supported=self.use_aiter_and_is_supported,
62
- )
63
- else:
64
- self.fp8_linear = Fp8LinearOp(
65
- act_quant_static=self.is_static_input_scheme,
66
- act_quant_group_shape=self.act_q_group_shape)
67
-
54
@classmethod
55
def get_min_capability(cls) -> int:
56
# lovelace and up
@@ -155,14 +141,13 @@ class CompressedTensorsW8A8Fp8(CompressedTensorsScheme):
141
x: torch.Tensor,
142
bias: Optional[torch.Tensor] = None) -> torch.Tensor:
143
158
- if self.weight_block_size is not None:
159
- return self.w8a8_block_fp8_linear.apply(
144
+ if layer.weight_block_size is not None:
145
+ return apply_fp8_block_linear(
146
+ layer,
147
input=x,
161
- weight=layer.weight,
162
- weight_scale=layer.weight_scale,
163
- input_scale=layer.input_scale,
148
bias=bias,
165
- )
149
+ cutlass_block_fp8_supported=self.cutlass_block_fp8_supported,
150
+ use_aiter_and_is_supported=self.use_aiter_and_is_supported)
151
152
return self.fp8_linear.apply(input=x,
153
weight=layer.weight,
vllm/model_executor/layers/quantization/deepgemm.py
+5
-5
@@ -42,7 +42,7 @@ def prepare_block_fp8_matmul_inputs(
42
return M, N, K, C
43
44
45
-def w8a8_deepgemm_block_scaled_mm(
45
+def w8a8_block_fp8_matmul_deepgemm(
46
A: torch.Tensor,
47
B: torch.Tensor,
48
As: torch.Tensor,
@@ -58,7 +58,7 @@ def w8a8_deepgemm_block_scaled_mm(
58
return C
59
60
61
-def w8a8_deepgemm_block_scaled_mm_fake(
61
+def w8a8_block_fp8_matmul_deepgemm_fake(
62
A: torch.Tensor,
63
B: torch.Tensor,
64
As: torch.Tensor,
@@ -72,7 +72,7 @@ def w8a8_deepgemm_block_scaled_mm_fake(
72
73
74
direct_register_custom_op(
75
- op_name="w8a8_deepgemm_block_scaled_mm",
76
- op_func=w8a8_deepgemm_block_scaled_mm,
77
- fake_impl=w8a8_deepgemm_block_scaled_mm_fake,
75
+ op_name="w8a8_block_fp8_matmul_deepgemm",
76
+ op_func=w8a8_block_fp8_matmul_deepgemm,
77
+ fake_impl=w8a8_block_fp8_matmul_deepgemm_fake,
78
)
vllm/model_executor/layers/quantization/fp8.py
+12
-28
@@ -33,7 +33,7 @@ from vllm.model_executor.layers.quantization.utils.flashinfer_utils import (
33
register_moe_scaling_factors, rotate_flashinfer_fp8_moe_weights,
34
select_cutlass_fp8_gemm_impl, swap_w13_to_w31)
35
from vllm.model_executor.layers.quantization.utils.fp8_utils import (
36
- W8A8BlockFp8LinearOp, check_aiter_fp8_linear_support,
36
+ apply_fp8_block_linear, check_aiter_fp8_linear_support,
37
create_fp8_input_scale, create_fp8_scale_parameter,
38
create_fp8_weight_parameter, expert_weight_is_col_major,
39
maybe_post_process_fp8_weight_block, process_fp8_weight_block_strategy,
@@ -242,28 +242,15 @@ class Fp8LinearMethod(LinearMethodBase):
242
self.weight_block_size = self.quant_config.weight_block_size
243
self.block_quant = self.weight_block_size is not None
244
self.act_q_static = self.quant_config.activation_scheme == "static"
245
- if self.weight_block_size:
246
- self.act_q_group_shape = GroupShape(1, self.weight_block_size[0])
245
+ # Use per-token quantization for better perf if dynamic and cutlass
246
+ if not self.act_q_static and cutlass_fp8_supported():
247
+ self.act_q_group_shape = GroupShape.PER_TOKEN
248
else:
248
- # Use per-token quantization for better perf if dynamic and cutlass
249
- if not self.act_q_static and cutlass_fp8_supported():
250
- self.act_q_group_shape = GroupShape.PER_TOKEN
251
- else:
252
- self.act_q_group_shape = GroupShape.PER_TENSOR
249
+ self.act_q_group_shape = GroupShape.PER_TENSOR
250
254
- if self.block_quant:
255
- assert not self.act_q_static
256
- assert self.weight_block_size is not None
257
- self.w8a8_block_fp8_linear = W8A8BlockFp8LinearOp(
258
- weight_group_shape=GroupShape(*self.weight_block_size),
259
- act_quant_group_shape=self.act_q_group_shape,
260
- cutlass_block_fp8_supported=self.cutlass_block_fp8_supported,
261
- use_aiter_and_is_supported=self.use_aiter_and_is_supported,
262
- )
263
- else:
264
- self.fp8_linear = Fp8LinearOp(
265
- act_quant_static=self.act_q_static,
266
- act_quant_group_shape=self.act_q_group_shape)
251
+ self.fp8_linear = Fp8LinearOp(
252
+ act_quant_static=self.act_q_static,
253
+ act_quant_group_shape=self.act_q_group_shape)
254
255
def create_weights(
256
self,
@@ -412,15 +399,12 @@ class Fp8LinearMethod(LinearMethodBase):
399
bias=bias)
400
401
if self.block_quant:
415
- assert self.weight_block_size is not None
416
-
417
- return self.w8a8_block_fp8_linear.apply(
402
+ return apply_fp8_block_linear(
403
+ layer,
404
input=x,
419
- weight=layer.weight,
420
- weight_scale=layer.weight_scale,
421
- input_scale=layer.input_scale,
405
bias=bias,
423
- )
406
+ cutlass_block_fp8_supported=self.cutlass_block_fp8_supported,
407
+ use_aiter_and_is_supported=self.use_aiter_and_is_supported)
408
409
return self.fp8_linear.apply(input=x,
410
weight=layer.weight,
vllm/model_executor/layers/quantization/input_quant_fp8.py
+8
-16
@@ -27,14 +27,11 @@ class QuantFP8(CustomOp):
27
This CustomOp supports both static and dynamic quantization.
28
"""
29
30
- def __init__(
31
- self,
32
- static: bool,
33
- group_shape: GroupShape,
34
- num_token_padding: Optional[int] = None,
35
- column_major_scales: bool = False,
36
- use_ue8m0: Optional[bool] = None, # for Torch compile
37
- ):
30
+ def __init__(self,
31
+ static: bool,
32
+ group_shape: GroupShape,
33
+ num_token_padding: Optional[int] = None,
34
+ column_major_scales: bool = False):
35
"""
36
:param static: static or dynamic quantization
37
:param group_shape: quantization group shape (PER_TOKEN, PER_TENSOR,
@@ -49,7 +46,6 @@ class QuantFP8(CustomOp):
46
self.group_shape = group_shape
47
self.num_token_padding = num_token_padding
48
self.column_major_scales = column_major_scales
52
- self.use_ue8m0 = use_ue8m0
49
50
self.is_group_quant = group_shape.is_per_group()
51
if self.is_group_quant:
@@ -74,8 +70,7 @@ class QuantFP8(CustomOp):
70
x,
71
group_size=self.group_size,
72
column_major_scales=self.column_major_scales,
77
- dtype=_FP8_DTYPE,
78
- use_ue8m0=self.use_ue8m0)
73
+ dtype=_FP8_DTYPE)
74
75
assert (scale is not None) == self.static
76
assert scale_ub is None or (not self.static and self.group_shape
@@ -142,10 +137,7 @@ class QuantFP8(CustomOp):
137
138
x_grouped = x.view(-1, num_groups, self.group_size)
139
absmax = x_grouped.abs().max(dim=-1, keepdim=True)[0].float()
145
- scales_raw = absmax / _FP8_MAX
146
- if self.use_ue8m0:
147
- scales_raw = torch.exp2(torch.ceil(torch.log2(scales_raw)))
148
- scales = (scales_raw).clamp(min=_FP8_MIN_SCALING_FACTOR)
140
+ scales = (absmax / _FP8_MAX).clamp(min=_FP8_MIN_SCALING_FACTOR)
141
142
x_scaled = x_grouped / scales
143
x_quant = x_scaled.clamp(_FP8_MIN, _FP8_MAX).to(_FP8_DTYPE)
@@ -159,6 +151,6 @@ class QuantFP8(CustomOp):
151
scales = scales.reshape(orig_shape[:-1] + (num_groups, ))
152
153
if self.column_major_scales:
162
- scales = scales.transpose(-2, -1).contiguous().transpose(-1, -2)
154
+ scales = scales.transpose(-2, -1).contiguous()
155
156
return x_quant, scales
vllm/model_executor/layers/quantization/utils/fp8_utils.py
+122
-177
@@ -13,9 +13,8 @@ import torch
13
import vllm.envs as envs
14
from vllm import _custom_ops as ops
15
from vllm.logger import init_logger
16
-from vllm.model_executor.layers.quantization.input_quant_fp8 import QuantFP8
16
from vllm.model_executor.layers.quantization.utils.quant_utils import (
18
- GroupShape, group_broadcast)
17
+ group_broadcast)
18
from vllm.model_executor.layers.quantization.utils.w8a8_utils import (
19
CUTLASS_BLOCK_FP8_SUPPORTED)
20
from vllm.model_executor.parameter import (BlockQuantScaleParameter,
@@ -25,7 +24,6 @@ from vllm.platforms import current_platform
24
from vllm.triton_utils import tl, triton
25
from vllm.utils import direct_register_custom_op
26
from vllm.utils.deep_gemm import (is_deep_gemm_e8m0_used,
28
- is_deep_gemm_supported,
27
should_use_deepgemm_for_fp8_linear)
28
29
logger = init_logger(__name__)
@@ -37,8 +35,6 @@ def is_fp8(x: Union[torch.dtype, torch.Tensor]) -> bool:
35
return x == torch.float8_e4m3fn or x == torch.float8_e4m3fnuz
36
37
40
-# We need to pass in the is_hopper flag as argument because the function
41
-# current_platform.is_device_capability() is not supported by Torch compiler.
38
def cutlass_scaled_mm(
39
A: torch.Tensor,
40
B: torch.Tensor,
@@ -46,17 +42,15 @@ def cutlass_scaled_mm(
42
Bs: torch.Tensor,
43
block_size: list[int],
44
output_dtype: torch.dtype = torch.float16,
49
- is_hopper: Optional[bool] = None,
45
) -> torch.Tensor:
51
- if is_hopper is None:
52
- is_hopper = current_platform.is_device_capability(90)
46
return ops.cutlass_scaled_mm(
47
A,
48
B.T,
49
out_dtype=output_dtype,
50
scale_a=As,
51
# SM90 block FP8 requires row-major scale_b, which we do ahead of time
59
- scale_b=Bs if block_size is not None and is_hopper else Bs.T)
52
+ scale_b=Bs if block_size is not None
53
+ and current_platform.is_device_capability(90) else Bs.T)
54
55
56
def rocm_aiter_gemm_w8a8_blockscale_impl(
@@ -102,190 +96,122 @@ if current_platform.is_rocm():
96
aiter_per1x128_quant = get_hip_quant(rocm_aiter.QuantType.per_1x128)
97
98
105
-# TODO we should be able to change the type of block_size to GroupShape
106
-# after we resolve GroupShape compilation issue
107
-# https://github.com/vllm-project/vllm/issues/25270
108
-def _w8a8_triton_block_scaled_mm_func(
109
- qx: torch.Tensor,
110
- weight: torch.Tensor,
111
- x_scale: torch.Tensor,
112
- weight_scale: torch.Tensor,
113
- block_size: list[int],
114
- output_dtype: torch.dtype,
115
-) -> torch.Tensor:
116
- return w8a8_triton_block_scaled_mm(qx, weight, x_scale, weight_scale,
117
- block_size, output_dtype)
99
+def dispatch_w8a8_blockscale_func(
100
+ use_cutlass: bool, use_aiter_and_is_supported: bool
101
+) -> Callable[[
102
+ torch.Tensor,
103
+ torch.Tensor,
104
+ torch.Tensor,
105
+ torch.Tensor,
106
+ list[int],
107
+ torch.dtype,
108
+], torch.Tensor]:
109
+ if use_cutlass:
110
+ return cutlass_scaled_mm
111
+ if (use_aiter_and_is_supported):
112
+ return torch.ops.vllm.rocm_aiter_gemm_w8a8_blockscale
113
+ return w8a8_block_fp8_matmul
114
115
120
-def _w8a8_triton_block_scaled_mm_fake(
121
- qx: torch.Tensor,
116
+# TODO fix ROCm->Triton custom path:
117
+# https://github.com/vllm-project/vllm/issues/14397
118
+def apply_w8a8_block_fp8_linear(
119
+ input: torch.Tensor,
120
weight: torch.Tensor,
123
- x_scale: torch.Tensor,
124
- weight_scale: torch.Tensor,
121
block_size: list[int],
126
- output_dtype: torch.dtype,
122
+ weight_scale: torch.Tensor,
123
+ input_scale: Optional[torch.Tensor] = None,
124
+ bias: Optional[torch.Tensor] = None,
125
+ cutlass_block_fp8_supported: bool = CUTLASS_BLOCK_FP8_SUPPORTED,
126
+ use_aiter_and_is_supported: bool = False,
127
) -> torch.Tensor:
128
- return torch.empty((qx.size(0), weight.size(0)),
129
- dtype=output_dtype,
130
- device=qx.device)
128
+ assert input_scale is None
129
+ # View input as 2D matrix for fp8 methods
130
+ input_2d = input.view(-1, input.shape[-1])
131
+ output_shape = [*input.shape[:-1], weight.shape[0]]
132
+ output_dtype = input.dtype
133
134
+ if should_use_deepgemm_for_fp8_linear(output_dtype, weight):
135
133
-# Note: the check can be removed when CPU torch > 2.7
134
-if not current_platform.is_cpu():
135
- direct_register_custom_op(
136
- "w8a8_triton_block_scaled_mm_func",
137
- _w8a8_triton_block_scaled_mm_func,
138
- fake_impl=_w8a8_triton_block_scaled_mm_fake,
139
- dispatch_key="CUDA",
140
- )
141
-
142
-
143
-# TODO fix ROCm->Triton custom path:
144
-# https://github.com/vllm-project/vllm/issues/14397
145
-class W8A8BlockFp8LinearOp:
146
- """
147
- This class executes a Blocked FP8 linear layer using cutlass if supported
148
- and torch.scaled_mm otherwise.
149
- """
150
-
151
- def __init__(
152
- self,
153
- weight_group_shape: GroupShape,
154
- act_quant_group_shape: GroupShape,
155
- cutlass_block_fp8_supported: bool = CUTLASS_BLOCK_FP8_SUPPORTED,
156
- use_aiter_and_is_supported: bool = False,
157
- ):
158
- self.weight_group_shape = weight_group_shape
159
- self.act_quant_group_shape = act_quant_group_shape
160
- self.is_deep_gemm_supported = is_deep_gemm_supported()
161
- self.is_hopper = current_platform.is_device_capability(90)
162
-
163
- # Get the correct blockscale mul and input quant operations.
164
- # We can't use _dispatch_w8a8_blockscale_op to figure out if we want
165
- # to use deepgemm because we don't know the shape of weights (and
166
- # whether deepgemm supports it) at the init time.
167
- self.w8a8_blockscale_op, self.input_quant_op = \
168
- self._dispatch_w8a8_blockscale_op(
169
- cutlass_block_fp8_supported, use_aiter_and_is_supported)
170
- self.deepgemm_input_quant_op = (QuantFP8(
171
- False,
172
- self.act_quant_group_shape,
173
- column_major_scales=True,
174
- use_ue8m0=is_deep_gemm_e8m0_used()) if self.is_deep_gemm_supported
175
- else None)
176
-
177
- def apply(
178
- self,
179
- input: torch.Tensor,
180
- weight: torch.Tensor,
181
- weight_scale: torch.Tensor,
182
- input_scale: Optional[torch.Tensor] = None,
183
- bias: Optional[torch.Tensor] = None,
184
- ) -> torch.Tensor:
185
- assert input_scale is None
186
- # View input as 2D matrix for fp8 methods
136
input_2d = input.view(-1, input.shape[-1])
137
output_shape = [*input.shape[:-1], weight.shape[0]]
189
- output_dtype = input.dtype
138
191
- if should_use_deepgemm_for_fp8_linear(output_dtype, weight,
192
- self.is_deep_gemm_supported):
193
- output = self._run_deepgemm(input, weight, weight_scale)
194
- if bias is not None:
195
- output = output + bias
196
- return output.to(dtype=input.dtype).view(*output_shape)
197
-
198
- output = self.w8a8_blockscale_op(input_2d, weight, weight_scale)
199
- if bias is not None:
200
- output = output + bias
201
- return output.to(dtype=input.dtype).view(*output_shape)
139
+ q_input, x_scale = per_token_group_quant_fp8(
140
+ input_2d,
141
+ block_size[1],
142
+ column_major_scales=True,
143
+ )
144
203
- def _run_deepgemm(
204
- self,
205
- input_2d: torch.Tensor,
206
- weight: torch.Tensor,
207
- weight_scale: torch.Tensor,
208
- ) -> torch.Tensor:
145
# ensure DeepGEMM-backed custom op is registered before use
146
import vllm.model_executor.layers.quantization.deepgemm # noqa: F401
147
212
- assert self.deepgemm_input_quant_op is not None
213
- q_input, x_scale = self.deepgemm_input_quant_op(input_2d)
214
- return torch.ops.vllm.w8a8_deepgemm_block_scaled_mm(
148
+ output = torch.ops.vllm.w8a8_block_fp8_matmul_deepgemm(
149
q_input,
150
weight,
151
x_scale,
152
weight_scale,
219
- self.weight_group_shape,
220
- output_dtype=input_2d.dtype)
221
-
222
- def _run_cutlass(
223
- self,
224
- input_2d: torch.Tensor,
225
- weight: torch.Tensor,
226
- weight_scale: torch.Tensor,
227
- ) -> torch.Tensor:
228
- assert self.input_quant_op is not None
229
- if self.is_hopper:
230
- # We pad unconditionally (even if shape is already divisible by 4)
231
- # to support dynamic shape for input_2d.shape[0] in torch.compile
232
- x = torch.nn.functional.pad(input_2d,
233
- (0, 0, 0, -input_2d.shape[0] % 4))
234
- else:
235
- x = input_2d
236
-
237
- q_input, x_scale = self.input_quant_op(x)
238
- output = cutlass_scaled_mm(q_input, weight, x_scale, weight_scale,
239
- list(self.weight_group_shape),
240
- input_2d.dtype, self.is_hopper)
241
- output = output[0:input_2d.shape[0], ...]
242
- return output
243
-
244
- def _run_aiter(
245
- self,
246
- input_2d: torch.Tensor,
247
- weight: torch.Tensor,
248
- weight_scale: torch.Tensor,
249
- ) -> torch.Tensor:
250
- assert self.act_quant_group_shape == GroupShape(1, 128)
251
- q_input, x_scale = aiter_per1x128_quant(
252
- input_2d.contiguous(), quant_dtype=rocm_aiter.dtypes.fp8)
253
- return torch.ops.vllm.rocm_aiter_gemm_w8a8_blockscale(
254
- q_input, weight, x_scale, weight_scale, self.weight_group_shape,
255
- input_2d.dtype)
256
-
257
- def _run_triton(
258
- self,
259
- input_2d: torch.Tensor,
260
- weight: torch.Tensor,
261
- weight_scale: torch.Tensor,
262
- ) -> torch.Tensor:
263
- assert self.input_quant_op is not None
264
- q_input, x_scale = self.input_quant_op(input_2d)
265
- return torch.ops.vllm.w8a8_triton_block_scaled_mm_func(
266
- q_input, weight, x_scale, weight_scale, self.weight_group_shape,
267
- input_2d.dtype)
268
-
269
- def _dispatch_w8a8_blockscale_op(
270
- self,
271
- use_cutlass: bool,
272
- use_aiter_and_is_supported: bool,
273
- ) -> tuple[Callable[[
274
- torch.Tensor,
275
- torch.Tensor,
276
- torch.Tensor,
277
- ], torch.Tensor], Optional[QuantFP8]]:
278
- if use_cutlass:
279
- return self._run_cutlass, (QuantFP8(False,
280
- self.act_quant_group_shape,
281
- column_major_scales=True,
282
- use_ue8m0=False))
153
+ block_size,
154
+ output_dtype=output_dtype)
155
+ if bias is not None:
156
+ output += bias
157
+ return output.to(dtype=output_dtype).view(*output_shape)
158
+
159
+ w8a8_blockscale_func = dispatch_w8a8_blockscale_func(
160
+ cutlass_block_fp8_supported, use_aiter_and_is_supported)
161
+ if cutlass_block_fp8_supported:
162
+ num_pad = 0
163
+ if current_platform.is_device_capability(90):
164
+ # pad first dimension to be divisible by 4 due to
165
+ # cutlass blockwise gemm limitation for hopper
166
+ num_pad = 4 - (input_2d.shape[0] % 4)
167
+ if num_pad > 0:
168
+ input_2d = torch.nn.functional.pad(input_2d,
169
+ (0, 0, 0, num_pad),
170
+ "constant", 0)
171
+ q_input, x_scale = per_token_group_quant_fp8(input_2d,
172
+ block_size[1],
173
+ column_major_scales=True)
174
+ output = w8a8_blockscale_func(q_input, weight, x_scale, weight_scale,
175
+ block_size, input.dtype)
176
+ if num_pad > 0:
177
+ output = output[:-num_pad]
178
+ else:
179
if use_aiter_and_is_supported:
284
- return self._run_aiter, None
285
- return self._run_triton, (QuantFP8(False,
286
- self.act_quant_group_shape,
287
- column_major_scales=False,
288
- use_ue8m0=False))
180
+ q_input, x_scale = aiter_per1x128_quant(
181
+ input_2d.contiguous(), quant_dtype=rocm_aiter.dtypes.fp8)
182
+ else:
183
+ q_input, x_scale = per_token_group_quant_fp8(
184
+ input_2d, block_size[1], column_major_scales=False)
185
+
186
+ output = w8a8_blockscale_func(q_input, weight, x_scale, weight_scale,
187
+ block_size, input.dtype)
188
+
189
+ if bias is not None:
190
+ output = output + bias
191
+ return output.to(dtype=input.dtype).view(*output_shape)
192
+
193
+
194
+def apply_w8a8_block_fp8_linear_fake(
195
+ input: torch.Tensor,
196
+ weight: torch.Tensor,
197
+ block_size: list[int],
198
+ weight_scale: torch.Tensor,
199
+ input_scale: Optional[torch.Tensor] = None,
200
+ bias: Optional[torch.Tensor] = None,
201
+ cutlass_block_fp8_supported: bool = CUTLASS_BLOCK_FP8_SUPPORTED,
202
+ use_aiter_and_is_supported: bool = False,
203
+) -> torch.Tensor:
204
+ output_shape = [*input.shape[:-1], weight.shape[0]]
205
+ return torch.empty(output_shape, dtype=input.dtype, device=input.device)
206
+
207
+
208
+if not current_platform.is_cpu():
209
+ direct_register_custom_op(
210
+ op_name="apply_w8a8_block_fp8_linear",
211
+ op_func=apply_w8a8_block_fp8_linear,
212
+ mutates_args=[],
213
+ fake_impl=apply_w8a8_block_fp8_linear_fake,
214
+ )
215
216
217
def input_to_float8(
@@ -537,7 +463,7 @@ def per_token_group_quant_fp8(
463
464
465
@triton.jit
540
-def _w8a8_triton_block_scaled_mm(
466
+def _w8a8_block_fp8_matmul(
467
# Pointers to inputs and output
468
A,
469
B,
@@ -662,7 +588,7 @@ def get_w8a8_block_fp8_configs(N: int, K: int, block_n: int,
588
return None
589
590
665
-def w8a8_triton_block_scaled_mm(
591
+def w8a8_block_fp8_matmul(
592
A: torch.Tensor,
593
B: torch.Tensor,
594
As: torch.Tensor,
@@ -722,7 +648,7 @@ def w8a8_triton_block_scaled_mm(
648
return (triton.cdiv(M, META["BLOCK_SIZE_M"]) *
649
triton.cdiv(N, META["BLOCK_SIZE_N"]), )
650
725
- _w8a8_triton_block_scaled_mm[grid](
651
+ _w8a8_block_fp8_matmul[grid](
652
A,
653
B,
654
C,
@@ -1005,6 +931,25 @@ def maybe_post_process_fp8_weight_block(layer: torch.nn.Module,
931
layer.weight_scale.data.T.contiguous(), requires_grad=False)
932
933
934
+def apply_fp8_block_linear(layer: torch.nn.Module, input: torch.Tensor,
935
+ bias: Optional[torch.Tensor],
936
+ cutlass_block_fp8_supported: bool,
937
+ use_aiter_and_is_supported: bool) -> torch.Tensor:
938
+ """Apply block-wise FP8 linear operation."""
939
+ assert layer.weight_block_size is not None
940
+
941
+ return torch.ops.vllm.apply_w8a8_block_fp8_linear(
942
+ input=input,
943
+ weight=layer.weight,
944
+ block_size=layer.weight_block_size,
945
+ weight_scale=layer.weight_scale,
946
+ input_scale=layer.input_scale,
947
+ bias=bias,
948
+ cutlass_block_fp8_supported=cutlass_block_fp8_supported,
949
+ use_aiter_and_is_supported=use_aiter_and_is_supported,
950
+ )
951
+
952
+
953
def expert_weight_is_col_major(x: torch.Tensor) -> bool:
954
assert x.dim() == 3
955
b, m, n = x.shape
vllm/utils/deep_gemm.py
+4
-8
@@ -9,7 +9,7 @@ from __future__ import annotations
9
import functools
10
import importlib
11
import os
12
-from typing import Any, Callable, NoReturn, Optional
12
+from typing import Any, Callable, NoReturn
13
14
import torch
15
@@ -184,13 +184,9 @@ def calc_diff(x: torch.Tensor, y: torch.Tensor):
184
return 1 - sim
185
186
187
-def should_use_deepgemm_for_fp8_linear(
188
- output_dtype: torch.dtype,
189
- weight: torch.Tensor,
190
- supports_deep_gemm: Optional[bool] = None):
191
- if supports_deep_gemm is None:
192
- supports_deep_gemm = is_deep_gemm_supported()
193
- return (supports_deep_gemm and output_dtype == torch.bfloat16
187
+def should_use_deepgemm_for_fp8_linear(output_dtype: torch.dtype,
188
+ weight: torch.Tensor):
189
+ return (is_deep_gemm_supported() and output_dtype == torch.bfloat16
190
and weight.shape[0] % 128 == 0 and weight.shape[1] % 128 == 0)
191
192