A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 2, 11:00 UTC
[ROCm][DSv4] Remove device pipeline stall in sparse attention (#43898)
Signed-off-by: kliuae <kuanfu.liu@embeddedllm.com>
kliuae committed
May 29, 2026 at 15:42 UTC
ab7521d77cf29ccaa6116e644ef0f4c6ff2a9abb
1 file changed
+4
-3
vllm/v1/attention/ops/rocm_aiter_mla_sparse.py
+4
-3
@@ -998,15 +998,16 @@ def build_ragged_indices_from_dense(
998
max_width = indices.shape[1] if indices.ndim == 2 else 0
999
lengths = lengths.clamp(min=0, max=max_width).contiguous()
1000
1001
- indptr = torch.empty(indices.shape[0] + 1, dtype=torch.int32, device=indices.device)
1002
- indptr[0] = 0
1001
+ indptr = torch.zeros(indices.shape[0] + 1, dtype=torch.int32, device=indices.device)
1002
torch.cumsum(lengths, dim=0, out=indptr[1:])
1003
1004
if indices.numel() == 0:
1005
flat = torch.empty(0, dtype=torch.int32, device=indices.device)
1006
else:
1007
flat = torch.empty(
1009
- int(indptr[-1].item()), dtype=torch.int32, device=indices.device
1008
+ indices.shape[0] * max_width,
1009
+ dtype=torch.int32,
1010
+ device=indices.device,
1011
)
1012
if flat.numel() > 0:
1013
block_size = 128