[Perf][DCP] Sparse indexer: key only significant bits in the DCP top-k merge

tests/v1/attention/test_indexer_dcp_localize.py

@@ -588,6 +588,58 @@ def test_cutedsl_dcp_candidate_pack_and_select_matches_reference(
         assert set(actual[row].cpu().tolist()) == set(expected[row].cpu().tolist())
 
 
+@pytest.mark.skipif(
+    not current_platform.is_cuda() or not has_cutedsl(),
+    reason="This test requires CUDA and CuteDSL",
+)
+@pytest.mark.parametrize("num_positions", [16384, 131072, 1 << 20])
+def test_cutedsl_gathered_select_compact_keys_match_reference(num_positions: int):
+    """BF16-valued scores (upcast SM100 MQA logits) tie at the top-k boundary.
+    Keying only their 16 high bits and the id bits below ``num_positions``
+    must select the same set as the full 64-bit key."""
+    from vllm.model_executor.kernels.attention.dsa.dcp_indexer_cutedsl import (
+        dcp_topk_id_bits,
+        stable_topk_from_gathered_candidates_cutedsl,
+    )
+
+    torch.manual_seed(7)
+    device = torch.device("cuda")
+    rows, world, topk = 5, 8, 2048
+    num_candidates = world * topk
+    scores = torch.randn((rows, num_candidates), device=device)
+    # Row 1: 64 distinct values, so the boundary tie spans hundreds of ids.
+    scores[1] = (scores[1] * 8).round().clamp(-32, 31)
+    # Row 2: all negative. Row 3: padding mixed in. Row 4: all tied.
+    scores[2] = -scores[2].abs()
+    scores[4] = 0.5
+    scores = scores.bfloat16().float()
+    ids = torch.stack(
+        [
+            torch.randperm(num_positions, device=device)[:num_candidates]
+            for _ in range(rows)
+        ]
+    ).to(torch.int32)
+    ids[:, 0] = num_positions - 1
+    ids[3, ::3] = -1
+    gathered = torch.stack((scores, ids.float()), dim=-1).contiguous()
+
+    expected = _ref_stable_topk_from_candidates_fp64(scores, ids, topk)
+    full = torch.empty((rows, topk), device=device, dtype=torch.int32)
+    stable_topk_from_gathered_candidates_cutedsl(gathered, topk, out=full)
+    compact = torch.empty_like(full)
+    stable_topk_from_gathered_candidates_cutedsl(
+        gathered,
+        topk,
+        out=compact,
+        score_bits=16,
+        id_bits=dcp_topk_id_bits(num_positions),
+    )
+    for row in range(rows):
+        want = sorted(expected[row].tolist())
+        assert sorted(full[row].tolist()) == want
+        assert sorted(compact[row].tolist()) == want
+
+
 @pytest.mark.skipif(not current_platform.is_cuda(), reason="This test requires CUDA")
 def test_sparse_prefill_dcp_metadata_localizes_causal_bounds():
     device = torch.device("cuda")

vllm/model_executor/kernels/attention/dsa/dcp_indexer_cutedsl.py

@@ -23,6 +23,7 @@
     VllmTritonJitKernel,
 )
 from vllm.triton_utils import tl, triton
+from vllm.utils.deep_gemm import mqa_logits_use_mx_scales
 
 
 def _dcp_merge_topk_and_interleave(vllm_config: Any) -> tuple[int, int]:
@@ -41,18 +42,29 @@ def _dcp_merge_topk_and_interleave(vllm_config: Any) -> tuple[int, int]:
     return topk, cp_interleave
 
 
+def dcp_topk_id_bits(num_positions: int) -> int:
+    """Bits that hold every candidate id (a position below num_positions)."""
+    return max(1, min(32, (num_positions - 1).bit_length()))
+
+
 def stable_topk_from_gathered_candidates_cutedsl(
     gathered: torch.Tensor,
     topk: int,
     out: torch.Tensor | None = None,
+    score_bits: int = 32,
+    id_bits: int = 32,
 ) -> torch.Tensor:
+    """`score_bits`: high score bits that rank (16 for upcast BF16 logits).
+    `id_bits`: bits that hold every token id."""
     if out is None:
         out = torch.empty(
             (gathered.shape[0], topk),
             dtype=torch.int32,
             device=gathered.device,
         )
-    _STABLE_TOPK_FROM_GATHERED_CANDIDATES_KERNEL(gathered, out, topk=topk)
+    _STABLE_TOPK_FROM_GATHERED_CANDIDATES_KERNEL(
+        gathered, out, topk=topk, score_bits=score_bits, id_bits=id_bits
+    )
     return out
 
 
@@ -274,24 +286,33 @@ class StableTopKFromGatheredCandidatesKernel(
     hist_bins = 2048
     radix_bits = (hist_bins - 1).bit_length()
     key_bits = Uint64.width
-    radix_passes = (key_bits + radix_bits - 1) // radix_bits
-    final_radix_bits = key_bits - radix_bits * (radix_passes - 1)
     hist_chunks = (hist_bins + tb_size - 1) // tb_size
     warps_per_block = tb_size // cute.arch.WARP_SIZE
 
     @dataclass(frozen=True)
     class CompileKey:
         topk: int
         num_candidates: int
+        # The key holds the top score_bits of the score (16 for upcast BF16
+        # logits, whose low half is zero) and the id_bits that hold every
+        # token id, so the radix never spends a pass on constant bits.
+        score_bits: int = 32
+        id_bits: int = 32
 
     @staticmethod
     def kernel(compile_key: CompileKey) -> Any:
         tb_size = StableTopKFromGatheredCandidatesKernel.tb_size
         hist_bins = StableTopKFromGatheredCandidatesKernel.hist_bins
         radix_bits = StableTopKFromGatheredCandidatesKernel.radix_bits
         key_bits = StableTopKFromGatheredCandidatesKernel.key_bits
-        radix_passes = StableTopKFromGatheredCandidatesKernel.radix_passes
-        final_radix_bits = StableTopKFromGatheredCandidatesKernel.final_radix_bits
+        score_bits = compile_key.score_bits
+        id_bits = compile_key.id_bits
+        assert 0 < score_bits <= 32 and 0 < id_bits <= 32
+        used_bits = score_bits + id_bits
+        radix_passes = (used_bits + radix_bits - 1) // radix_bits
+        final_radix_bits = used_bits - radix_bits * (radix_passes - 1)
+        id_shift = key_bits - used_bits
+        id_mask = (1 << id_bits) - 1
         hist_chunks = StableTopKFromGatheredCandidatesKernel.hist_chunks
         warps_per_block = StableTopKFromGatheredCandidatesKernel.warps_per_block
         topk = compile_key.topk
@@ -356,8 +377,10 @@ def stable_key(score: Float32, token_id: Int32) -> Uint64:
             mask = Uint32(0x80000000)
             if (bits & Uint32(0x80000000)) != Uint32(0):
                 mask = Uint32(0xFFFFFFFF)
-            score_key = Uint64(bits ^ mask) << Uint64(32)
-            id_key = Uint64(~Uint32(token_id))
+            score_key = Uint64((bits ^ mask) >> Uint32(32 - score_bits)) << Uint64(
+                key_bits - score_bits
+            )
+            id_key = Uint64(~Uint32(token_id) & Uint32(id_mask)) << Uint64(id_shift)
             key = score_key | id_key
             if token_id < Int32(0):
                 key = Uint64(0)
@@ -398,7 +421,10 @@ def radix_pass(
 
             prefix_bits = step * Int32(radix_bits)
             num_bins = 1 << bits
-            block_scan_iterations = (num_bins + tb_size - 1) // tb_size
+            # Narrow final passes still scan whole tb_size slices; bins at or
+            # above num_bins stay zero.
+            scan_bins = max(num_bins, tb_size)
+            block_scan_iterations = scan_bins // tb_size
             shift = Int32(key_bits) - prefix_bits - Int32(bits)
             bin_mask = Uint64(num_bins - 1)
             prefix = prefix_smem.load()
@@ -429,7 +455,7 @@ def radix_pass(
             iter = Int32(0)
             threshold_found = threshold_found_smem.load()
             while threshold_found == Int32(0) and iter < Int32(block_scan_iterations):
-                bin_idx = Int32(num_bins - 1) - (iter * Int32(tb_size) + tid)
+                bin_idx = Int32(scan_bins - 1) - (iter * Int32(tb_size) + tid)
                 count = hist_smem[bin_idx]
                 chunk_inclusive = block_scan_inclusive_i32(
                     count,
@@ -475,7 +501,12 @@ def radix_pass(
                             scope="cta",
                         )
                         if dst < Int32(topk):
-                            output[dst] = recast_val(~Uint32(key), Int32)
+                            # Padding keys are 0 and decode to -1.
+                            token_id = Int32(-1)
+                            if key != Uint64(0):
+                                id_key = Uint32(key >> Uint64(id_shift))
+                                token_id = recast_val(~id_key & Uint32(id_mask), Int32)
+                            output[dst] = token_id
             cute.arch.sync_threads()
 
             pass_finished = include_threshold_bin_smem.load()
@@ -563,9 +594,15 @@ def get_warmup_keys(self, vllm_config: Any) -> list[CompileKey]:
         topk, _ = _dcp_merge_topk_and_interleave(vllm_config)
         if topk <= 0:
             return []
+        model_config = vllm_config.model_config
+        kpool = getattr(model_config.hf_text_config, "index_kpool", None) or 1
         return self._trace_dispatch(self.dispatch)(
             topk=topk,
             num_candidates=topk * dcp_world_size,
+            # SM100 MQA logits are BF16, upcast for the top-k.
+            score_bits=16 if mqa_logits_use_mx_scales() else 32,
+            # Candidate ids are positions (pool ids for kpool) in max_model_len.
+            id_bits=dcp_topk_id_bits(model_config.max_model_len // kpool),
         )
 
     def warmup_inputs(self, compile_key: CompileKey) -> tuple[Any, ...]:
@@ -590,10 +627,14 @@ def __call__(
         out: torch.Tensor,
         *,
         topk: int,
+        score_bits: int = 32,
+        id_bits: int = 32,
     ) -> CuTeDSLLaunchSpec[CompileKey]:
         compile_key = self.dispatch(
             topk=topk,
             num_candidates=gathered.shape[1],
+            score_bits=score_bits,
+            id_bits=id_bits,
         )
         return compile_key, (gathered, out)
 

vllm/model_executor/layers/sparse_attn_indexer.py

@@ -100,6 +100,8 @@ def _merge_dcp_topk_global(
     dcp_world_size: int,
     cp_interleave: int,
     row_starts: torch.Tensor | None = None,
+    score_bits: int = 32,
+    num_positions: int | None = None,
 ) -> None:
     """Merge each DCP rank's local top-K into the global top-K.
 
@@ -112,6 +114,9 @@ def _merge_dcp_topk_global(
     score row. Overwrites ``topk_indices`` with global token ids (``-1`` for
     padding); the attention backend localizes them back to physical slots per
     rank.
+    The selector ranks only the high ``score_bits`` of the scores (16 when the
+    logits are upcast BF16) and the id bits that positions below
+    ``num_positions`` use.
     """
     if dcp_world_size <= 1:
         return
@@ -121,6 +126,7 @@ def _merge_dcp_topk_global(
     # stable-topk selector.
     _assert_cutedsl_dcp_merge_supported(logits, topk_indices, topk_tokens)
     from vllm.model_executor.kernels.attention.dsa.dcp_indexer_cutedsl import (
+        dcp_topk_id_bits,
         pack_dcp_topk_candidates_cutedsl,
         stable_topk_from_gathered_candidates_cutedsl,
     )
@@ -141,7 +147,11 @@ def _merge_dcp_topk_global(
     )
     gathered = get_dcp_group().all_gather(packed, dim=1)
     stable_topk_from_gathered_candidates_cutedsl(
-        gathered, topk_tokens, out=topk_indices
+        gathered,
+        topk_tokens,
+        out=topk_indices,
+        score_bits=score_bits,
+        id_bits=32 if num_positions is None else dcp_topk_id_bits(num_positions),
     )
 
 
@@ -690,6 +700,8 @@ def sparse_attn_indexer(
             # DeepGEMM returns compressed logits: row i starts at K position
             # cu_seqlen_ks[i]. None means every row starts at column zero.
             row_starts: torch.Tensor | None = None
+            # Empty chunks have no logits; match the other ranks and the warm-up.
+            score_bits = 16 if mqa_logits_use_mx_scales() else 32
             if chunk.local_total_seq_lens == 0 or q_slice.shape[0] == 0:
                 logits = q_slice.new_empty((q_slice.shape[0], 0), dtype=torch.float32)
                 topk_indices.fill_(-1)
@@ -724,17 +736,16 @@ def sparse_attn_indexer(
                         q_slice_cast, q_scale_slice, weights[row_start:row_end]
                     )
                     row_ends = cu_seqlen_ke - cu_seqlen_ks
-                    logits = as_fp32_logits(
-                        fp8_fp4_mqa_logits(
-                            (q_slice_cast, q_scale_slice),
-                            (k_quant_cast, k_scale_cast),
-                            chunk_weights,
-                            cu_seqlen_ks,
-                            cu_seqlen_ke,
-                            chunk.max_seqlen_k,
-                        ),
-                        row_ends,
+                    logits = fp8_fp4_mqa_logits(
+                        (q_slice_cast, q_scale_slice),
+                        (k_quant_cast, k_scale_cast),
+                        chunk_weights,
+                        cu_seqlen_ks,
+                        cu_seqlen_ke,
+                        chunk.max_seqlen_k,
                     )
+                    score_bits = 16 if logits.dtype == torch.bfloat16 else 32
+                    logits = as_fp32_logits(logits, row_ends)
                 num_rows = logits.shape[0]
                 if candidate_blocks is not None:
                     # Two-level selection (v4.1): the candidate source
@@ -782,6 +793,8 @@ def sparse_attn_indexer(
                     dcp_world_size,
                     cp_kv_cache_interleave_size,
                     row_starts=row_starts,
+                    score_bits=score_bits,
+                    num_positions=max_model_len,
                 )
 
         if row_shard is not None:
@@ -858,6 +871,7 @@ def sparse_attn_indexer(
             if use_fp4_cache
             else padded_q_quant_decode_tokens
         )
+        score_bits = 32
         if current_platform.is_xpu():
             if padded_q_scale is not None:
                 raise RuntimeError("XPU fp8_paged_mqa_logits does not support FP4 Q")
@@ -887,6 +901,8 @@ def sparse_attn_indexer(
                 max_model_len=max_model_len,
                 indices=decode_metadata.indices,
             )
+            if logits.dtype == torch.bfloat16:
+                score_bits = 16
             # Keeping bf16 halves top-k/candidate reads; other top-k backends and
             # DCP need fp32. SM100 decode is varlen: one context length per row.
             if topk_backend != "deep_select" or dcp_world_size != 1:
@@ -941,6 +957,8 @@ def sparse_attn_indexer(
                 dcp_rank,
                 dcp_world_size,
                 cp_kv_cache_interleave_size,
+                score_bits=score_bits,
+                num_positions=max_model_len,
             )
 
         if needs_padded_path: