diff --git a/docs/advanced_features/nsa_prefill_cp_accept_len1_garbage_output_debug_ledger.md b/docs/advanced_features/nsa_prefill_cp_accept_len1_garbage_output_debug_ledger.md index cecd58329..8c2938210 100644 --- a/docs/advanced_features/nsa_prefill_cp_accept_len1_garbage_output_debug_ledger.md +++ b/docs/advanced_features/nsa_prefill_cp_accept_len1_garbage_output_debug_ledger.md @@ -1781,3 +1781,254 @@ SSH 恢复后重新检查远端状态: - 本地 `py_compile` 通过:`decode.py`、`scheduler_output_processor_mixin.py`、`test_decode_queue_compaction.py`。 - 远端 g0034 容器 targeted tests 通过: `PYTHONPATH=python python -m pytest -q test/registered/unit/disaggregation/test_decode_queue_compaction.py test/registered/unit/managers/test_scheduler_health_check.py` → `17 passed`。 + +## 2026-06-06 full GSM8K prefill health-fail investigation + +Evidence from `/mnt/beegfs/cjy/log/sglang_cp_hicache_20260606_074842.log` during official GSM8K full run: + +- Official 200-question run passed before this: `Accuracy=0.970`, `Invalid=0.000`, no CP fail-fast/fallback/traceback in prefill log. +- Full run reached about 87% (`1150/1319`) before service failure. +- Prefill did not throw a Python exception. First prefill-side terminal symptom was health check: + - `08:13:32 Health check failed. Server couldn't get a response from detokenizer... last_heartbeat time: 08:12:18` + - `08:13:35 ... 128 requests still pending ... aborted 128 requests with 503` +- Last normal prefill batch was at `08:12:18`: + - `#new-seq: 7`, `#new-token: 832`, `#cached-token: 4480`, `#queue-req: 104`, `#inflight-req: 9`. + - write-through CP backups for node IDs `6271..6277` reached `final ack` and all write locks were released with `remaining_ongoing=0`. +- Decode errors begin after prefill health failure: + - decode logs start reporting `Could not fetch prefill parallel info from g0034:18992` at `08:13:48`, so decode failure is downstream of prefill bootstrap shutdown, not the first cause. +- Process state after shutdown: no active prefill process/port `17100`; only zombie python children remained. No container OOM evidence was found. + +Current hypothesis: + +- The scheduler likely stopped progressing after the last normal batch, before producing more detokenizer output. +- Two candidate blocking sites in prefill event loop both use `poll_and_all_reduce_attn_cp_tp_group(...)` and currently lack timing/status logs: + 1. `PrefillBootstrapQueue.pop_bootstrapped()` polls bootstrap queue. + 2. `SchedulerDisaggregationPrefillMixin.process_disagg_prefill_inflight_queue()` polls inflight KV transfer queue. +- Existing logs cannot distinguish whether the stall is in `recv_requests`, `pop_bootstrapped`, `get_next_disagg_prefill_batch_to_run`, `run_batch/process_batch_result`, or `process_disagg_prefill_inflight_queue`. + +Next diagnostic requirement: + +- Add low-frequency / env-gated timing logs around the prefill event-loop boundaries and around bootstrap/inflight poll-all-reduce. Log only when elapsed time exceeds a threshold or `SGLANG_CP_SHARED_KV_BS_GT1_DEBUG=1` is set. +- Do not change scheduling or transfer semantics until the blocking boundary is confirmed. + +Implemented diagnostic: + +- `prefill.py` now has env-gated `[CP_SHARED_KV_BS_GT1_TIMING]` markers/timings around: + - normal prefill event loop: `recv_requests`, `pop_bootstrapped`, `get_next_disagg_prefill_batch_to_run`, `run_batch`, `process_batch_result`, `process_disagg_prefill_inflight_queue`; + - bootstrap queue poll: `PrefillBootstrapQueue.pop_bootstrapped()` around `poll_and_all_reduce_attn_cp_tp_group`; + - inflight transfer poll: `process_disagg_prefill_inflight_queue()` around `poll_and_all_reduce_attn_cp_tp_group`. +- The diagnostic is off by default and should not change scheduling or transfer semantics. +- For the next reproduction, use: + - `SGLANG_CP_SHARED_KV_BS_GT1_TIMING=1` + - `SGLANG_CP_SHARED_KV_BS_GT1_TIMING_LIMIT=0` so the full GSM8K run does not exhaust the per-event log cap before the failure point; + - `SGLANG_CP_SHARED_KV_BS_GT1_TIMING_SLOW_MS=0` for the first reproduction so `*_start` markers are available. The last `*_start` without matching `*_done` is the blocking boundary. + +## 2026-06-06 restart with timing: current evidence + +Remote log: `/mnt/beegfs/cjy/log/sglang_cp_hicache_20260606_124320.log`. + +Observed state: + +- Prefill process is still alive after the new traffic batch; no current Python traceback or scheduler fatal. +- No `Health check failed` in this run so far. +- No `CP_SHARED_KV_FAIL_FAST`, no `RuntimeError`, no `Scheduler hit an exception`. +- Timing env is active, but the current run uses `SGLANG_CP_SHARED_KV_BS_GT1_TIMING_LIMIT=512` and `SGLANG_CP_SHARED_KV_BS_GT1_TIMING_SLOW_MS=1`; many timing events have already hit the per-rank cap (`4096` total across 8 ranks). If the service later hangs, the boundary marker may be exhausted and not show the final blocking point. For hang reproduction, use `SGLANG_CP_SHARED_KV_BS_GT1_TIMING_LIMIT=0`. + +Interesting non-fatal findings: + +- `CP_SHARED_KV_FALLBACK` appears 64 times, all seen as `tai_ipc_materialize reason=paged_start_slot_nonzero`; this means some current/prefix materialize paths still fall back from TAI IPC to local materialize plus collective when page-slot range starts from a nonzero slot. This is performance-relevant, not the fatal cause in this run. +- One request abort was observed at 13:00:28 across all CP ranks: + - `rid=76658e8841dc4db394a4df3052419e56` + - `prefill_bootstrap=4`, `prefill_inflight=2`, `abort_message=None` + - The scheduler continued normally afterward, so this abort is not the fatal condition. +- Batch size distribution from CP0 in the parsed window shows bs>1 is active: + - `bs=1: 411`, `bs=2: 132`, `bs=3: 128`, `bs=4: 62`, `bs=5: 22`, `bs=6: 8`, `bs=7: 4`. + +Timing hotspots: + +- `event_loop_run_batch_done`: avg ~1031 ms, p95 ~2025 ms, max ~5690 ms. This includes GPU forward and expected model work. +- `scheduler_prefill_prepare`: avg ~73 ms, p95 ~318 ms, max ~920 ms. +- Slowest `scheduler_prefill_prepare` cases are dominated by `prepare_for_extend`, not by `scan_waiting_queue`: + - example: `elapsed_ms=920.309`, `bs=3`, `extend_lens=[270, 2186, 1113]`, `prefix_lens=[51776, 79296, 94464]`, `prepare_for_extend=870.576 ms`. + - another example: `bs=1`, `extend_len=17363`, `prefix_len=137728`, `hicache_ready_to_load=43 ms`, `prepare_for_extend=716 ms`. +- `scan_waiting_queue` can still reach about 100 ms when queue is nontrivial, but it is not the largest hotspot in these slow samples. + +Current conclusion: + +- This restart has not reproduced the previous prefill-death symptom yet. +- The strongest evidence from this run is performance/control-path overhead in `prepare_for_extend` under large cached-prefix / bs>1 cases, plus residual TAI IPC fallback for nonzero page-slot materialize. +- For a future fatal reproduction, the current timing cap is insufficient; restart with unlimited timing limit before relying on the last `*_start`/`*_done` boundary. + +## 2026-06-06 full GSM8K accuracy verification + +Command run on g0034 container (`/sgl-workspace/sglang-tai`) against router `g0034:17100`: + +```bash +OPENAI_API_KEY=EMPTY PYTHONPATH=python python benchmark/gsm8k/bench_sglang.py \ + --host g0034 --port 17100 --backend srt \ + --data-path /mnt/beegfs/cjy/data/gsm8k_test.jsonl \ + --num-questions 1319 --parallel 64 \ + --temperature 0.0 --top-p 1.0 --max-new-tokens 512 \ + --result-file /mnt/beegfs/cjy/gsm8k_bench_sglang_full_20260606_132651.jsonl \ + --raw-result-file /mnt/beegfs/cjy/gsm8k_bench_sglang_full_raw_20260606_132651.jsonl +``` + +Result evidence: + +- Progress reached `1319/1319`. +- `Accuracy: 0.955`. +- `Invalid: 0.000`. +- `Latency: 311.296 s`. +- `Output throughput: 451.072 token/s`. +- Summary file: `/mnt/beegfs/cjy/gsm8k_bench_sglang_full_20260606_132651.jsonl`. +- Raw file: `/mnt/beegfs/cjy/gsm8k_bench_sglang_full_raw_20260606_132651.jsonl`. +- Log file: `/mnt/beegfs/cjy/gsm8k_bench_sglang_full_20260606_132651.log`. +- Post-run router health: `http://127.0.0.1:17100/health` returned `OK`. +- Post-run latest prefill log `/mnt/beegfs/cjy/log/sglang_cp_hicache_20260606_124320.log` had no matches for: + `Health check failed`, `Traceback`, `RuntimeError`, `Scheduler hit an exception`, `CP_SHARED_KV_FAIL_FAST`, `KVTransferError`, `SIGQUIT`. + +Interpretation: + +- Full GSM8K correctness is back to the expected official number (`0.955`) under current bs>1 CP shared-KV/HiCache configuration. +- The previous full-run prefill death was not reproduced in this run. +- This only verifies GSM8K accuracy and basic post-run service health; it does not close remaining performance issues such as `prepare_for_extend` CPU overhead or `tai_ipc_materialize reason=paged_start_slot_nonzero` fallback. + +Caveat: + +- `wc -l` on the raw result file reported `1318`, while the benchmark progress reached `1319/1319` and the summary result records `num_requests=1319`. Treat the printed benchmark summary as the accuracy source of truth unless raw-record completeness is needed for per-question diagnosis. + +## 2026-06-06 second full GSM8K cache-check prefill death + +Remote run under test: + +```bash +OPENAI_API_KEY=EMPTY PYTHONPATH=python python benchmark/gsm8k/bench_sglang.py \ + --host g0034 --port 17100 --backend srt \ + --data-path /mnt/beegfs/cjy/data/gsm8k_test.jsonl \ + --num-questions 1319 --parallel 64 \ + --temperature 0.0 --top-p 1.0 --max-new-tokens 512 \ + --result-file /mnt/beegfs/cjy/gsm8k_bench_sglang_full_cachecheck_20260606_135051.jsonl \ + --raw-result-file /mnt/beegfs/cjy/gsm8k_bench_sglang_full_cachecheck_raw_20260606_135051.jsonl +``` + +Observed failure: + +- Router health was `OK` before starting the second run. +- The benchmark progressed to about `77%` and then stalled; the benchmark process had to be killed manually. +- Prefill port `50001` was gone after the stall while router port `17100` was still listening. +- Latest prefill log remained `/mnt/beegfs/cjy/log/sglang_cp_hicache_20260606_124320.log`. +- No prefill-side Python traceback or explicit CP fail-fast was found immediately before death. +- Last normal prefill batch before watchdog was at `13:54:35` with `#new-seq: 7`, `#new-token: 448`, `#cached-token: 4928`, `#queue-req: 42`, `#inflight-req: 12`. +- Watchdog symptom: + - `13:55:17 Health check failed. Server couldn't get a response from detokenizer ... last_heartbeat time: 13:54:35`. + - `13:55:20 ... 65 requests still pending ... aborted 65 requests with 503`. +- Last timing evidence before death shows bootstrap queue did not fully drain: + - `bootstrap_poll_done elapsed_ms≈1.0 queue=7 polls_head=[2,2,2,2,2,2,1]` on multiple ranks. + - `event_loop_pop_bootstrapped_done ... bootstrapped=6 ... bootstrap_remaining=1`. +- Right after that, CP write locks for nodes `9803..9809` were released, followed by a burst of HiCache `load_back CP` logs for small host hits. These load_back entries report `SUCCESS`, so this is evidence of activity near the stall, not yet proof of a failing load_back transfer. +- Stale rid to investigate: `76658e8841dc4db394a4df3052419e56` was aborted much earlier at `13:00:28` with `prefill_bootstrap=4 prefill_inflight=2`, but appears again in the final watchdog `remaining_rids` list at `13:55:20`. + +Current investigation direction: + +- Root cause is not established yet. +- Two high-priority code paths to inspect: + 1. `PrefillBootstrapQueue.pop_bootstrapped()` ordering/cleanup behavior when one item remains pending while later items may be ready. + 2. Abort cleanup for requests already present in prefill bootstrap/inflight queues and tokenizer pending state. +- Do not treat the final HiCache load_back burst as root cause until a missing `*_done` timing boundary or code path proves it blocks scheduler progress. + +### 2026-06-06 root-cause finding: L1 owner-lane free-room was only diagnostic on success paths + +Evidence added after code inspection: + +- Final cache-check failure ended after a burst of CP HiCache load-back for tiny host hits: + - `host_hit_len=64` appeared `104` times in the final two-second window across CP ranks; + - `host_hit_len=128` appeared `8` times; + - owner-lane availability in those logs fell to `available_by_owner[0]=1` while `free_room_deficit_by_owner[0]≈443`. +- `hiradix_cache.load_back()` logged `free_room_deficit_by_owner`, but only evicted on `deficit_by_owner` (exact required pages), not on free-room deficit. Therefore successful tiny load-backs could drain owner lane 0 without replenishing L1 free room. +- `alloc_paged_token_slots_extend()` only called `_evict_for_compute_owner_lanes()` after the exact owner-lane allocation failed. Since `alloc_extend_compute_owner()` itself only enforces exact required pages, L1 free-room ratios were not maintained on successful extend allocations either. +- This explains why the second GSM8K run, unlike the first, can regress after cache is populated: many short cached prefixes map to the first CP owner lane, consume owner-0 pages, and leave the scheduler with little L1 room for subsequent cache-hit/extend work. + +Fix direction: + +- Preserve exact fail-fast semantics. +- Add proactive owner-lane free-room eviction before successful owner-aware extend allocation when L1 free-room trigger/target ratios are enabled. +- Add best-effort CP HiCache load-back owner-lane free-room eviction when exact capacity is sufficient but free room is below trigger. Failure to satisfy the free-room target should not reject the load-back; exact deficits still remain strict. + +### 2026-06-06 L1 owner-lane free-room fix implemented + +Implementation: + +- `alloc_paged_token_slots_extend()` now proactively calls CP owner-lane eviction before the first successful owner-aware allocation when L1 free-room target/trigger ratios are enabled. This keeps the configured free room from being only a failure-path mechanism. +- CP HiCache `load_back()` now performs best-effort owner-lane free-room eviction when exact `deficit_by_owner` is zero but `free_room_deficit_by_owner` is positive. Exact deficits still remain strict; free-room misses do not reject load-back. + +Regression coverage: + +- Added `test_compute_owner_alloc_maintains_l1_free_room_before_successful_alloc` to lock the case where exact allocation succeeds but lane free-room trigger requires proactive eviction. +- Remote g0034 container verification: + - `PYTHONPATH=python python -m py_compile python/sglang/srt/mem_cache/common.py python/sglang/srt/mem_cache/hiradix_cache.py test/registered/unit/mem_cache/test_cp_shared_kv_layout.py` passed. + - Targeted owner-lane tests: `3 passed`. + - Full `test_cp_shared_kv_layout.py`: `38 passed`. + +Open runtime validation: + +- No new prefill process has been started from this fix yet. The currently latest remote log remains the old failed run (`sglang_cp_hicache_20260606_124320.log`) with prefill port `50001` absent and router port `17100` still listening. + +### 2026-06-06 correction: free-room trigger/target semantics + +User correction applied to both L1 owner-lane and L2 host admission semantics: + +- Trigger condition must account for the pending allocation/reservation: + - trigger when `available < required + trigger_room`; + - equivalently, when `available - required < trigger_room`. +- Once triggered, `target_room` is the minimum eviction chunk. Evict + `max(required - available, target_room)`. +- Rationale: `required + target_room` leaves a full target watermark after the + current allocation and can waste too much L1/L2 cache residency under many + small cache-hit requests. Evicting only to `max(required, target_room)` + available is also wrong because it treats `target_room` as an available + watermark rather than a minimum eviction chunk. + +Tests added/updated: + +- L1: direct owner-lane helper test now covers `available=2, required=1, trigger=2`, which must trigger because the pending allocation would leave only one page. +- L2: `_free_room_deficit` and CP host admission tests now expect + `max(required - available, target_room)` after trigger. + +## 2026-06-07 repeat GSM8K precision run blocked by prefill death + +Remote run attempted for the requested first repeat: + +```bash +OPENAI_API_KEY=EMPTY PYTHONPATH=python python benchmark/gsm8k/bench_sglang.py \ + --host g0034 --port 17100 --backend srt \ + --data-path /mnt/beegfs/cjy/data/gsm8k_test.jsonl \ + --num-questions 1319 --parallel 64 \ + --temperature 0.0 --top-p 1.0 --max-new-tokens 512 \ + --result-file /mnt/beegfs/cjy/gsm8k_bench_sglang_full_repeat1_20260606_171214.jsonl \ + --raw-result-file /mnt/beegfs/cjy/gsm8k_bench_sglang_full_repeat1_raw_20260606_171214.jsonl +``` + +Observed failure: + +- Benchmark log: `/mnt/beegfs/cjy/gsm8k_bench_sglang_full_repeat1_20260606_171214.log`. +- Progress reached roughly the middle of the dataset, then failed with router/decode error: + `decode_internal_error: Could not fetch prefill parallel info from g0034:18992 after 20 attempts`. +- After failure, only router port `17100` was listening; prefill port `50001` and bootstrap port `18992` were gone. +- Prefill log: `/mnt/beegfs/cjy/log/sglang_cp_hicache_20260606_164817.log`. +- Last normal prefill batch before watchdog: + `17:14:54 #new-seq=5 #new-token=640 #cached-token=3200 #queue-req=47 #inflight-req=9`. +- CP write locks for nodes `5822..5826` were released immediately after that batch. +- Watchdog symptom: + - `17:15:37 Health check failed ... last_heartbeat time: 17:14:54`. + - `17:15:41 ... 65 requests still pending ... aborted 65 requests with 503`. +- No prefill-side Python traceback, explicit OOM, or CP fail-fast was found immediately before the watchdog shutdown. + +Important correction: + +- User confirmed this repeat was run on a new prefill process started after the free-room formula corrections. Therefore the failure is valid evidence that the latest free-room correction alone does not fully fix the repeat-GSM8K prefill death. +- The second requested repeat was not started because the first repeat took down prefill/bootstrap. + +Current next validation requirement: + +- Treat this as an active post-fix hang/regression, not an old-process artifact. +- If this hang needs deeper diagnosis, run with an uncapped or higher CP bs>1 timing log limit so the final scheduler/bootstrap boundary is visible; the current log reaches watchdog without a direct failing traceback. diff --git a/docs/advanced_features/nsa_prefill_cp_hicache_reactive_host_free_room_plan.md b/docs/advanced_features/nsa_prefill_cp_hicache_reactive_host_free_room_plan.md index 979084640..adf1d18ac 100644 --- a/docs/advanced_features/nsa_prefill_cp_hicache_reactive_host_free_room_plan.md +++ b/docs/advanced_features/nsa_prefill_cp_hicache_reactive_host_free_room_plan.md @@ -91,7 +91,7 @@ The new design should not add a per-tick maintenance loop. Add a **reactive host free-room watermark** to CP HiCache write admission. -When a reservation would fit, do nothing: +Initial draft, superseded below, only evicted when a reservation would not fit: ```text if required <= available: @@ -106,9 +106,9 @@ if required > available: deficit = required + free_room_target - available ``` -Because eviction happens before reservation, reserving `required` after this -eviction should leave approximately `free_room_target` host slots available on -the affected owner lanes. +This `required + free_room_target` target is now rejected because it can +over-reserve free room and waste cache residency under many small cache-hit +requests. ## Scope for the first implementation @@ -214,12 +214,15 @@ continues: if available >= required + trigger_room: deficit = 0 else: - deficit = max(0, required + target_room - available) + deficit = max(max(0, required - available), target_room) ``` -This keeps remaining cache usable when the current reservation fits, while still -making a triggered eviction release enough extra room to reduce repeated eviction -and give subsequent allocators a better chance of finding contiguous page runs. +The trigger is evaluated after accounting for the pending allocation +(`available - required < trigger_room`). Once triggered, `target_room` is the +minimum eviction chunk; if the current allocation needs more than that, evict the +larger exact deficit. It deliberately does **not** restore `required + +target_room`, because doing so would leave a full target room after the current +allocation and can evict useful cache too aggressively. Implications: diff --git a/docs/advanced_features/nsa_prefill_cp_shared_kv_bs_gt1_parallel_workstreams_zh.md b/docs/advanced_features/nsa_prefill_cp_shared_kv_bs_gt1_parallel_workstreams_zh.md index 6d0132f05..a31581e5a 100644 --- a/docs/advanced_features/nsa_prefill_cp_shared_kv_bs_gt1_parallel_workstreams_zh.md +++ b/docs/advanced_features/nsa_prefill_cp_shared_kv_bs_gt1_parallel_workstreams_zh.md @@ -1005,7 +1005,7 @@ scheduler admission 不做 bf16/fp8 字节级估算。原因是实际 KV pool 打开 scheduler bs>1 gate 后,CP HiCache write-through 的 radix node / host reservation 仍然必须保持 per request 独立;但 host 侧容量 admission 如果仍按 request 串行执行,会在同一个 prefill batch 内反复触发 L2 host free-room eviction: - correctness 没问题,因为 `reserve_write_cp()` 最终会消费真实 host allocator slot; -- 性能和碎片化较差,因为每个 request 都可能单独计算一次 `required + free_room_target - available`,在 host 接近满载时容易多次小额 eviction; +- 性能和碎片化较差,因为每个 request 都可能单独触发 free-room watermark admission;正确公式应在 `available - required < trigger_room` 时触发,并 evict `max(required - available, target_room)`,其中 `target_room` 是最小 eviction chunk,不是目标 available 水位; - 这与 L1/L2 free-room 的目标相反:应尽量一次 evict 到目标余量,降低 eviction 频率并给后续 allocate 更连续的空间。 ### 实现原则 diff --git a/python/sglang/srt/disaggregation/prefill.py b/python/sglang/srt/disaggregation/prefill.py index d9c01ac05..5ce64cf36 100644 --- a/python/sglang/srt/disaggregation/prefill.py +++ b/python/sglang/srt/disaggregation/prefill.py @@ -20,6 +20,7 @@ Life cycle of a request in the prefill server from __future__ import annotations import logging +import time from collections import deque from http import HTTPStatus from typing import TYPE_CHECKING, List, Optional @@ -86,6 +87,59 @@ def _cp_shared_kv_bs_gt1_prefill_debug( logger.info("[CP_SHARED_KV_BS_GT1_DEBUG] event=%s " + message, key, *args) +_CP_SHARED_KV_BS_GT1_PREFILL_TIMING_COUNTS = {} + + +def _cp_shared_kv_bs_gt1_prefill_timing_start() -> Optional[float]: + if not envs.SGLANG_CP_SHARED_KV_BS_GT1_TIMING.get(): + return None + return time.perf_counter() + + +def _cp_shared_kv_bs_gt1_prefill_timing( + key: str, + start_time: Optional[float], + message: str, + *args, +) -> None: + if start_time is None or not envs.SGLANG_CP_SHARED_KV_BS_GT1_TIMING.get(): + return + elapsed_ms = (time.perf_counter() - start_time) * 1000.0 + slow_ms = float(envs.SGLANG_CP_SHARED_KV_BS_GT1_TIMING_SLOW_MS.get()) + if slow_ms > 0 and elapsed_ms < slow_ms: + return + limit = int(envs.SGLANG_CP_SHARED_KV_BS_GT1_TIMING_LIMIT.get()) + count = _CP_SHARED_KV_BS_GT1_PREFILL_TIMING_COUNTS.get(key, 0) + if limit > 0 and count >= limit: + return + _CP_SHARED_KV_BS_GT1_PREFILL_TIMING_COUNTS[key] = count + 1 + logger.info( + "[CP_SHARED_KV_BS_GT1_TIMING] event=%s elapsed_ms=%.3f " + message, + key, + elapsed_ms, + *args, + ) + + +def _cp_shared_kv_bs_gt1_prefill_marker( + key: str, + message: str, + *args, +) -> None: + if not envs.SGLANG_CP_SHARED_KV_BS_GT1_TIMING.get(): + return + limit = int(envs.SGLANG_CP_SHARED_KV_BS_GT1_TIMING_LIMIT.get()) + count = _CP_SHARED_KV_BS_GT1_PREFILL_TIMING_COUNTS.get(key, 0) + if limit > 0 and count >= limit: + return + _CP_SHARED_KV_BS_GT1_PREFILL_TIMING_COUNTS[key] = count + 1 + logger.info( + "[CP_SHARED_KV_BS_GT1_TIMING] event=%s elapsed_ms=0.000 " + message, + key, + *args, + ) + + def _seq_summary(values) -> str: if values is None: return "None" @@ -458,11 +512,26 @@ class PrefillBootstrapQueue: else: return [], [] + poll_start = _cp_shared_kv_bs_gt1_prefill_timing_start() + _cp_shared_kv_bs_gt1_prefill_marker( + "bootstrap_poll_start", + "queue=%s return_failed=%s rids_to_check=%s", + len(self.queue), + return_failed_reqs, + len(rids_to_check) if rids_to_check is not None else None, + ) polls = poll_and_all_reduce_attn_cp_tp_group( [req.disagg_kv_sender for req in self.queue], self.scheduler.attn_cp_cpu_group, self.scheduler.attn_tp_cpu_group, ) + _cp_shared_kv_bs_gt1_prefill_timing( + "bootstrap_poll_done", + poll_start, + "queue=%s polls_head=%s", + len(self.queue), + polls[:8], + ) for i, (req, poll) in enumerate(zip(self.queue, polls)): if rids_to_check is not None: @@ -549,24 +618,114 @@ class SchedulerDisaggregationPrefillMixin: while True: # Receive requests - recv_reqs = self.recv_requests() - self.process_input_requests(recv_reqs) - self.waiting_queue.extend( - self.disagg_prefill_bootstrap_queue.pop_bootstrapped() + recv_start = _cp_shared_kv_bs_gt1_prefill_timing_start() + _cp_shared_kv_bs_gt1_prefill_marker( + "event_loop_recv_start", + "waiting=%s inflight=%s bootstrap=%s", + len(self.waiting_queue), + len(self.disagg_prefill_inflight_queue), + len(self.disagg_prefill_bootstrap_queue.queue), ) + recv_reqs = self.recv_requests() + _cp_shared_kv_bs_gt1_prefill_timing( + "event_loop_recv_done", + recv_start, + "recv=%s waiting=%s inflight=%s bootstrap=%s", + len(recv_reqs), + len(self.waiting_queue), + len(self.disagg_prefill_inflight_queue), + len(self.disagg_prefill_bootstrap_queue.queue), + ) + self.process_input_requests(recv_reqs) + pop_start = _cp_shared_kv_bs_gt1_prefill_timing_start() + bootstrapped_reqs = self.disagg_prefill_bootstrap_queue.pop_bootstrapped() + _cp_shared_kv_bs_gt1_prefill_timing( + "event_loop_pop_bootstrapped_done", + pop_start, + "bootstrapped=%s waiting_before=%s inflight=%s bootstrap_remaining=%s", + len(bootstrapped_reqs), + len(self.waiting_queue), + len(self.disagg_prefill_inflight_queue), + len(self.disagg_prefill_bootstrap_queue.queue), + ) + self.waiting_queue.extend(bootstrapped_reqs) # Get the next batch to run + batch_start = _cp_shared_kv_bs_gt1_prefill_timing_start() + _cp_shared_kv_bs_gt1_prefill_marker( + "event_loop_get_batch_start", + "waiting=%s inflight=%s bootstrap=%s", + len(self.waiting_queue), + len(self.disagg_prefill_inflight_queue), + len(self.disagg_prefill_bootstrap_queue.queue), + ) batch = self.get_next_disagg_prefill_batch_to_run() + _cp_shared_kv_bs_gt1_prefill_timing( + "event_loop_get_batch_done", + batch_start, + "has_batch=%s batch_size=%s waiting=%s inflight=%s bootstrap=%s", + batch is not None, + len(batch.reqs) if batch is not None else 0, + len(self.waiting_queue), + len(self.disagg_prefill_inflight_queue), + len(self.disagg_prefill_bootstrap_queue.queue), + ) self.cur_batch = batch # Launch the current batch if batch: + run_start = _cp_shared_kv_bs_gt1_prefill_timing_start() + _cp_shared_kv_bs_gt1_prefill_marker( + "event_loop_run_batch_start", + "batch_size=%s extend_lens=%s prefix_lens=%s inflight=%s", + len(batch.reqs), + _seq_summary(getattr(batch, "extend_lens", None)), + _seq_summary(getattr(batch, "prefix_lens", None)), + len(self.disagg_prefill_inflight_queue), + ) result = self.run_batch(batch) + _cp_shared_kv_bs_gt1_prefill_timing( + "event_loop_run_batch_done", + run_start, + "batch_size=%s inflight=%s", + len(batch.reqs), + len(self.disagg_prefill_inflight_queue), + ) + result_start = _cp_shared_kv_bs_gt1_prefill_timing_start() + _cp_shared_kv_bs_gt1_prefill_marker( + "event_loop_process_result_start", + "batch_size=%s inflight_before=%s", + len(batch.reqs), + len(self.disagg_prefill_inflight_queue), + ) self.process_batch_result(batch, result) + _cp_shared_kv_bs_gt1_prefill_timing( + "event_loop_process_result_done", + result_start, + "batch_size=%s inflight_after=%s", + len(batch.reqs), + len(self.disagg_prefill_inflight_queue), + ) else: self.self_check_during_idle() + inflight_start = _cp_shared_kv_bs_gt1_prefill_timing_start() + _cp_shared_kv_bs_gt1_prefill_marker( + "event_loop_process_inflight_start", + "inflight=%s waiting=%s bootstrap=%s", + len(self.disagg_prefill_inflight_queue), + len(self.waiting_queue), + len(self.disagg_prefill_bootstrap_queue.queue), + ) self.process_disagg_prefill_inflight_queue() + _cp_shared_kv_bs_gt1_prefill_timing( + "event_loop_process_inflight_done", + inflight_start, + "inflight=%s waiting=%s bootstrap=%s", + len(self.disagg_prefill_inflight_queue), + len(self.waiting_queue), + len(self.disagg_prefill_bootstrap_queue.queue), + ) # Update last_batch self.last_batch = batch @@ -766,11 +925,26 @@ class SchedulerDisaggregationPrefillMixin: done_reqs = [] + poll_start = _cp_shared_kv_bs_gt1_prefill_timing_start() + _cp_shared_kv_bs_gt1_prefill_marker( + "inflight_poll_start", + "inflight=%s rids_head=%s rids_to_check=%s", + len(self.disagg_prefill_inflight_queue), + [req.rid for req in self.disagg_prefill_inflight_queue[:8]], + len(rids_to_check) if rids_to_check is not None else None, + ) polls = poll_and_all_reduce_attn_cp_tp_group( [req.disagg_kv_sender for req in self.disagg_prefill_inflight_queue], self.attn_cp_cpu_group, self.attn_tp_cpu_group, ) + _cp_shared_kv_bs_gt1_prefill_timing( + "inflight_poll_done", + poll_start, + "inflight=%s polls_head=%s", + len(self.disagg_prefill_inflight_queue), + polls[:8], + ) undone_reqs: List[Req] = [] # Check .poll() for the reqs in disagg_prefill_inflight_queue. If Success, respond to the client and remove it from the queue diff --git a/python/sglang/srt/mem_cache/allocator.py b/python/sglang/srt/mem_cache/allocator.py index a00ab4ea3..bf2cceff8 100644 --- a/python/sglang/srt/mem_cache/allocator.py +++ b/python/sglang/srt/mem_cache/allocator.py @@ -57,10 +57,11 @@ def compute_owner_lane_free_room_deficits( if capacity > 0 and trigger_ratio > 0 else 0 ) - if int(avail) >= int(req) + trigger_room: + trigger_available = int(req) + trigger_room + if int(avail) >= trigger_available: deficits.append(0) else: - deficits.append(max(0, int(req) + target_room - int(avail))) + deficits.append(max(max(0, int(req) - int(avail)), target_room)) return deficits diff --git a/python/sglang/srt/mem_cache/common.py b/python/sglang/srt/mem_cache/common.py index b2c8778b7..71d4cfbde 100644 --- a/python/sglang/srt/mem_cache/common.py +++ b/python/sglang/srt/mem_cache/common.py @@ -444,6 +444,18 @@ def _evict_for_compute_owner_lanes( return +def _cp_owner_lane_free_room_enabled(tree_cache: BasePrefixCache | None) -> bool: + if tree_cache is None: + return False + return ( + float(getattr(tree_cache, "hicache_l1_free_room_ratio", 0.0) or 0.0) > 0.0 + or float( + getattr(tree_cache, "hicache_l1_free_room_trigger_ratio", 0.0) or 0.0 + ) + > 0.0 + ) + + def alloc_paged_token_slots_extend( tree_cache: BasePrefixCache, prefix_lens: torch.Tensor, @@ -540,6 +552,20 @@ def alloc_paged_token_slots_extend( if backup_state: state = allocator.backup_state() + # Maintain the configured L1 owner-lane free room on successful + # allocations as well. Previously this eviction path only ran after + # exact allocation failure, so cache hits and short extends could drain + # one CP owner lane down to a handful of pages while still "succeeding", + # leaving later batches to stall under skewed owner-lane pressure. + if _cp_owner_lane_free_room_enabled(tree_cache): + _evict_for_compute_owner_lanes( + tree_cache=tree_cache, + allocator=allocator, + page_compute_owners=page_compute_owners, + ) + if backup_state: + state = allocator.backup_state() + out_cache_loc = alloc_extend_compute_owner( prefix_lens, prefix_lens_cpu, diff --git a/python/sglang/srt/mem_cache/hiradix_cache.py b/python/sglang/srt/mem_cache/hiradix_cache.py index 895b5d8e4..78fac7e2a 100644 --- a/python/sglang/srt/mem_cache/hiradix_cache.py +++ b/python/sglang/srt/mem_cache/hiradix_cache.py @@ -152,7 +152,7 @@ def _free_room_deficit( trigger_room = _page_aligned_room_from_ratio(capacity, trigger_ratio, page_size) if int(available) >= int(required) + trigger_room: return 0 - return max(0, int(required) + target_room - int(available)) + return max(max(0, int(required) - int(available)), target_room) @dataclass @@ -3438,6 +3438,37 @@ class HiRadixCache(RadixCache): ) record_stage("owner_lane_evict") + if not any(v > 0 for v in load_back_plan.deficit_by_owner) and any( + v > 0 for v in load_back_plan.free_room_deficit_by_owner + ): + # Best-effort L1 free-room maintenance for CP load-back. Exact + # admission may still succeed with a nearly empty owner lane + # (common for one-page host hits that all map to owner 0), but + # letting those loads proceed without replenishing free room can + # starve the next extend batch. Free-room target misses are not + # fatal; exact deficits remain the strict condition below. + free_room_deficits = [ + int(v) for v in load_back_plan.free_room_deficit_by_owner + ] + free_room_pages = sum(max(0, v) for v in free_room_deficits) + if free_room_pages > 0: + evict_result = self.evict( + EvictParams( + num_tokens=free_room_pages * self.page_size, + owner_lane_deficits=free_room_deficits, + ) + ) + logger.debug( + "[HiCache-load] owner-lane free-room eviction before CP " + "load-back: node_id=%d free_room_deficit_by_owner=%s " + "num_tokens_evicted=%d", + last_hit_node.id, + free_room_deficits, + getattr(evict_result, "num_tokens_evicted", 0), + ) + load_back_plan = self._refresh_cp_load_back_plan(load_back_plan) + record_stage("owner_lane_free_room_evict") + if any(v > 0 for v in load_back_plan.deficit_by_owner): self.dec_lock_ref(ancester_node) logger.warning( diff --git a/test/registered/unit/mem_cache/test_cp_hicache_metadata.py b/test/registered/unit/mem_cache/test_cp_hicache_metadata.py index 40969ed87..3301f744f 100644 --- a/test/registered/unit/mem_cache/test_cp_hicache_metadata.py +++ b/test/registered/unit/mem_cache/test_cp_hicache_metadata.py @@ -631,7 +631,22 @@ class TestCpHiCacheFreeRoom(CustomTestCase): trigger_ratio=0.25, ) - self.assertEqual(deficit, 96) + self.assertEqual(deficit, 128) + + def test_free_room_deficit_trigger_accounts_for_pending_reservation(self): + deficit = hiradix_cache._free_room_deficit( + required=64, + available=96, + capacity=256, + page_size=64, + target_ratio=0.5, + trigger_ratio=0.25, + ) + + # available=96 is above trigger_room=64, but after reserving required=64 + # only 32 tokens would remain. Trigger checks available - required, + # then evicts at least target_room=128 tokens. + self.assertEqual(deficit, 128) def test_free_room_deficit_rounds_room_to_page(self): deficit = hiradix_cache._free_room_deficit( @@ -643,7 +658,7 @@ class TestCpHiCacheFreeRoom(CustomTestCase): trigger_ratio=0.0, ) - self.assertEqual(deficit, 128) + self.assertEqual(deficit, 64) def test_cp_host_write_admission_uses_trigger_target_room(self): cache = HiRadixCache.__new__(HiRadixCache) @@ -672,7 +687,7 @@ class TestCpHiCacheFreeRoom(CustomTestCase): ) self.assertEqual(admission.target_available_by_owner, (96, 256)) - self.assertEqual(admission.deficit_by_owner, (96, 0)) + self.assertEqual(admission.deficit_by_owner, (128, 0)) def test_cp_host_write_admission_does_not_evict_when_trigger_room_fits(self): cache = HiRadixCache.__new__(HiRadixCache) @@ -706,7 +721,7 @@ class TestCpHiCacheFreeRoom(CustomTestCase): def test_cp_host_write_admission_uses_draft_room_deficit(self): cache = HiRadixCache.__new__(HiRadixCache) cache.page_size = 64 - cache.hicache_host_free_room_ratio = 0.25 + cache.hicache_host_free_room_ratio = 0.5 cache.hicache_host_free_room_trigger_ratio = 0.25 cache._cp_required_host_tokens_by_rank = lambda _indices: (64, 0) cache._cp_host_capacity_snapshot = lambda: CpHiCacheCapacitySnapshot( @@ -731,7 +746,7 @@ class TestCpHiCacheFreeRoom(CustomTestCase): self.assertEqual(admission.target_available_by_owner, (256, 256)) self.assertEqual(admission.draft_available_by_owner, (64, 256)) - self.assertEqual(admission.deficit_by_owner, (64, 0)) + self.assertEqual(admission.deficit_by_owner, (128, 0)) def test_cp_host_write_batch_admission_uses_aggregate_required(self): cache = HiRadixCache.__new__(HiRadixCache) @@ -761,7 +776,7 @@ class TestCpHiCacheFreeRoom(CustomTestCase): self.assertEqual(admission.required_by_owner, (128, 0)) self.assertEqual(admission.target_available_by_owner, (96, 256)) # required=128, target_room=128, available=96. - self.assertEqual(admission.deficit_by_owner, (160, 0)) + self.assertEqual(admission.deficit_by_owner, (128, 0)) def test_prepare_write_backups_for_reqs_runs_one_batch_admission(self): cache = HiRadixCache.__new__(HiRadixCache) @@ -2873,7 +2888,7 @@ class TestHiRadixCacheCPLoadBack(CustomTestCase): self.assertEqual(plan.deficit_by_owner, [1, 0]) # The free-room target is still reported for observability/proactive policy. # required=1 page, available=0, target_room=ceil(8*0.5)=4 pages. - self.assertEqual(plan.free_room_deficit_by_owner, [5, 0]) + self.assertEqual(plan.free_room_deficit_by_owner, [4, 0]) def test_cp_load_back_uses_host_len_not_host_value(self): cache = HiRadixCache.__new__(HiRadixCache) diff --git a/test/registered/unit/mem_cache/test_cp_shared_kv_layout.py b/test/registered/unit/mem_cache/test_cp_shared_kv_layout.py index 8b0e419a5..cd2cef0bf 100644 --- a/test/registered/unit/mem_cache/test_cp_shared_kv_layout.py +++ b/test/registered/unit/mem_cache/test_cp_shared_kv_layout.py @@ -123,6 +123,26 @@ class TestCpSharedKVLayout(unittest.TestCase): class TestCPSharedPagedAllocator(CustomTestCase): + def test_compute_owner_free_room_trigger_accounts_for_pending_allocation(self): + from sglang.srt.mem_cache.allocator import ( + compute_owner_lane_free_room_deficits, + ) + + deficits = compute_owner_lane_free_room_deficits( + required=[1], + available=[2], + capacities=[8], + target_ratio=0.5, + trigger_ratio=0.25, + ) + + # available=2 would satisfy both exact required=1 and trigger_room=2 + # independently, but the pending allocation would leave only one free + # page. Therefore trigger condition must be available < required + + # trigger_room, and target_room is the minimum eviction chunk after + # trigger. + self.assertEqual(deficits, [4]) + def test_compute_owner_alloc_fallback_logs_every_event(self): from sglang.srt.mem_cache import common @@ -821,6 +841,102 @@ class TestCPSharedPagedAllocator(CustomTestCase): self.assertEqual(allocator.calls, 1) self.assertEqual(out.numel(), page_size * 8) + def test_compute_owner_alloc_maintains_l1_free_room_before_successful_alloc(self): + from types import SimpleNamespace + + from sglang.srt.mem_cache import common + from sglang.srt.mem_cache.base_prefix_cache import EvictResult + + page_size = 64 + + class FakeAllocator: + def __init__(self): + self.page_size = page_size + self.cp_size = 2 + self.calls = 0 + self.available_by_owner = [1, 8] + + def available_size(self): + return sum(self.available_by_owner) * page_size + + def compute_owner_lane_stats(self, page_compute_owners): + required = [ + sum(1 for owner in page_compute_owners if owner == lane) + for lane in range(self.cp_size) + ] + deficits = [ + max(0, req - avail) + for req, avail in zip(required, self.available_by_owner) + ] + return required, list(self.available_by_owner), deficits + + def compute_owner_lane_capacity_pages(self): + return [8, 8] + + def alloc_extend_compute_owner( + self, + _prefix_lens, + _prefix_lens_cpu, + _seq_lens, + _seq_lens_cpu, + _last_loc, + extend_num_tokens, + _page_compute_owners, + ): + self.calls += 1 + return torch.arange(extend_num_tokens, dtype=torch.int64) + + def alloc_extend(self, *_args, **_kwargs): + raise AssertionError("legacy allocation should not be used") + + class FakeTreeCache: + hicache_l1_free_room_ratio = 0.5 + hicache_l1_free_room_trigger_ratio = 0.25 + + def __init__(self, allocator): + self.token_to_kv_pool_allocator = allocator + self.owner_deficits = [] + + def is_chunk_cache(self): + return False + + def evictable_size(self): + return page_size * 8 + + def evict(self, params): + deficits = list(params.owner_lane_deficits) + self.owner_deficits.append(deficits) + for lane, deficit in enumerate(deficits): + allocator.available_by_owner[lane] += max(0, int(deficit)) + return EvictResult(num_tokens_evicted=sum(deficits) * page_size) + + allocator = FakeAllocator() + tree_cache = FakeTreeCache(allocator) + server_args = SimpleNamespace( + enable_nsa_prefill_cp_shared_kv=True, + enable_nsa_prefill_context_parallel=True, + nsa_prefill_cp_mode="in-seq-split", + ) + + with patch.object(common, "get_global_server_args", return_value=server_args): + out = common.alloc_paged_token_slots_extend( + tree_cache=tree_cache, + prefix_lens=torch.tensor([0], dtype=torch.int64), + prefix_lens_cpu=torch.tensor([0], dtype=torch.int64), + seq_lens=torch.tensor([page_size], dtype=torch.int64), + seq_lens_cpu=torch.tensor([page_size], dtype=torch.int64), + last_loc=torch.tensor([-1], dtype=torch.int64), + extend_num_tokens=page_size, + ) + + # required=1 page on owner lane 0, available=1, trigger=ceil(8*0.25)=2, + # target=ceil(8*0.5)=4. Exact allocation would succeed, but configured + # L1 free room requires pre-evicting at least target_room=4 pages. It + # must not evict required+target, but target_room is the minimum chunk. + self.assertEqual(tree_cache.owner_deficits, [[4, 0]]) + self.assertEqual(allocator.calls, 1) + self.assertEqual(out.numel(), page_size) + def test_compute_owner_alloc_uses_batch_owner_plan_for_multi_request(self): from types import SimpleNamespace @@ -1197,7 +1313,8 @@ class TestCPSharedPagedAllocator(CustomTestCase): ) # required=1 page, available=0, target_room=ceil(8*0.5)=4 pages. - self.assertEqual(tree_cache.owner_deficits[0], [5, 0]) + # Evict at least target_room=4 pages after trigger. + self.assertEqual(tree_cache.owner_deficits[0], [4, 0]) def test_compute_owner_capacity_wait_reports_owner_lane_deficits(self): from types import SimpleNamespace