Commit Graph
27 Commits
Author SHA1 Message Date
laoyao0822andOmX 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
laoyao0822andOmX 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 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
laoyao0822andOmX 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
laoyao0822andOmX 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
laoyao0822andOmX 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
laoyao0822andOmX 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
leaveletandClaude Opus 4.7 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 fa80c6278e 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:55 +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
ThomasX f38937723e fix: address CP HiCache review regressions 2026-05-08 04:11:25 +08:00
ThomasX 8844c303e2 fix: synchronize CP HiCache host eviction semantics 2026-05-08 03:20:47 +08:00
ThomasX 3e9152194b test: cover CP HiCache host integration 2026-05-08 02:59:27 +08:00
ThomasX 39f40c02b4 fix: handle empty HiCache match prefixes 2026-05-08 02:47:25 +08:00
ThomasX 829f3ebceb feat: restore CP HiCache host hits as logical locs 2026-05-08 02:34:16 +08:00
ThomasX 08d518c2ed fix: preserve CP host siblings during stale cleanup 2026-05-08 02:06:04 +08:00
ThomasX 2df832dec0 fix: unblock CP host eviction past stale leaves 2026-05-08 01:59:14 +08:00
ThomasX f1a7e605e8 feat: split and evict CP HiCache metadata 2026-05-08 01:50:07 +08:00
ThomasX bebbbfd0e8 fix: treat CP HiCache metadata as backup state 2026-05-08 01:41:53 +08:00
ThomasX 08a95a35a6 fix: implement CP HiCache write retry eviction 2026-05-08 01:35:55 +08:00
ThomasX 068aa2f910 feat: track CP HiCache write-back state 2026-05-08 01:29:39 +08:00
ThomasX eb0ec4299d test: cover CP HiCache metadata validation edges 2026-05-08 01:09:32 +08:00
ThomasX 1ede8bb999 test: harden CP HiCache metadata invariants 2026-05-08 01:06:58 +08:00
ThomasX 5fbe43381f feat: add CP HiCache node metadata 2026-05-08 01:00:09 +08:00