From e9198c57955888a920f6977fcf1299edd8b438c9 Mon Sep 17 00:00:00 2001 From: FujitsuPolycom <87842395+FujitsuPolycom@users.noreply.github.com> Date: Mon, 20 Jul 2026 12:26:59 -0500 Subject: [PATCH] feat(dcp): replicate target sparse-indexer cache --- .../layers/test_sparse_attn_indexer_b12x.py | 47 ++++- tests/models/test_dcp_shard_draft_defaults.py | 167 +++++++++++++++++- .../v1/attention/test_indexer_dcp_localize.py | 126 ++++++++++++- vllm/envs.py | 6 + .../layers/sparse_attn_indexer.py | 5 +- vllm/model_executor/models/deepseek_v2.py | 58 +++++- vllm/v1/attention/backends/mla/indexer.py | 35 +++- 7 files changed, 420 insertions(+), 24 deletions(-) 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 7f98f8979647..73b308f113f5 100644 --- a/tests/model_executor/layers/test_sparse_attn_indexer_b12x.py +++ b/tests/model_executor/layers/test_sparse_attn_indexer_b12x.py @@ -36,6 +36,44 @@ def get_simultaneous( return tensors +def test_replicated_constructor_uses_single_rank_without_dcp_group(monkeypatch): + def init_module(self): + torch.nn.Module.__init__(self) + + def fail_dcp_group(): + pytest.fail("replicated indexer construction must not query the DCP group") + + monkeypatch.setattr(indexer_mod.CustomOp, "__init__", init_module) + monkeypatch.setattr( + indexer_mod, + "get_current_vllm_config", + lambda: types.SimpleNamespace( + parallel_config=types.SimpleNamespace( + decode_context_parallel_size=4, + cp_kv_cache_interleave_size=1, + ) + ), + ) + monkeypatch.setattr(indexer_mod, "get_dcp_group", fail_dcp_group) + monkeypatch.setattr(indexer_mod, "use_b12x_sparse_indexer", lambda: True) + + indexer = indexer_mod.SparseAttnIndexer( + k_cache=torch.nn.Identity(), + quant_block_size=128, + scale_fmt="ue8m0", + topk_tokens=4, + head_dim=128, + max_model_len=4096, + max_total_seq_len=4096, + topk_indices_buffer=torch.empty((2, 4), dtype=torch.int32), + dcp_replicated=True, + ) + + assert indexer.dcp_replicated is True + assert indexer.dcp_world_size == 1 + assert indexer.dcp_rank == 0 + + def _install_fake_b12x_indexer( monkeypatch, calls: list[tuple], @@ -411,7 +449,7 @@ def test_b12x_prefill_indexer_requires_packed_contiguous_route(monkeypatch): 64 * 576, ], ) -def test_sparse_attn_indexer_decode_uses_non_shared_b12x_binding( +def test_replicated_decode_skips_dcp_merge_and_keeps_global_topk_ids( monkeypatch, page_stride0, ): @@ -437,6 +475,12 @@ def test_sparse_attn_indexer_decode_uses_non_shared_b12x_binding( lambda: torch.uint8, raising=False, ) + monkeypatch.setattr(indexer_mod, "_dcp_global_topk_requested", lambda: True) + + def fail_dcp_merge(**kwargs): + pytest.fail("replicated indexer must not all-gather top-k candidates") + + monkeypatch.setattr(indexer_mod, "_merge_b12x_dcp_topk", fail_dcp_merge) q_rows = 2 num_heads = 1 @@ -505,6 +549,7 @@ def test_sparse_attn_indexer_decode_uses_non_shared_b12x_binding( ) assert result is topk_indices_buffer + # B12X indexes the full replicated cache, so its logical IDs are already global. assert topk_indices_buffer.tolist() == [[123] * topk, [123] * topk] assert calls == [ ( diff --git a/tests/models/test_dcp_shard_draft_defaults.py b/tests/models/test_dcp_shard_draft_defaults.py index e6319599046c..143ae10ba1aa 100644 --- a/tests/models/test_dcp_shard_draft_defaults.py +++ b/tests/models/test_dcp_shard_draft_defaults.py @@ -3,16 +3,33 @@ from types import SimpleNamespace +import pytest import torch -from vllm.model_executor.models.deepseek_v2 import DeepseekV32IndexerCache +import vllm.model_executor.models.deepseek_v2 as deepseek_v2 +from vllm.model_executor.models.deepseek_v2 import ( + DeepseekV32IndexerCache, + _replicate_indexer_cache_under_dcp, +) -def _vllm_config(num_hidden_layers: int = 78): +def _vllm_config( + num_hidden_layers: int = 78, + dcp_size: int = 4, + pcp_size: int = 1, +): return SimpleNamespace( + use_v2_model_runner=True, cache_config=SimpleNamespace(block_size=256), + compilation_config=SimpleNamespace(static_forward_context={}), + attention_config=SimpleNamespace(backend="B12X_MLA_SPARSE"), model_config=SimpleNamespace( - hf_config=SimpleNamespace(num_hidden_layers=num_hidden_layers) + hf_config=SimpleNamespace(num_hidden_layers=num_hidden_layers), + max_model_len=4096, + ), + parallel_config=SimpleNamespace( + decode_context_parallel_size=dcp_size, + prefill_context_parallel_size=pcp_size, ), ) @@ -26,17 +43,157 @@ def _indexer_cache(layer_id: int = 78): return cache +def _get_indexer_spec(layer_id: int, config): + cache = _indexer_cache(layer_id) + cache.dcp_replicated = _replicate_indexer_cache_under_dcp(cache.prefix, config) + return cache, cache.get_kv_cache_spec(config) + + +def test_target_indexer_replication_defaults_off(monkeypatch): + monkeypatch.delenv("VLLM_DCP_REPLICATE_INDEXER_CACHE", raising=False) + + cache, spec = _get_indexer_spec(12, _vllm_config()) + + assert spec.dcp_replicated is False + assert spec.block_size == cache.cache_config.block_size + + def test_dcp_shard_draft_defaults_to_sharded(monkeypatch): monkeypatch.delenv("VLLM_DCP_SHARD_DRAFT", raising=False) + monkeypatch.setenv("VLLM_DCP_REPLICATE_INDEXER_CACHE", "1") - spec = _indexer_cache().get_kv_cache_spec(_vllm_config()) + cache, spec = _get_indexer_spec(78, _vllm_config()) assert spec.dcp_replicated is False + assert spec.block_size == cache.cache_config.block_size def test_dcp_shard_draft_can_restore_replicated_legacy_mode(monkeypatch): monkeypatch.setenv("VLLM_DCP_SHARD_DRAFT", "0") + monkeypatch.setenv("VLLM_DCP_REPLICATE_INDEXER_CACHE", "1") + + cache, spec = _get_indexer_spec(78, _vllm_config()) + + assert spec.dcp_replicated is True + assert spec.block_size == cache.cache_config.block_size + + +@pytest.mark.parametrize("dcp_size", [2, 4, 6, 8]) +def test_target_indexer_replication_equalizes_global_block_coverage( + monkeypatch, dcp_size: int +): + monkeypatch.setenv("VLLM_DCP_REPLICATE_INDEXER_CACHE", "1") + monkeypatch.setattr(deepseek_v2, "use_b12x_sparse_indexer", lambda: True) - spec = _indexer_cache().get_kv_cache_spec(_vllm_config()) + config = _vllm_config(dcp_size=dcp_size) + cache, spec = _get_indexer_spec(12, config) assert spec.dcp_replicated is True + assert spec.block_size == dcp_size * cache.cache_config.block_size + + +def test_target_indexer_replication_rejects_pcp(monkeypatch): + monkeypatch.setenv("VLLM_DCP_REPLICATE_INDEXER_CACHE", "1") + monkeypatch.setattr(deepseek_v2, "use_b12x_sparse_indexer", lambda: True) + + config = _vllm_config(dcp_size=4, pcp_size=2) + with pytest.raises(NotImplementedError, match="DCP2 through DCP8 with PCP1"): + _replicate_indexer_cache_under_dcp("model.layers.12.indexer.k_cache", config) + + +def test_target_indexer_replication_requires_v2_model_runner(monkeypatch): + monkeypatch.setenv("VLLM_DCP_REPLICATE_INDEXER_CACHE", "1") + + config = _vllm_config() + config.use_v2_model_runner = False + with pytest.raises(NotImplementedError, match="V2 model runner"): + _replicate_indexer_cache_under_dcp("model.layers.12.indexer.k_cache", config) + + +def test_real_cache_construction_replicates_target_but_not_draft(monkeypatch): + monkeypatch.setenv("VLLM_DCP_REPLICATE_INDEXER_CACHE", "1") + monkeypatch.delenv("VLLM_DCP_SHARD_DRAFT", raising=False) + monkeypatch.setattr(deepseek_v2, "use_b12x_sparse_indexer", lambda: True) + + config = _vllm_config() + monkeypatch.setattr(deepseek_v2, "get_current_vllm_config", lambda: config) + + target = DeepseekV32IndexerCache( + head_dim=132, + dtype=torch.uint8, + prefix="model.layers.12.self_attn.indexer.k_cache", + cache_config=config.cache_config, + ) + draft = DeepseekV32IndexerCache( + head_dim=132, + dtype=torch.uint8, + prefix="model.layers.78.self_attn.indexer.k_cache", + cache_config=config.cache_config, + ) + + assert target.dcp_replicated is True + assert target.get_kv_cache_spec(config).dcp_replicated is True + assert draft.dcp_replicated is False + assert draft.get_kv_cache_spec(config).dcp_replicated is False + assert config.compilation_config.static_forward_context == { + target.prefix: target, + draft.prefix: draft, + } + + +def test_indexer_constructor_forwards_cache_replication(monkeypatch): + captured: dict[str, object] = {} + + class FakeCache(torch.nn.Module): + def __init__(self, *args, **kwargs): + super().__init__() + self.dcp_replicated = True + + def fake_sparse_attn_indexer(*args, **kwargs): + captured["args"] = args + captured["kwargs"] = kwargs + return torch.nn.Identity() + + monkeypatch.setattr(deepseek_v2, "DeepseekV32IndexerCache", FakeCache) + monkeypatch.setattr(deepseek_v2, "SparseAttnIndexer", fake_sparse_attn_indexer) + monkeypatch.setattr( + deepseek_v2, + "ReplicatedLinear", + lambda *args, **kwargs: torch.nn.Identity(), + ) + monkeypatch.setattr( + deepseek_v2, + "MergedColumnParallelLinear", + lambda *args, **kwargs: torch.nn.Identity(), + ) + monkeypatch.setattr( + deepseek_v2, + "LayerNorm", + lambda *args, **kwargs: torch.nn.Identity(), + ) + monkeypatch.setattr(deepseek_v2.current_platform, "is_cuda", lambda: False) + + vllm_config = _vllm_config() + model_config = SimpleNamespace( + index_topk=4, + index_n_heads=1, + index_head_dim=128, + qk_rope_head_dim=64, + ) + indexer = deepseek_v2.Indexer( + vllm_config=vllm_config, + config=model_config, + hidden_size=256, + q_lora_rank=128, + quant_config=None, + cache_config=vllm_config.cache_config, + topk_indices_buffer=torch.empty((2, 4), dtype=torch.int32), + prefix="model.layers.12.self_attn.indexer", + ) + + sparse_args = captured["args"] + sparse_kwargs = captured["kwargs"] + assert isinstance(sparse_args, tuple) + assert isinstance(sparse_kwargs, dict) + assert sparse_args[0] is indexer.k_cache + assert sparse_kwargs["dcp_replicated"] is True diff --git a/tests/v1/attention/test_indexer_dcp_localize.py b/tests/v1/attention/test_indexer_dcp_localize.py index 3fafe2c4fdef..ab01ad01434d 100644 --- a/tests/v1/attention/test_indexer_dcp_localize.py +++ b/tests/v1/attention/test_indexer_dcp_localize.py @@ -1,24 +1,130 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +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_backend from vllm.platforms import current_platform from vllm.utils.import_utils import has_cutedsl -from vllm.v1.attention.backends.mla.indexer import build_prefill_chunk_metadata +from vllm.v1.attention.backend import CommonAttentionMetadata +from vllm.v1.attention.backends.mla.indexer import ( + DeepseekV32IndexerMetadataBuilder, + build_prefill_chunk_metadata, + get_indexer_max_num_blocks_per_req, +) from vllm.v1.attention.backends.mla.sparse_utils import ( triton_filter_and_convert_dcp_index, ) from vllm.v1.attention.backends.utils import get_dcp_local_seq_lens from vllm.v1.attention.ops.common import CPTritonContext, correct_attn_out +from vllm.v1.kv_cache_interface import MLAAttentionSpec def _local_count(length: int, rank: int, world: int, interleave: int) -> int: return sum(1 for pos in range(length) if (pos // interleave) % world == rank) +@pytest.mark.parametrize("dcp_size", [2, 4, 6, 8]) +def test_replicated_indexer_metadata_covers_full_context(dcp_size: int): + max_model_len = 131_071 + storage_block_size = 64 + manager_block_size = storage_block_size * dcp_size + + sharded_blocks = get_indexer_max_num_blocks_per_req( + max_model_len=max_model_len, + block_size=storage_block_size, + configured_cp_world_size=dcp_size, + dcp_replicated=False, + ) + replicated_blocks = get_indexer_max_num_blocks_per_req( + max_model_len=max_model_len, + block_size=manager_block_size, + configured_cp_world_size=dcp_size, + dcp_replicated=True, + ) + expected_blocks = (max_model_len + manager_block_size - 1) // manager_block_size + + assert sharded_blocks == expected_blocks + assert replicated_blocks == expected_blocks + + +def test_replicated_builder_uses_global_lengths_without_dcp_localization(monkeypatch): + def fail_dcp_path(*args, **kwargs): + pytest.fail("replicated indexer metadata must not enter DCP localization") + + monkeypatch.setattr(indexer_backend, "get_dcp_group", fail_dcp_path) + monkeypatch.setattr(indexer_backend, "num_compute_units", lambda device: 4) + monkeypatch.setattr(indexer_backend.envs, "VLLM_USE_B12X_SPARSE_INDEXER", True) + + spec = MLAAttentionSpec( + block_size=1024, + num_kv_heads=1, + head_size=132, + dtype=torch.uint8, + dcp_replicated=True, + ) + config = SimpleNamespace( + scheduler_config=SimpleNamespace( + max_num_batched_tokens=8, + max_num_seqs=4, + ), + parallel_config=SimpleNamespace( + decode_context_parallel_size=4, + prefill_context_parallel_size=1, + cp_kv_cache_interleave_size=1, + ), + speculative_config=None, + attention_config=SimpleNamespace(use_fp4_indexer_cache=False), + model_config=SimpleNamespace(max_model_len=4096), + ) + builder = DeepseekV32IndexerMetadataBuilder( + kv_cache_spec=spec, + layer_names=["model.layers.12.self_attn.indexer.k_cache"], + vllm_config=config, + device=torch.device("cpu"), + ) + monkeypatch.setattr(builder, "_dcp_localize_decode_seq_lens", fail_dcp_path) + monkeypatch.setattr( + builder, + "_maybe_build_b12x_schedule_metadata", + lambda *args, **kwargs: None, + ) + + global_seq_lens = torch.tensor([9], dtype=torch.int32) + local_seq_lens = torch.tensor([3], dtype=torch.int32) + common_metadata = CommonAttentionMetadata( + query_start_loc=torch.tensor([0, 1], dtype=torch.int32), + query_start_loc_cpu=torch.tensor([0, 1], dtype=torch.int32), + seq_lens=global_seq_lens, + num_reqs=1, + num_actual_tokens=1, + max_query_len=1, + max_seq_len=9, + block_table_tensor=torch.zeros((1, 1), dtype=torch.int32), + slot_mapping=torch.tensor([8], dtype=torch.int64), + dcp_local_seq_lens=local_seq_lens, + dcp_local_seq_lens_cpu=local_seq_lens.clone(), + seq_lens_cpu_upper_bound=global_seq_lens.clone(), + _seq_lens_cpu=global_seq_lens.clone(), + ) + + metadata = builder.build(0, common_metadata) + + assert builder.dcp_world_size == 1 + assert builder.dcp_rank == 0 + assert metadata.seq_lens.tolist() == [9] + assert metadata.slot_mapping.tolist() == [8] + assert metadata.decode is not None + assert metadata.decode.seq_lens.tolist() == [[9]] + assert metadata.decode.max_seq_len == 9 + assert metadata.decode.active_width.tolist() == [9] + assert metadata.decode.global_seq_lens is None + + def _global_to_local_indices( global_indices: torch.Tensor, rank: int, @@ -89,6 +195,16 @@ def _attention_from_indices( v: torch.Tensor, indices: torch.Tensor, ) -> tuple[torch.Tensor, torch.Tensor]: + if k.shape[0] == 0: + return ( + torch.zeros_like(q), + torch.full( + (q.shape[0],), + float("-inf"), + dtype=q.dtype, + device=q.device, + ), + ) valid = indices >= 0 safe_indices = indices.clamp_min(0) selected_k = k[safe_indices] @@ -253,7 +369,7 @@ def _merge_local_topks_global_with_fake_dcp( sparse_indexer.get_dcp_group = original_get_dcp_group -@pytest.mark.parametrize("world", [1, 2, 4]) +@pytest.mark.parametrize("world", [1, 2, 4, 6, 8]) @pytest.mark.parametrize("interleave", [1, 2, 4]) def test_get_dcp_local_seq_lens_matches_naive(world: int, interleave: int): seq_lens = torch.arange(0, 33, dtype=torch.int32) @@ -341,9 +457,11 @@ def test_get_dcp_local_seq_lens_must_run_after_decode_expansion(): @pytest.mark.parametrize("interleave", [1, 2]) -def test_sparse_dcp_attention_matches_global_topk_attention(interleave: int): +@pytest.mark.parametrize("world", [2, 4, 6, 8]) +def test_sparse_dcp_attention_matches_global_topk_attention( + interleave: int, world: int +): torch.manual_seed(0) - world = 2 topk = 3 num_queries = 4 max_seq_len = 13 diff --git a/vllm/envs.py b/vllm/envs.py index 770b861790bb..e3be1405db8b 100755 --- a/vllm/envs.py +++ b/vllm/envs.py @@ -73,6 +73,7 @@ VLLM_DCP_A2A_MAX_TOKENS: int = 0 VLLM_DCP_A2A_LARGE_BACKEND: Literal["ag_rs", "a2a"] = "ag_rs" VLLM_DCP_SHARD_DRAFT: str | None = None + VLLM_DCP_REPLICATE_INDEXER_CACHE: bool = False VLLM_DCP_GLOBAL_TOPK: bool = True VLLM_DCP_QUERY_SPLIT: bool = False VLLM_B12X_MLA_CKV_GATHER: bool = False @@ -1157,6 +1158,11 @@ def _resolve_rust_frontend_path() -> str | None: # target indexer cache and native MTP drafts, replicated for external # (Eagle-style) drafts. "VLLM_DCP_SHARD_DRAFT": lambda: os.getenv("VLLM_DCP_SHARD_DRAFT", None), + # Replicate the target model's sparse-indexer K cache on every DCP rank. + "VLLM_DCP_REPLICATE_INDEXER_CACHE": lambda: ( + os.getenv("VLLM_DCP_REPLICATE_INDEXER_CACHE", "0").lower() + in ("1", "true", "yes", "on") + ), # Under DCP, gather sparse-indexer logits across ranks and select a global # top-k instead of a per-rank local top-k. "VLLM_DCP_GLOBAL_TOPK": lambda: ( diff --git a/vllm/model_executor/layers/sparse_attn_indexer.py b/vllm/model_executor/layers/sparse_attn_indexer.py index b4fe13d67ad4..27c9922e91a4 100644 --- a/vllm/model_executor/layers/sparse_attn_indexer.py +++ b/vllm/model_executor/layers/sparse_attn_indexer.py @@ -2155,6 +2155,7 @@ def __init__( topk_scores_buffer: torch.Tensor | None = None, output_physical_slots: bool = False, num_q_heads: int | None = None, + dcp_replicated: bool = False, ): super().__init__() self.k_cache = k_cache @@ -2174,7 +2175,9 @@ def __init__( # during model construction) and pass them into the custom op, rather # than threading them through per-step metadata. parallel_config = get_current_vllm_config().parallel_config - self.dcp_world_size = parallel_config.decode_context_parallel_size + self.dcp_replicated = bool(dcp_replicated) + configured_dcp_world_size = parallel_config.decode_context_parallel_size + self.dcp_world_size = 1 if self.dcp_replicated else configured_dcp_world_size self.dcp_rank = get_dcp_group().rank_in_group if self.dcp_world_size > 1 else 0 self.cp_kv_cache_interleave_size = parallel_config.cp_kv_cache_interleave_size self.use_b12x_sparse_indexer = use_b12x_sparse_indexer() diff --git a/vllm/model_executor/models/deepseek_v2.py b/vllm/model_executor/models/deepseek_v2.py index 134500d172a6..a18f3ca1709a 100644 --- a/vllm/model_executor/models/deepseek_v2.py +++ b/vllm/model_executor/models/deepseek_v2.py @@ -617,6 +617,41 @@ def forward( return output +def _replicate_indexer_cache_under_dcp(prefix: str, vllm_config: VllmConfig) -> bool: + layer_id = extract_layer_index(prefix) + num_hidden_layers = getattr( + vllm_config.model_config.hf_config, "num_hidden_layers", None + ) + if layer_id is None or num_hidden_layers is None: + return False + + parallel_config = vllm_config.parallel_config + dcp_size = parallel_config.decode_context_parallel_size + pcp_size = parallel_config.prefill_context_parallel_size + if dcp_size * pcp_size <= 1: + return False + + if int(layer_id) >= int(num_hidden_layers): + return False + + requested = envs.VLLM_DCP_REPLICATE_INDEXER_CACHE + if requested and not vllm_config.use_v2_model_runner: + raise NotImplementedError( + "Replicated sparse-indexer KV requires the V2 model runner's " + "per-group context-parallel block tables." + ) + if requested and (not 2 <= dcp_size <= 8 or pcp_size != 1): + raise NotImplementedError( + "Replicated sparse-indexer KV currently supports DCP2 through " + "DCP8 with PCP1." + ) + if requested and not use_b12x_sparse_indexer(): + raise RuntimeError( + "VLLM_DCP_REPLICATE_INDEXER_CACHE requires the B12X sparse indexer." + ) + return requested + + class DeepseekV32IndexerCache(torch.nn.Module, AttentionLayerBase): def __init__( self, head_dim: int, dtype: torch.dtype, prefix: str, cache_config: CacheConfig @@ -627,30 +662,44 @@ def __init__( self.prefix = prefix self.cache_config = cache_config self.dtype = dtype - compilation_config = get_current_vllm_config().compilation_config + vllm_config = get_current_vllm_config() + self.dcp_replicated = _replicate_indexer_cache_under_dcp(prefix, vllm_config) + compilation_config = vllm_config.compilation_config if prefix in compilation_config.static_forward_context: raise ValueError(f"Duplicate layer name: {prefix}") compilation_config.static_forward_context[prefix] = self def get_kv_cache_spec(self, vllm_config: VllmConfig) -> KVCacheSpec: + block_size = self.cache_config.block_size + if self.dcp_replicated: + parallel_config = vllm_config.parallel_config + block_size *= ( + parallel_config.decode_context_parallel_size + * parallel_config.prefill_context_parallel_size + ) + logger.info_once( + "Using a DCP-replicated sparse-indexer K cache with %d-token " + "manager pages.", + block_size, + ) layer_id = extract_layer_index(self.prefix) num_hidden_layers = getattr( vllm_config.model_config.hf_config, "num_hidden_layers", None ) raw = envs.VLLM_DCP_SHARD_DRAFT shard_draft = ("1" if raw is None else raw).lower() in ("1", "true", "yes") - dcp_replicated = ( + draft_replicated = ( not shard_draft and layer_id is not None and num_hidden_layers is not None and int(layer_id) >= int(num_hidden_layers) ) return MLAAttentionSpec( # Only has one vector instead of K + V - block_size=self.cache_config.block_size, + block_size=block_size, num_kv_heads=1, head_size=self.head_dim, dtype=self.dtype, - dcp_replicated=dcp_replicated, + dcp_replicated=self.dcp_replicated or draft_replicated, ) def forward(self): ... @@ -751,6 +800,7 @@ def __init__( topk_scores_buffer=self.topk_scores_buffer, output_physical_slots=self.output_physical_slots, num_q_heads=self.n_head, + dcp_replicated=self.k_cache.dcp_replicated, ) self.is_inplace_rope = is_inplace_rope diff --git a/vllm/v1/attention/backends/mla/indexer.py b/vllm/v1/attention/backends/mla/indexer.py index f0172a10b95d..5fdcd378a1e5 100644 --- a/vllm/v1/attention/backends/mla/indexer.py +++ b/vllm/v1/attention/backends/mla/indexer.py @@ -31,7 +31,6 @@ split_decodes_and_prefills, ) from vllm.v1.kv_cache_interface import AttentionSpec, MLAAttentionSpec -from vllm.v1.worker.cp_utils import get_total_cp_world_size logger = init_logger(__name__) @@ -272,6 +271,17 @@ def get_max_prefill_buffer_size(vllm_config: VllmConfig): return max_model_len * 40 +def get_indexer_max_num_blocks_per_req( + max_model_len: int, + block_size: int, + configured_cp_world_size: int, + dcp_replicated: bool, +) -> int: + """Size indexer metadata for the cache group's effective DCP layout.""" + effective_cp_world_size = 1 if dcp_replicated else configured_cp_world_size + return cdiv(max_model_len, block_size * effective_cp_world_size) + + def _supports_varlen_paged_mqa_logits() -> bool: if ( envs.VLLM_USE_B12X_SPARSE_INDEXER @@ -347,7 +357,12 @@ def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) scheduler_config = self.vllm_config.scheduler_config parallel_config = self.vllm_config.parallel_config - self.dcp_world_size = parallel_config.decode_context_parallel_size + self.dcp_replicated = bool(getattr(self.kv_cache_spec, "dcp_replicated", False)) + configured_dcp_world_size = parallel_config.decode_context_parallel_size + configured_cp_world_size = ( + configured_dcp_world_size * parallel_config.prefill_context_parallel_size + ) + self.dcp_world_size = 1 if self.dcp_replicated else configured_dcp_world_size self.dcp_rank = get_dcp_group().rank_in_group if self.dcp_world_size > 1 else 0 self.cp_kv_cache_interleave_size = parallel_config.cp_kv_cache_interleave_size # The DCP sparse-indexer code is parameterized by interleave size, but @@ -444,9 +459,11 @@ def __init__(self, *args, **kwargs): dtype=torch.int32, device=self.device, ) - max_num_blocks_per_req = cdiv( + max_num_blocks_per_req = get_indexer_max_num_blocks_per_req( self.vllm_config.model_config.max_model_len, - self.kv_cache_spec.block_size * get_total_cp_world_size(), + self.kv_cache_spec.block_size, + configured_cp_world_size, + self.dcp_replicated, ) self.expanded_block_table_buffer = torch.zeros( ( @@ -986,7 +1003,7 @@ def build( # range and miss valid tokens. Keep the global seq_lens here and # localize the expanded bounds further down. global_seq_lens_for_decode: torch.Tensor | None = None - if dcp_local_seq_lens is not None: + if use_dcp_local_kv: global_seq_lens_for_decode = common_attn_metadata.seq_lens[:num_decodes] seq_lens = common_attn_metadata.seq_lens[:num_decodes] block_table = common_attn_metadata.block_table_tensor[:num_decodes, ...] @@ -1050,7 +1067,7 @@ def build( # Uncompressed DCP localizes after per-token expansion. Compressed # DCP must first convert global logical lengths to compressed-token # lengths and only then shard those retained tokens across ranks. - if dcp_local_seq_lens is not None and self.compress_ratio == 1: + if use_dcp_local_kv and self.compress_ratio == 1: seq_lens = self._dcp_localize_decode_seq_lens( seq_lens, num_decodes, seq_lens_is_buffer_view ) @@ -1060,7 +1077,7 @@ def build( if self.compress_ratio > 1: if seq_lens_is_buffer_view: seq_lens //= self.compress_ratio - if dcp_local_seq_lens is not None: + if use_dcp_local_kv: seq_lens.copy_( get_dcp_local_seq_lens( seq_lens, @@ -1072,7 +1089,7 @@ def build( else: # Copy to avoid mutating shared state; keeps CG address stable. compressed_decode_seq_lens = seq_lens // self.compress_ratio - if dcp_local_seq_lens is not None: + if use_dcp_local_kv: compressed_decode_seq_lens = get_dcp_local_seq_lens( compressed_decode_seq_lens, self.dcp_world_size, @@ -1123,7 +1140,7 @@ def build( use_native, ( common_attn_metadata.dcp_local_seq_lens_cpu[:num_decodes] - if dcp_local_seq_lens is not None + if use_dcp_local_kv and common_attn_metadata.dcp_local_seq_lens_cpu is not None else None ),