[ROCm] [Bugfix] Bugfix ROCm Sparse Indexer (#46222)
Signed-off-by: tjtanaa <tunjian.tan@embeddedllm.com>
TJian committed
Jun 21, 2026 at 04:38 UTC
1bdf9810aae30ae0b7002ac9f98bb9520b34e631
1 file changed
+6
-2
vllm/v1/attention/ops/rocm_aiter_mla_sparse.py
+6
-2
@@ -58,7 +58,9 @@ def _indexer_k_quant_and_cache_kernel(
58
slot_id = tl.load(slot_mapping_ptr + tid)
59
if slot_id < 0:
60
return
61
- block_id = slot_id // block_size
61
+ # The packed KV layout makes per-block strides large
62
+ # enough that block_id * stride can exceed 32-bit range.
63
+ block_id = (slot_id // block_size).to(tl.int64)
64
block_offset = slot_id % block_size
65
tile_block_id = block_offset // BLOCK_TILE_SIZE
66
tile_block_offset = block_offset % BLOCK_TILE_SIZE
@@ -179,7 +181,9 @@ def _cp_gather_indexer_quant_cache_kernel(
181
block_table_ptr + block_table_offset, mask=valid_block_table, other=-1
182
)
183
valid_block = valid_block_table & (block_id >= 0) & (block_id < NUM_BLOCKS)
182
- safe_block_id = tl.where(valid_block, block_id, 0)
184
+ # The packed KV layout makes per-block strides large
185
+ # enough that block_id * stride can exceed 32-bit range.
186
+ safe_block_id = tl.where(valid_block, block_id, 0).to(tl.int64)
187
safe_block_offset = tl.where(valid_block, block_offset, 0)
188
tiled_block_offset = safe_block_offset % BLOCK_TILE_SIZE
189
if LAYOUT == "SHUFFLE":