[Bugfix] Fix turboquant FP8 cast failure for BF16 models on Ampere GPUs (#39988)
Signed-off-by: Xu Zhou <xuzhou9417@163.com> Co-authored-by: Xu Zhou <xuzhou9417@163.com> Co-authored-by: Hoseung Kim <ghyutjik123@gmail.com>
XuZhou committed
Jul 10, 2026 at 21:55 UTC
b12cca6a2352ea6926864c925246eeb581be8caa
1 file changed
+1
-1
vllm/v1/attention/ops/triton_turboquant_store.py
+1
-1
@@ -188,7 +188,7 @@ def _tq_fused_store_fp8(
188
# ── FP8 KEY: cast to FP8 in-kernel and store ─────────────────
189
d_offs = tl.arange(0, BLOCK_D)
190
d_mask = d_offs < D
191
- k_vals = tl.load(Key_ptr + base + d_offs, mask=d_mask, other=0.0)
191
+ k_vals = tl.load(Key_ptr + base + d_offs, mask=d_mask, other=0.0).to(tl.float32)
192
k_fp8 = k_vals.to(tl.float8e4b15) if FP8_E4B15 else k_vals.to(tl.float8e4nv)
193
k_bytes = k_fp8.to(tl.uint8, bitcast=True)
194
tl.store(KV_cache_ptr + slot_base + d_offs, k_bytes, mask=d_mask)