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 13, 17:00 UTC
[ROCm] triton fp8 kernel (#27058)
Signed-off-by: Aleksandr Malyshev <maleksan@amd.com> Co-authored-by: Aleksandr Malyshev <maleksan@amd.com> Co-authored-by: Gregory Shtrasberg <156009573+gshtras@users.noreply.github.com>
Aleksandr Malyshev committed
Nov 6, 2025 at 11:46 UTC
449de9001af69592618516b298aa1c5f321ded34
1 file changed
+65
-34
vllm/model_executor/layers/quantization/utils/fp8_utils.py
+65
-34
@@ -69,30 +69,67 @@ def cutlass_scaled_mm(
69
70
71
def rocm_aiter_gemm_w8a8_blockscale_impl(
72
- A: torch.Tensor,
73
- B: torch.Tensor,
74
- As: torch.Tensor,
75
- Bs: torch.Tensor,
76
- block_size: list[int],
72
+ input_2d: torch.Tensor,
73
+ weight: torch.Tensor,
74
+ input_scale: torch.Tensor,
75
+ weight_scale: torch.Tensor,
76
+ group_size: int,
77
output_dtype: torch.dtype = torch.float16,
78
) -> torch.Tensor:
79
- import aiter as rocm_aiter
79
+ def is_aiter_triton_kernel_tuned(n, k):
80
+ return (n, k) in [
81
+ (1024, 8192),
82
+ (2112, 7168),
83
+ (3072, 1536),
84
+ (32768, 8192),
85
+ (4096, 7168),
86
+ (4608, 7168),
87
+ (512, 7168),
88
+ (7168, 2048),
89
+ (7168, 256),
90
+ (8192, 1024),
91
+ (8192, 32768),
92
+ ]
93
+
94
+ n, k = weight.shape
95
+ if input_scale is not None:
96
+ q_input = input_2d
97
+ elif not current_platform.is_fp8_fnuz() and is_aiter_triton_kernel_tuned(n, k):
98
+ from aiter.ops.triton.gemm_a8w8_blockscale import gemm_a8w8_blockscale
99
+
100
+ # MI350 case uses triton kernel
101
+ q_input, input_scale = per_token_group_quant_fp8(
102
+ input_2d,
103
+ group_size,
104
+ column_major_scales=False,
105
+ use_ue8m0=False,
106
+ )
107
+ else:
108
+ # MI300 uses tuned AITER ASM/C++ kernel
109
+ import aiter as rocm_aiter
110
+ from aiter import gemm_a8w8_blockscale, get_hip_quant
111
81
- return rocm_aiter.gemm_a8w8_blockscale(A, B, As, Bs, dtype=output_dtype)
112
+ aiter_per1x128_quant = get_hip_quant(rocm_aiter.QuantType.per_1x128)
113
+ q_input, input_scale = aiter_per1x128_quant(
114
+ input_2d.contiguous(), quant_dtype=rocm_aiter.dtypes.fp8
115
+ )
116
+
117
+ return gemm_a8w8_blockscale(
118
+ q_input, weight, input_scale, weight_scale, dtype=output_dtype
119
+ )
120
121
122
def rocm_aiter_gemm_w8a8_blockscale_fake(
85
- A: torch.Tensor,
86
- B: torch.Tensor,
87
- As: torch.Tensor,
88
- Bs: torch.Tensor,
89
- block_size: list[int],
123
+ input_2d: torch.Tensor,
124
+ weight: torch.Tensor,
125
+ input_scale: torch.Tensor,
126
+ weight_scale: torch.Tensor,
127
+ group_size: int,
128
output_dtype: torch.dtype = torch.float16,
129
) -> torch.Tensor:
92
- m = A.shape[0]
93
- n = B.shape[0]
94
- Y = torch.empty(m, n, dtype=output_dtype, device=A.device)
95
- return Y
130
+ m = input_2d.shape[0]
131
+ n = weight.shape[0]
132
+ return torch.empty(m, n, dtype=output_dtype, device=input_2d.device)
133
134
135
if current_platform.is_rocm():
@@ -101,15 +138,6 @@ if current_platform.is_rocm():
138
op_func=rocm_aiter_gemm_w8a8_blockscale_impl,
139
fake_impl=rocm_aiter_gemm_w8a8_blockscale_fake,
140
)
104
- if (
105
- envs.VLLM_ROCM_USE_AITER
106
- and envs.VLLM_ROCM_USE_AITER_LINEAR
107
- and current_platform.is_fp8_fnuz()
108
- ):
109
- import aiter as rocm_aiter
110
- from aiter import get_hip_quant
111
-
112
- aiter_per1x128_quant = get_hip_quant(rocm_aiter.QuantType.per_1x128)
141
142
143
# TODO we should be able to change the type of block_size to GroupShape
@@ -293,7 +321,9 @@ class W8A8BlockFp8LinearOp:
321
):
322
output = self._run_deepgemm(input_2d, weight, weight_scale)
323
else:
296
- output = self.w8a8_blockscale_op(input_2d, weight, weight_scale)
324
+ output = self.w8a8_blockscale_op(
325
+ input_2d, weight, weight_scale, input_scale
326
+ )
327
328
if bias is not None:
329
output = output + bias
@@ -322,7 +352,9 @@ class W8A8BlockFp8LinearOp:
352
input_2d: torch.Tensor,
353
weight: torch.Tensor,
354
weight_scale: torch.Tensor,
355
+ input_scale: torch.Tensor | None = None,
356
) -> torch.Tensor:
357
+ assert input_scale is None
358
assert self.input_quant_op is not None
359
q_input, input_scale = self.input_quant_op(input_2d)
360
if self.is_hopper:
@@ -350,17 +382,15 @@ class W8A8BlockFp8LinearOp:
382
input_2d: torch.Tensor,
383
weight: torch.Tensor,
384
weight_scale: torch.Tensor,
385
+ input_scale: torch.Tensor | None = None,
386
) -> torch.Tensor:
387
assert self.act_quant_group_shape == GroupShape(1, 128)
355
- q_input, input_scale = aiter_per1x128_quant(
356
- input_2d.contiguous(), quant_dtype=rocm_aiter.dtypes.fp8
357
- )
388
return torch.ops.vllm.rocm_aiter_gemm_w8a8_blockscale(
359
- q_input,
389
+ input_2d,
390
weight,
391
input_scale,
392
weight_scale,
363
- list(self.weight_group_shape),
393
+ self.act_quant_group_shape.col,
394
input_2d.dtype,
395
)
396
@@ -369,7 +399,9 @@ class W8A8BlockFp8LinearOp:
399
input_2d: torch.Tensor,
400
weight: torch.Tensor,
401
weight_scale: torch.Tensor,
402
+ input_scale: torch.Tensor | None = None,
403
) -> torch.Tensor:
404
+ assert input_scale is None
405
assert self.input_quant_op is not None
406
q_input, input_scale = self.input_quant_op(input_2d)
407
return torch.ops.vllm.w8a8_triton_block_scaled_mm_func(
@@ -391,6 +423,7 @@ class W8A8BlockFp8LinearOp:
423
torch.Tensor,
424
torch.Tensor,
425
torch.Tensor,
426
+ torch.Tensor | None,
427
],
428
torch.Tensor,
429
],
@@ -939,13 +972,11 @@ def requant_weight_ue8m0_inplace(
972
973
974
def check_aiter_fp8_linear_support() -> bool:
942
- """AITER is only supported on ROCm and only for FP8_FNUZ
943
- and at the moment are MI300 series"""
975
+ """AITER is only supported on ROCm for MI3XX"""
976
return (
977
current_platform.is_rocm()
978
and envs.VLLM_ROCM_USE_AITER
979
and envs.VLLM_ROCM_USE_AITER_LINEAR
948
- and current_platform.is_fp8_fnuz()
980
)
981
982