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
parent 2b1524bd8c
commit 2a9dfcca6f
9 changed files with 1067 additions and 48 deletions
@@ -1033,6 +1033,281 @@ Completed C14:
  123 passed, 5 warnings
  ```

+## C15: CP valid-tail radix nodes still need page-physical device accounting
+
+Finding:
+
+- Startup can fail immediately after disaggregated prefill warmup with:
+
+  ```text
+  token_to_kv_pool_allocator memory leak detected!
+  max_total_num_tokens=833024, available_size=833024,
+  evictable_size=3, protected_size=0, session_held=0
+  ```
+
+- The observed `evictable_size=3` is consistent with EAGLE/bigram warmup:
+  a 4-token request is cached as a 3-token bigram radix key.  Under CP HiCache
+  the radix key is valid-tail based, while the CP device allocator is
+  page-granular.
+- `RadixCache.cache_finished_req()` still freed `kv_indices[len(keys):]`.
+  For CP + page allocator, freeing the 4th token of a 64-token page releases the
+  whole physical page even though the radix node still owns the first 3 valid
+  tokens.  The allocator then reports the full pool available while radix still
+  reports an evictable 3-token node.
+- HiRadix device accounting also still used `len(node.key)` / `len(node.value)`
+  in several places.  With valid-tail nodes this is the scheduler-visible valid
+  length, not the allocator-visible physical page span.
+
+Correction:
+
+- CP HiCache device residency must account physical page spans while radix
+  matching and scheduler-visible prefix lengths remain valid-token based.
+- Tail free for CP valid-tail keys must start at the next page boundary, not at
+  `len(keys)`, so partial retained pages are never released out from under a
+  radix node.
+- Splits inside a CP physical page should be floored to the previous page
+  boundary (or rejected/deferred at zero) even when the node is not yet backed
+  by host metadata.  Page is the minimum ownership unit for both GPU and host
+  HiCache.
+
+Tests:
+
+- Add a regression for EAGLE/bigram CP finished-cache tail free: retained
+  partial pages must not call allocator free on the in-page tail.
+- Add CP page-physical accounting regressions for insert/lock/evict paths.
+- Add a split-floor regression for an unbacked CP valid-tail node.
+
+Implemented C15:
+
+- Added `ceil_to_page_len()` and made CP `cache_finished_req()` free suffix
+  pages only from the next page boundary after the valid radix key.  This keeps
+  a retained valid-tail node from releasing the physical page it still owns.
+- Added `HiRadixCache._node_device_resident_len()` and routed CP insert,
+  single-node write locks, path locks, load-back locks, regular eviction, and
+  backed demotion through allocator-visible physical page lengths.
+- Changed the CP split-floor policy to apply to unbacked valid-tail device
+  nodes too.  Splits inside a physical page now floor to the previous page
+  boundary independent of host-backup state.
+
+Completed C15 tests:
+
+```text
+remote g0034 container:
+  test_cp_hicache_metadata.py
+    90 passed, 5 warnings
+  test_cp_hicache_load_back_owner_lanes.py
+  test_cp_shared_kv_layout.py
+  test_cp_shared_kv_runtime.py
+    103 passed, 3 warnings
+```
+
+Not-tested C15:
+
+- `test_radix_cache_unit.py` collection currently fails in the remote container
+  before running tests because the `sgl_kernel::sgl_per_token_group_quant_8bit`
+  operator is not registered in that test import path.  This is an environment /
+  test bootstrap issue, not a C15 assertion failure.
+
+## C16: Exact valid-tail hits can make the next backup start mid-page
+
+Finding:
+
+- Current C7/C15 code keeps exact CP valid-tail cache hits token-precise.  That
+  means a request can hit, for example, 100 valid tokens on 64-token pages and
+  then start its next suffix at token 100.
+- `prepare_write_backup_for_req()` derives the backup start as
+  `max(req.cache_protected_len, existing_prefix_len)`.  If the protected prefix
+  is a valid tail, `start` can be inside a physical page.
+- `_cp_required_host_tokens_by_rank()` / `reserve_write_cp()` then call
+  `pad_token_locs_to_page_boundary()` on `device_indices[start:]`.  That helper
+  is only correct for spans that start at a page boundary; CPU tensors validate
+  this, but CUDA tensors intentionally skip the sync-heavy offset check.
+- Therefore the production CUDA path can silently build a padded physical span
+  from a mid-page start, producing incorrect page-owner accounting and host
+  reservation metadata.
+
+Risk:
+
+- This is a correctness risk for repeated prompts where the previous request
+  leaves a non-page-aligned valid-tail node and the next request extends beyond
+  that exact tail.
+- It also conflicts with the current design principle that CP HiCache owns cache
+  at page granularity and may sacrifice sub-page cache to keep management simple.
+
+Correction options:
+
+- Preferred under the current page-as-minimum-unit contract: expose only the
+  previous page boundary as `cache_protected_len` / `device_indices` for any hit
+  that will be extended, so new backup starts at a page boundary.
+- More complex alternative: implement explicit partial-page sharing/refcounting
+  between a cached prefix tail and the new current suffix.  This is intentionally
+  not the current direction.
+
+Tests needed:
+
+- Cache a 100-token valid-tail node, then issue a request that hits those 100
+  tokens and extends by another suffix.  Assert write reservation starts at 64,
+  not 100, or otherwise fails fast before mid-page padding.
+- Add a CUDA/remote regression or source-level guard proving CUDA hot path cannot
+  call `pad_token_locs_to_page_boundary()` with a non-page-start span.
+
+## C17: Duplicate/no-insert frees are still token-granular under CP pages
+
+Finding:
+
+- C15 fixed the final tail free in `RadixCache.cache_finished_req()` by starting
+  CP tail free at the next page boundary.
+- The duplicate/no-insert frees remain token-sliced:
+
+  ```text
+  cache_finished_req:
+    kv_indices[req.cache_protected_len : new_prefix_len]
+    kv_indices[req.cache_protected_len : len(keys)]
+
+  cache_unfinished_req:
+    kv_indices[req.cache_protected_len : new_prefix_len]
+  ```
+
+- With CP page allocators, freeing any loc inside a page can release the whole
+  page.  If the free range begins or ends inside a page, it can release a page
+  that still contains retained current/request KV or radix-owned valid-tail KV.
+
+Risk:
+
+- Remaining memory-accounting or allocator/radix divergence after the C15 startup
+  fix, especially on chunked/unfinished requests or insert-prefix duplicate
+  paths.
+
+Correction:
+
+- Audit every CP call to `token_to_kv_pool_allocator.free(kv_indices[a:b])`.
+  Under CP page mode, free ranges should either be page-aligned or deliberately
+  converted to whole pages that are no longer referenced by radix/request state.
+- If a range is not page-aligned and cannot be proven safe, fail fast in tests
+  first, then floor/ceil according to ownership semantics.
+
+Tests needed:
+
+- Finished insert where `req.cache_protected_len` and `new_prefix_len` differ
+  inside the same page.
+- `is_insert=False` finished path with a valid-tail key.
+- `cache_unfinished_req()` duplicate-free path with a non-page-aligned protected
+  prefix.
+
+## C18: `cache_protected_len` now mixes valid-token and physical-page semantics
+
+Finding:
+
+- CP valid-tail matching can expose `req.cache_protected_len` as a valid-token
+  length.
+- Some scheduler/runtime checker paths still assert or compute as if
+  `cache_protected_len` is page-aligned physical protection:
+
+  ```text
+  scheduler_runtime_checker_mixin.py::_get_batch_uncached_size
+  schedule_batch.py::_evict_swa
+  ```
+
+Risk:
+
+- Exact valid-tail cache hits may fail assertions or under/over-count uncached
+  tokens in memory self-check paths.
+- Even when asserts do not fire in the current CP/NSA configuration, this is an
+  API ambiguity that can reappear when batch/chunked/SWA paths are combined.
+
+Correction:
+
+- Split request-visible fields conceptually:
+  - valid protected length for scheduler/logical prompt progress;
+  - physical protected length for allocator/page accounting.
+- Until that refactor is done, CP HiCache should floor exposed protected length
+  to the nearest page boundary wherever allocator/page accounting consumes it.
+
+Tests needed:
+
+- Valid-tail CP hit through `Req.init_next_round_input()` and scheduler runtime
+  memory checker.
+- Guard that SWA/page-aligned paths either remain unreachable for CP NSA HiCache
+  or receive a page-aligned protected length.
+
+## C19: Prefetcher creation must not gate target partial-current reuse
+
+Finding:
+
+- Remote ETE run `/mnt/beegfs/cjy/sglang_cp_hicache_20260529_001129.log`
+  died at `2026-05-29 00:17:24` with:
+
+  ```text
+  [CP_SHARED_KV_FAIL_FAST][mla_partial_current_prefetch]
+  reason=missing_prefetcher prefix_lens=[320] extend_lens=[10808]
+  ```
+
+- The same log shows the async MLA prefetcher was intentionally not created:
+
+  ```text
+  create_skip reason=prefix_below_min prefix_pages=5 min_prefix_pages=8
+  ```
+
+- A later run showed the inverse shape: `prefix_lens=[16320]` and
+  `extend_lens=[16]` passed the prefix gate and created async prefetchers even
+  though the current suffix was sub-page.  So "short extend" did not request
+  prefetch directly; a large cache-hit prefix made the async prefetch gate pass.
+
+Corrected contract:
+
+- `CpSharedKVMlaPrefetcher` / `CpSharedKVIndexPrefetcher` mean async
+  next-layer prefetch only.  If async prefetch is disabled, below threshold, or
+  otherwise not worthwhile, do not create a prefetcher object.
+- Target MLA partial-current reuse must not depend on prefetcher existence.
+  When there is no prefetched prefix handle, `nsa_backend.py` synchronously
+  materializes the page-aligned prefix into the slot-dense buffer, all-reduces
+  that prefix range, then splices the current KV rows into their suffix page
+  slots with tail slack mapped to `-1`.
+- Draft/EAGLE cache-hit partial-current reuse remains disabled until the draft
+  KV lifetime has its own same-layer contract.  This does not affect target
+  partial-current reuse or cache-miss current-only reuse.
+- The disabled path is only the old compact materialize/current merge fallback;
+  it can expose padded tail slack.  The synchronous page-slot compose path is
+  the fallback-free current-layer contract when no async prefetcher exists.
+
+Tests:
+
+- Target cache-hit partial-current with `cp_shared_kv_mla_prefetcher=None` must
+  still return true from `should_reuse_current_extend_kv()`.
+- `materialize_prefix_and_reuse_current_kv_page_slots()` covers the no-prefetcher
+  synchronous compose path and verifies tail slack maps to `-1`.
+- Tiny-extend batches with large prefixes should skip async MLA/index prefetcher
+  creation before touching pool getters/remap builders.
+
+## C20: Prefetch disabled must not disable full current reuse
+
+Finding:
+
+- Cache-miss/current-only requests (`extend_prefix_lens_cpu=[0]` and
+  `seq_lens_cpu == extend_seq_lens_cpu`) do not need prefix materialization or a
+  prefetcher: they can reuse the full current forward KV directly.
+- Conflating "no prefetcher" with "no current reuse" regresses warmup,
+  cache-miss, and tiny-prefix traffic by forcing unnecessary full materialize or
+  rebuild work.
+
+Risk:
+
+- Performance regression on no-cache or tiny-prefix traffic.  This is especially
+  visible before HiCache warms up or when a prompt falls below the MLA prefetch
+  threshold (`prefix_pages < min_prefix_pages`, currently 8 pages by default for
+  `cp_size=8,page_size=64`).
+
+Correction:
+
+- Keep `should_reuse_current_extend_kv()` structured as:
+  - `current_only` is sufficient for full current reuse;
+  - target `partial_current and not current_only` is also sufficient, regardless
+    of prefetcher existence, because the backend has synchronous page-slot
+    compose;
+  - draft/EAGLE `partial_current and not current_only` remains rejected.
+- Tests include both target partial-current/no-prefetcher and target
+  current-only/no-prefetcher coverage.
+
 ## Testing requirements / status

 Targeted coverage needed for this contract:
@@ -1082,3 +1357,87 @@ Targeted coverage needed for this contract:
 - Do not add a hot-path CP collective to reconcile capacity.
 - Do not make draft KV manage an independent padded/valid contract.
 - Do not silently fall back from the page-aligned contract.
+
+## C21: Large prefix with tiny extend must not enter MLA/index prefetch
+
+Finding:
+
+- The current prefetch creation gate is prefix-sized: it skips MLA/index
+  prefetch only when `prefix_pages < min_prefix_pages`.
+- The 2026-05-29 remote hang did not violate that gate.  The last batch before
+  detokenizer health failure had:
+  - `prefix_lens=[16320]` (`255` pages at `page_size=64`), so the prefix gate
+    allowed prefetch;
+  - `extend_lens=[16]`, less than one page;
+  - `has_mla=True has_index=True`, followed by no further forward-layer progress
+    and detokenizer timeout.
+- Earlier batches with `prefix_lens=[320]` (`5` pages) correctly logged
+  `create_skip reason=prefix_below_min ... min_prefix_pages=8` and produced
+  `has_mla=False has_index=False`; those are not the observed hang pattern.
+- In other words, short extend did not directly request prefetch.  Cache hit
+  converted most of the request into a large cached prefix, and the large prefix
+  satisfied the prefetch gate while leaving only a sub-page current suffix.
+
+Root-cause analysis:
+
+- MLA/index prefetch is currently a next-layer optimization.  It starts prefix
+  materialization/reduce on a separate stream from layer prepare hooks, then the
+  consumer waits on the event later.
+- Partial-current reuse uses the MLA prefetcher as a page-slot compose contract:
+  prefetched prefix rows are reused and the fresh current suffix KV is inserted
+  into the padded tail page.
+- With `extend_len < page_size`, the suffix page has exactly one CP owner lane
+  and many rank-local zero-owned cases.  The current layer is also short enough
+  that async prefetch collective ordering and event waits are no longer hidden by
+  compute.  A mismatch in next-layer async collective order or event completion
+  can present as a CUDA/NCCL wait rather than a Python exception.
+- The observed symptom matches this: after `has_mla=True has_index=True` for
+  `prefix_lens=[16320], extend_lens=[16]`, no later forward probe appears and
+  the server only reports detokenizer heartbeat timeout.
+
+Risk:
+
+- Tiny cache-hit suffixes can pay prefetch launch/materialize/collective overhead
+  that cannot be amortized by the short extend.
+- More importantly, sub-page suffixes exercise the still-sensitive padded-tail
+  partial-current path.  Until that path is proven by ETE, allowing MLA/index
+  prefetch on `extend_len < page_size` is a correctness/hang risk.
+
+Correction:
+
+- Do not disable partial-current reuse just because `extend_len` is short.
+  Short cache-hit suffixes are exactly where avoiding suffix materialization is
+  valuable.
+- Add a separate minimum-extend gate only for the async next-layer prefetch
+  strategy.  When `extend_len` is below the async threshold, do **not** create an
+  MLA/index prefetcher object; a prefetcher object now means async next-layer
+  prefetch machinery exists.
+- Target partial-current reuse is independent of prefetcher existence.  Without
+  a prefetched prefix handle, the backend takes the synchronous page-slot compose
+  path: materialize/reduce the page-aligned prefix in the current layer and then
+  splice current KV rows into their padded suffix page slots.
+- Full current reuse is also independent of prefetcher existence.  Cache-miss
+  current-only batches keep using the forward KV directly even when async
+  prefetch is disabled, skipped by threshold, or not created.
+- If partial-current reuse is intentionally rejected, for example the current
+  conservative draft/EAGLE cache-hit guard, that rejection must not disable
+  current-only/full-current reuse for other request shapes.
+- Default the async gate to one page when `page_size` is known; allow an env
+  override to lower or raise the threshold for experiments.
+- This must not force tiny cache-hit suffixes back to compact full
+  materialization.
+
+Tests:
+
+- Add runtime helper coverage for the min-extend threshold default and override.
+- Add MLA/index prefetch creation tests where `prefix_pages >= 8` but
+  `extend_len < page_size`; both prefetchers should return `None` before
+  touching pool getters/remap builders, while current reuse remains available.
+- Add a synchronous no-prefetch compose test proving sub-page current suffix rows
+  are filled while tail slack in the current page is masked.
+
+Verified:
+
+- Remote container `/sgl-workspace/sglang-tai` on `g0034`:
+  - `test/registered/unit/mem_cache/test_cp_shared_kv_runtime.py`: `73 passed`
+  - `test/registered/unit/mem_cache/test_cp_hicache_metadata.py`: `90 passed`