fix(server): close residual paged-budget gaps left by #2077

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.

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%를 남긴다.