diff --git a/driver/cuda/CMakeLists.txt b/driver/cuda/CMakeLists.txt index 7fdbad989..b38b1d04f 100644 --- a/driver/cuda/CMakeLists.txt +++ b/driver/cuda/CMakeLists.txt @@ -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= 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 - $<$: - --extended-lambda - --expt-relaxed-constexpr - >) -set_target_properties(pie_flashinfer_cutlass_moe PROPERTIES - POSITION_INDEPENDENT_CODE ON) -target_sources(pie_driver_cuda_lib PRIVATE - $) +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= 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 + $<$: + --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_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() # ── Marlin (W4A16 / MXFP4 GEMM via vendored vLLM kernel) ────────────── # Built by default because GPT-OSS MXFP4 auto policy should use native diff --git a/driver/portable/src/graph_common.hpp b/driver/portable/src/graph_common.hpp index dccc14638..54081516c 100644 --- a/driver/portable/src/graph_common.hpp +++ b/driver/portable/src/graph_common.hpp @@ -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 diff --git a/inferlets/graph-of-thought/src/lib.rs b/inferlets/graph-of-thought/src/lib.rs index 41c4b3967..55df32f95 100644 --- a/inferlets/graph-of-thought/src/lib.rs +++ b/inferlets/graph-of-thought/src/lib.rs @@ -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, @@ -31,14 +42,54 @@ 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: \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> + '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: ", + 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. + } + Ok(tasks) +} /// Main logic for running the hierarchical aggregation workflow. async fn run_hierarchical_aggregation( @@ -46,13 +97,14 @@ async fn run_hierarchical_aggregation( question: &str, proposal_tokens: Vec, aggregation_tokens: usize, + model: &Model, ) -> Result> { - // --- 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()?; @@ -66,62 +118,35 @@ async fn run_hierarchical_aggregation( Ok::<_, String>((proposal_text, ctx)) }) }) - .collect::>>()? + .collect::>>()?; + + // 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::>(); - - // --- 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::>()?; + + // --- 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::>()?; - // --- 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 = future::join_all(second_agg_futures) + .await + .into_iter() + .collect::>>()? + .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::>>() + Ok(final_results) } #[inferlet::main] @@ -153,6 +178,7 @@ async fn main(input: Input) -> Result { &question, proposal_tokens, aggregation_tokens, + &model, ) .await?; diff --git a/server/build.rs b/server/build.rs index 75ef8490d..6195dded4 100644 --- a/server/build.rs +++ b/server/build.rs @@ -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