From 8ee40690331c25f19ecad49696dd54930fcab696 Mon Sep 17 00:00:00 2001 From: Zheng Gong Date: Wed, 30 Sep 2026 15:31:51 +0800 Subject: [PATCH] [ROCm][DSv4.1][Perf] Use the shared prefill chunk plan in the ROCm sparse prefill (#58539) Signed-off-by: Zheng Gong --- vllm/models/deepseek_v41/amd/rocm.py | 21 +++++++++++---------- 1 file changed, 11 insertions(+), 10 deletions(-) diff --git a/vllm/models/deepseek_v41/amd/rocm.py b/vllm/models/deepseek_v41/amd/rocm.py index fe020f4061..99fcfa6702 100644 --- a/vllm/models/deepseek_v41/amd/rocm.py +++ b/vllm/models/deepseek_v41/amd/rocm.py @@ -1379,7 +1379,6 @@ class DeepseekV41ROCMAiterMLAAttention(DeepseekV4Attention): swa_only = attn_metadata is None - num_prefills = swa_metadata.num_prefills num_prefill_tokens = swa_metadata.num_prefill_tokens num_decodes = swa_metadata.num_decodes num_decode_tokens = swa_metadata.num_decode_tokens @@ -1408,21 +1407,23 @@ class DeepseekV41ROCMAiterMLAAttention(DeepseekV4Attention): N = 0 M = N + self.window_size + self.max_num_batched_tokens - num_chunks = (num_prefills + self.PREFILL_CHUNK_SIZE - 1) // ( - self.PREFILL_CHUNK_SIZE + chunk_plan = swa_metadata.get_prefill_chunk_plan( + compress_ratio=self.compress_ratio, + prefill_chunk_size=self.PREFILL_CHUNK_SIZE, + has_compressed=not swa_only, ) + assert chunk_plan, "prefill chunk plan must be non-empty when num_prefills > 0" workspace = current_workspace_manager().get_simultaneous( *self._prefill_workspace_shapes(M, num_prefill_tokens, q) ) - kv = workspace[0] + kv_rows = workspace[0].view(-1, q.shape[-1]) if mxfp8_out is not None: output = workspace[1] assert output is not None - for chunk_idx in range(num_chunks): - chunk_start = chunk_idx * self.PREFILL_CHUNK_SIZE - chunk_end = min(chunk_start + self.PREFILL_CHUNK_SIZE, num_prefills) + for chunk_start, chunk_end, chunk_N, chunk_M in chunk_plan: chunk_size = chunk_end - chunk_start + kv = kv_rows[: chunk_size * chunk_M].view(chunk_size, chunk_M, -1) if not swa_only: assert attn_metadata is not None assert compressed_k_cache is not None @@ -1447,7 +1448,7 @@ class DeepseekV41ROCMAiterMLAAttention(DeepseekV4Attention): gather_lens=gather_lens[chunk_start:chunk_end], block_table=swa_block_table[chunk_start:chunk_end], block_size=swa_metadata.block_size, - offset=N, + offset=chunk_N, use_fnuz=current_platform.is_fp8_fnuz(), ) @@ -1468,8 +1469,8 @@ class DeepseekV41ROCMAiterMLAAttention(DeepseekV4Attention): self.window_size, self.compress_ratio, top_k, - M, - N, + chunk_M, + chunk_N, ) rocm_sparse_attn_prefill( q=q[query_start:query_end],