From d3715e81adb56e28295c0a105a98df957e1b94d2 Mon Sep 17 00:00:00 2001 From: Nyvo <75425811+Nyvo-io@users.noreply.github.com> Date: Sat, 18 Jul 2026 21:44:48 +0800 Subject: [PATCH] perf(qwen): reuse prefill plans for uncompiled GQA decode Signed-off-by: Nyvo <75425811+Nyvo-io@users.noreply.github.com> --- docs/index.md | 1 + .../kernels/reusable-prefill-plan.md | 259 ++++++++++ openinfer-core/src/ops.rs | 1 + openinfer-core/src/ops/paged_plan.rs | 143 +++++- openinfer-kernels/src/ops/attention.rs | 451 ++++++++++++++++-- openinfer-qwen3/src/executor.rs | 30 ++ openinfer-qwen3/src/unified_forward.rs | 74 ++- openinfer-qwen3/src/weights.rs | 79 ++- openinfer-qwen35-4b/src/batch_decode.rs | 56 ++- openinfer-qwen35-4b/src/batch_decode_graph.rs | 37 +- openinfer-qwen35-4b/src/prefill_buffers.rs | 108 ++++- openinfer-qwen35-4b/src/recurrent_state.rs | 59 ++- openinfer-qwen35-4b/src/tp_executor.rs | 4 +- openinfer-qwen35-4b/src/weights.rs | 75 ++- 14 files changed, 1236 insertions(+), 141 deletions(-) create mode 100644 docs/subsystems/kernels/reusable-prefill-plan.md diff --git a/docs/index.md b/docs/index.md index ade8d747e..6844ef97d 100644 --- a/docs/index.md +++ b/docs/index.md @@ -157,6 +157,7 @@ Organized by domain (model line / subsystem / playbook / lesson) instead of by l | `subsystems/kernels/openinfer-kernels-boundary.md` | Architecture decision: reusable frontend/runtime/data-plane layers plus per-model engines; `openinfer-kernels` keeps shared MoE/MLA substrate (`moe`: DeepEP/DeepGEMM/FlashMLA) separate from model-local surfaces such as the narrow GLM5.2 DeepGEMM/FlashMLA wrappers. | | `subsystems/kernels/build-rs-submodule-init.md` | `openinfer-kernels/build.rs` initializes missing git submodules automatically for first-time builds before checking vendored third-party kernel headers. | | `subsystems/kernels/kernel-op-reports.md` | Qwen3 kernel/report tooling is feature-gated: `qwen3_kernel_report` covers per-op kernel reports, and `qwen3_model_report` emits runtime-traced eager-DAG decode operator rollups with TensorSpec `KernelCall`s, latency stats, tables, and Graphviz DOT; measured FA2 `CTA_TILE_Q=64` prefill default in place. | +| `subsystems/kernels/reusable-prefill-plan.md` | Issue #711 / PR #714 reuses uncompiled-GQA prefill metadata with a scheduler-derived logical page-table bound and one checked allocation/accounting layout; Qwen3-14B Nsight A/B removes 22 allocation/free calls per step, while pointer stability, shared-prefix mixed load, both production golden suites, and Qwen3.5 scheduler E2E pass. | | `subsystems/kernels/typed-forward-pipeline.md` | Reusable typed tensor pipeline macro in `openinfer-kernels` so model crates can express common `typed_ops` chains without model-specific wrapper macros. | | `subsystems/kernels/tvm-ffi-mvp.md` | Optional `tvm-ffi-triton-cubin` bridge in `openinfer-kernels` plus a packed TVM wrapper for the Qwen3.5 GDR solve Triton AOT CUBIN launcher. | diff --git a/docs/subsystems/kernels/reusable-prefill-plan.md b/docs/subsystems/kernels/reusable-prefill-plan.md new file mode 100644 index 000000000..ede6e9f26 --- /dev/null +++ b/docs/subsystems/kernels/reusable-prefill-plan.md @@ -0,0 +1,259 @@ +# Reusable PrefillPagedPlan capacity and validation + +> **TL;DR:** Issue #711 / PR #714 removes per-step uncompiled-GQA plan allocation for Qwen3-14B and Qwen3.5-27B. Allocation and accounting share one checked layout, the scheduler-derived logical page-table bound covers shared prefixes, device pointers remain stable across changing shapes, and both production checkpoints pass their runtime gates. Five Qwen3-14B Nsight A/B pairs remove 22 allocation/free calls per decode step with a paired-median 19.27 us/step reduction in those CUDA API calls and no change in metadata uploads or kernels. +> +> **Last touched:** 2026-07 + +## Preparation + +- **Read**: + - `docs/index.md` - routed the cross-model attention-plan work to the kernels subsystem. + - `docs/models/qwen3/serving-perf-5090.md` - unified steps place prefill and decode rows in one `PrefillPagedPlan`. + - `docs/models/qwen3/prefix-cache.md` - prefix hits reuse physical KV blocks while each request retains its own logical page list. + - `docs/models/qwen3/model-crate.md` - Qwen3 owns scheduler and executor bounds while the reusable plan implementation lives in the shared kernel layer. + - `docs/models/qwen35/model-crate.md` - Qwen3.5 owns a separate decode graph and startup memory budget. + - `docs/subsystems/scheduler/scheduler.md` - admission bounds active plus prefilling requests by decode capacity and rejects sequences beyond the model context window. + - `docs/conventions/bench-regression.md` - performance evidence needs repeated runs; p99 is inspected rather than inferred from a five-iteration sample. + - `docs/conventions/coding-style.md` - tricky capacity invariants merit focused unit coverage, while runtime behavior should use an integration gate. +- **Relevant history**: + - `docs/models/qwen3/serving-perf-5090.md` records the original unified-attention fusion and why decode rows carry full logical page lists. + - PR #714's first RTX 5090 run reduced `cuMemAllocAsync` and `cuMemFreeAsync` calls from 708 to 642, but review found that physical pool size is not a valid logical page-table bound under shared prefixes. +- **Plan**: + 1. Add a checked `PrefillPagedPlan` capacity value derived from maximum step tokens, request rows, model context, page size, and GQA group size. + 2. Use `max_batch * ceil(max_context / page_size)` logical page references for Qwen3 and Qwen3.5, preserving startup memory accounting. + 3. Pass the same capacity and allocation layout through profiling, byte accounting, allocation, and serving ownership. + 4. Add CPU capacity coverage and a GPU gate that changes batch/context/page shapes while checking all eleven device pointers. + 5. Repeat the targeted tail benchmark and collect production Qwen3-14B allocation/API traces. + 6. Run Qwen3-14B golden and prefix-cache mixed gates plus Qwen3.5-27B short/long golden and scheduler E2E. +- **Acceptance risks**: + - Prefix sharing can make logical page references exceed unique physical pages. The capacity therefore uses the scheduler/model logical bound, not KV pool size. + - Persistent plan bytes compete with KV capacity. Every reserve is checked for overflow and charged before KV allocation. + - An allocation-count win could hide changed work. The production traces pin identical decode steps, metadata uploads, and kernel counts. + +## Execution Log + +### 1. Derive the reachable logical capacity + +Qwen3 admission subtracts both `active.len()` and `prefilling.len()` from `max_decode_batch_size`. Qwen3.5 likewise reserves prefilling promotion slots before admitting more work. Both schedulers reject request lifetimes beyond `max_position_embeddings`. + +The safe logical page-reference bound is therefore: + +```text +max_batch * ceil(max_position_embeddings / page_size) +``` + +This includes the worst shared-prefix case: every request keeps its own logical page list even when those lists refer to the same physical prefix pages. + +### 2. Share layout, allocation, and accounting + +- `PrefillPagedPlanCapacity` validates token, logical-page, batch, tile, and i32 ABI limits once. +- `PrefillPagedPlanLayout` is the single source for the eleven allocation lengths and their exact byte footprint. +- Qwen3 passes one capacity through temporary profile allocation, profile-byte subtraction, startup accounting, and the serving lane. +- Qwen3.5 charges the same capacity before KV allocation and passes it to `BatchDecodeGraphState`. +- Qwen3 and Qwen3.5 reserve arithmetic uses checked multiplication/addition through one checked total reused by admission and subtraction. +- Runtime updates return dimension-specific errors; they never fall back to the allocating constructor. + +The CPU layout tests cover the shared-prefix logical bound, invalid capacity, exact allocation layout, and i32 overflow. The GPU gate updates different page, token, and batch shapes, confirms all eleven device pointers remain unchanged, and checks page/token/batch/tile capacity errors. + +### 3. Validate shared-prefix mixed execution + +The Qwen3 group-size-5 fixture ran the complete `prefix_cache_behavior` integration test. It covered repeated warm prefixes, cold plus warm multi-request batching, the full-block prefix cap, and warm prefill plus active decode in one unified step. The test passed `1/1`. + +The final production rerun used Qwen3-14B itself and passed the same gate. Its mixed-batch rows were bit-identical to the cold reference; the unified-plan warm row stayed within the gate (`mean 0.0335`, `max 0.0592`). + +### 4. Repeat the targeted tail comparison + +The original five-iteration batch-16 result moved p99 from `1.384` to `1.737 ms`. The longer rerun fixed context at 512, batch at 16, decode at 256 steps, warmup at 32 steps, and distinct prompts at 16. CUDA Graphs were disabled. Five interleaved base/PR pairs each ran 100 iterations, producing 356,800 TPOT samples per run. + +The retained JSON records base `41b77566e48d46673cc0d8c7f279f6dbe275f4f5` and PR implementation `59a74963c5bf845befd55b8bd92b55681f1e72c1`. The collection loop was: + +```bash +run_decode() { + revision=$1 + binary=$2 + seed=$3 + "$binary" \ + --model-path /root/autodl-tmp/models/Qwen3-Group5-1L-fixture \ + --cuda-graph false \ + --format json \ + --label "${revision}-100iter-s${seed}" \ + --out "/root/autodl-tmp/pr714-100iter/${revision}-s${seed}.json" \ + decode \ + --ctxs 512 \ + --batches 16 \ + --decode-steps 256 \ + --warmup-steps 32 \ + --distinct-prompts 16 \ + --iters 100 \ + --seed "$seed" +} + +BASE_BIN=/root/autodl-tmp/target-base/release/bench_serving +PR_BIN=/root/autodl-tmp/target-pr714/release/bench_serving + +run_decode base "$BASE_BIN" 47; run_decode pr "$PR_BIN" 47 +run_decode pr "$PR_BIN" 48; run_decode base "$BASE_BIN" 48 +run_decode base "$BASE_BIN" 49; run_decode pr "$PR_BIN" 49 +run_decode pr "$PR_BIN" 50; run_decode base "$BASE_BIN" 50 +run_decode base "$BASE_BIN" 51; run_decode pr "$PR_BIN" 51 +``` + +| Seed | Average TPOT base -> PR | Delta | p99 TPOT base -> PR | Delta | Throughput base -> PR | Delta | +| ---: | ---: | ---: | ---: | ---: | ---: | ---: | +| 47 | 0.867688 -> 0.916159 ms | +5.59% | 1.013080 -> 1.964354 ms | +93.90% | 18,546.04 -> 18,153.58 tok/s | -2.12% | +| 48 | 0.864537 -> 0.852398 ms | -1.40% | 1.021063 -> 1.005956 ms | -1.48% | 18,581.68 -> 18,765.62 tok/s | +0.99% | +| 49 | 0.923673 -> 0.859485 ms | -6.95% | 1.867574 -> 1.003935 ms | -46.24% | 17,911.38 -> 18,711.22 tok/s | +4.47% | +| 50 | 0.917316 -> 0.858558 ms | -6.41% | 1.833100 -> 0.994906 ms | -45.73% | 18,023.70 -> 18,709.42 tok/s | +3.80% | +| 51 | 0.866872 -> 0.863014 ms | -0.45% | 1.033162 -> 1.061248 ms | +2.72% | 18,720.28 -> 18,778.51 tok/s | +0.31% | + +| Metric | Base median | PR median | Paired-delta median | +| --- | ---: | ---: | ---: | +| Average TPOT | 0.867688 ms | 0.859485 ms | -1.40% | +| p99 TPOT | 1.033162 ms | 1.005956 ms | -1.48% | +| Decode throughput | 18,546.04 tok/s | 18,711.22 tok/s | +0.99% | + +The isolated p99 spikes occur on both revisions. The repeated, order-interleaved medians do not reproduce the original short-run regression. + +### 5. Measure production allocation and plan API time + +The production A/B used Qwen3-14B group 5 on one NVIDIA RTX PRO 6000 Blackwell Server Edition (97,887 MiB), driver `580.119.02`, CUDA 12.8, `OPENINFER_CUDA_SM=120`, Rust nightly 2026-07-10, and Nsight Systems exporter 2024.6.2. The model revision was `40c069824f4251a91eefaf281ebe4c544efd3e18`; all eight weight shards matched the official Hugging Face SHA-256 values. + +The retained binaries are independently identifiable: + +```text +base a62416f12c6384fae2f4ce81474fb2f50a772df219808caf7462840b8432373a +PR fa60e425a59fef53066aa3ef3dc9fdac5b13a8390127f1b394d3c573376dee5a +``` + +Each trace captured exactly 32 steady eager decode steps at context 512 and batch 16. This is the exact pair-1 command preserved in the trace metadata; pairs 2-5 changed only the output pair number and selected base/PR binary: + +```bash +/opt/nvidia/nsight-compute/2025.1.1/host/target-linux-x64/nsys profile \ + --trace=cuda \ + --sample=none \ + --cpuctxsw=none \ + --capture-range=cudaProfilerApi \ + --capture-range-end=stop \ + --force-overwrite=true \ + --export=sqlite \ + -o /root/autodl-tmp/pr714-nsys-qwen3-14b/pair1-base \ + /root/autodl-tmp/pr714-ab/bin/qwen3_decode_context-base \ + --mode profile \ + --model-path /root/autodl-tmp/models/Qwen3-14B \ + --contexts 512 \ + --batch-size 16 \ + --profile-steps 32 \ + --capture-range \ + --disable-cuda-graph +``` + +CUDA Driver API counts and CPU durations were queried directly from the exported CUPTI table: + +```sql +SELECT + s.value AS api, + COUNT(*) AS calls, + ROUND(SUM(r.end - r.start) / 1000000.0, 6) AS cpu_api_ms +FROM CUPTI_ACTIVITY_KIND_RUNTIME AS r +JOIN StringIds AS s ON s.id = r.nameId +WHERE s.value IN ( + 'cuMemAllocAsync', + 'cuMemFreeAsync', + 'cuMemcpyHtoDAsync_v2', + 'cuEventSynchronize' +) +GROUP BY s.value +ORDER BY s.value; +``` + +Kernel and device-copy work were checked separately: + +```sql +SELECT + (SELECT COUNT(*) FROM CUPTI_ACTIVITY_KIND_KERNEL) AS kernels, + (SELECT COUNT(*) FROM CUPTI_ACTIVITY_KIND_MEMCPY) AS device_copies, + ROUND((SELECT SUM(end - start) FROM CUPTI_ACTIVITY_KIND_MEMCPY) / 1000000.0, 6) + AS device_copy_ms; +``` + +| Pair | Alloc+free calls base -> PR | Alloc+free API ms base -> PR | Saved per step | Alloc+free+H2D API ms base -> PR | Profile TPOT base -> PR | +| ---: | ---: | ---: | ---: | ---: | ---: | +| 1 | 1,728 -> 1,024 | 2.448186 -> 2.037476 | 12.83 us | 4.567320 -> 4.402517 | 24.5670 -> 24.6304 ms | +| 2 | 1,728 -> 1,024 | 2.887641 -> 2.253055 | 19.83 us | 5.270461 -> 4.902352 | 24.6975 -> 24.5400 ms | +| 3 | 1,728 -> 1,024 | 2.387319 -> 1.770831 | 19.27 us | 4.416813 -> 3.958116 | 24.6127 -> 24.5858 ms | +| 4 | 1,728 -> 1,024 | 2.750674 -> 2.764723 | -0.44 us | 5.269879 -> 6.438497 | 24.1020 -> 24.3742 ms | +| 5 | 1,728 -> 1,024 | 4.199736 -> 2.456119 | 54.49 us | 8.314309 -> 5.859004 | 24.8021 -> 24.5328 ms | + +Each API individually changes from 864 to 512 calls: `-352`, or exactly `-11` calls per step, for both allocation and free. The paired-median allocation/free API saving is `0.616488 ms / 32 = 19.265 us/step`. Including the unchanged H2D API calls, paired-median host API saving is `0.368109 ms / 32 = 11.503 us/step`. + +Every trace contains 448 `cuMemcpyHtoDAsync_v2` calls, 480 device-copy activities, 24,512 kernels, and 32 `cuEventSynchronize` calls on both revisions. Device-copy time stays within `0.172-0.181 ms` per trace. The optimization therefore removes allocation/free and pointer churn while retaining metadata uploads and computational work. Profile TPOT has a `24.6127 -> 24.5400 ms` median and a `-0.11%` paired-delta median, so the production claim is reduced plan-preparation API work, not a material TPOT speedup. + +### 6. Run production correctness and scheduler gates + +The production-gate source was rebased on `origin/main` `90440c5` and synced byte-for-byte to the GPU host. The final submission was then rebased onto `cb41b07` (#745); `git range-diff` shows that this only narrows the Qwen3.5 capacity field from `pub(super)` to private, leaving the runtime diff unchanged. Qwen3-14B used revision `40c069824f4251a91eefaf281ebe4c544efd3e18`. Qwen3.5-27B used revision `fc05daec18b0a78c049392ed2e771dde82bdf654`; all eleven weight shards and its tokenizer matched the official Hugging Face SHA-256 values. + +Environment and commands: + +```bash +export CUDA_HOME=/usr/local/cuda-12.8 +export CUDA_PATH=/usr/local/cuda-12.8 +export OPENINFER_CUDA_SM=120 +export OPENINFER_NVCC_JOBS=2 +export OPENINFER_TRITON_PYTHON=/root/miniconda3/bin/python +export CARGO_TARGET_DIR=/root/autodl-tmp/target-pr714 +export RUSTUP_TOOLCHAIN=nightly-2026-07-10-x86_64-unknown-linux-gnu + +cargo fmt --all -- --check +cargo check --release --locked -p openinfer-server --all-targets +cargo clippy --release --locked \ + -p openinfer-server -p openinfer-qwen3 -p openinfer-core \ + -p openinfer-kernels -p openinfer-kv-cache -p openinfer-kv-offload \ + -p openinfer-sample -p openinfer-bench \ + --all-targets -- -D warnings +cargo check --release --locked -p openinfer-qwen35-4b \ + --all-targets --features qwen35-4b + +cargo test --release --locked -p openinfer-core --lib \ + paged_plan::tests -- --nocapture +cargo test --release --locked -p openinfer-kernels --lib \ + preallocated_footprint -- --nocapture +cargo test --release --locked -p openinfer-kernels --lib \ + preallocated_update_preserves_all_pointers_and_reports_capacities \ + -- --ignored --nocapture + +OPENINFER_TEST_MODEL_PATH=/root/autodl-tmp/models/Qwen3-14B \ + cargo test --release --locked -p openinfer-qwen3 \ + --test hf_golden_gate -- --nocapture +OPENINFER_TEST_MODEL_PATH=/root/autodl-tmp/models/Qwen3-14B \ + cargo test --release --locked -p openinfer-qwen3 \ + --test prefix_cache -- --nocapture + +export OPENINFER_TEST_MODEL_PATH=/root/autodl-tmp/models/Qwen3.5-27B +export OPENINFER_TEST_MODEL_REVISION=fc05daec18b0a78c049392ed2e771dde82bdf654 +cargo test --release --locked -p openinfer-qwen35-4b --features qwen35-4b \ + --test hf_golden_gate \ + pega_logprobs_match_hf_golden_within_qwen35_tolerance -- --nocapture +cargo test --release --locked -p openinfer-qwen35-4b --features qwen35-4b \ + --test hf_golden_gate \ + pega_logprobs_match_hf_long_golden_within_qwen35_tolerance -- --nocapture +cargo test --release --locked -p openinfer-qwen35-4b --features qwen35-4b \ + --test e2e_scheduler test_e2e_qwen35_scheduler -- --nocapture +``` + +Final receipts: + +| Gate | Result | Key receipt | +| --- | --- | --- | +| Capacity/layout CPU tests | pass, 4 tests | shared-prefix bound, invalid page size, exact layout, i32 overflow | +| GPU pointer/capacity gate | pass, 1 test | all eleven pointers stable across changing shapes; dimension errors checked | +| Qwen3-14B HF golden | pass, 1 test | sequential mean `0.0291`, p99 `0.1087`; cached replay mean `0.0281`, p99 `0.1046` | +| Qwen3-14B prefix-cache mixed | pass, 1 test | warm/cold mixed and unified-plan paths pass | +| Qwen3.5-27B short golden | pass, 1 test | sequential mean `0.0201`, p99 `0.0694`; batch-5 p99 `0.0651`; slot-compaction p99 `0.0755` | +| Qwen3.5-27B long golden | pass, 1 test | prompt lengths 4097/8192; mean `0.0218`, p99 `0.0627` | +| Qwen3.5-27B scheduler E2E | pass, 1 test | full TP1 request-flow/liveness suite, `60.47 s` | + +## Debrief + +- **Outcome**: Both production uncompiled-GQA owners reuse fixed-capacity plan storage. Capacity follows reachable logical request state, allocation and byte accounting share one layout, startup arithmetic is checked end to end, and runtime overflow fails with a dimension-specific error. +- **Performance**: Targeted repeated medians show no tail regression. Production Nsight traces remove exactly eleven allocations and eleven frees per decode step, preserve every metadata upload and kernel, and reduce the paired-median allocation/free API time by 19.27 us/step. +- **Correctness**: Shared-prefix mixed execution, actual GPU pointer stability, Qwen3-14B golden/prefix-cache, and Qwen3.5-27B short/long golden plus scheduler E2E all pass on the final rebased source. diff --git a/openinfer-core/src/ops.rs b/openinfer-core/src/ops.rs index 057df9678..3c5a28a95 100644 --- a/openinfer-core/src/ops.rs +++ b/openinfer-core/src/ops.rs @@ -81,6 +81,7 @@ pub use openinfer_kernels::ops::single_prefill_nhd_causal_into; pub use openinfer_kernels::ops::single_prefill_nhd_noncausal_into; pub use openinfer_kernels::ops::write_vec_into; pub use paged_plan::PrefillPagedPlan; +pub use paged_plan::PrefillPagedPlanCapacity; #[cfg(feature = "kernel-call-trace")] pub use traced::embedding_batch; #[cfg(feature = "kernel-call-trace")] diff --git a/openinfer-core/src/ops/paged_plan.rs b/openinfer-core/src/ops/paged_plan.rs index 5d1861f86..2991f8f19 100644 --- a/openinfer-core/src/ops/paged_plan.rs +++ b/openinfer-core/src/ops/paged_plan.rs @@ -6,11 +6,108 @@ use cudarc::driver::CudaSlice; use crate::kv_pool::KvDesc; use crate::tensor::DeviceContext; +/// Checked dimensions for a reusable [`PrefillPagedPlan`]. +/// +/// `max_page_indices` counts logical page-table entries across every request, +/// not unique physical KV pages. Prefix-cached requests may reference the same +/// physical page, so the logical bound is `max_batch * pages_per_request`. +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub struct PrefillPagedPlanCapacity { + total_tokens: usize, + page_indices: usize, + batch: usize, + tiles: usize, +} + +impl PrefillPagedPlanCapacity { + pub fn for_context( + max_total_tokens: usize, + max_batch: usize, + max_context_tokens: usize, + page_size: usize, + gqa_group_size: usize, + ) -> Result { + anyhow::ensure!( + max_total_tokens > 0, + "prefill plan token capacity must be positive" + ); + anyhow::ensure!( + max_batch > 0, + "prefill plan batch capacity must be positive" + ); + anyhow::ensure!( + max_context_tokens > 0, + "prefill plan context capacity must be positive" + ); + anyhow::ensure!(page_size > 0, "prefill plan page size must be positive"); + anyhow::ensure!( + gqa_group_size > 0, + "prefill plan GQA group size must be positive" + ); + + let max_pages_per_request = max_context_tokens.div_ceil(page_size); + let max_page_indices = max_batch + .checked_mul(max_pages_per_request) + .ok_or_else(|| anyhow::anyhow!("prefill plan logical page capacity overflow"))?; + let max_tiles = max_total_tokens + .checked_mul(gqa_group_size) + .ok_or_else(|| anyhow::anyhow!("prefill plan tile capacity overflow"))?; + let capacity = Self { + total_tokens: max_total_tokens, + page_indices: max_page_indices, + batch: max_batch, + tiles: max_tiles, + }; + capacity.preallocated_bytes()?; + Ok(capacity) + } + + pub fn preallocated_bytes(self) -> Result { + PrefillPagedPlan::preallocated_bytes( + self.total_tokens, + self.page_indices, + self.batch, + self.tiles, + ) + } + + pub fn max_total_tokens(self) -> usize { + self.total_tokens + } + + pub fn max_page_indices(self) -> usize { + self.page_indices + } + + pub fn max_batch(self) -> usize { + self.batch + } + + pub fn max_tiles(self) -> usize { + self.tiles + } +} + pub struct PrefillPagedPlan { inner: openinfer_kernels::ops::PrefillPagedPlan, } impl PrefillPagedPlan { + /// Exact device bytes reserved by [`Self::new_preallocated`]. + pub fn preallocated_bytes( + max_total_tokens: usize, + max_page_indices: usize, + max_batch: usize, + max_tiles: usize, + ) -> Result { + openinfer_kernels::ops::PrefillPagedPlan::preallocated_bytes( + max_total_tokens, + max_page_indices, + max_batch, + max_tiles, + ) + } + pub fn new( ctx: &DeviceContext, desc: &KvDesc<'_>, @@ -95,7 +192,7 @@ impl PrefillPagedPlan { pub fn new_preallocated( ctx: &DeviceContext, max_total_tokens: usize, - max_total_pages: usize, + max_page_indices: usize, max_batch: usize, max_tiles: usize, ) -> Result { @@ -103,13 +200,28 @@ impl PrefillPagedPlan { inner: openinfer_kernels::ops::PrefillPagedPlan::new_preallocated( ctx, max_total_tokens, - max_total_pages, + max_page_indices, max_batch, max_tiles, )?, }) } + /// Allocate from one checked capacity value so allocation and memory + /// accounting cannot silently derive different dimensions. + pub fn new_preallocated_for_capacity( + ctx: &DeviceContext, + capacity: PrefillPagedPlanCapacity, + ) -> Result { + Self::new_preallocated( + ctx, + capacity.total_tokens, + capacity.page_indices, + capacity.batch, + capacity.tiles, + ) + } + /// Refill a pre-allocated plan in place (no allocation, pointers unchanged). #[allow(clippy::too_many_arguments)] pub fn update_batch_with_cta_tile_q( @@ -179,3 +291,30 @@ impl Deref for PrefillPagedPlan { &self.inner } } + +#[cfg(test)] +mod tests { + use std::collections::HashSet; + + use super::PrefillPagedPlanCapacity; + + #[test] + fn logical_page_capacity_counts_shared_prefix_for_each_request() { + let capacity = PrefillPagedPlanCapacity::for_context(16, 3, 32, 16, 5) + .expect("valid reusable plan capacity"); + let page_lists = [vec![7, 8], vec![7, 8], vec![7, 8]]; + let logical_pages: usize = page_lists.iter().map(Vec::len).sum(); + let physical_pages: HashSet = page_lists.into_iter().flatten().collect(); + + assert_eq!(physical_pages.len(), 2); + assert_eq!(logical_pages, 6); + assert_eq!(capacity.max_page_indices(), logical_pages); + } + + #[test] + fn capacity_rejects_zero_page_size() { + let error = PrefillPagedPlanCapacity::for_context(1, 1, 1, 0, 1) + .expect_err("zero page size must be rejected"); + assert!(error.to_string().contains("page size")); + } +} diff --git a/openinfer-kernels/src/ops/attention.rs b/openinfer-kernels/src/ops/attention.rs index d8967b935..c6e0aee5a 100644 --- a/openinfer-kernels/src/ops/attention.rs +++ b/openinfer-kernels/src/ops/attention.rs @@ -40,7 +40,98 @@ pub struct PrefillPagedPlan { cta_tile_q: i32, } +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +struct PrefillPagedPlanLayout { + page_indices: usize, + page_indptr: usize, + last_page_len: usize, + batch_indices: usize, + positions: usize, + q_indptr: usize, + request_indices: usize, + qo_tile_indices: usize, + kv_tile_indices: usize, + kv_chunk_size: usize, + total_num_rows: usize, +} + +impl PrefillPagedPlanLayout { + fn new( + max_total_tokens: usize, + max_page_indices: usize, + max_batch: usize, + max_tiles: usize, + ) -> Result { + anyhow::ensure!( + i32::try_from(max_total_tokens).is_ok(), + "prefill plan max_total_tokens capacity exceeds i32: {max_total_tokens}" + ); + anyhow::ensure!( + i32::try_from(max_page_indices).is_ok(), + "prefill plan max_page_indices capacity exceeds i32: {max_page_indices}" + ); + anyhow::ensure!( + i32::try_from(max_batch).is_ok(), + "prefill plan max_batch capacity exceeds i32: {max_batch}" + ); + anyhow::ensure!( + i32::try_from(max_tiles).is_ok(), + "prefill plan max_tiles capacity exceeds i32: {max_tiles}" + ); + let max_batch_plus_one = max_batch + .checked_add(1) + .ok_or_else(|| anyhow::anyhow!("prefill plan max_batch capacity overflows"))?; + Ok(Self { + page_indices: max_page_indices, + page_indptr: max_batch_plus_one, + last_page_len: max_batch, + batch_indices: max_total_tokens, + positions: max_total_tokens, + q_indptr: max_batch_plus_one, + request_indices: max_tiles, + qo_tile_indices: max_tiles, + kv_tile_indices: max_tiles, + kv_chunk_size: max_batch, + total_num_rows: 1, + }) + } + + fn preallocated_bytes(self) -> Result { + let elements = [ + self.page_indices, + self.page_indptr, + self.last_page_len, + self.batch_indices, + self.positions, + self.q_indptr, + self.request_indices, + self.qo_tile_indices, + self.kv_tile_indices, + self.kv_chunk_size, + self.total_num_rows, + ] + .into_iter() + .try_fold(0usize, usize::checked_add) + .ok_or_else(|| anyhow::anyhow!("prefill plan dimensions overflow footprint"))?; + elements + .checked_mul(std::mem::size_of::()) + .ok_or_else(|| anyhow::anyhow!("prefill plan footprint overflows usize")) + } +} + impl PrefillPagedPlan { + /// Exact device bytes reserved by [`Self::new_preallocated`]. Every + /// metadata array stores 32-bit integers (including `total_num_rows`). + pub fn preallocated_bytes( + max_total_tokens: usize, + max_page_indices: usize, + max_batch: usize, + max_tiles: usize, + ) -> Result { + PrefillPagedPlanLayout::new(max_total_tokens, max_page_indices, max_batch, max_tiles)? + .preallocated_bytes() + } + pub fn page_indices_d(&self) -> &CudaSlice { &self.page_indices_d } @@ -212,22 +303,25 @@ impl PrefillPagedPlan { pub fn new_preallocated( ctx: &DeviceContext, max_total_tokens: usize, - max_total_pages: usize, + max_page_indices: usize, max_batch: usize, max_tiles: usize, ) -> Result { + let layout = + PrefillPagedPlanLayout::new(max_total_tokens, max_page_indices, max_batch, max_tiles)?; + layout.preallocated_bytes()?; Ok(Self { - page_indices_d: ctx.stream.alloc_zeros(max_total_pages)?, - page_indptr_d: ctx.stream.alloc_zeros(max_batch + 1)?, - last_page_len_d: ctx.stream.alloc_zeros(max_batch)?, - batch_indices_d: ctx.stream.alloc_zeros(max_total_tokens)?, - positions_d: ctx.stream.alloc_zeros(max_total_tokens)?, - q_indptr_d: ctx.stream.alloc_zeros(max_batch + 1)?, - request_indices_d: ctx.stream.alloc_zeros(max_tiles)?, - qo_tile_indices_d: ctx.stream.alloc_zeros(max_tiles)?, - kv_tile_indices_d: ctx.stream.alloc_zeros(max_tiles)?, - kv_chunk_size_d: ctx.stream.alloc_zeros(max_batch)?, - total_num_rows_d: ctx.stream.alloc_zeros(1)?, + page_indices_d: ctx.stream.alloc_zeros(layout.page_indices)?, + page_indptr_d: ctx.stream.alloc_zeros(layout.page_indptr)?, + last_page_len_d: ctx.stream.alloc_zeros(layout.last_page_len)?, + batch_indices_d: ctx.stream.alloc_zeros(layout.batch_indices)?, + positions_d: ctx.stream.alloc_zeros(layout.positions)?, + q_indptr_d: ctx.stream.alloc_zeros(layout.q_indptr)?, + request_indices_d: ctx.stream.alloc_zeros(layout.request_indices)?, + qo_tile_indices_d: ctx.stream.alloc_zeros(layout.qo_tile_indices)?, + kv_tile_indices_d: ctx.stream.alloc_zeros(layout.kv_tile_indices)?, + kv_chunk_size_d: ctx.stream.alloc_zeros(layout.kv_chunk_size)?, + total_num_rows_d: ctx.stream.alloc_zeros(layout.total_num_rows)?, num_tiles: 0, batch_size: 0, total_tokens: 0, @@ -268,52 +362,86 @@ impl PrefillPagedPlan { anyhow::ensure!( host.all_page_indices.len() <= self.page_indices_d.len(), - "verify plan page_indices ({}) exceeds preallocated capacity ({})", + "prefill plan page_indices capacity exceeded: required={}, capacity={}", host.all_page_indices.len(), self.page_indices_d.len(), ); + anyhow::ensure!( + host.batch_size <= self.last_page_len_d.len(), + "prefill plan batch capacity exceeded: required={}, capacity={}", + host.batch_size, + self.last_page_len_d.len(), + ); + anyhow::ensure!( + host.total_tokens <= self.batch_indices_d.len(), + "prefill plan total_tokens capacity exceeded: required={}, capacity={}", + host.total_tokens, + self.batch_indices_d.len(), + ); + anyhow::ensure!( + host.request_indices_v.len() <= self.request_indices_d.len(), + "prefill plan tiles capacity exceeded: required={}, capacity={}", + host.request_indices_v.len(), + self.request_indices_d.len(), + ); anyhow::ensure!( host.page_indptr.len() <= self.page_indptr_d.len(), - "verify plan page_indptr ({}) exceeds preallocated capacity ({})", + "prefill plan page_indptr capacity exceeded: required={}, capacity={}", host.page_indptr.len(), self.page_indptr_d.len(), ); anyhow::ensure!( host.last_page_lens_i32.len() <= self.last_page_len_d.len(), - "verify plan last_page_lens ({}) exceeds preallocated capacity ({})", + "prefill plan last_page_lens capacity exceeded: required={}, capacity={}", host.last_page_lens_i32.len(), self.last_page_len_d.len(), ); anyhow::ensure!( host.batch_indices.len() <= self.batch_indices_d.len(), - "verify plan batch_indices ({}) exceeds preallocated capacity ({})", + "prefill plan batch_indices capacity exceeded: required={}, capacity={}", host.batch_indices.len(), self.batch_indices_d.len(), ); anyhow::ensure!( host.positions.len() <= self.positions_d.len(), - "verify plan positions ({}) exceeds preallocated capacity ({})", + "prefill plan positions capacity exceeded: required={}, capacity={}", host.positions.len(), self.positions_d.len(), ); anyhow::ensure!( host.q_indptr.len() <= self.q_indptr_d.len(), - "verify plan q_indptr ({}) exceeds preallocated capacity ({})", + "prefill plan q_indptr capacity exceeded: required={}, capacity={}", host.q_indptr.len(), self.q_indptr_d.len(), ); anyhow::ensure!( host.request_indices_v.len() <= self.request_indices_d.len(), - "verify plan tiles ({}) exceeds preallocated capacity ({})", + "prefill plan request_indices capacity exceeded: required={}, capacity={}", host.request_indices_v.len(), self.request_indices_d.len(), ); + anyhow::ensure!( + host.qo_tile_indices_v.len() <= self.qo_tile_indices_d.len(), + "prefill plan qo_tile_indices capacity exceeded: required={}, capacity={}", + host.qo_tile_indices_v.len(), + self.qo_tile_indices_d.len(), + ); + anyhow::ensure!( + host.kv_tile_indices_v.len() <= self.kv_tile_indices_d.len(), + "prefill plan kv_tile_indices capacity exceeded: required={}, capacity={}", + host.kv_tile_indices_v.len(), + self.kv_tile_indices_d.len(), + ); anyhow::ensure!( host.kv_chunk_sizes.len() <= self.kv_chunk_size_d.len(), - "verify plan kv_chunk_sizes ({}) exceeds preallocated capacity ({})", + "prefill plan kv_chunk_size capacity exceeded: required={}, capacity={}", host.kv_chunk_sizes.len(), self.kv_chunk_size_d.len(), ); + anyhow::ensure!( + !self.total_num_rows_d.is_empty(), + "prefill plan total_num_rows capacity is zero", + ); ctx.stream .memcpy_htod(&host.all_page_indices, &mut self.page_indices_d)?; @@ -379,10 +507,56 @@ impl BatchPlanHost { cta_tile_q_override: i32, ) -> Result { let batch_size = page_indices.len(); - assert_eq!(batch_size, last_page_lens.len()); - assert_eq!(batch_size, start_positions.len()); - assert_eq!(batch_size, seq_lens.len()); - let total_tokens: usize = seq_lens.iter().sum(); + anyhow::ensure!( + batch_size == last_page_lens.len(), + "prefill plan batch dimension mismatch: page_indices={}, last_page_lens={}", + batch_size, + last_page_lens.len() + ); + anyhow::ensure!( + batch_size == start_positions.len(), + "prefill plan batch dimension mismatch: page_indices={}, start_positions={}", + batch_size, + start_positions.len() + ); + anyhow::ensure!( + batch_size == seq_lens.len(), + "prefill plan batch dimension mismatch: page_indices={}, seq_lens={}", + batch_size, + seq_lens.len() + ); + anyhow::ensure!(num_q_heads > 0, "prefill plan num_q_heads must be positive"); + anyhow::ensure!( + num_kv_heads > 0, + "prefill plan num_kv_heads must be positive" + ); + anyhow::ensure!( + num_q_heads.is_multiple_of(num_kv_heads), + "prefill plan GQA dimensions incompatible: num_q_heads={}, num_kv_heads={}", + num_q_heads, + num_kv_heads + ); + anyhow::ensure!(head_dim > 0, "prefill plan head_dim must be positive"); + anyhow::ensure!( + i32::try_from(num_q_heads).is_ok(), + "prefill plan num_q_heads dimension exceeds i32: {}", + num_q_heads + ); + anyhow::ensure!( + i32::try_from(num_kv_heads).is_ok(), + "prefill plan num_kv_heads dimension exceeds i32: {}", + num_kv_heads + ); + anyhow::ensure!( + i32::try_from(head_dim).is_ok(), + "prefill plan head_dim dimension exceeds i32: {}", + head_dim + ); + let total_tokens: usize = seq_lens.iter().try_fold(0usize, |total, &len| { + total + .checked_add(len) + .ok_or_else(|| anyhow::anyhow!("prefill plan total_tokens overflows usize")) + })?; let group_size = num_q_heads / num_kv_heads; // Page metadata (concatenated across requests, CSR format) @@ -393,9 +567,27 @@ impl BatchPlanHost { for (i, pages) in page_indices.iter().enumerate() { all_page_indices.extend_from_slice(pages); + anyhow::ensure!( + i32::try_from(all_page_indices.len()).is_ok(), + "prefill plan page_indices dimension exceeds i32: {}", + all_page_indices.len() + ); page_indptr.push(all_page_indices.len() as i32); + anyhow::ensure!( + i32::try_from(last_page_lens[i]).is_ok(), + "prefill plan last_page_len dimension exceeds i32: {}", + last_page_lens[i] + ); last_page_lens_i32.push(last_page_lens[i] as i32); - kv_chunk_sizes.push((start_positions[i] + seq_lens[i]) as i32); + let kv_len = start_positions[i].checked_add(seq_lens[i]).ok_or_else(|| { + anyhow::anyhow!("prefill plan kv_chunk_size overflows usize for request {i}") + })?; + anyhow::ensure!( + i32::try_from(kv_len).is_ok(), + "prefill plan kv_chunk_size dimension exceeds i32: {}", + kv_len + ); + kv_chunk_sizes.push(kv_len as i32); } // Per-token metadata @@ -403,19 +595,39 @@ impl BatchPlanHost { let mut positions = Vec::with_capacity(total_tokens); for (i, &seq_len) in seq_lens.iter().enumerate() { let start = start_positions[i]; + anyhow::ensure!( + i32::try_from(i).is_ok(), + "prefill plan batch index exceeds i32: {i}" + ); batch_indices.extend(std::iter::repeat_n(i as i32, seq_len)); - positions.extend((start..start + seq_len).map(|p| p as i32)); + let end = start.checked_add(seq_len).ok_or_else(|| { + anyhow::anyhow!("prefill plan positions overflows usize for request {i}") + })?; + anyhow::ensure!( + i32::try_from(end).is_ok(), + "prefill plan positions dimension exceeds i32: {}", + end + ); + positions.extend((start..end).map(|p| p as i32)); } // Q token boundaries (CSR) let mut q_indptr = vec![0i32]; for &seq_len in seq_lens { let prev = *q_indptr.last().unwrap(); - q_indptr.push(prev + seq_len as i32); + let next = prev + .checked_add(seq_len as i32) + .ok_or_else(|| anyhow::anyhow!("prefill plan q_indptr dimension exceeds i32"))?; + q_indptr.push(next); } // Tile plan: use global cta_tile_q for consistent tiling - let cta_tile_q = unsafe { + anyhow::ensure!( + i32::try_from(total_tokens).is_ok(), + "prefill plan total_tokens dimension exceeds i32: {}", + total_tokens + ); + let cta_tile_q_i32 = unsafe { ffi::batch_prefill_cta_tile_q_with_override( total_tokens as i32, num_q_heads as i32, @@ -423,24 +635,40 @@ impl BatchPlanHost { head_dim as i32, cta_tile_q_override, ) - } as usize; + }; anyhow::ensure!( - cta_tile_q > 0, + cta_tile_q_i32 > 0, "invalid prefill CTA tile override {cta_tile_q_override}" ); + let cta_tile_q = cta_tile_q_i32 as usize; let mut request_indices_v = Vec::new(); let mut qo_tile_indices_v = Vec::new(); let mut kv_tile_indices_v = Vec::new(); for (req_idx, &seq_len) in seq_lens.iter().enumerate() { - let packed_qo_len = seq_len * group_size; + let packed_qo_len = seq_len.checked_mul(group_size).ok_or_else(|| { + anyhow::anyhow!("prefill plan packed query length overflows usize") + })?; let num_tiles_req = packed_qo_len.div_ceil(cta_tile_q); for tile in 0..num_tiles_req { + anyhow::ensure!( + i32::try_from(req_idx).is_ok(), + "prefill plan request index exceeds i32: {req_idx}" + ); + anyhow::ensure!( + i32::try_from(tile).is_ok(), + "prefill plan qo tile index exceeds i32: {tile}" + ); request_indices_v.push(req_idx as i32); qo_tile_indices_v.push(tile as i32); kv_tile_indices_v.push(0i32); } } + anyhow::ensure!( + i32::try_from(request_indices_v.len()).is_ok(), + "prefill plan num_tiles dimension exceeds i32: {}", + request_indices_v.len() + ); let num_tiles = request_indices_v.len() as i32; Ok(Self { @@ -1639,3 +1867,164 @@ pub fn paged_attention_batch_decode_via_prefill_hd256_into( Ok(()) } + +#[cfg(test)] +mod tests { + use anyhow::Result; + use cudarc::driver::DevicePtr; + + use super::PrefillPagedPlan; + use super::PrefillPagedPlanLayout; + use crate::tensor::DeviceContext; + + #[test] + fn preallocated_footprint_uses_allocation_layout() { + let layout = PrefillPagedPlanLayout::new(10, 7, 3, 5).unwrap(); + let bytes = PrefillPagedPlan::preallocated_bytes(10, 7, 3, 5).unwrap(); + + assert_eq!(bytes, layout.preallocated_bytes().unwrap()); + assert_eq!( + layout, + PrefillPagedPlanLayout { + page_indices: 7, + page_indptr: 4, + last_page_len: 3, + batch_indices: 10, + positions: 10, + q_indptr: 4, + request_indices: 5, + qo_tile_indices: 5, + kv_tile_indices: 5, + kv_chunk_size: 3, + total_num_rows: 1, + } + ); + } + + #[test] + fn preallocated_footprint_rejects_i32_capacity_overflow() { + let error = PrefillPagedPlan::preallocated_bytes((i32::MAX as usize) + 1, 1, 1, 1) + .expect_err("metadata dimensions must fit the i32 kernel ABI"); + assert!(error.to_string().contains("max_total_tokens")); + } + + #[test] + #[ignore = "requires a CUDA GPU"] + fn preallocated_update_preserves_all_pointers_and_reports_capacities() -> Result<()> { + let ctx = DeviceContext::new()?; + let layout = PrefillPagedPlanLayout::new(8, 8, 3, 64)?; + let mut plan = PrefillPagedPlan::new_preallocated(&ctx, 8, 8, 3, 64)?; + + assert_eq!(plan.page_indices_d.len(), layout.page_indices); + assert_eq!(plan.page_indptr_d.len(), layout.page_indptr); + assert_eq!(plan.last_page_len_d.len(), layout.last_page_len); + assert_eq!(plan.batch_indices_d.len(), layout.batch_indices); + assert_eq!(plan.positions_d.len(), layout.positions); + assert_eq!(plan.q_indptr_d.len(), layout.q_indptr); + assert_eq!(plan.request_indices_d.len(), layout.request_indices); + assert_eq!(plan.qo_tile_indices_d.len(), layout.qo_tile_indices); + assert_eq!(plan.kv_tile_indices_d.len(), layout.kv_tile_indices); + assert_eq!(plan.kv_chunk_size_d.len(), layout.kv_chunk_size); + assert_eq!(plan.total_num_rows_d.len(), layout.total_num_rows); + + macro_rules! pointers { + ($plan:expr) => { + [ + $plan.page_indices_d.device_ptr(&ctx.stream).0, + $plan.page_indptr_d.device_ptr(&ctx.stream).0, + $plan.last_page_len_d.device_ptr(&ctx.stream).0, + $plan.batch_indices_d.device_ptr(&ctx.stream).0, + $plan.positions_d.device_ptr(&ctx.stream).0, + $plan.q_indptr_d.device_ptr(&ctx.stream).0, + $plan.request_indices_d.device_ptr(&ctx.stream).0, + $plan.qo_tile_indices_d.device_ptr(&ctx.stream).0, + $plan.kv_tile_indices_d.device_ptr(&ctx.stream).0, + $plan.kv_chunk_size_d.device_ptr(&ctx.stream).0, + $plan.total_num_rows_d.device_ptr(&ctx.stream).0, + ] + }; + } + + let initial_pointers = pointers!(plan); + plan.update_batch_with_cta_tile_q( + &ctx, + &[vec![0, 1], vec![2]], + &[16, 1], + &[30, 0], + &[2, 1], + 24, + 4, + 128, + 0, + )?; + ctx.sync()?; + assert_eq!(pointers!(plan), initial_pointers); + + plan.update_batch_with_cta_tile_q( + &ctx, + &[vec![0, 1, 2], vec![3, 4], vec![5]], + &[16, 8, 1], + &[47, 21, 14], + &[1, 3, 2], + 24, + 4, + 128, + 0, + )?; + ctx.sync()?; + assert_eq!(pointers!(plan), initial_pointers); + + let page_error = plan + .update_batch_with_cta_tile_q( + &ctx, + &[vec![0, 1, 2], vec![3, 4, 5], vec![6, 7, 8]], + &[16, 16, 16], + &[47, 47, 47], + &[1, 1, 1], + 24, + 4, + 128, + 0, + ) + .expect_err("logical page capacity must be enforced"); + assert!(page_error.to_string().contains("page_indices capacity")); + + let token_error = plan + .update_batch_with_cta_tile_q( + &ctx, + &[vec![0], vec![1], vec![2]], + &[16, 16, 16], + &[16, 16, 16], + &[3, 3, 3], + 24, + 4, + 128, + 0, + ) + .expect_err("token capacity must be enforced"); + assert!(token_error.to_string().contains("total_tokens capacity")); + + let batch_error = plan + .update_batch_with_cta_tile_q( + &ctx, + &[vec![0], vec![1], vec![2], vec![3]], + &[16, 16, 16, 16], + &[16, 16, 16, 16], + &[1, 1, 1, 1], + 24, + 4, + 128, + 0, + ) + .expect_err("batch capacity must be enforced"); + assert!(batch_error.to_string().contains("batch capacity")); + + let mut tile_limited = PrefillPagedPlan::new_preallocated(&ctx, 8, 8, 3, 1)?; + let tile_error = tile_limited + .update_batch_with_cta_tile_q(&ctx, &[vec![0]], &[16], &[16], &[3], 24, 4, 128, 16) + .expect_err("tile capacity must be enforced"); + assert!(tile_error.to_string().contains("tiles capacity")); + + Ok(()) + } +} diff --git a/openinfer-qwen3/src/executor.rs b/openinfer-qwen3/src/executor.rs index 394557d57..3b98701ac 100644 --- a/openinfer-qwen3/src/executor.rs +++ b/openinfer-qwen3/src/executor.rs @@ -1250,6 +1250,7 @@ impl Qwen3Executor { total_blocks, padding_block_id, max_prefill_tokens, + budget.uncompiled_prefill_plan_capacity, )?, )?, workers: Vec::new(), @@ -1482,6 +1483,7 @@ impl Qwen3Executor { total_blocks, padding_block_id, max_prefill_tokens, + budget.uncompiled_prefill_plan_capacity, )?, )?; @@ -1496,6 +1498,7 @@ impl Qwen3Executor { total_blocks, padding_block_id, max_prefill_tokens, + budget.uncompiled_prefill_plan_capacity, )?; RankWorker::spawn(index + 1, lane) }) @@ -3137,6 +3140,9 @@ struct LocalQwen3Lane { padding_block_id: i32, /// Set by the sweep's `Finalize`; arms the TP replay fail-loud in serving. precapture_complete: bool, + /// Reusable metadata for the uncompiled-GQA unified decode fallback. + /// `None` for models with a compiled decode-attention kernel. + uncompiled_prefill_plan: Option, } /// Stored state for an async prefill that was launched but not yet synced. @@ -3160,6 +3166,7 @@ impl LocalQwen3Lane { total_blocks: usize, padding_block_id: i32, max_prefill_tokens: usize, + uncompiled_prefill_plan_capacity: Option, ) -> Result { let buf_layout = kv_buffer.layout(); let layout = KvLayout::new( @@ -3187,6 +3194,27 @@ impl LocalQwen3Lane { model.config().vocab_size, max_bucket, )?; + anyhow::ensure!( + uncompiled_prefill_plan_capacity.is_some() != model.config().decode_group_is_compiled(), + "Qwen3 reusable-plan capacity does not match decode route" + ); + let uncompiled_prefill_plan = if let Some(capacity) = uncompiled_prefill_plan_capacity { + let bytes = capacity.preallocated_bytes()?; + log::info!( + "Qwen3 uncompiled-GQA PrefillPagedPlan: tokens={}, page_indices={}, batch={}, tiles={}, footprint={} bytes", + capacity.max_total_tokens(), + capacity.max_page_indices(), + capacity.max_batch(), + capacity.max_tiles(), + bytes, + ); + Some(ops::PrefillPagedPlan::new_preallocated_for_capacity( + model.device_ctx(), + capacity, + )?) + } else { + None + }; Ok(Self { model, kv_buffer, @@ -3201,6 +3229,7 @@ impl LocalQwen3Lane { total_blocks, padding_block_id, precapture_complete: false, + uncompiled_prefill_plan, }) } @@ -3504,6 +3533,7 @@ impl LocalQwen3Lane { &mut self.bufs, self.kv_buffer.buffer(), &self.layout, + self.uncompiled_prefill_plan.as_mut(), ) } diff --git a/openinfer-qwen3/src/unified_forward.rs b/openinfer-qwen3/src/unified_forward.rs index 63df45af0..3c3f8e8ae 100644 --- a/openinfer-qwen3/src/unified_forward.rs +++ b/openinfer-qwen3/src/unified_forward.rs @@ -10,6 +10,7 @@ use half::bf16; use openinfer_core::kv_pool::KvLayout; use openinfer_core::ops; use openinfer_core::ops::PrefillPagedPlan; +use openinfer_core::ops::PrefillPagedPlanCapacity; use openinfer_core::sampler::SamplingParams; use openinfer_core::tensor::HiddenStates; use openinfer_kernels::ops::NumericPolicy; @@ -31,6 +32,7 @@ impl Qwen3Model { &self, max_prefill_tokens: usize, profile_decode_rows: usize, + uncompiled_prefill_plan_capacity: Option, kv_buffer: &KvBuffer, decode_bufs: &mut BatchDecodeBuffers, sample_scratch: &mut openinfer_sample::SampleScratch, @@ -99,6 +101,17 @@ impl Qwen3Model { )]; let prefill_adapters: Vec> = vec![None; num_prefill_reqs]; + // The uncompiled-GQA profile exercises the same in-place metadata path + // used by serving. Keep this temporary plan graph-stable as well, so + // the measured peak includes its real persistent footprint. + anyhow::ensure!( + uncompiled_prefill_plan_capacity.is_some() != self.config.decode_group_is_compiled(), + "Qwen3 profile reusable-plan capacity does not match decode route" + ); + let mut uncompiled_prefill_plan = uncompiled_prefill_plan_capacity + .map(|capacity| PrefillPagedPlan::new_preallocated_for_capacity(&self.ctx, capacity)) + .transpose()?; + let logits = self.unified_step_with_peak( &prefill_tokens_list, &prefill_single_views, @@ -109,6 +122,7 @@ impl Qwen3Model { decode_bufs, kv_buffer.buffer(), &layout, + uncompiled_prefill_plan.as_mut(), mark_peak, )?; mark_peak()?; @@ -145,6 +159,7 @@ impl Qwen3Model { decode_bufs: &mut BatchDecodeBuffers, kv_buffer: &CudaSlice, layout: &KvLayout, + uncompiled_prefill_plan: Option<&mut PrefillPagedPlan>, ) -> Result { let mut mark_peak = || Ok(()); self.unified_step_with_peak( @@ -157,6 +172,7 @@ impl Qwen3Model { decode_bufs, kv_buffer, layout, + uncompiled_prefill_plan, &mut mark_peak, ) } @@ -172,6 +188,7 @@ impl Qwen3Model { decode_bufs: &mut BatchDecodeBuffers, kv_buffer: &CudaSlice, layout: &KvLayout, + uncompiled_prefill_plan: Option<&mut PrefillPagedPlan>, mark_peak: &mut dyn FnMut() -> Result<()>, ) -> Result { let num_prefill_reqs = prefill_prompts.len(); @@ -236,7 +253,7 @@ impl Qwen3Model { NumericPolicy::Pin | NumericPolicy::PerToken ) && num_decode_reqs > 0 && self.config.decode_group_is_compiled(); - let plan = if split_decode_attention { + let split_prefill_plan = if split_decode_attention { let positions: Vec = decode_positions.iter().map(|&pos| pos as i32).collect(); self.ctx .stream @@ -267,6 +284,12 @@ impl Qwen3Model { } else { None } + } else { + None + }; + let fresh_prefill_plan; + let plan = if split_decode_attention { + split_prefill_plan.as_ref() } else { // One attention plan over prefill requests + decode rows (qo_len=1, // start at the decode position so the row attends its full history). @@ -284,17 +307,42 @@ impl Qwen3Model { start_positions.extend_from_slice(&decode_positions); let mut seq_lens = prefill_seq_lens.clone(); seq_lens.extend(std::iter::repeat_n(1, num_decode_reqs)); - Some(PrefillPagedPlan::from_raw_batch_with_cta_tile_q( - &self.ctx, - &page_indices, - &last_page_lens, - &start_positions, - &seq_lens, - self.local_num_attention_heads(), - self.local_num_key_value_heads(), - self.config.head_dim, - PREFILL_ATTENTION_CTA_TILE_Q, - )?) + if self.config.decode_group_is_compiled() { + // Tuned keeps compiled GQA decode rows on the unified + // BatchPrefill path. Those models do not reserve the + // uncompiled fallback plan, so retain the original fresh-plan + // behavior for their startup profile and mixed steps. + fresh_prefill_plan = PrefillPagedPlan::from_raw_batch_with_cta_tile_q( + &self.ctx, + &page_indices, + &last_page_lens, + &start_positions, + &seq_lens, + self.local_num_attention_heads(), + self.local_num_key_value_heads(), + self.config.head_dim, + PREFILL_ATTENTION_CTA_TILE_Q, + )?; + Some(&fresh_prefill_plan) + } else { + let plan = uncompiled_prefill_plan.ok_or_else(|| { + anyhow::anyhow!( + "uncompiled GQA decode requires a preallocated PrefillPagedPlan" + ) + })?; + plan.update_batch_with_cta_tile_q( + &self.ctx, + &page_indices, + &last_page_lens, + &start_positions, + &seq_lens, + self.local_num_attention_heads(), + self.local_num_key_value_heads(), + self.config.head_dim, + PREFILL_ATTENTION_CTA_TILE_Q, + )?; + Some(&*plan) + } }; mark_peak()?; @@ -302,7 +350,7 @@ impl Qwen3Model { let hidden = self.unified_layers_with_peak( hidden, total_tokens, - plan.as_ref(), + plan, decode_bufs, split_decode_attention, total_prefill, diff --git a/openinfer-qwen3/src/weights.rs b/openinfer-qwen3/src/weights.rs index 380c77995..4abb82ae3 100644 --- a/openinfer-qwen3/src/weights.rs +++ b/openinfer-qwen3/src/weights.rs @@ -10,6 +10,7 @@ use cudarc::nccl::safe::ReduceOp; use half::bf16; use log::debug; use log::info; +use openinfer_core::ops::PrefillPagedPlanCapacity; use openinfer_core::tensor::DeviceContext; use openinfer_core::tensor::DeviceMatrix; use openinfer_core::tensor::DeviceVec; @@ -103,6 +104,7 @@ pub(crate) struct KvBudget { pub(crate) head_dim: usize, pub(crate) block_size: usize, pub(crate) num_blocks: usize, + pub(crate) uncompiled_prefill_plan_capacity: Option, } #[derive(Clone, Copy, Debug)] @@ -986,6 +988,11 @@ impl Qwen3Model { "Qwen3 memory profile requires decode capacity above prefill rows" ); let profile_rows = profile_prefill_rows + profile_decode_rows; + let uncompiled_prefill_plan_capacity = self.uncompiled_prefill_plan_capacity( + max_prefill_tokens, + max_decode_batch_size, + geometry.block_size, + )?; let profile_blocks = profile_temp_blocks(max_prefill_tokens, profile_decode_rows, geometry.block_size); let profile_kv_bytes = profile_blocks * bytes_per_block; @@ -1034,6 +1041,7 @@ impl Qwen3Model { self.profile_unified_step_memory( max_prefill_tokens, profile_decode_rows, + uncompiled_prefill_plan_capacity, &profile_kv, &mut decode_bufs, &mut sample_scratch, @@ -1045,19 +1053,33 @@ impl Qwen3Model { // the dummy step legal. The final KV pool is sized separately below, so // remove that profile-only backing store from the measured non-KV peak. let profile_peak_increase = peak_used_bytes.saturating_sub(initial_used_bytes); - let non_kv_peak_increase = profile_peak_increase.saturating_sub(profile_kv_bytes); - let non_kv_bytes = initial_used_bytes - .saturating_add(non_kv_peak_increase) - .saturating_add(memory_options.kv_cache_memory_margin_bytes); + let profile_plan_bytes = uncompiled_prefill_plan_capacity + .map(PrefillPagedPlanCapacity::preallocated_bytes) + .transpose()? + .unwrap_or(0); + let non_kv_peak_increase = profile_peak_increase + .saturating_sub(profile_kv_bytes) + .saturating_sub(profile_plan_bytes); + let base_non_kv_bytes = initial_used_bytes + .checked_add(non_kv_peak_increase) + .and_then(|bytes| bytes.checked_add(memory_options.kv_cache_memory_margin_bytes)) + .ok_or_else(|| anyhow::anyhow!("Qwen3 non-KV memory reserve overflows usize"))?; + let persistent_plan_bytes = profile_plan_bytes; + let total_non_kv_bytes = base_non_kv_bytes + .checked_add(persistent_plan_bytes) + .ok_or_else(|| anyhow::anyhow!("Qwen3 total non-KV reserve overflows usize"))?; anyhow::ensure!( - requested_bytes > non_kv_bytes, - "Qwen3 memory profile leaves no room for KV cache: requested={} MiB, \ - non_kv={} MiB, margin={} MiB", + requested_bytes > total_non_kv_bytes, + "Qwen3 memory profile leaves no room for KV cache and persistent uncompiled-GQA plan: requested={} MiB, base_non_kv={} MiB, plan={} MiB", requested_bytes / (1024 * 1024), - non_kv_bytes / (1024 * 1024), - memory_options.kv_cache_memory_margin_bytes / (1024 * 1024) + base_non_kv_bytes / (1024 * 1024), + persistent_plan_bytes / (1024 * 1024), ); - let kv_budget_bytes = requested_bytes - non_kv_bytes; + let kv_budget_bytes = requested_bytes + .checked_sub(total_non_kv_bytes) + .ok_or_else(|| { + anyhow::anyhow!("Qwen3 total non-KV reserve exceeds requested device memory") + })?; let min_kv_bytes = 64 * bytes_per_block; anyhow::ensure!( kv_budget_bytes >= min_kv_bytes, @@ -1067,23 +1089,26 @@ impl Qwen3Model { ); log::info!( "memory profile: total={} MiB requested={} MiB ({:.0}%) initial_used={} MiB \ - peak_non_kv_increase={} MiB margin={} MiB -> KV budget={} MiB", + peak_non_kv_increase={} MiB margin={} MiB persistent_plan={} MiB -> KV budget={} MiB", total_bytes / (1024 * 1024), requested_bytes / (1024 * 1024), memory_options.gpu_memory_utilization * 100.0, initial_used_bytes / (1024 * 1024), non_kv_peak_increase / (1024 * 1024), memory_options.kv_cache_memory_margin_bytes / (1024 * 1024), + persistent_plan_bytes / (1024 * 1024), kv_budget_bytes / (1024 * 1024), ); - Ok(Self::kv_budget_from_bytes( + let mut budget = Self::kv_budget_from_bytes( geometry, bytes_per_block, dflash_kv_bytes_per_token, kv_budget_bytes, initial_free_bytes, "profiled", - )) + ); + budget.uncompiled_prefill_plan_capacity = uncompiled_prefill_plan_capacity; + Ok(budget) } fn kv_budget_geometry(&self, page_size: usize) -> KvBudget { @@ -1094,7 +1119,35 @@ impl Qwen3Model { head_dim: self.config.head_dim, block_size: page_size, num_blocks: 0, + uncompiled_prefill_plan_capacity: None, + } + } + + fn uncompiled_prefill_plan_capacity( + &self, + max_prefill_tokens: usize, + max_batch: usize, + page_size: usize, + ) -> Result> { + if self.config.decode_group_is_compiled() { + return Ok(None); } + // A step with prefill work can use at most max_batch - 1 decode rows; + // a decode-only step has only max_batch tokens. Admission reserves a + // slot for every active or prefilling request. + let max_total_tokens = max_prefill_tokens + .checked_add(max_batch.saturating_sub(1)) + .ok_or_else(|| anyhow::anyhow!("Qwen3 prefill plan token capacity overflow"))? + .max(max_batch); + let group_size = self.local_num_attention_heads() / self.local_num_key_value_heads(); + PrefillPagedPlanCapacity::for_context( + max_total_tokens, + max_batch, + self.config.max_position_embeddings, + page_size, + group_size, + ) + .map(Some) } fn kv_bytes_per_block(geometry: &KvBudget) -> usize { diff --git a/openinfer-qwen35-4b/src/batch_decode.rs b/openinfer-qwen35-4b/src/batch_decode.rs index e0381ba39..aa3606325 100644 --- a/openinfer-qwen35-4b/src/batch_decode.rs +++ b/openinfer-qwen35-4b/src/batch_decode.rs @@ -429,33 +429,18 @@ impl Qwen35Model { start_positions.push(pos); } - let bufs = &mut graph_state.buffers; - bufs.set_batch_size(bs); - self.ctx - .stream - .memcpy_htod(token_ids, &mut bufs.token_ids_d) - .map_err(|e| { - anyhow::anyhow!( - "hybrid decode H2D token_ids bs={bs}, cap={}: {e}", - bufs.max_batch_size - ) - })?; - self.ctx - .stream - .memcpy_htod(&positions_i32, &mut bufs.positions_d) - .map_err(|e| { - anyhow::anyhow!( - "hybrid decode H2D positions bs={bs}, cap={}: {e}", - bufs.max_batch_size - ) - })?; - let page_indices: Vec> = kv_states.iter().map(|kv| kv.page_indices_i32()).collect(); let last_page_lens: Vec = kv_states.iter().map(|kv| kv.last_page_len()).collect(); let seq_lens = vec![1usize; bs]; - // cta_tile_q 0 = the kernel's own FA2 derivation; the hd256 FFI takes no override. - let plan = ops::PrefillPagedPlan::from_raw_batch_with_cta_tile_q( + + let plan = graph_state + .uncompiled_prefill_plan + .as_mut() + .ok_or_else(|| { + anyhow::anyhow!("uncompiled GQA decode requires a preallocated PrefillPagedPlan") + })?; + plan.update_batch_with_cta_tile_q( &self.ctx, &page_indices, &last_page_lens, @@ -468,7 +453,7 @@ impl Qwen35Model { ) .with_context(|| { format!( - "hybrid decode build PrefillPagedPlan bs={bs}, pages={}, heads={}/{}, head_dim={}", + "hybrid decode update PrefillPagedPlan bs={bs}, pages={}, heads={}/{}, head_dim={}", page_indices.iter().map(Vec::len).sum::(), self.config.num_attention_heads, self.config.num_key_value_heads, @@ -476,6 +461,27 @@ impl Qwen35Model { ) })?; + let bufs = &mut graph_state.buffers; + bufs.set_batch_size(bs); + self.ctx + .stream + .memcpy_htod(token_ids, &mut bufs.token_ids_d) + .map_err(|e| { + anyhow::anyhow!( + "hybrid decode H2D token_ids bs={bs}, cap={}: {e}", + bufs.max_batch_size + ) + })?; + self.ctx + .stream + .memcpy_htod(&positions_i32, &mut bufs.positions_d) + .map_err(|e| { + anyhow::anyhow!( + "hybrid decode H2D positions bs={bs}, cap={}: {e}", + bufs.max_batch_size + ) + })?; + let kv_buffer = kv_states[0].buffer(); let layout = *kv_states[0].layout(); anyhow::ensure!( @@ -490,7 +496,7 @@ impl Qwen35Model { self.batch_decode_batched_hybrid_kernels( kv_buffer, &layout, - &plan, + plan, bs, &graph_state.linear_pointer_tables.state_ptrs, &graph_state.linear_pointer_tables.conv_state_ptrs, diff --git a/openinfer-qwen35-4b/src/batch_decode_graph.rs b/openinfer-qwen35-4b/src/batch_decode_graph.rs index 8148046b3..34564dff6 100644 --- a/openinfer-qwen35-4b/src/batch_decode_graph.rs +++ b/openinfer-qwen35-4b/src/batch_decode_graph.rs @@ -3,6 +3,8 @@ use anyhow::Result; use openinfer_core::cuda_graph::CudaGraphState; use openinfer_core::kv_pool::KvPool; +use openinfer_core::ops::PrefillPagedPlan; +use openinfer_core::ops::PrefillPagedPlanCapacity; use openinfer_core::tensor::DeviceContext; use super::config::Config35; @@ -54,6 +56,9 @@ pub(crate) struct BatchDecodeGraphState { pub(crate) linear_pointer_tables: LinearStatePointerTables, /// One `CudaGraphState` per BATCH_BUCKETS entry (indexed by position). pub(crate) graphs: Vec, + /// Reusable metadata for the uncompiled-GQA hybrid decode fallback. + /// `None` for models with a compiled decode-attention kernel. + pub(crate) uncompiled_prefill_plan: Option, } impl BatchDecodeGraphState { @@ -64,16 +69,17 @@ impl BatchDecodeGraphState { tensor_parallel: TensorParallelConfig, kv_pool: &KvPool, max_batch: usize, + uncompiled_prefill_plan_capacity: Option, ) -> Result { let padding_page_id = kv_pool.padding_page_id(); - let max_total_pages = kv_pool.capacity_pages(); + let kv_page_capacity = kv_pool.capacity_pages(); let buffers = BatchDecodeBuffers35::new( ctx, config, tensor_parallel, max_batch, - max_total_pages, + kv_page_capacity, padding_page_id, )?; @@ -97,11 +103,38 @@ impl BatchDecodeGraphState { .map(|_| CudaGraphState::new()) .collect(); + anyhow::ensure!( + uncompiled_prefill_plan_capacity.is_some() != config.decode_group_is_compiled(), + "Qwen3.5 reusable-plan capacity does not match decode route" + ); + let uncompiled_prefill_plan = if let Some(capacity) = uncompiled_prefill_plan_capacity { + anyhow::ensure!( + max_batch <= capacity.max_batch(), + "Qwen3.5 graph batch {max_batch} exceeds reusable-plan batch capacity {}", + capacity.max_batch() + ); + let bytes = capacity.preallocated_bytes()?; + log::info!( + "Qwen3.5 uncompiled-GQA PrefillPagedPlan: tokens={}, page_indices={}, batch={}, tiles={}, footprint={} bytes", + capacity.max_total_tokens(), + capacity.max_page_indices(), + capacity.max_batch(), + capacity.max_tiles(), + bytes, + ); + Some(PrefillPagedPlan::new_preallocated_for_capacity( + ctx, capacity, + )?) + } else { + None + }; + Ok(Self { buffers, slot_states, linear_pointer_tables, graphs, + uncompiled_prefill_plan, }) } diff --git a/openinfer-qwen35-4b/src/prefill_buffers.rs b/openinfer-qwen35-4b/src/prefill_buffers.rs index 0a26863c0..c130ee53b 100644 --- a/openinfer-qwen35-4b/src/prefill_buffers.rs +++ b/openinfer-qwen35-4b/src/prefill_buffers.rs @@ -8,6 +8,21 @@ use openinfer_core::tensor::HiddenStates; use super::config::Config35; +fn checked_product(factors: &[usize], label: &str) -> Result { + factors.iter().try_fold(1usize, |product, &factor| { + product + .checked_mul(factor) + .ok_or_else(|| anyhow::anyhow!("Qwen3.5 {label} size overflows usize")) + }) +} + +fn checked_sum(terms: &[usize], label: &str) -> Result { + terms.iter().try_fold(0usize, |sum, &term| { + sum.checked_add(term) + .ok_or_else(|| anyhow::anyhow!("Qwen3.5 {label} size overflows usize")) + }) +} + /// Scratch buffers for a single Qwen3.5 linear-attention chunk-wise GDR prefill call. /// /// The first implementation target is intentionally narrow: @@ -120,7 +135,7 @@ impl GdrChunkwiseScratch35 { /// /// Direct-paged prefill writes full-attention K/V into the paged pool, so /// HND KVCache staging buffers are no longer part of the prefill scratch. - pub(crate) fn estimate_bytes(config: &Config35, max_seq_len: usize) -> usize { + pub(crate) fn estimate_bytes(config: &Config35, max_seq_len: usize) -> Result { let num_vh = config.linear_num_value_heads; let key_dim = config.linear_key_head_dim; let val_dim = config.linear_value_head_dim; @@ -128,23 +143,40 @@ impl GdrChunkwiseScratch35 { let num_chunks = max_seq_len.div_ceil(chunk_sz); let seq = max_seq_len; - let kv_hidden = num_vh * key_dim; - let vv_hidden = num_vh * val_dim; + let kv_hidden = checked_product(&[num_vh, key_dim], "prefill KV hidden")?; + let vv_hidden = checked_product(&[num_vh, val_dim], "prefill value hidden")?; // 1. GDR scratch (bf16 = 2 bytes, f32 = 4 bytes) let gdr_bytes = { - let f32_elems = seq * num_vh // g_cumsum - + seq * num_vh // beta - + seq * num_vh * chunk_sz // a_tril - + num_chunks * num_vh * val_dim * key_dim; // chunk_state - let bf16_elems = seq * num_vh * chunk_sz // a_inv - + kv_hidden * seq // q_expanded - + kv_hidden * seq // k_expanded - + vv_hidden * seq // v_raw - + kv_hidden * seq // w - + vv_hidden * seq // u - + vv_hidden * seq; // v_new - f32_elems * 4 + bf16_elems * 2 + let seq_vh = checked_product(&[seq, num_vh], "prefill per-head scratch")?; + let seq_vh_chunk = checked_product(&[seq, num_vh, chunk_sz], "prefill chunk matrix")?; + let chunk_state = checked_product( + &[num_chunks, num_vh, val_dim, key_dim], + "prefill chunk state", + )?; + let kv_seq = checked_product(&[kv_hidden, seq], "prefill KV sequence")?; + let vv_seq = checked_product(&[vv_hidden, seq], "prefill value sequence")?; + let f32_elems = checked_sum( + &[seq_vh, seq_vh, seq_vh_chunk, chunk_state], + "prefill f32 scratch", + )?; + let bf16_elems = checked_sum( + &[seq_vh_chunk, kv_seq, kv_seq, vv_seq, kv_seq, vv_seq, vv_seq], + "prefill bf16 scratch", + )?; + checked_sum( + &[ + checked_product( + &[f32_elems, std::mem::size_of::()], + "prefill f32 bytes", + )?, + checked_product( + &[bf16_elems, std::mem::size_of::()], + "prefill bf16 bytes", + )?, + ], + "prefill GDR scratch bytes", + )? }; // 2. Per-layer transient peak (all bf16 = 2 bytes). @@ -153,20 +185,44 @@ impl GdrChunkwiseScratch35 { let intermediate = config.intermediate_size; // Shared: hidden_batch + normed + hidden_plus_attn + normed_for_mlp - let shared_layer = hidden_dim * seq * 4; + let shared_layer = checked_product(&[hidden_dim, seq, 4], "prefill shared layer scratch")?; // Full attention: q_full(with gate) + k + v + attn_out + q_prepped - let full_qkv = config.num_attention_heads * config.head_dim * 2; - let full_kv = config.num_key_value_heads * config.head_dim; - let full_out = config.num_attention_heads * config.head_dim; - let full_attn_temps = (full_qkv + full_kv * 2 + full_out * 2) * seq; + let full_qkv = checked_product( + &[config.num_attention_heads, config.head_dim, 2], + "prefill full-attention Q", + )?; + let full_kv = checked_product( + &[config.num_key_value_heads, config.head_dim], + "prefill full-attention KV", + )?; + let full_out = checked_product( + &[config.num_attention_heads, config.head_dim], + "prefill full-attention output", + )?; + let full_attn_width = checked_sum( + &[ + full_qkv, + checked_product(&[full_kv, 2], "prefill full-attention KV pair")?, + checked_product(&[full_out, 2], "prefill full-attention outputs")?, + ], + "prefill full-attention width", + )?; + let full_attn_temps = + checked_product(&[full_attn_width, seq], "prefill full-attention scratch")?; // MLP: gate_up_out + act_out (same peak footprint as separate gate/up) - let mlp_temps = intermediate * seq * 3; - - let peak_layer = shared_layer + full_attn_temps.max(mlp_temps); - let per_layer_bytes = peak_layer * 2; // bf16 - - gdr_bytes + per_layer_bytes + let mlp_temps = checked_product(&[intermediate, seq, 3], "prefill MLP scratch")?; + + let peak_layer = checked_sum( + &[shared_layer, full_attn_temps.max(mlp_temps)], + "prefill layer peak", + )?; + let per_layer_bytes = checked_product( + &[peak_layer, std::mem::size_of::()], + "prefill layer bytes", + )?; + + checked_sum(&[gdr_bytes, per_layer_bytes], "prefill scratch reserve") } } diff --git a/openinfer-qwen35-4b/src/recurrent_state.rs b/openinfer-qwen35-4b/src/recurrent_state.rs index 8a78a8ef5..c1013ea1b 100644 --- a/openinfer-qwen35-4b/src/recurrent_state.rs +++ b/openinfer-qwen35-4b/src/recurrent_state.rs @@ -43,18 +43,44 @@ pub(crate) struct LinearStatePointerTables { /// Per-layer element counts shared by allocation and reservation: /// (linear layers, f32 state elements, bf16 conv elements). -fn per_layer_dims(config: &Config35) -> (usize, usize, usize) { - let num_linear_layers = config.num_hidden_layers - config.num_full_attention_layers(); - let state_size = - config.linear_num_value_heads * config.linear_key_head_dim * config.linear_value_head_dim; - let conv_state_size = config.linear_attn_qkv_dim() * (config.linear_conv_kernel_dim - 1); - (num_linear_layers, state_size, conv_state_size) +fn per_layer_dims(config: &Config35) -> Result<(usize, usize, usize)> { + let num_linear_layers = config + .num_hidden_layers + .checked_sub(config.num_full_attention_layers()) + .ok_or_else(|| { + anyhow::anyhow!("Qwen3.5 full-attention layer count exceeds total layers") + })?; + let state_size = config + .linear_num_value_heads + .checked_mul(config.linear_key_head_dim) + .and_then(|size| size.checked_mul(config.linear_value_head_dim)) + .ok_or_else(|| anyhow::anyhow!("Qwen3.5 recurrent state dimensions overflow usize"))?; + let linear_q = config + .linear_num_key_heads + .checked_mul(config.linear_key_head_dim) + .ok_or_else(|| anyhow::anyhow!("Qwen3.5 linear Q dimension overflows usize"))?; + let linear_v = config + .linear_num_value_heads + .checked_mul(config.linear_value_head_dim) + .ok_or_else(|| anyhow::anyhow!("Qwen3.5 linear V dimension overflows usize"))?; + let linear_qkv = linear_q + .checked_mul(2) + .and_then(|qk| qk.checked_add(linear_v)) + .ok_or_else(|| anyhow::anyhow!("Qwen3.5 linear QKV dimension overflows usize"))?; + let conv_tail = config + .linear_conv_kernel_dim + .checked_sub(1) + .ok_or_else(|| anyhow::anyhow!("Qwen3.5 linear conv kernel must be positive"))?; + let conv_state_size = linear_qkv + .checked_mul(conv_tail) + .ok_or_else(|| anyhow::anyhow!("Qwen3.5 conv state dimensions overflow usize"))?; + Ok((num_linear_layers, state_size, conv_state_size)) } impl RecurrentState { /// Allocate zeroed recurrent state for all linear attention layers. pub(crate) fn new(ctx: &DeviceContext, config: &Config35) -> Result { - let (num_linear_layers, state_size, conv_state_size) = per_layer_dims(config); + let (num_linear_layers, state_size, conv_state_size) = per_layer_dims(config)?; let mut layers = Vec::with_capacity(num_linear_layers); for _ in 0..num_linear_layers { @@ -145,15 +171,24 @@ impl LinearStatePointerTables { } /// Device bytes of one request's recurrent state. -pub(crate) fn bytes_per_request(config: &Config35) -> usize { - let (num_linear_layers, state_size, conv_state_size) = per_layer_dims(config); +pub(crate) fn bytes_per_request(config: &Config35) -> Result { + let (num_linear_layers, state_size, conv_state_size) = per_layer_dims(config)?; + let state_bytes = state_size + .checked_mul(std::mem::size_of::()) + .ok_or_else(|| anyhow::anyhow!("Qwen3.5 recurrent state bytes overflow usize"))?; + let conv_bytes = conv_state_size + .checked_mul(std::mem::size_of::()) + .ok_or_else(|| anyhow::anyhow!("Qwen3.5 conv state bytes overflow usize"))?; + let per_layer_bytes = state_bytes + .checked_add(conv_bytes) + .ok_or_else(|| anyhow::anyhow!("Qwen3.5 per-layer recurrent bytes overflow usize"))?; num_linear_layers - * (state_size * std::mem::size_of::() - + conv_state_size * std::mem::size_of::()) + .checked_mul(per_layer_bytes) + .ok_or_else(|| anyhow::anyhow!("Qwen3.5 per-request recurrent bytes overflow usize")) } impl RecurrentState { - pub(crate) fn allocation_bytes(config: &Config35) -> usize { + pub(crate) fn allocation_bytes(config: &Config35) -> Result { bytes_per_request(config) } } diff --git a/openinfer-qwen35-4b/src/tp_executor.rs b/openinfer-qwen35-4b/src/tp_executor.rs index 422f852ba..5895806b1 100644 --- a/openinfer-qwen35-4b/src/tp_executor.rs +++ b/openinfer-qwen35-4b/src/tp_executor.rs @@ -790,10 +790,10 @@ impl TpWorkerPrepared { .ctx .mem_get_info() .map_err(|err| anyhow::anyhow!("failed to query TP rank {rank} memory: {err}"))?; - let recurrent_bytes = RecurrentState::allocation_bytes(model.config()); + let recurrent_bytes = RecurrentState::allocation_bytes(model.config())?; let prefill_scratch_tokens = prefill_scratch_tokens(max_prefill_tokens); let prefill_scratch_bytes = - GdrChunkwiseScratch35::estimate_bytes(model.config(), prefill_scratch_tokens); + GdrChunkwiseScratch35::estimate_bytes(model.config(), prefill_scratch_tokens)?; let max_batch = effective_recurrent_capacity( requested_max_batch, free_bytes, diff --git a/openinfer-qwen35-4b/src/weights.rs b/openinfer-qwen35-4b/src/weights.rs index cc0cecaa8..ebef0f8f3 100644 --- a/openinfer-qwen35-4b/src/weights.rs +++ b/openinfer-qwen35-4b/src/weights.rs @@ -7,6 +7,7 @@ use cudarc::nccl::safe::Comm; use cudarc::nccl::safe::ReduceOp; use log::debug; use log::info; +use openinfer_core::ops::PrefillPagedPlanCapacity; use openinfer_core::tensor::DeviceContext; use openinfer_core::tensor::DeviceMatrix; use openinfer_core::tensor::DeviceVec; @@ -124,6 +125,8 @@ pub struct Qwen35Model { /// sit below `reserved_decode_slots` when the request is not a bucket /// (e.g. `--max-batch 5` allocates bucket 8 but admits at most 5). See #470. pub(super) decode_admission_batch: usize, + /// Capacity charged at startup and later consumed by the graph state. + uncompiled_prefill_plan_capacity: Option, tp_comm: Option, } @@ -207,6 +210,11 @@ impl Qwen35Model { let mut config = Config35::from_file(model_path)?; let tensor_parallel = runtime.tensor_parallel.unwrap_or_default(); + // TP currently uses the eager decode kernel path, which does not own a + // hybrid PrefillPagedPlan. Only the single-GPU graph owner needs the + // persistent uncompiled-GQA metadata reservation. + let needs_uncompiled_prefill_plan = + tensor_parallel.world_size == 1 && !config.decode_group_is_compiled(); tensor_parallel.validate_for(&config, runtime.enable_cuda_graph)?; debug!( "Config: hidden_size={}, num_layers={}, full_attn={}, linear_attn={}, max_position_embeddings={}, tp_rank={}, tp_world_size={}", @@ -491,36 +499,71 @@ impl Qwen35Model { config.head_dim, page_size, ); - let bytes_per_page = layout.page_stride * std::mem::size_of::(); + let bytes_per_page = layout + .page_stride + .checked_mul(std::mem::size_of::()) + .ok_or_else(|| anyhow::anyhow!("Qwen3.5 KV page bytes overflow usize"))?; let (free_bytes, _total_bytes) = cudarc::driver::result::mem_get_info() .map_err(|e| anyhow::anyhow!("cuMemGetInfo failed: {e}"))?; // Reserve space for prefill scratch (GDR chunkwise + per-layer transients) // before allocating KV pool, so prefill doesn't OOM. let max_prefill_len = super::prefill::SCRATCH_ESTIMATE_SEQ; - let scratch_reserve = - super::prefill_buffers::GdrChunkwiseScratch35::estimate_bytes(&config, max_prefill_len); - let recurrent_reserve = - STATES_PER_DECODE_SLOT * max_batch * super::recurrent_state::bytes_per_request(&config); - let min_kv_bytes = MIN_KV_PAGES * bytes_per_page; + let scratch_reserve = super::prefill_buffers::GdrChunkwiseScratch35::estimate_bytes( + &config, + max_prefill_len, + )?; + let recurrent_bytes_per_request = super::recurrent_state::bytes_per_request(&config)?; + let recurrent_reserve = STATES_PER_DECODE_SLOT + .checked_mul(max_batch) + .and_then(|slots| slots.checked_mul(recurrent_bytes_per_request)) + .ok_or_else(|| anyhow::anyhow!("Qwen3.5 recurrent reserve overflows usize"))?; + let min_kv_bytes = MIN_KV_PAGES + .checked_mul(bytes_per_page) + .ok_or_else(|| anyhow::anyhow!("Qwen3.5 minimum KV reserve overflows usize"))?; + let uncompiled_prefill_plan_capacity = if needs_uncompiled_prefill_plan { + let group_size = config.num_attention_heads / config.num_key_value_heads; + Some(PrefillPagedPlanCapacity::for_context( + max_batch, + max_batch, + config.max_position_embeddings, + page_size, + group_size, + )?) + } else { + None + }; + let prefill_plan_reserve = uncompiled_prefill_plan_capacity + .map(PrefillPagedPlanCapacity::preallocated_bytes) + .transpose()? + .unwrap_or(0); + let base_reserve = scratch_reserve + .checked_add(recurrent_reserve) + .ok_or_else(|| anyhow::anyhow!("Qwen3.5 scratch/recurrent reserves overflow usize"))?; + let startup_reserve = base_reserve + .checked_add(prefill_plan_reserve) + .ok_or_else(|| anyhow::anyhow!("Qwen3.5 startup reserves overflow usize"))?; + let required_bytes = startup_reserve + .checked_add(min_kv_bytes) + .ok_or_else(|| anyhow::anyhow!("Qwen3.5 total memory requirement overflows usize"))?; anyhow::ensure!( - free_bytes >= scratch_reserve + recurrent_reserve + min_kv_bytes, - "insufficient device memory for Qwen3.5: {} MB free, but prefill scratch needs {} MB, \ - recurrent state needs {} MB ({STATES_PER_DECODE_SLOT} x {max_batch} decode slots), \ - and the minimal KV pool needs {} MB; lower the decode batch capacity (--max-batch) \ - or use a smaller model", + free_bytes >= required_bytes, + "insufficient device memory for Qwen3.5: free={} MB, scratch/recurrent={} MB, persistent plan={} MB, minimum KV={} MB; lower --max-batch or use a smaller model", free_bytes / (1024 * 1024), - scratch_reserve / (1024 * 1024), - recurrent_reserve / (1024 * 1024), + base_reserve / (1024 * 1024), + prefill_plan_reserve / (1024 * 1024), min_kv_bytes / (1024 * 1024), ); - let available = free_bytes - scratch_reserve - recurrent_reserve; + let available = free_bytes + .checked_sub(startup_reserve) + .ok_or_else(|| anyhow::anyhow!("Qwen3.5 startup reserves exceed free device memory"))?; let kv_budget = (available as f64 * 0.85) as usize; let num_pages = (kv_budget / bytes_per_page).max(MIN_KV_PAGES); let kv_mb = num_pages * bytes_per_page / (1024 * 1024); let scratch_mb = scratch_reserve / (1024 * 1024); let recurrent_mb = recurrent_reserve / (1024 * 1024); info!( - "Qwen3.5 KV cache: {num_pages} pages ({kv_mb} MB), prefill scratch reserve: {scratch_mb} MB, recurrent-state reserve: {recurrent_mb} MB ({STATES_PER_DECODE_SLOT} x {max_batch} slots), {:.0}% of {:.0} MB free", + "Qwen3.5 KV cache: {num_pages} pages ({kv_mb} MB), prefill scratch reserve: {scratch_mb} MB, recurrent-state reserve: {recurrent_mb} MB ({STATES_PER_DECODE_SLOT} x {max_batch} slots), persistent plan reserve: {} MB, {:.0}% of {:.0} MB free", + prefill_plan_reserve / (1024 * 1024), kv_budget as f64 / free_bytes as f64 * 100.0, free_bytes as f64 / 1024.0 / 1024.0 ); @@ -546,6 +589,7 @@ impl Qwen35Model { kv_pool, reserved_decode_slots: max_batch, decode_admission_batch, + uncompiled_prefill_plan_capacity, tp_comm: None, }) } @@ -727,6 +771,7 @@ impl Qwen35Model { self.tensor_parallel, &self.kv_pool, max_batch, + self.uncompiled_prefill_plan_capacity, ) }