chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 11 Packages Projects Releases Wiki Activity

Stabilize CP HiCache page-tail ownership under EAGLE reuse

Browse Source 
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
This commit is contained in:
laoyao0822
2026-05-30 01:20:01 +08:00
parent 21065cdfdf
commit b56a4f2e6b
16 changed files with 964 additions and 17 deletions
Download Patch File Download Diff File Expand all files Collapse all files
docs/advanced_features/nsa_prefill_cp_page_aligned_cache_contract.md
+408
View File
@@ -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.
python/sglang/srt/disaggregation/decode.py
+16
View File
@@ -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)
python/sglang/srt/disaggregation/prefill.py
+16
View File
@@ -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,
python/sglang/srt/environ.py
+2
View File
@@ -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)
python/sglang/srt/layers/attention/nsa/cp_shared_kv_runtime.py
+5 -1
View File
@@ -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:
python/sglang/srt/layers/attention/nsa/utils.py
+49
View File
@@ -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
python/sglang/srt/layers/attention/nsa_backend.py
+23 -5
View File
@@ -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),
python/sglang/srt/managers/scheduler_dp_attn_mixin.py
+11 -1
View File
@@ -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()
python/sglang/srt/managers/scheduler_output_processor_mixin.py
+77
View File
@@ -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
python/sglang/srt/mem_cache/common.py
+5 -4
View File
@@ -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",
python/sglang/srt/mem_cache/hiradix_cache.py
+106
View File
@@ -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
python/sglang/srt/mem_cache/radix_cache.py
+40 -2
View File
@@ -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))
test/registered/unit/layers/test_nsa_cp_utils.py
+28 -3
View File
@@ -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,
)
test/registered/unit/managers/test_scheduler_dp_attn_mixin.py
+37
View File
@@ -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)
test/registered/unit/mem_cache/test_cp_hicache_load_back_owner_lanes.py
+133 -1
View File
@@ -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()
test/registered/unit/mem_cache/test_cp_shared_kv_runtime.py
+8
View File
@@ -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