Commit Graph

7298 Commits

Author SHA1 Message Date
laoyao0822
c2d25ff591 Prevent CP HiCache pending-backup splits from killing chunked prefill
Chunked prefill can revisit a sub-page CP HiCache tail while a per-layer backup is still in flight. The old insert path split the radix node first and only then tried to prune the stale tail, so an unprunable pending backup raised after tree mutation and propagated to the scheduler.

This makes split/prune atomic from the radix-tree perspective: drain completed write acks only on the conflict path, preflight pending/unprunable state before split, and return a deferred insert result when the backup is still in flight. Unfinished requests keep their KV ownership for transfer/release instead of rematching or freeing pages under a stale tree state.

Constraint: CP HiCache backup metadata is node/page owned and cannot be repartitioned while per-layer D2H is pending
Rejected: Split the pending backup node | would require repartitioning in-flight backup metadata and host reservations
Rejected: Delete the stale tail unconditionally | risks freeing device/host pages still owned by pending backup
Confidence: high
Scope-risk: moderate
Directive: Do not move stale-tail prune after split without a preflight; pending backup split conflicts must remain non-mutating
Tested: remote g0034 container py_compile for mem_cache files
Tested: remote g0034 PYTHONPATH=python pytest -q test/registered/unit/mem_cache/test_cp_hicache_metadata.py test/registered/unit/mem_cache/test_cp_hicache_load_back_owner_lanes.py => 115 passed
Tested: local git diff --check
Not-tested: full chunked prefill ETE after service restart
Co-authored-by: OmX <omx@oh-my-codex.dev>
2026-06-02 04:51:39 +08:00
laoyao0822
6ef4face89 Preserve FP8 CP shared-KV page contracts
NSA FP8 CP shared-KV reuse must operate on packed page-slot rows, not bf16 compact rows. The change keeps current-only and partial-current reuse inside the page-aligned materialization contract, fails fast for non-page-aligned CP split inputs, and prevents FP8 FlashMLA-KV prefill from reaching incompatible in-seq CP metadata.

Constraint: NSA FP8 persistent MLA KV rows are packed 656-byte records and CP shared KV cache management is page-granular.\nConstraint: FlashMLA-KV prefill metadata is not CP-local after NSA in-seq splitting.\nRejected: Silently splice bf16 current rows into FP8 materialized cache | corrupts the packed cache layout.\nRejected: Keep FP8 FlashMLA-KV auto prefill under NSA CP | reaches num_splits shape errors after q-row splitting.\nConfidence: medium\nScope-risk: moderate\nDirective: Do not re-enable FP8 FlashMLA-KV prefill for NSA in-seq CP until metadata is rebuilt after CP splitting or made CP-local.\nTested: Local git diff --check and py_compile for touched SGLang files.\nTested: Remote g0034 related unit sweep recorded in docs: test_nsa_cp_utils.py, test_cp_shared_kv_layout.py, test_cp_shared_kv_runtime.py, test_cp_hicache_metadata.py passed.\nNot-tested: Full FP8 ETE startup and performance run after this commit.
2026-06-01 03:33:44 +08:00
laoyao0822
46be97adc0 Align CP shared-KV prefetch with the attention overlap window
Launching CP shared-KV prefetch from MLA prepare or the indexer can make
next-layer prefix work overlap current-layer MQA/materialization instead of
the attention window. Centralize the launch in the NSA backend after
current-layer materialization and before attention, and leave the early
indexer hook inert so the call site cannot regress silently.

The accompanying notes capture the draft-as-forward-layer follow-up plan and
the latest OOM diagnosis: the observed 178945-token failure matches the CP
in-seq MQA logits allocation, so the follow-up fix is q-dimension chunking in
_get_topk_ragged_with_cp(), not a max-prefetch-size gate.

Constraint: CP shared-KV prefetch must avoid overlapping current-layer MQA/materialization and must not add silent fallback behavior.
Rejected: Limit maximum prefetch size | hides the CP logits peak and can reduce cache/prefetch effectiveness.
Rejected: Keep prepare/indexer launch sites | they place next-layer collectives in the wrong overlap window.
Confidence: medium
Scope-risk: moderate
Directive: Do not reintroduce early CP prefetch launch without checking Nsight overlap and CP MQA memory peaks.
Tested: Local git diff --check and py_compile for touched Python files.
Tested: Remote container py_compile plus targeted pytest: 3 passed, 5 warnings.
Not-tested: Full ETE under production traffic after this commit.
Not-tested: CP in-seq MQA logits chunking; documented as follow-up.
2026-06-01 01:09:43 +08:00
laoyao0822
4342de0463 Avoid redundant CP collectives on sync shared-KV materialize
Synchronous CP shared KV full-hit and partial-current paths can now use the tai-kernel CUDA IPC slot-dense materialize path instead of first copying local owner pages and then running a dense CP all-reduce. This keeps the async prefetch pipeline unchanged while routing the safer synchronous runtime path through descriptor-driven owner-page reads.

The runtime builds descriptors from the existing page-aligned slot contract, caches peer pointer tables for long-lived KV/index buffers, and falls back with explicit warnings if the tai-kernel IPC capability is missing. Unit coverage locks offset handle exchange, descriptor construction, and both MLA/index full and partial-current calls.

Constraint: Async prefetch currently has unresolved scheduling contention with GEMM/MoE and current-layer KV collectives, so this commit intentionally does not wire IPC into prefetch.

Constraint: tai-kernel must provide offset-aware CUDA IPC symbols from af9fb67.

Rejected: Replace prefetch all-reduce in this slice | Nsight shows prefetch timing and communicator contention need a separate scheduling design.

Rejected: Fail-fast on missing tai IPC immediately | remote ETE still needs to validate production deployment capability before removing the warning fallback.

Confidence: medium

Scope-risk: moderate

Directive: Keep prefetch and synchronous materialize decisions separate until prefetch has a low-SM/copy-engine schedule.

Related: tai-kernel af9fb67

Tested: git diff --check; remote runtime/unit evidence recorded in docs/advanced_features/nsa_prefill_cp_page_aligned_cache_contract.md.

Not-tested: Fresh GLM5 ETE after this commit; async prefetch IPC path.

Co-authored-by: OmX <omx@oh-my-codex.dev>
2026-06-01 00:00:15 +08:00
laoyao0822
71872bb851 Expose empty chat messages behind conversion failures
The replay failure returns HTTP 400 before generation, so the next diagnostic point must be the chat-template conversion boundary rather than HiCache or scheduler paths. This adds a failure-only summary that records message roles, content shapes, empty indices, tool metadata, bootstrap room, and template mode without logging raw user content.

Constraint: Production replay traffic can contain large/private messages, so diagnostics must avoid full content dumps.

Rejected: Enable full request logging | too noisy and exposes complete prompt payloads.

Rejected: Normalize empty content immediately | would hide whether the malformed message is user, assistant, tool, or template-derived.

Confidence: high

Scope-risk: narrow

Directive: Keep this log failure-only; do not move it to the hot path or print raw content.

Tested: python -m py_compile python/sglang/srt/entrypoints/openai/serving_chat.py

Tested: remote docker py_compile after scp to g0034:/mnt/beegfs/cjy/sglang-dev

Not-tested: Full replay reproduction; requires user-driven traffic after service restart.
2026-05-31 04:06:50 +08:00
laoyao0822
0fc95b6439 Keep CP HiCache reuse page-safe without eviction log churn
Page-aligned CP shared KV can pad out_cache_loc beyond valid current rows, so current reuse now gates MLA composition on the valid extend rows and permits draft partial-current reuse once the TAI sparse-page capability check passes. The TAI current-slot path self-tests sparse pages before use and falls back to the torch reference when the installed kernel is stale.

Eviction success and no-op diagnostics were also moved from INFO to DEBUG so owner-lane and host-admission churn does not flood production logs; true write failures remain WARNING.

Constraint: CP shared KV uses page-aligned physical reservations where valid suffix rows can be shorter than padded out_cache_loc.

Constraint: Production failure/fallback logs must remain visible, but hot successful eviction paths should not emit INFO per victim/rank.

Rejected: Keep draft partial-current reuse disabled | would preserve avoidable full materialization on draft cache-hit suffixes.

Rejected: Trust the TAI current-slot kernel unconditionally | stale kernels can corrupt sparse current-page composition.

Confidence: medium

Scope-risk: moderate

Directive: Do not reintroduce INFO logging in eviction hot paths without rate limiting and runtime evidence.

Tested: local py_compile for touched Python files

Tested: local git diff --check

Tested: remote container py_compile for touched Python files

Tested: remote PYTHONPATH=python python -m pytest -q test/registered/unit/mem_cache/test_cp_hicache_metadata.py::TestHiCacheEvictLoggingLevels::test_evict_hot_path_success_logs_are_debug_only test/registered/unit/mem_cache/test_cp_shared_kv_runtime.py -> 82 passed, 5 warnings, 2 subtests passed

Not-tested: full ETE traffic after this commit; draft partial-current accept length still needs user-driven runtime validation
2026-05-31 03:31:06 +08:00
laoyao0822
3c14b1f127 Enable index partial-current reuse without replaying prefix materialize
The index path now mirrors the target MLA partial-current contract: prefetched or synchronously materialized prefix pages are composed with valid current index K/scale rows in slot-dense page buffers. Current-only batches keep the compact current-index fast path, while partial cache-hit batches share one composed dense index buffer across the in-seq prev/next topk pair.\n\nThe prefetch consume path remaps through the slot page inverse instead of treating the slot-dense buffer as physical-pool capacity, and current-index quantization uses valid extend rows so padded out_cache_loc does not disable reuse.\n\nConstraint: CP shared KV remains page-slot based; padding rows must stay invisible to attention/index semantics\nConstraint: Draft/EAGLE partial-current reuse remains guarded by should_reuse_current_extend_kv\nRejected: Replace prefix all-reduce with all-gather | NCCL all-gather still uses SM and would require an additional compose/scatter step\nConfidence: medium\nScope-risk: moderate\nDirective: Do not reintroduce current-only gating for index reuse; partial target cache hits must compose prefix + valid current rows\nTested: Local py_compile for touched Python files\nTested: g0034 sglang-glm5-dev-2 PYTHONPATH=python python -m pytest -q test/registered/unit/mem_cache/test_cp_shared_kv_runtime.py -> 77 passed, 5 warnings, 2 subtests passed\nNot-tested: Full ETE traffic with latest commit; CUDA perf impact of index partial-current prefetch under production load
2026-05-31 02:31:09 +08:00
laoyao0822
251a48fb0a Stabilize CP HiCache page-first direct transfers on CUDA 13
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.
2026-05-31 01:26:07 +08:00
laoyao0822
b328baec7c Stabilize EAGLE draft cache hits under CP HiCache
The failing runs showed EAGLE accept length collapsing when draft cache-hit suffixes used the new partial-current splice path.  This keeps target partial-current reuse enabled, but returns EAGLE/NextN draft cache-hit suffixes to the previous full-materialize path with an explicit fallback warning until the draft splice path has value-level ETE proof.\n\nThe same change set also tightens the page-granular CP HiCache contract for scheduler-visible hits and makes the prefill-to-decode EAGLE handoff observable without cloning hot-path metadata.  Exact non-page CP hits are floored to a page boundary for new scheduling decisions, while internal unfinished-request refresh keeps its exact accounting.\n\nConstraint: CP shared KV and HiCache operate at page granularity; exposing token-precise CP tails to scheduler-visible cache hits can force non-page partial materialization.\nConstraint: EAGLE/NextN draft has only one executable layer, so draft prefetch and draft partial-current splice need a separate correctness contract from target layers.\nRejected: Keep draft partial-current splice enabled | remote logs correlate it with avg accept length around 0.068 and median 0.\nRejected: Clone decode metadata tensors on transfer | slot ownership until process_prebuilt consumes them avoids extra hot-path copies.\nConfidence: medium\nScope-risk: moderate\nDirective: Do not re-enable draft partial-current reuse without metadata/draft-KV value checks and ETE accept-length evidence.\nTested: g0034 container py_compile for touched modules.\nTested: g0034 container PYTHONPATH=python python -m pytest -q test/registered/unit/disaggregation/test_decode_queue_compaction.py test/registered/unit/mem_cache/test_cp_hicache_metadata.py test/registered/unit/mem_cache/test_cp_shared_kv_runtime.py -> 183 passed, 5 warnings, 2 subtests passed.\nNot-tested: Fresh ETE accept-length run after this exact commit; requires user-driven traffic restart.
2026-05-30 22:31:43 +08:00
laoyao0822
10296a5fef Prevent stale CP HiCache tails from overlapping new page owners
CP HiCache owns KV at page granularity, but exact valid-tail extension and backed partial-tail split could leave an old sub-page tail child beside a new suffix that reuses the same physical page. That makes radix residency ambiguous across device, host, and draft mirrors. The insert/match split paths now prune stale floored tails when safe, and defer/fail through the existing pending-split path when the subtree is protected or has in-flight backup state.\n\nThis also keeps a temporary scheduler boundary warning for externally observed zero-output responses so future ETE runs can classify whether zero visible output reaches SGLang's output processor.\n\nConstraint: CP shared KV and HiCache manage physical KV by page, while radix keys retain valid-token lengths.\nRejected: Keep overlapping old tail nodes after page-floor split | leaves two independent cache states for one physical tail page.\nRejected: Force-prune protected or in-flight backup tails | can mutate cache state still used by active transfer or inference.\nConfidence: medium\nScope-risk: moderate\nDirective: Do not remove the stale-tail prune without replacing it with another page-granular ownership rule for CP HiCache radix splits.\nTested: Remote py_compile for hiradix_cache.py and scheduler_output_processor_mixin.py in g0034 container.\nTested: Remote PYTHONPATH=python python -m pytest -q test/registered/unit/mem_cache/test_cp_hicache_metadata.py -> 97 passed, 5 warnings.\nNot-tested: Full ETE recovery of EAGLE accept length; latest ETE still shows accept collapse, documented in C55.\nNot-tested: Router/client-side output_len=0 correlation when server-side OUTPUT_ZERO_DEBUG does not fire.
2026-05-30 05:51:53 +08:00
laoyao0822
e9c341afe8 Hold decode EAGLE metadata until prebuilt consumption
Transferred EAGLE metadata buffers are reusable slot views. The previous clone-based mitigation protected correctness but added copies on the transfer hot path and hid the actual lifetime contract. This change makes the transferred request own the metadata slot while it sits in the decode waiting queue, then releases it immediately after process_prebuilt has consumed top-k and hidden state into the prebuilt batch. Abort paths also release any held decode metadata slot.

Constraint: Decode disaggregation metadata buffers are reusable slot views consumed later by process_prebuilt.

Rejected: Clone transferred EAGLE tensors at commit time | correct but less efficient and masks the ownership contract.

Rejected: Release in process_batch_result_prebuilt | holds slots across forward longer than needed.

Confidence: medium

Scope-risk: moderate

Directive: Do not free successful EAGLE transfer metadata in pop_transferred unless process_prebuilt consumption is also moved earlier.

Tested: Remote py_compile for decode.py, scheduler.py, scheduler_output_processor_mixin.py, and test_decode_queue_compaction.py.

Tested: Remote focused lifecycle tests passed: 3 passed.

Tested: Remote full test_decode_queue_compaction.py passed: 11 passed, 5 warnings.

Not-tested: Fresh ETE runtime validation of EAGLE accept-length recovery after the C48 sync.
2026-05-30 04:51:50 +08:00
laoyao0822
07c9544737 Preserve transferred EAGLE state past metadata-slot reuse
Decode committed EAGLE top-k and hidden-state tensors as views into reusable metadata-buffer rows. The metadata index is freed immediately after transfer commit, while the request may wait before process_prebuilt consumes the draft state. Under concurrent cache-hit traffic a later transfer can overwrite the same row, leaving output_id copied correctly but EAGLE draft state corrupted, which matches low accept length despite successful KV/state registration.

Constraint: Metadata slots are intentionally recycled right after transfer commit for throughput.

Rejected: Hold metadata slots until process_prebuilt | larger lifetime change and reduces transfer capacity; cloning the small prebuilt EAGLE state is narrower.

Confidence: high

Scope-risk: narrow

Directive: Do not store reusable metadata-buffer views on Req unless the slot lifetime is extended through all consumers.

Tested: Local py_compile for decode.py and test_decode_queue_compaction.py.

Tested: Remote g0034 container py_compile for decode.py and test_decode_queue_compaction.py.

Tested: Remote g0034 focused clone-lifetime test: 1 passed.

Tested: Remote g0034 test_decode_queue_compaction.py: 10 passed, 5 warnings.

Not-tested: ETE cache-hit accept-length validation after restarting prefill/decode with this synced code.
2026-05-30 03:36:42 +08:00
laoyao0822
b7364d23f9 Keep draft cache-hit KV on current-suffix compose
Cache-hit EAGLE/NextN draft was falling back to full materialization while the target path used page-aligned prefix materialize plus fresh current-suffix splice. That creates a target/draft KV source asymmetry exactly on the high-cache-hit path where decode accept length collapsed. The draft model still does not get next-layer async prefetch; only the same-layer current suffix compose contract is made role-agnostic.

Constraint: EAGLE/NextN has no real next decoder layer, so draft async prefetch remains disabled.

Rejected: Restore draft full-materialize fallback | recreates the observed target/draft cache-hit asymmetry and hides stale-current-suffix bugs.

Confidence: medium

Scope-risk: moderate

Directive: Do not reintroduce draft cache-hit current-reuse fallback without proving the draft persistent pool has fully fresh suffix rows before attention.

Tested: Local py_compile for cp_shared_kv_runtime.py and test_cp_shared_kv_runtime.py.

Tested: Remote g0034 container py_compile for changed runtime/test files.

Tested: Remote g0034 pytest test_cp_shared_kv_runtime.py: 73 passed, 5 warnings, 2 subtests passed.

Tested: Remote g0034 related suite test_nsa_cp_utils.py test_cp_shared_kv_layout.py test_cp_shared_kv_runtime.py: 128 passed, 5 warnings, 2 subtests passed.

Not-tested: ETE cache-hit accept-length validation after restarting prefill/decode with this synced code.
2026-05-30 03:06:17 +08:00
laoyao0822
5f343c65ca Keep EAGLE diagnostics from breaking disagg startup
The C32 startup audit was diagnostic-only, but it incorrectly read spec_algorithm from helper queue objects. PrefillBootstrapQueue and DecodePreallocQueue keep scheduler-owned runtime state instead, so the debug path failed before KV manager initialization could complete.\n\nUse scheduler.spec_algorithm for both startup audits and record the ownership rule in the page-aligned cache ledger.\n\nConstraint: Queue helpers do not copy every scheduler field onto self.\nRejected: Disable EAGLE accept diagnostics | would lose the draft-transfer evidence needed for the accept-len investigation.\nConfidence: high\nScope-risk: narrow\nDirective: Disaggregation helper diagnostics should read scheduler-owned runtime state unless the constructor explicitly copied the field.\nTested: Local py_compile for prefill.py and decode.py.\nTested: Local git diff --check.\nTested: Synced prefill.py/decode.py to g0034 and remote container py_compile passed.\nNot-tested: Full disaggregated startup and ETE traffic after this fix; user controls runtime launch and load generation.
2026-05-30 01:37:05 +08:00
laoyao0822
b56a4f2e6b Stabilize CP HiCache page-tail ownership under EAGLE reuse
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
2026-05-30 01:20:01 +08:00
laoyao0822
21065cdfdf Keep tiny CP cache-hit suffixes off async prefetch
The repeated-request hang reproduced on g0034 after the cached-prefix
request created both MLA and index async prefetchers with
prefix_lens=[40320] and extend_lens=[65]. The existing one-page default
only blocked sub-page suffixes, so a barely-over-one-page suffix still
entered the next-layer collective path before any later forward progress
was logged.

Raise the default async prefetch extend gate to one page per CP lane while
keeping the env override. This only gates async prefetcher object
creation; target partial-current reuse still uses the synchronous
page-slot compose/current-splice path when no prefetcher exists.

Constraint: cp_size=8,page_size=64 repeated prompt had extend_len=65 and hung immediately after has_mla=True has_index=True create_result logs
Rejected: Disable partial-current reuse for short extends | that would lose the cache-hit benefit and regress current/full reuse
Rejected: Disable all async prefetch by default | broader performance impact than the observed tiny-suffix failure
Confidence: medium
Scope-risk: moderate
Directive: Do not lower the default below one page per CP lane without ETE proof that repeated cache-hit tiny suffixes no longer hang
Tested: Remote g0034 container py_compile for touched runtime/prefetch/test files; targeted C22 tests passed 2 tests plus 2 subtests; full test_cp_shared_kv_runtime.py passed 73 tests plus 2 subtests
Not-tested: Full multi-node ETE repeated-request run after this threshold change
2026-05-29 22:04:23 +08:00
laoyao0822
40cf691c78 Keep CP HiCache valid tails page-owned during extension
CP shared KV HiCache now treats non-page-aligned valid-tail nodes as page-owned when a later request extends beyond them. The prefix probe, match path, insert path, and prepared backup start now agree on flooring the reusable prefix to the previous physical page boundary, so prepared backup metadata cannot start mid-page or fail to attach at insertion.

Duplicate frees under CP HiCache now go through a page-safe free helper. Insert and unfinished duplicate ranges free only fully unprotected pages; no-insert completion still releases the right tail owned only by the finishing request.

Constraint: CP HiCache allocator frees whole physical pages even when called with token-granular locs.

Rejected: Partial-page sharing/refcounting | too complex for the current page-as-minimum-unit contract.

Rejected: Fix only prepare_write_backup_for_req | match_prefix and insert would still expose exact valid-tail hits and desynchronize prepared backup length.

Confidence: medium

Scope-risk: moderate

Directive: Do not expose non-page-aligned CP valid-tail hits to extending requests unless partial-page ownership is explicitly implemented end-to-end.

Tested: remote g0034 py_compile for touched files

Tested: remote g0034 test_cp_hicache_metadata.py 97 passed

Tested: remote g0034 test_cp_shared_kv_runtime.py 73 passed

Not-tested: test_cp_shared_kv_layout.py aborts during installed sgl_kernel architecture-specific op loading before assertions
2026-05-29 21:49:17 +08:00
laoyao0822
2a9dfcca6f Keep current reuse independent of async CP prefetch
Async MLA/index prefetch is a scheduling optimization, not the correctness contract for target current reuse. Tiny cache-hit suffixes can skip async prefetcher creation while target partial-current reuse still composes page-slot prefix materialization with current KV rows synchronously. CP HiCache radix/device accounting now treats retained valid-tail pages as physical page spans so allocator state stays consistent when logical cache keys are shorter than the retained page.

Constraint: CP shared KV ownership and HiCache residency are page-granular while request-visible cache lengths remain valid-token lengths.
Constraint: Async prefetch can hang or regress on large-prefix tiny-extend traffic and must not be required for current reuse.
Rejected: Treat missing prefetcher as fail-fast for target partial-current reuse | disabled useful current reuse and broke tiny-prefix/tiny-suffix traffic.
Rejected: Keep async prefetcher object with synchronous consume mode | conflates prefetch object existence with current-layer correctness and hides fallback semantics.
Confidence: medium
Scope-risk: moderate
Directive: Do not make current-only or target partial-current reuse depend on MLA/index prefetcher creation; prefetcher objects mean async next-layer work exists.
Tested: Remote g0034 container py_compile for touched modules.
Tested: Remote g0034 PYTHONPATH=python python -m pytest -q test/registered/unit/mem_cache/test_cp_shared_kv_runtime.py -> 73 passed.
Tested: Remote g0034 PYTHONPATH=python python -m pytest -q test/registered/unit/mem_cache/test_cp_hicache_metadata.py -> 90 passed.
Not-tested: Latest full ETE traffic run with GLM-5.1 CP HiCache after this commit.
Not-tested: CUDA kernel-level performance impact of synchronous no-prefetch partial-current compose.
2026-05-29 19:34:27 +08:00
laoyao0822
2b1524bd8c Apply CP HiCache page-floor policy before backup and insert
Backed CP HiCache tail splits were floored during match, but write preparation and radix insertion could still observe the raw sub-page prefix. That could skip one token during pre-forward reservation or call CpHiCacheNodeMetadata.split() with a non-page boundary. The shared helper now floors backed partial-tail prefixes before probe, match, and insert.

Constraint: CP HiCache uses page-granular physical ownership while radix keys may keep valid-tail lengths.

Rejected: Let probe and insert keep token-precise tail prefixes | it reintroduces half-page ownership through a different entry point.

Confidence: high

Scope-risk: moderate

Directive: Any CP radix path that can split or reserve against backed HiCache metadata must use the same page-floor policy.

Tested: Remote red tests first for probe and insert failures in g0034 container.

Tested: Remote py_compile for hiradix_cache.py and test_cp_hicache_metadata.py.

Tested: Remote targeted C7 tests: 4 passed, 3 warnings.

Tested: Remote pytest test_cp_hicache_metadata.py test_cp_hicache_load_back_owner_lanes.py: 94 passed, 5 warnings.

Not-tested: Live ETE traffic with concurrent divergent tail prompts.

Co-authored-by: OmX <omx@oh-my-codex.dev>
2026-05-29 07:10:53 +08:00
laoyao0822
8001c4ae8e Floor CP HiCache tail splits to page boundaries
CP HiCache ownership is page-granular, so a backed radix node must not be split inside a padded physical tail page. When a shorter hit would require an interior tail-page boundary, matching now floors to the previous page boundary and sacrifices the sub-page cache prefix instead of splitting ownership metadata.

Constraint: Host/device/draft CP HiCache metadata tracks page owners and padded physical spans.

Rejected: Split one padded tail page across two radix nodes | it would require half-page ownership semantics and risks double-counting capacity.

Confidence: high

Scope-risk: moderate

Directive: Keep exact valid-tail hits, but floor partial backed-node splits to page boundaries unless metadata gains explicit sub-page ownership.

Tested: Remote py_compile for hiradix_cache.py and test_cp_hicache_metadata.py in g0034 container.

Tested: Remote pytest targeted backed-tail split tests plus exact valid-tail hit test: 3 passed.

Tested: Remote pytest test_cp_hicache_metadata.py test_cp_hicache_load_back_owner_lanes.py: 92 passed, 5 warnings.

Not-tested: Live ETE traffic under divergent short-prefix prompts.

Co-authored-by: OmX <omx@oh-my-codex.dev>
2026-05-29 06:57:16 +08:00
laoyao0822
c9f790cde9 Fail fast when CP MLA partial-current compose is unavailable
Mixed prefix/current MLA reuse depends on the page-slot prefetch compose path to preserve padded tail semantics. The compact materialize/current merge path can re-expose suffix slack as valid dense rows, so the backend now raises with an explicit fail-fast marker instead of silently falling back.

Constraint: Page-aligned CP shared KV contract requires suffix tail slack to remain invalid.

Rejected: Keep compact merge as a fallback | it has a different dense-row contract and can hide correctness bugs.

Confidence: high

Scope-risk: moderate

Directive: Do not reintroduce merge_materialized_and_current_kv on MLA partial-current reuse without a page-slot correctness proof.

Tested: Remote py_compile for nsa_backend.py and test_nsa_cp_utils.py in g0034 container.

Tested: Remote pytest test_nsa_cp_utils.py::TestNSAInSeqCPUtils::test_mla_partial_current_path_fails_fast_instead_of_compact_fallback.

Tested: Remote pytest test_nsa_cp_utils.py test_cp_shared_kv_layout.py test_cp_shared_kv_runtime.py: 124 passed, 5 warnings.

Not-tested: Live ETE traffic and CUDA kernel execution for this fail-fast path.

Co-authored-by: OmX <omx@oh-my-codex.dev>
2026-05-29 06:46:39 +08:00
laoyao0822
dd69c5970c Standardize CP shared-KV runtime fallback markers
Current reuse, TAI materialize, and TAI fused MLA store are optimized CP shared-KV paths. Their fallback helpers already emitted warnings, but the messages were not consistently grep-able with the standard CP shared-KV fallback marker.

This moves the marker into the helper layer so individual fallback call sites do not have to remember to include it, and keeps the existing per-reason rate limiting.

Constraint: Runtime fallback logs must be visible without enabling debug logging.

Rejected: Prefix only selected call sites | helper-level prefixing avoids future unmarked fallback messages.

Confidence: high

Scope-risk: narrow

Directive: New CP shared-KV runtime fallback helpers should use [CP_SHARED_KV_FALLBACK] and a component name.

Tested: Local py_compile for cp_shared_kv_runtime.py and test_cp_shared_kv_runtime.py.

Tested: Local extracted assertLogs check for all three fallback helper prefixes.

Tested: Remote g0034 pytest exact runtime helper tests: 2 passed, 3 warnings.

Not-tested: Full cp_shared_kv_runtime.py suite in this slice.

Co-authored-by: OmX <omx@oh-my-codex.dev>
2026-05-29 06:02:13 +08:00
laoyao0822
e09c8256d7 Expose CP shared-KV allocation fallback as a warning
The compute-owner allocation path is an intended hot path for CP shared KV. Falling back to legacy page allocation changes the behavior profile enough that INFO/no-marker logging is too easy to miss during profiling.

This makes the fallback warning-visible and gives it the standard CP shared-KV fallback marker while preserving the existing behavior of logging every fallback event.

Constraint: Production fallbacks from intended CP shared-KV hot paths must be visible and grep-able.

Rejected: Keep INFO logging | historical silent or low-visibility fallbacks hid inactive optimized paths.

Confidence: medium

Scope-risk: narrow

Directive: Keep fallback logs warning-level unless the path is proven to be expected steady state and separately observable.

Tested: Local py_compile for common.py and test_cp_shared_kv_layout.py.

Tested: Local extracted assertLogs check for _log_cp_shared_kv_alloc_fallback warning prefix and reason formatting.

Tested: Remote g0034 py_compile for common.py and test_cp_shared_kv_layout.py.

Not-tested: Full remote pytest for test_cp_shared_kv_layout.py; current g0034 container aborts while importing installed sgl_kernel common_ops before test execution.

Co-authored-by: OmX <omx@oh-my-codex.dev>
2026-05-29 05:59:02 +08:00
laoyao0822
38532a1fc9 Prevent incompatible CP shared-KV transfer mapping
Mooncake is the only disaggregation transfer backend in this branch with CP shared-KV owner filtering plus logical-page-position destination selection. NIXL still slices destination pages by the original chunk slice, so allowing CP shared-KV prefill on NIXL can silently pair filtered prefill pages with the wrong decode pages.

This keeps the supported path narrow while preserving the page-aligned transfer contract: non-page-aligned valid tails transfer their physical tail page, but do not get padded to CP-size pages.

Constraint: CP shared-KV transfer remaps prefill logical pages to per-rank physical pages while decode metadata remains request-position based.

Rejected: Let NIXL continue through the generic slice path | it lacks logical-page-position selection and can silently corrupt CP shared-KV transfers.

Confidence: high

Scope-risk: narrow

Directive: Do not enable CP shared-KV on another PD transfer backend until its sender filters owner pages and selects decode pages by logical request-page position.

Tested: Local py_compile for server_args and touched tests.

Tested: Remote g0034 pytest test_cp_shared_kv_transfer_mapping.py test_req_to_token_pool.py TestHiCacheArgs: 22 passed, 8 subtests passed.

Not-tested: End-to-end PD transfer with a live non-page-aligned prompt.

Co-authored-by: OmX <omx@oh-my-codex.dev>
2026-05-29 05:53:37 +08:00
laoyao0822
0043037f78 Keep CP HiCache owner lanes padded for valid-tail victims
Load-back owner-lane eviction can need to evict a device-resident node whose radix value is a valid tail shorter than the physical tail page. The eviction planner now pads that value only to the page boundary before deriving owner counts, matching the existing write/load capacity contract without exposing padding to radix or scheduler lengths.

Constraint: CP HiCache capacity remains page-owner based while radix node values remain valid-token based.
Constraint: Avoid collectives; owner-lane capacity must be deterministic from local metadata and logical page ids.
Rejected: Require device-resident victim values to be page-aligned | valid-tail cache nodes are now an intentional supported state.
Rejected: Pad to cp_size pages | this would waste KV and violate the page-boundary-only contract.
Confidence: medium
Scope-risk: narrow
Directive: If split-inside-tail support is added later, preserve page ownership/refcount semantics before sharing one padded physical page across radix nodes.
Tested: local py_compile for hiradix_cache.py and touched CP HiCache tests.
Tested: remote g0034 new C8 exact tests: 3 passed, 3 warnings.
Tested: remote g0034 CP HiCache impacted suites: 146 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 starlette.
Not-tested: CUDA E2E runtime for this commit.
Co-authored-by: OmX <omx@oh-my-codex.dev>
2026-05-29 05:39:25 +08:00
laoyao0822
7cfc3c1324 Preserve CP HiCache valid tails while padding physical pages
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>
2026-05-29 05:28:53 +08:00
laoyao0822
c551623ca8 Keep CP shared-KV cache hits page-aligned without short fallback
CP shared KV and HiCache need a stable contract where physical cache coverage is page-aligned, while scheduler/radix-visible hit length remains the valid token length. This records the contract, adds page-aligned extent metadata, keeps owner assignment on actual tail pages instead of short-prefix fallback, and updates partial current reuse tests around tail-page masking.

Constraint: CP owner lanes operate on page units while scheduler and radix hit accounting must remain token-valid.

Rejected: Pad short suffixes to cp_size or 2*cp_size pages | wastes KV capacity and can turn a small tail into a much larger physical span.

Rejected: Silent direct-write or prefetch fallback | production fallback must be warning-visible for diagnosis.

Confidence: medium

Scope-risk: moderate

Directive: Do not reintroduce replicated short-radix fallback without checking docs/advanced_features/nsa_prefill_cp_page_aligned_cache_contract.md.

Tested: local py_compile for touched runtime, utility, owner, and unit-test files.

Tested: remote g0034 container three-file suite: 122 passed, 5 warnings.

Not-tested: full local pytest, blocked by missing runtime dependencies such as orjson.

Not-tested: CUDA E2E runtime for this commit.

Co-authored-by: OmX <omx@oh-my-codex.dev>
2026-05-29 05:04:50 +08:00
laoyao0822
1e1807f2bc reset prefetch default env 2026-05-29 02:27:05 +08:00
laoyao0822
c3fc3ff752 Stabilize CP shared-KV prefetch around draft cache hits
Cache-hit EAGLE/NextN draft extends can enter the draft DeepEP MoE immediately after CP shared-KV attention. The partial current-reuse path is kept for target layers, but draft cache-hit suffixes now use full materialization until draft has an explicit same-layer reuse contract. Next-layer MLA/index prefetch is also gated by the actual model depth, so the single-layer draft model does not enqueue unused next-layer async work.

The temporary stage traces used to isolate the hang are removed. The retained draft current-reuse fallback is a bounded warning because it changes the runtime path intentionally.

Constraint: EAGLE/NextN has one executable draft layer and mirrors target KV state.

Rejected: Keep partial current reuse for draft cache-hit suffixes | reproduced hangs at draft layer0 before DeepEP MoE completion.

Rejected: Keep temporary stage traces | useful for diagnosis but too noisy for normal runs.

Confidence: medium

Scope-risk: moderate

Directive: Do not re-enable draft cache-hit partial current reuse without an explicit draft same-layer reuse contract and ETE validation with CP shared KV + HiCache + EAGLE.

Tested: py_compile on edited Python files; git diff --check; temp trace grep returned no matches.

Not-tested: Local targeted pytest is blocked by missing pybase64 in this environment; full ETE after log cleanup not run.
2026-05-29 00:33:41 +08:00
laoyao0822
26c792939d Gate tiny CP shared-KV prefetch by token threshold
Prefix materialize prefetch has fixed launch and coordination overhead, so tiny cache-hit prefixes should stay on the simpler synchronous path. The default threshold is now expressed as a cached import-time token threshold while the existing page override remains available for workload-specific tuning.

Both MLA and index prefix prefetchers use the same page-size-aware threshold, so CP shared-KV prefix prefetch is enabled only when the prefix covers at least one page per CP lane and at least the configured token threshold.

Constraint: The default token threshold must be read once to avoid hot-path env lookups

Constraint: Existing page threshold override must continue to disable or force the gate

Rejected: Hard-code 1024 tokens | tuning needs to be possible without code changes

Rejected: Read token threshold on every maybe_create call | unnecessary hot-path env lookup

Confidence: medium

Scope-risk: narrow

Directive: Keep MLA and index prefetch threshold logic shared; do not let their gates diverge

Tested: Not rerun per user request; prior py_compile, diff check, and isolated env-cache check had passed before commit

Not-tested: full pytest in local environment due missing optional dependencies/kernel import constraints

Not-tested: GLM5 E2E throughput after this threshold change
2026-05-28 08:57:45 +08:00
laoyao0822
25f2147677 Reduce CP HiCache capacity synchronization to owner-lane logic
CP shared KV and HiCache now use owner-lane metadata as the
authoritative capacity view for host write admission and GPU load-back
planning. This removes the debug scalar capacity env and keeps CP load-back
from relying on a rank-wide scalar collective when per-owner availability is
already known. The load-back planner also accounts for evicting child leaves
that unlock ancestor device residency, which fixes small lane deficits despite
large aggregate evictable capacity.

The commit also adds gated CPU timing logs for CP shared-KV MLA/index
prefetch and a CUDA microbenchmark for comparing dense all-reduce with
owner-packed all-gather layouts. The timing logs are intentionally behind the
existing MLA prefetch log env and should not be enabled for throughput
measurements.

Constraint: CP shared KV owner lanes require target/draft capacity decisions to preserve page_owners rather than total-token scalars
Constraint: CUDA collective benchmarks must run on target GPU hosts, not locally
Rejected: Keep SGLANG_CP_HICACHE_CAPACITY_DEBUG observer env | owner-lane admission now replaces that scalar debug path
Rejected: Add a silent scalar-allreduce fallback | unexpected owner-lane mismatch should fail fast or log loudly
Confidence: medium
Scope-risk: moderate
Directive: Do not reintroduce CP capacity collectives on the scheduler hot path without proving the owner-lane metadata is insufficient
Directive: Disable SGLANG_CP_SHARED_KV_LOG_MLA_PREFETCH for end-to-end performance runs; it is diagnostic and high-volume
Tested: git diff --check
Tested: python -m py_compile on changed runtime/test/benchmark Python files
Tested: remote pytest -q test/registered/unit/mem_cache/test_cp_hicache_load_back_owner_lanes.py test/registered/unit/mem_cache/test_cp_hicache_metadata.py (81 passed, 5 warnings)
Not-tested: CUDA benchmark benchmark/hicache/bench_cp_shared_kv_prefetch_collective.py
Not-tested: full GLM5 E2E throughput after this commit
2026-05-28 08:31:49 +08:00
laoyao0822
ff33446787 Make CP HiCache residency owner-lane deterministic
CP shared KV cannot treat capacity as a scalar token count: cache-hit load-back and fresh extend allocation both have to preserve the logical page owner pattern or later direct writes, HiCache reload, and prefix materialization can read the wrong lane. This change moves the critical paths to owner-lane plans, makes owner-lane exhaustion recoverable during prefill scheduling, and routes shared-KV prefix prefetch through prefetch-stream-safe KV getters so HiCache layer-load waits do not attach to the forward stream.

Constraint: CP shared KV correctness depends on page owner lane preservation across allocation, backup, load, eviction, and prefix materialization.

Constraint: Avoid adding CP/global collectives for capacity agreement; derive capacity from deterministic local owner-lane state.

Rejected: Keep SGLANG_DISABLE_TAI_OWNER_SELECT fallback | legacy allocation can silently break owner-lane invariants.

Rejected: Scalar total-token eviction for CP HiCache load-back | total capacity can be sufficient while the required owner lane is exhausted.

Confidence: medium

Scope-risk: broad

Directive: Do not reintroduce silent legacy fallback in owner-lane paths; unexpected owner-lane failure must be warning-level fail-closed or recoverable capacity wait.

Tested: Remote g0034 container PYTHONPATH=python python -m pytest test/registered/unit/mem_cache/test_alloc_pages_with_owners.py test/registered/unit/mem_cache/test_cp_shared_kv_layout.py test/registered/unit/mem_cache/test_cp_hicache_load_back_owner_lanes.py test/registered/unit/mem_cache/test_cp_shared_kv_runtime.py -q -> 95 passed.

Tested: Local py_compile for modified runtime/cache/scheduler modules.

Not-tested: Full CUDA ETE performance trace for cache-hit overlap and MTP accept-rate impact.

Co-authored-by: OmX <omx@oh-my-codex.dev>
2026-05-28 05:54:23 +08:00
laoyao0822
67d52346de Route page-first direct HiCache loads through TAI
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
2026-05-28 02:37:56 +08:00
laoyao0822
40a8de5fd1 Make CP HiCache backup admission deterministic
CP HiCache write-through under shared KV was still using rank-wide collectives to decide host reservation eviction, and per-layer backup registration could be bypassed before the final forward boundary. This moves backup registration to the final run_batch pre-forward boundary, forwards it through SessionAwareCache, exposes fallback paths as explicit warnings, and introduces deterministic owner-lane capacity planning for CP host reservation.

Constraint: CP shared-KV ranks must keep target and draft host reservations owner-lane consistent without adding hot-path collective synchronization

Constraint: Remote CUDA validation must run in the g0034 container, not locally

Rejected: Keep reserve_slots_max all_reduce as the default admission path | observed reserve collectives reaching double-digit and occasional 100ms+ latency

Rejected: Silent post-forward catch-up backup | hides when per-layer forward-overlap backup is not actually active

Confidence: medium

Scope-risk: broad

Directive: Do not reintroduce CP HiCache hot-path collectives without a measured mismatch case and explicit fallback warning

Tested: py_compile for modified Python modules and CP HiCache metadata test file in remote g0034 container

Tested: python3 -m pytest test/registered/unit/mem_cache/test_cp_hicache_metadata.py -q in remote g0034 container (75 passed, 5 warnings)

Tested: git diff --check HEAD~1..HEAD

Not-tested: Local pytest blocked by missing pybase64 in the local environment

Not-tested: Full CP HiCache + MTP E2E after the no-collective reservation change

Co-authored-by: OmX <omx@oh-my-codex.dev>
2026-05-27 23:11:23 +08:00
laoyao0822
f355fdd39e Overlap CP HiCache backup without exposing partial host state
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
2026-05-27 09:50:47 +08:00
laoyao0822
03529319a1 Keep CP HiCache host visibility behind per-layer target and draft backup
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.
2026-05-27 07:45:16 +08:00
laoyao0822
367dff06f3 Keep CP HiCache draft KV invisible until joint readiness
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
2026-05-27 05:46:34 +08:00
laoyao0822
71c4f66968 Enforce draft KV strong-sync for CP HiCache hits
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>
2026-05-27 05:23:31 +08:00
laoyao0822
8571fe0cd9 Share CP HiCache host budget across target and draft KV
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>
2026-05-27 04:28:50 +08:00
laoyao0822
d14c02b0dc Count evictable device cache when gating HiCache load-back
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>
2026-05-27 02:41:45 +08:00
e5982dcceb fix: call _update_leaf_status in inc/dec_node_lock_ref to prevent phantom evictable nodes
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>
2026-05-27 01:52:41 +08:00
3d65944a22 Remove stale layer_transfer_counter prefetch guards (port from e293d4a39 fix #3)
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>
2026-05-27 00:40:39 +08:00
3186382be1 Drop env-gated CP HiCache page-index validator
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>
2026-05-27 00:15:35 +08:00
1d630def95 Fix CP HiCache load_cp owner-pattern mismatch (cache-hit corruption)
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>
2026-05-27 00:01:39 +08:00
laoyao0822
f2834b3403 Preserve draft NSA state during CP disaggregated transfer
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.
2026-05-26 23:59:28 +08:00
laoyao0822
315eaaff56 Preserve draft KV across CP HiCache hits
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.
2026-05-26 23:58:40 +08:00
d655fad040 Dispatch EAGLE radix bigram builder to tai-kernel
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>
2026-05-26 23:41:17 +08:00
laoyao0822
ec7e9fbc57 Avoid full-prompt embedding in CP MTP prefill
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.
2026-05-23 15:48:36 +08:00
bacad1d498 fix: update leaf status for hicache 2026-05-22 14:09:48 +00:00
bd6e28f8ce fix: embed before CP split in nextn to prevent TP all_reduce shape mismatch
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>
2026-05-22 11:21:09 +00:00