[ROCm][Kernel][DSV4] Remove tl.constexpr to avoid cold-compile churn in indexer gather kernel (#56153)
Signed-off-by: xawang <xawang@amd.com>
amd-xavierwang committed
Sep 11, 2026 at 17:34 UTC
2d75e586fcaf88231f7a75f482dc8bfb5ad9da10
1 file changed
+8
-8
vllm/v1/attention/ops/rocm_aiter_mla_sparse.py
+8
-8
@@ -292,16 +292,16 @@ def _cp_gather_indexer_quant_cache_kernel(
292
HEAD_DIM: tl.constexpr,
293
BLOCK_TILE_SIZE: tl.constexpr,
294
HEAD_TILE_SIZE: tl.constexpr,
295
- NUM_TOKENS: tl.constexpr,
296
- NUM_BATCHES: tl.constexpr,
297
- BLOCK_TABLE_WIDTH: tl.constexpr,
298
- NUM_BLOCKS: tl.constexpr,
295
+ num_tokens,
296
+ num_batches,
297
+ block_table_width,
298
+ num_blocks,
299
):
300
tid = tl.program_id(0)
301
offset = tl.arange(0, HEAD_DIM)
302
- valid_tid = tid < NUM_TOKENS
302
+ valid_tid = tid < num_tokens
303
batch_id = tl.load(token_to_seq_ptr + tid, mask=valid_tid, other=-1)
304
- valid_batch = (batch_id >= 0) & (batch_id < NUM_BATCHES)
304
+ valid_batch = (batch_id >= 0) & (batch_id < num_batches)
305
safe_batch_id = tl.where(valid_batch, batch_id, 0)
306
batch_start = tl.load(cu_seqlen_ptr + safe_batch_id, mask=valid_batch, other=0)
307
batch_end = tl.load(cu_seqlen_ptr + safe_batch_id + 1, mask=valid_batch, other=0)
@@ -314,7 +314,7 @@ def _cp_gather_indexer_quant_cache_kernel(
314
valid_block_table = (
315
valid_token
316
& (block_table_id >= 0)
317
- & (block_table_id < BLOCK_TABLE_WIDTH)
317
+ & (block_table_id < block_table_width)
318
& (block_offset >= 0)
319
& (block_offset < block_size)
320
)
@@ -323,7 +323,7 @@ def _cp_gather_indexer_quant_cache_kernel(
323
block_id = tl.load(
324
block_table_ptr + block_table_offset, mask=valid_block_table, other=-1
325
)
326
- valid_block = valid_block_table & (block_id >= 0) & (block_id < NUM_BLOCKS)
326
+ valid_block = valid_block_table & (block_id >= 0) & (block_id < num_blocks)
327
# The packed KV layout makes per-block strides large
328
# enough that block_id * stride can exceed 32-bit range.
329
safe_block_id = tl.where(valid_block, block_id, 0).to(tl.int64)