fix(qwen3): cap PerToken CUDA Graph buckets - #848
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: 4c832603de
ℹ️ 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".
Signed-off-by: RicardoMin <17879681016@163.com>
4c83260 to
5b77c70
Compare
|
The CI workflow for this fork PR is still awaiting maintainer approval. |
|
sure |
Signed-off-by: RicardoMin <17879681016@163.com>
|
@codex review |
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 1fd04d8111
ℹ️ 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".
Signed-off-by: RicardoMin <17879681016@163.com>
|
Could a maintainer please approve the pending CI workflow so the checks can run? Thank you! |
FeathBow
left a comment
FeathBow
left a comment
There was a problem hiding this comment.
Thank you for working on #780. The cap-and-eager direction is reasonable, but the current head does not yet close the issue.
The profile accounts for bucket 32 only. The other seven allowed buckets can still be captured lazily after KV sizing and remain resident alongside it, so the complete eight-graph set is not budgeted. Please account for the retained set and provide an exact-head base-fails/head-passes receipt for the original 16 GB, utilization-0.90 failure. The 32/33 counters prove dispatch mode, not memory closure.
PerToken+TP is not reachable through the public server contract: --batch-invariant selects Pin and rejects TP. Unless this low-level combination is intentionally supported, please remove its dedicated implementation and 199-line test. Most of the 50-line helper-return block tests local implementation details. Please remove the supplemental-profile matrix and retain, at most, a compact policy-boundary check for PerToken versus Tuned and Pin.
I’d appreciate it if you could give the full diff a careful self-review pass before requesting re-review. That would help keep the next review focused on the substantive remaining issues :)
| // instantiate. Capture that allowed bucket first on the single-GPU | ||
| // path so its executable memory is included in the KV budget. TP | ||
| // keeps this profile eager because ranks are not coordinated here. | ||
| let policy = numeric_policy(); |
There was a problem hiding this comment.
Bucket 32 is not a memory upper bound: BatchDecodeBuffers retains independent graph state for every allowed bucket, so buckets 1, 2, 4, 8, 16, 20, and 24 can still be captured lazily after KV sizing and remain live alongside bucket 32. #780 is specifically a cumulative-residency OOM; please profile or reserve the complete retained set and attach the original 16 GB base-fails/head-passes receipt.
There was a problem hiding this comment.
At the time of testing, I only had access to a 22 GB-class NVIDIA L4, so the results above were collected on that GPU with gpu_memory_utilization=0.90. If validation must be performed on a physical 16 GB GPU, I can rerun the same base/candidate tests on a 16 GB GPU when I am back home and provide the results.
| }; | ||
| run_phase(PrecapturePhase::Warmup)?; | ||
| for bucket_idx in 0..BATCH_BUCKETS.len() { | ||
| let graph_policy = pegainfer_kernels::ops::numeric_policy(); |
There was a problem hiding this comment.
The public Qwen3 launch path maps --batch-invariant to Pin and rejects tensor parallelism for that mode; #780 is also explicitly scoped as verifiable on one consumer GPU. This TP branch is therefore exercised only by callers that directly set the diagnostic PerToken policy and construct the low-level executor, as the new test does.
If that combination is not a supported contract, the TP filtering and 199-line TP test are unnecessary scope for this fix. If it is intended to be supported, please state that contract and retain the following per-rank memory verification as part of the PR evidence.
| 160, 168, 176, 184, 192, 200, 208, 216, 224, 232, 240, 248, 256, | ||
| ]; | ||
| /// Largest batch bucket for which PerToken keeps a CUDA Graph | ||
| pub(crate) const PERTOKEN_GRAPH_MAX_BUCKET: usize = 32; |
There was a problem hiding this comment.
Bucket 32 is a plausible conservative boundary because #780 first observed visible drain around bucket 56, but the current counters only prove graph-versus-eager routing. They do not show that the complete retained allowed set [1, 2, 4, 8, 16, 20, 24, 32] closes the original 16 GB, utilization-0.90 OOM.
Please attach the exact-head command, GPU/model/CUDA configuration, and a base-fails/head-passes receipt from the same full bucket sweep. A post-sweep cuMemGetInfo reading would make the remaining headroom auditable. Please also report the decode-step or TPOT cost at the 32/33 boundary.
| .filter(|&&bucket| per_token_graph_bucket_allowed(graph_policy, bucket)) | ||
| .count(); | ||
|
|
||
| for (bucket_idx, &bucket) in BATCH_BUCKETS.iter().enumerate() { |
There was a problem hiding this comment.
TP profiling is eager, and the KV pool has already been sized and allocated before this loop captures the eight allowed PerToken graphs on every rank. The new TP counter test proves which path executes, but it does not measure this post-budget residency.
Capping may be sufficient; it just needs evidence. Please either reserve the complete per-rank graph footprint or report the exact cumulative per-rank allocation/headroom under the configuration used for the TP claim.
|
|
||
| /// Number of CUDA devices visible to this test process | ||
| fn cuda_device_count() -> usize { | ||
| cudarc::driver::CudaContext::device_count().map_or(0, |count| count.max(0) as usize) |
There was a problem hiding this comment.
map_or(0, ...) makes a CUDA driver/context failure indistinguishable from a host with zero GPUs. The test then returns successfully, so an explicitly requested TP acceptance run can be green without executing any model code.
Please propagate or panic on device_count() errors. Only a successful query that reports fewer than two visible devices should take the availability-skip path.
| let use_cuda_graph = graphs_available && graph_use != DecodeGraphUse::Eager; | ||
| // PerToken records one GEMM node per now; cap Large buckets to bound graph memory. | ||
| // Above the cap, keep the PetToken arithmetic but execute the kernels eagerly | ||
| let mut use_cuda_graph = graphs_available && graph_use != DecodeGraphUse::Eager; |
There was a problem hiding this comment.
The cap is enforced here only for Serve; CaptureOnly and Replay still accept a disallowed PerToken bucket and rely on separate controller/lane/Finalize filters. The lane currently turns such a request into a silent successful no-op, which can hide a protocol mismatch.
Please give graph eligibility one owner. Serve may intentionally choose eager above the cap, while an explicit capture/replay request for a disallowed bucket should be impossible by construction or fail loudly.
| NumericPolicy::Tuned, | ||
| *super::BATCH_BUCKETS.last().unwrap() | ||
| )); | ||
| assert_eq!( |
There was a problem hiding this comment.
These assertions mostly restate a local helper and remain green despite the cumulative-accounting bug. The retained 32/33 integration probe already covers the serving boundary. Please remove the supplemental-profile matrix; if the unaffected-policy contract needs direct coverage, keep only a compact check for PerToken, Tuned, and Pin. A full-plan unit test is useful only if profiling and serving consume one canonical graph plan.
| @@ -0,0 +1,199 @@ | |||
| //! TP=2 regression for the PerToken CUDA-Graph bucket cap. | |||
There was a problem hiding this comment.
Please remove this test unless direct PerToken+TP is a supported low-level contract. If it is supported, retain only the TP-specific pre-capture/barrier contract that the single-GPU gate cannot cover.
|
Dear Reviewer,
Thank you very much for your careful review and valuable feedback on this PR. I sincerely apologize for the oversights in the code. Your feedback has pointed out clear areas for improvement, which I truly appreciate.
I fully understand and accept your comments, and will immediately rectify the issues according to your guidance to ensure they are properly resolved. Once again, apologies for any inconvenience caused, and thanks for your patience and instruction!
Best regards,
RicardoMin
---- Replied Message ----
| From | Salt ***@***.***> |
| Date | 08/07/2026 01:07 |
| To | ***@***.***> |
| Cc | ***@***.***>***@***.***> |
| Subject | Re: [openinfer-project/pegainfer] fix(qwen3): cap PerToken CUDA Graph buckets (PR #848) |
@FeathBow requested changes on this pull request.
Thank you for working on #780. The cap-and-eager direction is reasonable, but the current head does not yet close the issue.
The profile accounts for bucket 32 only. The other seven allowed buckets can still be captured lazily after KV sizing and remain resident alongside it, so the complete eight-graph set is not budgeted. Please account for the retained set and provide an exact-head base-fails/head-passes receipt for the original 16 GB, utilization-0.90 failure. The 32/33 counters prove dispatch mode, not memory closure.
PerToken+TP is not reachable through the public server contract: --batch-invariant selects Pin and rejects TP. Unless this low-level combination is intentionally supported, please remove its dedicated implementation and 199-line test. Most of the 50-line helper-return block tests local implementation details. Please remove the supplemental-profile matrix and retain, at most, a compact policy-boundary check for PerToken versus Tuned and Pin.
I’d appreciate it if you could give the full diff a careful self-review pass before requesting re-review. That would help keep the next review focused on the substantive remaining issues :)
In pegainfer-qwen3/src/unified_forward.rs:
@@ -73,6 +74,28 @@ impl Qwen3Model { } else { crate::batch_decode::DecodeGraphUse::Serve }; + // PerToken buckets above 32 intentionally run eager, so the usual + // max-row probe would miss the largest graph that serving can still + // instantiate. Capture that allowed bucket first on the single-GPU + // path so its executable memory is included in the KV budget. TP + // keeps this profile eager because ranks are not coordinated here. + let policy = numeric_policy();
Bucket 32 is not a memory upper bound: BatchDecodeBuffers retains independent graph state for every allowed bucket, so buckets 1, 2, 4, 8, 16, 20, and 24 can still be captured lazily after KV sizing and remain live alongside bucket 32. #780 is specifically a cumulative-residency OOM; please profile or reserve the complete retained set and attach the original 16 GB base-fails/head-passes receipt.
In pegainfer-qwen3/src/executor.rs:
@@ -1544,15 +1545,24 @@ impl Qwen3Executor { Ok(()) }; run_phase(PrecapturePhase::Warmup)?; - for bucket_idx in 0..BATCH_BUCKETS.len() { + let graph_policy = pegainfer_kernels::ops::numeric_policy();
The public Qwen3 launch path maps --batch-invariant to Pin and rejects tensor parallelism for that mode; #780 is also explicitly scoped as verifiable on one consumer GPU. This TP branch is therefore exercised only by callers that directly set the diagnostic PerToken policy and construct the low-level executor, as the new test does.
If that combination is not a supported contract, the TP filtering and 199-line TP test are unnecessary scope for this fix. If it is intended to be supported, please state that contract and retain the following per-rank memory verification as part of the PR evidence.
In pegainfer-qwen3/src/batch_decode_buffers.rs:
@@ -26,6 +26,31 @@ pub(crate) const BATCH_BUCKETS: &[usize] = &[ 1, 2, 4, 8, 16, 20, 24, 32, 40, 48, 56, 64, 72, 80, 88, 96, 104, 112, 120, 128, 136, 144, 152, 160, 168, 176, 184, 192, 200, 208, 216, 224, 232, 240, 248, 256, ]; +/// Largest batch bucket for which PerToken keeps a CUDA Graph +pub(crate) const PERTOKEN_GRAPH_MAX_BUCKET: usize = 32;
Bucket 32 is a plausible conservative boundary because #780 first observed visible drain around bucket 56, but the current counters only prove graph-versus-eager routing. They do not show that the complete retained allowed set [1, 2, 4, 8, 16, 20, 24, 32] closes the original 16 GB, utilization-0.90 OOM.
Please attach the exact-head command, GPU/model/CUDA configuration, and a base-fails/head-passes receipt from the same full bucket sweep. A post-sweep cuMemGetInfo reading would make the remaining headroom auditable. Please also report the decode-step or TPOT cost at the 32/33 boundary.
In pegainfer-qwen3/src/executor.rs:
@@ -1544,15 +1545,24 @@ impl Qwen3Executor { Ok(()) }; run_phase(PrecapturePhase::Warmup)?; - for bucket_idx in 0..BATCH_BUCKETS.len() { + let graph_policy = pegainfer_kernels::ops::numeric_policy(); + let captured_buckets = BATCH_BUCKETS + .iter() + .filter(|&&bucket| per_token_graph_bucket_allowed(graph_policy, bucket)) + .count(); + + for (bucket_idx, &bucket) in BATCH_BUCKETS.iter().enumerate() {
TP profiling is eager, and the KV pool has already been sized and allocated before this loop captures the eight allowed PerToken graphs on every rank. The new TP counter test proves which path executes, but it does not measure this post-budget residency.
Capping may be sufficient; it just needs evidence. Please either reserve the complete per-rank graph footprint or report the exact cumulative per-rank allocation/headroom under the configuration used for the TP claim.
In pegainfer-qwen3/tests/tp_pertoken_graph_cap.rs:
+const PROMPT_LEN: usize = 8; +const AT_CAP_BATCH: usize = 32; +const ABOVE_CAP_BATCH: usize = 33; + +/// Reslove the model path or skip this GPU integration test +fn model_path_or_skip() -> Option<String> { + let Ok(path) = std::env::var("PEGAINFER_TEST_MODEL_PATH") else { + eprintln!("skipping tp_pertoken_graph_cap: set PEGAINFER_TEST_MODEL_PATH"); + return None; + }; + Some(path) +} + +/// Number of CUDA devices visible to this test process +fn cuda_device_count() -> usize { + cudarc::driver::CudaContext::device_count().map_or(0, |count| count.max(0) as usize)
map_or(0, ...) makes a CUDA driver/context failure indistinguishable from a host with zero GPUs. The test then returns successfully, so an explicitly requested TP acceptance run can be green without executing any model code.
Please propagate or panic on device_count() errors. Only a successful query that reports fewer than two visible devices should take the availability-skip path.
In pegainfer-qwen3/src/batch_decode.rs:
@@ -94,13 +95,32 @@ impl Qwen3Model { graphs_available || matches!(graph_use, DecodeGraphUse::Serve | DecodeGraphUse::Eager), "batch_decode {graph_use:?} requires CUDA graphs enabled and no LoRA rows" ); - let use_cuda_graph = graphs_available && graph_use != DecodeGraphUse::Eager; + // PerToken records one GEMM node per now; cap Large buckets to bound graph memory. + // Above the cap, keep the PetToken arithmetic but execute the kernels eagerly + let mut use_cuda_graph = graphs_available && graph_use != DecodeGraphUse::Eager;
The cap is enforced here only for Serve; CaptureOnly and Replay still accept a disallowed PerToken bucket and rely on separate controller/lane/Finalize filters. The lane currently turns such a request into a silent successful no-op, which can hide a protocol mismatch.
Please give graph eligibility one owner. Serve may intentionally choose eager above the cap, while an explicit capture/replay request for a disallowed bucket should be impossible by construction or fail loudly.
In pegainfer-qwen3/src/batch_decode_buffers.rs:
+ NumericPolicy::PerToken, + PERTOKEN_GRAPH_MAX_BUCKET + )); + assert!(!per_token_graph_bucket_allowed( + NumericPolicy::PerToken, + PERTOKEN_GRAPH_MAX_BUCKET + 1 + )); + assert!(per_token_graph_bucket_allowed( + NumericPolicy::Pin, + *super::BATCH_BUCKETS.last().unwrap() + )); + assert!(per_token_graph_bucket_allowed( + NumericPolicy::Tuned, + *super::BATCH_BUCKETS.last().unwrap() + )); + assert_eq!(
These assertions mostly restate a local helper and remain green despite the cumulative-accounting bug. The retained 32/33 integration probe already covers the serving boundary. Please remove the supplemental-profile matrix; if the unaffected-policy contract needs direct coverage, keep only a compact check for PerToken, Tuned, and Pin. A full-plan unit test is useful only if profiling and serving consume one canonical graph plan.
In pegainfer-qwen3/tests/tp_pertoken_graph_cap.rs:
@@ -0,0 +1,199 @@ +//! TP=2 regression for the PerToken CUDA-Graph bucket cap.
Please remove this test unless direct PerToken+TP is a supported low-level contract. If it is supported, retain only the TP-specific pre-capture/barrier contract that the single-GPU gate cannot cover.
—
Reply to this email directly, view it on GitHub, or unsubscribe.
You are receiving this because you authored the thread.Message ID: ***@***.***>
|
|
@codex review |
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 3420aaa8a4
ℹ️ 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".
Qwen3 PerToken CUDA Graph memory validationThank you for the detailed review. I revised the implementation to account for the complete retained Graph set and validated the memory and performance behavior with a same-GPU base/candidate comparison. Current PR head: Test environment
Implementation
Full-bucket memory sweepThe same diagnostic exercised all 36 decode buckets twice in one executor and sampled
Candidate dispatch was: The base failure surfaced during the bucket-136 prefill allocation after the preceding Graph captures had reduced free memory to 40.12 MiB. The candidate completed all 36 buckets with a final retained delta of only 16.00 MiB. Batch 32/33 boundaryThe direct
Across 100 paired rounds, the mean latency delta was Validation
ConclusionThe base run confirms that the OOM is caused by cumulative PerToken Graph-executable residency, not by one largest Graph. The revision accounts for the complete eight-Graph retained set before KV sizing and prevents additional persistent Graph allocation above bucket 32. All 36 buckets complete successfully, while the measured aggregate throughput change at the 32/33 boundary is PerToken+TP remains outside the supported public contract and is not claimed by this PR. |
Signed-off-by: RicardoMin <17879681016@163.com>
3420aaa to
824f5d8
Compare
|
I've fixed the DCO sign-off issue. The CI workflow is currently awaiting maintainer approval. Thanks! |
FeathBow
left a comment
FeathBow
left a comment
There was a problem hiding this comment.
Thanks for the rework.
Imo, two blockers remain.
PerToken+TP is described as unsupported but is not rejected at the public low-level constructor. It loads and profiles the ranks, allocates KV, spawns workers and arms the TP watchdog before failing at bucket 40. Returning that error drops the success sender, while the watchdog remains armed and aborts the process at its 600-second deadline. Please reject this combination before model loading and watchdog creation.
Profiling reserves only the full-SM graphs cache. Decode overlap captures into the independent graphs_split cache, while enable_decode_overlap() rejects Pin but still allows PerToken. Please reject PerToken+overlap at that guard or account for both graph sets.
Btw, the remaining documentation and cleanup items are non-blocking, but it would be great if you could address them in the same pass.
| ); | ||
| Ok(None) | ||
| } | ||
| DecodeGraphUse::CaptureOnly | DecodeGraphUse::Replay => { |
There was a problem hiding this comment.
This part is reached by the unconditional TP sweep at bucket 40, after model/KV allocation, worker startup and watchdog creation. Returning Err drops sweep_done_tx; the watchdog deliberately remains armed and later calls std::process::abort(), so even a caller that handles the error can lose the process ten minutes later.
If PerToken+TP is unsupported, please reject it at the start of from_runtime_with_lora_options(). This does not require TP-specific cap handling or restoring the deleted TP test.
| && self.tensor_parallel.world_size == 1 | ||
| && graph_plan.requires_cumulative_profile() | ||
| { | ||
| for graph_rows in graph_plan.retained_buckets() { |
There was a problem hiding this comment.
This loop reserves only the full-SM graphs cache. SplitConcurrent captures into the independent graphs_split cache; enable_decode_overlap() rejects Pin but allows PerToken, and the PerToken GEMM accepts the override stream. One executor can therefore retain two eight-graph sets while profiling reserves one.
Please reject PerToken+overlap at the existing guard or profile both caches. The guard is the smaller fix if this diagnostic combination is unsupported.
| fn assert_pertoken_graph_cap_behavior(model_path: &str) { | ||
| set_numeric_policy(NumericPolicy::PerToken); | ||
|
|
||
| let mut ex = Qwen3Executor::from_runtime(model_path, true, &[0]).expect("build probe executor"); |
There was a problem hiding this comment.
The test already constructs a PerToken executor in run_policy(), drops it, then loads a fourth model for this probe. Please run the 32/33 counter check on the existing PerToken executor. Keep the counter transition; the extra model startup, LOGPROBS=64, diagnostic print and helper-name comments add no coverage.
3 participants
Description
Fixes #780
PerToken routes projection GEMMs through a row-wise numerical path. CUDA Graph
capture records one GEMM node per token row, so graph executable memory grows
with larger decode buckets.
This change limits PerToken CUDA Graph capture to bucket 32. Batches above the
cap keep the PerToken numerical path but fall back to eager decode. Tensor
parallel startup uses the same rule when performing graph pre-capture and
finalization.
Type of Change
Validation
cargo fmt --all --checkgit diff --checkcargo check --release -p pegainfer-qwen3 --libbatch_invariance_decode_gemm_graphtp_pertoken_graph_captp_concurrent_decodeObserved:
bs=32 served 0->0;bs=33 served 506->1012.Checklist