perf(qwen): reuse PrefillPagedPlan for uncompiled-GQA decode

[Nyvo-io]

Summary

• reuse one owner-scoped PrefillPagedPlan for uncompiled-GQA decode in Qwen3 and Qwen3.5
• size logical page metadata as max_batch * ceil(max_position_embeddings / page_size), including shared-prefix workloads
• derive all eleven allocation lengths and exact byte accounting from one checked layout
• use checked startup reserve arithmetic end to end for both model loaders
• retain the original fresh-plan path for compiled-GQA Qwen3 models

Implementation

PrefillPagedPlanCapacity validates token, logical-page, batch, tile, and i32 ABI limits. PrefillPagedPlanLayout is the single source for allocation lengths and memory accounting. Runtime updates refill the eleven device buffers in place and return dimension-specific capacity errors instead of allocating a fallback plan.

Qwen3 passes one capacity through profile allocation, profile-byte subtraction, startup accounting, and every serving lane. Qwen3.5 charges the same capacity before KV allocation and gives it to BatchDecodeGraphState. Prefix sharing is bounded by logical request page lists rather than unique physical KV pages.

Validation

Final source: d3715e8, rebased on origin/main cb41b07.

The production-gate source was initially based on 90440c5. The final rebase onto cb41b07 (#745) only narrows the Qwen3.5 capacity field from pub(super) to private; git range-diff confirms that the runtime diff is unchanged.

• capacity/layout CPU tests: 4 passed
• GPU pointer/capacity gate: all eleven pointers stable across changing shapes; page/token/batch/tile errors verified
• Qwen3-14B HF golden: passed
• Qwen3-14B prefix-cache mixed prefill/decode: passed
• Qwen3.5-27B short golden: passed
• Qwen3.5-27B long golden (4097/8192-token prompts): passed
• Qwen3.5-27B scheduler E2E: passed
• release checks, eight-package Clippy with -D warnings, and Qwen3.5 all-target feature check: passed
• all 15 GitHub CI checks on final head d3715e8: passed

Production model revisions:

• Qwen3-14B: 40c069824f4251a91eefaf281ebe4c544efd3e18
• Qwen3.5-27B: fc05daec18b0a78c049392ed2e771dde82bdf654

Performance Evidence

Five interleaved 100-iteration fixture pairs produced 356,800 TPOT samples per run:

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%

Five production Qwen3-14B Nsight pairs each captured 32 steady eager decode steps at context 512 and batch 16:

• cuMemAllocAsync: 864 -> 512 calls per trace, exactly -11 per step
• cuMemFreeAsync: 864 -> 512 calls per trace, exactly -11 per step
• paired-median allocation/free API saving: 19.265 us/step
• H2D calls: unchanged at 448 per trace
• kernels: unchanged at 24,512 per trace
• profile TPOT paired-delta median: -0.11%

The production claim is reduced plan-preparation allocation/free API work without changed metadata uploads or kernel work, not a material TPOT speedup.

Exact commands, environment, per-pair tables, CUPTI SQL, runtime receipts, and model revisions are retained in docs/subsystems/kernels/reusable-prefill-plan.md.

Fixes #711

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: 91a8ffd05f

ℹ️ About Codex in GitHub

Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".

[chatgpt-codex-connector]

Keep fresh plans for compiled Tuned unified steps

With the default NumericPolicy::Tuned, split_decode_attention is false even when the Qwen3 GQA group is compiled, so stock Qwen3-4B/8B enters this branch during the startup memory profile and any mixed prefill+decode unified step. Callers only pass a reusable uncompiled_prefill_plan when !decode_group_is_compiled(), so compiled default runs now fail with uncompiled GQA decode requires a preallocated PrefillPagedPlan instead of building the fresh plan this branch used before.

Useful? React with 👍 / 👎.

[Nyvo-io]

Fixed in 59a7496. The unified path now uses the reusable plan only for an uncompiled GQA group; compiled GQA retains the original fresh-plan behavior for default NumericPolicy::Tuned startup profiling and mixed prefill+decode steps.

I validated both sides of the branch on an RTX 5090 with CUDA 12.8 using unified_within_path_gemm_n_invariant_under_pin: stock Qwen3-4B (compiled GQA, default Tuned) passed, and a one-layer Qwen3 fixture with 20 Q heads / 4 KV heads (uncompiled group size 5) passed. Both runs reported 1 passed, 0 failed.

The current head is 8a6c73a; the commit after 59a7496 is only a behavior-preserving Qwen3.5 Clippy branch-order fix.

[FeathBow]

Passing CI alone isn't enough to rule out performance regressions. Since the GPU allocation tracing and Qwen3 measurements are pending, I will convert this PR to a Draft for now. Please take your time to add and run those actual performance tests, and we can look into it together once the data is ready.

[github-actions]

All commit attributions previously flagged on this pull request are resolved.

[Nyvo-io]

GPU validation update for current head 8a6c73a.

Environment

• NVIDIA GeForce RTX 5090 (32 GB), driver 595.71.05
• CUDA 12.8, OPENINFER_CUDA_SM=120
• Rust nightly 2026-07-10

Build and runtime gates

• cargo fmt --all -- --check: pass
• cargo check --release --locked -p openinfer-server --all-targets: pass
• Qwen3 CUDA Clippy set with -D warnings: pass
• cargo check --release --locked -p openinfer-qwen35-4b --all-targets --features qwen35-4b: pass, including the Qwen3.5 Triton AOT build and test targets
• Stock Qwen3-4B compiled-GQA/default-Tuned unified runtime gate: 1 passed, 0 failed
• One-layer Qwen3 uncompiled-GQA fixture (20 Q heads / 4 KV heads, group size 5) on the same gate: 1 passed, 0 failed

The crate-wide Qwen3.5 nightly Clippy gate currently reports existing lint debt on origin/main. It also found one new if_not_else warning from this PR; that warning was fixed in 8a6c73a, and the focused lint gate passes.

Nsight Systems allocation A/B

Both revisions ran the same uncompiled-GQA runtime test. Counts come from the exported CUPTI SQLite trace:

CUDA Driver API	origin/main	PR	Delta
cuMemAllocAsync	708	642	-66 (-9.3%)
cuMemFreeAsync	708	642	-66 (-9.3%)
cuLaunchKernel	8466	8466	0

The unchanged kernel count supports that the comparison removed allocation/free churn without reducing the test workload.

Decode benchmark A/B

Same fixture, context 512, 256 decode steps, 32 warmups, 5 iterations, CUDA graphs disabled:

Batch	Avg TPOT base -> PR	Delta	Decode tok/s base -> PR	Delta
1	0.814142 -> 0.788851 ms	-3.11%	1301.69 -> 1328.19	+2.04%
8	0.863782 -> 0.848747 ms	-1.74%	9731.16 -> 9898.55	+1.72%
16	0.917572 -> 0.902274 ms	-1.67%	17795.63 -> 18177.62	+2.15%

Average TPOT, p50, p95, and throughput moved in the expected direction for all three batches. I do not want to overstate the short run: batch-16 p99 regressed from 1.384099 ms to 1.736591 ms, and this fixture is targeted path evidence rather than a full Qwen3-14B or Qwen3.5-27B benchmark.

I am keeping the PR as Draft. Please let me know whether this targeted runtime/allocation evidence is sufficient for review or whether you would prefer full-model validation first.

[FeathBow]

The targeted allocation evidence is useful and supports the churn claim. Before undrafting, could the batch-16 p99 result (1.384 → 1.737 ms) be checked with a longer or repeated run? Five iterations are not enough to determine whether this tail movement is noise or a regression.
The capacity question above also remains uncovered by the one-layer, single-request fixture. A real uncompiled-GQA run with prefix caching and mixed prefill/decode would be the most direct validation of the new fixed-capacity path.

[FeathBow]

The plan's page_indices capacity is set to the physical KV-pool size (total_blocks here, max_total_pages in the Qwen3.5 path), but the plan stores concatenated page lists for all requests. With prefix caching, multiple requests can reference the same physical prefix pages, so the logical page-table length can exceed the number of unique pool pages.

The old path allocated a right-sized plan per step; this path now returns page_indices capacity exceeded on a valid prefix-cache workload. This is directly relevant to the uncompiled-GQA target, while the current single-request fixture cannot exercise it. Could the plan use a bound for the maximum logical page-table size, with the same capacity included in startup accounting, or explicitly establish and enforce an invariant that makes the pool-size bound valid?

The scheduler appears to reserve active + prefilling against max_bucket (openinfer-qwen3/src/scheduler.rs:1123-1165); please confirm that this invariant covers every unified-step route.

[Nyvo-io]

Fixed in 6197e92. The reusable plan now sizes page_indices from the logical request bound:

max_batch * ceil(max_position_embeddings / page_size)

It no longer uses the number of unique physical KV-pool pages. The scheduler admission bound covers active plus prefilling requests, and context-length rejection bounds the pages contributed by each request, so this capacity covers every unified-step route even when several requests reference the same cached prefix pages.

I also added and ran a group-size-5 prefix-cache integration case with a shared prefix and mixed prefill/decode; it passed on the RTX 5090.

[FeathBow]

The profile plan dimensions are re-derived independently from the dimensions used when the temporary plan is created in profile_unified_step_memory: pages use max_prefill_tokens.div_ceil(block_size) + profile_decode_rows here versus prefill_pages + profile_decode_rows there, and batch uses profile_rows versus num_prefill_reqs + profile_decode_rows.

These expressions happen to be equal today, but the exact-reservation claim depends on them remaining identical. Could the capacity dimensions be computed once and shared by both allocation and subtraction, or at least checked for equality? This is a drift-prevention issue rather than an observed mismatch in the current head.

[Nyvo-io]

Fixed in 6197e92. A single PrefillPagedPlanCapacity is now computed once and shared instead of re-deriving equivalent dimensions.

For Qwen3, that same value drives profile allocation, profile subtraction, startup reserve accounting, and serving allocation. Qwen3.5 likewise shares one capacity between startup accounting and graph-state allocation. This removes the drift risk between the temporary profile plan and the reserved persistent plan.

[Nyvo-io]

@FeathBow I reran the batch-16 benchmark as five interleaved base/PR pairs with 100 iterations per run (356,800 TPOT samples in each run). The medians were:

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%
Throughput	18,546.04 tok/s	18,711.22 tok/s	+0.99%

The original five-iteration +25% p99 movement was not reproducible in these repeated interleaved medians. There were isolated tail spikes on both revisions: PR seed 47 reached 1.964 ms p99, while base seeds 49 and 50 reached 1.868 ms and 1.833 ms.

I also replaced the physical-pool capacity with the logical bound max_batch * ceil(max_context / page_size) and added a real group-size-5 shared-prefix mixed prefill/decode integration run. That test passed on the RTX 5090. The fixture exercises the target path, but I am not treating it as full Qwen3-14B or Qwen3.5-27B model validation.

[FeathBow]

Since this PR is intended to close #711 for both production uncompiled-GQA models, I think the final acceptance should also include the existing Qwen3-14B runtime/golden gates and the Qwen3.5-27B short/long golden plus scheduler E2E. The one-layer group-size-5 fixture is good mechanism coverage, but it does not replace end-to-end validation of either production checkpoint.

[FeathBow]

The PR changes Qwen3.5 loader reservation, graph-state ownership, and the complete batch_decode_batched_hybrid plan construction/update boundary, but every reported runtime test uses the Qwen3 group-size-5 fixture. A Qwen3 fixture cannot execute Qwen3.5's hybrid recurrent/full-attention path. The committed doc nevertheless says the outcome is validated and relegates full Qwen3.5-27B coverage to a follow-up. At minimum, a group-6 Qwen3.5 fixture must execute consecutive decode shapes and the intended model-level gates must remain open before Fixes #711 is accepted.

[FeathBow]

Qwen3.5 computes recurrent and base reserves with ordinary multiplication/addition, then repeats unchecked sums in the admission predicate and inside checked_sub. Qwen3 similarly evaluates base_non_kv_bytes + persistent_plan_bytes before checking the result. A sufficiently large external model/config dimension can wrap these expressions in release and make the loader admit an oversized KV pool, contradicting the PR's explicit checked-accounting contract. Derive every reserve with checked_mul/checked_add once and subtract that shared checked total.

[FeathBow]

The footprint test manually repeats the same element-count formula but never compares it with the allocation layout, so changing one of the eleven allocation lengths without changing preallocated_bytes leaves the test green. The logical-page test initializes a capacity and asserts the expected formula; it does not update a real plan. No test checks device-pointer stability or capacity errors across changing batch/context/page shapes. The external group-size fixture is not committed, and CI only compiles these test targets. Replace formula restatements with one shared layout derivation used by allocation and accounting, plus a focused GPU update/pointer gate or retain a reproducible target fixture.

[FeathBow]

Retain the per-pair table and exact collection/query commands at minimum; the original acceptance also requested plan-preparation CPU/API time, which remains absent.

[FeathBow]

Could we add one runtime receipt for this Qwen3.5 group-6 path before treating #711 as closed? The current group-size-5 fixture is useful evidence for the shared plan mechanism, but it exercises the Qwen3 owner rather than this hybrid recurrent/full-attention path, its loader reservation, or its graph-state allocation. A Qwen3.5 group-6 fixture covering consecutive batch/context/page shapes, or the existing Qwen3.5-27B short/long golden and scheduler E2E gates, would make the cross-model acceptance claim verifiable.

[FeathBow]

Would you mind deriving the complete reserve with checked multiplication/addition once, then reusing that checked total in both the admission check and checked_sub? recurrent_reserve, base_reserve, and base_reserve + prefill_plan_reserve are still evaluated with ordinary arithmetic here, so a sufficiently large external model configuration can wrap before checked_sub sees it in a release build. The Qwen3 branch has the same smaller gap at base_non_kv_bytes + persistent_plan_bytes. Closing both would make the PR's checked-accounting claim hold end to end.

[FeathBow]

Could this test be tied to the allocation layout rather than restating the footprint formula? As written, changing one of the eleven new_preallocated buffer lengths without updating preallocated_bytes would still leave this test green. Deriving both allocation lengths and byte accounting from one small layout value would remove that drift class; a focused GPU gate that updates different batch/context/page shapes and compares the eleven device pointers would then cover the pointer-stability part of the contract. The existing logical-page formula test is useful for the shared-prefix bound, but it does not exercise an actual plan update.

[Nyvo-io]

@FeathBow All requested follow-ups are complete on final head d3715e8 (rebased on origin/main cb41b07). I also addressed every open line-level concern before rerunning the gates.

Code and invariant fixes

• Logical page capacity is now max_batch * ceil(max_position_embeddings / page_size), so shared physical prefix pages cannot under-size concatenated request page lists.
• PrefillPagedPlanLayout drives all eleven allocation lengths and exact byte accounting; Qwen3 profile/accounting/serving and Qwen3.5 accounting/graph ownership consume the same checked capacity.
• Qwen3 and Qwen3.5 startup reserves use checked multiplication/addition through one total reused by admission and subtraction.
• The compiled-GQA Qwen3 route retains fresh plans. The uncompiled route updates the persistent owner-scoped plan and fails with dimension-specific errors rather than allocating a fallback.
• The focused CUDA gate passed across changing page/token/batch shapes, verified all eleven device pointers remain stable, and exercised page/token/batch/tile capacity failures.

Requested production gates

• Qwen3-14B HF golden: 1 passed; sequential mean/p99 0.0291/0.1087, cached replay 0.0281/0.1046.
• Qwen3-14B prefix-cache mixed gate: 1 passed; warm/cold mixed rows were bit-identical to cold reference, and the unified warm row passed at mean/max 0.0335/0.0592.
• Qwen3.5-27B short golden: 1 passed; sequential mean/p99 0.0201/0.0694, batch-5 p99 0.0651, slot-compaction p99 0.0755.
• Qwen3.5-27B long golden: 1 passed; 4097/8192-token prompts, mean/p99 0.0218/0.0627.
• Qwen3.5-27B scheduler E2E: 1 passed in 60.47 s.

Performance follow-ups

• The requested longer tail check used five interleaved base/PR pairs, 100 iterations each. Median p99 was 1.033162 -> 1.005956 ms (-1.48% paired median); the earlier five-iteration movement did not reproduce.
• Five production Qwen3-14B Nsight pairs each captured 32 eager decode steps. Both cuMemAllocAsync and cuMemFreeAsync changed 864 -> 512 calls per trace, exactly -11 each per step. Paired-median allocation/free API time saved was 19.265 us/step.
• Every production trace retained 448 H2D calls and 24,512 kernels. Profile TPOT paired-delta median was -0.11%, so the claim is limited to reduced plan-preparation API work.

The production-gate source was initially based on 90440c5. The final rebase onto cb41b07 (#745) only narrows the Qwen3.5 capacity field from pub(super) to private; git range-diff confirms that the runtime diff is unchanged.

The exact environment, commands, five benchmark pairs, five Nsight pairs, CUPTI SQL, model revisions, and all receipts are in reusable-prefill-plan.md. The final commit is DCO-signed, two final full-diff reviews found no actionable defects, and all 15 GitHub CI checks on d3715e8 are green.

[FeathBow]

Please rebase onto main before review.