CP HiCache now treats draft KV as a strict target-owned payload through pending write visibility, host eviction, and state-buffer registration. Host metadata created before async D2H ack is no longer request-visible, so match_prefix cannot select an in-flight host node. Draft host eviction now fails before target cleanup when draft metadata is missing, and prefill/decode share one helper for draft NSA state buffers so shared-KV mode cannot silently skip mismatched draft state.
Constraint: CP shared KV + HiCache + EAGLE/MTP must not expose target-only host hits or skipped draft state as valid cache hits
Rejected: Rely on event-loop ordering and lock_ref to hide in-flight writes | match_prefix does not consult lock_ref and can observe host_len/cp_hicache directly
Rejected: Keep draft state mismatch as debug-only skip | it can poison speculative acceptance while looking like a successful cache hit
Confidence: high
Scope-risk: moderate
Directive: Do not reintroduce silent draft/target fallback in CP shared-KV HiCache paths; malformed strong-sync metadata should fail fast
Tested: python -m py_compile targeted modified files
Tested: remote g0034 container pytest test/registered/unit/mem_cache/test_cp_hicache_metadata.py test/registered/unit/managers/test_hicache_controller_cp.py test/registered/unit/disaggregation/test_cp_shared_kv_transfer_mapping.py -q (91 passed)
Not-tested: Full CP shared KV + HiCache + EAGLE/MTP ETE server run after this commit
CP HiCache host hits must not advertise target residency unless the draft payload is also valid when EAGLE/MTP draft HiCache is attached. This closes target-only metadata paths by making the CP host-valid predicate and load replay fail fast, resets draft host storage with the target host pool, and records the P1-P3 strong-sync plan state.
The page-index validator is restored for CPU/fake-test tensors only, preserving unit-test coverage for malformed page spans without reintroducing CUDA hot-path host sync.
Constraint: CP shared KV + HiCache + EAGLE/MTP cannot safely demote malformed target/draft metadata to an ordinary cache miss
Rejected: keep permissive fallback for missing draft_host_indices | it can look like a successful cache hit while poisoning speculative acceptance
Rejected: re-enable generic CUDA tensor page validation | it can force host sync in the HiCache transfer hot path
Confidence: high
Scope-risk: moderate
Reversibility: clean
Directive: Do not add silent fallback around CP draft HiCache metadata; unexpected target/draft divergence should fail fast with node/rank context
Tested: remote container targeted tests: 5 passed
Tested: remote container files test_cp_hicache_metadata.py and test_hicache_controller_cp.py: 77 passed
Tested: remote container test_page_index_utils.py: 8 passed
Tested: local git diff --check and py_compile for modified Python files
Not-tested: full CP shared KV + HiCache + EAGLE/MTP ETE
Co-authored-by: OmX <omx@oh-my-codex.dev>
CP HiCache previously let the target host pool and the draft/MTP host pool each consume the full --hicache-size budget. With EAGLE/MTP enabled this doubled per-rank host allocation and could kill scheduler ranks during startup before Python emitted a traceback. The cache now treats target KV and draft KV as one logical host-cache object: target and draft capacities are computed from one per-rank byte budget, draft may receive more token capacity when its per-token footprint is smaller, and draft attachment remains tied to target residency.
Constraint: --hicache-size is a per-rank host budget and must not be multiplied by attaching draft KV.
Rejected: Give draft another independent --hicache-size allocation | repeats the observed host OOM failure mode.
Rejected: Disable draft HiCache attachment under CP | avoids OOM but breaks target/draft cache-hit consistency for MTP.
Confidence: medium
Scope-risk: moderate
Directive: Keep target and draft KV as one logical HiCache object; do not let draft host allocation consume an independent full hicache-size budget.
Tested: python -m py_compile on modified scheduler/cache/test files
Tested: remote g0034 container PYTHONPATH=python python -m pytest test/registered/unit/mem_cache/test_cp_hicache_metadata.py -q (45 passed)
Not-tested: full multi-rank GLM5 server restart after clearing existing remote router/defunct process state
Co-authored-by: OmX <omx@oh-my-codex.dev>
CP shared KV already registered the draft model main KV buffer with the
prefill/decode Mooncake managers, but NSA draft state buffers were not part
of the state registration set. HiCache/cache-hit traffic could then transfer
pages from the draft pool without transferring the matching draft
index/scale state, which is a plausible cause of the EAGLE/MTP accept-length
collapse after cache hits.
This appends compatible draft NSA state buffers to the existing state
transfer registration on both prefill and decode, and extends transfer-side
diagnostics so source/destination state-buffer counts are visible. The
mismatch guard degrades to the common prefix of registered state buffers
instead of crashing if a rolling deployment exposes asymmetric registration.
Constraint: Scope is intentionally limited to target_state_type=nsa and draft_state_type=nsa.
Rejected: Treat draft main KV transfer as sufficient | NSA attention also needs draft index/scale state for transferred pages.
Rejected: Add Mamba/SWA draft-state semantics now | those state layouts need separate correctness analysis.
Confidence: medium
Scope-risk: moderate
Directive: Do not remove the draft_state_buffer_start/count fields without checking Mooncake source/destination registration symmetry.
Tested: PYTHONDONTWRITEBYTECODE=1 python3 -m py_compile python/sglang/srt/disaggregation/prefill.py python/sglang/srt/disaggregation/decode.py python/sglang/srt/disaggregation/mooncake/conn.py
Tested: git diff --check
Tested: Remote prefill log showed registered_state_bufs=79 and maybe_send_extra_state src_state_bufs=79 dst_state_bufs=79 with no state-buffer mismatch.
Not-tested: Full accept-length recovery; latest remote run hit an unrelated prefill KV allocator idle-check leak after transfer registration succeeded.
Prefill CP only needs the local hidden shard for DeepSeek NextN draft extend. The change adds a draft shared-KV path that captures target hidden locally, feeds only the CP-local slice into the draft model, and keeps draft KV writes/transfers on the same shared logical-to-physical page mapping as target KV.\n\nDebug logs are gated behind SGLANG_CP_DRAFT_SHARED_KV_DEBUG and cover scheduler pool selection, KV manager buffer registration, local physical writes, prefill sender filtering, transfer pages, and decode commit metadata so ETE runs can prove draft KV is sharded rather than full-concatenated on a prefill rank.\n\nConstraint: Prefill runs CP while decode remains DP, so prefill must avoid full hidden/KV materialization but decode still receives full logical KV pages.\nRejected: Keep draft extend on full hidden state | preserves correctness but wastes prefill memory and defeats CP shared-KV intent.\nRejected: Transfer draft KV with a separate mapping | target and draft pools share req_to_token logical indices, so duplicating mapping adds risk without benefit.\nConfidence: medium\nScope-risk: moderate\nDirective: Do not remove the debug logs until ETE evidence confirms draft MLA/index writes and transfer pages are CP-sharded on all ranks.\nTested: Remote compileall for changed CP draft, transfer, scheduler, NSA index, MLA write, and EAGLE files.\nNot-tested: Full GLM-5 EAGLE ETE with SGLANG_CP_DRAFT_SHARED_KV_DEBUG=1 after this logging addition; local pytest intentionally not run.
The CP shared-KV path now has a gated tai-kernel replacement for NSA index
K/scale plus MQA range preparation, and Phase8 prefetch can skip tiny prefixes
that do not cover all CP lanes. The Phase9 plan documents the next scheduler
work for overlapping CP communication with peer-request attention windows.
Temporary diagnostic logs added while validating prefetch ownership and fused
index prepare routing were removed before committing so the runtime path does
not add log-only synchronization, log counters, or shape-reporting overhead.
Constraint: Production profiling showed small per-request CPU/GPU overhead from diagnostic logging and sync-prone debug counters.
Rejected: Keep fused-index prepare fallback/used logs behind a new env var | it leaves another runtime branch and logging surface for a path that should be benchmarked with profiler evidence instead.
Rejected: Keep owned page-count prefetch logs | they require sync-prone tensor reductions and were only useful for one-off diagnosis.
Confidence: medium
Scope-risk: moderate
Directive: Reintroduce CP shared-KV diagnostics only behind explicit debug paths, and avoid .item()/shape-heavy logging in hot prefill paths.
Tested: git diff --check for staged sglang-dev changes.
Tested: AST parse for environ.py, cp_shared_kv_prefetch.py, cp_shared_kv_runtime.py, nsa_indexer.py, and test_cp_shared_kv_runtime.py.
Not-tested: Full unit test suite.
Not-tested: Multi-node GLM5 prefill/decode/router runtime after this exact commit.
CP shared KV materialization repeatedly rebuilt the same logical-page slot remaps and page inverse metadata for each layer. Cache the token and paged remap metadata on the forward batch so MLA KV, index K/scale, and prefetch paths can reuse the layer-independent mapping while still materializing layer-specific data through the existing tai/torch runtime paths.
Constraint: Only mapping metadata is batch-scoped; dense KV/index contents remain layer-specific and are not reused.
Rejected: Cache fully materialized dense KV/index buffers | would add large per-layer memory residency and invalidation complexity.
Confidence: medium
Scope-risk: moderate
Directive: Do not assume this removes materialize or CP all-reduce cost; profile tai fallback logs and Nsight kernels before attributing E2E gains or losses.
Tested: git diff --check
Tested: remote g0034 container PYTHONPATH=python python3 -m pytest test/registered/unit/mem_cache/test_cp_shared_kv_runtime.py -q (52 passed, 5 warnings)
Not-tested: Full GLM-5 disaggregated E2E performance run
The shared-KV path now keeps more CP metadata on-device and reuses
physical out-cache locations across MLA and NSA index writes, so each
layer avoids repeating logical-to-physical remaps. The in-seq CP
all-gather rerange path now delegates to tai-kernel when available and
falls back to the existing torch split/cat path with an explicit log.
This also extends the Phase8 prefetch machinery to cover shared KV
materialization metadata and keeps debug/fallback behavior gated so the
fast path is not polluted by diagnostic checks.
Constraint: Custom CP kernels must live in tai-kernel and be imported lazily from SGLang
Constraint: Decode does not use CP; these changes target NSA prefill CP in-seq-split shared KV
Rejected: Recompute physical local cache locations separately for MLA and index writes | repeats the same remap work every layer
Rejected: Keep the in-seq rerange Triton code inline in SGLang | duplicates kernel ownership and blocks tai-kernel reuse
Confidence: medium
Scope-risk: moderate
Directive: Keep CP collective ordering identical across ranks; do not add rank-local fallback decisions inside shared KV materialize paths
Tested: Remote g0034 container py_compile for modified SGLang/tai-kernel files; remote pytest test/registered/unit/layers/test_nsa_cp_utils.py passed with 24 tests
Not-tested: Full multi-node GLM5 prefill/decode throughput after the final commit boundary
The shared-KV prefill path now optionally calls tai_kernel.nsa_prefill.fused_store_mla_kv before falling back to logical_locs_to_physical plus set_mla_kv_buffer. The fast path supports packed FP8 and BF16/FP16 direct KV buffers, while debug mode and kernel failures still preserve the existing fallback behavior. Success logging was removed after path verification because per-layer/per-rank logs are too noisy in normal server runs.
Constraint: Runtime must remain safe when tai-kernel is absent or debug checks are enabled
Rejected: Keep success logs permanently | floods prefill logs once every rank/layer starts using the fast path
Confidence: high
Scope-risk: moderate
Directive: Keep fallback warnings; do not re-add per-layer success logs outside explicit debug instrumentation
Tested: g0034 container python -m py_compile python/sglang/srt/layers/attention/nsa/cp_shared_kv_runtime.py
Tested: g0034 container PYTHONPATH=python pytest -q test/registered/unit/mem_cache/test_cp_shared_kv_runtime.py -q (40 passed)
Not-tested: Full multi-node PD server throughput after log removal
CP shared KV now avoids the PyTorch sort/search remap for the single-request current-only path by deriving compact rows from page-level inverse mapping. The same change keeps sort NVTX attribution gated and splits high-frequency MoE sort markers behind a separate env var so profiling does not perturb normal runs.
Decode-side disaggregation prealloc also avoids rebuilding large token index tensors and records finer allocation timing, while compute-owner allocation/free tests cover the shared-KV page-lane behavior.
Constraint: The runtime tree used for validation is the remote /sgl-workspace/sglang-tai mount, which is not itself a Git repository, so these tracked files were synchronized into the local repo before commit.
Rejected: Keep torch.sort/searchsorted for current remap | it emits ATen/CCCL radixSortKVInPlace kernels in the attention hot path.
Rejected: Enable MoE sort NVTX under the generic sort env | the MoE preprocess sort is too frequent and can make profiling look like a hang.
Confidence: medium
Scope-risk: moderate
Directive: Do not reintroduce token-level torch.sort/searchsorted in CP shared-KV current remap without profiling the attention hot path under Nsight.
Tested: Remote container py_compile for modified runtime files; git diff --cached --check.
Not-tested: Full multi-node GLM5 PD throughput/profile rerun after the page-inverse current remap.
Phase8 only gains useful overlap when the next-layer MLA prefix prefetch is allowed to run until the next layer actually consumes the prefetched buffer. The old wait-after-attention switch let runtime configuration collapse the optimization back into current-layer tail latency, so the prefetch path now has one wait policy and the documentation records the implemented behavior.
Constraint: Phase8 should keep the production environment surface minimal while preserving the existing enable and debug-log knobs
Rejected: SGLANG_CP_SHARED_KV_MLA_PREFETCH_WAIT_AFTER_ATTENTION | it reintroduced current-layer synchronous waiting and made profiling behavior depend on a nonessential policy knob
Confidence: medium
Scope-risk: narrow
Directive: Do not add another Phase8 wait policy knob without first proving the added policy improves end-to-end prefill latency under CP shared KV
Tested: Python AST parse for touched Python files
Tested: git diff --check
Not-tested: Full pytest and remote server integration were not run in this commit
Shared KV now relies on page-aligned CP metadata and compute-owner page allocation so persistent MLA KV and NSA index shards can be written by the rank that computed them. The compatibility read path keeps the dense full-view contract for existing topk and attention kernels, but removes duplicated prev/next index materialize, adds optional tai materialize integration, and tightens tests/docs around the fallback boundaries.
Constraint: Decode remains non-CP while prefill CP owns the shared-KV changes
Constraint: Existing attention/topk kernels still expect dense full-view KV/index inputs
Rejected: Change attention kernels to read owner-sharded KV directly | larger semantic change reserved for later phases
Rejected: Merge index K/scale storage with MLA KV storage | would couple topk and attention cache lifecycles before materialize overhead is isolated
Confidence: medium
Scope-risk: broad
Directive: Do not remove fallback logging or debug-gated assertions without reproducing long-context chunked/radix-hit paths
Tested: git diff --check --cached
Not-tested: Local pytest/runtime server verification not run in this commit step per current workflow constraints
Phase 8 needs a narrow boundary because previous shared-KV phases already changed persistent layout and materialize behavior. The plan records an MLA-only, one-layer-ahead prefix prefetch path that targets chunked-prefill/radix-hit sync materialize overhead without adding index prefetch or bandwidth scheduling yet.
Constraint: Phase 8 must not introduce SGLANG_CP_SHARED_KV_LAYER_PREFETCH_KIND
Constraint: Current request is documentation only; do not modify runtime code
Rejected: Add index K/scale prefetch to the same phase | widens topk correctness and collective-order risk before MLA overlap is validated
Rejected: Persist dense KV across chunks | reintroduces large full-context memory pressure
Confidence: high
Scope-risk: narrow
Directive: Keep Phase 8 v1 MLA-only unless profiling proves index materialize is the next blocker
Tested: git diff --check -- docs/advanced_features/nsa_prefill_cp_phase8_mla_prefix_prefetch_plan.md
Not-tested: Runtime/server execution; documentation-only commit
Phase 5 needs each current KV page to have exactly one CP compute owner before local KV/index direct writes can be safe. This change teaches in-seq NSA prefill CP to produce page-aligned split metadata under shared-KV mode, threads page size into the metadata builders, and fixes local pair splitting so unequal page-aligned zigzag segments do not corrupt topk inputs.
Constraint: Phase 5 direct-write layout requires page ownership to be expressible at page granularity
Constraint: Short page-unit batches remain on the token-balanced fallback to avoid zero-page segment risk
Rejected: Split local q/weights by half | page-aligned zigzag segments can have unequal token counts
Confidence: medium
Scope-risk: moderate
Directive: Do not enable compute-owner direct writes unless nsa_cp_metadata.page_aligned is true and local loc ownership is verified
Tested: python3 -m py_compile python/sglang/srt/layers/attention/nsa/utils.py python/sglang/srt/layers/attention/nsa/nsa_indexer.py python/sglang/srt/models/deepseek_v2.py python/sglang/srt/models/deepseek_nextn.py test/registered/unit/layers/test_nsa_cp_utils.py
Not-tested: Local pytest collection is blocked in this environment by missing pybase64; container/runtime tests were not rerun during this commit step
Phase 4 is now documented as the page-aligned CP split step, leaving the existing shared KV layout untouched while establishing the page ownership invariant needed for direct writes. Phase 5 is documented as the follow-up compute-owner layout change that can route local MLA KV and NSA index writes into the owning rank's physical pool.
Constraint: Current Phase 2/3 layout is page-interleaved and still relies on compatibility materialize paths.
Rejected: Combine split alignment, allocator changes, and shard-aware attention into one phase | too broad to debug safely.
Confidence: high
Scope-risk: narrow
Directive: Do not start compute-owner KV layout work before preserving the Phase 4 page-aligned split invariant and fallback semantics.
Tested: git diff --cached --check
Not-tested: Runtime server startup; documentation-only commit.
Phase 2 expanded prefill CP persistent KV capacity by sharding KV at rest, but its compatibility runtime can still rebuild full-view KV/index buffers from sharded storage. This documents a narrower Phase 3 target: reuse current chunk KV/index that was already CP all-gathered and reranged, while leaving deeper history shard-aware work to Phase 4.
The new plan separates current-only and mixed current/history cases, defines fallback requirements to the Phase 2 full-view compatibility path, and records validation criteria before any runtime implementation starts.
Constraint: Phase 3 must reduce avoidable duplicate materialization without rewriting NSA topk or sparse attention kernels
Constraint: History KV/index access remains Phase 2 compatibility materialization until Phase 4
Rejected: Start with shard-aware topk/distributed attention | too broad for Phase 3 and harder to validate against the current Phase 2 baseline
Rejected: Remove all shared-KV runtime materialization immediately | history/prefix/chunked-prefill paths still need a compatibility fallback
Confidence: high
Scope-risk: narrow
Directive: Keep Phase 3 scoped to current chunk reuse; put selected-history and shard-aware algorithms in Phase 4 unless the design is explicitly revised
Tested: git diff --check for the two documentation files
Not-tested: Runtime behavior; documentation-only change
Prefill CP previously replicated NSA/MLA persistent KV on every CP rank, so CP8 consumed eight copies of KV memory while exposing only one rank of logical cache capacity. This change splits logical KV locs from per-rank physical storage, shards MLA latent KV and NSA index K/scale by deterministic page ownership, and keeps existing NSA attention kernels working through a full-view runtime materialization layer.
Mooncake PD transfer now sends each prefill CP rank's owned physical pages with explicit logical page positions so non-CP decode can reconstruct full-layout KV. The implementation is guarded by an explicit server flag and startup checks, and the design documentation records the implemented scope, debug environment, and Phase 3 boundary.
Constraint: Phase 2 must preserve existing NSA attention/index kernels via runtime full-view materialization
Constraint: Decode side remains non-CP and receives full KV through Mooncake
Rejected: Shard-aware NSA attention in this change | belongs to Phase 3 because it requires distributed topk/softmax/output contracts
Rejected: Request-contiguous CP ownership | unstable under chunked prefill and tied to attention split mode
Confidence: medium
Scope-risk: broad
Directive: Do not enable round-robin CP shared KV without wiring runtime materialization/PD transfer contracts for that split mode
Directive: Keep SGLANG_DEBUG_CP_SHARED_KV disabled for perf measurements; it intentionally enables CUDA-syncing diagnostics
Tested: Remote py_compile for shared-KV touched Python files in g0034 container
Tested: Remote pytest selected cp_shared/shared_kv/nsa suite: 37 passed, 34 deselected
Not-tested: Full GLM5 multi-node throughput/regression run after final doc update
Not-tested: Phase 3 shard-aware runtime, round-robin CP mode, and non-Mooncake PD backends
- nsa_prefill_cp_all_gather.md: explains why there are two all-gathers
in current NSA prefill CP (per-layer KV/index gather vs tail output
hidden gather), with data flow diagrams and involved file lists
- nsa_prefill_cp_phase1_narrow_output_collection.md: detailed Phase 1
design plan covering batch-level eligibility, fallback conditions,
affected modules, KV transfer impact, risks, and test matrix
