From fcd317cf33b5113768a8ba1570f5211627f6fd63 Mon Sep 17 00:00:00 2001 From: Zheng Gong Date: Wed, 30 Sep 2026 15:26:57 +0800 Subject: [PATCH] [ROCm][DSv4][Perf] Use the shared prefill chunk plan in the ROCm sparse prefill (#58405) Signed-off-by: Zheng Gong Co-authored-by: Cursor --- vllm/models/deepseek_v4/amd/rocm.py | 27 +++++++++++---------------- 1 file changed, 11 insertions(+), 16 deletions(-) diff --git a/vllm/models/deepseek_v4/amd/rocm.py b/vllm/models/deepseek_v4/amd/rocm.py index 292901dfe9..c4c4d416b8 100644 --- a/vllm/models/deepseek_v4/amd/rocm.py +++ b/vllm/models/deepseek_v4/amd/rocm.py @@ -1426,7 +1426,6 @@ class DeepseekV4ROCMAiterMLAAttention(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 @@ -1458,26 +1457,22 @@ class DeepseekV4ROCMAiterMLAAttention(DeepseekV4Attention): topk_indices = attn_metadata.c128a_prefill_topk_indices assert topk_indices is not None top_k = topk_indices.shape[-1] - N = (self.max_model_len + self.compress_ratio - 1) // self.compress_ratio else: assert self.topk_indices_buffer is not None topk_indices = self.topk_indices_buffer[num_decode_tokens:] top_k = 0 - 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, ) - + assert chunk_plan, "prefill chunk plan must be non-empty when num_prefills > 0" workspace_manager = current_workspace_manager() - kv = workspace_manager.get_simultaneous( - ((self.PREFILL_CHUNK_SIZE, M, q.shape[-1]), torch.bfloat16), - )[0] - 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 = workspace_manager.get_simultaneous( + ((chunk_size, chunk_M, q.shape[-1]), torch.bfloat16), + )[0] if not swa_only: assert attn_metadata is not None assert compressed_k_cache is not None @@ -1502,7 +1497,7 @@ class DeepseekV4ROCMAiterMLAAttention(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(), ) @@ -1523,8 +1518,8 @@ class DeepseekV4ROCMAiterMLAAttention(DeepseekV4Attention): self.window_size, self.compress_ratio, top_k, - M, - N, + chunk_M, + chunk_N, max_image_tokens=self.max_image_tokens, left_visible=( left_visible[query_start:query_end]