CP shared KV and HiCache need a stable contract where physical cache coverage is page-aligned, while scheduler/radix-visible hit length remains the valid token length. This records the contract, adds page-aligned extent metadata, keeps owner assignment on actual tail pages instead of short-prefix fallback, and updates partial current reuse tests around tail-page masking.
Constraint: CP owner lanes operate on page units while scheduler and radix hit accounting must remain token-valid.
Rejected: Pad short suffixes to cp_size or 2*cp_size pages | wastes KV capacity and can turn a small tail into a much larger physical span.
Rejected: Silent direct-write or prefetch fallback | production fallback must be warning-visible for diagnosis.
Confidence: medium
Scope-risk: moderate
Directive: Do not reintroduce replicated short-radix fallback without checking docs/advanced_features/nsa_prefill_cp_page_aligned_cache_contract.md.
Tested: local py_compile for touched runtime, utility, owner, and unit-test files.
Tested: remote g0034 container three-file suite: 122 passed, 5 warnings.
Not-tested: full local pytest, blocked by missing runtime dependencies such as orjson.
Not-tested: CUDA E2E runtime for this commit.
Co-authored-by: OmX <omx@oh-my-codex.dev>
The CP draft shared-KV path only needs this rank's local draft tokens, but the previous compatibility path embedded the full prompt before CP-splitting. For long MTP/EAGLE prefill this recreates the large hidden activation that CP shared KV is trying to avoid.\n\nThis pads local draft input ids to the per-rank max token count recorded in NSA CP metadata, embeds the padded local tensor, then trims back to the true local length. That keeps rank shapes compatible while avoiding full-prompt embedding on every rank. Missing or stale metadata keeps the existing full-embedding fallback.\n\nConstraint: CP ranks can own uneven token counts, so the local embedding path needs a rank-uniform padded shape.\nRejected: Pad local ids to the full prompt length | this preserves compatibility but loses the intended memory reduction.\nConfidence: medium\nScope-risk: moderate\nDirective: Do not remove the full-embedding fallback unless all CP draft metadata producers guarantee max_rank_len for every prefill path.\nTested: g0034 container py_compile for utils.py and deepseek_nextn.py; g0034 container pytest -q test/registered/unit/layers/test_nsa_cp_utils.py => 25 passed, 5 warnings.\nNot-tested: Full distributed E2E with HiCache cache-hit MTP accept-length recovery.
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
Prefill CP shared KV uses the PAGED fused topk path, but topk_transform still built RAGGED topk offsets before dispatching by method. That introduced cumsum/repeat_interleave work after MQA, showing up as DeviceScanInitKernel and host/device traffic in profiles even though PAGED topk only needs cu_seqlens_q. Move metadata construction into the selected branch and pass precomputed single-segment CP cu_seqlens overrides from NSA CP metadata.
Constraint: PAGED fused topk needs cu_seqlens_q but does not consume topk_indices_offset.\nRejected: Add a kernel to fuse repeat_interleave for PAGED | the offset is unused in the current path, so avoiding it is cheaper and safer.\nConfidence: high\nScope-risk: narrow\nDirective: Do not reintroduce unconditional topk_indices_offset construction in topk_transform; keep RAGGED-only metadata on the RAGGED branch.\nTested: python -m py_compile for modified files locally; g0034 container python3 -m py_compile for modified files; g0034 container python3 test/registered/unit/layers/test_nsa_cp_utils.py ran 23 tests OK.\nNot-tested: Full server profile after restart; full SGLang test suite.
The CP in-seq NSA indexer path was constructing a CUDA validity mask around each fp8_mqa_logits call, synchronizing it back with .item(), and then rebuilding small CUDA tensors before topk. The valid rows are always a contiguous prefix because ke_offset is monotonic, so the count can be derived from existing CP metadata and the path can use slices instead of boolean gather/scatter.
RAGGED topk now reuses the zero row-start vector as the per-row offset override for the single-segment case to avoid cumsum/repeat_interleave metadata kernels after MQA.
Constraint: CP shared KV prefill still needs to preserve prev/next causal visibility and topk transform semantics.
Rejected: Add a new Triton range-builder first | the Python-side synchronization and boolean indexing were the larger immediate risk, and kernel work should be profiled after this smaller change.
Confidence: medium
Scope-risk: moderate
Directive: Do not reintroduce GPU mask .item() or CPU-created CUDA tensors in the fp8_mqa_logits hot path without profiler evidence.
Tested: python3 -m py_compile python/sglang/srt/layers/attention/nsa/nsa_indexer.py test/registered/unit/layers/test_nsa_cp_utils.py
Tested: g0034 docker /sgl-workspace/sglang-tai python3 test/registered/unit/layers/test_nsa_cp_utils.py -> Ran 22 tests OK
Not-tested: full GLM5 PD server throughput/profiler run
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 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
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
