[GG] perf(indexer): shard sparse prefill queries and halve result traffic

tests/distributed/test_indexer_parallel_groups.py

@@ -21,6 +21,39 @@ def test_build_indexer_two_by_four_groups_for_tp8():
     assert query_split_groups == [[0, 4], [1, 5], [2, 6], [3, 7]]
 
 
+@pytest.mark.parametrize(
+    ("indexer_shards", "expected_dcp", "expected_query_split"),
+    [
+        (
+            1,
+            [[0], [1], [2], [3], [4], [5], [6], [7]],
+            [list(range(8))],
+        ),
+        (
+            2,
+            [[0, 1], [2, 3], [4, 5], [6, 7]],
+            [[0, 2, 4, 6], [1, 3, 5, 7]],
+        ),
+        (
+            8,
+            [list(range(8))],
+            [[0], [1], [2], [3], [4], [5], [6], [7]],
+        ),
+    ],
+)
+def test_build_indexer_groups_cover_dcp1_partial_and_full(
+    indexer_shards,
+    expected_dcp,
+    expected_query_split,
+):
+    dcp_groups, query_split_groups = _build_indexer_replica_group_ranks(
+        [list(range(8))], indexer_shards
+    )
+
+    assert dcp_groups == expected_dcp
+    assert query_split_groups == expected_query_split
+
+
 def test_build_indexer_replica_groups_stay_inside_each_tp_group():
     dcp_groups, query_split_groups = _build_indexer_replica_group_ranks(
         [list(range(8)), list(range(8, 16))], 4
@@ -76,6 +109,18 @@ def test_indexer_group_selector_supports_partial_target_and_full_draft(monkeypat
     assert parallel_state.get_indexer_query_split_group(4) is full_query_split
 
 
+def test_indexer_group_selector_uses_tp_query_split_for_dcp1(monkeypatch):
+    dcp = SimpleNamespace(world_size=1)
+    query_split = SimpleNamespace(world_size=8)
+    monkeypatch.setattr(parallel_state, "_INDEXER_DCP", None)
+    monkeypatch.setattr(parallel_state, "_INDEXER_QUERY_SPLIT", None)
+    monkeypatch.setattr(parallel_state, "_DCP", dcp)
+    monkeypatch.setattr(parallel_state, "_QUERY_SPLIT", query_split)
+
+    assert parallel_state.get_indexer_dcp_group(1) is dcp
+    assert parallel_state.get_indexer_query_split_group(1) is query_split
+
+
 def test_indexer_group_selector_rejects_unknown_shard_count(monkeypatch):
     monkeypatch.setattr(parallel_state, "_INDEXER_DCP", SimpleNamespace(world_size=2))
     monkeypatch.setattr(parallel_state, "_DCP", SimpleNamespace(world_size=4))

tests/model_executor/layers/test_sparse_attn_indexer_b12x.py

@@ -288,6 +288,94 @@ def gather_topk_ids_by_position(candidate_ids, positions, out):
     )
 
 
+def test_query_split_gathers_indices_in_place_without_scores():
+    calls: list[tuple[int, int]] = []
+
+    class FakePyNccl:
+        disabled = False
+
+        def all_gather(self, output, input_):
+            calls.append((output.data_ptr(), input_.data_ptr()))
+            rows = input_.shape[0]
+            output[:rows].fill_(10)
+            output[rows:].fill_(20)
+
+    group = types.SimpleNamespace(
+        world_size=2,
+        rank_in_group=1,
+        device_communicator=types.SimpleNamespace(pynccl_comm=FakePyNccl()),
+    )
+    gathered_indices = torch.full((4, 3), -1, dtype=torch.int32)
+    local_indices = gathered_indices[2:]
+    local_indices.fill_(20)
+    scores = torch.arange(12, dtype=torch.float32).reshape(4, 3)
+    scores_before = scores.clone()
+
+    indexer_mod._query_split_all_gather_indices(
+        group,
+        local_indices,
+        gathered_indices,
+    )
+
+    assert calls == [(gathered_indices.data_ptr(), local_indices.data_ptr())]
+    assert gathered_indices.tolist() == [[10] * 3, [10] * 3, [20] * 3, [20] * 3]
+    assert torch.equal(scores, scores_before)
+
+
+def test_query_split_copies_aliased_input_for_torch_distributed_fallback(
+    monkeypatch,
+):
+    calls: list[tuple[int, int]] = []
+
+    def fake_all_gather_into_tensor(output, input_, *, group):
+        assert group == "device-group"
+        calls.append((output.data_ptr(), input_.data_ptr()))
+        rows = input_.shape[0]
+        output[:rows].fill_(10)
+        output[rows:].copy_(input_)
+
+    monkeypatch.setattr(
+        torch.distributed,
+        "all_gather_into_tensor",
+        fake_all_gather_into_tensor,
+    )
+    group = types.SimpleNamespace(
+        world_size=2,
+        rank_in_group=1,
+        device_group="device-group",
+        device_communicator=types.SimpleNamespace(
+            pynccl_comm=types.SimpleNamespace(disabled=True),
+        ),
+    )
+    gathered_indices = torch.full((4, 3), -1, dtype=torch.int32)
+    local_indices = gathered_indices[2:]
+    local_indices.fill_(20)
+    local_ptr = local_indices.data_ptr()
+
+    indexer_mod._query_split_all_gather_indices(
+        group,
+        local_indices,
+        gathered_indices,
+    )
+
+    assert calls[0][0] == gathered_indices.data_ptr()
+    assert calls[0][1] != local_ptr
+    assert gathered_indices.tolist() == [[10] * 3, [10] * 3, [20] * 3, [20] * 3]
+
+
+def test_query_split_rejects_non_aliasing_local_indices():
+    group = types.SimpleNamespace(world_size=2, rank_in_group=0)
+    gathered_indices = torch.empty((4, 3), dtype=torch.int32)
+    local_indices = torch.empty((2, 3), dtype=torch.int32)
+
+    with pytest.raises(RuntimeError, match="must alias"):
+        indexer_mod._query_split_all_gather_indices(
+            group,
+            local_indices,
+            gathered_indices,
+        )
+
+
 @pytest.mark.parametrize(
     "page_stride0",
     [

vllm/distributed/parallel_state.py

@@ -2073,17 +2073,17 @@ def initialize_model_parallel(
         group_name="dcp",
     )
 
-    # Build the query-split groups for the indexer query split (Fix A).
-    # Ranks sharing the same dcp_rank (position within their DCP group)
-    # form a query-split group.  At TP=8/DCP=2 the DCP groups are
-    # {0,1},{2,3},{4,5},{6,7} and the query-split groups are
-    # {0,2,4,6} (dcp_rank=0) and {1,3,5,7} (dcp_rank=1).
+    # Build the full-indexer query-split groups from the same topology helper
+    # used by partially replicated indexers below. At TP8/DCP2 this produces
+    # {0,2,4,6}/{1,3,5,7}. DCP1 intentionally produces one TP-wide group:
+    # every rank has the replicated indexer inputs and can process a query-row
+    # shard before restoring the exact int32 top-k indices.
     global _QUERY_SPLIT
     assert _QUERY_SPLIT is None, "query split group is already initialized"
-    if decode_context_model_parallel_size > 1 and envs.VLLM_DCP_QUERY_SPLIT:
-        query_split_ranks: list[list[int]] = []
-        for dcp_rank_idx in range(decode_context_model_parallel_size):
-            query_split_ranks.append([grp[dcp_rank_idx] for grp in group_ranks])
+    if envs.VLLM_DCP_QUERY_SPLIT:
+        _, query_split_ranks = _build_indexer_replica_group_ranks(
+            tp_group_ranks, decode_context_model_parallel_size
+        )
         _QUERY_SPLIT = init_model_parallel_group(
             query_split_ranks,
             get_world_group().local_rank,

vllm/model_executor/layers/sparse_attn_indexer.py

@@ -407,17 +407,68 @@ def _dcp_all_gather_first_dim_into(
         pynccl_comm.all_gather(output_tensor, input_tensor)
         return
 
-    device_group = getattr(communicator, "device_group", None)
+    device_group = getattr(group, "device_group", None)
+    if device_group is None:
+        device_group = getattr(communicator, "device_group", None)
     if device_group is not None:
         import torch.distributed as dist
 
-        dist.all_gather_into_tensor(output_tensor, input_tensor, group=device_group)
+        gather_input = input_tensor
+        input_start = input_tensor.data_ptr()
+        input_end = input_start + input_tensor.numel() * input_tensor.element_size()
+        output_start = output_tensor.data_ptr()
+        output_end = output_start + output_tensor.numel() * output_tensor.element_size()
+        if input_start < output_end and output_start < input_end:
+            # PyTorch does not expose NCCL's in-place all-gather contract.
+            # Keep the generic fallback portable; the deployed PyNCCL path
+            # above retains the rank-local zero-copy layout.
+            gather_input = input_tensor.clone()
+        dist.all_gather_into_tensor(output_tensor, gather_input, group=device_group)
         return
 
     gathered = group.all_gather(input_tensor, dim=0)
     output_tensor.copy_(gathered)
 
 
+def _query_split_all_gather_indices(
+    group,
+    local_indices: torch.Tensor,
+    gathered_indices: torch.Tensor,
+) -> None:
+    """Restore query-split rows directly into their shared index buffer.
+
+    Args:
+        group: Query-split process group and rank metadata.
+        local_indices: Contiguous rank-local rows aliasing their output slot.
+        gathered_indices: Contiguous output containing every rank's rows.
+
+    Raises:
+        RuntimeError: If the buffers are noncontiguous, have incompatible
+            shapes, or do not satisfy the rank-local alias contract.
+    """
+    if group.world_size <= 1:
+        return
+    if not local_indices.is_contiguous() or not gathered_indices.is_contiguous():
+        raise RuntimeError("query-split top-k buffers must be contiguous")
+    if gathered_indices.shape[0] != local_indices.shape[0] * group.world_size:
+        raise RuntimeError("query-split output has an invalid row count")
+    if gathered_indices.shape[1:] != local_indices.shape[1:]:
+        raise RuntimeError("query-split output has an invalid trailing shape")
+
+    rank = int(group.rank_in_group)
+    expected_local = gathered_indices.narrow(
+        0, rank * local_indices.shape[0], local_indices.shape[0]
+    )
+    if expected_local.data_ptr() != local_indices.data_ptr():
+        raise RuntimeError(
+            "query-split local indices must alias their rank slot in the output"
+        )
+
+    # NCCL supports in-place all-gather when the send buffer aliases the
+    # rank-local receive slice. The generic fallback copies only that slice.
+    _dcp_all_gather_first_dim_into(group, local_indices, gathered_indices)
+
+
 def _unpack_b12x_dcp_gathered_candidates(
     gathered_candidates: torch.Tensor,
     candidate_indices: torch.Tensor,
@@ -1994,19 +2045,17 @@ def sparse_attn_indexer(
                 if used_owner_merge:
                     continue
                 if qs_active:
-                    gathered_indices = qs_group.all_gather(
-                        topk_indices.contiguous(), dim=0
-                    )
-                    topk_indices_buffer[
+                    gathered_indices = topk_indices_buffer[
                         chunk.token_start : chunk.token_end, :topk_tokens
-                    ].copy_(gathered_indices)
-                    if topk_scores is not None:
-                        gathered_scores = qs_group.all_gather(
-                            topk_scores.contiguous(), dim=0
-                        )
-                        topk_scores_buffer[
-                            chunk.token_start : chunk.token_end, :topk_tokens
-                        ].copy_(gathered_scores)
+                    ]
+                    _query_split_all_gather_indices(
+                        qs_group,
+                        topk_indices,
+                        gathered_indices,
+                    )
+                    # Scores are consumed only by the DCP-local candidate
+                    # merge above. Sparse attention needs the restored indices,
+                    # so gathering scores here would double result traffic.
                 continue
 
             cu_seqlen_ks = chunk.cu_seqlen_ks