--- baked-ours/v1/attention/backends/mla/indexer.py 2026-09-10 19:56:14.000000000 +0200 +++ baked/v1/attention/backends/mla/indexer.py 2026-09-11 10:13:47.000000000 +0200 @@ -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( @@ -1235,16 +1260,12 @@ # batches. Keep its address stable across varlen graph replays. seq_lens_is_buffer_view = not use_native or next_n > 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 - ) - # 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 + ), )