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
[Perf] Fix and reapply move apply w8a8 block fp8 linear to class (#25696)
Signed-off-by: ElizaWszola <ewszola@redhat.com> Signed-off-by: ElizaWszola <elizaw.9289@gmail.com> Signed-off-by: Luka Govedič <lgovedic@redhat.com> Signed-off-by: Luka Govedič <ProExpertProg@users.noreply.github.com> Co-authored-by: Luka Govedič <ProExpertProg@users.noreply.github.com> Co-authored-by: Michael Goin <mgoin64@gmail.com> Co-authored-by: Luka Govedič <lgovedic@redhat.com>
ElizaWszola committed
Oct 2, 2025 at 21:35 UTC
502640c3f944f8b9702989bd45f8add970b67ea0
13 files changed
+412
-200
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_block_fp8_matmul,
20
+ w8a8_triton_block_scaled_mm,
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_block_fp8_matmul(
161
+ "triton_fp8_fp8_fp16_scaled_mm_blockwise": lambda: w8a8_triton_block_scaled_mm(
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_block_fp8_matmul,
12
+ w8a8_triton_block_scaled_mm,
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_block_fp8_matmul(A_vllm,
66
+ return w8a8_triton_block_scaled_mm(A_vllm,
67
B_vllm,
68
A_scale_vllm,
69
B_scale_vllm,
tests/kernels/quantization/test_block_fp8.py
+3
-2
@@ -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_block_fp8_matmul)
14
+ cutlass_scaled_mm, per_token_group_quant_fp8, w8a8_triton_block_scaled_mm)
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,7 +91,8 @@ 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_block_fp8_matmul(A_fp8, B_fp8, As, Bs, block_size, out_dtype)
94
+ out = w8a8_triton_block_scaled_mm(A_fp8, B_fp8, As, Bs, block_size,
95
+ out_dtype)
96
97
rel_diff = (torch.mean(
98
torch.abs(out.to(torch.float32) - ref_out.to(torch.float32))) /
tests/kernels/quantization/test_fp8_quant_group.py
+21
-9
@@ -20,9 +20,11 @@ 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])
24
@torch.inference_mode()
25
def test_quantfp8_group_functionality(batch_size: int, hidden_dim: int,
25
- group_size: int, seed: int) -> None:
26
+ group_size: int, seed: int,
27
+ use_ue8m0: bool) -> None:
28
"""Test QuantFP8 group quantization with various configurations.
29
30
Tests both CUDA and native implementations, column-major scales,
@@ -38,7 +40,8 @@ def test_quantfp8_group_functionality(batch_size: int, hidden_dim: int,
40
group_shape = GroupShape(1, group_size)
41
quant_op = QuantFP8(static=False,
42
group_shape=group_shape,
41
- column_major_scales=False)
43
+ column_major_scales=False,
44
+ use_ue8m0=use_ue8m0)
45
46
# 1. Test native implementation (always available)
47
x_quant_native, scales_native = quant_op.forward_native(x.clone())
@@ -48,9 +51,15 @@ def test_quantfp8_group_functionality(batch_size: int, hidden_dim: int,
51
# 2. Test column-major scales configuration
52
quant_op_col = QuantFP8(static=False,
53
group_shape=group_shape,
51
- column_major_scales=True)
54
+ column_major_scales=True,
55
+ use_ue8m0=use_ue8m0)
56
_, scales_col = quant_op_col.forward_native(x.clone())
53
- assert scales_col.shape == (expected_num_groups, batch_size)
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)
63
64
# 3. Test CUDA implementation (only for divisible dimensions)
65
if is_divisible:
@@ -68,21 +77,23 @@ def test_quantfp8_group_functionality(batch_size: int, hidden_dim: int,
77
78
79
@pytest.mark.parametrize("seed", [42])
80
+@pytest.mark.parametrize("use_ue8m0", [True, False])
81
@torch.inference_mode()
72
-def test_quantfp8_group_multidimensional(seed: int) -> None:
82
+def test_quantfp8_group_multidimensional(seed: int, use_ue8m0: bool) -> None:
83
current_platform.seed_everything(seed)
84
85
group_size = 64
86
87
# Test with 3D input
78
- batch1, batch2, hidden_dim = 4, 8, 512
88
+ batch1, batch2, hidden_dim = 4, 8, 1024
89
x_3d = torch.randn(
90
(batch1, batch2, hidden_dim), dtype=torch.bfloat16, device="cuda") * 8
91
92
group_shape = GroupShape(1, group_size)
93
quant_op = QuantFP8(static=False,
94
group_shape=group_shape,
85
- column_major_scales=False)
95
+ column_major_scales=False,
96
+ use_ue8m0=use_ue8m0)
97
98
x_quant, scales = quant_op.forward_native(x_3d.clone())
99
assert x_quant.shape == x_3d.shape
@@ -91,9 +102,10 @@ def test_quantfp8_group_multidimensional(seed: int) -> None:
102
# Test column_major_scales with multi-dim
103
quant_op_col = QuantFP8(static=False,
104
group_shape=group_shape,
94
- column_major_scales=True)
105
+ column_major_scales=True,
106
+ use_ue8m0=use_ue8m0)
107
_, scales_col = quant_op_col.forward_native(x_3d.clone())
96
- assert scales_col.shape == (batch1, hidden_dim // group_size, batch2)
108
+ assert scales_col.shape == (batch1, batch2, hidden_dim // group_size)
109
110
# Test with 4D input
111
batch1, batch2, batch3, hidden_dim = 2, 3, 4, 256
tests/model_executor/test_enabled_custom_ops.py
-30
@@ -17,8 +17,6 @@ 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)
20
from vllm.platforms import current_platform
21
22
RMS_NORM_SUPPORTED_DTYPES = [torch.float16, torch.bfloat16]
@@ -111,34 +109,6 @@ def test_enabled_ops_invalid(env: str):
109
RMSNorm(1024).enabled()
110
111
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
-
112
@pytest.mark.parametrize("use_rocm_aiter", ["0", "1"])
113
def test_topk_dispatch(use_rocm_aiter: str, monkeypatch):
114
monkeypatch.setenv("VLLM_ROCM_USE_AITER", use_rocm_aiter)
tests/quantization/test_compressed_tensors.py
+35
@@ -18,6 +18,9 @@ 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)
24
from vllm.model_executor.layers.quantization.utils.quant_utils import (
25
cutlass_fp4_supported)
26
from vllm.model_executor.layers.quantization.utils.w8a8_utils import (
@@ -742,3 +745,35 @@ def test_compressed_tensors_transforms_perplexity(vllm_runner, model, prompt,
745
perplexity = llm.generate_prompt_perplexity([prompt])[0]
746
print(perplexity)
747
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/vllm.py
+17
@@ -516,6 +516,23 @@ class VllmConfig:
516
" by VLLM_DEBUG_DUMP_PATH to %s", env_path)
517
self.compilation_config.debug_dump_path = env_path
518
519
+ def has_blocked_weights():
520
+ if self.quant_config is not None:
521
+ if hasattr(self.quant_config, "weight_block_size"):
522
+ return self.quant_config.weight_block_size is not None
523
+ elif hasattr(self.quant_config, "has_blocked_weights"):
524
+ return self.quant_config.has_blocked_weights()
525
+ return False
526
+
527
+ # Enable quant_fp8 CUDA ops (TODO disable in follow up)
528
+ # On H100 the CUDA kernel is faster than
529
+ # native implementation
530
+ # https://github.com/vllm-project/vllm/issues/25094
531
+ if has_blocked_weights():
532
+ custom_ops = self.compilation_config.custom_ops
533
+ if "none" not in custom_ops and "-quant_fp8" not in custom_ops:
534
+ custom_ops.append("+quant_fp8")
535
+
536
def update_sizes_for_sequence_parallelism(self,
537
possible_sizes: list) -> list:
538
# remove the sizes that not multiple of tp_size when
vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors.py
+8
@@ -644,6 +644,14 @@ 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
+
655
@staticmethod
656
def supports_cutlass_24(
657
weight_quant: Optional[QuantizationArgs],
vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_w8a8_fp8.py
+26
-11
@@ -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
- apply_fp8_block_linear, check_aiter_fp8_linear_support,
14
+ W8A8BlockFp8LinearOp, 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,16 +41,30 @@ 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)
44
45
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
+
52
self.cutlass_block_fp8_supported = cutlass_block_fp8_supported()
53
self.use_aiter_and_is_supported = check_aiter_fp8_linear_support()
54
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
+
68
@classmethod
69
def get_min_capability(cls) -> int:
70
# lovelace and up
@@ -142,13 +156,14 @@ class CompressedTensorsW8A8Fp8(CompressedTensorsScheme):
156
x: torch.Tensor,
157
bias: Optional[torch.Tensor] = None) -> torch.Tensor:
158
145
- if layer.weight_block_size is not None:
146
- return apply_fp8_block_linear(
147
- layer,
159
+ if self.weight_block_size is not None:
160
+ return self.w8a8_block_fp8_linear.apply(
161
input=x,
162
+ weight=layer.weight,
163
+ weight_scale=layer.weight_scale,
164
+ input_scale=layer.input_scale,
165
bias=bias,
150
- cutlass_block_fp8_supported=self.cutlass_block_fp8_supported,
151
- use_aiter_and_is_supported=self.use_aiter_and_is_supported)
166
+ )
167
168
return self.fp8_linear.apply(input=x,
169
weight=layer.weight,
vllm/model_executor/layers/quantization/fp8.py
+28
-12
@@ -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
- apply_fp8_block_linear, check_aiter_fp8_linear_support,
36
+ W8A8BlockFp8LinearOp, 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,15 +242,28 @@ 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
- # 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
245
+ if self.weight_block_size:
246
+ self.act_q_group_shape = GroupShape(1, self.weight_block_size[0])
247
else:
249
- self.act_q_group_shape = GroupShape.PER_TENSOR
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
253
251
- self.fp8_linear = Fp8LinearOp(
252
- act_quant_static=self.act_q_static,
253
- act_quant_group_shape=self.act_q_group_shape)
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)
267
268
def create_weights(
269
self,
@@ -399,12 +412,15 @@ class Fp8LinearMethod(LinearMethodBase):
412
bias=bias)
413
414
if self.block_quant:
402
- return apply_fp8_block_linear(
403
- layer,
415
+ assert self.weight_block_size is not None
416
+
417
+ return self.w8a8_block_fp8_linear.apply(
418
input=x,
419
+ weight=layer.weight,
420
+ weight_scale=layer.weight_scale,
421
+ input_scale=layer.input_scale,
422
bias=bias,
406
- cutlass_block_fp8_supported=self.cutlass_block_fp8_supported,
407
- use_aiter_and_is_supported=self.use_aiter_and_is_supported)
423
+ )
424
425
return self.fp8_linear.apply(input=x,
426
weight=layer.weight,
vllm/model_executor/layers/quantization/input_quant_fp8.py
+16
-8
@@ -27,11 +27,14 @@ class QuantFP8(CustomOp):
27
This CustomOp supports both static and dynamic quantization.
28
"""
29
30
- def __init__(self,
31
- static: bool,
32
- group_shape: GroupShape,
33
- num_token_padding: Optional[int] = None,
34
- column_major_scales: bool = False):
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
+ ):
38
"""
39
:param static: static or dynamic quantization
40
:param group_shape: quantization group shape (PER_TOKEN, PER_TENSOR,
@@ -46,6 +49,7 @@ class QuantFP8(CustomOp):
49
self.group_shape = group_shape
50
self.num_token_padding = num_token_padding
51
self.column_major_scales = column_major_scales
52
+ self.use_ue8m0 = use_ue8m0
53
54
self.is_group_quant = group_shape.is_per_group()
55
if self.is_group_quant:
@@ -70,7 +74,8 @@ class QuantFP8(CustomOp):
74
x,
75
group_size=self.group_size,
76
column_major_scales=self.column_major_scales,
73
- dtype=_FP8_DTYPE)
77
+ dtype=_FP8_DTYPE,
78
+ use_ue8m0=self.use_ue8m0)
79
80
assert (scale is not None) == self.static
81
assert scale_ub is None or (not self.static and self.group_shape
@@ -137,7 +142,10 @@ class QuantFP8(CustomOp):
142
143
x_grouped = x.view(-1, num_groups, self.group_size)
144
absmax = x_grouped.abs().max(dim=-1, keepdim=True)[0].float()
140
- scales = (absmax / _FP8_MAX).clamp(min=_FP8_MIN_SCALING_FACTOR)
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)
149
150
x_scaled = x_grouped / scales
151
x_quant = x_scaled.clamp(_FP8_MIN, _FP8_MAX).to(_FP8_DTYPE)
@@ -151,6 +159,6 @@ class QuantFP8(CustomOp):
159
scales = scales.reshape(orig_shape[:-1] + (num_groups, ))
160
161
if self.column_major_scales:
154
- scales = scales.transpose(-2, -1).contiguous()
162
+ scales = scales.transpose(-2, -1).contiguous().transpose(-1, -2)
163
164
return x_quant, scales
vllm/model_executor/layers/quantization/utils/fp8_utils.py
+240
-117
@@ -13,8 +13,9 @@ 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
17
from vllm.model_executor.layers.quantization.utils.quant_utils import (
17
- group_broadcast)
18
+ GroupShape, group_broadcast)
19
from vllm.model_executor.layers.quantization.utils.w8a8_utils import (
20
CUTLASS_BLOCK_FP8_SUPPORTED)
21
from vllm.model_executor.parameter import (BlockQuantScaleParameter,
@@ -24,6 +25,7 @@ from vllm.platforms import current_platform
25
from vllm.triton_utils import tl, triton
26
from vllm.utils import direct_register_custom_op
27
from vllm.utils.deep_gemm import (fp8_gemm_nt, is_deep_gemm_e8m0_used,
28
+ is_deep_gemm_supported,
29
should_use_deepgemm_for_fp8_linear)
30
31
logger = init_logger(__name__)
@@ -35,6 +37,8 @@ def is_fp8(x: Union[torch.dtype, torch.Tensor]) -> bool:
37
return x == torch.float8_e4m3fn or x == torch.float8_e4m3fnuz
38
39
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.
42
def cutlass_scaled_mm(
43
A: torch.Tensor,
44
B: torch.Tensor,
@@ -42,15 +46,17 @@ def cutlass_scaled_mm(
46
Bs: torch.Tensor,
47
block_size: list[int],
48
output_dtype: torch.dtype = torch.float16,
49
+ is_hopper: Optional[bool] = None,
50
) -> torch.Tensor:
51
+ if is_hopper is None:
52
+ is_hopper = current_platform.is_device_capability(90)
53
return ops.cutlass_scaled_mm(
54
A,
55
B.T,
56
out_dtype=output_dtype,
57
scale_a=As,
58
# SM90 block FP8 requires row-major scale_b, which we do ahead of time
52
- scale_b=Bs if block_size is not None
53
- and current_platform.is_device_capability(90) else Bs.T)
59
+ scale_b=Bs if block_size is not None and is_hopper else Bs.T)
60
61
62
def rocm_aiter_gemm_w8a8_blockscale_impl(
@@ -96,115 +102,251 @@ if current_platform.is_rocm():
102
aiter_per1x128_quant = get_hip_quant(rocm_aiter.QuantType.per_1x128)
103
104
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
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)
118
119
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
+def _w8a8_triton_block_scaled_mm_fake(
121
+ qx: torch.Tensor,
122
weight: torch.Tensor,
121
- block_size: list[int],
123
+ x_scale: torch.Tensor,
124
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,
125
+ block_size: list[int],
126
+ output_dtype: torch.dtype,
127
) -> torch.Tensor:
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
128
+ return torch.empty((qx.size(0), weight.size(0)),
129
+ dtype=output_dtype,
130
+ device=qx.device)
131
134
- if should_use_deepgemm_for_fp8_linear(output_dtype, weight):
132
136
- input_2d = input.view(-1, input.shape[-1])
137
- output_shape = [*input.shape[:-1], weight.shape[0]]
133
+direct_register_custom_op(
134
+ "w8a8_triton_block_scaled_mm_func",
135
+ _w8a8_triton_block_scaled_mm_func,
136
+ fake_impl=_w8a8_triton_block_scaled_mm_fake,
137
+)
138
139
- q_input, x_scale = per_token_group_quant_fp8(
140
- input_2d,
141
- block_size[1],
142
- column_major_scales=True,
143
- )
144
- output = torch.empty((q_input.shape[0], weight.shape[0]),
145
- dtype=torch.bfloat16,
146
- device=q_input.device)
147
- fp8_gemm_nt((q_input, x_scale), (weight, weight_scale), output)
148
- if bias is not None:
149
- output += bias
150
- return output.to(dtype=output_dtype).view(*output_shape)
151
-
152
- w8a8_blockscale_func = dispatch_w8a8_blockscale_func(
153
- cutlass_block_fp8_supported, use_aiter_and_is_supported)
154
- if cutlass_block_fp8_supported:
155
- num_pad = 0
156
- if current_platform.is_device_capability(90):
157
- # pad first dimension to be divisible by 4 due to
158
- # cutlass blockwise gemm limitation for hopper
159
- num_pad = 4 - (input_2d.shape[0] % 4)
160
- if num_pad > 0:
161
- input_2d = torch.nn.functional.pad(input_2d,
162
- (0, 0, 0, num_pad),
163
- "constant", 0)
164
- q_input, x_scale = per_token_group_quant_fp8(input_2d,
165
- block_size[1],
166
- column_major_scales=True)
167
- output = w8a8_blockscale_func(q_input, weight, x_scale, weight_scale,
168
- block_size, input.dtype)
169
- if num_pad > 0:
170
- output = output[:-num_pad]
171
- else:
172
- if use_aiter_and_is_supported:
173
- q_input, x_scale = aiter_per1x128_quant(
174
- input_2d.contiguous(), quant_dtype=rocm_aiter.dtypes.fp8)
175
- else:
176
- q_input, x_scale = per_token_group_quant_fp8(
177
- input_2d, block_size[1], column_major_scales=False)
139
179
- output = w8a8_blockscale_func(q_input, weight, x_scale, weight_scale,
180
- block_size, input.dtype)
140
+def _padded_cutlass(
141
+ qx: torch.Tensor,
142
+ weight: torch.Tensor,
143
+ x_scale: torch.Tensor,
144
+ weight_scale: torch.Tensor,
145
+ block_size: list[int],
146
+ output_dtype: torch.dtype,
147
+) -> torch.Tensor:
148
+ pad_multiple = 4
149
+ dim = qx.shape[0]
150
+ padded = dim if dim % pad_multiple == 0 else dim + pad_multiple - (
151
+ dim % pad_multiple)
152
182
- if bias is not None:
183
- output = output + bias
184
- return output.to(dtype=input.dtype).view(*output_shape)
153
+ padded_shape = [padded, *qx.shape[1:]]
154
+ padded_qx = torch.zeros(padded_shape, device=qx.device, dtype=qx.dtype)
155
+ padded_qx[0:qx.shape[0], ...].copy_(qx)
156
157
+ padded_x_scale_shape = [*x_scale.shape[1:], padded]
158
+ padded_x_scale = torch.ones(padded_x_scale_shape,
159
+ device=x_scale.device,
160
+ dtype=x_scale.dtype).permute(-1, -2)
161
+ padded_x_scale[0:x_scale.shape[0], ...].copy_(x_scale)
162
187
-def apply_w8a8_block_fp8_linear_fake(
188
- input: torch.Tensor,
163
+ output = cutlass_scaled_mm(padded_qx, weight, padded_x_scale, weight_scale,
164
+ block_size, output_dtype, True)
165
+ return output[0:qx.shape[0], ...]
166
+
167
+
168
+def _padded_cutlass_fake(
169
+ qx: torch.Tensor,
170
weight: torch.Tensor,
190
- block_size: list[int],
171
+ x_scale: torch.Tensor,
172
weight_scale: torch.Tensor,
192
- input_scale: Optional[torch.Tensor] = None,
193
- bias: Optional[torch.Tensor] = None,
194
- cutlass_block_fp8_supported: bool = CUTLASS_BLOCK_FP8_SUPPORTED,
195
- use_aiter_and_is_supported: bool = False,
173
+ block_size: list[int],
174
+ output_dtype: torch.dtype,
175
) -> torch.Tensor:
197
- output_shape = [*input.shape[:-1], weight.shape[0]]
198
- return torch.empty(output_shape, dtype=input.dtype, device=input.device)
176
+ return torch.empty((qx.size(0), weight.size(0)),
177
+ dtype=output_dtype,
178
+ device=qx.device)
179
180
201
-if not current_platform.is_cpu():
202
- direct_register_custom_op(
203
- op_name="apply_w8a8_block_fp8_linear",
204
- op_func=apply_w8a8_block_fp8_linear,
205
- mutates_args=[],
206
- fake_impl=apply_w8a8_block_fp8_linear_fake,
207
- )
181
+direct_register_custom_op(
182
+ "padded_cutlass",
183
+ _padded_cutlass,
184
+ fake_impl=_padded_cutlass_fake,
185
+)
186
+
187
+
188
+def _fp8_gemm_nt_op(q_input: torch.Tensor, input_scale: torch.Tensor,
189
+ weight: torch.Tensor, weight_scale: torch.Tensor,
190
+ output: torch.Tensor, use_deep_gemm_e8m0: bool) -> None:
191
+ fp8_gemm_nt((q_input, input_scale), (weight, weight_scale),
192
+ output,
193
+ is_deep_gemm_e8m0_used=use_deep_gemm_e8m0)
194
+
195
+
196
+def _fp8_gemm_nt_op_fake(q_input: torch.Tensor, input_scale: torch.Tensor,
197
+ weight: torch.Tensor, weight_scale: torch.Tensor,
198
+ output: torch.Tensor,
199
+ use_deep_gemm_e8m0: bool) -> None:
200
+ return None
201
+
202
+
203
+direct_register_custom_op(
204
+ "fp8_gemm_nt_op",
205
+ _fp8_gemm_nt_op,
206
+ mutates_args=["output"],
207
+ fake_impl=_fp8_gemm_nt_op_fake,
208
+)
209
+
210
+
211
+# TODO fix ROCm->Triton custom path:
212
+# https://github.com/vllm-project/vllm/issues/14397
213
+class W8A8BlockFp8LinearOp:
214
+ """
215
+ This class executes a Blocked FP8 linear layer using cutlass if supported
216
+ and torch.scaled_mm otherwise.
217
+ """
218
+
219
+ def __init__(
220
+ self,
221
+ weight_group_shape: GroupShape,
222
+ act_quant_group_shape: GroupShape,
223
+ cutlass_block_fp8_supported: bool = CUTLASS_BLOCK_FP8_SUPPORTED,
224
+ use_aiter_and_is_supported: bool = False,
225
+ ):
226
+ self.weight_group_shape = weight_group_shape
227
+ self.act_quant_group_shape = act_quant_group_shape
228
+ self.is_deep_gemm_supported = is_deep_gemm_supported()
229
+ self.is_hopper = current_platform.is_device_capability(90)
230
+ self.use_deep_gemm_e8m0 = is_deep_gemm_e8m0_used()
231
+
232
+ # Get the correct blockscale mul and input quant operations.
233
+ # We can't use _dispatch_w8a8_blockscale_op to figure out if we want
234
+ # to use deepgemm because we don't know the shape of weights (and
235
+ # whether deepgemm supports it) at the init time.
236
+ self.w8a8_blockscale_op, self.input_quant_op = \
237
+ self._dispatch_w8a8_blockscale_op(
238
+ cutlass_block_fp8_supported, use_aiter_and_is_supported)
239
+ self.deepgemm_input_quant_op = (QuantFP8(
240
+ False,
241
+ self.act_quant_group_shape,
242
+ column_major_scales=True,
243
+ use_ue8m0=self.use_deep_gemm_e8m0) if self.is_deep_gemm_supported
244
+ else None)
245
+
246
+ def apply(
247
+ self,
248
+ input: torch.Tensor,
249
+ weight: torch.Tensor,
250
+ weight_scale: torch.Tensor,
251
+ input_scale: Optional[torch.Tensor] = None,
252
+ bias: Optional[torch.Tensor] = None,
253
+ ) -> torch.Tensor:
254
+ assert input_scale is None
255
+ # View input as 2D matrix for fp8 methods
256
+ input_2d = input.view(-1, input.shape[-1])
257
+ output_shape = [*input.shape[:-1], weight.shape[0]]
258
+ output_dtype = input.dtype
259
+
260
+ if should_use_deepgemm_for_fp8_linear(output_dtype, weight,
261
+ self.is_deep_gemm_supported):
262
+ output = self._run_deepgemm(input_2d, weight, weight_scale)
263
+ else:
264
+ output = self.w8a8_blockscale_op(input_2d, weight, weight_scale)
265
+
266
+ if bias is not None:
267
+ output = output + bias
268
+ return output.to(dtype=input.dtype).view(*output_shape)
269
+
270
+ def _run_deepgemm(
271
+ self,
272
+ input_2d: torch.Tensor,
273
+ weight: torch.Tensor,
274
+ weight_scale: torch.Tensor,
275
+ ) -> torch.Tensor:
276
+ assert self.deepgemm_input_quant_op is not None
277
+ q_input, input_scale = self.deepgemm_input_quant_op(input_2d)
278
+ output = torch.empty((q_input.shape[0], weight.shape[0]),
279
+ dtype=torch.bfloat16,
280
+ device=q_input.device)
281
+ torch.ops.vllm.fp8_gemm_nt_op(q_input, input_scale, weight,
282
+ weight_scale, output,
283
+ self.use_deep_gemm_e8m0)
284
+ return output
285
+
286
+ def _run_cutlass(
287
+ self,
288
+ input_2d: torch.Tensor,
289
+ weight: torch.Tensor,
290
+ weight_scale: torch.Tensor,
291
+ ) -> torch.Tensor:
292
+ assert self.input_quant_op is not None
293
+ q_input, input_scale = self.input_quant_op(input_2d)
294
+ if self.is_hopper:
295
+ return torch.ops.vllm.padded_cutlass(q_input, weight, input_scale,
296
+ weight_scale,
297
+ list(self.weight_group_shape),
298
+ input_2d.dtype)
299
+ else:
300
+ return cutlass_scaled_mm(q_input, weight,
301
+ input_scale, weight_scale,
302
+ list(self.weight_group_shape),
303
+ input_2d.dtype, False)
304
+
305
+ def _run_aiter(
306
+ self,
307
+ input_2d: torch.Tensor,
308
+ weight: torch.Tensor,
309
+ weight_scale: torch.Tensor,
310
+ ) -> torch.Tensor:
311
+ assert self.act_quant_group_shape == GroupShape(1, 128)
312
+ q_input, input_scale = aiter_per1x128_quant(
313
+ input_2d.contiguous(), quant_dtype=rocm_aiter.dtypes.fp8)
314
+ return torch.ops.vllm.rocm_aiter_gemm_w8a8_blockscale(
315
+ q_input, weight, input_scale, weight_scale,
316
+ self.weight_group_shape, input_2d.dtype)
317
+
318
+ def _run_triton(
319
+ self,
320
+ input_2d: torch.Tensor,
321
+ weight: torch.Tensor,
322
+ weight_scale: torch.Tensor,
323
+ ) -> torch.Tensor:
324
+ assert self.input_quant_op is not None
325
+ q_input, input_scale = self.input_quant_op(input_2d)
326
+ return torch.ops.vllm.w8a8_triton_block_scaled_mm_func(
327
+ q_input, weight, input_scale, weight_scale,
328
+ self.weight_group_shape, input_2d.dtype)
329
+
330
+ def _dispatch_w8a8_blockscale_op(
331
+ self,
332
+ use_cutlass: bool,
333
+ use_aiter_and_is_supported: bool,
334
+ ) -> tuple[Callable[[
335
+ torch.Tensor,
336
+ torch.Tensor,
337
+ torch.Tensor,
338
+ ], torch.Tensor], Optional[QuantFP8]]:
339
+ if use_cutlass:
340
+ return self._run_cutlass, (QuantFP8(False,
341
+ self.act_quant_group_shape,
342
+ column_major_scales=True,
343
+ use_ue8m0=False))
344
+ if use_aiter_and_is_supported:
345
+ return self._run_aiter, None
346
+ return self._run_triton, (QuantFP8(False,
347
+ self.act_quant_group_shape,
348
+ column_major_scales=False,
349
+ use_ue8m0=False))
350
351
352
def input_to_float8(
@@ -456,7 +598,7 @@ def per_token_group_quant_fp8(
598
599
600
@triton.jit
459
-def _w8a8_block_fp8_matmul(
601
+def _w8a8_triton_block_scaled_mm(
602
# Pointers to inputs and output
603
A,
604
B,
@@ -581,7 +723,7 @@ def get_w8a8_block_fp8_configs(N: int, K: int, block_n: int,
723
return None
724
725
584
-def w8a8_block_fp8_matmul(
726
+def w8a8_triton_block_scaled_mm(
727
A: torch.Tensor,
728
B: torch.Tensor,
729
As: torch.Tensor,
@@ -641,7 +783,7 @@ def w8a8_block_fp8_matmul(
783
return (triton.cdiv(M, META["BLOCK_SIZE_M"]) *
784
triton.cdiv(N, META["BLOCK_SIZE_N"]), )
785
644
- _w8a8_block_fp8_matmul[grid](
786
+ _w8a8_triton_block_scaled_mm[grid](
787
A,
788
B,
789
C,
@@ -924,25 +1066,6 @@ def maybe_post_process_fp8_weight_block(layer: torch.nn.Module,
1066
layer.weight_scale.data.T.contiguous(), requires_grad=False)
1067
1068
927
-def apply_fp8_block_linear(layer: torch.nn.Module, input: torch.Tensor,
928
- bias: Optional[torch.Tensor],
929
- cutlass_block_fp8_supported: bool,
930
- use_aiter_and_is_supported: bool) -> torch.Tensor:
931
- """Apply block-wise FP8 linear operation."""
932
- assert layer.weight_block_size is not None
933
-
934
- return torch.ops.vllm.apply_w8a8_block_fp8_linear(
935
- input=input,
936
- weight=layer.weight,
937
- block_size=layer.weight_block_size,
938
- weight_scale=layer.weight_scale,
939
- input_scale=layer.input_scale,
940
- bias=bias,
941
- cutlass_block_fp8_supported=cutlass_block_fp8_supported,
942
- use_aiter_and_is_supported=use_aiter_and_is_supported,
943
- )
944
-
945
-
1069
def expert_weight_is_col_major(x: torch.Tensor) -> bool:
1070
assert x.dim() == 3
1071
b, m, n = x.shape
vllm/utils/deep_gemm.py
+14
-7
@@ -9,7 +9,7 @@ from __future__ import annotations
9
import functools
10
import importlib
11
import os
12
-from typing import Any, Callable, NoReturn
12
+from typing import Any, Callable, NoReturn, Optional
13
14
import torch
15
@@ -136,9 +136,12 @@ def fp8_gemm_nt(*args, **kwargs):
136
_lazy_init()
137
if _fp8_gemm_nt_impl is None:
138
return _missing(*args, **kwargs)
139
- return _fp8_gemm_nt_impl(*args,
140
- disable_ue8m0_cast=not is_deep_gemm_e8m0_used(),
141
- **kwargs)
139
+ if "is_deep_gemm_e8m0_used" in kwargs:
140
+ use_ue8m0 = kwargs["is_deep_gemm_e8m0_used"]
141
+ del kwargs["is_deep_gemm_e8m0_used"]
142
+ else:
143
+ use_ue8m0 = is_deep_gemm_e8m0_used()
144
+ return _fp8_gemm_nt_impl(*args, disable_ue8m0_cast=not use_ue8m0, **kwargs)
145
146
147
def m_grouped_fp8_gemm_nt_contiguous(*args, **kwargs):
@@ -301,9 +304,13 @@ def calc_diff(x: torch.Tensor, y: torch.Tensor):
304
return 1 - sim
305
306
304
-def should_use_deepgemm_for_fp8_linear(output_dtype: torch.dtype,
305
- weight: torch.Tensor):
306
- return (is_deep_gemm_supported() and output_dtype == torch.bfloat16
307
+def should_use_deepgemm_for_fp8_linear(
308
+ output_dtype: torch.dtype,
309
+ weight: torch.Tensor,
310
+ supports_deep_gemm: Optional[bool] = None):
311
+ if supports_deep_gemm is None:
312
+ supports_deep_gemm = is_deep_gemm_supported()
313
+ return (supports_deep_gemm and output_dtype == torch.bfloat16
314
and weight.shape[0] % 128 == 0 and weight.shape[1] % 128 == 0)
315
316