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, 23:45 UTC
[FEAT] [ROCm] [AITER]: Add AITER HIP block quant kernel (#21242)
TJian committed
Jul 27, 2025 at 22:07 UTC
e626d286f5ac997bca5bd635c3db9e457fe91df9
1 file changed
+13
-2
vllm/model_executor/layers/quantization/utils/fp8_utils.py
+13
-2
@@ -82,6 +82,13 @@ if current_platform.is_rocm():
82
fake_impl=rocm_aiter_gemm_w8a8_blockscale_fake,
83
dispatch_key=current_platform.dispatch_key,
84
)
85
+ if (envs.VLLM_ROCM_USE_AITER and envs.VLLM_ROCM_USE_AITER_LINEAR
86
+ and current_platform.is_fp8_fnuz()):
87
+
88
+ import aiter as rocm_aiter
89
+ from aiter import get_hip_quant
90
+
91
+ aiter_per1x128_quant = get_hip_quant(rocm_aiter.QuantType.per_1x128)
92
93
94
def dispatch_w8a8_blockscale_func(
@@ -178,8 +185,12 @@ def apply_w8a8_block_fp8_linear(
185
block_size, input.dtype)
186
187
else:
181
- q_input, x_scale = per_token_group_quant_fp8(
182
- input_2d, block_size[1], column_major_scales=use_cutlass)
188
+ if use_aiter_and_is_supported:
189
+ q_input, x_scale = aiter_per1x128_quant(
190
+ input_2d.contiguous(), quant_dtype=rocm_aiter.dtypes.fp8)
191
+ else:
192
+ q_input, x_scale = per_token_group_quant_fp8(
193
+ input_2d, block_size[1], column_major_scales=use_cutlass)
194
195
output = w8a8_blockscale_func(q_input, weight, x_scale, weight_scale,
196
block_size, input.dtype)