fix(server): disable GGML_CUDA_GRAPHS to avoid ggml CUDA graph key collision crash#460
Open
zatchbell1311-wq wants to merge 6 commits into
Open
Conversation
…llision crash Pie's portable driver allocates a fresh ggml context arena per request. ggml's CUDA graph cache keys captured graphs by the raw pointer address of the graph's first node (cgraph->nodes[0]) — when a new request's first-node address happens to coincide with a previous (now-freed) request's address, the cache incorrectly treats the new, differently- shaped graph as identical to the cached one, skips rebuilding it, and replays a stale CUDA graph against mismatched tensor shapes. This crashed pie-server reliably (via different GGML_ASSERT failures depending on which shapes collided) after 1-2 successful requests to the same inferlet, or under moderate concurrency. This is a known upstream ggml/llama.cpp issue (ggml-org/llama.cpp #20315, #19708) — the maintainers' sanctioned workaround is disabling CUDA graph capture via GGML_CUDA_GRAPHS=OFF, which this commit forces at the Pie build.rs level for the portable driver's CUDA backend. Verified: before this fix, 6 identical repeated requests to the same inferlet crashed the server on request 2 every time. After, 6 repeated requests (self-consistency, 4-8 concurrent proposal contexts) and 3 repeated requests (graph-of-thought, 14 generation calls each) all completed cleanly with consistent correct outputs and no crashes. Tradeoff: disables a ~10-15% raw decode throughput optimization in exchange for correctness. Given crashes make the backend unusable for any sustained workload, this is a clear net win; a more surgical fix (incorporating cgraph->n_nodes or another shape-derived value into the cache key, matching the graph_recompute fix suggested upstream) could restore graph-reuse performance for genuinely-repeated shapes as a follow-up.
WalkthroughThis PR lowers the portable driver's graph node budget, rewrites the graph-of-thought example to batch proposal and aggregation work, disables GGML_CUDA_GRAPHS in the portable CUDA build, and adds an opt-in CMake switch for FlashInfer/CUTLASS MoE kernels. ChangesDriver Graph Node Budget
Graph-of-Thought Batched Aggregation
Server CUDA Graphs Build Flag
CUDA MoE Build Gate
Sequence Diagram(s)sequenceDiagram
participant Main
participant Model
participant ProposalTasks as Proposal futures
participant AggTasks as Aggregation futures
Main->>Model: fork base_context per proposal
Model-->>ProposalTasks: generate proposals concurrently
ProposalTasks-->>Main: join_all collects proposals
Main->>AggTasks: launch first-level aggregation pairs
AggTasks-->>Main: join_all collects aggregated texts and contexts
Main->>AggTasks: launch second-level aggregation pairs
AggTasks-->>Main: join_all collects final texts
🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (3)
server/build.rs (1)
130-140: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winStale comment contradicts the new
OFFsetting.The comment directly above (lines 135-139) explains the rationale for turning CUDA graphs ON ("Enable explicitly" for perf reasons), but Line 140 now sets
GGML_CUDA_GRAPHSto"OFF". This leaves misleading documentation in place that contradicts the actual behavior and the PR's stated fix rationale (stale graph-cache collisions causing crashes). Update the comment to explain the new rationale (crash avoidance) instead of the old perf-enabling rationale.📝 Proposed comment fix
- // ggml's CUDA-graph capture/replay is OFF by default at the ggml- - // subproject level; llama.cpp's parent CMakeLists.txt flips it - // ON. Embedding ggml directly (as we do here) misses that override - // and ends up issuing ~30 kernel launches per decode step, which - // costs ~1 ms vs ~0.3 ms with capture. Enable explicitly. + // ggml-cuda's CUDA graph cache keys captured graphs by the raw + // address of the first node. Since Pie allocates a fresh ggml + // context arena per request, a new graph can collide with a + // previously freed one and replay against mismatched tensor + // shapes, tripping GGML_ASSERT failures under repeated/ + // concurrent requests. Disable explicitly until upstream + // resolves the caching issue; costs ~1 ms vs ~0.3 ms per + // decode step (more kernel launches instead of a captured + // graph replay). .define("GGML_CUDA_GRAPHS", "OFF")🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@server/build.rs` around lines 130 - 140, The comment above the GGML_CUDA_GRAPHS setting is stale and still explains enabling CUDA graphs for performance, which now conflicts with the actual OFF value. Update the inline explanation in the cuda_enabled block within build.rs to match the new behavior and rationale (crash avoidance from stale graph-cache collisions), keeping the comment aligned with the GGML_CUDA_GRAPHS configuration.inferlets/graph-of-thought/src/lib.rs (2)
130-147: 🚀 Performance & Scalability | 🔵 TrivialBound concurrency for large proposal sets.
future::join_alllaunches every proposal and then every first-level aggregation pair at once, each holding its ownContextand KV pages simultaneously. This is the intended batching win, but with a largeproposal_tokensit can spike KV/page usage. Consider a chunked launch (e.g. bounded concurrency) if inputs can grow, so peak context/KV allocation stays within budget.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@inferlets/graph-of-thought/src/lib.rs` around lines 130 - 147, The unbounded use of future::join_all in the first and second aggregation stages can launch too many proposal and Context allocations at once, spiking KV/page usage for large inputs. Update the launch flow in lib.rs around launch_aggregation_pairs and the two join_all calls to use bounded concurrency or chunked batching so only a limited number of proposal pairs are active at a time while preserving the batching behavior.
49-49: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winRemove the dead
AGGREGATE_PROMPTconstant and its stale comment.The constant is empty and explicitly marked unused, and the comment points to
format_aggregate_prompt, which no longer exists (the prompt is now inlined inlaunch_aggregation_pairs). This is misleading dead code.♻️ Proposed cleanup
-const AGGREGATE_PROMPT: &str = ""; // unused - see format_aggregate_prompt below -🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@inferlets/graph-of-thought/src/lib.rs` at line 49, Remove the dead AGGREGATE_PROMPT constant and its stale inline comment from lib.rs. The prompt is now inlined in launch_aggregation_pairs, so delete the unused empty constant and make sure any references to the old format_aggregate_prompt path are also cleaned up so the module only contains the active prompt source.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@inferlets/graph-of-thought/src/lib.rs`:
- Around line 65-89: The pairing loop in launch_aggregation_pairs is dropping
the leftover proposal when iterating with iter.next(), which can leave
final_results empty on the next pass. Preserve the unpaired item by carrying it
forward into the next aggregation level instead of discarding it, and make sure
the logic around the aggregation task creation still handles the existing
Context/new prompt flow for paired items.
---
Nitpick comments:
In `@inferlets/graph-of-thought/src/lib.rs`:
- Around line 130-147: The unbounded use of future::join_all in the first and
second aggregation stages can launch too many proposal and Context allocations
at once, spiking KV/page usage for large inputs. Update the launch flow in
lib.rs around launch_aggregation_pairs and the two join_all calls to use bounded
concurrency or chunked batching so only a limited number of proposal pairs are
active at a time while preserving the batching behavior.
- Line 49: Remove the dead AGGREGATE_PROMPT constant and its stale inline
comment from lib.rs. The prompt is now inlined in launch_aggregation_pairs, so
delete the unused empty constant and make sure any references to the old
format_aggregate_prompt path are also cleaned up so the module only contains the
active prompt source.
In `@server/build.rs`:
- Around line 130-140: The comment above the GGML_CUDA_GRAPHS setting is stale
and still explains enabling CUDA graphs for performance, which now conflicts
with the actual OFF value. Update the inline explanation in the cuda_enabled
block within build.rs to match the new behavior and rationale (crash avoidance
from stale graph-cache collisions), keeping the comment aligned with the
GGML_CUDA_GRAPHS configuration.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro
Run ID: cd1b10e0-796f-4701-bdfa-495991c4bf10
📒 Files selected for processing (3)
driver/portable/src/graph_common.hppinferlets/graph-of-thought/src/lib.rsserver/build.rs
Comment on lines
+65
to
+89
| while let Some((text_a, _ctx_a)) = iter.next() { | ||
| if let Some((text_b, _ctx_b)) = iter.next() { | ||
| // Both proposal contexts are dropped here — all their KV pages | ||
| // are freed immediately. The aggregation uses a fresh context | ||
| // with only the question + both proposals, so the model always | ||
| // knows what problem it is solving. | ||
| let mut ctx = Context::new(model)?; | ||
| let prompt = format!( | ||
| "You are solving this math problem: {}\n\nHere are two proposed solutions:\n\nSolution A:\n{}\n\nSolution B:\n{}\n\nPick the correct answer. Show your reasoning briefly. End with: ANSWER: <number>", | ||
| question.clone(), text_a, text_b | ||
| ); | ||
| ctx.system(SYSTEM_PROMPT); | ||
| ctx.user(&prompt); | ||
| ctx.cue(); | ||
| tasks.push(async move { | ||
| let text = ctx | ||
| .generate(Sampler::TopP { temperature: 0.6, p: 0.95 }) | ||
| .max_tokens(aggregation_tokens) | ||
| .collect_text() | ||
| .await?; | ||
| Ok((text, ctx)) | ||
| }); | ||
| } | ||
| // Odd item out: context dropped, KV pages freed. | ||
| } |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
# Confirm the default proposal count to gauge how often N<4 is hit.
rg -nP 'fn default_proposal_tokens' -A6 inferlets/graph-of-thought/src/lib.rsRepository: pie-project/pie
Length of output: 541
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Inspect the aggregation logic and its caller.
nl -ba inferlets/graph-of-thought/src/lib.rs | sed -n '1,220p'
echo
echo "---- call sites ----"
rg -n "run_hierarchical_aggregation|final_results|default_proposal_tokens" inferlets/graph-of-thought/src/lib.rsRepository: pie-project/pie
Length of output: 193
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Read the relevant slice with line numbers using available tools.
sed -n '1,220p' inferlets/graph-of-thought/src/lib.rs | cat -n
echo
echo "---- call sites ----"
rg -n "run_hierarchical_aggregation|final_results|default_proposal_tokens" inferlets/graph-of-thought/src/lib.rsRepository: pie-project/pie
Length of output: 9088
Carry the unpaired item into the next aggregation level
With custom proposal_tokens lengths of 2 or 3, the leftover item is dropped in launch_aggregation_pairs, and the second pass can return an empty final_results, so main prints nothing. Forward the unpaired item instead of discarding it.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@inferlets/graph-of-thought/src/lib.rs` around lines 65 - 89, The pairing loop
in launch_aggregation_pairs is dropping the leftover proposal when iterating
with iter.next(), which can leave final_results empty on the next pass. Preserve
the unpaired item by carrying it forward into the next aggregation level instead
of discarding it, and make sure the logic around the aggregation task creation
still handles the existing Context/new prompt flow for paired items.
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@driver/cuda/CMakeLists.txt`:
- Around line 435-438: The PIE_CUDA_BUILD_MOE=OFF branch in CMake adds a source
file that does not exist, causing default OFF builds to fail. Either add the
missing src/ops/flashinfer_moe_stub.cpp implementation with the three
ops::flashinfer_cutlass_moe_* stub symbols, or remove/update the
target_sources(pie_driver_cuda_lib ...) reference in driver/cuda/CMakeLists.txt
so it no longer points to a nonexistent file.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro
Run ID: 39ec1d44-da4a-4b11-b827-dee662dff609
📒 Files selected for processing (1)
driver/cuda/CMakeLists.txt
Comment on lines
+435
to
+438
| else() | ||
| message(STATUS "PIE_CUDA_BUILD_MOE=OFF: skipping FlashInfer/CUTLASS MoE kernels (Nemotron-H only)") | ||
| target_sources(pie_driver_cuda_lib PRIVATE src/ops/flashinfer_moe_stub.cpp) | ||
| endif() |
There was a problem hiding this comment.
🗄️ Data Integrity & Integration | 🟠 Major | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
# Verify stub exists and implements all three interface functions
fd flashinfer_moe_stub.cpp
cat driver/cuda/src/ops/flashinfer_moe_stub.cppRepository: pie-project/pie
Length of output: 229
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo '== files =='
git ls-files | rg 'flashinfer_moe|nemotron_h_forward|CMakeLists.txt$' || true
echo
echo '== search symbols =='
rg -n 'flashinfer_cutlass_moe_enabled|flashinfer_cutlass_moe_bf16_relu2|flashinfer_cutlass_moe_workspace_bytes|flashinfer_moe_stub' driver . || true
echo
echo '== relevant CMake excerpt =='
sed -n '420,450p' driver/cuda/CMakeLists.txtRepository: pie-project/pie
Length of output: 4843
Add the missing src/ops/flashinfer_moe_stub.cpp file
driver/cuda/CMakeLists.txt wires in src/ops/flashinfer_moe_stub.cpp when PIE_CUDA_BUILD_MOE=OFF, but that file is not present. Default OFF builds will fail as soon as CMake tries to compile it. Either add the stub with the three ops::flashinfer_cutlass_moe_* symbols or stop referencing it here.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@driver/cuda/CMakeLists.txt` around lines 435 - 438, The
PIE_CUDA_BUILD_MOE=OFF branch in CMake adds a source file that does not exist,
causing default OFF builds to fail. Either add the missing
src/ops/flashinfer_moe_stub.cpp implementation with the three
ops::flashinfer_cutlass_moe_* stub symbols, or remove/update the
target_sources(pie_driver_cuda_lib ...) reference in driver/cuda/CMakeLists.txt
so it no longer points to a nonexistent file.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Problem
pie servecrashes reliably on the CUDA portable driver backend after 1-2 successful requests to the same inferlet, or under moderate concurrency.The crash signature varies (different
GGML_ASSERTfailure each time —ggml.c:6965graph-size overflow,ggml.c:3637tensor shape mismatch),which initially made this look like several unrelated bugs.Root cause
ggml's CUDA graph cache (
ggml_cuda_graph_get_key,ggml-cuda.cu) keys cached graphs by the raw memory address of the graph's first node(
cgraph->nodes[0]):Pie's portable driver allocates a fresh ggml context arena per request. When a new request's first-node address happens to land at the same address as a previous (now-freed) request's, the cache treats the new — and possibly differently-shaped — graph as identical to the cached one, skips rebuilding it, and replays a stale captured CUDA graph against mismatched tensor
shapes. This is a known upstream issue: ggml-org/llama.cpp#20315, ggml-org/llama.cpp#19708. The maintainers' sanctioned workaround is disabling CUDA graph capture entirely via
GGML_CUDA_GRAPHS=OFF.Pie's
server/build.rswas explicitly forcingGGML_CUDA_GRAPHS=ONwhen configuring the portable driver's CUDA CMake build, which is what exposes this bug.Fix
Flip
GGML_CUDA_GRAPHSfrom"ON"to"OFF"inserver/build.rs.Verification
Before this fix: 6 identical repeated requests to the same inferlet (self-consistency, 8 concurrent proposal contexts) crashed the server on request 2, every time, reproducibly.
After this fix: 6 repeated self-consistency requests and 3 repeated graph-of-thought requests (14 generation calls each) all completed cleanly with consistent, correct outputs. Zero crashes.
Tradeoff
Disabling CUDA graphs removes a ~10-15% raw decode throughput optimization (per upstream benchmarks) in exchange for correctness. Given the backend is otherwise unusable for any sustained workload, this is a clear net improvement.
Possible follow-up
A more surgical fix — incorporating
cgraph->n_nodesor another shape-derived value into the cache key (similar to thegraph_recomputeapproach suggested in ggml-org/llama.cpp#20315) — could restore CUDA graph reuse performance for genuinely-repeated shapes while preserving correctness. Happy to explore this as a separate PR if there's interest; opted for the minimal, verified-safe fix here given the crash makes the backend fully unusable as-is.Summary by CodeRabbit