[kv-shard 2/4] Sharded pools (#37615)

Co-authored-by: Zhangheng <hzh0425@apache.org>

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

@@ -22,6 +22,10 @@ from sglang.kernels.ops.kvcache.trtllm_mha_page_table import (
 )
 from sglang.srt.configs.model_config import AttentionArch
 from sglang.srt.layers.attention.base_attn_backend import AttentionBackend
+from sglang.srt.layers.attention.kv_shard_hooks import (
+    get_kv_shard_pool,
+    prepare_kv_shard_forward,
+)
 from sglang.srt.layers.attention.verify_mask import VerifyMask, maybe_create_verify_mask
 from sglang.srt.layers.cp.base import CPAttentionBackendKind, get_cp_strategy
 from sglang.srt.layers.cp.utils import is_cp_active
@@ -46,6 +50,7 @@ from sglang.srt.utils.common import get_device_capability
 
 if TYPE_CHECKING:
     from sglang.srt.layers.radix_attention import RadixAttention
+    from sglang.srt.mem_cache.page_interleave_pool import PageInterleaveKVPoolMixin
     from sglang.srt.model_executor.model_runner import ModelRunner
 
 from sgl_kernel import merge_state_v2
@@ -145,6 +150,10 @@ class FlashAttentionBackend(AttentionBackend):
     needs_cpu_seq_lens: bool = False
     supports_ragged_verify_graph: bool = True
 
+    # Set from the pool type in __init__; the class default keeps the extend
+    # metadata guard readable on instances built with __new__ (test stubs).
+    _kv_shard_pool: Optional[PageInterleaveKVPoolMixin] = None
+
     # Chunked-prefix attention reads the stable ForwardBatch cu-seqlens and
     # KV-index buffers directly, so it needs no backend-private replay state.
     supports_full_cuda_graph_chunked_prefix = True
@@ -205,6 +214,12 @@ class FlashAttentionBackend(AttentionBackend):
             and model_runner.token_to_kv_pool.swa_layer_nums > 0
         )
 
+        self._kv_shard_pool = get_kv_shard_pool(self.token_to_kv_pool)
+        # begin_shard_extend builds the owner-major gather plan from host-side
+        # prefix/final lengths. Normal FA3 metadata is device-only, so opt the
+        # sharded variant back into FutureMap's CPU mirror publication.
+        self.needs_cpu_seq_lens = self._kv_shard_pool is not None
+
         self.topk = get_spec().speculative_eagle_topk or 0
         self.speculative_num_steps = speculative_num_steps
         self.speculative_num_draft_tokens = get_spec().speculative_num_draft_tokens
@@ -1154,6 +1169,23 @@ class FlashAttentionBackend(AttentionBackend):
                     )
                 )
 
+        # Logical-page KV sharding: capture the batch's gather plan and swap the
+        # page table to scratch rows. During a sharded extend, attention reads
+        # the assembled [prefix | chunk] scratch, never the striped pool rows;
+        # the plan capture also kicks the first layer's prefix gather.
+        #
+        # Runs after KVIndexTranslator and before the `// page_size` reduction.
+        # The unified-memory UnifiedKVPool and page-interleaved pools are
+        # alternatives, so at most one translation fires.
+        if self._kv_shard_pool is not None and prepare_kv_shard_forward(
+            self._kv_shard_pool,
+            self.req_to_token,
+            forward_batch,
+        ):
+            metadata.page_table = self._kv_shard_pool.translate_loc_to_scratch(
+                metadata.page_table
+            ).to(torch.int32)
+
         # Convert the page table to a strided format which is needed by FA3 API
         if self.page_size > 1 and not _unified_read:
             self.strided_indices = torch.arange(

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

@@ -0,0 +1,67 @@
+# 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.
+# ==============================================================================
+
+"""Logical-page KV-shard helpers shared by attention backends.
+
+Page-interleaved pools build one gather plan per extend batch. Keeping pool
+detection and that begin/end lifecycle here prevents each compatible attention
+backend from implementing a subtly different version of the contract.
+"""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Optional
+
+import torch
+
+from sglang.srt.mem_cache.page_interleave_pool import PageInterleaveKVPoolMixin
+
+if TYPE_CHECKING:
+    from sglang.srt.model_executor.forward_batch_info import ForwardBatch
+
+
+def get_kv_shard_pool(token_to_kv_pool) -> Optional[PageInterleaveKVPoolMixin]:
+    """Return the page-interleaved pool, if logical-page sharding is active."""
+    return (
+        token_to_kv_pool
+        if isinstance(token_to_kv_pool, PageInterleaveKVPoolMixin)
+        else None
+    )
+
+
+def prepare_kv_shard_forward(
+    pool: PageInterleaveKVPoolMixin,
+    req_to_token: torch.Tensor,
+    forward_batch: ForwardBatch,
+) -> bool:
+    """Update the pool's gather plan and report whether this is an extend."""
+    if not forward_batch.forward_mode.is_extend_or_draft_extend_or_mixed():
+        pool.end_shard_extend()
+        return False
+
+    req_pool_indices = forward_batch.req_pool_indices
+    prefix_lens = forward_batch.extend_prefix_lens_cpu
+    seq_lens = forward_batch.seq_lens_cpu
+    if req_pool_indices is None or prefix_lens is None or seq_lens is None:
+        raise RuntimeError(
+            "KV-sharded attention requires request indices and CPU length metadata"
+        )
+
+    pool.begin_shard_extend(
+        req_to_token,
+        req_pool_indices,
+        prefix_lens,
+        seq_lens,
+    )
+    return True

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

@@ -50,6 +50,10 @@ from sglang.srt.layers.attention.flashinfer_mla_backend import (
     FlashInferMLAAttnBackend,
     FlashInferMLAMultiStepDraftBackend,
 )
+from sglang.srt.layers.attention.kv_shard_hooks import (
+    get_kv_shard_pool,
+    prepare_kv_shard_forward,
+)
 from sglang.srt.layers.attention.verify_mask import VerifyMask, maybe_create_verify_mask
 from sglang.srt.layers.dcp.layout import get_dcp_lens
 from sglang.srt.layers.logits_processor import get_in_autotune_dummy_run
@@ -239,6 +243,8 @@ class TRTLLMMLABackend(FlashInferMLAAttnBackend):
         self.q_data_type = model_runner.dtype
         self.page_size = model_runner.page_size
         self.req_to_token = model_runner.req_to_token_pool.req_to_token
+        self._kv_shard_pool = get_kv_shard_pool(model_runner.token_to_kv_pool)
+        self.needs_cpu_seq_lens |= self._kv_shard_pool is not None
 
         # Workspace allocation
         self.workspace_size = DEFAULT_WORKSPACE_SIZE_MB * 1024 * 1024
@@ -806,6 +812,13 @@ class TRTLLMMLABackend(FlashInferMLAAttnBackend):
 
     def init_forward_metadata(self, forward_batch: ForwardBatch):
         """Initialize the metadata for a forward pass."""
+        if self._kv_shard_pool is not None:
+            prepare_kv_shard_forward(
+                self._kv_shard_pool,
+                self.req_to_token,
+                forward_batch,
+            )
+
         self._decode_kernel_loc = None
         # Delegate to parent for non-decode modes.
         if (

python/sglang/srt/mem_cache/page_interleave.py

@@ -23,8 +23,8 @@ sees only its own physical pages; the boundary is the pure bijection below.
 The shard group is the group across which KV storage is replicated today and
 therefore can be striped without extra compute-time communication:
 
-- GQA/MHA models: the **attention CP group** — prefill CP already allgathers
-  the full chunk's K/V to every CP rank (``cp_allgather_and_save_kv_cache``).
+- GQA/MHA models: the **attention CP group** — the prefill CP strategy
+  materializes the full chunk's K/V on every CP rank before the pool write.
 - MLA models: the **attention TP group** — the latent KV projection is
   ``ReplicatedLinear``, so every attn-TP rank computes identical latent KV.
 """
@@ -32,10 +32,16 @@ therefore can be striped without extra compute-time communication:
 from __future__ import annotations
 
 import logging
+from typing import TYPE_CHECKING
 
 import msgspec
 import torch
 
+from sglang.srt.runtime_context import get_parallel
+
+if TYPE_CHECKING:
+    from sglang.srt.distributed.parallel_state import GroupCoordinator
+
 logger = logging.getLogger(__name__)
 
 
@@ -84,3 +90,24 @@ class PageInterleavePlacement:
     def filter_local(self, loc: torch.Tensor, rank: int) -> torch.Tensor:
         """Logical slots -> this rank's physical pool rows, order-preserving."""
         return self.local_index(loc[self.local_mask(loc, rank)])
+
+
+def get_kv_shard_group(use_mla_backend: bool) -> GroupCoordinator:
+    """The group KV pages are striped across — the axis that replicates KV
+    at rest, chosen by topology:
+
+    - An active attention-CP group takes precedence: prefill CP replicates
+      KV storage across CP ranks for every attention type (GQA via the
+      full-chunk allgather, MLA via the CP latent-KV rebuild).
+    - Without CP, MLA latent KV is still replicated across attention-TP
+      (ReplicatedLinear projection), so the attn-TP group is the shard axis.
+    - GQA without CP has no replicated axis (KV is head-sharded across TP);
+      the returned trivial CP group has world_size 1, which disables
+      sharding in get_kv_shard_group_info.
+    """
+    cp_group = get_parallel().attn_cp_group
+    if cp_group.world_size > 1:
+        return cp_group
+    if use_mla_backend:
+        return get_parallel().attn_tp_group
+    return cp_group

python/sglang/srt/mem_cache/page_interleave_pool.py

@@ -0,0 +1,703 @@
+# 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.
+# ==============================================================================
+
+"""KV pools striped at logical-page granularity across a shard group.
+
+Per-layer pool tensors keep the stock shape and size — identical rows on every
+rank; only the written rows differ: each rank persists the rows it owns under
+the pure placement bijection (``mem_cache/page_interleave.py``).
+
+Central consequence for the forward pass: during a sharded prefill, extend
+attention never reads the KV pool — a rank's pool holds only its stripe of
+the prefix *and* of the current chunk. Attention reads a per-layer assembled
+scratch slot laid out ``[prefix region | chunk region | trash page]``:
+
+- The prefix is NCCL-allgathered from the shard group into the slot one layer
+  ahead, on a dedicated side stream with a dedicated ``PyNcclCommunicator``
+  ahead. Under rotated
+  owner-classed allocation the owners of a cached prefix's pages are exactly
+  cyclic, so per-rank owned counts differ by <= 1: each rank sends its own
+  prefix pages in local-page order, padded to ``ceil(prefix_pages / N)`` pages
+  with the (never referenced) trash page — a *regular* in-place allgather,
+  owner-major output, no reorder pass.
+- The chunk region is staged locally on the compute stream where the write
+  path already holds the full chunk (GQA CP allgathers K/V before the pool
+  write; MLA latent is replicated across the shard group at write time).
+- The trash page absorbs padded/dummy locations (the reserved logical pages
+  covering loc < N*ps, and any location outside the current batch's plan).
+
+Scratch addressing goes through a per-batch ``logical page -> scratch
+page`` lookup (``_page_pos``, one int32 per logical page, rebuilt in
+``begin_shard_extend``): purely local, mirrored by construction, and valid
+for any consumer index vector (attention page tables, MLA prefix
+``kv_indices``), not just whole-prefix position arithmetic.
+
+Layer-ahead pipelining needs no model changes: acquiring layer ``l``'s slot
+for reading kicks layer ``l+1``'s gather. Reads happen in SPMD lockstep on
+every rank, so the collective order is symmetric by construction.
+"""
+
+from __future__ import annotations
+
+import logging
+from contextlib import nullcontext
+from typing import TYPE_CHECKING, Dict, List, Optional, Tuple
+
+import torch
+
+from sglang.srt.distributed.device_communicators.pynccl import PyNcclCommunicator
+from sglang.srt.mem_cache.memory_pool import (
+    GPU_MEMORY_TYPE_KV_CACHE,
+    MHATokenToKVPool,
+    MLATokenToKVPool,
+    RadixAttention,
+    unwrap_write_loc,
+)
+from sglang.srt.mem_cache.page_interleave import (
+    PageInterleavePlacement,
+    PageShardSpec,
+)
+from sglang.srt.mem_cache.utils import (
+    get_mla_kv_buffer_triton,
+    set_mla_kv_buffer_triton,
+)
+from sglang.srt.utils import ceil_div, get_bool_env_var
+
+if TYPE_CHECKING:
+    from sglang.srt.distributed.parallel_state import GroupCoordinator
+
+logger = logging.getLogger(__name__)
+
+
+class _ScratchSlot:
+    """One half of the double buffer: scratch tensors + ready event.
+
+    ``resident_key`` identifies what the slot currently holds — ``(layer_id,
+    epoch)`` — so repeated reads of the same layer skip the re-gather and a
+    new batch (epoch bump) retires stale residency without any invalidation
+    walk.
+    """
+
+    def __init__(self, tensors: Dict[str, torch.Tensor], device_module):
+        self.tensors = tensors
+        self.ready = device_module.Event()
+        self.resident_key: Optional[Tuple[int, int]] = None
+
+
+class PageInterleaveKVPoolMixin:
+    """Shard-generic state and mechanics; mixed into concrete pools below.
+
+    Subclasses call ``_init_page_shard_state`` after the base pool has created
+    its buffers, and implement ``_scratch_tensor_specs`` (per-slot tensors) and
+    ``_gather_pairs`` (pool buffer -> scratch tensor pairs of one layer).
+    """
+
+    # ---- init ---------------------------------------------------------------
+
+    def _init_page_shard_state(
+        self, shard_spec: PageShardSpec, shard_group: GroupCoordinator
+    ):
+        spec = shard_spec
+        assert spec.shard_size > 1, "page-interleave sharding needs shard_size > 1"
+        assert spec.shard_size == shard_group.world_size
+        assert spec.shard_rank == shard_group.rank_in_group
+        assert spec.page_size == self.page_size
+        # The prefix region must fit N * ceil(prefix_pages / N) pages for any
+        # prefix, i.e. be a multiple of the full-group span; the chunk region
+        # is per-page.
+        assert spec.max_prefix_tokens % spec.logical_page_size == 0
+        assert spec.chunk_tokens % spec.page_size == 0
+
+        self.shard_spec = spec
+        self.placement = PageInterleavePlacement(spec)
+        self.shard_rank = spec.shard_rank
+        self.shard_size = spec.shard_size
+        self.device_module = torch.get_device_module(self.device)
+
+        # Dedicated communicator + stream so the layer-ahead gathers never
+        # interleave with the group's main collectives.
+        self.kv_gather_comm: PyNcclCommunicator = PyNcclCommunicator(
+            group=shard_group.cpu_group, device=shard_group.device
+        )
+        self.kv_gather_stream = self.device_module.Stream()
+
+        # Scratch: [prefix | chunk | trash page], double-buffered.
+        scratch_rows = spec.max_prefix_tokens + spec.chunk_tokens + spec.page_size
+        self._chunk_base = spec.max_prefix_tokens
+        self._trash_base = spec.max_prefix_tokens + spec.chunk_tokens
+        with self.memory_saver_adapter.region(GPU_MEMORY_TYPE_KV_CACHE):
+            with (
+                torch.cuda.use_mem_pool(self.custom_mem_pool)
+                if self.custom_mem_pool
+                else nullcontext()
+            ):
+                self._slots = [
+                    _ScratchSlot(
+                        self._scratch_tensor_specs(scratch_rows), self.device_module
+                    )
+                    for _ in range(2)
+                ]
+                # Per-batch absolute scratch-page index of every logical page.
+                # Unplanned pages initially target the trash page. Logical page
+                # ids run [0, N * (size/ps + 1)) — N per physical page of one
+                # rank's pool incl. the reserved padded page.
+                self._page_pos = torch.full(
+                    (spec.shard_size * (self.size // spec.page_size + 2),),
+                    self._trash_base // spec.page_size,
+                    dtype=torch.int32,
+                    device=self.device,
+                )
+
+        self._epoch = 0
+        self._shard_extend_active = False
+        self._block_pages = 0
+        # Strictly larger than any local physical page id: the owner-major
+        # sort key of logical page l is (l % N) * stride + l // N.
+        self._local_page_stride = self.size // spec.page_size + 2
+        self._debug_plan_checks = get_bool_env_var("SGLANG_DEBUG_MEMORY_POOL")
+        self._send_rows: Optional[torch.Tensor] = None
+        self._write_plan_key = None
+        self._write_plan: Optional[Tuple[torch.Tensor, torch.Tensor]] = None
+        self._translate_cache: Dict[Tuple[int, int], torch.Tensor] = {}
+
+        logger.info(
+            "Page-interleave KV sharding enabled: shard_rank=%d shard_size=%d "
+            "page_size=%d scratch_rows=%d x2",
+            self.shard_rank,
+            self.shard_size,
+            spec.page_size,
+            scratch_rows,
+        )
+
+    def set_kv_buffer_prefix_valid(self, *args, **kwargs):
+        raise NotImplementedError(
+            "prefix-valid commit is unsupported under logical-page KV sharding "
+            "(it writes pool rows directly, bypassing the ownership filter)"
+        )
+
+    # ---- subclass hooks -------------------------------------------------------
+
+    def _scratch_tensor_specs(self, rows: int) -> Dict[str, torch.Tensor]:
+        raise NotImplementedError
+
+    def _gather_pairs(self, local_layer: int) -> List[Tuple[torch.Tensor, str]]:
+        """(pool buffer of ``local_layer``, scratch tensor name) pairs."""
+        raise NotImplementedError
+
+    # ---- per-batch plan -------------------------------------------------------
+
+    def begin_shard_extend(
+        self,
+        req_to_token: torch.Tensor,
+        req_pool_indices: torch.Tensor,
+        prefix_lens_cpu: List[int],
+        seq_lens_cpu: List[int],
+    ) -> None:
+        """Capture the batch's gather plan and kick the first layer's gather.
+
+        Called once per scheduled extend batch, at attention-metadata build
+        time (after ``req_to_token`` holds the batch's allocation). All inputs
+        are mirrored across ranks, so every rank derives the identical plan.
+        """
+        spec = self.shard_spec
+        ps, shard_size = spec.page_size, spec.shard_size
+
+        # Per-request page samples. One logical page per position-page
+        # (whole-page draws + ps-aligned chunk starts), so stride-ps sampling
+        # enumerates each request's pages in position order — no run
+        # collapsing. block_pages (the padded allgather block) is the
+        # sync-free host bound sum_i ceil(K_i / N): each request's prefix is
+        # ONE cyclic rotation run (graft-declined inserts + per-request
+        # bases), so a rank owns at most ceil(K_i / N) of its pages, and the
+        # cross-request dedup below only shrinks per-rank counts.
+        empty = torch.empty((0,), dtype=torch.int64, device=self.device)
+        prefix_parts = []
+        chunk_parts = []
+        block_pages = 0
+        for i in range(len(prefix_lens_cpu)):
+            prefix_len, seq_len = int(prefix_lens_cpu[i]), int(seq_lens_cpu[i])
+            assert prefix_len % ps == 0, (
+                f"sharded prefill requires physical-page-aligned prefixes "
+                f"(the radix-tree match quantum), got prefix_len="
+                f"{prefix_len}, page_size={ps}"
+            )
+            row = req_to_token[req_pool_indices[i]]
+            if prefix_len:
+                pages_i = row[:prefix_len:ps].long() // ps
+                prefix_parts.append(pages_i)
+                block_pages += ceil_div(prefix_len // ps, shard_size)
+                if self._debug_plan_checks:
+                    # Owner-congruence: each request's prefix owners must be
+                    # exactly cyclic — the block bound above relies on it. A
+                    # broken rotation would overflow a rank's block and the
+                    # translation would silently address the wrong scratch
+                    # rows. Fail loud instead (debug mode syncs).
+                    owners = pages_i % shard_size
+                    expected = (
+                        int(owners[0])
+                        + torch.arange(
+                            owners.numel(), dtype=torch.int64, device=self.device
+                        )
+                    ) % shard_size
+                    assert torch.equal(owners, expected), (
+                        f"sharded prefix owners of request {i} are not "
+                        f"cyclic — rotation base out of sync "
+                        f"(owners[:16]={owners[:16].tolist()})"
+                    )
+            if seq_len > prefix_len:
+                chunk_parts.append(row[prefix_len:seq_len:ps].long() // ps)
+
+        # A prefix page shared by several requests gathers into ONE scratch
+        # slot (torch.unique sorts — deterministic, mirrored); chunk pages
+        # are per-request fresh allocations, disjoint by construction.
+        prefix_pages = torch.unique(torch.cat(prefix_parts)) if prefix_parts else empty
+        chunk_pages = torch.cat(chunk_parts) if chunk_parts else empty
+        n_prefix = prefix_pages.numel()
+        n_chunk = chunk_pages.numel()
+        n_prefix_slots = shard_size * block_pages
+        assert n_prefix_slots * ps <= spec.max_prefix_tokens, (
+            f"prefix ({n_prefix} pages, padded gather span "
+            f"{n_prefix_slots * ps} tokens) exceeds the scratch prefix "
+            f"capacity ({spec.max_prefix_tokens}) — the PrefillAdder scratch "
+            f"reservation should have deferred this batch"
+        )
+        assert n_chunk * ps <= spec.chunk_tokens, (
+            f"chunk ({n_chunk * ps} tokens) exceeds the scratch chunk "
+            f"capacity ({spec.chunk_tokens})"
+        )
+
+        self._page_pos.fill_(self._trash_base // ps)
+        if n_prefix:
+            # Owner-major slot assignment: sort the batch's unique pages by
+            # (owner, local page) so rank r's pages are contiguous at
+            # r * block_pages in local-page order — the same order rank r
+            # packs its send block below, on every rank (mirrored).
+            owners = prefix_pages % shard_size
+            local_pages = prefix_pages // shard_size
+            order = torch.argsort(owners * self._local_page_stride + local_pages)
+            sorted_pages = prefix_pages[order]
+            sorted_owners = owners[order]
+            counts = torch.bincount(sorted_owners, minlength=shard_size)
+            starts = torch.cumsum(counts, 0) - counts
+            within = (
+                torch.arange(n_prefix, dtype=torch.int64, device=self.device)
+                - starts[sorted_owners]
+            )
+            self._page_pos[sorted_pages] = (sorted_owners * block_pages + within).to(
+                torch.int32
+            )
+        if n_chunk:
+            self._page_pos[chunk_pages] = torch.arange(
+                self._chunk_base // ps,
+                self._chunk_base // ps + n_chunk,
+                dtype=torch.int32,
+                device=self.device,
+            )
+
+        self._epoch += 1
+        self._block_pages = block_pages
+        self._shard_extend_active = True
+        self._translate_cache.clear()
+        if block_pages:
+            # This rank's owned pages in slot order (already local-sorted) —
+            # logical page l is its local physical page l // N — padded to
+            # the regular allgather block with the reserved trash page
+            # (local page 0), whose rows the plan never references.
+            own_local = sorted_pages[sorted_owners == self.shard_rank] // shard_size
+            n_pad = block_pages - own_local.numel()
+            if n_pad:
+                own_local = torch.cat(
+                    [
+                        own_local,
+                        torch.zeros((n_pad,), dtype=torch.int64, device=self.device),
+                    ]
+                )
+            self._send_rows = (
+                own_local[:, None] * ps
+                + torch.arange(ps, dtype=torch.int64, device=self.device)
+            ).reshape(-1)
+            self._prefetch_layer(self.start_layer)
+        else:
+            self._send_rows = None
+
+    def end_shard_extend(self) -> None:
+        """Mark no sharded extend in flight (non-extend forward modes)."""
+        self._shard_extend_active = False
+
+    # ---- translation ----------------------------------------------------------
+
+    def translate_loc_to_scratch(self, loc: torch.Tensor) -> torch.Tensor:
+        """Logical token slots -> rows of the current batch's scratch slots.
+
+        ``_page_pos`` holds every logical page's absolute scratch page slot:
+        prefix pages owner-major in ``[0, N * block_pages)`` (rank ``l % N``'s
+        pages contiguous at ``rank * block_pages``, local-page order), chunk
+        pages in batch-sequence order at ``_chunk_base``, and anything outside
+        the plan at the trash page.
+        """
+        ps = self.shard_spec.page_size
+        loc64 = loc.long()
+        scratch_page = self._page_pos[loc64 // ps].long()
+        return scratch_page * ps + loc64 % ps
+
+    # ---- the layer-ahead gather -------------------------------------------------
+
+    def _prefetch_layer(self, layer_id: int) -> None:
+        """Kick the allgather assembling ``layer_id``'s prefix into its slot,
+        on the gather stream. Idempotent per ``(layer_id, epoch)``."""
+        local_layer = layer_id - self.start_layer
+        if local_layer >= self.layer_num:
+            return
+        slot = self._slots[layer_id % 2]
+        key = (layer_id, self._epoch)
+        if slot.resident_key == key:
+            return
+        block = self._block_pages * self.shard_spec.page_size
+        # Order the gather after all prior compute-stream work: the
+        # previous tenant's reads (attention of layer_id - 2) and the pool
+        # writes that produced the prefix rows.
+        self.kv_gather_stream.wait_stream(self.device_module.current_stream())
+        with self.device_module.stream(self.kv_gather_stream):
+            for pool_buf, name in self._gather_pairs(local_layer):
+                scratch = slot.tensors[name]
+                send = scratch[self.shard_rank * block : (self.shard_rank + 1) * block]
+                torch.index_select(pool_buf, 0, self._send_rows, out=send)
+                # In-place regular allgather: send is exactly the rank's
+                # block of the output, every rank contributes `block` rows.
+                with self.kv_gather_comm.change_state(enable=True):
+                    self.kv_gather_comm.all_gather(
+                        scratch[: self.shard_size * block], send
+                    )
+            slot.ready.record(self.kv_gather_stream)
+        slot.resident_key = key
+
+    def _acquire_slot_for_read(self, layer_id: int) -> _ScratchSlot:
+        """Block the compute stream on the slot's ready event (a no-op when
+        the layer-ahead gather already landed) and kick the next layer's
+        gather. A residency miss is a caller bug — the batch plan captured by
+        ``begin_shard_extend`` must have prefetched this layer."""
+        assert self._shard_extend_active, (
+            "sharded pool read outside an active sharded extend batch "
+            "(begin_shard_extend not called?)"
+        )
+        slot = self._slots[layer_id % 2]
+        if self._block_pages:
+            assert slot.resident_key == (layer_id, self._epoch), (
+                f"prefix scratch miss for layer {layer_id} "
+                f"(resident={slot.resident_key}, epoch={self._epoch})"
+            )
+            self.device_module.current_stream().wait_event(slot.ready)
+            self._prefetch_layer(layer_id + 1)
+        return slot
+
+    def _translate_loc_cached(self, loc: torch.Tensor) -> torch.Tensor:
+        """Per-batch memoized ``translate_loc_to_scratch`` for the per-layer
+        callers: the same loc tensors (``out_cache_loc`` at every layer's
+        ``set_kv_buffer``; the prefix ``kv_indices`` at every layer's
+        ``get_mla_kv_buffer``) arrive at all layers of a forward, and the
+        plan is frozen per batch — so each distinct loc tensor is translated
+        once per batch instead of once per layer. ``begin_shard_extend`` clears
+        the cache when it installs a new plan."""
+        key = (loc.data_ptr(), loc.numel())
+        rows = self._translate_cache.get(key)
+        if rows is None:
+            rows = self.translate_loc_to_scratch(loc)
+            self._translate_cache[key] = rows
+        return rows
+
+    # ---- write plan (owner filter), cached per (loc tensor, epoch) --------------
+
+    def _get_write_plan(self, loc: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
+        """(owned positions into ``loc``, their local physical rows).
+
+        The same ``out_cache_loc`` tensor is passed for every layer of a
+        forward, so the owner filter is computed once per batch, not per
+        layer.
+        """
+        key = (loc.data_ptr(), loc.numel(), self._epoch)
+        if self._write_plan_key != key:
+            owned_idx = torch.nonzero(
+                self.placement.local_mask(loc, self.shard_rank)
+            ).squeeze(1)
+            local_rows = self.placement.local_index(loc[owned_idx])
+            self._write_plan_key = key
+            self._write_plan = (owned_idx, local_rows)
+        return self._write_plan
+
+
+class PageInterleaveMHATokenToKVPool(PageInterleaveKVPoolMixin, MHATokenToKVPool):
+    """MHA/GQA pool striped across the attention CP group.
+
+    The write path relies on the prefill CP strategy materializing the full
+    chunk's K/V before ``set_kv_buffer`` receives the full (unsplit)
+    ``out_cache_loc`` on every rank. Each rank persists only its stripe and
+    stages the full chunk into the current layer's scratch chunk region
+    (extend attention reads prefix *and* chunk through the scratch page table).
+    """
+
+    def __init__(
+        self,
+        *args,
+        shard_spec: PageShardSpec,
+        shard_group: GroupCoordinator,
+        **kwargs,
+    ):
+        super().__init__(*args, **kwargs)
+        assert not self.use_hnd, "KV sharding supports the NHD layout only"
+        assert self.kv_cache_layout == "nhd", (
+            f"KV sharding supports the NHD layout only, got {self.kv_cache_layout}"
+        )
+        assert not self.post_capture_active
+        self._init_page_shard_state(shard_spec, shard_group)
+
+    def _scratch_tensor_specs(self, rows: int) -> Dict[str, torch.Tensor]:
+        return {
+            "k": torch.zeros(
+                (rows, self.head_num, self.head_dim),
+                dtype=self.store_dtype,
+                device=self.device,
+            ),
+            "v": torch.zeros(
+                (rows, self.head_num, self.v_head_dim),
+                dtype=self.store_dtype,
+                device=self.device,
+            ),
+        }
+
+    def _gather_pairs(self, local_layer: int):
+        return [
+            (self.k_buffer[local_layer], "k"),
+            (self.v_buffer[local_layer], "v"),
+        ]
+
+    def set_kv_buffer(
+        self,
+        layer: RadixAttention,
+        loc_info,
+        cache_k: torch.Tensor,
+        cache_v: torch.Tensor,
+        k_scale: Optional[float] = None,
+        v_scale: Optional[float] = None,
+        layer_id_override: Optional[int] = None,
+        dcp_kv_mask: Optional[torch.Tensor] = None,
+    ):
+        assert dcp_kv_mask is None, "DCP is mutually exclusive with KV sharding"
+        loc, _, _ = unwrap_write_loc(loc_info)
+        layer_id = (
+            layer_id_override if layer_id_override is not None else layer.layer_id
+        )
+
+        if cache_k.dtype != self.dtype:
+            if k_scale is not None:
+                cache_k = cache_k / k_scale
+            if v_scale is not None:
+                cache_v = cache_v / v_scale
+            cache_k = cache_k.to(self.dtype)
+            cache_v = cache_v.to(self.dtype)
+        if self.store_dtype != self.dtype:
+            cache_k = cache_k.view(self.store_dtype)
+            cache_v = cache_v.view(self.store_dtype)
+
+        if self._shard_extend_active:
+            # Stage the full chunk into this layer's slot on the compute
+            # stream: ordered before this layer's attention read for free
+            # and disjoint from the gather stream, which only writes the
+            # prefix region. Padded locations translate to the trash page.
+            slot = self._slots[layer_id % 2]
+            rows = self._translate_loc_cached(loc)
+            slot.tensors["k"][rows] = cache_k
+            slot.tensors["v"][rows] = cache_v
+
+        owned_idx, local_rows = self._get_write_plan(loc)
+        if owned_idx.numel() > 0:
+            self._store_kv_layer(
+                layer_id - self.start_layer,
+                local_rows,
+                cache_k.index_select(0, owned_idx),
+                cache_v.index_select(0, owned_idx),
+            )
+
+    def get_key_buffer(self, layer_id: int):
+        if self._shard_extend_active:
+            slot = self._acquire_slot_for_read(layer_id)
+            k = slot.tensors["k"]
+            return k.view(self.dtype) if self.store_dtype != self.dtype else k
+        return super().get_key_buffer(layer_id)
+
+    def get_value_buffer(self, layer_id: int):
+        if self._shard_extend_active:
+            slot = self._acquire_slot_for_read(layer_id)
+            v = slot.tensors["v"]
+            return v.view(self.dtype) if self.store_dtype != self.dtype else v
+        return super().get_value_buffer(layer_id)
+
+
+class PageInterleaveMLATokenToKVPool(PageInterleaveKVPoolMixin, MLATokenToKVPool):
+    """MLA latent pool striped across its shard group (the attn-CP group when
+    prefill CP is active, the attn-TP group otherwise — see
+    ``page_interleave.get_kv_shard_group``).
+
+    The latent KV reaching the write path is identical on every shard-group
+    rank (``ReplicatedLinear`` projection across attn-TP; the prefill CP
+    strategy's latent-KV rebuild across attn-CP), so the write filter needs no
+    compute-time communication. Read consumers, all served from the assembled
+    scratch: the chunked-prefix MHA path fetches prefix rows through
+    ``get_mla_kv_buffer``; the absorbed-MLA and one-shot paths read
+    ``get_key_buffer``/``get_value_buffer`` through the translated page
+    table. The current chunk is staged into the slot at write time so the
+    page-table consumers cover ``[prefix | chunk]`` uniformly.
+    """
+
+    def __init__(
+        self,
+        *args,
+        shard_spec: PageShardSpec,
+        shard_group: GroupCoordinator,
+        **kwargs,
+    ):
+        super().__init__(*args, **kwargs)
+        assert not self.use_dsa, "DSA models are not supported by KV sharding yet"
+        self._init_page_shard_state(shard_spec, shard_group)
+
+    def _scratch_tensor_specs(self, rows: int) -> Dict[str, torch.Tensor]:
+        return {
+            "kv": torch.zeros(
+                (rows, 1, self.kv_cache_dim),
+                dtype=self.store_dtype,
+                device=self.device,
+            ),
+        }
+
+    def _gather_pairs(self, local_layer: int):
+        return [(self.kv_buffer[local_layer], "kv")]
+
+    def _scratch_kv(self, slot: _ScratchSlot) -> torch.Tensor:
+        kv = slot.tensors["kv"]
+        if self.store_dtype != self.dtype:
+            return kv.view(self.dtype)
+        return kv
+
+    def set_kv_buffer(
+        self,
+        layer: RadixAttention,
+        loc_info,
+        cache_k: torch.Tensor,
+        cache_v: torch.Tensor,
+        layer_id_override: Optional[int] = None,
+    ):
+        loc, _, _ = unwrap_write_loc(loc_info)
+        layer_id = (
+            layer_id_override if layer_id_override is not None else layer.layer_id
+        )
+        if self._shard_extend_active:
+            # Stage the full chunk into this layer's slot (compute
+            # stream): the absorbed / one-shot readers cover the current
+            # chunk through the translated page table too.
+            slot = self._slots[layer_id % 2]
+            rows = self._translate_loc_cached(loc)
+            staged_k = cache_k
+            if staged_k.dtype != self.dtype:
+                staged_k = staged_k.to(self.dtype)
+            self._scratch_kv(slot)[rows] = staged_k
+        owned_idx, local_rows = self._get_write_plan(loc)
+        if owned_idx.numel() == 0:
+            return
+        super().set_kv_buffer(
+            layer,
+            local_rows,
+            cache_k.index_select(0, owned_idx),
+            cache_v,
+            layer_id_override=layer_id_override,
+        )
+
+    def set_mla_kv_buffer(
+        self,
+        layer: RadixAttention,
+        loc: torch.Tensor,
+        cache_k_nope: torch.Tensor,
+        cache_k_rope: torch.Tensor,
+        layer_id_override: Optional[int] = None,
+    ):
+        layer_id = (
+            layer_id_override if layer_id_override is not None else layer.layer_id
+        )
+        if self._shard_extend_active:
+            slot = self._slots[layer_id % 2]
+            rows = self._translate_loc_cached(loc)
+            staged_nope, staged_rope = cache_k_nope, cache_k_rope
+            if staged_nope.dtype != self.dtype:
+                staged_nope = staged_nope.to(self.dtype)
+                staged_rope = staged_rope.to(self.dtype)
+            if self.store_dtype != self.dtype:
+                staged_nope = staged_nope.view(self.store_dtype)
+                staged_rope = staged_rope.view(self.store_dtype)
+            set_mla_kv_buffer_triton(slot.tensors["kv"], rows, staged_nope, staged_rope)
+        owned_idx, local_rows = self._get_write_plan(loc)
+        if owned_idx.numel() == 0:
+            return
+        super().set_mla_kv_buffer(
+            layer,
+            local_rows,
+            cache_k_nope.index_select(0, owned_idx),
+            cache_k_rope.index_select(0, owned_idx),
+            layer_id_override=layer_id_override,
+        )
+
+    def get_kv_buffer_shape(self):
+        # Shape probes (e.g. the eager runner's DCP-metadata prep) must not
+        # route through the attention getters below — they may run before
+        # this batch's metadata build while the previous batch's shard-extend
+        # flag is still set.
+        k = self.kv_buffer[0]
+        return k.shape, k[..., : self.kv_lora_rank].shape
+
+    def get_key_buffer(self, layer_id: int):
+        # During a sharded extend the pool holds only this rank's stripe;
+        # page-table readers (absorbed MLA, incl. the CP zigzag wrapper) get
+        # the assembled scratch — metadata.page_table is already translated
+        # to scratch rows.
+        if self._shard_extend_active:
+            return self._scratch_kv(self._acquire_slot_for_read(layer_id))
+        return super().get_key_buffer(layer_id)
+
+    def get_value_buffer(self, layer_id: int):
+        if self._shard_extend_active:
+            kv = self._scratch_kv(self._acquire_slot_for_read(layer_id))
+            return kv[..., : self.kv_lora_rank]
+        return super().get_value_buffer(layer_id)
+
+    def get_mla_kv_buffer(
+        self,
+        layer: RadixAttention,
+        loc: torch.Tensor,
+        dst_dtype: Optional[torch.dtype] = None,
+    ):
+        slot = self._acquire_slot_for_read(layer.layer_id)
+        rows = self._translate_loc_cached(loc)
+        kv_buffer = slot.tensors["kv"]
+        if self.store_dtype != self.dtype:
+            kv_buffer = kv_buffer.view(self.dtype)
+        dst_dtype = dst_dtype or self.dtype
+        cache_k_nope = torch.empty(
+            (loc.shape[0], 1, self.kv_lora_rank),
+            dtype=dst_dtype,
+            device=kv_buffer.device,
+        )
+        cache_k_rope = torch.empty(
+            (loc.shape[0], 1, self.qk_rope_head_dim),
+            dtype=dst_dtype,
+            device=kv_buffer.device,
+        )
+        get_mla_kv_buffer_triton(kv_buffer, rows, cache_k_nope, cache_k_rope)
+        return cache_k_nope, cache_k_rope