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>
HiCache host hits can be skipped before load-back when the quota gate only counts immediately free KV allocator space. Under CP shared-KV pressure most reusable capacity may be represented as evictable radix-cache leaves, so the gate can incorrectly reject a host hit and leave prefill with cached-token zero despite host residency. Count device evictable cache in the quota estimate while leaving actual owner-lane allocation and eviction checks in the load path.
Constraint: CP HiCache load-back still has to respect owner-lane allocation and allocator eviction semantics.
Rejected: Force load-back regardless of quota | would bypass the scheduler pressure signal and increase OOM risk.
Rejected: Treat cache-hit zero as a transfer issue | logs showed host hits were found but skipped by quota before transfer.
Confidence: medium
Scope-risk: moderate
Directive: Do not remove evictable cache from load-back capacity accounting without checking CP HiCache host-hit behavior under device pressure.
Tested: git diff --check
Tested: remote g0034 container pytest -q test/registered/unit/managers/test_prefill_adder.py test/registered/unit/managers/test_hicache_controller_cp.py test/registered/unit/mem_cache/test_cp_hicache_metadata.py test/registered/unit/mem_cache/test_alloc_pages_with_owners.py (90 passed, 3 warnings)
Not-tested: Full ETE GLM5 CP+HiCache+EAGLE pressure run after this quota change
Co-authored-by: OmX <omx@oh-my-codex.dev>
inc_node_lock_ref/dec_node_lock_ref adjusted evictable_size_ but did
not call _update_leaf_status, so nodes could become "phantoms" —
counted in evictable_size_ but missing from evictable_leaves. Under
load with write-behind, this caused eviction to return 0 despite
large evictable_size_, leading to OOM:
Available tokens: 778624 (evictable_size=706880)
evict_result=(num_tokens_evicted=0)
The race: write-behind locks node X via inc_node_lock_ref (X stays
in evictable_leaves). A request path touches X via inc_lock_ref,
which calls _update_leaf_status and removes X from evictable_leaves.
Request finishes, dec_lock_ref keeps X out (lock_ref still >0).
Write ack calls dec_node_lock_ref dropping lock_ref to 0, but never
calls _update_leaf_status — X is permanently lost.
Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
CpSharedKVMlaPrefetcher.create and CpSharedKVIndexPrefetcher.create
returned None whenever ``token_to_kv_pool.layer_transfer_counter`` was
set, which permanently disables NSA indexer prefetch whenever HiCache
is active — even when no H2D transfer is in progress.
The guard is unnecessary: buffer getters synchronize via wait_until(),
and the prefetcher's stream calls wait_stream(current_stream) before
materialization. Removing it restores intended prefetch parallelism
under HiCache + NSA.
Only fix#3 of e293d4a39 is portable to this branch — fixes#1 (SHM
timeout) and #2 (SHM exception safety) target shm_allreduce.py and the
batched-CP coordination path in hiradix_cache.py, neither of which
exists here.
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
The ported 44ba832 added _validate_cp_hicache_page_indices gated on
envs.SGLANG_DEBUG_HICACHE_VALIDATE, but that env was introduced by
97a9f850c (not on this branch). Calling _write_cp / load_cp crashed:
AttributeError: 'Envs' object has no attribute 'SGLANG_DEBUG_HICACHE_VALIDATE'
The validator is purely defensive (page-alignment holds by construction
in HostKVCache.alloc and CpSharedKVLayout.logical_locs_to_physical), so
remove the function entirely along with its 4 call sites in _write_cp
and load_cp. The wrapping try/except blocks existed solely to free
allocations when validation raised; with no raising call inside, they
become dead and are removed too. Stale imports (envs,
validate_page_aligned_token_indices) dropped.
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
In CP=8 + NSA-shared-KV + HiCache disagg-prefill, cache-hit prefill produced
incoherent decode output. Cold prefill on CP was correct; pure CP without
HiCache was correct. The bug lived at the HiCache load_cp / device-alloc
interface.
Root cause: cache_controller.load_cp called the plain
mem_pool_device_allocator.alloc(logical_len), which returns logical pages
with no CP owner-pattern preservation. Cold prefill instead uses
alloc_extend_compute_owner with a zigzag owner pattern from
build_in_seq_page_compute_owners. The saved CpHiCacheNodeMetadata.owned_positions
records WHICH POSITIONS in the write-time alloc were owned by this rank. At
load time, those same positions are applied to a new alloc whose per-position
owner pattern is arbitrary -- each rank loads its host bytes into physical
slots whose corresponding logical page is owned by a DIFFERENT rank.
Attention's materialize_shared_token_kv_buffer reads from the owner's
physical slot, which was never loaded. Result: garbage.
Fix:
- CpHiCacheNodeMetadata gains two required fields: page_owners (int8 per
logical page, identical on all CP ranks) and page_size. __post_init__
validates; split() bisects page_owners by page index with a page-alignment
check.
- _write_cp derives page_owners from device_indices (page-first slot of each
page -> logical page id -> layout.owner_for_logical_pages) and stores in
both metadata-construction sites (zero-owned and normal).
- New CPSharedPagedTokenToKVPoolAllocator.alloc_pages_with_owners() reuses
_select_compute_owner_pages (with its tai-kernel fast path) and returns
page-contiguous token locs whose per-page owner sequence equals the input.
- load_cp now concats page_owners across nodes_to_load and calls
alloc_pages_with_owners. On None (lane exhausted) the caller hits the
retry-with-eviction path; further failure returns None and degrades to
cache miss. No silent fallback to plain alloc -- that recreated the bug.
- load_back retry path now calls _evict_for_compute_owner_lanes (module-top
import) instead of plain evict(); this targets the deficit lane and gives
the next alloc attempt a chance to satisfy it.
- envs import moved to module top in cache_controller.py per code-review
feedback. Removed an over-defensive owned_check.all().item() in load_cp
that would have re-introduced the host-sync anti-pattern 97a9f850c
removed -- the invariant is already guaranteed by alloc_pages_with_owners.
Tests: 40 existing CpHiCacheNodeMetadata constructions migrated to pass the
new required fields. 9 new metadata tests (validators + split page-alignment).
10 new allocator tests in test_alloc_pages_with_owners.py covering input-order
preservation, lane exhaustion, release_pages fallback, debug-mode invariant.
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
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.
Cache-hit prefill can skip draft forward for the prefix while PD transfer still reads draft KV for that same prefix. CP HiCache therefore needs to persist draft/MTP KV alongside target KV instead of relying on whatever remains in the draft GPU pool.
Constraint: CP HiCache is host-only here; storage backends remain unsupported for CP shared KV.
Constraint: CP shared KV must keep owner-page semantics and avoid falling back to full KV on every rank.
Rejected: Recompute cached-prefix draft KV during prefill | loses the HiCache benefit and reintroduces the large hidden/KV footprint.
Rejected: Change PD transfer to skip draft prefix KV | decode still needs draft cache continuity for MTP acceptance.
Confidence: medium
Scope-risk: moderate
Directive: Keep target and draft CP HiCache metadata/load/write/evict paths in lockstep; changing one without the other can silently reduce MTP accept length.
Tested: Remote g0034 container /sgl-workspace/sglang-tai: python3 -m pytest -q test/registered/unit/managers/test_hicache_controller_cp.py test/registered/unit/mem_cache/test_cp_hicache_metadata.py => 58 passed, 3 warnings
Not-tested: Full multi-node HiCache+MTP serving benchmark and accept-length recovery.
The pure-Python convert_to_bigram_key list comprehension runs on every
cache_finished_req / cache_unfinished_req of every EAGLE radix variant
(radix_cache, hiradix_cache, swa_radix_cache), with token-list lengths
that scale with prompt + output. Scheduler profiles consistently flag it
as the largest mem_cache-side CPU hotspot.
This commit wires sglang.srt.mem_cache.utils.convert_to_bigram_key to
tai_kernel.radix.convert_to_bigram_key when the extension is importable,
falling back to the pure-Python implementation otherwise. The tai-kernel
path uses a pybind11 module that calls the CPython C API directly
(PyTuple_New + PyTuple_SET_ITEM with ref-stealing) rather than rebuilding
the list comprehension's bytecode-level tuple allocations. Measured 1.4x
at n=131k and up to 2.5x for n=1k on g0034 Python 3.11; allocator-bound
at large n because CPython's 2-tuple freelist already amortises the
construction. The int64-packing follow-up that bypasses tuple allocation
entirely is parked as a separate work item.
Runtime safety:
- The dispatcher catches any first-call JIT compile / runtime failure,
logs once, and falls through to the pure-Python path for the rest of
the process — JIT failures must degrade rather than crash a serving
loop.
- SGLANG_DISABLE_TAI_BIGRAM forces the Python path for bisecting.
Constraint: Output must be a real Python List[Tuple[int, int]] because
downstream radix dicts use the tuples as hashable keys.
Rejected: Pre-allocated tuple slab pool | CPython's per-interpreter 2-tuple
freelist already serves this case, and we cannot recycle tuples that
become radix-tree keys without changing the consumer.
Rejected: int64-packed keys this round | requires changes to RadixKey,
get_child_key_fn, key_match_fn, and EAGLE bigram detection; deserves its
own plan.
Confidence: high
Scope-risk: low
Directive: Keep _python_convert_to_bigram_key reachable; if the tai-kernel
path is ever removed, the EAGLE radix cache must continue to work
unchanged.
Tested: tai-kernel side validated on the cluster
(``python benchmark/radix/benchmark_convert_to_bigram_key.py --check``
prints byte-exact correctness then 1.4-2.5x speedup across sizes
128..131072 on g0034). The new
``test/registered/unit/mem_cache/test_convert_to_bigram_key.py``
exercises both dispatch paths via SGLANG_DISABLE_TAI_BIGRAM patching.
Not-tested: End-to-end EAGLE serving accuracy + scheduler-time delta on
a real workload.
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
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.
With CP=8 and dp_size=1, enable_dp_attention gets reset to False, so
VocabParallelEmbedding uses tensor_model_parallel_all_reduce (tp_size=8).
The CP local draft path was splitting tokens before embedding, giving each
rank a different local_tokens count. This caused an NCCL all_reduce shape
mismatch and a collective hang.
Move embed_tokens() before the CP split: embed on the full input (all ranks
see the same shape), then cp_split_and_rebuild_data the result. The decoder
layer still runs on CP-local tokens, preserving the CP performance benefit.
Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
nsa_indexer.py imports log_cp_draft_shared_kv_debug from
cp_shared_kv_runtime, but ad358d164 defined it in nsa/utils.py.
The missing symbol caused deepseek_v2.py import to fail silently
(registry.py catches and warns), falling back to transformers.py
which loads unsharded 256-expert MoE parameters → OOM on H200.
Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
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.
Phase8 prefetch was starting local materialization and the CP all-reduce from the same late hook, after rank-skewed MQA/topk/materialize work had already separated CP ranks. This commit separates local prefix materialization from reduce launch, starts next-layer prefetch earlier in MLA prepare, and explicitly launches pending MLA/index reduces at attention boundaries so the collectives can overlap the intended window instead of drifting behind skewed ranks.
Constraint: Prefetch remains optional and gated by the existing CP shared-KV prefetch environment controls.
Constraint: Profiling required per-source materialize all-reduce ranges, so NVTX labels are gated behind SGLANG_CP_SHARED_KV_MATERIALIZE_NVTX.
Rejected: Launch all-reduce immediately in start_next_layer_prefix | this reproduced the old skew-sensitive timing and could put prefetch behind current-layer work.
Rejected: Remove later hooks entirely | they are still needed as fallbacks when the early MLA prepare hook is bypassed.
Confidence: medium
Scope-risk: moderate
Directive: Preserve the consume-time fallback launch; otherwise missed launch paths silently lose correctness or overlap.
Tested: Remote g0034 docker py_compile for changed SGLang files; python -m pytest -q test/registered/unit/mem_cache/test_cp_shared_kv_runtime.py -q (50 passed).
Not-tested: Full GLM5 multi-node prefill/decode profile after this exact commit.
The existing CP MQA prepare integration only used the TAI GetK/GetS+range kernel when current_index_kv was absent. Current-index reuse and full-kernel fallback still launched torch arange and zeros_like before fp8_mqa_logits. Route both cases through the new tai-kernel range-only API while preserving the torch fallback when the kernel is disabled or unavailable.
Constraint: SGLANG_CP_SHARED_KV_FUSED_INDEX_MQA_PREPARE remains the gate for this optimization.
Rejected: Force the full prepare kernel in current_index_kv path | it would redo index KV gathering that current reuse intentionally avoids.
Confidence: high
Scope-risk: narrow
Directive: Do not remove the torch fallback; mixed deployments may run without the updated tai-kernel package.
Related: tai-kernel 34cb7a8
Tested: Remote container g0034 docker py_compile for cp_shared_kv_runtime.py and nsa_indexer.py; tai-kernel range unit tests passed remotely (4 passed).
Not-tested: Full GLM5 prefill/decode server profile after this exact commit.
Shared-KV prefill already persists each rank's MLA shard and reconstructs the dense attention KV from the shared pool before attention. Keeping the legacy CP rebuild all-gather after direct write duplicated communication on the hot MQA-to-attention path. The rebuild remains enabled only for current-only reuse, where the backend intentionally consumes full current KV tensors instead of materializing history from the pool. Index and MLA prefetch now share one FIFO CUDA stream so their CP collectives preserve local launch order.
Constraint: CP shared-KV materialize is the authoritative KV source for prefix/cache-hit MLA attention.
Rejected: Gate prefetch by prefix owner-lane coverage | owner skew is not the root cause and would add an extra collective plus CPU sync.
Confidence: medium
Scope-risk: moderate
Directive: Do not reintroduce rebuild_cp_kv_cache for shared-KV direct-write unless the backend consumes k_nope/k_pe directly.
Tested: git diff --check; py_compile for cp_shared_kv_prefetch.py, nsa_backend.py, forward_mla.py; remote container py_compile after scp sync.
Not-tested: Full multi-node GLM5 performance run after this exact commit.
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
In disaggregated prefill with --enable-single-batch-overlap,
batch.copy() was used to defer process_batch_result, but the copy
omitted spec_info. This caused req.hidden_states_tensor to be None
when using Eagle/spec-v2 speculative decoding on the decode server,
silently degrading draft quality for the first decode round.
The CpSharedKVLayout initialization was gated behind
`if not self.is_draft_worker`, leaving the draft model runner
with `uses_cp_shared_kv=True` but `cp_shared_kv_layout=None`.
This caused an AssertionError in NSAIndexer._filter_shared_index_write
during DRAFT_EXTEND under CP shared KV + MTP.
Move the layout init after the `is_draft_worker` guard so both
target and draft model runners have it available.
Add logger.info statements with [CacheCtrl-write], [MemCache-evict], [MemCache-alloc], [HiCache-write], [HiCache-evict], [HiCache-load] tags across cache_controller.py, common.py, and hiradix_cache.py to improve observability of cache operations.
Add validate_page_aligned_token_indices utility and apply it across
HiCache write/load paths, NSA indexer transfers, and CUDA direct copy
kernels to reject malformed (partial, misaligned, non-contiguous) page
groups before they reach native transfer code. Also validate supported
CP HiCache backend/layout combinations at server startup.
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.
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
