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, 05:30 UTC
Aiter mha fp8 fix (#24991)
Signed-off-by: Doug Lehr <douglehr@amd.com> Co-authored-by: Doug Lehr <douglehr@amd.com>
Douglas Lehr committed
Sep 17, 2025 at 17:29 UTC
1a456c7c90afdb534d1203d7e4ea5747aada801c
2 files changed
+4
-4
vllm/attention/ops/rocm_aiter_paged_attn.py
+2
-2
@@ -81,8 +81,8 @@ class AITERPagedAttention(PagedAttention):
81
blocksparse_head_sliding_step=blocksparse_head_sliding_step)
82
83
if "fp8" in kv_cache_dtype:
84
- key_cache = key_cache.view(torch.float8_e4m3fnuz)
85
- value_cache = value_cache.view(torch.float8_e4m3fnuz)
84
+ key_cache = key_cache.view(current_platform.fp8_dtype())
85
+ value_cache = value_cache.view(current_platform.fp8_dtype())
86
87
if blocksparse_vert_stride is not None and blocksparse_vert_stride > 1:
88
# use blocksparse paged attention
vllm/v1/attention/backends/rocm_aiter_fa.py
+2
-2
@@ -479,8 +479,8 @@ class AiterFlashAttentionImpl(AttentionImpl):
479
)
480
481
if self.kv_cache_dtype.startswith("fp8"):
482
- key_cache = key_cache.view(torch.float8_e4m3fnuz)
483
- value_cache = value_cache.view(torch.float8_e4m3fnuz)
482
+ key_cache = key_cache.view(current_platform.fp8_dtype())
483
+ value_cache = value_cache.view(current_platform.fp8_dtype())
484
485
if not attn_metadata.use_cascade:
486
cu_seqlens_q = attn_metadata.query_start_loc