karylab_agents/vllm
1
0
Fork 0
forked from Karylab-cklius/vllm
Code Pull Requests 2 Actions 1 Packages Activity

[ROCm][DSv4] Remove device pipeline stall in sparse attention (#43898)

Browse Source 
Signed-off-by: kliuae <kuanfu.liu@embeddedllm.com>
This commit is contained in:
kliuae
2026-05-29 15:42:40 +08:00
committed by GitHub
parent 94d3f4d205
commit ab7521d77c
1 changed files with 4 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files
vllm/v1/attention/ops/rocm_aiter_mla_sparse.py
+4 -3
View File
@@ -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