[ROCm][Perf] Optimize sparse attention prefill kernel for DeepSeek-V4 (#48519)
Signed-off-by: kliuae <kuanfu.liu@embeddedllm.com> Co-authored-by: TJian <tunjian.tan@embeddedllm.com>
kliuae committed
Jul 16, 2026 at 08:58 UTC
3c1bc1fc0dd5384c92d608824f03e6e60e98d592
1 file changed
+10
-3
vllm/v1/attention/ops/rocm_aiter_mla_sparse.py
+10
-3
@@ -1187,10 +1187,12 @@ def _sparse_attn_prefill_ragged_kernel(
1187
kv_len = kv_end - kv_start
1188
1189
k_offsets = tl.arange(0, BLOCK_K)
1190
+ slot = tl.load(
1191
+ kv_indices_ptr + kv_start + k_offsets, mask=k_offsets < kv_len, other=-1
1192
+ )
1193
for k_start in tl.range(0, kv_len, BLOCK_K):
1194
k_pos = k_start + k_offsets
1195
in_range = k_pos < kv_len
1193
- slot = tl.load(kv_indices_ptr + kv_start + k_pos, mask=in_range, other=-1)
1196
valid = in_range & (slot >= 0) & (slot < num_kv)
1197
safe_slot = tl.where(valid, slot, 0)
1198
@@ -1201,7 +1203,11 @@ def _sparse_attn_prefill_ragged_kernel(
1203
mask=valid[:, None] & dim_mask[None, :],
1204
other=0.0,
1205
)
1204
- kv = tl.where(valid[:, None] & dim_mask[None, :], kv, 0.0)
1206
+
1207
+ next_k_pos = k_start + BLOCK_K + k_offsets
1208
+ slot = tl.load(
1209
+ kv_indices_ptr + kv_start + next_k_pos, mask=next_k_pos < kv_len, other=-1
1210
+ )
1211
1212
scores = tl.dot(q, tl.trans(kv)) * scale
1213
scores = tl.where(head_mask[:, None] & valid[None, :], scores, neg_large)
@@ -1865,6 +1871,7 @@ def _rocm_sparse_attn_prefill_ragged_triton(
1871
block_h = 16
1872
block_d = triton.next_power_of_2(head_dim)
1873
block_k = 16 if head_dim >= 256 else 32
1874
+ num_warps = 4
1875
out = torch.empty_like(q, dtype=torch.bfloat16)
1876
_sparse_attn_prefill_ragged_kernel[(num_queries, triton.cdiv(num_heads, block_h))](
1877
q,
@@ -1889,7 +1896,7 @@ def _rocm_sparse_attn_prefill_ragged_triton(
1896
BLOCK_H=block_h,
1897
BLOCK_D=block_d,
1898
BLOCK_K=block_k,
1892
- num_warps=8,
1899
+ num_warps=num_warps,
1900
)
1901
return out
1902