From a563e578b1054cb604bc75a03856d592ed950601 Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Thu, 30 Jul 2026 22:00:17 +0000 Subject: [PATCH] perf(indexer): restore query-split indices without score traffic --- .../test_indexer_parallel_groups.py | 45 ++++++++++ .../layers/test_sparse_attn_indexer_b12x.py | 88 +++++++++++++++++++ vllm/distributed/parallel_state.py | 18 ++-- .../layers/sparse_attn_indexer.py | 77 +++++++++++++--- 4 files changed, 205 insertions(+), 23 deletions(-) diff --git a/tests/distributed/test_indexer_parallel_groups.py b/tests/distributed/test_indexer_parallel_groups.py index 98b9ca8b4ca6..8414eb255eec 100644 --- a/tests/distributed/test_indexer_parallel_groups.py +++ b/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)) diff --git a/tests/model_executor/layers/test_sparse_attn_indexer_b12x.py b/tests/model_executor/layers/test_sparse_attn_indexer_b12x.py index cfc785c8ecff..71c9e6ad2fd5 100644 --- a/tests/model_executor/layers/test_sparse_attn_indexer_b12x.py +++ b/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", [ diff --git a/vllm/distributed/parallel_state.py b/vllm/distributed/parallel_state.py index 0f2ef8564b65..0364580062f8 100644 --- a/vllm/distributed/parallel_state.py +++ b/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, diff --git a/vllm/model_executor/layers/sparse_attn_indexer.py b/vllm/model_executor/layers/sparse_attn_indexer.py index 3c1c1cd0e70f..8733a5b79948 100644 --- a/vllm/model_executor/layers/sparse_attn_indexer.py +++ b/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