From ab7521d77cf29ccaa6116e644ef0f4c6ff2a9abb Mon Sep 17 00:00:00 2001 From: kliuae <17350011+kliuae@users.noreply.github.com> Date: Fri, 29 May 2026 15:42:40 +0800 Subject: [PATCH] [ROCm][DSv4] Remove device pipeline stall in sparse attention (#43898) Signed-off-by: kliuae --- vllm/v1/attention/ops/rocm_aiter_mla_sparse.py | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/vllm/v1/attention/ops/rocm_aiter_mla_sparse.py b/vllm/v1/attention/ops/rocm_aiter_mla_sparse.py index 7ddcc493449..332350d8380 100644 --- a/vllm/v1/attention/ops/rocm_aiter_mla_sparse.py +++ b/vllm/v1/attention/ops/rocm_aiter_mla_sparse.py @@ -998,15 +998,16 @@ def build_ragged_indices_from_dense( max_width = indices.shape[1] if indices.ndim == 2 else 0 lengths = lengths.clamp(min=0, max=max_width).contiguous() - indptr = torch.empty(indices.shape[0] + 1, dtype=torch.int32, device=indices.device) - indptr[0] = 0 + indptr = torch.zeros(indices.shape[0] + 1, dtype=torch.int32, device=indices.device) torch.cumsum(lengths, dim=0, out=indptr[1:]) if indices.numel() == 0: flat = torch.empty(0, dtype=torch.int32, device=indices.device) else: flat = torch.empty( - int(indptr[-1].item()), dtype=torch.int32, device=indices.device + indices.shape[0] * max_width, + dtype=torch.int32, + device=indices.device, ) if flat.numel() > 0: block_size = 128