[ROCm][DSv4][Perf] Use the shared prefill chunk plan in the ROCm sparse prefill (#58405)
Signed-off-by: Zheng Gong <zgong@amd.com> Co-authored-by: Cursor <cursoragent@cursor.com>
Zheng Gong committed
Sep 30, 2026 at 15:26 UTC
fcd317cf33b5113768a8ba1570f5211627f6fd63
1 file changed
+11
-16
vllm/models/deepseek_v4/amd/rocm.py
+11
-16
@@ -1426,7 +1426,6 @@ class DeepseekV4ROCMAiterMLAAttention(DeepseekV4Attention):
1426
1427
swa_only = attn_metadata is None
1428
1429
- num_prefills = swa_metadata.num_prefills
1429
num_prefill_tokens = swa_metadata.num_prefill_tokens
1430
num_decodes = swa_metadata.num_decodes
1431
num_decode_tokens = swa_metadata.num_decode_tokens
@@ -1458,26 +1457,22 @@ class DeepseekV4ROCMAiterMLAAttention(DeepseekV4Attention):
1457
topk_indices = attn_metadata.c128a_prefill_topk_indices
1458
assert topk_indices is not None
1459
top_k = topk_indices.shape[-1]
1461
- N = (self.max_model_len + self.compress_ratio - 1) // self.compress_ratio
1460
else:
1461
assert self.topk_indices_buffer is not None
1462
topk_indices = self.topk_indices_buffer[num_decode_tokens:]
1463
top_k = 0
1466
- N = 0
1464
1468
- M = N + self.window_size + self.max_num_batched_tokens
1469
- num_chunks = (num_prefills + self.PREFILL_CHUNK_SIZE - 1) // (
1470
- self.PREFILL_CHUNK_SIZE
1465
+ chunk_plan = swa_metadata.get_prefill_chunk_plan(
1466
+ compress_ratio=self.compress_ratio,
1467
+ prefill_chunk_size=self.PREFILL_CHUNK_SIZE,
1468
)
1472
-
1469
+ assert chunk_plan, "prefill chunk plan must be non-empty when num_prefills > 0"
1470
workspace_manager = current_workspace_manager()
1474
- kv = workspace_manager.get_simultaneous(
1475
- ((self.PREFILL_CHUNK_SIZE, M, q.shape[-1]), torch.bfloat16),
1476
- )[0]
1477
- for chunk_idx in range(num_chunks):
1478
- chunk_start = chunk_idx * self.PREFILL_CHUNK_SIZE
1479
- chunk_end = min(chunk_start + self.PREFILL_CHUNK_SIZE, num_prefills)
1471
+ for chunk_start, chunk_end, chunk_N, chunk_M in chunk_plan:
1472
chunk_size = chunk_end - chunk_start
1473
+ kv = workspace_manager.get_simultaneous(
1474
+ ((chunk_size, chunk_M, q.shape[-1]), torch.bfloat16),
1475
+ )[0]
1476
if not swa_only:
1477
assert attn_metadata is not None
1478
assert compressed_k_cache is not None
@@ -1502,7 +1497,7 @@ class DeepseekV4ROCMAiterMLAAttention(DeepseekV4Attention):
1497
gather_lens=gather_lens[chunk_start:chunk_end],
1498
block_table=swa_block_table[chunk_start:chunk_end],
1499
block_size=swa_metadata.block_size,
1505
- offset=N,
1500
+ offset=chunk_N,
1501
use_fnuz=current_platform.is_fp8_fnuz(),
1502
)
1503
@@ -1523,8 +1518,8 @@ class DeepseekV4ROCMAiterMLAAttention(DeepseekV4Attention):
1518
self.window_size,
1519
self.compress_ratio,
1520
top_k,
1526
- M,
1527
- N,
1521
+ chunk_M,
1522
+ chunk_N,
1523
max_image_tokens=self.max_image_tokens,
1524
left_visible=(
1525
left_visible[query_start:query_end]