From 60f6f0bc822b4f648525aa3419a24a464666bd20 Mon Sep 17 00:00:00 2001 From: lanqinghuan Date: Sat, 29 Aug 2026 15:27:43 +0800 Subject: [PATCH] [Perf][Attention][DSA] Shard prefill indexer rows across TP ranks The CUDA DSA prefill indexer recomputes identical logits and Top-K on every tensor-parallel rank. Assign each prefill query row to exactly one rank and publish the finished rows with one all_gatherv per indexer layer. Rows are independent, so this is a layout-preserving concatenation rather than a Top-K merge. The partition is contiguous and balances the exact number of scored keys, derived from CPU scheduler metadata with no device sync. Assisted-by: OpenAI Codex Assisted-by: Claude Code (Opus) Signed-off-by: lanqinghuan --- .../v1/attention/test_indexer_tp_row_shard.py | 524 ++++++++++++++++++ .../layers/sparse_attn_indexer.py | 68 ++- vllm/v1/attention/backends/mla/indexer.py | 133 ++++- 3 files changed, 705 insertions(+), 20 deletions(-) create mode 100644 tests/v1/attention/test_indexer_tp_row_shard.py diff --git a/tests/v1/attention/test_indexer_tp_row_shard.py b/tests/v1/attention/test_indexer_tp_row_shard.py new file mode 100644 index 000000000000..bc4ed413b568 --- /dev/null +++ b/tests/v1/attention/test_indexer_tp_row_shard.py @@ -0,0 +1,524 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""TP query-row sharding of the DSA indexer prefill. + +The indexer projections are replicated across TP, so every rank recomputes the +same prefill logits and the same top-k. Rows are independent (one block per row +over that row's ``[ks, ke)``), so each rank can own a disjoint slice and the +group exchanges ``index_topk`` int32s per row instead of the logits. +""" + +from types import SimpleNamespace + +import pytest +import torch + +import vllm.model_executor.layers.sparse_attn_indexer as sparse_indexer +import vllm.v1.attention.backends.mla.indexer as indexer +from vllm.config import CUDAGraphMode +from vllm.v1.attention.backends.mla.indexer import ( + DeepseekV32IndexerMetadata, + DeepseekV32IndexerPrefillChunkMetadata, +) + +INDEXER_LAYER = "model.layers.0.self_attn.indexer.k_cache" + +_TOPK = 8 +_NUM_KV = 64 +# Leading decode rows of the shared buffer. The prefill shard must address the +# window [num_decode_tokens, num_decode_tokens + num_prefill_tokens) and leave +# these alone; mixed decode+prefill batches are the concurrency > 1 case. +_DECODE_ROWS = 5 + + +def _ref_top_k_per_row_prefill(logits, cu_ks, cu_ke, out, num_rows, _s0, _s1, topk): + """Row-independent stand-in for ``ops.top_k_per_row_prefill``. + + Mirrors the real kernel's contract (one block per row, reading only that + row's ``[ks, ke)``), which is the property the row shard rests on. + """ + positions = torch.arange(logits.shape[1]).unsqueeze(0) + lo = cu_ks[:num_rows].long().unsqueeze(1) + hi = cu_ke[:num_rows].long().unsqueeze(1) + scores = ( + logits[:num_rows] + .float() + .masked_fill((positions < lo) | (positions >= hi), -float("inf")) + ) + k = min(topk, logits.shape[1]) + picked = scores.topk(k, dim=1).indices.int() + keep = torch.arange(k).unsqueeze(0) < (hi - lo).clamp(min=0, max=k) + out[:num_rows, :k] = torch.where(keep, picked, torch.full_like(picked, -1)) + out[:num_rows, k:] = -1 + + +def _build_chunks(row_counts): + """Ragged chunks: uneven row counts, per-row causal bounds that leave some + rows short of topk, a continuation chunk that must reuse the gathered K, + and a trailing empty-KV chunk.""" + chunks, token_start = [], _DECODE_ROWS + for idx, num_rows in enumerate(row_counts): + empty_kv = idx == len(row_counts) - 1 + ke = ((torch.arange(num_rows) * 7 + idx * 3) % (_NUM_KV + 1)).int() + chunks.append( + DeepseekV32IndexerPrefillChunkMetadata( + block_table=torch.zeros(1, 1, dtype=torch.int32), + cu_seqlen_ks=torch.zeros(num_rows, dtype=torch.int32), + cu_seqlen_ke=torch.zeros_like(ke) if empty_kv else ke, + cu_seq_lens=torch.zeros(2, dtype=torch.int32), + token_to_seq=torch.zeros(1, dtype=torch.int32), + total_seq_lens=0 if empty_kv else _NUM_KV, + token_start=token_start, + token_end=token_start + num_rows, + num_reqs=1, + skip_kv_gather=idx % 2 == 1, + local_cu_seq_lens=torch.zeros(2, dtype=torch.int32), + local_total_seq_lens=0 if empty_kv else _NUM_KV, + max_local_total_seq_lens=_NUM_KV, + ) + ) + token_start += num_rows + return chunks, token_start + + +def _bound_tables(chunks, num_tokens): + """Per-global-row causal bounds, zero for the leading decode rows.""" + ks = torch.zeros(num_tokens, dtype=torch.int32) + ke = torch.zeros(num_tokens, dtype=torch.int32) + chunk_ks = torch.cat([c.cu_seqlen_ks for c in chunks]) + chunk_ke = torch.cat([c.cu_seqlen_ke for c in chunks]) + rows = slice(_DECODE_ROWS, _DECODE_ROWS + chunk_ks.numel()) + ks[rows] = chunk_ks + ke[rows] = chunk_ke + return ks, ke + + +def _run_rank( + monkeypatch, *, world, rank, chunks, num_tokens, logits, exchange, split=None +): + """Drive the real ``sparse_attn_indexer`` prefill path for one TP rank. + + ``exchange(local_rows, sizes)`` stands in for the group's all_gatherv. + Returns ``(topk_buffer, gather_call_count)``. + """ + ks_table, ke_table = _bound_tables(chunks, num_tokens) + gathers = [] + + metadata = DeepseekV32IndexerMetadata( + seq_lens=torch.empty(0, dtype=torch.int32), + max_seq_len=2048, + slot_mapping=torch.zeros(num_tokens, dtype=torch.long), + num_decodes=0, + # num_decodes=0 keeps the decode path out of this test while still + # placing the prefill rows behind a decode offset in the buffer. + num_decode_tokens=_DECODE_ROWS, + num_prefills=len(chunks), + num_prefill_tokens=num_tokens - _DECODE_ROWS, + prefill=SimpleNamespace(chunks=chunks, row_shard_sizes=split), + ) + + set_ = monkeypatch.setattr + set_( + sparse_indexer, + "get_forward_context", + lambda: SimpleNamespace( + attn_metadata={INDEXER_LAYER: metadata}, + cudagraph_runtime_mode=CUDAGraphMode.PIECEWISE, + ), + ) + set_(sparse_indexer.current_platform, "fp8_dtype", lambda: torch.float16) + set_(sparse_indexer.current_platform, "is_xpu", lambda: False) + set_(sparse_indexer, "get_tensor_model_parallel_rank", lambda: rank) + set_( + sparse_indexer, + "get_tp_group", + lambda: SimpleNamespace(all_gatherv=lambda t, dim, sizes: exchange(t, sizes)), + ) + set_( + sparse_indexer, + "current_workspace_manager", + lambda: SimpleNamespace( + get_simultaneous=lambda *specs: tuple( + torch.zeros(shape, dtype=dtype) for shape, dtype in specs + ) + ), + ) + set_(sparse_indexer.ops, "top_k_per_row_prefill", _ref_top_k_per_row_prefill) + set_( + sparse_indexer.ops, + "cp_gather_indexer_k_quant_cache", + lambda *args: gathers.append(1), + ) + + def fake_mqa_logits(q, _k, weights, cu_ks, cu_ke, clean_logits=True): + rows = q[0][:, 0, 0].long() + # The q slice, the weight slice and the causal bounds must all name the + # same global rows; an off-by-one anywhere in the shard breaks this. + torch.testing.assert_close(weights[:, 0].long(), rows) + torch.testing.assert_close(cu_ks, ks_table[rows]) + torch.testing.assert_close(cu_ke, ke_table[rows]) + return logits[rows] + + set_(sparse_indexer, "fp8_fp4_mqa_logits", fake_mqa_logits) + + row_ids = torch.arange(num_tokens, dtype=torch.float32) + buffer = torch.full((num_tokens, _TOPK), 17, dtype=torch.int32) + sparse_indexer.sparse_attn_indexer( + torch.zeros(num_tokens, 1), # hidden_states + INDEXER_LAYER, + torch.empty(1), # kv_cache + row_ids.reshape(num_tokens, 1, 1), # q_quant carries its global row id + None, # q_scale + None, # k + row_ids.reshape(num_tokens, 1), # weights carry it too + 128, + "ue8m0", + _TOPK, + 4, + 4096, + _NUM_KV, + buffer, + True, # skip_k_cache_insert + False, # use_pcp + "", # dense_mha_metadata_layer_name + ) + return buffer, len(gathers) + + +def _run_group(monkeypatch, world, chunks, num_tokens, logits, split=None): + """Collect every rank's slice, then replay the concatenation each rank + would receive. There is exactly one exchange per forward.""" + rows = sum(chunk.token_end - chunk.token_start for chunk in chunks) + split = split or [rows // world + int(rank < rows % world) for rank in range(world)] + slices: dict[int, torch.Tensor] = {} + results: list[tuple[torch.Tensor, int]] = [] + for replay in (False, True): + results = [] + for rank in range(world): + + def exchange(local, sizes, rank=rank, replay=replay): + assert local.is_contiguous() + assert sizes == split + assert local.shape[0] == sizes[rank] + slices[rank] = local.clone() + if not replay: + return torch.zeros(sum(sizes), _TOPK, dtype=torch.int32) + return torch.cat([slices[r] for r in range(world)]) + + with monkeypatch.context() as m: + results.append( + _run_rank( + m, + world=world, + rank=rank, + chunks=chunks, + num_tokens=num_tokens, + logits=logits, + exchange=exchange, + split=split, + ) + ) + return results + + +def test_tp_row_shard_prefill_ignores_padded_tail( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """The gathered real rows must not be assigned across padded tokens.""" + chunks, unpadded_end = _build_chunks([2049, 849, 651, 549]) + num_tokens = unpadded_end + 13 + logits = torch.randn( + num_tokens, _NUM_KV, generator=torch.Generator().manual_seed(53691) + ) + + def no_exchange(local, sizes): + raise AssertionError("the replicated reference must not exchange") + + with monkeypatch.context() as m: + baseline, baseline_gathers = _run_rank( + m, + world=1, + rank=0, + chunks=chunks, + num_tokens=num_tokens, + logits=logits, + exchange=no_exchange, + ) + + for rank, (buffer, gathers) in enumerate( + _run_group( + monkeypatch, + world=4, + chunks=chunks, + num_tokens=num_tokens, + logits=logits, + split=[2049, 849, 651, 549], + ) + ): + torch.testing.assert_close( + buffer[:unpadded_end], baseline[:unpadded_end], msg=f"rank {rank} diverged" + ) + assert torch.all(buffer[unpadded_end:] == -1) + assert gathers == baseline_gathers + + +@pytest.mark.parametrize("world", [2, 3, 8]) +@pytest.mark.parametrize("ties", [False, True]) +@pytest.mark.parametrize("uneven", [False, True]) +def test_tp_row_shard_prefill_matches_row_independent_reference( + monkeypatch: pytest.MonkeyPatch, world: int, ties: bool, uneven: bool +) -> None: + """Every TP rank receives the reference result for every completed row. + + Covers non-power-of-two tp_size, ragged chunks that straddle the shard + boundary, rows with fewer than topk candidates, an empty-KV chunk, and a + dense-tie logits table (ties are where a row-order-dependent merge would + diverge), and both the equal-row and a lopsided cost-balanced split. + """ + # Two chunks deliberately have the same local query range length. The + # gather decision is positional metadata, not a set keyed by row bounds. + chunks, num_tokens = _build_chunks([3300, 2500, 3300, 7, 1100]) + logits = torch.randn( + num_tokens, _NUM_KV, generator=torch.Generator().manual_seed(1003) + ) + if ties: + logits = (logits * 2).round() / 2 + + def no_exchange(local, sizes): + raise AssertionError("tp_size == 1 must not exchange") + + baseline, baseline_gathers = _run_rank( + monkeypatch, + world=1, + rank=0, + chunks=chunks, + num_tokens=num_tokens, + logits=logits, + exchange=no_exchange, + ) + expected_gathers = sum( + chunk.total_seq_lens > 0 and not chunk.skip_kv_gather for chunk in chunks + ) + assert baseline_gathers == expected_gathers + # No row is silently dropped: each holds exactly min(ke - ks, topk) valid + # slots, padded with -1. + ks_table, ke_table = _bound_tables(chunks, num_tokens) + torch.testing.assert_close( + (baseline >= 0).sum(dim=1).int(), + (ke_table - ks_table).clamp(min=0, max=_TOPK).int(), + ) + assert torch.all(baseline[:_DECODE_ROWS] == -1) + + rows = num_tokens - _DECODE_ROWS + split = None + if uneven: + # a lopsided but valid partition, as cost balancing produces + split = [rows - (world - 1) * _TOPK] + [_TOPK] * (world - 1) + for rank, (buffer, gathers) in enumerate( + _run_group(monkeypatch, world, chunks, num_tokens, logits, split) + ): + torch.testing.assert_close(buffer, baseline, msg=f"rank {rank} diverged") + assert torch.all(buffer[:_DECODE_ROWS] == -1), "exchange clobbered decode rows" + # The K gather is a workspace side effect that later chunks reuse via + # `skip_kv_gather`; narrowing the query rows must not change it. + assert gathers == baseline_gathers + + +def _scored_keys(seq_lens, query_lens, compress_ratio): + """Per-row ke - ks, spelled out the way the Triton metadata kernel does.""" + out = [] + for seq_len, query_len in zip(seq_lens, query_lens): + context = seq_len - query_len + out += [(context + 1 + j) // compress_ratio for j in range(query_len)] + return out + + +def _rank_costs(sizes, per_row): + costs, off = [], 0 + for size in sizes: + costs.append(sum(per_row[off : off + size])) + off += size + return costs + + +@pytest.mark.parametrize("tp_size", [2, 3, 4, 8]) +@pytest.mark.parametrize( + "shape", + [ + pytest.param("fresh", id="fresh_prompt"), + pytest.param("prefix", id="prefix_context"), + pytest.param("ragged", id="mixed_ragged"), + pytest.param("chunk_tail", id="chunked_prefill_tail"), + pytest.param("tight", id="just_above_floor"), + pytest.param("deep_first", id="deep_context_first"), + ], +) +def test_balanced_row_shard_equalises_scored_keys(tp_size: int, shape: str) -> None: + """The split beats equal rows on cost balance, or declines to split.""" + rows_needed = indexer.MIN_TP_SHARD_ROWS_PER_RANK * tp_size + per_req = 4 * rows_needed + if shape == "fresh": # no context: the full causal ramp + query_lens, seq_lens = [per_req], [per_req] + elif shape == "prefix": # long shared prefix already in cache + query_lens, seq_lens = [per_req], [per_req + 200_000] + elif shape == "ragged": # uneven requests at different depths + query_lens = [per_req // 2, 7, per_req, per_req // 3 + 5] + seq_lens = [q + c for q, c in zip(query_lens, [0, 9_000, 500, 60_000])] + elif shape == "deep_first": + # An expensive deep-context request ahead of a cheap fresh one, so the + # cost profile falls and the early ranks own fewer rows. + query_lens = [rows_needed // 2, 2 * rows_needed - rows_needed // 2] + seq_lens = [query_lens[0] + 500_000, query_lens[1]] + elif shape == "tight": + # Just above the total-size gate, where sub-1024 inexpensive shards are + # still beneficial. + query_lens, seq_lens = [rows_needed + 2], [rows_needed + 2] + else: # a mid-prompt chunk: high, nearly flat cost + query_lens, seq_lens = [per_req], [per_req + 15 * per_req] + + compress_ratio = 4 + sizes = indexer.balanced_prefill_row_shard( + torch.tensor(seq_lens, dtype=torch.int32), + torch.tensor(query_lens, dtype=torch.int32), + compress_ratio, + tp_size, + ) + assert sizes is not None + num_rows = sum(query_lens) + assert len(sizes) == tp_size + assert sum(sizes) == num_rows, "the split must cover every row exactly once" + assert min(sizes) >= 1 + + per_row = _scored_keys(seq_lens, query_lens, compress_ratio) + assert per_row == [ + int(x) for x in _reference_ke_minus_ks(seq_lens, query_lens, compress_ratio) + ] + base, rem = divmod(num_rows, tp_size) + equal_sizes = [base + int(r < rem) for r in range(tp_size)] + + def imbalance(split): + costs = _rank_costs(split, per_row) + return max(costs) / (sum(costs) / len(costs)) + + assert imbalance(sizes) <= imbalance(equal_sizes) + 1e-9 + assert imbalance(sizes) < 1.02 + + +def _reference_ke_minus_ks(seq_lens, query_lens, compress_ratio): + """Independent restatement of the kernel formula, vectorised.""" + out = [] + for seq_len, query_len in zip(seq_lens, query_lens): + pos = torch.arange(query_len) + (seq_len - query_len) + 1 + out += (pos // compress_ratio).tolist() + return out + + +def test_balanced_row_shard_declines_below_the_floor() -> None: + """Too few rows to give every rank the floor -> no split at all.""" + tp_size = 4 + rows = indexer.MIN_TP_SHARD_ROWS_PER_RANK * tp_size - 1 + assert ( + indexer.balanced_prefill_row_shard( + torch.tensor([rows], dtype=torch.int32), + torch.tensor([rows], dtype=torch.int32), + 4, + tp_size, + ) + is None + ) + assert ( + indexer.balanced_prefill_row_shard( + torch.tensor([rows + 1], dtype=torch.int32), + torch.tensor([rows + 1], dtype=torch.int32), + 4, + 1, + ) + is None + ) + + +@pytest.mark.parametrize( + "max_seq_len,index_topk,force_mqa,expected", + [ + (2047, 2048, False, False), + (2048, 2048, False, False), + (2049, 2048, False, True), + (2047, 2048, True, True), + (2048, 2048, True, True), + ], +) +def test_prefill_uses_mqa_matches_attention_routing( + max_seq_len: int, index_topk: int, force_mqa: bool, expected: bool +) -> None: + """Row sharding only runs when sparse MQA consumes the Top-K output.""" + assert indexer._prefill_uses_mqa(max_seq_len, index_topk, force_mqa) is expected + + +def _sharding_config(cudagraph_mode=CUDAGraphMode.PIECEWISE): + return SimpleNamespace( + compilation_config=SimpleNamespace(cudagraph_mode=cudagraph_mode) + ) + + +def _set_sharding_env( + monkeypatch: pytest.MonkeyPatch, overrides: dict[str, bool] | None = None +) -> None: + """Set every lazy environment flag read by the sharding gate.""" + values = { + "VLLM_DISABLE_PYNCCL": False, + "VLLM_USE_NCCL_SYMM_MEM": False, + "VLLM_BATCH_INVARIANT": False, + **(overrides or {}), + } + for name, value in values.items(): + monkeypatch.setenv(name, str(int(value))) + + +@pytest.mark.parametrize( + "kwargs,env,expected", + [ + ({}, {}, True), + ({"tp_size": 1}, {}, False), + ({"dcp_world_size": 2}, {}, False), + ({"use_pcp": True}, {}, False), + ({}, {"VLLM_DISABLE_PYNCCL": True}, False), + ({}, {"VLLM_USE_NCCL_SYMM_MEM": True}, False), + ({}, {"VLLM_BATCH_INVARIANT": True}, False), + ], +) +def test_row_sharding_gate_rejects_unsupported_configurations( + monkeypatch: pytest.MonkeyPatch, kwargs: dict, env: dict, expected: bool +) -> None: + """The gate is the whole safety envelope; nothing else guards the exchange.""" + monkeypatch.setattr(indexer.current_platform, "is_cuda", lambda: True) + _set_sharding_env(monkeypatch, env) + args = {"dcp_world_size": 1, "use_pcp": False, "tp_size": 4, **kwargs} + supported = indexer.tp_prefill_row_sharding_supported(_sharding_config(), **args) + assert supported is expected + + +@pytest.mark.parametrize( + "cudagraph_mode,expected", + [ + (None, True), + (CUDAGraphMode.NONE, True), + (CUDAGraphMode.PIECEWISE, True), + # Mixed batches run under PIECEWISE, so prefill is never captured whole. + (CUDAGraphMode.FULL_AND_PIECEWISE, True), + (CUDAGraphMode.FULL_DECODE_ONLY, True), + # Mixed batches would be captured whole; the exchange must stay out. + (CUDAGraphMode.FULL, False), + ], +) +def test_row_sharding_gate_follows_the_mixed_batch_cudagraph_mode( + monkeypatch: pytest.MonkeyPatch, cudagraph_mode, expected: bool +) -> None: + monkeypatch.setattr(indexer.current_platform, "is_cuda", lambda: True) + _set_sharding_env(monkeypatch) + supported = indexer.tp_prefill_row_sharding_supported( + _sharding_config(cudagraph_mode), + dcp_world_size=1, + use_pcp=False, + tp_size=4, + ) + assert supported is expected diff --git a/vllm/model_executor/layers/sparse_attn_indexer.py b/vllm/model_executor/layers/sparse_attn_indexer.py index 488eb3899c4d..8ade2c5adfaf 100644 --- a/vllm/model_executor/layers/sparse_attn_indexer.py +++ b/vllm/model_executor/layers/sparse_attn_indexer.py @@ -9,7 +9,12 @@ from vllm._aiter_ops import rocm_aiter_ops from vllm.compilation.breakable_cudagraph import eager_break_during_capture from vllm.config import CUDAGraphMode, get_current_vllm_config -from vllm.distributed import get_dcp_group, get_pcp_group +from vllm.distributed import ( + get_dcp_group, + get_pcp_group, + get_tensor_model_parallel_rank, + get_tp_group, +) from vllm.forward_context import get_forward_context from vllm.logger import init_logger from vllm.model_executor.custom_op import CustomOp @@ -446,6 +451,17 @@ def sparse_attn_indexer( values_spec, scales_spec, ) + # Contiguous window of prefill rows this rank scores. The builder emits + # row_shard_sizes from replicated scheduler metadata, so every TP rank + # agrees on whether the exchange below runs. + shard_sizes = prefill_metadata.row_shard_sizes + shard_start = shard_stop = 0 + if shard_sizes is not None: + assert dcp_world_size == 1 and not use_pcp + assert forward_context.cudagraph_runtime_mode != CUDAGraphMode.FULL + tp_rank = get_tensor_model_parallel_rank() + shard_start = num_decode_tokens + sum(shard_sizes[:tp_rank]) + shard_stop = shard_start + shard_sizes[tp_rank] for chunk in prefill_metadata.chunks: cu_seqlen_ks = chunk.cu_seqlen_ks cu_seqlen_ke = chunk.cu_seqlen_ke @@ -461,15 +477,23 @@ def sparse_attn_indexer( chunk.local_cu_seq_lens, ) - q_slice = q_quant[chunk.token_start : chunk.token_end] - q_scale_slice = ( - q_scale[chunk.token_start : chunk.token_end] - if q_scale is not None - else None - ) - topk_indices = topk_indices_buffer[ - chunk.token_start : chunk.token_end, :topk_tokens - ] + # Narrow the scoring to this rank's rows. The gather above stays + # unconditional: later chunks of the same request reuse that + # workspace via `skip_kv_gather`, so every rank must fill it. + row_start, row_end = chunk.token_start, chunk.token_end + if shard_sizes is not None: + row_start = max(row_start, shard_start) + row_end = min(row_end, shard_stop) + if row_start >= row_end: + continue + lo = row_start - chunk.token_start + hi = row_end - chunk.token_start + cu_seqlen_ks = cu_seqlen_ks[lo:hi] + cu_seqlen_ke = cu_seqlen_ke[lo:hi] + + q_slice = q_quant[row_start:row_end] + q_scale_slice = q_scale[row_start:row_end] if q_scale is not None else None + topk_indices = topk_indices_buffer[row_start:row_end, :topk_tokens] if chunk.local_total_seq_lens == 0: logits = q_slice.new_empty((q_slice.shape[0], 0), dtype=torch.float32) @@ -492,7 +516,7 @@ def sparse_attn_indexer( q_slice_cast, k_quant_cast, k_scale_cast, - weights[chunk.token_start : chunk.token_end], + weights[row_start:row_end], cu_seqlen_ks, cu_seqlen_ke, ) @@ -500,7 +524,7 @@ def sparse_attn_indexer( logits = fp8_fp4_mqa_logits( (q_slice_cast, q_scale_slice), (k_quant_cast, k_scale_cast), - weights[chunk.token_start : chunk.token_end], + weights[row_start:row_end], cu_seqlen_ks, cu_seqlen_ke, clean_logits=False, @@ -524,7 +548,25 @@ def sparse_attn_indexer( dcp_rank, dcp_world_size, cp_kv_cache_interleave_size, - row_starts=chunk.cu_seqlen_ks, + row_starts=cu_seqlen_ks, + ) + + if shard_sizes is not None: + # Every row was scored and ranked end to end by one rank, so this is + # a layout-preserving concatenation, not a top-k merge. all_gatherv + # allocates its output, so the source may alias the destination. + # ``num_prefill_tokens`` follows the padded attention metadata, + # while shard_sizes is built from the real scheduler query lengths. + # Match the destination to the unpadded all_gatherv result. + prefill_end = num_decode_tokens + sum(shard_sizes) + topk_indices_buffer[num_decode_tokens:prefill_end, :topk_tokens] = ( + get_tp_group().all_gatherv( + topk_indices_buffer[ + shard_start:shard_stop, :topk_tokens + ].contiguous(), + dim=0, + sizes=shard_sizes, + ) ) if has_decode: diff --git a/vllm/v1/attention/backends/mla/indexer.py b/vllm/v1/attention/backends/mla/indexer.py index 6198c14eb738..f4e11ed63a83 100644 --- a/vllm/v1/attention/backends/mla/indexer.py +++ b/vllm/v1/attention/backends/mla/indexer.py @@ -6,8 +6,12 @@ import torch import vllm.envs as envs -from vllm.config import VllmConfig -from vllm.distributed import get_dcp_group, get_pcp_group +from vllm.config import CUDAGraphMode, VllmConfig +from vllm.distributed import ( + get_dcp_group, + get_pcp_group, + get_tensor_model_parallel_world_size, +) from vllm.logger import init_logger from vllm.model_executor.warmup.jit_warmup_triton_helper import ( LaunchSpec, @@ -226,6 +230,88 @@ def split_indexer_prefill_chunks( return chunks +# Conservative floor for amortizing the exchange latency. An H20 TP=4 sweep +# regressed at 512 rows/rank and became net-positive at 1024; the exact +# crossover is hardware-specific. +MIN_TP_SHARD_ROWS_PER_RANK = 1024 + + +def _prefill_uses_mqa(max_seq_len: int, index_topk: int, force_mqa: bool) -> bool: + """Whether sparse MQA consumes prefill Top-K indices.""" + return max_seq_len > index_topk or force_mqa + + +def balanced_prefill_row_shard( + seq_lens_cpu: torch.Tensor, + query_lens_cpu: torch.Tensor, + compress_ratio: int, + tp_size: int, +) -> list[int] | None: + """Contiguous per-rank row counts that equalise scored keys across TP. + + Row ``j`` of a prefill request scores + ``(seq_len - query_len + 1 + j) // compress_ratio`` compressed keys, the + same value ``BuildPrefillChunkMetadataKernel`` writes as + ``cu_seqlen_ke - cu_seqlen_ks``, so the cost profile is reproducible from + CPU scheduler metadata with no device sync. + + Returns: + Per-rank row counts, or None to keep the replicated path. + """ + num_rows = int(query_lens_cpu.sum()) + if tp_size < 2 or num_rows < MIN_TP_SHARD_ROWS_PER_RANK * tp_size: + return None + + query_lens = query_lens_cpu.to(torch.int64) + first_key = torch.repeat_interleave( + seq_lens_cpu.to(torch.int64) - query_lens + 1, query_lens + ) + row_in_request = torch.arange(num_rows) - torch.repeat_interleave( + torch.cumsum(query_lens, 0) - query_lens, query_lens + ) + cost = torch.cumsum((first_key + row_in_request) // compress_ratio, 0) + total = int(cost[-1]) + if total <= 0: + return None + + targets = torch.arange(1, tp_size) * total // tp_size + bounds = [0, *torch.searchsorted(cost, targets).tolist(), num_rows] + # Force strictly increasing boundaries so every rank owns at least one row. + # The floor above guarantees there is room. A rank whose share is small is + # also cheap, so it is not on the critical path. + for i in range(1, tp_size): + bounds[i] = max(bounds[i], bounds[i - 1] + 1) + for i in range(tp_size - 1, 0, -1): + bounds[i] = min(bounds[i], bounds[i + 1] - 1) + return [bounds[i + 1] - bounds[i] for i in range(tp_size)] + + +def tp_prefill_row_sharding_supported( + vllm_config: VllmConfig, + dcp_world_size: int, + use_pcp: bool, + tp_size: int, +) -> bool: + """Whether the prefill row-shard exchange may run at all. + + ``row_shard_sizes`` doubles as the collective's consensus bit, so this may + only depend on state that is identical on every TP rank. In particular, do + not consult the lazily-mutated symmetric-memory compiler state or + per-forward capture state. + """ + cudagraph_mode = vllm_config.compilation_config.cudagraph_mode or CUDAGraphMode.NONE + return ( + current_platform.is_cuda() + and dcp_world_size == 1 + and not use_pcp + and tp_size > 1 + and not envs.VLLM_DISABLE_PYNCCL + and not envs.VLLM_USE_NCCL_SYMM_MEM + and not envs.VLLM_BATCH_INVARIANT + and cudagraph_mode.mixed_mode() != CUDAGraphMode.FULL + ) + + class DeepseekV32IndexerBackend(AttentionBackend): @classmethod def supports_pcp(cls) -> bool: @@ -516,6 +602,9 @@ def __call__( class DeepseekV32IndexerPrefillMetadata: chunks: list[DeepseekV32IndexerPrefillChunkMetadata] max_prefill_seq_len: int = -1 + # Contiguous per-TP-rank row counts for the replicated indexer prefill, or + # None to keep the replicated path. See balanced_prefill_row_shard. + row_shard_sizes: list[int] | None = None @dataclass @@ -735,6 +824,19 @@ def __init__(self, *args, block_table_width: int, **kwargs) -> None: ) self.use_fp4_indexer_cache = dsa_indexer_uses_fp4(self.vllm_config) + self.enable_tp_prefill_row_sharding = tp_prefill_row_sharding_supported( + self.vllm_config, + self.dcp_world_size, + self.use_pcp, + get_tensor_model_parallel_world_size(), + ) + if self.enable_tp_prefill_row_sharding: + logger.info_once( + "DSA indexer TP prefill row sharding enabled " + "(engages at >= %d prefill rows per rank)", + MIN_TP_SHARD_ROWS_PER_RANK, + ) + next_n = self.num_speculative_tokens + 1 self.decode_threshold = next_n self.reorder_batch_threshold = None @@ -1126,13 +1228,30 @@ def build( # Skip when total_seq_lens is 0 (i.e., no compressed token). if metadata is not None: chunks.append(metadata) + row_shard_sizes = None + # Sparse MLA skips the prefill Top-K entirely when every prefill + # request is short enough for dense MHA, so sharding those rows + # would only add an exchange. Mirrors SparseMLAAttention's + # `use_dense_mha`; both read hf_config.index_topk. + prefill_max_seq_len = int( + seq_lens_cpu[num_decodes : num_decodes + num_prefills].max() + ) + prefill_uses_mqa = _prefill_uses_mqa( + prefill_max_seq_len, + self.vllm_config.model_config.hf_config.index_topk, + self.vllm_config.attention_config.sparse_mla_force_mqa, + ) + if self.enable_tp_prefill_row_sharding and prefill_uses_mqa: + row_shard_sizes = balanced_prefill_row_shard( + seq_lens_cpu[num_decodes : num_decodes + num_prefills], + prefill_query_lens_cpu, + self.compress_ratio, + get_tensor_model_parallel_world_size(), + ) prefill_metadata = DeepseekV32IndexerPrefillMetadata( chunks, - max_prefill_seq_len=( - int(seq_lens_cpu[num_decodes:].max().item()) - if num_prefills > 0 - else 0 - ), + max_prefill_seq_len=prefill_max_seq_len, + row_shard_sizes=row_shard_sizes, ) decode_metadata = None