Keep CP HiCache free room from collapsing under cache reuse
Repeated GSM8K/cache-hit traffic can drain a CP owner lane while exact allocation still succeeds. This changes L1/L2 free-room logic so the trigger accounts for the pending allocation, then evicts a minimum target chunk or the exact shortage, whichever is larger. CP load-back also performs best-effort owner-lane free-room eviction before admitting tiny cache hits that would otherwise skew one lane. The commit also adds gated bs>1 prefill timing markers around recv/bootstrap/batch/inflight polling so future hangs expose the scheduler boundary instead of only surfacing as a watchdog shutdown. The post-fix repeat GSM8K failure is recorded as an active regression to continue investigating, not as old-process noise. Constraint: Free-room policy must reduce repeated eviction and owner-lane starvation without reserving required+target pages after every allocation Rejected: Evict to required+target availability | wastes L1/L2 residency under many short cache hits Rejected: Treat free-room misses as fatal on load-back | exact capacity should remain the strict admission condition Confidence: medium Scope-risk: moderate Directive: Do not remove the gated bs>1 timing markers until repeat-GSM8K hangs have a direct failing boundary Tested: git diff --check Tested: python -m py_compile for touched Python files Tested: remote earlier test_cp_shared_kv_layout.py and test_cp_hicache_metadata.py passed 157 tests after this free-room formula Not-tested: Two full repeat GSM8K runs after this commit; latest repeat still killed prefill and requires follow-up root-cause work
This commit is contained in:
laoyao0822
@@ -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.
|
||||
|
||||
@@ -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:
|
||||
|
||||
|
||||
@@ -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 更连续的空间。
|
||||
|
||||
### 实现原则
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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(
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user