[Bugfix][ROCm] Drop -1 sentinels when building the ragged sparse-MLA indices (#58058)
Signed-off-by: Mehmet Cagri Kaymak <mehmet.kaymak@amd.com> Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>
Mehmet Cagri committed
Sep 28, 2026 at 18:40 UTC
75fad5bbef7b2c0264c4b7ce6c4c10e32233d523
2 files changed
+124
-27
tests/kernels/attention/test_rocm_triton_attn_dsv4.py
+32
@@ -8,6 +8,7 @@ import pytest
8
import torch
9
10
from vllm.platforms import current_platform
11
+from vllm.utils.torch_utils import set_random_seed
12
13
pytestmark = pytest.mark.skipif(
14
not current_platform.is_rocm(), reason="Only used by ROCm"
@@ -424,6 +425,37 @@ def test_compute_global_topk_ragged_indices_and_indptr() -> None:
425
torch.testing.assert_close(actual_lens, expected_lens)
426
427
428
+@pytest.mark.parametrize("width", [16, 640, 2176])
429
+@pytest.mark.parametrize("num_queries", [1, 70, 1000])
430
+@torch.inference_mode()
431
+def test_build_ragged_indices_from_dense_drops_invalid_entries(
432
+ width: int, num_queries: int
433
+) -> None:
434
+ from vllm.v1.attention.ops.rocm_aiter_mla_sparse import (
435
+ build_ragged_indices_from_dense,
436
+ )
437
+
438
+ device = torch.device("cuda")
439
+ set_random_seed(width + num_queries)
440
+ num_rows = 5000
441
+ indices = torch.randint(0, num_rows + 100, (num_queries, width))
442
+ indices[torch.rand(num_queries, width) < 0.3] = -1
443
+ lengths = torch.randint(0, width + 8, (num_queries,))
444
+
445
+ flat, indptr = build_ragged_indices_from_dense(
446
+ indices.to(device, torch.int32),
447
+ lengths.to(device, torch.int32),
448
+ num_rows=num_rows,
449
+ )
450
+
451
+ # Each row keeps its in-range entries within its length, in order.
452
+ expected = [
453
+ [x for x in row[:n] if 0 <= x < num_rows]
454
+ for row, n in zip(indices.tolist(), lengths.tolist())
455
+ ]
456
+ assert _rows_from_ragged(flat, indptr) == expected
457
+
458
+
459
@torch.inference_mode()
460
def test_combine_topk_swa_indices_adds_image_visibility() -> None:
461
from vllm.models.deepseek_v4.amd.rocm import combine_topk_swa_indices
vllm/v1/attention/ops/rocm_aiter_mla_sparse.py
+92
-27
@@ -2017,36 +2017,74 @@ def _validate_dsv4_sparse_dims(
2017
2018
2019
@triton.jit
2020
-def _pack_dense_prefix_to_ragged_kernel(
2020
+def _count_valid_row_entries_kernel(
2021
indices_ptr,
2022
lengths_ptr,
2023
+ counts_ptr,
2024
indptr_ptr,
2024
- out_ptr,
2025
indices_stride0,
2026
num_rows_limit,
2027
+ num_queries,
2028
row_width,
2028
- BLOCK_SIZE: tl.constexpr,
2029
+ BLOCK_R: tl.constexpr,
2030
+ BLOCK_W: tl.constexpr,
2031
):
2030
- row_idx = tl.program_id(0)
2031
- block_idx = tl.program_id(1)
2032
- offsets = block_idx * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
2032
+ pid = tl.program_id(0)
2033
+ # The scan fills indptr[1:], so write indptr[0] here and skip a zero fill.
2034
+ if pid == 0:
2035
+ tl.store(indptr_ptr, 0)
2036
+
2037
+ rows = pid * BLOCK_R + tl.arange(0, BLOCK_R)
2038
+ row_ok = rows < num_queries
2039
+ row_len = tl.load(lengths_ptr + rows, mask=row_ok, other=0)
2040
+
2041
+ acc = tl.zeros((BLOCK_R,), dtype=tl.int32)
2042
+ for start in range(0, row_width, BLOCK_W):
2043
+ cols = start + tl.arange(0, BLOCK_W)
2044
+ in_row = cols[None, :] < row_width
2045
+ vals = tl.load(
2046
+ indices_ptr + rows[:, None] * indices_stride0 + cols[None, :],
2047
+ mask=row_ok[:, None] & in_row,
2048
+ other=-1,
2049
+ ).to(tl.int32)
2050
+ keep = (
2051
+ row_ok[:, None] & in_row & (cols[None, :] < row_len[:, None]) & (vals >= 0)
2052
+ )
2053
+ if num_rows_limit >= 0:
2054
+ keep = keep & (vals < num_rows_limit)
2055
+ acc += tl.sum(keep.to(tl.int32), axis=1)
2056
2034
- row_len = tl.load(lengths_ptr + row_idx)
2035
- if block_idx * BLOCK_SIZE >= row_len:
2036
- return
2057
+ tl.store(counts_ptr + rows, acc, mask=row_ok)
2058
2038
- mask = offsets < row_len
2039
- safe_offsets = tl.where(offsets < row_width, offsets, 0)
2040
- vals = tl.load(
2041
- indices_ptr + row_idx * indices_stride0 + safe_offsets,
2042
- mask=mask & (offsets < row_width),
2043
- other=-1,
2044
- ).to(tl.int32)
2045
- if num_rows_limit >= 0:
2046
- vals = tl.where((vals >= 0) & (vals < num_rows_limit), vals, -1)
2059
2060
+@triton.jit
2061
+def _compact_dense_row_to_ragged_kernel(
2062
+ indices_ptr,
2063
+ lengths_ptr,
2064
+ indptr_ptr,
2065
+ out_ptr,
2066
+ indices_stride0,
2067
+ num_rows_limit,
2068
+ row_width,
2069
+ BLOCK_W: tl.constexpr,
2070
+):
2071
+ row_idx = tl.program_id(0)
2072
+ row_len = tl.load(lengths_ptr + row_idx)
2073
out_start = tl.load(indptr_ptr + row_idx)
2049
- tl.store(out_ptr + out_start + offsets, vals, mask=mask)
2074
+
2075
+ written = 0
2076
+ for start in range(0, row_width, BLOCK_W):
2077
+ cols = start + tl.arange(0, BLOCK_W)
2078
+ in_row = cols < row_width
2079
+ vals = tl.load(
2080
+ indices_ptr + row_idx * indices_stride0 + cols, mask=in_row, other=-1
2081
+ ).to(tl.int32)
2082
+ keep = in_row & (cols < row_len) & (vals >= 0)
2083
+ if num_rows_limit >= 0:
2084
+ keep = keep & (vals < num_rows_limit)
2085
+ rank = tl.cumsum(keep.to(tl.int32), axis=0) - 1
2086
+ tl.store(out_ptr + out_start + written + rank, vals, mask=keep)
2087
+ written += tl.sum(keep.to(tl.int32), axis=0)
2088
2089
2090
def build_ragged_indices_from_dense(
@@ -2061,10 +2099,40 @@ def build_ragged_indices_from_dense(
2099
)
2100
2101
max_width = indices.shape[1] if indices.ndim == 2 else 0
2064
- lengths = lengths.clamp(min=0, max=max_width).contiguous()
2102
+ # Both kernels bound the row by row_width and row_len, so no clamp needed.
2103
+ lengths = lengths.contiguous()
2104
+
2105
+ # Dense rows reserve slots they may not fill, leaving -1 anywhere in the
2106
+ # first `lengths` entries. Keep only the valid ones, since the consumers
2107
+ # index the KV pool without checking the sign.
2108
+ num_queries = indices.shape[0]
2109
+ # max_width is whatever the caller packed (align(topk + window, 128) for the
2110
+ # V4.1 indexer) and nothing caps it, so loop the row rather than size the
2111
+ # block by it.
2112
+ block = min(triton.next_power_of_2(max_width), 1024) if max_width > 0 else 1
2113
+ counts = torch.empty(num_queries, dtype=torch.int32, device=indices.device)
2114
+ indptr = torch.empty(num_queries + 1, dtype=torch.int32, device=indices.device)
2115
+ if num_queries > 0 and max_width > 0:
2116
+ # A program per row only gets row_width elements, which is launch
2117
+ # bound for narrow rows. Tile rows to fill the block.
2118
+ block_r = max(1, 1024 // block)
2119
+ _count_valid_row_entries_kernel[(triton.cdiv(num_queries, block_r),)](
2120
+ indices,
2121
+ lengths,
2122
+ counts,
2123
+ indptr,
2124
+ indices.stride(0),
2125
+ int(num_rows),
2126
+ num_queries,
2127
+ max_width,
2128
+ BLOCK_R=block_r,
2129
+ BLOCK_W=block,
2130
+ )
2131
+ else:
2132
+ counts.zero_()
2133
+ indptr.zero_()
2134
2066
- indptr = torch.zeros(indices.shape[0] + 1, dtype=torch.int32, device=indices.device)
2067
- torch.cumsum(lengths, dim=0, out=indptr[1:])
2135
+ torch.cumsum(counts, dim=0, out=indptr[1:])
2136
2137
if indices.numel() == 0:
2138
flat = torch.empty(0, dtype=torch.int32, device=indices.device)
@@ -2075,10 +2143,7 @@ def build_ragged_indices_from_dense(
2143
device=indices.device,
2144
)
2145
if flat.numel() > 0:
2078
- block_size = 128
2079
- _pack_dense_prefix_to_ragged_kernel[
2080
- (indices.shape[0], triton.cdiv(max_width, block_size))
2081
- ](
2146
+ _compact_dense_row_to_ragged_kernel[(indices.shape[0],)](
2147
indices,
2148
lengths,
2149
indptr,
@@ -2086,7 +2151,7 @@ def build_ragged_indices_from_dense(
2151
indices.stride(0),
2152
int(num_rows),
2153
max_width,
2089
- BLOCK_SIZE=block_size,
2154
+ BLOCK_W=block,
2155
)
2156
2157
return flat, indptr