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, 13:15 UTC
[ROCm] Add missing gemm_a8w8_blockscale import (#28378)
Signed-off-by: Yong Hoon Shin <yhshin@meta.com>
Yong Hoon Shin committed
Nov 10, 2025 at 13:13 UTC
021143561fcffa9bee133d0b3bd311bc5cb3703c
1 file changed
+21
-20
vllm/model_executor/layers/quantization/utils/fp8_utils.py
+21
-20
@@ -316,38 +316,39 @@ class W8A8BlockFp8LinearOp:
316
assert self.act_quant_group_shape == GroupShape(1, 128)
317
318
n, k = weight.shape
319
- if input_scale is not None:
320
- q_input = input_2d
319
322
- # MI350 case uses triton kernel
323
- if (
320
+ use_triton = (
321
not current_platform.is_fp8_fnuz()
322
and rocm_aiter_ops.is_triton_gemm_w8a8_tuned(n, k)
326
- ):
323
+ )
324
+
325
+ if use_triton:
326
+ gemm_a8w8_blockscale_op = rocm_aiter_ops.triton_gemm_a8w8_blockscale
327
+ else:
328
+ gemm_a8w8_blockscale_op = rocm_aiter_ops.gemm_w8a8_blockscale
329
+
330
+ if input_scale is not None:
331
+ q_input = input_2d
332
+ # MI350 case uses triton kernel
333
+ elif use_triton:
334
q_input, input_scale = per_token_group_quant_fp8(
335
input_2d,
336
self.act_quant_group_shape.col,
337
column_major_scales=False,
338
use_ue8m0=False,
339
)
333
- return rocm_aiter_ops.triton_gemm_a8w8_blockscale(
334
- q_input,
335
- weight,
336
- input_scale,
337
- weight_scale,
338
- input_2d.dtype,
339
- )
340
-
340
# MI300 uses tuned AITER ASM/C++ kernel
341
else:
342
q_input, input_scale = rocm_aiter_ops.per_1x128_fp8_quant(input_2d)
344
- return rocm_aiter_ops.gemm_w8a8_blockscale(
345
- q_input,
346
- weight,
347
- input_scale,
348
- weight_scale,
349
- input_2d.dtype,
350
- )
343
+
344
+ return gemm_a8w8_blockscale_op(
345
+ q_input,
346
+ weight,
347
+ input_scale,
348
+ weight_scale,
349
+ list(self.weight_group_shape),
350
+ output_dtype=input_2d.dtype,
351
+ )
352
353
def _run_triton(
354
self,