merge: main into refactor/issue-2169-row-sampler

Brings in #2205 (one prefill plan and one trim hook). The only conflict was the import block in generate.rs (prefill_plan alongside RowSampler, the replaced sampler-state import dropped). Post-merge clippy, the sampling, finish-step and prefill-plan test modules, and the real-checkpoint parity runs match the baseline.

Author: inureyes

TECHNICAL_REPORTS/2205-one-prefill-plan-trim-hook-20261007.en.md

@@ -0,0 +1,91 @@
+# Technical Report: PR #2205 - One prefill plan and one trim hook for the CLI and the server
+
+**Date**: 2026-10-07
+
+**Status**: Implemented and verified on GB10; pending merge.
+
+**Languages**: Rust, Markdown
+
+**Risk Level**: Medium.
+- Every CLI and server prefill now executes one plan, and the server's default prefill chunk changes from 512 to 2048.
+- Greedy output on three real checkpoints is unchanged before and after.
+- The new default's effect on concurrent decode latency is not measured yet; the epic's end-of-run benchmark covers it.
+
+## Executive Summary
+
+Phase 3 of epic #2166 (#2170).
+- **One plan.** `mlxcel_core::prefill_plan::PrefillPlan` now decides, in one place, how a prompt is split into forwards: the adopted prompt-cache prefix, the history-boundary split, chunking, tile padding and the pad trim. Every CLI and server prefill site only executes its pieces.
+- **One trim hook.** `LanguageModel::trim_state(Option<SequenceId>, excess)` replaces the two separate hooks.
+- **One chunk size.** The chunk is a single policy, 2048 per ADR 0007's measurement.
+- **Prompt-cache divergence explained.** The hit vs miss divergence found in Phase 0 is root-caused (the prefill partition) and bounded by an explicit invariant.
+
+## 1. Problem Statement
+
+- **Prefill was decided twice.** The CLI chunked at 2048, the server at 512, and padding and trimming were coded separately on each side.
+- **Two trim hooks.** The model API had `trim_internal_caches` for the CLI and `trim_sequence_state` for the server, so a family had to implement both. #2140 was that failure.
+- **Prompt-cache divergence.** The Phase 0 parity baseline showed a server prompt-cache hit diverging from the uncached run of the same prompt; Qwen3-1.7B greedy diverged at token 0.
+
+## 2. Change Summary
+
+- **`PrefillPlan`:**
+  - `new`/`with_prefix(prompt_len, adopted, boundary, chunk, PrefillCaps)` produce ordered `PrefillPiece { range, padded_len }`.
+  - The boundary segment is one unpadded piece, and chunks restart at the boundary.
+  - The final piece is tile-padded where the model and input allow it.
+  - `reproduces` and `split_points` state when a cache hit can reproduce a miss.
+- **CLI:** `prefill_prompt_last_logits`, the `generate_with_stats` copy and both embedding variants run plan pieces through one executor.
+- **Server:**
+  - `planned_prefill.rs` runs pieces for full, chunked, batched and handoff prefill. The plan is rebuilt from the sequence each tick, with `prefill_offset` as the cursor.
+  - The history-boundary snapshot is inserted after the boundary piece, with no extra forward.
+  - Gemma 4's `mtp_prefill_ranges` and the paged block reservation use the plan.
+- **`trim_state`:** seven model-owned families implement it, the VLM wrappers forward it, and `LoadedModel` delegates to it. `rewind_decode_appends` stays a separate exact-rewind contract.
+- **Chunk policy:** `prefill_chunk_len()` reads `MLXCEL_PREFILL_CHUNK`, else 2048. It is now also the server's `--prefill-chunk-size` default, on both server binaries.
+  - The memory estimate's activation term follows it (`activation_prefill_tokens()`).
+  - The batched-prefill token budget caps each row's share at 512, so the cohort budget stays 4096.
+
+## 3. Technical Decisions
+
+**Root cause of the cache divergence: the partition, not adoption.** Under `MLXCEL_SDPA_DETERMINISTIC=1`, chunking the miss at the cached length made miss and hit identical:
+- Qwen3-1.7B: chunk at 45.
+- Llama-3.2-1B: chunk at 71.
+
+The mechanism: the dense-KV miss is one 52- or 75-row forward (tiled qmm). The hit forwards only a 7- or 4-row suffix, which takes the per-row qmv kernel (`M*B < 8`) and a different attention tiling. The two reduce in a different order and flip a near-tied token.
+
+**Documented bound instead of forcing the split.** A hit reproduces the miss exactly when its adopted prefix ends on a split point of the miss's plan, and the adopted rows were written by those same pieces. Splitting every family at the history boundary would make the harness rows identical, but it was rejected for three reasons:
+- It adds a forward launch to every cold chat prefill.
+- It would make the prompt cache change single-turn output.
+- It does not fix whole-prompt replay anyway.
+
+**One chunk size: 2048.** Phase 0 measured 8 to 29 percent lower TTFT at 8192 tokens with 2048. Review found three places that silently assumed 512, and they were fixed:
+- `mlx_server`'s flag default.
+- A `clamp(1, 0)` panic in the memory estimate.
+- The batched-prefill budget.
+
+## 4. Validation
+
+- **Unit tests:**
+  - `prefill_plan` 11, mlxcel-core `prefill` 76, `generate::` 46, `decode_finish` 9.
+  - Root modules run one at a time: `scheduler_prompt_cache_plan_tests` 3 (including the chunk-counter and progress-frame cases), `server::batch::scheduler::tests` 94, `scheduler_prompt_cache_tests` 31, `prefill_span_coverage_tests` 3, `finish_step_tests`, `speculative_burst_tests`, `prefill_cohort`, `cli_input`, `memory_estimate`, `gemma4_mtp_target`.
+  - `vlm_wrapper_capability_delegation`, including the tile-padding guard audit.
+- **Real checkpoints:** `mlxcel-engine-parity` with `MLXCEL_SDPA_DETERMINISTIC=1` on qwen3-1.7b-4bit, llama-3.2-1b-instruct-4bit and lfm2-350m-8bit.
+  - Every pair result was unchanged before, after and after the merge with `main`.
+  - CLI and server token streams were identical to the base commit.
+- **Known flakiness:** the full `server::batch::scheduler` filter aborts intermittently in CUDA graph capture. Every module run alone is clean.
+
+## 5. Residual Risks
+
+- **Concurrent decode latency is unmeasured.** Inter-token latency of concurrent decode streams under the 2048 default is left to the end-of-epic benchmark. A longer prefill piece delays decode grants by up to four times.
+- **History segments are not chunked.** A snapshot family's history segment is one unchunked forward regardless of chunk size. This is pre-existing (#1143).
+- **One environment variable now drives both paths.** `MLXCEL_PREFILL_CHUNK` now also sets the server default, and the memory estimate does not see an explicit `--prefill-chunk-size`.
+- **No numeric bound test.** Nothing asserts a numeric bound on hit-vs-miss divergence. The bound is the documented invariant plus the measured table.
+
+## 6. Learning Points
+
+- **Treat a partition as part of the numerics.** Two correct prefills of the same tokens differ in their low bits when they are split differently, because the kernel dispatch depends on row count.
+- **Changing a default needs a sweep for its implicit dependents.** Three of the four review fixes were places that assumed the old 512 without naming it.
+
+## 7. Related
+
+- Epic #2166, issue #2170.
+- ADR 0007, ADR 0005.
+- #2201 (finish step, merged into this branch).
+- #2140, #1143, #2185.

TECHNICAL_REPORTS/2205-one-prefill-plan-trim-hook-20261007.ko.md

@@ -0,0 +1,91 @@
+# 기술 보고서: PR #2205 - CLI와 서버가 함께 쓰는 prefill 계획 하나와 trim 훅 하나
+
+**작성일**: 2026-10-07
+
+**상태**: GB10에서 구현하고 검증했으며 머지 대기 중이다.
+
+**언어**: Rust, Markdown
+
+**위험도**: 중간.
+- 이제 CLI와 서버의 모든 prefill이 계획 하나를 실행하고, 서버의 기본 prefill 청크가 512에서 2048로 바뀐다.
+- 실제 체크포인트 세 개에서 greedy 출력은 변경 전후가 같다.
+- 새 기본값이 동시 decode 지연에 주는 영향은 아직 측정하지 않았고, epic의 마지막 측정에서 다룬다.
+
+## 요약
+
+epic #2166의 Phase 3(#2170)이다.
+- **계획 하나**: 이제 `mlxcel_core::prefill_plan::PrefillPlan` 한 곳에서 프롬프트를 어떻게 forward로 나눌지 정한다. 재사용할 prompt-cache 접두사, 대화 이력 경계에서의 분할, 청크 분할, 타일 패딩, 패딩 trim이 모두 여기서 정해진다. CLI와 서버의 prefill 경로는 그 계획의 조각을 실행하기만 한다.
+- **trim 훅 하나**: 따로 있던 두 훅을 `LanguageModel::trim_state(Option<SequenceId>, excess)` 하나로 바꿨다.
+- **청크 크기 하나**: ADR 0007의 측정에 따라 2048 하나로 정했다.
+- **prompt-cache 차이 설명**: Phase 0에서 발견한 캐시 적중과 미적중의 출력 차이를 원인(prefill 분할)까지 밝히고, 명시적인 불변식으로 범위를 정했다.
+
+## 1. 문제 정의
+
+- **prefill 결정이 두 곳에 있었다.** CLI는 2048, 서버는 512로 청크를 나눴고, 패딩과 trim도 양쪽에 따로 짜여 있었다.
+- **trim 훅이 두 개였다.** 모델 API에 CLI용 `trim_internal_caches`와 서버용 `trim_sequence_state`가 따로 있어서, 모델 계열마다 둘 다 구현해야 했다. #2140이 그 때문에 생긴 문제다.
+- **prompt-cache 차이**: Phase 0 패리티 기준선에서, 서버의 prompt-cache 적중 실행이 같은 프롬프트를 캐시 없이 돌린 결과와 달랐다. Qwen3-1.7B greedy는 0번째 토큰부터 달랐다.
+
+## 2. 변경 요약
+
+- **`PrefillPlan`**
+  - `new`/`with_prefix(prompt_len, adopted, boundary, chunk, PrefillCaps)`가 순서대로 `PrefillPiece { range, padded_len }`를 만든다.
+  - 대화 이력 경계 구간은 패딩 없는 조각 하나로 두고, 청크는 경계부터 다시 센다.
+  - 마지막 조각은 모델과 입력이 허용할 때만 타일에 맞춰 패딩한다.
+  - `reproduces`와 `split_points`는 캐시 적중이 미적중 결과를 재현할 수 있는 조건을 나타낸다.
+- **CLI**: `prefill_prompt_last_logits`, `generate_with_stats` 안의 복사본, 임베딩 변형 두 개가 모두 하나의 실행기로 계획의 조각을 실행한다.
+- **서버**
+  - `planned_prefill.rs`가 전체 prefill, 청크 prefill, 배치 prefill, handoff prefill의 조각을 실행한다. 계획은 tick마다 시퀀스에서 다시 만들고, 진행 위치는 `prefill_offset`이 가리킨다.
+  - 대화 이력 경계의 snapshot은 경계 조각을 실행한 직후에 forward를 더 돌리지 않고 넣는다.
+  - Gemma 4의 `mtp_prefill_ranges`와 paged 블록 예약도 계획을 쓴다.
+- **`trim_state`**: 자체 상태를 가진 모델 계열 일곱 개가 구현하고, VLM wrapper는 전달하며, `LoadedModel`은 위임한다. `rewind_decode_appends`는 별개의 정확한 되감기 계약으로 그대로 둔다.
+- **청크 정책**: `prefill_chunk_len()`은 `MLXCEL_PREFILL_CHUNK`를 읽고, 없으면 2048을 쓴다. 이제 서버 바이너리 두 개 모두에서 `--prefill-chunk-size`의 기본값이기도 하다.
+  - 메모리 추정의 활성화 항(`activation_prefill_tokens()`)도 이 값을 따른다.
+  - 배치 prefill 토큰 예산은 행마다 512로 상한을 두어 묶음 예산을 4096으로 유지한다.
+
+## 3. 기술적 선택과 그 이유
+
+**캐시 차이의 원인은 재사용이 아니라 분할이다.** `MLXCEL_SDPA_DETERMINISTIC=1`에서 미적중 실행을 캐시된 길이에서 청크로 나누자 두 실행이 같아졌다.
+- Qwen3-1.7B: 45에서 분할
+- Llama-3.2-1B: 71에서 분할
+
+원리는 이렇다. dense KV를 쓰는 미적중 실행은 52행 또는 75행을 forward 한 번으로 처리한다(타일형 qmm). 적중 실행은 7행 또는 4행 꼬리만 forward하는데, 이 크기는 행 단위 qmv 커널(`M*B < 8`)과 다른 attention 타일링을 탄다. 두 실행은 합산 순서가 달라서, 점수가 거의 같은 토큰이 뒤바뀐다.
+
+**분할을 강제하지 않고 경계를 문서로 정했다.** 적중 실행은 재사용하는 접두사가 미적중 실행의 분할 지점에서 끝나고, 그 행들을 같은 조각이 썼을 때 미적중 결과를 정확히 재현한다. 모든 모델 계열을 대화 이력 경계에서 강제로 나누면 하네스 결과는 같아지지만, 세 가지 이유로 택하지 않았다.
+- 모든 cold chat prefill에 forward가 한 번 더 든다.
+- prompt cache가 한 턴짜리 출력까지 바꾸게 된다.
+- 프롬프트 전체를 다시 보내는 경우는 어차피 해결되지 않는다.
+
+**청크 크기는 2048 하나다.** Phase 0에서 2048이 8192토큰 기준 TTFT를 8~29% 낮췄다. 리뷰에서 512를 말없이 가정하던 곳 세 군데를 찾아 고쳤다.
+- `mlx_server`의 플래그 기본값
+- 메모리 추정의 `clamp(1, 0)` panic
+- 배치 prefill 예산
+
+## 4. 검증
+
+- **단위 테스트**
+  - `prefill_plan` 11, mlxcel-core `prefill` 76, `generate::` 46, `decode_finish` 9.
+  - 루트 크레이트는 모듈을 하나씩 돌렸다: `scheduler_prompt_cache_plan_tests` 3(청크 카운터와 진행 프레임 경우 포함), `server::batch::scheduler::tests` 94, `scheduler_prompt_cache_tests` 31, `prefill_span_coverage_tests` 3, `finish_step_tests`, `speculative_burst_tests`, `prefill_cohort`, `cli_input`, `memory_estimate`, `gemma4_mtp_target`.
+  - 타일 패딩 보호 감사를 포함한 `vlm_wrapper_capability_delegation`.
+- **실제 체크포인트**: qwen3-1.7b-4bit, llama-3.2-1b-instruct-4bit, lfm2-350m-8bit에서 `MLXCEL_SDPA_DETERMINISTIC=1`로 `mlxcel-engine-parity`를 돌렸다.
+  - 모든 비교 쌍의 결과가 변경 전, 변경 후, main 병합 후 모두 같았다.
+  - CLI와 서버의 토큰 열은 변경 전 커밋과 같았다.
+- **알려진 불안정성**: `server::batch::scheduler` 필터 전체를 한 번에 돌리면 CUDA graph capture에서 가끔 중단된다. 모듈을 하나씩 돌리면 모두 깨끗하다.
+
+## 5. 남은 위험
+
+- **동시 decode 지연을 재지 않았다.** 2048 기본값에서 동시 decode 스트림의 토큰 간 지연은 epic 마지막 측정으로 넘겼다. prefill 조각이 길어지면 decode 차례가 최대 네 배까지 늦어진다.
+- **대화 이력 구간은 청크로 나뉘지 않는다.** snapshot 계열의 이력 구간은 청크 크기와 상관없이 forward 한 번으로 처리된다. 이전부터 있던 동작이다(#1143).
+- **환경 변수 하나가 양쪽을 움직인다.** `MLXCEL_PREFILL_CHUNK`가 이제 서버 기본값도 정하고, 메모리 추정은 명시적인 `--prefill-chunk-size`를 보지 못한다.
+- **수치 경계 테스트가 없다.** 적중과 미적중의 차이에 수치 경계를 거는 테스트는 없다. 경계는 문서화한 불변식과 측정표가 대신한다.
+
+## 6. 학습 포인트
+
+- **분할 방식도 수치 계산의 일부로 봐야 한다.** 같은 토큰을 올바르게 처리한 두 prefill도 나누는 방식이 다르면 낮은 비트가 달라진다. 커널 선택이 행 수에 따라 달라지기 때문이다.
+- **기본값을 바꿀 때는 그 값에 암묵적으로 기대던 곳을 모두 찾아야 한다.** 리뷰에서 고친 네 건 가운데 세 건이 이름을 밝히지 않고 512를 가정하던 곳이었다.
+
+## 7. 관련 항목
+
+- epic #2166, 이슈 #2170.
+- ADR 0007, ADR 0005.
+- #2201(종료 처리 통합, 이 브랜치에 병합됨).
+- #2140, #1143, #2185.