sglang/docs/superpowers/plans/2026-04-24-nsa-prefill-cp-shared-kv-phase2.md

# NSA Prefill CP Shared KV Phase 2 Implementation Plan

> **For agentic workers:** REQUIRED SUB-SKILL: Use superpowers:subagent-driven-development (recommended) or superpowers:executing-plans to implement this plan task-by-task. Steps use checkbox (`- [ ]`) syntax for tracking.

**Goal:** Implement Phase 2 shared/sharded persistent KV for NSA prefill CP so each prefill CP rank keeps its physical KV pool size while the CP group exposes an expanded logical KV capacity.

**Architecture:** Introduce a CP shared KV layout layer that separates logical KV locations from per-rank physical KV locations. Keep scheduler, radix, and req-to-token state in logical loc space; translate to physical locs only at KV pool writes, compatibility reads, and Mooncake PD transfer. Phase 2 intentionally keeps an attention full-view compatibility path and defers shard-aware runtime attention to Phase 3.

**Tech Stack:** Python, PyTorch, Triton-backed SGLang KV pools, SGLang scheduler/radix cache, NSA/MLA attention backend, Mooncake PD transfer, pytest/unittest.

---

## 0. Review Mapping From Design To Current Code

| Design area | Current code | Required change |
| --- | --- | --- |
| Config gate | `python/sglang/srt/server_args.py` fields around `enable_nsa_prefill_context_parallel`; parser near `--enable-nsa-prefill-context-parallel` | Add `--enable-nsa-prefill-cp-shared-kv`; validate NSA+MLA, prefill CP, `in-seq-split`, `page_size=64`, prefill disagg only, Mooncake/all-CP-transfer when PD is used. |
| Physical KV sizing | `python/sglang/srt/model_executor/model_runner_kv_cache_mixin.py::profile_max_num_token`, `_resolve_memory_pool_config`, `_apply_memory_pool_config`, `_init_pools` | Preserve physical profiled token count for `NSATokenToKVPool(size=...)`; expose logical capacity as `physical * attn_cp_size` to scheduler. |
| Persistent KV allocation | `python/sglang/srt/mem_cache/memory_pool.py::NSATokenToKVPool`, `MLATokenToKVPool._create_buffers` | Allocate physical per-rank buffers only; log physical tokens and logical tokens separately. |
| Logical/physical loc mapping | No dedicated file today | Add `python/sglang/srt/mem_cache/cp_shared_kv_layout.py`. |
| Allocator | `python/sglang/srt/mem_cache/allocator.py::PagedTokenToKVPoolAllocator` | Add CP shared wrapper/allocator returning logical locs and tracking logical free pages while checking deterministic physical ownership. |
| Extend allocation | `python/sglang/srt/mem_cache/common.py::alloc_for_extend`, `write_cache_indices` | Keep writing logical locs into `req_to_token_pool`; do not pass logical locs directly to physical KV writes. |
| Forward batch metadata | `python/sglang/srt/model_executor/forward_batch_info.py::ForwardBatch` | Add layout and logical/physical helper fields so attention code can translate locs. |
| KV writes | `python/sglang/srt/layers/attention/nsa_backend.py` calls to `set_mla_kv_buffer`; `python/sglang/srt/layers/attention/nsa/nsa_indexer.py::set_index_k_scale_buffer` | Filter current full chunk to owner tokens and write only physical locs on this rank. |
| Attention compatibility reads | `nsa_backend.py` `dequantize_k_cache_paged`; `nsa_indexer.py::get_index_k_continuous`; page tables from `req_to_token_pool` | Add explicit Phase2 compatibility helper that can materialize full logical view; keep runtime workspace risk visible. |
| Mooncake transfer chunk | `python/sglang/srt/disaggregation/mooncake/conn.py::TransferKVChunk`, `CommonKVSender.send`, Mooncake worker uses `req.dst_kv_indices[kv_chunk.index_slice]` | Replace `index_slice` with explicit `logical_page_positions`; send source physical pages; choose decode dst pages by absolute request page positions. |
| Prefill transfer source | `python/sglang/srt/disaggregation/prefill.py::send_kv_chunk` | Interpret `req_to_token` as logical locs; filter owner pages; convert to source physical page ids. |
| Decode transfer dst | `python/sglang/srt/disaggregation/decode.py` preallocation | Keep decode full/non-CP physical layout; receiver gets full dst page list; Mooncake selects dst pages via `logical_page_positions`. |
| Metrics/logs | `scheduler.py`, `model_runner_kv_cache_mixin.py`, existing KV allocation log in `memory_pool.py` | Log physical tokens, logical tokens, cp size, shard policy, full-view materialization bytes. |

## File Structure

Create:

- `python/sglang/srt/mem_cache/cp_shared_kv_layout.py`
  - Pure mapping helper. No CUDA allocation. Unit-testable on CPU.
- `test/registered/unit/mem_cache/test_cp_shared_kv_layout.py`
  - Tests dummy page handling, owner mapping, logical-to-physical conversion, NumPy page filtering.
- `test/registered/unit/disaggregation/test_cp_shared_kv_transfer_mapping.py`
  - Tests Mooncake-style logical page positions and source/destination page mapping without network.

Modify:

- `python/sglang/srt/server_args.py`
  - Add flag and validation.
- `python/sglang/srt/model_executor/model_runner.py`
  - Carry shared KV enablement and layout to the runner.
- `python/sglang/srt/model_executor/model_runner_kv_cache_mixin.py`
  - Split physical/logical memory pool capacity and initialize CP shared allocator.
- `python/sglang/srt/model_executor/forward_batch_info.py`
  - Attach layout and keep `out_cache_loc` documented as logical when shared KV is enabled.
- `python/sglang/srt/mem_cache/allocator.py`
  - Add `CPSharedPagedTokenToKVPoolAllocator`.
- `python/sglang/srt/mem_cache/common.py`
  - Keep allocation/write paths logical; add assertions when shared KV is enabled.
- `python/sglang/srt/mem_cache/memory_pool.py`
  - Add log path for physical/logical tokens; avoid changing buffer shape beyond physical size.
- `python/sglang/srt/layers/attention/nsa_backend.py`
  - Translate and filter locs for MLA KV writes; call compatibility helper before physical reads where needed.
- `python/sglang/srt/layers/attention/nsa/nsa_indexer.py`
  - Translate and filter locs for NSA index K writes; use compatibility path for prefix index reads.
- `python/sglang/srt/disaggregation/utils.py`
  - Add CP shared page filtering utilities. Keep old contiguous CP filtering for replicated mode.
- `python/sglang/srt/disaggregation/prefill.py`
  - Source logical pages from `req_to_token`; send only owner physical pages with absolute logical page positions.
- `python/sglang/srt/disaggregation/mooncake/conn.py`
  - Replace chunk `index_slice` semantics for shared KV; keep old behavior when disabled.
- `python/sglang/srt/disaggregation/common/conn.py`
  - Enforce all CP ranks transfer for shared KV and expose shared KV mode to sender.
- `python/sglang/srt/disaggregation/decode.py`
  - Keep full decode layout; ensure preallocated `dst_kv_indices` are selected by explicit positions.

---

### Task 1: Add Pure CP Shared KV Layout Helper

**Files:**
- Create: `python/sglang/srt/mem_cache/cp_shared_kv_layout.py`
- Create: `test/registered/unit/mem_cache/test_cp_shared_kv_layout.py`

- [ ] **Step 1: Write failing layout tests**

Create `test/registered/unit/mem_cache/test_cp_shared_kv_layout.py`:

```python
import unittest

import numpy as np
import torch

from sglang.srt.mem_cache.cp_shared_kv_layout import CpSharedKVLayout
from sglang.test.ci.ci_register import register_cpu_ci

register_cpu_ci(est_time=1, suite="stage-a-test-cpu")


class TestCpSharedKVLayout(unittest.TestCase):
    def test_page_owner_skips_dummy_page(self):
        layout = CpSharedKVLayout(page_size=64, cp_size=4, cp_rank=0)
        pages = torch.tensor([0, 1, 2, 3, 4, 5, 8, 9], dtype=torch.int64)
        owners = layout.owner_for_logical_pages(pages)
        self.assertEqual(owners.tolist(), [-1, 0, 1, 2, 3, 0, 3, 0])

    def test_logical_to_physical_pages_keeps_dummy_zero(self):
        layout = CpSharedKVLayout(page_size=64, cp_size=4, cp_rank=0)
        pages = torch.tensor([0, 1, 2, 3, 4, 5, 8, 9], dtype=torch.int64)
        physical = layout.logical_pages_to_physical(pages)
        self.assertEqual(physical.tolist(), [0, 1, 1, 1, 1, 2, 2, 3])

    def test_owned_mask_and_loc_translation(self):
        layout = CpSharedKVLayout(page_size=64, cp_size=4, cp_rank=2)
        locs = torch.tensor([0, 64, 128, 192, 256, 320, 384], dtype=torch.int64)
        mask = layout.owned_by_this_rank(locs)
        self.assertEqual(mask.tolist(), [False, False, False, True, False, False, False])
        physical = layout.logical_locs_to_physical(locs[mask])
        self.assertEqual(physical.tolist(), [64])

    def test_numpy_filter_returns_request_absolute_positions(self):
        layout = CpSharedKVLayout(page_size=64, cp_size=4, cp_rank=1)
        logical_pages = np.array([1, 2, 3, 4, 5, 6, 7, 8, 9], dtype=np.int32)
        request_positions = np.arange(10, 19, dtype=np.int32)
        src_physical, positions = layout.filter_owned_pages_np(
            logical_pages, request_positions
        )
        self.assertEqual(src_physical.tolist(), [1, 2])
        self.assertEqual(positions.tolist(), [11, 15])


if __name__ == "__main__":
    unittest.main()
```

- [ ] **Step 2: Run tests and verify failure**

Run:

```bash
python3 -m pytest test/registered/unit/mem_cache/test_cp_shared_kv_layout.py -q
```

Expected: FAIL with `ModuleNotFoundError: No module named 'sglang.srt.mem_cache.cp_shared_kv_layout'`.

- [ ] **Step 3: Implement layout helper**

Create `python/sglang/srt/mem_cache/cp_shared_kv_layout.py`:

```python
from __future__ import annotations

from dataclasses import dataclass

import numpy as np
import torch


@dataclass(frozen=True)
class CpSharedKVLayout:
    """Maps CP-group logical KV locations to per-rank physical KV locations.

    Page 0 is the existing dummy/padding page and is never owned by a real CP
    rank. Usable pages start at page 1.
    """

    page_size: int
    cp_size: int
    cp_rank: int

    def __post_init__(self):
        if self.page_size <= 0:
            raise ValueError(f"page_size must be positive, got {self.page_size}")
        if self.cp_size <= 0:
            raise ValueError(f"cp_size must be positive, got {self.cp_size}")
        if not 0 <= self.cp_rank < self.cp_size:
            raise ValueError(
                f"cp_rank must be in [0, {self.cp_size}), got {self.cp_rank}"
            )

    def owner_for_logical_pages(self, logical_pages: torch.Tensor) -> torch.Tensor:
        owners = torch.remainder(logical_pages - 1, self.cp_size)
        return torch.where(logical_pages == 0, torch.full_like(owners, -1), owners)

    def owned_pages_mask(self, logical_pages: torch.Tensor) -> torch.Tensor:
        return self.owner_for_logical_pages(logical_pages) == self.cp_rank

    def owned_by_this_rank(self, logical_locs: torch.Tensor) -> torch.Tensor:
        logical_pages = torch.div(logical_locs, self.page_size, rounding_mode="floor")
        return self.owned_pages_mask(logical_pages)

    def logical_pages_to_physical(self, logical_pages: torch.Tensor) -> torch.Tensor:
        physical_pages = torch.div(
            logical_pages - 1, self.cp_size, rounding_mode="floor"
        ) + 1
        return torch.where(logical_pages == 0, torch.zeros_like(physical_pages), physical_pages)

    def logical_locs_to_physical(self, logical_locs: torch.Tensor) -> torch.Tensor:
        logical_pages = torch.div(logical_locs, self.page_size, rounding_mode="floor")
        offsets = torch.remainder(logical_locs, self.page_size)
        return self.logical_pages_to_physical(logical_pages) * self.page_size + offsets

    def owner_for_logical_pages_np(self, logical_pages: np.ndarray) -> np.ndarray:
        pages = np.asarray(logical_pages, dtype=np.int64)
        owners = (pages - 1) % self.cp_size
        owners = np.where(pages == 0, -1, owners)
        return owners.astype(np.int32, copy=False)

    def logical_pages_to_physical_np(self, logical_pages: np.ndarray) -> np.ndarray:
        pages = np.asarray(logical_pages, dtype=np.int64)
        physical = (pages - 1) // self.cp_size + 1
        physical = np.where(pages == 0, 0, physical)
        return physical.astype(np.int32, copy=False)

    def filter_owned_pages_np(
        self,
        logical_pages: np.ndarray,
        request_page_positions: np.ndarray,
    ) -> tuple[np.ndarray, np.ndarray]:
        pages = np.asarray(logical_pages, dtype=np.int64)
        positions = np.asarray(request_page_positions, dtype=np.int64)
        if pages.shape != positions.shape:
            raise ValueError(
                f"logical_pages and request_page_positions must have the same shape, "
                f"got {pages.shape} and {positions.shape}"
            )
        mask = self.owner_for_logical_pages_np(pages) == self.cp_rank
        return (
            self.logical_pages_to_physical_np(pages[mask]),
            positions[mask].astype(np.int32, copy=False),
        )
```

- [ ] **Step 4: Run layout tests**

Run:

```bash
python3 -m pytest test/registered/unit/mem_cache/test_cp_shared_kv_layout.py -q
```

Expected: PASS.

- [ ] **Step 5: Commit**

```bash
git add python/sglang/srt/mem_cache/cp_shared_kv_layout.py test/registered/unit/mem_cache/test_cp_shared_kv_layout.py
git commit -m "feat(cp): add shared KV layout mapping"
```

---

### Task 2: Add Server Flag And Validation

**Files:**
- Modify: `python/sglang/srt/server_args.py`
- Test: `test/registered/unit/server_args/test_server_args.py`

- [ ] **Step 1: Write failing parser test**

Append to `test/registered/unit/server_args/test_server_args.py`:

```python
def test_enable_nsa_prefill_cp_shared_kv_parser_flag():
    import argparse

    from sglang.srt.server_args import ServerArgs

    parser = argparse.ArgumentParser()
    ServerArgs.add_cli_args(parser)
    raw_args = parser.parse_args(
        [
            "--model-path",
            "dummy",
            "--enable-nsa-prefill-context-parallel",
            "--enable-nsa-prefill-cp-shared-kv",
            "--nsa-prefill-cp-mode",
            "in-seq-split",
        ]
    )
    args = ServerArgs.from_cli_args(raw_args)
    assert args.enable_nsa_prefill_cp_shared_kv is True
```

- [ ] **Step 2: Run test and verify failure**

Run:

```bash
python3 -m pytest test/registered/unit/server_args/test_server_args.py -q -k shared_kv
```

Expected: FAIL because the CLI flag or dataclass field does not exist.

- [ ] **Step 3: Add dataclass field and parser argument**

Modify `python/sglang/srt/server_args.py` near the NSA CP fields:

```python
# Context parallelism used in the long sequence prefill phase of DeepSeek v3.2
enable_nsa_prefill_context_parallel: bool = False
nsa_prefill_cp_mode: str = "round-robin-split"
enable_nsa_prefill_cp_shared_kv: bool = False
```

Add parser argument near `--enable-nsa-prefill-context-parallel`:

```python
parser.add_argument(
    "--enable-nsa-prefill-cp-shared-kv",
    action="store_true",
    help=(
        "Enable Phase 2 shared/sharded persistent KV pool for NSA prefill CP. "
        "Only prefill CP with NSA+MLA is supported; decode CP remains disabled."
    ),
)
```

- [ ] **Step 4: Add validation gate in NSA model-specific adjustment**

Inside the DeepSeek DSA branch after existing `enable_nsa_prefill_context_parallel` checks, add:

```python
if self.enable_nsa_prefill_cp_shared_kv:
    assert self.enable_nsa_prefill_context_parallel, (
        "--enable-nsa-prefill-cp-shared-kv requires "
        "--enable-nsa-prefill-context-parallel."
    )
    assert self.nsa_prefill_cp_mode == "in-seq-split", (
        "Phase 2 shared KV is initially validated only with "
        "--nsa-prefill-cp-mode in-seq-split. The layout is mode-neutral, "
        "but round-robin runtime wiring is not enabled yet."
    )
    assert self.disaggregation_mode != "decode", (
        "Phase 2 shared KV supports prefill CP only; decode CP/shared KV is not supported."
    )
    assert self.page_size == 64, "Phase 2 shared KV requires page_size=64 for NSA."
    if self.disaggregation_mode == "prefill":
        from sglang.srt.environ import envs

        assert envs.SGLANG_DISAGGREGATION_ALL_CP_RANKS_TRANSFER.get(), (
            "Phase 2 shared KV with PD disaggregation requires "
            "SGLANG_DISAGGREGATION_ALL_CP_RANKS_TRANSFER=1 so all prefill CP ranks transfer shards."
        )
```

- [ ] **Step 5: Run parser test**

Run:

```bash
python3 -m pytest test/registered/unit/server_args/test_server_args.py -q -k shared_kv
```

Expected: PASS.

- [ ] **Step 6: Commit**

```bash
git add python/sglang/srt/server_args.py test/registered/unit/server_args/test_server_args.py
git commit -m "feat(cp): add NSA prefill shared KV flag"
```

---

### Task 3: Split Physical And Logical KV Capacity In ModelRunner

**Files:**
- Modify: `python/sglang/srt/model_executor/model_runner_kv_cache_mixin.py`
- Modify: `python/sglang/srt/model_executor/model_runner.py`
- Modify: `python/sglang/srt/managers/tp_worker.py`
- Modify: `python/sglang/srt/managers/scheduler.py`

- [ ] **Step 1: Add capacity fields to `MemoryPoolConfig`**

Modify `MemoryPoolConfig` in `model_runner_kv_cache_mixin.py`:

```python
@dataclass
class MemoryPoolConfig:
    max_total_num_tokens: int
    max_running_requests: int
    physical_max_total_num_tokens: Optional[int] = None
    full_max_total_num_tokens: Optional[int] = None
    swa_max_total_num_tokens: Optional[int] = None
    mem_fraction_static: Optional[float] = None

    def __post_init__(self):
        if self.physical_max_total_num_tokens is None:
            self.physical_max_total_num_tokens = self.max_total_num_tokens
        if self.max_total_num_tokens <= 0 or self.physical_max_total_num_tokens <= 0:
            msg = "Not enough memory. Please try to increase --mem-fraction-static."
            if self.mem_fraction_static is not None:
                msg += f" Current value: mem_fraction_static={self.mem_fraction_static}"
            raise RuntimeError(msg)
```

- [ ] **Step 2: Resolve logical capacity when shared KV is enabled**

In `_resolve_memory_pool_config`, after `token_capacity = self._resolve_token_capacity(profiled_tokens)`, add:

```python
physical_token_capacity = token_capacity
if self.server_args.enable_nsa_prefill_cp_shared_kv:
    assert self.server_args.page_size == self.page_size
    physical_pages = physical_token_capacity // self.page_size
    logical_token_capacity = physical_pages * self.server_args.attn_cp_size * self.page_size
    token_capacity = logical_token_capacity
```

Return both values:

```python
return MemoryPoolConfig(
    max_total_num_tokens=token_capacity,
    physical_max_total_num_tokens=physical_token_capacity,
    max_running_requests=self._resolve_max_num_reqs(token_capacity),
    full_max_total_num_tokens=full_tokens,
    swa_max_total_num_tokens=swa_tokens,
    mem_fraction_static=self.server_args.mem_fraction_static,
)
```

- [ ] **Step 3: Apply both capacities**

In `_apply_memory_pool_config`, set:

```python
self.max_total_num_tokens = config.max_total_num_tokens
self.physical_max_total_num_tokens = config.physical_max_total_num_tokens
```

In `ModelRunner.__init__`, initialize fallback:

```python
self.physical_max_total_num_tokens = None
```

- [ ] **Step 4: Ensure scheduler sees logical capacity**

Keep these existing paths using `self.model_runner.max_total_num_tokens`:

- `tp_worker.py` sets `self.max_total_num_tokens`
- `scheduler.py` logs `max_total_num_tokens`
- metrics use scheduler `max_total_num_tokens`

Add a log in `init_memory_pool` after `_apply_memory_pool_config`:

```python
if self.server_args.enable_nsa_prefill_cp_shared_kv:
    logger.info(
        "CP shared KV enabled. physical_tokens_per_rank=%s, logical_tokens=%s, cp_size=%s, shard_policy=page_interleaved",
        self.physical_max_total_num_tokens,
        self.max_total_num_tokens,
        self.server_args.attn_cp_size,
    )
```

- [ ] **Step 5: Run import smoke test**

Run:

```bash
python3 -m py_compile python/sglang/srt/model_executor/model_runner_kv_cache_mixin.py python/sglang/srt/model_executor/model_runner.py python/sglang/srt/managers/tp_worker.py python/sglang/srt/managers/scheduler.py
```

Expected: no output, exit code 0.

- [ ] **Step 6: Commit**

```bash
git add python/sglang/srt/model_executor/model_runner_kv_cache_mixin.py python/sglang/srt/model_executor/model_runner.py python/sglang/srt/managers/tp_worker.py python/sglang/srt/managers/scheduler.py
git commit -m "feat(cp): split logical and physical KV capacity"
```

---

### Task 4: Add CP Shared Paged Allocator Returning Logical Locs

**Files:**
- Modify: `python/sglang/srt/mem_cache/allocator.py`
- Modify: `python/sglang/srt/model_executor/model_runner_kv_cache_mixin.py`
- Test: `test/registered/unit/mem_cache/test_cp_shared_kv_layout.py`

- [ ] **Step 1: Extend tests for logical allocation capacity**

Append to `test_cp_shared_kv_layout.py`:

```python
class TestCPSharedPagedAllocator(unittest.TestCase):
    def test_shared_allocator_exposes_logical_capacity(self):
        from sglang.srt.mem_cache.allocator import CPSharedPagedTokenToKVPoolAllocator

        allocator = CPSharedPagedTokenToKVPoolAllocator(
            logical_size=64 * 8,
            physical_size=64 * 2,
            page_size=64,
            dtype=torch.bfloat16,
            device="cpu",
            kvcache=None,
            need_sort=False,
            cp_size=4,
            cp_rank=0,
        )
        self.assertEqual(allocator.available_size(), 64 * 8)
        locs = allocator.alloc(64 * 2)
        self.assertEqual(locs.numel(), 64 * 2)
        self.assertEqual(allocator.available_size(), 64 * 6)
        allocator.free(locs)
        self.assertEqual(allocator.available_size(), 64 * 8)
```

- [ ] **Step 2: Run test and verify failure**

Run:

```bash
python3 -m pytest test/registered/unit/mem_cache/test_cp_shared_kv_layout.py -q -k shared_allocator
```

Expected: FAIL because `CPSharedPagedTokenToKVPoolAllocator` does not exist.

- [ ] **Step 3: Implement allocator**

Add to `python/sglang/srt/mem_cache/allocator.py` after `PagedTokenToKVPoolAllocator`:

```python
class CPSharedPagedTokenToKVPoolAllocator(PagedTokenToKVPoolAllocator):
    """Paged allocator that returns CP-group logical KV locs.

    It tracks logical pages. The physical KV pool is smaller and is addressed by
    CpSharedKVLayout only at actual KV buffer access time.
    """

    def __init__(
        self,
        logical_size: int,
        physical_size: int,
        page_size: int,
        dtype: torch.dtype,
        device: str,
        kvcache: KVCache,
        need_sort: bool,
        cp_size: int,
        cp_rank: int,
    ):
        if logical_size % page_size != 0:
            raise ValueError("logical_size must be page aligned")
        if physical_size % page_size != 0:
            raise ValueError("physical_size must be page aligned")
        if logical_size != physical_size * cp_size:
            raise ValueError(
                f"logical_size must equal physical_size * cp_size, got "
                f"{logical_size=} {physical_size=} {cp_size=}"
            )
        super().__init__(logical_size, page_size, dtype, device, kvcache, need_sort)
        self.physical_size = physical_size
        self.cp_size = cp_size
        self.cp_rank = cp_rank
```

This initial implementation relies on `PagedTokenToKVPoolAllocator` tracking logical pages. Physical safety comes from deterministic owner mapping and the exact `logical_size == physical_size * cp_size` invariant.

- [ ] **Step 4: Wire allocator into `_init_pools`**

In `model_runner_kv_cache_mixin.py`, import the class:

```python
from sglang.srt.mem_cache.allocator import (
    CPSharedPagedTokenToKVPoolAllocator,
    PagedTokenToKVPoolAllocator,
    TokenToKVPoolAllocator,
)
```

In `_init_pools`, compute pool size:

```python
physical_kv_pool_size = (
    self.physical_max_total_num_tokens
    if self.server_args.enable_nsa_prefill_cp_shared_kv
    else self.max_total_num_tokens
)
```

Pass `physical_kv_pool_size` to `NSATokenToKVPool(size=...)` while keeping `self.max_total_num_tokens` logical.

For allocator creation when `self.page_size != 1`, use:

```python
if self.server_args.enable_nsa_prefill_cp_shared_kv:
    self.token_to_kv_pool_allocator = CPSharedPagedTokenToKVPoolAllocator(
        logical_size=self.max_total_num_tokens,
        physical_size=self.physical_max_total_num_tokens,
        page_size=self.page_size,
        dtype=self.kv_cache_dtype,
        device=self.device,
        kvcache=self.token_to_kv_pool,
        need_sort=need_sort,
        cp_size=self.server_args.attn_cp_size,
        cp_rank=self.tp_rank % self.server_args.attn_cp_size,
    )
else:
    self.token_to_kv_pool_allocator = PagedTokenToKVPoolAllocator(
        self.max_total_num_tokens,
        page_size=self.page_size,
        dtype=self.kv_cache_dtype,
        device=self.device,
        kvcache=self.token_to_kv_pool,
        need_sort=need_sort,
    )
```

Use `get_attention_cp_rank()` from `sglang.srt.distributed.parallel_state` for the `cp_rank` argument. Do not derive it from `tp_rank` manually.

- [ ] **Step 5: Run allocator tests**

Run:

```bash
python3 -m pytest test/registered/unit/mem_cache/test_cp_shared_kv_layout.py -q
```

Expected: PASS.

- [ ] **Step 6: Commit**

```bash
git add python/sglang/srt/mem_cache/allocator.py python/sglang/srt/model_executor/model_runner_kv_cache_mixin.py test/registered/unit/mem_cache/test_cp_shared_kv_layout.py
git commit -m "feat(cp): add shared KV logical allocator"
```

---

### Task 5: Attach CP Shared KV Layout To ModelRunner And ForwardBatch

**Files:**
- Modify: `python/sglang/srt/model_executor/model_runner.py`
- Modify: `python/sglang/srt/model_executor/forward_batch_info.py`

- [ ] **Step 1: Add fields to `ForwardBatch`**

Modify `ForwardBatch` dataclass in `forward_batch_info.py`:

```python
from sglang.srt.mem_cache.cp_shared_kv_layout import CpSharedKVLayout
```

Add fields near `nsa_cp_metadata`:

```python
uses_cp_shared_kv: bool = False
cp_shared_kv_layout: Optional[CpSharedKVLayout] = None
```

- [ ] **Step 2: Initialize layout in `ModelRunner`**

In `ModelRunner.__init__`, add:

```python
self.uses_cp_shared_kv = server_args.enable_nsa_prefill_cp_shared_kv
self.cp_shared_kv_layout = None
```

After distributed groups are initialized and before forward batches are created, initialize:

```python
if self.uses_cp_shared_kv:
    from sglang.srt.distributed.parallel_state import get_attention_cp_rank
    from sglang.srt.mem_cache.cp_shared_kv_layout import CpSharedKVLayout

    self.cp_shared_kv_layout = CpSharedKVLayout(
        page_size=self.page_size,
        cp_size=self.server_args.attn_cp_size,
        cp_rank=get_attention_cp_rank(),
    )
```

- [ ] **Step 3: Copy layout into `ForwardBatch.init_new`**

In `ForwardBatch.init_new`, after `ret = cls(...)`:

```python
ret.uses_cp_shared_kv = model_runner.uses_cp_shared_kv
ret.cp_shared_kv_layout = model_runner.cp_shared_kv_layout
```

- [ ] **Step 4: Run compile check**

Run:

```bash
python3 -m py_compile python/sglang/srt/model_executor/model_runner.py python/sglang/srt/model_executor/forward_batch_info.py
```

Expected: no output, exit code 0.

- [ ] **Step 5: Commit**

```bash
git add python/sglang/srt/model_executor/model_runner.py python/sglang/srt/model_executor/forward_batch_info.py
git commit -m "feat(cp): attach shared KV layout to forward batches"
```

---

### Task 6: Convert MLA Persistent KV Writes To Owner-Only Physical Writes

**Files:**
- Modify: `python/sglang/srt/layers/attention/nsa_backend.py`
- Optional helper modify: `python/sglang/srt/layers/attention/nsa/utils.py`

- [ ] **Step 1: Add local helper in `nsa_backend.py`**

Add near imports or class helpers:

```python
def _shared_kv_owner_write_args(forward_batch, logical_locs, *token_tensors):
    if not getattr(forward_batch, "uses_cp_shared_kv", False):
        return logical_locs, token_tensors
    layout = forward_batch.cp_shared_kv_layout
    assert layout is not None
    mask = layout.owned_by_this_rank(logical_locs)
    physical_locs = layout.logical_locs_to_physical(logical_locs[mask])
    return physical_locs.contiguous(), tuple(t[mask].contiguous() for t in token_tensors)
```

- [ ] **Step 2: Update all NSA backend `set_mla_kv_buffer` sites**

For each `set_mla_kv_buffer` in `nsa_backend.py`, replace:

```python
forward_batch.token_to_kv_pool.set_mla_kv_buffer(layer, cache_loc, k, k_rope)
```

with:

```python
cache_loc, (k_to_store, k_rope_to_store) = _shared_kv_owner_write_args(
    forward_batch, cache_loc, k, k_rope
)
if cache_loc.numel() > 0:
    forward_batch.token_to_kv_pool.set_mla_kv_buffer(
        layer, cache_loc, k_to_store, k_rope_to_store
    )
```

Apply to the known sites around current lines:

- `nsa_backend.py:1303`
- `nsa_backend.py:1501`
- `nsa_backend.py:1961`

- [ ] **Step 3: Run compile check**

Run:

```bash
python3 -m py_compile python/sglang/srt/layers/attention/nsa_backend.py
```

Expected: no output, exit code 0.

- [ ] **Step 4: Commit**

```bash
git add python/sglang/srt/layers/attention/nsa_backend.py
git commit -m "feat(cp): shard NSA MLA KV writes"
```

---

### Task 7: Convert NSA Index K Persistent Writes To Owner-Only Physical Writes

**Files:**
- Modify: `python/sglang/srt/layers/attention/nsa/nsa_indexer.py`

- [ ] **Step 1: Add shared KV write filter helper**

In `nsa_indexer.py`, add:

```python
def _shared_kv_owner_index_write_args(forward_batch, key, k_fp8=None, k_scale=None):
    if not getattr(forward_batch, "uses_cp_shared_kv", False):
        out_loc = forward_batch.out_cache_loc
        return out_loc.contiguous() if not out_loc.is_contiguous() else out_loc, key, k_fp8, k_scale
    layout = forward_batch.cp_shared_kv_layout
    assert layout is not None
    logical_locs = forward_batch.out_cache_loc
    mask = layout.owned_by_this_rank(logical_locs)
    physical_locs = layout.logical_locs_to_physical(logical_locs[mask]).contiguous()
    key = key[mask].contiguous() if key is not None else None
    k_fp8 = k_fp8[mask].contiguous() if k_fp8 is not None else None
    k_scale = k_scale[mask].contiguous() if k_scale is not None else None
    return physical_locs, key, k_fp8, k_scale
```

- [ ] **Step 2: Update fused store path**

At current fused `fused_store_index_k_cache(...)` path, replace use of `forward_batch.out_cache_loc` with filtered physical locs:

```python
out_loc, key_to_store, _, _ = _shared_kv_owner_index_write_args(
    forward_batch, key
)
if out_loc.numel() > 0:
    fused_store_index_k_cache(
        key_to_store,
        buf,
        out_loc,
        forward_batch.token_to_kv_pool.page_size,
    )
return
```

- [ ] **Step 3: Update fallback `set_index_k_scale_buffer` path**

Replace:

```python
out_loc = forward_batch.out_cache_loc
if not out_loc.is_contiguous():
    out_loc = out_loc.contiguous()
forward_batch.token_to_kv_pool.set_index_k_scale_buffer(...)
```

with:

```python
out_loc, _, k_fp8, k_scale = _shared_kv_owner_index_write_args(
    forward_batch, key=None, k_fp8=k_fp8, k_scale=k_scale
)
if out_loc.numel() > 0:
    forward_batch.token_to_kv_pool.set_index_k_scale_buffer(
        layer_id=layer_id,
        loc=out_loc,
        index_k=k_fp8,
        index_k_scale=k_scale,
    )
```

- [ ] **Step 4: Run compile check**

Run:

```bash
python3 -m py_compile python/sglang/srt/layers/attention/nsa/nsa_indexer.py
```

Expected: no output, exit code 0.

- [ ] **Step 5: Commit**

```bash
git add python/sglang/srt/layers/attention/nsa/nsa_indexer.py
git commit -m "feat(cp): shard NSA index KV writes"
```

---

### Task 8: Add Mooncake Shared KV Page Mapping Utilities

**Files:**
- Modify: `python/sglang/srt/disaggregation/utils.py`
- Create: `test/registered/unit/disaggregation/test_cp_shared_kv_transfer_mapping.py`

- [ ] **Step 1: Write failing utility tests**

Create `test/registered/unit/disaggregation/test_cp_shared_kv_transfer_mapping.py`:

```python
import unittest

import numpy as np

from sglang.srt.disaggregation.utils import filter_kv_pages_for_cp_shared_kv
from sglang.srt.mem_cache.cp_shared_kv_layout import CpSharedKVLayout
from sglang.test.ci.ci_register import register_cpu_ci

register_cpu_ci(est_time=1, suite="stage-a-test-cpu")


class TestCPSharedKVTransferMapping(unittest.TestCase):
    def test_filter_uses_absolute_request_page_positions(self):
        layout = CpSharedKVLayout(page_size=64, cp_size=4, cp_rank=2)
        logical_pages = np.array([9, 10, 11, 12, 13, 14], dtype=np.int32)
        chunk_page_start = 8
        src_pages, positions = filter_kv_pages_for_cp_shared_kv(
            layout=layout,
            logical_pages=logical_pages,
            chunk_page_start=chunk_page_start,
        )
        self.assertEqual(src_pages.tolist(), [3, 4])
        self.assertEqual(positions.tolist(), [10, 14])


if __name__ == "__main__":
    unittest.main()
```

- [ ] **Step 2: Run test and verify failure**

Run:

```bash
python3 -m pytest test/registered/unit/disaggregation/test_cp_shared_kv_transfer_mapping.py -q
```

Expected: FAIL because `filter_kv_pages_for_cp_shared_kv` does not exist.

- [ ] **Step 3: Implement utility**

Add to `python/sglang/srt/disaggregation/utils.py`:

```python
def filter_kv_pages_for_cp_shared_kv(layout, logical_pages: np.ndarray, chunk_page_start: int):
    """Return source physical pages and request-absolute page positions for this CP rank."""
    logical_pages = np.asarray(logical_pages, dtype=np.int32)
    request_positions = (
        np.arange(len(logical_pages), dtype=np.int32) + np.int32(chunk_page_start)
    )
    return layout.filter_owned_pages_np(logical_pages, request_positions)
```

- [ ] **Step 4: Run mapping tests**

Run:

```bash
python3 -m pytest test/registered/unit/disaggregation/test_cp_shared_kv_transfer_mapping.py -q
```

Expected: PASS.

- [ ] **Step 5: Commit**

```bash
git add python/sglang/srt/disaggregation/utils.py test/registered/unit/disaggregation/test_cp_shared_kv_transfer_mapping.py
git commit -m "feat(cp): add shared KV transfer page mapping"
```

---

### Task 9: Update Mooncake Transfer Chunk From Slice To Explicit Positions

**Files:**
- Modify: `python/sglang/srt/disaggregation/mooncake/conn.py`

- [ ] **Step 1: Extend `TransferKVChunk` dataclass**

Change:

```python
class TransferKVChunk:
    room: int
    prefill_kv_indices: npt.NDArray[np.int32]
    index_slice: slice
    is_last_chunk: bool
    prefill_aux_index: Optional[int]
    state_indices: Optional[List[int]]
```

To:

```python
class TransferKVChunk:
    room: int
    prefill_kv_indices: npt.NDArray[np.int32]
    index_slice: Optional[slice]
    logical_page_positions: Optional[npt.NDArray[np.int32]]
    is_last_chunk: bool
    prefill_aux_index: Optional[int]
    state_indices: Optional[List[int]]
```

- [ ] **Step 2: Extend `add_transfer_request` signature**

Change signature to:

```python
def add_transfer_request(
    self,
    bootstrap_room: int,
    kv_indices: npt.NDArray[np.int32],
    index_slice: Optional[slice],
    is_last_chunk: bool,
    aux_index: Optional[int] = None,
    state_indices: Optional[List[int]] = None,
    logical_page_positions: Optional[npt.NDArray[np.int32]] = None,
):
```

When constructing `TransferKVChunk`, pass `logical_page_positions=logical_page_positions`.

- [ ] **Step 3: Select decode dst pages by explicit positions when present**

Replace in Mooncake transfer worker:

```python
chunked_dst_kv_indice = req.dst_kv_indices[kv_chunk.index_slice]
```

with:

```python
if kv_chunk.logical_page_positions is not None:
    chunked_dst_kv_indice = req.dst_kv_indices[kv_chunk.logical_page_positions]
else:
    assert kv_chunk.index_slice is not None
    chunked_dst_kv_indice = req.dst_kv_indices[kv_chunk.index_slice]
```

- [ ] **Step 4: Keep old sender path intact**

In `CommonKVSender.send` implementation inside Mooncake conn, keep existing replicated behavior when shared KV is disabled:

```python
self.kv_mgr.add_transfer_request(
    self.bootstrap_room,
    kv_indices,
    index_slice,
    is_last_chunk,
    aux_index=self.aux_index,
    state_indices=state_indices,
)
```

Shared KV path will call this method with `logical_page_positions` from Task 10.

- [ ] **Step 5: Run compile check**

Run:

```bash
python3 -m py_compile python/sglang/srt/disaggregation/mooncake/conn.py
```

Expected: no output, exit code 0.

- [ ] **Step 6: Commit**

```bash
git add python/sglang/srt/disaggregation/mooncake/conn.py
git commit -m "feat(cp): support explicit Mooncake KV page positions"
```

---

### Task 10: Send Only Owned CP Shared KV Pages From Prefill

**Files:**
- Modify: `python/sglang/srt/disaggregation/prefill.py`
- Modify: `python/sglang/srt/disaggregation/common/conn.py`
- Modify: `python/sglang/srt/disaggregation/mooncake/conn.py`

- [ ] **Step 1: Enforce all CP ranks transfer in CommonKVManager**

In `CommonKVManager.__init__`, after `self.enable_all_cp_ranks_for_transfer` is set:

```python
if (
    server_args.enable_nsa_prefill_cp_shared_kv
    and disaggregation_mode == DisaggregationMode.PREFILL
    and self.attn_cp_size > 1
):
    assert self.enable_all_cp_ranks_for_transfer, (
        "CP shared KV requires SGLANG_DISAGGREGATION_ALL_CP_RANKS_TRANSFER=1."
    )
```

- [ ] **Step 2: Add prefill page filtering in `send_kv_chunk`**

In `prefill.py::send_kv_chunk`, after `page_indices = kv_to_page_indices(kv_indices, page_size)`, add shared branch:

```python
logical_page_positions = None
if self.server_args.enable_nsa_prefill_cp_shared_kv:
    from sglang.srt.disaggregation.utils import filter_kv_pages_for_cp_shared_kv

    layout = self.tp_worker.model_runner.cp_shared_kv_layout
    assert layout is not None
    chunk_page_start = start_idx // page_size
    page_indices, logical_page_positions = filter_kv_pages_for_cp_shared_kv(
        layout=layout,
        logical_pages=page_indices,
        chunk_page_start=chunk_page_start,
    )
```

- [ ] **Step 3: Pass explicit positions to sender**

Update the sender call. If `req.disagg_kv_sender.send` signature cannot be changed globally, add a Mooncake-specific keyword-compatible method. Preferred signature:

```python
req.disagg_kv_sender.send(
    page_indices,
    state_indices,
    logical_page_positions=logical_page_positions,
)
```

Update Mooncake sender `send` signature:

```python
def send(
    self,
    kv_indices: npt.NDArray[np.int32],
    state_indices: Optional[List[int]] = None,
    logical_page_positions: Optional[npt.NDArray[np.int32]] = None,
):
```

When `logical_page_positions is not None`, do not call old `filter_kv_indices_for_cp_rank`; `kv_indices` are already source physical page ids for this rank. Call:

```python
self.kv_mgr.add_transfer_request(
    self.bootstrap_room,
    kv_indices,
    None,
    is_last_chunk,
    aux_index=self.aux_index if is_last_chunk else None,
    state_indices=state_indices if is_last_chunk else None,
    logical_page_positions=logical_page_positions,
)
```

- [ ] **Step 4: Handle empty owner chunk correctly**

If `len(page_indices) == 0` in shared KV mode, do not treat the request as fully skipped forever. For non-last chunks, return. For last chunk, still enqueue an empty final chunk or update success only after aux/state semantics are satisfied. Implement:

```python
if len(page_indices) == 0:
    if not last_chunk:
        return
    if self.server_args.enable_nsa_prefill_cp_shared_kv:
        req.disagg_kv_sender.send(
            page_indices,
            state_indices,
            logical_page_positions=logical_page_positions,
        )
        return
    logger.info(
        "Skip sending kv chunk for request %s room=%s because page_indices is empty",
        req.rid,
        req.bootstrap_room,
    )
    return
```

- [ ] **Step 5: Run compile check**

Run:

```bash
python3 -m py_compile python/sglang/srt/disaggregation/prefill.py python/sglang/srt/disaggregation/common/conn.py python/sglang/srt/disaggregation/mooncake/conn.py
```

Expected: no output, exit code 0.

- [ ] **Step 6: Commit**

```bash
git add python/sglang/srt/disaggregation/prefill.py python/sglang/srt/disaggregation/common/conn.py python/sglang/srt/disaggregation/mooncake/conn.py
git commit -m "feat(cp): transfer shared KV shards from all prefill CP ranks"
```

---

### Task 11: Wire NSA State/Index Cache Transfer Positions

**Files:**
- Modify: `python/sglang/srt/disaggregation/prefill.py`
- Modify: `python/sglang/srt/disaggregation/mooncake/conn.py`

- [ ] **Step 1: Extend chunk state fields**

Change `TransferKVChunk.state_indices` from only list of page ids to a structure that can carry positions:

```python
@dataclasses.dataclass
class TransferStatePages:
    src_indices: npt.NDArray[np.int32]
    logical_page_positions: Optional[npt.NDArray[np.int32]]
```

Use a backward-compatible type because existing Mamba/SWA/NSA paths pass lists:

```python
from typing import Union

state_indices: Optional[Union[List[int], TransferStatePages]]
```

- [ ] **Step 2: Build NSA state src pages and positions in prefill last chunk**

In `prefill.py`, inside the existing `elif isinstance(self.token_to_kv_pool_allocator.get_kvcache(), NSATokenToKVPool):` branch, use this shared-KV sub-branch:

```python
state_logical_pages = kv_to_page_indices(kv_indices_full.cpu().numpy(), page_size)
state_request_positions = np.arange(len(state_logical_pages), dtype=np.int32)
state_src_pages, state_positions = layout.filter_owned_pages_np(
    state_logical_pages,
    state_request_positions,
)
state_indices = TransferStatePages(
    src_indices=state_src_pages,
    logical_page_positions=state_positions,
)
```

- [ ] **Step 3: Use state positions in Mooncake extra send**

Where Mooncake currently calls `maybe_send_extra(req, kv_chunk.state_indices, target_rank_registration_info.dst_state_data_ptrs, executor, target_rank_registration_info)`, branch:

```python
state_indices = kv_chunk.state_indices
if isinstance(state_indices, TransferStatePages):
    dst_state_indices = target_rank_registration_info.dst_state_indices[
        state_indices.logical_page_positions
    ]
    self.maybe_send_extra(
        req,
        state_indices.src_indices,
        target_rank_registration_info.dst_state_data_ptrs,
        executor,
        target_rank_registration_info,
        dst_state_indices_override=dst_state_indices,
    )
else:
    self.maybe_send_extra(
        req,
        state_indices,
        target_rank_registration_info.dst_state_data_ptrs,
        executor,
        target_rank_registration_info,
    )
```

Change `maybe_send_extra` signature to accept the override explicitly:

```python
def maybe_send_extra(
    self,
    req: TransferInfo,
    prefill_state_indices: list[int],
    dst_state_data_ptrs: list[int],
    executor: concurrent.futures.ThreadPoolExecutor,
    target_rank_registration_info: Optional[KVArgsRegisterInfo] = None,
    dst_state_indices_override: Optional[npt.NDArray[np.int32]] = None,
):
```

Inside the `state_type in ["swa", "nsa"]` branch, replace `dst_state_indices = np.array(req.dst_state_indices, dtype=np.int32)` with:

```python
dst_state_indices = (
    np.asarray(dst_state_indices_override, dtype=np.int32)
    if dst_state_indices_override is not None
    else np.array(req.dst_state_indices, dtype=np.int32)
)
```

- [ ] **Step 4: Run compile check**

Run:

```bash
python3 -m py_compile python/sglang/srt/disaggregation/prefill.py python/sglang/srt/disaggregation/mooncake/conn.py
```

Expected: no output, exit code 0.

- [ ] **Step 5: Commit**

```bash
git add python/sglang/srt/disaggregation/prefill.py python/sglang/srt/disaggregation/mooncake/conn.py
git commit -m "feat(cp): transfer NSA index state shards with positions"
```

---

### Task 12: Add Attention Runtime Full-View Compatibility Hooks

**Files:**
- Create: `python/sglang/srt/layers/attention/nsa/cp_shared_kv_runtime.py`
- Modify: `python/sglang/srt/layers/attention/nsa_backend.py`
- Modify: `python/sglang/srt/layers/attention/nsa/nsa_indexer.py`
- Test: `test/registered/unit/mem_cache/test_cp_shared_kv_layout.py`

- [ ] **Step 1: Add CPU-testable page remapping helper tests**

Append to `test/registered/unit/mem_cache/test_cp_shared_kv_layout.py`:

```python
class TestCPSharedRuntimePageRemap(unittest.TestCase):
    def test_build_dense_page_remap(self):
        from sglang.srt.layers.attention.nsa.cp_shared_kv_runtime import (
            build_dense_page_remap_np,
        )

        logical_pages = np.array([9, 10, 13, 18], dtype=np.int32)
        remapped = build_dense_page_remap_np(logical_pages)
        self.assertEqual(remapped.tolist(), [0, 1, 2, 3])

    def test_build_dense_page_remap_rejects_duplicates(self):
        from sglang.srt.layers.attention.nsa.cp_shared_kv_runtime import (
            build_dense_page_remap_np,
        )

        with self.assertRaises(ValueError):
            build_dense_page_remap_np(np.array([9, 9], dtype=np.int32))
```

- [ ] **Step 2: Run test and verify failure**

Run:

```bash
python3 -m pytest test/registered/unit/mem_cache/test_cp_shared_kv_layout.py -q -k RuntimePageRemap
```

Expected: FAIL because `cp_shared_kv_runtime.py` does not exist.

- [ ] **Step 3: Create runtime compatibility module**

Create `python/sglang/srt/layers/attention/nsa/cp_shared_kv_runtime.py`:

```python
from __future__ import annotations

import logging

import numpy as np
import torch
import torch.distributed as dist

from sglang.srt.distributed.parallel_state import get_attention_cp_group

logger = logging.getLogger(__name__)


def build_dense_page_remap_np(logical_pages: np.ndarray) -> np.ndarray:
    pages = np.asarray(logical_pages, dtype=np.int64)
    if len(np.unique(pages)) != len(pages):
        raise ValueError(f"logical_pages must be unique, got {pages.tolist()}")
    return np.arange(len(pages), dtype=np.int32)


def materialize_shared_kv_pages_dense(
    *,
    forward_batch,
    layer_id: int,
    logical_pages: torch.Tensor,
) -> torch.Tensor:
    """Materialize dense MLA KV pages for Phase 2 compatibility.

    Every CP rank allocates a dense scratch tensor for the requested logical pages.
    The owner rank copies its physical pages into the matching dense positions;
    an all-reduce SUM makes the full dense view visible to every CP rank.
    This is intentionally a Phase 2 compatibility path and may consume maxlen-scale
    workspace.
    """
    layout = forward_batch.cp_shared_kv_layout
    assert layout is not None
    kv_pool = forward_batch.token_to_kv_pool
    key_buffer = kv_pool.get_key_buffer(layer_id)
    page_size = kv_pool.page_size
    kv_cache_dim = key_buffer.shape[-1]

    logical_pages = logical_pages.to(device=key_buffer.device, dtype=torch.int64)
    scratch = torch.zeros(
        (logical_pages.numel(), page_size, kv_cache_dim),
        dtype=key_buffer.dtype,
        device=key_buffer.device,
    )
    owned_mask = layout.owned_pages_mask(logical_pages)
    if owned_mask.any():
        physical_pages = layout.logical_pages_to_physical(logical_pages[owned_mask])
        src = key_buffer.view(-1, page_size, kv_cache_dim)[physical_pages]
        scratch[owned_mask] = src

    dist.all_reduce(scratch, op=dist.ReduceOp.SUM, group=get_attention_cp_group().device_group)
    logger.info(
        "CP shared KV full-view materialization: layer=%s, num_pages=%s, bytes=%s",
        layer_id,
        int(logical_pages.numel()),
        int(scratch.nbytes),
    )
    return scratch


def materialize_shared_index_pages_dense(
    *,
    forward_batch,
    layer_id: int,
    logical_pages: torch.Tensor,
) -> torch.Tensor:
    layout = forward_batch.cp_shared_kv_layout
    assert layout is not None
    kv_pool = forward_batch.token_to_kv_pool
    index_buffer = kv_pool.get_index_k_with_scale_buffer(layer_id)

    logical_pages = logical_pages.to(device=index_buffer.device, dtype=torch.int64)
    scratch = torch.zeros(
        (logical_pages.numel(), index_buffer.shape[-1]),
        dtype=index_buffer.dtype,
        device=index_buffer.device,
    )
    owned_mask = layout.owned_pages_mask(logical_pages)
    if owned_mask.any():
        physical_pages = layout.logical_pages_to_physical(logical_pages[owned_mask])
        scratch[owned_mask] = index_buffer[physical_pages]

    dist.all_reduce(scratch, op=dist.ReduceOp.SUM, group=get_attention_cp_group().device_group)
    logger.info(
        "CP shared KV index full-view materialization: layer=%s, num_pages=%s, bytes=%s",
        layer_id,
        int(logical_pages.numel()),
        int(scratch.nbytes),
    )
    return scratch


def extract_index_k_and_scale_from_dense_pages(
    dense_pages: torch.Tensor,
    seq_len: int,
    page_size: int,
    index_head_dim: int,
) -> tuple[torch.Tensor, torch.Tensor]:
    scale_bytes_per_token = index_head_dim // 128 * 4
    k_bytes = page_size * index_head_dim
    k_fp8 = dense_pages[:, :k_bytes].reshape(-1, index_head_dim)[:seq_len]
    k_scale = dense_pages[:, k_bytes:].reshape(-1, scale_bytes_per_token)[:seq_len]
    return k_fp8.contiguous(), k_scale.contiguous()
```

- [ ] **Step 4: Wire MLA dense materialization at physical-read sites**

In `nsa_backend.py`, import:

```python
from sglang.srt.layers.attention.nsa.cp_shared_kv_runtime import (
    materialize_shared_kv_pages_dense,
)
```

Before existing physical KV reads that use logical page tables, branch on shared KV. For the `flashmla_sparse` ragged path before `dequantize_k_cache_paged`, replace the original `kv_cache` source with dense scratch pages:

```python
if getattr(forward_batch, "uses_cp_shared_kv", False):
    logical_pages = page_table_1_flattened.to(dtype=torch.int64)
    dense_pages = materialize_shared_kv_pages_dense(
        forward_batch=forward_batch,
        layer_id=layer.layer_id,
        logical_pages=logical_pages,
    )
    kv_cache = dense_pages.view(-1, 1, dense_pages.shape[-1])
    page_table_1_flattened = torch.arange(
        dense_pages.shape[0], device=dense_pages.device, dtype=page_table_1_flattened.dtype
    )
kv_cache = dequantize_k_cache_paged(kv_cache, page_table_1_flattened)
```

- [ ] **Step 5: Wire index dense materialization at index prefix-read sites**

In `nsa_indexer.py`, import:

```python
from sglang.srt.layers.attention.nsa.cp_shared_kv_runtime import (
    extract_index_k_and_scale_from_dense_pages,
    materialize_shared_index_pages_dense,
)
```

Before calls to `get_index_k_continuous(layer_id, seq_len, block_tables[i])` and `get_index_k_continuous(layer_id, end_seq_position, block_tables[batch_idx])`, branch:

```python
if getattr(forward_batch, "uses_cp_shared_kv", False):
    logical_pages = block_tables[batch_idx].to(dtype=torch.int64)
    dense_index_pages = materialize_shared_index_pages_dense(
        forward_batch=forward_batch,
        layer_id=layer_id,
        logical_pages=logical_pages,
    )
    k_fp8, k_scale = extract_index_k_and_scale_from_dense_pages(
        dense_index_pages,
        seq_len=seq_len,
        page_size=forward_batch.token_to_kv_pool.page_size,
        index_head_dim=self.index_head_dim,
    )
else:
    k_fp8 = forward_batch.token_to_kv_pool.get_index_k_continuous(
        layer_id,
        seq_len,
        block_tables[i],
    )
    k_scale = forward_batch.token_to_kv_pool.get_index_k_scale_continuous(
        layer_id,
        seq_len,
        block_tables[i],
    )
```

For the `cp_index` loop that uses `batch_idx`, use the same branch but pass `seq_len=end_seq_position` and `logical_pages=block_tables[batch_idx]`.

- [ ] **Step 6: Run tests and compile checks**

Run:

```bash
python3 -m pytest test/registered/unit/mem_cache/test_cp_shared_kv_layout.py -q -k RuntimePageRemap
python3 -m py_compile python/sglang/srt/layers/attention/nsa/cp_shared_kv_runtime.py python/sglang/srt/layers/attention/nsa_backend.py python/sglang/srt/layers/attention/nsa/nsa_indexer.py
```

Expected: tests PASS and compile check exits 0.

- [ ] **Step 7: Commit**

```bash
git add python/sglang/srt/layers/attention/nsa/cp_shared_kv_runtime.py python/sglang/srt/layers/attention/nsa_backend.py python/sglang/srt/layers/attention/nsa/nsa_indexer.py test/registered/unit/mem_cache/test_cp_shared_kv_layout.py
git commit -m "feat(cp): add shared KV runtime full-view compatibility"
```


### Task 13: Make Scheduler And Metrics Explicitly Report Logical Capacity

**Files:**
- Modify: `python/sglang/srt/managers/scheduler.py`
- Modify: `python/sglang/srt/observability/scheduler_metrics_mixin.py`
- Modify: `python/sglang/srt/observability/metrics_collector.py` if a new metric is desired

- [ ] **Step 1: Keep existing `max_total_num_tokens` logical**

No code should reinterpret scheduler `max_total_num_tokens` as physical after Task 3. Audit these locations and keep them logical:

```text
scheduler.py: max_total_num_tokens log
schedule_policy.py: PrefillAdder.rem_total_tokens via allocator.available_size()
scheduler_runtime_checker_mixin.py: token usage denominator
scheduler_metrics_mixin.py: token_capacity
```

- [ ] **Step 2: Add physical capacity to scheduler log when available**

In scheduler init debug info, add:

```python
physical_tokens = getattr(
    self.tp_worker.model_runner,
    "physical_max_total_num_tokens",
    self.max_total_num_tokens,
)
if self.server_args.enable_nsa_prefill_cp_shared_kv:
    logger.info(
        "CP shared KV scheduler capacity: logical_tokens=%s, physical_tokens_per_rank=%s, cp_size=%s",
        self.max_total_num_tokens,
        physical_tokens,
        self.attn_cp_size,
    )
```

- [ ] **Step 3: Run compile check**

Run:

```bash
python3 -m py_compile python/sglang/srt/managers/scheduler.py python/sglang/srt/observability/scheduler_metrics_mixin.py
```

Expected: no output, exit code 0.

- [ ] **Step 4: Commit**

```bash
git add python/sglang/srt/managers/scheduler.py python/sglang/srt/observability/scheduler_metrics_mixin.py
git commit -m "feat(cp): report shared KV logical scheduler capacity"
```

---

### Task 14: End-To-End Smoke Verification On Target Launch

**Files:**
- No source files required unless verification exposes defects.

- [ ] **Step 1: Start prefill with shared KV flag**

Use the target command plus:

```bash
export SGLANG_DISAGGREGATION_ALL_CP_RANKS_TRANSFER=1
```

and add:

```bash
--enable-nsa-prefill-cp-shared-kv
```

Expected startup logs:

```text
CP shared KV enabled. physical_tokens_per_rank=237312, logical_tokens=1898496, cp_size=8, shard_policy=page_interleaved
CP shared KV scheduler capacity: logical_tokens=1898496, physical_tokens_per_rank=237312, cp_size=8
```

- [ ] **Step 2: Verify physical KV allocation is not multiplied**

Expected:

```text
KV Cache is allocated. #tokens: approximately previous physical count, KV size: approximately previous 22.14 GB
```

The physical KV size must not become `22.14 * 8 GB` on each rank.

- [ ] **Step 3: Verify logical capacity is multiplied**

Expected:

```text
logical_tokens ≈ physical_tokens_per_rank * 8
```

For the observed example:

```text
physical_tokens_per_rank=237312
logical_tokens=1898496
```

- [ ] **Step 4: Verify all prefill CP ranks participate in Mooncake transfer**

Expected logs or counters:

```text
CP shared KV transfer: rank=0 sent pages > 0
CP shared KV transfer: rank=1 sent pages > 0
CP shared KV transfer: rank=2 sent pages > 0
CP shared KV transfer: rank=3 sent pages > 0
CP shared KV transfer: rank=4 sent pages > 0
CP shared KV transfer: rank=5 sent pages > 0
CP shared KV transfer: rank=6 sent pages > 0
CP shared KV transfer: rank=7 sent pages > 0
```

- [ ] **Step 5: Run one small correctness request**

Send a prompt shorter than one physical rank capacity. Expected:

```text
prefill succeeds
decode receives complete KV
first token output is produced
no missing NSA index K cache error
```

- [ ] **Step 6: Run admission test above old single-rank capacity**

Send or simulate a prefill whose logical prompt/cache pressure exceeds old `237312` but is below `237312 * 8`, while decode capacity is sufficient for the chosen request.

Expected:

```text
prefill scheduler admits based on logical capacity
per-rank physical KV allocation remains bounded
runtime may hit compatibility workspace limit; if it does, log identifies full-view materialization as the source
```

- [ ] **Step 7: Commit verification fixes**

If verification requires code fixes:

```bash
git add <fixed files>
git commit -m "fix(cp): stabilize shared KV smoke path"
```

---

## Self-Review

- Spec coverage:
  - Config/guard: Task 2.
  - logical/physical loc split: Tasks 1, 3, 4, 5.
  - KV pool physical allocation: Tasks 3, 4.
  - scheduler logical capacity: Tasks 3, 13.
  - KV writes owner-only: Tasks 6, 7.
  - Mooncake PD transfer explicit positions: Tasks 8, 9, 10, 11.
  - attention full-view compatibility and Phase3 boundary: Task 12.
  - validation: Task 14.
- Placeholder scan: no placeholder markers, no unconstrained future-work steps, no unbounded edge-condition instructions.
- Type consistency:
  - `CpSharedKVLayout` is defined in Task 1 and used consistently in later tasks.
  - `logical_page_positions` is request-absolute everywhere.
  - `out_cache_loc` is logical when shared KV is enabled; physical locs are only passed to KV pool writes after layout conversion.