Files

T

laoyao0822 76c830dda8 docs(cp): plan NSA prefill CP shared KV phase2

2026-04-24 16:31:03 +08:00

56 KiB

Raw Blame History

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:

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:

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:

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:

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

Expected: PASS.

Step 5: Commit

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:

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:

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:

# 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:

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:

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:

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

Expected: PASS.

Step 6: Commit

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:

@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:

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:

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:

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:

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:

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:

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

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:

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:

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:

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:

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

In _init_pools, compute pool size:

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:

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:

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

Expected: PASS.

Step 6: Commit

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:

from sglang.srt.mem_cache.cp_shared_kv_layout import CpSharedKVLayout

Add fields near nsa_cp_metadata:

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

Step 2: Initialize layout in ModelRunner

In ModelRunner.__init__, add:

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:

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(...):

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:

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

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:

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:

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

with:

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:

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

Expected: no output, exit code 0.

Step 4: Commit

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:

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:

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:

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:

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:

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

Expected: no output, exit code 0.

Step 5: Commit

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:

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:

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:

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:

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

Expected: PASS.

Step 5: Commit

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:

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:

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:

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:

chunked_dst_kv_indice = req.dst_kv_indices[kv_chunk.index_slice]

with:

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:

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:

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

Expected: no output, exit code 0.

Step 6: Commit

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:

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:

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:

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

Update Mooncake sender send signature:

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:

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:

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:

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

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:

@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:

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:

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:

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:

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:

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:

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

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:

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:

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:

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:

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:

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:

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:

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:

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

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:

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:

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:

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

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:

export SGLANG_DISAGGREGATION_ALL_CP_RANKS_TRANSFER=1

and add:

--enable-nsa-prefill-cp-shared-kv

Expected startup logs:

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:

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:

logical_tokens ≈ physical_tokens_per_rank * 8

For the observed example:

physical_tokens_per_rank=237312
logical_tokens=1898496

Step 4: Verify all prefill CP ranks participate in Mooncake transfer

Expected logs or counters:

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:

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:

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:

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.

56 KiB Raw Blame History

NSA Prefill CP Shared KV Phase 2 Implementation Plan

0. Review Mapping From Design To Current Code

File Structure

Task 1: Add Pure CP Shared KV Layout Helper

Task 2: Add Server Flag And Validation

Task 3: Split Physical And Logical KV Capacity In ModelRunner

Task 4: Add CP Shared Paged Allocator Returning Logical Locs

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

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

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

Task 8: Add Mooncake Shared KV Page Mapping Utilities

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

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

Task 11: Wire NSA State/Index Cache Transfer Positions

Task 12: Add Attention Runtime Full-View Compatibility Hooks

Task 13: Make Scheduler And Metrics Explicitly Report Logical Capacity

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

Self-Review

56 KiB

Raw Blame History