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 4, 12:30 UTC
[ROCm] Use quant_dtype in per_token_quant instead of hardcoded FP8 (#39121)
Signed-off-by: Bortlesboat <bortstheboat@gmail.com>
Andrew Barnes committed
Apr 29, 2026 at 16:46 UTC
169988a3c0e0912fc20be2d104a4b76a51ad9fa4
1 file changed
+2
-2
vllm/_aiter_ops.py
+2
-2
@@ -778,7 +778,7 @@ def _rocm_aiter_per_token_quant_impl(
778
assert quant_dtype in [torch.int8, FP8_DTYPE]
779
780
out_shape = x.shape
781
- out = torch.empty(x.shape, dtype=FP8_DTYPE, device=x.device)
781
+ out = torch.empty(x.shape, dtype=quant_dtype, device=x.device)
782
if scale is None:
783
scale = torch.empty((*out_shape[:-1], 1), dtype=torch.float32, device=x.device)
784
dynamic_per_token_scaled_quant(
@@ -798,7 +798,7 @@ def _rocm_aiter_per_token_quant_fake(
798
) -> tuple[torch.Tensor, torch.Tensor]:
799
out_shape = x.shape
800
return (
801
- torch.empty(x.shape, dtype=FP8_DTYPE, device=x.device),
801
+ torch.empty(x.shape, dtype=quant_dtype, device=x.device),
802
torch.empty((*out_shape[:-1], 1), dtype=torch.float32, device=x.device),
803
)
804