[Perf][DCP] Sparse indexer: read the rank-major DCP top-k gather in place

tests/v1/attention/test_indexer_dcp_localize.py

@@ -120,7 +120,7 @@ def _dcp_lse_merge(
 
 
 class _FakeDCPGroup(GroupCoordinator):
-    """Single-process stand-in: ``all_gather`` returns the pre-built
+    """Single-process stand-in: ``all_gather`` returns the pre-built dim-0
     concatenation of every rank's packed ``(score, global_id)`` candidates,
     mirroring the one packed all-gather the merge issues.
 
@@ -130,7 +130,7 @@ def __init__(self, gathered_packed: torch.Tensor) -> None:
         self.gathered_packed = gathered_packed
 
     def all_gather(self, input_: torch.Tensor, dim: int = -1) -> torch.Tensor:
-        assert dim == 1
+        assert dim == 0
         return self.gathered_packed.clone()
 
 
@@ -234,7 +234,7 @@ def _merge_local_topks_global_with_fake_dcp(
             torch.stack((scores.float(), global_ids.to(torch.float32)), dim=-1)
         )
 
-    fake_group = _FakeDCPGroup(torch.cat(packed_per_rank, dim=1).contiguous())
+    fake_group = _FakeDCPGroup(torch.cat(packed_per_rank, dim=0).contiguous())
     original_get_dcp_group = sparse_indexer.get_dcp_group
     sparse_indexer.get_dcp_group = lambda: fake_group
     try:
@@ -574,13 +574,15 @@ def test_cutedsl_dcp_candidate_pack_and_select_matches_reference(
         torch.testing.assert_close(packed[..., 1], expected_ids)
         packed_by_rank.append(packed)
 
-    gathered = torch.cat(packed_by_rank, dim=1).contiguous()
+    # Rank-major [world, rows, topk, 2], as a dim-0 all-gather produces.
+    gathered = torch.stack(packed_by_rank)
     actual = torch.empty((rows, topk), device=device, dtype=torch.int32)
     returned = stable_topk_from_gathered_candidates_cutedsl(gathered, topk, out=actual)
     assert returned is actual
+    all_candidates = torch.cat(packed_by_rank, dim=1)
     expected = _ref_stable_topk_from_candidates_fp64(
-        gathered[..., 0],
-        gathered[..., 1].to(torch.int32),
+        all_candidates[..., 0],
+        all_candidates[..., 1].to(torch.int32),
         topk,
     )
 

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

@@ -46,9 +46,23 @@ def stable_topk_from_gathered_candidates_cutedsl(
     topk: int,
     out: torch.Tensor | None = None,
 ) -> torch.Tensor:
+    """Select each row's global top-``topk`` from all ranks' candidates.
+
+    ``gathered`` is rank-major ``[world_size, rows, topk, 2]``, as produced by a
+    dim-0 all-gather of each rank's ``[rows, topk, 2]`` candidates, and is read
+    in place: a dim-1 gather into ``[rows, world_size * topk, 2]`` would cost a
+    full reorder copy.
+    """
+    assert gathered.dim() == 4 and gathered.shape[2:] == (topk, 2), (
+        f"Expected [world_size, rows, {topk}, 2] candidates, "
+        f"got {tuple(gathered.shape)}"
+    )
+    assert gathered.stride(-1) == 1 and gathered.stride(-2) == 2, (
+        f"Candidate pairs must be contiguous, got strides {gathered.stride()}"
+    )
     if out is None:
         out = torch.empty(
-            (gathered.shape[0], topk),
+            (gathered.shape[1], topk),
             dtype=torch.int32,
             device=gathered.device,
         )
@@ -301,6 +315,10 @@ def kernel(compile_key: CompileKey) -> Any:
             f"to be a multiple of {tb_size}, got {compile_key.num_candidates}"
         )
         keys_per_thread = compile_key.num_candidates // tb_size
+        # Each rank contributes topk candidates (rank-major input). Whole
+        # tb_size-wide slices then fall in one rank, so each thread key's rank
+        # and offset are compile-time constants plus tid.
+        assert topk % tb_size == 0, f"Expected {topk=} to be a multiple of {tb_size=}"
 
         @cute.struct
         class SharedStorage:
@@ -486,9 +504,9 @@ def radix_pass(
 
         @cute.kernel
         def device_kernel(input: cute.Tensor, out: cute.Tensor):
+            # input is rank-major [world_size, rows, topk, 2].
             row, _, _ = cute.arch.block_idx()
             tid, _, _ = cute.arch.thread_idx()
-            input_row = input[row, None, None]
             output_row = out[row, None]
             keys = cute.make_rmem_tensor((keys_per_thread,), Uint64)
 
@@ -500,9 +518,10 @@ def device_kernel(input: cute.Tensor, out: cute.Tensor):
                 output_row[i] = Int32(-1)
 
             for key_idx in cutlass.range_constexpr(keys_per_thread):
-                col = tid + Int32(key_idx * tb_size)
-                score = Float32(input_row[col, 0])
-                token_id = Int32(input_row[col, 1])
+                src_rank = Int32((key_idx * tb_size) // topk)
+                cand = tid + Int32((key_idx * tb_size) % topk)
+                score = Float32(input[src_rank, row, cand, 0])
+                token_id = Int32(input[src_rank, row, cand, 1])
                 keys[key_idx] = stable_key(score, token_id)
 
             if tid == Int32(0):
@@ -541,7 +560,7 @@ def host_entrypoint(
             out: cute.Tensor,
             stream: CUstream,
         ):
-            grid = (gathered.shape[0], 1, 1)
+            grid = (gathered.shape[1], 1, 1)  # rows (rank-major input)
             device_kernel(gathered, out).launch(
                 grid=grid,
                 block=(tb_size, 1, 1),
@@ -561,7 +580,9 @@ def get_warmup_keys(self, vllm_config: Any) -> list[CompileKey]:
         if dcp_world_size <= 1:
             return []
         topk, _ = _dcp_merge_topk_and_interleave(vllm_config)
-        if topk <= 0:
+        # The merge only runs with topk in (512, 1024, 2048): whole tb_size
+        # slices of each rank's candidates.
+        if topk <= 0 or topk % self.tb_size != 0:
             return []
         return self._trace_dispatch(self.dispatch)(
             topk=topk,
@@ -570,10 +591,17 @@ def get_warmup_keys(self, vllm_config: Any) -> list[CompileKey]:
 
     def warmup_inputs(self, compile_key: CompileKey) -> tuple[Any, ...]:
         num_rows = cute.sym_int()
+        # Rank-major [world_size, rows, topk, 2] candidates.
+        world_size = compile_key.num_candidates // compile_key.topk
         gathered = cute.runtime.make_fake_tensor(
             Float32,
-            (num_rows, compile_key.num_candidates, 2),
-            stride=(cute.sym_int64(divisibility=2), 2, 1),
+            (world_size, num_rows, compile_key.topk, 2),
+            stride=(
+                cute.sym_int64(divisibility=2),
+                cute.sym_int64(divisibility=2),
+                2,
+                1,
+            ),
             assumed_align=8,
         )
         out = make_fake_tensor(
@@ -593,7 +621,7 @@ def __call__(
     ) -> CuTeDSLLaunchSpec[CompileKey]:
         compile_key = self.dispatch(
             topk=topk,
-            num_candidates=gathered.shape[1],
+            num_candidates=gathered.shape[0] * topk,
         )
         return compile_key, (gathered, out)
 

vllm/model_executor/layers/sparse_attn_indexer.py

@@ -152,9 +152,14 @@ def _merge_dcp_topk_global(
         cp_interleave,
         row_starts,
     )
-    gathered = get_dcp_group().all_gather(packed, dim=1)
+    # Gather along dim 0 and let the selector read the rank-major
+    # [world_size, rows, topk, 2] result: a dim-1 gather would reorder it into
+    # [rows, world_size * topk, 2] with a full extra copy.
+    gathered = get_dcp_group().all_gather(packed, dim=0)
     stable_topk_from_gathered_candidates_cutedsl(
-        gathered, topk_tokens, out=topk_indices
+        gathered.view(dcp_world_size, *packed.shape),
+        topk_tokens,
+        out=topk_indices,
     )