[ROCm] Resolve the indexer fp8 cache dtype once at import (#53792)
Signed-off-by: Sriram Kumar <sriramkumar.kishorekumar@amd.com> Co-authored-by: Douglas Lehr <91553416+dllehr-amd@users.noreply.github.com>
Sriram Kumar committed
Sep 17, 2026 at 20:51 UTC
ff6b5808c4571a37ba3d27f74d74ef43c3a14245
1 file changed
+9
-11
vllm/v1/attention/ops/rocm_aiter_mla_sparse.py
+9
-11
@@ -29,6 +29,8 @@ else:
29
30
logger = init_logger(__name__)
31
32
+FP8_DTYPE = current_platform.fp8_dtype()
33
+
34
35
@functools.cache
36
def _get_aiter_topk_ops() -> tuple[Callable[..., None], Callable[..., None]] | None:
@@ -260,8 +262,7 @@ def indexer_k_quant_and_cache_triton(
262
# In real layout, we store the first portion as kv cache value
263
# and second portion as kv cache scale
264
kv_cache = kv_cache.view(num_blocks, -1)
263
- fp8_dtype = current_platform.fp8_dtype()
264
- kv_cache_value = kv_cache[:, : block_size * head_dim].view(fp8_dtype)
265
+ kv_cache_value = kv_cache[:, : block_size * head_dim].view(FP8_DTYPE)
266
kv_cache_scale = kv_cache[:, block_size * head_dim :].view(torch.float32)
267
head_tile_size = head_tile_size // kv_cache.element_size()
268
layout = "NORMAL" if block_size == 1 else "SHUFFLE"
@@ -279,7 +280,7 @@ def indexer_k_quant_and_cache_triton(
280
layout,
281
block_tile_size,
282
head_tile_size,
282
- IS_FNUZ=current_platform.fp8_dtype() == torch.float8_e4m3fnuz,
283
+ IS_FNUZ=torch.float8_e4m3fnuz == FP8_DTYPE,
284
USE_UE8M0=scale_fmt == "ue8m0",
285
)
286
@@ -465,8 +466,7 @@ def cp_gather_indexer_k_quant_cache_triton(
466
num_blocks = k_cache.shape[0]
467
# we assume the kv cache already been split to 2 portion
468
k_cache = k_cache.view(num_blocks, -1)
468
- fp8_dtype = current_platform.fp8_dtype()
469
- k_cache_value = k_cache[:, : block_size * head_dim].view(fp8_dtype)
469
+ k_cache_value = k_cache[:, : block_size * head_dim].view(FP8_DTYPE)
470
k_cache_scale = k_cache[:, block_size * head_dim :].view(torch.float32)
471
grid = (num_tokens,)
472
k_fp8_scale = k_fp8_scale.view(torch.float32)
@@ -516,7 +516,6 @@ def fp8_paged_mqa_logits_torch(
516
):
517
from vllm.utils.math_utils import cdiv
518
519
- fp8_dtype = current_platform.fp8_dtype()
519
batch_size, next_n, _, dim = q.size()
520
if next_n == 1:
521
block_size = kv_cache.shape[1]
@@ -540,7 +539,7 @@ def fp8_paged_mqa_logits_torch(
539
cache = kv_cache_flat[pages]
540
scale_offset = block_size * dim
541
cache_value = (
543
- cache[..., :scale_offset].view(dtype=fp8_dtype).to(torch.float32)
542
+ cache[..., :scale_offset].view(dtype=FP8_DTYPE).to(torch.float32)
543
)
544
cache_scale = (
545
cache[..., scale_offset:].view(dtype=torch.float32).contiguous()
@@ -558,7 +557,7 @@ def fp8_paged_mqa_logits_torch(
557
kv_cache, scale = kv_cache[..., :dim], kv_cache[..., dim:]
558
scale = scale.contiguous().view(torch.float)
559
q = q.float()
561
- kv_cache = kv_cache.view(fp8_dtype).float() * scale
560
+ kv_cache = kv_cache.view(FP8_DTYPE).float() * scale
561
num_block, block_size, _, dim = kv_cache.size()
562
logits = torch.full(
563
[batch_size * next_n, max_model_len],
@@ -1021,7 +1020,6 @@ def rocm_aiter_sparse_attn_indexer(
1020
# careful! this will be None in dummy run
1021
forward_context = get_forward_context()
1022
attn_metadata = forward_context.attn_metadata
1024
- fp8_dtype = current_platform.fp8_dtype()
1023
from vllm.utils.torch_utils import _resolve_layer_name
1024
1025
k_cache_prefix = _resolve_layer_name(k_cache_prefix)
@@ -1037,7 +1035,7 @@ def rocm_aiter_sparse_attn_indexer(
1035
# Prefill k_fp8 and k_scale buffers, used by
1036
# rocm_aiter_sparse_attn_indexer's prefill path
1037
workspace_manager.get_simultaneous(
1040
- ((total_seq_lens, head_dim), fp8_dtype),
1038
+ ((total_seq_lens, head_dim), FP8_DTYPE),
1039
((total_seq_lens, 4), torch.uint8),
1040
)
1041
@@ -1117,7 +1115,7 @@ def rocm_aiter_sparse_attn_indexer(
1115
1116
workspace_manager = current_workspace_manager()
1117
k_fp8_full, k_scale_full = workspace_manager.get_simultaneous(
1120
- ((total_seq_lens, head_dim), fp8_dtype),
1118
+ ((total_seq_lens, head_dim), FP8_DTYPE),
1119
((total_seq_lens, 4), torch.uint8),
1120
)
1121
for chunk in prefill_metadata.chunks: