--- a/vllm/distributed/communication_op.py 2026-09-12 07:48:35.828568398 +0200 +++ b/vllm/distributed/communication_op.py 2026-09-11 10:13:47.000000000 +0200 @@ -9,8 +9,17 @@ from .parallel_state import get_tp_group +_DSV41_PREFILL_TIMING = __import__("os").environ.get("DSV41_PREFILL_TIMING", "0") == "1" + + def tensor_model_parallel_all_reduce(input_: torch.Tensor) -> torch.Tensor: """All-reduce the input tensor across model parallel group.""" + if _DSV41_PREFILL_TIMING and input_.dim() >= 2 and input_.shape[0] >= 1024: + # Stage 7: phase "allreduce" of the prefill budget. + from vllm.model_executor.kernels.attention.dsa.candidate_blocks import _PrefillTimer + + with _PrefillTimer("allreduce", int(input_.shape[0])): + return get_tp_group().all_reduce(input_) return get_tp_group().all_reduce(input_) --- a/vllm/model_executor/layers/fused_moe/b12x.py 2026-09-12 07:48:35.776439710 +0200 +++ b/vllm/model_executor/layers/fused_moe/b12x.py 2026-09-11 10:13:47.000000000 +0200 @@ -755,16 +755,20 @@ ) scratch = _workspace_as_b12x_scratch(workspace2, plan) - _run_b12x_moe_plan( - plan=plan, - scratch=scratch, - hidden_states=hidden_states, - prepared=prepared, - topk_weights=topk_weights, - topk_ids=topk_ids, - output=output, - unit_scale_contract=self._quant_mode == "w4a16", - ) + # Stage 7: phase "moe" of the prefill budget (inert unless DSV41_PREFILL_TIMING=1). + from vllm.model_executor.kernels.attention.dsa.candidate_blocks import _PrefillTimer + + with _PrefillTimer("moe", int(hidden_states.shape[0])): + _run_b12x_moe_plan( + plan=plan, + scratch=scratch, + hidden_states=hidden_states, + prepared=prepared, + topk_weights=topk_weights, + topk_ids=topk_ids, + output=output, + unit_scale_contract=self._quant_mode == "w4a16", + ) def moe_sum(self, input: torch.Tensor, output: torch.Tensor) -> None: raise NotImplementedError("LoRA is not supported for B12xExperts") --- a/vllm/model_executor/layers/sparse_attn_indexer.py 2026-09-12 07:48:35.313947794 +0200 +++ b/vllm/model_executor/layers/sparse_attn_indexer.py 2026-09-11 10:13:47.000000000 +0200 @@ -13,6 +13,15 @@ from vllm.forward_context import get_forward_context from vllm.logger import init_logger from vllm.model_executor.custom_op import CustomOp +from vllm.model_executor.kernels.attention.dsa.candidate_blocks import ( + apply_candidate_mask as _apply_candidate_mask, +) +from vllm.model_executor.kernels.attention.dsa.candidate_blocks import ( + _PrefillTimer as _DSV41PrefillTimer, +) +from vllm.model_executor.kernels.attention.dsa.candidate_blocks import ( + select_candidate_blocks as _select_candidate_blocks, +) from vllm.model_executor.layers.quantization.utils.quant_utils import ( get_fp8_min_max, ) @@ -316,6 +325,10 @@ dcp_world_size: int = 1, cp_kv_cache_interleave_size: int = 1, skip_topk_buffer_clear: bool = False, + candidate_blocks: torch.Tensor | None = None, + candidate_block_size: int = 0, + candidate_write: bool = False, + compress_ratio: int = 1, ) -> torch.Tensor: # careful! this will be None in dummy run forward_context = get_forward_context() @@ -323,6 +336,26 @@ fp8_dtype = current_platform.fp8_dtype() k_cache_prefix = _resolve_layer_name(k_cache_prefix) + # v4.1 two-level candidate filtering under DCP: the logits columns are this + # rank's local compressed states; with interleave 1 and world | block, local + # block b (size block // world) is exactly this rank's shard of global block + # b, so candidate ids need no translation. The producer MAX-reduces block + # scores across ranks and pins the global newest block (candidate_blocks.py). + local_cand_bs = candidate_block_size + dcp_group = None + if candidate_blocks is not None: + assert candidate_block_size > 0 + if dcp_world_size > 1: + assert cp_kv_cache_interleave_size == 1, ( + "v4.1 candidate filtering under DCP needs cp_kv_cache_interleave_size=1" + ) + assert candidate_block_size % dcp_world_size == 0, ( + f"candidate_block_size ({candidate_block_size}) must be a multiple " + f"of dcp_world_size ({dcp_world_size})" + ) + local_cand_bs = candidate_block_size // dcp_world_size + dcp_group = get_dcp_group() + # assert isinstance(attn_metadata, dict) if not isinstance(attn_metadata, dict): # Reserve workspace for indexer during profiling run @@ -361,6 +394,9 @@ use_pcp, dense_mha_metadata_layer_name, use_fp4_cache, + candidate_blocks=candidate_blocks, + candidate_block_size=candidate_block_size, + candidate_write=candidate_write, ) attn_metadata_narrowed = attn_metadata[k_cache_prefix] assert isinstance(attn_metadata_narrowed, DeepseekV32IndexerMetadata) @@ -474,6 +510,23 @@ if chunk.local_total_seq_lens == 0: logits = q_slice.new_empty((q_slice.shape[0], 0), dtype=torch.float32) topk_indices.fill_(-1) + if candidate_blocks is not None and candidate_write and dcp_group is not None: + # Other ranks may hold states for this chunk: join the MAX + # reduce with an empty local contribution. + _select_candidate_blocks( + logits, + cu_seqlen_ks, + cu_seqlen_ke, + candidate_blocks.shape[1], + local_cand_bs, + candidate_blocks[chunk.token_start : chunk.token_end], + dcp_group=dcp_group, + global_row_ke=chunk.global_ctx_lens, + global_block_size=candidate_block_size, + max_global_blocks=triton.cdiv( + chunk.max_global_ctx, candidate_block_size + ), + ) else: # DeepGEMM scalar-type tags (zero-copy): MXFP4 values → int8 # (kPackedFP4), scales → int32 squeezed to 1-D kv_sf / 2-D q_sf. @@ -506,6 +559,38 @@ clean_logits=False, ) num_rows = logits.shape[0] + if candidate_blocks is not None: + # Two-level selection (v4.1): the candidate source + # publishes its top blocks; later indexers mask their + # scores to them. Both before the row top-k. + chunk_candidates = candidate_blocks[ + chunk.token_start : chunk.token_end + ] + if candidate_write: + _select_candidate_blocks( + logits, + cu_seqlen_ks, + cu_seqlen_ke, + chunk_candidates.shape[1], + local_cand_bs, + chunk_candidates, + dcp_group=dcp_group, + global_row_ke=chunk.global_ctx_lens, + global_block_size=candidate_block_size, + max_global_blocks=( + triton.cdiv(chunk.max_global_ctx, candidate_block_size) + if dcp_group is not None + else None + ), + ) + else: + _apply_candidate_mask( + logits, + cu_seqlen_ks, + cu_seqlen_ke, + chunk_candidates, + local_cand_bs, + ) ops.top_k_per_row_prefill( logits, cu_seqlen_ks, @@ -612,6 +697,64 @@ indices=decode_metadata.indices, ) num_rows = logits.shape[0] + if candidate_blocks is not None: + # Two-level selection (v4.1) on the decode logits; columns are + # request-local compressed positions. seq_lens is (B, next_n) + # for native spec decode (per-row effective lens) and (B, 1) + # otherwise. + vis = seq_lens.reshape(-1) + row_repeat = next_n if vis.numel() != num_rows else 1 + vis = vis[:num_rows] + decode_candidates = candidate_blocks[:num_rows] + if candidate_write: + global_row_ke = None + max_global_blocks = None + if dcp_group is not None: + # Global (un-sharded) per-row context in STATES for the + # newest-block pin: per token when flattened / next_n == 1, + # per request (row_repeat) or expanded per row otherwise. + gl = decode_metadata.global_seq_lens + assert gl is not None + gl = gl.reshape(-1) + if row_repeat > 1: + gl = gl[:batch_size] + elif gl.numel() != num_rows: + # native spec decode: request lens -> per-row lens + gl = ( + gl[:batch_size].unsqueeze(1) + - next_n + + 1 + + torch.arange(next_n, device=gl.device, dtype=gl.dtype) + ).clamp_(min=0).reshape(-1) + else: + gl = gl[:num_rows] + global_row_ke = gl // compress_ratio + max_global_blocks = triton.cdiv( + attn_metadata_narrowed.max_seq_len // compress_ratio, + candidate_block_size, + ) + _select_candidate_blocks( + logits, + None, + vis, + decode_candidates.shape[1], + local_cand_bs, + decode_candidates, + row_repeat, + dcp_group=dcp_group, + global_row_ke=global_row_ke, + global_block_size=candidate_block_size, + max_global_blocks=max_global_blocks, + ) + else: + _apply_candidate_mask( + logits, + None, + vis, + decode_candidates, + local_cand_bs, + row_repeat, + ) topk_indices = topk_indices_buffer[:num_padded_tokens, :topk_tokens] use_cooperative_topk = ( @@ -622,10 +765,13 @@ and current_platform.has_device_capability(90) and not current_platform.is_device_capability_family(120) ) - use_persistent_topk = current_platform.is_cuda() and topk_tokens in ( - 512, - 1024, - 2048, + # Tech2Wild 2026-09-10 (SM12x): GB10 has 48 SMs and 99 KB smem per block. On long + # rows persistent_topk oversubscribes and its FilteredTopK fallback needs >=128 KB + # smem, so the launch fails (32K context: 90 CTAs > 48). Use top_k_per_row_decode. + use_persistent_topk = ( + current_platform.is_cuda() + and topk_tokens in (512, 1024, 2048) + and not current_platform.is_device_capability_family(120) ) if use_cooperative_topk: workspace_manager = current_workspace_manager() @@ -712,6 +858,10 @@ dcp_world_size: int = 1, cp_kv_cache_interleave_size: int = 1, skip_topk_buffer_clear: bool = False, + candidate_blocks: torch.Tensor | None = None, + candidate_block_size: int = 0, + candidate_write: bool = False, + compress_ratio: int = 1, ) -> torch.Tensor: return topk_indices_buffer @@ -719,7 +869,7 @@ direct_register_custom_op( op_name="sparse_attn_indexer", op_func=sparse_attn_indexer, - mutates_args=["topk_indices_buffer"], + mutates_args=["topk_indices_buffer", "candidate_blocks"], fake_impl=sparse_attn_indexer_fake, dispatch_key=current_platform.dispatch_key, ) @@ -751,6 +901,9 @@ skip_k_cache_insert: bool = False, use_fp4_cache: bool = False, compress_ratio: int = 1, + candidate_blocks: torch.Tensor | None = None, + candidate_block_size: int = 0, + candidate_write: bool = False, ): super().__init__() self.k_cache = k_cache @@ -764,6 +917,11 @@ self.skip_k_cache_insert = skip_k_cache_insert self.use_fp4_cache = use_fp4_cache self.compress_ratio = compress_ratio + # v4.1 two-level selection: the candidate source indexer writes the + # top candidate blocks here; later indexers mask their scores with it. + self.candidate_blocks = candidate_blocks + self.candidate_block_size = candidate_block_size + self.candidate_write = candidate_write self.dense_mha_metadata_layer_name = "" # DCP scalars are constant for the run; resolve them here (config is set # during model construction) and pass them into the custom op, rather @@ -818,6 +976,17 @@ k: torch.Tensor | None, weights: torch.Tensor, ): + # Stage 7: phase "indexer" of the prefill timing (inert unless enabled). + with _DSV41PrefillTimer("indexer", int(hidden_states.shape[0])): + return self._forward_cuda_impl(hidden_states, q_quant, k, weights) + + def _forward_cuda_impl( + self, + hidden_states: torch.Tensor, + q_quant: torch.Tensor | tuple[torch.Tensor, torch.Tensor], + k: torch.Tensor | None, + weights: torch.Tensor, + ): # FP8 path: single tensor (per-token scale is folded into `weights`). # FP4 path: (values, scales) tuple with scales required by the kernel. if isinstance(q_quant, tuple): @@ -846,6 +1015,10 @@ self.dcp_rank, self.dcp_world_size, self.cp_kv_cache_interleave_size, + candidate_blocks=self.candidate_blocks, + candidate_block_size=self.candidate_block_size, + candidate_write=self.candidate_write, + compress_ratio=self.compress_ratio, ) def forward_xpu( @@ -868,12 +1041,16 @@ assert isinstance(q_quant, torch.Tensor), ( "AMD sparse_attn_indexer expects a single FP8 q_quant tensor" ) - from vllm.platforms.rocm import on_gfx11 + from vllm.platforms.rocm import on_gfx11, on_gfx950 if ( rocm_aiter_ops.is_enabled() or rocm_aiter_ops.is_rdna_aiter_enabled() or on_gfx11() + # The so-called AITER sparse indexer op has a native gfx950 path: + # its cache insert, MQA logits, and top-k fallbacks are implemented + # by local Triton/C++ kernels and do not require the aiter package. + or on_gfx950() ): return torch.ops.vllm.rocm_aiter_sparse_attn_indexer( hidden_states, @@ -891,8 +1068,11 @@ self.topk_indices_buffer, skip_k_cache_insert=self.skip_k_cache_insert, compress_ratio=self.compress_ratio, + candidate_blocks=self.candidate_blocks, + candidate_block_size=self.candidate_block_size, + candidate_write=self.candidate_write, ) raise RuntimeError( - "Sparse attention indexer ROCm path is only supported on AITER. " - "Please enable aiter with VLLM_ROCM_USE_AITER=1" + "Sparse attention indexer ROCm path requires AITER or a supported " + "native architecture (gfx950/gfx11)." ) --- a/vllm/model_executor/model_loader/weight_utils.py 2026-09-12 07:48:35.320443782 +0200 +++ b/vllm/model_executor/model_loader/weight_utils.py 2026-09-11 10:13:47.000000000 +0200 @@ -9,6 +9,8 @@ import hashlib import json import os + +_DSV41_ENGRAM_DISK = os.environ.get("DSV41_ENGRAM_DISK", "0") == "1" import tempfile import threading import time @@ -964,6 +966,10 @@ for name in f.keys(): # noqa: SIM118 if should_skip_weight(name, local_expert_ids): continue + # Tech2Wild/Kai 2026-09-10: DSV41_ENGRAM_DISK=1 leaves the + # Engram tables on disk (see deepseek_v4_1/common/engram.py) + if _DSV41_ENGRAM_DISK and name.endswith((".engram.embed.weight", ".engram.embed.scale")): + continue param = f.get_tensor(name) yield name, param --- a/vllm/v1/attention/backend.py 2026-09-12 07:48:35.616608046 +0200 +++ b/vllm/v1/attention/backend.py 2026-09-11 10:13:47.000000000 +0200 @@ -611,6 +611,26 @@ self.vllm_config = vllm_config self.device = device self.kernel_block_size: int | None = None + # The DCP geometry this builder's group actually runs at: sharded + # groups use the process DCP size; DCP-replicated groups (every rank + # holds the full KV) behave as dcp=1 for metadata purposes. + from vllm.v1.kv_cache_interface import is_dcp_replicated_spec + + process_dcp_size = vllm_config.parallel_config.decode_context_parallel_size + self.dcp_replicated = is_dcp_replicated_spec(kv_cache_spec) + if process_dcp_size > 1 and not self.dcp_replicated: + from vllm.distributed.parallel_state import get_dcp_group + + try: + self.dcp_world_size = get_dcp_group().world_size + self.dcp_rank = get_dcp_group().rank_in_group + except AssertionError: + # DCP group not initialized (unit tests). + self.dcp_world_size = process_dcp_size + self.dcp_rank = 0 + else: + self.dcp_world_size = 1 + self.dcp_rank = 0 def set_kernel_block_size(self, kernel_block_size: int) -> None: self.kernel_block_size = kernel_block_size @@ -650,10 +670,7 @@ max_num_queries_for_spec, ) - if ( - self.vllm_config.parallel_config.decode_context_parallel_size > 1 - and not supports_dcp_with_varlen - ): + if getattr(self, "dcp_world_size", 1) > 1 and not supports_dcp_with_varlen: self.reorder_batch_threshold = 1 @abstractmethod --- a/vllm/v1/attention/backends/mla/compressor_utils.py 2026-09-12 07:48:35.555787153 +0200 +++ b/vllm/v1/attention/backends/mla/compressor_utils.py 2026-09-11 10:13:47.000000000 +0200 @@ -38,6 +38,9 @@ compress_ratio: int triton_block_size: int block_size: int + dcp_world: int = 1 + dcp_rank: int = 0 + dcp_interleave: int = 1 @staticmethod @triton.jit(do_not_specialize=["block_table_stride"]) @@ -51,10 +54,14 @@ # [num_reqs, max_num_blocks] block_table_ptr, block_table_stride, + # states per page on THIS rank (kv_cache_spec.num_states) block_size, COMPRESS_RATIO: tl.constexpr, PAD_ID: tl.constexpr, TRITON_BLOCK_SIZE: tl.constexpr, + DCP_WORLD: tl.constexpr, + DCP_RANK: tl.constexpr, + DCP_INTERLEAVE: tl.constexpr, ): batch_idx = tl.program_id(0) @@ -71,14 +78,32 @@ pos = start_pos + i + tl.arange(0, TRITON_BLOCK_SIZE) is_valid = (pos + 1) % COMPRESS_RATIO == 0 - pos_after_compress = pos // COMPRESS_RATIO + # Global compressed state index of the group this token closes. + state = pos // COMPRESS_RATIO + if DCP_WORLD > 1: + # DCP shards the compressed cache over STATES with the runner's + # token-space formula (worker/gpu/block_table.py) applied in + # state space: one block-table entry spans block_size * DCP_WORLD + # states, this rank owns the states whose interleave group is + # congruent to DCP_RANK, and stores them contiguously per entry. + virtual_block_size = block_size * DCP_WORLD + virtual_block_idx = state // virtual_block_size + virtual_off = state % virtual_block_size + is_local = (virtual_off // DCP_INTERLEAVE) % DCP_WORLD == DCP_RANK + local_off = (virtual_off // (DCP_INTERLEAVE * DCP_WORLD)) * DCP_INTERLEAVE + ( + virtual_off % DCP_INTERLEAVE + ) + local_state = virtual_block_idx * block_size + local_off + is_valid = is_valid & is_local + else: + local_state = state - block_ids = pos_after_compress // block_size + block_ids = local_state // block_size block_numbers = tl.load( block_table_ptr + batch_idx * block_table_stride + block_ids, mask=mask & is_valid, ) - slot_ids = block_numbers * block_size + pos_after_compress % block_size + slot_ids = block_numbers * block_size + local_state % block_size # NOTE slot_ids = tl.where(is_valid, slot_ids, PAD_ID) @@ -89,15 +114,29 @@ *, compress_ratio: int, block_size: int, + dcp_world_size: int = 1, + dcp_rank: int = 0, + cp_kv_cache_interleave_size: int = 1, ) -> CompileKey: return self.CompileKey( compress_ratio=compress_ratio, triton_block_size=self.TRITON_BLOCK_SIZE, block_size=triton_scalar_specialization_rep(block_size), + dcp_world=dcp_world_size, + dcp_rank=dcp_rank, + dcp_interleave=cp_kv_cache_interleave_size, ) def get_warmup_keys(self, vllm_config: Any) -> list[CompileKey]: hf_config = vllm_config.model_config.hf_config + parallel_config = vllm_config.parallel_config + dcp_world_size = parallel_config.decode_context_parallel_size + cp_interleave = parallel_config.cp_kv_cache_interleave_size + dcp_rank = 0 + if dcp_world_size > 1: + from vllm.distributed import get_dcp_group + + dcp_rank = get_dcp_group().rank_in_group configured_ratios = ( *(getattr(hf_config, "compress_ratios", None) or ()), getattr(hf_config, "index_kpool", 1) or 1, @@ -113,6 +152,9 @@ dict( compress_ratio=ratio, block_size=vllm_config.cache_config.block_size // ratio, + dcp_world_size=dcp_world_size, + dcp_rank=dcp_rank, + cp_kv_cache_interleave_size=cp_interleave, ) for ratio in compress_ratios ) @@ -128,6 +170,9 @@ block_table=int32_ptr, block_size=compile_key.block_size, compress_ratio=compile_key.compress_ratio, + dcp_world_size=compile_key.dcp_world, + dcp_rank=compile_key.dcp_rank, + cp_kv_cache_interleave_size=compile_key.dcp_interleave, ) @kernel_launcher @@ -139,12 +184,22 @@ block_table: torch.Tensor, block_size: int, compress_ratio: int, + dcp_world_size: int = 1, + dcp_rank: int = 0, + cp_kv_cache_interleave_size: int = 1, ) -> LaunchSpec: + assert block_size % cp_kv_cache_interleave_size == 0, ( + f"states per page ({block_size}) must be a multiple of " + f"cp_kv_cache_interleave_size ({cp_kv_cache_interleave_size})" + ) return (block_table.shape[0],), dict( block_table_stride=block_table.stride(0), COMPRESS_RATIO=compress_ratio, PAD_ID=-1, TRITON_BLOCK_SIZE=self.TRITON_BLOCK_SIZE, + DCP_WORLD=dcp_world_size, + DCP_RANK=dcp_rank, + DCP_INTERLEAVE=cp_kv_cache_interleave_size, ) @@ -156,7 +211,17 @@ block_size: int, compress_ratio: int, out: torch.Tensor | None = None, + dcp_world_size: int = 1, + dcp_rank: int = 0, + cp_kv_cache_interleave_size: int = 1, ) -> torch.Tensor: + """Slot of the compressed state each token closes (PAD otherwise). + + ``block_size`` is the number of states per page on this rank. Under DCP + (``dcp_world_size > 1``) states are sharded across ranks and only this + rank's states get a slot; the block table is the shared (manager) table + whose entries each span ``block_size * dcp_world_size`` states. + """ if out is not None: # Guard: for padded / invalid sequences. # Negative positions produce bogus block indices that lead to illegal memory @@ -176,6 +241,9 @@ block_table, block_size, compress_ratio, + dcp_world_size, + dcp_rank, + cp_kv_cache_interleave_size, ) return slot_mapping --- a/vllm/v1/attention/backends/mla/indexer.py 2026-09-12 07:48:35.764402198 +0200 +++ b/vllm/v1/attention/backends/mla/indexer.py 2026-09-11 10:13:47.000000000 +0200 @@ -294,7 +294,7 @@ @staticmethod def get_supported_kernel_block_sizes() -> list[int | MultipleOf]: - return [256] + return [64 if current_platform.is_device_capability_family(90) else 128] @dataclass @@ -314,6 +314,11 @@ local_cu_seq_lens: torch.Tensor | None = None local_total_seq_lens: int = 0 max_local_total_seq_lens: int = 0 + # Per query token: GLOBAL compressed context length (states), i.e. the + # un-sharded `cu_seqlen_ke - cu_seqlen_ks`. Under DCP the candidate-block + # selection pins each row's newest global block from this. + global_ctx_lens: torch.Tensor | None = None + max_global_ctx: int = 0 class BuildPrefillChunkMetadataKernel( @@ -346,6 +351,8 @@ token_to_seq_ptr, cu_compressed_seq_len_ks_ptr, cu_compressed_seq_len_ke_ptr, + # [output_query_len] global (un-sharded) compressed context per token + global_ctx_len_ptr, query_slice_start, query_slice_stop, DCP_RANK, @@ -388,6 +395,7 @@ # global per-token length is sharded across ranks. global_ctx = start_pos + 1 + offset len_per_token = global_ctx // COMPRESS_RATIO + tl.store(global_ctx_len_ptr + out_pos, len_per_token, mask=mask) if DCP_WORLD > 1: # Per-rank local context length under interleave-aware DCP, matching # get_dcp_local_seq_lens. K == 1 reduces to (len + world-1-rank)//world. @@ -478,6 +486,7 @@ token_to_seq=int32_ptr, cu_compressed_seq_len_ks=int32_ptr, cu_compressed_seq_len_ke=int32_ptr, + global_ctx_len=int32_ptr, query_slice_start=compile_key.query_slice_start, query_slice_stop=compile_key.query_slice_stop, DCP_RANK=compile_key.dcp_rank, @@ -497,6 +506,7 @@ token_to_seq: torch.Tensor, cu_compressed_seq_len_ks: torch.Tensor, cu_compressed_seq_len_ke: torch.Tensor, + global_ctx_len: torch.Tensor, query_slice_start: int, query_slice_stop: int, DCP_RANK: int, @@ -697,6 +707,9 @@ # so this is None; its own split uses self.decode_threshold. reorder_batch_threshold: int | None = None requires_block_table_width = True + # DeepSeek-V4.1 DSpark under DCP: the draft is SWA-only (no indexer); the + # target's non-causal verification rows use the DCP indexer path. + supports_non_causal_multi_token_dcp = True @classmethod def get_cudagraph_support( @@ -812,9 +825,14 @@ assert isinstance(self.kv_cache_spec.tokens_per_state, int) self.compress_ratio = self.kv_cache_spec.tokens_per_state if self.dcp_world_size > 1 and self.compress_ratio > 1: - raise NotImplementedError( - "DCP is not supported with sparse indexer KV compression " - f"(compress_ratio={self.compress_ratio})." + # DCP + compressed indexer cache (DeepSeek-V4.1 ratio-2 index + # sources): the K cache is sharded over compressed STATES; the + # compressed slot mapping and the decode bounds below are computed + # in state coordinates (DCP2-DESIGN-stage23.md §A.1). + logger.info_once( + "DSA indexer: DCP%d with compress_ratio=%d (state-sharded K cache)", + self.dcp_world_size, + self.compress_ratio, ) # Pre-allocate buffers for CUDA graph compatibility when @@ -1057,6 +1075,10 @@ and self.kv_cache_spec.block_size % kernel_block_size == 0 ): factor = self.kv_cache_spec.block_size // kernel_block_size + assert self.dcp_world_size == 1, ( + "DCP with a compressed indexer cache needs kernel_block_size == " + f"kv_cache_spec.block_size (got factor {factor})." + ) indexer_block_table = (block_table[:, ::factor] // factor).contiguous() padded_num_tokens = num_tokens if self.pcp_world_size > 1: @@ -1069,6 +1091,9 @@ self.kv_cache_spec.num_states, self.compress_ratio, out=self.compressed_slot_mapping_buffer, + dcp_world_size=self.dcp_world_size, + dcp_rank=self.dcp_rank, + cp_kv_cache_interleave_size=self.cp_kv_cache_interleave_size, ) if self.pcp_world_size > 1: compressed_slot_mapping = get_pcp_group().all_gather( @@ -1231,20 +1256,16 @@ ) block_table = self.indexer_decode_block_table_buffer[:rows, :cols] - seq_lens_is_buffer_view = (use_native and next_n > 1) or ( - not use_native and max_decode_len > 1 - ) - - # DCP: localize the now-expanded per-token global bounds to this - # rank's owned KV. Done here (after expansion) so each token's global - # causal length is localized individually; see the comment above. - if dcp_local_seq_lens is not None: - seq_lens = self._dcp_localize_decode_seq_lens( - seq_lens, num_decodes, seq_lens_is_buffer_view - ) + # Flattening always returns a buffer view, including single-token + # batches. Keep its address stable across varlen graph replays. + seq_lens_is_buffer_view = not use_native or next_n > 1 # For DeepseekV4 (compress_ratio > 1), the indexer KV cache stores # compressed tokens. Convert uncompressed seq_lens to compressed. + # This must happen BEFORE the DCP localization below: the cache is + # sharded over compressed states, so this rank's share is + # localize(len // R), not localize(len) // R (len=10, R=2, W=2: + # rank 0 owns states {0, 2, 4} -> 3, but 5 // 2 == 2). if self.compress_ratio > 1: if seq_lens_is_buffer_view: seq_lens //= self.compress_ratio @@ -1255,6 +1276,17 @@ ) self.expanded_seq_lens_buffer[num_decodes:num_decode_tokens] = 0 seq_lens = self.expanded_seq_lens_buffer[:num_decode_tokens] + # The copy is private: localize it in place below. + seq_lens_is_buffer_view = True + + # DCP: localize the now-expanded (and now-compressed) per-token + # global bounds to this rank's owned KV. Done here (after expansion) + # so each token's global causal length is localized individually; + # see the comment above. + if dcp_local_seq_lens is not None: + seq_lens = self._dcp_localize_decode_seq_lens( + seq_lens, num_decodes, seq_lens_is_buffer_view + ) # Non-MTP: deep_gemm paged MQA logits requires 2D context_lens # (csrc/apis/attention.hpp). Unsqueeze to (B, 1) so downstream @@ -1367,6 +1399,7 @@ cu_seq_len_ks = torch.empty(output_query_len, dtype=torch.int32, device=device) cu_seq_len_ke = torch.empty(output_query_len, dtype=torch.int32, device=device) + global_ctx_len = torch.empty(output_query_len, dtype=torch.int32, device=device) # Under DCP the kernel writes this rank's local row bounds into # cu_seq_len_ks/ke; otherwise local_cu_seq_lens aliases cu_seq_lens. @@ -1378,6 +1411,7 @@ token_to_seq, cu_seq_len_ks, cu_seq_len_ke, + global_ctx_len, qs_start, qs_stop, dcp_rank, @@ -1409,6 +1443,12 @@ local_cu_seq_lens=local_cu_seq_lens, local_total_seq_lens=local_total_seq_lens, max_local_total_seq_lens=max_local_total_seq_lens, + global_ctx_lens=global_ctx_len if dcp_world_size > 1 else None, + max_global_ctx=( + int(compressed_seq_lens_cpu[start_idx:end_idx].max().item()) + if dcp_world_size > 1 + else 0 + ), ) --- a/vllm/v1/attention/backends/mla/sparse_swa.py 2026-09-12 07:48:35.334168175 +0200 +++ b/vllm/v1/attention/backends/mla/sparse_swa.py 2026-09-11 10:13:47.000000000 +0200 @@ -46,18 +46,26 @@ _LAYER_TYPE_SWAONLY = "swaonly" _LAYER_TYPE_C4A = "c4a" _LAYER_TYPE_C128A = "c128a" +# v4.1 ratio-1 / ratio-2 indexer layers: same indexer-path shape as C4A, but +# their compressed page block is block_size // ratio, so each gets its own plan. +# v4.1 builders classify layer types themselves (ratio 0 is SWA-only there); see +# deepseek_v4_1/sparse_mla.py. +_LAYER_TYPE_C1A = "c1a" +_LAYER_TYPE_C2A = "c2a" def _layer_type_for(compress_ratio: int) -> str: if compress_ratio <= 1: return _LAYER_TYPE_SWAONLY + if compress_ratio == 2: + return _LAYER_TYPE_C2A if compress_ratio == 4: return _LAYER_TYPE_C4A if compress_ratio == 128: return _LAYER_TYPE_C128A raise ValueError( f"Unsupported DeepseekV4 compress_ratio={compress_ratio}; " - "expected 1, 4, or 128." + "expected 1, 2, 4, or 128." ) @@ -70,6 +78,7 @@ prefix: str, cache_config: CacheConfig, backend_cls: "type[AttentionBackend] | None" = None, + block_size: int = 64, ): super().__init__() self.backend_cls = backend_cls or DeepseekSparseSWABackend @@ -84,12 +93,9 @@ raise ValueError(f"Duplicate layer name: {prefix}") compilation_config.static_forward_context[prefix] = self - # Block size is constrained by tensor sharing between SWA and C4A KV blocks. - # Since both block types share the same physical tensor, they must use the - # same page size. The C4A KV block shape [256//4, head_dim] = [64, head_dim] - # determines the SWA block size of 64 tokens per block. - # TODO(yifan): make SWA block size automatically determined and configurable. - self.block_size = 64 + # Any multiple of 32; the sparse decode kernels take the page size at + # runtime. + self.block_size = block_size # uint8: fp8_ds_mla UE8M0 paged layout. bfloat16 / float8_e4m3fn: # contiguous full-cache layout. assert self.dtype in (torch.uint8, torch.bfloat16, torch.float8_e4m3fn) @@ -115,6 +121,9 @@ alignment=576 if uses_fp8_ds_mla_layout else 512, model_version="deepseek_v4", kv_quant_mode=get_kv_quant_mode(self.cache_config.cache_dtype), + # DCP: the 128-token window is tiny; every rank keeps all of it + # (no sharding), so SWA needs no cross-rank combine. + dcp_replicated=True, ) def forward(self): ... @@ -130,7 +139,14 @@ @staticmethod def get_supported_kernel_block_sizes() -> list[int | MultipleOf]: - return [MultipleOf(64)] + return [MultipleOf(32)] + + @classmethod + def get_swa_block_size(cls) -> int: + """Tech2Wild/Kai 2026-09-10: tokens per SWA cache page handed to + DeepseekV4SWACache. FlashMLA / TRTLLM-gen take the page size at + runtime (32); kernels instantiated for one page size override this.""" + return 32 @classmethod def get_preferred_block_size(cls, default_block_size: int) -> int: @@ -207,12 +223,17 @@ tile_sched_swaonly: "FlashMLASchedMeta | None" = None tile_sched_c4a: "FlashMLASchedMeta | None" = None tile_sched_c128a: "FlashMLASchedMeta | None" = None + tile_sched_c1a: "FlashMLASchedMeta | None" = None + tile_sched_c2a: "FlashMLASchedMeta | None" = None flashinfer_sparse_index_cache: dict[str, tuple[torch.Tensor, torch.Tensor]] = field( default_factory=dict ) def get_prefill_chunk_plan( - self, compress_ratio: int, prefill_chunk_size: int + self, + compress_ratio: int, + prefill_chunk_size: int, + has_compressed: bool | None = None, ) -> list[tuple[int, int, int, int]]: if self.num_prefills == 0: return [] @@ -220,17 +241,20 @@ assert self.prefill_seq_lens_cpu is not None assert self.prefill_query_lens_cpu is not None + # Whether the layer gathers a compressed-KV region into the prefill + # workspace. v4.0 callers keep the legacy default (ratio <= 1 means + # SWA-only); v4.1 passes has_compressed explicitly because its + # compress_ratio==1 layers DO have a full-length compressed cache. + if has_compressed is None: + has_compressed = compress_ratio > 1 + # query_len <= max_num_batched_tokens and # gather_len = query_len + min(prefix_len, window_size - 1), so the # worst-case gathered width is bounded by # max_num_batched_tokens + window_size - 1. The compressed prefix pool # is bounded by ceil(max_model_len / compress_ratio). max_workspace_area = prefill_chunk_size * ( - ( - 0 - if compress_ratio <= 1 - else cdiv(self.prefill_max_model_len, compress_ratio) - ) + (cdiv(self.prefill_max_model_len, compress_ratio) if has_compressed else 0) + self.prefill_window_size + self.prefill_max_num_batched_tokens ) @@ -239,13 +263,13 @@ prefix_lens_cpu, min=0, max=self.prefill_window_size - 1 ) compressed_lens_cpu = ( - torch.zeros_like(self.prefill_seq_lens_cpu) - if compress_ratio <= 1 - else torch.div( + torch.div( self.prefill_seq_lens_cpu, compress_ratio, rounding_mode="floor", ) + if has_compressed + else torch.zeros_like(self.prefill_seq_lens_cpu) ) chunk_plan: list[tuple[int, int, int, int]] = [] @@ -390,6 +414,9 @@ reorder_batch_threshold: int | None = None _cudagraph_support: ClassVar[AttentionCGSupport] = AttentionCGSupport.UNIFORM_BATCH supports_draft_decode_metadata_update = True + # DSpark's non-causal multi-token decode only touches the replicated SWA + # caches (no cross-rank combine), so it is supported under DCP. + supports_non_causal_multi_token_dcp: ClassVar[bool] = True def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) @@ -709,6 +736,8 @@ tile_sched_swaonly=tile_sched[_LAYER_TYPE_SWAONLY], tile_sched_c4a=tile_sched[_LAYER_TYPE_C4A], tile_sched_c128a=tile_sched[_LAYER_TYPE_C128A], + tile_sched_c1a=tile_sched[_LAYER_TYPE_C1A], + tile_sched_c2a=tile_sched[_LAYER_TYPE_C2A], **deepseek_v4_fields, # type: ignore[arg-type] ) @@ -789,6 +818,8 @@ metadata.tile_sched_swaonly = tile_sched[_LAYER_TYPE_SWAONLY] metadata.tile_sched_c4a = tile_sched[_LAYER_TYPE_C4A] metadata.tile_sched_c128a = tile_sched[_LAYER_TYPE_C128A] + metadata.tile_sched_c1a = tile_sched[_LAYER_TYPE_C1A] + metadata.tile_sched_c2a = tile_sched[_LAYER_TYPE_C2A] metadata.flashinfer_sparse_index_cache.clear() def build_tile_scheduler( @@ -809,6 +840,8 @@ _LAYER_TYPE_SWAONLY: None, _LAYER_TYPE_C4A: None, _LAYER_TYPE_C128A: None, + _LAYER_TYPE_C1A: None, + _LAYER_TYPE_C2A: None, } if ( num_decode_tokens == 0 --- a/vllm/v1/core/kv_cache_coordinator.py 2026-09-12 07:48:35.835361219 +0200 +++ b/vllm/v1/core/kv_cache_coordinator.py 2026-09-11 10:13:47.000000000 +0200 @@ -24,6 +24,9 @@ KVCacheSpec, MambaSpec, SlidingWindowSpec, + is_dcp_replicated_spec, + is_dcp_sharded_spec, + iter_layer_specs, ) from vllm.v1.request import Request @@ -619,13 +622,19 @@ assert pcp_world_size == 1, "PCP not support hybrid attn now." if dcp_world_size > 1: # DCP shards full-attention KV across ranks and replicates Mamba - # state; other spec types (e.g. sliding window) have no DCP-aware - # handling yet, so reject them explicitly. + # state and attention groups marked ``dcp_replicated`` (their + # managers run at dcp=1, see dcp_world_size_for_kv_cache_spec); + # any other spec type has no DCP-aware handling, so reject it. for g in kv_cache_config.kv_cache_groups: - assert isinstance(g.kv_cache_spec, (FullAttentionSpec, MambaSpec)), ( - "DCP with hybrid KV cache layouts only supports " - "full-attention and Mamba groups, got: " - f"{type(g.kv_cache_spec).__name__}." + spec = g.kv_cache_spec + assert ( + is_dcp_sharded_spec(spec) + or is_dcp_replicated_spec(spec) + or all(isinstance(s, MambaSpec) for s in iter_layer_specs(spec)) + ), ( + "DCP with hybrid KV cache layouts only supports sharded " + "full-attention, dcp_replicated attention and Mamba groups, " + f"got: {type(spec).__name__}." ) # Fine-grained hash hits require Mamba "align" and compatible cache # managers in every group. TP needs hashing finer than the Mamba block; --- a/vllm/v1/core/kv_cache_utils.py 2026-09-12 07:48:35.531214294 +0200 +++ b/vllm/v1/core/kv_cache_utils.py 2026-09-11 10:13:47.000000000 +0200 @@ -22,6 +22,7 @@ from vllm.v1.kv_cache_interface import ( AttentionSpec, ChunkedLocalAttentionSpec, + CircularBufferSpec, FullAttentionSpec, HiddenStateCacheSpec, KpoolTailSpec, @@ -36,6 +37,7 @@ SlidingWindowSpec, UniformTypeKVCacheSpecs, compute_layout_strides, + is_dcp_sharded_spec, iter_layer_specs, replace_as, ) @@ -650,13 +652,15 @@ def resolve_dcp_kv_block_size(spec: KVCacheSpec, dcp_world_size: int) -> int: - """Return the token span of a cache block under DCP.""" - layer_specs = iter_layer_specs(spec) - if len(layer_specs) > 0 and all( - isinstance(layer_spec, AttentionSpec) for layer_spec in layer_specs - ): - return spec.block_size * dcp_world_size - return spec.block_size + """Return the token span of a cache block under DCP. + + Only DCP-sharded groups span ``dcp_world_size`` blocks of tokens per + manager block; replicated attention groups (``dcp_replicated``), Mamba and + other per-rank state keep their plain block size. Must agree with + ``dcp_world_size_for_kv_cache_spec`` (the coordinator scales manager block + sizes with that helper). + """ + return spec.block_size * dcp_world_size_for_kv_cache_spec(spec, dcp_world_size) def resolve_dcp_kv_cache_spec(spec: KVCacheSpec, dcp_world_size: int) -> KVCacheSpec: @@ -685,15 +689,13 @@ must keep ``dcp_world_size=1`` even when the process runs with DCP > 1. Draft MLA groups on the sharded DSpark path are ``FullAttentionSpec`` / - ``MLAAttentionSpec`` and therefore keep the process DCP size. A replicated - draft group would need a different spec, not this helper. + ``MLAAttentionSpec`` and therefore keep the process DCP size. Attention + groups whose layers set ``dcp_replicated`` (DeepSeek-V4.1 SWA cache and + compressor ring) hold the full KV on every rank and keep ``1``. """ if dcp_world_size <= 1: return 1 - inner = spec - if isinstance(spec, UniformTypeKVCacheSpecs): - inner = next(iter(spec.kv_cache_specs.values())) - if isinstance(inner, FullAttentionSpec): + if is_dcp_sharded_spec(spec): return dcp_world_size return 1 @@ -2010,15 +2012,22 @@ cdiv(len(names), n) * page for page, names in page_size_layers.items() ) - # Bytes a block must hold however the mamba buckets end up split: a mamba - # bucket can go down to one state per group, every other bucket's split is - # already fixed by the repeat pattern. + # Bytes a block must hold however the state buckets end up split: mamba, + # circular-buffer and unsplit sliding-window buckets can go down to one + # state per group, every other bucket's split is fixed by the repeat pattern. + def is_state_bucket(spec: UniformTypeKVCacheSpecs) -> bool: + if isinstance(spec.first_spec, (MambaSpec, CircularBufferSpec)): + return True + return repeats_per_group is None and isinstance( + spec.first_spec, SlidingWindowSpec + ) + anchor_bytes = max( ( widest_group_bytes( page_size_layers, len(spec.kv_cache_specs) - if isinstance(spec.first_spec, MambaSpec) + if is_state_bucket(spec) else num_groups_for(spec, balanced), ) for spec, page_size_layers, balanced in bucketed @@ -2029,10 +2038,9 @@ groups = [] for spec, page_size_layers, balanced in bucketed: num_groups = num_groups_for(spec, balanced) - # `_align_hybrid_block_size` pads a mamba state up to one attention - # page, so cap a mamba group at the states a block already fits rather - # than let it widen the block. - if anchor_bytes and isinstance(spec.first_spec, MambaSpec): + # Cap a state group at the states a block already fits rather than let + # it widen the block. + if anchor_bytes and is_state_bucket(spec): states_per_block = max(anchor_bytes // spec.first_spec.page_size_bytes, 1) num_groups = max( num_groups, cdiv(len(spec.kv_cache_specs), states_per_block) @@ -2068,8 +2076,9 @@ if spec_config is None or not spec_config.use_eagle(): return False model_config = vllm_config.model_config - return ( - model_config is not None and model_config.hf_config.model_type == "deepseek_v4" + return model_config is not None and model_config.hf_config.model_type in ( + "deepseek_v4", + "deepseek_v41", ) @@ -2090,9 +2099,9 @@ spec merging, wherever grouping happens to land. It is sufficient but not necessary: a drafter whose spec is indistinguishable from the target's cannot be found this way. - 2. Model-scoped positional fallback for DeepseekV4, whose MTP block reuses - the target's own decoder layer and so carries no spec marker. Its draft - attention layer is always the last registered layer, so flag whichever + 2. Model-scoped positional fallback for DeepseekV4/V4.1, whose MTP block + reuses the target's own decoder layer and so carries no spec marker. Its + draft attention layer is always the last registered layer, so flag whichever group holds it. This rule is only valid where the groups partition exactly the layers of ``kv_cache_spec``, which is true on the packed grouping path and not in general; other callers must leave @@ -2105,7 +2114,8 @@ kv_cache_spec: The kv cache spec of each attention layer, in layer registration order. Only read by rule 2. kv_cache_groups: Groups to annotate in place. - use_deepseek_v4_fallback: Enable rule 2 for a DeepseekV4 packed group. + use_deepseek_v4_fallback: Enable rule 2 for a DeepseekV4/V4.1 packed + group. """ spec_config = vllm_config.speculative_config if spec_config is None or not spec_config.use_eagle_block_drop(): --- a/vllm/v1/kv_cache_interface.py 2026-09-12 07:48:35.351087500 +0200 +++ b/vllm/v1/kv_cache_interface.py 2026-09-11 10:13:47.000000000 +0200 @@ -240,8 +240,11 @@ f"Unsupported KV cache spec type: {type(self)}. " "Please register it using @register_kv_cache_spec decorator." ) + dcp_replicated = getattr(self, "dcp_replicated", None) return all( - isinstance(spec, uniform_type_base_spec) for spec in kv_cache_specs.values() + isinstance(spec, uniform_type_base_spec) + and getattr(spec, "dcp_replicated", None) == dcp_replicated + for spec in kv_cache_specs.values() ) @@ -400,6 +403,13 @@ """Tokens covered by one stored state. Ints > 1 compress multiple tokens into one state (DSv4 sparse MLA); fractions < 1 store multiple states per token (Whisper block pooling: ``Fraction(1, block_pool_size)``).""" + dcp_replicated: bool = False + """DCP: every rank holds this layer's full KV instead of a 1/dcp shard. + A replicated group keeps its plain block geometry under DCP (no block-size + scaling, no per-rank memory divide, block table sized for the whole + context) and every rank writes every position (DeepSeek-V4.1 SWA cache and + compressor ring). Sharded groups (the default) round-robin positions across + ranks at ``cp_kv_cache_interleave_size`` granularity.""" def __post_init__(self): if self.head_size_v is None: @@ -438,7 +448,9 @@ def max_num_blocks_per_req(self, vllm_config: VllmConfig, max_len: int) -> int: parallel_config = vllm_config.parallel_config - kv_shard_count = parallel_config.decode_context_parallel_size + kv_shard_count = ( + 1 if self.dcp_replicated else parallel_config.decode_context_parallel_size + ) return cdiv(max_len, self.block_size * kv_shard_count) @@ -471,7 +483,7 @@ def max_memory_usage_bytes(self, vllm_config: VllmConfig) -> int: max_model_len = vllm_config.model_config.max_model_len dcp_world_size = vllm_config.parallel_config.decode_context_parallel_size - if dcp_world_size > 1: + if dcp_world_size > 1 and not self.dcp_replicated: max_model_len = cdiv(max_model_len, dcp_world_size) return cdiv(max_model_len, self.block_size) * self.page_size_bytes @@ -519,6 +531,7 @@ num_head_slots=specs[0].num_head_slots, state_content_bytes=specs[0].state_content_bytes, tokens_per_state=specs[0].tokens_per_state, + dcp_replicated=specs[0].dcp_replicated, sliding_window=cls.merge_window_sizes(sliding_window), attention_chunk_size=cls.merge_window_sizes(attention_chunk_size), # If any layer in the group is non-causal, treat the group as @@ -602,6 +615,7 @@ state_content_bytes=specs[0].state_content_bytes, cache_dtype_str=cache_dtype_str_set.pop(), tokens_per_state=tokens_per_state_set.pop(), + dcp_replicated=specs[0].dcp_replicated, model_version=model_version_set.pop(), storage_block_size=storage_block_size_set.pop(), non_causal_multi_token_decode=non_causal_mtd_set.pop(), @@ -657,6 +671,7 @@ num_head_slots=base.num_head_slots, state_content_bytes=base.state_content_bytes, tokens_per_state=base.tokens_per_state, + dcp_replicated=base.dcp_replicated, sliding_window=base.sliding_window, attention_chunk_size=base.attention_chunk_size, non_causal=base.non_causal, @@ -700,6 +715,7 @@ return all( isinstance(spec, ChunkedLocalAttentionSpec) and spec.attention_chunk_size == self.attention_chunk_size + and spec.dcp_replicated == self.dcp_replicated for spec in kv_cache_specs.values() ) @@ -743,9 +759,10 @@ return cdiv(num_tokens, self.block_size) + 1 def max_memory_usage_bytes(self, vllm_config: VllmConfig) -> int: - assert vllm_config.parallel_config.decode_context_parallel_size == 1, ( - "DCP not support sliding window." - ) + assert ( + vllm_config.parallel_config.decode_context_parallel_size == 1 + or self.dcp_replicated + ), "DCP not support sliding window (sharded); mark the spec dcp_replicated." max_blocks = self.max_admission_blocks_per_request( max_in_flight_tokens=vllm_config.max_in_flight_tokens, max_model_len=vllm_config.model_config.max_model_len, @@ -758,6 +775,7 @@ return all( isinstance(spec, SlidingWindowSpec) and spec.sliding_window == self.sliding_window + and spec.dcp_replicated == self.dcp_replicated for spec in kv_cache_specs.values() ) @@ -786,7 +804,9 @@ self, kv_cache_specs: dict[str, KVCacheSpec] ) -> bool: return all( - isinstance(spec, CircularBufferSpec) for spec in kv_cache_specs.values() + isinstance(spec, CircularBufferSpec) + and spec.dcp_replicated == self.dcp_replicated + for spec in kv_cache_specs.values() ) @property @@ -847,6 +867,7 @@ extra_retained_tokens=extra_retained_set.pop(), cache_dtype_str=cache_dtype_str_set.pop(), tokens_per_state=tokens_per_state_set.pop(), + dcp_replicated=specs[0].dcp_replicated, model_version=model_version_set.pop(), ) @@ -856,6 +877,7 @@ return all( isinstance(spec, SlidingWindowMLASpec) and spec.sliding_window == self.sliding_window + and spec.dcp_replicated == self.dcp_replicated for spec in kv_cache_specs.values() ) @@ -1013,6 +1035,7 @@ page_size_padded=specs[0].page_size_padded, num_head_slots=specs[0].num_head_slots, state_content_bytes=specs[0].state_content_bytes, + dcp_replicated=specs[0].dcp_replicated, sliding_window=cls.merge_window_sizes(sliding_window), attention_chunk_size=cls.merge_window_sizes(attention_chunk_size), non_causal=any(spec.non_causal for spec in specs), @@ -1147,6 +1170,28 @@ ) +def is_dcp_replicated_spec(kv_cache_spec: KVCacheSpec) -> bool: + """Whether a KV cache group spec is (or wraps) DCP-replicated attention + layers: every DCP rank holds the full KV, so the group keeps its plain + block geometry and every rank writes every position.""" + layer_specs = iter_layer_specs(kv_cache_spec) + return len(layer_specs) > 0 and all( + isinstance(spec, AttentionSpec) and spec.dcp_replicated + for spec in layer_specs + ) + + +def is_dcp_sharded_spec(kv_cache_spec: KVCacheSpec) -> bool: + """Whether a KV cache group spec is (or wraps) DCP-sharded full attention + (incl. MLA): positions are round-robined across DCP ranks, so manager block + size, prefix hashing and per-rank memory scale by the DCP size.""" + layer_specs = iter_layer_specs(kv_cache_spec) + return len(layer_specs) > 0 and all( + isinstance(spec, FullAttentionSpec) and not spec.dcp_replicated + for spec in layer_specs + ) + + def get_kv_cache_spec_kind(kv_cache_spec: KVCacheSpec) -> KVCacheSpecKind: if isinstance(kv_cache_spec, UniformTypeKVCacheSpecs): inner_kinds = { --- a/vllm/v1/worker/cp_utils.py 2026-09-12 07:48:35.759086508 +0200 +++ b/vllm/v1/worker/cp_utils.py 2026-09-11 10:13:47.000000000 +0200 @@ -35,6 +35,10 @@ f"but {backend.get_name()} does not support PCP." ) layer_impl = getattr(layer, "impl", None) + if layer_impl is None and hasattr(layer, "need_to_return_lse_for_decode"): + # DeepSeek-V4.1 attention runs its own forward (no + # AttentionImpl) and exposes the DCP contract on the layer. + layer_impl = layer if layer_impl is None: continue if vllm_config.speculative_config is not None and interleave_size > 1: --- a/vllm/v1/worker/gpu/attn_utils.py 2026-09-12 07:48:35.523821429 +0200 +++ b/vllm/v1/worker/gpu/attn_utils.py 2026-09-11 10:13:47.000000000 +0200 @@ -20,6 +20,7 @@ KVCacheConfig, KVCacheSpec, UniformTypeKVCacheSpecs, + is_dcp_replicated_spec, ) from vllm.v1.worker.gpu.model_states.interface import ModelSpecificAttnMetadata from vllm.v1.worker.utils import ( @@ -283,6 +284,13 @@ for i in range(num_kv_cache_groups): block_table = block_tables[i] slot_mapping = slot_mappings[i] + # DCP-replicated groups hold every position on every rank: their + # builders must see the full seq_lens, not this rank's shard. + group_dcp_local_seq_lens = ( + None + if is_dcp_replicated_spec(kv_cache_config.kv_cache_groups[i].kv_cache_spec) + else dcp_local_seq_lens + ) # Per-group causal for hybrid drafters (mixed SWA/full attention). group_causal = ( causal if isinstance(causal, (bool, torch.Tensor)) else causal.get(i, True) @@ -310,7 +318,7 @@ block_table_tensor=block_table, slot_mapping=slot_mapping, causal=group_causal, - dcp_local_seq_lens=dcp_local_seq_lens, + dcp_local_seq_lens=group_dcp_local_seq_lens, positions=positions, is_prefilling=group_is_prefilling, mm_req_doc_ranges=mm_req_doc_ranges, --- a/vllm/v1/worker/gpu/block_table.py 2026-09-12 07:48:35.514464887 +0200 +++ b/vllm/v1/worker/gpu/block_table.py 2026-09-11 10:13:47.000000000 +0200 @@ -27,6 +27,7 @@ cp_rank: int = 0, cp_interleave: int = 1, slot_mapping_enabled: list[bool] | None = None, + cp_sharded: list[bool] | None = None, ): self.block_sizes = block_sizes self.kernel_block_sizes = kernel_block_sizes @@ -44,6 +45,13 @@ slot_mapping_enabled = [True] * self.num_kv_cache_groups assert len(slot_mapping_enabled) == self.num_kv_cache_groups self._slot_mapping_enabled = slot_mapping_enabled + # Per group: True = positions are round-robined across the cp_size + # ranks (this rank stores 1/cp_size of them); False = DCP-replicated + # group, every rank stores every position at its plain block geometry. + if cp_sharded is None: + cp_sharded = [True] * self.num_kv_cache_groups + assert len(cp_sharded) == self.num_kv_cache_groups + self._cp_sharded = cp_sharded self.blocks_per_kv_block = [ bs // kbs for bs, kbs in zip(block_sizes, kernel_block_sizes) @@ -108,6 +116,9 @@ self.slot_mapping_enabled = torch.tensor( self._slot_mapping_enabled, dtype=torch.bool, device=self.device ) + self.cp_sharded = torch.tensor( + self._cp_sharded, dtype=torch.bool, device=self.device + ) self.input_block_table_ptrs = self._make_ptr_tensor(self.input_block_tables) def append_block_ids( @@ -211,6 +222,7 @@ self.block_sizes_tensor, self.kernel_block_sizes_tensor, self.slot_mapping_enabled, + self.cp_sharded, slot_mappings, slot_mappings.stride(0), self.cp_rank, @@ -284,6 +296,7 @@ block_sizes, # [num_kv_cache_groups] kernel_block_sizes, # [num_kv_cache_groups] slot_mapping_enabled, # [num_kv_cache_groups] + cp_sharded, # [num_kv_cache_groups] slot_mappings_ptr, # [num_kv_cache_groups, max_num_tokens] slot_mappings_stride, cp_rank, @@ -313,6 +326,7 @@ kv_block_size = tl.load(block_sizes + group_id) kernel_block_size = tl.load(kernel_block_sizes + group_id) mapping_enabled = tl.load(slot_mapping_enabled + group_id) + group_sharded = tl.load(cp_sharded + group_id) req_state_idx = tl.load(idx_mapping + batch_idx) start_idx = tl.load(query_start_loc + batch_idx) @@ -326,15 +340,19 @@ local_positions = positions is_local = True else: - # Context parallelism is used. + # Context parallelism is used. Sharded groups round-robin positions + # across the CP ranks; DCP-replicated groups (cp_sharded False) + # keep every position on every rank at the plain block geometry. virtual_block_size = kv_block_size * CP_SIZE virtual_block_indices = positions // virtual_block_size virtual_block_offsets = positions % virtual_block_size - is_local = virtual_block_offsets // CP_INTERLEAVE % CP_SIZE == cp_rank + owner_is_me = virtual_block_offsets // CP_INTERLEAVE % CP_SIZE == cp_rank rounds = virtual_block_offsets // (CP_INTERLEAVE * CP_SIZE) remainder = virtual_block_offsets % CP_INTERLEAVE local_offsets = rounds * CP_INTERLEAVE + remainder - local_positions = virtual_block_indices * kv_block_size + local_offsets + sharded_positions = virtual_block_indices * kv_block_size + local_offsets + local_positions = tl.where(group_sharded, sharded_positions, positions) + is_local = tl.where(group_sharded, owner_is_me, True) block_indices = tl.where( mapping_enabled, local_positions // kernel_block_size, 0 --- a/vllm/v1/worker/gpu/model_runner.py 2026-09-12 07:48:35.558729666 +0200 +++ b/vllm/v1/worker/gpu/model_runner.py 2026-09-11 10:13:47.000000000 +0200 @@ -70,6 +70,7 @@ KVCacheConfig, MambaSpec, UniformTypeKVCacheSpecs, + is_dcp_replicated_spec, ) from vllm.v1.outputs import ( DraftTokenIds, @@ -170,6 +171,82 @@ logger = init_logger(__name__) +# Stage 7 (DSV41_PREFILL_TIMING=1): per-phase budget of eager prefill steps. +import os as _dsv41_os # noqa: E402 + +_DSV41_PREFILL_TIMING = _dsv41_os.environ.get("DSV41_PREFILL_TIMING", "0") == "1" + + +def _dsv41_prefill_timer_start(input_batch, batch_desc, dummy_run): + if ( + not _DSV41_PREFILL_TIMING + or dummy_run + or input_batch.num_tokens < 1024 + or batch_desc.cg_mode == CUDAGraphMode.FULL + or torch.cuda.is_current_stream_capturing() + ): + return None + try: # drop phase time accumulated outside this forward (e.g. the DSpark drafter of the previous step) + from vllm.model_executor.kernels.attention.dsa.candidate_blocks import prefill_acc_take + + prefill_acc_take() + except Exception: # noqa: BLE001 + pass + e0 = torch.cuda.Event(enable_timing=True) + e0.record() + return e0 + + +def _dsv41_prefill_timer_end(e0, input_batch, model_state): + if e0 is None: + return + e1 = torch.cuda.Event(enable_timing=True) + e1.record() + e1.synchronize() + fwd = e0.elapsed_time(e1) + try: + from vllm.model_executor.kernels.attention.dsa.candidate_blocks import ( + prefill_acc_take, + ) + + acc = prefill_acc_take() + except Exception: # noqa: BLE001 + acc = {} + g = acc.get("attn_gather", 0.0) + k = acc.get("attn_kernel", 0.0) + c = acc.get("attn_combine", 0.0) + idx = acc.get("indexer", 0.0) + cand = acc.get("candidates", 0.0) + stager = getattr(model_state, "engram_stager", None) + eng = float(getattr(stager, "_t_last_total_ms", 0.0) or 0.0) if stager is not None else 0.0 + n = int(input_batch.num_tokens) + moe = acc.get("moe", 0.0) + ar = acc.get("allreduce", 0.0) + egf = acc.get("engram_fwd", 0.0) + rest = fwd - (g + k + c) - idx + logger.info( + "Prefill step timing: %d tokens / %d reqs: forward %.0f ms (%.0f tok/s) = " + "attention %.0f (gather %.0f, kernel %.0f, combine %.0f) + indexer %.0f " + "(incl. candidates %.0f) + rest[MoE/dense/other] %.0f [moe %.0f, allreduce %.0f, " + "engram_fwd %.0f, other %.0f]; Engram staging before the step %.0f ms", + n, + int(input_batch.num_reqs), + fwd, + 1e3 * n / max(fwd, 1e-6), + g + k + c, + g, + k, + c, + idx, + cand, + rest, + moe, + ar, + egf, + rest - moe - ar - egf, + eng, + ) + class GPUModelRunner(LoRAModelRunnerMixin): def __init__(self, vllm_config: VllmConfig, device: torch.device): @@ -558,9 +635,13 @@ block_sizes = [] max_num_blocks_per_group = [] slot_mapping_enabled = [] + cp_sharded = [] for kv_cache_group in kv_cache_config.kv_cache_groups: spec = kv_cache_group.kv_cache_spec block_sizes.append(spec.block_size) + # DCP-replicated groups (V4.1 SWA cache, compressor ring) store + # every position on every rank; only sharded groups round-robin. + cp_sharded.append(not is_dcp_replicated_spec(spec)) layer_spec = ( spec.first_spec if isinstance(spec, UniformTypeKVCacheSpecs) else spec ) @@ -623,7 +704,15 @@ cp_size=self.dcp_size, cp_rank=self.dcp_rank, cp_interleave=self.cp_interleave, + cp_sharded=cp_sharded, ) + if self.use_dcp: + logger.info( + "DCP%d KV cache groups: sharded=%s (block sizes %s)", + self.dcp_size, + cp_sharded, + block_sizes, + ) self.pcp_manager = pcp.maybe_build_pcp_manager( self.vllm_config, self.device, @@ -1752,6 +1841,7 @@ input_batch, batch_desc.cg_mode == CUDAGraphMode.FULL ) self.step_timing.forward_start() + _dsv41_pf = _dsv41_prefill_timer_start(input_batch, batch_desc, dummy_run) # Run model. if batch_desc.cg_mode == CUDAGraphMode.FULL: @@ -1794,6 +1884,7 @@ else: # Eager (NONE): call the raw model directly. model_output = self.model(**model_inputs) + _dsv41_prefill_timer_end(_dsv41_pf, input_batch, self.model_state) if self.is_last_pp_rank: if self.use_aux_hidden_state_outputs: --- a/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py 2026-09-12 07:48:35.759099133 +0200 +++ b/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py 2026-09-11 10:13:47.000000000 +0200 @@ -377,6 +377,19 @@ assert self.draft_kv_cache_group_id >= 0 # Support multiple draft KV cache groups by preparing inputs once for each for i, gid in enumerate(self.draft_kv_cache_group_ids): + # The draft's context/query slots must follow the GROUP's CP view: + # a DCP-replicated group (DeepSeek-V4.1 draft SWA caches) stores + # every position on every rank at the plain block geometry, so it + # gets the DCP1 mapping; only sharded groups round-robin. + group_sharded = getattr(self.block_tables, "_cp_sharded", None) + if group_sharded is None or group_sharded[gid]: + cp_rank, cp_size, cp_interleave = ( + self.block_tables.cp_rank, + self.block_tables.cp_size, + self.block_tables.cp_interleave, + ) + else: + cp_rank, cp_size, cp_interleave = 0, 1, 1 prepare_dflash_inputs( self.input_buffers, self.block_tables.slot_mappings[gid], @@ -396,9 +409,9 @@ seeds, self.block_tables.input_block_tables[gid], self.block_tables.kernel_block_sizes[gid], - self.block_tables.cp_rank, - self.block_tables.cp_size, - self.block_tables.cp_interleave, + cp_rank, + cp_size, + cp_interleave, self.parallel_drafting_token_id, self.num_query_per_req, self.num_speculative_steps,