PR4: on-demand CUDA graph capture extension

apps/cli/main.cpp

@@ -283,6 +283,7 @@ int main(int argc, char** argv) {
         engine_options.speculative    = cli.speculative;
         engine_options.enable_vision  = cli.enable_vision;
         engine_options.use_cuda_graph = cli.use_cuda_graph;
+        engine_options.graph_capture_ceiling = cli.graph_capture_ceiling;
         // One CLI invocation owns exactly one request, so retained cross-request context has no
         // consumer and must not reserve an extra Device StateImage or run terminal capture.
         engine_options.context_cache.enabled                = false;

apps/cli/options.cpp

@@ -85,7 +85,7 @@ std::string usage_text(const char* argv0) {
            "       [--stop-token-id N]... [--stop <text>]... [--reasoning-stop <text>]...\n"
            "       [--raw-output] [--print-token-ids] [--no-thinking] [--thinking-budget N]\n"
            "       [--reasoning-effort low|medium|xhigh] [--vision]\n"
-           "       [--no-cuda-graph]\n"
+           "       [--no-cuda-graph] [--graph-capture-ceiling N]\n"
            "\n"
            "Streams answer content to stdout and reasoning plus diagnostics to stderr.\n"
            "Structured message content accepts text, image/image_url, and video/video_url parts;\n"
@@ -152,6 +152,8 @@ Options parse_options(int argc, char** argv) {
             options.reasoning_effort = parse_reasoning_effort(value(arg));
         } else if (arg == "--vision") {
             options.enable_vision = true;
+        } else if (arg == "--graph-capture-ceiling") {
+            options.graph_capture_ceiling = parse_u32(value(arg), "graph-capture-ceiling");
         } else if (arg == "--no-cuda-graph") {
             options.use_cuda_graph = false;
         } else if (arg == "--stop-token-id") {

apps/cli/options.h

@@ -27,6 +27,7 @@ struct Options {
     SpeculativeOptions speculative;
     bool enable_vision  = false;
     bool use_cuda_graph = true;
+    std::uint32_t graph_capture_ceiling = 0;
 
     bool raw_output      = false;
     bool print_token_ids = false;

include/ninfer/types.h

@@ -122,6 +122,11 @@ struct EngineOptions {
     std::uint32_t media_preprocess_threads = 0;
     bool enable_vision                     = false;
     bool use_cuda_graph                    = true;
+    // On-demand graph capture: 0 (default) captures the full decode ladder at
+    // startup; a positive value captures only segments fully below it and the
+    // runtime extends the family on demand as the decode frontier grows past
+    // each captured segment (one capture per growth crossing).
+    std::uint32_t graph_capture_ceiling = 0;
     ContextCacheOptions context_cache;
     ContextCostOptions context_cost;
     LoadProgress load_progress;

src/targets/qwen3_6/impl/runtime/layouts.h

@@ -79,6 +79,7 @@ struct SequencePlanningInputs {
     ProposalHead proposal_head             = ProposalHead::Full;
     StartupFeatures features;
     bool use_cuda_graph = true;
+    std::uint32_t graph_capture_ceiling = 0;
     bool causal_scoring = false;
     int device          = 0;
     ContextCacheOptions context_cache;
@@ -103,6 +104,7 @@ struct SequencePlanImpl<NINFER_QWEN36_VARIANT> {
     ProposalHead proposal_head             = ProposalHead::Full;
     StartupFeatures features;
     bool use_cuda_graph = true;
+    std::uint32_t graph_capture_ceiling = 0;
     bool causal_scoring = false;
     int device          = 0;
     ContextCacheOptions context_cache;

src/targets/qwen3_6/impl/runtime/layouts_impl.h

@@ -137,7 +137,7 @@ PersistentLayout persistent_layout(const SequencePlanImpl& plan) {
                      .kv_dtype                  = plan.kv_dtype,
                      .kv_quant_group            = plan.kv_quant_group,
                      .enable_mtp                = plan.features.mtp(),
-                     .kv_table_rows             = static_cast<std::int32_t>(plan.max_concurrency),
+                     .kv_table_rows             = static_cast<std::int32_t>(plan.max_concurrency + 1),
                      .text_physical_page_groups = physical_pages,
                      .mtp_physical_page_groups  = mtp_physical_pages,
                  });
@@ -199,7 +199,7 @@ PersistentLayout persistent_layout(const SequencePlanImpl& plan) {
                     builder,
                     KVExecutionTableSpec{
                         .logical_page_capacity = logical_pages,
-                        .table_rows            = static_cast<std::int32_t>(plan.max_concurrency),
+                        .table_rows            = static_cast<std::int32_t>(plan.max_concurrency + 1),
                     }),
                 .layers      = 1,
                 .max_context = plan.capacity,
@@ -671,6 +671,7 @@ std::unique_ptr<SequencePlanImpl> build_sequence_candidate(const SequencePlannin
     impl->features            = inputs.features;
     impl->use_cuda_graph      = inputs.use_cuda_graph;
     impl->causal_scoring      = inputs.causal_scoring;
+    impl->graph_capture_ceiling = inputs.graph_capture_ceiling;
     impl->device              = inputs.device;
     impl->context_cache       = inputs.context_cache;
     impl->kv_dtype            = inputs.kv_dtype;
@@ -745,6 +746,7 @@ make_sequence_planner_impl(DeviceContext& device, const EngineOptions& options,
         .proposal_head       = options.speculative.proposal_head,
         .features            = qwen3_6::startup_features(options),
         .use_cuda_graph      = options.use_cuda_graph,
+        .graph_capture_ceiling = options.graph_capture_ceiling,
         .causal_scoring      = options.purpose == EnginePurpose::CausalScoring,
         .device              = options.device,
         .context_cache       = options.context_cache,

src/targets/qwen3_6/impl/runtime/program.h

@@ -12,6 +12,7 @@
 #include <ninfer/targets/qwen3_6/prepared_prompt.h>
 
 #include "targets/qwen3_6/impl/runtime/layouts.h"
+#include "targets/qwen3_6/impl/runtime/schedule.h"
 #include "targets/qwen3_6/impl/runtime/dflash_context.h"
 #include "targets/qwen3_6/impl/runtime/host_kv_extent_store.h"
 #include "targets/qwen3_6/impl/runtime/logical_kv_store.h"
@@ -687,6 +688,11 @@ class ProgramImplCore {
     qwen3_6::DFlashDecodeEgress* dflash_host_egress   = nullptr;
 
     std::size_t workspace_logical_peak_bytes = 0;
+
+    // On-demand graph capture state (see DecodeGraphFamily comment).
+    std::uint32_t graph_capture_ceiling = 0;
+    void extend_ordinary_graphs(std::uint32_t batch_size, std::uint32_t frontier);
+    schedule::ExecutionCore make_execution_core();
     std::size_t vision_handoff_peak_bytes    = 0;
 
 private:

src/targets/qwen3_6/impl/runtime/program_impl.h

@@ -616,6 +616,19 @@ DecodeGraphProfile& select_graph_profile(DecodeGraphFamily& family, std::uint32_
     return *it;
 }
 
+// True when no profile of this batch covers the frontier — the caller's cue
+// to extend the family on demand (on-demand graph capture).
+[[nodiscard]] inline bool graph_profile_missing(const DecodeGraphFamily& family,
+                                                std::uint32_t batch_size,
+                                                std::uint32_t frontier) noexcept {
+    return std::none_of(family.profiles.begin(), family.profiles.end(),
+                        [&](const DecodeGraphProfile& profile) {
+                            return profile.batch_size == batch_size &&
+                                   profile.min_execution_frontier <= frontier &&
+                                   frontier <= profile.max_execution_frontier;
+                        });
+}
+
 void validate_graph_profiles(const std::vector<GraphExecutionProfile>& profiles,
                              std::uint32_t max_frontier, const char* label) {
     if (profiles.empty() || profiles.front().min != 0 || profiles.back().max != max_frontier) {
@@ -727,6 +740,7 @@ ProgramImplCore::ProgramImplCore(const LoadedModelData& model_in, const Sequence
       speculative_backend(plan.speculative_backend), kv_dtype(plan.kv_dtype),
       kv_quant_group(plan.kv_quant_group), proposal_head(plan.proposal_head),
       vision_enabled(plan.features.vision), use_cuda_graph(plan.use_cuda_graph),
+      graph_capture_ceiling(plan.graph_capture_ceiling),
       causal_scoring(plan.causal_scoring), kv_payload_bytes(plan.persistent.kv_payload_bytes),
       graph_allowance_bytes(plan.graph_allowance_bytes), workspace_plan(plan.workspace),
       persistent(plan.persistent.bytes), workspace_storage(plan.workspace.capacity),
@@ -10337,6 +10351,18 @@ void ProgramImplCore::prepare_graphs() {
         const auto ordinary_profiles = ordinary_graph_profiles(capacity);
         validate_graph_profiles(ordinary_profiles, capacity - 1, "ordinary");
         const std::uint32_t ordinary_batch_limit = max_concurrency;
+        // On-demand capture: with a positive ceiling, startup captures only
+        // the segments fully below it; decode extends on growth crossings and
+        // full coverage is revalidated after each extension.
+        std::vector<GraphExecutionProfile> startup_profiles;
+        if (graph_capture_ceiling == 0) {
+            startup_profiles = ordinary_profiles;
+        } else {
+            for (const GraphExecutionProfile& planned : ordinary_profiles) {
+                if (planned.max <= graph_capture_ceiling) { startup_profiles.push_back(planned); }
+            }
+            if (startup_profiles.empty()) { startup_profiles.push_back(ordinary_profiles.front()); }
+        }
         schedule::OrdinaryBatchContext ordinary_state{
             execution_core(),      decoder->text_kv,
             *io.ordinary,          *ordinary_host_ingress,
@@ -10348,9 +10374,9 @@ void ProgramImplCore::prepare_graphs() {
                                         nullptr);
         device.synchronize();
 
-        ordinary_graphs.profiles.reserve(ordinary_profiles.size() * ordinary_batch_limit);
+        ordinary_graphs.profiles.reserve(startup_profiles.size() * ordinary_batch_limit);
         for (std::uint32_t batch_size = 1; batch_size <= ordinary_batch_limit; ++batch_size) {
-            for (const GraphExecutionProfile planned : ordinary_profiles) {
+            for (const GraphExecutionProfile planned : startup_profiles) {
                 ordinary_graphs.profiles.emplace_back();
                 DecodeGraphProfile& profile    = ordinary_graphs.profiles.back();
                 profile.batch_size             = batch_size;
@@ -10365,6 +10391,9 @@ void ProgramImplCore::prepare_graphs() {
                                                         envelope, profile.definition);
             }
         }
+        if (graph_capture_ceiling == 0) {
+            validate_graph_profiles(ordinary_profiles, capacity - 1, "ordinary");
+        }
     }
 
     if (speculative_backend == SpeculativeBackend::Mtp) {
@@ -10497,6 +10526,76 @@ void ProgramImplCore::prepare_graphs() {
     release_capture_rows(*text_kv_addresses, text_capture_allocations);
 }
 
+schedule::ExecutionCore ProgramImplCore::make_execution_core() {
+    return schedule::ExecutionCore{device,
+                                   model,
+                                   work,
+                                   state_images->linear(),
+                                   replay_records ? &*replay_records : nullptr,
+                                   io,
+                                   prefill_hidden,
+                                   prefill_chunk,
+                                   proposal_head};
+}
+
+// On-demand graph capture: capture the missing ordinary-family segments that
+// cover `frontier` for one batch size. One segment per growth crossing; each
+// is captured exactly once and the family coverage check revalidates. The
+// capture reuses prepare_graphs' dummy-page machinery through the address
+// store: one transient address space on a dedicated execution row whose
+// private page is repeated across the whole table, so arbitrary envelopes
+// read/write valid addresses without disturbing any live request.
+void ProgramImplCore::extend_ordinary_graphs(std::uint32_t batch_size,
+                                             std::uint32_t frontier) {
+    const auto ordinary_profiles = ordinary_graph_profiles(capacity);
+    std::vector<GraphExecutionProfile> missing;
+    for (const GraphExecutionProfile& planned : ordinary_profiles) {
+        const bool covered =
+            std::any_of(ordinary_graphs.profiles.begin(), ordinary_graphs.profiles.end(),
+                        [&](const DecodeGraphProfile& profile) {
+                            return profile.batch_size == batch_size &&
+                                   profile.min_execution_frontier == planned.min &&
+                                   profile.max_execution_frontier == planned.max;
+                        });
+        if (!covered && planned.min <= frontier) { missing.push_back(planned); }
+    }
+    if (missing.empty()) { return; }
+    // Row max_concurrency is the dedicated never-bound capture row (the
+    // address store is sized max_concurrency + 1); request rows 0..C-1 stay
+    // untouched during runtime capture.
+    std::optional<KVAddressSpaceHandle> allocation =
+        text_kv_addresses->create_active(1, static_cast<std::int32_t>(max_concurrency));
+    if (!allocation) { throw std::bad_alloc(); }
+    text_kv_addresses->materialize_to_tokens(*allocation, 1, device.stream);
+    decoder->text_kv.execution_tables().publish_repeated(
+        text_kv_addresses->execution_row(*allocation).handle(),
+        text_kv_addresses->physical_page(*allocation, 0),
+        decoder->text_kv.execution_tables().logical_page_capacity(), device.stream);
+    device.synchronize();
+
+    schedule::OrdinaryBatchContext ordinary_state{
+        make_execution_core(), decoder->text_kv,
+        *io.ordinary,          *ordinary_host_ingress,
+        *ordinary_host_egress, state_images->continuation_hidden_store()};
+    for (const GraphExecutionProfile& planned : missing) {
+        ordinary_graphs.profiles.emplace_back();
+        DecodeGraphProfile& profile    = ordinary_graphs.profiles.back();
+        profile.batch_size             = batch_size;
+        profile.min_execution_frontier = planned.min;
+        profile.max_execution_frontier = planned.max;
+        profile.topology_class         = planned.topology_class * max_concurrency + (batch_size - 1U);
+        const ops::CausalAttentionExecutionEnvelope envelope{planned.min + 1, planned.max + 1};
+        schedule::capture_ordinary_decode_batch(ordinary_state,
+                                                static_cast<std::int32_t>(batch_size), envelope,
+                                                profile.definition);
+    }
+    if (text_kv_addresses->active(*allocation)) { text_kv_addresses->deactivate(*allocation); }
+    (void)text_kv_addresses->release(*allocation);
+    std::fprintf(stderr, "[graphs] extended ordinary family: batch %u +%zu segments through "
+                         "frontier %u\n",
+                 batch_size, missing.size(), frontier);
+}
+
 void ProgramImplCore::install_sampling(SequenceState& sequence, RequestControl& request,
                                        const ops::SamplingConfig& config) {
     Tensor counts = token_counts.slice(1, static_cast<std::int32_t>(sequence.lane), 1)
@@ -10958,6 +11057,16 @@ ProgramImplCore::decode_ordinary_batch(std::span<const std::uint32_t> lanes,
         DecodeGraphExecutable* executable = nullptr;
         ops::CausalAttentionExecutionEnvelope envelope{maximum_frontier + 1, maximum_frontier + 1};
         if (use_cuda_graph) {
+            // On-demand capture: with a startup ceiling, growth past the
+            // captured segments extends the family once per crossing here.
+            if (graph_capture_ceiling != 0 &&
+                graph_profile_missing(ordinary_graphs,
+                                      static_cast<std::uint32_t>(lanes.size()),
+                                      maximum_frontier)) {
+                (void)cudaStreamSynchronize(device.stream);
+                extend_ordinary_graphs(static_cast<std::uint32_t>(lanes.size()),
+                                       maximum_frontier);
+            }
             DecodeGraphProfile& profile =
                 select_graph_profile(ordinary_graphs, static_cast<std::uint32_t>(lanes.size()),
                                      maximum_frontier, "ordinary batch");