CP shared-KV bs>1 cache-hit loads already merge request load ops, but the host pool still rebuilt layer-invariant mapping work from the same host/device indices. Introduce a PreparedLoadDescriptor lifecycle around begin/end load, wire MLA KV and NSA index H2D loads through tai-kernel prepared submit when available, and add timing hooks plus regression coverage for descriptor reuse and explicit fallback logging. Record the P4/P6b design and benchmark results in the advanced feature notes.
Constraint: Radix residency and allocator decisions remain synchronous; only the data-transfer descriptor is prepared for per-layer async submit.
Constraint: Production fast path must not silently fall back when tai prepared H2D support is missing.
Rejected: Cross-batch descriptor reuse | descriptor lifetime and tensor ownership are only safe within one load operation.
Rejected: Change L2->L1 scheduling to layer-ahead prefetch in this commit | that is a separate lifecycle change after descriptor reuse is stable.
Confidence: medium
Scope-risk: moderate
Directive: Keep LayerDoneCounter per-layer readiness semantics; do not replace with all-layer waits.
Tested: python -m py_compile python/sglang/srt/mem_cache/memory_pool_host.py python/sglang/srt/managers/cache_controller.py
Tested: Remote g0034:cjy-glm5-new PYTHONPATH=python python -m pytest -q test/registered/unit/managers/test_hicache_controller_cp.py (88 passed)
Tested: Remote tai-kernel prepared descriptor CUDA test (6 passed) and P4 benchmark full matrix (90 rows)
Not-tested: ETE replay/GSM8K cache-hit correctness after this commit
Not-tested: Layer-ahead L2->L1 prefetch scheduling
Co-authored-by: OmX <omx@oh-my-codex.dev>
When chunked prefill is active, CP shared-KV bs>1 cannot consume more extend
tokens than the current chunk budget. If the CP-specific extend-token limit is
omitted, default it to rem_chunk_tokens so scheduler admission reflects the
reachable chunk capacity. The request-count and cached-token knobs keep their
None-as-unlimited behavior.
Constraint: CP bs>1 batching must not advertise a larger extend batch than chunked prefill can execute.
Rejected: Require users to always set --cp-shared-kv-prefill-max-total-extend-tokens | the safe default is already available from chunked prefill state.
Rejected: Default batch request or cached-token limits | those are policy knobs and None should remain unlimited.
Confidence: high
Scope-risk: narrow
Directive: Keep --cp-shared-kv-prefill-max-total-extend-tokens as min(user_limit, chunk_budget) when both exist.
Tested: Local py_compile for schedule_policy.py and test_prefill_adder.py.
Tested: Remote g0034 cjy-glm5-new targeted prefill_adder tests: 2 passed.
Not-tested: Full ETE scheduler batching distribution after defaulting the extend limit.
Co-authored-by: OmX <omx@oh-my-codex.dev>
CP shared-KV batching previously estimated MQA logits from full request
extend/context rows, which overstated memory because CP in-seq split only
computes each rank's two zigzag segments. Add CP-size aware row accounting
that mirrors the fused CP MQA materialization path and take the worst local
rank peak for scheduler admission.
Expose SGLANG_NSA_MQA_LOGITS_CHUNK_MAX_GB as a more direct cap for one fp32
MQA logits chunk. Runtime and scheduler now both translate this GB cap into
chunk rows from the actual K rows, while keeping the old row cap as a
mutually-exclusive expert override.
Constraint: Scheduler admission must stay CUDA-sync-free and use static budget information only.
Rejected: Keep full-request q*k admission | it over-gates CP bs>1 batches because CP splits q rows per rank.
Rejected: Let rows and GB caps both apply | precedence would be ambiguous during tuning.
Confidence: medium
Scope-risk: moderate
Directive: Keep MQA logits admission tied to the fused CP MQA segment shape; do not revert to full request token counts.
Tested: Local py_compile for touched runtime, scheduler, estimator, and tests.
Tested: Local pytest test_cp_shared_kv_prefill_buffer_estimator.py: 8 passed.
Tested: Remote g0034 cjy-glm5-new py_compile and targeted estimator/runtime tests: 13 passed.
Not-tested: Full ETE high-cache-hit CP bs>1 load with SGLANG_NSA_MQA_LOGITS_CHUNK_MAX_GB.
Co-authored-by: OmX <omx@oh-my-codex.dev>
CP shared-KV bs>1 admission already bounds request count, extend tokens,
cached tokens, and an estimated temporary buffer size. The estimate missed
the fp32 MQA logits temporary, whose peak grows with query rows times
context rows and can dominate high-cache-hit multi-request batches.
Add an MQA logits peak term to the CPU-only estimator and include it in
the layer-forward peak enforced by --cp-shared-kv-prefill-max-buffer-size.
When SGLANG_NSA_MQA_LOGITS_CHUNK_MAX_ROWS is set, admission estimates the
post-chunk peak using that row cap; otherwise it remains conservative and
assumes the full extend-row count.
Constraint: Scheduler admission must stay CPU-only and cannot query CUDA free memory.
Rejected: Add a separate scheduler limit for MQA logits | the existing max-buffer-size knob is the right aggregate admission budget.
Rejected: Use SGLANG_NSA_MQA_LOGITS_FREE_MEM_FRACTION in scheduler | that depends on runtime CUDA free memory and would make admission host-sync or stale.
Confidence: medium
Scope-risk: moderate
Directive: Keep the estimator conservative when chunk max rows is unset; do not rely on CUDA free-memory queries in scheduler admission.
Tested: Local py_compile for estimator, scheduler, schedule_policy, and estimator tests.
Tested: Local pytest test_cp_shared_kv_prefill_buffer_estimator.py: 5 passed.
Tested: Remote g0034 cjy-glm5-new py_compile and estimator pytest: 5 passed.
Not-tested: ETE scheduler admission under high-cache-hit bs>1 traffic.
Co-authored-by: OmX <omx@oh-my-codex.dev>
CP shared-KV bs>1 batching was only bounded by request count, extend tokens, and cached tokens. That left temporary GPU buffers such as MLA/index materialization, remap metadata, logits windows, and transfer descriptors implicit, and raw extend-token limits could exceed the active chunked-prefill budget.\n\nThis adds an explicit max-buffer-size admission gate with a CPU-only stream-aware estimator, wires it through PrefillAdder/Scheduler, performs a startup CUDA smoke allocation when configured, and reports the estimate in the scheduler admission benchmark. When chunked prefill is active, the effective CP extend-token limit is capped by the current chunk budget so the CP path does not advertise unreachable batch capacity or lift max-prefill-tokens too far.\n\nConstraint: Admission estimation must stay CPU-only on the scheduler hot path; CUDA allocation is limited to startup smoke checking.\nConstraint: Single oversized requests must still be allowed to run alone to avoid scheduler deadlock.\nRejected: Rely only on --max-prefill-tokens | it does not reliably bound the first oversized request and does not model cache-hit/load-back pressure.\nRejected: Let CP extend limit exceed chunked-prefill size | it creates an unreachable effective capacity and misleading budget lift.\nConfidence: medium\nScope-risk: moderate\nDirective: If bs>1 L1 prefetch is enabled later, update CPSharedKVPrefillBufferEstimatorContext.bs_gt1_l1_prefetch_enabled and include the live prefetch dense buffers in overlap windows.\nTested: local py_compile for touched files\nTested: local PYTHONPATH=python pytest -q test/registered/unit/managers/test_cp_shared_kv_prefill_buffer_estimator.py (4 passed)\nTested: remote cjy-glm5-new targeted pytest for new server_args, PrefillAdder, estimator, and benchmark cases (10 passed)\nTested: remote cjy-glm5-new PYTHONPATH=python pytest -q test/registered/unit/managers/test_cp_shared_kv_prefill_buffer_estimator.py test/registered/unit/managers/test_prefill_adder.py test/registered/unit/managers/test_prefill_scheduler_admission_bench.py (29 passed before chunk cap, then test_prefill_adder.py 21 passed after chunk cap)\nNot-tested: full server_args suite because existing TestPrepareServerArgs tries to reach HuggingFace and fails under container DNS/network\nNot-tested: GLM5 ETE smoke with --cp-shared-kv-prefill-max-buffer-size
Add an offline benchmark that reuses PrefillAdder to model how L1 cached tokens, L2 HiCache hits, and extend tokens shape CP shared-KV prefill batch admission. The tool makes scheduler stop reasons and fake L2 load-back capacity pressure observable without starting a model.
Constraint: The benchmark must stay CPU/offline and avoid depending on CUDA execution or live services.
Constraint: L2 cached tokens are modeled as host_hit_length, so successful load-back both increases prefix_len and consumes fake L1 capacity.
Rejected: Build an ETE benchmark first | too slow for isolating scheduler admission behavior.
Rejected: Reimplement scheduler logic from scratch | would drift from PrefillAdder semantics.
Confidence: high
Scope-risk: narrow
Directive: Treat duration_us as Python admission overhead only; it is not an ETE latency metric.
Tested: Remote pytest test/registered/unit/managers/test_prefill_scheduler_admission_bench.py: 4 passed as part of 6 targeted tests.
Tested: Remote synthetic benchmark run with --cp-max-total-cached-tokens showed second 4096-token cached request stopped with OTHER.
Not-tested: Real traffic trace import from production logs.
High cache-hit CP shared-KV batches can have small extend tokens while still carrying substantial prefix/load-back work. Add a CP-specific cached-token admission limit so operators can bound that pressure independently from extend-token batching.
Constraint: The limit must not deadlock a single high-cache-hit request; it only stops adding additional requests to a non-empty batch.
Constraint: Cached tokens are counted after L2 load-back planning via prefix_len, so L1 hits and successful L2 hits share one scheduler budget.
Rejected: Reuse max_prefill_tokens | it limits generic input budget and does not represent cached-token work.
Rejected: Count only L1 prefix before load-back | would miss L2 hit pressure, which is one of the target cases.
Confidence: high
Scope-risk: moderate
Directive: Keep cached-token and extend-token limits separate; they bound different scheduler costs.
Tested: Remote pytest targeted cached-token PrefillAdder cases: 2 passed.
Tested: Remote pytest test/registered/unit/managers/test_prefill_adder.py: 18 passed.
Tested: Remote ServerArgs CP validation smoke: SERVER_ARGS_CACHED_LIMIT_OK.
Not-tested: Full ETE replay with a production cached-token limit value.
CP shared-KV bs>1 uses cp_shared_kv_prefill_max_total_extend_tokens as its grouping admission limit, but the generic max_prefill_tokens budget could still stop batching earlier. Raise only the legacy input-token admission budget for this CP path while keeping allocator-owned capacity checks unchanged.
Constraint: CP shared-KV bs>1 needs large cache-hit batches without relying on the generic max_prefill_tokens default.
Constraint: Allocator capacity must remain enforced by rem_total_tokens, cur_rem_tokens, and prepare_for_extend().
Rejected: Increase server-wide max_prefill_tokens | would change generic scheduler behavior and non-CP paths.
Confidence: high
Scope-risk: narrow
Directive: Do not use max_prefill_tokens as the CP shared-KV bs>1 grouping limit; use the CP-specific total extend token knob.
Tested: Local py_compile for schedule_policy.py and test_prefill_adder.py.
Tested: Remote pytest targeted CP PrefillAdder cases: 5 passed.
Tested: Remote pytest test/registered/unit/managers/test_prefill_adder.py: 16 passed.
Not-tested: Full ETE replay after this scheduler-only change.
CP shared-KV HiCache transfers are per-layer, so the host layout should match the access pattern instead of forcing page-major strides through every layer. This adds a direct-only layer_page_first layout, routes per-layer KV and NSA index backup/load through the TAI LF<->LPF direct kernels, and keeps storage/page-buffer metadata paths fail-fast until their page-level contract is redesigned.\n\nThe direct controller keeps host indices in caller order for both page_first_direct and layer_page_first because the TAI direct path requires CPU index descriptors and owns descriptor coalescing. All-layer backup intentionally loops over per-layer direct kernels rather than using the sgl-kernel all-layer direct ABI.\n\nConstraint: layer_page_first is currently host-only CP HiCache; storage backends assume page-major contiguous page metadata.\nConstraint: TAI LPF direct kernels require CPU int64 page indices and complete page spans.\nRejected: silently fallback to SM copy when TAI LPF kernels are missing | that hides production performance regressions.\nRejected: support storage page metadata in this commit | LPF requires a layer-page-level storage contract, not a one-pointer-per-page contract.\nConfidence: medium\nScope-risk: moderate\nDirective: Do not enable storage or kernel backend for layer_page_first without redesigning page-buffer metadata and adding remote ETE coverage.\nTested: local py_compile for touched runtime files.\nTested: remote py_compile in g0034 container for touched runtime files.\nTested: remote targeted pytest: 5 passed for parser/storage/layout/move_indices smoke coverage.\nNot-tested: full CP HiCache ETE with --hicache-mem-layout layer_page_first after this commit step.\nNot-tested: combined CUDA roundtrip tests in one pytest process; previous independent runs passed but combined run exposed a host-memory registration lifecycle issue.
Burst decode can sit in disaggregation prealloc/transfer queues while health probes still need an immediate scheduler-alive response, and EAGLE prebuilt metadata must outlive the synthetic prebuilt step until the first real decode consume. The transfer path also needs to send final aux metadata even when there are no new KV pages in the final chunk.
This keeps decode metadata slot ownership narrow instead of cloning EAGLE tensors, adds fail-safe cleanup for prebuilt exceptions/finished prebuilt requests, fixes final empty chunk metadata transfer, and records the investigation ledger for future debugging.
Constraint: bs=1 historically worked, so decode compute/sampling behavior must not be changed without direct evidence.
Rejected: Clone EAGLE metadata tensors on transfer commit | avoids lifetime issues but adds hot-path CPU/GPU memory traffic.
Rejected: Treat prefill AbortReq logs as root cause | current evidence shows they can be downstream of decode/router aborts.
Confidence: medium
Scope-risk: moderate
Directive: Do not move EAGLE metadata slot release earlier than first real decode result processing without proving the H2D/spec_info consume point changed.
Tested: g0034 docker py_compile for touched scheduler/disagg/test files; PYTHONPATH=python python -m pytest -q test/registered/unit/disaggregation/test_decode_queue_compaction.py test/registered/unit/mem_cache/test_req_to_token_pool.py test/registered/unit/managers/test_scheduler_health_check.py -> 24 passed
Not-tested: Fresh end-to-end burst traffic after restart; decode node_rank=1 persistent log capture.
The scheduler can now admit multi-request NSA in-seq CP shared-KV prefill batches only when the shared-KV bs>1 flag is explicitly enabled. The gate is still disabled by default and is scoped to CP shared-KV so ordinary CP is not widened accidentally.
Batch admission is bounded by optional request-count and page-aligned extend-token limits while real memory capacity remains allocator-owned. This keeps bf16 and fp8 on the same scheduler path because dtype differences are already reflected in KV pool token/page capacity.
Constraint: bs>1 runtime paths remain guarded by existing CP shared-KV fail-fast checks.
Constraint: Scheduler must not duplicate bf16/fp8 byte-level capacity estimation.
Rejected: Open the old CP gate unconditionally | ordinary CP would inherit an unverified shared-KV-specific batching path.
Rejected: Treat the extend-token cap as a hard per-request limit | a single large request could deadlock the scheduler.
Confidence: medium
Scope-risk: moderate
Directive: Keep CP shared-KV batching gated until ETE validates EAGLE accept length, output length, and HiCache load/backup behavior under real traffic.
Tested: local py_compile for server_args, schedule_policy, scheduler, prefill_adder tests, and server_args tests.
Tested: remote g0034 py_compile for the same files.
Tested: remote g0034 pytest target set: 5 passed for parser, parameter validation, default single-request CP gate, enabled bs>1 gate, and page-aligned extend cap.
Tested: remote g0034 pytest test_prefill_adder.py => 13 passed.
Not-tested: full server_args test file has an unrelated HuggingFace DNS/config-download failure in TestPrepareServerArgs.test_prepare_server_args.
Not-tested: ETE production traffic with --enable-cp-shared-kv-prefill-bs-gt1.
CP HiCache write reservations must stay per radix node, but the transfer descriptor does not need to be per request. This changes the layer-end hook to group pending write states for the same source and layer, so bs>1 prefill emits one target D2H descriptor and one draft D2H descriptor per layer while preserving per-node metadata, rollback, and ack semantics.\n\nConstraint: CP shared-KV HiCache metadata, host slots, and radix acknowledgements remain per request/node.\nConstraint: TAI direct transfer kernels already accept flattened page descriptors, so no tai-kernel change is required.\nRejected: Merge HiCache reservations or radix nodes | would complicate rollback and split handling.\nRejected: Add collective synchronization for grouped backup | grouping is local descriptor construction and must not add rank-level sync.\nConfidence: high\nScope-risk: moderate\nDirective: Keep target and draft source notifications separate; final ack must wait for both when draft HiCache is attached.\nTested: local py_compile for cache_controller.py and test_hicache_controller_cp.py\nTested: local git diff --check\nTested: remote pytest test/registered/unit/managers/test_hicache_controller_cp.py test/registered/unit/mem_cache/test_cp_hicache_load_back_owner_lanes.py => 85 passed, 3 warnings\nNot-tested: full ETE bs>1 CP HiCache replay with admission gate removed\nNot-tested: Nsight/throughput validation of reduced D2H submit count
Host HiCache reservations were paying token-level free-slot scans when trying to preserve page contiguity. The allocator now keeps a lazy page-extent index so availability checks and contiguous-preferred allocations avoid materializing the full 220GB-equivalent free-slot metadata path.
The companion benchmark models steady-state L2 churn near full occupancy, including burn-in and historical node-size effects, so LPF/RDMA descriptor quality can be separated from ETE noise.
Constraint: CP HiCache host allocations are page-shaped, but existing callers may still read free_slots directly.
Rejected: Sort and scan free_slots on each alloc_contiguous_preferred call | measured ms-level CPU overhead on 220GB-equivalent metadata.
Rejected: Remove free_slots compatibility | storage/tests still rely on the public tensor surface.
Confidence: medium
Scope-risk: moderate
Directive: Do not reintroduce per-allocation full free_slots scans on HostKVCache; preserve page-extent metadata or benchmark before changing allocator shape.
Tested: Local py_compile for memory_pool_host.py, allocator benchmark, and related tests.
Tested: Local test_cp_hicache_allocator_bench.py 10 passed.
Tested: Remote g0034 test_hicache_controller_cp.py 67 passed; test_cp_hicache_allocator_bench.py 10 passed.
Tested: Remote 220GB-equivalent host_churn benchmark: contiguous path reduced from ms-level to ~30-292us p50 depending on fragmentation.
Not-tested: Full CUDA ETE run after allocator change.
Not-tested: Production long-run fragmentation behavior under live traffic.
CP shared KV now keeps explicit L1 and host free-room targets so pressure is handled by planned eviction instead of repeated capacity-edge retries. The host allocator gains contiguous-preferred page reservation, L1 owner-lane allocation prefers contiguous physical pages, and CP HiCache metadata preserves pending backup safety for page-granular radix updates. Mooncake transfer stats and allocator microbenchmarks are included to make the remaining transfer bottlenecks measurable rather than inferred.
Constraint: CP shared KV uses decode CP size 1 with all prefill CP ranks participating in transfer, so L1/L2 cache residency must remain page-granular and avoid extra collectives.\nConstraint: Production HiCache can be hundreds of GB, so allocator metadata overhead must be visible before enabling aggressive contiguous allocation broadly.\nRejected: Evict only the exact deficit | this keeps the cache at the cliff and causes repeated evict/allocate pressure.\nRejected: Rely on allocator scans alone for contiguity | remote microbenchmarks show fragmented 220GB-equivalent host metadata can make contiguous-preferred scans multi-ms.\nConfidence: medium\nScope-risk: moderate\nDirective: Do not increase L1/L2 free-room defaults or add new CP collectives without ETE evidence and transfer/allocator measurements.\nTested: python -m py_compile on touched runtime/test/benchmark files.\nTested: PYTHONPATH=. python -m pytest -q test/registered/unit/benchmark/test_cp_hicache_allocator_bench.py => 4 passed, 1 warning.\nTested: Remote g0034 log /mnt/beegfs/cjy/log/sglang_cp_hicache_20260601_233723.log shows active prefill process with L1/L2 free-room args, 702 HTTP 200 chat completions, 6272 prefill batches, and no fatal scheduler traceback in latest scan.\nTested: User-reported L1/L2 cache ETE validation passed on remote run.\nNot-tested: Full local pytest suite; local environment is missing several runtime dependencies.\nNot-tested: CUDA allocator microbenchmark during active production prefill process.\nNot-tested: Mooncake straggler fix; stats show transfer tail latency remains a separate bottleneck.
CP HiCache direct/page_first_direct all-layer backup was still able to enter sgl-kernel's stale cudaMemcpyBatchAsync path, which segfaults under CUDA 13 before Python can surface an error. The SGLang route now avoids that all-layer sgl-kernel op for page_first_direct backup and uses the TAI per-layer direct LF->PF op for MHA, MLA, and NSA indexer data.\n\nThe load-back path also prepares NSA indexer page indices once per load op and reuses them across per-layer H2D loads, preserving per-layer overlap while removing redundant page-index derivation.\n\nConstraint: Remote runtime is CUDA 13.0 where sgl-kernel's all-layer direct LF->PF op uses the wrong cudaMemcpyBatchAsync ABI.\nRejected: Patch sgl-kernel in this branch | we are converging production HiCache direct/page_first_direct paths onto tai-kernel and do not want to maintain another CUDA-ABI-sensitive copy path here.\nRejected: Collapse H2D load-back into one all-layer op | that would reduce submit count but lose per-layer completion visibility and forward overlap.\nConfidence: medium\nScope-risk: moderate\nDirective: Do not reintroduce sgl_kernel.transfer_kv_all_layer_direct_lf_pf for direct/page_first_direct HiCache backup without CUDA 13 ABI verification.\nTested: g0034 container: PYTHONPATH=python python -m pytest -q -s test/registered/unit/mem_cache/test_nsa_pool_host_unit.py -> 10 passed, 3 warnings.\nTested: g0034 container: PYTHONPATH=python python -m pytest -q test/registered/unit/managers/test_hicache_controller_cp.py -> 61 passed, 3 warnings.\nTested: python -m py_compile python/sglang/srt/mem_cache/memory_pool_host.py python/sglang/srt/managers/cache_controller.py test/registered/unit/mem_cache/test_nsa_pool_host_unit.py test/registered/unit/managers/test_hicache_controller_cp.py\nNot-tested: Full ETE prefill/decode traffic after this commit.\nNot-tested: sgl-kernel implementation itself remains unchanged.
CP shared KV and HiCache now keep page-aligned physical ownership while preserving valid-token radix semantics. Repeated tiny EAGLE exact hits free duplicate tail pages instead of leaking one allocator page, owner-lane load-back uses page-vector admission/eviction, and single-DP idle schedulers avoid entering an unnecessary MLP-sync collective.
The commit also records the current page-aligned cache contract and adds gated decode-side EAGLE accept diagnostics so future accept-length collapses can be tied to draft KV/state transfer evidence instead of more prefill cache speculation.
Constraint: CP HiCache allocator ownership is page-granular while radix matching remains valid-token based.
Constraint: New diagnostics must be gated and must not alter normal EAGLE, transfer, or cache behavior.
Rejected: Padding short requests to cp_size or 2*cp_size pages | wastes KV capacity and still hides valid-tail lifecycle bugs.
Rejected: Adding more unconditional collectives to prove CP consistency | hot-path collectives previously caused severe performance risk.
Confidence: medium
Scope-risk: broad
Directive: Do not reintroduce silent fallback for CP shared KV/HiCache paths; warning-level fallback or fail-fast is intentional.
Tested: git diff --check
Tested: local py_compile for all modified Python files
Tested: remote g0034 container py_compile for modified Python/test files
Tested: remote g0034 container PYTHONPATH=python python -m pytest -q test/registered/unit/layers/test_nsa_cp_utils.py test/registered/unit/mem_cache/test_cp_shared_kv_runtime.py test/registered/unit/mem_cache/test_cp_hicache_load_back_owner_lanes.py test/registered/unit/managers/test_scheduler_dp_attn_mixin.py => 114 passed, 5 warnings, 2 subtests passed
Not-tested: full ETE traffic rerun after this commit
Not-tested: CUDA/TAI kernel benchmark coverage for all production shapes
Draft HiCache already piggybacks target CP reservations, but the controller coverage only exercised page-aligned nodes. The new tests pin the intended contract for valid-tail nodes: target and draft reserve/load the same padded physical page while scheduler-visible load results keep only valid logical locs.
Constraint: Draft KV is a target mirror; it must not choose an independent valid-only host path.
Rejected: Re-enable draft partial-current reuse in this slice | current safe contract keeps EAGLE/NextN cache-hit draft on full materialization until same-layer padded visibility is proven.
Confidence: medium
Scope-risk: narrow
Directive: Keep draft and target host metadata coupled; do not add draft-only capacity or prefetch decisions.
Tested: local py_compile for test_hicache_controller_cp.py.
Tested: remote g0034 test_hicache_controller_cp.py: 59 passed, 3 warnings.
Not-tested: CUDA E2E runtime for this commit.
Co-authored-by: OmX <omx@oh-my-codex.dev>
CP HiCache now keeps radix and scheduler-visible lengths as valid tokens while host/device transfers reserve and replay the padded physical page span. Exact valid-tail write, insertion, and match paths no longer fall back to page-flooring; the physical owner-lane contract still uses padded page metadata.
Constraint: Scheduler prefix indices must never include padded tail locs.
Constraint: Host/device transfer and owner-lane admission remain page-based.
Rejected: Pad to cp_size or 2*cp_size pages | wastes KV and recreates short-tail fallback behavior.
Rejected: Expose padded locs through load_cp return | would leak fake tokens into req.prefix_indices.
Confidence: medium
Scope-risk: moderate
Directive: Do not implement split-inside-tail by duplicating page_owners without a page-sharing/refcount design.
Tested: local py_compile for touched CP HiCache/radix/controller files and tests.
Tested: remote g0034 CP HiCache impacted suites: 143 passed, 5 warnings.
Tested: remote g0034 CP shared KV C1-C5 suite: 122 passed, 5 warnings.
Not-tested: full local pytest, blocked by missing runtime dependencies such as orjson/starlette.
Not-tested: CUDA E2E runtime for this commit.
Co-authored-by: OmX <omx@oh-my-codex.dev>
Use tai-kernel for direct page_first_direct per-layer H2D load across MHA, MLA, and NSA indexer pools. This keeps SGLang off the sgl-kernel cudaMemcpyBatchAsync path that crashes on CUDA 13 while preserving fail-fast behavior when the required TAI op is unavailable.
Constraint: remote CUDA 13 stack crashes in sgl-kernel PF->LF direct load via cuMemcpyBatchAsync_v2
Rejected: Silent fallback to sgl-kernel or Python loop | fallbacks would hide either a crash-prone ABI path or a large performance regression
Confidence: high
Scope-risk: moderate
Directive: page_first_direct direct load must remain fail-fast if tai_kernel.nsa_prefill.transfer_kv_per_layer_direct_pf_lf is missing
Tested: remote g0034 PYTHONPATH=python pytest -q test/registered/unit/managers/test_hicache_controller_cp.py: 55 passed, 3 warnings
Tested: remote g0034 CUDA smoke for MLATokenToKVPoolHost.load_to_device_per_layer with direct/page_first_direct passed
Not-tested: full SGLang ETE server after the final commit
CP shared KV with HiCache and EAGLE needs host backup to overlap forward while keeping radix visibility synchronous. The change reserves host slots before forward, drives target and draft backup from explicit layer-end hooks, and commits host visibility only after the final target/draft ack. It also probes the final insertion prefix before early reservation so repeated EAGLE prompts do not prepare duplicate suffix backups that later rollback as insert_miss.
Constraint: CP ranks use independent shared-KV pools, so target/draft host state must remain atomically visible at the radix boundary.
Constraint: Fused MLA and NSA store paths can bypass store-side notifier hooks, so layer end is the safer backup progress boundary.
Rejected: Store-side backup notifier as the primary trigger | fused store and zero-local paths made notifier coverage fragile.
Rejected: Reserve from cache_protected_len alone | EAGLE bigram/page alignment can make final insertion find a longer existing prefix and force duplicate rollback work.
Confidence: medium
Scope-risk: moderate
Directive: Do not add per-layer CP collectives here; keep radix state synchronous and data transfer asynchronous/local-event driven.
Tested: local git diff --check
Tested: local py_compile for touched CP HiCache/cache-controller/deepseek/test files
Tested: remote pytest test/registered/unit/mem_cache/test_cp_hicache_metadata.py test/registered/unit/managers/test_hicache_controller_cp.py -q (115 passed, 5 warnings)
Not-tested: full GLM5 ETE server rerun after this commit
CP shared KV needs HiCache backup to overlap with layer execution without exposing partially copied host state. Split CP backup into reservation, pending radix state, per-layer target/draft D2H submission, and one final ack-driven visibility commit. The all-layer path remains available only as an explicit fallback and now logs a warning when used.
Constraint: CP shared KV owner-lane metadata and draft/MTP KV must stay strongly synchronized with target KV.
Constraint: Local CUDA tests are disallowed; CUDA verification was run only in the g0034 container.
Rejected: Let target layer hooks copy draft KV too | draft may not have stored that layer yet, which can corrupt MTP accept behavior.
Rejected: Silent all-layer fallback | it hides performance regressions and makes ETE logs ambiguous.
Confidence: medium
Scope-risk: broad
Directive: Reserved or partially copied host payloads must remain invisible until final ack commits pending_host_backups.
Tested: g0034 docker /sgl-workspace/sglang-tai PYTHONPATH=/mnt/beegfs/cjy/tai-kernel/python:python python -m pytest test/registered/unit/managers/test_hicache_controller_cp.py -q -> 49 passed.
Tested: g0034 docker /sgl-workspace/sglang-tai PYTHONPATH=/mnt/beegfs/cjy/tai-kernel/python:python python -m pytest test/registered/unit/mem_cache/test_cp_hicache_metadata.py -q -> 58 passed.
Tested: g0034 docker /mnt/beegfs/cjy/tai-kernel PYTHONPATH=python python -m pytest tests/nsa_prefill/test_kvcacheio_lf_pf.py -q -> 7 passed.
Not-tested: Long-running GLM5 CP+HiCache+MTP ETE throughput and host-pressure soak.
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>
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>
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>
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.
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.
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.
