From 6229c7da60052b0640f76f410cd6ede12b64e9fa Mon Sep 17 00:00:00 2001 From: laoyao0822 Date: Wed, 10 Jun 2026 04:28:26 +0800 Subject: [PATCH] Plan index-skip cache and transfer reduction Record the phased design for turning NSA index-topk sharing from a compute-only optimization into HiCache and PD-transfer savings. The plan deliberately separates metadata validation from allocation compaction so prefill/decode state-slot mismatches fail before any bandwidth reduction can corrupt cache contents. Constraint: CP shared-KV cache corruption previously appeared only on cache-hit paths, so state-layer mapping must be explicit before compact transfer or allocation. Rejected: Compact L1/L2 index buffers first | without state_layer_ids this can silently map different logical layers to the same transfer slot across prefill and decode. Confidence: high Scope-risk: narrow Directive: Do not skip the state_layer_ids phase before compacting NSA index state buffers. Tested: Documentation placeholder scan and referenced source-path existence check. Not-tested: Runtime behavior; this commit is documentation only. --- ..._cp_index_skip_cache_communication_plan.md | 1108 +++++++++++++++++ 1 file changed, 1108 insertions(+) create mode 100644 docs/advanced_features/nsa_prefill_cp_index_skip_cache_communication_plan.md diff --git a/docs/advanced_features/nsa_prefill_cp_index_skip_cache_communication_plan.md b/docs/advanced_features/nsa_prefill_cp_index_skip_cache_communication_plan.md new file mode 100644 index 000000000..23b840976 --- /dev/null +++ b/docs/advanced_features/nsa_prefill_cp_index_skip_cache_communication_plan.md @@ -0,0 +1,1108 @@ +# 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= total_layers= state_bufs= draft_state_bufs= saved_index_layers= +``` + +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. +```