diff --git a/TECHNICAL_REPORTS/2088-paged-budget-gaps-20261002.en.md b/TECHNICAL_REPORTS/2088-paged-budget-gaps-20261002.en.md new file mode 100644 index 000000000..fd13ad62f --- /dev/null +++ b/TECHNICAL_REPORTS/2088-paged-budget-gaps-20261002.en.md @@ -0,0 +1,58 @@ +# Technical Report: PR (issue #2088), residual paged-budget gaps after #2077 + +**Date**: 2026-10-02 + +**Status**: Implemented and validated on GB10 with meta-llama-3.1-8b-instruct-4bit under a forced 4444-block budget; not yet merged. + +**Languages**: Rust + +**Risk level**: Medium (changes paged admission and reclaim decisions; adds a default-on admission watermark of 0.01) + +## Executive summary + +PR #2077 left three gaps in the paged KV block budget: the free-block figure was pool-wide while blocks are reused per layer, a parked chunked prefill's next chunk had no reservation of its own, and nothing kept decode headroom at admission. This PR closes all three and adds `--kv-admission-watermark` (default 0.01). Measuring the watermark on GB10 exposed two more accounting gaps from #2077: adopted prompt-cache prefixes were charged twice at admission (which wedged the queue), and blocks pinned by queued adoptions could not be reclaimed (which shed a running row). Both are fixed here as well. + +## Problem statement + +1. `PagedBlockPool::free_block_budget` returned `budget - live`. A freed block goes back to its own layer's free list, and `acquire_block` for layer `L` reuses only `free_lists[L]` before minting under the pool-wide cap. Once the free lists diverge, one layer can fail to acquire while the figure is positive. +2. A parked chunked prefill is protected only by subtracting its remaining blocks from the figure other consumers read. That figure saturates at 0, so when the pool fell below the reservation every other consumer deferred but nothing restored it, and the chunk's forward could hit an exhausted pool. The reservation also ignored the tile padding a padded chunk writes before trimming. +3. Under a tight budget, admitting into exactly the free blocks leaves running rows no room to grow, so decode preempts them and they re-prefill (12 preemptions for 9 turns in the #2077 pressure run). + +## Change summary + +- Per-layer accounting (`mlxcel-core` `paged.rs`): `free_block_budget` returns `n * num_layers` for the largest `n` such that every layer can acquire `n` blocks from its own free list plus a share of the mint headroom `budget - allocated`, capped at `budget - live`. Every scheduler demand is uniform across layers, so with balanced free lists the figure is unchanged. Re-tagging free blocks across layers was rejected: pool rows are per-layer slabs, so a moved block needs a new row and the budget would stop bounding memory. Layers that have never minted a block are left out of the constraint. +- Chunk reservation (`block_reclaim.rs`, `prefill.rs`): `continue_chunked_prefill` calls `reserve_prefill_chunk_blocks` before the forward. It checks the pool figure (not the saturating set-aside figure), evicts cold prefixes, then preempts with floor 0 and no priority limit, since the chunk's admission already checked the whole prompt. If even that leaves too few blocks, the request fails with "KV cache budget exhausted" instead of panicking the worker. This covers the `MLXCEL_MIXED_STEP` tick and the #1011 grant path. `chunked_prefill_reserved_blocks` now includes the next chunk's padding. +- Admission watermark (`admission_watermark.rs`, `--kv-admission-watermark`, env `MLXCEL_KV_ADMISSION_WATERMARK`, range 0.0 to 0.5): while a decode batch is live, admission requires free blocks minus `ceil(F * budget)` to cover the request and reclaims toward that target before deferring: cold prefixes first, then (for a higher-priority request) strictly lower-priority rows, as admission already did for the request itself. With an empty batch it does not apply, so a request that fits the budget is never refused. `paged_prefill_room` applies it, so the batched window and the single-sequence gate charge the same figure. +- Adopted prefix charge (`admission.rs`, `CachePool::paged_blocks_to_reach`): admission charges the blocks the suffix will mint, measured from the sequence's paged state, and at least one block per layer after an adoption (a full-prompt hit re-runs its last token into a shared tail block, which forks). Before this, an adopted prefix larger than half the budget could never be admitted, even into an empty batch. +- Queued adoption reclaim (`queued_adoption.rs`, `PrefillQueue::find_lowest_first`/`get_mut`/`remove`): adoption happens at enqueue and pins the prefix's blocks while the request waits. Decode and chunk reclaim now drop one queued text request's adoption (lowest priority, newest first) after evicting cold prefixes and before preempting; the request re-prefills cold and keeps its prompt-cache context for donate-back. Admission does not do this; a queued request waits instead of costing another one its hit. Multimodal requests are skipped because an adopted VLM request no longer holds its full-prompt embeddings. + +## Validation + +- Unit tests: `paged_detach_tests.rs` adds 3 per-layer cases; on a diverged state (layer 0 owns 2 free blocks, layer 1 none, budget fully minted) the old figure is 2 while a layer-1 append fails, and the new figure is 0; with 2 blocks of headroom the figure is 4 and exactly that much is acquirable. `paged_append_need_tests.rs` adds the suffix charge. `block_reclaim_tests.rs` adds MixedStep tick-order cases (decode defers and the chunk fits; mid-block rows leave the pool below the reservation, the chunk append fails without the reservation and succeeds after it preempts one row; a cold prefix is evicted before any row), the watermark deferral and eviction cases, and the queued-adoption cases (a lone boundary row survives by dropping a queued adoption; order is prefix, adoption, row; admission never drops one). `scheduler_prompt_cache_tests.rs` adds the adopted-prefix charge on a real pool. +- `server::batch`, `server::prompt_cache`, `runtime_settings`, `cli_input`, `config_tests`, `commands::serve` filters: 822 passed. mlxcel-core `paged` and `budget` filters: 282 passed. `cargo clippy --release --features cuda --lib --tests -- -D warnings` (root and `-p mlxcel-core`) and `cargo fmt --all` are clean. +- GB10 pressure run (`--kv-cache-budget 600000000`, 4444 blocks, paged, `--max-batch-size 4`, a prefix parked first, then three concurrent 3-turn conversations of 600-token generations): + +| Build | Watermark | Preemptions | Decode reclaims | Turns | Errors | +|-------|-----------|-------------|-----------------|-------|--------| +| main (#2077 report) | n/a | 12 | 16 | 9 | none | +| gaps 1-3 only | 0 | 12 | 16 | 9 | none | +| gaps 1-3 only | 0.15 | 2 | 4 | 4 | 2 shed rows, 1 request hung 900 s | +| this PR (x2) | 0 | 10 | 43 | 9 | none | +| this PR (x2) | 0.01 | 6 | 6 | 9 | none | +| this PR (x2) | 0.02 | 6 | 6 | 9 | none | +| this PR (x2) | 0.05 | 5 | 5 | 9 | none | +| this PR | 0.10 | 1 | 2 | 9 | none | +| this PR | 0.15 | 1 | 3 | 9 | none | + + The 0.15 failure before the last two fixes: both sheds had `preempted=0` with one row in the batch and nothing evictable, so the remaining blocks were held by a queued adoption, and the hang was a request whose adopted prefix exceeded half the budget, deferred forever with an empty batch (an aborted 0.20 run showed the same pattern with the worker spinning at 96% CPU). The default 0.01 is the smallest grid value that lowered preemptions without losing a turn. At 0 the branch's 43 reclaim passes are 33 one-tick deferrals of a row short only by a parked chunk's reservation. +- `MLXCEL_MIXED_STEP=1` at 0.01, `--metrics`: 21 mixed steps and 55 prefill chunks, 9 of 9 turns, no panics, 7 preemptions. +- Default launch (auto budget, 753620 blocks) at 0.01, two 3-turn conversations: no reclaim or preemption, turns 2 and 3 hit the prompt cache (1280 and 1504 tokens). + +## Not validated + +- The per-layer figure trusts that every refcount-0 block sits on exactly one layer free list. Its `debug_assert!` never ran (all tests were release builds); instead every mutation site was read: `release_block` pushes on reaching 0, `acquire_block` and `restore_block` pop or remove before setting refcount 1, `retain_block` refuses a refcount-0 block, and no path removes a record from `blocks`. + +- The chunk reservation's reclaim path did not trigger in any real-server run (`Chunked prefill reclaimed` never logged); it is covered by unit tests only. +- The queued-adoption drop triggered once each at 0.10 and 0.15; its re-prefill of a dropped request was checked only through those runs completing all turns. +- The padded-chunk reservation applies only on Metal M5+ (`should_align_prefill`), which this host cannot run. +- The watermark default was measured on one model and one budget; under the auto budget it holds back 1% of the pool while rows decode. diff --git a/TECHNICAL_REPORTS/2088-paged-budget-gaps-20261002.ko.md b/TECHNICAL_REPORTS/2088-paged-budget-gaps-20261002.ko.md new file mode 100644 index 000000000..98cbc430b --- /dev/null +++ b/TECHNICAL_REPORTS/2088-paged-budget-gaps-20261002.ko.md @@ -0,0 +1,58 @@ +# 기술 보고서: PR (이슈 #2088), #2077 이후 남은 paged 예산 결함 + +**날짜**: 2026-10-02 + +**상태**: GB10에서 meta-llama-3.1-8b-instruct-4bit, 4444블록 예산으로 검증 완료, 아직 머지되지 않음 + +**언어**: Rust + +**위험도**: 중간 (paged 승인 및 회수 결정이 바뀌고, 기본값 0.01의 승인 워터마크가 켜짐) + +## 요약 + +PR #2077은 paged KV 블록 예산에 세 가지 결함을 남겼다. 빈 블록 수치는 풀 전체 기준인데 블록 재사용은 레이어 단위였고, 대기 중인 청크 프리필의 다음 청크에는 자체 예약이 없었으며, 승인 시점에 디코드 여유분을 남기는 장치가 없었다. 이 PR은 세 가지를 모두 해결하고 `--kv-admission-watermark`(기본값 0.01)를 추가한다. GB10에서 워터마크를 측정하는 과정에서 #2077의 회계 결함 두 가지가 더 드러났다. 채택한 프롬프트 캐시 접두사가 승인 시 이중으로 계산되어 큐가 멈췄고, 대기 중인 요청이 채택으로 고정한 블록은 회수할 수 없어 실행 중인 행이 탈락했다. 두 문제도 이 PR에서 함께 고친다. + +## 문제 + +1. `PagedBlockPool::free_block_budget`은 `budget - live`를 반환했다. 해제된 블록은 자기 레이어의 free list로 돌아가고, 레이어 `L`의 `acquire_block`은 `free_lists[L]`만 재사용한 뒤 풀 전체 상한 아래에서 새 블록을 만든다. free list가 레이어마다 달라지면 수치가 양수인데도 한 레이어가 블록을 얻지 못한다. +2. 대기 중인 청크 프리필은 다른 소비자가 읽는 수치에서 남은 블록을 빼는 방식으로만 보호되었다. 이 수치는 0에서 포화되므로, 풀이 예약량 아래로 떨어지면 다른 소비자는 모두 미루지만 예약을 복구하는 경로가 없어 청크 forward가 고갈된 풀에 부딪힐 수 있었다. 패딩된 청크가 잘라내기 전에 쓰는 타일 패딩도 예약에 포함되지 않았다. +3. 예산이 빡빡할 때 빈 블록에 꼭 맞게 승인하면 실행 중인 행이 자랄 공간이 없어 디코드가 이를 선점하고, 선점된 행은 프리필을 다시 한다(#2077 압박 실행에서 9턴에 선점 12회). + +## 변경 사항 + +- 레이어별 회계(`mlxcel-core` `paged.rs`): `free_block_budget`은 모든 레이어가 자기 free list와 mint 여유분 `budget - allocated`의 몫으로 `n`블록씩 얻을 수 있는 최대 `n`에 대해 `n * num_layers`를 반환하고, `budget - live`로 상한을 둔다. 스케줄러의 수요는 모두 레이어 간 균일하므로 free list가 균형 잡혀 있으면 값은 그대로다. 빈 블록을 다른 레이어로 옮기는 방식은 기각했다. 풀 행은 레이어별 slab이라 옮긴 블록은 새 레이어에서 새 행이 필요하고, 그러면 예산이 메모리를 제한하지 못한다. 아직 블록을 만든 적이 없는 레이어는 제약에서 제외한다. +- 청크 예약(`block_reclaim.rs`, `prefill.rs`): `continue_chunked_prefill`은 forward 전에 `reserve_prefill_chunk_blocks`를 호출한다. 포화되는 예약 차감 수치가 아니라 풀 수치를 확인하고, 콜드 접두사를 축출한 뒤 floor 0, 우선순위 제한 없이 선점한다. 청크의 승인은 이미 프롬프트 전체를 확인했기 때문이다. 그래도 블록이 부족하면 워커를 panic시키는 대신 "KV cache budget exhausted"로 요청을 실패시킨다. `MLXCEL_MIXED_STEP` 틱과 #1011 grant 경로를 모두 덮는다. `chunked_prefill_reserved_blocks`는 이제 다음 청크의 패딩을 포함한다. +- 승인 워터마크(`admission_watermark.rs`, `--kv-admission-watermark`, 환경 변수 `MLXCEL_KV_ADMISSION_WATERMARK`, 범위 0.0~0.5): 디코드 배치가 살아 있는 동안 승인은 빈 블록에서 `ceil(F * budget)`을 뺀 값이 요청을 덮어야 하며, 미루기 전에 그 목표까지 회수한다. 콜드 접두사를 먼저 축출하고, 더 높은 우선순위 요청이면 기존 승인과 같이 엄격히 낮은 우선순위의 행을 선점한다. 배치가 비어 있으면 적용되지 않으므로 예산에 들어가는 요청이 거부되는 일은 없다. `paged_prefill_room`이 이를 적용하므로 배치 윈도와 단일 시퀀스 게이트가 같은 수치로 계산한다. +- 채택 접두사 계산(`admission.rs`, `CachePool::paged_blocks_to_reach`): 승인은 시퀀스의 paged 상태 기준으로 접미사가 만들 블록만 계산하고, 채택 후에는 레이어당 최소 1블록을 계산한다(프롬프트 전체가 적중하면 마지막 토큰을 공유 꼬리 블록에 다시 써서 fork가 일어난다). 이전에는 예산의 절반보다 큰 접두사를 채택한 요청은 배치가 비어 있어도 승인될 수 없었다. +- 대기 채택 회수(`queued_adoption.rs`, `PrefillQueue::find_lowest_first`/`get_mut`/`remove`): 채택은 enqueue 시점에 일어나 요청이 기다리는 동안 접두사 블록을 고정한다. 디코드와 청크 회수는 이제 콜드 접두사 축출 다음, 선점 전에 대기 중인 텍스트 요청 하나(가장 낮은 우선순위, 가장 최근 것부터)의 채택을 버린다. 그 요청은 콜드 프리필을 하고 donate-back을 위한 프롬프트 캐시 컨텍스트는 유지한다. 승인 경로는 이를 하지 않는다. 대기 요청은 다른 요청의 캐시 적중을 빼앗는 대신 기다린다. 채택한 VLM 요청은 전체 프롬프트 임베딩을 더 이상 갖고 있지 않으므로 멀티모달 요청은 건너뛴다. + +## 검증 + +- 단위 테스트: `paged_detach_tests.rs`에 레이어별 사례 3개를 추가했다. 어긋난 상태(레이어 0만 빈 블록 2개, 레이어 1은 0개, 예산 전부 mint됨)에서 이전 수치는 2인데 레이어 1 append가 실패하고, 새 수치는 0이다. 여유분 2블록이 있으면 수치는 4이고 정확히 그만큼 얻을 수 있다. `paged_append_need_tests.rs`에 접미사 계산을 추가했다. `block_reclaim_tests.rs`에 MixedStep 틱 순서 사례(디코드가 미루고 청크가 들어감, 블록 중간 행들이 풀을 예약 아래로 떨어뜨리면 예약 없이는 청크 append가 실패하고 행 하나를 선점한 뒤 성공함, 행보다 콜드 접두사를 먼저 축출함), 워터마크 미루기 및 축출 사례, 대기 채택 사례(경계의 단독 행이 대기 채택을 버려 살아남음, 순서는 접두사, 채택, 행, 승인은 채택을 버리지 않음)를 추가했다. `scheduler_prompt_cache_tests.rs`에 실제 풀에서의 채택 접두사 계산을 추가했다. +- `server::batch`, `server::prompt_cache`, `runtime_settings`, `cli_input`, `config_tests`, `commands::serve` 필터: 822개 통과. mlxcel-core `paged`, `budget` 필터: 282개 통과. `cargo clippy --release --features cuda --lib --tests -- -D warnings`(루트와 `-p mlxcel-core`)와 `cargo fmt --all` 통과. +- GB10 압박 실행(`--kv-cache-budget 600000000`, 4444블록, paged, `--max-batch-size 4`, 접두사를 먼저 하나 고정한 뒤 600토큰 생성의 3턴 대화 3개를 동시에): + +| 빌드 | 워터마크 | 선점 | 디코드 회수 | 턴 | 오류 | +|------|----------|------|-------------|----|------| +| main (#2077 보고서) | 해당 없음 | 12 | 16 | 9 | 없음 | +| 결함 1-3만 | 0 | 12 | 16 | 9 | 없음 | +| 결함 1-3만 | 0.15 | 2 | 4 | 4 | 행 2개 탈락, 요청 1개 900초 정지 | +| 이 PR (2회) | 0 | 10 | 43 | 9 | 없음 | +| 이 PR (2회) | 0.01 | 6 | 6 | 9 | 없음 | +| 이 PR (2회) | 0.02 | 6 | 6 | 9 | 없음 | +| 이 PR (2회) | 0.05 | 5 | 5 | 9 | 없음 | +| 이 PR | 0.10 | 1 | 2 | 9 | 없음 | +| 이 PR | 0.15 | 1 | 3 | 9 | 없음 | + + 마지막 두 수정 전 0.15 실패: 두 탈락 모두 배치에 행이 하나뿐이고 축출할 것이 없는 상태에서 `preempted=0`이었으므로 남은 블록은 대기 채택이 쥐고 있었다. 정지는 예산의 절반보다 큰 접두사를 채택한 요청이 빈 배치에서 영원히 미뤄진 것이다(중단한 0.20 실행에서도 같은 양상으로 워커가 CPU 96%로 공회전했다). 기본값 0.01은 턴을 잃지 않고 선점을 줄인 측정 격자의 가장 작은 값이다. 0에서 이 브랜치의 회수 43회 중 33회는 대기 청크의 예약만큼만 모자란 행의 한 틱 미루기다. +- `MLXCEL_MIXED_STEP=1`, 0.01, `--metrics`: mixed step 21회, 프리필 청크 55개, 9턴 모두 완료, panic 없음, 선점 7회. +- 기본 실행(auto 예산, 753620블록), 0.01, 3턴 대화 2개: 회수와 선점 없음, 2턴과 3턴이 프롬프트 캐시에 적중(1280, 1504토큰). + +## 검증하지 못한 것 + +- 레이어별 수치는 refcount가 0인 블록이 정확히 하나의 레이어 free list에 있다는 불변식에 의존한다. 이를 확인하는 `debug_assert!`는 테스트가 모두 release 빌드여서 실행되지 않았다. 대신 모든 변경 지점을 읽어 확인했다. `release_block`은 0이 되면 push하고, `acquire_block`과 `restore_block`은 refcount를 1로 만들기 전에 꺼내거나 제거하며, `retain_block`은 refcount 0 블록을 거부하고, `blocks`에서 레코드를 지우는 경로는 없다. + +- 청크 예약의 회수 경로는 실제 서버 실행에서 한 번도 동작하지 않았다(`Chunked prefill reclaimed` 로그 없음). 단위 테스트로만 확인했다. +- 대기 채택 버리기는 0.10과 0.15에서 각각 한 번 동작했다. 버려진 요청의 재프리필은 해당 실행이 모든 턴을 마친 것으로만 확인했다. +- 패딩 청크 예약은 Metal M5+(`should_align_prefill`)에서만 적용되며 이 호스트에서는 실행할 수 없다. +- 워터마크 기본값은 모델 하나, 예산 하나로 측정했다. auto 예산에서는 행이 디코드하는 동안 풀의 1%를 남긴다. diff --git a/docs/CONTINUOUS_BATCHING.md b/docs/CONTINUOUS_BATCHING.md index 3aa379e9e..4c46d3fcb 100644 --- a/docs/CONTINUOUS_BATCHING.md +++ b/docs/CONTINUOUS_BATCHING.md @@ -22,6 +22,7 @@ token, then streams the new tokens out. Relevant flags: | `--max-queue-depth N` | 32 | Maximum queued (not yet admitted) requests. | | `--prefill-chunk-size N` | 512 | Token chunk size for prefill; bounds prefill's effect on decode latency. | | `--prefill-grant-interval N` | 16 | Decode ticks a parked chunked prefill yields before it is granted one; bounds an admitted long prompt's time to first token. `0` disables the grant (unbounded wait). | +| `--kv-admission-watermark F` | 0.01 | Fraction of the paged KV block budget (`0.0` to `0.5`) a prefill admission keeps free while rows are decoding, so they can grow without being preempted. `0` disables it. See below. | | `--enable-preemption` | off | Allow evicting a lower-priority sequence to admit a waiting one. | | `--no-batch` | off | Disable batching and serve sequentially (the legacy single worker). | | `--no-prompt-cache` | off | Disable the prompt-prefix KV cache (it is on by default). | @@ -139,6 +140,46 @@ budget is inert. Disable the guard with `--kv-cache-budget none`. Memory- constrained hosts can also lower `--parallel` or cap `--ctx-size` (see the context-sizing note in [environment-variables.md](environment-variables.md)). +### Decode headroom under a tight budget (`--kv-admission-watermark`) + +When the budget is small relative to the concurrent work, admitting a request +into exactly the free blocks leaves the running rows nowhere to grow. The next +decode ticks then preempt the least-progressed row, which re-prefills from +scratch when it is admitted again. `--kv-admission-watermark F` keeps the +fraction `F` of the block budget free at admission: a request is admitted only +when the free blocks minus `F * budget` cover it. To make that room it evicts +cold prompt-cache prefixes and, for a higher-priority request, preempts +strictly lower-priority rows (as admission already does for the request +itself) before it waits. The watermark applies +only while a decode batch is live, so a request that fits the budget is never +refused; it waits for running rows to finish instead. `0` disables it, and it +is inert without a block budget. + +The default, 0.01, is the smallest value on the measured grid that lowered +preemptions without losing a turn. On GB10 with meta-llama-3.1-8b-instruct-4bit, +`--kv-cache-budget 600000000` (4444 blocks), `--max-batch-size 4`, and three +concurrent 3-turn conversations of 600-token generations, each value run twice +with identical results: + +| Watermark | Preemptions | Decode reclaim passes | Turns completed | Wall time | +|-----------|-------------|-----------------------|-----------------|-----------| +| 0 | 10 | 43 | 9 / 9 | 103-104 s | +| 0.01 | 6 | 6 | 9 / 9 | 90-91 s | +| 0.02 | 6 | 6 | 9 / 9 | 89-90 s | +| 0.05 | 5 | 5 | 9 / 9 | 89-90 s | +| 0.10 | 1 | 2 | 9 / 9 | 86 s (one run) | +| 0.15 | 1 | 3 | 9 / 9 | 88 s (one run) | + +Raise it when a tight budget preempts often and the extra queueing latency +for new requests is acceptable. + +Two more rules keep the budget accounting honest (issue #2088). A request +that adopted a cached prefix is charged only for the blocks its suffix will +mint, since the adopted blocks are already in use. When a running row needs a +block and none is free, the reclaim order is: evict a cold prompt-cache +prefix, then drop a queued request's adopted prefix (it re-prefills cold when +admitted), then preempt a running row. + ### How a chunked prefill shares ticks with decode (`--prefill-grant-interval`) The flag table above says `--prefill-chunk-size` "bounds prefill's effect on diff --git a/docs/environment-variables.md b/docs/environment-variables.md index 655d1673f..2e03ac38e 100644 --- a/docs/environment-variables.md +++ b/docs/environment-variables.md @@ -124,6 +124,7 @@ Both server entry points implement llama-server b10621's Vertex AI custom-contai | `MLXCEL_DEFAULT_ORG` | HuggingFace org/user name | `mlx-community` | Org prepended to a bare, prefix-less model name (no `/`) by the `-m/--model` resolver (`generate` / `serve` / `inspect` / `run`), the `mlxcel-server -m` resolver, and the `download` verb (`mlxcel download` / `mlx-server download`), so `mlxcel run Qwen3-4B-4bit` resolves to `mlx-community/Qwen3-4B-4bit` and `mlxcel download Qwen3-4B-4bit` downloads that same repo. An explicit `owner/name` repo-id and an existing local path are unaffected. An empty/whitespace value falls back to `mlx-community`. | | `MLXCEL_SERVER_DECODE_STORAGE` | `auto`, `dense`, `paged` | `auto` | Server continuous-batching decode storage. `--decode-storage-backend` takes precedence. Invalid values warn and fall back to `auto`. | | `MLXCEL_KV_CACHE_BUDGET` | `auto`, unsigned integer bytes, or `none`/`0` | `auto` | Paged KV block-pool budget for continuous batching. `--kv-cache-budget` takes precedence. Defaults to `auto` (#628): it pairs with the batched-decode default so admission caps KV for the concurrent batch and returns backpressure instead of an OOM abort. Applies to pool-backed Fp16 caches under the paged decode backend (the `--parallel > 1` default); inert on the dense backend. `none` / `0` leaves the pool unbounded. | +| `MLXCEL_KV_ADMISSION_WATERMARK` | fraction `0.0` to `0.5` | `0.01` | Paged KV admission watermark (issue #2088): the share of the block budget a prefill admission keeps free while rows are decoding, so they can grow without preemption. Admission evicts cold prompt-cache prefixes, and preempts strictly lower-priority rows for a higher-priority request, to make that room before it waits. Applies only while a decode batch is live and only with a block budget; `0` disables it. Out-of-range or non-numeric values are rejected at startup. `--kv-admission-watermark` takes precedence. See [CONTINUOUS_BATCHING.md](CONTINUOUS_BATCHING.md). | | `MLXCEL_PREFILL_CHUNK` | unsigned integer tokens, `0` disables | `2048` | Cache-level chunked prefill for the single-sequence `mlxcel generate` / `mlxcel-bench-decode` path: the prompt is fed through the model in chunks of this many tokens, so sliding-window KV caches trim between chunks and prefill peak memory stays near the chunk size instead of scaling with the whole prompt (issues #672/#674). `0` forces the previous single-pass prefill. Models can opt out of multi-call prefill via `LanguageModel::supports_chunked_prefill`. The server path is controlled by `--prefill-chunk-size` instead. | | `MLXCEL_MAX_BATCH_PREFILL_TOKENS` | unsigned integer tokens, `0` disables | derived (`2 * max_batch_prefill * prefill_chunk_size`) | Padded-token budget bounding one server batched prefill's transient memory (issue #715). The batched path pads a cohort to its longest prompt `L` and materializes a `[B, L, L]` FP32 mask, an `O(B*L^2)` transient; this caps the drained window by total padded tokens (`rows * L`) so the mask stays within `~2*N^2` bytes, with rows past the budget spilling to the chunked single-sequence path. `--max-batch-prefill-tokens` takes precedence. Unset derives `2 * max_batch_prefill * prefill_chunk_size` (the shipped `2 * 4 * 512 = 4096`; the 2x headroom keeps a full batch of slightly-over-chunk-sized prompts in one window); `0` disables the cap (pre-#715 unbounded behavior). Only affects families with `supports_batched_prefill()` under `--max-batch-prefill > 1`. | | `MLXCEL_PREFILL_GRANT_INTERVAL` | unsigned integer decode ticks, `0` disables | `16` | Prefill fairness grant for the continuous-batching scheduler (issue #1011). A prompt longer than `--prefill-chunk-size` admitted next to a busy decode batch runs one chunk and is then parked; this is how many consecutive decode ticks it yields before the scheduler grants it one. A `C`-chunk prompt therefore reaches its first token within `C * (N + 1)` ticks however long the batch keeps decoding, where before #1011 it waited for the batch to drain and its time to first token had no bound at all. The decoding streams pay for that bound: over one grant cycle they get `N` tokens per `N * D + P` of wall clock for a decode step `D` and a chunk forward `P`, so their mean inter-token latency during the admission window is `D + P / N`. Lower is faster to first token and noisier for everyone else. `0` restores the pre-#1011 arbitration exactly, unbounded wait included. An unparseable value falls back to the default rather than to `0`, so a typo cannot silently reinstate the starvation. `--prefill-grant-interval` takes precedence. Read once at scheduler construction. Attribute a run with `mlxcel_batch_prefill_grants_total` on `/metrics`; see `scripts/bench_mixed_step_admission.py` and `scripts/bench/starvation_probe.sh`. | diff --git a/src/bin/mlx_server.rs b/src/bin/mlx_server.rs index 17e9be459..e68adc44a 100644 --- a/src/bin/mlx_server.rs +++ b/src/bin/mlx_server.rs @@ -976,6 +976,27 @@ struct ServerArgs { )] kv_cache_budget: Option, + /// Fraction of the paged KV block budget a prefill admission keeps free + /// for running rows to keep decoding (issue #2088), in `0.0..=0.5`. + /// + /// Under a tight `--kv-cache-budget`, admitting a request into exactly the + /// free blocks leaves the decoding rows no room to grow, so the next decode + /// ticks preempt them and they re-prefill from scratch. With a watermark, + /// admission waits until the free blocks cover the request plus this share + /// of the budget, evicting cold prompt-cache prefixes (and, for a + /// higher-priority request, preempting strictly lower-priority rows) to + /// make that room. It applies only + /// while rows are decoding, so a request that fits the budget is never + /// refused. `0` disables it. Ignored without a block budget. Defaults to + /// 0.01. Also reads `MLXCEL_KV_ADMISSION_WATERMARK`. + #[arg( + long = "kv-admission-watermark", + env = "MLXCEL_KV_ADMISSION_WATERMARK", + value_name = "FRACTION", + value_parser = mlxcel::server::parse_kv_admission_watermark + )] + kv_admission_watermark: Option, + /// Maximum number of responses persisted by the OpenAI /// `/v1/responses` store (in-memory). `0` disables persistence /// entirely. Also reads `LLAMA_ARG_RESPONSES_STORE_MAX_ENTRIES`. @@ -2707,6 +2728,8 @@ fn build_startup_input(mut args: ServerArgs) -> anyhow::Result anyhow::Result crate::ServeArgs { ctx_size: 8192, max_kv_size: 0, kv_cache_budget: None, + kv_admission_watermark: None, enable_vlm_prefix_cache: false, allowed_origins: Vec::new(), n_predict: 128, diff --git a/src/lib/mlxcel-core/src/cache.rs b/src/lib/mlxcel-core/src/cache.rs index 0f6ff9774..932b75ca4 100644 --- a/src/lib/mlxcel-core/src/cache.rs +++ b/src/lib/mlxcel-core/src/cache.rs @@ -7008,8 +7008,9 @@ impl CachePool { self.paged_slab_blocks } - /// Blocks still **acquirable** before the paged budget is hit - /// (`budget − live`), or `None` when unbounded / no paged pool. `Some(0)` + /// Blocks still **acquirable** before the paged budget is hit, counted + /// per layer (see [`PagedBlockPool::free_block_budget`]), or `None` when + /// unbounded / no paged pool. `Some(0)` /// means every budgeted block is in use; the admission gate must reclaim /// (evict cold prefixes, then preempt) or defer. Eviction raises this even /// though allocated rows are retained, because freed blocks are reusable. @@ -7033,6 +7034,18 @@ impl CachePool { pool.borrow().blocks_to_append(&state.borrow(), token_count) } + /// Pool blocks a prefill growing every layer of sequence `id` to + /// `total_len` positions would acquire (see + /// [`PagedBlockPool::blocks_to_reach`]). A sequence that adopted a cached + /// prefix is charged only for its suffix. `None` when `id` has no paged + /// state or no paged pool exists, so callers can fall back to a + /// whole-prompt estimate (issue #2088). + pub fn paged_blocks_to_reach(&self, id: SequenceId, total_len: usize) -> Option { + let pool = self.paged_pool.as_ref()?; + let state = self.active.get(&id).and_then(|seq| seq.paged.as_ref())?; + Some(pool.borrow().blocks_to_reach(&state.borrow(), total_len)) + } + /// Read-only access to the underlying [`PagedBlockPool`]. /// /// Returns a [`Ref`] guard; release it before any path that mutates the diff --git a/src/lib/mlxcel-core/src/cache/paged.rs b/src/lib/mlxcel-core/src/cache/paged.rs index f8a26b8f6..08817e406 100644 --- a/src/lib/mlxcel-core/src/cache/paged.rs +++ b/src/lib/mlxcel-core/src/cache/paged.rs @@ -801,19 +801,61 @@ impl PagedBlockPool { } /// Physical blocks still **acquirable** before the budget is hit, or `None` - /// when unbounded. This is `budget − live_block_count`: a fresh - /// `acquire_block` succeeds whenever fewer than `budget` blocks are live, - /// because it can either reuse a freed (allocated-but-not-live) block or - /// mint a new one while `allocated < budget`. Eviction / preemption that - /// drops a block's refcount to 0 therefore *raises* this figure even though - /// `allocated_block_count` is unchanged (the freed row is retained for - /// reuse). The scheduler gates prefill admission on this value and reclaims - /// (evict cold prefixes, then preempt) until it covers the sequence's need. - /// `Some(0)` means every budgeted block is in use — admission must reclaim - /// or defer. + /// when unbounded, counted the way the scheduler spends them: the same + /// number of blocks on every layer (issue #2088). + /// + /// A freed block keeps its layer: [`Self::release_block`] pushes it onto + /// that layer's free list, and [`Self::acquire_block`] for layer `L` reuses + /// only `free_lists[L]` before minting a new block, which the budget caps + /// pool-wide. The pool-wide `budget - live` figure therefore over-reports + /// as soon as the free lists diverge: with every free block on layer 0 and + /// nothing left to mint, a layer-1 acquire fails while `budget - live` is + /// positive. Re-tagging a free block to another layer is not an option, + /// because pool rows are per-layer slabs and the moved block would need a + /// new row on its new layer, so the budget would stop bounding memory. + /// + /// Every scheduler demand is uniform across layers (an admission charges + /// `ceil(len / block_size) * num_layers`, a decode tick sums + /// [`Self::blocks_to_append`] over equal-length layers), so this returns + /// `n * num_layers` for the largest `n` such that every layer can acquire + /// `n` blocks: the free blocks a layer already owns, plus a share of the + /// mint headroom `budget - allocated`. With uniform free lists it equals + /// `budget - live` rounded down to a multiple of `num_layers`, which + /// changes no uniform comparison. Layers that have never minted a block + /// (none yet, or a layer the model never writes) are left out of the + /// per-layer constraint, and the result never exceeds `budget - live`. + /// + /// Eviction / preemption that drops a block's refcount to 0 raises this + /// figure even though `allocated_block_count` is unchanged (the freed row + /// is retained for reuse). The scheduler gates prefill admission and the + /// decode reservation on this value and reclaims (evict cold prefixes, then + /// preempt) until it covers the need. `Some(0)` means no further uniform + /// step fits; admission must reclaim or defer. pub fn free_block_budget(&self) -> Option { - self.block_budget - .map(|max| max.saturating_sub(self.live_block_count())) + let max = self.block_budget?; + let num_layers = self.layout.num_layers; + let mut minted = vec![false; num_layers]; + let mut live = 0usize; + for record in self.blocks.values() { + if let Some(flag) = minted.get_mut(record.layer_idx) { + *flag = true; + } + if record.is_in_use() { + live += 1; + } + } + debug_assert_eq!( + live + self.free_lists.iter().map(Vec::len).sum::(), + self.blocks.len(), + "every refcount-0 block sits on exactly one layer free list" + ); + let pool_wide = max.saturating_sub(live); + let free_per_layer: Vec = (0..num_layers) + .filter(|&layer| minted[layer]) + .map(|layer| self.free_lists[layer].len()) + .collect(); + let headroom = max.saturating_sub(self.blocks.len()); + Some(uniform_acquirable_blocks(&free_per_layer, headroom, num_layers).min(pool_wide)) } /// Blocks appending `token_count` tokens to every layer of `state` would @@ -827,25 +869,40 @@ impl PagedBlockPool { /// a decode tick's rows and reclaims budget before the forward runs, since /// a failed acquire inside the forward is fatal. pub fn blocks_to_append(&self, state: &PagedSequenceState, token_count: usize) -> usize { - if token_count == 0 { - return 0; - } - let block_size = self.layout.block_size.max(1); state .layers .iter() - .map(|layer| { - let required = (layer.len + token_count).div_ceil(block_size); - let fresh = required.saturating_sub(layer.block_ids.len()); - let cow = layer - .block_ids - .get(layer.len / block_size) - .is_some_and(|&tail| self.refcount(tail) > 1); - fresh + usize::from(cow) - }) + .map(|layer| self.layer_blocks_to_append(layer, token_count)) .sum() } + /// Blocks growing every layer of `state` to `total_len` positions would + /// acquire (issue #2088): [`Self::blocks_to_append`] with each layer's own + /// remaining count. Prefill admission charges this instead of the whole + /// prompt, because a sequence that adopted a cached prefix already holds + /// those blocks (they are live) and only mints the suffix. + pub fn blocks_to_reach(&self, state: &PagedSequenceState, total_len: usize) -> usize { + state + .layers + .iter() + .map(|layer| self.layer_blocks_to_append(layer, total_len.saturating_sub(layer.len))) + .sum() + } + + fn layer_blocks_to_append(&self, layer: &PagedLayerState, token_count: usize) -> usize { + if token_count == 0 { + return 0; + } + let block_size = self.layout.block_size.max(1); + let required = (layer.len + token_count).div_ceil(block_size); + let fresh = required.saturating_sub(layer.block_ids.len()); + let cow = layer + .block_ids + .get(layer.len / block_size) + .is_some_and(|&tail| self.refcount(tail) > 1); + fresh + usize::from(cow) + } + /// Whether the pool's cache mode requires Turbo4 sidecar storage. /// /// Used by: paged detach/adopt to decide whether to round-trip sidecars, @@ -2550,6 +2607,46 @@ impl PagedBlockPool { } } +/// Blocks a uniform per-layer demand can acquire (issue #2088): `n * +/// num_layers` for the largest `n` such that the layers in `free_per_layer` +/// can each take `n` blocks, using their own free blocks first and minting +/// the rest from the shared `headroom`. An empty `free_per_layer` (no layer +/// has minted a block yet) leaves only the headroom as the limit. +fn uniform_acquirable_blocks( + free_per_layer: &[usize], + headroom: usize, + num_layers: usize, +) -> usize { + if free_per_layer.is_empty() { + return headroom; + } + // Blocks to mint so that every constrained layer holds `n` free blocks. + let mint_cost = |n: usize| -> usize { + free_per_layer + .iter() + .map(|&free| n.saturating_sub(free)) + .fold(0usize, usize::saturating_add) + }; + // `mint_cost(hi + 1) > headroom` because the layer with the most free + // blocks alone needs `headroom + 1` mints there. + let mut lo = 0usize; + let mut hi = free_per_layer + .iter() + .copied() + .max() + .unwrap_or(0) + .saturating_add(headroom); + while lo < hi { + let mid = lo + (hi - lo).div_ceil(2); + if mint_cost(mid) <= headroom { + lo = mid; + } else { + hi = mid - 1; + } + } + lo.saturating_mul(num_layers) +} + /// Normalize a write block into the layout-A slot update shape /// `[1, n_slots, n_kv_heads, head_dim]` and return its geometry /// `(n_slots, n_kv_heads, head_dim)`. diff --git a/src/lib/mlxcel-core/src/cache/paged_append_need_tests.rs b/src/lib/mlxcel-core/src/cache/paged_append_need_tests.rs index 8496eb82a..feb6ca8c5 100644 --- a/src/lib/mlxcel-core/src/cache/paged_append_need_tests.rs +++ b/src/lib/mlxcel-core/src/cache/paged_append_need_tests.rs @@ -81,3 +81,21 @@ fn shared_partial_tail_counts_the_copy_on_write_fork() { pool.retain_block(aligned.layers[0].block_ids[1]).unwrap(); assert_eq!(pool.blocks_to_append(&aligned, 1), 2); } + +#[test] +fn reaching_a_prompt_length_charges_only_the_suffix() { + // Issue #2088: a sequence that adopted a cached prefix of 9 tokens (3 + // blocks per layer, the last one partial) holds those blocks already, so + // prefilling a 13-token prompt mints one block per layer, not the 4 per + // layer the whole prompt would need. + let layout = layout(); + let mut pool = PagedBlockPool::new(layout.clone()); + let state = state_with_len(&mut pool, &layout, 9); + assert_eq!(pool.blocks_to_reach(&state, 13), 2); + assert_eq!(pool.blocks_to_reach(&state, 12), 0); + assert_eq!(pool.blocks_to_reach(&state, 9), 0); + assert_eq!(pool.blocks_to_reach(&state, 4), 0, "shorter than held"); + // A fresh sequence is charged the whole prompt. + let fresh = PagedSequenceState::new(&layout); + assert_eq!(pool.blocks_to_reach(&fresh, 13), 8); +} diff --git a/src/lib/mlxcel-core/src/cache/paged_detach_tests.rs b/src/lib/mlxcel-core/src/cache/paged_detach_tests.rs index e42009ea7..11597ea30 100644 --- a/src/lib/mlxcel-core/src/cache/paged_detach_tests.rs +++ b/src/lib/mlxcel-core/src/cache/paged_detach_tests.rs @@ -345,6 +345,74 @@ fn free_block_budget_rises_when_blocks_are_freed() { ); } +/// Two layers whose free lists diverge (issue #2088): sequence `a` writes 2 +/// blocks on layer 0 only and is released, sequence `b` keeps 2 live blocks +/// on layer 1. Layer 0 then owns 2 free blocks and layer 1 owns none. +fn diverged_free_lists(budget: usize) -> (PagedBlockPool, PagedKvLayout) { + let layout = PagedKvLayout::uniform(2, 4, 128).unwrap(); + let mut pool = PagedBlockPool::new(layout.clone()); + pool.set_block_budget(Some(budget)); + let mut a = PagedSequenceState::new(&layout); + pool.append_tokens(&mut a, 0, 8).unwrap(); + let mut b = PagedSequenceState::new(&layout); + pool.append_tokens(&mut b, 1, 8).unwrap(); + pool.release_sequence(&mut a).unwrap(); + assert_eq!(pool.allocated_block_count(), 4); + assert_eq!(pool.live_block_count(), 2); + (pool, layout) +} + +#[test] +fn free_block_budget_counts_per_layer_free_lists() { + // Budget 4 is fully minted, so layer 1 can neither reuse a free block (all + // free ones belong to layer 0) nor mint one. The pool-wide `budget - live` + // figure is 2 here, which made admission and decode reservation believe a + // layer-1 write would succeed; the per-layer figure reports no room. + let (mut pool, layout) = diverged_free_lists(4); + let mut fresh = PagedSequenceState::new(&layout); + let err = pool.append_tokens(&mut fresh, 1, 1).unwrap_err(); + assert!(err.contains("budget exhausted"), "got: {err}"); + assert_eq!( + pool.free_block_budget(), + Some(0), + "no uniform per-layer step fits while layer 1 cannot acquire" + ); +} + +#[test] +fn free_block_budget_on_diverged_free_lists_is_acquirable() { + // Budget 6 leaves 2 blocks of mint headroom: layer 0 reuses its 2 free + // blocks and layer 1 mints 2, so 2 blocks per layer (4 total) are + // acquirable. The figure must match exactly what the pool hands out. + let (mut pool, layout) = diverged_free_lists(6); + assert_eq!(pool.free_block_budget(), Some(4)); + let mut fresh = PagedSequenceState::new(&layout); + for layer in 0..2 { + pool.append_tokens(&mut fresh, layer, 8) + .expect("the reported blocks are acquirable on every layer"); + } + assert_eq!(pool.free_block_budget(), Some(0)); + assert!(pool.append_tokens(&mut fresh, 1, 1).is_err()); +} + +#[test] +fn free_block_budget_on_uniform_free_lists_matches_budget_minus_live() { + // Uniform layers: the per-layer figure equals the historical + // `budget - live` whenever that is a multiple of the layer count. + let layout = PagedKvLayout::uniform(2, 4, 128).unwrap(); + let mut pool = PagedBlockPool::new(layout.clone()); + pool.set_block_budget(Some(10)); + let mut a = PagedSequenceState::new(&layout); + let mut b = PagedSequenceState::new(&layout); + for layer in 0..2 { + pool.append_tokens(&mut a, layer, 8).unwrap(); + pool.append_tokens(&mut b, layer, 4).unwrap(); + } + pool.release_sequence(&mut a).unwrap(); + assert_eq!(pool.live_block_count(), 2); + assert_eq!(pool.free_block_budget(), Some(8)); +} + #[test] fn cache_pool_budget_set_before_pool_creation_applies_on_creation() { // The scheduler may set the budget before the first paged allocation lazily diff --git a/src/main.rs b/src/main.rs index 2cd8c1caa..7ee2df372 100644 --- a/src/main.rs +++ b/src/main.rs @@ -1794,6 +1794,27 @@ pub(crate) struct ServeArgs { )] kv_cache_budget: Option, + /// Fraction of the paged KV block budget a prefill admission keeps free + /// for running rows to keep decoding (issue #2088), in `0.0..=0.5`. + /// + /// Under a tight `--kv-cache-budget`, admitting a request into exactly the + /// free blocks leaves the decoding rows no room to grow, so the next decode + /// ticks preempt them and they re-prefill from scratch. With a watermark, + /// admission waits until the free blocks cover the request plus this share + /// of the budget, evicting cold prompt-cache prefixes (and, for a + /// higher-priority request, preempting strictly lower-priority rows) to + /// make that room. It applies only + /// while rows are decoding, so a request that fits the budget is never + /// refused. `0` disables it. Ignored without a block budget. Defaults to + /// 0.01. Also reads `MLXCEL_KV_ADMISSION_WATERMARK`. + #[arg( + long = "kv-admission-watermark", + env = "MLXCEL_KV_ADMISSION_WATERMARK", + value_name = "FRACTION", + value_parser = mlxcel::server::parse_kv_admission_watermark + )] + kv_admission_watermark: Option, + /// Maximum number of responses persisted by the OpenAI /// `/v1/responses` store (in-memory). `0` disables persistence /// entirely, in which case `GET /v1/responses/:id` and diff --git a/src/server/batch/queue.rs b/src/server/batch/queue.rs index 3beb5933c..5d6f4e679 100644 --- a/src/server/batch/queue.rs +++ b/src/server/batch/queue.rs @@ -23,6 +23,8 @@ use std::collections::VecDeque; use std::sync::atomic::Ordering; +use mlxcel_core::cache::SequenceId; + use super::sequence::{RequestPriority, SequenceInfo}; /// Default maximum queue depth when no explicit limit is given. @@ -226,6 +228,41 @@ impl PrefillQueue { self.max_size } + /// Id of the first queued sequence matching `predicate`, scanning from the + /// lowest-priority lane's newest entry toward the highest-priority lane's + /// oldest, so the request served last is found first (issue #2088). + pub fn find_lowest_first(&self, mut predicate: F) -> Option + where + F: FnMut(&SequenceInfo) -> bool, + { + self.low + .iter() + .rev() + .chain(self.normal.iter().rev()) + .chain(self.high.iter().rev()) + .find(|seq| predicate(seq)) + .map(|seq| seq.seq_id) + } + + /// Mutable access to the queued sequence `id`, if it is queued. + pub fn get_mut(&mut self, id: SequenceId) -> Option<&mut SequenceInfo> { + self.high + .iter_mut() + .chain(self.normal.iter_mut()) + .chain(self.low.iter_mut()) + .find(|seq| seq.seq_id == id) + } + + /// Remove and return the queued sequence `id`, if it is queued. + pub fn remove(&mut self, id: SequenceId) -> Option { + for lane in [&mut self.high, &mut self.normal, &mut self.low] { + if let Some(pos) = lane.iter().position(|seq| seq.seq_id == id) { + return lane.remove(pos); + } + } + None + } + /// Remove and return all queued sequences whose cancellation flag is set. /// /// Scans all priority lanes and removes sequences where diff --git a/src/server/batch/scheduler/admission.rs b/src/server/batch/scheduler/admission.rs index 2ea67757f..85cf5442f 100644 --- a/src/server/batch/scheduler/admission.rs +++ b/src/server/batch/scheduler/admission.rs @@ -688,12 +688,44 @@ impl BatchScheduler { } } + /// Pool blocks `seq`'s prefill will mint (issue #2088). A sequence that + /// adopted a cached prefix at enqueue already holds that prefix's blocks, + /// and they count as live, so charging the whole prompt counted them twice: + /// a request whose adopted prefix exceeded half the budget could never be + /// admitted, even into an empty batch. The charge is the suffix measured + /// from the sequence's paged state, and at least one block per layer after + /// an adoption, because a hit covering the whole prompt re-runs its last + /// token into the shared tail block, which forks it. Falls back to the + /// whole-prompt estimate for a sequence without paged state. The batched + /// window keeps the peek-based estimate: over-charging there only ends the + /// window early and routes the head through this gate. + pub(super) fn prefill_block_need(&self, seq: &SequenceInfo) -> usize { + let estimate = self.estimate_prefill_blocks(seq.prompt_tokens.len()); + if estimate == 0 { + return 0; + } + let Some(suffix) = self + .cache_pool + .paged_blocks_to_reach(seq.seq_id, seq.prompt_tokens.len()) + else { + return estimate; + }; + if seq.already_cached_tokens > 0 { + suffix.max(self.model.num_layers()) + } else { + suffix + } + } + /// Paged pool blocks a new prefill may still pin without reclaiming, or /// `None` when no block budget is configured. Before the pool is created - /// nothing is allocated, so the whole budget is free. + /// nothing is allocated, so the whole budget is free. The admission + /// watermark (issue #2088) is already taken off, so the single-sequence + /// gate and the batched-prefill window charge against the same figure. pub(super) fn paged_prefill_room(&self) -> Option { let total = self.cache_pool.paged_block_budget()?; - Some(self.available_paged_blocks().unwrap_or(total)) + let free = self.available_paged_blocks().unwrap_or(total); + Some(free.saturating_sub(self.paged_admission_headroom())) } /// Paged block-budget admission gate. Returns `Some(seq)` to proceed with @@ -706,9 +738,9 @@ impl BatchScheduler { Some(t) => t, None => return Some(seq), }; - let need = self.estimate_prefill_blocks(seq.prompt_tokens.len()); + let need = self.prefill_block_need(&seq); if need == 0 { - return Some(seq); // model does not use the paged pool + return Some(seq); // no paged pool, or nothing left to write } // If it cannot fit the entire budget, reject — deferring forever would // wedge the queue behind a request that can never run. @@ -721,17 +753,16 @@ impl BatchScheduler { ); return None; } - // Acquirable blocks (budget − live). `None` means the pool is not yet - // created (nothing allocated ⇒ the whole budget is free). - let free = self.available_paged_blocks().unwrap_or(total); - if need <= free { + // Acquirable blocks minus the admission watermark (issue #2088). + if need <= self.paged_prefill_room().unwrap_or(total) { return Some(seq); } if self.reclaim_paged_blocks(need, seq.priority) { return Some(seq); } - // Still no room — defer to a later tick. Decodes in flight will free - // blocks as their sequences finish; this request retries then. + // Still no room (or no room beyond the watermark): defer to a later + // tick. Decodes in flight free blocks as their sequences finish, and + // the watermark lifts once the batch is empty; this request retries. if let Err(rejected) = self.prefill_queue.enqueue(seq) { self.prompt_cache_seq_ctx.remove(&rejected.seq_id); self.release_sequence_caches(rejected.seq_id); diff --git a/src/server/batch/scheduler/admission_watermark.rs b/src/server/batch/scheduler/admission_watermark.rs new file mode 100644 index 000000000..02cfe148c --- /dev/null +++ b/src/server/batch/scheduler/admission_watermark.rs @@ -0,0 +1,80 @@ +// Copyright 2025-2026 Lablup Inc. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +//! Paged KV admission watermark (issue #2088). +//! +//! Under a tight `--kv-cache-budget`, admitting a prefill into exactly the +//! free blocks leaves the running rows no room to grow, so the next decode +//! ticks preempt them and they re-prefill. The watermark keeps a fraction of +//! the budget free at admission while a decode batch is live, trading a later +//! admission for fewer preemptions. + +use super::block_reclaim::{PagedBlockReclaimer, reclaim_blocks}; +use super::*; + +/// Blocks the admission watermark keeps free: `fraction` of the `total` +/// budget, rounded up, but only while a decode batch is live. With no +/// running row there is no decode growth to protect, and holding headroom +/// back would wedge a request that fits the budget but not the budget minus +/// the watermark. +pub(super) fn admission_watermark_blocks(total: usize, fraction: f64, decode_live: bool) -> usize { + if !decode_live || fraction <= 0.0 { + return 0; + } + // `fraction` is clamped to `0.0..=MAX_KV_ADMISSION_WATERMARK`, so the + // product fits `total`. + ((total as f64) * fraction).ceil() as usize +} + +/// Whether a prefill needing `need` blocks may be admitted now, keeping +/// `watermark` blocks free for decode growth (issue #2088). Admits when the +/// free blocks minus the watermark cover the request; otherwise reclaims +/// toward `need + watermark` (cold prefixes first, then rows of strictly +/// lower priority than `priority`) and admits only if that much is free. +/// `false` means the caller defers the request until running rows finish. +pub(super) fn admit_with_watermark( + r: &mut R, + need: usize, + watermark: usize, + priority: RequestPriority, +) -> bool { + let target = need.saturating_add(watermark); + if r.free_blocks().is_none_or(|free| free >= target) { + return true; + } + reclaim_blocks(r, target, 0, Some(priority), false).fits +} + +impl BatchScheduler { + /// Blocks the admission watermark currently keeps free (issue #2088); 0 + /// without a budget or while no row is decoding. + pub(super) fn paged_admission_headroom(&self) -> usize { + self.cache_pool.paged_block_budget().map_or(0, |total| { + admission_watermark_blocks( + total, + self.paged_admission_watermark, + !self.active_batch.is_empty(), + ) + }) + } + + /// Admit a prefill needing `need` blocks at `priority` under the + /// admission watermark, reclaiming if needed. Only strictly lower-priority + /// rows may be preempted; an equal-priority request waits for running rows + /// to finish instead. See [`admit_with_watermark`]. + pub(super) fn reclaim_paged_blocks(&mut self, need: usize, priority: RequestPriority) -> bool { + let watermark = self.paged_admission_headroom(); + admit_with_watermark(self, need, watermark, priority) + } +} diff --git a/src/server/batch/scheduler/block_reclaim.rs b/src/server/batch/scheduler/block_reclaim.rs index 75ab39251..8225d4310 100644 --- a/src/server/batch/scheduler/block_reclaim.rs +++ b/src/server/batch/scheduler/block_reclaim.rs @@ -30,11 +30,28 @@ use super::*; /// The scheduler operations the block reclaim loop drives. pub(super) trait PagedBlockReclaimer { - /// Blocks still acquirable before the budget is hit, `None` when unbounded. - fn free_blocks(&self) -> Option; + /// Blocks the pool can still hand out before the budget is hit, before + /// anything is set aside for a parked chunked prefill. `None` when + /// unbounded. + fn pool_free_blocks(&self) -> Option; + /// Blocks still acquirable by anything other than the parked chunked + /// prefill: [`Self::pool_free_blocks`] minus [`Self::reserved_blocks`], + /// `None` when unbounded. Saturates at 0, so it cannot show how far the + /// pool has fallen below the reservation; [`reserve_chunk_blocks`] reads + /// the pool figure for that reason. + fn free_blocks(&self) -> Option { + self.pool_free_blocks() + .map(|free| free.saturating_sub(self.reserved_blocks())) + } /// Evict the least recently used prompt-cache entry and return its pinned /// blocks to the pool. `false` when there was nothing to evict. fn evict_cold_prefix(&mut self) -> bool; + /// Drop the cached prefix one queued (not yet admitted) request adopted at + /// enqueue, so it re-prefills cold when admitted. Adoption pins the + /// prefix's blocks, and once the prompt-cache entry is evicted those pins + /// are the only owner, so neither eviction nor preemption can reach them. + /// `false` when no queued request holds an adoption it can give up. + fn drop_queued_prefix(&mut self) -> bool; /// Preempt one running sequence (it re-prefills on resume), restricted to /// rows of priority strictly below `below` when set. `false` when no /// eligible victim exists. @@ -59,13 +76,18 @@ pub(super) struct ReclaimOutcome { pub fits: bool, /// Prompt-cache entries evicted. pub evicted: usize, + /// Queued requests whose adopted prefix was dropped. + pub dropped_adoptions: usize, /// Running sequences preempted. pub preempted: usize, } /// Reclaim pool blocks until at least `need` are acquirable, or no further /// reclamation is possible. Evicts cold prompt-cache prefixes first, then +/// (when `drop_queued` is set) drops prefixes queued requests adopted, then /// preempts running sequences while more than `preempt_floor` remain active. +/// The order is by what each step throws away: an unused cache entry, a +/// queued request's cache hit, a running request's work. /// The scheduler picks preemption victims with /// [`select_block_reclaim_victim_from`] (least progress first), so the row /// closest to finishing is never the one discarded. @@ -75,12 +97,15 @@ pub(super) struct ReclaimOutcome { /// strictly lower-priority rows. Decode passes a floor of 1 and no priority /// limit: the rows asking for blocks are themselves in the batch, and /// preempting the last one to make room for itself would only re-prefill it -/// back to the same point. +/// back to the same point. Decode and the chunk reservation set `drop_queued` +/// (the alternative is shedding or preempting); admission does not, since a +/// queued request can simply wait rather than cost another one its hit. pub(super) fn reclaim_blocks( r: &mut R, need: usize, preempt_floor: usize, preempt_below: Option, + drop_queued: bool, ) -> ReclaimOutcome { let room = |r: &R| r.free_blocks().is_none_or(|free| free >= need); let mut outcome = ReclaimOutcome::default(); @@ -88,7 +113,15 @@ pub(super) fn reclaim_blocks( while !room(r) && r.evict_cold_prefix() { outcome.evicted += 1; } - // 2. Preempt running sequences (drop their KV; they re-prefill on resume). + // 2. Drop queued requests' adopted prefixes (issue #2088). + while drop_queued && !room(r) && r.drop_queued_prefix() { + outcome.dropped_adoptions += 1; + // The entry a dropped adoption came from may now be evictable. + while !room(r) && r.evict_cold_prefix() { + outcome.evicted += 1; + } + } + // 3. Preempt running sequences (drop their KV; they re-prefill on resume). while !room(r) && r.active_len() > preempt_floor && r.preempt_one(preempt_below) { outcome.preempted += 1; } @@ -145,7 +178,7 @@ pub(super) fn reserve_decode_blocks( ..DecodeReservation::default() }; } - let outcome = reclaim_blocks(r, need, 1, None); + let outcome = reclaim_blocks(r, need, 1, None, true); let mut run: Vec = ids.iter().copied().filter(|&id| r.is_active(id)).collect(); let mut shed = Vec::new(); let mut deferred = Vec::new(); @@ -173,6 +206,66 @@ pub(super) fn reserve_decode_blocks( } } +/// Views a reclaimer with nothing set aside, so [`reclaim_blocks`] measures +/// room against the pool figure. Used when the blocks being reserved are the +/// parked chunked prefill's own. +struct Unreserved<'a, R: ?Sized>(&'a mut R); + +impl PagedBlockReclaimer for Unreserved<'_, R> { + fn pool_free_blocks(&self) -> Option { + self.0.pool_free_blocks() + } + fn evict_cold_prefix(&mut self) -> bool { + self.0.evict_cold_prefix() + } + fn drop_queued_prefix(&mut self) -> bool { + self.0.drop_queued_prefix() + } + fn preempt_one(&mut self, below: Option) -> bool { + self.0.preempt_one(below) + } + fn active_len(&self) -> usize { + self.0.active_len() + } + fn blocks_to_append(&self, id: SequenceId, tokens: usize) -> usize { + self.0.blocks_to_append(id, tokens) + } + fn is_active(&self, id: SequenceId) -> bool { + self.0.is_active(id) + } + fn reserved_blocks(&self) -> usize { + 0 + } +} + +/// Reserve the `need` blocks the parked chunked prefill's next chunk will +/// write, right before that chunk's forward (issue #2088). +/// +/// The chunk is otherwise protected only by the set-aside in +/// [`PagedBlockReclaimer::free_blocks`]. That figure saturates at 0, so once +/// the pool has fallen below the reservation (a decode row that needed no new +/// block still ran, a write the estimate undercounted, a padded chunk) every +/// other consumer sees "no room" and defers, but nothing restores the +/// reservation, and the chunk then hits an exhausted pool inside the forward. +/// This runs the shared reclaim loop against the pool figure: cold prefixes +/// first, then running rows with a preemption floor of 0 and no priority +/// limit. The chunk's admission already checked the whole prompt, so its +/// blocks are owed to it, and a lower floor is what keeps a single deferred +/// decode row and the chunk from waiting on each other forever. +pub(super) fn reserve_chunk_blocks( + r: &mut R, + need: usize, +) -> ReclaimOutcome { + let mut view = Unreserved(r); + if view.free_blocks().is_none_or(|free| free >= need) { + return ReclaimOutcome { + fits: true, + ..ReclaimOutcome::default() + }; + } + reclaim_blocks(&mut view, need, 0, None, true) +} + /// Charge `need` blocks against a batched-prefill window's remaining paged /// budget. `room == None` means no budget is configured and always admits. /// Returns `false`, leaving `room` untouched, when the row does not fit. @@ -188,8 +281,8 @@ pub(super) fn take_paged_room(room: &mut Option, need: usize) -> bool { } impl PagedBlockReclaimer for BatchScheduler { - fn free_blocks(&self) -> Option { - self.available_paged_blocks() + fn pool_free_blocks(&self) -> Option { + self.cache_pool.free_paged_block_budget() } fn evict_cold_prefix(&mut self) -> bool { @@ -203,6 +296,10 @@ impl PagedBlockReclaimer for BatchScheduler { true } + fn drop_queued_prefix(&mut self) -> bool { + self.drop_queued_adoption() + } + fn preempt_one(&mut self, below: Option) -> bool { match select_block_reclaim_victim_from(self.active_batch.iter_sequences(), below) { Some(victim_id) => self.preempt_sequence(victim_id), @@ -232,32 +329,82 @@ impl BatchScheduler { /// checked the whole prompt against the budget, but the blocks are only /// acquired chunk by chunk, so decode growth and later admissions between /// chunks must leave them free or the next chunk hits an exhausted pool. + /// Counts the next chunk's tile padding too (issue #2088): a padded chunk + /// writes its pad positions before trimming them, so a terminal chunk can + /// mint a block the unpadded remainder would not. pub(super) fn chunked_prefill_reserved_blocks(&self) -> usize { self.chunked_prefill_seq.as_ref().map_or(0, |seq| { - let remaining = seq.prompt_tokens.len().saturating_sub(seq.prefill_offset); + let total = seq.prompt_tokens.len(); + let remaining = total.saturating_sub(seq.prefill_offset); + let pad = + next_chunked_prefill_range(total, seq.prefill_offset, self.prefill_chunk_size) + .map_or(0, |range| { + let len = range.end - range.start; + self.prefill_write_len(len) - len + }); self.cache_pool - .paged_blocks_to_append(seq.seq_id, remaining) + .paged_blocks_to_append(seq.seq_id, remaining + pad) }) } + /// Positions a prefill chunk of `len` prompt tokens writes into the KV + /// cache: `len` rounded up to the neural-accelerator tile when the chunk + /// is padded (M5+ Metal), `len` otherwise. + pub(super) fn prefill_write_len(&self, len: usize) -> usize { + if self.model.supports_padded_prefill() && should_align_prefill() { + align_to_na_tile(len) + } else { + len + } + } + /// Paged blocks acquirable without reclaim once the parked chunked /// prefill's remaining blocks are set aside, or `None` when no budget is /// configured or the pool does not exist yet. Every block-budget decision /// (prefill admission, decode reservation, lookahead prime) reads this /// rather than the raw pool figure (#1982). pub(super) fn available_paged_blocks(&self) -> Option { - self.cache_pool - .free_paged_block_budget() - .map(|free| free.saturating_sub(self.chunked_prefill_reserved_blocks())) + PagedBlockReclaimer::free_blocks(self) } - /// Reclaim paged pool blocks until at least `need` are acquirable, or no - /// further reclamation is possible, for admitting a prefill of - /// `priority`. Only strictly lower-priority rows may be preempted; an - /// equal-priority request waits for running rows to finish instead. See - /// [`reclaim_blocks`]. Returns whether `need` blocks are now acquirable. - pub(super) fn reclaim_paged_blocks(&mut self, need: usize, priority: RequestPriority) -> bool { - reclaim_blocks(self, need, 0, Some(priority)).fits + /// Reserve the blocks the parked chunked prefill `seq_id` writes for a + /// chunk of `write_len` positions, right before its forward (issue + /// #2088). Returns whether the chunk may run. + /// + /// Called from `continue_chunked_prefill`, so it covers both ways a chunk + /// runs beside a live decode batch: the `MLXCEL_MIXED_STEP` tick, where + /// the decode step has just run and may have left the pool below the + /// chunk's reservation, and the #1011 prefill grant. A no-op returning + /// `true` when no budget is configured or the pool has room. Under + /// pressure it tears down any prebuilt lookahead (preemption changes batch + /// membership) and runs [`reserve_chunk_blocks`]. `false` means even + /// preempting every running row and evicting every cold prefix left too + /// few blocks; the caller fails the request instead of letting the forward + /// hit an exhausted pool. + pub(super) fn reserve_prefill_chunk_blocks( + &mut self, + seq_id: SequenceId, + write_len: usize, + ) -> bool { + let need = self.cache_pool.paged_blocks_to_append(seq_id, write_len); + let free_before = self.cache_pool.free_paged_block_budget(); + if free_before.is_none_or(|free| free >= need) { + return true; + } + self.discard_lookahead(); + let outcome = reserve_chunk_blocks(self, need); + tracing::info!( + %seq_id, + need, + free_before = free_before.unwrap_or_default(), + free_after = self.cache_pool.free_paged_block_budget().unwrap_or_default(), + evicted_prefixes = outcome.evicted, + dropped_adoptions = outcome.dropped_adoptions, + preempted = outcome.preempted, + fits = outcome.fits, + "Chunked prefill reclaimed paged KV blocks" + ); + outcome.fits } /// Reserve pool blocks for the decode tick over `seq_ids` (#1982). Returns @@ -290,6 +437,7 @@ impl BatchScheduler { free_before = free_before.unwrap_or_default(), free_after = self.available_paged_blocks().unwrap_or_default(), evicted_prefixes = outcome.evicted, + dropped_adoptions = outcome.dropped_adoptions, preempted = outcome.preempted, shed = reservation.shed.len(), deferred = reservation.deferred.len(), diff --git a/src/server/batch/scheduler/block_reclaim_tests.rs b/src/server/batch/scheduler/block_reclaim_tests.rs index 454f4b42f..5f0a2d212 100644 --- a/src/server/batch/scheduler/block_reclaim_tests.rs +++ b/src/server/batch/scheduler/block_reclaim_tests.rs @@ -28,9 +28,12 @@ use mlxcel_core::cache::{CachePool, KVCache, PagedKvLayout, SequenceId, Sequence use mlxcel_core::generate::LanguageModel; use mlxcel_core::{MlxArray, UniquePtr}; +use super::super::admission_watermark::{admission_watermark_blocks, admit_with_watermark}; use super::{ - PagedBlockReclaimer, RequestPriority, reclaim_blocks, reserve_decode_blocks, take_paged_room, + PagedBlockReclaimer, RequestPriority, reclaim_blocks, reserve_chunk_blocks, + reserve_decode_blocks, take_paged_room, }; +use crate::server::config::clamp_kv_admission_watermark; use crate::server::prompt_cache::{ CacheEntry, DetachedKvSet, PromptCacheConfig, PromptCacheStore, key::{MultimodalDigest, PromptCacheKey}, @@ -79,6 +82,10 @@ struct Harness { store: Arc, active: Vec, preempted: Vec, + /// Queued requests holding an adopted prefix's blocks, oldest first. + queued: Vec, + /// Queued requests whose adopted prefix was dropped. + dropped: Vec, /// Blocks set aside for a parked chunked prefill, like /// `BatchScheduler::chunked_prefill_reserved_blocks`. reserved: usize, @@ -102,6 +109,8 @@ impl Harness { ))), active: Vec::new(), preempted: Vec::new(), + queued: Vec::new(), + dropped: Vec::new(), reserved: 0, } } @@ -152,10 +161,8 @@ impl Harness { } impl PagedBlockReclaimer for Harness { - fn free_blocks(&self) -> Option { - self.pool - .free_paged_block_budget() - .map(|free| free.saturating_sub(self.reserved)) + fn pool_free_blocks(&self) -> Option { + self.pool.free_paged_block_budget() } fn evict_cold_prefix(&mut self) -> bool { @@ -168,6 +175,16 @@ impl PagedBlockReclaimer for Harness { true } + fn drop_queued_prefix(&mut self) -> bool { + // Newest queued request first, like `PrefillQueue::find_lowest_first`. + let Some(id) = self.queued.pop() else { + return false; + }; + self.pool.release(id); + self.dropped.push(id); + true + } + fn preempt_one(&mut self, below: Option) -> bool { // Every harness row is Normal priority, so a limit at or below Normal // leaves no eligible victim. @@ -328,7 +345,7 @@ fn admission_floor_zero_may_preempt_every_lower_priority_row() { let mut h = Harness::new(4); h.decoding(BLOCK); h.decoding(BLOCK); - let outcome = reclaim_blocks(&mut h, 4, 0, Some(RequestPriority::High)); + let outcome = reclaim_blocks(&mut h, 4, 0, Some(RequestPriority::High), false); assert!(outcome.fits); assert_eq!(outcome.preempted, 2); assert!(h.active.is_empty()); @@ -339,12 +356,15 @@ fn admission_floor_zero_may_preempt_every_lower_priority_row() { struct NoPreempt<'a>(&'a mut Harness); impl PagedBlockReclaimer for NoPreempt<'_> { - fn free_blocks(&self) -> Option { - self.0.free_blocks() + fn pool_free_blocks(&self) -> Option { + self.0.pool_free_blocks() } fn evict_cold_prefix(&mut self) -> bool { self.0.evict_cold_prefix() } + fn drop_queued_prefix(&mut self) -> bool { + self.0.drop_queued_prefix() + } fn preempt_one(&mut self, _below: Option) -> bool { false } @@ -402,7 +422,7 @@ fn admission_never_preempts_rows_of_equal_priority() { let mut h = Harness::new(4); h.decoding(BLOCK); h.decoding(BLOCK); - let outcome = reclaim_blocks(&mut h, 4, 0, Some(RequestPriority::Normal)); + let outcome = reclaim_blocks(&mut h, 4, 0, Some(RequestPriority::Normal), false); assert!(!outcome.fits); assert_eq!(outcome.preempted, 0); assert_eq!(h.active.len(), 2); @@ -424,3 +444,213 @@ fn a_row_short_only_by_the_chunked_prefill_reservation_waits_instead_of_failing( assert!(reservation.shed.is_empty()); assert_eq!(reservation.deferred, vec![row]); } + +/// A parked chunked prefill `seq` that has written `written` tokens on every +/// layer and still has `remaining` to go, with its reservation set aside the +/// way `BatchScheduler::chunked_prefill_reserved_blocks` does. +fn park_chunked_prefill(h: &mut Harness, written: usize, remaining: usize) -> SequenceId { + let id = h.pool.allocate(&h.model).unwrap(); + h.append_all_layers(id, written).unwrap(); + h.reserved = h.pool.paged_blocks_to_append(id, remaining); + id +} + +#[test] +fn mixed_step_decode_then_chunk_share_the_last_blocks() { + // MixedStep tick order (issue #2088): decode first, then the chunk. Budget + // 6: the decode row holds 2 blocks at a boundary, the parked prefill holds + // 2 and needs 2 more, and nothing else is free. The decode row may not take + // the chunk's 2 blocks (it waits a tick), and the chunk reservation then + // finds them without reclaiming anything. + let mut h = Harness::new(6); + let row = h.decoding(BLOCK); + let chunked = park_chunked_prefill(&mut h, BLOCK, BLOCK); + assert_eq!(h.reserved, 2); + + let decode = reserve_decode_blocks(&mut h, &[row], 1); + assert_eq!( + decode.deferred, + vec![row], + "decode leaves the chunk's blocks" + ); + assert!(decode.shed.is_empty()); + + let need = h.pool.paged_blocks_to_append(chunked, BLOCK); + let chunk = reserve_chunk_blocks(&mut h, need); + assert!(chunk.fits); + assert_eq!((chunk.evicted, chunk.preempted), (0, 0)); + h.append_all_layers(chunked, BLOCK) + .expect("the chunk writes all of its blocks"); +} + +#[test] +fn mixed_step_chunk_reclaims_when_decode_left_the_pool_below_its_reservation() { + // Budget 8: the parked prefill holds 2 blocks and needs 2 more, but two + // mid-block decode rows hold the other 6 (a row that needs no new block + // always runs, so nothing stopped them). The set-aside figure saturates at + // 0 and hides the shortfall; without the chunk reservation the chunk's + // write would hit "block budget exhausted" inside the forward. + let mut h = Harness::new(8); + let chunked = park_chunked_prefill(&mut h, BLOCK, BLOCK); + let first = h.decoding(BLOCK + 1); + let second = h.decoding(BLOCK - 1); + assert_eq!(h.pool.free_paged_block_budget(), Some(0)); + assert_eq!(h.free_blocks(), Some(0)); + + let decode = reserve_decode_blocks(&mut h, &[first, second], 1); + assert_eq!( + decode.run, + vec![first, second], + "mid-block rows still decode" + ); + + let need = h.pool.paged_blocks_to_append(chunked, BLOCK); + let err = h.pool.append_paged_tokens(chunked, 0, BLOCK).unwrap_err(); + assert!(err.contains("block budget exhausted"), "{err}"); + + let chunk = reserve_chunk_blocks(&mut h, need); + assert!( + chunk.fits, + "preempting a decode row restores the reservation" + ); + assert_eq!(chunk.preempted, 1); + assert_eq!(h.preempted, vec![second]); + h.append_all_layers(chunked, BLOCK) + .expect("the chunk runs instead of being dropped"); +} + +#[test] +fn chunk_reservation_evicts_cold_prefixes_before_preempting() { + // Budget 8: decode row (2) + parked prefill (2) + a cached prefix (4) fill + // the pool while the prefill still needs 2 blocks. Evicting the prefix is + // enough; the decode row keeps running. + let mut h = Harness::new(8); + let row = h.decoding(BLOCK); + let chunked = park_chunked_prefill(&mut h, BLOCK, BLOCK); + h.cache_prefix(2 * BLOCK, 1); + assert_eq!(h.pool.free_paged_block_budget(), Some(0)); + + let need = h.pool.paged_blocks_to_append(chunked, BLOCK); + let chunk = reserve_chunk_blocks(&mut h, need); + assert!(chunk.fits); + assert_eq!((chunk.evicted, chunk.preempted), (1, 0)); + assert!(h.is_active(row)); + h.append_all_layers(chunked, BLOCK).unwrap(); +} + +#[test] +fn admission_defers_when_free_blocks_minus_the_watermark_cannot_cover_it() { + // Budget 12, one decoding row holding 2 blocks: 10 free. A 4-block + // request fits the free blocks, but not with a 7-block watermark kept for + // decode growth, and nothing is reclaimable, so admission defers. + let mut h = Harness::new(12); + h.decoding(BLOCK); + assert_eq!(h.free_blocks(), Some(10)); + assert!(admit_with_watermark(&mut h, 4, 0, RequestPriority::Normal)); + assert!(!admit_with_watermark(&mut h, 4, 7, RequestPriority::Normal)); + assert!( + h.preempted.is_empty(), + "equal-priority rows are never preempted" + ); + assert!(admit_with_watermark(&mut h, 4, 6, RequestPriority::Normal)); +} + +#[test] +fn admission_watermark_evicts_cold_prefixes_to_restore_headroom() { + // Budget 12: a decoding row (2) and a cached prefix (4) leave 6 free. A + // 4-block request with a 4-block watermark needs 8 free, so the cold + // prefix is evicted and the request is admitted. + let mut h = Harness::new(12); + h.decoding(BLOCK); + h.cache_prefix(2 * BLOCK, 1); + assert!(admit_with_watermark(&mut h, 4, 4, RequestPriority::Normal)); + assert_eq!(h.store.len(), 0); + assert!(h.preempted.is_empty()); +} + +#[test] +fn admission_watermark_applies_only_while_rows_decode() { + assert_eq!(admission_watermark_blocks(4444, 0.1, true), 445); + assert_eq!( + admission_watermark_blocks(4444, 0.1, false), + 0, + "an empty batch has no decode growth to protect" + ); + assert_eq!(admission_watermark_blocks(4444, 0.0, true), 0); + assert_eq!(clamp_kv_admission_watermark(-1.0), 0.0); + assert_eq!(clamp_kv_admission_watermark(0.9), 0.5); + assert_eq!(clamp_kv_admission_watermark(f64::NAN), 0.0); +} + +/// A queued request that adopted a `tokens`-long prefix whose prompt-cache +/// entry has since been evicted: its pins are the blocks' only owner. +fn queued_adoption(h: &mut Harness, tokens: usize) -> SequenceId { + let id = h.pool.allocate(&h.model).unwrap(); + h.append_all_layers(id, tokens).unwrap(); + h.queued.push(id); + id +} + +#[test] +fn decode_drops_a_queued_adoption_before_shedding_the_last_row() { + // Budget 6: the only running row sits at a block boundary (2 blocks) and a + // queued request's adopted prefix holds the other 4. Nothing is cached and + // the floor keeps the row from preempting itself, so before #2088 the row + // was shed (GB10 run at a 0.15 watermark). Dropping the queued adoption + // frees the blocks; the queued request re-prefills cold when admitted. + let mut h = Harness::new(6); + let row = h.decoding(BLOCK); + let queued = queued_adoption(&mut h, 2 * BLOCK); + assert_eq!(h.free_blocks(), Some(0)); + + let reservation = reserve_decode_blocks(&mut h, &[row], 1); + assert_eq!(reservation.run, vec![row]); + assert!(reservation.shed.is_empty()); + let outcome = reservation.reclaim.expect("reclaim ran"); + assert_eq!( + ( + outcome.evicted, + outcome.dropped_adoptions, + outcome.preempted + ), + (0, 1, 0) + ); + assert_eq!(h.dropped, vec![queued]); + h.append_all_layers(row, 1).unwrap(); +} + +#[test] +fn queued_adoptions_go_after_cold_prefixes_and_before_running_rows() { + // Budget 10: two rows at a boundary (4 blocks), a cached prefix (2) and a + // queued adoption (4). One boundary crossing needs 4 blocks: evicting the + // prefix gives 2, dropping the adoption gives the rest, and no row is + // preempted. + let mut h = Harness::new(10); + let first = h.decoding(BLOCK); + let second = h.decoding(BLOCK); + h.cache_prefix(BLOCK, 1); + queued_adoption(&mut h, 2 * BLOCK); + + let reservation = reserve_decode_blocks(&mut h, &[first, second], 1); + assert_eq!(reservation.run, vec![first, second]); + let outcome = reservation.reclaim.expect("reclaim ran"); + assert_eq!( + ( + outcome.evicted, + outcome.dropped_adoptions, + outcome.preempted + ), + (1, 1, 0) + ); + assert!(h.preempted.is_empty()); +} + +#[test] +fn admission_never_drops_another_requests_adoption() { + // A queued request waits rather than cost another one its cache hit. + let mut h = Harness::new(6); + h.decoding(BLOCK); + queued_adoption(&mut h, 2 * BLOCK); + assert!(!admit_with_watermark(&mut h, 2, 0, RequestPriority::Normal)); + assert!(h.dropped.is_empty()); +} diff --git a/src/server/batch/scheduler/config.rs b/src/server/batch/scheduler/config.rs index cb9c2878e..73c9ad929 100644 --- a/src/server/batch/scheduler/config.rs +++ b/src/server/batch/scheduler/config.rs @@ -124,6 +124,7 @@ impl BatchScheduler { // `with_prefill_grant_interval` overrides this later with an // explicit CLI value when one was passed. prefill_grant_interval: resolve_prefill_grant_interval(None), + paged_admission_watermark: DEFAULT_KV_ADMISSION_WATERMARK, decode_ticks_since_prefill_grant: 0, shutdown_requested: false, consecutive_decode_eval_failures: 0, @@ -220,6 +221,15 @@ impl BatchScheduler { self } + /// Install the paged KV admission watermark (issue #2088, + /// `--kv-admission-watermark`): the fraction of the block budget a prefill + /// admission keeps free for decode growth while rows are decoding. Clamped + /// to `0.0..=MAX_KV_ADMISSION_WATERMARK`; inert without a block budget. + pub fn with_paged_admission_watermark(mut self, fraction: f64) -> Self { + self.paged_admission_watermark = clamp_kv_admission_watermark(fraction); + self + } + /// Returns the resolved prefill fairness interval (0 = grant disabled). /// Exposed for tests. pub fn prefill_grant_interval(&self) -> u32 { diff --git a/src/server/batch/scheduler/mod.rs b/src/server/batch/scheduler/mod.rs index 821166235..634a14c5a 100644 --- a/src/server/batch/scheduler/mod.rs +++ b/src/server/batch/scheduler/mod.rs @@ -62,7 +62,8 @@ use crate::models::gemma4_mtp_target::{ use crate::server::ServerGenerateOptions; use crate::server::batch::observability::{BatchObservability, apc_trace_enabled}; use crate::server::config::{ - DecodeStorageBackend, PreemptionPolicy, PromptCacheRequestContext, ReasoningBudgetOverride, + DEFAULT_KV_ADMISSION_WATERMARK, DecodeStorageBackend, PreemptionPolicy, + PromptCacheRequestContext, ReasoningBudgetOverride, clamp_kv_admission_watermark, }; use crate::server::model_provider::model_worker::{ StreamingDecodeState, build_generation_result_with_cache, merge_config_stop_tokens, @@ -396,6 +397,10 @@ pub struct BatchScheduler { /// step with the arbitration it feeds. Nothing else in the scheduler needs /// to remember to reset it. decode_ticks_since_prefill_grant: u32, + /// Fraction of the paged KV block budget a prefill admission leaves free + /// for the running rows' decode growth while a decode batch is live + /// (issue #2088, `--kv-admission-watermark`). 0 disables it. + paged_admission_watermark: f64, // -- Shutdown flag -- shutdown_requested: bool, @@ -966,6 +971,7 @@ fn validate_dense_detached_kv_modes_against_table( } mod admission; +mod admission_watermark; mod block_reclaim; mod config; mod decode_tick; @@ -975,6 +981,7 @@ mod mtp_dispatch; mod paged_layout; mod prefill; mod prompt_cache; +mod queued_adoption; mod run_loop; mod shared_budget; mod speculative_finalize; diff --git a/src/server/batch/scheduler/prefill.rs b/src/server/batch/scheduler/prefill.rs index 18095d04a..dada4230c 100644 --- a/src/server/batch/scheduler/prefill.rs +++ b/src/server/batch/scheduler/prefill.rs @@ -1209,6 +1209,20 @@ impl BatchScheduler { (chunk.to_vec(), None) }; + // Reserve this chunk's blocks before its forward (issue #2088). A + // MixedStep tick has just run a decode step, and nothing else restores + // the set-aside once the pool falls below it; a failed acquire inside + // the forward is fatal for the worker. + if !self.reserve_prefill_chunk_blocks(seq.seq_id, eff_chunk.len()) { + let total = self.cache_pool.paged_block_budget().unwrap_or_default(); + self.abort_sequence( + seq, + &format!( + "KV cache budget exhausted: no free blocks in the {total}-block KV cache budget to continue the chunked prefill" + ), + ); + return false; + } let eff_len = eff_chunk.len() as i32; let input = mlxcel_core::from_slice_i32(&eff_chunk, &[1, eff_len]); let logits = { diff --git a/src/server/batch/scheduler/queued_adoption.rs b/src/server/batch/scheduler/queued_adoption.rs new file mode 100644 index 000000000..85c1f07ea --- /dev/null +++ b/src/server/batch/scheduler/queued_adoption.rs @@ -0,0 +1,77 @@ +// Copyright 2025-2026 Lablup Inc. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +//! Giving back a queued request's adopted prompt-cache prefix (issue #2088). +//! +//! A request adopts a cached prefix when it is enqueued, which pins the +//! prefix's paged blocks for as long as it waits for admission. Under a tight +//! `--kv-cache-budget` (and more often with an admission watermark, which +//! lengthens the wait) those pins can be the only reclaimable blocks left +//! when a running row needs one. The block reclaim loop drops them before it +//! preempts running work. + +use super::*; + +impl BatchScheduler { + /// Drop the cached prefix the lowest-priority, newest queued text request + /// adopted at enqueue (issue #2088): release its cache slot, allocate a + /// fresh one, and reset its offsets so it prefills cold when admitted, the + /// way [`Self::preempt_sequence`] resets a running row. Its prompt-cache + /// context moves to the new id, so it still donates its prefix back when + /// it finishes. Multimodal requests are skipped: an adopted VLM request + /// dropped its full-prompt embeddings and could not re-prefill cold. + /// Returns whether an adoption was dropped. + pub(super) fn drop_queued_adoption(&mut self) -> bool { + let Some(old_id) = self.prefill_queue.find_lowest_first(|seq| { + seq.already_cached_tokens > 0 + && !seq.is_vlm_request() + && seq.audio.is_empty() + && self + .cache_pool + .paged_blocks_to_reach(seq.seq_id, 0) + .is_some() + }) else { + return false; + }; + self.release_sequence_caches(old_id); + let new_id = match self.allocate_sequence_state() { + Ok(id) => id, + Err(err) => { + self.prompt_cache_seq_ctx.remove(&old_id); + if let Some(seq) = self.prefill_queue.remove(old_id) { + let _ = seq.response_tx.send(GenerateEvent::Error(format!( + "Server busy: re-allocation after dropping the cached prefix failed: {err}" + ))); + } + return true; + } + }; + if let Some(ctx) = self.prompt_cache_seq_ctx.remove(&old_id) { + self.prompt_cache_seq_ctx.insert(new_id, ctx); + } + if let Some(seq) = self.prefill_queue.get_mut(old_id) { + tracing::info!( + old = %old_id, + new = %new_id, + cached_tokens = seq.already_cached_tokens, + "Dropped a queued request's adopted prefix to free paged KV blocks" + ); + seq.seq_id = new_id; + seq.prefill_offset = 0; + seq.prefill_start_offset = 0; + seq.already_cached_tokens = 0; + } + true + } +} diff --git a/src/server/batch/scheduler_prompt_cache_tests.rs b/src/server/batch/scheduler_prompt_cache_tests.rs index dbddeef4a..1ba8b5d7e 100644 --- a/src/server/batch/scheduler_prompt_cache_tests.rs +++ b/src/server/batch/scheduler_prompt_cache_tests.rs @@ -1471,3 +1471,36 @@ fn evicting_cold_prefix_restores_paged_block_budget() { "evicting the cold prefix reclaimed its 2 blocks into the budget" ); } + +#[test] +fn adopted_prefix_is_charged_only_for_its_suffix() { + // Issue #2088: admission used to charge an adopted request for its whole + // prompt while the adopted blocks already counted as live, so a prefix + // larger than half the budget could never be admitted, even into an empty + // batch (the GB10 run at a 0.15 watermark wedged on one). Budget 8 blocks + // (1 layer, block 4): a 20-token adopted prefix holds 5 blocks and leaves 3 + // free. The 24-token prompt mints 1 block, not the 6 the whole prompt + // would, so it fits. + let model = PagedStub { + layout: PagedKvLayout::uniform(1, 4, 128).unwrap(), + }; + let mut pool = CachePool::new(4); + pool.set_paged_block_budget(Some(8)); + let (set, _blocks) = mint_paged_set(&mut pool, &model, 20); + let DetachedKvSet::Paged(paged) = set else { + panic!("paged set"); + }; + let adopted = pool.adopt_paged(&model, paged).expect("adopt"); + let free = pool.free_paged_block_budget().expect("budgeted"); + assert_eq!(free, 3); + + let whole_prompt = 24usize.div_ceil(4); + assert!(whole_prompt > free, "the old charge could never fit"); + assert_eq!(pool.paged_blocks_to_reach(adopted, 24), Some(1)); + pool.append_paged_tokens(adopted, 0, 4) + .expect("the suffix fits the free blocks"); + + // No paged state: callers fall back to the whole-prompt estimate. + let dense = SequenceId::from_raw(u64::MAX); + assert_eq!(pool.paged_blocks_to_reach(dense, 24), None); +} diff --git a/src/server/cli_input.rs b/src/server/cli_input.rs index 7d8c3b8c7..a30b11b0f 100644 --- a/src/server/cli_input.rs +++ b/src/server/cli_input.rs @@ -541,6 +541,10 @@ pub struct ServerStartupInput { /// resolved to a concrete block count on the worker thread. pub kv_cache_budget: Option, + /// `--kv-admission-watermark` (issue #2088): fraction of the paged block + /// budget kept free at admission while rows decode. `None` = the default. + pub kv_admission_watermark: Option, + /// `--enable-vlm-prefix-cache` (#124 step c). Default off. Enables /// experimental VLM prompt-prefix cache sharing for multi-turn same-image /// conversations; forwarded verbatim to the scheduler. @@ -1179,6 +1183,7 @@ impl ServerStartupInput { // forward the paged KV block-budget directive verbatim (resolved // to a block count on the worker thread). kv_cache_budget: self.kv_cache_budget, + kv_admission_watermark: self.kv_admission_watermark, // forward the experimental VLM prefix-cache toggle (#124 step c). enable_vlm_prefix_cache: self.enable_vlm_prefix_cache, // forward the resolved CORS policy: the b10621 flags, or the diff --git a/src/server/cli_input_tests.rs b/src/server/cli_input_tests.rs index 52a22c4a2..b17eb7c38 100644 --- a/src/server/cli_input_tests.rs +++ b/src/server/cli_input_tests.rs @@ -215,6 +215,7 @@ fn sample_input() -> ServerStartupInput { max_kv_size: 0, // paged KV pool block-budget directive (None = unbounded, the default). kv_cache_budget: None, + kv_admission_watermark: None, // experimental VLM prefix-cache toggle off in tests (#124 step c). enable_vlm_prefix_cache: false, // CORS allow-list unset in tests (#244): permissive default. diff --git a/src/server/config.rs b/src/server/config.rs index 7883034ec..b9d7049be 100644 --- a/src/server/config.rs +++ b/src/server/config.rs @@ -29,6 +29,44 @@ use crate::server::prompt_cache::key::MultimodalDigest; use mlxcel_core::lang_analyzer::LangBiasConfig; use mlxcel_core::sampling::LogprobsConfig; +/// Default `--kv-admission-watermark` (issue #2088): the fraction of the +/// paged KV block budget a prefill admission keeps free for decode growth +/// while rows are decoding. The smallest value on the measured grid (0.01 to +/// 0.15) that lowered preemptions without losing a turn: on GB10 with +/// meta-llama-3.1-8b-instruct-4bit at a 4444-block budget, three concurrent +/// 3-turn conversations went from 10 preemptions to 6 in both repeats. See +/// `docs/CONTINUOUS_BATCHING.md`. +pub const DEFAULT_KV_ADMISSION_WATERMARK: f64 = 0.01; + +/// Largest accepted `--kv-admission-watermark`. Above half the budget the +/// server would mostly admit one sequence at a time. +pub const MAX_KV_ADMISSION_WATERMARK: f64 = 0.5; + +/// Clamp an admission watermark into `0.0..=MAX_KV_ADMISSION_WATERMARK`; NaN +/// disables it. +#[must_use] +pub fn clamp_kv_admission_watermark(fraction: f64) -> f64 { + if fraction.is_nan() { + return 0.0; + } + fraction.clamp(0.0, MAX_KV_ADMISSION_WATERMARK) +} + +/// clap value parser for `--kv-admission-watermark`: a fraction of the block +/// budget in `0.0..=MAX_KV_ADMISSION_WATERMARK`. +pub fn parse_kv_admission_watermark(s: &str) -> Result { + let value: f64 = s + .trim() + .parse() + .map_err(|_| format!("`{s}` is not a number; expected a fraction such as 0.1"))?; + if !(0.0..=MAX_KV_ADMISSION_WATERMARK).contains(&value) { + return Err(format!( + "`{s}` is out of range; expected a fraction in 0.0..={MAX_KV_ADMISSION_WATERMARK}" + )); + } + Ok(value) +} + /// Storage backend used by the server batch scheduler for decode-time state. #[derive(Debug, Clone, Copy, PartialEq, Eq, Default)] pub enum DecodeStorageBackend { @@ -980,6 +1018,11 @@ pub struct ServerConfig { /// [`crate::server::batch::BatchScheduler::with_paged_block_budget`]. Only /// meaningful for pool-backed (Fp16, dense-natural-backend) sequences. pub kv_cache_budget: Option, + /// Paged KV admission watermark (issue #2088, `--kv-admission-watermark`): + /// the fraction of the block budget a prefill admission keeps free for + /// the running rows' decode growth while a decode batch is live. `0.0` + /// disables it. Inert without a block budget. + pub kv_admission_watermark: f64, /// `--enable-vlm-prefix-cache` (#124 step c). Default off. When on, the /// scheduler permits VLM (image/audio) chat requests to adopt and donate /// KV prefixes for multi-turn same-image conversations; text-only and @@ -1157,6 +1200,7 @@ impl Default for ServerConfig { // default with an `auto` paged KV budget so admission sheds load // instead of OOMing. Disable with `--kv-cache-budget none`. kv_cache_budget: Some(crate::memory_estimate::PagedBudgetDirective::Auto), + kv_admission_watermark: DEFAULT_KV_ADMISSION_WATERMARK, enable_vlm_prefix_cache: false, cors_policy: crate::server::CorsPolicy::default(), serving_mode: crate::distributed::disaggregated::ServingMode::Hybrid, diff --git a/src/server/config_tests.rs b/src/server/config_tests.rs index 32f61d399..da59c0c3c 100644 --- a/src/server/config_tests.rs +++ b/src/server/config_tests.rs @@ -12,7 +12,10 @@ // See the License for the specific language governing permissions and // limitations under the License. -use super::{DecodeStorageBackend, ReasoningAliasField, ServerConfig}; +use super::{ + DEFAULT_KV_ADMISSION_WATERMARK, DecodeStorageBackend, ReasoningAliasField, ServerConfig, + parse_kv_admission_watermark, +}; use crate::memory_estimate::PagedBudgetDirective; #[test] @@ -73,4 +76,19 @@ fn server_config_default_matches_llama_server_compatibility_defaults() { // Serving-throughput default guard: `auto` paged KV budget pairs with the // batched-decode default so admission sheds load instead of OOMing (#628). assert_eq!(config.kv_cache_budget, Some(PagedBudgetDirective::Auto)); + assert_eq!( + config.kv_admission_watermark, + DEFAULT_KV_ADMISSION_WATERMARK + ); +} + +#[test] +fn kv_admission_watermark_parser_accepts_fractions_up_to_half() { + assert_eq!(parse_kv_admission_watermark("0"), Ok(0.0)); + assert_eq!(parse_kv_admission_watermark(" 0.1 "), Ok(0.1)); + assert_eq!(parse_kv_admission_watermark("0.5"), Ok(0.5)); + assert!(parse_kv_admission_watermark("0.6").is_err()); + assert!(parse_kv_admission_watermark("-0.1").is_err()); + assert!(parse_kv_admission_watermark("NaN").is_err()); + assert!(parse_kv_admission_watermark("ten").is_err()); } diff --git a/src/server/mod.rs b/src/server/mod.rs index 4cc2b4ad1..2fd957e66 100644 --- a/src/server/mod.rs +++ b/src/server/mod.rs @@ -120,6 +120,10 @@ pub use cli_input::{ resolve_batch_kv_quant_config, resolve_compat_toggle, resolve_kv_cache_mode, resolve_n_parallel, }; +pub use config::{ + DEFAULT_KV_ADMISSION_WATERMARK, MAX_KV_ADMISSION_WATERMARK, clamp_kv_admission_watermark, + parse_kv_admission_watermark, +}; pub use config::{ DecodeStorageBackend, LiveSettings, PipelineParallelRuntimeConfig, PreemptionPolicy, ReasoningAliasField, RemotePipelineStageConfig, ServerConfig, ServerGenerateOptions, diff --git a/src/server/model_provider.rs b/src/server/model_provider.rs index c55bbe027..ae2f190b1 100644 --- a/src/server/model_provider.rs +++ b/src/server/model_provider.rs @@ -775,6 +775,8 @@ impl ModelProvider { }, // forward the --kv-cache-budget directive to the worker. config.kv_cache_budget, + // forward the --kv-admission-watermark fraction (#2088). + config.kv_admission_watermark, // experimental VLM prompt-prefix cache toggle (#124 step c). config.enable_vlm_prefix_cache, // disaggregated serving role from `--node-role` (#126 B2). @@ -1161,6 +1163,7 @@ impl ModelProvider { // legacy wrapper: b10621 defaults (shift disabled, keep 0). Default::default(), kv_cache_budget, + crate::server::DEFAULT_KV_ADMISSION_WATERMARK, enable_vlm_prefix_cache, serving_mode, decode_peers, @@ -1219,6 +1222,7 @@ impl ModelProvider { // b10621 context-retention policy at the KV bound (#1472). context_retention: crate::server::batch::ContextRetentionPolicy, kv_cache_budget: Option, + kv_admission_watermark: f64, enable_vlm_prefix_cache: bool, serving_mode: crate::distributed::disaggregated::ServingMode, decode_peers: Vec, @@ -1288,6 +1292,7 @@ impl ModelProvider { // paged KV pool block-budget directive; resolved to a block count // on the worker thread once the model geometry is known. kv_cache_budget, + kv_admission_watermark, // experimental VLM prompt-prefix cache toggle (#124 step c). enable_vlm_prefix_cache, // disaggregated serving role from `--node-role` (#126 B2). The @@ -1414,6 +1419,7 @@ impl ModelProvider { context_size_total: 0, // model default in minimal test path explicit_max_kv_size: None, // unset in minimal test path kv_cache_budget: None, // unbounded in minimal test path + kv_admission_watermark: 0.0, // inert without a budget enable_vlm_prefix_cache: false, // off in minimal test path // minimal test path is single-node. serving_mode: crate::distributed::disaggregated::ServingMode::Hybrid, diff --git a/src/server/model_worker.rs b/src/server/model_worker.rs index 4047a36af..ca237eb64 100644 --- a/src/server/model_worker.rs +++ b/src/server/model_worker.rs @@ -154,6 +154,11 @@ pub(crate) struct WorkerSchedulerConfig { /// the scheduler's pool through /// [`crate::server::batch::BatchScheduler::with_paged_block_budget`]. pub kv_cache_budget: Option, + /// paged KV admission watermark (issue #2088, `--kv-admission-watermark`): + /// the fraction of the block budget a prefill admission keeps free while + /// rows decode. Installed through + /// [`crate::server::batch::BatchScheduler::with_paged_admission_watermark`]. + pub kv_admission_watermark: f64, /// experimental VLM prompt-prefix cache toggle (#124 step c, /// `--enable-vlm-prefix-cache`). /// @@ -759,6 +764,8 @@ pub(crate) fn spawn_model_worker_with_batch_config( .with_context_retention(sched_config.context_retention) // install the resolved paged KV block budget (epic #116 #122 b3). .with_paged_block_budget(paged_block_budget) + // keep decode headroom at admission under that budget (#2088). + .with_paged_admission_watermark(sched_config.kv_admission_watermark) // install the resolved paged KV slab size (#899). .with_paged_slab_blocks(paged_slab_blocks) // experimental VLM prompt-prefix cache sharing (#124 step c). diff --git a/src/server/runtime_settings.rs b/src/server/runtime_settings.rs index 87213303f..95a28d54a 100644 --- a/src/server/runtime_settings.rs +++ b/src/server/runtime_settings.rs @@ -143,6 +143,7 @@ pub const CLASSIFIED_SERVER_CONFIG_FIELDS: &[&str] = &[ "context_shift", "n_keep", "kv_cache_budget", + "kv_admission_watermark", "enable_vlm_prefix_cache", "cors_policy", "serving_mode", @@ -512,6 +513,7 @@ fn read_only_reason(field: &str) -> &'static str { | "context_shift" | "n_keep" | "kv_cache_budget" + | "kv_admission_watermark" | "enable_vlm_prefix_cache" => SCHEDULER_REASON, // b10621 `--sleep-idle-seconds` (#1440): the serving worker reads the // window when it builds its scheduler, so changing it live would only @@ -648,6 +650,7 @@ fn read_only_value(config: &ServerConfig, field: &str) -> Value { "context_shift" => json!(config.context_shift), "n_keep" => json!(config.n_keep), "kv_cache_budget" => debug(&config.kv_cache_budget), + "kv_admission_watermark" => json!(config.kv_admission_watermark), "enable_vlm_prefix_cache" => json!(config.enable_vlm_prefix_cache), "cors_policy" => debug(&config.cors_policy), "serving_mode" => debug(&config.serving_mode), @@ -733,7 +736,8 @@ fn read_only_kind(field: &str) -> KnobKind { | "default_dynatemp_range" | "default_dynatemp_exponent" | "default_adaptive_target" - | "default_adaptive_decay" => KnobKind::Float, + | "default_adaptive_decay" + | "kv_admission_watermark" => KnobKind::Float, "api_keys" | "lora_adapters" | "lora_runtime" => KnobKind::Object, "reasoning_budget_message" | "model_alias" diff --git a/src/server/runtime_settings_tests.rs b/src/server/runtime_settings_tests.rs index 5a2c710e6..13fc3b305 100644 --- a/src/server/runtime_settings_tests.rs +++ b/src/server/runtime_settings_tests.rs @@ -59,9 +59,9 @@ fn rejection_reason<'a>(result: &'a ApplyResult, name: &str) -> &'a str { } #[test] -fn server_config_schema_classifies_all_110_fields() { +fn server_config_schema_classifies_all_111_fields() { let declared = declared_server_config_fields(); - assert_eq!(declared.len(), 110, "ServerConfig field count changed"); + assert_eq!(declared.len(), 111, "ServerConfig field count changed"); assert_eq!( declared.as_slice(), CLASSIFIED_SERVER_CONFIG_FIELDS, @@ -69,7 +69,7 @@ fn server_config_schema_classifies_all_110_fields() { ); let specs = schema(&ServerConfig::default()); - assert_eq!(specs.len(), 110); + assert_eq!(specs.len(), 111); let expected_names: Vec<_> = CLASSIFIED_SERVER_CONFIG_FIELDS .iter() .copied() @@ -79,7 +79,7 @@ fn server_config_schema_classifies_all_110_fields() { assert_eq!(actual_names, expected_names); assert_eq!( actual_names.iter().copied().collect::>().len(), - 110, + 111, "every management API name must be unique" ); diff --git a/src/server/startup.rs b/src/server/startup.rs index f8ad92b8e..785bf45da 100644 --- a/src/server/startup.rs +++ b/src/server/startup.rs @@ -565,6 +565,10 @@ pub struct ServerStartupConfig { /// semantics. `None` (the default) keeps the pool unbounded. pub kv_cache_budget: Option, + /// Paged KV admission watermark (`--kv-admission-watermark`, issue + /// #2088). `None` keeps [`crate::server::DEFAULT_KV_ADMISSION_WATERMARK`]. + pub kv_admission_watermark: Option, + /// maximum number of responses kept in the /// [`crate::server::responses_store::ResponsesStore`]. `0` disables /// response persistence entirely, in which case `GET /v1/responses/:id` @@ -812,6 +816,7 @@ impl Default for ServerStartupConfig { // pairs with the batched-decode default. Disable with // `--kv-cache-budget none`. kv_cache_budget: Some(crate::memory_estimate::PagedBudgetDirective::Auto), + kv_admission_watermark: None, responses_store_max_entries: 1024, responses_store_max_bytes: super::responses_store::DEFAULT_RESPONSES_STORE_MAX_BYTES, responses_store_ttl_secs: 3600, @@ -1764,6 +1769,12 @@ pub(super) fn build_server_config( // forward the paged KV block-budget directive verbatim; the worker // resolves it to a concrete block count once the model is loaded. kv_cache_budget: startup.kv_cache_budget, + // issue #2088: resolve the admission watermark to its default here. + kv_admission_watermark: crate::server::clamp_kv_admission_watermark( + startup + .kv_admission_watermark + .unwrap_or(crate::server::DEFAULT_KV_ADMISSION_WATERMARK), + ), // forward the experimental VLM prefix-cache toggle (#124 step c). enable_vlm_prefix_cache: startup.enable_vlm_prefix_cache, // forward the validated CORS allow-list (#244); `None` keeps permissive.