[unified-memory] Support MLA-hybrid-Mamba (Kimi-Linear) on the Triton backend (#32971)

python/sglang/kernels/ops/attention/fla/chunk_delta_h.py

@@ -60,6 +60,7 @@ def chunk_gated_delta_rule_fwd_kernel_h_blockdim64(
     h,
     initial_state,
     initial_state_indices,
+    stride_init_state,
     cu_seqlens,
     chunk_offsets,
     T,
@@ -113,9 +114,13 @@ def chunk_gated_delta_rule_fwd_kernel_h_blockdim64(
     stride_k = Hg * K
     stride_w = H * K
 
-    index = tl.load(initial_state_indices + i_n).to(tl.int32)
-    h0 = initial_state + index * stride_h
-    ht = initial_state + index * stride_h
+    # Slot stride comes from the caller (initial_state.stride(0)): the state pool
+    # may be an envelope-strided view (page-major / unified memory), where the
+    # per-slot pitch spans ALL layers' state, not H*V*K. int64: envelope pitches
+    # overflow an int32 index product.
+    index = tl.load(initial_state_indices + i_n).to(tl.int64)
+    h0 = initial_state + index * stride_init_state
+    ht = initial_state + index * stride_init_state
     if USE_INITIAL_STATE:
         h0 = h0 + i_h * V * K
     if INPLACE_UPDATE:
@@ -355,6 +360,9 @@ def chunk_gated_delta_rule_fwd_h(
         h=h,
         initial_state=initial_state,
         initial_state_indices=initial_state_indices,
+        # Envelope-strided state pools (page-major / unified memory) have a
+        # per-slot pitch != H*V*K; contiguous pools pass exactly H*V*K.
+        stride_init_state=(initial_state.stride(0) if initial_state is not None else 0),
         cu_seqlens=cu_seqlens,
         chunk_offsets=chunk_offsets,
         T=T,

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

@@ -160,10 +160,13 @@ class TritonAttnBackend(AttentionBackend):
         # byte-identical to the slot-based envelope.
         self.page_size = getattr(model_runner, "page_size", 1) or 1
         # Unified pool v2p hook (None = no-op): req_to_token holds VIRTUAL ids but
-        # kernels need PHYSICAL. Applied eagerly so the captured graph has no translate.
+        # kernels need the kernel-facing id space — PHYSICAL for MHA, DENSE for the
+        # dense-view MLA pool (translate_kv_loc_dense falls back to the physical
+        # translate when kernel_page_multiplier == 1, so preferring it is exact for
+        # both). Applied eagerly so the captured graph has no translate.
         self._translate_kv_loc = getattr(
-            self.token_to_kv_pool_allocator, "translate_kv_loc", None
-        )
+            self.token_to_kv_pool_allocator, "translate_kv_loc_dense", None
+        ) or getattr(self.token_to_kv_pool_allocator, "translate_kv_loc", None)
         self.num_draft_tokens = model_runner.server_args.speculative_num_draft_tokens
         self.speculative_num_steps = model_runner.server_args.speculative_num_steps
         self.topk = model_runner.server_args.speculative_eagle_topk or 0
@@ -1243,6 +1246,9 @@ class TritonAttnBackend(AttentionBackend):
             cache_loc = forward_batch.out_cache_loc
             if isinstance(pool, SWAKVPool) and pool.layers_mapping[layer.layer_id][1]:
                 cache_loc = pool.translate_loc_from_full_to_swa(cache_loc)
+            elif self._translate_kv_loc is not None:
+                # Unified pool: buffers are indexed in the kernel-facing id space.
+                cache_loc = self._translate_kv_loc(cache_loc)
             k_buffer, v_buffer = pool.get_kv_buffer(layer.layer_id)
             k = k_buffer[cache_loc]
             v = v_buffer[cache_loc]
@@ -1710,7 +1716,15 @@ class TritonAttnBackend(AttentionBackend):
                     k.div_(layer.k_scale)
                 self.token_to_kv_pool.set_kv_buffer(
                     layer,
-                    forward_batch.out_cache_loc,
+                    # `full_loc` carries the pre-translated loc under the unified
+                    # pool, refreshed into a capture-stable buffer before replay —
+                    # translating inside set_kv_buffer would be captured and replay
+                    # a stale v2p. None (-> raw loc) for static pools.
+                    KVWriteLoc(
+                        forward_batch.out_cache_loc,
+                        self.forward_metadata.swa_out_cache_loc,
+                        full_loc=self.forward_metadata.out_cache_loc_full_physical,
+                    ),
                     k,
                     v,
                 )

python/sglang/srt/mem_cache/kv_cache_configurator.py

@@ -430,9 +430,6 @@ class KVCacheConfigurator:
 
         config = self.mambaish_config
         assert config is not None
-        assert (
-            not self.use_mla_backend
-        ), "unified memory pool does not support MLA-hybrid-Mamba yet"
         # The full sub-pool is page-aware (via `MultiEndedAllocator(page_size=...)`);
         # the mamba sub-pool stays page=1.
         assert self.page_size >= 1, f"page_size must be >= 1, got {self.page_size}"
@@ -462,6 +459,12 @@ class KVCacheConfigurator:
             end_layer=self.layer_info.end_layer,
             is_draft_worker=self.is_draft_worker,
             use_mla_backend=self.use_mla_backend,
+            kv_lora_rank=(
+                self.model_config.kv_lora_rank if self.use_mla_backend else None
+            ),
+            qk_rope_head_dim=(
+                self.model_config.qk_rope_head_dim if self.use_mla_backend else None
+            ),
             mamba_layer_ids=mamba_layer_ids,
             full_attention_layer_ids=full_attention_layer_ids,
             mamba2_cache_params=config.mamba2_cache_params,

python/sglang/srt/mem_cache/layout/page_major.py

@@ -108,6 +108,74 @@ def build_page_major_mha_views(
     return k_buffer, v_buffer
 
 
+def mla_entry_bytes(*, layer_num: int, kv_cache_dim: int, itemsize: int) -> int:
+    """Bytes occupied by one MLA slot across all layers (single latent row, no V)."""
+    return layer_num * kv_cache_dim * itemsize
+
+
+def build_dense_mla_views(
+    raw: torch.Tensor,
+    *,
+    layer_num: int,
+    kv_cache_dim: int,
+    store_dtype: torch.dtype,
+    page_size: int,
+    num_pages: int,
+    anchor_bytes: int = 0,
+) -> List[torch.Tensor]:
+    """Per-layer DENSE views over ``raw`` for MLA in the page-major layout.
+
+    The page envelope is ``[L0_latent * ps | L1_latent * ps | ...]``. Because all
+    MLA layers share one uniform row size (``kv_cache_dim``), the envelope is
+    itself a valid dense paged pool under a re-numbered index space: folding the
+    layer offset ``l * ps * kv_cache_dim`` into each view's storage_offset makes
+    every per-layer view a plain CONTIGUOUS ``(num_pages * layer_num * ps, 1,
+    kv_cache_dim)`` tensor, addressed by the layer-independent dense id
+
+        dense(t) = (t // ps) * (ps * layer_num) + t % ps      (t = physical token)
+
+    so one shared block table (entry = page * layer_num) serves every layer, and
+    kernels that require ``.view(-1, page_size, kv_cache_dim)`` (trtllm/cutlass/
+    flashmla) work on the views natively.
+
+    The views overlap each other (view ``l+1`` is view ``l`` shifted by ``ps``
+    rows); that is safe because layer ``l`` is only ever indexed at dense ids,
+    which always resolve to layer-``l`` bytes relative to view ``l``'s origin.
+    Layer ``layer_num-1``'s view extends ``(layer_num-1) * ps`` rows past the
+    last page envelope, so ``raw`` must carry at least one extra page envelope
+    of tail padding (``UnifiedKVPool``'s ``view_tail_pad_bytes``).
+    """
+    itemsize = store_dtype.itemsize
+    row_bytes = kv_cache_dim * itemsize
+    page_bytes = page_size * layer_num * row_bytes
+    n_dense = num_pages * layer_num * page_size
+    assert anchor_bytes % itemsize == 0
+    last_view_end = (
+        anchor_bytes + (layer_num - 1) * page_size * row_bytes + (n_dense * row_bytes)
+    )
+    assert last_view_end <= raw.numel() * raw.itemsize, (
+        f"build_dense_mla_views: layer {layer_num - 1}'s view ends at byte "
+        f"{last_view_end} but the raw buffer holds only "
+        f"{raw.numel() * raw.itemsize} bytes; allocate the tail pad "
+        f"(one page envelope = {page_bytes} B) via view_tail_pad_bytes"
+    )
+
+    as_dtype_view = raw.view(store_dtype)
+    views: List[torch.Tensor] = []
+    for layer in range(layer_num):
+        base_bytes = anchor_bytes + layer * page_size * row_bytes
+        assert base_bytes % itemsize == 0
+        views.append(
+            torch.as_strided(
+                as_dtype_view,
+                size=(n_dense, 1, kv_cache_dim),
+                stride=(kv_cache_dim, kv_cache_dim, 1),
+                storage_offset=base_bytes // itemsize,
+            )
+        )
+    return views
+
+
 def mamba_entry_bytes(
     *,
     layer_num: int,

python/sglang/srt/mem_cache/memory_pool.py

@@ -3551,6 +3551,10 @@ class HybridLinearKVPool(KVCache):
         # virtual->physical mamba-slot translate for the HiCache offload path;
         # identity for a static pool, the allocator's `translate` for the unified pool.
         self._mamba_translate = lambda ids: ids
+        # virtual->dense full-KV translate for the model-level MLA entry points
+        # (`set_mla_kv_buffer` / `get_mla_kv_buffer` receive VIRTUAL locs);
+        # identity for a static pool, `translate_kv_loc_dense` for the unified pool.
+        self._full_translate = lambda ids: ids
         self.use_mla = use_mla
         if full_kv_pool is not None:
             # Shared-KV-pool path: the caller built a UnifiedMHATokenToKVPool
@@ -3791,10 +3795,13 @@ class HybridLinearKVPool(KVCache):
                 dcp_kv_mask=dcp_kv_mask,
             )
         else:
+            # Mirror the MHA branch: `full_loc` is the unified pool's
+            # pre-translated (dense) loc; None for a static pool.
+            write_loc = full_loc if full_loc is not None else loc
             with self._transfer_id_context(layer):
                 self.full_kv_pool.set_kv_buffer(
                     layer,
-                    loc,
+                    write_loc,
                     cache_k,
                     cache_v,
                 )
@@ -3831,6 +3838,10 @@ class HybridLinearKVPool(KVCache):
         cache_k_rope: torch.Tensor,
     ):
         assert self.use_mla, "set_mla_kv_buffer called when use_mla is False"
+        # Model-level MLA entry point: `loc` is a VIRTUAL loc under the unified
+        # pool (eager prefill only; the decode write goes through set_kv_buffer's
+        # pre-translated `full_loc`), so translate to the dense id space here.
+        loc = self._full_translate(loc)
         with self._transfer_id_context(layer):
             self.full_kv_pool.set_mla_kv_buffer(layer, loc, cache_k_nope, cache_k_rope)
 
@@ -3841,6 +3852,7 @@ class HybridLinearKVPool(KVCache):
         dst_dtype: Optional[torch.dtype] = None,
     ):
         assert self.use_mla, "get_mla_kv_buffer called when use_mla is False"
+        loc = self._full_translate(loc)
         with self._transfer_id_context(layer):
             return self.full_kv_pool.get_mla_kv_buffer(layer, loc, dst_dtype)
 

python/sglang/srt/mem_cache/multi_ended_allocator.py

@@ -111,6 +111,7 @@ class MultiEndedAllocator(BaseTokenToKVPoolAllocator):
         need_sort: bool = False,
         forward_stream: Optional[torch.cuda.Stream] = None,
         lazy_compaction: bool = False,
+        kernel_page_multiplier: int = 1,
     ):
         spec = unified_buffer.spec(sub_pool_name)
         max_slots = unified_buffer.max_slots(sub_pool_name)
@@ -130,6 +131,11 @@ class MultiEndedAllocator(BaseTokenToKVPoolAllocator):
         self.entry_bytes = spec.entry_bytes()
         self.min_slot_index = unified_buffer.min_slot_index(sub_pool_name)
         self.is_id_owner = is_id_owner
+        # Dense (kernel-facing) index space scale: the page-major envelope of a
+        # multi-layer uniform-entry sub-pool (MLA) is a valid dense paged pool
+        # once page ids are scaled by layer_num — `translate_kv_loc_dense` emits
+        # that space. 1 for sub-pools whose kernels take real physical ids.
+        self.kernel_page_multiplier = kernel_page_multiplier
         # Overlap mode: `free` drops a wait_stream(forward_stream) barrier so its
         # v2p writes + move kernel serialize after the in-flight forward.
         self.forward_stream = forward_stream
@@ -681,6 +687,59 @@ class MultiEndedAllocator(BaseTokenToKVPoolAllocator):
         result = phys_pages * self.page_size + offsets
         return torch.clamp_min(result, 0)
 
+    def translate_kv_loc_dense(
+        self,
+        virt_tokens: torch.Tensor,
+        *,
+        out: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        """Translate virtual token ids to DENSE (kernel-facing) ids.
+
+        dense(t) = (t // ps) * (ps * kernel_page_multiplier) + t % ps for the
+        physical token t — i.e. `translate_kv_loc` with the page stride scaled by
+        `kernel_page_multiplier` (= layer_num for a dense-view MLA sub-pool; see
+        `build_dense_mla_views`). Internal machinery (compaction, in-flight write
+        sets) MUST keep using `translate_kv_loc`: dense ids are for kernels only.
+
+        The tombstone clamp routes -1 entries to dense id 0 — inside the page-0
+        reserved sink for every layer view. Supports ``out=`` like
+        `translate_kv_loc` for cuda-graph buffer stability.
+        """
+        if self.kernel_page_multiplier == 1:
+            return self.translate_kv_loc(virt_tokens, out=out)
+        if out is not None:
+            assert out.dtype == torch.int64, (
+                f"translate_kv_loc_dense: out= dtype must be int64 (matches v2p), "
+                f"got {out.dtype}"
+            )
+            assert out.shape == virt_tokens.shape, (
+                f"translate_kv_loc_dense: out= shape {tuple(out.shape)} must "
+                f"match virt_tokens shape {tuple(virt_tokens.shape)}"
+            )
+        with record_function("MultiEndedAlloc.translate_kv_loc_dense"):
+            dense_page_stride = self.page_size * self.kernel_page_multiplier
+            if self.page_size == 1:
+                # dense = phys * multiplier; tombstone -1 scales negative → clamp 0.
+                if out is not None:
+                    tmp = torch.index_select(self.virtual_to_physical, 0, virt_tokens)
+                    tmp = torch.clamp_min(tmp * dense_page_stride, 0)
+                    out.copy_(tmp)
+                    return out
+                result = torch.index_select(self.virtual_to_physical, 0, virt_tokens)
+                return torch.clamp_min(result * dense_page_stride, 0)
+            virt_pages = virt_tokens // self.page_size
+            offsets = virt_tokens % self.page_size
+            if out is not None:
+                torch.index_select(self.virtual_to_physical, 0, virt_pages, out=out)
+                out.mul_(dense_page_stride)
+                out.add_(offsets)
+                # tombstoned page: -1*dense_page_stride + offset < 0
+                out.clamp_(min=0)
+                return out
+            phys_pages = self.virtual_to_physical[virt_pages]
+            result = phys_pages * dense_page_stride + offsets
+            return torch.clamp_min(result, 0)
+
     # -- alloc --
 
     def alloc(self, need_size: int) -> Optional[torch.Tensor]:
@@ -1654,12 +1713,13 @@ class UnifiedMambaTokenToKVPoolAllocator(BaseTokenToKVPoolAllocator):
         need_sort: bool = False,
         forward_stream: Optional[torch.cuda.Stream] = None,
         lazy_compaction: bool = False,
+        full_kernel_page_multiplier: int = 1,
     ):
         full_max = unified_buffer.max_slots("full")
         super().__init__(
             size=full_max - 1,
             page_size=page_size,
-            dtype=unified_buffer.mha_spec("full").store_dtype,
+            dtype=unified_buffer.spec("full").get_dtype(),
             device=device,
             kvcache=kvcache,
             need_sort=need_sort,
@@ -1681,6 +1741,7 @@ class UnifiedMambaTokenToKVPoolAllocator(BaseTokenToKVPoolAllocator):
             need_sort=need_sort,
             forward_stream=forward_stream,
             lazy_compaction=lazy_compaction,
+            kernel_page_multiplier=full_kernel_page_multiplier,
         )
         self.mamba_allocator = MultiEndedAllocator(
             kvcache=kvcache.mamba_pool,
@@ -1833,6 +1894,20 @@ class UnifiedMambaTokenToKVPoolAllocator(BaseTokenToKVPoolAllocator):
         result = self.full_attn_allocator.translate_kv_loc(loc, out=out)
         return result
 
+    @property
+    def kernel_page_multiplier(self) -> int:
+        return self.full_attn_allocator.kernel_page_multiplier
+
+    def translate_kv_loc_dense(
+        self,
+        loc: torch.Tensor,
+        *,
+        out: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        """Full-pool virtual TOKEN ids -> DENSE (kernel-facing) ids. Falls back
+        to the physical translate when `kernel_page_multiplier == 1` (MHA)."""
+        return self.full_attn_allocator.translate_kv_loc_dense(loc, out=out)
+
     def is_slot_allocated(self, slot: int) -> bool:
         return self.full_attn_allocator.is_slot_allocated(slot)
 

python/sglang/srt/mem_cache/unified_memory_pool.py

@@ -35,6 +35,7 @@ from torch.profiler import record_function
 from sglang.kernels.ops.kvcache.cache_move import store_cache_4d_kernel
 from sglang.srt.constants import GPU_MEMORY_TYPE_KV_CACHE
 from sglang.srt.mem_cache.layout.page_major import (
+    build_dense_mla_views,
     build_page_major_mamba_views,
     build_page_major_mha_views,
 )
@@ -42,6 +43,7 @@ from sglang.srt.mem_cache.memory_pool import (
     HybridReqToTokenPool,
     MambaPool,
     MHATokenToKVPool,
+    MLATokenToKVPool,
     move_kv_cache_native,
     unwrap_write_loc,
 )
@@ -139,6 +141,40 @@ class MHASubPoolSpec(SubPoolSpec):
         return self.store_dtype
 
 
+@dataclass(frozen=True, kw_only=True)
+class MLASubPoolSpec(SubPoolSpec):
+    """Per-slot layout of one MLA-shaped sub-pool.
+
+    One latent row (``kv_lora_rank + qk_rope_head_dim``) per token per layer; V
+    is a prefix slice of the same row, so there is no separate V region. Not a
+    subclass of ``MHASubPoolSpec`` — the K+V byte math and the ``v_head_dim > 0``
+    invariant there do not apply.
+    """
+
+    kv_lora_rank: int
+    qk_rope_head_dim: int
+    store_dtype: torch.dtype
+
+    def __post_init__(self):
+        super().__post_init__()
+        assert (
+            self.kv_lora_rank > 0
+        ), f"kv_lora_rank must be positive; got {self.kv_lora_rank}"
+        assert (
+            self.qk_rope_head_dim > 0
+        ), f"qk_rope_head_dim must be positive; got {self.qk_rope_head_dim}"
+
+    @property
+    def kv_cache_dim(self) -> int:
+        return self.kv_lora_rank + self.qk_rope_head_dim
+
+    def entry_bytes(self) -> int:
+        return self.layer_num * self.kv_cache_dim * self.store_dtype.itemsize
+
+    def get_dtype(self) -> torch.dtype:
+        return self.store_dtype
+
+
 @dataclass(frozen=True, kw_only=True)
 class MambaSubPoolSpec(SubPoolSpec):
     """Per-slot layout of one Mamba-shaped sub-pool."""
@@ -188,6 +224,7 @@ class UnifiedKVPool:
         device: str,
         enable_memory_saver: bool,
         page_size: int = 1,
+        view_tail_pad_bytes: int = 0,
     ):
         assert page_size >= 1, f"page_size must be >= 1; got {page_size}"
         assert len(sub_pool_specs) == 2, (
@@ -213,25 +250,44 @@ class UnifiedKVPool:
         self.memory_saver_adapter = TorchMemorySaverAdapter.create(
             enable=enable_memory_saver
         )
+        # `view_tail_pad_bytes` extends the ALLOCATION only (dense MLA views are
+        # per-layer shifted, so the last layer's view reaches past the final page
+        # envelope); all slot/watermark math stays on the unpadded `total_bytes`.
+        self.view_tail_pad_bytes = view_tail_pad_bytes
         with self.memory_saver_adapter.region(GPU_MEMORY_TYPE_KV_CACHE):
-            self._raw = torch.empty(total_bytes, dtype=torch.uint8, device=device)
+            self._raw = torch.empty(
+                total_bytes + view_tail_pad_bytes, dtype=torch.uint8, device=device
+            )
         self._raw.zero_()  # unset slots must read as zeros (matches non-shared)
 
         self._max_slots: Dict[str, int] = {}
         self._anchor_bytes: Dict[str, int] = {}
         self._min_slot_index: Dict[str, int] = {}
-        # MHA: (k_buffer, v_buffer); Mamba: (conv_state_list, temporal_state)
+        # MHA: (k_buffer, v_buffer); MLA: [per-layer dense views];
+        # Mamba: (conv_state_list, temporal_state)
         self._mha_views: Dict[str, Tuple[List[torch.Tensor], List[torch.Tensor]]] = {}
+        self._mla_views: Dict[str, List[torch.Tensor]] = {}
         self._mamba_views: Dict[str, Tuple[List[torch.Tensor], torch.Tensor]] = {}
 
-        # Slot-0 dummy writes for both pools land in [0, entry_max); each pool's
-        # first allocatable slot is chosen so real data starts at >= entry_max.
+        # Slot-0 dummy writes for both pools land in the reserved low-byte sink;
+        # each pool's first allocatable slot is chosen so real data starts past it.
+        # For a page-aware sub-pool the slot-0 write touches layer blocks spread
+        # across the WHOLE page-0 envelope (up to page_size * entry_bytes), not
+        # just one slot envelope — reserve the max of both.
         entry_max = max(s.entry_bytes() for s in sub_pool_specs)
+        reserved_floor = max(
+            [entry_max]
+            + [
+                page_size * s.entry_bytes()
+                for s in sub_pool_specs
+                if not isinstance(s, MambaSubPoolSpec)  # mamba is page_size=1
+            ]
+        )
 
         for spec in sub_pool_specs:
             entry_bytes = spec.entry_bytes()
             max_slots = total_bytes // entry_bytes
-            min_slot_index = (entry_max + entry_bytes - 1) // entry_bytes  # ceil
+            min_slot_index = (reserved_floor + entry_bytes - 1) // entry_bytes  # ceil
             if max_slots <= min_slot_index:
                 raise RuntimeError(
                     f"UnifiedKVPool: sub-pool {spec.name!r} fits only {max_slots} "
@@ -249,6 +305,13 @@ class UnifiedKVPool:
                     max_slots,
                     page_size=page_size,
                 )
+            elif isinstance(spec, MLASubPoolSpec):
+                self._mla_views[spec.name] = self._build_mla_views(
+                    spec,
+                    anchor,
+                    max_slots,
+                    page_size=page_size,
+                )
             elif isinstance(spec, MambaSubPoolSpec):
                 self._mamba_views[spec.name] = self._build_mamba_views(
                     spec, anchor, max_slots
@@ -289,6 +352,13 @@ class UnifiedKVPool:
         ), f"sub-pool {name!r} is {type(s).__name__}, expected MHASubPoolSpec"
         return s
 
+    def mla_spec(self, name: str) -> MLASubPoolSpec:
+        s = self._specs_by_name[name]
+        assert isinstance(
+            s, MLASubPoolSpec
+        ), f"sub-pool {name!r} is {type(s).__name__}, expected MLASubPoolSpec"
+        return s
+
     def mamba_spec(self, name: str) -> MambaSubPoolSpec:
         s = self._specs_by_name[name]
         assert isinstance(
@@ -310,6 +380,9 @@ class UnifiedKVPool:
     def mha_views_for(self, name: str) -> Tuple[List[torch.Tensor], List[torch.Tensor]]:
         return self._mha_views[name]
 
+    def mla_views_for(self, name: str) -> List[torch.Tensor]:
+        return self._mla_views[name]
+
     def mamba_views_for(self, name: str) -> Tuple[List[torch.Tensor], torch.Tensor]:
         return self._mamba_views[name]
 
@@ -332,6 +405,23 @@ class UnifiedKVPool:
             anchor_bytes=anchor_bytes,
         )
 
+    def _build_mla_views(
+        self,
+        spec: MLASubPoolSpec,
+        anchor_bytes: int,
+        max_slots: int,
+        page_size: int,
+    ) -> List[torch.Tensor]:
+        return build_dense_mla_views(
+            self._raw,
+            layer_num=spec.layer_num,
+            kv_cache_dim=spec.kv_cache_dim,
+            store_dtype=spec.store_dtype,
+            page_size=page_size,
+            num_pages=max_slots // page_size,
+            anchor_bytes=anchor_bytes,
+        )
+
     def _build_mamba_views(
         self, spec: MambaSubPoolSpec, anchor_bytes: int, max_slots: int
     ) -> Tuple[List[torch.Tensor], torch.Tensor]:
@@ -497,6 +587,91 @@ class UnifiedMHATokenToKVPool(MHATokenToKVPool):
             )
 
 
+class UnifiedMLATokenToKVPool(MLATokenToKVPool):
+    """MLA KV pool whose per-layer `kv_buffer` entries are DENSE views into a
+    `UnifiedKVPool` (see `build_dense_mla_views`).
+
+    Loc-space contract: every loc this pool receives through the KVCache API
+    (`set_kv_buffer` / `set_mla_kv_buffer` / `get_mla_kv_buffer`, and the
+    kv_indices consumed by attention kernels reading `get_key_buffer` /
+    `get_value_buffer`) is a DENSE id — the `translate_kv_loc_dense` output
+
+        dense(t) = (t // ps) * (ps * layer_num) + t % ps
+
+    which is layer-independent (the layer offset is folded into each view's
+    storage_offset), so the stock `MLATokenToKVPool` read/write methods work on
+    the views unmodified. The ONE exception is `move_kv_cache`: the allocator's
+    compaction calls it with REAL physical token ids, and it is overridden to
+    relocate whole page envelopes on the raw buffer.
+    """
+
+    def __init__(
+        self,
+        *,
+        unified_buffer: UnifiedKVPool,
+        sub_pool_name: str,
+        kv_cache_dtype: torch.dtype,
+        page_size: int = 1,
+    ):
+        spec = unified_buffer.mla_spec(sub_pool_name)
+        store_dtype = _store_dtype_for(kv_cache_dtype)
+        assert spec.store_dtype == store_dtype, (
+            f"sub-pool {sub_pool_name!r} store dtype {spec.store_dtype} does not "
+            f"match kv cache dtype {kv_cache_dtype} (store {store_dtype})"
+        )
+
+        self._unified_buffer = unified_buffer
+        self._sub_pool_name = sub_pool_name
+        self._kv_views = unified_buffer.mla_views_for(sub_pool_name)
+        max_slots = unified_buffer.max_slots(sub_pool_name)
+        self._num_pages = max_slots // page_size
+        self._page_bytes = page_size * spec.entry_bytes()
+        # Dense row count per view; also the OOB bound for dense locs.
+        self._dense_size = self._num_pages * spec.layer_num * page_size
+
+        super().__init__(
+            # OOB checks bound locs by `size + page_size`; dense ids run to
+            # `_dense_size` (page 0 is the reserved padding sink).
+            size=self._dense_size - page_size,
+            page_size=page_size,
+            dtype=kv_cache_dtype,
+            kv_lora_rank=spec.kv_lora_rank,
+            qk_rope_head_dim=spec.qk_rope_head_dim,
+            layer_num=spec.layer_num,
+            device=unified_buffer.device,
+            enable_memory_saver=False,  # buffer owned by UnifiedKVPool
+        )
+
+    def _create_buffers(self):
+        self.kv_buffer = self._kv_views
+
+    def _clear_buffers(self):
+        # Lifetime owned by UnifiedKVPool; do not delete the views.
+        pass
+
+    def get_kv_size_bytes(self):
+        return 0  # UnifiedKVPool logs the total; per-sub-pool would double-count
+
+    def move_kv_cache(self, tgt_loc: torch.Tensor, src_loc: torch.Tensor):
+        """Relocate whole page envelopes.
+
+        `tgt_loc`/`src_loc` are REAL physical token ids (NOT dense ids): both
+        compaction paths expand page ids into page-major-ordered token runs
+        (`pages[:, None] * ps + offsets`), relied on here to recover the page
+        lists. One contiguous envelope copy replaces the per-layer strided moves.
+        """
+        if tgt_loc.numel() == 0:
+            return
+        ps = self.page_size
+        tgt_pages = tgt_loc.view(-1, ps)[:, 0] // ps
+        src_pages = src_loc.view(-1, ps)[:, 0] // ps
+        with record_function("UnifiedMLA.move_kv_cache"):
+            env = self._unified_buffer._raw[: self._num_pages * self._page_bytes].view(
+                self._num_pages, self._page_bytes
+            )
+            env[tgt_pages] = env[src_pages]
+
+
 class UnifiedMambaPool(MambaPool):
     """Mamba state pool whose conv/temporal state are strided views into a `UnifiedKVPool`.
 
@@ -766,10 +941,12 @@ class UnifiedHybridReqToTokenPool(HybridReqToTokenPool):
         mamba_envelope_layout: bool = False,
         enable_linear_replayssm: bool = False,
         linear_replayssm_cache_len: int = 16,
+        enable_gdn_replayssm_spec: bool = False,
     ):
         # mamba_envelope_layout / speculative_eagle_topk / enable_linear_replayssm /
-        # linear_replayssm_cache_len: accepted to match the parent signature but NOT
-        # forwarded — the shared pool's conv/temporal state are fixed-shape views.
+        # linear_replayssm_cache_len / enable_gdn_replayssm_spec: accepted to match
+        # the parent signature but NOT forwarded — the shared pool's conv/temporal
+        # state are fixed-shape views (replayssm/spec are gated off under unified).
         assert mamba_size == self._shared_mamba_size, (
             f"UnifiedHybridReqToTokenPool._init_mamba_pool: mamba_size={mamba_size} "
             f"!= unified_buffer.max_slots({self._mamba_sub_pool_name!r}) - 1 "
@@ -837,6 +1014,8 @@ def init_unified_mamba_pools(
     end_layer: int,
     is_draft_worker: bool,
     use_mla_backend: bool,
+    kv_lora_rank: Optional[int] = None,
+    qk_rope_head_dim: Optional[int] = None,
     mamba_layer_ids: List[int],
     full_attention_layer_ids: List[int],
     mamba2_cache_params,
@@ -860,22 +1039,37 @@ def init_unified_mamba_pools(
         UnifiedMambaTokenToKVPoolAllocator,
     )
 
-    assert (
-        not use_mla_backend
-    ), "unified memory pool does not support MLA-hybrid-Mamba yet"
     # Full sub-pool is page-aware; mamba stays page=1 (state is per-request).
     assert page_size >= 1, f"page_size must be >= 1, got {page_size}"
 
     store_dtype = _store_dtype_for(kv_cache_dtype)
     # full-attn at the high-byte end (grow-down), mamba at the low-byte end (grow-up).
-    full_spec = MHASubPoolSpec(
-        name="full",
-        layer_num=len(full_attention_layer_ids),
-        head_num=head_num,
-        head_dim=head_dim,
-        store_dtype=store_dtype,
-        grow_direction="down",
-    )
+    if use_mla_backend:
+        assert kv_lora_rank and qk_rope_head_dim, (
+            "init_unified_mamba_pools: MLA-hybrid-Mamba needs kv_lora_rank and "
+            f"qk_rope_head_dim; got {kv_lora_rank} / {qk_rope_head_dim}"
+        )
+        assert not is_draft_worker, (
+            "init_unified_mamba_pools: draft workers (speculative decoding) are "
+            "not supported with the MLA unified pool"
+        )
+        full_spec = MLASubPoolSpec(
+            name="full",
+            layer_num=len(full_attention_layer_ids),
+            kv_lora_rank=kv_lora_rank,
+            qk_rope_head_dim=qk_rope_head_dim,
+            store_dtype=store_dtype,
+            grow_direction="down",
+        )
+    else:
+        full_spec = MHASubPoolSpec(
+            name="full",
+            layer_num=len(full_attention_layer_ids),
+            head_num=head_num,
+            head_dim=head_dim,
+            store_dtype=store_dtype,
+            grow_direction="down",
+        )
     cp = mamba2_cache_params
     mamba_spec = MambaSubPoolSpec(
         name="mamba",
@@ -891,12 +1085,16 @@ def init_unified_mamba_pools(
         max_total_num_tokens * full_spec.entry_bytes()
         + max_mamba_cache_size * mamba_spec.entry_bytes()
     )
+    # Dense MLA views are per-layer shifted, so the last layer's view reaches one
+    # page envelope past the final page — allocation-only tail pad (~page bytes).
+    view_tail_pad_bytes = page_size * full_spec.entry_bytes() if use_mla_backend else 0
     shared_pool = UnifiedKVPool(
         total_bytes=total_bytes,
         sub_pool_specs=[full_spec, mamba_spec],
         device=device,
         enable_memory_saver=enable_memory_saver,
         page_size=page_size,
+        view_tail_pad_bytes=view_tail_pad_bytes,
     )
     req_to_token_pool = UnifiedHybridReqToTokenPool(
         unified_buffer=shared_pool,
@@ -913,13 +1111,23 @@ def init_unified_mamba_pools(
         enable_overlap_schedule=not disable_overlap_schedule,
         start_layer=start_layer,
     )
-    unified_full_kv_pool = UnifiedMHATokenToKVPool(
-        unified_buffer=shared_pool,
-        sub_pool_name="full",
-        page_size=page_size,
-        start_layer=start_layer,
-        end_layer=end_layer,
-    )
+    if use_mla_backend:
+        # start_layer stays 0: HybridLinearKVPool patches layer ids to the dense
+        # 0..N-1 index via _transfer_id_context before every MLA pool call.
+        unified_full_kv_pool = UnifiedMLATokenToKVPool(
+            unified_buffer=shared_pool,
+            sub_pool_name="full",
+            kv_cache_dtype=kv_cache_dtype,
+            page_size=page_size,
+        )
+    else:
+        unified_full_kv_pool = UnifiedMHATokenToKVPool(
+            unified_buffer=shared_pool,
+            sub_pool_name="full",
+            page_size=page_size,
+            start_layer=start_layer,
+            end_layer=end_layer,
+        )
     full_attn_layer_ids_for_pool = (
         [0] if is_draft_worker else list(full_attention_layer_ids)
     )
@@ -945,6 +1153,9 @@ def init_unified_mamba_pools(
         need_sort=need_sort,
         forward_stream=forward_stream,
         lazy_compaction=lazy_compaction,
+        full_kernel_page_multiplier=(
+            len(full_attention_layer_ids) if use_mla_backend else 1
+        ),
     )
 
     # Wrap the composite's mamba MultiEndedAllocator in a slot allocator (PHYSICAL view).
@@ -956,23 +1167,43 @@ def init_unified_mamba_pools(
     # `_mamba_translate` feeds the HiCache offload path, GATED OFF here — wired but inert.
     req_to_token_pool.mamba_allocator = mamba_slot_allocator
     token_to_kv_pool._mamba_translate = mamba_slot_allocator.translate
+    if use_mla_backend:
+        # Model-level MLA entry points (`set_mla_kv_buffer` / `get_mla_kv_buffer`)
+        # receive VIRTUAL locs and translate to the dense space internally
+        # (eager-prefill-only paths; never captured in a cuda graph).
+        token_to_kv_pool._full_translate = allocator.translate_kv_loc_dense
 
     logger.info(
         "[unified-memory-pool] ============================================================"
     )
     logger.info(
-        "[unified-memory-pool] UNIFIED MEMORY POOL ENABLED -- path=Mamba hybrid"
-    )
-    logger.info(
-        "[unified-memory-pool]   full_layers=%d, mamba_layers=%d, head_num=%d, head_dim=%d, "
-        "page_size=%d, is_draft_worker=%s",
-        len(full_attention_layer_ids),
-        len(mamba_layer_ids),
-        head_num,
-        head_dim,
-        page_size,
-        is_draft_worker,
+        "[unified-memory-pool] UNIFIED MEMORY POOL ENABLED -- path=Mamba hybrid (%s full side)",
+        "MLA" if use_mla_backend else "MHA",
     )
+    if use_mla_backend:
+        logger.info(
+            "[unified-memory-pool]   full_layers=%d, mamba_layers=%d, kv_lora_rank=%d, "
+            "qk_rope_head_dim=%d, page_size=%d (dense views, kernel_page_multiplier=%d, "
+            "view_tail_pad=%d B)",
+            len(full_attention_layer_ids),
+            len(mamba_layer_ids),
+            kv_lora_rank,
+            qk_rope_head_dim,
+            page_size,
+            len(full_attention_layer_ids),
+            view_tail_pad_bytes,
+        )
+    else:
+        logger.info(
+            "[unified-memory-pool]   full_layers=%d, mamba_layers=%d, head_num=%d, head_dim=%d, "
+            "page_size=%d, is_draft_worker=%s",
+            len(full_attention_layer_ids),
+            len(mamba_layer_ids),
+            head_num,
+            head_dim,
+            page_size,
+            is_draft_worker,
+        )
     logger.info(
         "[unified-memory-pool]   total_bytes=%d, max_total_num_tokens=%d, max_mamba_cache_size=%d, "
         "max_num_reqs=%d, speculative_num_draft_tokens=%s",

test/registered/unit/mem_cache/test_full_loc_fast_path.py

@@ -202,5 +202,115 @@ class TestHybridLinearFullLocRouting(unittest.TestCase):
         self.assertNotIn("already_physical", kwargs)
 
 
+class _RecordingMLAPool(_RecordingPool):
+    """Also records the model-level MLA entry points."""
+
+    def __init__(self):
+        super().__init__()
+        self.mla_set_calls = []
+        self.mla_get_calls = []
+
+    def set_mla_kv_buffer(self, layer, loc, cache_k_nope, cache_k_rope):
+        self.mla_set_calls.append(loc)
+
+    def get_mla_kv_buffer(self, layer, loc, dst_dtype=None):
+        self.mla_get_calls.append(loc)
+        return None, None
+
+
+class TestHybridLinearMLARouting(unittest.TestCase):
+    """MLA-side routing contracts of `HybridLinearKVPool`:
+
+    - `set_kv_buffer` (MLA branch) mirrors the MHA branch — write the
+      pre-translated `KVWriteLoc.full_loc` when present (unified pool, where it
+      carries the DENSE loc), else the raw `loc` (static pool, already physical).
+    - `set_mla_kv_buffer` / `get_mla_kv_buffer` receive VIRTUAL locs and apply
+      `_full_translate` exactly once (identity for a static pool)."""
+
+    def _make_bare_pool(self, translate=None):
+        from sglang.srt.mem_cache.memory_pool import HybridLinearKVPool
+
+        pool = object.__new__(HybridLinearKVPool)
+        pool.full_kv_pool = _RecordingMLAPool()
+        pool.use_mla = True
+        pool.full_attention_layer_id_mapping = {0: 0}
+        pool._full_translate = translate if translate is not None else (lambda x: x)
+        return pool
+
+    def test_mla_writes_full_loc_from_write_loc(self):
+        pool = self._make_bare_pool()
+        virtual_loc = torch.tensor([7, 8, 9], dtype=torch.int64)
+        dense_phys = torch.tensor([21, 24, 27], dtype=torch.int64)
+
+        layer = types.SimpleNamespace(layer_id=0)
+        pool.set_kv_buffer(
+            layer,
+            _loc_info(virtual_loc, full_phys=dense_phys),
+            torch.zeros(3, 1, 8),
+            None,
+        )
+
+        self.assertEqual(len(pool.full_kv_pool.calls), 1)
+        forwarded, _ = pool.full_kv_pool.calls[0]
+        self.assertIs(forwarded, dense_phys)
+        self.assertIsNot(forwarded, virtual_loc)
+
+    def test_mla_falls_back_to_loc_when_absent(self):
+        pool = self._make_bare_pool()
+        phys_loc = torch.tensor([7, 8, 9], dtype=torch.int64)
+
+        layer = types.SimpleNamespace(layer_id=0)
+        pool.set_kv_buffer(
+            layer,
+            _loc_info(phys_loc),
+            torch.zeros(3, 1, 8),
+            None,
+        )
+
+        self.assertEqual(len(pool.full_kv_pool.calls), 1)
+        forwarded, _ = pool.full_kv_pool.calls[0]
+        self.assertIs(forwarded, phys_loc)
+
+    def test_set_mla_kv_buffer_translates_exactly_once(self):
+        calls = []
+
+        def translate(ids):
+            calls.append(ids)
+            return ids + 100
+
+        pool = self._make_bare_pool(translate=translate)
+        virtual_loc = torch.tensor([7, 8, 9], dtype=torch.int64)
+        layer = types.SimpleNamespace(layer_id=0)
+
+        pool.set_mla_kv_buffer(
+            layer, virtual_loc, torch.zeros(3, 1, 6), torch.zeros(3, 1, 2)
+        )
+
+        self.assertEqual(len(calls), 1)
+        self.assertEqual(len(pool.full_kv_pool.mla_set_calls), 1)
+        self.assertTrue(
+            torch.all(pool.full_kv_pool.mla_set_calls[0] == virtual_loc + 100)
+        )
+
+    def test_get_mla_kv_buffer_translates_exactly_once(self):
+        calls = []
+
+        def translate(ids):
+            calls.append(ids)
+            return ids + 100
+
+        pool = self._make_bare_pool(translate=translate)
+        virtual_loc = torch.tensor([4, 5], dtype=torch.int64)
+        layer = types.SimpleNamespace(layer_id=0)
+
+        pool.get_mla_kv_buffer(layer, virtual_loc)
+
+        self.assertEqual(len(calls), 1)
+        self.assertEqual(len(pool.full_kv_pool.mla_get_calls), 1)
+        self.assertTrue(
+            torch.all(pool.full_kv_pool.mla_get_calls[0] == virtual_loc + 100)
+        )
+
+
 if __name__ == "__main__":
     unittest.main()

test/registered/unit/mem_cache/test_unified_mla_gpu_parity.py

@@ -0,0 +1,203 @@
+# Copyright 2023-2026 SGLang Team
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""GPU parity of the dense-view `UnifiedMLATokenToKVPool` against the stock
+`MLATokenToKVPool` on real K3 MLA geometry (L=24, D=512+64).
+
+The unified pool receives DENSE locs (dense(t) = (t//ps)*(ps*L) + t%ps); the
+reference pool receives the raw token ids. Every (layer, token) cell must hold
+identical bytes afterwards. Covers:
+
+  - `set_mla_kv_buffer` under BOTH kernel paths — the Triton fallback
+    (n_loc < 768) and the TMA JIT fast path (n_loc >= 768, which flattens the
+    buffer via `.view(shape[0], -1)`, only legal because dense views are
+    contiguous);
+  - `set_kv_buffer` (combined pre-concatenated write, the Triton-backend path);
+  - `get_mla_kv_buffer` roundtrip;
+  - page_size 1 and 64.
+
+    python -m pytest test/registered/unit/mem_cache/test_unified_mla_gpu_parity.py -v
+"""
+
+import types
+import unittest
+
+import torch
+
+from sglang.test.ci.ci_register import register_cuda_ci
+
+register_cuda_ci(est_time=60, stage="base-b", runner_config="1-gpu-small")
+
+_HAS_CUDA = torch.cuda.is_available()
+_DEV = "cuda"
+
+_L = 24
+_LORA = 512
+_ROPE = 64
+_D = _LORA + _ROPE
+_DTYPE = torch.bfloat16
+
+
+def _dense(t: torch.Tensor, ps: int) -> torch.Tensor:
+    return (t // ps) * (ps * _L) + t % ps
+
+
+def _make_pools(ps: int, n_tokens: int = 4096):
+    from sglang.srt.mem_cache.memory_pool import MLATokenToKVPool
+    from sglang.srt.mem_cache.unified_memory_pool import (
+        MambaSubPoolSpec,
+        MLASubPoolSpec,
+        UnifiedKVPool,
+        UnifiedMLATokenToKVPool,
+    )
+
+    full = MLASubPoolSpec(
+        name="full",
+        layer_num=_L,
+        kv_lora_rank=_LORA,
+        qk_rope_head_dim=_ROPE,
+        store_dtype=_DTYPE,
+        grow_direction="down",
+    )
+    mamba = MambaSubPoolSpec(
+        name="mamba",
+        layer_num=2,
+        conv_state_shapes=((8, 16),),
+        conv_dtype=torch.bfloat16,
+        temporal_state_shape=(4, 8, 8),
+        temporal_dtype=torch.float32,
+        grow_direction="up",
+    )
+    total = full.entry_bytes() * n_tokens + mamba.entry_bytes() * 16
+    pool = UnifiedKVPool(
+        total_bytes=total,
+        sub_pool_specs=[full, mamba],
+        device=_DEV,
+        enable_memory_saver=False,
+        page_size=ps,
+        view_tail_pad_bytes=ps * full.entry_bytes(),
+    )
+    unified = UnifiedMLATokenToKVPool(
+        unified_buffer=pool,
+        sub_pool_name="full",
+        kv_cache_dtype=_DTYPE,
+        page_size=ps,
+    )
+    max_tokens = pool.max_slots("full")
+    ref = MLATokenToKVPool(
+        size=max_tokens - ps,
+        page_size=ps,
+        dtype=_DTYPE,
+        kv_lora_rank=_LORA,
+        qk_rope_head_dim=_ROPE,
+        layer_num=_L,
+        device=_DEV,
+        enable_memory_saver=False,
+    )
+    return unified, ref, max_tokens
+
+
+def _rand_locs(max_tokens: int, ps: int, n: int) -> torch.Tensor:
+    # distinct physical token ids clear of the reserved page 0
+    g = torch.Generator(device="cpu").manual_seed(1234 + n + ps)
+    perm = torch.randperm(max_tokens - ps, generator=g)[:n] + ps
+    return perm.to(_DEV)
+
+
+@unittest.skipUnless(_HAS_CUDA, "requires CUDA")
+class TestUnifiedMLAPoolGPUParity(unittest.TestCase):
+    def _assert_parity(self, unified, ref, locs, ps, layers=range(_L)):
+        for l in layers:
+            got = unified.get_key_buffer(l)[_dense(locs, ps)]
+            want = ref.get_key_buffer(l)[locs]
+            torch.testing.assert_close(got, want, rtol=0, atol=0)
+
+    def _run_set_mla(self, ps: int, n_loc: int):
+        unified, ref, max_tokens = _make_pools(ps)
+        locs = _rand_locs(max_tokens, ps, n_loc)
+        torch.manual_seed(7)
+        for l in range(_L):
+            layer = types.SimpleNamespace(layer_id=l)
+            nope = torch.randn(n_loc, 1, _LORA, dtype=_DTYPE, device=_DEV)
+            rope = torch.randn(n_loc, 1, _ROPE, dtype=_DTYPE, device=_DEV)
+            unified.set_mla_kv_buffer(layer, _dense(locs, ps), nope, rope)
+            ref.set_mla_kv_buffer(layer, locs, nope, rope)
+        torch.cuda.synchronize()
+        self._assert_parity(unified, ref, locs, ps)
+
+    def test_set_mla_kv_buffer_triton_fallback_ps1(self):
+        self._run_set_mla(ps=1, n_loc=256)  # < 768 -> Triton fallback kernel
+
+    def test_set_mla_kv_buffer_tma_jit_ps1(self):
+        self._run_set_mla(ps=1, n_loc=1024)  # >= 768 -> TMA JIT fast path
+
+    def test_set_mla_kv_buffer_triton_fallback_ps64(self):
+        self._run_set_mla(ps=64, n_loc=256)
+
+    def test_set_mla_kv_buffer_tma_jit_ps64(self):
+        self._run_set_mla(ps=64, n_loc=1024)
+
+    def test_set_kv_buffer_combined_write(self):
+        for ps in (1, 64):
+            unified, ref, max_tokens = _make_pools(ps)
+            n_loc = 512
+            locs = _rand_locs(max_tokens, ps, n_loc)
+            torch.manual_seed(11)
+            for l in (0, _L // 2, _L - 1):
+                layer = types.SimpleNamespace(layer_id=l)
+                k = torch.randn(n_loc, 1, _D, dtype=_DTYPE, device=_DEV)
+                unified.set_kv_buffer(layer, _dense(locs, ps), k, None)
+                ref.set_kv_buffer(layer, locs, k, None)
+            torch.cuda.synchronize()
+            self._assert_parity(unified, ref, locs, ps, layers=(0, _L // 2, _L - 1))
+
+    def test_get_mla_kv_buffer_roundtrip(self):
+        for ps in (1, 64):
+            unified, ref, max_tokens = _make_pools(ps)
+            n_loc = 1024  # exercise both get paths against the same bytes
+            locs = _rand_locs(max_tokens, ps, n_loc)
+            torch.manual_seed(13)
+            layer = types.SimpleNamespace(layer_id=3)
+            nope = torch.randn(n_loc, 1, _LORA, dtype=_DTYPE, device=_DEV)
+            rope = torch.randn(n_loc, 1, _ROPE, dtype=_DTYPE, device=_DEV)
+            unified.set_mla_kv_buffer(layer, _dense(locs, ps), nope, rope)
+            got_nope, got_rope = unified.get_mla_kv_buffer(layer, _dense(locs, ps))
+            torch.cuda.synchronize()
+            torch.testing.assert_close(got_nope, nope, rtol=0, atol=0)
+            torch.testing.assert_close(got_rope, rope, rtol=0, atol=0)
+
+    def test_move_kv_cache_page_envelope_gpu(self):
+        for ps in (1, 64):
+            unified, ref, max_tokens = _make_pools(ps)
+            num_pages = max_tokens // ps
+            n_loc = ps  # one full page of tokens
+            src_page, dst_page = num_pages - 2, 2
+            src_t = torch.arange(ps, device=_DEV, dtype=torch.int64) + src_page * ps
+            dst_t = torch.arange(ps, device=_DEV, dtype=torch.int64) + dst_page * ps
+            torch.manual_seed(17)
+            for l in range(_L):
+                layer = types.SimpleNamespace(layer_id=l)
+                k = torch.randn(n_loc, 1, _D, dtype=_DTYPE, device=_DEV)
+                unified.set_kv_buffer(layer, _dense(src_t, ps), k, None)
+            before = [
+                unified.get_key_buffer(l)[_dense(src_t, ps)].clone() for l in range(_L)
+            ]
+            unified.move_kv_cache(dst_t, src_t)
+            torch.cuda.synchronize()
+            for l in range(_L):
+                got = unified.get_key_buffer(l)[_dense(dst_t, ps)]
+                torch.testing.assert_close(got, before[l], rtol=0, atol=0)
+
+
+if __name__ == "__main__":
+    unittest.main()

test/registered/unit/mem_cache/test_unified_mla_views.py

@@ -0,0 +1,391 @@
+# Copyright 2023-2026 SGLang Team
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Dense MLA views for the unified memory pool (MLA-hybrid-Mamba, Kimi K3).
+
+Covers, CPU-only (pure torch — no GPU / Triton kernels):
+  - `MLASubPoolSpec` byte math;
+  - `build_dense_mla_views` addressing: view_l[dense(t)] must land exactly at
+    the page-major envelope byte offset `p*(L*ps*D) + l*(ps*D) + s*D`, the
+    overlapping per-layer views must not alias at equal dense ids, and the
+    missing-tail-pad case must fail loud;
+  - `UnifiedKVPool` MLA plumbing: `view_tail_pad_bytes` extends the allocation
+    only, and the reserved sink floor covers the whole page-0 envelope;
+  - `UnifiedMLATokenToKVPool`: buffer wiring, V-as-prefix-slice, and the
+    page-envelope `move_kv_cache` (REAL physical token ids, page-major runs);
+  - `MultiEndedAllocator.translate_kv_loc_dense`: dense = v2p-page * (ps*L) +
+    offset, tombstone clamp to the sink, `out=` contract, multiplier-1
+    fallback, and correctness across eager compaction.
+
+GPU parity of the actual read/write kernels (set_mla_kv_buffer TMA path etc.)
+lives in the server-level tests, not here.
+
+    python -m pytest test/registered/unit/mem_cache/test_unified_mla_views.py -v
+"""
+
+from sglang.test.ci.ci_register import register_cpu_ci
+
+register_cpu_ci(est_time=8, suite="base-a-test-cpu")
+
+import unittest
+
+import torch
+
+from sglang.srt.mem_cache.layout.page_major import (
+    build_dense_mla_views,
+    mla_entry_bytes,
+)
+from sglang.srt.mem_cache.multi_ended_allocator import MultiEndedAllocator
+from sglang.srt.mem_cache.unified_memory_pool import (
+    MambaSubPoolSpec,
+    MLASubPoolSpec,
+    UnifiedKVPool,
+    UnifiedMLATokenToKVPool,
+)
+
+_DEV = "cpu"
+
+# Small-but-nontrivial MLA geometry: L=3 layers, D=8 (=6+2), so every byte
+# offset is hand-checkable. Real K3 is L=24, D=576 (=512+64).
+_L = 3
+_LORA = 6
+_ROPE = 2
+_D = _LORA + _ROPE
+_DTYPE = torch.bfloat16
+_ITEM = _DTYPE.itemsize
+
+
+def _mla_spec(grow="down", layer_num=_L):
+    return MLASubPoolSpec(
+        name="full",
+        layer_num=layer_num,
+        kv_lora_rank=_LORA,
+        qk_rope_head_dim=_ROPE,
+        store_dtype=_DTYPE,
+        grow_direction=grow,
+    )
+
+
+def _mamba_spec(grow="up", layer_num=2):
+    return MambaSubPoolSpec(
+        name="mamba",
+        layer_num=layer_num,
+        conv_state_shapes=((4, 3),),
+        conv_dtype=torch.float32,
+        temporal_state_shape=(2, 2, 2),
+        temporal_dtype=torch.float32,
+        grow_direction=grow,
+    )
+
+
+def _make_unified(page_size=1, n_full_tokens=64, n_mamba_slots=8):
+    full = _mla_spec()
+    mamba = _mamba_spec()
+    total = full.entry_bytes() * n_full_tokens + mamba.entry_bytes() * n_mamba_slots
+    pool = UnifiedKVPool(
+        total_bytes=total,
+        sub_pool_specs=[full, mamba],
+        device=_DEV,
+        enable_memory_saver=False,
+        page_size=page_size,
+        view_tail_pad_bytes=page_size * full.entry_bytes(),
+    )
+    return pool, full, mamba
+
+
+def _dense(t, ps, layer_num):
+    return (t // ps) * (ps * layer_num) + t % ps
+
+
+class TestMLASubPoolSpec(unittest.TestCase):
+    def test_entry_bytes_and_dim(self):
+        spec = _mla_spec()
+        self.assertEqual(spec.kv_cache_dim, _D)
+        self.assertEqual(spec.entry_bytes(), _L * _D * _ITEM)
+        self.assertEqual(
+            spec.entry_bytes(),
+            mla_entry_bytes(layer_num=_L, kv_cache_dim=_D, itemsize=_ITEM),
+        )
+        self.assertEqual(spec.get_dtype(), _DTYPE)
+
+    def test_rejects_nonpositive_dims(self):
+        with self.assertRaises(AssertionError):
+            MLASubPoolSpec(
+                name="full",
+                layer_num=_L,
+                kv_lora_rank=0,
+                qk_rope_head_dim=_ROPE,
+                store_dtype=_DTYPE,
+                grow_direction="down",
+            )
+
+
+class TestDenseMLAViews(unittest.TestCase):
+    def _make_raw(self, ps, num_pages, pad_pages=1):
+        page_bytes = ps * _L * _D * _ITEM
+        raw = torch.zeros(
+            (num_pages + pad_pages) * page_bytes, dtype=torch.uint8, device=_DEV
+        )
+        return raw, page_bytes
+
+    def test_view_addressing_matches_envelope_formula(self):
+        for ps in (1, 4):
+            num_pages = 6
+            raw, _ = self._make_raw(ps, num_pages)
+            views = build_dense_mla_views(
+                raw,
+                layer_num=_L,
+                kv_cache_dim=_D,
+                store_dtype=_DTYPE,
+                page_size=ps,
+                num_pages=num_pages,
+            )
+            self.assertEqual(len(views), _L)
+            n_dense = num_pages * _L * ps
+            for v in views:
+                self.assertEqual(tuple(v.shape), (n_dense, 1, _D))
+                # contiguous in the (row, dim) sense — .view(-1, ps, D) legality
+                self.assertEqual(v.stride(0), _D)
+                self.assertEqual(v.stride(2), 1)
+            flat = raw.view(_DTYPE)
+            for p, l, s in [(0, 0, 0), (1, 2, ps - 1), (4, 1, ps // 2), (5, 2, 0)]:
+                t = p * ps + s
+                marker = float(p * 100 + l * 10 + s + 1)
+                views[l][_dense(t, ps, _L)] = marker
+                # envelope formula, in elements
+                elem = p * (_L * ps * _D) + l * (ps * _D) + s * _D
+                self.assertTrue(
+                    torch.all(flat[elem : elem + _D] == marker),
+                    f"(p={p}, l={l}, s={s}, ps={ps}) landed off-formula",
+                )
+
+    def test_views_do_not_alias_across_layers(self):
+        ps = 4
+        num_pages = 4
+        raw, _ = self._make_raw(ps, num_pages)
+        views = build_dense_mla_views(
+            raw,
+            layer_num=_L,
+            kv_cache_dim=_D,
+            store_dtype=_DTYPE,
+            page_size=ps,
+            num_pages=num_pages,
+        )
+        t = 2 * ps + 1  # page 2, slot 1
+        d = _dense(t, ps, _L)
+        for l in range(_L):
+            views[l][d] = float(l + 1)
+        for l in range(_L):
+            self.assertTrue(torch.all(views[l][d] == float(l + 1)))
+
+    def test_missing_tail_pad_fails_loud(self):
+        ps = 2
+        num_pages = 4
+        raw, _ = self._make_raw(ps, num_pages, pad_pages=0)
+        with self.assertRaises(AssertionError):
+            build_dense_mla_views(
+                raw,
+                layer_num=_L,
+                kv_cache_dim=_D,
+                store_dtype=_DTYPE,
+                page_size=ps,
+                num_pages=num_pages,
+            )
+
+
+class TestUnifiedKVPoolMLA(unittest.TestCase):
+    def test_max_slots_ignore_tail_pad(self):
+        pool, full, mamba = _make_unified(page_size=4)
+        total = full.entry_bytes() * 64 + mamba.entry_bytes() * 8
+        self.assertEqual(pool.max_slots("full"), total // full.entry_bytes())
+        self.assertEqual(pool.max_slots("mamba"), total // mamba.entry_bytes())
+        # allocation actually carries the pad
+        self.assertEqual(pool._raw.numel(), total + 4 * full.entry_bytes())
+
+    def test_reserved_floor_covers_page0_envelope(self):
+        ps = 4
+        pool, full, mamba = _make_unified(page_size=ps)
+        floor = max(
+            max(full.entry_bytes(), mamba.entry_bytes()), ps * full.entry_bytes()
+        )
+        for spec in (full, mamba):
+            self.assertGreaterEqual(
+                pool.min_slot_index(spec.name) * spec.entry_bytes(), floor
+            )
+
+    def test_mla_views_accessor(self):
+        pool, full, _ = _make_unified(page_size=1)
+        views = pool.mla_views_for("full")
+        self.assertEqual(len(views), _L)
+        self.assertIs(pool.mla_spec("full"), full)
+
+
+class TestUnifiedMLATokenToKVPool(unittest.TestCase):
+    def _make(self, ps=1):
+        pool, full, mamba = _make_unified(page_size=ps)
+        kv_pool = UnifiedMLATokenToKVPool(
+            unified_buffer=pool,
+            sub_pool_name="full",
+            kv_cache_dtype=_DTYPE,
+            page_size=ps,
+        )
+        return pool, kv_pool
+
+    def test_buffers_and_prefix_value_slice(self):
+        pool, kv_pool = self._make(ps=1)
+        self.assertEqual(len(kv_pool.kv_buffer), _L)
+        self.assertEqual(kv_pool.get_kv_size_bytes(), 0)
+        k = kv_pool.get_key_buffer(1)
+        v = kv_pool.get_value_buffer(1)
+        self.assertEqual(k.shape[-1], _D)
+        self.assertEqual(v.shape[-1], _LORA)
+        # V is a prefix slice of K's storage: writing K shows up in V
+        k[7] = 2.5
+        self.assertTrue(torch.all(v[7] == 2.5))
+
+    def test_move_kv_cache_moves_page_envelopes(self):
+        for ps in (1, 4):
+            pool, kv_pool = self._make(ps=ps)
+            num_pages = pool.max_slots("full") // ps
+            page_bytes = ps * _L * _D * _ITEM
+            env = pool._raw[: num_pages * page_bytes].view(num_pages, page_bytes)
+            src_pages = torch.tensor([num_pages - 2, num_pages - 4])
+            dst_pages = torch.tensor([2, 3])
+            env[src_pages[0]] = 7
+            env[src_pages[1]] = 9
+            # page-major token runs, exactly how compaction expands pages
+            offsets = torch.arange(ps, dtype=torch.int64)
+            src_t = (src_pages[:, None] * ps + offsets).reshape(-1)
+            dst_t = (dst_pages[:, None] * ps + offsets).reshape(-1)
+            kv_pool.move_kv_cache(dst_t, src_t)
+            self.assertTrue(torch.all(env[dst_pages[0]] == 7), f"ps={ps}")
+            self.assertTrue(torch.all(env[dst_pages[1]] == 9), f"ps={ps}")
+
+    def test_move_then_dense_readback(self):
+        ps = 4
+        pool, kv_pool = self._make(ps=ps)
+        num_pages = pool.max_slots("full") // ps
+        src_page, dst_page = num_pages - 3, 5
+        # write through the views at src, expect it at dst after the move
+        for l in range(_L):
+            for s in range(ps):
+                kv_pool.kv_buffer[l][_dense(src_page * ps + s, ps, _L)] = float(
+                    l * ps + s + 1
+                )
+        offsets = torch.arange(ps, dtype=torch.int64)
+        kv_pool.move_kv_cache(
+            (torch.tensor([dst_page])[:, None] * ps + offsets).reshape(-1),
+            (torch.tensor([src_page])[:, None] * ps + offsets).reshape(-1),
+        )
+        for l in range(_L):
+            for s in range(ps):
+                got = kv_pool.kv_buffer[l][_dense(dst_page * ps + s, ps, _L)]
+                self.assertTrue(
+                    torch.all(got == float(l * ps + s + 1)), f"(l={l}, s={s})"
+                )
+
+
+class _FakeKVCache:
+    def __init__(self, max_slots: int):
+        self.buf = torch.full((max_slots,), -1, dtype=torch.int64)
+
+    def move_kv_cache(self, dst_loc: torch.Tensor, src_loc: torch.Tensor):
+        self.buf[dst_loc] = self.buf[src_loc].clone()
+
+
+class TestTranslateKvLocDense(unittest.TestCase):
+    def _build(self, ps=1, n_full_tokens=64, multiplier=_L):
+        pool, full, mamba = _make_unified(page_size=ps, n_full_tokens=n_full_tokens)
+        full_alloc = MultiEndedAllocator(
+            kvcache=_FakeKVCache(pool.max_slots("full")),
+            unified_buffer=pool,
+            sub_pool_name="full",
+            device=_DEV,
+            is_id_owner=True,
+            page_size=ps,
+            kernel_page_multiplier=multiplier,
+        )
+        mamba_alloc = MultiEndedAllocator(
+            kvcache=_FakeKVCache(pool.max_slots("mamba")),
+            unified_buffer=pool,
+            sub_pool_name="mamba",
+            device=_DEV,
+            is_id_owner=True,
+        )
+        full_alloc.bind_peer(mamba_alloc)
+        mamba_alloc.bind_peer(full_alloc)
+        return full_alloc
+
+    def test_dense_matches_formula_ps1(self):
+        alloc = self._build(ps=1)
+        v = alloc.alloc(8)
+        self.assertIsNotNone(v)
+        phys = alloc.translate_kv_loc(v)
+        dense = alloc.translate_kv_loc_dense(v)
+        self.assertTrue(torch.all(dense == phys * _L))
+
+    def test_dense_matches_formula_paged(self):
+        ps = 4
+        alloc = self._build(ps=ps)
+        v = alloc.alloc(3 * ps)
+        self.assertIsNotNone(v)
+        phys = alloc.translate_kv_loc(v)
+        dense = alloc.translate_kv_loc_dense(v)
+        expected = (phys // ps) * (ps * _L) + phys % ps
+        self.assertTrue(torch.all(dense == expected))
+
+    def test_tombstone_clamps_to_sink(self):
+        alloc = self._build(ps=1)
+        # never-allocated virtual ids -> v2p == -1 -> dense id 0
+        virt = torch.tensor([alloc.min_slot_index + 1], dtype=torch.int64)
+        dense = alloc.translate_kv_loc_dense(virt)
+        self.assertTrue(torch.all(dense == 0))
+
+    def test_out_matches_and_aliases(self):
+        for ps in (1, 4):
+            alloc = self._build(ps=ps)
+            v = alloc.alloc(2 * ps)
+            self.assertIsNotNone(v)
+            no_out = alloc.translate_kv_loc_dense(v)
+            out = torch.empty_like(v)
+            ret = alloc.translate_kv_loc_dense(v, out=out)
+            self.assertIs(ret, out)
+            self.assertTrue(torch.all(out == no_out))
+            # canonical in-place aliasing: translate(x, out=x)
+            x = v.clone()
+            alloc.translate_kv_loc_dense(x, out=x)
+            self.assertTrue(torch.all(x == no_out))
+
+    def test_multiplier_one_falls_back_to_physical(self):
+        alloc = self._build(ps=1, multiplier=1)
+        v = alloc.alloc(4)
+        self.assertIsNotNone(v)
+        self.assertTrue(
+            torch.all(alloc.translate_kv_loc_dense(v) == alloc.translate_kv_loc(v))
+        )
+
+    def test_dense_follows_compaction(self):
+        alloc = self._build(ps=1)
+        a = alloc.alloc(4)
+        b = alloc.alloc(4)
+        c = alloc.alloc(4)
+        self.assertIsNotNone(c)
+        alloc.free(b)  # eager compaction relocates survivors
+        phys_a = alloc.translate_kv_loc(a)
+        phys_c = alloc.translate_kv_loc(c)
+        self.assertTrue(torch.all(alloc.translate_kv_loc_dense(a) == phys_a * _L))
+        self.assertTrue(torch.all(alloc.translate_kv_loc_dense(c) == phys_c * _L))
+
+
+if __name__ == "__main__":
+    unittest.main()