Skip to content

fix(server): disable GGML_CUDA_GRAPHS to avoid ggml CUDA graph key collision crash #460

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
zatchbell1311-wq wants to merge 6 commits into
base: main
Choose a base branch
from
Open

fix(server): disable GGML_CUDA_GRAPHS to avoid ggml CUDA graph key collision crash #460

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
6 commits
Select commit Hold shift + click to select a range
109f9ad
perf(got): batch aggregation pairs to reduce scheduler round-trips
zatchbell1311-wq Jun 26, 2026
da32422
fix(got): include question in aggregation prompt to prevent context loss
dhruvdubey1311-crypto Jun 27, 2026
621a0f3
Merge remote-tracking branch 'origin/main' into feat/got-batched-aggr…
dhruvdubey1311-crypto Jul 3, 2026
2092fb1
WIP: GRAPH_MAX_NODES fix for CPU backend
dhruvdubey1311-crypto Jul 3, 2026
f712de7
fix(server): disable GGML_CUDA_GRAPHS to avoid ggml CUDA graph key co…
dhruvdubey1311-crypto Jul 3, 2026
03d6683
fix(driver/cuda): remove Python3 build dependency, handle CUDA 13 CCC…
dhruvdubey1311-crypto Jul 5, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
126 changes: 69 additions & 57 deletions driver/cuda/CMakeLists.txt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -367,63 +367,75 @@ target_include_directories(pie_driver_cuda_lib
# third_party/flashinfer_generated/ (the SM80 set; byte-identical for SM80/86/89
# and the only set the build globs). This removes the build-time Python codegen
# and its incidental tvm_ffi / pynvml / tqdm imports -- no external Python deps.
# Refresh after bumping the vendored FlashInfer revision with:
# FLASHINFER_SRC=<src> driver/cuda/third_party/flashinfer_generated/regenerate.sh
set(PIE_FLASHINFER_MOE_GEN_DIR
"${CMAKE_CURRENT_SOURCE_DIR}/third_party/flashinfer_generated")
file(GLOB PIE_FLASHINFER_MOE_GENERATED_SOURCES CONFIGURE_DEPENDS
"${PIE_FLASHINFER_MOE_GEN_DIR}/gemm_grouped/80/*.generated.cu")
list(LENGTH PIE_FLASHINFER_MOE_GENERATED_SOURCES PIE_FLASHINFER_MOE_GENERATED_COUNT)
message(STATUS "Using ${PIE_FLASHINFER_MOE_GENERATED_COUNT} FlashInfer generated MoE launchers")

# FlashInfer/TensorRT-LLM CUTLASS MoE runner. Keep this isolated from the
# regular driver TUs because the dependency slice is large and should remain a
# clear opt-in boundary for Nemotron experimentation.
add_library(pie_flashinfer_cutlass_moe OBJECT
src/ops/flashinfer_moe.cu
${flashinfer_SOURCE_DIR}/csrc/nv_internal/cpp/common/envUtils.cpp
${flashinfer_SOURCE_DIR}/csrc/nv_internal/cpp/common/logger.cpp
${flashinfer_SOURCE_DIR}/csrc/nv_internal/cpp/common/stringUtils.cpp
${flashinfer_SOURCE_DIR}/csrc/nv_internal/cpp/common/tllmException.cpp
${flashinfer_SOURCE_DIR}/csrc/nv_internal/cpp/common/memoryUtils.cu
${flashinfer_SOURCE_DIR}/csrc/nv_internal/tensorrt_llm/kernels/preQuantScaleKernel.cu
${flashinfer_SOURCE_DIR}/csrc/nv_internal/tensorrt_llm/kernels/cutlass_kernels/cutlass_heuristic.cpp
${flashinfer_SOURCE_DIR}/csrc/nv_internal/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/moe_gemm_tma_warp_specialized_input.cu
${flashinfer_SOURCE_DIR}/csrc/nv_internal/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/moe_gemm_kernels_bf16_bf16.cu
${flashinfer_SOURCE_DIR}/csrc/nv_internal/tensorrt_llm/kernels/cutlass_kernels/fp8_blockscale_gemm/fp8_blockscale_gemm.cu
${flashinfer_SOURCE_DIR}/csrc/nv_internal/tensorrt_llm/kernels/lora/lora.cpp
${PIE_FLASHINFER_MOE_GENERATED_SOURCES})
target_include_directories(pie_flashinfer_cutlass_moe PRIVATE
${CMAKE_CURRENT_SOURCE_DIR}/src
${CUDAToolkit_INCLUDE_DIRS}
${PIE_CCCL_INCLUDE_DIRS}
${flashinfer_SOURCE_DIR}/include
${flashinfer_SOURCE_DIR}/csrc
${flashinfer_SOURCE_DIR}/csrc/fused_moe/cutlass_backend
${flashinfer_SOURCE_DIR}/csrc/nv_internal
${flashinfer_SOURCE_DIR}/csrc/nv_internal/include
${flashinfer_SOURCE_DIR}/csrc/nv_internal/tensorrt_llm/cutlass_extensions/include
${flashinfer_SOURCE_DIR}/csrc/nv_internal/tensorrt_llm/kernels/cutlass_kernels/include
${flashinfer_SOURCE_DIR}/csrc/nv_internal/tensorrt_llm/kernels/cutlass_kernels
${flashinfer_SOURCE_DIR}/csrc/nv_internal/tensorrt_llm/kernels/cutlass_kernels/moe_gemm
${flashinfer_SOURCE_DIR}/csrc/nv_internal/tensorrt_llm/kernels/cutlass_kernels/fp8_blockscale_gemm
${PIE_FLASHINFER_MOE_GEN_DIR}
${flashinfer_SOURCE_DIR}/3rdparty/cutlass/include
${flashinfer_SOURCE_DIR}/3rdparty/cutlass/tools/util/include
${flashinfer_SOURCE_DIR}/3rdparty/spdlog/include)
target_compile_definitions(pie_flashinfer_cutlass_moe PRIVATE
ENABLE_BF16=1
ENABLE_FP8=1
USING_OSS_CUTLASS_MOE_GEMM=1)
target_compile_options(pie_flashinfer_cutlass_moe PRIVATE
$<$<COMPILE_LANGUAGE:CUDA>:
--extended-lambda
--expt-relaxed-constexpr
>)
set_target_properties(pie_flashinfer_cutlass_moe PROPERTIES
POSITION_INDEPENDENT_CODE ON)
target_sources(pie_driver_cuda_lib PRIVATE
$<TARGET_OBJECTS:pie_flashinfer_cutlass_moe>)
option(PIE_CUDA_BUILD_MOE
"Build the FlashInfer/CUTLASS MoE GEMM kernels (required only for Nemotron-H). \
OFF by default: each generated kernel file can need several GB RAM to compile, \
which OOMs on small dev machines (e.g. WSL2 with <8GB RAM)."
OFF)

if(PIE_CUDA_BUILD_MOE)
# Refresh after bumping the vendored FlashInfer revision with:
# FLASHINFER_SRC=<src> driver/cuda/third_party/flashinfer_generated/regenerate.sh
set(PIE_FLASHINFER_MOE_GEN_DIR
"${CMAKE_CURRENT_SOURCE_DIR}/third_party/flashinfer_generated")
file(GLOB PIE_FLASHINFER_MOE_GENERATED_SOURCES CONFIGURE_DEPENDS
"${PIE_FLASHINFER_MOE_GEN_DIR}/gemm_grouped/80/*.generated.cu")
list(LENGTH PIE_FLASHINFER_MOE_GENERATED_SOURCES PIE_FLASHINFER_MOE_GENERATED_COUNT)
message(STATUS "Using ${PIE_FLASHINFER_MOE_GENERATED_COUNT} FlashInfer generated MoE launchers")

# FlashInfer/TensorRT-LLM CUTLASS MoE runner. Keep this isolated from the
# regular driver TUs because the dependency slice is large and should remain a
# clear opt-in boundary for Nemotron experimentation.
add_library(pie_flashinfer_cutlass_moe OBJECT
src/ops/flashinfer_moe.cu
${flashinfer_SOURCE_DIR}/csrc/nv_internal/cpp/common/envUtils.cpp
${flashinfer_SOURCE_DIR}/csrc/nv_internal/cpp/common/logger.cpp
${flashinfer_SOURCE_DIR}/csrc/nv_internal/cpp/common/stringUtils.cpp
${flashinfer_SOURCE_DIR}/csrc/nv_internal/cpp/common/tllmException.cpp
${flashinfer_SOURCE_DIR}/csrc/nv_internal/cpp/common/memoryUtils.cu
${flashinfer_SOURCE_DIR}/csrc/nv_internal/tensorrt_llm/kernels/preQuantScaleKernel.cu
${flashinfer_SOURCE_DIR}/csrc/nv_internal/tensorrt_llm/kernels/cutlass_kernels/cutlass_heuristic.cpp
${flashinfer_SOURCE_DIR}/csrc/nv_internal/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/moe_gemm_tma_warp_specialized_input.cu
${flashinfer_SOURCE_DIR}/csrc/nv_internal/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/moe_gemm_kernels_bf16_bf16.cu
${flashinfer_SOURCE_DIR}/csrc/nv_internal/tensorrt_llm/kernels/cutlass_kernels/fp8_blockscale_gemm/fp8_blockscale_gemm.cu
${flashinfer_SOURCE_DIR}/csrc/nv_internal/tensorrt_llm/kernels/lora/lora.cpp
${PIE_FLASHINFER_MOE_GENERATED_SOURCES})
target_include_directories(pie_flashinfer_cutlass_moe PRIVATE
${CMAKE_CURRENT_SOURCE_DIR}/src
${CUDAToolkit_INCLUDE_DIRS}
${PIE_CCCL_INCLUDE_DIRS}
${flashinfer_SOURCE_DIR}/include
${flashinfer_SOURCE_DIR}/csrc
${flashinfer_SOURCE_DIR}/csrc/fused_moe/cutlass_backend
${flashinfer_SOURCE_DIR}/csrc/nv_internal
${flashinfer_SOURCE_DIR}/csrc/nv_internal/include
${flashinfer_SOURCE_DIR}/csrc/nv_internal/tensorrt_llm/cutlass_extensions/include
${flashinfer_SOURCE_DIR}/csrc/nv_internal/tensorrt_llm/kernels/cutlass_kernels/include
${flashinfer_SOURCE_DIR}/csrc/nv_internal/tensorrt_llm/kernels/cutlass_kernels
${flashinfer_SOURCE_DIR}/csrc/nv_internal/tensorrt_llm/kernels/cutlass_kernels/moe_gemm
${flashinfer_SOURCE_DIR}/csrc/nv_internal/tensorrt_llm/kernels/cutlass_kernels/fp8_blockscale_gemm
${PIE_FLASHINFER_MOE_GEN_DIR}
${flashinfer_SOURCE_DIR}/3rdparty/cutlass/include
${flashinfer_SOURCE_DIR}/3rdparty/cutlass/tools/util/include
${flashinfer_SOURCE_DIR}/3rdparty/spdlog/include)
target_compile_definitions(pie_flashinfer_cutlass_moe PRIVATE
ENABLE_BF16=1
ENABLE_FP8=1
USING_OSS_CUTLASS_MOE_GEMM=1)
target_compile_options(pie_flashinfer_cutlass_moe PRIVATE
$<$<COMPILE_LANGUAGE:CUDA>:
--extended-lambda
--expt-relaxed-constexpr
>)
set_target_properties(pie_flashinfer_cutlass_moe PROPERTIES
POSITION_INDEPENDENT_CODE ON)
target_sources(pie_driver_cuda_lib PRIVATE
$<TARGET_OBJECTS:pie_flashinfer_cutlass_moe>)
target_compile_definitions(pie_driver_cuda_lib PRIVATE PIE_CUDA_HAVE_FLASHINFER_MOE=1)
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()
Comment on lines +435 to +438

@coderabbitai coderabbitai Bot Jul 5, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🗄️ 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.cpp

Repository: 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.txt

Repository: 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.


# ── Marlin (W4A16 / MXFP4 GEMM via vendored vLLM kernel) ──────────────
# Built by default because GPT-OSS MXFP4 auto policy should use native
Expand Down
2 changes: 1 addition & 1 deletion driver/portable/src/graph_common.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -21,7 +21,7 @@ namespace pie_portable_driver {

// Per-call ggml graph node budget. Sized to comfortably fit MoE + spec
// decode batches (every layer's MoE op count is ~10-20 nodes).
inline constexpr std::size_t GRAPH_MAX_NODES = 1ull << 19;
inline constexpr std::size_t GRAPH_MAX_NODES = 1ull << 14; // reduced from 1<<19 (524288) to 1<<14 (16384) — the original value requests a ~10GB scheduler buffer regardless of model size, which fails malloc on memory-constrained systems for small models like Qwen3-0.6B.

// ggml's CUDA flash_attn supports head_dim ≤ 256 across all GQA ratios.
// Gemma 4's full-attention layers use head_dim=512 — those layers fall
Expand Down
160 changes: 93 additions & 67 deletions inferlets/graph-of-thought/src/lib.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,12 +1,23 @@
//! Demonstrates Graph-of-Thought (GoT) for hierarchical aggregation.
//!
//! This example generates multiple initial proposals concurrently, then
//! progressively aggregates them in pairs across multiple levels. The streaming
//! nature allows aggregation to begin as soon as pairs of proposals are ready,
//! maximizing parallelism.
//! progressively aggregates them in pairs across multiple levels.
//!
//! # Batched aggregation
//!
//! The original implementation processed proposals one at a time in a
//! `while let Some` loop, which meant aggregation pairs were launched
//! sequentially. This prevented the runtime scheduler from seeing all
//! aggregation requests simultaneously and coalescing them into a single
//! batched forward pass.
//!
//! This version collects all proposals first via `future::join_all`, then
//! launches all aggregation pairs in one shot. The scheduler now receives
//! all N/2 aggregation `flush()` calls concurrently and can batch them
//! together, reducing the number of GPU kernel invocations per level from
//! N/2 sequential firings to ideally 1 batched firing.

use futures::stream::FuturesUnordered;
use futures::{StreamExt, future};
use futures::future;
use inferlet::{
Context, sample::Sampler, model::Model,
runtime, Result,
Expand All @@ -31,28 +42,69 @@ fn default_aggregation_tokens() -> usize { 256 }
const SYSTEM_PROMPT: &str = "You are a helpful, respectful and honest assistant.";

const PROPOSAL_PROMPT_TEMPLATE: &str = "\
Could you suggest a method or approach to solve the following question? \
Please provide a high-level plan without doing the actual calculation. \
Keep it concise, around 80 words. Question: {}";

const AGGREGATE_PROMPT: &str = "\
Please compare the following solution with the one you just provided \
and aggregate their ideas into a single, improved solution:\n";
Solve the following math problem step by step. Show your work and give the final numeric answer. \
End your response with: ANSWER: <number>\n\
Question: {}";

const AGGREGATE_PROMPT: &str = ""; // unused - see format_aggregate_prompt below

/// Pair up a flat list of (text, context) results into aggregation tasks,
/// launching all pairs simultaneously so the scheduler can batch them.
///
/// If the input has an odd number of items, the last unpaired item is
/// dropped (its context is released immediately, freeing KV pages).
fn launch_aggregation_pairs(
items: Vec<(String, Context)>,
aggregation_tokens: usize,
question: String,
model: &Model,
) -> Result<Vec<impl std::future::Future<Output = Result<(String, Context)>> + 'static>> {
let mut tasks = Vec::new();
let mut iter = items.into_iter();

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.
}
Comment on lines +65 to +89

@coderabbitai coderabbitai Bot Jul 3, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🎯 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.rs

Repository: 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.rs

Repository: 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.rs

Repository: 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.


Ok(tasks)
}

/// Main logic for running the hierarchical aggregation workflow.
async fn run_hierarchical_aggregation(
base_context: &mut Context,
question: &str,
proposal_tokens: Vec<usize>,
aggregation_tokens: usize,
model: &Model,
) -> Result<Vec<String>> {
// --- Stage 1: Generate Initial Proposals ---
// --- Stage 1: Generate Initial Proposals (all concurrent) ---
let propose_prompt = PROPOSAL_PROMPT_TEMPLATE.replace("{}", question);
base_context.user(&propose_prompt);
base_context.flush().await?;

let mut proposal_tasks = proposal_tokens
let proposal_futures = proposal_tokens
.into_iter()
.map(|max_tokens| {
let mut ctx = base_context.fork()?;
Expand All @@ -66,62 +118,35 @@ async fn run_hierarchical_aggregation(
Ok::<_, String>((proposal_text, ctx))
})
})
.collect::<Result<Vec<_>>>()?
.collect::<Result<Vec<_>>>()?;

// Collect ALL proposals before pairing. This lets Stage 2 launch all
// aggregation pairs simultaneously rather than as each proposal dribbles in.
let proposals: Vec<(String, Context)> = future::join_all(proposal_futures)
.await
.into_iter()
.collect::<FuturesUnordered<_>>();

// --- Stage 2: First-Level Aggregation (Pairing Proposals) ---
let mut first_aggregation_tasks = FuturesUnordered::new();
let mut pending_proposal: Option<(String, Context)> = None;

while let Some(result) = proposal_tasks.next().await {
let (proposal_text, mut proposal_ctx) = result?;
if pending_proposal.is_none() {
pending_proposal = Some((proposal_text, proposal_ctx));
} else {
let (previous_proposal_text, _) = pending_proposal.take().unwrap();
let aggregation_prompt = format!("{}{}", AGGREGATE_PROMPT, previous_proposal_text);
proposal_ctx.user(&aggregation_prompt);
proposal_ctx.cue();

first_aggregation_tasks.push(async move {
let aggregation_text = proposal_ctx
.generate(Sampler::TopP { temperature: 0.6, p: 0.95 })
.max_tokens(aggregation_tokens)
.collect_text()
.await?;
Ok::<_, String>((aggregation_text, proposal_ctx))
});
}
}
.collect::<Result<_>>()?;

// --- Stage 2: First-Level Aggregation (all pairs launched at once) ---
// All N/2 aggregation flush() calls hit the scheduler simultaneously,
// which can coalesce them into a single batched forward pass.
let first_agg_futures = launch_aggregation_pairs(proposals, aggregation_tokens, question.to_string(), &model)?;
let first_aggregations: Vec<(String, Context)> = future::join_all(first_agg_futures)
.await
.into_iter()
.collect::<Result<_>>()?;

// --- Stage 3: Second-Level Aggregation (Pairing Aggregations) ---
let mut second_aggregation_tasks = Vec::new();
let mut pending_aggregation: Option<(String, Context)> = None;

while let Some(result) = first_aggregation_tasks.next().await {
let (aggregation_text, mut aggregation_ctx) = result?;
if pending_aggregation.is_none() {
pending_aggregation = Some((aggregation_text, aggregation_ctx));
} else {
let (previous_aggregation_text, _) = pending_aggregation.take().unwrap();
let final_prompt = format!("{}{}", AGGREGATE_PROMPT, previous_aggregation_text);
aggregation_ctx.user(&final_prompt);
aggregation_ctx.cue();

second_aggregation_tasks.push(async move {
aggregation_ctx
.generate(Sampler::TopP { temperature: 0.6, p: 0.95 })
.max_tokens(aggregation_tokens)
.collect_text()
.await
});
}
}
// --- Stage 3: Second-Level Aggregation (all pairs launched at once) ---
let second_agg_futures = launch_aggregation_pairs(first_aggregations, aggregation_tokens, question.to_string(), &model)?;
let final_results: Vec<String> = future::join_all(second_agg_futures)
.await
.into_iter()
.collect::<Result<Vec<(String, Context)>>>()?
.into_iter()
.map(|(text, _ctx)| text)
.collect();

// --- Stage 4: Collect Final Results ---
let results = future::join_all(second_aggregation_tasks).await;
results.into_iter().collect::<Result<Vec<_>>>()
Ok(final_results)
}

#[inferlet::main]
Expand Down Expand Up @@ -153,6 +178,7 @@ async fn main(input: Input) -> Result<String> {
&question,
proposal_tokens,
aggregation_tokens,
&model,
)
.await?;

Expand Down
2 changes: 1 addition & 1 deletion server/build.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -137,7 +137,7 @@ fn build_portable() {
// 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.
.define("GGML_CUDA_GRAPHS", "ON")
.define("GGML_CUDA_GRAPHS", "OFF")
.define("GGML_STATIC", "ON")
// llama.cpp b8994 added a multi-GPU NCCL allreduce path inside
// ggml-cuda (`GGML_CUDA_NCCL`, default ON when NCCL is present
Expand Down