From b56a4f2e6b460593d3caffcc1b89d29def877f9e Mon Sep 17 00:00:00 2001 From: laoyao0822 Date: Sat, 30 May 2026 01:20:01 +0800 Subject: [PATCH] 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 --- ..._prefill_cp_page_aligned_cache_contract.md | 408 ++++++++++++++++++ python/sglang/srt/disaggregation/decode.py | 16 + python/sglang/srt/disaggregation/prefill.py | 16 + python/sglang/srt/environ.py | 2 + .../attention/nsa/cp_shared_kv_runtime.py | 6 +- .../sglang/srt/layers/attention/nsa/utils.py | 49 +++ .../srt/layers/attention/nsa_backend.py | 28 +- .../srt/managers/scheduler_dp_attn_mixin.py | 12 +- .../scheduler_output_processor_mixin.py | 77 ++++ python/sglang/srt/mem_cache/common.py | 9 +- python/sglang/srt/mem_cache/hiradix_cache.py | 106 +++++ python/sglang/srt/mem_cache/radix_cache.py | 42 +- .../unit/layers/test_nsa_cp_utils.py | 31 +- .../managers/test_scheduler_dp_attn_mixin.py | 37 ++ .../test_cp_hicache_load_back_owner_lanes.py | 134 +++++- .../mem_cache/test_cp_shared_kv_runtime.py | 8 + 16 files changed, 964 insertions(+), 17 deletions(-) create mode 100644 test/registered/unit/managers/test_scheduler_dp_attn_mixin.py diff --git a/docs/advanced_features/nsa_prefill_cp_page_aligned_cache_contract.md b/docs/advanced_features/nsa_prefill_cp_page_aligned_cache_contract.md index 7fe3ab187..c7b3b19f4 100644 --- a/docs/advanced_features/nsa_prefill_cp_page_aligned_cache_contract.md +++ b/docs/advanced_features/nsa_prefill_cp_page_aligned_cache_contract.md @@ -1558,3 +1558,411 @@ Verified C22: `73 passed, 2 subtests passed`. - Local pytest remains blocked by missing optional deps such as `orjson`; remote container is the authoritative verification environment for this path. + +## C23: Repeated EAGLE/bigram requests can expose 65-token tiny suffixes on 64-token pages + +Finding: + +- A repeated request can still show `extend_lens=[65]` even when the logical + prompt is identical and page flooring is only supposed to lose at most + `page_size - 1` tokens. +- The extra token comes from EAGLE/bigram key semantics plus SGLang's existing + one-token prefix-cache holdback: + - request sequence length in the observed log: `seq_len=40385`; + - scheduler prefix match caps token keys at `input_len - 1`; + - EAGLE radix keys are bigrams, so `convert_to_bigram_key()` reduces that key + length by one more token; + - effective matchable key length becomes `40385 - 2 = 40383`; + - page-flooring to 64-token pages gives `40320`; + - forward therefore sees `40385 - 40320 = 65` extend tokens. + +Implication: + +- For non-EAGLE page-granular reuse, the repeated-request suffix can be up to + `1 + (page_size - 1)` tokens because the last token is held back. +- For EAGLE/bigram reuse, it can be up to `2 + (page_size - 1)` tokens. With + `page_size=64`, the maximum tiny suffix is therefore `65`, matching the + observed hang log. +- This does not mean the user prompt changed; it is an artifact of the current + cache key contract. + +## C24: Tiny page-count cache-hit suffixes still enter NSA in-seq CP split + +Finding: + +- Latest remote log `/mnt/beegfs/cjy/sglang_cp_hicache_20260529_140440.log` + shows C22 is active: the repeated cache-hit request with + `prefix_lens=[40320]`, `extend_lens=[65]`, `page_size=64`, `cp_size=8` + skips both MLA and index async prefetchers: + `create_skip reason=extend_below_min ... min_async_extend_tokens=512`, + followed by `create_result ... has_mla=False has_index=False`. +- The same request then hangs before any `forward_probe` / `forward_layer` + line for layer 2; the next relevant lines are detokenizer heartbeat timeout. +- The earlier log `/mnt/beegfs/cjy/sglang_cp_hicache_20260529_134616.log` + showed the same 65-token repeated suffix hanging when async prefetchers were + created. Therefore async prefetch creation is not the root cause. +- `can_cp_split()` still allows NSA in-seq CP split whenever + `sum(extend_seq_lens_cpu) >= cp_size`. For the observed repeated request, + page-aligned physical coverage is only `ceil(65 / 64) = 2` pages, far below + `cp_size=8`. +- `build_page_aligned_in_seq_split_list(total_len=128, extend_len=65, + cp_size=8)` produces a degenerate split where only two page units are real and + most CP lanes get zero valid suffix tokens. This is the common condition + across both hang logs. + +Correction: + +- Async prefetch gating alone is insufficient. CP shared-KV page-aligned + in-seq split itself should require at least one physical suffix page per CP + rank: `ceil(extend_len / page_size) >= cp_size`. +- Below that threshold, run the request without NSA in-seq CP split rather than + generating mostly-zero CP segments. This keeps the cache page contract + intact while avoiding the fragile zero-lane distributed path. +- This must not disable target current reuse. With no async prefetcher, the + backend should still use synchronous prefix materialize + current splice for + cache-hit suffixes, and full current reuse for cache-miss/current-only + requests. + +Implementation: + +- `can_cp_split()` now skips CP shared-KV in-seq split when a page-aligned + cache-hit suffix has fewer padded suffix pages than CP lanes. For example + `extend_len=65,page_size=64,cp_size=8` has two physical suffix pages and no + longer creates mostly-zero CP segments. +- `can_reuse_current_extend_kv()` now treats tail-padded `out_cache_loc` as + eligible when it still contains all valid current suffix rows. The backend + slices partial-current `current_kv_cache/current_locs` down to the valid + suffix length before synchronous prefix compose, so padding rows do not become + current page-slot data. +- Current-only/full-current reuse remains unsliced in the existing path because + padded current-only CP requests already rely on the padded query/cache shape. + +Verified C24: + +- Remote container `/sgl-workspace/sglang-tai` on `g0034`: + - `python -m py_compile` for `utils.py`, `cp_shared_kv_runtime.py`, and + `nsa_backend.py` passed. + - Red tests first failed on the old code: + `test_can_cp_split_skips_cp_when_radix_hit_suffix_has_too_few_pages`, + `test_can_cp_split_skips_cp_when_page_units_do_not_cover_all_lanes`, and + padded-current-reuse coverage. + - After the fix, the targeted C24 tests passed: `4 passed`. + - Related unit suite passed: + `test_nsa_cp_utils.py`, `test_cp_shared_kv_runtime.py`, + `test_cp_shared_kv_layout.py` => `128 passed, 5 warnings, 2 subtests passed`. + +## C25: 2026-05-29 remote "hang" snapshot is idle prefill plus L1 owner-lane eviction churn + +Finding: + +- Latest remote log `/mnt/beegfs/cjy/sglang_cp_hicache_20260529_143924.log` + did not show a prefill scheduler deadlock at the time of inspection: + - `/health` responded `200 OK`; + - metrics reported `num_running_reqs=0`, `num_queue_reqs=0`, + `num_prefill_prealloc_queue_reqs=0`; + - detokenizer `py-spy` stack was idle in `recv_pyobj`; + - scheduler ranks were in the normal idle receive path: + CP0 sampling at `recv_requests()` nonblocking ZMQ receive and CP1-7 + sampling at `broadcast_pyobj()` from `recv_requests()`. +- The last completed user POST in that log was at `2026-05-29 14:52:08`. + After that, the log only showed periodic health/internal tiny one-token + prefill batches, so the prefill process itself had no visible user request + pending. +- The performance collapse immediately before that idle period coincided with + repeated L1/device owner-lane eviction and load-back churn, not host-L2 + eviction: + - no `deterministic CP host eviction before write`, + `_evict_host_for_physical_slots`, or host-full fallback appeared in the + log after startup; + - `HiCache-load` repeatedly had low per-owner availability and evicted device + victims before CP load-back, e.g. node `873` loaded `101632` host tokens + but first evicted victims `[627,604]` for `102016` device tokens; + - generic device eviction then also evicted large backed nodes for much + smaller capacity deficits, e.g. `num_tokens=10752` evicted `79488`, and + `num_tokens=19968` evicted `91200`. + +Implication: + +- The current evidence points to L1/device cache thrashing under CP owner-lane + constraints once the working set is saturated. Host cache may be nearly full + externally, but this log's hot path is not host eviction; it is repeated + host-hit load-back requiring device-owner-lane free pages, followed by + coarse whole-node device eviction. +- The existing CP load-back owner-lane planner avoids random lane mismatches, + but its eviction granularity is still whole radix nodes. When leaves are + large, a small lane deficit can evict a much larger cached prefix, destroying + future L1 hits and increasing H2D load-back work. +- The generic compute-owner allocation eviction path still calls + `tree_cache.evict(EvictParams(... owner_lane_deficits=...))`, but + `owner_lane_deficits` is not consumed by `HiRadixCache.evict()`. It therefore + remains a coarse total-token eviction path and can over-evict when only a few + owner lanes are short. + +Next correction target: + +- Do not treat this snapshot as an async-prefetch hang. The next fix should + reduce saturated-L1 owner-lane churn: + 1. add an owner-lane-aware device eviction path for compute-owner allocation, + reusing the CP load-back victim selection logic where possible; + 2. cap or skip load-back when the required owner-lane eviction cost is far + larger than the host-hit benefit, unless the cache-hit prefix is large + enough to amortize the churn; + 3. longer term, split/schedule radix leaves at page-aligned boundaries before + async backup/load so eviction can free page-sized or small-node units + instead of whole long-request leaves. + +## C26: 2026-05-29 hang is a scheduler collective-order divergence, not only idle L1 churn + +Finding: + +- A later remote stack sample on `g0034` for + `/mnt/beegfs/cjy/sglang_cp_hicache_20260529_143924.log` showed a real + distributed scheduler hang: + - CP0 / CP6 were in `recv_requests()->broadcast_pyobj()`; + - CP1 / CP2 / CP5 / CP7 were in + `scheduler_dp_attn_mixin.all_gather()` from + `prepare_mlp_sync_batch_raw()->maybe_prepare_mlp_sync_batch()`; + - CP3 was in idle memory check and CP4 in the prefill event-loop boundary; + - detokenizer was idle in `recv_pyobj`. +- This means ranks disagreed on the next collective: some entered the DP/attn + MLP-sync `all_gather`, while others entered scheduler request broadcast. + That is a collective-order divergence and can deadlock even when queues and + metrics look empty. + +Immediate debugging target: + +- Read `event_loop_normal_disagg_prefill`, + `get_next_disagg_prefill_batch_to_run`, and + `scheduler_dp_attn_mixin.maybe_prepare_mlp_sync_batch` to identify how one CP + rank can decide it has a batch requiring MLP sync while another decides it is + idle and returns to `recv_requests`. +- The earlier owner-lane eviction churn remains a performance issue, but it is + not sufficient to explain this collective-order mismatch. + +Correction C26: + +- Added a single-DP idle fast path in `prepare_mlp_sync_batch_raw`: when + `local_batch is None`, `dp_size == 1`, and `SGLANG_SCHEDULER_SKIP_ALL_GATHER` + is not enabled, return `None` without building an idle batch or entering the + MLP-sync all-gather. +- Rationale: CP in-seq split runs inside one DP scheduling domain in the current + deployment. All TP/CP ranks should have identical batch presence. An idle + all-gather only exchanges zero metadata; after tiny internal health requests + it can overlap with the next request broadcast and produce the observed + split: some ranks in Gloo `broadcast_pyobj`, others in NCCL/CPU MLP-sync + `all_gather`. +- This does not remove MLP sync for real batches or for multi-DP cases where a + rank may need an idle batch because another DP worker has work. + +## C27: 2026-05-29 request failure evidence is not in the prefill log + +Finding: + +- In the stopped run using `/mnt/beegfs/cjy/sglang_cp_hicache_20260529_152501.log`, + the prefill log does not contain the request-failure traceback: + - no `Traceback`, `RuntimeError`, scheduler-dead, health-failure, OOM, or + Python exception lines after startup; + - HTTP status counts in the prefill log are all `200`; + - the last visible prefill log line is a successful health check at + `2026-05-29 15:45:43`. +- All prefill/decode Python workers were already stopped/zombie when inspected. + Decode on `g0035`/`g0036` had no persisted stdout/stderr for this run; the + process output was attached to a TTY rather than a log file. Therefore the + concrete decode/router-side request failure cannot be reconstructed from the + available files. +- Decode workers on both decode hosts show zombie creation around + `2026-05-29 15:27:26-15:27:39`, while the prefill log continued accepting + successful `/v1/chat/completions` requests until `15:43:39` and health checks + until `15:45:43`. This strongly suggests the user-visible request failure is + downstream of decode/router availability or external process stop, not a + prefill-side HTTP 4xx/5xx captured in the HiCache log. + +Implication: + +- Do not infer a CP HiCache correctness failure from this prefill log alone. + The available evidence is an observability gap: decode/router failure output + was not persisted. +- For the next ETE run, decode/router stdout must be persisted with `tee` or an + equivalent log sink; otherwise a stopped decode process leaves only zombie + metadata and no Python traceback. + +Next correction target: + +- Continue code fixes that are already evidenced in prefill logs (owner-lane + device eviction churn and excessive tiny/internal-request CP fallback logging), + but treat this specific request failure as unrooted until decode/router logs + are captured. + +C25 correction update: + +- Implemented owner-lane-aware `HiRadixCache.evict()` handling when + `EvictParams.owner_lane_deficits` is set. The eviction path now plans device + victims using owner-lane contribution, so a low-priority non-contributing + leaf is no longer evicted before a higher-priority leaf that actually frees + the deficit lane. +- Reduced `_evict_for_compute_owner_lanes()` requested eviction from + `deficit_pages * page_size * cp_size` to `deficit_pages * page_size` because + the deficit vector is already counted in owner-lane pages. This removes the + previous full-CP-stripe over-eviction for a single-lane deficit. +- Remote verification on `g0034` container: + - RED before fix: + `test_owner_lane_evict_params_choose_deficit_contributing_victim` evicted + the non-contributing `[8,9,10,11]` page instead of the deficit owner page. + - GREEN after fix: + targeted test passed. + - Related suite passed: + `test_cp_hicache_load_back_owner_lanes.py` => `9 passed`. + +## C28: 2026-05-29 accept length collapse means EAGLE draft rejection, not a HiCache hit-rate counter + +Finding: + +- Decode `accept len` is reported from `spec_num_accepted_tokens / spec_num_forward_ct`. + `update_spec_metrics()` adds `num_accepted_tokens + batch_size`, while + `_resolve_spec_overlap_token_ids()` defines `num_accepted_tokens` as + `sum(accept_lens) - len(batch.reqs)`. Therefore a logged accept length near + `1.0` means EAGLE accepted zero draft tokens and only produced the mandatory + target token per verify step. +- The current remote decode processes are running with EAGLE, but not with + `SGLANG_CP_DRAFT_SHARED_KV=1` in their environment. Prefill has + `SGLANG_CP_DRAFT_SHARED_KV=1`, CP shared KV, HiCache, current reuse, and MLA + prefetch enabled. Decode still includes draft KV buffers through its draft + worker, so this is not alone proof of missing transfer, but it is an important + deployment asymmetry to keep visible. +- Current live metrics on the new run already show low/uneven speculative + acceptance on some decode DP ranks (for example `spec_accept_length` around + `1.05` on DP4/DP9 and higher on DP8/DP13), while several ranks are still at + zero because they have not reported a recent decode window. + +Implication: + +- The accept-length symptom should be debugged as target/draft state divergence + or sampling/draft-quality degradation on decode, not as a direct HiCache host + capacity issue. The key invariant to verify is: after disaggregated prefill + and any cache-hit/current-reuse/load-back path, target KV, draft KV, sequence + lengths, and draft hidden state used by EAGLE are still aligned for the same + request positions. +- Next evidence should be low-frequency decode-side EAGLE acceptance diagnostics + keyed by request/cache-hit status and whether draft KV buffers were received; + per-layer prefill logs are not sufficient to prove why decode rejects drafts. + +### C29. 2026-05-30 live run evidence: accept-len collapse is decode-side EAGLE rejection, not missing CP current reuse + +Evidence from `/mnt/beegfs/cjy/log/sglang_cp_hicache_20260529_160555.log`, `decode0_20260529_160608.log`, and `decode1_20260529_160611.log`: + +- Decode-side `accept len: 1.00` recurs across both decode nodes. Metrics at 2026-05-30 00:16 CST showed examples: DP4=1.0 on g0035 and DP9=1.0 on g0036. This means speculative draft tokens are rejected; the accepted token is only the mandatory target token. +- The low-accept events are not limited to requests with active disaggregation transfer. Log parsing showed many `accept len <= 1.001` events with `#transfer-req: 0`, so the symptom persists after transfer/preallocation, not only during the transfer window. +- Prefill CP shared KV current reuse is active in the same run. Counts: `forward_partial_current_reuse=2032`, with `used_prefetch=True=1104` and `used_prefetch=False=928`. Small-prefix cases below the async-prefetch threshold still use synchronous current compose, e.g. `reason=missing_prefetcher prefix_lens=[320] ... forward_partial_current_reuse used_prefetch=False`. +- Therefore `prefix_below_min` disables async prefix prefetcher creation only; it does not disable partial/full current reuse. It should not by itself explain accept length collapsing to 1. +- Decode startup remains asymmetric: g0035/node-rank0 command lacks `--speculative-attention-mode prefill`, while g0036/node-rank1 includes it. Both decode launch environments only expose `SGLANG_DISAGGREGATION_ALL_CP_RANKS_TRANSFER=1` and `SGLANG_DISAGGREGATION_QUEUE_SIZE=8` among CP/disagg flags. + +Current hypothesis to test next: accept-len collapse comes from decode-side EAGLE draft/target state divergence or decode launch/config asymmetry, not from CP shared KV partial current reuse being disabled. Add low-frequency decode-side EAGLE diagnostics around accept-lens / draft KV availability / transferred request state before changing prefill cache logic again. + +C29 correction note: + +- `--speculative-attention-mode prefill` is the default in `ServerArgs`; therefore g0035 omitting the explicit flag and g0036 spelling it out is not enough to prove a runtime mismatch. Keep checking actual parsed server args if needed, but do not treat the command-line spelling difference alone as root cause. + +### C30. 2026-05-30 router failures are downstream of prefill one-page cache accounting failure + +Evidence from `/mnt/beegfs/cjy/log/sglang_cp_hicache_20260529_160555.log` and router errors around `2026-05-29 16:19:08`: + +- Router `Prefill server failed (CRITICAL)` errors align with the prefill scheduler exiting, not with an independent router/decode routing bug. +- The direct prefill failure is the idle memory checker: + `token_to_kv_pool_allocator memory leak detected! self.max_total_num_tokens=833024, available_size=16384, evictable_size=816576, protected_size=0, session_held=0`. +- The accounting gap is exactly one page: `833024 - (16384 + 816576 + 0 + 0) = 64` tokens, matching the configured page size. +- Immediately before the failure, logs show tiny internal/prefill batches where scheduler allocation is page-granular (`#new-token: 64`) but CP shared KV metadata has valid lengths such as `seq_lens=[1]` or `seq_lens=[6]`, `real_pages=1`. +- Therefore the next root-cause target is a one-page lifecycle/accounting mismatch around tiny CP+EAGLE disagg prefill requests, especially the path `cache_unfinished_req(req)` before KV transfer and `release_kv_cache(req, tree_cache)` after transfer completion. + +Implication: + +- The router error is a symptom: once prefill rank workers die, decode/router cannot fetch prefill KV and reports failures. +- This is separate from the low EAGLE `accept len` symptom. Low accept length remains a decode-side draft rejection/state divergence issue; the one-page accounting failure explains prefill death and router errors. + +Current hypothesis to verify before fixing: + +- Disaggregated prefill inserts and locks an unfinished radix node before transfer. On transfer completion, the generic `release_kv_cache()` path calls `cache_finished_req()` again, which can re-insert/split a finished-key view. For tiny valid lengths under page granularity, this can leave exactly one physical page outside `available + evictable + protected` accounting. + +### C31. Repeated tiny EAGLE requests leak one CP HiCache page when exact matched tail is not page aligned + +New root-cause evidence: + +- A minimal remote repro with CP HiCache, `page_size=64`, EAGLE/bigram keys, and repeated `seq_len=6` requests shows the first request is balanced after release, but every repeated exact hit loses one 64-token page: + - after first release: `avail=8128 ev=64 prot=0 sum=8192`; + - after second release: `avail=8064 ev=64 prot=0 sum=8128`; + - after third release: `avail=8000 ev=64 prot=0 sum=8064`. +- This matches the live crash where `max_total_num_tokens - (available + evictable + protected + session_held) = 64`. +- Cause: CP HiCache uses valid-token radix keys but page-granular allocator accounting. For repeated tiny EAGLE requests, the full key is an exact non-page-aligned hit, e.g. `len(keys)=5` for `seq_len=6`. +- `cache_unfinished_req()` / `cache_finished_req()` call `_free_kv_indices_range(kv_indices, cache_protected_len, new_prefix_len)` for duplicate prefixes. In CP mode `_free_kv_indices_range` ceil-aligns `start` and floor-aligns `end` unless `include_partial_tail=True`. With `start=0,end=5,page_size=64`, the duplicate range floors to empty, so the newly allocated physical page is neither freed nor represented by a new radix node. The existing cached node accounts only its own page, leaving the new duplicate page leaked. + +Correction target: + +- When an insertion result is an exact full-key duplicate (`new_prefix_len == len(keys)`) under CP HiCache, free the duplicate request KV range with `include_partial_tail=True`. This frees the newly allocated physical tail page while preserving the existing cached node. Keep normal partial-split/new-node behavior page-floored to avoid freeing pages retained by radix nodes. + +C31 correction: + +- `RadixCache.cache_finished_req()` and `cache_unfinished_req()` now free duplicate KV ranges with `include_partial_tail=True` when CP HiCache insertion reports an exact full-key duplicate (`new_prefix_len == len(keys)`). +- This is intentionally limited to exact duplicate hits. Partial/new-node insertions still use the page-floored free path so a newly inserted radix tail page is not accidentally returned to the allocator. +- Added regression coverage for repeated tiny EAGLE exact hits: three repeated `seq_len=2` requests must keep `available + evictable + protected == allocator.size` after each release. + +Verified C31: + +- Local compile: `python -m py_compile python/sglang/srt/mem_cache/radix_cache.py test/registered/unit/mem_cache/test_cp_hicache_load_back_owner_lanes.py` passed. +- Remote `g0034` container `/sgl-workspace/sglang-tai`: + - targeted regression passed: `test_repeated_tiny_eagle_exact_hit_frees_duplicate_tail_page`; + - related owner-lane/CP HiCache suite passed: `test_cp_hicache_load_back_owner_lanes.py => 10 passed`. + +### C32. Decode accept-length collapse needs draft-transfer evidence, not more prefill cache changes + +Finding: + +- The latest persisted decode logs (`decode0_20260529_160608.log`, + `decode1_20260529_160611.log`) show low EAGLE acceptance after transfer has + already drained: parsing 970 decode-stat lines gives 855 lines with + `#transfer-req: 0`, average `accept len=1.342`, and 134 exact `accept len=1.00` + windows. Therefore the symptom is not limited to the prefill-to-decode + transfer polling window. +- Prefill did attach the CP draft HiCache mirror in the same run. All eight CP + ranks logged `[HiCache-draft] attached CP draft KV host pool`, and the host + budget was shared across target+draft. This rules out the simple hypothesis + that prefill had no draft host mirror. +- The logs do not prove whether the disaggregation transfer actually registered + and copied draft KV/state buffers, because `SGLANG_CP_DRAFT_SHARED_KV_DEBUG` + was not enabled; `draft_kv_buffer_count`, `draft_state_buffer_count`, and + per-room draft transfer records were therefore absent. +- Decode `accept len ~= 1` still means EAGLE draft rejection: the decode metric + adds the mandatory target token, so `1.00` is zero accepted draft tokens. + +Implication: + +- Do not continue changing CP partial current reuse or HiCache load-back based + solely on this accept-length symptom. The next evidence must come from + decode-side EAGLE state and disaggregation draft-transfer metadata: + transferred draft KV/state buffer counts, metadata hidden/topk presence, and + low-frequency per-request accept-lens after transfer. +- A safe next step is gated diagnostics, not an unconditional behavior change: + log startup/transfer draft-buffer availability once per rank and optionally + sample low-accept decode batches under an explicit debug env. + +C32 correction: + +- Added gated decode-side EAGLE acceptance diagnostics controlled by + `SGLANG_EAGLE_ACCEPT_DEBUG=1` and sampled by + `SGLANG_EAGLE_ACCEPT_DEBUG_INTERVAL` (default 128). The log is warning-level + only when a sampled batch has zero accepted draft tokens; otherwise it is + interval-sampled info. It records per-request accepted draft count, prompt + length, output length, cached token count, PD metadata tensor presence/shape, + and current `spec_info` topk/hidden shapes. +- This is diagnostic-only. It does not change EAGLE sampling, KV transfer, + cache load-back, or CP shared KV behavior. +- Synced `environ.py` and `scheduler_output_processor_mixin.py` to remote + `g0034:/mnt/beegfs/cjy/sglang-dev`; remote container py_compile passed. + +C32 correction update: + +- Extended `SGLANG_EAGLE_ACCEPT_DEBUG=1` to print one startup KV-manager audit per + rank on both prefill and decode. The audit includes target/draft KV buffer + counts and target/draft state buffer counts. This avoids enabling the much + noisier `SGLANG_CP_DRAFT_SHARED_KV_DEBUG` just to answer whether draft KV/state + transfer surfaces are registered. +- Synced `prefill.py` and `decode.py` to remote `g0034:/mnt/beegfs/cjy/sglang-dev`; + remote container py_compile passed for the updated files. diff --git a/python/sglang/srt/disaggregation/decode.py b/python/sglang/srt/disaggregation/decode.py index 9d03c645c..cfb9495b6 100644 --- a/python/sglang/srt/disaggregation/decode.py +++ b/python/sglang/srt/disaggregation/decode.py @@ -509,6 +509,22 @@ class DecodePreallocQueue: len(kv_args.state_data_ptrs), ) + if envs.SGLANG_EAGLE_ACCEPT_DEBUG.get() and self.spec_algorithm.is_eagle(): + logger.info( + "[EAGLE_ACCEPT_DEBUG] decode_kv_manager cp_rank=%s " + "target_kv_bufs=%s draft_kv_bufs=%s total_kv_bufs=%s " + "target_state_type=%s registered_state_bufs=%s " + "draft_state_type=%s draft_state_bufs=%s", + self.tp_rank, + target_kv_buffer_count, + draft_kv_buffer_count, + len(kv_args.kv_data_ptrs), + kv_args.state_type, + len(kv_args.state_data_ptrs), + kv_args.draft_state_type, + kv_args.draft_state_buffer_count, + ) + kv_args.ib_device = self.scheduler.server_args.disaggregation_ib_device kv_args.gpu_id = self.scheduler.gpu_id kv_manager_class = get_kv_class(self.transfer_backend, KVClassType.MANAGER) diff --git a/python/sglang/srt/disaggregation/prefill.py b/python/sglang/srt/disaggregation/prefill.py index a972c7d92..21bc2407b 100644 --- a/python/sglang/srt/disaggregation/prefill.py +++ b/python/sglang/srt/disaggregation/prefill.py @@ -350,6 +350,22 @@ class PrefillBootstrapQueue: len(kv_args.state_data_ptrs), ) + if envs.SGLANG_EAGLE_ACCEPT_DEBUG.get() and self.spec_algorithm.is_eagle(): + logger.info( + "[EAGLE_ACCEPT_DEBUG] prefill_kv_manager cp_rank=%s " + "target_kv_bufs=%s draft_kv_bufs=%s total_kv_bufs=%s " + "target_state_type=%s registered_state_bufs=%s " + "draft_state_type=%s draft_state_bufs=%s", + self.tp_rank, + target_kv_buffer_count, + draft_kv_buffer_count, + len(kv_args.kv_data_ptrs), + kv_args.state_type, + len(kv_args.state_data_ptrs), + kv_args.draft_state_type, + kv_args.draft_state_buffer_count, + ) + kv_manager_class = get_kv_class(self.transfer_backend, KVClassType.MANAGER) kv_manager = kv_manager_class( kv_args, diff --git a/python/sglang/srt/environ.py b/python/sglang/srt/environ.py index 84f787bca..434afe091 100644 --- a/python/sglang/srt/environ.py +++ b/python/sglang/srt/environ.py @@ -217,6 +217,8 @@ class Envs: SGLANG_CP_SHARED_KV_MLA_PREFETCH_MIN_EXTEND_TOKENS = EnvInt(-1) SGLANG_CP_DRAFT_SHARED_KV = EnvBool(False) SGLANG_CP_DRAFT_SHARED_KV_DEBUG = EnvBool(False) + SGLANG_EAGLE_ACCEPT_DEBUG = EnvBool(False) + SGLANG_EAGLE_ACCEPT_DEBUG_INTERVAL = EnvInt(128) SGLANG_DISABLE_TAI_BIGRAM = EnvBool(False) SGLANG_TEST_REQUEST_TIME_STATS = EnvBool(False) SGLANG_DISABLE_TP_MEMORY_INBALANCE_CHECK = EnvBool(False) diff --git a/python/sglang/srt/layers/attention/nsa/cp_shared_kv_runtime.py b/python/sglang/srt/layers/attention/nsa/cp_shared_kv_runtime.py index 28aeebb46..a3b972447 100644 --- a/python/sglang/srt/layers/attention/nsa/cp_shared_kv_runtime.py +++ b/python/sglang/srt/layers/attention/nsa/cp_shared_kv_runtime.py @@ -968,7 +968,11 @@ def can_reuse_current_extend_kv(forward_batch) -> bool: seq_len = int(seq_lens_cpu[0].item()) if extend_len <= 0 or seq_len < extend_len: return False - return int(out_cache_loc.numel()) == extend_len + # ForwardBatch pads tensors such as out_cache_loc at the tail for CUDA graph + # and CP alignment. The first extend_len rows still cover the valid current + # suffix, so padded batches remain eligible for current reuse as long as the + # valid suffix is fully present. + return int(out_cache_loc.numel()) >= extend_len def should_reuse_current_extend_kv(forward_batch) -> bool: diff --git a/python/sglang/srt/layers/attention/nsa/utils.py b/python/sglang/srt/layers/attention/nsa/utils.py index 08bc31cb4..29e33d703 100644 --- a/python/sglang/srt/layers/attention/nsa/utils.py +++ b/python/sglang/srt/layers/attention/nsa/utils.py @@ -403,6 +403,51 @@ def should_use_replicated_compute_for_short_radix_hit( return False +def should_skip_cp_shared_kv_cp_split_for_short_page_extent( + forward_batch: "ForwardBatch", + cp_size: int, +) -> bool: + """Avoid in-seq CP split when a shared-KV suffix has too few pages. + + CP shared KV is page-owned. A cache-hit suffix with fewer physical pages + than CP lanes creates mostly-zero in-seq segments; the distributed NSA path + has repeatedly shown hangs on that shape. Keep the page cache contract by + running those tiny suffixes without NSA in-seq CP split instead of falling + back to token-balanced page-splitting. + """ + + if ( + forward_batch is None + or not getattr(forward_batch, "uses_cp_shared_kv", False) + or cp_size <= 1 + ): + return False + + extend_seq_lens_cpu = getattr(forward_batch, "extend_seq_lens_cpu", None) + extend_prefix_lens_cpu = getattr(forward_batch, "extend_prefix_lens_cpu", None) + token_to_kv_pool = getattr(forward_batch, "token_to_kv_pool", None) + page_size = int(getattr(token_to_kv_pool, "page_size", 0) or 0) + if ( + extend_seq_lens_cpu is None + or len(extend_seq_lens_cpu) != 1 + or extend_prefix_lens_cpu is None + or len(extend_prefix_lens_cpu) != 1 + or page_size <= 0 + ): + return False + + extend_len = int(extend_seq_lens_cpu[0]) + if extend_len <= 0: + return False + + prefix_len = int(extend_prefix_lens_cpu[0]) + if prefix_len <= 0 or prefix_len % page_size != 0: + return False + + padded_pages = ceil_div(extend_len, page_size) + return padded_pages < cp_size + + def can_cp_split(seq_len: int, cp_size: int, use_nsa: bool, forward_batch): if is_nsa_prefill_cp_round_robin_split(): cur_cp_seq_len = seq_len // cp_size @@ -415,6 +460,10 @@ def can_cp_split(seq_len: int, cp_size: int, use_nsa: bool, forward_batch): # the seq data needs to be divided and recombined at twice the size of cp_size. if should_use_replicated_compute_for_short_radix_hit(forward_batch, cp_size): return False + if should_skip_cp_shared_kv_cp_split_for_short_page_extent( + forward_batch, cp_size + ): + return False cur_cp_seq_len = seq_len // (cp_size * 2) if ( cur_cp_seq_len != 0 diff --git a/python/sglang/srt/layers/attention/nsa_backend.py b/python/sglang/srt/layers/attention/nsa_backend.py index bbdf2a0b8..7963ed774 100644 --- a/python/sglang/srt/layers/attention/nsa_backend.py +++ b/python/sglang/srt/layers/attention/nsa_backend.py @@ -1782,6 +1782,7 @@ class NativeSparseAttnBackend( ) if can_reuse_current_kv: current_kv_cache = _cat([k, k_rope], dim=-1) + current_locs_for_reuse = forward_batch.out_cache_loc logical_page_table_1 = page_table_1 current_remap_page_size, current_remap_logical_page_capacity = ( current_loc_remap_fast_path_args(forward_batch) @@ -1809,13 +1810,30 @@ class NativeSparseAttnBackend( "MLA current reuse cp_rank=%s layer=%s current_locs=%s remapped=%s kv_ck=%s rope_ck=%s", forward_batch.cp_shared_kv_layout.cp_rank, layer.layer_id, - tensor_debug_summary(forward_batch.out_cache_loc), + tensor_debug_summary(current_locs_for_reuse), tensor_debug_summary(page_table_1), tensor_debug_checksum(k), tensor_debug_checksum(k_rope), ) kv_cache = current_kv_cache else: + extend_lens_cpu_for_current = getattr( + forward_batch, "extend_seq_lens_cpu", None + ) + if ( + extend_lens_cpu_for_current is not None + and len(extend_lens_cpu_for_current) == 1 + ): + valid_current_rows = int(extend_lens_cpu_for_current[0]) + if ( + valid_current_rows > 0 + and valid_current_rows < int(current_locs_for_reuse.numel()) + and valid_current_rows <= int(current_kv_cache.shape[0]) + ): + current_kv_cache = current_kv_cache[:valid_current_rows] + current_locs_for_reuse = current_locs_for_reuse[ + :valid_current_rows + ] prefetched_kv = None if mla_prefetcher is not None: prefetched_kv = mla_prefetcher.consume_prefix_with_current( @@ -1823,7 +1841,7 @@ class NativeSparseAttnBackend( kv_cache=kv_cache, logical_locs=logical_page_table_1, current_kv_cache=current_kv_cache, - current_locs=forward_batch.out_cache_loc, + current_locs=current_locs_for_reuse, current_remap_page_size=current_remap_page_size, current_remap_logical_page_capacity=current_remap_logical_page_capacity, ) @@ -1869,7 +1887,7 @@ class NativeSparseAttnBackend( f"extend_lens={extend_lens} " f"current_rows={int(current_kv_cache.shape[0])} " f"logical_page_table_shape={tuple(logical_page_table_1.shape)} " - f"current_locs_shape={tuple(forward_batch.out_cache_loc.shape)} " + f"current_locs_shape={tuple(current_locs_for_reuse.shape)} " f"page_size={page_size}" ) prefix_pages = int(prefix_lens_cpu[0]) // page_size @@ -1885,7 +1903,7 @@ class NativeSparseAttnBackend( kv_cache=kv_cache, logical_locs=logical_page_table_1, current_kv_cache=current_kv_cache, - current_locs=forward_batch.out_cache_loc, + current_locs=current_locs_for_reuse, slot_remap=slot_remap, layout=forward_batch.cp_shared_kv_layout, page_size=page_size, @@ -1946,7 +1964,7 @@ class NativeSparseAttnBackend( "MLA partial current reuse cp_rank=%s layer=%s current_locs=%s remapped=%s kv_ck=%s rope_ck=%s", forward_batch.cp_shared_kv_layout.cp_rank, layer.layer_id, - tensor_debug_summary(forward_batch.out_cache_loc), + tensor_debug_summary(current_locs_for_reuse), tensor_debug_summary(page_table_1), tensor_debug_checksum(k), tensor_debug_checksum(k_rope), diff --git a/python/sglang/srt/managers/scheduler_dp_attn_mixin.py b/python/sglang/srt/managers/scheduler_dp_attn_mixin.py index 9733cd3e1..6a1a4729c 100644 --- a/python/sglang/srt/managers/scheduler_dp_attn_mixin.py +++ b/python/sglang/srt/managers/scheduler_dp_attn_mixin.py @@ -154,6 +154,17 @@ def prepare_mlp_sync_batch_raw( disable_overlap_schedule: bool, offload_tags: set[str], ): + skip_all_gather = envs.SGLANG_SCHEDULER_SKIP_ALL_GATHER.get() + + # With a single DP scheduling domain, all TP/CP ranks are expected to have + # the same local batch presence. If there is no local batch, an MLP-sync + # collective would only gather all-zero idle metadata. More importantly, + # in CP disaggregated prefill this idle NCCL/Gloo sync can race with the + # next request-broadcast collective after tiny health/internal requests and + # leave some ranks in recv broadcast while others enter this all-gather. + if local_batch is None and dp_size == 1 and not skip_all_gather: + return None + # Check if other DP workers have running batches if local_batch is None or local_batch.forward_mode.is_prebuilt(): num_tokens = 0 @@ -176,7 +187,6 @@ def prepare_mlp_sync_batch_raw( or num_tokens_for_logprob == local_batch.batch_size() ) - skip_all_gather = envs.SGLANG_SCHEDULER_SKIP_ALL_GATHER.get() can_cuda_graph = ( local_batch is None or local_batch.forward_mode.is_decode_or_idle() diff --git a/python/sglang/srt/managers/scheduler_output_processor_mixin.py b/python/sglang/srt/managers/scheduler_output_processor_mixin.py index 15d6c3cf5..db9e8077b 100644 --- a/python/sglang/srt/managers/scheduler_output_processor_mixin.py +++ b/python/sglang/srt/managers/scheduler_output_processor_mixin.py @@ -51,6 +51,82 @@ class SchedulerOutputProcessorMixin: storage_backend_type = type(storage_backend).__name__ return storage_backend_type + def _maybe_log_eagle_accept_debug( + self: Scheduler, + batch: ScheduleBatch, + result: GenerationBatchResult, + ) -> None: + """Low-frequency EAGLE accept diagnostics for PD/debug runs. + + This is intentionally gated. The hot path normally only reports an + aggregate accept length; when accept length collapses we need enough + per-request state to distinguish draft rejection from transfer/cache + lifecycle issues without adding per-token or per-layer logging. + """ + + if not envs.SGLANG_EAGLE_ACCEPT_DEBUG.get(): + return + if batch.spec_algorithm.is_none(): + return + + accept_per_req = getattr(result, "accept_length_per_req_cpu", None) + if accept_per_req is None: + return + + counter = getattr(self, "_eagle_accept_debug_counter", 0) + 1 + self._eagle_accept_debug_counter = counter + + interval = envs.SGLANG_EAGLE_ACCEPT_DEBUG_INTERVAL.get() + interval = max(1, int(interval or 1)) + has_zero_accept = any(int(x) <= 0 for x in accept_per_req) + if not has_zero_accept and counter % interval != 0: + return + + spec_info = getattr(batch, "spec_info", None) + + def _shape(obj): + shape = getattr(obj, "shape", None) + return tuple(shape) if shape is not None else None + + samples = [] + for i, req in enumerate(batch.reqs[: min(4, len(batch.reqs))]): + if not has_zero_accept and counter % interval != 0: + break + if has_zero_accept and int(accept_per_req[i]) > 0: + continue + samples.append( + { + "rid": str(getattr(req, "rid", ""))[:8], + "accepted_draft": int(accept_per_req[i]), + "origin": len(getattr(req, "origin_input_ids", []) or []), + "output": len(getattr(req, "output_ids", []) or []), + "cached": int(getattr(req, "cached_tokens", 0) or 0), + "spec_verify_ct": int(getattr(req, "spec_verify_ct", 0) or 0), + "has_pd_hidden": getattr(req, "hidden_states_tensor", None) + is not None, + "pd_hidden_shape": _shape(getattr(req, "hidden_states_tensor", None)), + "pd_topk_shape": _shape(getattr(req, "output_topk_p", None)), + } + ) + + log_fn = logger.warning if has_zero_accept else logger.info + log_fn( + "[EAGLE_ACCEPT_DEBUG] step=%d mode=%s bs=%d avg_accept=%.3f " + "min_accept=%d max_accept=%d num_zero=%d spec_v2=%s " + "spec_topk_shape=%s spec_hidden_shape=%s samples=%s", + counter, + getattr(self, "disaggregation_mode", None), + len(batch.reqs), + sum(int(x) for x in accept_per_req) / max(1, len(accept_per_req)), + min(int(x) for x in accept_per_req), + max(int(x) for x in accept_per_req), + sum(1 for x in accept_per_req if int(x) <= 0), + getattr(batch, "is_spec_v2", None), + _shape(getattr(spec_info, "topk_p", None)), + _shape(getattr(spec_info, "hidden_states", None)), + samples, + ) + def _get_cached_tokens_details(self, req: Req) -> Optional[dict]: """Get detailed cache breakdown for a request, if available. @@ -410,6 +486,7 @@ class SchedulerOutputProcessorMixin: self.num_generated_tokens += len(batch.reqs) if not batch.spec_algorithm.is_none(): self.update_spec_metrics(batch.batch_size(), result.num_accepted_tokens) + self._maybe_log_eagle_accept_debug(batch, result) if self.enable_metrics: self.metrics_collector.increment_decode_cuda_graph_pass( value=can_run_cuda_graph diff --git a/python/sglang/srt/mem_cache/common.py b/python/sglang/srt/mem_cache/common.py index 3d551a382..b5bb8d6ca 100644 --- a/python/sglang/srt/mem_cache/common.py +++ b/python/sglang/srt/mem_cache/common.py @@ -349,10 +349,11 @@ def _evict_for_compute_owner_lanes( ) return - evict_tokens = max( - allocator.page_size, - deficit_pages * allocator.page_size * int(getattr(allocator, "cp_size", 1)), - ) + # ``deficits`` is already expressed in physical owner-lane pages. The + # old ``* cp_size`` multiplier turned a one-page lane deficit into one + # full CP stripe of eviction and caused large L1 churn under HiCache + # load-back pressure. + evict_tokens = max(allocator.page_size, deficit_pages * allocator.page_size) before_available = allocator.available_size() logger.info( "[MemCache-evict] _evict_for_compute_owner_lanes attempt=%d: deficit_pages=%d evict_tokens=%d before_available=%d evictable_size=%d", diff --git a/python/sglang/srt/mem_cache/hiradix_cache.py b/python/sglang/srt/mem_cache/hiradix_cache.py index 46328d8b4..53adf4313 100644 --- a/python/sglang/srt/mem_cache/hiradix_cache.py +++ b/python/sglang/srt/mem_cache/hiradix_cache.py @@ -1238,6 +1238,105 @@ class HiRadixCache(RadixCache): ) return refreshed + def _evict_cp_owner_lane_deficit_nodes( + self, params: EvictParams + ) -> EvictResult: + deficits = [max(0, int(v)) for v in (params.owner_lane_deficits or [])] + if not deficits or all(v <= 0 for v in deficits): + return EvictResult() + + plan = CpLoadBackPlan( + page_owners=[], + required_by_owner=[], + available_by_owner=[], + deficit_by_owner=deficits, + host_hit_len=0, + ) + eviction_plan = self._plan_cp_load_back_owner_lane_evictions(plan) + if len(eviction_plan.victims) == 0: + logger.info( + "[HiCache-evict] owner-lane evict found no contributing victims: " + "num_tokens=%d deficits=%s evictable_size=%d available_size=%d", + params.num_tokens, + deficits, + self.evictable_size_, + self.token_to_kv_pool_allocator.available_size(), + ) + return EvictResult() + + num_evicted = 0 + num_locked_skipped = 0 + write_back_nodes: List[TreeNode] = [] + for victim in eviction_plan.victims: + if victim.lock_ref > 0: + num_locked_skipped += 1 + continue + + if victim.pin_expiry > 0 and time.monotonic() > victim.pin_expiry: + self._clear_pin(victim) + + if not self._cp_device_leaf_is_load_back_victim(victim): + logger.info( + "[HiCache-evict] owner-lane evict victim no longer evictable: " + "victim_id=%s deficits=%s", + getattr(victim, "id", None), + deficits, + ) + continue + + if self._is_pinned(victim): + if self._node_backuped(victim): + num_evicted += self._evict_backuped(victim) + else: + written = self.write_backup(victim, write_back=True) + if written > 0: + write_back_nodes.append(victim) + else: + self._clear_pin(victim) + continue + + if not self._node_backuped(victim): + if self.cache_controller.write_policy == "write_back": + written = self.write_backup(victim, write_back=True) + if written > 0: + write_back_nodes.append(victim) + continue + num_evicted += self._evict_regular(victim) + else: + num_evicted += self._evict_backuped(victim) + + for child in victim.parent.children.values(): + if child in write_back_nodes: + continue + if not child.evicted: + break + else: + self._update_leaf_status(victim.parent) + + if write_back_nodes: + self.writing_check(write_back=True) + for victim in write_back_nodes: + if self._node_backuped(victim): + num_evicted += self._evict_backuped(victim) + + logger.info( + "[HiCache-evict] owner-lane evict END: requested_tokens=%d " + "deficits=%s victims=%s planned_freed_by_owner=%s " + "remaining_deficit_by_owner=%s num_evicted=%d " + "num_locked_skipped=%d evictable_size_after=%d " + "available_size_after=%d", + params.num_tokens, + deficits, + [getattr(node, "id", None) for node in eviction_plan.victims], + eviction_plan.planned_freed_by_owner, + eviction_plan.remaining_deficit_by_owner, + num_evicted, + num_locked_skipped, + self.evictable_size_, + self.token_to_kv_pool_allocator.available_size(), + ) + return EvictResult(num_tokens_evicted=num_evicted) + def _cp_build_write_admission( self, device_indices: torch.Tensor, *, node_id: int, phase: str ) -> CpWriteAdmission: @@ -2575,6 +2674,13 @@ class HiRadixCache(RadixCache): def evict(self, params: EvictParams) -> EvictResult: start_time = time.perf_counter() num_tokens = params.num_tokens + if params.owner_lane_deficits is not None and any( + int(v) > 0 for v in params.owner_lane_deficits + ): + result = self._evict_cp_owner_lane_deficit_nodes(params) + self.update_eviction_metrics(result.num_tokens_evicted, start_time) + return result + leaves = list(self.evictable_leaves) eviction_heap = [ (self.eviction_strategy.get_priority(node), node) for node in leaves diff --git a/python/sglang/srt/mem_cache/radix_cache.py b/python/sglang/srt/mem_cache/radix_cache.py index 66ab4da1a..f5f79e364 100644 --- a/python/sglang/srt/mem_cache/radix_cache.py +++ b/python/sglang/srt/mem_cache/radix_cache.py @@ -494,6 +494,28 @@ class RadixCache(BasePrefixCache): end = start self.token_to_kv_pool_allocator.free(kv_indices[start:end]) + def _should_free_cp_exact_match_tail( + self, *, start: int, end: int, key_len: int + ) -> bool: + """Return whether duplicate-prefix freeing must include a CP tail page. + + CP HiCache keeps radix keys at scheduler-visible valid lengths while the + allocator owns whole physical pages. For an exact duplicate hit whose + valid key ends inside a page (typical tiny EAGLE/bigram requests), the + duplicate request page is not retained by a new radix node. Flooring + ``end`` to the previous page boundary would therefore leak that newly + allocated page. Partial/new-node insertions still use the conservative + page-floored free path because their tail page is represented by the + radix node being inserted. + """ + + return ( + getattr(self, "_uses_cp_hicache", False) + and key_len > 0 + and end == key_len + and end > start + ) + def cache_finished_req(self, req: Req, is_insert: bool = True): """Cache request when it finishes.""" # In deterministic mode, disable finished request insertion to radix cache @@ -541,7 +563,14 @@ class RadixCache(BasePrefixCache): new_prefix_len = result.prefix_len # Free the duplicates that were already in the tree self._free_kv_indices_range( - kv_indices, req.cache_protected_len, new_prefix_len + kv_indices, + req.cache_protected_len, + new_prefix_len, + include_partial_tail=self._should_free_cp_exact_match_tail( + start=req.cache_protected_len, + end=new_prefix_len, + key_len=len(keys), + ), ) else: if prepared_cp_backup is not None and hasattr( @@ -612,7 +641,16 @@ class RadixCache(BasePrefixCache): req.cp_hicache_prepared_backup = None new_prefix_len = result.prefix_len - self._free_kv_indices_range(kv_indices, req.cache_protected_len, new_prefix_len) + self._free_kv_indices_range( + kv_indices, + req.cache_protected_len, + new_prefix_len, + include_partial_tail=self._should_free_cp_exact_match_tail( + start=req.cache_protected_len, + end=new_prefix_len, + key_len=len(keys), + ), + ) # The prefix indices could be updated, reuse it match_result = self.match_prefix(MatchPrefixParams(key=radix_key)) diff --git a/test/registered/unit/layers/test_nsa_cp_utils.py b/test/registered/unit/layers/test_nsa_cp_utils.py index a84795e96..1028f1e94 100644 --- a/test/registered/unit/layers/test_nsa_cp_utils.py +++ b/test/registered/unit/layers/test_nsa_cp_utils.py @@ -233,7 +233,32 @@ class TestNSAInSeqCPUtils(unittest.TestCase): split_list, extend_prefix_len=54464, extend_len=256, page_size=64 ) - def test_can_cp_split_keeps_cp_for_short_radix_hit_suffix(self): + def test_can_cp_split_skips_cp_when_radix_hit_suffix_has_too_few_pages(self): + class Mode: + def is_context_parallel_extend(self): + return True + + forward_batch = SimpleNamespace( + uses_cp_shared_kv=True, + extend_seq_lens_cpu=[65], + extend_prefix_lens_cpu=[54464], + token_to_kv_pool=SimpleNamespace(page_size=64), + forward_mode=Mode(), + ) + + with ( + patch( + "sglang.srt.layers.attention.nsa.utils.is_nsa_prefill_cp_round_robin_split", + return_value=False, + ), + patch( + "sglang.srt.layers.attention.nsa.utils.is_nsa_enable_prefill_cp", + return_value=True, + ), + ): + self.assertFalse(can_cp_split(128, 8, True, forward_batch)) + + def test_can_cp_split_skips_cp_when_page_units_do_not_cover_all_lanes(self): class Mode: def is_context_parallel_extend(self): return True @@ -256,7 +281,7 @@ class TestNSAInSeqCPUtils(unittest.TestCase): return_value=True, ), ): - self.assertTrue(can_cp_split(256, 8, True, forward_batch)) + self.assertFalse(can_cp_split(256, 8, True, forward_batch)) def test_can_cp_split_keeps_cp_for_radix_hit_suffix_with_one_page_per_rank(self): class Mode: @@ -847,7 +872,7 @@ class TestNSAInSeqCPUtils(unittest.TestCase): "materialize/current merge when page-slot prefetch compose is unavailable.", ) self.assertIn( - "[CP_SHARED_KV_FAIL_FAST][mla_partial_current_prefetch]", + "[CP_SHARED_KV_FAIL_FAST][mla_partial_current_sync]", source, ) diff --git a/test/registered/unit/managers/test_scheduler_dp_attn_mixin.py b/test/registered/unit/managers/test_scheduler_dp_attn_mixin.py new file mode 100644 index 000000000..a7a38607f --- /dev/null +++ b/test/registered/unit/managers/test_scheduler_dp_attn_mixin.py @@ -0,0 +1,37 @@ +from unittest import TestCase +from unittest.mock import patch + +from sglang.srt.managers.scheduler_dp_attn_mixin import ( + MLPSyncBatchInfo, + prepare_mlp_sync_batch_raw, +) + + +class _FakeTPGroup: + device_group = object() + cpu_group = object() + device = "cuda" + + +class TestSchedulerDPAttnMixin(TestCase): + def test_single_dp_idle_batch_does_not_enter_mlp_sync_collective(self): + def fail_all_gather(self, *args, **kwargs): + raise AssertionError("idle single-DP scheduler must not all-gather") + + with patch.object(MLPSyncBatchInfo, "all_gather", fail_all_gather): + result = prepare_mlp_sync_batch_raw( + local_batch=None, + dp_size=1, + attn_tp_size=1, + attn_cp_size=8, + tp_group=_FakeTPGroup(), + get_idle_batch=lambda: (_ for _ in ()).throw( + AssertionError("idle single-DP scheduler must not build idle batch") + ), + disable_cuda_graph=False, + require_mlp_tp_gather=True, + disable_overlap_schedule=True, + offload_tags=set(), + ) + + self.assertIsNone(result) diff --git a/test/registered/unit/mem_cache/test_cp_hicache_load_back_owner_lanes.py b/test/registered/unit/mem_cache/test_cp_hicache_load_back_owner_lanes.py index 90286e72f..d94628afd 100644 --- a/test/registered/unit/mem_cache/test_cp_hicache_load_back_owner_lanes.py +++ b/test/registered/unit/mem_cache/test_cp_hicache_load_back_owner_lanes.py @@ -89,11 +89,17 @@ for _schema in ( raise from sglang.srt.mem_cache.allocator import CPSharedPagedTokenToKVPoolAllocator +import sglang.srt.mem_cache.common as mem_cache_common from sglang.srt.mem_cache.hiradix_cache import ( CpHiCacheNodeMetadata, HiRadixCache, ) -from sglang.srt.mem_cache.radix_cache import RadixKey, TreeNode, get_child_key +from sglang.srt.mem_cache.radix_cache import ( + RadixKey, + TreeNode, + _key_match_paged, + get_child_key, +) from sglang.test.ci.ci_register import register_cpu_ci from sglang.test.test_utils import CustomTestCase @@ -216,6 +222,57 @@ def _make_cache(allocator): return cache +class _TinyReqToTokenPool: + def __init__(self, size=8, max_context_len=32): + self.req_to_token = torch.zeros((size, max_context_len), dtype=torch.int64) + + def write(self, indices, values): + self.req_to_token[indices] = values + + def free(self, req): + req.req_pool_idx = None + + +def _make_tiny_eagle_req(cache, allocator, *, seq_len, req_pool_idx): + page_size = allocator.page_size + alloc_len = ((seq_len + page_size - 1) // page_size) * page_size + locs = allocator.alloc(alloc_len) + cache.req_to_token_pool.write( + (req_pool_idx, slice(0, seq_len)), locs[:seq_len] + ) + + req = types.SimpleNamespace( + req_pool_idx=req_pool_idx, + mamba_pool_idx=None, + fill_ids=list(range(seq_len)), + origin_input_ids=list(range(seq_len)), + output_ids=[], + extra_key=None, + cache_protected_len=0, + last_node=cache.root_node, + priority=0, + kv_committed_len=seq_len, + kv_allocated_len=seq_len, + kv_committed_freed=False, + kv_overallocated_freed=False, + cp_hicache_prepared_backup=None, + ) + + def pop_committed_kv_cache(): + assert not req.kv_committed_freed + req.kv_committed_freed = True + return req.kv_committed_len + + def pop_overallocated_kv_cache(): + assert not req.kv_overallocated_freed + req.kv_overallocated_freed = True + return req.kv_committed_len, req.kv_allocated_len + + req.pop_committed_kv_cache = pop_committed_kv_cache + req.pop_overallocated_kv_cache = pop_overallocated_kv_cache + return req + + class TestCpHiCacheLoadBackOwnerLanes(CustomTestCase): def test_load_back_plan_reports_owner_lane_vectors(self): allocator = _make_allocator() @@ -302,6 +359,46 @@ class TestCpHiCacheLoadBackOwnerLanes(CustomTestCase): self.assertEqual(target.value.tolist(), loaded.tolist()) self.assertIn(target.id, cache.ongoing_load_back) + def test_owner_lane_evict_params_choose_deficit_contributing_victim(self): + from sglang.srt.mem_cache.base_prefix_cache import EvictParams + + allocator = _make_allocator() + cache = _make_cache(allocator) + + non_contributing = _make_node( + 50, + 900, + [1], + value=torch.tensor([8, 9, 10, 11], dtype=torch.int64), + priority=0, + ) + contributing = _make_node( + 51, + 1000, + [0], + value=torch.tensor([4, 5, 6, 7], dtype=torch.int64), + priority=10, + ) + _attach_child(cache, cache.root_node, non_contributing) + _attach_child(cache, cache.root_node, contributing) + cache.evictable_leaves.update({non_contributing, contributing}) + cache.evictable_size_ = len(non_contributing.value) + len(contributing.value) + + result = cache.evict( + EvictParams( + num_tokens=allocator.page_size, + owner_lane_deficits=[1, 0, 0, 0], + ) + ) + + self.assertEqual(result.num_tokens_evicted, allocator.page_size) + self.assertEqual( + [indices.tolist() for indices in cache.cache_controller.evicted_device_indices], + [[4, 5, 6, 7]], + ) + self.assertIsNone(contributing.value) + self.assertIsNotNone(non_contributing.value) + def test_load_back_evicts_blocking_leaf_to_unlock_parent_owner_lane(self): allocator = _make_allocator() # Only owner lane 1 is initially available. The target needs owner lane @@ -384,6 +481,41 @@ class TestCpHiCacheLoadBackOwnerLanes(CustomTestCase): self.assertNotIn(target.id, cache.ongoing_load_back) self.assertEqual(cache.cache_controller.ack_load_queue, []) + def test_repeated_tiny_eagle_exact_hit_frees_duplicate_tail_page(self): + allocator = _make_allocator(page_size=4, cp_size=4) + cache = _make_cache(allocator) + cache.req_to_token_pool = _TinyReqToTokenPool() + cache.key_match_fn = functools.partial( + _key_match_paged, page_size=allocator.page_size + ) + cache.device = "cpu" + cache.is_eagle = True + cache.disable_finished_insert = False + cache.enable_storage = False + cache.write_through_threshold = 10**9 + + old_get_global_server_args = mem_cache_common.get_global_server_args + mem_cache_common.get_global_server_args = lambda: types.SimpleNamespace( + page_size=allocator.page_size, + speculative_algorithm="EAGLE", + ) + try: + for req_pool_idx in (1, 2, 3): + req = _make_tiny_eagle_req( + cache, allocator, seq_len=2, req_pool_idx=req_pool_idx + ) + cache.cache_unfinished_req(req) + mem_cache_common.release_kv_cache(req, cache) + + accounted = ( + allocator.available_size() + + cache.evictable_size() + + cache.protected_size() + ) + self.assertEqual(accounted, allocator.size) + finally: + mem_cache_common.get_global_server_args = old_get_global_server_args + if __name__ == "__main__": unittest.main() diff --git a/test/registered/unit/mem_cache/test_cp_shared_kv_runtime.py b/test/registered/unit/mem_cache/test_cp_shared_kv_runtime.py index c22955623..d1bd19e27 100644 --- a/test/registered/unit/mem_cache/test_cp_shared_kv_runtime.py +++ b/test/registered/unit/mem_cache/test_cp_shared_kv_runtime.py @@ -561,6 +561,14 @@ class TestCpSharedKVRuntimeHelpers(unittest.TestCase): forward_batch.out_cache_loc = torch.arange(127, dtype=torch.int64) self.assertFalse(can_reuse_current_extend_kv(forward_batch)) + forward_batch.extend_seq_lens_cpu = [65] + forward_batch.seq_lens_cpu = torch.tensor([40320 + 65], dtype=torch.int32) + forward_batch.out_cache_loc = torch.arange(128, dtype=torch.int64) + self.assertTrue(can_reuse_current_extend_kv(forward_batch)) + + forward_batch.out_cache_loc = torch.arange(64, dtype=torch.int64) + self.assertFalse(can_reuse_current_extend_kv(forward_batch)) + def test_should_reuse_current_extend_kv_disables_draft_cache_hit_suffix(self): from sglang.srt.environ import envs from sglang.srt.layers.attention.nsa import cp_shared_kv_runtime as runtime