diff --git a/docs/CONTINUOUS_BATCHING.md b/docs/CONTINUOUS_BATCHING.md index 6a26d52cd..8167b2312 100644 --- a/docs/CONTINUOUS_BATCHING.md +++ b/docs/CONTINUOUS_BATCHING.md @@ -181,6 +181,70 @@ savings, the decode throughput, and `--kv-cache-budget` are documented in byte-identical to the dense backend; it is the storage backend the disaggregated roles below build on. +#### The fused decode kernel + +Since v0.4.4 a batched paged decode step runs as **one fused attention launch +over the whole batch** (issue #899) rather than as a per-sequence loop that +copied each sequence's visible KV out of the pool before attending. The kernel +reads the pool's scattered pages in place through a CSR page table and splits +the KV range across thread blocks, so its parallelism grows with context +instead of saturating. Measured against the previous path on an Apple M1 Ultra: +1.4x to 3.1x at batch 4 across 1K to 32K of context, and 1.3x to 1.5x for a +single sequence at 16K and 32K +(`docs/benchmark_results/paged-decode-v2-m1ultra-2026-07-31.md`). + +What it does *not* serve falls back to the previous gather-then-SDPA path, per +launch, with identical output: + +| case | why | +|---|---| +| a lone request under 4096 visible KV tokens | measured slower there (0.91x at batch 1 with 1K of context) | +| a multi-request launch under 512 visible tokens per request | below the measured evidence, and gather is trivially cheap at that size | +| a layer whose pool rows span more than one slab | the kernel reads one contiguous buffer per side | +| multi-token steps: batched prefill, speculative / MTP verify | the kernel serves a single query token per sequence | +| logit-softcap families, explicit attention masks | no soft-cap or mask term in the kernel | +| Int8 / Turbo KV modes | those sequences keep dense caches and are never pool-backed | + +The floor is two-regime because the measured loss is a property of batch 1, not +of total work: at 1024 tokens of context per request the kernel runs 0.91x at +batch 1 but 1.41x at batch 4 and 1.47x at batch 8, since a batched launch +spreads the same chunk count over more requests and amortizes the merge pass. + +#### Seeing which path ran + +The server announces the first occurrence of each distinct dispatch outcome at +**info**, with no `RUST_LOG` needed: + +```text +INFO ...: paged decode v2: fused v2 launch (batch 4, 3828 visible KV tokens, 16 chunks, merge on) +INFO ...: paged decode v2: gather: the layer spans 3 K / 3 V pool slabs and the fused kernels read one buffer per side (slab_blocks=32); raise --ctx-size or MLXCEL_PAGED_SLAB_BLOCKS so a layer's rows fit one slab +INFO ...: paged decode v2: gather: 1024 visible KV tokens across 1 request(s) is below the 4096-token dispatch floor +INFO ...: paged decode v2: gather: pinned by MLXCEL_PAGED_ATTENTION_NATIVE +``` + +A run that never prints `fused v2 launch` never used the fused kernel, whatever +its throughput looks like. Check for that line before drawing any conclusion +from a before/after comparison. + +Two knobs, neither normally needed: + +- `MLXCEL_PAGED_ATTENTION_NATIVE=0` is the kill switch. It restores the + pre-#899 gather behaviour end to end, for bisecting a suspected kernel + regression. `=1` forces the fused path for every servable shape, bypassing + the token floor; that is the supported way to benchmark the declined corner. + The floors themselves move with `MLXCEL_PAGED_V2_MIN_KV_TOKENS` (lone + request, default 4096) and `MLXCEL_PAGED_V2_MIN_KV_TOKENS_PER_REQUEST` + (multi-request, default 512). +- `MLXCEL_PAGED_SLAB_BLOCKS` overrides the pool's slab size in blocks. The + server derives it from `--ctx-size` and `--parallel`, clamped by the KV + budget, and logs the resolved value at startup (`Paged KV slab size: N blocks + per layer`). A layer's first write allocates its whole slab, so this + front-loads exactly the KV bytes the startup memory estimate already reports + for that batch and context. **Serving contexts longer than `--ctx-size` + implies makes layers outgrow the slab and silently returns them to the gather + path**; raise `--ctx-size` or set this variable if the startup line looks too + small. `0` pins the historical 32-block default. + Recurrent and hybrid SSM / linear-attention families cannot safely reuse arbitrary KV blocks, so they keep the hybrid-SSM/APC exclusion. Families that opt into `supports_snapshot_reuse()` instead use a separate exact-prefix diff --git a/docs/adr/0001-paged-attention-gather-vs-fused-kernel.md b/docs/adr/0001-paged-attention-gather-vs-fused-kernel.md index c90ef452a..51ecb392b 100644 --- a/docs/adr/0001-paged-attention-gather-vs-fused-kernel.md +++ b/docs/adr/0001-paged-attention-gather-vs-fused-kernel.md @@ -1,6 +1,8 @@ # ADR 0001: Paged-attention gather strategy and KV pool tensor layout -**Status:** Accepted (2026-05-31). Part of epic #116 (unified paged KV cache), Phase 0 (#117). +**Status:** Superseded in part by issue #899 (2026-07-31). Accepted (2026-05-31), part of epic #116 (unified paged KV cache), Phase 0 (#117). + +The layout and block-size decisions below still stand. The *attention strategy* decision does not: this ADR chose gather-then-SDPA and named the trigger that would justify building a fused kernel, that trigger has now fired and been measured, and since #899 the fused path is what serves production batched paged decode. See the 2026-07-31 addendum on issue #899 at the end of this document; the historical reasoning is kept intact above it because the measurements are what justified the switch. ## Context @@ -188,6 +190,20 @@ Scope of #898: library only. `MLXCEL_PAGED_ATTENTION_V2=1` selects it inside `Pa Correctness on an M1 Ultra: 68 configurations against the gather-then-SDPA reference across head dims 64/128, GQA 1/4/8, page sizes 16/32/64, contexts 512 to 32768, and batches 1/4/8, worst relative deviation 5.2e-4 against a 2e-2 tolerance. +## Addendum (2026-07-31): the trigger fired, the fused path is production (#899) + +This ADR set the trigger for building a fused path at "single-sequence context past ~16384 tokens, or any sustained batched decode". Since v0.4 the server defaults to `--parallel 4`, which is that second condition, and #898's v2 kernel was measured against the production gather path on an Apple M1 Ultra (`docs/benchmark_results/paged-decode-v2-m1ultra-2026-07-31.md`). Both trigger points are met: batch 4 runs 1.41x to 3.08x, and single-sequence decode runs 1.29x at 16384 and 1.47x at 32768. The 2x-3x figure this ADR predicted for batch 4 past 4096 tokens is reproduced almost exactly (2.04x, 2.78x, 3.08x). + +**The decision is therefore reversed for batched paged decode.** `PagedBlockPool::paged_decode_batched`, reached from `cache::paged_batch_decode::paged_batch_decode_attention`, is the production path for `--decode-storage paged`; `gather_visible` + SDPA remains as the fallback and the parity reference. + +Three qualifications, each with its own reason: + +**The switch is not unconditional.** Exactly one measured cell loses: batch 1 at 1024 tokens, 0.91x median and below parity in all three repetitions, where the plan degenerates to two pages per chunk and the merge pass plus the workspace round trip costs more than the attention. `paged_v2::dispatch` keeps that regime on gather with a two-regime floor: a lone request needs 4096 visible KV tokens, a multi-request launch needs 512 per request. The first formulation used a single 4096-token total and had to be corrected: the loss is a property of batch 1, not of total work (the same 1024-token context runs 0.91x at batch 1 and 1.41x at batch 4), and a total-token floor separated the two only by accident, declining a real batch-4 launch that delivered 956 tokens per request. + +**The single-slab constraint from #235 had to be relaxed, not removed.** The fused kernels read one contiguous pool buffer per side, and 32-row slabs capped them at 1024 tokens across the whole batch, below the dispatch floor. `PagedBlockPool` now carries its own slab size (default unchanged) and the server sizes it from the configuration the KV budget was already computed from. A layer that outgrows its slab still appends without copying and simply stops being eligible, so the degradation is to this ADR's original strategy rather than to a failure. + +**The library-only entry point retired in #710 stays retired.** `paged_decode_attention_pooled` and `select_pooled_paged_dispatch` are untouched: they still serve external `mlxcel-core` consumers with the v1 kernel and the #331 selector. #899 adds a separate production entry point rather than reviving that one, because the two have different inputs (a whole batch versus one sequence at a time) and different dispatch evidence. + ## References - Epic #116, unified KV cache. diff --git a/docs/benchmark_results/paged-decode-production-m1ultra-2026-08-02.md b/docs/benchmark_results/paged-decode-production-m1ultra-2026-08-02.md new file mode 100644 index 000000000..99c8dcdd8 --- /dev/null +++ b/docs/benchmark_results/paged-decode-production-m1ultra-2026-08-02.md @@ -0,0 +1,121 @@ +# Production paged decode through fused v2: Apple M1 Ultra, 2026-08-02 + +Validation run for issue #899 (epic #909), which routes the server's batched +paged decode through the fused v2 kernel from #898 and retires gather-then-SDPA +as the default. + +**Result: decode throughput improves in four of five scenarios, 1.15x to 1.57x, +with parity in the fifth and no regression anywhere.** Aggregate throughput moves +much less, for a reason explained below that is a property of the metric rather +than of the kernel. + +## Environment + +| Field | Value | +|---|---| +| Hardware | Apple M1 Ultra, 20 cores (16P + 4E), 128 GB unified memory | +| OS | macOS 26.5.2 (Darwin 25.5.0) | +| Backend | Metal | +| mlxcel | 0.4.3, branch `feature/issue-899-production-paged-v2-dispatch` | +| Model | `mlx-community/Llama-3.2-1B-Instruct-4bit` (2048 hidden, 32 q heads, 8 kv heads, head_dim 64, 16 layers) | +| Harness | `scripts/benchmark_paged_decode_production.sh`, `--parallel 4`, 128 decode tokens | +| Arms | `before` = `MLXCEL_PAGED_ATTENTION_NATIVE=0` (gather), `after` = default (fused v2) | +| Load average | 5.0 to 5.5 throughout, from unrelated concurrent work | +| Repetitions | 3 full sweeps, fresh server per arm per sweep | + +CUDA was not available. The CUDA JIT bodies inherited from #898 remain +uncompiled and unrun. + +## Dispatch was verified, not assumed + +This is recorded first because the first attempt at this benchmark produced a +completely null result (0.98x to 1.07x everywhere) that looked like a legitimate +finding and was not. **Both arms had run the gather path**: the fused kernel never +launched once across the entire sweep, for two independent reasons (a token floor +whose shape excluded the measured scenarios, and a single-sequence decode path +that was never wired). A before/after comparison of gather against gather agrees +to within noise, which is indistinguishable from "the kernel does not help". + +Every run below therefore carries proof of which kernel it executed, taken from +the server logs: + +``` +before: paged decode v2: gather: pinned by MLXCEL_PAGED_ATTENTION_NATIVE +after: paged decode v2: fused v2 launch (batch 4, 3828 visible KV tokens, 16 chunks, merge on) +``` + +The harness now fails the run if the after arm never logs a fused launch, or if +the before arm does. A null sweep can no longer be produced silently. + +## Decode throughput, after / before + +The metric the kernel actually affects. Three independent sweeps. + +| scenario | rep 1 | rep 2 | rep 3 | median | +|---|---|---|---|---| +| 4 clients, ctx 1024 | 1.33x | 1.82x | 1.57x | **1.57x** | +| 4 clients, ctx 4096 | 0.98x | 1.00x | 1.02x | 1.00x | +| 4 clients, ctx 16384 | 1.15x | 1.13x | 1.15x | **1.15x** | +| 1 client, ctx 16384 | 1.24x | 1.42x | 1.22x | **1.24x** | +| 1 client, ctx 32768 | 1.41x | 1.41x | 1.44x | **1.41x** | + +Two cells are tight enough to quote without hedging: 4 clients at 16384 spans +1.13x to 1.15x, and 1 client at 32768 spans 1.41x to 1.44x. Both are large +relative to their own spread and both reproduce in every sweep. + +The 4-client / 1024 cell is the noisiest (1.33x to 1.82x) but never drops below +1.33x, so the direction is not in doubt even though the magnitude is. + +## Aggregate throughput, after / before + +| scenario | rep 1 | rep 2 | rep 3 | +|---|---|---|---| +| 4 clients, ctx 1024 | 1.21x | 1.20x | 1.23x | +| 4 clients, ctx 4096 | 1.00x | 1.00x | 1.01x | +| 4 clients, ctx 16384 | 1.00x | 1.00x | 1.00x | +| 1 client, ctx 16384 | 1.09x | 1.26x | 1.18x | +| 1 client, ctx 32768 | 1.04x | 1.05x | 1.06x | + +Aggregate includes time to first token, so at long context it is dominated by +prefill, which this change does not touch. At 4 clients and 16384 tokens the +server spends several seconds prefilling four prompts before decoding 128 tokens +each, so a 1.15x decode improvement is diluted to 1.00x. Read the decode column +for the kernel's effect and this column for what a user of that particular +workload shape would feel. + +## Against the issue's required outcomes + +- **"No scenario regresses beyond noise (3%)."** Met. The weakest cell is 4 + clients at 4096, whose median is exactly 1.00x (0.98x, 1.00x, 1.02x across + sweeps), which is parity rather than regression. +- **"Batched decode aggregate throughput improves at 4K+ context."** Partially + met, and stated precisely rather than rounded up: batched *decode* improves at + 16384 (1.15x, reproducible) but is at parity at 4096, and batched *aggregate* + does not improve at either because prefill dominates it. ADR-0001's prediction + that removing gather overhead is worth 2x-3x of SDPA time is an op-level claim, + and #898 reproduced it at the op level (2.78x at this shape); it does not follow + that end-to-end serving throughput moves by the same factor, because attention + is one component of a decode step. + +## The 4096 cell + +Batched decode at 4096 tokens is the one scenario that gains nothing, sitting +between a 1.57x gain at 1024 and a 1.15x gain at 16384. That non-monotonicity is +not explained. It is reproducible (three sweeps within 4%), so it is a real +property of that shape rather than noise, and it is worth understanding before +the dispatch floor is tuned further. Recorded as an open question rather than +smoothed over. + +## Open items + +- The non-monotonic 4096 cell above. +- Batch 2 and 3 are interpolated in the dispatch floor: every multi-request cell + #898 measured was batch 4 or 8. +- The 1-client / 4096 crossover cell has the least margin of any win (1.08x at + the op level). If single-sequence 4K regresses in a future run, + `MIN_SINGLE_REQUEST_KV_TOKENS` can be raised to 8192, which gives up only that + cell. +- CUDA remains entirely unvalidated across #898 and #899. +- One model family measured end to end. Greedy parity covers three families + (Llama3, Qwen2.5, Qwen3), all exact against the gather baseline, but throughput + was measured only on Llama-3.2-1B. diff --git a/scripts/benchmark_paged_decode_production.sh b/scripts/benchmark_paged_decode_production.sh new file mode 100755 index 000000000..1b0735df2 --- /dev/null +++ b/scripts/benchmark_paged_decode_production.sh @@ -0,0 +1,210 @@ +#!/usr/bin/env bash +# +# Server-level before/after benchmark for the production paged decode v2 path +# (issue #899). +# +# Runs the issue's mandatory matrix twice against the same binary and the same +# model, changing only which decode path serves the batch: +# +# after = the fused v2 kernel (the new default) +# before = gather-then-SDPA, pinned with MLXCEL_PAGED_ATTENTION_NATIVE=0 +# (the kill switch, which restores the pre-#899 behaviour end to end) +# +# Matrix: +# * 4 concurrent clients at 1K / 4K / 16K prompt tokens (the documented v0.4 +# --parallel default), reporting aggregate decode throughput and TTFT. +# * single-sequence long-context decode at 16K and 32K. +# +# The two arms are run serially against a freshly started server each time, so +# no measurement overlaps another. Results land as CSV plus a markdown skeleton +# ready to become docs/benchmark_results/paged-decode-production--.md. +# +# Usage: +# MODEL=models/qwen2.5-7b-instruct-4bit ./scripts/benchmark_paged_decode_production.sh +# +# Environment: +# MODEL model directory (required) +# SERVER_BIN default target/release/mlxcel-server +# PARALLEL concurrent slots, default 4 +# MAX_TOKENS decode tokens per request, default 128 +# CTX_SIZE total --ctx-size; default is PARALLEL * 32768 so the per-slot +# window covers the 32K single-sequence case and the paged slab +# is sized for it. Undersizing this is the single most likely +# way to accidentally measure the gather path in both arms. +# PORT default 18991 +# OUT_DIR default benchmarks/ + +set -euo pipefail + +ROOT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)" +MODEL="${MODEL:-}" +SERVER_BIN="${SERVER_BIN:-$ROOT_DIR/target/release/mlxcel-server}" +PARALLEL="${PARALLEL:-4}" +MAX_TOKENS="${MAX_TOKENS:-128}" +CTX_SIZE="${CTX_SIZE:-$((PARALLEL * 32768))}" +PORT="${PORT:-18991}" +HOST="${HOST:-127.0.0.1}" +OUT_DIR="${OUT_DIR:-$ROOT_DIR/benchmarks}" +DATE_TAG="${DATE_TAG:-$(date +%F)}" +HW_TAG="${HW_TAG:-$(uname -m)}" + +if [[ -z "$MODEL" ]]; then + echo "MODEL is required, e.g. MODEL=models/qwen2.5-7b-instruct-4bit $0" >&2 + exit 2 +fi +if [[ ! -x "$SERVER_BIN" ]]; then + echo "server binary not found at $SERVER_BIN; build with:" >&2 + echo " cargo build --release --features metal,accelerate" >&2 + exit 2 +fi + +mkdir -p "$OUT_DIR" +CSV="$OUT_DIR/paged_decode_production_${HW_TAG}_${DATE_TAG}.csv" +LOG_DIR="$OUT_DIR/paged_decode_production_logs_${DATE_TAG}" +mkdir -p "$LOG_DIR" +echo "arm,scenario,concurrency,prompt_tokens,max_tokens,raw_log" > "$CSV" + +SERVER_PID="" +cleanup() { + if [[ -n "$SERVER_PID" ]]; then + kill "$SERVER_PID" >/dev/null 2>&1 || true + wait "$SERVER_PID" >/dev/null 2>&1 || true + SERVER_PID="" + fi +} +trap cleanup EXIT + +wait_for_health() { + for _ in $(seq 1 300); do + if curl -fsS "http://$HOST:$PORT/health" >/dev/null 2>&1; then + return 0 + fi + sleep 1 + done + echo "server did not become healthy on port $PORT" >&2 + return 1 +} + +start_server() { + local arm="$1" + local log="$LOG_DIR/server_${arm}.log" + echo "--- starting server (arm: $arm), log $log" + if [[ "$arm" == "before" ]]; then + MLXCEL_PAGED_ATTENTION_NATIVE=0 "$SERVER_BIN" \ + --model "$MODEL" --host "$HOST" --port "$PORT" \ + --parallel "$PARALLEL" --ctx-size "$CTX_SIZE" \ + >"$log" 2>&1 & + else + "$SERVER_BIN" \ + --model "$MODEL" --host "$HOST" --port "$PORT" \ + --parallel "$PARALLEL" --ctx-size "$CTX_SIZE" \ + >"$log" 2>&1 & + fi + SERVER_PID=$! + wait_for_health + # Surface the resolved slab size: if this line is missing, the fused path is + # not eligible and the "after" arm will silently measure gather. + grep -m1 "Paged KV slab size" "$log" || \ + echo "NOTE: no slab-size line in the server log (see MLXCEL_PAGED_SLAB_BLOCKS)" +} + +# Confirm each arm ran the kernel it was supposed to. The first sweep for this +# issue produced a full matrix of null results because both arms were on the +# gather path and nothing said so; a benchmark that cannot tell you which kernel +# it measured is worse than no benchmark, so fail loudly instead. +check_arm_dispatch() { + local arm="$1" + local log="$LOG_DIR/server_${arm}.log" + local fused + fused=$(grep -c "paged decode v2: fused v2 launch" "$log" || true) + echo "--- $arm: fused-launch announcements in the server log: $fused" + grep -m3 "paged decode v2:" "$log" || true + if [[ "$arm" == "after" && "$fused" -eq 0 ]]; then + echo "" >&2 + echo "ERROR: the 'after' arm never logged a fused v2 launch, so it measured the" >&2 + echo "gather path and the comparison is meaningless. The decline reason is in" >&2 + echo " $log" >&2 + echo "on a 'paged decode v2: gather: ...' line, which names what to change." >&2 + return 1 + fi + if [[ "$arm" == "before" && "$fused" -ne 0 ]]; then + echo "" >&2 + echo "ERROR: the 'before' arm logged a fused v2 launch, so the kill switch did" >&2 + echo "not take effect and both arms ran the same kernel." >&2 + return 1 + fi + return 0 +} + +run_case() { + local arm="$1" scenario="$2" concurrency="$3" prompt_tokens="$4" + local log="$LOG_DIR/${arm}_${scenario}.txt" + echo "--- $arm / $scenario: concurrency=$concurrency prompt_tokens=$prompt_tokens" + python3 "$ROOT_DIR/scripts/bench_serving_concurrency.py" \ + --host "$HOST" --port "$PORT" \ + --concurrency "$concurrency" \ + --prompt-tokens "$prompt_tokens" \ + --max-tokens "$MAX_TOKENS" \ + | tee "$log" + echo "$arm,$scenario,$concurrency,$prompt_tokens,$MAX_TOKENS,$log" >> "$CSV" +} + +run_arm() { + local arm="$1" + start_server "$arm" + # 4 concurrent clients, the v0.4 default scenario. + run_case "$arm" "batch${PARALLEL}_ctx1k" "$PARALLEL" 1024 + run_case "$arm" "batch${PARALLEL}_ctx4k" "$PARALLEL" 4096 + run_case "$arm" "batch${PARALLEL}_ctx16k" "$PARALLEL" 16384 + # Single-sequence long context. + run_case "$arm" "batch1_ctx16k" 1 16384 + run_case "$arm" "batch1_ctx32k" 1 32768 + check_arm_dispatch "$arm" + cleanup + # Let the allocator settle before the next arm so the two see the same + # starting conditions. + sleep 5 +} + +echo "model : $MODEL" +echo "server : $SERVER_BIN" +echo "parallel : $PARALLEL" +echo "ctx-size : $CTX_SIZE (per slot $((CTX_SIZE / PARALLEL)))" +echo "max-tokens : $MAX_TOKENS" +echo "csv : $CSV" +echo + +run_arm "before" +run_arm "after" + +MD="$OUT_DIR/paged-decode-production-${HW_TAG}-${DATE_TAG}.md" +{ + echo "# Production paged decode v2: ${HW_TAG}, ${DATE_TAG}" + echo + echo "Issue #899. Before = gather-then-SDPA (\`MLXCEL_PAGED_ATTENTION_NATIVE=0\`)," + echo "after = the fused v2 kernel (the new default). Same binary, same model," + echo "same host, arms run serially." + echo + echo "| Field | Value |" + echo "|---|---|" + echo "| Model | \`$MODEL\` |" + echo "| Parallel slots | $PARALLEL |" + echo "| \`--ctx-size\` | $CTX_SIZE (per slot $((CTX_SIZE / PARALLEL))) |" + echo "| Decode tokens per request | $MAX_TOKENS |" + echo "| Raw logs | \`$LOG_DIR\` |" + echo + echo "## Results" + echo + echo "Fill from the per-case logs listed in \`$CSV\`." + echo + echo "| scenario | before aggregate tok/s | after aggregate tok/s | ratio | before TTFT p50 | after TTFT p50 |" + echo "|---|---|---|---|---|---|" + for scenario in "batch${PARALLEL}_ctx1k" "batch${PARALLEL}_ctx4k" "batch${PARALLEL}_ctx16k" batch1_ctx16k batch1_ctx32k; do + echo "| $scenario | | | | | |" + done +} > "$MD" + +echo +echo "CSV : $CSV" +echo "Markdown : $MD (skeleton; fill from the logs)" +echo "Logs : $LOG_DIR" diff --git a/src/execution/memory_estimate.rs b/src/execution/memory_estimate.rs index 5ccd25ccf..e3cbc8030 100644 --- a/src/execution/memory_estimate.rs +++ b/src/execution/memory_estimate.rs @@ -870,9 +870,177 @@ pub fn resolve_paged_block_budget( auto_kv_budget_bytes(&est) } }; + // #899: the fused decode v2 workspace lives outside the block pool but + // inside the same memory, so charge it to the KV budget before converting + // the remainder into blocks. Admission then reserves blocks it can actually + // back with memory. + let workspace = paged_v2_workspace_reserve_bytes(model_dir, num_layers, batch); + let budget_bytes = budget_bytes.saturating_sub(workspace); Some(usize::try_from(budget_bytes / per_block).unwrap_or(usize::MAX)) } +/// Concurrent v2 launches whose workspaces can be alive at once. +/// +/// The decode lookahead pipeline (#632) can have the current step's forward and +/// the next step's prime in flight together, and each holds one layer's +/// workspace; nothing holds more. +const PAGED_V2_CONCURRENT_LAUNCHES: u64 = 2; + +/// Upper bound on the GQA replication factor used for the workspace reserve. +/// +/// The workspace scales with the query-head count, which +/// [`mlxcel_core::hardware::KvCacheParams`] does not carry (it only needs the +/// KV side). Bounding `Hq / Hkv` at 16 covers every family in +/// `docs/supported-models.md` (Llama 3 is 4, Qwen 2.5 7B is 7, the widest MQA +/// ports are 8) with headroom, and over-reserving a few megabytes is the safe +/// direction for an admission bound. +const PAGED_V2_MAX_N_REP: u64 = 16; + +/// Bytes to reserve for the fused decode v2 workspace (issue #899). +/// +/// The workspace is the partial kernel's `(partial_v, lse)` output pair, sized +/// `num_chunks * Hq * (head_dim + 1) * 4` bytes +/// ([`mlxcel_core::paged_v2::PagedDecodePlan::workspace_bytes`]). The plan's +/// binary search stops at the largest chunk size still reaching the device CTA +/// target, so `num_chunks * ctas_per_chunk` lands within a factor of two of +/// that target and +/// +/// ```text +/// num_chunks <= 2 * target_ctas / Hkv + batch +/// ``` +/// +/// since `ctas_per_chunk = Hkv * q_groups >= Hkv`. Substituting +/// `Hq = Hkv * n_rep` makes the dominant term independent of the head counts: +/// `2 * target_ctas * n_rep * (head_dim + 1) * 4`. +/// +/// [`PAGED_V2_MAX_N_REP`] makes the result a few times larger than any real +/// plan needs, which is the safe direction for an admission bound. Even so it +/// lands in the low tens of megabytes at most, well under a tenth of a percent +/// of a serving KV budget, which is why it was not accounted for before v2 +/// became the production path. It is charged now so the accounting is truthful, +/// not because it is large. +#[must_use] +pub fn paged_v2_workspace_reserve_bytes(model_dir: &Path, num_layers: usize, batch: u64) -> u64 { + let Some(params) = kv_cache_params_from_path(model_dir, DEFAULT_CTX_LEN, false, 1) else { + return 0; + }; + if num_layers == 0 || params.num_kv_heads == 0 || params.head_dim == 0 { + return 0; + } + let target = mlxcel_core::paged_v2::device_target_ctas() as u64; + let batch = batch.max(1); + let chunks = target + .saturating_mul(2) + .div_ceil(params.num_kv_heads) + .saturating_add(batch); + let hq = params.num_kv_heads.saturating_mul(PAGED_V2_MAX_N_REP); + let per_launch = chunks + .saturating_mul(hq) + .saturating_mul(params.head_dim.saturating_add(1)) + .saturating_mul(std::mem::size_of::() as u64); + per_launch.saturating_mul(PAGED_V2_CONCURRENT_LAUNCHES) +} + +// ── Paged pool slab sizing (issue #899) ────────────────────────────────────── + +/// Environment override for the paged pool's slab size, in blocks. +/// +/// `0` pins the pool default ([`mlxcel_core::cache::POOL_SLAB_BLOCKS`]), which +/// keeps the pre-#899 allocation behaviour and, as a side effect, keeps the +/// fused decode path unreachable for anything past one slab. Any other positive +/// integer is used verbatim, which is how a benchmark sweeps the setting. +pub const PAGED_SLAB_BLOCKS_ENV: &str = "MLXCEL_PAGED_SLAB_BLOCKS"; + +/// Resolve the paged pool's slab size in blocks (issue #899). +/// +/// ## Why this exists +/// +/// The fused paged-attention decode kernels read **one contiguous pool buffer +/// per side**, so they can only serve a layer whose physical rows all live in +/// the pool's first slab. With the historical 32-block slab that caps them at +/// 1024 tokens across the entire batch (at `block_size` 32), which is below +/// every context the #899 dispatch policy would pick them for. Sizing the slab +/// to the workload is what makes the fused path reachable. +/// +/// ## The policy +/// +/// One slab per layer big enough for the batch the server was configured to +/// run at the context it was configured to serve: +/// +/// ```text +/// blocks = ceil(per_slot_ctx / block_size) * batch +/// ``` +/// +/// clamped below by the pool default and above by the per-layer share of the +/// paged block budget, so the eager allocation can never exceed what +/// `--kv-cache-budget` already reserved. `per_slot_ctx` is the effective +/// per-slot `--ctx-size` (the same figure `--max-kv-size` is resolved from); +/// [`DEFAULT_CTX_LEN`] stands in when the operator did not set one, which is +/// also what [`estimate_total_memory`] assumes. +/// +/// ## What it costs +/// +/// The slab is allocated lazily, per layer, on that layer's first write, and +/// the whole slab is allocated at once. So the first prefill front-loads +/// `blocks * paged_block_bytes * num_layers` bytes: exactly the KV cache for +/// `batch` sequences at `per_slot_ctx` tokens, which is the figure the startup +/// memory estimate already reports. A layer that outgrows its slab appends a +/// second one exactly as before (no copy, #235) and simply stops being eligible +/// for the fused path, so an under-sized slab degrades to the pre-#899 gather +/// behaviour rather than failing. +/// +/// Returns `None` when the geometry is unavailable or the operator pinned the +/// pool default, in which case the caller should not touch the pool's slab +/// size. +#[must_use] +pub fn resolve_paged_slab_blocks( + model_dir: &Path, + num_layers: usize, + block_size: usize, + batch: u64, + per_slot_ctx: u64, + kv_dtype_int8: bool, + block_budget: Option, +) -> Option { + if num_layers == 0 || block_size == 0 { + return None; + } + if let Ok(raw) = std::env::var(PAGED_SLAB_BLOCKS_ENV) { + return match raw.trim().parse::() { + Ok(0) => None, + Ok(n) => Some(n), + Err(_) => { + tracing::warn!( + "{PAGED_SLAB_BLOCKS_ENV}={raw:?} is not a non-negative integer; \ + using the derived slab size" + ); + None + } + }; + } + // The geometry probe doubles as the "is this model pool-eligible at all" + // check: without it the byte clamp below would be meaningless. + paged_block_bytes(model_dir, num_layers, block_size, kv_dtype_int8)?; + + let ctx = if per_slot_ctx == 0 { + DEFAULT_CTX_LEN + } else { + per_slot_ctx + }; + let per_seq_blocks = ctx.div_ceil(block_size as u64); + let want = per_seq_blocks.saturating_mul(batch.max(1)); + let want = usize::try_from(want).unwrap_or(usize::MAX); + + let capped = match block_budget { + // The budget is a global block count across every layer, so a layer may + // claim at most its even share without the eager allocation being able + // to exceed the budget the operator approved. + Some(budget) => want.min((budget / num_layers).max(1)), + None => want, + }; + Some(capped.max(mlxcel_core::cache::POOL_SLAB_BLOCKS)) +} + // ── Output formatting ───────────────────────────────────────────────────────── /// Format a byte count as a human-readable string (GiB, MiB, or exact bytes). @@ -1525,6 +1693,12 @@ mod tests { let tmp = tempfile::tempdir().unwrap(); write_minimal_config(tmp.path()); let per_block = paged_block_bytes(tmp.path(), 32, 32, false).unwrap(); // 131072 + // The paged decode v2 workspace is charged to the byte budget before it + // is divided into blocks (#899), so a caller asking for exactly N blocks + // of KV has to ask for the workspace too. + let workspace = paged_v2_workspace_reserve_bytes(tmp.path(), 32, 1); + assert!(workspace > 0, "the minimal config has a derivable geometry"); + // Exactly 100 blocks. assert_eq!( resolve_paged_block_budget( @@ -1533,7 +1707,7 @@ mod tests { 32, 1, false, - PagedBudgetDirective::Bytes(per_block * 100), + PagedBudgetDirective::Bytes(per_block * 100 + workspace), ), Some(100), ); @@ -1545,7 +1719,7 @@ mod tests { 32, 1, false, - PagedBudgetDirective::Bytes(per_block * 100 + per_block / 2), + PagedBudgetDirective::Bytes(per_block * 100 + per_block / 2 + workspace), ), Some(100), ); @@ -1558,10 +1732,44 @@ mod tests { 32, 1, false, - PagedBudgetDirective::Bytes(per_block - 1), + PagedBudgetDirective::Bytes(per_block - 1 + workspace), ), Some(0), ); + // A budget that did not account for the workspace comes back short by + // exactly the reserve, which is the point: the blocks it hands out are + // blocks it can actually back. + let shortfall = usize::try_from((per_block * 100 - workspace) / per_block).unwrap(); + assert!(shortfall < 100); + assert_eq!( + resolve_paged_block_budget( + tmp.path(), + 32, + 32, + 1, + false, + PagedBudgetDirective::Bytes(per_block * 100), + ), + Some(shortfall), + ); + } + + #[test] + fn the_v2_workspace_reserve_is_small_and_scales_with_the_head_dim() { + let tmp = tempfile::tempdir().unwrap(); + write_minimal_config(tmp.path()); + let reserve = paged_v2_workspace_reserve_bytes(tmp.path(), 32, 4); + // Single-digit megabytes: the reserve exists so admission is truthful, + // not because it is a material share of a KV budget. + assert!(reserve > 0, "expected a non-zero reserve, got {reserve}"); + assert!( + reserve < 64 * 1024 * 1024, + "the reserve should stay small, got {reserve} bytes" + ); + // A model with no derivable geometry reserves nothing rather than + // guessing. + let empty = tempfile::tempdir().unwrap(); + assert_eq!(paged_v2_workspace_reserve_bytes(empty.path(), 32, 4), 0); } #[test] diff --git a/src/lib/mlxcel-core/src/cache.rs b/src/lib/mlxcel-core/src/cache.rs index 9d28496a5..4155f81c2 100644 --- a/src/lib/mlxcel-core/src/cache.rs +++ b/src/lib/mlxcel-core/src/cache.rs @@ -78,6 +78,9 @@ pub mod batch_quant; mod detach; mod paged; +/// Whole-batch decode over pool-backed KV caches, the production paged decode +/// path (issue #899). +pub mod paged_batch_decode; /// CSR page-table view of a decode batch, consumed by the paged decode v2 /// kernels (issue #898). pub mod paged_csr; @@ -103,8 +106,11 @@ pub use batch_quant::{ }; pub use detach::{DetachedCacheSet, DetachedHandle, DetachedKVCache, DetachedRotatingKVCache}; pub use paged::{ - GatheredKv, PagedBlockId, PagedBlockPool, PagedCacheStats, PagedKvLayout, PagedLayerState, - PagedSequenceState, + GatheredKv, POOL_SLAB_BLOCKS, PagedBlockId, PagedBlockPool, PagedCacheStats, PagedKvLayout, + PagedLayerState, PagedSequenceState, +}; +pub use paged_batch_decode::{ + PagedBatchDecodeStats, paged_batch_decode_attention, paged_batch_decode_stats, }; pub use paged_csr::{PagedCsrView, build_paged_csr_view}; pub use paged_detach::DetachedPagedCacheSet; @@ -5874,6 +5880,12 @@ pub struct CachePool { /// to a live pool in [`CachePool::set_paged_block_budget`]). `None` = /// unbounded (the default). paged_block_budget: Option, + /// Paged slab size (physical block rows per pool slab tensor) remembered + /// across lazy pool creation, exactly like [`Self::paged_block_budget`]. + /// `None` leaves [`PagedBlockPool`] on its default, which is the pre-#899 + /// behaviour. The scheduler sets it when decode storage is paged, so the + /// fused decode kernels can serve a layer whose rows all live in one slab. + paged_slab_blocks: Option, } impl CachePool { @@ -5886,6 +5898,7 @@ impl CachePool { paged_pool: None, detached: HashMap::new(), paged_block_budget: None, + paged_slab_blocks: None, } } @@ -6191,6 +6204,28 @@ impl CachePool { self.paged_block_budget } + /// Set the paged pool's slab size in physical block rows (issue #899). + /// + /// Stored intent: the pool is created lazily on the first paged + /// allocation, so the value is applied there. Setting it after the pool + /// exists is rejected by [`PagedBlockPool::set_slab_blocks`] once storage + /// has been allocated, because the slab shape is baked into the row + /// mapping; this returns that error rather than silently ignoring it. + /// + /// `None` leaves the pool's default in place. + pub fn set_paged_slab_blocks(&mut self, slab_blocks: Option) -> Result<(), String> { + self.paged_slab_blocks = slab_blocks; + match (self.paged_pool.as_ref(), slab_blocks) { + (Some(pool), Some(n)) => pool.borrow_mut().set_slab_blocks(n), + _ => Ok(()), + } + } + + /// The configured paged slab size, or `None` when the pool default applies. + pub fn paged_slab_blocks(&self) -> Option { + self.paged_slab_blocks + } + /// Blocks still **acquirable** before the paged budget is hit /// (`budget − live`), or `None` when unbounded / no paged pool. `Some(0)` /// means every budgeted block is in use; the admission gate must reclaim @@ -6654,6 +6689,11 @@ impl CachePool { // Apply the budget remembered before the pool existed (set via // `set_paged_block_budget` before the first paged allocation). pool.set_block_budget(self.paged_block_budget); + // Same for the slab size (#899). It must be applied before the first + // write, which is exactly here: the pool has no storage yet. + if let Some(slab_blocks) = self.paged_slab_blocks { + pool.set_slab_blocks(slab_blocks)?; + } self.paged_pool = Some(Rc::new(RefCell::new(pool))); Ok(()) } diff --git a/src/lib/mlxcel-core/src/cache/paged.rs b/src/lib/mlxcel-core/src/cache/paged.rs index 0e1a5b2bc..7f52e2d27 100644 --- a/src/lib/mlxcel-core/src/cache/paged.rs +++ b/src/lib/mlxcel-core/src/cache/paged.rs @@ -471,13 +471,42 @@ pub struct PagedBlockPool { /// "presize allocates once, incremental growth appends once" and /// observability can surface growth churn. grow_events: u64, + /// Monotonic counter bumped by every mutation that can move a block: which + /// blocks a sequence holds, and which physical row a block occupies. It is + /// the invalidation signal for [`Self::decode_v2_cache`] (issue #899); see + /// [`crate::paged_v2::plan_cache`] for the reuse argument it underwrites. + /// + /// Deliberately **not** bumped by a plain token append that fits in the + /// current block, because that is the decode hot path and it changes no + /// page list. + block_epoch: u64, + /// Cached CSR page tables (per layer) and chunk plan for the active decode + /// batch (issue #899). Purely derived state: dropping it is always safe and + /// only costs a rebuild. + decode_v2_cache: crate::paged_v2::PagedDecodeV2Cache, + /// Physical block rows per pool slab tensor (issue #899). + /// + /// Defaults to [`POOL_SLAB_BLOCKS`], which is byte-for-byte the pre-#899 + /// behaviour. The fused decode kernels read **one contiguous pool buffer + /// per side**, so they can only serve a layer whose rows all live in the + /// first slab; with the 32-row default that caps the fused path at 1024 + /// tokens across the whole batch at `block_size` 32, which is below every + /// context the production dispatch would choose it for. Raising this to the + /// working set the server was configured for is what makes the fused path + /// reachable in production. See [`Self::set_slab_blocks`]. + slab_blocks: usize, } -/// Number of physical block rows per pool slab tensor. Growth appends whole -/// slabs of this size (#235), so a freshly assigned row never triggers a +/// Default number of physical block rows per pool slab tensor. Growth appends +/// whole slabs of this size (#235), so a freshly assigned row never triggers a /// reallocation or copy of existing storage; this also keeps the allocation /// granularity of the previous single-tensor chunked growth (32 blocks). -const POOL_SLAB_BLOCKS: usize = 32; +/// +/// Since issue #899 this is a *default* rather than a fixed constant: the pool +/// carries its own [`PagedBlockPool::slab_blocks`], because the fused decode +/// kernels can only serve a layer whose rows live in one slab. See +/// [`PagedBlockPool::set_slab_blocks`]. +pub const POOL_SLAB_BLOCKS: usize = 32; /// Gathered visible K and V tensors for one layer, each shaped /// `[1, n_kv_heads, visible_len, head_dim]` (SDPA-ready). Returned by @@ -488,6 +517,36 @@ pub type GatheredKv = (UniquePtr, UniquePtr); /// paired with its `(n_slots, n_kv_heads, head_dim)` geometry. type NormalizedBlock = (UniquePtr, (i32, i32, i32)); +/// Run a resolved v2 plan and return the result in `q`'s dtype. +/// +/// The kernels read Q and emit their output in f32 deterministically (they +/// never re-specialise by dtype), so Q is cast in and the result cast back. +/// That is the same round trip [`PagedBlockPool::paged_decode_fused`] performs, +/// which is what keeps v1, v2, and gather byte-comparable. +fn launch_v2( + q: &MlxArray, + pool_k: &MlxArray, + pool_v: &MlxArray, + view: &super::paged_csr::PagedCsrView, + plan: &crate::paged_v2::PagedDecodePlan, + scale: f32, +) -> Result, String> { + let q_dtype = ffi::array_dtype(q); + let q_f32 = if q_dtype == crate::dtype::FLOAT32 { + None + } else { + Some(ffi::astype(q, crate::dtype::FLOAT32)) + }; + let q_in: &MlxArray = q_f32.as_deref().unwrap_or(q); + let ctx = crate::paged_v2::V2Context::build(q_in, pool_k, pool_v, view, plan.geometry, scale)?; + let out_f32 = ctx.launch(plan)?; + Ok(if q_dtype == crate::dtype::FLOAT32 { + out_f32 + } else { + ffi::astype(&out_f32, q_dtype) + }) +} + impl PagedBlockPool { pub fn new(layout: PagedKvLayout) -> Self { let num_layers = layout.num_layers; @@ -505,9 +564,54 @@ impl PagedBlockPool { next_row: vec![0; num_layers], block_budget: None, grow_events: 0, + block_epoch: 0, + decode_v2_cache: crate::paged_v2::PagedDecodeV2Cache::new(), + slab_blocks: POOL_SLAB_BLOCKS, } } + /// Physical block rows per pool slab tensor. + #[must_use] + pub fn slab_blocks(&self) -> usize { + self.slab_blocks + } + + /// Set the slab size, in physical block rows (issue #899). + /// + /// Must be called before any layer's pool storage is allocated (that is, + /// before the first write), because the slab shape is baked into every + /// tensor the layer owns and into the row-to-slab arithmetic + /// [`Self::slab_coords`] performs. Returns an error rather than silently + /// corrupting the row mapping if storage already exists. + /// + /// ## What it costs and what it buys + /// + /// Slabs are allocated lazily, one layer at a time, on that layer's first + /// write, and each allocation is `slab_blocks` whole rows. Raising this + /// therefore front-loads memory: a layer's first token allocates the whole + /// slab instead of growing into it. In exchange, every row of that layer + /// lives in one contiguous buffer, which is the precondition for the fused + /// decode kernels; a layer that outgrows its first slab appends a second + /// one exactly as before (no copy, #235) and simply stops being eligible + /// for the fused path, falling back to gather. + /// + /// The server sizes this from the configuration it already reserved KV + /// memory for (context length times batch size, clamped by the block + /// budget), so the front-loaded bytes are bytes the budget already + /// promised. Callers with no such figure should leave the default alone. + pub fn set_slab_blocks(&mut self, slab_blocks: usize) -> Result<(), String> { + if slab_blocks == 0 { + return Err("PagedBlockPool::set_slab_blocks: slab_blocks must be > 0".to_string()); + } + if self.pool_meta.iter().any(Option::is_some) { + return Err( + "PagedBlockPool::set_slab_blocks: pool storage is already allocated".to_string(), + ); + } + self.slab_blocks = slab_blocks; + Ok(()) + } + /// Number of pool growth episodes since construction (see /// [`Self::ensure_layer_capacity`]). A long prefill should presize into /// one episode instead of accumulating these. @@ -515,10 +619,49 @@ impl PagedBlockPool { self.grow_events } + /// Current block-topology epoch (issue #899). + /// + /// Bumped by every mutation that can change which blocks a sequence holds + /// or which physical row a block occupies. Stable across a decode step that + /// only appends tokens inside already-allocated blocks, which is what makes + /// the CSR page table cacheable. See [`crate::paged_v2::plan_cache`]. + #[must_use] + pub fn block_epoch(&self) -> u64 { + self.block_epoch + } + + /// Invalidate every cached CSR page table and chunk plan. + /// + /// Called internally on any block-topology change. Public so a caller that + /// mutates block tables through a path this pool cannot see (there is none + /// today) can still force a rebuild, and so tests can assert the + /// invalidation is what makes a stale plan impossible rather than luck. + pub fn invalidate_decode_plan_cache(&mut self) { + self.block_epoch = self.block_epoch.wrapping_add(1); + self.decode_v2_cache.clear(); + } + + /// Bump the topology epoch without dropping the cached views. + /// + /// The views carry the epoch they were built at, so a bump alone is enough + /// to make every one of them unusable; keeping the allocations lets the + /// next rebuild reuse their `Vec` capacity. + #[inline] + fn bump_block_epoch(&mut self) { + self.block_epoch = self.block_epoch.wrapping_add(1); + } + + /// Counters for the decode plan cache (issue #899), for tests and tracing. + #[must_use] + pub fn decode_plan_cache_stats(&self) -> crate::paged_v2::PlanCacheStats { + self.decode_v2_cache.stats() + } + /// Map a physical row index to its slab and the row within that slab. #[inline] - fn slab_coords(row: usize) -> (usize, i32) { - (row / POOL_SLAB_BLOCKS, (row % POOL_SLAB_BLOCKS) as i32) + fn slab_coords(&self, row: usize) -> (usize, i32) { + let per_slab = self.slab_blocks.max(1); + (row / per_slab, (row % per_slab) as i32) } /// Physical row currently assigned to `block_id` on `layer_idx`, if any. @@ -669,6 +812,9 @@ impl PagedBlockPool { self.release_block(block_id)?; } } + // A back-trim moves `len` (and may reset `logical_start`), which moves + // the visible page range even when no block was released (#899). + self.bump_block_epoch(); Ok(trimmed) } @@ -689,6 +835,9 @@ impl PagedBlockPool { layer.len = 0; layer.logical_start = 0; } + // A sequence leaving the batch (finish, eviction, preemption) is one of + // the invalidation events issue #899 names for the decode plan cache. + self.bump_block_epoch(); Ok(()) } @@ -810,6 +959,7 @@ impl PagedBlockPool { )); } record.refcount = record.refcount.saturating_add(1); + self.bump_block_epoch(); Ok(()) } @@ -847,6 +997,7 @@ impl PagedBlockPool { self.free_rows[layer_idx].push(row); } self.free_lists[layer_idx].push(block_id); + self.bump_block_epoch(); Ok(()) } @@ -1090,13 +1241,14 @@ impl PagedBlockPool { } } None => { - self.create_layer_pool(layer_idx, POOL_SLAB_BLOCKS, n_kv_heads, head_dim, k_dtype); + let initial = self.slab_blocks; + self.create_layer_pool(layer_idx, initial, n_kv_heads, head_dim, k_dtype); } } // Resolve (and grow for) the physical row for this block. let row = self.assign_row(block_id, layer_idx)?; - let (slab_i, row_in_slab) = Self::slab_coords(row); + let (slab_i, row_in_slab) = self.slab_coords(row); // Reassign the row's slab with the slot update so MLX donates the // buffer (in-place append, O(block)). The null placeholder makes the @@ -1129,9 +1281,10 @@ impl PagedBlockPool { dtype: i32, ) { debug_assert!(self.pool_meta[layer_idx].is_none()); - let slabs = min_capacity_blocks.div_ceil(POOL_SLAB_BLOCKS).max(1); + let per_slab = self.slab_blocks.max(1); + let slabs = min_capacity_blocks.div_ceil(per_slab).max(1); let block_size = self.layout.block_size as i32; - let shape = [POOL_SLAB_BLOCKS as i32, block_size, n_kv_heads, head_dim]; + let shape = [per_slab as i32, block_size, n_kv_heads, head_dim]; for _ in 0..slabs { self.pool_k[layer_idx].push(ffi::zeros(&shape, dtype)); self.pool_v[layer_idx].push(ffi::zeros(&shape, dtype)); @@ -1140,7 +1293,7 @@ impl PagedBlockPool { n_kv_heads, head_dim, dtype, - capacity_blocks: slabs * POOL_SLAB_BLOCKS, + capacity_blocks: slabs * per_slab, }); } @@ -1365,7 +1518,7 @@ impl PagedBlockPool { "PagedBlockPool::copy_on_write_block: shared block {src_block_id} on layer {layer_idx} has no pool row (was it written?)" ) })?; - let (slab_i, row_in_slab) = Self::slab_coords(src_row); + let (slab_i, row_in_slab) = self.slab_coords(src_row); let block_size = self.layout.block_size as i32; // Slice the source row's [block_size, H, D] slab out of K and V (the @@ -1415,7 +1568,7 @@ impl PagedBlockPool { "PagedBlockPool::read_block_contents: block {block_id} on layer {layer_idx} has no pool row (never written)" ) })?; - let (slab_i, row_in_slab) = Self::slab_coords(src_row); + let (slab_i, row_in_slab) = self.slab_coords(src_row); let block_size = self.layout.block_size as i32; // Slice the source row's [block_size, H, D] slab out of K and V (the bare @@ -1519,15 +1672,16 @@ impl PagedBlockPool { // monotonically, so a sequence spans about one run per slab; the // concat output copy is the same magnitude as the gather copy this // path always materialized. + let per_slab = self.slab_blocks.max(1); let gather_rows = |slabs: &[UniquePtr]| -> UniquePtr { let mut parts: Vec> = Vec::new(); let mut i = 0usize; while i < rows.len() { - let slab_i = rows[i] / POOL_SLAB_BLOCKS; + let slab_i = rows[i] / per_slab; let mut local: Vec = Vec::new(); let mut j = i; - while j < rows.len() && rows[j] / POOL_SLAB_BLOCKS == slab_i { - local.push((rows[j] % POOL_SLAB_BLOCKS) as i32); + while j < rows.len() && rows[j] / per_slab == slab_i { + local.push((rows[j] % per_slab) as i32); j += 1; } let idx = ffi::from_slice_i32(&local, &[local.len() as i32]); @@ -1680,6 +1834,184 @@ impl PagedBlockPool { Ok(Some(out)) } + /// Production paged decode for a whole batch (issue #899). + /// + /// This is the entry point the server's pool-backed batched decode takes, + /// replacing the per-sequence `gather_visible` + dense SDPA loop that ADR + /// 0001 measured as the paged decode hot path. It differs from + /// [`Self::paged_decode_fused_v2`] in three ways, all of which are what + /// makes it usable on a hot path rather than in a harness: + /// + /// 1. **It caches.** The CSR page table is kept per layer and the chunk + /// plan per batch, keyed by [`Self::block_epoch`] plus a per-request + /// fingerprint (see [`crate::paged_v2::plan_cache`]). A decode step that + /// does not cross a page boundary refreshes O(batch) scalars instead of + /// resolving O(pages) block ids. + /// 2. **It applies the dispatch policy.** [`crate::paged_v2::dispatch`] + /// keeps small launches on the gather path, because #898 measured v2 + /// losing to gather for a single request at 1024 tokens of context. + /// 3. **It says what it did.** The returned + /// [`PagedDecodeOutcome`](crate::paged_v2::PagedDecodeOutcome) carries + /// the decision and the numbers behind it, so the caller can surface it + /// at whatever level makes it visible. It is a value rather than a log + /// line because the first cut of this used `tracing::debug!` and a whole + /// production benchmark sweep then compared gather against gather + /// without anything saying so. + /// + /// `q` is `[B, Hq, 1, head_dim]`, `states[b]` is sequence `b`'s per-layer + /// state, and `Ok((Some(out), _))` is `[B, Hq, 1, head_dim]` in `q`'s dtype. + /// `Ok((None, outcome))` is the caller's signal to run the gather path, and + /// `outcome` says why. + pub fn paged_decode_batched( + &mut self, + q: &MlxArray, + states: &[&PagedSequenceState], + layer_idx: usize, + scale: f32, + ) -> Result< + ( + Option>, + crate::paged_v2::PagedDecodeOutcome, + ), + String, + > { + use crate::paged_v2::{ + PagedDecodeOutcome, PagedV2Dispatch, active_required_visible_tokens, + select_paged_v2_dispatch, + }; + + // Same single-slab restriction as v1 and as `paged_decode_fused_v2`: + // both kernels read one contiguous pool buffer per side, so a layer + // grown past one slab (#235) is declined rather than stitched. This is + // the decline that silently disables the whole fused path when the slab + // is sized below the served working set, so it names the knob. + if self.single_slab_tensors(layer_idx).is_none() { + return Ok(( + None, + PagedDecodeOutcome::MultiSlab { + k_slabs: self.pool_k.get(layer_idx).map_or(0, Vec::len), + v_slabs: self.pool_v.get(layer_idx).map_or(0, Vec::len), + slab_blocks: self.slab_blocks(), + }, + )); + } + + let mut layers: Vec<&PagedLayerState> = Vec::with_capacity(states.len()); + for state in states { + layers.push(state.layer(layer_idx).ok_or_else(|| { + format!( + "PagedBlockPool::paged_decode_batched: layer {layer_idx} out of range for {} layers", + state.layers.len() + ) + })?); + } + + let (batch, geometry) = { + let (pool_k, _) = self.single_slab_tensors(layer_idx).expect("checked above"); + crate::paged_v2::geometry_from_shapes(q, pool_k)? + }; + if batch != states.len() { + return Err(format!( + "PagedBlockPool::paged_decode_batched: q batch {batch} disagrees with {} states", + states.len() + )); + } + if let Err(reason) = geometry.check() { + return Ok((None, PagedDecodeOutcome::UnservableGeometry(reason))); + } + + let epoch = self.block_epoch; + let page_size = self.layout.block_size; + let target = crate::paged_v2::device_target_ctas(); + + // The cache is a disjoint field, but the view builder needs `&self`. + // Move it out for the duration so both borrows are independent, and put + // it back unconditionally (including on the error path) so a failed + // build never loses the other layers' entries. + let mut cache = std::mem::take(&mut self.decode_v2_cache); + let key = crate::paged_v2::DecodeBatchKey { + layer_idx, + epoch, + page_size, + layers: &layers, + geometry: &geometry, + }; + // Resolve the page table first and decide from it, then build the plan + // only if the fused path is actually going to run. A declined launch + // must not pay for the chunk search, and keeping the two steps separate + // is what makes the plan-rebuild counter mean "the fused path ran". + let view_built = cache.ensure_view(&key, || self.paged_csr_view(states, layer_idx)); + + let decision = match &view_built { + Ok(()) => { + let view = cache + .view(layer_idx) + .expect("ensure_view populated the entry"); + let visible_tokens: usize = + view.seq_lens.iter().map(|&len| len.max(0) as usize).sum(); + if !view.any_visible() { + Err(PagedDecodeOutcome::NoVisibleTokens) + } else { + let floor = active_required_visible_tokens(batch); + let policy = select_paged_v2_dispatch(visible_tokens, floor); + match crate::layers::resolve_paged_v2_dispatch(|| policy) { + PagedV2Dispatch::V2 => Ok((visible_tokens, view.page_counts())), + // A gather decision the policy also made is the floor; + // a gather decision it did not make is the kill switch, + // which is the only other thing + // `resolve_paged_v2_dispatch` can do. + PagedV2Dispatch::Gather if policy == PagedV2Dispatch::Gather => { + Err(PagedDecodeOutcome::BelowFloor { + batch, + visible_tokens, + floor, + }) + } + PagedV2Dispatch::Gather => Err(PagedDecodeOutcome::KillSwitch), + } + } + } + Err(reason) => Err(PagedDecodeOutcome::ViewFailed(reason.clone())), + }; + + let launched: Result>, String>; + let outcome: PagedDecodeOutcome; + match decision { + Ok((visible_tokens, page_counts)) => { + cache.ensure_plan(&geometry, &page_counts, || { + crate::paged_v2::PagedDecodePlan::heuristic(geometry, &page_counts, target) + }); + // Both accessors are shared borrows of the same cache, so they + // coexist; the two mutations above have already finished. + let view = cache.view(layer_idx).expect("populated above"); + let plan = cache.plan().expect("populated above"); + match plan.validate() { + Ok(()) => { + let (pool_k, pool_v) = + self.single_slab_tensors(layer_idx).expect("checked above"); + outcome = PagedDecodeOutcome::Fused { + batch, + visible_tokens, + chunks: plan.num_chunks, + merged: plan.needs_merge, + }; + launched = launch_v2(q, pool_k, pool_v, view, plan, scale).map(Some); + } + Err(reason) => { + outcome = PagedDecodeOutcome::PlanRejected(reason); + launched = Ok(None); + } + } + } + Err(declined) => { + outcome = declined; + launched = Ok(None); + } + } + self.decode_v2_cache = cache; + launched.map(|out| (out, outcome)) + } + /// Fused paged-attention decode over the pool via the native Metal kernel /// (epic #116 Phase 6, #123). /// @@ -1895,6 +2227,9 @@ impl PagedBlockPool { self.ensure_layer_capacity(layer_idx, row + 1)?; } self.block_rows[layer_idx].insert(block_id, row); + // A block that previously had no row now has one, so any cached CSR + // page table built before this point may name a different row (#899). + self.bump_block_epoch(); Ok(row) } @@ -1916,20 +2251,16 @@ impl PagedBlockPool { if min_capacity <= meta.capacity_blocks { return Ok(()); } - let target_slabs = min_capacity.div_ceil(POOL_SLAB_BLOCKS); + let per_slab = self.slab_blocks.max(1); + let target_slabs = min_capacity.div_ceil(per_slab); let block_size = self.layout.block_size as i32; - let shape = [ - POOL_SLAB_BLOCKS as i32, - block_size, - meta.n_kv_heads, - meta.head_dim, - ]; + let shape = [per_slab as i32, block_size, meta.n_kv_heads, meta.head_dim]; while self.pool_k[layer_idx].len() < target_slabs { self.pool_k[layer_idx].push(ffi::zeros(&shape, meta.dtype)); self.pool_v[layer_idx].push(ffi::zeros(&shape, meta.dtype)); } self.pool_meta[layer_idx] = Some(PagedPoolMeta { - capacity_blocks: target_slabs * POOL_SLAB_BLOCKS, + capacity_blocks: target_slabs * per_slab, ..meta }); self.grow_events += 1; @@ -1958,6 +2289,7 @@ impl PagedBlockPool { "free-list block must have zero refcount" ); record.refcount = 1; + self.bump_block_epoch(); return Ok(block_id); } @@ -1976,6 +2308,7 @@ impl PagedBlockPool { self.next_block_id += 1; self.blocks .insert(block_id, PagedBlockRecord::new(layer_idx)); + self.bump_block_epoch(); Ok(block_id) } @@ -2007,6 +2340,7 @@ impl PagedBlockPool { } self.next_block_id = self.next_block_id.max(block_id.as_u64().saturating_add(1)); + self.bump_block_epoch(); Ok(()) } diff --git a/src/lib/mlxcel-core/src/cache/paged_batch_decode.rs b/src/lib/mlxcel-core/src/cache/paged_batch_decode.rs new file mode 100644 index 000000000..885d70dd1 --- /dev/null +++ b/src/lib/mlxcel-core/src/cache/paged_batch_decode.rs @@ -0,0 +1,367 @@ +// Copyright 2025-2026 Lablup Inc. and Jeongkyu Shin +// +// 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. + +//! Whole-batch decode over pool-backed KV caches (issue #899). +//! +//! ## What this replaces +//! +//! Before #899, a model's batched decode handed pool-backed caches to the +//! per-sequence loop: for each sequence, `KVCache::update_and_fetch` appended +//! the new token to the pool and called `gather_visible` to materialize a dense +//! `[1, Hkv, visible_len, D]` copy, then the model ran an ordinary fused SDPA +//! against it. That is ADR 0001 strategy A, and the ADR measures its overhead +//! at 2x-3x of SDPA time past 4096 tokens. +//! +//! [`paged_batch_decode_attention`] takes the whole batch instead: it performs +//! the same pool appends, then issues **one** fused v2 launch across every +//! sequence and every KV head, with no gather copy anywhere. +//! +//! ## Contract +//! +//! The function either owns the step or has not touched anything: +//! +//! - `None` means it declined **before writing**, so the caller runs its +//! unchanged per-sequence loop. Every decline check is evaluated first, for +//! this reason. +//! - `Some(out)` means the pool appends happened here and `out` is +//! `[B, Hq, 1, D]` attention output. If v2 was not the right choice for the +//! shape, or declined it, the gather-then-SDPA fallback ran *inside* this +//! function; the caller must not run its own loop, which would double-write +//! the pool. +//! +//! ## Variant coverage +//! +//! | family shape | handling | +//! |---|---| +//! | full attention, f16 pool | v2 above the token floor, gather below | +//! | sliding window (`logical_start > 0`) | v2; the window is expressed as a CSR range, see below | +//! | logit softcap | declined, gather (`softcap != 0.0`) | +//! | explicit attention mask | declined (`mask` is not `None` at the call site) | +//! | speculative / MTP verify | never reached: those steps have `seq_len > 1` | +//! | Int8 / Turbo KV modes | declined: those sequences are never pool-backed | +//! +//! ### Sliding window +//! +//! v2 applies **no windowing mask of its own**. A trimmed window is expressed +//! entirely in the page table: [`build_paged_csr_view`] emits only the pages +//! from `logical_start / page_size` onward and sets `first_page_offset` to +//! `logical_start % page_size`, so the kernel's token `i` of request `r` +//! resolves to absolute position `logical_start + i` and the retired prefix is +//! never addressed. That is the same `[logical_start, len)` window +//! `gather_visible` slices, which is why the two paths agree. +//! +//! Two consequences worth stating explicitly, because both are silent if +//! violated: +//! +//! 1. **The window must be contiguous.** A future attention-sink retention on +//! the paged path (`[0, keep) ++ [start, len)`, the shape +//! `KVCache::trim_front_keep_sink` produces on dense caches) is *not* +//! expressible in one CSR range and must not be routed here. Today it cannot +//! be: `trim_front_keep_sink` returns `0` for pool-backed caches, so +//! `logical_start` on a served sequence only ever moves through +//! `PagedBlockPool::trim_tokens`, which keeps the window contiguous. +//! 2. **RoPE is applied upstream.** The model rotates Q and K from the +//! scheduler's `BatchedAttentionMetadata::rope_offsets` before this function +//! is called, so [`PagedCsrView::rope_offsets`] is not consumed here. It +//! records the same quantity (the absolute next-token position, which is the +//! written length and not the visible length) and exists for callers that +//! have to derive it themselves. +//! +//! [`build_paged_csr_view`]: super::paged_csr::build_paged_csr_view +//! [`PagedCsrView::rope_offsets`]: super::paged_csr::PagedCsrView::rope_offsets + +use std::cell::Ref; +use std::rc::Rc; +use std::sync::atomic::{AtomicBool, AtomicU64, Ordering}; + +use cxx::UniquePtr; + +use super::{KVCache, KVCacheMode, PagedSequenceState}; +use crate::ffi; +use crate::ffi::MlxArray; +use crate::paged_v2::{PAGED_DECODE_OUTCOME_KINDS, PagedDecodeOutcome}; + +/// Counters for what the production paged decode actually did. +/// +/// Process-wide and relaxed: they exist so tests can assert that v2 really took +/// the step (rather than silently falling back) and so an operator can see the +/// split without a profiler. They are never read on the hot path. +#[derive(Debug, Default)] +pub struct PagedBatchDecodeCounters { + /// Layer steps served by a fused v2 launch. + pub v2_launches: AtomicU64, + /// Layer steps served by the in-function gather-then-SDPA fallback. + pub gather_fallbacks: AtomicU64, + /// Calls that declined before writing, leaving the caller's own loop to run. + pub declines: AtomicU64, +} + +static COUNTERS: PagedBatchDecodeCounters = PagedBatchDecodeCounters { + v2_launches: AtomicU64::new(0), + gather_fallbacks: AtomicU64::new(0), + declines: AtomicU64::new(0), +}; + +/// One flag per [`PagedDecodeOutcome`] kind, so the *first* launch of each +/// distinct outcome is announced and the rest are silent. +/// +/// Per kind rather than one global flag on purpose. A single one-shot reports +/// only whatever happened first, which in a real server is a short warmup +/// request; a later, permanent decline for an entirely different reason then +/// never surfaces. That is exactly how the #899 production benchmark ran a full +/// sweep on the gather path without any line saying so. +static REPORTED: [AtomicBool; PAGED_DECODE_OUTCOME_KINDS] = + [const { AtomicBool::new(false) }; PAGED_DECODE_OUTCOME_KINDS]; + +/// Announce an outcome the first time its kind occurs. +/// +/// **Info, not debug.** A decode-path diagnostic that only an operator who +/// already suspects a problem can enable is worthless for the case it exists +/// for: knowing which kernel a benchmark actually measured. There are at most +/// [`PAGED_DECODE_OUTCOME_KINDS`] of these lines in a process lifetime, so the +/// cost of always emitting them is bounded and tiny. +fn report_once(outcome: &PagedDecodeOutcome) { + let slot = &REPORTED[outcome.kind_index()]; + if slot + .compare_exchange(false, true, Ordering::Relaxed, Ordering::Relaxed) + .is_ok() + { + tracing::info!("paged decode v2: {}", outcome.describe()); + } +} + +/// Forget which outcomes have been announced. Test-only: the flags are +/// process-wide, and a test that asserts on logging needs a clean slate. +#[cfg(test)] +fn reset_reported() { + for slot in &REPORTED { + slot.store(false, Ordering::Relaxed); + } +} + +/// Snapshot of [`COUNTERS`] as plain integers. +#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)] +pub struct PagedBatchDecodeStats { + pub v2_launches: u64, + pub gather_fallbacks: u64, + pub declines: u64, +} + +/// Read the production paged decode counters. +#[must_use] +pub fn paged_batch_decode_stats() -> PagedBatchDecodeStats { + PagedBatchDecodeStats { + v2_launches: COUNTERS.v2_launches.load(Ordering::Relaxed), + gather_fallbacks: COUNTERS.gather_fallbacks.load(Ordering::Relaxed), + declines: COUNTERS.declines.load(Ordering::Relaxed), + } +} + +/// Whether every cache in the batch can be served by the whole-batch path. +/// +/// Pure and MLX-free apart from the shape reads, and evaluated in full before +/// any pool mutation: a `false` here means the caller's per-sequence loop runs +/// with the pool untouched. +fn batch_is_servable(caches: &[&mut KVCache], softcap: f32) -> Result<(), &'static str> { + if caches.is_empty() { + return Err("empty batch"); + } + // A logit soft-cap is applied inside SDPA; the v2 kernel has no soft-cap + // term and the gather fallback in this function does not thread one either, + // so the whole batch goes back to the caller (issue #899 scopes soft-cap + // families to a follow-up). + if softcap != 0.0 { + return Err("logit soft-cap requested"); + } + let Some(first) = caches[0].paged_backing.as_ref() else { + return Err("caches are not pool-backed"); + }; + for cache in caches { + let Some(backing) = cache.paged_backing.as_ref() else { + return Err("caches are not pool-backed"); + }; + if cache.mode != KVCacheMode::Fp16 { + return Err("cache mode is not Fp16"); + } + if backing.layer_idx != first.layer_idx { + return Err("caches name different layers"); + } + // One launch reads one pool; sequences backed by different pools cannot + // share a page table. + if !Rc::ptr_eq(&backing.pool, &first.pool) { + return Err("caches are backed by different pools"); + } + } + Ok(()) +} + +/// `[B, H, 1, D]` shape check for the decode tensors. +fn is_single_token_batch(shape: &[i32], batch: usize) -> bool { + shape.len() == 4 && shape[0] as usize == batch && shape[2] == 1 +} + +/// Whole-batch pooled paged decode: append this step's K/V for every sequence, +/// then run attention once over the batch. +/// +/// See the module docs for the decline contract. `q_batched` is `[B, Hq, 1, D]`, +/// `k_batched` / `v_batched` are `[B, Hkv, 1, D]`, and `caches[b]` is sequence +/// `b`'s cache **for one layer** (all entries must name the same `layer_idx` +/// of the same pool). +/// +/// Used by: Llama3 / Qwen2 / Qwen2.5 / Helium `forward_split_attention`, Qwen3 +/// `forward_split_attention`, and every VLM whose text backbone is one of +/// those. +pub fn paged_batch_decode_attention( + q_batched: &MlxArray, + k_batched: &MlxArray, + v_batched: &MlxArray, + caches: &mut [&mut KVCache], + scale: f32, + softcap: f32, +) -> Option> { + let batch = caches.len(); + let servable = batch_is_servable(caches, softcap).and_then(|()| { + if !is_single_token_batch(&ffi::array_shape(q_batched), batch) + || !is_single_token_batch(&ffi::array_shape(k_batched), batch) + || !is_single_token_batch(&ffi::array_shape(v_batched), batch) + { + // Batched prefill and speculative / MTP verify land here: more than + // one query token per sequence, which the kernel does not serve. + return Err("not a single-token decode step"); + } + Ok(()) + }); + if let Err(reason) = servable { + COUNTERS.declines.fetch_add(1, Ordering::Relaxed); + report_once(&PagedDecodeOutcome::NotServable(reason)); + return None; + } + + let backing = caches[0] + .paged_backing + .clone() + .expect("batch_is_servable established a paged backing"); + let layer_idx = backing.layer_idx; + + // Past this point the pool is mutated, so every exit must return `Some`. + for (b, cache) in caches.iter_mut().enumerate() { + let k_i = slice_row(k_batched, b); + let v_i = slice_row(v_batched, b); + cache.write_paged(&k_i, &v_i); + } + + let states: Vec> = caches + .iter() + .map(|cache| { + cache + .paged_backing + .as_ref() + .expect("checked above") + .state + .borrow() + }) + .collect(); + let state_refs: Vec<&PagedSequenceState> = states.iter().map(|state| &**state).collect(); + + let mut pool = backing.pool.borrow_mut(); + match pool.paged_decode_batched(q_batched, &state_refs, layer_idx, scale) { + Ok((launched, outcome)) => { + // Which kernel a server actually ran has to be visible without a + // profiler and without an operator opting into debug logging; a + // before/after benchmark is uninterpretable otherwise, and the + // first #899 sweep measured gather against gather because this was + // not true. + report_once(&outcome); + match launched { + Some(out) => { + COUNTERS.v2_launches.fetch_add(1, Ordering::Relaxed); + Some(out) + } + None => { + COUNTERS.gather_fallbacks.fetch_add(1, Ordering::Relaxed); + Some(gather_fallback( + q_batched, + &pool, + &state_refs, + layer_idx, + scale, + )) + } + } + } + Err(reason) => { + // A hard error here means the batch's own bookkeeping is + // inconsistent (a length outrunning its blocks, a block with no + // pool row). The gather path applies the same guards and would fail + // the same way, so surface it rather than papering over it: the + // pre-#899 path panicked on exactly these conditions too. + panic!("PagedBlockPool::paged_decode_batched failed on layer {layer_idx}: {reason}"); + } + } +} + +/// Per-sequence `gather_visible` + fused SDPA, concatenated on the batch axis. +/// +/// Byte-identical to what the caller's per-sequence loop produced before #899: +/// the same gather, the same `attention_from_ptr` call, and a batch-axis concat +/// that commutes with the transpose/reshape the caller applies afterward +/// (neither touches axis 0). +fn gather_fallback( + q_batched: &MlxArray, + pool: &super::PagedBlockPool, + states: &[&PagedSequenceState], + layer_idx: usize, + scale: f32, +) -> UniquePtr { + let mut outputs: Vec> = Vec::with_capacity(states.len()); + for (b, state) in states.iter().enumerate() { + let q_i = slice_row(q_batched, b); + let (key_visible, value_visible) = pool + .gather_visible(state, layer_idx) + .expect("PagedBlockPool::gather_visible failed for pool-backed cache") + .expect("gather_visible returned None for a pool-backed cache"); + outputs.push(unsafe { + crate::layers::attention_from_ptr( + &q_i, + &key_visible, + &value_visible, + scale, + std::ptr::null(), + 0.0, + 0, + ) + }); + } + let mut outputs = outputs.into_iter(); + let mut result = outputs + .next() + .expect("a servable batch has at least one sequence"); + for output in outputs { + result = crate::concatenate(&result, &output, 0); + } + result +} + +/// Slice row `b` out of a `[B, H, T, D]` tensor, keeping the batch axis. +fn slice_row(arr: &MlxArray, b: usize) -> UniquePtr { + ffi::slice( + arr, + &[b as i32, 0, 0, 0], + &[b as i32 + 1, i32::MAX, i32::MAX, i32::MAX], + ) +} + +#[cfg(test)] +#[path = "paged_batch_decode_tests.rs"] +mod paged_batch_decode_tests; diff --git a/src/lib/mlxcel-core/src/cache/paged_batch_decode_tests.rs b/src/lib/mlxcel-core/src/cache/paged_batch_decode_tests.rs new file mode 100644 index 000000000..aa8456de0 --- /dev/null +++ b/src/lib/mlxcel-core/src/cache/paged_batch_decode_tests.rs @@ -0,0 +1,536 @@ +// Copyright 2025-2026 Lablup Inc. and Jeongkyu Shin +// +// 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. + +//! Tests for the production whole-batch paged decode (issue #899). +//! +//! Two halves. The decline checks are pure and run anywhere: they pin the +//! contract that the function touches nothing before it has decided it can +//! serve the batch. The parity checks drive a real [`PagedBlockPool`] through +//! real [`KVCache::new_paged`] caches, run the batched entry point, and compare +//! against the gather-then-SDPA path computed from the *same* post-write pool +//! state, which is the exact baseline #899 has to preserve. + +use std::cell::RefCell; + +use super::*; +use crate::cache::{PagedBlockPool, PagedKvLayout, PagedSequenceState}; +use crate::dtype; + +const PAGE: usize = 32; + +// ── decline contract ──────────────────────────────────────────────────────── + +#[test] +fn a_dense_batch_is_declined() { + let mut a = KVCache::new(); + let mut b = KVCache::new(); + let caches: Vec<&mut KVCache> = vec![&mut a, &mut b]; + assert!(batch_is_servable(&caches, 0.0).is_err()); +} + +#[test] +fn an_empty_batch_is_declined() { + let caches: Vec<&mut KVCache> = Vec::new(); + assert!(batch_is_servable(&caches, 0.0).is_err()); +} + +#[test] +fn a_softcap_batch_is_declined() { + let (pool, states) = fresh_pool(2, 2, 64); + let mut a = KVCache::new_paged(pool.clone(), states[0].clone(), 0); + let mut b = KVCache::new_paged(pool, states[1].clone(), 0); + let caches: Vec<&mut KVCache> = vec![&mut a, &mut b]; + assert!(batch_is_servable(&caches, 0.0).is_ok()); + assert!(batch_is_servable(&caches, 30.0).is_err()); +} + +#[test] +fn a_mixed_dense_and_paged_batch_is_declined() { + let (pool, states) = fresh_pool(1, 2, 64); + let mut paged = KVCache::new_paged(pool, states[0].clone(), 0); + let mut dense = KVCache::new(); + let caches: Vec<&mut KVCache> = vec![&mut paged, &mut dense]; + assert!(batch_is_servable(&caches, 0.0).is_err()); +} + +#[test] +fn caches_from_two_pools_are_declined() { + let (pool_a, states_a) = fresh_pool(1, 2, 64); + let (pool_b, states_b) = fresh_pool(1, 2, 64); + let mut a = KVCache::new_paged(pool_a, states_a[0].clone(), 0); + let mut b = KVCache::new_paged(pool_b, states_b[0].clone(), 0); + let caches: Vec<&mut KVCache> = vec![&mut a, &mut b]; + assert!(batch_is_servable(&caches, 0.0).is_err()); +} + +#[test] +fn caches_from_two_layers_are_declined() { + let (pool, states) = fresh_pool(2, 2, 64); + let mut a = KVCache::new_paged(pool.clone(), states[0].clone(), 0); + let mut b = KVCache::new_paged(pool, states[1].clone(), 1); + let caches: Vec<&mut KVCache> = vec![&mut a, &mut b]; + assert!(batch_is_servable(&caches, 0.0).is_err()); +} + +#[test] +fn shape_checks_require_a_single_query_token() { + assert!(is_single_token_batch(&[4, 8, 1, 64], 4)); + assert!(!is_single_token_batch(&[4, 8, 2, 64], 4)); + assert!(!is_single_token_batch(&[2, 8, 1, 64], 4)); + assert!(!is_single_token_batch(&[8, 1, 64], 4)); +} + +#[test] +fn a_multi_token_step_declines_without_touching_the_pool() { + // Batched prefill and speculative / MTP verify both arrive with more than + // one query token; the decline must happen before any pool write. + let (pool, states) = fresh_pool(1, 2, 64); + let mut a = KVCache::new_paged(pool.clone(), states[0].clone(), 0); + let mut b = KVCache::new_paged(pool.clone(), states[1].clone(), 0); + let mut caches: Vec<&mut KVCache> = vec![&mut a, &mut b]; + + let q = ffi::zeros(&[2, 8, 3, 64], dtype::FLOAT16); + let k = ffi::zeros(&[2, 2, 3, 64], dtype::FLOAT16); + let v = ffi::zeros(&[2, 2, 3, 64], dtype::FLOAT16); + assert!(paged_batch_decode_attention(&q, &k, &v, &mut caches, 0.125, 0.0).is_none()); + + assert_eq!(states[0].borrow().layer(0).unwrap().len, 0); + assert_eq!(states[1].borrow().layer(0).unwrap().len, 0); + assert_eq!(pool.borrow().allocated_block_count(), 0); +} + +// ── parity against the gather path ────────────────────────────────────────── + +/// A pool sized for the fused path plus one `PagedSequenceState` per sequence. +/// +/// `set_slab_blocks` is what makes the fused kernels reachable at all: they +/// read one contiguous pool buffer per side, so every row the batch touches has +/// to live in the first slab. +fn fresh_pool( + num_layers: usize, + kv_heads: usize, + head_dim: usize, +) -> ( + Rc>, + Vec>>, +) { + let bytes_per_block = PAGE * kv_heads * head_dim * 2; + let layout = PagedKvLayout::uniform(num_layers, PAGE, bytes_per_block).expect("valid layout"); + let mut pool = PagedBlockPool::new(layout.clone()); + pool.set_slab_blocks(512) + .expect("fresh pool has no storage"); + let states = (0..4) + .map(|_| Rc::new(RefCell::new(PagedSequenceState::new(&layout)))) + .collect(); + (Rc::new(RefCell::new(pool)), states) +} + +/// xorshift64* in [-1, 1), deterministic so a failure reproduces exactly. +struct Rng(u64); + +impl Rng { + fn new(seed: u64) -> Self { + Self(seed | 1) + } + fn next_f32(&mut self) -> f32 { + self.0 ^= self.0 << 13; + self.0 ^= self.0 >> 7; + self.0 ^= self.0 << 17; + ((self.0 >> 40) as f32 / (1u32 << 24) as f32) * 2.0 - 1.0 + } +} + +fn random_array(rng: &mut Rng, shape: &[i32]) -> UniquePtr { + let n: usize = shape.iter().map(|d| *d as usize).product(); + let data: Vec = (0..n).map(|_| rng.next_f32()).collect(); + let f32_arr = ffi::from_slice_f32(&data, shape); + ffi::astype(&f32_arr, dtype::FLOAT16) +} + +fn to_vec_f32(arr: &MlxArray) -> Vec { + let as_f32 = ffi::astype(arr, dtype::FLOAT32); + ffi::eval(&as_f32); + ffi::array_to_raw_bytes(&as_f32) + .chunks_exact(4) + .map(|c| f32::from_ne_bytes(c.try_into().expect("4-byte chunk"))) + .collect() +} + +/// Relative RMS between two flattened outputs. +fn relative_rms(a: &[f32], b: &[f32]) -> f32 { + assert_eq!(a.len(), b.len(), "outputs must have the same length"); + let mut num = 0.0f64; + let mut den = 0.0f64; + for (x, y) in a.iter().zip(b.iter()) { + num += f64::from(x - y) * f64::from(x - y); + den += f64::from(*y) * f64::from(*y); + } + if den == 0.0 { + return num.sqrt() as f32; + } + (num / den).sqrt() as f32 +} + +/// Drive one whole-batch decode step and return `(batched output, gather +/// reference computed from the same post-write pool state, cache stats)`. +fn run_step( + prompt_lens: &[usize], + q_heads: i32, + kv_heads: i32, + head_dim: i32, + seed: u64, +) -> (Vec, Vec, crate::paged_v2::PlanCacheStats) { + let batch = prompt_lens.len(); + let (pool, states) = fresh_pool(1, kv_heads as usize, head_dim as usize); + let mut rng = Rng::new(seed); + let mut caches: Vec = (0..batch) + .map(|b| KVCache::new_paged(pool.clone(), states[b].clone(), 0)) + .collect(); + + // Prefill each sequence into the pool. + for (b, &len) in prompt_lens.iter().enumerate() { + let k = random_array(&mut rng, &[1, kv_heads, len as i32, head_dim]); + let v = random_array(&mut rng, &[1, kv_heads, len as i32, head_dim]); + caches[b].update(k, v); + } + + let q = random_array(&mut rng, &[batch as i32, q_heads, 1, head_dim]); + let k_step = random_array(&mut rng, &[batch as i32, kv_heads, 1, head_dim]); + let v_step = random_array(&mut rng, &[batch as i32, kv_heads, 1, head_dim]); + let scale = 1.0 / (head_dim as f32).sqrt(); + + let out = { + let mut refs: Vec<&mut KVCache> = caches.iter_mut().collect(); + paged_batch_decode_attention(&q, &k_step, &v_step, &mut refs, scale, 0.0) + .expect("the batch is servable") + }; + + // Reference: the pre-#899 path, run against the pool state the batched + // entry point just left behind, so the two see identical K/V. + let reference = { + let pool_ref = pool.borrow(); + let borrowed: Vec> = + states.iter().take(batch).map(|s| s.borrow()).collect(); + let state_refs: Vec<&PagedSequenceState> = borrowed.iter().map(|s| &**s).collect(); + gather_fallback(&q, &pool_ref, &state_refs, 0, scale) + }; + + let stats = pool.borrow().decode_plan_cache_stats(); + (to_vec_f32(&out), to_vec_f32(&reference), stats) +} + +#[test] +fn batched_decode_matches_the_gather_path_above_the_floor() { + // 4 x 1024 = 4096 visible tokens, exactly the dispatch floor, so v2 runs. + let (out, reference, stats) = run_step(&[1024, 1024, 1024, 1024], 8, 2, 64, 0xA11CE); + assert!( + stats.view_rebuilds >= 1, + "the CSR view should have been built" + ); + assert!(stats.plan_rebuilds >= 1, "the plan should have been built"); + let rms = relative_rms(&out, &reference); + assert!(rms < 5e-3, "relative RMS {rms} exceeds 5e-3"); +} + +#[test] +fn batched_decode_matches_the_gather_path_with_ragged_lengths() { + // Mixed lengths, none a multiple of the page size, still above the floor. + let (out, reference, _) = run_step(&[1500, 900, 1201, 777], 8, 2, 64, 0xBEEF); + let rms = relative_rms(&out, &reference); + assert!(rms < 5e-3, "relative RMS {rms} exceeds 5e-3"); +} + +#[test] +fn batched_decode_matches_the_gather_path_at_gqa_one() { + let (out, reference, _) = run_step(&[2048, 2048], 4, 4, 64, 0xC0FFEE); + let rms = relative_rms(&out, &reference); + assert!(rms < 5e-3, "relative RMS {rms} exceeds 5e-3"); +} + +#[test] +fn below_the_floor_the_batch_still_answers_correctly() { + // 2 x 256 = 512 visible tokens against a 2 x 512 = 1024 batched floor, so + // the entry point takes its own gather fallback. It must still produce the + // reference answer exactly, not merely within tolerance, and must still + // have written the step's K/V once. + let (out, reference, stats) = run_step(&[256, 256], 8, 2, 64, 0xD00D); + assert_eq!( + stats.plan_rebuilds, 0, + "a below-floor launch must not build a plan" + ); + let rms = relative_rms(&out, &reference); + assert!( + rms < 1e-6, + "the gather fallback must reproduce the reference" + ); +} + +#[test] +fn a_single_sequence_above_its_floor_takes_the_fused_path() { + // The single-sequence decode path (`decode_single_step` -> the model's + // one-sequence `forward`) reaches this entry point with a batch of one. Two + // of the five scenarios in the issue's benchmark matrix are that shape, so + // a batch-1 launch above the single-request floor has to fuse. + let (out, reference, stats) = run_step(&[4096], 8, 2, 64, 0x1EAF); + assert!( + stats.plan_rebuilds >= 1, + "a 4096-token single-sequence launch should have built a plan, got {stats:?}" + ); + let rms = relative_rms(&out, &reference); + assert!(rms < 5e-3, "relative RMS {rms} exceeds 5e-3"); +} + +#[test] +fn a_short_single_sequence_stays_on_gather() { + // The one measured loss in #898 is batch 1 at 1024 tokens, so a lone + // request below 4096 must not fuse even though a batched launch of the same + // per-request size would. + let (_, _, stats) = run_step(&[1024], 8, 2, 64, 0x1EAF); + assert_eq!( + stats.plan_rebuilds, 0, + "a 1024-token single-sequence launch must stay on gather, got {stats:?}" + ); +} + +#[test] +fn a_batched_launch_just_under_a_nominal_1k_prompt_fuses() { + // The exact shape the production benchmark delivered for its nominal 1K + // scenario: 4 requests of 956 tokens. The first floor formulation summed to + // 3824 and declined it; the two-regime floor requires 4 x 512 = 2048. + let (out, reference, stats) = run_step(&[956, 956, 956, 956], 8, 2, 64, 0x956); + assert!( + stats.plan_rebuilds >= 1, + "4 x 956 tokens should fuse, got {stats:?}" + ); + let rms = relative_rms(&out, &reference); + assert!(rms < 5e-3, "relative RMS {rms} exceeds 5e-3"); +} + +#[test] +fn the_step_is_written_exactly_once() { + let (pool, states) = fresh_pool(1, 2, 64); + let mut rng = Rng::new(0xF00D); + let mut caches: Vec = (0..2) + .map(|b| KVCache::new_paged(pool.clone(), states[b].clone(), 0)) + .collect(); + for cache in caches.iter_mut() { + let k = random_array(&mut rng, &[1, 2, 100, 64]); + let v = random_array(&mut rng, &[1, 2, 100, 64]); + cache.update(k, v); + } + assert_eq!(states[0].borrow().layer(0).unwrap().len, 100); + + let q = random_array(&mut rng, &[2, 8, 1, 64]); + let k = random_array(&mut rng, &[2, 2, 1, 64]); + let v = random_array(&mut rng, &[2, 2, 1, 64]); + { + let mut refs: Vec<&mut KVCache> = caches.iter_mut().collect(); + paged_batch_decode_attention(&q, &k, &v, &mut refs, 0.125, 0.0).expect("servable"); + } + for state in states.iter().take(2) { + assert_eq!(state.borrow().layer(0).unwrap().len, 101); + } + // `offset` drives RoPE for the next step and must advance in lockstep. + assert_eq!(caches[0].offset, 101); + assert_eq!(caches[1].offset, 101); +} + +#[test] +fn a_steady_step_reuses_the_page_table() { + let (pool, states) = fresh_pool(1, 2, 64); + let mut rng = Rng::new(0x5EED); + let mut caches: Vec = (0..4) + .map(|b| KVCache::new_paged(pool.clone(), states[b].clone(), 0)) + .collect(); + for cache in caches.iter_mut() { + // 1024 = exactly 32 pages, so the next 31 decode tokens all land in a + // fresh page 32 opened by the first of them. + let k = random_array(&mut rng, &[1, 2, 1024, 64]); + let v = random_array(&mut rng, &[1, 2, 1024, 64]); + cache.update(k, v); + } + + let scale = 0.125; + let mut rebuilds = Vec::new(); + for _ in 0..4 { + let q = random_array(&mut rng, &[4, 8, 1, 64]); + let k = random_array(&mut rng, &[4, 2, 1, 64]); + let v = random_array(&mut rng, &[4, 2, 1, 64]); + let mut refs: Vec<&mut KVCache> = caches.iter_mut().collect(); + paged_batch_decode_attention(&q, &k, &v, &mut refs, scale, 0.0).expect("servable"); + drop(refs); + rebuilds.push(pool.borrow().decode_plan_cache_stats().view_rebuilds); + } + + // Step 1 opens page 32 for every request (epoch bump, rebuild). Steps 2, 3 + // and 4 stay inside it, so at most one further rebuild is allowed for the + // trailing epoch bump the first step left behind. + assert!( + rebuilds[3] - rebuilds[0] <= 1, + "steady steps should reuse the page table, rebuild counts were {rebuilds:?}" + ); + let stats = pool.borrow().decode_plan_cache_stats(); + assert!( + stats.view_reuses >= 2, + "expected page-table reuse, got {stats:?}" + ); +} + +#[test] +fn a_multi_slab_layer_falls_back_to_gather() { + // The pool default (32 rows) cannot hold 4 x 1024 tokens in one slab, so + // the fused path declines and the entry point answers from gather. + let bytes_per_block = PAGE * 2 * 64 * 2; + let layout = PagedKvLayout::uniform(1, PAGE, bytes_per_block).expect("valid layout"); + let pool = Rc::new(RefCell::new(PagedBlockPool::new(layout.clone()))); + let states: Vec>> = (0..4) + .map(|_| Rc::new(RefCell::new(PagedSequenceState::new(&layout)))) + .collect(); + let mut rng = Rng::new(0x51AB); + let mut caches: Vec = (0..4) + .map(|b| KVCache::new_paged(pool.clone(), states[b].clone(), 0)) + .collect(); + for cache in caches.iter_mut() { + let k = random_array(&mut rng, &[1, 2, 1024, 64]); + let v = random_array(&mut rng, &[1, 2, 1024, 64]); + cache.update(k, v); + } + assert!( + pool.borrow().slab_count(0) > 1, + "the default slab size should not hold this batch" + ); + + let q = random_array(&mut rng, &[4, 8, 1, 64]); + let k = random_array(&mut rng, &[4, 2, 1, 64]); + let v = random_array(&mut rng, &[4, 2, 1, 64]); + let scale = 0.125; + let out = { + let mut refs: Vec<&mut KVCache> = caches.iter_mut().collect(); + paged_batch_decode_attention(&q, &k, &v, &mut refs, scale, 0.0).expect("servable") + }; + let reference = { + let pool_ref = pool.borrow(); + let borrowed: Vec> = + states.iter().map(|s| s.borrow()).collect(); + let state_refs: Vec<&PagedSequenceState> = borrowed.iter().map(|s| &**s).collect(); + gather_fallback(&q, &pool_ref, &state_refs, 0, scale) + }; + let rms = relative_rms(&to_vec_f32(&out), &to_vec_f32(&reference)); + assert!( + rms < 1e-6, + "the gather fallback must reproduce the reference" + ); + // Nothing was cached, because the fused path was never reached. + assert_eq!(pool.borrow().decode_plan_cache_stats().plan_rebuilds, 0); +} + +// ── outcome reporting ─────────────────────────────────────────────────────── + +/// Drive one launch through the pool directly and return only its outcome. +fn outcome_for(prompt_len: usize, slabbed: bool) -> PagedDecodeOutcome { + let (pool, states) = if slabbed { + fresh_pool(1, 2, 64) + } else { + let bytes_per_block = PAGE * 2 * 64 * 2; + let layout = PagedKvLayout::uniform(1, PAGE, bytes_per_block).expect("valid layout"); + let pool = Rc::new(RefCell::new(PagedBlockPool::new(layout.clone()))); + let states = (0..4) + .map(|_| Rc::new(RefCell::new(PagedSequenceState::new(&layout)))) + .collect(); + (pool, states) + }; + let mut rng = Rng::new(0xD1A6); + let mut cache = KVCache::new_paged(pool.clone(), states[0].clone(), 0); + let k = random_array(&mut rng, &[1, 2, prompt_len as i32, 64]); + let v = random_array(&mut rng, &[1, 2, prompt_len as i32, 64]); + cache.update(k, v); + + let q = random_array(&mut rng, &[1, 8, 1, 64]); + let state_ref = states[0].borrow(); + let refs: Vec<&PagedSequenceState> = vec![&state_ref]; + let mut pool_mut = pool.borrow_mut(); + pool_mut + .paged_decode_batched(&q, &refs, 0, 0.125) + .expect("no hard error") + .1 +} + +#[test] +fn a_multi_slab_layer_reports_the_knob_that_would_fix_it() { + // The decline that silently disabled the whole fused path in the first #899 + // production benchmark. It has to name the slab counts and the knob, not + // just return `None`. + let outcome = outcome_for(2048, false); + match &outcome { + PagedDecodeOutcome::MultiSlab { + k_slabs, + slab_blocks, + .. + } => { + assert!(*k_slabs > 1, "expected several slabs, got {k_slabs}"); + assert_eq!(*slab_blocks, crate::cache::POOL_SLAB_BLOCKS); + } + other => panic!("expected a multi-slab decline, got {other:?}"), + } + assert!(!outcome.is_fused()); + assert!(outcome.describe().contains("MLXCEL_PAGED_SLAB_BLOCKS")); +} + +#[test] +fn a_below_floor_launch_reports_both_numbers() { + match outcome_for(64, true) { + PagedDecodeOutcome::BelowFloor { + batch, + visible_tokens, + floor, + } => { + assert_eq!(batch, 1); + assert_eq!(visible_tokens, 64); + assert_eq!(floor, crate::paged_v2::MIN_SINGLE_REQUEST_KV_TOKENS); + } + other => panic!("expected a below-floor decline, got {other:?}"), + } +} + +#[test] +fn a_fused_launch_reports_its_shape() { + match outcome_for(8192, true) { + PagedDecodeOutcome::Fused { + batch, + visible_tokens, + chunks, + .. + } => { + assert_eq!(batch, 1); + assert_eq!(visible_tokens, 8192); + assert!(chunks >= 1); + } + other => panic!("expected a fused launch, got {other:?}"), + } +} + +#[test] +fn an_unservable_batch_declines_without_touching_the_pool() { + reset_reported(); + let (pool, states) = fresh_pool(1, 2, 64); + let mut a = KVCache::new_paged(pool.clone(), states[0].clone(), 0); + let mut dense = KVCache::new(); + let mut caches: Vec<&mut KVCache> = vec![&mut a, &mut dense]; + let q = ffi::zeros(&[2, 8, 1, 64], dtype::FLOAT16); + let k = ffi::zeros(&[2, 2, 1, 64], dtype::FLOAT16); + let v = ffi::zeros(&[2, 2, 1, 64], dtype::FLOAT16); + assert!(paged_batch_decode_attention(&q, &k, &v, &mut caches, 0.125, 0.0).is_none()); + assert_eq!(states[0].borrow().layer(0).unwrap().len, 0); + assert_eq!(pool.borrow().allocated_block_count(), 0); +} diff --git a/src/lib/mlxcel-core/src/layers.rs b/src/lib/mlxcel-core/src/layers.rs index 8b796c390..01362accf 100644 --- a/src/lib/mlxcel-core/src/layers.rs +++ b/src/lib/mlxcel-core/src/layers.rs @@ -4054,12 +4054,21 @@ pub fn select_pooled_paged_dispatch( } /// Process-wide `MLXCEL_PAGED_ATTENTION_NATIVE` override for the fused -/// paged-attention kernel (#123, extended to tri-state in #331). +/// paged-attention kernel (#123, extended to tri-state in #331, extended to the +/// production paged decode path in #899). /// -/// Library-only control: this variable steers only the library-only -/// [`paged_decode_attention_pooled`] entry point (ADR 0001 #710), not the live -/// `mlxcel-server` block-table decode path (which uses the separate -/// `DecodeBatchContext::use_native_paged_kernel` request). +/// Two consumers: +/// +/// - The library-only [`paged_decode_attention_pooled`] entry point (ADR 0001 +/// #710), through [`resolve_dispatch_decision`]. Unchanged. +/// - The server's pool-backed batched paged decode, through +/// [`resolve_paged_v2_dispatch`]. Issue #899 made the fused v2 kernel the +/// production path for that decode, and named this variable's force-off +/// values (`0` / `false` / `off` / `no`) as its kill switch: setting it +/// restores the pre-#899 gather-then-SDPA behaviour end to end. +/// +/// The separate `DecodeBatchContext::use_native_paged_kernel` request still +/// governs the *dense-compat* block-table decode, which is a different path. #[derive(Debug, Clone, Copy, PartialEq, Eq)] enum NativePagedOverride { /// Force the fused kernel, bypassing the adaptive selector (the original @@ -4127,6 +4136,34 @@ fn resolve_dispatch_decision( } } +/// Apply the `MLXCEL_PAGED_ATTENTION_NATIVE` override to the production paged +/// decode v2 dispatch (issue #899). +/// +/// The override is checked *before* the selector so the kill switch cannot be +/// out-thought by a shape the selector likes: +/// +/// - force-off (`0` / `false` / `off` / `no`) pins the gather path, restoring +/// pre-#899 behaviour end to end; +/// - force-on (`1` / `true` / `on` / `yes`) pins v2 for every shape the kernel +/// can actually serve, bypassing the measured token floor (this is how the +/// floor itself gets re-measured, and how the low-context regime is +/// benchmarked deliberately); +/// - unset, or any other value, lets `selector` decide. +/// +/// A forced dispatch still goes through the kernel's own structural declines +/// (single-slab layer, servable geometry, non-empty batch), so forcing it on +/// cannot produce a launch the kernel would reject. +pub(crate) fn resolve_paged_v2_dispatch( + selector: impl FnOnce() -> crate::paged_v2::PagedV2Dispatch, +) -> crate::paged_v2::PagedV2Dispatch { + use crate::paged_v2::PagedV2Dispatch; + match native_paged_override() { + NativePagedOverride::ForceNative => PagedV2Dispatch::V2, + NativePagedOverride::ForceGather => PagedV2Dispatch::Gather, + NativePagedOverride::Auto => selector(), + } +} + /// The backend the fused kernel would run on, cached for the decode hot path. /// /// The fused kernel has a Metal JIT body and a CUDA JIT body (#634), so it is a diff --git a/src/lib/mlxcel-core/src/paged_v2/dispatch.rs b/src/lib/mlxcel-core/src/paged_v2/dispatch.rs new file mode 100644 index 000000000..4d183fada --- /dev/null +++ b/src/lib/mlxcel-core/src/paged_v2/dispatch.rs @@ -0,0 +1,273 @@ +// Copyright 2025-2026 Lablup Inc. and Jeongkyu Shin +// +// 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. + +//! Production dispatch policy for paged decode v2 (issue #899). +//! +//! Issue #898 landed v2 as a library capability and measured it against the +//! production gather-then-SDPA path on an Apple M1 Ultra +//! (`docs/benchmark_results/paged-decode-v2-m1ultra-2026-07-31.md`, medians of +//! three repetitions, head_dim 128 / 32 q heads / 8 kv heads / block 32): +//! +//! | batch | ctx 1024 | ctx 4096 | ctx 16384 | ctx 32768 | +//! |---|---|---|---|---| +//! | 1 | **0.91x** | 1.08x | 1.29x | 1.47x | +//! | 4 | 1.41x | 2.04x | 2.78x | 3.08x | +//! | 8 | 1.47x | 2.00x | 2.27x | 2.32x | +//! +//! Exactly one cell loses: batch 1 at 1024 tokens, 0.91x median and below +//! parity in all three repetitions. The plan degenerates to two pages per chunk +//! there, and the merge pass plus the workspace round trip costs more than the +//! whole attention does at that size. So v2 must not be dispatched +//! unconditionally. +//! +//! ## Why the floor is two-regime, and not one number +//! +//! The first cut of this policy summed the launch's visible tokens and required +//! 4096 of them, on the reasoning that `batch * ctx` is what the kernel reads +//! and therefore what the fixed merge and workspace cost has to amortize +//! against. That fits the table (the loss has 1024 total tokens, every win has +//! 4096 or more) but it is the wrong *shape*, and the server benchmark for #899 +//! showed it plainly: **the loss is a property of batch 1, not of total +//! tokens.** +//! +//! Read the ctx-1024 column: 0.91x at batch 1, 1.41x at batch 4, 1.47x at batch +//! 8. Same per-request context, same plan degeneracy, opposite outcome. What +//! changes is that at batch 4 the same chunk count is spread over four +//! requests, so the plan picks a larger chunk size, emits fewer chunks per +//! request, and amortizes the merge over four times the useful work. A +//! total-token floor separates those two cells only by accident, and with +//! almost no margin: a batched launch at 1024 tokens per request sits exactly +//! on 4096. The production benchmark's nominal "1K" scenario delivered 956 +//! tokens per request, `4 * 956 = 3824`, and was declined despite being the +//! same shape as the measured 1.41x win. +//! +//! So the floor is now stated the way the measurements are: +//! +//! | launch | floor | evidence | +//! |---|---|---| +//! | one request | [`MIN_SINGLE_REQUEST_KV_TOKENS`] visible tokens | 1024 loses (0.91x), 4096 wins (1.08x) | +//! | more than one | [`MIN_BATCHED_KV_TOKENS_PER_REQUEST`] per request | 1024 per request wins at batch 4 (1.41x) and batch 8 (1.47x) | +//! +//! The batched floor sits at 512 rather than 1024 deliberately. 1024 is the +//! *lowest measured* batched context and it wins comfortably; putting a floor +//! on top of a measured point means a real workload landing just under it (a +//! tokenizer delivering 956 tokens for a nominal 1K prompt, exactly what +//! happened) is declined for no measured reason. 512 keeps margin below the +//! evidence while still refusing launches small enough that the gather path is +//! trivially cheap. Batch 2 and 3 are interpolated, not measured; the trend +//! across batch 1, 4, and 8 at fixed context is monotone and the mechanism +//! above explains why, so the interpolation is stated rather than hidden. +//! +//! ## Overrides +//! +//! Both floors have an environment override for re-measurement on new hardware. +//! To force the fused path for every servable shape, including the regime this +//! policy declines, use `MLXCEL_PAGED_ATTENTION_NATIVE=1`, which bypasses the +//! selector entirely; that is the supported way to benchmark the declined +//! corner, rather than driving a floor to zero. + +use std::sync::OnceLock; + +/// Where a pooled paged decode batch runs. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum PagedV2Dispatch { + /// The fused v2 kernel: CSR page table, cross-CTA split-KV, merge. + V2, + /// Per-sequence `gather_visible` + dense SDPA, the pre-#899 production + /// path. + Gather, +} + +/// Visible tokens a **single-request** launch needs before v2 is dispatched. +/// +/// From the table in this module's docs: batch 1 at 1024 tokens is the only +/// measured loss (0.91x), and batch 1 at 4096 is the weakest measured win +/// (1.08x median, 1.02x to 1.21x across repetitions, so a win in every +/// repetition). +pub const MIN_SINGLE_REQUEST_KV_TOKENS: usize = 4096; + +/// Visible tokens **per request** a multi-request launch needs. +/// +/// Every measured multi-request cell wins, the smallest being batch 4 at 1024 +/// tokens per request (1.41x). This sits below that with margin rather than on +/// top of it; see the module docs for why the margin matters. +pub const MIN_BATCHED_KV_TOKENS_PER_REQUEST: usize = 512; + +/// Environment override for [`MIN_SINGLE_REQUEST_KV_TOKENS`]. +pub const MIN_KV_TOKENS_ENV: &str = "MLXCEL_PAGED_V2_MIN_KV_TOKENS"; + +/// Environment override for [`MIN_BATCHED_KV_TOKENS_PER_REQUEST`]. +pub const MIN_KV_TOKENS_PER_REQUEST_ENV: &str = "MLXCEL_PAGED_V2_MIN_KV_TOKENS_PER_REQUEST"; + +/// Pure parse of a floor override. +/// +/// Anything that is not a non-negative integer falls back to the default, so a +/// typo degrades to the measured policy rather than to an unmeasured one. +#[must_use] +pub fn parse_min_total_kv_tokens(value: Option<&str>, default: usize) -> usize { + match value.map(str::trim) { + Some(raw) if !raw.is_empty() => raw.parse::().unwrap_or(default), + _ => default, + } +} + +/// The active single-request floor, read once per process. +#[must_use] +pub fn min_total_kv_tokens() -> usize { + static FLOOR: OnceLock = OnceLock::new(); + *FLOOR.get_or_init(|| { + let raw = std::env::var(MIN_KV_TOKENS_ENV).ok(); + parse_min_total_kv_tokens(raw.as_deref(), MIN_SINGLE_REQUEST_KV_TOKENS) + }) +} + +/// The active per-request batched floor, read once per process. +#[must_use] +pub fn min_kv_tokens_per_request() -> usize { + static FLOOR: OnceLock = OnceLock::new(); + *FLOOR.get_or_init(|| { + let raw = std::env::var(MIN_KV_TOKENS_PER_REQUEST_ENV).ok(); + parse_min_total_kv_tokens(raw.as_deref(), MIN_BATCHED_KV_TOKENS_PER_REQUEST) + }) +} + +/// Visible tokens this launch must reach, given how many requests it serves. +/// +/// Pure, so the number that decided a dispatch can be logged next to the number +/// that was measured against it. +#[must_use] +pub fn required_visible_tokens(batch: usize, single: usize, per_request: usize) -> usize { + if batch <= 1 { + single + } else { + batch.saturating_mul(per_request) + } +} + +/// The floor this launch faces under the process's active configuration. +#[must_use] +pub fn active_required_visible_tokens(batch: usize) -> usize { + required_visible_tokens(batch, min_total_kv_tokens(), min_kv_tokens_per_request()) +} + +/// Pure regime selector: v2 once the launch reads at least `floor` visible KV +/// tokens in total, gather below that. +/// +/// Allocation-free and MLX-free, so it is trivially unit-testable and cheap +/// enough to evaluate per layer without memoization. +#[must_use] +pub fn select_paged_v2_dispatch(total_visible_tokens: usize, floor: usize) -> PagedV2Dispatch { + if total_visible_tokens >= floor { + PagedV2Dispatch::V2 + } else { + PagedV2Dispatch::Gather + } +} + +#[cfg(test)] +mod tests { + use super::*; + + /// Apply the shipped policy to a `(batch, per-request context)` cell. + fn decide(batch: usize, ctx: usize) -> PagedV2Dispatch { + let floor = required_visible_tokens( + batch, + MIN_SINGLE_REQUEST_KV_TOKENS, + MIN_BATCHED_KV_TOKENS_PER_REQUEST, + ); + select_paged_v2_dispatch(batch * ctx, floor) + } + + #[test] + fn the_only_measured_loss_is_declined() { + assert_eq!(decide(1, 1024), PagedV2Dispatch::Gather); + } + + #[test] + fn every_measured_win_is_dispatched() { + for (batch, ctx) in [ + (1, 4096), + (1, 16384), + (1, 32768), + (4, 1024), + (4, 4096), + (4, 16384), + (4, 32768), + (8, 1024), + (8, 4096), + (8, 16384), + (8, 32768), + ] { + assert_eq!( + decide(batch, ctx), + PagedV2Dispatch::V2, + "batch {batch} / ctx {ctx} should dispatch v2" + ); + } + } + + #[test] + fn a_batched_launch_just_under_a_nominal_1k_prompt_still_dispatches() { + // The regression that motivated the two-regime floor: the production + // benchmark's nominal 1K scenario delivered 956 tokens per request, and + // a flat 4096-total floor declined it even though batch 4 at ~1K is a + // measured 1.41x win. + assert_eq!(decide(4, 956), PagedV2Dispatch::V2); + assert_eq!(decide(2, 956), PagedV2Dispatch::V2); + } + + #[test] + fn a_single_request_is_not_helped_by_the_batched_floor() { + // The batched floor must never leak into batch 1, which is the only + // regime with a measured loss. + assert_eq!(decide(1, 512), PagedV2Dispatch::Gather); + assert_eq!(decide(1, 2048), PagedV2Dispatch::Gather); + assert_eq!(decide(1, 4095), PagedV2Dispatch::Gather); + assert_eq!(decide(1, 4096), PagedV2Dispatch::V2); + } + + #[test] + fn a_tiny_batched_launch_is_still_declined() { + // Not measured, and the gather path is trivially cheap at this size. + assert_eq!(decide(4, 128), PagedV2Dispatch::Gather); + assert_eq!(decide(2, 256), PagedV2Dispatch::Gather); + } + + #[test] + fn an_empty_batch_stays_on_gather() { + assert_eq!(decide(0, 0), PagedV2Dispatch::Gather); + assert_eq!(decide(4, 0), PagedV2Dispatch::Gather); + } + + #[test] + fn the_floor_scales_with_the_request_count() { + assert_eq!(required_visible_tokens(1, 4096, 512), 4096); + assert_eq!(required_visible_tokens(2, 4096, 512), 1024); + assert_eq!(required_visible_tokens(4, 4096, 512), 2048); + assert_eq!(required_visible_tokens(8, 4096, 512), 4096); + // Saturating, so an absurd request count cannot wrap the floor down. + assert_eq!(required_visible_tokens(usize::MAX, 4096, 512), usize::MAX); + } + + #[test] + fn the_env_overrides_parse_only_non_negative_integers() { + assert_eq!(parse_min_total_kv_tokens(None, 4096), 4096); + assert_eq!(parse_min_total_kv_tokens(Some(""), 4096), 4096); + assert_eq!(parse_min_total_kv_tokens(Some(" "), 4096), 4096); + assert_eq!(parse_min_total_kv_tokens(Some("nope"), 4096), 4096); + assert_eq!(parse_min_total_kv_tokens(Some("-1"), 4096), 4096); + assert_eq!(parse_min_total_kv_tokens(Some("0"), 4096), 0); + assert_eq!(parse_min_total_kv_tokens(Some(" 8192 "), 4096), 8192); + } +} diff --git a/src/lib/mlxcel-core/src/paged_v2/mod.rs b/src/lib/mlxcel-core/src/paged_v2/mod.rs index 39c1f26e4..9c619552b 100644 --- a/src/lib/mlxcel-core/src/paged_v2/mod.rs +++ b/src/lib/mlxcel-core/src/paged_v2/mod.rs @@ -45,34 +45,50 @@ //! for an occupancy-derived CTA target and emits the flat index arrays as //! plain data that can be cached across decode steps. //! -//! ## Default off +//! ## Two entry points, two gates //! -//! This issue lands v2 as a library capability. `MLXCEL_PAGED_ATTENTION_V2=1` -//! selects it inside -//! [`crate::cache::PagedBlockPool::paged_decode_fused`]; with the variable -//! unset the entry point is one `OnceLock` read away from the exact v1 code it -//! ran before, no v2 array is built, and no v2 kernel is JIT-compiled. v1 is -//! left fully intact for comparison. Production wiring (scheduler, server -//! defaults, dispatch thresholds) is issue #899. +//! **The v1 entry point** ([`crate::cache::PagedBlockPool::paged_decode_fused`], +//! reached only from the library-only [`crate::layers::paged_decode_attention_pooled`]) +//! tries v2 first under `MLXCEL_PAGED_ATTENTION_V2=1`, default off. That is +//! issue #898's comparison gate and is unchanged. +//! +//! **The production entry point** +//! ([`crate::cache::PagedBlockPool::paged_decode_batched`], reached from the +//! server's pool-backed batched decode) is issue #899 and is **default on**. It +//! is governed by [`dispatch`] (a measured token floor below which the gather +//! path is kept) and by `MLXCEL_PAGED_ATTENTION_NATIVE`, whose force-off values +//! are its kill switch. //! //! ## Not covered //! -//! Sliding-window and softcap variants keep their documented fallbacks, and -//! speculative multi-token verify stays on its current path: both are out of -//! scope per the issue. `first_page_offset` means the *page table* copes with a -//! trimmed window, but v2 applies no windowing mask of its own. +//! Logit-softcap families and speculative / MTP multi-token verify steps stay +//! on their existing paths; see [`crate::cache::paged_batch_decode`] for where +//! each is declined. `first_page_offset` means the *page table* copes with a +//! trimmed sliding window, but v2 applies no windowing mask of its own, so the +//! visible range has to be expressed entirely in the CSR range. use std::sync::OnceLock; +pub mod dispatch; pub mod launch; +pub mod outcome; pub mod plan; +pub mod plan_cache; +pub use dispatch::{ + MIN_BATCHED_KV_TOKENS_PER_REQUEST, MIN_KV_TOKENS_ENV, MIN_KV_TOKENS_PER_REQUEST_ENV, + MIN_SINGLE_REQUEST_KV_TOKENS, PagedV2Dispatch, active_required_visible_tokens, + min_kv_tokens_per_request, min_total_kv_tokens, parse_min_total_kv_tokens, + required_visible_tokens, select_paged_v2_dispatch, +}; pub use launch::{V2Context, geometry_from_shapes, resolve_plan, run_decode_v2}; +pub use outcome::{PAGED_DECODE_OUTCOME_KINDS, PagedDecodeOutcome}; pub use plan::{ MAX_CHUNKS, PagedDecodeGeometry, PagedDecodePlan, TARGET_CTAS_ENV, chunks_for_batch, chunks_for_request, device_target_ctas, max_pages_per_chunk, min_pages_per_chunk, search_pages_per_chunk, }; +pub use plan_cache::{DecodeBatchKey, PagedDecodeV2Cache, PlanCacheStats, RequestFingerprint}; /// Environment variable selecting the v2 decode path. Default off. pub const V2_ENV: &str = "MLXCEL_PAGED_ATTENTION_V2"; diff --git a/src/lib/mlxcel-core/src/paged_v2/outcome.rs b/src/lib/mlxcel-core/src/paged_v2/outcome.rs new file mode 100644 index 000000000..7e7ef142f --- /dev/null +++ b/src/lib/mlxcel-core/src/paged_v2/outcome.rs @@ -0,0 +1,229 @@ +// Copyright 2025-2026 Lablup Inc. and Jeongkyu Shin +// +// 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. + +//! What a production paged decode launch did, and why (issue #899). +//! +//! ## Why this is a value and not a log line +//! +//! The first cut of #899 reported declines with `tracing::debug!` from inside +//! the pool. The production benchmark then measured gather against gather for a +//! full sweep without anything saying so, because a real `mlxcel-server` run +//! emits no `DEBUG` from this crate: an operator has to opt into it, and +//! `RUST_LOG="info,mlxcel_core=debug"` did not produce a single line in the run +//! that was supposed to validate the change. A diagnostic that is invisible in +//! the one situation it exists for is not a diagnostic. +//! +//! So the decision is now a **return value**. The pool reports what it did, the +//! caller decides how loudly to say it, and the reason carries the numbers that +//! produced it rather than a bare "declined". Nothing on the hot path formats a +//! string unless a genuinely new outcome occurs. + +/// The outcome of one production paged decode launch, for one layer. +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum PagedDecodeOutcome { + /// The fused v2 kernel ran. + Fused { + /// Requests in the launch. + batch: usize, + /// Visible KV tokens the launch read. + visible_tokens: usize, + /// Chunks the plan emitted. + chunks: usize, + /// Whether the merge pass ran. + merged: bool, + }, + /// The layer's rows do not fit one contiguous pool buffer, so neither fused + /// kernel can address them. This is the decline that silently disables the + /// whole fused path when the slab is sized too small. + MultiSlab { + k_slabs: usize, + v_slabs: usize, + slab_blocks: usize, + }, + /// The kernel cannot serve this head geometry. + UnservableGeometry(String), + /// No request in the launch has a visible token. + NoVisibleTokens, + /// The launch is smaller than the dispatch policy's floor. + BelowFloor { + batch: usize, + visible_tokens: usize, + floor: usize, + }, + /// `MLXCEL_PAGED_ATTENTION_NATIVE` pinned the gather path. + KillSwitch, + /// The plan failed its structural check. + PlanRejected(String), + /// The page table could not be built. + ViewFailed(String), + /// The caller declined before any pool write: not all caches are + /// pool-backed on one pool and layer, a soft-cap was requested, or the step + /// carries more than one query token (batched prefill, speculative verify). + NotServable(&'static str), +} + +/// Number of distinct outcome kinds, for the caller's one-shot log table. +pub const PAGED_DECODE_OUTCOME_KINDS: usize = 9; + +impl PagedDecodeOutcome { + /// Whether the fused kernel actually ran. + #[must_use] + pub fn is_fused(&self) -> bool { + matches!(self, Self::Fused { .. }) + } + + /// Stable index for this kind, so a caller can keep a fixed-size table of + /// "already reported" flags without allocating or hashing. + #[must_use] + pub fn kind_index(&self) -> usize { + match self { + Self::Fused { .. } => 0, + Self::MultiSlab { .. } => 1, + Self::UnservableGeometry(_) => 2, + Self::NoVisibleTokens => 3, + Self::BelowFloor { .. } => 4, + Self::KillSwitch => 5, + Self::PlanRejected(_) => 6, + Self::ViewFailed(_) => 7, + Self::NotServable(_) => 8, + } + } + + /// One-line human summary carrying the numbers that produced the decision. + /// + /// Allocates, so callers format it only when they are about to emit it. + #[must_use] + pub fn describe(&self) -> String { + match self { + Self::Fused { + batch, + visible_tokens, + chunks, + merged, + } => format!( + "fused v2 launch (batch {batch}, {visible_tokens} visible KV tokens, \ + {chunks} chunks, merge {})", + if *merged { "on" } else { "skipped" } + ), + Self::MultiSlab { + k_slabs, + v_slabs, + slab_blocks, + } => format!( + "gather: the layer spans {k_slabs} K / {v_slabs} V pool slabs and the fused \ + kernels read one buffer per side (slab_blocks={slab_blocks}); raise --ctx-size \ + or MLXCEL_PAGED_SLAB_BLOCKS so a layer's rows fit one slab" + ), + Self::UnservableGeometry(reason) => { + format!("gather: the kernel cannot serve this geometry ({reason})") + } + Self::NoVisibleTokens => "gather: no request has a visible token".to_string(), + Self::BelowFloor { + batch, + visible_tokens, + floor, + } => format!( + "gather: {visible_tokens} visible KV tokens across {batch} request(s) is below \ + the {floor}-token dispatch floor" + ), + Self::KillSwitch => "gather: pinned by MLXCEL_PAGED_ATTENTION_NATIVE".to_string(), + Self::PlanRejected(reason) => format!("gather: the chunk plan was rejected ({reason})"), + Self::ViewFailed(reason) => { + format!("gather: the page table failed to build ({reason})") + } + Self::NotServable(reason) => format!("gather: batch not servable ({reason})"), + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn kind_indices_are_distinct_and_in_range() { + let all = [ + PagedDecodeOutcome::Fused { + batch: 1, + visible_tokens: 1, + chunks: 1, + merged: false, + }, + PagedDecodeOutcome::MultiSlab { + k_slabs: 2, + v_slabs: 2, + slab_blocks: 32, + }, + PagedDecodeOutcome::UnservableGeometry("x".to_string()), + PagedDecodeOutcome::NoVisibleTokens, + PagedDecodeOutcome::BelowFloor { + batch: 1, + visible_tokens: 1, + floor: 2, + }, + PagedDecodeOutcome::KillSwitch, + PagedDecodeOutcome::PlanRejected("x".to_string()), + PagedDecodeOutcome::ViewFailed("x".to_string()), + PagedDecodeOutcome::NotServable("x"), + ]; + assert_eq!(all.len(), PAGED_DECODE_OUTCOME_KINDS); + let mut seen = [false; PAGED_DECODE_OUTCOME_KINDS]; + for outcome in &all { + let i = outcome.kind_index(); + assert!(i < PAGED_DECODE_OUTCOME_KINDS, "{outcome:?} index {i}"); + assert!(!seen[i], "duplicate index for {outcome:?}"); + seen[i] = true; + assert!(!outcome.describe().is_empty()); + } + } + + #[test] + fn only_the_fused_variant_reports_a_fused_launch() { + assert!( + PagedDecodeOutcome::Fused { + batch: 4, + visible_tokens: 4096, + chunks: 16, + merged: true, + } + .is_fused() + ); + assert!(!PagedDecodeOutcome::NoVisibleTokens.is_fused()); + assert!(!PagedDecodeOutcome::KillSwitch.is_fused()); + } + + #[test] + fn the_multi_slab_message_names_the_knob_that_fixes_it() { + let text = PagedDecodeOutcome::MultiSlab { + k_slabs: 3, + v_slabs: 3, + slab_blocks: 32, + } + .describe(); + assert!(text.contains("MLXCEL_PAGED_SLAB_BLOCKS"), "{text}"); + assert!(text.contains("--ctx-size"), "{text}"); + } + + #[test] + fn the_floor_message_carries_both_numbers() { + let text = PagedDecodeOutcome::BelowFloor { + batch: 4, + visible_tokens: 3824, + floor: 4096, + } + .describe(); + assert!(text.contains("3824"), "{text}"); + assert!(text.contains("4096"), "{text}"); + } +} diff --git a/src/lib/mlxcel-core/src/paged_v2/plan_cache.rs b/src/lib/mlxcel-core/src/paged_v2/plan_cache.rs new file mode 100644 index 000000000..8f5ef0c7d --- /dev/null +++ b/src/lib/mlxcel-core/src/paged_v2/plan_cache.rs @@ -0,0 +1,369 @@ +// Copyright 2025-2026 Lablup Inc. and Jeongkyu Shin +// +// 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. + +//! Per-batch caching of the CSR page table and the chunk plan (issue #899). +//! +//! ## What is expensive and what is not +//! +//! Building a [`PagedCsrView`] costs one hash lookup per **visible page**: at +//! 32768 tokens with 32-token pages that is 1024 lookups per request per layer, +//! and a 32-layer model at batch 4 would pay ~131k lookups every decode step. +//! Refreshing the per-request scalars costs one pass over the **batch**: four +//! `i32` writes per request, so 16 writes for that same batch. +//! +//! A decode step appends exactly one token per request, which changes only the +//! scalars unless the request crosses a page boundary. That is the asymmetry +//! this cache exploits: keep `indices` / `indptr` (the expensive half), +//! recompute `last_page_len` / `first_page_offset` / `seq_lens` / +//! `rope_offsets` (the cheap half) from the live layer states every step. +//! +//! The cheap half is **recomputed, never incremented**. A delta ("bump +//! `last_page_len` by one") would be one missed invalidation away from feeding +//! the kernel a length that outruns its pages, which is an out-of-bounds read +//! rather than a wrong answer. Recomputing from [`PagedLayerState`] costs the +//! same O(batch) and cannot drift. +//! +//! ## Reuse predicate +//! +//! `indices` and `indptr` are a function of exactly two things: +//! +//! 1. Each request's visible **page range**, `logical_start / page_size ..= +//! (len - 1) / page_size`, which is derived from the per-request +//! fingerprint below. +//! 2. The block-id to physical-row mapping and the block table contents, which +//! change only through a `&mut PagedBlockPool` method. +//! +//! (1) is checked directly against the live layer states. (2) is covered by +//! [`PagedBlockPool::block_epoch`], a counter the pool bumps on every mutation +//! that can move a block: acquire, release, row assignment, row forget, block +//! restore, sequence restore, sequence release, and token trim. A cached view +//! is reusable only when the epoch it was built at is still current. +//! +//! The epoch is what makes the issue's named invalidation events fall out for +//! free rather than needing scheduler plumbing: +//! +//! | event | what bumps the epoch | +//! |---|---| +//! | admission of a new sequence | its first `write_prefill` acquires blocks | +//! | eviction / preemption | `release_sequence` releases them | +//! | sequence finish | same | +//! | crossing a page boundary | `append_tokens` acquires the new block | +//! | prompt-cache adopt / detach | `restore_sequence` / `forget` | +//! +//! The epoch is global rather than per-layer, so a boundary crossing costs two +//! fully rebuilt steps rather than one (the crossing step, plus the following +//! step whose per-layer entries were cached at intermediate epoch values). With +//! 32-token pages that is 2 rebuilt steps in 32, or ~6% of the uncached cost. +//! +//! ## Why the pool owns it +//! +//! The issue asks for the plan to be cached "on the active batch". The pool is +//! where that lives in practice: it is the object the decode path already holds +//! mutably, it is the object every invalidation event already goes through, and +//! putting the cache anywhere else would need the scheduler to forward events +//! that the pool observes first hand. See the PR body for the trade-off. +//! +//! [`PagedBlockPool::block_epoch`]: crate::cache::PagedBlockPool::block_epoch + +use crate::cache::{PagedCsrView, PagedLayerState}; +use crate::paged_v2::plan::{PagedDecodeGeometry, PagedDecodePlan}; + +/// Everything about one request that can change `indices` / `indptr`. +/// +/// Deliberately not a pointer or an id: two different sequences cannot present +/// the same fingerprint at the same epoch, because a sequence entering or +/// leaving the batch moves blocks and therefore moves the epoch. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct RequestFingerprint { + /// Blocks retained by the layer (the block table length). + pub blocks: usize, + /// First visible page index, `logical_start / page_size`. + pub first_page: usize, + /// Last written page index, `(len - 1) / page_size`. Meaningless when the + /// request has no visible tokens, in which case `visible` is false. + pub last_page: usize, + /// Whether the request has any visible token at all. + pub visible: bool, +} + +impl RequestFingerprint { + /// Fingerprint one layer state at a given page size. + #[must_use] + pub fn of(layer: &PagedLayerState, page_size: usize) -> Self { + let page_size = page_size.max(1); + let visible = layer.visible_len() > 0; + Self { + blocks: layer.block_ids.len(), + first_page: layer.logical_start / page_size, + last_page: if layer.len == 0 { + 0 + } else { + (layer.len - 1) / page_size + }, + visible, + } + } +} + +/// Everything [`PagedDecodeV2Cache::view_and_plan`] needs to identify the batch +/// it is being asked about: which layer, at which block-topology epoch, over +/// which live layer states, for which launch geometry. +/// +/// A struct rather than five parameters because these five always travel +/// together and are meaningless apart. +#[derive(Debug, Clone, Copy)] +pub struct DecodeBatchKey<'a> { + /// Layer whose page table is wanted. + pub layer_idx: usize, + /// The pool's current block-topology epoch. + pub epoch: u64, + /// Tokens per pool page. + pub page_size: usize, + /// One live layer state per request, in batch order. + pub layers: &'a [&'a PagedLayerState], + /// Launch geometry the plan is built against. + pub geometry: &'a PagedDecodeGeometry, +} + +/// A cached page table plus the epoch and fingerprints that justify reusing it. +#[derive(Debug, Clone)] +struct CachedView { + epoch: u64, + page_size: usize, + fingerprints: Vec, + view: PagedCsrView, +} + +/// Counters for the cache's behaviour, surfaced for tests and tracing. +#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)] +pub struct PlanCacheStats { + /// Views served by refreshing only the per-request scalars. + pub view_reuses: u64, + /// Views rebuilt from the block tables. + pub view_rebuilds: u64, + /// Plans served from the cached plan. + pub plan_reuses: u64, + /// Plans rebuilt (chunk search re-run). + pub plan_rebuilds: u64, +} + +/// Per-pool cache of CSR views (one entry per layer) and the chunk plan. +#[derive(Debug, Default)] +pub struct PagedDecodeV2Cache { + views: Vec>, + plan: Option, + stats: PlanCacheStats, +} + +impl PagedDecodeV2Cache { + /// An empty cache. + #[must_use] + pub fn new() -> Self { + Self::default() + } + + /// Drop everything. Called by the pool whenever it cannot cheaply reason + /// about what changed (layout replacement, explicit scheduler request). + pub fn clear(&mut self) { + self.views.clear(); + self.plan = None; + } + + /// Cache counters since construction. + #[must_use] + pub fn stats(&self) -> PlanCacheStats { + self.stats + } + + /// The CSR view for one layer, reusing the cached page list when the batch + /// has only advanced within its current pages. + /// + /// `build` is invoked only on a miss and must produce the view for exactly + /// these layer states. + pub fn view_for( + &mut self, + layer_idx: usize, + epoch: u64, + page_size: usize, + layers: &[&PagedLayerState], + build: F, + ) -> Result<&PagedCsrView, String> + where + F: FnOnce() -> Result, + { + self.ensure_view_inner(layer_idx, epoch, page_size, layers, build)?; + Ok(self + .view(layer_idx) + .expect("ensure_view_inner populated the entry")) + } + + /// The cached view for a layer, if one is present. + #[must_use] + pub fn view(&self, layer_idx: usize) -> Option<&PagedCsrView> { + self.views + .get(layer_idx) + .and_then(Option::as_ref) + .map(|cached| &cached.view) + } + + /// The cached chunk plan, if one is present. + #[must_use] + pub fn plan(&self) -> Option<&PagedDecodePlan> { + self.plan.as_ref() + } + + /// Populate the layer's CSR view, reusing the cached page list when the + /// batch has only advanced within its current pages. + /// + /// Separate from [`Self::ensure_plan`] on purpose: the dispatch policy + /// needs the view's visible-token count to decide whether the fused path + /// runs at all, and a declined launch must not pay for the plan's chunk + /// search. Splitting the two also makes the plan-rebuild counter an honest + /// "the fused path ran" signal rather than "the fused path was considered". + pub fn ensure_view(&mut self, key: &DecodeBatchKey<'_>, build: F) -> Result<(), String> + where + F: FnOnce() -> Result, + { + self.ensure_view_inner(key.layer_idx, key.epoch, key.page_size, key.layers, build) + } + + /// Populate the chunk plan for these page counts, reusing the cached plan + /// when it still emits the same chunk grouping. + /// + /// [`PagedDecodePlan::matches`] is the plan's own O(batch) predicate: it + /// checks that every request still needs exactly the number of chunks the + /// cached `o_indptr` allotted it, which is what makes the plan's flat index + /// arrays still correct. + pub fn ensure_plan( + &mut self, + geometry: &PagedDecodeGeometry, + page_counts: &[usize], + build: F, + ) where + F: FnOnce() -> PagedDecodePlan, + { + let hit = self + .plan + .as_ref() + .is_some_and(|plan| plan.matches(geometry, page_counts)); + if hit { + self.stats.plan_reuses += 1; + } else { + self.plan = Some(build()); + self.stats.plan_rebuilds += 1; + } + } + + fn ensure_view_inner( + &mut self, + layer_idx: usize, + epoch: u64, + page_size: usize, + layers: &[&PagedLayerState], + build: F, + ) -> Result<(), String> + where + F: FnOnce() -> Result, + { + let fingerprints: Vec = layers + .iter() + .map(|layer| RequestFingerprint::of(layer, page_size)) + .collect(); + + if self.views.len() <= layer_idx { + self.views.resize_with(layer_idx + 1, || None); + } + let reusable = match &self.views[layer_idx] { + Some(cached) => { + cached.epoch == epoch + && cached.page_size == page_size + && cached.fingerprints == fingerprints + } + None => false, + }; + + if reusable { + let cached = self.views[layer_idx] + .as_mut() + .expect("reusable implies the entry is present"); + refresh_scalars(&mut cached.view, layers, page_size); + // Cheap (one pass over the batch, not over the pages) and the same + // assertion the builder runs, so a reused view is held to exactly + // the invariants a freshly built one is. + cached.view.validate()?; + self.stats.view_reuses += 1; + return Ok(()); + } + + let view = build()?; + self.stats.view_rebuilds += 1; + self.views[layer_idx] = Some(CachedView { + epoch, + page_size, + fingerprints, + view, + }); + Ok(()) + } + + /// The chunk plan for these page counts, reusing the cached plan when it + /// still emits the same chunk grouping. + /// + /// [`PagedDecodePlan::matches`] is the plan's own O(batch) predicate: it + /// checks that every request still needs exactly the number of chunks the + /// cached `o_indptr` allotted it, which is what makes the plan's flat index + /// arrays still correct. + pub fn plan_for( + &mut self, + geometry: &PagedDecodeGeometry, + page_counts: &[usize], + build: F, + ) -> &PagedDecodePlan + where + F: FnOnce() -> PagedDecodePlan, + { + self.ensure_plan(geometry, page_counts, build); + self.plan.as_ref().expect("plan is present after the miss") + } +} + +/// Recompute the per-request scalars of a reused view from the live layer +/// states. The page list (`indices` / `indptr`) is untouched; the caller has +/// already established that it is still correct. +fn refresh_scalars(view: &mut PagedCsrView, layers: &[&PagedLayerState], page_size: usize) { + let page_size = page_size.max(1); + for (r, layer) in layers.iter().enumerate() { + let visible = layer.visible_len(); + view.rope_offsets[r] = clamp_i32(layer.len); + if visible == 0 { + view.last_page_len[r] = 0; + view.first_page_offset[r] = 0; + view.seq_lens[r] = 0; + continue; + } + view.last_page_len[r] = clamp_i32((layer.len - 1) % page_size + 1); + view.first_page_offset[r] = clamp_i32(layer.logical_start % page_size); + view.seq_lens[r] = clamp_i32(visible); + } +} + +/// Saturating `usize -> i32`, matching `build_paged_csr_view`'s conversion so a +/// refreshed view and a rebuilt view agree element for element. +fn clamp_i32(v: usize) -> i32 { + i32::try_from(v).unwrap_or(i32::MAX) +} + +#[cfg(test)] +#[path = "plan_cache_tests.rs"] +mod plan_cache_tests; diff --git a/src/lib/mlxcel-core/src/paged_v2/plan_cache_tests.rs b/src/lib/mlxcel-core/src/paged_v2/plan_cache_tests.rs new file mode 100644 index 000000000..4c79ab131 --- /dev/null +++ b/src/lib/mlxcel-core/src/paged_v2/plan_cache_tests.rs @@ -0,0 +1,333 @@ +// Copyright 2025-2026 Lablup Inc. and Jeongkyu Shin +// +// 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. + +//! Unit tests for the CSR view / chunk plan cache (issue #899). +//! +//! MLX-free: every test drives synthetic [`PagedLayerState`]s and the pure +//! [`build_paged_csr_view`] builder, so the reuse predicate is exercised +//! without a GPU. + +use super::*; +use crate::cache::{PagedBlockId, build_paged_csr_view}; + +const PAGE: usize = 32; + +/// A layer state with `len` written tokens, `logical_start` trimmed from the +/// front, and enough blocks to cover `len`. Block ids are `base + page` so +/// distinct requests get distinct blocks. +fn layer(base: u64, len: usize, logical_start: usize) -> PagedLayerState { + let blocks = len.div_ceil(PAGE); + PagedLayerState { + block_ids: (0..blocks) + .map(|p| PagedBlockId::from_raw(base + p as u64)) + .collect(), + len, + logical_start, + } +} + +/// Row assignment mirroring a real pool closely enough for the builder: block +/// id `n` lives on row `n`. +fn build(layers: &[&PagedLayerState]) -> Result { + build_paged_csr_view(PAGE, layers, |id| Some(id.as_u64() as usize)) +} + +fn geometry() -> PagedDecodeGeometry { + PagedDecodeGeometry { + q_heads: 32, + kv_heads: 8, + head_dim: 128, + page_size: PAGE as i32, + } +} + +#[test] +fn a_step_inside_the_current_pages_reuses_the_page_list() { + let mut cache = PagedDecodeV2Cache::new(); + let a = layer(0, 100, 0); + let b = layer(100, 200, 0); + let layers: Vec<&PagedLayerState> = vec![&a, &b]; + let first = cache + .view_for(0, 7, PAGE, &layers, || build(&layers)) + .expect("build") + .clone(); + assert_eq!(cache.stats().view_rebuilds, 1); + + // One decode token each: 100 -> 101 and 200 -> 201, neither crosses a page + // boundary (100 % 32 = 4, 200 % 32 = 8). + let a2 = layer(0, 101, 0); + let b2 = layer(100, 201, 0); + let layers2: Vec<&PagedLayerState> = vec![&a2, &b2]; + let reused = cache + .view_for(0, 7, PAGE, &layers2, || panic!("must not rebuild")) + .expect("reuse") + .clone(); + + assert_eq!(cache.stats().view_reuses, 1); + assert_eq!(cache.stats().view_rebuilds, 1); + assert_eq!(reused.indices, first.indices); + assert_eq!(reused.indptr, first.indptr); + assert_eq!(reused.seq_lens, vec![101, 201]); + assert_eq!(reused.rope_offsets, vec![101, 201]); + assert_eq!(reused.last_page_len, vec![101 - 96, 201 - 192]); +} + +#[test] +fn a_reused_view_is_identical_to_a_fresh_build() { + let mut cache = PagedDecodeV2Cache::new(); + let a = layer(0, 100, 0); + let b = layer(100, 200, 5); + let layers: Vec<&PagedLayerState> = vec![&a, &b]; + cache + .view_for(0, 1, PAGE, &layers, || build(&layers)) + .expect("build"); + + let a2 = layer(0, 101, 0); + let b2 = layer(100, 201, 5); + let layers2: Vec<&PagedLayerState> = vec![&a2, &b2]; + let reused = cache + .view_for(0, 1, PAGE, &layers2, || panic!("must not rebuild")) + .expect("reuse") + .clone(); + let fresh = build(&layers2).expect("fresh build"); + assert_eq!(reused, fresh); +} + +#[test] +fn crossing_a_page_boundary_rebuilds() { + let mut cache = PagedDecodeV2Cache::new(); + // 96 = 3 full pages; the next token opens page 3. + let a = layer(0, 96, 0); + let layers: Vec<&PagedLayerState> = vec![&a]; + cache + .view_for(0, 1, PAGE, &layers, || build(&layers)) + .expect("build"); + + let a2 = layer(0, 97, 0); + let layers2: Vec<&PagedLayerState> = vec![&a2]; + let rebuilt = cache + .view_for(0, 1, PAGE, &layers2, || build(&layers2)) + .expect("rebuild") + .clone(); + assert_eq!(cache.stats().view_rebuilds, 2); + assert_eq!(cache.stats().view_reuses, 0); + assert_eq!(rebuilt.indices.len(), 4); +} + +#[test] +fn an_epoch_bump_rebuilds_even_when_the_fingerprints_match() { + let mut cache = PagedDecodeV2Cache::new(); + let a = layer(0, 100, 0); + let layers: Vec<&PagedLayerState> = vec![&a]; + cache + .view_for(0, 1, PAGE, &layers, || build(&layers)) + .expect("build"); + cache + .view_for(0, 2, PAGE, &layers, || build(&layers)) + .expect("rebuild after epoch bump"); + assert_eq!(cache.stats().view_rebuilds, 2); + assert_eq!(cache.stats().view_reuses, 0); +} + +#[test] +fn a_front_trim_that_retires_a_page_rebuilds() { + let mut cache = PagedDecodeV2Cache::new(); + let a = layer(0, 200, 0); + let layers: Vec<&PagedLayerState> = vec![&a]; + let before = cache + .view_for(0, 1, PAGE, &layers, || build(&layers)) + .expect("build") + .clone(); + assert_eq!(before.indices.len(), 7); + + // logical_start 64 retires pages 0 and 1: same epoch, different first_page. + let trimmed = layer(0, 200, 64); + let layers2: Vec<&PagedLayerState> = vec![&trimmed]; + let after = cache + .view_for(0, 1, PAGE, &layers2, || build(&layers2)) + .expect("rebuild") + .clone(); + assert_eq!(cache.stats().view_rebuilds, 2); + assert_eq!(after.indices.len(), 5); + assert_eq!(after.indices[0], 2); + assert_eq!(after.first_page_offset, vec![0]); + assert_eq!(after.seq_lens, vec![136]); +} + +#[test] +fn a_front_trim_inside_the_same_page_reuses_and_moves_the_offset() { + let mut cache = PagedDecodeV2Cache::new(); + let a = layer(0, 200, 3); + let layers: Vec<&PagedLayerState> = vec![&a]; + cache + .view_for(0, 1, PAGE, &layers, || build(&layers)) + .expect("build"); + + // Still inside page 0, so the page list is unchanged. + let a2 = layer(0, 200, 7); + let layers2: Vec<&PagedLayerState> = vec![&a2]; + let reused = cache + .view_for(0, 1, PAGE, &layers2, || panic!("must not rebuild")) + .expect("reuse") + .clone(); + assert_eq!(cache.stats().view_reuses, 1); + assert_eq!(reused.first_page_offset, vec![7]); + assert_eq!(reused.seq_lens, vec![193]); + assert_eq!(reused, build(&layers2).expect("fresh")); +} + +#[test] +fn a_batch_size_change_rebuilds() { + let mut cache = PagedDecodeV2Cache::new(); + let a = layer(0, 100, 0); + let b = layer(100, 100, 0); + let two: Vec<&PagedLayerState> = vec![&a, &b]; + cache + .view_for(0, 1, PAGE, &two, || build(&two)) + .expect("build"); + let one: Vec<&PagedLayerState> = vec![&a]; + cache + .view_for(0, 1, PAGE, &one, || build(&one)) + .expect("rebuild"); + assert_eq!(cache.stats().view_rebuilds, 2); +} + +#[test] +fn layers_are_cached_independently() { + let mut cache = PagedDecodeV2Cache::new(); + let l0 = layer(0, 100, 0); + let l1 = layer(500, 100, 0); + let a: Vec<&PagedLayerState> = vec![&l0]; + let b: Vec<&PagedLayerState> = vec![&l1]; + cache.view_for(0, 1, PAGE, &a, || build(&a)).expect("l0"); + cache.view_for(5, 1, PAGE, &b, || build(&b)).expect("l5"); + assert_eq!(cache.stats().view_rebuilds, 2); + + let l0b = layer(0, 101, 0); + let l1b = layer(500, 101, 0); + let a2: Vec<&PagedLayerState> = vec![&l0b]; + let b2: Vec<&PagedLayerState> = vec![&l1b]; + let v0 = cache + .view_for(0, 1, PAGE, &a2, || panic!("l0 must reuse")) + .expect("l0 reuse") + .clone(); + let v5 = cache + .view_for(5, 1, PAGE, &b2, || panic!("l5 must reuse")) + .expect("l5 reuse") + .clone(); + assert_eq!(cache.stats().view_reuses, 2); + assert_eq!(v0.indices, vec![0, 1, 2, 3]); + assert_eq!(v5.indices, vec![500, 501, 502, 503]); +} + +#[test] +fn clear_drops_every_layer() { + let mut cache = PagedDecodeV2Cache::new(); + let a = layer(0, 100, 0); + let layers: Vec<&PagedLayerState> = vec![&a]; + cache + .view_for(0, 1, PAGE, &layers, || build(&layers)) + .expect("build"); + cache.clear(); + cache + .view_for(0, 1, PAGE, &layers, || build(&layers)) + .expect("rebuild after clear"); + assert_eq!(cache.stats().view_rebuilds, 2); + assert_eq!(cache.stats().view_reuses, 0); +} + +#[test] +fn the_plan_is_reused_while_the_chunk_grouping_holds() { + let mut cache = PagedDecodeV2Cache::new(); + let geo = geometry(); + let counts = vec![64usize, 64]; + let plan = cache + .plan_for(&geo, &counts, || { + PagedDecodePlan::heuristic(geo, &counts, 512) + }) + .clone(); + assert_eq!(cache.stats().plan_rebuilds, 1); + + // Same page counts: the plan is bit-for-bit reusable. + let again = cache + .plan_for(&geo, &counts, || panic!("must not rebuild")) + .clone(); + assert_eq!(cache.stats().plan_reuses, 1); + assert_eq!(again, plan); +} + +#[test] +fn the_plan_is_rebuilt_when_a_request_needs_another_chunk() { + let mut cache = PagedDecodeV2Cache::new(); + let geo = geometry(); + let counts = vec![64usize, 64]; + let plan = cache + .plan_for(&geo, &counts, || { + PagedDecodePlan::heuristic(geo, &counts, 512) + }) + .clone(); + + // Grow one request by a whole chunk's worth of pages so `matches` fails. + let grown = vec![64usize + plan.pages_per_chunk as usize, 64]; + let rebuilt = cache + .plan_for(&geo, &grown, || { + PagedDecodePlan::heuristic(geo, &grown, 512) + }) + .clone(); + assert_eq!(cache.stats().plan_rebuilds, 2); + assert!(rebuilt.num_chunks > plan.num_chunks); +} + +#[test] +fn the_plan_is_rebuilt_when_the_geometry_changes() { + let mut cache = PagedDecodeV2Cache::new(); + let geo = geometry(); + let counts = vec![64usize]; + cache.plan_for(&geo, &counts, || { + PagedDecodePlan::heuristic(geo, &counts, 512) + }); + let other = PagedDecodeGeometry { + head_dim: 64, + ..geo + }; + cache.plan_for(&other, &counts, || { + PagedDecodePlan::heuristic(other, &counts, 512) + }); + assert_eq!(cache.stats().plan_rebuilds, 2); + assert_eq!(cache.stats().plan_reuses, 0); +} + +#[test] +fn fingerprints_capture_the_visible_page_range() { + let empty = layer(0, 0, 0); + let fp = RequestFingerprint::of(&empty, PAGE); + assert!(!fp.visible); + assert_eq!(fp.blocks, 0); + + let one = layer(0, 1, 0); + let fp = RequestFingerprint::of(&one, PAGE); + assert!(fp.visible); + assert_eq!((fp.first_page, fp.last_page), (0, 0)); + + let spanning = layer(0, 100, 40); + let fp = RequestFingerprint::of(&spanning, PAGE); + assert_eq!((fp.first_page, fp.last_page), (1, 3)); + + // A fully trimmed layer has no visible tokens even though `len` is + // non-zero, and must not compare equal to a visible one. + let trimmed = layer(0, 100, 100); + let fp_trimmed = RequestFingerprint::of(&trimmed, PAGE); + assert!(!fp_trimmed.visible); + assert_ne!(fp_trimmed, fp); +} diff --git a/src/models/llama3.rs b/src/models/llama3.rs index 18e6aee42..6b33bff03 100644 --- a/src/models/llama3.rs +++ b/src/models/llama3.rs @@ -440,7 +440,30 @@ impl Attention { mlxcel_core::cache::turbo::sparse_v::turbo4_delegated_compressed_attention_enabled(); let use_turbo4_dequant_sdpa = mlxcel_core::cache::turbo::sparse_v::turbo4_dequant_sdpa_enabled(); - let attn_out = if l == 1 + // Single-sequence pooled paged decode (#899). The server takes this + // path whenever the active batch is one request: `decode_single_step` + // calls the single-sequence `forward`, never `forward_split_attention`. + // Without this branch the two single-sequence scenarios in the issue's + // benchmark matrix (16K and 32K) could never reach the fused kernel, and + // neither could any moment in a batched run where only one request is + // still decoding. The batch-1 launch is exactly the `[1, H, 1, D]` shape + // the whole-batch entry point already serves. + let pooled_single = if l == 1 && mask.is_none() && cache.is_paged_backed() { + mlxcel_core::cache::paged_batch_decode_attention( + &q, + &k, + &v, + &mut [&mut *cache], + self.scale, + 0.0, + ) + } else { + None + }; + + let attn_out = if let Some(out) = pooled_single { + out + } else if l == 1 && use_turbo4_asym_dequant_sdpa && cache.turbo4_asym_dequant_sdpa_available() { @@ -546,6 +569,29 @@ impl Attention { &metadata.rope_offsets, ); + // Pool-backed paged decode (#899): the whole batch in one fused launch. + // This is the production path for `--decode-storage paged`; it appends + // every sequence's new K/V to the shared `PagedBlockPool` and runs + // attention once over the batch, replacing the per-sequence + // `update_and_fetch` + `gather_visible` + SDPA loop below. It declines + // (leaving the pool untouched) for anything it cannot serve, including + // dense-backed batches, so the loop below is still reached unchanged. + // Multi-token steps (batched prefill, speculative / MTP verify) never + // enter it: `seq_len != 1`. + if seq_len == 1 + && mask.is_none() + && decode_context.is_some_and(|context| context.is_paged_decode()) + && let Some(attn_out) = mlxcel_core::cache::paged_batch_decode_attention( + &q_batched, &k_batched, &v_batched, caches, self.scale, 0.0, + ) + { + let attn_out = mlxcel_core::transpose_axes(&attn_out, &[0, 2, 1, 3]); + return mlxcel_core::reshape( + &attn_out, + &[b as i32, seq_len, self.num_heads * self.head_dim], + ); + } + let paged_decode = decode_context.and_then(|context| { if seq_len != 1 || mask.is_some() || !context.is_paged_decode() { return None; @@ -553,12 +599,12 @@ impl Attention { if caches.iter().any(|cache| cache.mode != KVCacheMode::Fp16) { return None; } - // Pool-backed caches (scheduler paged decode, #121) keep no dense - // `keys`/`values` buffers for the native paged kernel to read. - // Route them through the per-sequence `update_and_fetch` loop below, - // whose transparent pool intercept writes new K/V into the shared - // `PagedBlockPool` (`write_prefill`) and gathers the visible window - // back (`gather_visible`) — the #152-validated single-stream path. + // Pool-backed caches keep no dense `keys`/`values` buffers for the + // dense-compat paged kernel to read. Since #899 they are normally + // served by `paged_batch_decode_attention` above; reaching here + // means that path declined, so fall through to the per-sequence + // `update_and_fetch` loop below, whose pool intercept writes into + // the shared `PagedBlockPool` and gathers the visible window back. if caches.iter().any(|cache| cache.is_paged_backed()) { return None; } diff --git a/src/models/qwen3.rs b/src/models/qwen3.rs index 6c1e2d9f3..fe8b8c376 100644 --- a/src/models/qwen3.rs +++ b/src/models/qwen3.rs @@ -175,7 +175,30 @@ impl Attention { mlxcel_core::cache::turbo::sparse_v::turbo4_delegated_compressed_attention_enabled(); let use_turbo4_dequant_sdpa = mlxcel_core::cache::turbo::sparse_v::turbo4_dequant_sdpa_enabled(); - let attn_out = if l == 1 + // Single-sequence pooled paged decode (#899). The server takes this + // path whenever the active batch is one request: `decode_single_step` + // calls the single-sequence `forward`, never `forward_split_attention`. + // Without this branch the two single-sequence scenarios in the issue's + // benchmark matrix (16K and 32K) could never reach the fused kernel, and + // neither could any moment in a batched run where only one request is + // still decoding. The batch-1 launch is exactly the `[1, H, 1, D]` shape + // the whole-batch entry point already serves. + let pooled_single = if l == 1 && mask.is_none() && cache.is_paged_backed() { + mlxcel_core::cache::paged_batch_decode_attention( + &q, + &k, + &v, + &mut [&mut *cache], + self.scale, + 0.0, + ) + } else { + None + }; + + let attn_out = if let Some(out) = pooled_single { + out + } else if l == 1 && use_turbo4_asym_dequant_sdpa && cache.turbo4_asym_dequant_sdpa_available() { @@ -282,6 +305,29 @@ impl Attention { &metadata.rope_offsets, ); + // Pool-backed paged decode (#899): the whole batch in one fused launch. + // This is the production path for `--decode-storage paged`; it appends + // every sequence's new K/V to the shared `PagedBlockPool` and runs + // attention once over the batch, replacing the per-sequence + // `update_and_fetch` + `gather_visible` + SDPA loop below. It declines + // (leaving the pool untouched) for anything it cannot serve, including + // dense-backed batches, so the loop below is still reached unchanged. + // Multi-token steps (batched prefill, speculative / MTP verify) never + // enter it: `seq_len != 1`. + if seq_len == 1 + && mask.is_none() + && decode_context.is_some_and(|context| context.is_paged_decode()) + && let Some(attn_out) = mlxcel_core::cache::paged_batch_decode_attention( + &q_batched, &k_batched, &v_batched, caches, self.scale, 0.0, + ) + { + let attn_out = mlxcel_core::transpose_axes(&attn_out, &[0, 2, 1, 3]); + return mlxcel_core::reshape( + &attn_out, + &[b as i32, seq_len, self.num_heads * self.head_dim], + ); + } + let paged_decode = decode_context.and_then(|context| { if seq_len != 1 || mask.is_some() || !context.is_paged_decode() { return None; @@ -289,12 +335,12 @@ impl Attention { if caches.iter().any(|cache| cache.mode != KVCacheMode::Fp16) { return None; } - // Pool-backed caches (scheduler paged decode, #121) keep no dense - // `keys`/`values` buffers for the native paged kernel to read. - // Route them through the per-sequence `update_and_fetch` loop below, - // whose transparent pool intercept writes new K/V into the shared - // `PagedBlockPool` (`write_prefill`) and gathers the visible window - // back (`gather_visible`) — the #152-validated single-stream path. + // Pool-backed caches keep no dense `keys`/`values` buffers for the + // dense-compat paged kernel to read. Since #899 they are normally + // served by `paged_batch_decode_attention` above; reaching here + // means that path declined, so fall through to the per-sequence + // `update_and_fetch` loop below, whose pool intercept writes into + // the shared `PagedBlockPool` and gathers the visible window back. if caches.iter().any(|cache| cache.is_paged_backed()) { return None; } diff --git a/src/server/batch/scheduler.rs b/src/server/batch/scheduler.rs index 78a9b622c..7ba676966 100644 --- a/src/server/batch/scheduler.rs +++ b/src/server/batch/scheduler.rs @@ -1445,6 +1445,26 @@ impl BatchScheduler { self } + /// Install the paged KV slab size in blocks (issue #899). + /// + /// `Some(n)` makes each layer's pool storage one contiguous `n`-row slab, + /// which is the precondition for the fused paged-attention decode kernels: + /// they read one pool buffer per side, so a layer spread across several + /// slabs is declined and falls back to gather-then-SDPA. `None` leaves the + /// pool's own default (32 rows), which is the pre-#899 behaviour. + /// + /// Resolved from the operator's `--ctx-size` / `--parallel` and the KV + /// budget by [`crate::memory_estimate::resolve_paged_slab_blocks`] on the + /// worker thread. Applied to the pool when it is lazily created; a failure + /// (a pool that already has storage) is logged and ignored, because an + /// unsized slab costs performance, not correctness. + pub fn with_paged_slab_blocks(mut self, slab_blocks: Option) -> Self { + if let Err(reason) = self.cache_pool.set_paged_slab_blocks(slab_blocks) { + tracing::warn!("could not install the paged KV slab size: {reason}"); + } + self + } + /// Attach the resolved speculative-decoding dispatch. /// /// Default (constructed by [`Self::with_config`]) is diff --git a/src/server/model_worker.rs b/src/server/model_worker.rs index f0ae3f56c..4803d51e2 100644 --- a/src/server/model_worker.rs +++ b/src/server/model_worker.rs @@ -521,6 +521,26 @@ pub(crate) fn spawn_model_worker_with_batch_config( sched_config.kv_cache_budget, ); + // #899: size the paged pool's slab so a layer's rows stay in one + // contiguous buffer, which is what the fused decode kernels + // require. Derived from the same configuration the KV budget was, + // and clamped by it. `None` leaves the pool default alone. + let paged_slab_blocks = crate::memory_estimate::resolve_paged_slab_blocks( + &model_path, + model.num_layers(), + crate::server::batch::scheduler::DEFAULT_PAGED_BLOCK_SIZE, + effective_max_batch_size.max(1) as u64, + sched_config.max_kv_size.unwrap_or(0) as u64, + false, + paged_block_budget, + ); + if let Some(blocks) = paged_slab_blocks { + tracing::info!( + "Paged KV slab size: {blocks} blocks per layer \ + (fused decode serves a layer only while its rows fit one slab)" + ); + } + let mut scheduler = super::super::batch::BatchScheduler::with_config( model, tokenizer, @@ -549,6 +569,8 @@ pub(crate) fn spawn_model_worker_with_batch_config( .with_max_kv_size(sched_config.max_kv_size) // install the resolved paged KV block budget (epic #116 #122 b3). .with_paged_block_budget(paged_block_budget) + // install the resolved paged KV slab size (#899). + .with_paged_slab_blocks(paged_slab_blocks) // experimental VLM prompt-prefix cache sharing (#124 step c). .with_vlm_prefix_cache(sched_config.enable_vlm_prefix_cache) // attach the resolved speculative dispatch so the diff --git a/tests/paged_decode_v2_greedy_parity.rs b/tests/paged_decode_v2_greedy_parity.rs new file mode 100644 index 000000000..8da7bf627 --- /dev/null +++ b/tests/paged_decode_v2_greedy_parity.rs @@ -0,0 +1,342 @@ +// Copyright 2025-2026 Lablup Inc. and Jeongkyu Shin +// +// 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. + +//! Greedy token parity for the production paged decode v2 path (issue #899). +//! +//! This is the regression guard the issue makes mandatory: on a pinned prompt +//! set, batch 4 with mixed lengths and long prompts, the greedy token stream +//! produced by the fused v2 decode must match the gather-then-SDPA baseline +//! exactly. +//! +//! # How both paths run in one process +//! +//! `MLXCEL_PAGED_ATTENTION_NATIVE` is read once per process behind a +//! `OnceLock`, so it cannot be toggled between two runs of the same test +//! binary. The test therefore selects the path structurally instead, using the +//! one property the fused kernels genuinely require: **every physical row of a +//! layer must live in one contiguous pool buffer**. +//! +//! - The **v2 arm** sets `CachePool::set_paged_slab_blocks` to a slab large +//! enough for the whole batch, which is exactly what +//! `resolve_paged_slab_blocks` does on the server. The fused path is +//! eligible and, above the dispatch token floor, taken. +//! - The **gather arm** leaves the pool on its 32-row default, so every layer +//! is multi-slab, the fused path declines, and +//! `paged_batch_decode_attention` answers from `gather_visible` + SDPA. That +//! is byte-for-byte the pre-#899 production path. +//! +//! Both arms run identical scheduler wiring, identical prompts, and identical +//! greedy sampling, so a token difference can only be the kernel. The test +//! asserts afterwards that the v2 arm really did build a plan, so a silent +//! fallback cannot pass as parity. +//! +//! # Choosing a checkpoint +//! +//! `MLXCEL_PARITY_MODEL` names a model directory (absolute, or relative to the +//! repository root). Without it, each family case looks for its default +//! directory under `models/` and soft-skips when absent, so the matrix can be +//! run family by family as checkpoints arrive. +//! +//! # Running +//! +//! ```text +//! # One explicit checkpoint (any pool-backed family): +//! MLXCEL_PARITY_MODEL=$HOME/.cache/mlxcel/models/Qwen2.5-0.5B-Instruct-4bit \ +//! cargo test --release --test paged_decode_v2_greedy_parity \ +//! --features metal,accelerate -- --ignored --nocapture +//! +//! # The default matrix (each case skips when its directory is absent): +//! cargo test --release --test paged_decode_v2_greedy_parity \ +//! --features metal,accelerate -- --ignored --nocapture +//! ``` +//! +//! `MLXCEL_PARITY_PROMPT_LEN` (default 8192) sets the shortest prompt; the four +//! sequences use that length plus staggered offsets so the batch has mixed +//! lengths and the requests cross page boundaries on different steps. +//! `MLXCEL_PARITY_STEPS` (default 32) sets the number of compared decode steps. + +mod common; +use common::repo_model_dir; + +use std::path::PathBuf; + +use mlxcel::{DecodeBatchContext, LanguageModel, initialize_runtime, load_model}; +use mlxcel_core::cache::{CachePool, PagedKvLayout, SequenceStateLayout}; + +/// Paged block size, matching the scheduler's `DEFAULT_PAGED_BLOCK_SIZE`. +const BLOCK_SIZE: usize = 32; + +/// Sequences in the batch. Four is the documented v0.4 `--parallel` default and +/// the batch size the issue's benchmark matrix uses. +const BATCH: usize = 4; + +/// Token offsets added to the base prompt length, so the four sequences have +/// mixed lengths and open new pages on different decode steps. +const LENGTH_STAGGER: [usize; BATCH] = [0, 37, 96, 151]; + +/// Deterministic prompt vocabulary. Every id is < 128k, valid for the qwen and +/// llama vocabularies alike, so no tokenizer is needed and the two arms are fed +/// identical bytes. +const PROMPT_ALPHABET: &[i32] = &[ + 9707, 11, 358, 1079, 264, 4128, 1614, 13, 5209, 3291, 752, 911, 697, 7990, 13, 358, 1079, + 21815, 911, 1246, 498, 3705, 1293, 37597, 323, 1128, 8573, 979, 279, 6193, 27715, 13, +]; + +fn env_usize(key: &str, default: usize) -> usize { + std::env::var(key) + .ok() + .and_then(|raw| raw.trim().parse::().ok()) + .filter(|value| *value > 0) + .unwrap_or(default) +} + +fn prompt_len() -> usize { + env_usize("MLXCEL_PARITY_PROMPT_LEN", 8192) +} + +fn decode_steps() -> usize { + env_usize("MLXCEL_PARITY_STEPS", 32) +} + +/// The prompt for sequence `slot`: a deterministic repeating pattern, offset by +/// the slot so no two sequences share a prompt. +fn prompt_for(slot: usize, len: usize) -> Vec { + (0..len) + .map(|i| PROMPT_ALPHABET[(i + slot * 7) % PROMPT_ALPHABET.len()]) + .collect() +} + +/// Greedy argmax over the vocab at position `pos` of a `[batch, seq, vocab]` +/// logits tensor, for batch row `row`. +fn greedy_token(logits: &mlxcel_core::MlxArray, row: i32, pos: i32) -> i32 { + let shape = mlxcel_core::array_shape(logits); + let vocab = shape[2]; + let at_pos = mlxcel_core::slice(logits, &[row, pos, 0], &[row + 1, pos + 1, vocab]); + let flat = mlxcel_core::reshape(&at_pos, &[vocab]); + mlxcel_core::eval(&flat); + mlxcel_core::item_i32(&mlxcel_core::argmax_last_axis(&flat)) +} + +/// The paged sequence-state layout the scheduler builds for the default Fp16 KV +/// mode. +fn scheduler_paged_layout(num_layers: usize) -> SequenceStateLayout { + SequenceStateLayout::paged_kv_cache( + PagedKvLayout::uniform(num_layers, BLOCK_SIZE, BLOCK_SIZE).expect("valid paged layout"), + ) +} + +/// Slab size that holds the whole batch's rows for one layer, which is what +/// makes the fused path eligible. Mirrors +/// `crate::memory_estimate::resolve_paged_slab_blocks` at this batch and +/// context. +fn slab_blocks_for(prompt: usize, steps: usize) -> usize { + let longest = prompt + LENGTH_STAGGER[BATCH - 1] + steps + 1; + longest.div_ceil(BLOCK_SIZE) * BATCH +} + +/// Prefill every sequence, then run `steps` batched greedy decode steps through +/// the paged batched path, returning one token stream per sequence. +/// +/// `slab_blocks` is the only difference between the two arms: `Some(n)` makes +/// each layer single-slab and the fused path eligible, `None` leaves the pool +/// default so it declines and the gather fallback answers. +fn run_batched( + model: &mlxcel::LoadedModel, + slab_blocks: Option, + prompt: usize, + steps: usize, +) -> (Vec>, bool) { + let num_layers = model.num_layers(); + let mut pool = CachePool::new(BATCH + 2); + pool.set_paged_slab_blocks(slab_blocks) + .expect("a fresh pool has no storage yet"); + + let ids: Vec<_> = (0..BATCH) + .map(|slot| { + pool.allocate_with_layout(model, Some(scheduler_paged_layout(num_layers))) + .unwrap_or_else(|e| panic!("paged allocate {slot}: {e}")) + }) + .collect(); + + // Prefill each sequence individually, as `execute_full_prefill` does. + let mut next = vec![0i32; BATCH]; + for (slot, &id) in ids.iter().enumerate() { + let len = prompt + LENGTH_STAGGER[slot]; + let tokens = prompt_for(slot, len); + let input = mlxcel_core::from_slice_i32(&tokens, &[1, len as i32]); + let mask = mlxcel_core::utils::create_causal_mask(len as i32, 0); + let caches = pool.get_caches_mut(id).expect("caches"); + let logits = model.forward(&input, caches, Some(&mask)); + mlxcel_core::eval(&logits); + next[slot] = greedy_token(&logits, 0, len as i32 - 1); + } + + // Batched greedy decode, exactly the dispatch `execute_batched_decode` + // performs for the paged storage backend. + let context = DecodeBatchContext::paged_with_native(BLOCK_SIZE as i32, true); + let mut streams: Vec> = (0..BATCH).map(|_| Vec::with_capacity(steps)).collect(); + for _ in 0..steps { + for (slot, stream) in streams.iter_mut().enumerate() { + stream.push(next[slot]); + } + let input = mlxcel_core::from_slice_i32(&next, &[BATCH as i32, 1]); + let logits = { + let mut batch_caches = pool.get_batch_caches_mut(&ids).expect("batch caches"); + model.forward_batched_with_context_and_ids( + &input, + Some(&ids), + &mut batch_caches, + None, + Some(&context), + ) + }; + mlxcel_core::eval(&logits); + for (slot, entry) in next.iter_mut().enumerate() { + *entry = greedy_token(&logits, slot as i32, 0); + } + } + + // Did the fused path actually run? The plan cache is per pool, so this is + // immune to any other test in the binary. + let used_v2 = pool + .paged_pool_ref() + .is_some_and(|paged| paged.decode_plan_cache_stats().plan_rebuilds > 0); + (streams, used_v2) +} + +/// Resolve the checkpoint for a case: `MLXCEL_PARITY_MODEL` wins, then the +/// family's default directory under `models/`, then the downloader's cache at +/// `~/.cache/mlxcel/models/`. +fn resolve_model_dir(default_dir: &str, hf_repo: &str) -> Option { + if let Some(raw) = std::env::var_os("MLXCEL_PARITY_MODEL") { + let path = PathBuf::from(raw); + let path = if path.is_absolute() { + path + } else { + PathBuf::from(env!("CARGO_MANIFEST_DIR")).join(path) + }; + if path.exists() { + return Some(path); + } + eprintln!( + "MLXCEL_PARITY_MODEL points at {}, which does not exist", + path.display() + ); + return None; + } + let path = repo_model_dir(default_dir); + if path.exists() { + return Some(path); + } + // Fall back to the downloader's cache, which is where `mlxcel download` + // puts a checkpoint when no `--local-dir` was given. + hf_cache_dir(hf_repo) +} + +/// `~/.cache/mlxcel/models//` for a `org/name` repository id, when +/// it exists. +fn hf_cache_dir(hf_repo: &str) -> Option { + let home = std::env::var_os("HOME")?; + let path = PathBuf::from(home) + .join(".cache") + .join("mlxcel") + .join("models") + .join(hf_repo); + path.exists().then_some(path) +} + +/// Run the parity comparison for one checkpoint, soft-skipping when it is +/// absent. +fn assert_parity_for(default_dir: &str, fetch_repo: &str) { + let Some(model_dir) = resolve_model_dir(default_dir, fetch_repo) else { + eprintln!( + "Skipping {default_dir}: checkpoint not found.\n\ + Fetch with: ./target/release/mlxcel download {fetch_repo}\n\ + Or point the test at one: MLXCEL_PARITY_MODEL=/path/to/model" + ); + return; + }; + let (model, _tokenizer) = + load_model(&model_dir).unwrap_or_else(|e| panic!("load {}: {e:?}", model_dir.display())); + if !model.supports_paged_decode_backend() { + eprintln!( + "Skipping {}: the family does not support the paged decode backend", + model_dir.display() + ); + return; + } + + let prompt = prompt_len(); + let steps = decode_steps(); + let slab = slab_blocks_for(prompt, steps); + eprintln!( + "\n=== paged decode v2 greedy parity: {} ===\n\ + batch {BATCH}, prompts {prompt}{LENGTH_STAGGER:?}, {steps} decode steps, \ + slab {slab} blocks", + model_dir.display() + ); + + let (gather_streams, gather_used_v2) = run_batched(&model, None, prompt, steps); + assert!( + !gather_used_v2, + "the gather arm must not have reached the fused path; the slab default changed" + ); + + let (v2_streams, v2_used_v2) = run_batched(&model, Some(slab), prompt, steps); + assert!( + v2_used_v2, + "the v2 arm never built a decode plan, so it silently fell back to gather; \ + parity would be vacuous. Check the slab size ({slab} blocks) and the \ + dispatch token floor." + ); + + for slot in 0..BATCH { + assert_eq!( + v2_streams[slot], gather_streams[slot], + "sequence {slot} diverged.\ngather: {:?}\nv2: {:?}", + gather_streams[slot], v2_streams[slot] + ); + } + eprintln!( + "OK: {BATCH} sequences x {steps} greedy decode steps identical between the fused v2 \ + path and the gather baseline." + ); +} + +#[test] +#[ignore = "loads a real checkpoint and runs long real GPU forwards; run with --ignored"] +fn paged_decode_v2_greedy_parity_qwen3() { + let _runtime = initialize_runtime(); + assert_parity_for("qwen3-0.6b-4bit", "mlx-community/Qwen3-0.6B-4bit"); +} + +#[test] +#[ignore = "loads a real checkpoint and runs long real GPU forwards; run with --ignored"] +fn paged_decode_v2_greedy_parity_llama3() { + let _runtime = initialize_runtime(); + assert_parity_for( + "llama-3.2-1b-4bit", + "mlx-community/Llama-3.2-1B-Instruct-4bit", + ); +} + +#[test] +#[ignore = "loads a real checkpoint and runs long real GPU forwards; run with --ignored"] +fn paged_decode_v2_greedy_parity_qwen25() { + let _runtime = initialize_runtime(); + assert_parity_for( + "qwen2.5-0.5b-instruct-4bit", + "mlx-community/Qwen2.5-0.5B-Instruct-4bit", + ); +}