[ROCm][DSv4.1][Perf] Use the shared prefill chunk plan in the ROCm sparse prefill (#58539)
Signed-off-by: Zheng Gong <zgong@amd.com>
Zheng Gong committed
Sep 30, 2026 at 15:31 UTC
8ee40690331c25f19ecad49696dd54930fcab696
1 file changed
+11
-10
vllm/models/deepseek_v41/amd/rocm.py
+11
-10
@@ -1379,7 +1379,6 @@ class DeepseekV41ROCMAiterMLAAttention(DeepseekV4Attention):
1379
1380
swa_only = attn_metadata is None
1381
1382
- num_prefills = swa_metadata.num_prefills
1382
num_prefill_tokens = swa_metadata.num_prefill_tokens
1383
num_decodes = swa_metadata.num_decodes
1384
num_decode_tokens = swa_metadata.num_decode_tokens
@@ -1408,21 +1407,23 @@ class DeepseekV41ROCMAiterMLAAttention(DeepseekV4Attention):
1407
N = 0
1408
1409
M = N + self.window_size + self.max_num_batched_tokens
1411
- num_chunks = (num_prefills + self.PREFILL_CHUNK_SIZE - 1) // (
1412
- self.PREFILL_CHUNK_SIZE
1410
+ chunk_plan = swa_metadata.get_prefill_chunk_plan(
1411
+ compress_ratio=self.compress_ratio,
1412
+ prefill_chunk_size=self.PREFILL_CHUNK_SIZE,
1413
+ has_compressed=not swa_only,
1414
)
1415
+ assert chunk_plan, "prefill chunk plan must be non-empty when num_prefills > 0"
1416
1417
workspace = current_workspace_manager().get_simultaneous(
1418
*self._prefill_workspace_shapes(M, num_prefill_tokens, q)
1419
)
1418
- kv = workspace[0]
1420
+ kv_rows = workspace[0].view(-1, q.shape[-1])
1421
if mxfp8_out is not None:
1422
output = workspace[1]
1423
assert output is not None
1422
- for chunk_idx in range(num_chunks):
1423
- chunk_start = chunk_idx * self.PREFILL_CHUNK_SIZE
1424
- chunk_end = min(chunk_start + self.PREFILL_CHUNK_SIZE, num_prefills)
1424
+ for chunk_start, chunk_end, chunk_N, chunk_M in chunk_plan:
1425
chunk_size = chunk_end - chunk_start
1426
+ kv = kv_rows[: chunk_size * chunk_M].view(chunk_size, chunk_M, -1)
1427
if not swa_only:
1428
assert attn_metadata is not None
1429
assert compressed_k_cache is not None
@@ -1447,7 +1448,7 @@ class DeepseekV41ROCMAiterMLAAttention(DeepseekV4Attention):
1448
gather_lens=gather_lens[chunk_start:chunk_end],
1449
block_table=swa_block_table[chunk_start:chunk_end],
1450
block_size=swa_metadata.block_size,
1450
- offset=N,
1451
+ offset=chunk_N,
1452
use_fnuz=current_platform.is_fp8_fnuz(),
1453
)
1454
@@ -1468,8 +1469,8 @@ class DeepseekV41ROCMAiterMLAAttention(DeepseekV4Attention):
1469
self.window_size,
1470
self.compress_ratio,
1471
top_k,
1471
- M,
1472
- N,
1472
+ chunk_M,
1473
+ chunk_N,
1474
)
1475
rocm_sparse_attn_prefill(
1476
q=q[query_start:query_end],