fix(ds4): enforce sparse top-k row-length contract

csrc/libtorch_stable/cooperative_topk.cuh

@@ -496,7 +496,7 @@ __device__ void large_topk(const float* __restrict__ row_input,
 template <uint32_t TopK, uint32_t CS>
 __device__ void cooperative_topk_body(CooperativeTopKParams<TopK> params) {
   const auto rank = blockIdx.y, row = blockIdx.x, tx = threadIdx.x;
-  const auto sl = params.lengths[row];
+  const int32_t sl = params.lengths[row] > 0 ? params.lengths[row] : 0;
   int32_t* out = params.output + row * TopK;
   const float* in = params.input + row * params.stride;
 

csrc/libtorch_stable/persistent_topk.cuh

@@ -661,7 +661,7 @@ __device__ void radix_topk(const float* __restrict__ row_input,
                            uint32_t* shared_scalars, uint32_t* shared_ordered,
                            RadixRowState* state, uint32_t cta_in_group,
                            uint32_t ctas_per_group, int& barrier_phase,
-                           uint32_t iter, uint32_t tx) {
+                           uint32_t radix_iter, uint32_t tx) {
   const uint32_t my_chunk_end = (my_chunk_start + chunk_size < seq_len)
                                     ? my_chunk_start + chunk_size
                                     : seq_len;
@@ -718,7 +718,7 @@ __device__ void radix_topk(const float* __restrict__ row_input,
 
   // -- Stage 2: 4 rounds of radix select --
   for (uint32_t round = 0; round < 4; round++) {
-    const uint32_t global_round = iter * 4 + round;
+    const uint32_t global_round = radix_iter * 4 + round;
     const uint32_t shift = 24 - round * 8;
     const uint32_t prefix = shared_scalars[0];
     const uint32_t remaining_k = shared_scalars[1];
@@ -898,14 +898,24 @@ __global__ void __launch_bounds__(kThreadsPerBlock, 2)
   RadixRowState* state = &params.row_states[group_id];
 
   int barrier_phase = 0;
+  uint32_t radix_iter = 0;
   const uint32_t total_iters = (params.num_rows + num_groups - 1) / num_groups;
 
   for (uint32_t iter = 0; iter < total_iters; iter++) {
     // Static round-robin: all CTAs in the group implicitly agree on the row
     uint32_t row_idx = group_id + iter * num_groups;
     if (row_idx >= params.num_rows) break;
 
-    const uint32_t seq_len = params.lengths[row_idx];
+    // A row may only expose scores inside both its physical stride and the
+    // host-declared active sequence bound. Padding rows use non-positive
+    // lengths and therefore select no scores.
+    const int32_t raw_len = params.lengths[row_idx];
+    const uint32_t row_bound =
+        params.stride < params.max_seq_len ? params.stride : params.max_seq_len;
+    const uint32_t non_negative_len =
+        raw_len > 0 ? static_cast<uint32_t>(raw_len) : 0u;
+    const uint32_t seq_len =
+        non_negative_len < row_bound ? non_negative_len : row_bound;
     int32_t* row_output = params.output + row_idx * params.top_k;
     const float* row_input = params.input + row_idx * params.stride;
 
@@ -930,7 +940,8 @@ __global__ void __launch_bounds__(kThreadsPerBlock, 2)
     radix_topk<TopK, VEC_SIZE>(
         row_input, row_output, seq_len, my_chunk_start, chunk_size,
         local_histogram, suffix_sum, shared_scalars, shared_ordered, state,
-        cta_in_group, ctas_per_group, barrier_phase, iter, tx);
+        cta_in_group, ctas_per_group, barrier_phase, radix_iter, tx);
+    radix_iter++;
   }
 }
 

tests/kernels/test_top_k_per_row.py

@@ -926,6 +926,90 @@ def test_workspace_topk(test_config: dict, top_k: int, backend: str) -> None:
     )
 
 
+@pytest.mark.skipif(not current_platform.is_cuda(), reason="This test requires CUDA")
+@torch.inference_mode()
+def test_persistent_topk_reused_group_after_short_row() -> None:
+    """A short row must not advance a group's radix histogram ring."""
+    torch.set_default_device("cuda:0")
+    set_random_seed(0)
+
+    top_k = 2048
+    long_seq_len = 32769
+    radix = 256
+    fixed_smem = ((radix + radix + 5) * 4 + 15) & ~15
+    props = torch.cuda.get_device_properties(0)
+    max_smem = props.shared_memory_per_block_optin
+    if max_smem <= props.shared_memory_per_multiprocessor // 2:
+        pytest.skip("Cannot force one persistent_topk CTA per SM")
+
+    max_chunk = ((max_smem - fixed_smem) // 4 // 4) * 4
+    ctas_per_group = max(
+        (props.multi_processor_count - 1 + 9) // 10,
+        (long_seq_len + max_chunk - 1) // max_chunk,
+    )
+    if ctas_per_group >= props.multi_processor_count:
+        pytest.skip("Not enough SMs to construct a reused CTA group")
+
+    stride = ctas_per_group * max_chunk
+    num_groups = max(1, (props.multi_processor_count - 1) // ctas_per_group)
+    num_rows = 3 * num_groups
+    lengths = torch.full((num_rows,), top_k, dtype=torch.int32, device="cuda")
+    target_row = 2 * num_groups
+    lengths[0] = long_seq_len
+    lengths[num_groups] = long_seq_len - 1
+    lengths[target_row] = long_seq_len
+
+    logits = torch.randn(num_rows, stride, dtype=torch.float32, device="cuda")
+    indices = torch.empty((num_rows, top_k), dtype=torch.int32, device="cuda")
+    _run_topk_backend("persistent_topk", logits, lengths, indices, top_k, stride)
+    torch.accelerator.synchronize()
+
+    expected = logits[target_row, :long_seq_len].topk(top_k).indices
+    assert set(indices[target_row].cpu().tolist()) == set(expected.cpu().tolist())
+
+
+@pytest.mark.skipif(not current_platform.is_cuda(), reason="This test requires CUDA")
+@pytest.mark.parametrize("backend", WORKSPACE_TOPK_BACKENDS)
+@torch.inference_mode()
+def test_workspace_topk_negative_length_is_empty(backend: str) -> None:
+    """A graph-padding row with a negative length selects no scores."""
+    torch.set_default_device("cuda:0")
+    top_k = 2048
+    stride = 4096
+    logits = torch.randn((1, stride), dtype=torch.float32, device="cuda")
+    lengths = torch.tensor([-1], dtype=torch.int32, device="cuda")
+    indices = torch.empty((1, top_k), dtype=torch.int32, device="cuda")
+
+    _run_topk_backend(backend, logits, lengths, indices, top_k, stride)
+    torch.accelerator.synchronize()
+
+    assert torch.all(indices == -1)
+
+
+@pytest.mark.skipif(not current_platform.is_cuda(), reason="This test requires CUDA")
+@torch.inference_mode()
+def test_persistent_topk_bounds_length_to_active_row() -> None:
+    """Persistent top-k must not select mapped scores beyond the active row."""
+    torch.set_default_device("cuda:0")
+    set_random_seed(0)
+
+    top_k = 2048
+    stride = 4096
+    active_length = 3072
+    backing = torch.empty(2 * stride, dtype=torch.float32, device="cuda")
+    backing[:active_length] = torch.randn(active_length, device="cuda")
+    backing[active_length:] = 1000.0
+    logits = backing[:stride].view(1, stride)
+    lengths = torch.tensor([2 * stride], dtype=torch.int32, device="cuda")
+    indices = torch.empty((1, top_k), dtype=torch.int32, device="cuda")
+
+    _run_topk_backend("persistent_topk", logits, lengths, indices, top_k, active_length)
+    torch.accelerator.synchronize()
+
+    expected = logits[0, :active_length].topk(top_k).indices
+    assert set(indices[0].cpu().tolist()) == set(expected.cpu().tolist())
+
+
 @pytest.mark.skipif(not current_platform.is_cuda(), reason="This test requires CUDA")
 @pytest.mark.parametrize("top_k", [512, 2048])
 @pytest.mark.parametrize("backend", WORKSPACE_TOPK_BACKENDS)

tests/model_executor/layers/test_sparse_attn_indexer_b12x.py

@@ -1650,3 +1650,67 @@ def test_mtp_variable_decode_preserves_block_table_alignment_padding():
             next_n=4,
             max_decode_len=3,
         )
+
+
+def test_native_mtp_padding_lengths_are_non_negative():
+    from vllm.v1.attention.backends.mla import indexer as mla_indexer_mod
+
+    builder = object.__new__(mla_indexer_mod.DeepseekV32IndexerMetadataBuilder)
+    builder.decode_seq_lens_buffer = torch.empty(8, dtype=torch.int32)
+    builder.offsets_buffer = torch.arange(4, dtype=torch.int32)
+
+    seq_lens = torch.tensor([0, 5], dtype=torch.int32)
+    block_table = torch.zeros((2, 2), dtype=torch.int32)
+    decode_lens = torch.tensor([2, 2], dtype=torch.int32)
+
+    prepared, _, _, batch_size, requires_padding = builder._prepare_decode_tensors(
+        seq_lens=seq_lens,
+        block_table=block_table,
+        decode_lens=decode_lens,
+        decode_lens_cpu=decode_lens,
+        query_start_loc=torch.tensor([0, 2], dtype=torch.int32),
+        num_decodes=2,
+        num_decode_tokens=4,
+        use_native=True,
+        next_n=2,
+        max_decode_len=2,
+    )
+
+    assert prepared.tolist() == [[0, 0], [4, 5]]
+    assert batch_size == 2
+    assert not requires_padding
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="This test requires CUDA")
+def test_uniform_decode_kernel_clamps_graph_padding_lengths():
+    from vllm.v1.attention.backends.mla.indexer import (
+        _prepare_uniform_decode_kernel,
+    )
+
+    device = torch.device("cuda")
+    seq_lens = torch.tensor([0, 5], dtype=torch.int32, device=device)
+    block_table = torch.tensor([[11, 12], [21, 22]], dtype=torch.int32, device=device)
+    prepared = torch.empty(4, dtype=torch.int32, device=device)
+    expanded_block_table = torch.empty((4, 2), dtype=torch.int32, device=device)
+    decode_lens = torch.empty(4, dtype=torch.int32, device=device)
+
+    _prepare_uniform_decode_kernel[(4,)](
+        seq_lens,
+        prepared,
+        block_table,
+        block_table.stride(0),
+        expanded_block_table,
+        expanded_block_table.stride(0),
+        decode_lens,
+        2,
+        BLOCK_SIZE=1024,
+    )
+
+    assert prepared.cpu().tolist() == [0, 0, 4, 5]
+    assert expanded_block_table.cpu().tolist() == [
+        [11, 12],
+        [11, 12],
+        [21, 22],
+        [21, 22],
+    ]
+    assert decode_lens.cpu().tolist() == [1, 1, 1, 1]

vllm/v1/attention/backends/mla/indexer.py

@@ -63,9 +63,10 @@ def _prepare_uniform_decode_kernel(
     req_id = idx // max_decode_len
     local_idx = idx % max_decode_len
 
-    # Compute number of KVs attended to by this token.
+    # Graph-padding requests have no KV entries. Clamp their per-token lengths
+    # before downstream kernels consume the signed values as unsigned bounds.
     seq_len = tl.load(seq_lens_ptr + req_id)
-    per_token_seq_len = seq_len - max_decode_len + local_idx + 1
+    per_token_seq_len = tl.maximum(seq_len - max_decode_len + local_idx + 1, 0)
     tl.store(decode_seq_lens_ptr + idx, per_token_seq_len)
 
     # Copy block table row.
@@ -956,7 +957,7 @@ def _prepare_decode_tensors(
                     - max_decode_len
                     + 1
                     + self.offsets_buffer[:max_decode_len]
-                )
+                ).clamp_(min=0)
                 seq_lens = seq_lens_buffer
             return seq_lens, block_table, decode_lens, num_decodes, requires_padding