# NSA Prefill CP Index Skip Cache/Communication Optimization 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:** Extend NSA `index_topk_freq` / IndexCache from compute-only savings to HiCache index load/backup savings, PD state-transfer savings, and eventually device/host index-cache capacity savings without changing model semantics.

**Architecture:** Treat index-cache layers as an explicit logical contract derived from the same skip formula used by `DeepseekV2AttentionMLA`. First centralize the active-layer map and use it to skip unnecessary per-layer indexer load/backup while keeping memory layout unchanged. Only after that is verified, add transfer metadata and compact device/host state buffers with fail-fast mapping checks between prefill and decode.

**Tech Stack:** Python runtime (`sglang.srt`), NSA/MLA attention, CP shared KV, HiCache host/device pools, Mooncake/NIXL PD transfer metadata, pytest, remote CUDA validation on `g0034`.

---

## 0. Current evidence from code reading

### 0.1 Index skip is currently a model-forward compute optimization only

Files:
- `python/sglang/srt/models/deepseek_v2.py:1215-1254`
- `python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla.py:206-235`
- `python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla.py:571-578`

Current semantics:

```python
# deepseek_v2.py, target layers
if index_topk_pattern is None and index_skip_topk_offset is not None:
    skip_topk = max(layer_id - index_skip_topk_offset + 1, 0) % index_topk_freq != 0
    next_skip_topk = max(layer_id - index_skip_topk_offset + 2, 0) % index_topk_freq != 0
elif index_topk_pattern is None:
    skip_topk = max(layer_id - 1, 0) % index_topk_freq != 0
    next_skip_topk = layer_id % index_topk_freq != 0
else:
    skip_topk = index_topk_pattern[layer_id] == "S"
    next_skip_topk = layer_id + 1 < len(index_topk_pattern) and index_topk_pattern[layer_id + 1] == "S"
```

`forward_mla.py` only calls `self.indexer(...)` when the layer is not skipped, except the intentional `is_nextn and prev_topk_indices is None` draft fallback:

```python
if not self.skip_topk or (self.is_nextn and prev_topk_indices is None):
    topk_indices = self.indexer(...)
else:
    topk_indices = prev_topk_indices
```

Implication: skipped target layers do not need their own index-K cache content. They reuse top-k indices from an active producer layer.

### 0.2 Device NSA index buffers are still allocated and registered for every layer

File: `python/sglang/srt/mem_cache/memory_pool.py:1853-2019`

Current behavior:

```python
self.index_k_with_scale_buffer = [
    torch.zeros((num_pages, page_size * (index_head_dim + index_head_dim // 128 * 4)), ...)
    for _ in range(layer_num)
]

def _get_index_k_with_scale_buffer(self, layer_id: int) -> torch.Tensor:
    return self.index_k_with_scale_buffer[layer_id - self.start_layer]

def get_state_buf_infos(self):
    data_ptrs = [self.index_k_with_scale_buffer[i].data_ptr() for i in range(self.layer_num)]
    data_lens = [self.index_k_with_scale_buffer[i].nbytes for i in range(self.layer_num)]
    item_lens = [self.index_k_with_scale_buffer[i][0].nbytes for i in range(self.layer_num)]
    return data_ptrs, data_lens, item_lens
```

Implication: PD state registration and transfer still see all index layers even when `index_topk_freq=4` means most layers are semantically skipped.

### 0.3 Host HiCache NSA index buffers are still full-layer

File: `python/sglang/srt/mem_cache/memory_pool_host.py:1833-1931`

Current behavior:

```python
self.indexer_size_per_token = index_head_dim + index_head_dim // quant_block_size * 4
self.indexer_page_stride_size = self.indexer_size_per_token * page_size * dtype.itemsize
self.indexer_layout_dim = self.indexer_page_stride_size * self.layer_num

# page_first/page_first_direct
self.index_k_with_scale_buffer = alloc_func(
    (self.indexer_page_num, self.layer_num, 1, self.indexer_page_stride_size),
    ...,
)

# layer_page_first
self.index_k_with_scale_buffer = alloc_func(
    (self.layer_num, self.indexer_page_num, 1, self.indexer_page_stride_size),
    ...,
)
```

For `index_head_dim=128`, `quant_block_size=128`, `page_size=64`, the indexer stride is `8448` bytes per page per layer. A 78-layer model stores/transfers about `78 * 8448 = 658,944` bytes per page of NSA indexer state before considering skipped layers.

### 0.4 HiCache index load/backup always follows every model layer

File: `python/sglang/srt/mem_cache/memory_pool_host.py:1959-2186`

Current behavior:

```python
def load_to_device_per_layer(..., layer_id, io_backend):
    super().load_to_device_per_layer(..., layer_id, io_backend)
    self._load_indexer_to_device_per_layer(..., layer_id, io_backend)

def backup_from_device_all_layer(...):
    super().backup_from_device_all_layer(...)
    self._backup_indexer_from_device_all_layer(...)

def backup_from_device_per_layer(..., layer_id, io_backend):
    super().backup_from_device_per_layer(..., layer_id, io_backend)
    self._backup_indexer_from_device_per_layer(..., layer_id, io_backend)
```

The MLA KV transfer must still happen for every layer. Only NSA indexer transfer can be skipped for skipped top-k layers.

### 0.5 PD state transfer has no logical layer IDs today

Files:
- `python/sglang/srt/disaggregation/prefill.py:372-420`
- `python/sglang/srt/disaggregation/decode.py:425-473`
- `python/sglang/srt/disaggregation/utils.py:222-263`
- `python/sglang/srt/disaggregation/mooncake/conn.py:164-202`
- `python/sglang/srt/disaggregation/mooncake/conn.py:890-1014`
- `python/sglang/srt/disaggregation/mooncake/conn.py:1820-1884`

Current behavior:

```python
state_data_ptrs, state_data_lens, state_item_lens = pool.get_state_buf_infos()
kv_args.state_data_ptrs = state_data_ptrs
kv_args.state_data_lens = state_data_lens
kv_args.state_item_lens = state_item_lens
```

Mooncake registers only state pointers and lengths. It does not register logical index layer IDs. Therefore compacting `get_state_buf_infos()` before adding layer-ID metadata can silently map prefill layer slot `i` to a different decode layer slot `i`.

### 0.6 Draft/EAGLE must remain conservative initially

File: `python/sglang/srt/disaggregation/utils.py:222-263`

`append_cp_draft_state_buffers()` appends draft NSA state buffers to target state buffers. Draft (`is_nextn`) sets `skip_topk=True`, but `forward_mla.py` intentionally allows indexer execution when `is_nextn and prev_topk_indices is None`. The first implementation must not compact or skip draft index state unless prefill and decode attach a separate, validated draft state-layer namespace.

---

## 1. Target behavior

With `--nsa-index-topk-freq 4` and no offset, target active index layers are:

```text
0, 1, 5, 9, 13, ...
```

With `--nsa-index-topk-freq 4 --nsa-index-skip-topk-offset 1`, target active index layers are:

```text
0, 4, 8, 12, ...
```

For an `index_topk_pattern`, active target index layers are exactly layers whose pattern character is not `S`.

Required properties:

1. The formula used by model forward, HiCache, memory estimation, and PD transfer must be one shared helper.
2. Skipped target layers must not run indexer, must not backup indexer state, and must not load indexer state.
3. MLA KV cache is still per-layer and must not be skipped.
4. State-buffer compaction requires logical-layer metadata. Mismatch must fail fast; truncation is not allowed.
5. Draft/EAGLE state remains full and explicit until a separate draft mapping is implemented and tested.

---

## 2. File structure and responsibilities

Create:
- `python/sglang/srt/configs/nsa_index_layers.py`
  - Owns the single source of truth for index skip flags, active index layers, and logical-layer-to-slot mapping.

Create:
- `test/registered/unit/configs/test_nsa_index_layers.py`
  - Tests the helper against the existing model skip formula.

Modify:
- `python/sglang/srt/models/deepseek_v2.py`
  - Replaces local skip formula with helper calls.

Modify:
- `python/sglang/srt/mem_cache/memory_pool.py`
  - Adds active-layer metadata to `NSATokenToKVPool`.
  - Later phases map logical layer IDs to compact state slots.

Modify:
- `python/sglang/srt/mem_cache/memory_pool_host.py`
  - Skips indexer H2D/D2H for inactive target layers.
  - Later phases compact host index buffers.

Modify:
- `python/sglang/srt/model_executor/model_runner_kv_cache_mixin.py`
  - Later phases estimate NSA index cache using active layer count.

Modify:
- `python/sglang/srt/disaggregation/base/conn.py`
- `python/sglang/srt/disaggregation/prefill.py`
- `python/sglang/srt/disaggregation/decode.py`
- `python/sglang/srt/disaggregation/utils.py`
- `python/sglang/srt/disaggregation/mooncake/conn.py`
- `python/sglang/srt/disaggregation/nixl/conn.py`
  - Adds optional state layer IDs to PD registration and fail-fast validation before compact state transfer.

Modify tests:
- `test/registered/unit/mem_cache/test_nsa_pool_host_unit.py`
- `test/registered/unit/mem_cache/test_cp_shared_kv_runtime.py`
- `test/registered/unit/disaggregation/test_mooncake_conn.py` if present; otherwise create `test/registered/unit/disaggregation/test_pd_state_layer_ids.py`.

---

## 3. Phase P0: central active-index-layer helper

### Goal

Create a shared helper that exactly reproduces the current `DeepseekV2AttentionMLA` skip formula and exposes active index layer IDs.

### Implementation steps

- [ ] Create `python/sglang/srt/configs/nsa_index_layers.py` with this interface:

```python
from dataclasses import dataclass
from typing import Dict, Iterable, List, Optional, Sequence


@dataclass(frozen=True)
class NSAIndexLayerPlan:
    start_layer: int
    end_layer: int
    active_layer_ids: tuple[int, ...]
    layer_to_slot: Dict[int, int]

    def is_active(self, layer_id: int) -> bool:
        return layer_id in self.layer_to_slot

    def slot_for_layer(self, layer_id: int) -> int:
        try:
            return self.layer_to_slot[layer_id]
        except KeyError as exc:
            raise RuntimeError(
                "[CP_SHARED_KV_FAIL_FAST][index_cache_layer] "
                f"inactive index layer requested: layer_id={layer_id} "
                f"active_layer_ids={list(self.active_layer_ids)}"
            ) from exc


def nsa_index_skip_flags(config, layer_id: int, *, is_nextn: bool = False) -> tuple[bool, bool]:
    if is_nextn:
        return True, True

    index_topk_freq = getattr(config, "index_topk_freq", 1)
    index_topk_pattern = getattr(config, "index_topk_pattern", None)
    index_skip_topk_offset = getattr(config, "index_skip_topk_offset", None)

    if index_topk_freq < 1:
        raise ValueError(f"index_topk_freq must be >= 1, got {index_topk_freq}")

    if index_topk_pattern is None and index_skip_topk_offset is not None:
        if index_skip_topk_offset <= 0:
            raise ValueError(
                "index_skip_topk_offset must be positive when configured; "
                f"got {index_skip_topk_offset}"
            )
        skip_topk = max(layer_id - index_skip_topk_offset + 1, 0) % index_topk_freq != 0
        next_skip_topk = max(layer_id - index_skip_topk_offset + 2, 0) % index_topk_freq != 0
        return skip_topk, next_skip_topk

    if index_topk_pattern is None:
        skip_topk = max(layer_id - 1, 0) % index_topk_freq != 0
        next_skip_topk = layer_id % index_topk_freq != 0
        return skip_topk, next_skip_topk

    if layer_id < 0 or layer_id >= len(index_topk_pattern):
        raise ValueError(
            f"layer_id={layer_id} outside index_topk_pattern length={len(index_topk_pattern)}"
        )
    skip_topk = index_topk_pattern[layer_id] == "S"
    next_skip_topk = layer_id < len(index_topk_pattern) - 1 and index_topk_pattern[layer_id + 1] == "S"
    return skip_topk, next_skip_topk


def build_nsa_index_layer_plan(config, start_layer: int, end_layer: int, *, is_nextn: bool = False) -> NSAIndexLayerPlan:
    if end_layer < start_layer:
        raise ValueError(f"end_layer={end_layer} must be >= start_layer={start_layer}")
    if is_nextn:
        active = tuple(range(start_layer, end_layer))
    else:
        active = tuple(
            layer_id
            for layer_id in range(start_layer, end_layer)
            if not nsa_index_skip_flags(config, layer_id, is_nextn=False)[0]
        )
    return NSAIndexLayerPlan(
        start_layer=start_layer,
        end_layer=end_layer,
        active_layer_ids=active,
        layer_to_slot={layer_id: slot for slot, layer_id in enumerate(active)},
    )
```

- [ ] Add unit tests in `test/registered/unit/configs/test_nsa_index_layers.py`:

```python
from types import SimpleNamespace

import pytest

from sglang.srt.configs.nsa_index_layers import (
    build_nsa_index_layer_plan,
    nsa_index_skip_flags,
)


def test_default_freq_one_marks_all_target_layers_active():
    cfg = SimpleNamespace(index_topk_freq=1)
    plan = build_nsa_index_layer_plan(cfg, 0, 6)
    assert plan.active_layer_ids == (0, 1, 2, 3, 4, 5)
    assert [nsa_index_skip_flags(cfg, i)[0] for i in range(6)] == [False] * 6


def test_freq_four_without_offset_matches_current_model_formula():
    cfg = SimpleNamespace(index_topk_freq=4)
    plan = build_nsa_index_layer_plan(cfg, 0, 12)
    assert plan.active_layer_ids == (0, 1, 5, 9)
    assert [nsa_index_skip_flags(cfg, i)[0] for i in range(10)] == [
        False, False, True, True, True, False, True, True, True, False
    ]


def test_freq_four_with_offset_one_uses_layers_zero_four_eight():
    cfg = SimpleNamespace(index_topk_freq=4, index_skip_topk_offset=1)
    plan = build_nsa_index_layer_plan(cfg, 0, 12)
    assert plan.active_layer_ids == (0, 4, 8)
    assert [nsa_index_skip_flags(cfg, i)[0] for i in range(9)] == [
        False, True, True, True, False, True, True, True, False
    ]


def test_pattern_marks_non_shared_layers_active():
    cfg = SimpleNamespace(index_topk_freq=1, index_topk_pattern="CSSSCSS")
    plan = build_nsa_index_layer_plan(cfg, 0, len(cfg.index_topk_pattern))
    assert plan.active_layer_ids == (0, 4)


def test_nextn_keeps_all_draft_layers_active_for_state_safety():
    cfg = SimpleNamespace(index_topk_freq=4, index_skip_topk_offset=1)
    plan = build_nsa_index_layer_plan(cfg, 0, 1, is_nextn=True)
    assert plan.active_layer_ids == (0,)
    assert nsa_index_skip_flags(cfg, 0, is_nextn=True) == (True, True)


def test_inactive_slot_lookup_fails_fast():
    cfg = SimpleNamespace(index_topk_freq=4, index_skip_topk_offset=1)
    plan = build_nsa_index_layer_plan(cfg, 0, 8)
    with pytest.raises(RuntimeError, match="inactive index layer requested"):
        plan.slot_for_layer(1)
```

- [ ] Run local non-CUDA test:

```bash
cd /root/sglang-work/sglang-dev
PYTHONPATH=python python -m pytest -q test/registered/unit/configs/test_nsa_index_layers.py
```

Expected output:

```text
6 passed
```

- [ ] Replace the local formula in `python/sglang/srt/models/deepseek_v2.py` with the helper:

```python
from sglang.srt.configs.nsa_index_layers import nsa_index_skip_flags

...
self.skip_topk, self.next_skip_topk = nsa_index_skip_flags(
    config, layer_id, is_nextn=is_nextn
)
```

- [ ] Run targeted syntax check:

```bash
cd /root/sglang-work/sglang-dev
python -m py_compile \
  python/sglang/srt/configs/nsa_index_layers.py \
  python/sglang/srt/models/deepseek_v2.py
```

Expected output: no output and exit code `0`.

---

## 4. Phase P1: wire active-layer metadata into `NSATokenToKVPool` without compacting allocation

### Goal

Expose the active target index layer plan from the device pool while keeping `index_k_with_scale_buffer` full-layer. This makes the first runtime change low-risk and gives later phases a stable API.

### Implementation steps

- [ ] Extend `NSATokenToKVPool.__init__` in `python/sglang/srt/mem_cache/memory_pool.py` with optional constructor arguments:

```python
index_active_layer_ids: Optional[Sequence[int]] = None,
compact_index_layers: bool = False,
```

- [ ] Store identity slot mapping for this phase:

```python
if index_active_layer_ids is None:
    index_active_layer_ids = tuple(range(self.start_layer, self.end_layer))
else:
    index_active_layer_ids = tuple(int(x) for x in index_active_layer_ids)

self.index_active_layer_ids = index_active_layer_ids
self.index_active_layer_id_set = frozenset(index_active_layer_ids)
self.index_compact_layers = bool(compact_index_layers)
self.index_layer_to_slot = {
    layer_id: layer_id - self.start_layer for layer_id in index_active_layer_ids
}
```

- [ ] Add helpers in `NSATokenToKVPool`:

```python
def is_index_layer_active(self, layer_id: int) -> bool:
    return layer_id in self.index_active_layer_id_set


def get_index_layer_slot(self, layer_id: int) -> int:
    try:
        return self.index_layer_to_slot[layer_id]
    except KeyError as exc:
        raise RuntimeError(
            "[CP_SHARED_KV_FAIL_FAST][index_cache_layer] "
            f"inactive index layer requested: layer_id={layer_id} "
            f"active_layer_ids={list(self.index_active_layer_ids)}"
        ) from exc
```

- [ ] Keep raw allocation full-layer in this phase:

```python
index_buffer_layer_num = layer_num
```

- [ ] Update index buffer accessors to use `get_index_layer_slot(layer_id)` only for public operations that require active index state:

```python
def _get_index_k_with_scale_buffer(self, layer_id: int) -> torch.Tensor:
    return self.index_k_with_scale_buffer[self.get_index_layer_slot(layer_id)]
```

- [ ] Run unit tests that touch memory pool helpers. If constructing the real CUDA pool is not possible locally, run only py_compile locally and the real tests remotely:

```bash
cd /root/sglang-work/sglang-dev
python -m py_compile python/sglang/srt/mem_cache/memory_pool.py
```

Remote command:

```bash
ssh g0034 "docker exec sglang-glm5-dev-2 bash -lc 'cd /sgl-workspace/sglang-tai && PYTHONPATH=python python -m pytest -q test/registered/unit/mem_cache/test_nsa_pool_host_unit.py'"
```

Expected output for remote: test process exits `0`.

---

## 5. Phase P2: skip inactive target indexer HiCache load/backup while preserving MLA KV transfer

### Goal

Reduce H2D/D2H indexer transfer immediately for `index_topk_freq > 1` without changing host/device allocation shape or PD state registration.

### Implementation steps

- [ ] In `python/sglang/srt/mem_cache/memory_pool_host.py`, add an inactive check at the top of `_load_indexer_to_device_per_layer`:

```python
if not device_pool.is_index_layer_active(layer_id):
    return
```

- [ ] Add the same check at the top of `_backup_indexer_from_device_per_layer`:

```python
if not device_pool.is_index_layer_active(layer_id):
    return
```

- [ ] Change `_backup_indexer_from_device_all_layer` direct loops from `range(self.layer_num)` to the active layer IDs:

```python
active_layer_ids = tuple(getattr(device_pool, "index_active_layer_ids", range(self.layer_num)))
for logical_layer_id in active_layer_ids:
    slot = device_pool.get_index_layer_slot(logical_layer_id)
    tai_transfer(
        src_ptrs=[device_pool.index_k_with_scale_buffer[slot]],
        dst_ptrs=[self.index_k_with_scale_buffer],
        src_indices=device_page_indices,
        dst_indices=host_page_indices,
        layer_id=logical_layer_id,
        page_size=1,
    )
```

For `layer_first`, `dst_layers` remains a list and must use matching slots on both sides:

```python
src_layers=[device_pool.index_k_with_scale_buffer[device_pool.get_index_layer_slot(i)] for i in active_layer_ids]
dst_layers=[self.index_k_with_scale_buffer[i - self.start_layer] for i in active_layer_ids]
```

- [ ] Add a CPU-side unit test with monkeypatched transfer functions so the test does not require CUDA kernels. The test should assert that inactive layers are not passed to indexer transfer while `super().backup_from_device_per_layer` is still called for MLA KV. If an existing host-pool CUDA fixture is easier, add this assertion to `test/registered/unit/mem_cache/test_nsa_pool_host_unit.py`.

Test shape:

```python
def test_index_skip_does_not_transfer_inactive_index_layers(monkeypatch):
    calls = []

    class FakeDevicePool:
        index_active_layer_ids = (0, 4)
        index_k_with_scale_buffer = [object() for _ in range(8)]
        def is_index_layer_active(self, layer_id):
            return layer_id in self.index_active_layer_ids
        def get_index_layer_slot(self, layer_id):
            return layer_id

    def fake_tai_transfer():
        def inner(**kwargs):
            calls.append(kwargs["layer_id"])
        return inner

    monkeypatch.setattr(memory_pool_host, "_load_tai_transfer_kv_per_layer_direct_lf_pf", fake_tai_transfer)
    host_pool = make_fake_nsa_host_pool(layout="page_first_direct", layer_num=8)
    host_pool._backup_indexer_from_device_per_layer(FakeDevicePool(), host_indices, device_indices, 1, "direct")
    host_pool._backup_indexer_from_device_per_layer(FakeDevicePool(), host_indices, device_indices, 4, "direct")
    assert calls == [4]
```

- [ ] Run targeted tests:

```bash
cd /root/sglang-work/sglang-dev
PYTHONPATH=python python -m pytest -q test/registered/unit/mem_cache/test_nsa_pool_host_unit.py -k 'index_skip or indexer'
```

Expected output: selected tests pass. If CUDA allocation prevents local execution, run the same command in the remote container.

- [ ] ETE validation after P2:

Start prefill with:

```bash
--nsa-index-topk-freq 4
```

Verify logs show no `CP_SHARED_KV_FAIL_FAST` and no index-related fallback warnings. Run two GSM8K samples back-to-back; both first-run and cache-hit-run accuracy must stay near the known baseline (`0.953-0.955` for the current GLM5 setup).

---

## 6. Phase P3: add state layer IDs to PD transfer before compacting state buffers

### Goal

Make PD state transfer layer-aware so compact state registration can be verified instead of relying on positional list order.

### Implementation steps

- [ ] Extend `KVArgs` state metadata where `state_data_ptrs`, `state_data_lens`, and `state_item_lens` are assigned. Add:

```python
kv_args.state_layer_ids = pool.get_state_layer_ids() if hasattr(pool, "get_state_layer_ids") else []
```

Apply in both:
- `python/sglang/srt/disaggregation/prefill.py`
- `python/sglang/srt/disaggregation/decode.py`

- [ ] Add `get_state_layer_ids()` to `NSATokenToKVPool`:

```python
def get_state_layer_ids(self):
    return list(range(self.start_layer, self.end_layer))
```

This returns full logical IDs in P3. It will switch to `index_active_layer_ids` in P4.

- [ ] Update `append_cp_draft_state_buffers()` in `python/sglang/srt/disaggregation/utils.py` to append draft layer IDs with an explicit namespace offset. Use negative IDs for draft state to avoid collision with target target-layer IDs:

```python
draft_state_layer_ids = getattr(kv_args, "draft_state_layer_ids", None)
if draft_state_layer_ids is None:
    draft_state_layer_ids = [-(i + 1) for i in range(draft_state_bufs)]
kv_args.state_layer_ids += draft_state_layer_ids
```

- [ ] Extend Mooncake registration in `python/sglang/srt/disaggregation/mooncake/conn.py`:

Packing side:

```python
packed_state_layer_ids = b"".join(
    struct.pack("i", int(layer_id)) for layer_id in getattr(self.kv_mgr.kv_args, "state_layer_ids", [])
)
```

Add the frame after `packed_state_dim_per_tensor`.

Unpacking side:

```python
dst_state_layer_ids=(
    list(struct.unpack(f"{len(msg[12])//4}i", msg[12]))
    if len(msg) > 12 and len(msg[12]) > 0
    else []
),
```

- [ ] Add `dst_state_layer_ids: list[int]` to `KVArgsRegisterInfo`.

- [ ] Validate before state transfer in `maybe_send_extra()`:

```python
src_state_layer_ids = list(getattr(self.kv_args, "state_layer_ids", []) or [])
dst_state_layer_ids = (
    list(target_rank_registration_info.dst_state_layer_ids)
    if target_rank_registration_info is not None
    else []
)
if src_state_layer_ids or dst_state_layer_ids:
    if src_state_layer_ids != dst_state_layer_ids:
        raise RuntimeError(
            "[CP_SHARED_KV_FAIL_FAST][state_layer_ids] "
            f"prefill={src_state_layer_ids} decode={dst_state_layer_ids} "
            f"state_type={state_type} room={req.room} session={req.mooncake_session_id}"
        )
```

- [ ] Replace current warning-and-truncate behavior for state-buffer count mismatch with fail-fast when either side provides `state_layer_ids`:

```python
if len(src_state_data_ptrs) != len(dst_state_ptrs):
    if src_state_layer_ids or dst_state_layer_ids:
        raise RuntimeError(
            "[CP_SHARED_KV_FAIL_FAST][state_buffer_count] "
            f"src={len(src_state_data_ptrs)} dst={len(dst_state_ptrs)} "
            f"src_layers={src_state_layer_ids} dst_layers={dst_state_layer_ids}"
        )
    # legacy warning path remains only for non-layer-aware state types
```

- [ ] Mirror the metadata addition in NIXL if the same state transfer path is used in the deployment. If NIXL is not used for this target, keep the NIXL change limited to optional packing/unpacking and no behavior change.

- [ ] Add unit tests in `test/registered/unit/disaggregation/test_pd_state_layer_ids.py`:

```python
def test_mooncake_register_info_roundtrips_state_layer_ids():
    msg = make_register_msg_with_state_layer_ids([0, 4, 8, -1])
    info = KVArgsRegisterInfo.from_zmq(msg)
    assert info.dst_state_layer_ids == [0, 4, 8, -1]


def test_mooncake_state_layer_id_mismatch_fails_fast():
    manager = make_fake_mooncake_manager(state_layer_ids=[0, 4, 8])
    info = make_fake_register_info(dst_state_layer_ids=[0, 5, 8])
    with pytest.raises(RuntimeError, match="state_layer_ids"):
        manager.maybe_send_extra(req, prefill_state_indices, dst_state_ptrs, executor, info)
```

- [ ] Run targeted tests:

```bash
cd /root/sglang-work/sglang-dev
PYTHONPATH=python python -m pytest -q test/registered/unit/disaggregation/test_pd_state_layer_ids.py
```

Expected output: all tests pass.

---

## 7. Phase P4: compact PD state registration to active target index layers

### Goal

Reduce PD NSA state transfer buffer count and transfer bytes by registering only active target index layers.

### Implementation steps

- [ ] Change `NSATokenToKVPool.get_state_buf_infos()` to return active slots:

```python
def get_state_buf_infos(self):
    slots = [self.get_index_layer_slot(layer_id) for layer_id in self.index_active_layer_ids]
    data_ptrs = [self.index_k_with_scale_buffer[slot].data_ptr() for slot in slots]
    data_lens = [self.index_k_with_scale_buffer[slot].nbytes for slot in slots]
    item_lens = [self.index_k_with_scale_buffer[slot][0].nbytes for slot in slots]
    return data_ptrs, data_lens, item_lens


def get_state_layer_ids(self):
    return list(self.index_active_layer_ids)
```

- [ ] Keep draft full-state append enabled by `append_cp_draft_state_buffers()`. The combined `state_layer_ids` must look like this for target active layers `[0, 4, 8]` and one draft layer:

```text
[0, 4, 8, -1]
```

- [ ] Add a unit test:

```python
def test_nsa_state_buf_infos_returns_active_layers_only():
    pool = make_fake_nsa_pool(layer_num=8, active_layer_ids=(0, 4))
    ptrs, lens, item_lens = pool.get_state_buf_infos()
    assert len(ptrs) == 2
    assert pool.get_state_layer_ids() == [0, 4]
```

- [ ] Add a PD registration test showing prefill/decode active IDs match and count mismatch fails instead of truncating.

- [ ] Run remote ETE:

```bash
# First run fills L1/L2 cache
python benchmark/gsm8k/bench_sglang.py \
  --host g0034 --port 17100 --backend srt \
  --data-path /mnt/beegfs/cjy/data/gsm8k_test.jsonl \
  --num-questions 200 --parallel 64 \
  --temperature 0.0 --top-p 1.0 --max-new-tokens 512

# Second run should be cache-hit heavy and must not lose accuracy
python benchmark/gsm8k/bench_sglang.py \
  --host g0034 --port 17100 --backend srt \
  --data-path /mnt/beegfs/cjy/data/gsm8k_test.jsonl \
  --num-questions 200 --parallel 64 \
  --temperature 0.0 --top-p 1.0 --max-new-tokens 512
```

Expected result: second run accuracy stays within normal sample variance of the first run; no large cache-hit drop.

---

## 8. Phase P5: compact device NSA index allocation

### Goal

Reduce L1/device index-cache memory allocation by allocating buffers only for active target index layers.

### Implementation steps

- [ ] Pass active target index layers from model config to pool construction. The construction site must use:

```python
from sglang.srt.configs.nsa_index_layers import build_nsa_index_layer_plan

index_plan = build_nsa_index_layer_plan(
    self.model_config.hf_config,
    start_layer,
    end_layer,
    is_nextn=self.model_config.is_draft_model,
)
```

- [ ] Pass `index_active_layer_ids=index_plan.active_layer_ids` to `NSATokenToKVPool`.

- [ ] Change `NSATokenToKVPool.__init__` allocation:

```python
index_buffer_layer_num = len(self.index_active_layer_ids)
self.index_layer_to_slot = {
    layer_id: slot for slot, layer_id in enumerate(self.index_active_layer_ids)
}
self.index_k_with_scale_buffer = [
    torch.zeros(...)
    for _ in range(index_buffer_layer_num)
]
```

- [ ] Update all direct indexing sites from `layer_id - self.start_layer` to `self.get_index_layer_slot(layer_id)`:

```python
buf = self.index_k_with_scale_buffer[self.get_index_layer_slot(layer_id)]
```

Affected methods:
- `_get_index_k_with_scale_buffer`
- `get_index_k_continuous`
- `get_index_k_scale_continuous`
- `get_index_k_scale_buffer`
- `set_index_k_scale_buffer`

- [ ] Update `get_kv_size_bytes()` to sum compact buffers. The existing loop over `self.index_k_with_scale_buffer` remains correct after allocation compaction.

- [ ] Update `python/sglang/srt/model_executor/model_runner_kv_cache_mixin.py` estimate:

```python
from sglang.srt.configs.nsa_index_layers import build_nsa_index_layer_plan

index_plan = build_nsa_index_layer_plan(
    self.model_config.hf_config,
    start_layer,
    end_layer,
    is_nextn=self.model_config.is_draft_model,
)
cell_size += indexer_size_per_token * len(index_plan.active_layer_ids) * element_size
```

- [ ] Add tests:

```python
def test_compact_device_index_buffers_use_active_layer_slots():
    pool = make_fake_nsa_pool(layer_num=8, active_layer_ids=(0, 4), compact=True)
    assert len(pool.index_k_with_scale_buffer) == 2
    assert pool.get_index_layer_slot(0) == 0
    assert pool.get_index_layer_slot(4) == 1
    with pytest.raises(RuntimeError, match="inactive index layer requested"):
        pool.get_index_layer_slot(1)
```

- [ ] Run remote CUDA unit tests for NSA pool and runtime:

```bash
ssh g0034 "docker exec sglang-glm5-dev-2 bash -lc 'cd /sgl-workspace/sglang-tai && PYTHONPATH=python python -m pytest -q test/registered/unit/mem_cache/test_nsa_pool_host_unit.py test/registered/unit/mem_cache/test_cp_shared_kv_runtime.py'"
```

Expected output: all selected tests pass.

---

## 9. Phase P6: compact host HiCache NSA index allocation

### Goal

Reduce L2/host index-cache memory and direct H2D/D2H transfer stride by allocating host index buffers over active index slots instead of all logical layers.

### Implementation steps

- [ ] In `NSATokenToKVPoolHost.__init__`, derive:

```python
self.index_active_layer_ids = tuple(device_pool.index_active_layer_ids)
self.index_active_layer_num = len(self.index_active_layer_ids)
self.index_logical_to_slot = {
    layer_id: slot for slot, layer_id in enumerate(self.index_active_layer_ids)
}
```

- [ ] Change index layout dimensions:

```python
self.indexer_layout_dim = self.indexer_page_stride_size * self.index_active_layer_num
```

- [ ] Change `get_size_per_token()`:

```python
return base + self.indexer_size_per_token * self.index_active_layer_num * self.indexer_dtype.itemsize
```

- [ ] Change host allocation shapes:

```python
# layer_first
for _ in range(self.index_active_layer_num)

# page_first/page_first_direct
(self.indexer_page_num, self.index_active_layer_num, 1, self.indexer_page_stride_size)

# layer_page_first
(self.index_active_layer_num, self.indexer_page_num, 1, self.indexer_page_stride_size)
```

- [ ] Add a host helper:

```python
def get_index_host_layer_slot(self, layer_id: int) -> int:
    try:
        return self.index_logical_to_slot[layer_id]
    except KeyError as exc:
        raise RuntimeError(
            "[CP_SHARED_KV_FAIL_FAST][host_index_cache_layer] "
            f"inactive host index layer requested: layer_id={layer_id} "
            f"active_layer_ids={list(self.index_active_layer_ids)}"
        ) from exc
```

- [ ] Update all index transfer calls to pass slot IDs to host layout kernels and logical IDs only to model/device APIs:

For page-first-direct load:

```python
host_slot = self.get_index_host_layer_slot(layer_id)
device_slot = device_pool.get_index_layer_slot(layer_id)
_load_tai_transfer_kv_per_layer_direct_pf_lf()(
    src_ptrs=[self.index_k_with_scale_buffer],
    dst_ptrs=[device_pool.index_k_with_scale_buffer[device_slot]],
    src_indices=host_page_indices,
    dst_indices=device_page_indices,
    layer_id=host_slot,
    page_size=1,
)
```

For page-first-direct backup:

```python
host_slot = self.get_index_host_layer_slot(layer_id)
device_slot = device_pool.get_index_layer_slot(layer_id)
_load_tai_transfer_kv_per_layer_direct_lf_pf()(
    src_ptrs=[device_pool.index_k_with_scale_buffer[device_slot]],
    dst_ptrs=[self.index_k_with_scale_buffer],
    src_indices=device_page_indices,
    dst_indices=host_page_indices,
    layer_id=host_slot,
    page_size=1,
)
```

- [ ] Add shape tests for all supported layouts:

```python
@pytest.mark.parametrize("layout,expected_shape", [
    ("page_first_direct", (page_num, 2, 1, stride)),
    ("layer_page_first", (2, page_num, 1, stride)),
])
def test_host_index_cache_compacts_to_active_layers(layout, expected_shape):
    host_pool = make_fake_nsa_host_pool(layout=layout, layer_num=8, active_layer_ids=(0, 4))
    assert tuple(host_pool.index_k_with_scale_buffer.shape) == expected_shape
```

- [ ] Run remote CUDA roundtrip for direct page-first-direct and layer-page-first layouts. The roundtrip must write an active layer, load it back, and verify byte equality for the active layer only.

---

## 10. Phase P7: ETE validation and performance acceptance

### Required launch coverage

Run prefill with index skip enabled:

```bash
--enable-nsa-prefill-context-parallel \
--nsa-prefill-cp-mode in-seq-split \
--enable-nsa-prefill-cp-shared-kv \
--enable-hierarchical-cache \
--hicache-io-backend direct \
--hicache-mem-layout page_first_direct \
--kv-cache-dtype fp8_e4m3 \
--enable-cp-shared-kv-prefill-bs-gt1 \
--nsa-index-topk-freq 4
```

Run two validation modes:

1. `--nsa-index-topk-freq 4` without offset.
2. `--nsa-index-topk-freq 4 --nsa-index-skip-topk-offset 1`.

### Correctness commands

Use 200-question fast checks before full GSM8K:

```bash
python benchmark/gsm8k/bench_sglang.py \
  --host g0034 \
  --port 17100 \
  --backend srt \
  --data-path /mnt/beegfs/cjy/data/gsm8k_test.jsonl \
  --num-questions 200 \
  --parallel 64 \
  --temperature 0.0 \
  --top-p 1.0 \
  --max-new-tokens 512
```

Then run full GSM8K twice without service restart:

```bash
python benchmark/gsm8k/bench_sglang.py \
  --host g0034 \
  --port 17100 \
  --backend srt \
  --data-path /mnt/beegfs/cjy/data/gsm8k_test.jsonl \
  --num-questions 1319 \
  --parallel 64 \
  --temperature 0.0 \
  --top-p 1.0 \
  --max-new-tokens 512
```

Expected acceptance:

```text
Accuracy: 0.953 - 0.955 on first run
Accuracy: within normal variance on second cache-hit-heavy run
Invalid: 0.000 or near 0.000
No repeated gibberish / accept-len collapse attributable to prefill cache corruption
```

### Performance counters to collect

Add one temporary debug line gated by an existing CP shared-KV debug env, not a new env:

```text
[CP_SHARED_KV_INDEX_SKIP] active_layers=<n> total_layers=<m> state_bufs=<k> draft_state_bufs=<d> saved_index_layers=<m-n>
```

Collect from prefill logs:

```bash
grep -E "CP_SHARED_KV_INDEX_SKIP|state_bufs|maybe_send_extra_state|CP_SHARED_KV_FAIL_FAST|CP_SHARED_KV_FALLBACK" /mnt/beegfs/cjy/log/sglang_cp_hicache_*.log | tail -n 200
```

Expected after P4/P6 with `freq=4`: target `state_bufs` and host index layer dimension drop approximately from `78` to about `21` target index buffers, plus draft buffers when enabled.

---

## 11. Explicit risks and mitigations

1. **Prefill/decode state slot mismatch corrupts cache.**
   - Mitigation: P3 adds `state_layer_ids` and fails fast before compact registration.

2. **Draft/EAGLE semantics differ from target.**
   - Mitigation: keep draft state active/full unless the draft mapping is explicitly represented with negative IDs and tested.

3. **Model forward formula divergence.**
   - Mitigation: P0 replaces model-local logic with the shared helper and tests default, offset, pattern, and nextn cases.

4. **Skipping MLA KV by mistake would break attention.**
   - Mitigation: P2 only guards `_load_indexer_*` and `_backup_indexer_*`; `super().load_to_device_per_layer()` and `super().backup_from_device_per_layer()` still run for MLA KV on every layer.

5. **Mooncake/NIXL backward compatibility.**
   - Mitigation: new `state_layer_ids` frame is optional; compact state transfer requires both sides to provide matching IDs. If IDs are absent and buffer counts differ, compact mode fails instead of truncating.

6. **Memory estimate undercount.**
   - Mitigation: compact allocation is not enabled until `model_runner_kv_cache_mixin.py` uses active layer count.

7. **Pattern and PP slices.**
   - Mitigation: helper accepts `start_layer/end_layer` and validates pattern indexing. Unit tests include non-zero `start_layer` before compact allocation is enabled.

---

## 12. Implementation order and commit boundaries

Use these commit boundaries to keep regressions bisectable:

1. `Centralize NSA index skip layer planning`
   - P0 helper + tests + model formula replacement.

2. `Expose NSA active index layers in KV pools`
   - P1 metadata only, no behavior change except fail-fast inactive direct access if hit.

3. `Skip HiCache index transfer for inactive NSA layers`
   - P2 transfer skip only, no allocation compaction.

4. `Attach logical layer IDs to PD state registration`
   - P3 metadata + fail-fast validation, no state-buffer compaction.

5. `Transfer only active NSA index state buffers`
   - P4 compact PD state registration.

6. `Compact NSA device index cache by active layers`
   - P5 L1/device allocation reduction + memory estimate.

7. `Compact NSA host index cache by active layers`
   - P6 L2/host allocation reduction + direct transfer slot mapping.

Each commit message must follow the Lore Commit Protocol in `/root/sglang-work/AGENTS.md`.

---

## 13. Verification matrix

| Phase | Local verification | Remote verification | ETE required |
| --- | --- | --- | --- |
| P0 | `test_nsa_index_layers.py`, py_compile | same tests in container | no |
| P1 | py_compile | NSA pool tests | no |
| P2 | host pool transfer-skip tests | CUDA host pool tests | 200 GSM8K x2 |
| P3 | PD registration unit tests | same tests in container | no |
| P4 | state buf active-layer tests | Mooncake transfer unit tests | 200 GSM8K x2, then full GSM8K x2 |
| P5 | memory estimate tests | CUDA NSA pool tests | full GSM8K x2 |
| P6 | host shape tests | CUDA direct roundtrip | full GSM8K x2 + replay perf |

Completion criteria:

```text
- No CP shared-KV fallback warnings on the hot path.
- No state-buffer truncation warning in Mooncake/NIXL.
- GSM8K first run and cache-hit run maintain baseline accuracy.
- Replay throughput does not regress relative to current bs>1 baseline.
- Logs show active index state buffers are reduced when index_topk_freq > 1.
```

## 14. Implementation status ledger

- P0-P2: implemented in `d21952b90 Reduce inactive NSA index-cache transfer safely`.
  - Central helper now owns target/draft skip formula.
  - `NSATokenToKVPool` exposes active-layer metadata without allocation compaction.
  - HiCache indexer load/backup skips inactive target index layers while MLA KV still transfers every layer.
- P3: implemented locally after this plan revision.
  - `KVArgs.state_layer_ids` is populated for NSA state buffers.
  - CP draft NSA state is appended under negative layer IDs (`-1`, `-2`, ...).
  - Mooncake and NIXL registration roundtrip optional signed `state_layer_ids`.
  - Mooncake/NIXL layer-aware state transfer fails fast on layer-id mismatch; Mooncake also fails fast on layer-aware state-buffer count mismatch instead of truncating.
  - Verification: remote container `test_cp_shared_kv_transfer_mapping.py`, `test_pd_state_layer_ids.py`, and `test_cp_per_layer_transfer.py` passed (`39 passed`).
- P4: implemented locally after this plan revision.
  - `NSATokenToKVPool.get_state_buf_infos()` returns only active index-layer slots.
  - `NSATokenToKVPool.get_state_layer_ids()` returns active logical layer IDs.
  - Draft/EAGLE remains appended under P3 negative layer IDs.
  - Verification: remote container `TestNSAIndexerPageIndices`, `test_cp_shared_kv_transfer_mapping.py`, `test_pd_state_layer_ids.py`, and `test_cp_per_layer_transfer.py` passed (`52 passed`).
- P5: implemented locally after this plan revision.
  - Device/L1 NSA index cache allocation uses `len(index_active_layer_ids)` when `compact_index_layers=True`.
  - Model runner enables compact device index allocation for NSA pools.
  - Capacity estimation counts active index layers instead of all local layers.
  - Verification: remote container `TestNSAIndexerPageIndices`, `test_model_runner_kv_cache_mixin.py`, and `test_cp_shared_kv_runtime.py` passed (`141 passed`, `2 subtests passed`).
- P6: implemented locally after this plan revision.
  - `NSATokenToKVPoolHost` mirrors the device pool's active index-layer IDs and allocates host/L2 index buffers with active-layer depth only.
  - Host index transfer maps logical layer IDs to compact host slots for `layer_first`, `page_first_direct`, and `layer_page_first`.
  - Test-only/legacy host stubs without compact metadata keep dense-slot behavior; real compact host pools still fail fast if an inactive logical index layer is requested.
  - Direct HiCache tests now follow the production CPU-index contract before calling TAI `cudaMemcpyBatchAsync` paths.
  - Host-register allocation now checks `cudaHostRegister` errors and unregisters host tensors on object finalization; this prevents repeated CUDA unit tests from poisoning later CUDA ops with stale `cudaErrorHostMemoryAlreadyRegistered`.
  - Follow-up correction: target+draft shared `--hicache-size` budgeting must also use active index-layer count in `_estimate_hicache_size_per_token`; otherwise the host/L2 buffer is physically compact but the L2 token capacity remains conservatively sized as if all target layers still had index cache.
  - Verification: remote container P3-P6 suite passed:
    `test_nsa_pool_host_unit.py`, `test_model_runner_kv_cache_mixin.py`,
    `test_cp_shared_kv_transfer_mapping.py`, `test_pd_state_layer_ids.py`,
    `test_cp_per_layer_transfer.py`, and `test_cp_shared_kv_runtime.py`
    (`200 passed`, `2 subtests passed`).
- P7: not rerun after P3. Needs GSM8K/replay after P4+ compact changes.
