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
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 CP KV materialization remained on the critical path for cached
NSA/MLA prefill batches. This change introduces a one-layer-ahead
prefetcher that materializes the cached prefix for the next layer on a
separate CUDA stream and consumes it when that layer reaches attention.
The prefetch path keeps the existing dense page-table semantics, defers
waiting until the prefetched buffer is actually consumed, and uses the
TAI optimized materialize/remap helpers when enabled before falling back
to the torch implementation.
The implementation is intentionally gated by environment variables and
keeps layer-2-only probe logging for functional confirmation without
making normal profiling noisy.
Constraint: Prefill CP shared KV must preserve existing page-table and dense KV semantics for NSA paged topk attention
Constraint: The production performance path requires SGLANG_CP_SHARED_KV_USE_TAI_MATERIALIZE=1 and logging disabled
Rejected: Wait immediately after the producer layer attention | this truncated the overlap window and hid less work
Rejected: Torch-only prefetch materialize | it bypassed the optimized TAI materialize/remap path and could erase the expected win
Confidence: medium
Scope-risk: moderate
Directive: Do not evaluate Phase8 throughput with SGLANG_CP_SHARED_KV_LOG_MLA_PREFETCH=1; use it only to confirm create/start/consume_hit behavior
Tested: Local AST parse for modified Python files
Tested: Local git diff --check
Tested: Remote g0034 container AST parse for modified files under /sgl-workspace/sglang-tai
Tested: Remote g0034 container pytest target covering Phase8 log env, TAI range materialize, optimized slot inverse/remap, and existing token TAI path
Not-tested: Full prefill/decode/router throughput after the TAI prefetch-path fix
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
The shared-KV materialize path was spending time in Python-observed CUDA tensor predicates and dynamic-shape remap helpers. Keep the runtime changes that move the hot paged path to slot-based device remapping, while dropping the NVTX experiment from this commit so profiling annotations do not become part of the runtime surface yet.\n\nThe MLA read path now passes the real page table as the page remap source, which keeps paged topk indices tied to the same logical page-table domain used to build the dense materialized KV view.\n\nConstraint: CP shared KV still needs a dense per-call view before deeper Phase4/Phase5 layout changes remove the materialize cost.\nRejected: Keep NVTX ranges in this commit | user requested reverting NVTX instrumentation before commit\nRejected: Restore compact unique-page remap everywhere | it reintroduces CUDA sync-prone dynamic-shape ops on the hot paged materialize path\nConfidence: medium\nScope-risk: moderate\nDirective: Benchmark slot-remap buffer size against compact unique-page remap before treating this as the final performance path; Phase4/5 should reduce materialize instead of relying on this aggregation path.\nTested: git diff --check on changed files; python -m py_compile on changed runtime/backend/test files; grep confirmed NVTX symbols removed\nNot-tested: pytest blocked locally by missing pybase64 dependency; multi-node PD runtime not rerun
Phase 2 shared KV keeps persistent KV/index sharded across CP ranks but
uses a full-view compatibility layer before NSA topk and MLA attention.
For current-only prefill chunks, the current KV/index tensors have already
been CP all-gathered and reranged before being written to the sharded
persistent pool. This change adds a guarded current-reuse path that remaps
logical current locs to compact tensor rows and skips the shared-KV
materialize path for current-only MLA and NSA indexer reads.
Constraint: Existing NSA/MLA kernels still consume full-view/compact page tables; history and mixed current/history batches must keep the Phase 2 fallback.
Rejected: Make all history attention shard-aware in this patch | that requires global topk merge and distributed sparse attention and belongs to a later phase.
Confidence: medium
Scope-risk: moderate
Reversibility: clean
Directive: Do not remove the Phase 2 fallback until mixed/history shared-KV paths have correctness and performance coverage.
Tested: python -m py_compile on modified Python files
Tested: git diff --check on staged modified files
Not-tested: local pytest collection is blocked by missing pybase64 in this environment.
Not-tested: full long-context chunked prefill/decode performance in 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
Replace full hidden all-gather at prefill tail with per-request last-token
hidden collection, reducing communication from total_tokens x hidden_size
to cp_size x bs x hidden_size for both in-seq-split and round-robin modes.
- nsa/utils.py: add cp_collect_last_token_hidden() with mode-specific
narrow collection helpers that only gather the last token hidden
- deepseek_v2.py: add _should_use_narrow_output_path() gate on
DeepseekV2Model, fallback to full gather for EAGLE/return_logprob/
capture_hidden batches
- logits_processor.py: add _is_compact_hidden_states() to bypass
_get_pruned_states() when hidden is already compact
