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, 18:15 UTC
[Perf] Create TMA-aligned input scale tensor for DeepGemm on Hopper (#32619)
Signed-off-by: Xin Yang <xyangx@amazon.com>
Xin Yang committed
Jan 22, 2026 at 12:47 UTC
d08b356ee08dc8e0b0dd052ed7bef5245e21d7af
7 files changed
+75
-17
benchmarks/kernels/deepgemm/benchmark_fp8_block_dense_gemm.py
+3
-3
@@ -14,7 +14,6 @@ from vllm.triton_utils import triton
14
from vllm.utils.deep_gemm import (
15
calc_diff,
16
fp8_gemm_nt,
17
- get_col_major_tma_aligned_tensor,
17
per_block_cast_to_fp8,
18
)
19
@@ -48,8 +47,9 @@ def benchmark_shape(
47
block_size = [128, 128]
48
49
# Pre-quantize A for all implementations
51
- A_deepgemm, A_scale_deepgemm = per_token_group_quant_fp8(A, block_size[1])
52
- A_scale_deepgemm = get_col_major_tma_aligned_tensor(A_scale_deepgemm)
50
+ A_deepgemm, A_scale_deepgemm = per_token_group_quant_fp8(
51
+ A, block_size[1], column_major_scales=True, tma_aligned_scales=True
52
+ )
53
C_deepgemm = torch.empty((m, n), device="cuda", dtype=torch.bfloat16)
54
A_vllm, A_scale_vllm = per_token_group_quant_fp8(A, block_size[1])
55
A_vllm_cutlass, A_scale_vllm_cutlass = per_token_group_quant_fp8(
tests/kernels/quant_utils.py
+3
@@ -9,6 +9,7 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import (
9
group_broadcast,
10
)
11
from vllm.platforms import current_platform
12
+from vllm.utils.deep_gemm import _ceil_to_ue8m0, is_deep_gemm_e8m0_used
13
from vllm.utils.math_utils import round_up
14
15
FP8_DTYPE = current_platform.fp8_dtype()
@@ -170,6 +171,8 @@ def native_per_token_group_quant_fp8(
171
x_ = x.reshape(x.numel() // group_size, group_size)
172
amax = x_.abs().max(dim=-1, keepdim=True)[0].clamp(min=eps).to(torch.float32)
173
x_s = amax / fp8_max
174
+ if is_deep_gemm_e8m0_used():
175
+ x_s = _ceil_to_ue8m0(x_s)
176
x_q = (x_ / x_s).clamp(min=fp8_min, max=fp8_max).to(dtype)
177
x_q = x_q.reshape(x.shape)
178
x_s = x_s.reshape(x.shape[:-1] + (x.shape[-1] // group_size,))
tests/kernels/quantization/test_block_fp8.py
+30
-9
@@ -20,7 +20,7 @@ from vllm.model_executor.layers.quantization.utils.fp8_utils import (
20
from vllm.platforms import current_platform
21
from vllm.utils.deep_gemm import (
22
fp8_gemm_nt,
23
- get_col_major_tma_aligned_tensor,
23
+ get_tma_aligned_size,
24
per_block_cast_to_fp8,
25
should_use_deepgemm_for_fp8_linear,
26
)
@@ -40,6 +40,8 @@ DTYPES = [torch.bfloat16] # [torch.half, torch.bfloat16, torch.float32]
40
NUM_TOKENS = [7, 2050]
41
D = [512, 4096, 5120, 13824]
42
GROUP_SIZE = [64, 128, 512]
43
+COLUMN_MAJOR_SCALES = [True, False]
44
+TMA_ALIGNED_SCALES = [True, False]
45
M = [1, 7, 8, 83, 84, 4096]
46
N = [128, 512, 7168, 7748, 13824]
47
K = [256, 3884, 4096, 13824, 16384]
@@ -63,20 +65,40 @@ def setup_cuda():
65
reason="This platform supports e4m3fnuz, not e4m3fn.",
66
)
67
@pytest.mark.parametrize(
66
- "num_tokens,d,dtype,group_size,seed",
67
- itertools.product(NUM_TOKENS, D, DTYPES, GROUP_SIZE, SEEDS),
68
+ "num_tokens,d,dtype,group_size,column_major_scales,tma_aligned_scales,seed",
69
+ itertools.product(
70
+ NUM_TOKENS,
71
+ D,
72
+ DTYPES,
73
+ GROUP_SIZE,
74
+ COLUMN_MAJOR_SCALES,
75
+ TMA_ALIGNED_SCALES,
76
+ SEEDS,
77
+ ),
78
)
79
@torch.inference_mode()
70
-def test_per_token_group_quant_fp8(num_tokens, d, dtype, group_size, seed):
80
+def test_per_token_group_quant_fp8(
81
+ num_tokens, d, dtype, group_size, column_major_scales, tma_aligned_scales, seed
82
+):
83
torch.manual_seed(seed)
84
x = torch.rand(num_tokens, d, dtype=dtype)
85
86
ref_out, ref_scale = native_per_token_group_quant_fp8(x, group_size)
75
- out, scale = per_token_group_quant_fp8(x, group_size)
87
+ out, scale = per_token_group_quant_fp8(
88
+ x,
89
+ group_size,
90
+ column_major_scales=column_major_scales,
91
+ tma_aligned_scales=tma_aligned_scales,
92
+ )
93
94
assert torch.allclose(out.to(torch.float32), ref_out.to(torch.float32), rtol=0.15)
95
assert torch.allclose(scale, ref_scale)
96
97
+ if column_major_scales:
98
+ assert scale.stride()[-2] == 1
99
+ if tma_aligned_scales:
100
+ assert scale.stride()[-1] == get_tma_aligned_size(num_tokens, 4)
101
+
102
103
@pytest.mark.parametrize(
104
"M,N,K,block_size,out_dtype,seed",
@@ -186,7 +208,9 @@ def test_w8a8_block_fp8_deep_gemm_matmul(M, N, K, block_size, out_dtype, seed):
208
):
209
pytest.skip(f"Skipping test; invalid size {M}, {N}, {K}")
210
189
- A_fp8, As_fp8 = per_token_group_quant_fp8(A_fp32, block_size[1])
211
+ A_fp8, As_fp8 = per_token_group_quant_fp8(
212
+ A_fp32, block_size[1], column_major_scales=True, tma_aligned_scales=True
213
+ )
214
B_fp8, Bs_fp8 = per_block_cast_to_fp8(B_fp32, block_size=block_size)
215
216
As = As_fp8.to(torch.float32)
@@ -194,9 +218,6 @@ def test_w8a8_block_fp8_deep_gemm_matmul(M, N, K, block_size, out_dtype, seed):
218
219
ref_out = native_w8a8_block_matmul(A_fp8, B_fp8, As, Bs, block_size, out_dtype)
220
197
- # Transpose earlier so that the testing will not trigger transposing kernels
198
- As_fp8 = get_col_major_tma_aligned_tensor(As_fp8)
199
-
221
out = torch.zeros((M, N), device="cuda", dtype=out_dtype)
222
223
assert As_fp8.shape == (M, (K + 127) // 128), (
tests/kernels/quantization/test_per_token_group_quant.py
+6
-2
@@ -8,13 +8,16 @@ import torch
8
from vllm.model_executor.layers.quantization.utils import fp8_utils, int8_utils
9
10
11
-@pytest.mark.parametrize("shape", [(32, 128), (64, 256), (16, 512)])
11
+@pytest.mark.parametrize(
12
+ "shape", [(31, 128), (32, 128), (63, 256), (64, 256), (16, 512)]
13
+)
14
@pytest.mark.parametrize("column_major", [False, True])
15
+@pytest.mark.parametrize("tma_aligned", [False, True])
16
@pytest.mark.parametrize("scale_ue8m0", [False, True])
17
@pytest.mark.parametrize("group_size", [64, 128])
18
@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available")
19
def test_per_token_group_quant_fp8(
17
- shape, column_major: bool, scale_ue8m0: bool, group_size: int
20
+ shape, column_major: bool, tma_aligned: bool, scale_ue8m0: bool, group_size: int
21
):
22
device = "cuda"
23
@@ -28,6 +31,7 @@ def test_per_token_group_quant_fp8(
31
x,
32
group_size,
33
column_major_scales=column_major,
34
+ tma_aligned_scales=tma_aligned,
35
use_ue8m0=scale_ue8m0,
36
)
37
vllm/model_executor/layers/quantization/input_quant_fp8.py
+5
@@ -36,6 +36,7 @@ class QuantFP8(CustomOp):
36
group_shape: GroupShape,
37
num_token_padding: int | None = None,
38
column_major_scales: bool = False,
39
+ tma_aligned_scales: bool = False,
40
use_ue8m0: bool | None = None, # for Torch compile
41
):
42
"""
@@ -44,6 +45,8 @@ class QuantFP8(CustomOp):
45
PER_CHANNEL, or arbitrary block size)
46
:param num_token_padding: Pad the token dimension of output to this
47
size
48
+ :param tma_aligned_scales: For group quantization, output scales in
49
+ TMA-aligned layout
50
:param column_major_scales: For group quantization, output scales in
51
column major format
52
"""
@@ -53,6 +56,7 @@ class QuantFP8(CustomOp):
56
self.use_per_token_if_dynamic = group_shape == GroupShape.PER_TOKEN
57
self.num_token_padding = num_token_padding
58
self.column_major_scales = column_major_scales
59
+ self.tma_aligned_scales = tma_aligned_scales
60
self.use_ue8m0 = use_ue8m0
61
62
self.use_aiter = rocm_aiter_ops.is_linear_fp8_enabled()
@@ -82,6 +86,7 @@ class QuantFP8(CustomOp):
86
x,
87
group_size=self.group_size,
88
column_major_scales=self.column_major_scales,
89
+ tma_aligned_scales=self.tma_aligned_scales,
90
dtype=_FP8_DTYPE,
91
use_ue8m0=self.use_ue8m0,
92
)
vllm/model_executor/layers/quantization/utils/fp8_utils.py
+23
-3
@@ -35,6 +35,7 @@ from vllm.triton_utils import tl, triton
35
from vllm.utils.deep_gemm import (
36
DeepGemmQuantScaleFMT,
37
fp8_gemm_nt,
38
+ get_tma_aligned_size,
39
is_deep_gemm_e8m0_used,
40
is_deep_gemm_supported,
41
should_use_deepgemm_for_fp8_linear,
@@ -378,6 +379,7 @@ class W8A8BlockFp8LinearOp:
379
False,
380
self.act_quant_group_shape,
381
column_major_scales=True,
382
+ tma_aligned_scales=True,
383
use_ue8m0=self.use_deep_gemm_e8m0,
384
)
385
if self.is_deep_gemm_supported
@@ -868,6 +870,7 @@ def per_token_group_quant_fp8(
870
eps: float = 1e-10,
871
dtype: torch.dtype | None = None,
872
column_major_scales: bool = False,
873
+ tma_aligned_scales: bool = False,
874
out_q: torch.Tensor | None = None,
875
use_ue8m0: bool | None = None,
876
) -> tuple[torch.Tensor, torch.Tensor]:
@@ -878,9 +881,10 @@ def per_token_group_quant_fp8(
881
x: The input tensor with ndim >= 2.
882
group_size: The group size used for quantization.
883
eps: The minimum to avoid dividing zero.
881
- dtype: The dype of output tensor. Note that only `torch.float8_e4m3fn`
884
+ dtype: The dtype of output tensor. Note that only `torch.float8_e4m3fn`
885
is supported for now.
886
column_major_scales: Outputs scales in column major.
887
+ tma_aligned_scales: Outputs scales in TMA-aligned layout.
888
out_q: Optional output tensor. If not provided, function will create.
889
Returns:
890
tuple[torch.Tensor, torch.Tensor]: The quantized tensor and the
@@ -904,8 +908,24 @@ def per_token_group_quant_fp8(
908
909
# Allocate the scale tensor in either row- or column-major format.
910
if column_major_scales:
907
- shape = (x.shape[-1] // group_size,) + x.shape[:-1]
908
- x_s = torch.empty(shape, device=x.device, dtype=torch.float32).permute(-1, -2)
911
+ if tma_aligned_scales:
912
+ m = x.shape[-2]
913
+ sf_k = x.shape[-1] // group_size
914
+ tma_aligned_m = get_tma_aligned_size(m, 4)
915
+ shape = x.shape[:-2] + (m, sf_k)
916
+ stride = (
917
+ (1, tma_aligned_m)
918
+ if x.dim() == 2
919
+ else (tma_aligned_m * sf_k, 1, tma_aligned_m)
920
+ )
921
+ x_s = torch.empty_strided(
922
+ shape, stride, device=x.device, dtype=torch.float32
923
+ )
924
+ else:
925
+ shape = x.shape[:-2] + (x.shape[-1] // group_size, x.shape[-2])
926
+ x_s = torch.empty(shape, device=x.device, dtype=torch.float32).permute(
927
+ -1, -2
928
+ )
929
else:
930
shape = x.shape[:-1] + (x.shape[-1] // group_size,)
931
x_s = torch.empty(shape, device=x.device, dtype=torch.float32)
vllm/utils/deep_gemm.py
+5
@@ -340,6 +340,11 @@ def _align(x: int, y: int) -> int:
340
return cdiv(x, y) * y
341
342
343
+# Taken from https://github.com/deepseek-ai/DeepGEMM/blob/v2.1.1/csrc/utils/math.hpp#L19
344
+def get_tma_aligned_size(x: int, element_size: int):
345
+ return _align(x, 16 // element_size)
346
+
347
+
348
DEFAULT_BLOCK_SIZE = [128, 128]
349
350