perf(qwen): reuse PrefillPagedPlan for uncompiled-GQA decode#714
perf(qwen): reuse PrefillPagedPlan for uncompiled-GQA decode#714Nyvo-io wants to merge 1 commit into
Conversation
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 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".
| let plan = uncompiled_prefill_plan.as_deref_mut().ok_or_else(|| { | ||
| anyhow::anyhow!( | ||
| "uncompiled GQA decode requires a preallocated PrefillPagedPlan" | ||
| ) | ||
| })?; |
There was a problem hiding this comment.
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 👍 / 👎.
There was a problem hiding this comment.
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.
|
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. |
|
All commit attributions previously flagged on this pull request are resolved. |
067e4d5 to
59a7496
Compare
|
GPU validation update for current head Environment
Build and runtime gates
The crate-wide Qwen3.5 nightly Clippy gate currently reports existing lint debt on Nsight Systems allocation A/B Both revisions ran the same uncompiled-GQA runtime test. Counts come from the exported CUPTI SQLite trace:
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:
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
left a comment
FeathBow
left a comment
There was a problem hiding this comment.
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.
| Some(ops::PrefillPagedPlan::new_preallocated( | ||
| model.device_ctx(), | ||
| max_total_tokens, | ||
| total_blocks, |
There was a problem hiding this comment.
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.
There was a problem hiding this comment.
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.
| })?; | ||
| openinfer_core::ops::PrefillPagedPlan::preallocated_bytes( | ||
| plan_tokens, | ||
| max_prefill_tokens.div_ceil(geometry.block_size) + profile_decode_rows, |
There was a problem hiding this comment.
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.
There was a problem hiding this comment.
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.
Signed-off-by: Nyvo <75425811+Nyvo-io@users.noreply.github.com>
8a6c73a to
6197e92
Compare
|
@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:
The original five-iteration I also replaced the physical-pool capacity with the logical bound |
FeathBow
left a comment
FeathBow
left a comment
There was a problem hiding this comment.
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.
| // 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 |
There was a problem hiding this comment.
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.
| @@ -517,25 +525,42 @@ impl Qwen35Model { | |||
| let recurrent_reserve = | |||
There was a problem hiding this comment.
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.
|
|
||
| use super::PrefillPagedPlanCapacity; | ||
|
|
||
| #[test] |
There was a problem hiding this comment.
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.
|
|
||
| The test passed `1/1`. Focused capacity tests also passed `2/2`, including a regression case where three logical request page lists refer to only two unique physical pages. | ||
|
|
||
| ### Step 4: Repeat the tail-latency comparison |
There was a problem hiding this comment.
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.
| .ok_or_else(|| { | ||
| anyhow::anyhow!("uncompiled GQA decode requires a preallocated PrefillPagedPlan") | ||
| })?; | ||
| plan.update_batch_with_cta_tile_q( |
There was a problem hiding this comment.
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.
| .map(PrefillPagedPlanCapacity::preallocated_bytes) | ||
| .transpose()? | ||
| .unwrap_or(0); | ||
| let base_reserve = scratch_reserve + recurrent_reserve; |
There was a problem hiding this comment.
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.
|
|
||
| #[test] | ||
| fn preallocated_footprint_matches_all_metadata_arrays() { | ||
| let bytes = PrefillPagedPlan::preallocated_bytes(10, 7, 3, 5).unwrap(); |
There was a problem hiding this comment.
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.
2 participants
Summary
PrefillPagedPlanin each Qwen3 lane and Qwen3.5 graph-state owner for uncompiled GQA groupsPrefillPagedPlanCapacityImplementation
Qwen3 group-5 decode now reuses an owner-scoped plan across unified decode steps. Qwen3.5 group-6 hybrid decode does the same in
BatchDecodeGraphState. Compiled GQA groups preserve the existing fresh-plan behavior and do not reserve a persistent plan.The logical
page_indicescapacity is:This bounds the number of page references even when requests share physical prefix pages. Scheduler admission bounds active plus prefilling requests, while context-length rejection bounds pages per request.
For Qwen3, one capacity object drives profile allocation, profile subtraction, startup accounting, and serving allocation. Qwen3.5 likewise shares one capacity between startup accounting and graph-state allocation. The plan API validates batch, token, page, tile, head, position, and i32 ABI dimensions.
Validation
cargo fmt --all -- --check-D warningscuMemAllocAsync: 708 -> 642cuMemFreeAsync: 708 -> 642Five interleaved 100-iteration benchmark pairs produced 356,800 TPOT samples per run:
The earlier five-iteration
+25%p99 result was not reproducible in repeated interleaved medians. Isolated tail spikes occurred on both revisions. These tests use a one-layer Qwen3 fixture with the same group-size-5 path; they are targeted path evidence, not full Qwen3-14B or Qwen3.5-27B model validation.Fixes #711