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 5, 09:45 UTC
[ROCm] Remove unnecessary fp8 roundtrip in gather cache NHD dequant (#39122)
Signed-off-by: Bortlesboat <bortstheboat@gmail.com>
Andrew Barnes committed
Apr 9, 2026 at 03:12 UTC
8a34c5087aa723eafd9995a3af814fcae8334c4d
1 file changed
+6
-4
vllm/v1/attention/backends/rocm_aiter_fa.py
+6
-4
@@ -112,10 +112,12 @@ if current_platform.is_rocm():
112
if DEQUANT:
113
k_scale = tl.load(k_scale_ptr)
114
v_scale = tl.load(v_scale_ptr)
115
- k_dtype = k_reg.dtype
116
- v_dtype = v_reg.dtype
117
- k_reg = (k_reg.to(tl.float32) * k_scale).to(k_dtype)
118
- v_reg = (v_reg.to(tl.float32) * v_scale).to(v_dtype)
115
+ k_reg = (k_reg.to(tl.float32) * k_scale).to(
116
+ key_ptr_offset.dtype.element_ty
117
+ )
118
+ v_reg = (v_reg.to(tl.float32) * v_scale).to(
119
+ value_ptr_offset.dtype.element_ty
120
+ )
121
tl.store(key_ptr_offset + col_offsets, k_reg)
122
tl.store(value_ptr_offset + col_offsets, v_reg)
123