[Feature] Unified memory: support decode context parallelism for Kimi-Linear (#36890)

python/sglang/kernels/ops/attention/dcp_kernels.py

@@ -174,11 +174,7 @@ def update_kv_lens_and_indices(
     local_kv_indices_offsets = local_kv_indices_start + offsets
 
     kv_values = tl.load(kv_indices + kv_indice_offsets, mask=mask)
-    tl.store(
-        local_kv_indices + local_kv_indices_offsets,
-        kv_values // dcp_world_size,
-        mask=mask,
-    )
+    tl.store(local_kv_indices + local_kv_indices_offsets, kv_values, mask=mask)
 
 
 # ---------------------------------------------------------------------------

python/sglang/kernels/ops/kvcache/mla_buffer.py

@@ -88,13 +88,15 @@ def set_mla_kv_buffer_kernel(
 _TMA_BULK_STORE_MIN_LOCS = 768
 
 
-def set_mla_kv_buffer_triton(
+def _set_mla_kv_buffer_impl(
     kv_buffer: torch.Tensor,
     loc: torch.Tensor,
     cache_k_nope: torch.Tensor,
     cache_k_rope: torch.Tensor,
     *,
-    reserved_skip_index: int = 0,
+    reserved_skip_index: int,
+    dcp_world_size: int,
+    dcp_rank: int,
 ):
     """Dispatch MLA paged-KV scatter writes to the fastest available path.
 
@@ -121,6 +123,9 @@ def set_mla_kv_buffer_triton(
 
     Writes targeting ``reserved_skip_index`` are skipped. Slot 0 is reserved
     for CUDA-graph padding by default; pass -1 to disable skipping.
+
+    Shared body of the two entry points below; the owner rule reaches it as
+    ``1, 0`` (nothing to select) or as the live topology.
     """
     from sglang.kernels.ops.kvcache.set_mla_kv_buffer import (
         can_use_set_mla_kv_buffer,
@@ -136,7 +141,7 @@ def set_mla_kv_buffer_triton(
         n_loc >= _TMA_BULK_STORE_MIN_LOCS
         and is_arch_support_pdl()
         and can_use_set_mla_kv_buffer(nope_bytes, rope_bytes)
-        and not get_parallel().dcp_enabled
+        and dcp_world_size == 1
     ):
         jit_set_mla_kv_buffer(
             kv_buffer,
@@ -170,12 +175,54 @@ def set_mla_kv_buffer_triton(
         nope_dim,
         rope_dim,
         BLOCK=BLOCK,
-        DCP_RANK=get_parallel().attn_dcp_rank,
-        DCP_WORLD_SIZE=get_parallel().attn_dcp_size,
+        DCP_RANK=dcp_rank,
+        DCP_WORLD_SIZE=dcp_world_size,
         **pdl_kwargs,
     )
 
 
+def set_mla_kv_buffer_triton(
+    kv_buffer: torch.Tensor,
+    loc: torch.Tensor,
+    cache_k_nope: torch.Tensor,
+    cache_k_rope: torch.Tensor,
+    *,
+    reserved_skip_index: int = 0,
+):
+    """Scatter at locs already addressing this rank's rows (widened ->
+    `set_mla_kv_buffer_dcp_sharded_triton`)."""
+    _set_mla_kv_buffer_impl(
+        kv_buffer,
+        loc,
+        cache_k_nope,
+        cache_k_rope,
+        reserved_skip_index=reserved_skip_index,
+        dcp_world_size=1,
+        dcp_rank=0,
+    )
+
+
+def set_mla_kv_buffer_dcp_sharded_triton(
+    kv_buffer: torch.Tensor,
+    loc: torch.Tensor,
+    cache_k_nope: torch.Tensor,
+    cache_k_rope: torch.Tensor,
+    *,
+    reserved_skip_index: int = 0,
+):
+    """Scatter at DCP-WIDENED locs: select this rank's ids and collapse them."""
+    parallel = get_parallel()
+    _set_mla_kv_buffer_impl(
+        kv_buffer,
+        loc,
+        cache_k_nope,
+        cache_k_rope,
+        reserved_skip_index=reserved_skip_index,
+        dcp_world_size=parallel.attn_dcp_size,
+        dcp_rank=parallel.attn_dcp_rank,
+    )
+
+
 @triton.jit
 def set_mla_kv_buffer_fp8_quant_kernel(
     kv_buffer_fp8_ptr,

python/sglang/srt/arg_groups/kv_cache_hook.py

@@ -253,12 +253,8 @@ def handle_unified_memory_pool(server_args: Any) -> None:
         "full-attention slots are VIRTUAL — the host-offload path does not "
         "translate them to physical."
     )
-    assert cfg.dcp_size == 1, (
-        "--enable-unified-memory is not yet compatible with decode context "
-        "parallelism (--dcp-size > 1): the pool has no DCP-aware masked write "
-        "path (UnifiedMHATokenToKVPool.set_kv_buffer asserts dcp_kv_mask is None), "
-        "so a DCP run would boot and then fail on the first KV write."
-    )
+    if cfg.dcp_size > 1:
+        _validate_unified_memory_dcp(server_args)
     # Only monolithic decode cuda-graph capture is wired; piecewise prefill
     # capture is not. Guard when the user opts into it.
     _cg_cfg = cfg.cuda_graph_config
@@ -280,6 +276,53 @@ def handle_unified_memory_pool(server_args: Any) -> None:
         )
 
 
+def _validate_unified_memory_dcp(server_args: Any) -> None:
+    """Gate --enable-unified-memory + --dcp-size > 1 to the audited path.
+
+    Under DCP the unified allocator hands out a WIDENED virtual id space
+    (dcp_size logical ids per stored row) and every read index reaches
+    `translate_kv_loc*` already collapsed by a DCP index kernel. Only the
+    pieces below have been converted to that two-stage contract.
+    """
+    assert use_mla_backend(server_args), (
+        "--enable-unified-memory with decode context parallelism "
+        "(--dcp-size > 1) supports MLA models only (e.g. kimi-linear): the "
+        "MHA unified pool has no DCP-aware masked write path "
+        "(UnifiedMHATokenToKVPool.set_kv_buffer asserts dcp_kv_mask is None)."
+    )
+    assert not model_config_of(server_args).is_hybrid_swa, (
+        "--enable-unified-memory with decode context parallelism "
+        "(--dcp-size > 1) does not support hybrid sliding-window models: "
+        "UnifiedSWATokenToKVPoolAllocator does not widen its virtual id "
+        "space, and the full->swa mapping is not DCP-sharded."
+    )
+    cfg = resolving_view(server_args)
+    assert cfg.disaggregation_mode == "null", (
+        "--enable-unified-memory with decode context parallelism "
+        "(--dcp-size > 1) does not support PD disaggregation: the transfer "
+        "ships whole page envelopes, which under DCP hold only this rank's "
+        "shard of each widened page. Rejected here rather than at the first KV "
+        "transfer, where translate_kv_indices_for_transfer would abort a "
+        "server that had already booted."
+    )
+    # trtllm_mla (and its cutedsl_mla / tokenspeed_mla subclasses) build the
+    # MLA block table straight from req_to_token with
+    # create_flashmla_kv_indices_triton, whose v2p gather assumes UNWIDENED
+    # page ids; the DCP variant (create_mla_kv_page_table_for_dcp) has no v2p
+    # gather at all. Wire one of them through the other to add those here.
+    dcp_allowed = {"flashinfer"}
+    backends = set(attention_backends_of(resolved_view(server_args)))
+    backends.discard(None)
+    assert backends <= dcp_allowed, (
+        "--enable-unified-memory with decode context parallelism "
+        f"(--dcp-size > 1) requires {sorted(dcp_allowed)} for the "
+        f"full-attention layers; got {sorted(backends)}. The other paged MLA "
+        "backends build their block table from raw (widened) req_to_token "
+        "page ids and do not translate them through the unified pool's "
+        "virtual->physical page table."
+    )
+
+
 def handle_page_major_kv_layout(server_args: Any):
     # The unified pool stores state in the page-major envelope-strided layout, so
     # enabling it implies --enable-page-major-kv-layout — routing it through the

python/sglang/srt/layers/attention/flashinfer_mla_backend.py

@@ -930,8 +930,10 @@ class FlashInferMLAIndicesUpdaterDecode:
                 ENTRY_PAGE_SIZE=kv_view.entry_page_size,
             )
 
+            # The table above is deliberately VIRTUAL under DCP.
+            n_kernel_ids = paged_kernel_lens_sum
             if get_parallel().dcp_enabled:
-                plan_dcp_decode_metadata(
+                n_kernel_ids = plan_dcp_decode_metadata(
                     kv_lens,
                     kv_indptr,
                     kv_indices,
@@ -939,6 +941,18 @@ class FlashInferMLAIndicesUpdaterDecode:
                     fast_decode_kwargs,
                     bs,
                 )
+            # Written back IN PLACE: on cuda-graph replay `kv_indices` IS the
+            # capture-stable buffer the captured wrapper reads, so rebinding the
+            # local name would leave the graph on virtual ids. Only the prefix
+            # just filled is translated; the stale tail never indexes v2p.
+            translator = self.attn_backend.kv_index_translator
+            if (
+                not kv_view.is_translated
+                and n_kernel_ids > 0
+                and translator.needs_read_translate
+            ):
+                valid = kv_indices[:n_kernel_ids]
+                valid.copy_(translator.translate_dcp_read_ids(valid))
         else:
             kv_indptr, kv_indices = spec_info.kv_indptr, spec_info.kv_indices
 

python/sglang/srt/layers/dcp/layout.py

@@ -42,12 +42,13 @@ def get_dcp_lens(
 
 
 def filter_dcp_local_kv_indices(kv_indices: torch.Tensor):
+    """Keep this rank's share of a read-index tensor, still WIDENED.
+
+    Selection only; the caller collapses via translate_dcp_read_ids.
+    """
     parallel = get_parallel()
     if parallel.dcp_enabled:
-        kv_indices = (
-            kv_indices[kv_indices % parallel.dcp_size == parallel.dcp_rank]
-            // parallel.dcp_size
-        )
+        kv_indices = kv_indices[kv_indices % parallel.dcp_size == parallel.dcp_rank]
     return kv_indices
 
 
@@ -67,7 +68,7 @@ def filter_dcp_local_chunk_kv_indices(
         first = (parallel.dcp_rank - start) % dcp_size
         parts.append(kv_indices[offset + first : offset + length : dcp_size])
         offset += length
-    return torch.cat(parts) // dcp_size
+    return torch.cat(parts)
 
 
 def update_local_kv_lens_for_dcp(kv_len_arr):

python/sglang/srt/layers/dcp/planner.py

@@ -26,6 +26,7 @@ from sglang.kernels.ops.attention.dcp_kernels import (
 )
 from sglang.srt.layers.dcp.layout import update_local_kv_lens_for_dcp
 from sglang.srt.layers.dcp.metadata import DecodeContextParallelMetadata
+from sglang.srt.model_executor.forward_context import get_attn_backend
 from sglang.srt.runtime_context import get_device, get_parallel
 
 
@@ -110,9 +111,10 @@ def prepare_decode_context_parallel_metadata(
         parallel.dcp_size,
     )
     # Prefix lengths are dcp_size-aligned (widened allocator page), so no nonzero().
-    dcp_local_prefix_kv_indices = (
+    # `get_mla_kv_buffer` is a read door with the caller-translates contract.
+    translator = get_attn_backend().kv_index_translator
+    dcp_local_prefix_kv_indices = translator.translate_dcp_read_ids(
         dcp_prefix_kv_indices[parallel.dcp_rank :: parallel.dcp_size]
-        // parallel.dcp_size
     )
     dcp_kv_buffer = torch.empty(
         (
@@ -139,7 +141,14 @@ def plan_dcp_decode_metadata(
     init_metadata_replay: bool,
     fast_decode_kwargs: dict,
     bs: int,
-):
+) -> int:
+    """Shard `kv_indices` to this DCP rank in place; return the shard's length.
+
+    `kv_lens` / `kv_indptr` are rewritten to the per-rank lengths, and this
+    rank's ids (`loc % dcp_size == dcp_rank`) are compacted into
+    `kv_indices[:total_local_len]`, still WIDENED. The returned length bounds the
+    prefix the caller hands to `KVIndexTranslator.translate_dcp_read_ids`.
+    """
     parallel = get_parallel()
     local_kv_lens = kv_lens.clone()
     update_local_kv_lens_for_dcp(local_kv_lens)
@@ -185,3 +194,4 @@ def plan_dcp_decode_metadata(
     kv_indices[:total_local_len] = local_kv_indices[:total_local_len]
     kv_lens.copy_(local_kv_lens)
     kv_indptr[: bs + 1] = local_kv_lens_cumsum[: bs + 1]
+    return total_local_len

python/sglang/srt/mem_cache/kv_index_translator.py

@@ -63,6 +63,7 @@ from sglang.srt.mem_cache.multi_ended_allocator import (
     UnifiedSWATokenToKVPoolAllocator,
 )
 from sglang.srt.mem_cache.swa_memory_pool import SWAKVPool
+from sglang.srt.runtime_context import get_parallel
 
 
 class KVReadTables(msgspec.Struct, frozen=True):
@@ -122,6 +123,14 @@ class KVIndexTranslator:
             self._full_p2v_table = alloc.full_p2v_page_table
             self._full_page_multiplier = alloc.kernel_page_multiplier
             self._translate_full = alloc.translate_kv_loc_for_kernel
+            # The WRITE loc is the one id that arrives DCP-WIDENED: read indices
+            # are collapsed by the DCP index kernels, `out_cache_loc` still
+            # carries the owner rule in `loc % dcp_size`. Identity with the read
+            # translate when dcp_size == 1.
+            self._translate_write_full = alloc.translate_write_loc_for_kernel
+            # DCP read ids stay WIDENED to the consumer: selecting this rank's
+            # share changes the length, so only the production site can do it.
+            self.defer_read_translate = get_parallel().attn_dcp_size > 1
             if isinstance(alloc, UnifiedSWATokenToKVPoolAllocator):
                 self._swa_v2p_table = alloc.swa_v2p_page_table
                 self._swa_page_multiplier = alloc.swa_kernel_page_multiplier
@@ -135,6 +144,8 @@ class KVIndexTranslator:
             self._full_p2v_table = None
             self._full_page_multiplier = 1
             self._translate_full = None
+            self._translate_write_full = None
+            self.defer_read_translate = False
             self._swa_v2p_table = None
             self._swa_page_multiplier = 1
             self._swa_write_loc_from_full = (
@@ -192,7 +203,7 @@ class KVIndexTranslator:
         captured graph bakes it) passes its own tables in ``into``;
         ``into=None`` allocates of width ``max_pages`` instead.
         """
-        if not self.is_translating:
+        if not self.is_translating or self.defer_read_translate:
             return KVIndexTable(
                 ids=self.req_to_token,
                 row_ids=req_pool_indices,
@@ -308,20 +319,22 @@ class KVIndexTranslator:
         self._index_table_memo = (weakref.ref(forward_batch), view)
         return view
 
-    def assert_backends_carry_translator(self, backends) -> None:
-        """Boot guard: under the unified pool every backend a forward can reach
-        must carry THIS translator."""
-        if not self.is_translating:
-            return
+    def bind_and_verify_backends(self, backends) -> None:
+        """Boot: make every reachable backend carry THIS translator.
+
+        Model-layer producers read it off `get_attn_backend()`, so an unset
+        attribute is an unreachable hook, not "no translation needed".
+        """
         for backend in backends:
             if backend is None:
                 continue
+            if backend.kv_index_translator is None:
+                backend.kv_index_translator = self
+                continue
             assert backend.kv_index_translator is self, (
-                f"{type(backend).__name__} does not carry the runner's "
-                "KVIndexTranslator. A backend (or wrapper) reachable under "
-                "--enable-unified-memory must forward `kv_index_translator`, or "
-                "read-index producers silently skip the virtual->kernel-facing "
-                "translation."
+                f"{type(backend).__name__} carries a KVIndexTranslator that is "
+                "not this runner's. A wrapper must forward the inner backend's "
+                "copy, not build its own."
             )
 
     # -- write loc (phase 1; phase 2 lives in build_index_table) ----------------
@@ -338,7 +351,9 @@ class KVIndexTranslator:
         self._index_table_memo = None
         if not self.is_translating or forward_batch.out_cache_loc is None:
             return
-        forward_batch.out_cache_loc = self._translate_full(forward_batch.out_cache_loc)
+        forward_batch.out_cache_loc = self._translate_write_full(
+            forward_batch.out_cache_loc
+        )
 
     def sliding_window_write_loc_for(
         self, out_cache_loc: Optional[torch.Tensor]
@@ -363,6 +378,23 @@ class KVIndexTranslator:
 
     # -- token-level translate surface (the mixin / local-attn consumers) ------
 
+    @property
+    def needs_read_translate(self) -> bool:
+        """Whether `translate_dcp_read_ids` is anything but the identity, so a
+        hot path can skip the call rather than round-trip a no-op copy."""
+        return self.is_translating or get_parallel().attn_dcp_size > 1
+
+    def translate_dcp_read_ids(self, widened_ids: torch.Tensor) -> torch.Tensor:
+        """Widened logical READ ids -> kernel-facing ids, for either pool.
+
+        The one hook every DCP read-index production site calls; on a static
+        pool `widened // dcp_size` IS the whole virtual->physical translation.
+        """
+        dcp_size = get_parallel().attn_dcp_size
+        if dcp_size > 1:
+            widened_ids = widened_ids // dcp_size
+        return self.translate_full_attn_ids(widened_ids)
+
     def translate_full_attn_ids(
         self, kv_indices: torch.Tensor, *, out: Optional[torch.Tensor] = None
     ) -> torch.Tensor:

python/sglang/srt/mem_cache/memory_pool.py

@@ -66,6 +66,7 @@ from sglang.srt.mem_cache.layout.page_major import (
 from sglang.srt.mem_cache.utils import (
     get_mla_kv_buffer_triton,
     maybe_init_custom_mem_pool,
+    set_mla_kv_buffer_dcp_sharded_triton,
     set_mla_kv_buffer_triton,
     set_mla_kv_buffer_triton_fp8_quant,
     set_mla_kv_scale_buffer_triton,
@@ -3657,8 +3658,9 @@ class HybridLinearKVPool(KVCache):
         self.head_num = head_num
         self.head_dim = head_dim
         self.mamba_pool = mamba_pool
-        # virtual->physical mamba-slot translate for the HiCache offload path;
-        # identity for a static pool, the allocator's `translate` for the unified pool.
+        # Identity even though the unified pool holds VIRTUAL mamba ids: its
+        # composite allocator implements neither `get_cpu_copy` nor
+        # `load_cpu_copy`, the only readers, so those ids never arrive here.
         self._mamba_translate = lambda ids: ids
         self.use_mla = use_mla
         if full_kv_pool is not None:
@@ -4078,6 +4080,32 @@ class MLATokenToKVPool(KVCache):
     def get_kv_buffer(self, layer_id: int):
         return self.get_key_buffer(layer_id), self.get_value_buffer(layer_id)
 
+    # Has the WRITE loc arriving here already had the DCP owner rule resolved?
+    # False: this pool takes a WIDENED loc. The unified pool resolves it in
+    # `KVIndexTranslator.rebind_write_loc` and flips this. Not derivable from
+    # `kernel_page_blocks`: that is `layer_num`, so a rank owning one
+    # full-attention layer is translated with blocks_per_page 1.
+    write_loc_is_dcp_resolved = False
+
+    @property
+    def _write_loc_dcp_span(self) -> int:
+        """How many logical ids one stored row spans in the write-loc space."""
+        return 1 if self.write_loc_is_dcp_resolved else get_parallel().attn_dcp_size
+
+    def _scatter_mla_rows(
+        self,
+        dst_buffer: torch.Tensor,
+        loc: torch.Tensor,
+        cache_k_nope: torch.Tensor,
+        cache_k_rope: torch.Tensor,
+    ) -> None:
+        if self.write_loc_is_dcp_resolved:
+            set_mla_kv_buffer_triton(dst_buffer, loc, cache_k_nope, cache_k_rope)
+        else:
+            set_mla_kv_buffer_dcp_sharded_triton(
+                dst_buffer, loc, cache_k_nope, cache_k_rope
+            )
+
     def set_kv_buffer(
         self,
         layer: RadixAttention,
@@ -4095,12 +4123,15 @@ class MLATokenToKVPool(KVCache):
             layer_id_override if layer_id_override is not None else layer.layer_id
         )
         assert not self.dsa_kv_cache_store_fp8
-        parallel = get_parallel()
-        if parallel.dcp_enabled:
-            valid_mask = loc % parallel.attn_dcp_size == parallel.attn_dcp_rank
-            if not valid_mask.all():
-                loc = loc[valid_mask]
-                cache_k = cache_k[valid_mask]
+        # No DCP-aware variant is possible: the two backends reaching this door
+        # disagree on the loc space (flashinfer-MLA widened, Triton collapsed).
+        assert self.write_loc_is_dcp_resolved or not get_parallel().dcp_enabled, (
+            "MLATokenToKVPool.set_kv_buffer has no DCP-aware write path. Under "
+            "--dcp-size > 1 the MLA write must go through set_mla_kv_buffer, "
+            "whose kernel resolves the owner rule; reaching the combined-row "
+            "door means an attention backend took a write path that never "
+            "declared which loc space it emits."
+        )
         if cache_k.dtype != self.dtype:
             cache_k = cache_k.to(self.dtype)
 
@@ -4118,6 +4149,10 @@ class MLATokenToKVPool(KVCache):
         cache_k_nope: torch.Tensor,
         cache_k_rope: torch.Tensor,
     ) -> None:
+        assert not (
+            self.write_loc_is_dcp_resolved
+            and (self.use_dsa or self.dsa_kv_cache_store_fp8)
+        ), "the DSA write paths have no resolved-loc variant"
         if _is_hip and self.use_dsa and self.dtype == fp8_dtype:
             # HIP FP8 path uses raw MLA KV layout (nope + rope) without per-block scales.
             # Fuse BF16/FP16 -> FP8 cast with paged KV write.
@@ -4139,12 +4174,7 @@ class MLATokenToKVPool(KVCache):
             # Reuse existing two-tensor write kernel (works with FP8 byte layout)
             # cache_k_nope_fp8: (num_tokens, 1, 528) uint8 [nope_fp8(512) | scales(16)]
             # cache_k_rope_fp8: (num_tokens, 1, 128) uint8 [rope_bf16_bytes(128)]
-            set_mla_kv_buffer_triton(
-                dst_buffer,
-                loc,
-                cache_k_nope_fp8,
-                cache_k_rope_fp8,
-            )
+            self._scatter_mla_rows(dst_buffer, loc, cache_k_nope_fp8, cache_k_rope_fp8)
         else:
             if cache_k_nope.dtype != self.dtype:
                 cache_k_nope = cache_k_nope.to(self.dtype)
@@ -4153,12 +4183,7 @@ class MLATokenToKVPool(KVCache):
                 cache_k_nope = cache_k_nope.view(self.store_dtype)
                 cache_k_rope = cache_k_rope.view(self.store_dtype)
 
-            set_mla_kv_buffer_triton(
-                dst_buffer,
-                loc,
-                cache_k_nope,
-                cache_k_rope,
-            )
+            self._scatter_mla_rows(dst_buffer, loc, cache_k_nope, cache_k_rope)
 
     def set_mla_kv_buffer(
         self,
@@ -4168,11 +4193,11 @@ class MLATokenToKVPool(KVCache):
         cache_k_rope: torch.Tensor,
         layer_id_override: Optional[int] = None,
     ):
-        # loc is widened under DCP; the kernel divides by the world size itself.
+        # loc is widened under DCP unless the pool declares it resolved.
         maybe_detect_oob(
             loc,
             0,
-            (self.size + self.page_size) * get_parallel().attn_dcp_size,
+            (self.size + self.page_size) * self._write_loc_dcp_span,
             "set_mla_kv_buffer (MLA)",
         )
         maybe_detect_kernel_facing_loc(
@@ -4379,7 +4404,7 @@ class MLATokenToKVPoolFP4(MLATokenToKVPool):
                 cache_k_nope = cache_k_nope.view(self.store_dtype)
                 cache_k_rope = cache_k_rope.view(self.store_dtype)
 
-            set_mla_kv_buffer_triton(
+            self._scatter_mla_rows(
                 self.kv_buffer[layer_id - self.start_layer],
                 loc,
                 cache_k_nope_fp4,

python/sglang/srt/mem_cache/multi_ended_allocator.py

@@ -52,6 +52,7 @@ from sglang.srt.mem_cache.unified_memory_pool import (
     UnifiedKVPool,
     UnifiedMLATokenToKVPool,
 )
+from sglang.srt.runtime_context import get_parallel
 from sglang.srt.utils.common import get_num_new_pages, next_power_of_2
 
 logger = logging.getLogger(__name__)
@@ -263,6 +264,7 @@ class MultiEndedAllocator(BaseTokenToKVPoolAllocator):
         device: str,
         is_id_owner: bool,
         page_size: int = 1,
+        shards_under_dcp: bool = False,
         need_sort: bool = False,
         forward_stream: Optional[torch.cuda.Stream] = None,
         lazy_compaction: bool = False,
@@ -270,9 +272,13 @@ class MultiEndedAllocator(BaseTokenToKVPoolAllocator):
     ):
         spec = unified_buffer.spec(sub_pool_name)
         max_slots = unified_buffer.max_slots(sub_pool_name)
+        # DCP shards KV tokens only. Mamba state and the SWA rows are
+        # replicated, so they stay slot-granular whatever the process width is.
+        self.shards_under_dcp = shards_under_dcp
+        dcp_size = get_parallel().attn_dcp_size if shards_under_dcp else 1
         super().__init__(
-            size=max_slots,
-            page_size=page_size,
+            size=max_slots * dcp_size,
+            page_size=page_size * dcp_size,
             dtype=spec.get_dtype(),
             device=device,
             kvcache=kvcache,
@@ -301,12 +307,26 @@ class MultiEndedAllocator(BaseTokenToKVPoolAllocator):
         self.forward_stream = forward_stream
 
         # --- Page-aware bookkeeping ---
-        # `min_page_index` = ceil(min_slot_index / page_size), keeping the
+        # Two page sizes, equal unless decode context parallelism is on:
+        # `page_size` is VIRTUAL (what the scheduler, the tree cache and the
+        # alloc/free surface speak, matching PagedTokenToKVPoolAllocator's
+        # widened DCP contract), `pool_page_size` is the PHYSICAL rows one page
+        # occupies here. Under DCP a virtual page holds dcp_size logical ids per
+        # stored row, of which this rank owns `loc % dcp_size == dcp_rank`;
+        # `KVIndexTranslator.translate_dcp_read_ids` collapses `loc // dcp_size`
+        # before reaching `translate_kv_loc*`, so everything at or below the v2p
+        # table -- byte budget, compaction moves, translate -- stays on
+        # `pool_page_size`.
+        # Page ids are invariant under the widening, so v2p/p2v are unchanged.
+        self.pool_page_size = page_size
+        self.page_size = page_size * dcp_size
+        self.num_pages = max_slots // self.pool_page_size
+        # `min_page_index` = ceil(min_slot_index / pool_page_size), keeping the
         # reserved-sink invariant (min_page_index * entry_bytes_per_page >= entry_max).
-        self.page_size = page_size
-        self.num_pages = max_slots // page_size
-        self.min_page_index = (self.min_slot_index + page_size - 1) // page_size
-        self.entry_bytes_per_page = self.entry_bytes * page_size
+        self.min_page_index = (
+            self.min_slot_index + self.pool_page_size - 1
+        ) // self.pool_page_size
+        self.entry_bytes_per_page = self.entry_bytes * self.pool_page_size
 
         # v2p / p2v sized by PAGES. Page 0 is the padding anchor; trailing row is
         # the -1 sentinel.
@@ -982,6 +1002,10 @@ class MultiEndedAllocator(BaseTokenToKVPoolAllocator):
     ) -> torch.Tensor:
         """Translate token-granular virtual ids to physical ids.
 
+        Under DCP the input is the DCP-collapsed id (`widened // dcp_size`, what
+        `KVIndexTranslator.translate_dcp_read_ids` hands down), so this works on
+        `pool_page_size`.
+
         ``out=`` writes in-place into a caller-owned buffer — required under
         cuda-graph capture for buffer-stability (the captured graph records the
         gather against a fixed ``data_ptr``).
@@ -1008,7 +1032,8 @@ class MultiEndedAllocator(BaseTokenToKVPoolAllocator):
         # routes any tombstoned read/write to physical slot 0 — reserved
         # padding-sink space by the `min_slot_index` invariant (bytes [0, entry_max)
         # across all sub-pools hold no real data).
-        if self.page_size == 1:
+        ps = self.pool_page_size
+        if ps == 1:
             if out is not None:
                 # `index_select(out=out)` forbids index/out aliasing, but the
                 # canonical caller does in-place `translate(kv_indices, out=kv_indices)`.
@@ -1019,18 +1044,18 @@ class MultiEndedAllocator(BaseTokenToKVPoolAllocator):
                 return out
             result = torch.index_select(self.virtual_to_physical, 0, virt_tokens)
             return torch.clamp_min(result, 0)
-        # page_size > 1: page math. `virt_pages`/`offsets` are fresh, so they
+        # ps > 1: page math. `virt_pages`/`offsets` are fresh, so they
         # cannot alias `out` — `index_select(out=out)` is safe.
-        virt_pages = virt_tokens // self.page_size
-        offsets = virt_tokens % self.page_size
+        virt_pages = virt_tokens // ps
+        offsets = virt_tokens % ps
         if out is not None:
             torch.index_select(self.virtual_to_physical, 0, virt_pages, out=out)
-            out.mul_(self.page_size)
+            out.mul_(ps)
             out.add_(offsets)
             out.clamp_(min=0)  # tombstoned page: -1*ps + offset in [-ps, -1]
             return out
         phys_pages = self.virtual_to_physical[virt_pages]
-        result = phys_pages * self.page_size + offsets
+        result = phys_pages * ps + offsets
         return torch.clamp_min(result, 0)
 
     def translate_kv_loc_for_kernel(
@@ -1048,7 +1073,7 @@ class MultiEndedAllocator(BaseTokenToKVPoolAllocator):
         clamp to kernel-facing id 0, the page-0 sink. int64 out; a consumer whose
         kernel ABI wants int32 narrows where it fills that buffer.
         """
-        ps = self.page_size
+        ps = self.pool_page_size
         stride = ps * self.kernel_page_multiplier
         with record_function("MultiEndedAlloc.translate_kv_loc_for_kernel"):
             pages = virt_tokens if ps == 1 else virt_tokens // ps
@@ -1076,6 +1101,33 @@ class MultiEndedAllocator(BaseTokenToKVPoolAllocator):
                 out.add_(offsets)
             return out.clamp_(min=0)
 
+    def translate_write_loc_for_kernel(
+        self,
+        widened_loc: torch.Tensor,
+        *,
+        out: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        """Widened virtual WRITE loc (`out_cache_loc`) -> kernel-facing id.
+
+        Reads arrive already DCP-collapsed (every DCP index kernel divides), but
+        `out_cache_loc` does not: it still carries the owner rule in
+        `loc % dcp_size`. Resolve ownership, collapse, translate; ids this rank
+        does not own go to kernel id 0, the padding sink every write kernel
+        skips. Identity with `translate_kv_loc_for_kernel` at dcp_size == 1.
+        """
+        parallel = get_parallel()
+        dcp_size = parallel.attn_dcp_size if self.shards_under_dcp else 1
+        if dcp_size == 1:
+            return self.translate_kv_loc_for_kernel(widened_loc, out=out)
+        with record_function("MultiEndedAlloc.translate_write_loc_for_kernel"):
+            owned = (widened_loc % dcp_size) == parallel.attn_dcp_rank
+            dense = self.translate_kv_loc_for_kernel(widened_loc // dcp_size)
+            dense = torch.where(owned, dense, torch.zeros_like(dense))
+            if out is not None:
+                out.copy_(dense)
+                return out
+            return dense
+
     # -- alloc --
 
     def alloc(self, need_size: int) -> Optional[torch.Tensor]:
@@ -1503,15 +1555,15 @@ class MultiEndedAllocator(BaseTokenToKVPoolAllocator):
         """
         v_moved = self.physical_to_virtual[src_pages].clone()  # read pre-wipe
 
-        # Expand page ids to token ids for the token-granular move kernel.
-        if self.page_size == 1:
+        # Expand to PHYSICAL token granularity (the move kernel is
+        # token-granular over pool rows).
+        if self.pool_page_size == 1:
             src_t, dst_t = src_pages, dst_pages
         else:
-            offsets = torch.arange(
-                self.page_size, dtype=torch.int64, device=self.device
-            )
-            src_t = (src_pages[:, None] * self.page_size + offsets).reshape(-1)
-            dst_t = (dst_pages[:, None] * self.page_size + offsets).reshape(-1)
+            ps = self.pool_page_size
+            offsets = torch.arange(ps, dtype=torch.int64, device=self.device)
+            src_t = (src_pages[:, None] * ps + offsets).reshape(-1)
+            dst_t = (dst_pages[:, None] * ps + offsets).reshape(-1)
 
         # Un-translated copy: the public copy_from translates virtual ids,
         # which we must NOT do here.
@@ -1571,9 +1623,15 @@ class MultiEndedAllocator(BaseTokenToKVPoolAllocator):
             return None
         # `oclv` is non-None here (set_inflight_forward clears the slot otherwise).
         with record_function("MultiEndedAlloc._materialize_inflight_write_set"):
+            # `oclv` is a WIDENED virtual id under DCP; collapse to the id space
+            # translate speaks. The write set is a page set, and a widened page
+            # covers exactly the same page, so the non-owned ids fold in harmlessly.
+            dcp_size = get_parallel().attn_dcp_size if self.shards_under_dcp else 1
+            if dcp_size > 1:
+                oclv = oclv // dcp_size
             phys_tokens = self.translate_kv_loc(oclv)
-            if self.page_size > 1:
-                phys_pages = (phys_tokens // self.page_size).unique()
+            if self.pool_page_size > 1:
+                phys_pages = (phys_tokens // self.pool_page_size).unique()
             else:
                 phys_pages = phys_tokens
             return set(phys_pages.tolist())  # .tolist() syncs schedule_stream
@@ -1999,17 +2057,15 @@ class MultiEndedAllocator(BaseTokenToKVPoolAllocator):
                 (v_moveds_t >= 0).all(),
                 "invalid p2v mapping in MultiEndedAllocator._flush",
             )
-            # Expand to token granularity (the move kernel is token-granular).
-            if self.page_size == 1:
+            # Expand to PHYSICAL token granularity (the move kernel is
+            # token-granular over pool rows).
+            if self.pool_page_size == 1:
                 src_t, dst_t = src_pages_t, dst_pages_t
             else:
-                offsets = torch.arange(
-                    self.page_size,
-                    dtype=torch.int64,
-                    device=self.device,
-                )
-                src_t = (src_pages_t[:, None] * self.page_size + offsets).reshape(-1)
-                dst_t = (dst_pages_t[:, None] * self.page_size + offsets).reshape(-1)
+                ps = self.pool_page_size
+                offsets = torch.arange(ps, dtype=torch.int64, device=self.device)
+                src_t = (src_pages_t[:, None] * ps + offsets).reshape(-1)
+                dst_t = (dst_pages_t[:, None] * ps + offsets).reshape(-1)
             self._kvcache.move_kv_cache(dst_t, src_t)
             # ONE bulk remap (single-writer on schedule_stream).
             self.virtual_to_physical[v_moveds_t] = dst_pages_t
@@ -2711,9 +2767,10 @@ class UnifiedMambaTokenToKVPoolAllocator(BaseTokenToKVPoolAllocator):
         lazy_compaction: bool = False,
     ):
         full_max = unified_buffer.max_slots("full")
+        dcp_size = get_parallel().attn_dcp_size
         super().__init__(
-            size=full_max - 1,
-            page_size=page_size,
+            size=(full_max - 1) * dcp_size,
+            page_size=page_size * dcp_size,
             dtype=unified_buffer.spec("full").get_dtype(),
             device=device,
             kvcache=kvcache,
@@ -2721,11 +2778,13 @@ class UnifiedMambaTokenToKVPoolAllocator(BaseTokenToKVPoolAllocator):
         )
         self.unified_buffer = unified_buffer
         self._kvcache = kvcache
-        self.page_size = page_size
+        # Widened under DCP, matching the full sub-allocator; see its __init__.
+        self.page_size = page_size * dcp_size
         self.lazy_compaction = lazy_compaction
 
         # FULL is page-aware; MAMBA stays page_size=1 (state is per-request,
-        # orthogonal to the full side's per-token paging).
+        # orthogonal to the full side's per-token paging), and only FULL shards
+        # under DCP: mamba state is replicated on every rank.
         self.full_attn_allocator = MultiEndedAllocator(
             kvcache=kvcache.full_kv_pool,
             unified_buffer=unified_buffer,
@@ -2733,6 +2792,7 @@ class UnifiedMambaTokenToKVPoolAllocator(BaseTokenToKVPoolAllocator):
             device=device,
             is_id_owner=True,
             page_size=page_size,
+            shards_under_dcp=True,
             need_sort=need_sort,
             forward_stream=forward_stream,
             lazy_compaction=lazy_compaction,
@@ -2813,15 +2873,22 @@ class UnifiedMambaTokenToKVPoolAllocator(BaseTokenToKVPoolAllocator):
         stays inside the JOINT budget. = mamba bytes/slot ÷ full bytes/token, rounded
         UP (conservative). Only on the shared composite (non-shared pools are separate,
         so the planner sources this via `getattr(..., None)`).
+
+        The planner charges this against `rem_total_tokens`, which is fed by
+        `available_size()` -- widened under DCP. One widened token is
+        `entry_bytes / dcp_size` local bytes, so the conversion carries the same
+        `dcp_size`; leaving it out under-reserves the shared gap by that factor.
         """
         return -(
             -self.mamba_allocator.entry_bytes_per_page
+            * get_parallel().attn_dcp_size
             // self.full_attn_allocator.entry_bytes
         )
 
     @property
     def size_full(self) -> int:
-        return self.full_attn_allocator.max_slots - 1
+        # Widened like `size`: a logical token capacity, not a row count.
+        return (self.full_attn_allocator.max_slots - 1) * get_parallel().attn_dcp_size
 
     @property
     def size_mamba(self) -> int:
@@ -2915,6 +2982,15 @@ class UnifiedMambaTokenToKVPoolAllocator(BaseTokenToKVPoolAllocator):
         """Full-pool virtual TOKEN ids -> kernel-facing ids."""
         return self.full_attn_allocator.translate_kv_loc_for_kernel(loc, out=out)
 
+    def translate_write_loc_for_kernel(
+        self,
+        loc: torch.Tensor,
+        *,
+        out: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        """Widened virtual WRITE loc -> DENSE id; see the sub-allocator's copy."""
+        return self.full_attn_allocator.translate_write_loc_for_kernel(loc, out=out)
+
     def translate_kv_indices_for_transfer(
         self, kv_indices: torch.Tensor
     ) -> torch.Tensor:
@@ -2923,6 +2999,13 @@ class UnifiedMambaTokenToKVPoolAllocator(BaseTokenToKVPoolAllocator):
         PHYSICAL, not kernel-facing: the transfer registers page ENVELOPES (see
         `UnifiedMLATokenToKVPool.get_contiguous_buf_infos`).
         """
+        # Defensive: `_validate_unified_memory_dcp` rejects this pairing at
+        # argument validation, so reaching it means a config path got past that.
+        assert get_parallel().attn_dcp_size == 1, (
+            "PD-disaggregation transfer with the unified memory pool does not "
+            "support decode context parallelism: the transfer ships whole page "
+            "envelopes, which hold only this rank's shard of each widened page."
+        )
         return self.full_attn_allocator.translate_kv_loc(kv_indices.to(torch.int64))
 
     def set_disagg_move_gate(self, gate: Callable[[], bool]) -> None:
@@ -3361,6 +3444,17 @@ class UnifiedSWATokenToKVPoolAllocator(SWATokenToKVPoolAllocator):
         """Full-pool virtual TOKEN ids -> kernel-facing ids."""
         return self.full_attn_allocator.translate_kv_loc_for_kernel(loc, out=out)
 
+    def translate_write_loc_for_kernel(
+        self,
+        loc: torch.Tensor,
+        *,
+        out: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        """Widened virtual WRITE loc -> kernel-facing id; see the sub-allocator's
+        copy. DCP is rejected for this composite at argument validation, so this
+        is the dcp_size == 1 identity with the read translate."""
+        return self.full_attn_allocator.translate_write_loc_for_kernel(loc, out=out)
+
     @property
     def swa_kernel_page_multiplier(self) -> int:
         return self.swa_attn_allocator.kernel_page_multiplier

python/sglang/srt/mem_cache/unified_memory_pool.py

@@ -708,6 +708,10 @@ class UnifiedMLATokenToKVPool(MLATokenToKVPool):
         # Lifetime owned by UnifiedKVPool; do not delete the views.
         pass
 
+    # `rebind_write_loc` already collapsed the widened id and sent the rows this
+    # rank does not own to the padding sink.
+    write_loc_is_dcp_resolved = True
+
     def get_kv_size_bytes(self):
         return 0  # UnifiedKVPool logs the total; per-sub-pool would double-count
 
@@ -1332,7 +1336,7 @@ def init_unified_mamba_pools(
         max_size=req_to_token_pool._shared_mamba_size,
         device=device,
     )
-    # `_mamba_translate` feeds the HiCache offload path, GATED OFF here — wired but inert.
+    # Inert: this allocator implements neither reader (see HybridLinearKVPool).
     req_to_token_pool.mamba_allocator = mamba_slot_allocator
     token_to_kv_pool._mamba_translate = mamba_slot_allocator.translate
     # No full-KV translate hook is wired: both MLA doors now receive

python/sglang/srt/mem_cache/utils.py

@@ -22,6 +22,9 @@ from sglang.kernels.ops.kvcache.mla_buffer import (
 from sglang.kernels.ops.kvcache.mla_buffer import (
     get_mla_kv_buffer_triton as get_mla_kv_buffer_triton,
 )
+from sglang.kernels.ops.kvcache.mla_buffer import (
+    set_mla_kv_buffer_dcp_sharded_triton as set_mla_kv_buffer_dcp_sharded_triton,
+)
 from sglang.kernels.ops.kvcache.mla_buffer import (
     set_mla_kv_buffer_fp8_quant_kernel as set_mla_kv_buffer_fp8_quant_kernel,
 )

python/sglang/srt/model_executor/forward_batch_deepseek_mha_mixin.py

@@ -91,10 +91,8 @@ class ForwardBatchDeepSeekMHAMixin:
                 self.prefix_chunk_starts_cpu[idx],
                 self.prefix_chunk_seq_lens_cpu[idx],
             )
-            # None on a backend that never bound a translator.
-            src = get_attn_backend().kv_index_translator
-            if src is not None:
-                chunk_kv_indices = src.translate_full_attn_ids(chunk_kv_indices)
+            translator = get_attn_backend().kv_index_translator
+            chunk_kv_indices = translator.translate_dcp_read_ids(chunk_kv_indices)
             self.prefix_chunk_kv_indices.append(chunk_kv_indices)
 
     # Here we suppose the length of each chunk is equal

python/sglang/srt/model_executor/model_runner.py

@@ -1020,7 +1020,7 @@ class ModelRunner:
         self.attn_backend = backends.attn_backend
         self.decode_attn_backend = backends.decode_attn_backend
         self.decode_attn_backend_group = backends.decode_attn_backend_group
-        self.kv_index_translator.assert_backends_carry_translator(
+        self.kv_index_translator.bind_and_verify_backends(
             [self.attn_backend, self.decode_attn_backend]
         )
 

python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mha.py

@@ -499,6 +499,10 @@ class DeepseekMHAForwardMixin:
             # Without this, a chunked-prefill split (extend_prefix_lens != 0) that
             # reads cached prefix KV crashes with "576 != 656".
             kv_indices = filter_dcp_local_kv_indices(kv_indices=kv_indices)
+            # Read door: the pool never translates, so the production site does.
+            kv_indices = get_attn_backend().kv_index_translator.translate_dcp_read_ids(
+                kv_indices
+            )
             kv_a, k_pe = get_token_to_kv_pool().get_mla_kv_buffer(
                 self.attn_mha, kv_indices, torch.bfloat16
             )

python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mha_rocm.py

@@ -22,7 +22,10 @@ from sglang.srt.layers.quantization.fp8_utils import (
     materialize_bpreshuffle_fp8_scale_tuple,
 )
 from sglang.srt.model_executor.forward_batch_info import ForwardBatch
-from sglang.srt.model_executor.forward_context import get_token_to_kv_pool
+from sglang.srt.model_executor.forward_context import (
+    get_attn_backend,
+    get_token_to_kv_pool,
+)
 from sglang.srt.models.deepseek_common.attention_forward_methods.forward_mha import (
     forward_dsa_indexer_for_mha,
     resolve_attn_backend,
@@ -306,6 +309,10 @@ class DeepseekMHARocmForwardMixin:
     ):
         if _use_aiter_gfx95:
             kv_indices = filter_dcp_local_kv_indices(kv_indices=kv_indices)
+            # Read door: the pool never translates, so the production site does.
+            kv_indices = get_attn_backend().kv_index_translator.translate_dcp_read_ids(
+                kv_indices
+            )
             kv_a, k_pe = get_token_to_kv_pool().get_mla_kv_buffer(
                 self.attn_mha, kv_indices, dst_dtype
             )