launcher/dcp2/diffs/13-v1_attention_backends_mla_compressor_utils.py.diff

--- baked-ours/v1/attention/backends/mla/compressor_utils.py	2026-09-10 19:56:14.000000000 +0200
+++ baked/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