chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 11 Packages Projects Releases Wiki Activity

Compact skipped NSA index-cache state safely

Browse Source 
Index skip reduces the number of target layers that own NSA index state,
but PD transfer and HiCache still assumed dense full-layer state buffers.
This change carries explicit state layer IDs through prefill/decode
registration, compacts device and host index buffers to active layers,
and maps logical layer IDs to compact slots on transfer paths.

The PD side fails fast when prefill/decode disagree on NSA state layer
identity instead of silently truncating or copying mismatched buffers.
Host direct tests now use the same CPU-index descriptor contract required
by the TAI cudaMemcpyBatchAsync path, and host registered memory is
unregistered on tensor finalization to avoid stale cudaHostRegister state
across CUDA tests.

Constraint: CP shared-KV with index_topk skip must keep target/draft state identity explicit before compacting buffers
Constraint: Direct HiCache TAI transfer rejects CUDA indices to avoid hidden D2H copies on the control path
Rejected: Keep full-layer L1/L2 index buffers | wastes the memory/bandwidth that index skip is meant to save
Rejected: Infer state buffer order by count only | can silently corrupt cache when active layer sets differ
Confidence: high
Scope-risk: moderate
Directive: Do not compact or reorder NSA state buffers without carrying logical layer IDs through PD registration and validating both sides
Tested: Remote container py_compile for touched runtime files
Tested: Remote container pytest: 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, test_cp_shared_kv_runtime.py -> 200 passed, 2 subtests passed
Not-tested: Full ETE GSM8K/replay after compacted P3-P6 changes
Co-authored-by: OmX <omx@oh-my-codex.dev>
This commit is contained in:
laoyao0822
2026-06-10 05:37:06 +08:00
parent d21952b903
commit 1ebde44e59
14 changed files with 611 additions and 16 deletions
Download Patch File Download Diff File Expand all files Collapse all files

35
docs/advanced_features/nsa_prefill_cp_index_skip_cache_communication_plan.md
View File

@@ -1106,3 +1106,38 @@ Completion criteria:
- 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`.
- 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.