diff --git a/apps/cli/main.cpp b/apps/cli/main.cpp index 933192da9b..a0928131ba 100644 --- a/apps/cli/main.cpp +++ b/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; diff --git a/apps/cli/options.cpp b/apps/cli/options.cpp index b5c5798d9a..112b5d1e26 100644 --- a/apps/cli/options.cpp +++ b/apps/cli/options.cpp @@ -85,7 +85,7 @@ std::string usage_text(const char* argv0) { " [--stop-token-id N]... [--stop ]... [--reasoning-stop ]...\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") { diff --git a/apps/cli/options.h b/apps/cli/options.h index 3c0c2e0960..1fc0b9ac04 100644 --- a/apps/cli/options.h +++ b/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; diff --git a/include/ninfer/types.h b/include/ninfer/types.h index 677c381e88..95c2a938f8 100644 --- a/include/ninfer/types.h +++ b/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; diff --git a/src/targets/qwen3_6/impl/runtime/layouts.h b/src/targets/qwen3_6/impl/runtime/layouts.h index 859f4f5d7e..f7514c6bb2 100644 --- a/src/targets/qwen3_6/impl/runtime/layouts.h +++ b/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 { 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; diff --git a/src/targets/qwen3_6/impl/runtime/layouts_impl.h b/src/targets/qwen3_6/impl/runtime/layouts_impl.h index b21165ac86..8d86deeb74 100644 --- a/src/targets/qwen3_6/impl/runtime/layouts_impl.h +++ b/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(plan.max_concurrency), + .kv_table_rows = static_cast(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(plan.max_concurrency), + .table_rows = static_cast(plan.max_concurrency + 1), }), .layers = 1, .max_context = plan.capacity, @@ -671,6 +671,7 @@ std::unique_ptr 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, diff --git a/src/targets/qwen3_6/impl/runtime/program.h b/src/targets/qwen3_6/impl/runtime/program.h index 6c0f8d33a0..cd45452354 100644 --- a/src/targets/qwen3_6/impl/runtime/program.h +++ b/src/targets/qwen3_6/impl/runtime/program.h @@ -12,6 +12,7 @@ #include #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: diff --git a/src/targets/qwen3_6/impl/runtime/program_impl.h b/src/targets/qwen3_6/impl/runtime/program_impl.h index 36f04c8bdb..bb02b3e4e1 100644 --- a/src/targets/qwen3_6/impl/runtime/program_impl.h +++ b/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& 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 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 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 allocation = + text_kv_addresses->create_active(1, static_cast(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(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(sequence.lane), 1) @@ -10958,6 +11057,16 @@ ProgramImplCore::decode_ordinary_batch(std::span 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(lanes.size()), + maximum_frontier)) { + (void)cudaStreamSynchronize(device.stream); + extend_ordinary_graphs(static_cast(lanes.size()), + maximum_frontier); + } DecodeGraphProfile& profile = select_graph_profile(ordinary_graphs, static_cast(lanes.size()), maximum_frontier, "ordinary batch");