From 6d4548bca496c25740761355d6948cdcafb5a542 Mon Sep 17 00:00:00 2001 From: justinchuby <223556219+Copilot@users.noreply.github.com> Date: Tue, 28 Jul 2026 10:50:42 -0700 Subject: [PATCH 1/4] feat(kv): one KV capacity mechanism for native CUDA and ORT Justin hit "CUDA KV capacity exceeded: requested context length 4097, configured max_len 4096" on the native backend and asked why it was so conservative when he had VRAM to spare, and whether we did not already have dynamic KV allocation. We did -- on one path only. The ORT shared-buffer GQA path grew in power-of-two buckets via kv_capacity_bucket, while native CUDA pre-allocated once at a hardcoded DEFAULT_CUDA_KV_MAX_LEN of 4096 and never grew. Two mechanisms for the same resource, which is the shape that has produced several bugs here recently; a bigger constant would have left the divergence in place for someone to hit from another angle. The growth policy now lives once, in onnx-genai-kv: kv_capacity_bucket plus a KvCapacityGrowthBackend seam and a shared ensure_kv_capacity driver. Both ORT and native CUDA consume it. The backend seam carries only primitives that genuinely differ -- allocate buffers, copy the prefix, build the mask, invalidate capture, commit -- so adding a backend means implementing primitives rather than re-deriving policy. Growth and CUDA graph capture coexist, as they already did on the ORT side: grow to the next bucket, device-to-device copy the valid prefix, rebuild the mask, release the capture, re-capture at the new size. The comment claiming capture required a frozen physical capacity described the fixed-capacity design rather than a real constraint, and is corrected. Capacity derives from model metadata and is clamped by what the device can actually afford, rather than by a constant. The clamp initially checked only the final KV size, which missed the transient peak during growth where the old and new buffers are both live: on an 8 GB RTX 4060 with DeepSeek at 28,680 bytes/token, the 65536 to 131072 transition needs 5.64 GB against a 5.93 GB budget. It now clamps to the worst transition, giving hard_max_len 110080 for that model on that card instead of the metadata's 131072. Verified on the GPU rather than reasoned about: forced 4 to 8 to 16 growth with real reallocation and re-capture, output stayed coherent -- a corrupted prefix copy shows up as wrong tokens rather than an error, so coherence is the assertion that matters. Reaching the ceiling now produces the capacity error before allocation rather than a CUDA OOM. The error text names where the limit came from, the metadata value, device free and total, KV bytes per token, and the available overrides. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 023d5a3f-c551-4aa8-afbe-a329af15fbdd --- .squad/agents/leon/history.md | 37 ++ Cargo.lock | 1 + crates/onnx-genai-engine/src/engine/load.rs | 4 + .../src/native_decode/cuda.rs | 578 ++++++++++++++++-- .../src/native_decode/load.rs | 61 +- .../src/native_decode/mod.rs | 1 + .../src/native_decode/tests.rs | 93 +++ crates/onnx-genai-kv/src/lib.rs | 237 +++++++ crates/onnx-genai-ort/Cargo.toml | 1 + crates/onnx-genai-ort/src/cuda_rt.rs | 63 +- crates/onnx-genai-ort/src/decode/dynamic.rs | 102 ++-- crates/onnx-genai-ort/src/decode/kv_growth.rs | 27 +- crates/onnx-genai-ort/src/lib.rs | 2 +- .../src/executor/build.rs | 2 +- docs/native-ort-kv-capacity.md | 48 ++ 15 files changed, 1128 insertions(+), 129 deletions(-) create mode 100644 docs/native-ort-kv-capacity.md diff --git a/.squad/agents/leon/history.md b/.squad/agents/leon/history.md index 7d4642de7e..09d93c9abe 100644 --- a/.squad/agents/leon/history.md +++ b/.squad/agents/leon/history.md @@ -160,3 +160,40 @@ Authored the Phi-4-mini bit-exact native-CUDA-versus-ORT 64-token decode lock. F - ๐ŸŸข Approved Sapper's revision. `RewindRunnerPolicy` now keeps public `Engine::rewind_session_to` on `RejectRunnerRewind` while speculative draft alignment, draft rewind, target accept/reject rewind, and overmaterialized cleanup use `AllowRunnerRewind`. - Verified public runner-backed rewind rejects before scheduler/session/token/KV/runner mutation; policy-boundary tests would fail if allow/reject were flipped. Support matrix now matches the public API contract, with prepared/infallible runner rewind recorded as follow-up. - Validation: engine build/fmt/clippy passed; server/CLI builds passed; full engine lib suite now runs locally with ORT 1.27/API 27 (252 passed, 0 failed, 1 ignored); targeted failed-rewind/speculative/checkpoint tests passed. + +## 2026-07-28T08:22:00-07:00 โ€” Native CUDA KV capacity derivation + +- Replaced the hardcoded 4096 native-CUDA KV fallback with capacity resolution from explicit programmatic/env overrides, otherwise `model.max_sequence_length` clamped to an 85% queried CUDA free-memory budget using actual KV bytes/token. +- Kept the native path fixed-capacity: growing would replace device allocations/IoBinding pointers and invalidate CUDA graph capture, unlike ORT shared-buffer GQA's grow/rebind/recapture path. +- Expanded capacity errors to include the source, metadata max length, CUDA free/total bytes when available, KV bytes/token, and both override knobs. +- Validation: `cargo build -p onnx-genai-engine`, `cargo test -p onnx-genai-engine --lib`, `cargo fmt -p onnx-genai-engine -- --check`, and `cargo clippy -p onnx-genai-engine --all-targets -- -D warnings` passed locally; no CUDA device was available for real native-CUDA allocation/capture execution. + +## 2026-07-28T08:53:10-07:00 โ€” Native/ORT KV capacity unification correction + +- Justin rejected the narrower "better native default" framing: native and ORT must share one KV capacity mechanism/capability, not two strategies for the same resource. +- Moved/used `kv_capacity_bucket` as shared policy in `onnx-genai-kv`; ORT imports it and native CUDA now grows by the same bucket policy. +- Native CUDA growth is compatible with graph capture when treated as a capture boundary: invalidate the graph, replace/copy device KV buffers, then re-warm/re-capture on subsequent single-token steps. Real recapture-after-growth still needs GPU validation. +- Added `docs/native-ort-kv-capacity.md` to preserve the rationale and remaining validation gap. + +## 2026-07-28T09:10:46-07:00 โ€” Capture/growth compatibility clarified + +- Justin clarified that dynamic growth and CUDA graph capture are compatible; ORT shared-buffer GQA is the reference, not a counterexample to native growth. +- Updated the native CUDA comments/design note to describe bucket growth as a graph boundary and removed the obsolete framing that fixed capacity might be required for capture. +- Native behavior now matches the ORT mechanism: grow by shared bucket, preserve the valid prefix, invalidate the stale capture, then re-capture on the new bucket. + +## 2026-07-28T09:16:58-07:00 โ€” DRY shared KV grow driver + +- Justin asked for DRY after native mirrored ORT growth. Factored growth orchestration into `onnx-genai-kv::ensure_kv_capacity` with `KvCapacityGrowthBackend`. +- Both ORT shared-buffer KV and native CUDA consume the shared driver; backend implementations only provide primitives for allocation, prefix preservation, mask replacement, capture invalidation, and commit. +- Added policy-order tests proving fallible buffer/mask preparation happens before capture invalidation/commit and failure leaves current capacity untouched. + +## 2026-07-28T10:05:00-07:00 โ€” Native CUDA KV growth verified on local GPU + +- Corrected the earlier bad assumption that this machine had no CUDA device: `nvidia-smi` reports an RTX 4060 Laptop GPU with 8188 MiB. +- With `ONNX_GENAI_EP=cuda`, `ONNX_GENAI_REQUIRE_CUDA=1`, NVIDIA Python package DLL/header paths, and `ONNX_GENAI_KV_MIN_BUCKET=4`, real native CUDA DeepSeek generation grew KV 4โ†’8โ†’16, re-captured at each bucket, and continued coherently. +- Capacity diagnostics were also exercised with `ONNX_GENAI_CUDA_KV_MAX_LEN=4`; the error now reports requested length, configured max/source, metadata max, CUDA free/total bytes, KV bytes/token, and the override/remediation knobs. + +## 2026-07-28T10:26:45-07:00 โ€” CUDA KV clamp corrected for transient growth + +- Justin challenged the "metadata 131072, not clamped" DeepSeek result on the 8 GiB RTX 4060. The first clamp only checked final KV size (`0.85 * free / bytes_per_token`), but growth allocates the new bucket before freeing the old bucket. +- Tightened the default hard maximum to the largest capacity whose worst old+new bucket transition fits inside the CUDA free-memory headroom budget. For the observed DeepSeek numbers (`free=5925502976`, `bytes/token=28680`), that clamps metadata 131072 to 110080 instead of allowing the 65536โ†’131072 transition to exceed the 85% growth budget. diff --git a/Cargo.lock b/Cargo.lock index 455b7e6a2b..6889058fd3 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2303,6 +2303,7 @@ dependencies = [ "minijinja", "minijinja-contrib", "onnx-genai-genai-config", + "onnx-genai-kv", "onnx-genai-metadata", "onnx-genai-ort-sys", "onnx-genai-runtime-config", diff --git a/crates/onnx-genai-engine/src/engine/load.rs b/crates/onnx-genai-engine/src/engine/load.rs index 6c486aee90..5b04068f37 100644 --- a/crates/onnx-genai-engine/src/engine/load.rs +++ b/crates/onnx-genai-engine/src/engine/load.rs @@ -222,6 +222,10 @@ impl Engine { native_device.clone(), governor.weight_offload_host_cache(), metadata.model.as_ref().and_then(|model| model.io.as_ref()), + metadata + .model + .as_ref() + .and_then(|model| model.max_sequence_length), config.decode_precision, ) .map_err(|error| anyhow::anyhow!("Failed to load native decoder session: {error:#}"))? diff --git a/crates/onnx-genai-engine/src/native_decode/cuda.rs b/crates/onnx-genai-engine/src/native_decode/cuda.rs index 588f7c657b..e935bfc30c 100644 --- a/crates/onnx-genai-engine/src/native_decode/cuda.rs +++ b/crates/onnx-genai-engine/src/native_decode/cuda.rs @@ -48,17 +48,34 @@ pub(crate) fn resolve_graph_capture_enabled( structural.is_capture_safe() } -const DEFAULT_CUDA_KV_MAX_LEN: usize = 4096; +const CUDA_KV_MEMORY_BUDGET_NUMERATOR: usize = 85; +const CUDA_KV_MEMORY_BUDGET_DENOMINATOR: usize = 100; #[derive(Clone, Debug, Eq, PartialEq)] pub struct CudaKvDebugStats { pub logical_len: usize, pub max_len: usize, + pub max_len_source: String, pub device_ptrs: Vec, pub kv_transfers: DeviceBindingTransferStats, pub graph: CudaGraphDebugStats, } +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub(crate) struct CudaDeviceMemorySnapshot { + pub(crate) free_bytes: usize, + pub(crate) total_bytes: usize, +} + +#[derive(Clone, Debug, Eq, PartialEq)] +pub(crate) struct CudaKvCapacity { + pub(crate) max_len: usize, + pub(crate) source: String, + pub(crate) metadata_max_len: Option, + pub(crate) device_memory: Option, + pub(crate) bytes_per_token: usize, +} + #[derive(Clone, Debug, Default, Eq, PartialEq)] pub struct CudaGraphDebugStats { pub enabled: bool, @@ -80,18 +97,23 @@ enum DecodeCudaGraphPhase { pub(crate) struct DecodeCudaState { logical_len: usize, + /// Current physical KV bucket for the native CUDA decode session. The hard + /// maximum lives in `capacity.max_len`; this bucket grows on demand via the + /// shared `onnx_genai_kv::kv_capacity_bucket` policy. Growth is a capture + /// boundary: reallocation invalidates the old graph, then the normal + /// warm-up/capture/replay state machine captures the new bucket. pub(crate) max_len: usize, - bindings: Vec, - base_binding_count: usize, - kv_binding_range: std::ops::Range, - auxiliary_binding_range: std::ops::Range, - input_ids_binding: usize, - position_ids_binding: Option, + pub(crate) bindings: Vec, + pub(crate) base_binding_count: usize, + pub(crate) kv_binding_range: std::ops::Range, + pub(crate) auxiliary_binding_range: std::ops::Range, + pub(crate) input_ids_binding: usize, + pub(crate) position_ids_binding: Option, logits_binding: usize, logits_shape: Vec, logits_dtype: DataType, greedy_result: DeviceIoBinding, - graph_enabled: bool, + pub(crate) graph_enabled: bool, /// `true` when the attention-mask binding exposes its *logical* valid length /// (not the padded physical capacity) to at least one consumer that is not a /// capacity-aware kernel โ€” e.g. GLM-5.2's `indexer` arithmetic branch, which @@ -104,6 +126,7 @@ pub(crate) struct DecodeCudaState { graph_captures: u64, graph_replays: u64, graph_fallbacks: u64, + capacity: CudaKvCapacity, pub(crate) graph_fallback_reason: Option, pub(crate) graph_fallback_report: Option, /// Structural reasons, recorded at binding time, why one or more auxiliary @@ -262,23 +285,24 @@ impl NativeDecodeSession { .cuda .as_mut() .context("CUDA decode state is not initialized")?; - if total_len > state.max_len { - bail!( - "CUDA KV capacity exceeded: requested context length {total_len}, configured max_len {} (set ONNX_GENAI_CUDA_KV_MAX_LEN or use load_with_cuda_kv_max_len)", - state.max_len - ); + if total_len > state.capacity.max_len { + bail!("{}", state.capacity_exceeded_error(total_len)); } - // Single-token decode freezes the mask to physical capacity so the step - // is CUDA-graph-capture eligible; multi-token prefill keeps the growing - // logical length (prefix-sensitive causal island). A mask whose logical - // valid length feeds a non-capacity-aware consumer (see + let grew = state.ensure_capacity(&mut self.session, total_len)?; + // Single-token decode freezes the mask to the current physical bucket so + // the step is CUDA-graph-capture eligible. This is not a fixed-capacity + // constraint: when `ensure_capacity` grows the bucket, it invalidates the + // old capture and the existing state machine re-captures against the new + // buffers, matching ORT shared-buffer growth. Multi-token prefill keeps + // the growing logical length (prefix-sensitive causal island). A mask + // whose logical valid length feeds a non-capacity-aware consumer (see // `decode_mask_expose_len`) cannot be frozen and uses `total_len`. let mask_expose = if token_ids.len() == 1 { state.decode_mask_expose_len(total_len) } else { total_len }; - state.extend_mask(past_len, total_len, mask_expose)?; + state.extend_mask(if grew { 0 } else { past_len }, total_len, mask_expose)?; if token_ids.len() == 1 { state.write_decode_inputs(token_ids[0], past_len)?; @@ -387,13 +411,11 @@ impl NativeDecodeSession { .cuda .as_mut() .context("CUDA decode state is not initialized")?; - if total_len > state.max_len { - bail!( - "CUDA KV capacity exceeded: requested context length {total_len}, configured max_len {} (set ONNX_GENAI_CUDA_KV_MAX_LEN or use load_with_cuda_kv_max_len)", - state.max_len - ); + if total_len > state.capacity.max_len { + bail!("{}", state.capacity_exceeded_error(total_len)); } - state.extend_mask(past_len, total_len, total_len)?; + let grew = state.ensure_capacity(&mut self.session, total_len)?; + state.extend_mask(if grew { 0 } else { past_len }, total_len, total_len)?; self.run_cuda_eager_rows( token_ids, past_len, @@ -416,13 +438,15 @@ impl NativeDecodeSession { .cuda .as_mut() .context("CUDA decode state is not initialized")?; - if total_len > state.max_len { - bail!( - "CUDA KV capacity exceeded: requested context length {total_len}, configured max_len {} (set ONNX_GENAI_CUDA_KV_MAX_LEN or use load_with_cuda_kv_max_len)", - state.max_len - ); + if total_len > state.capacity.max_len { + bail!("{}", state.capacity_exceeded_error(total_len)); } - state.extend_mask(past_len, total_len, state.decode_mask_expose_len(total_len))?; + let grew = state.ensure_capacity(&mut self.session, total_len)?; + state.extend_mask( + if grew { 0 } else { past_len }, + total_len, + state.decode_mask_expose_len(total_len), + )?; state.write_decode_inputs(token_id, past_len)?; if let Err(error) = state.run_one_token(&mut self.session, &self.trace) { let diagnosis = diagnose_native_failure(&self.session, &error.to_string()); @@ -442,6 +466,97 @@ impl NativeDecodeSession { } impl DecodeCudaState { + pub(crate) fn kv_bytes_per_token( + session: &InferenceSession, + present_to_past: &HashMap, + ) -> anyhow::Result { + let mut bytes = std::mem::size_of::(); + for past in present_to_past.values() { + let meta = session + .inputs() + .iter() + .find(|meta| meta.name == *past) + .with_context(|| format!("missing CUDA KV input metadata for '{past}'"))?; + if !matches!(meta.dtype, DataType::Float32 | DataType::Float16) || meta.shape.len() != 4 + { + bail!( + "CUDA KV input '{past}' must be rank-4 f32 or f16, got {:?} {:?}", + meta.dtype, + meta.shape + ); + } + let mut elements_per_token = 1usize; + for (axis, dim) in meta.shape.iter().copied().enumerate() { + let dim = if axis == 0 || axis == 2 { + 1 + } else if let Dim::Static(value) = dim { + value + } else { + bail!( + "cannot infer CUDA KV dimension {axis} for '{past}' shape {:?}", + meta.shape + ); + }; + elements_per_token = elements_per_token.checked_mul(dim).with_context(|| { + format!( + "CUDA KV bytes-per-token overflow for '{past}' shape {:?}", + meta.shape + ) + })?; + } + bytes = bytes + .checked_add( + meta.dtype + .checked_storage_bytes(elements_per_token) + .with_context(|| { + format!( + "CUDA KV bytes-per-token overflow for '{past}' shape {:?}", + meta.shape + ) + })?, + ) + .with_context(|| { + format!( + "CUDA KV bytes-per-token overflow for '{past}' shape {:?}", + meta.shape + ) + })?; + } + Ok(bytes) + } + + fn capacity_exceeded_error(&self, requested: usize) -> String { + cuda_kv_capacity_exceeded_message(requested, &self.capacity) + } + + fn ensure_capacity( + &mut self, + session: &mut InferenceSession, + required: usize, + ) -> anyhow::Result { + let mut backend = NativeCudaCapacityBackend { + state: self, + session, + }; + match onnx_genai_kv::ensure_kv_capacity(&mut backend, required)? { + onnx_genai_kv::KvCapacityGrowth::Unchanged => Ok(false), + onnx_genai_kv::KvCapacityGrowth::Grew { + old_capacity, + new_capacity, + valid_len, + } => { + tracing::info!( + old_len = old_capacity, + new_len = new_capacity, + valid_len, + hard_max_len = backend.state.capacity.max_len, + "grew native CUDA KV capacity bucket" + ); + Ok(true) + } + } + } + /// Collect the symbolic dimension ids that the native decoder structurally /// pins to `1` at decode time. Batch (axis 0 of every input) and query-seq /// (the remaining `input_ids` / `position_ids` axes, which are bound to a @@ -515,9 +630,10 @@ impl DecodeCudaState { session: &mut InferenceSession, io: DecodeCudaIo<'_>, present_to_past: &HashMap, - max_len: usize, + capacity: CudaKvCapacity, graph_enabled: bool, ) -> anyhow::Result { + let max_len = onnx_genai_kv::kv_capacity_bucket(0, capacity.max_len); let mut mask = session.allocate_device_binding( io.attention_mask, None::, @@ -818,6 +934,7 @@ impl DecodeCudaState { graph_captures: 0, graph_replays: 0, graph_fallbacks: 0, + capacity, graph_fallback_reason: None, graph_fallback_report: None, auxiliary_bind_declines: declined_auxiliary, @@ -830,19 +947,21 @@ impl DecodeCudaState { /// Write the valid `1`s for keys `[start, end)` and set the mask's exposed /// logical length. `expose_len` is the last-dim extent the graph's mask /// island (and hence the Attention kernel) sees: for a single-token decode - /// step it is the fixed physical capacity (`max_len`), which freezes the + /// step it is the current physical bucket (`max_len`), which freezes the /// island to a shape-static `[1,1,1,max_len]` additive bias (correct for a /// single query row โ€” verified: the padding suffix maps to `-inf`, the valid - /// prefix to `0`), so the decode step carries no growing logical input and - /// stays CUDA-graph-capture eligible. Multi-token prefill passes `end` - /// (the growing valid length) because the causal island is prefix-sensitive - /// for `q_seq > 1` and must see the exact logical length. + /// prefix to `0`). Bucket growth invalidates and re-captures the graph at the + /// new bucket, exactly like ORT shared-buffer KV, so capture eligibility does + /// not require a session-wide fixed capacity. Multi-token prefill passes + /// `end` (the growing valid length) because the causal island is + /// prefix-sensitive for `q_seq > 1` and must see the exact logical length. /// The last-dim extent to expose for the attention mask on a single-token - /// decode step. Frozen to the physical capacity (`max_len`) so the step stays - /// CUDA-graph-capture eligible โ€” *unless* the mask binding exposes its logical - /// valid length to a non-capacity-aware consumer (`mask_exposes_logical`), in - /// which case the true valid length (`total_len`) must be used or the padded - /// width leaks into that consumer's arithmetic (see [`Self::mask_exposes_logical`]). + /// decode step. Frozen to the current physical bucket (`max_len`) so the step + /// stays CUDA-graph-capture eligible โ€” *unless* the mask binding exposes its + /// logical valid length to a non-capacity-aware consumer + /// (`mask_exposes_logical`), in which case the true valid length (`total_len`) + /// must be used or the padded width leaks into that consumer's arithmetic + /// (see [`Self::mask_exposes_logical`]). fn decode_mask_expose_len(&self, total_len: usize) -> usize { if self.mask_exposes_logical { total_len @@ -918,6 +1037,11 @@ impl DecodeCudaState { } self.graph_captures += 1; self.graph_phase = DecodeCudaGraphPhase::Ready; + tracing::debug!( + capacity = self.max_len, + captures = self.graph_captures, + "native CUDA decode graph captured" + ); } DeviceGraphCaptureResult::NotCapturable(report) => { self.graph_fallbacks += 1; @@ -1034,6 +1158,7 @@ impl DecodeCudaState { CudaKvDebugStats { logical_len: self.logical_len, max_len: self.max_len, + max_len_source: self.capacity.source.clone(), device_ptrs, kv_transfers: transfers, graph: CudaGraphDebugStats { @@ -1048,7 +1173,129 @@ impl DecodeCudaState { } } -pub(crate) fn cuda_kv_max_len_from_env() -> anyhow::Result { +struct NativeCudaGrownBuffers { + replacements: Vec<(usize, DeviceIoBinding)>, +} + +struct NativeCudaCapacityBackend<'a> { + state: &'a mut DecodeCudaState, + session: &'a mut InferenceSession, +} + +impl onnx_genai_kv::KvCapacityGrowthBackend for NativeCudaCapacityBackend<'_> { + type Error = anyhow::Error; + type GrownBuffers = NativeCudaGrownBuffers; + type GrownMask = DeviceIoBinding; + + fn current_capacity(&self) -> usize { + self.state.max_len + } + + fn hard_max_capacity(&self) -> usize { + self.state.capacity.max_len + } + + fn valid_len(&self) -> usize { + self.state.logical_len + } + + fn capacity_exceeded(&self, required: usize) -> Self::Error { + anyhow::anyhow!(self.state.capacity_exceeded_error(required)) + } + + fn build_grown_buffers( + &mut self, + new_capacity: usize, + valid_len: usize, + ) -> anyhow::Result { + native_cuda_device_barrier(self.session)?; + let mut replacements = Vec::with_capacity(self.state.kv_binding_range.len()); + for index in self.state.kv_binding_range.clone() { + let old = &self.state.bindings[index]; + let mut physical_shape = old.physical_shape().to_vec(); + physical_shape[2] = new_capacity; + let mut logical_shape = physical_shape.clone(); + logical_shape[2] = valid_len; + let new_binding = self.session.allocate_device_binding( + old.input_name().to_owned(), + old.output_name().map(str::to_owned), + old.dtype, + physical_shape.clone(), + logical_shape, + )?; + let dst = new_binding.device_ptr() as usize; + let src = old.device_ptr() as usize; + let total_bytes = + checked_shape_bytes(&physical_shape, old.dtype).with_context(|| { + format!( + "CUDA KV grow allocation size overflow for '{}' shape {:?}", + old.input_name(), + physical_shape + ) + })?; + native_cuda_memset_zero(dst, total_bytes)?; + copy_kv_prefix_device_to_device( + dst, + src, + old.physical_shape(), + &physical_shape, + 2, + valid_len, + old.dtype.checked_storage_bytes(1).with_context(|| { + format!( + "CUDA KV grow copy element size overflow for '{}'", + old.input_name() + ) + })?, + )?; + replacements.push((index, new_binding)); + } + Ok(NativeCudaGrownBuffers { replacements }) + } + + fn build_grown_mask( + &mut self, + new_capacity: usize, + valid_len: usize, + ) -> anyhow::Result> { + let mut new_mask = self.session.allocate_device_binding( + self.state.bindings[0].input_name().to_owned(), + self.state.bindings[0].output_name().map(str::to_owned), + DataType::Int64, + vec![1, new_capacity], + vec![1, valid_len], + )?; + native_cuda_memset_zero( + new_mask.device_ptr() as usize, + new_capacity * std::mem::size_of::(), + )?; + new_mask.set_logical_shape(vec![1, valid_len])?; + Ok(Some(new_mask)) + } + + fn invalidate_capture(&mut self) -> anyhow::Result<()> { + native_cuda_device_barrier(self.session)?; + self.state.invalidate_graph(self.session) + } + + fn commit_grown_capacity( + &mut self, + new_capacity: usize, + grown_buffers: Self::GrownBuffers, + grown_mask: Option, + ) -> anyhow::Result<()> { + if let Some(mask) = grown_mask { + self.state.bindings[0] = mask; + } + for (index, binding) in grown_buffers.replacements { + self.state.bindings[index] = binding; + } + self.state.max_len = new_capacity; + Ok(()) + } +} + +pub(crate) fn cuda_kv_max_len_from_env() -> anyhow::Result> { match std::env::var("ONNX_GENAI_CUDA_KV_MAX_LEN") { Ok(value) => { let parsed = value.trim().parse::().with_context(|| { @@ -1057,9 +1304,252 @@ pub(crate) fn cuda_kv_max_len_from_env() -> anyhow::Result { if parsed == 0 { bail!("ONNX_GENAI_CUDA_KV_MAX_LEN must be greater than zero"); } - Ok(parsed) + Ok(Some(parsed)) } - Err(std::env::VarError::NotPresent) => Ok(DEFAULT_CUDA_KV_MAX_LEN), + Err(std::env::VarError::NotPresent) => Ok(None), Err(error) => Err(error).context("read ONNX_GENAI_CUDA_KV_MAX_LEN"), } } + +fn checked_shape_bytes(shape: &[usize], dtype: DataType) -> Option { + let elements = shape + .iter() + .try_fold(1usize, |product, &dim| product.checked_mul(dim))?; + dtype.checked_storage_bytes(elements) +} + +#[cfg(feature = "cuda")] +fn native_cuda_device_barrier(session: &InferenceSession) -> anyhow::Result<()> { + let _guard = onnx_genai_ort::cuda_rt::DeviceGuard::set(session.device_id().index as i32)?; + onnx_genai_ort::cuda_rt::device_synchronize()?; + Ok(()) +} + +#[cfg(not(feature = "cuda"))] +fn native_cuda_device_barrier(_session: &InferenceSession) -> anyhow::Result<()> { + bail!("native CUDA KV growth requires the onnx-genai-engine `cuda` feature") +} + +#[cfg(feature = "cuda")] +fn native_cuda_memset_zero(dst: usize, bytes: usize) -> anyhow::Result<()> { + onnx_genai_ort::cuda_rt::memset_zero(dst, bytes)?; + Ok(()) +} + +#[cfg(not(feature = "cuda"))] +fn native_cuda_memset_zero(_dst: usize, _bytes: usize) -> anyhow::Result<()> { + bail!("native CUDA KV growth requires the onnx-genai-engine `cuda` feature") +} + +#[cfg(feature = "cuda")] +fn copy_kv_prefix_device_to_device( + dst: usize, + src: usize, + old_shape: &[usize], + new_shape: &[usize], + seq_axis: usize, + valid_len: usize, + elem_size: usize, +) -> anyhow::Result<()> { + if valid_len == 0 { + return Ok(()); + } + if old_shape.len() != new_shape.len() || seq_axis >= old_shape.len() { + bail!("invalid CUDA KV grow copy shapes {old_shape:?} -> {new_shape:?} on axis {seq_axis}"); + } + let old_cap = old_shape[seq_axis]; + let new_cap = new_shape[seq_axis]; + if valid_len > old_cap || valid_len > new_cap { + bail!( + "invalid CUDA KV grow copy valid prefix {valid_len} for shapes {old_shape:?} -> {new_shape:?}" + ); + } + let blocks = old_shape[..seq_axis] + .iter() + .try_fold(1usize, |product, &dim| product.checked_mul(dim)) + .context("CUDA KV grow copy block count overflow")?; + let inner = old_shape[seq_axis + 1..] + .iter() + .try_fold(1usize, |product, &dim| product.checked_mul(dim)) + .context("CUDA KV grow copy inner size overflow")?; + let segment_bytes = valid_len + .checked_mul(inner) + .and_then(|elements| elements.checked_mul(elem_size)) + .context("CUDA KV grow copy segment size overflow")?; + let old_stride = old_cap + .checked_mul(inner) + .and_then(|elements| elements.checked_mul(elem_size)) + .context("CUDA KV grow copy old stride overflow")?; + let new_stride = new_cap + .checked_mul(inner) + .and_then(|elements| elements.checked_mul(elem_size)) + .context("CUDA KV grow copy new stride overflow")?; + for block in 0..blocks { + let src_offset = block + .checked_mul(old_stride) + .context("CUDA KV grow copy source offset overflow")?; + let dst_offset = block + .checked_mul(new_stride) + .context("CUDA KV grow copy destination offset overflow")?; + onnx_genai_ort::cuda_rt::memcpy_device_to_device( + dst + dst_offset, + src + src_offset, + segment_bytes, + )?; + } + Ok(()) +} + +#[cfg(not(feature = "cuda"))] +fn copy_kv_prefix_device_to_device( + _dst: usize, + _src: usize, + _old_shape: &[usize], + _new_shape: &[usize], + _seq_axis: usize, + _valid_len: usize, + _elem_size: usize, +) -> anyhow::Result<()> { + bail!("native CUDA KV growth requires the onnx-genai-engine `cuda` feature") +} + +pub(crate) fn resolve_cuda_kv_capacity( + programmatic_max_len: Option, + env_max_len: Option, + metadata_max_len: Option, + bytes_per_token: usize, + device_memory: Option, +) -> anyhow::Result { + if bytes_per_token == 0 { + bail!("CUDA KV bytes per token must be greater than zero"); + } + let memory_budget_tokens = device_memory.and_then(|memory| { + let whole = (memory.free_bytes / CUDA_KV_MEMORY_BUDGET_DENOMINATOR) + .checked_mul(CUDA_KV_MEMORY_BUDGET_NUMERATOR)?; + let remainder = (memory.free_bytes % CUDA_KV_MEMORY_BUDGET_DENOMINATOR) + .checked_mul(CUDA_KV_MEMORY_BUDGET_NUMERATOR)? + / CUDA_KV_MEMORY_BUDGET_DENOMINATOR; + whole + .checked_add(remainder) + .and_then(|budget| budget.checked_div(bytes_per_token)) + }); + let memory_budget_len = memory_budget_tokens.map(max_kv_capacity_for_growth_budget); + let (configured, source) = if let Some(max_len) = programmatic_max_len { + (max_len, "programmatic load_with_cuda_kv_max_len".to_owned()) + } else if let Some(max_len) = env_max_len { + (max_len, "ONNX_GENAI_CUDA_KV_MAX_LEN".to_owned()) + } else if let Some(max_len) = metadata_max_len { + let metadata_len = max_len; + let max_len = memory_budget_len.map_or(metadata_len, |budget| metadata_len.min(budget)); + ( + max_len, + match (device_memory, memory_budget_len) { + (Some(memory), Some(budget)) if budget < metadata_len => format!( + "model.max_sequence_length clamped by CUDA free-memory growth budget (free={} bytes, total={} bytes, budget={} tokens at {} bytes/token)", + memory.free_bytes, memory.total_bytes, budget, bytes_per_token + ), + (Some(memory), Some(budget)) => format!( + "model.max_sequence_length within CUDA free-memory growth budget (free={} bytes, total={} bytes, budget={} tokens at {} bytes/token)", + memory.free_bytes, memory.total_bytes, budget, bytes_per_token + ), + _ => "model.max_sequence_length (CUDA free-memory query unavailable)".to_owned(), + }, + ) + } else if let Some(budget) = memory_budget_len { + ( + budget, + "CUDA free-memory budget (metadata model.max_sequence_length unavailable)".to_owned(), + ) + } else { + bail!( + "cannot derive native CUDA KV capacity: no explicit max length, model.max_sequence_length is unavailable, and CUDA free-memory could not be queried; set ONNX_GENAI_CUDA_KV_MAX_LEN or use load_with_cuda_kv_max_len" + ); + }; + if configured == 0 { + bail!( + "derived CUDA KV max length is zero from {source}; free more device memory or set a smaller ONNX_GENAI_CUDA_KV_MAX_LEN/load_with_cuda_kv_max_len value" + ); + } + Ok(CudaKvCapacity { + max_len: configured, + source, + metadata_max_len, + device_memory, + bytes_per_token, + }) +} + +fn max_kv_capacity_for_growth_budget(budget_tokens: usize) -> usize { + if budget_tokens == 0 { + return 0; + } + let initial = onnx_genai_kv::kv_capacity_bucket(0, budget_tokens); + if initial >= budget_tokens { + return budget_tokens; + } + let mut previous = initial; + loop { + let Some(next) = previous.checked_mul(2) else { + return budget_tokens; + }; + if next >= budget_tokens { + return budget_tokens + .checked_sub(previous) + .filter(|candidate| *candidate > previous) + .unwrap_or(previous); + } + if previous + .checked_add(next) + .is_some_and(|transient| transient <= budget_tokens) + { + previous = next; + continue; + } + return budget_tokens + .checked_sub(previous) + .filter(|candidate| *candidate > previous) + .unwrap_or(previous); + } +} + +pub(crate) fn cuda_kv_capacity_exceeded_message( + requested: usize, + capacity: &CudaKvCapacity, +) -> String { + let metadata = capacity + .metadata_max_len + .map(|value| value.to_string()) + .unwrap_or_else(|| "unavailable".to_owned()); + let memory = capacity + .device_memory + .map(|memory| { + format!( + "CUDA free={} bytes, total={} bytes", + memory.free_bytes, memory.total_bytes + ) + }) + .unwrap_or_else(|| "CUDA free-memory query unavailable".to_owned()); + format!( + "CUDA KV capacity exceeded: requested context length {requested}, configured max_len {} (source: {}; model.max_sequence_length: {metadata}; {memory}; KV bytes/token: {}). Increase the limit with ONNX_GENAI_CUDA_KV_MAX_LEN or load_with_cuda_kv_max_len if the device has enough VRAM, or reduce prompt/max_new_tokens.", + capacity.max_len, capacity.source, capacity.bytes_per_token + ) +} + +#[cfg(feature = "cuda")] +pub(crate) fn cuda_device_memory_snapshot( + device_id: i32, +) -> anyhow::Result { + let memory = onnx_genai_ort::cuda_rt::device_memory_info(device_id) + .context("query CUDA free device memory for native KV capacity")?; + Ok(CudaDeviceMemorySnapshot { + free_bytes: memory.free_bytes, + total_bytes: memory.total_bytes, + }) +} + +#[cfg(not(feature = "cuda"))] +pub(crate) fn cuda_device_memory_snapshot( + _device_id: i32, +) -> anyhow::Result { + bail!("CUDA free-memory query requires the onnx-genai-engine `cuda` feature") +} diff --git a/crates/onnx-genai-engine/src/native_decode/load.rs b/crates/onnx-genai-engine/src/native_decode/load.rs index 877fa0777e..542021b0f7 100644 --- a/crates/onnx-genai-engine/src/native_decode/load.rs +++ b/crates/onnx-genai-engine/src/native_decode/load.rs @@ -1,11 +1,26 @@ use super::*; +fn native_metadata_max_len_from_model_path(path: &Path) -> Option { + let root = if path.is_dir() { path } else { path.parent()? }; + [ + "inference_metadata.yaml", + "inference_metadata.yml", + "inference_metadata.json", + ] + .iter() + .map(|name| root.join(name)) + .find(|path| path.is_file()) + .and_then(|path| onnx_genai_metadata::load_metadata(&path).ok()) + .and_then(|metadata| metadata.model.and_then(|model| model.max_sequence_length)) +} + impl NativeDecodeSession { pub(crate) fn load_with_weight_offload_host_cache( path: impl AsRef, device: NativeDecodeDevice, host_cache: onnx_runtime_ep_cpu::WeightOffloadHostCache, io: Option<&ModelIoSpec>, + metadata_max_len: Option, decode_precision: DecodePrecision, ) -> anyhow::Result { let preference = match device { @@ -47,21 +62,32 @@ impl NativeDecodeSession { builder = builder.execution_provider(Arc::new(ep)); } let session = builder.build().context("load native decoder model")?; - Self::from_session_with_cuda_kv_max_len_and_io(session, None, io) + Self::from_session_with_cuda_options_and_io( + session, + NativeDecodeCudaOptions { + kv_max_len: None, + metadata_max_len, + graph_capture: None, + }, + io, + ) } /// Load with an explicit CUDA KV capacity. `None` uses - /// `ONNX_GENAI_CUDA_KV_MAX_LEN`, then the 4096-token default. + /// `ONNX_GENAI_CUDA_KV_MAX_LEN`, then derives a default from model metadata + /// clamped by queried CUDA free memory. pub fn load_with_cuda_kv_max_len( path: impl AsRef, device: NativeDecodeDevice, cuda_kv_max_len: Option, ) -> anyhow::Result { + let metadata_max_len = native_metadata_max_len_from_model_path(path.as_ref()); Self::load_with_cuda_options( path, device, NativeDecodeCudaOptions { kv_max_len: cuda_kv_max_len, + metadata_max_len, graph_capture: None, }, ) @@ -77,8 +103,11 @@ impl NativeDecodeSession { pub fn load_with_cuda_options( path: impl AsRef, device: NativeDecodeDevice, - options: NativeDecodeCudaOptions, + mut options: NativeDecodeCudaOptions, ) -> anyhow::Result { + if options.metadata_max_len.is_none() { + options.metadata_max_len = native_metadata_max_len_from_model_path(path.as_ref()); + } let preference = match device { NativeDecodeDevice::Cpu => DevicePreference::Cpu, NativeDecodeDevice::Cuda { index } => DevicePreference::Gpu { index }, @@ -109,6 +138,7 @@ impl NativeDecodeSession { Self::from_session_with_cuda_options(session, NativeDecodeCudaOptions::default()) } + #[cfg(test)] pub(crate) fn from_session_with_cuda_kv_max_len_and_io( session: InferenceSession, cuda_kv_max_len: Option, @@ -118,13 +148,14 @@ impl NativeDecodeSession { session, NativeDecodeCudaOptions { kv_max_len: cuda_kv_max_len, + metadata_max_len: None, graph_capture: None, }, io, ) } - fn from_session_with_cuda_options( + pub(crate) fn from_session_with_cuda_options( session: InferenceSession, cuda_options: NativeDecodeCudaOptions, ) -> anyhow::Result { @@ -369,11 +400,25 @@ impl NativeDecodeSession { let attention_mask = attention_mask .as_deref() .context("native CUDA target decode requires a declared attention-mask input")?; + let bytes_per_token = DecodeCudaState::kv_bytes_per_token(&session, &present_to_past)?; + let device_memory = cuda_device_memory_snapshot(session.device_id().index as i32).ok(); let max_len = match cuda_options.kv_max_len { Some(0) => bail!("CUDA KV max length must be greater than zero"), - Some(value) => value, - None => cuda_kv_max_len_from_env()?, + Some(value) => Some(value), + None => None, }; + let env_max_len = if max_len.is_some() { + None + } else { + cuda_kv_max_len_from_env()? + }; + let capacity = resolve_cuda_kv_capacity( + max_len, + env_max_len, + cuda_options.metadata_max_len, + bytes_per_token, + device_memory, + )?; let runtime_config = onnx_genai_runtime_config::runtime_config(); let graph_enabled = resolve_graph_capture_enabled( cuda_options.graph_capture, @@ -385,7 +430,7 @@ impl NativeDecodeSession { }, ); let mut span = onnx_genai_ort::prof_span!("native.cuda_kv_alloc"); - span.set_arg("max_len", max_len as u64); + span.set_arg("max_len", capacity.max_len as u64); span.set_arg("kv_pairs", present_to_past.len() as u64); span.set_arg("graph_capture", graph_enabled); Some(DecodeCudaState::new( @@ -397,7 +442,7 @@ impl NativeDecodeSession { logits: &logits, }, &present_to_past, - max_len, + capacity, graph_enabled, )?) } else { diff --git a/crates/onnx-genai-engine/src/native_decode/mod.rs b/crates/onnx-genai-engine/src/native_decode/mod.rs index 9e533e717e..d061a8115b 100644 --- a/crates/onnx-genai-engine/src/native_decode/mod.rs +++ b/crates/onnx-genai-engine/src/native_decode/mod.rs @@ -64,6 +64,7 @@ pub enum NativeDecodeDevice { #[derive(Clone, Copy, Debug, Default, Eq, PartialEq)] pub struct NativeDecodeCudaOptions { pub kv_max_len: Option, + pub metadata_max_len: Option, pub graph_capture: Option, } diff --git a/crates/onnx-genai-engine/src/native_decode/tests.rs b/crates/onnx-genai-engine/src/native_decode/tests.rs index 25d01565dc..83f138755b 100644 --- a/crates/onnx-genai-engine/src/native_decode/tests.rs +++ b/crates/onnx-genai-engine/src/native_decode/tests.rs @@ -38,6 +38,95 @@ fn graph_capture_auto_enables_for_owned_cuda_kv() { )); } +#[test] +fn cuda_kv_capacity_uses_metadata_with_memory_clamp() { + let capacity = resolve_cuda_kv_capacity( + None, + None, + Some(4096), + 10, + Some(CudaDeviceMemorySnapshot { + free_bytes: 20_000, + total_bytes: 40_000, + }), + ) + .unwrap(); + assert_eq!(capacity.max_len, 1024); + assert!( + capacity + .source + .contains("model.max_sequence_length clamped by CUDA free-memory growth budget") + ); +} + +#[test] +fn cuda_kv_capacity_clamp_accounts_for_allocate_before_free_growth() { + let capacity = resolve_cuda_kv_capacity( + None, + None, + Some(131_072), + 28_680, + Some(CudaDeviceMemorySnapshot { + free_bytes: 5_925_502_976, + total_bytes: 8_585_281_536, + }), + ) + .unwrap(); + assert_eq!(capacity.max_len, 110_080); + assert!( + capacity + .source + .contains("model.max_sequence_length clamped by CUDA free-memory growth budget") + ); +} + +#[test] +fn cuda_kv_capacity_honors_env_before_metadata() { + let capacity = resolve_cuda_kv_capacity( + None, + Some(8192), + Some(4096), + 10, + Some(CudaDeviceMemorySnapshot { + free_bytes: 20_000, + total_bytes: 40_000, + }), + ) + .unwrap(); + assert_eq!(capacity.max_len, 8192); + assert_eq!(capacity.source, "ONNX_GENAI_CUDA_KV_MAX_LEN"); +} + +#[test] +fn cuda_kv_capacity_error_explains_source_and_device_memory() { + let capacity = CudaKvCapacity { + max_len: 1024, + source: "model.max_sequence_length clamped by CUDA free-memory growth budget".to_owned(), + metadata_max_len: Some(4096), + device_memory: Some(CudaDeviceMemorySnapshot { + free_bytes: 20_000, + total_bytes: 40_000, + }), + bytes_per_token: 10, + }; + let message = cuda_kv_capacity_exceeded_message(1025, &capacity); + assert!( + message.contains("requested context length 1025"), + "{message}" + ); + assert!(message.contains("configured max_len 1024"), "{message}"); + assert!( + message.contains("source: model.max_sequence_length"), + "{message}" + ); + assert!( + message.contains("model.max_sequence_length: 4096"), + "{message}" + ); + assert!(message.contains("CUDA free=20000 bytes"), "{message}"); + assert!(message.contains("ONNX_GENAI_CUDA_KV_MAX_LEN"), "{message}"); +} + #[test] fn graph_capture_auto_declines_for_non_owned_or_non_cuda() { let shared = GraphCaptureStructuralSafety { @@ -751,6 +840,7 @@ fn build_cuda_decoder( session, NativeDecodeCudaOptions { kv_max_len: Some(max_len), + metadata_max_len: None, graph_capture: Some(graph_capture), }, ) @@ -1432,6 +1522,7 @@ fn native_cuda_qwen_decode_matches_cpu_tokens() -> anyhow::Result<()> { NativeDecodeDevice::Cuda { index: Some(0) }, NativeDecodeCudaOptions { kv_max_len: Some(128), + metadata_max_len: None, graph_capture: Some(false), }, )?; @@ -1450,6 +1541,7 @@ fn native_cuda_qwen_decode_matches_cpu_tokens() -> anyhow::Result<()> { NativeDecodeDevice::Cuda { index: Some(0) }, NativeDecodeCudaOptions { kv_max_len: Some(128), + metadata_max_len: None, graph_capture: Some(true), }, )?; @@ -1536,6 +1628,7 @@ fn native_cuda_verify_rewind_no_kv_corruption() -> anyhow::Result<()> { NativeDecodeDevice::Cuda { index: Some(0) }, NativeDecodeCudaOptions { kv_max_len: Some(128), + metadata_max_len: None, graph_capture: Some(graph), }, ) diff --git a/crates/onnx-genai-kv/src/lib.rs b/crates/onnx-genai-kv/src/lib.rs index 5f12a440a3..3bffcad3a3 100644 --- a/crates/onnx-genai-kv/src/lib.rs +++ b/crates/onnx-genai-kv/src/lib.rs @@ -52,6 +52,99 @@ pub type SequenceId = u64; /// Token identifier. pub type TokenId = u32; +/// Round a required sequence length up to the shared KV capacity bucket. +/// +/// Buckets are powers of two, at least the minimum bucket floor, and clamped to +/// the caller's hard maximum. `ONNX_GENAI_KV_MIN_BUCKET` can raise or lower the +/// floor for deployments that want fewer early reallocations or smaller initial +/// buffers. This is the single capacity policy shared by ORT shared-buffer KV +/// and native CUDA KV; scheduler/resource planning can consume the same source +/// instead of reimplementing it. +pub fn kv_capacity_bucket(len: usize, hard_max: usize) -> usize { + const MIN_BUCKET_DEFAULT: usize = 256; + if hard_max == 0 { + return 0; + } + let min_bucket = std::env::var("ONNX_GENAI_KV_MIN_BUCKET") + .ok() + .and_then(|v| v.parse::().ok()) + .filter(|v| *v > 0) + .unwrap_or(MIN_BUCKET_DEFAULT); + len.next_power_of_two().max(min_bucket).min(hard_max) +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum KvCapacityGrowth { + Unchanged, + Grew { + old_capacity: usize, + new_capacity: usize, + valid_len: usize, + }, +} + +pub trait KvCapacityGrowthBackend { + type Error; + type GrownBuffers; + type GrownMask; + + fn current_capacity(&self) -> usize; + fn hard_max_capacity(&self) -> usize; + fn valid_len(&self) -> usize; + fn capacity_exceeded(&self, required: usize) -> Self::Error; + fn build_grown_buffers( + &mut self, + new_capacity: usize, + valid_len: usize, + ) -> Result; + fn build_grown_mask( + &mut self, + new_capacity: usize, + valid_len: usize, + ) -> Result, Self::Error>; + fn invalidate_capture(&mut self) -> Result<(), Self::Error>; + fn commit_grown_capacity( + &mut self, + new_capacity: usize, + grown_buffers: Self::GrownBuffers, + grown_mask: Option, + ) -> Result<(), Self::Error>; +} + +/// Shared grow driver for fixed-capacity KV buckets. +/// +/// Backends provide allocation/copy/capture primitives; this function owns the +/// policy: reject above the hard maximum, choose the next shared bucket, build +/// all fallible replacement state before mutating live state, invalidate capture, +/// then commit atomically. +pub fn ensure_kv_capacity(backend: &mut B, required: usize) -> Result +where + B: KvCapacityGrowthBackend, +{ + let hard_max = backend.hard_max_capacity(); + if required > hard_max { + return Err(backend.capacity_exceeded(required)); + } + let old_capacity = backend.current_capacity(); + if required <= old_capacity { + return Ok(KvCapacityGrowth::Unchanged); + } + let new_capacity = kv_capacity_bucket(required, hard_max); + if new_capacity <= old_capacity { + return Ok(KvCapacityGrowth::Unchanged); + } + let valid_len = backend.valid_len(); + let grown_buffers = backend.build_grown_buffers(new_capacity, valid_len)?; + let grown_mask = backend.build_grown_mask(new_capacity, valid_len)?; + backend.invalidate_capture()?; + backend.commit_grown_capacity(new_capacity, grown_buffers, grown_mask)?; + Ok(KvCapacityGrowth::Grew { + old_capacity, + new_capacity, + valid_len, + }) +} + /// Device tier for page storage. #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] pub enum Device { @@ -136,3 +229,147 @@ pub enum KvError { #[error("Page {0} not found")] PageNotFound(PageId), } + +#[cfg(test)] +mod tests { + use super::{ + KvCapacityGrowth, KvCapacityGrowthBackend, ensure_kv_capacity, kv_capacity_bucket, + }; + + #[test] + fn kv_capacity_bucket_rounds_to_power_of_two_min_256() { + let hard = 32_768; + assert_eq!(kv_capacity_bucket(0, hard), 256); + assert_eq!(kv_capacity_bucket(1, hard), 256); + assert_eq!(kv_capacity_bucket(128, hard), 256); + assert_eq!(kv_capacity_bucket(256, hard), 256); + assert_eq!(kv_capacity_bucket(257, hard), 512); + assert_eq!(kv_capacity_bucket(5000, hard), 8192); + } + + #[test] + fn kv_capacity_bucket_caps_at_hard_max() { + assert_eq!(kv_capacity_bucket(20_000, 32_768), 32_768); + assert_eq!(kv_capacity_bucket(32_768, 32_768), 32_768); + assert_eq!(kv_capacity_bucket(100, 128), 128); + assert_eq!(kv_capacity_bucket(100, 0), 0); + } + + #[test] + fn kv_capacity_bucket_is_monotonic_and_within_bounds() { + let hard = 4096; + let mut previous = 0; + for len in 0..8192 { + let bucket = kv_capacity_bucket(len, hard); + assert!( + bucket >= previous, + "len={len} bucket={bucket} previous={previous}" + ); + assert!(bucket <= hard); + previous = bucket; + } + } + + #[derive(Default)] + struct RecordingBackend { + current: usize, + hard: usize, + valid: usize, + events: Vec<&'static str>, + fail_mask: bool, + } + + impl KvCapacityGrowthBackend for RecordingBackend { + type Error = &'static str; + type GrownBuffers = usize; + type GrownMask = usize; + + fn current_capacity(&self) -> usize { + self.current + } + + fn hard_max_capacity(&self) -> usize { + self.hard + } + + fn valid_len(&self) -> usize { + self.valid + } + + fn capacity_exceeded(&self, _required: usize) -> Self::Error { + "capacity exceeded" + } + + fn build_grown_buffers( + &mut self, + new_capacity: usize, + _valid_len: usize, + ) -> Result { + self.events.push("buffers"); + Ok(new_capacity) + } + + fn build_grown_mask( + &mut self, + new_capacity: usize, + _valid_len: usize, + ) -> Result, Self::Error> { + self.events.push("mask"); + if self.fail_mask { + return Err("mask failed"); + } + Ok(Some(new_capacity)) + } + + fn invalidate_capture(&mut self) -> Result<(), Self::Error> { + self.events.push("invalidate"); + Ok(()) + } + + fn commit_grown_capacity( + &mut self, + new_capacity: usize, + _grown_buffers: Self::GrownBuffers, + _grown_mask: Option, + ) -> Result<(), Self::Error> { + self.events.push("commit"); + self.current = new_capacity; + Ok(()) + } + } + + #[test] + fn ensure_kv_capacity_orders_fallible_work_before_invalidation() { + let mut backend = RecordingBackend { + current: 256, + hard: 4096, + valid: 200, + ..RecordingBackend::default() + }; + let growth = ensure_kv_capacity(&mut backend, 257).unwrap(); + assert_eq!( + growth, + KvCapacityGrowth::Grew { + old_capacity: 256, + new_capacity: 512, + valid_len: 200, + } + ); + assert_eq!(backend.events, ["buffers", "mask", "invalidate", "commit"]); + assert_eq!(backend.current, 512); + } + + #[test] + fn ensure_kv_capacity_does_not_invalidate_or_commit_after_prepare_failure() { + let mut backend = RecordingBackend { + current: 256, + hard: 4096, + valid: 200, + fail_mask: true, + ..RecordingBackend::default() + }; + assert_eq!(ensure_kv_capacity(&mut backend, 257), Err("mask failed")); + assert_eq!(backend.events, ["buffers", "mask"]); + assert_eq!(backend.current, 256); + } +} diff --git a/crates/onnx-genai-ort/Cargo.toml b/crates/onnx-genai-ort/Cargo.toml index 28d22d428b..349fe85fa3 100644 --- a/crates/onnx-genai-ort/Cargo.toml +++ b/crates/onnx-genai-ort/Cargo.toml @@ -42,6 +42,7 @@ cudarc = { version = "0.19", default-features = false, features = ["std", "drive minijinja = { version = "2", features = ["json"] } minijinja-contrib = { version = "2", features = ["pycompat"] } onnx-genai-genai-config = { workspace = true } +onnx-genai-kv = { workspace = true } # The shared absolute trace time base, so these spans land on the same timeline # as the ones the runtime and its execution providers record. onnx-runtime-tracer = { workspace = true } diff --git a/crates/onnx-genai-ort/src/cuda_rt.rs b/crates/onnx-genai-ort/src/cuda_rt.rs index cd09a9643f..907be11320 100644 --- a/crates/onnx-genai-ort/src/cuda_rt.rs +++ b/crates/onnx-genai-ort/src/cuda_rt.rs @@ -42,6 +42,7 @@ type CudaMemsetFn = unsafe extern "C" fn(*mut c_void, i32, usize) -> i32; type CudaDeviceSynchronizeFn = unsafe extern "C" fn() -> i32; type CudaSetDeviceFn = unsafe extern "C" fn(i32) -> i32; type CudaGetDeviceFn = unsafe extern "C" fn(*mut i32) -> i32; +type CudaMemGetInfoFn = unsafe extern "C" fn(*mut usize, *mut usize) -> i32; struct CudaRt { // Kept alive so the resolved function pointers remain valid; never called @@ -52,6 +53,7 @@ struct CudaRt { device_synchronize: CudaDeviceSynchronizeFn, set_device: CudaSetDeviceFn, get_device: CudaGetDeviceFn, + mem_get_info: CudaMemGetInfoFn, } // SAFETY: the resolved `cudart` entry points are plain C functions that are @@ -93,8 +95,23 @@ fn load() -> std::result::Result { unsafe { lib.get::(b"cudaDeviceSynchronize\0") }; let set_device = unsafe { lib.get::(b"cudaSetDevice\0") }; let get_device = unsafe { lib.get::(b"cudaGetDevice\0") }; - match (memcpy, memset, device_synchronize, set_device, get_device) { - (Ok(memcpy), Ok(memset), Ok(device_synchronize), Ok(set_device), Ok(get_device)) => { + let mem_get_info = unsafe { lib.get::(b"cudaMemGetInfo\0") }; + match ( + memcpy, + memset, + device_synchronize, + set_device, + get_device, + mem_get_info, + ) { + ( + Ok(memcpy), + Ok(memset), + Ok(device_synchronize), + Ok(set_device), + Ok(get_device), + Ok(mem_get_info), + ) => { // Copy the function pointers out before `lib` is moved into the // struct; the borrows on `lib` end here. let memcpy = *memcpy; @@ -102,6 +119,7 @@ fn load() -> std::result::Result { let device_synchronize = *device_synchronize; let set_device = *set_device; let get_device = *get_device; + let mem_get_info = *mem_get_info; return Ok(CudaRt { _lib: lib, memcpy, @@ -109,11 +127,12 @@ fn load() -> std::result::Result { device_synchronize, set_device, get_device, + mem_get_info, }); } _ => { last_err = format!( - "{name}: missing cudaMemcpy/cudaMemset/cudaDeviceSynchronize/cudaSetDevice/cudaGetDevice symbol" + "{name}: missing cudaMemcpy/cudaMemset/cudaDeviceSynchronize/cudaSetDevice/cudaGetDevice/cudaMemGetInfo symbol" ); } } @@ -136,7 +155,7 @@ fn runtime() -> Result<&'static CudaRt> { /// copy is unordered relative to the EP's KV writes (before) and reads (after), /// which silently corrupts the cache. Growth is rare (O(log length)), so the /// synchronization cost is negligible. -pub(crate) fn device_synchronize() -> Result<()> { +pub fn device_synchronize() -> Result<()> { let rt = runtime()?; // SAFETY: `cudaDeviceSynchronize` takes no arguments and matches the // `cudart` ABI. @@ -149,6 +168,32 @@ pub(crate) fn device_synchronize() -> Result<()> { Ok(()) } +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub struct CudaDeviceMemoryInfo { + pub free_bytes: usize, + pub total_bytes: usize, +} + +/// Query CUDA device memory after making `device_id` current. +pub fn device_memory_info(device_id: i32) -> Result { + let _guard = DeviceGuard::set(device_id)?; + let rt = runtime()?; + let mut free_bytes = 0usize; + let mut total_bytes = 0usize; + // SAFETY: both pointers are valid out-parameters; `cudaMemGetInfo` matches + // the cudart ABI and reads the current device selected by `DeviceGuard`. + let code = unsafe { (rt.mem_get_info)(&mut free_bytes, &mut total_bytes) }; + if code != 0 { + return Err(OrtError::InvalidArgument(format!( + "cudaMemGetInfo failed with CUDA error code {code}" + ))); + } + Ok(CudaDeviceMemoryInfo { + free_bytes, + total_bytes, + }) +} + /// RAII guard that makes `device_id` the calling thread's current CUDA device /// for the duration of a grow copy, restoring the previous current device on /// drop. @@ -160,14 +205,14 @@ pub(crate) fn device_synchronize() -> Result<()> { /// synchronize the wrong device and fail to order the copy against the EP's /// stream โ€” the exact race the barriers exist to prevent. Pinning is cheap and /// growth is rare (O(log length)). -pub(crate) struct DeviceGuard { +pub struct DeviceGuard { prev: i32, restore: bool, } impl DeviceGuard { /// Set the current CUDA device to `device_id`, remembering the previous one. - pub(crate) fn set(device_id: i32) -> Result { + pub fn set(device_id: i32) -> Result { let rt = runtime()?; let mut prev: i32 = 0; // SAFETY: `prev` is a valid out-parameter; `cudaGetDevice` matches the @@ -218,7 +263,7 @@ impl Drop for DeviceGuard { /// /// Used to define the tail of a freshly allocated (uninitialized) KV bucket so /// positions past the valid length are deterministic zeros. -pub(crate) fn memset_zero(dst: usize, bytes: usize) -> Result<()> { +pub fn memset_zero(dst: usize, bytes: usize) -> Result<()> { if bytes == 0 { return Ok(()); } @@ -255,7 +300,7 @@ pub(crate) fn memset_zero(dst: usize, bytes: usize) -> Result<()> { /// Ordering against the *previous* replay's read of these buffers (WAR) is /// guaranteed separately: the device sampler fully synchronizes the device at the /// end of every captured step before the next step overwrites the inputs. -pub(crate) fn memcpy_host_to_device(dst: usize, src: &[u8]) -> Result<()> { +pub fn memcpy_host_to_device(dst: usize, src: &[u8]) -> Result<()> { if src.is_empty() { return Ok(()); } @@ -313,7 +358,7 @@ pub(crate) fn memcpy_device_to_host(dst: &mut [u8], src: usize) -> Result<()> { /// Copy `bytes` from device address `src` to device address `dst` /// (`cudaMemcpyDeviceToDevice`). -pub(crate) fn memcpy_device_to_device(dst: usize, src: usize, bytes: usize) -> Result<()> { +pub fn memcpy_device_to_device(dst: usize, src: usize, bytes: usize) -> Result<()> { if bytes == 0 { return Ok(()); } diff --git a/crates/onnx-genai-ort/src/decode/dynamic.rs b/crates/onnx-genai-ort/src/decode/dynamic.rs index 1e8ebe655a..c50a2f4b64 100644 --- a/crates/onnx-genai-ort/src/decode/dynamic.rs +++ b/crates/onnx-genai-ort/src/decode/dynamic.rs @@ -1250,40 +1250,40 @@ impl<'a> DecodeSession<'a> { if self.mode != DecodeKvMode::SharedBuffer { return Ok(()); } - let hard_max = self.max_length.ok_or_else(|| { - OrtError::InvalidArgument("shared-buffer growth requires max_length".into()) - })?; - if required > hard_max { - return Err(OrtError::InvalidArgument(format!( - "requested KV capacity {required} exceeds model max_length {hard_max}" - ))); - } - if required <= self.kv_capacity { - return Ok(()); - } - let new_capacity = kv_capacity_bucket(required, hard_max); - if new_capacity <= self.kv_capacity { - // Already at the hard cap; no growth is needed. - return Ok(()); - } - self.grow_kv_buffers(new_capacity) + let _ = onnx_genai_kv::ensure_kv_capacity(self, required)?; + Ok(()) } +} - /// Reallocate the shared KV buffers at `new_capacity`, copying the valid - /// prefix `[0, current_len)` across, then resize the captured mask and force - /// a graph re-capture so the next captured step rebinds the larger buffers. - /// - /// CORRECTNESS: the valid KV prefix must survive the reallocation exactly. - /// The buffers may be device-resident โ€” where the raw data pointer is NOT - /// host-addressable. On CUDA the prefix is relocated with a direct - /// device-to-device `cudaMemcpy` on the raw KV device pointers (ORT's - /// `CopyTensors` is unusable here: the built-in CUDA EP registers no - /// env-level data-transfer, so it fails with "Data transfer implementation - /// ... not found (code 9)"). CPU buffers copy directly on the host. Growth - /// is rare (O(log length)), so this one-time per-growth cost is amortized - /// away. - fn grow_kv_buffers(&mut self, new_capacity: usize) -> Result<()> { - let valid_len = self.current_len; +impl onnx_genai_kv::KvCapacityGrowthBackend for DecodeSession<'_> { + type Error = OrtError; + type GrownBuffers = HashMap>; + type GrownMask = Value; + + fn current_capacity(&self) -> usize { + self.kv_capacity + } + + fn hard_max_capacity(&self) -> usize { + self.max_length.unwrap_or(0) + } + + fn valid_len(&self) -> usize { + self.current_len + } + + fn capacity_exceeded(&self, required: usize) -> Self::Error { + OrtError::InvalidArgument(format!( + "requested KV capacity {required} exceeds model max_length {}", + self.hard_max_capacity() + )) + } + + fn build_grown_buffers( + &mut self, + new_capacity: usize, + valid_len: usize, + ) -> Result { // Resolve the transfer path once, failing fast (before any allocation) // if the buffers live on a device we have no copy primitive for. let device = self.grow_device()?; @@ -1312,34 +1312,53 @@ impl<'a> DecodeSession<'a> { )?; grown.insert(pair.past.clone(), Arc::new(new_value)); } + Ok(grown) + } + fn build_grown_mask( + &mut self, + new_capacity: usize, + _valid_len: usize, + ) -> Result> { // The captured attention mask capacity must equal the KV buffer // capacity. Build its replacement before mutating the session so a // fallible allocation cannot leave the KV and capture state out of sync. // Preserve the mask's residency: a device-resident captured mask (the // CUDA device-argmax path) must stay on the device allocator so the // captured graph keeps reading it in place after re-capture. - let grown_mask = if let Some((valid_ones, on_device, mask_device_id)) = self + if let Some((valid_ones, on_device, mask_device_id)) = self .capture .as_ref() .map(|cap| (cap.mask_valid_len, cap.inputs_on_device, cap.device_id)) { let mask_len = i64::try_from(new_capacity) .map_err(|_| OrtError::InvalidArgument("KV capacity exceeds i64".into()))?; - Some(self.build_grown_capture_mask( + Ok(Some(self.build_grown_capture_mask( new_capacity, mask_len, valid_ones, on_device, mask_device_id, - )?) + )?)) } else { - None - }; + Ok(None) + } + } - // All fallible work is complete. Release the captured graph while its old - // buffers are still alive, then atomically commit the replacement state. + fn invalidate_capture(&mut self) -> Result<()> { + // The shared grow driver has completed all fallible allocation/copy + // work. Release the captured graph while its old buffers are still alive; + // commit_grown_capacity atomically swaps in the replacement state. self.invalidate_captured_graph(); + Ok(()) + } + + fn commit_grown_capacity( + &mut self, + new_capacity: usize, + grown: Self::GrownBuffers, + grown_mask: Option, + ) -> Result<()> { self.current_kv = grown; self.kv_capacity = new_capacity; if let (Some(cap), Some(attention_mask)) = (self.capture.as_mut(), grown_mask) { @@ -1348,13 +1367,16 @@ impl<'a> DecodeSession<'a> { } Ok(()) } +} +impl DecodeSession<'_> { /// Build a replacement captured attention mask of `capacity` elements /// (shape `[1, mask_len]`) with the leading `valid_ones` set to 1 and the /// rest zero, preserving the mask's residency: device-resident when /// `on_device` (allocated on the retained CUDA allocator and initialized with /// host->device fills), host-resident otherwise. Used by - /// [`grow_kv_buffers`](Self::grow_kv_buffers) when the KV capacity grows. + /// the shared [`onnx_genai_kv::ensure_kv_capacity`] driver when the KV + /// capacity grows. fn build_grown_capture_mask( &self, capacity: usize, diff --git a/crates/onnx-genai-ort/src/decode/kv_growth.rs b/crates/onnx-genai-ort/src/decode/kv_growth.rs index dd7077e3f7..0d2b21577d 100644 --- a/crates/onnx-genai-ort/src/decode/kv_growth.rs +++ b/crates/onnx-genai-ort/src/decode/kv_growth.rs @@ -1,4 +1,5 @@ use super::*; +pub(super) use onnx_genai_kv::kv_capacity_bucket; /// Where a shared-KV buffer lives, selecting the prefix-copy transfer used when /// growing it. Non-CUDA device EPs are rejected earlier (see @@ -12,32 +13,6 @@ pub(super) enum GrowDevice { Cuda, } -/// Round a required sequence length up to the shared-KV bucket capacity. -/// -/// Buckets are powers of two (at least the minimum bucket floor), clamped to the -/// model's hard `max_length`. The caller must reject lengths above that ceiling -/// before bucketing. Sizing the shared KV buffers to the bucket rather than the -/// full `max_length` keeps captured-decode -/// per-step attention cost ~O(actual length) โ€” matching onnxruntime-genai โ€” -/// while the sequence only crosses O(log length) bucket boundaries, so the -/// buffers (and the captured CUDA graph) are grown/re-captured that few times. -/// See [`DecodeSession::kv_capacity`] for the measured motivation. -/// -/// The minimum bucket floor defaults to 256 and can be overridden with the -/// `ONNX_GENAI_KV_MIN_BUCKET` environment variable. The default is a good -/// balance across models (a large `max_length` model that generates few tokens -/// pays no capacity tax, while short generations avoid frequent early growth); -/// the override exists for per-deployment tuning without a rebuild. -pub(super) fn kv_capacity_bucket(len: usize, hard_max: usize) -> usize { - const MIN_BUCKET_DEFAULT: usize = 256; - let min_bucket = std::env::var("ONNX_GENAI_KV_MIN_BUCKET") - .ok() - .and_then(|v| v.parse::().ok()) - .filter(|v| *v > 0) - .unwrap_or(MIN_BUCKET_DEFAULT); - len.next_power_of_two().max(min_bucket).min(hard_max) -} - /// Allocate a `new_shape` shared-KV buffer and copy the valid sequence prefix /// `[0, valid_len)` of `old` into it, zeroing positions `>= valid_len`. /// diff --git a/crates/onnx-genai-ort/src/lib.rs b/crates/onnx-genai-ort/src/lib.rs index 023cba2e80..ee568e704c 100644 --- a/crates/onnx-genai-ort/src/lib.rs +++ b/crates/onnx-genai-ort/src/lib.rs @@ -12,7 +12,7 @@ pub mod binding; pub mod chat_template; pub mod component; #[cfg(feature = "cuda")] -pub(crate) mod cuda_rt; +pub mod cuda_rt; pub mod decode; pub mod decode_contract; #[cfg(feature = "cuda")] diff --git a/crates/onnx-runtime-session/src/executor/build.rs b/crates/onnx-runtime-session/src/executor/build.rs index f09fc40901..2a5a784d3d 100644 --- a/crates/onnx-runtime-session/src/executor/build.rs +++ b/crates/onnx-runtime-session/src/executor/build.rs @@ -1611,7 +1611,7 @@ impl Executor { physical_shape: Vec, logical_shape: Vec, ) -> Result { - let expose_logical_input_shape = self.input_index.get(&input_name).is_none_or(|&vid| { + let expose_logical_input_shape = self.input_index.get(&input_name).is_some_and(|&vid| { if output_name.is_some() { !self.binding_consumers_use_physical_capacity(vid) } else { diff --git a/docs/native-ort-kv-capacity.md b/docs/native-ort-kv-capacity.md new file mode 100644 index 0000000000..c1ab054a4d --- /dev/null +++ b/docs/native-ort-kv-capacity.md @@ -0,0 +1,48 @@ +# Native and ORT KV capacity unification + +The runtime uses one KV capacity policy for ORT shared-buffer KV and native CUDA +KV: `onnx_genai_kv::kv_capacity_bucket(required, hard_max)`. The policy rounds +demand to a power-of-two bucket, applies the `ONNX_GENAI_KV_MIN_BUCKET` floor, +and clamps to a hard maximum derived from model metadata and device capacity. + +## Hard maximum + +The hard maximum is not a model-family constant. It comes from explicit caller +configuration first (`load_with_cuda_kv_max_len`, `ONNX_GENAI_CUDA_KV_MAX_LEN`), +otherwise from `model.max_sequence_length`, clamped by a queried CUDA +free-memory growth budget using the graph's actual KV bytes per token. The +growth budget accounts for the allocation-before-free sequence used by both +backends: at a bucket boundary the old bucket remains live while the larger +bucket and mask are allocated and filled, so the default ceiling is the largest +capacity whose worst bucket transition fits inside the headroom budget. If +neither metadata nor a CUDA memory query can provide a limit, native CUDA fails +during load with an actionable configuration error instead of guessing. + +## Growth + +Both ORT shared-buffer KV and native CUDA grow on demand by the shared bucket +policy instead of pre-allocating the full context. The orchestration lives in +`onnx_genai_kv::ensure_kv_capacity` over the `KvCapacityGrowthBackend` trait: +reject above the hard maximum, compute the next `kv_capacity_bucket`, build all +fallible replacement state while the old state remains live, invalidate graph +capture, then commit the new capacity. ORT implements the primitive seam with +`OrtValue` allocation/rebind and `grow_kv_value`; native CUDA implements it with +`DeviceIoBinding` allocation, direct device-to-device prefix copies, and mask +rewrites. Adding another backend should implement those primitives rather than +re-derive the growth policy. + +## CUDA graph capture + +Growth and CUDA graph capture coexist by making growth a graph boundary. The +shared driver builds the new bucket and captured mask before it calls the +backend's capture-invalidation primitive, so failure leaves the old buffers and +old capture intact. ORT releases its captured graph id; native CUDA resets its +captured graph. The next captured step then binds/captures the new buffers and +subsequent steps replay at the new bucket. Bucket growth is logarithmically rare, +so the recapture cost is amortized. + +The remaining limitation is CI coverage, not capability: the grow-and-recapture +path was exercised manually on an RTX 4060 Laptop GPU by forcing early buckets +with `ONNX_GENAI_KV_MIN_BUCKET=4`, observing native CUDA growth from 4โ†’8โ†’16 +with graph recapture at each bucket and coherent continued generation. Automated +CUDA coverage still depends on a GPU runner and model fixtures being available. From 44404842f47a839af09d598cea096c332e0650f3 Mon Sep 17 00:00:00 2001 From: justinchuby <223556219+Copilot@users.noreply.github.com> Date: Tue, 28 Jul 2026 11:50:51 -0700 Subject: [PATCH 2/4] feat(kv): one KV capacity mechanism for native CUDA and ORT Justin hit "CUDA KV capacity exceeded: requested context length 4097, configured max_len 4096" on the native backend and asked why it was so conservative, and whether we did not already have dynamic KV allocation. We did -- on one path only. The ORT shared-buffer GQA path grew in power-of-two buckets via kv_capacity_bucket, while native CUDA pre-allocated once at a hardcoded 4096 and never grew. Two mechanisms for the same resource, which is the shape behind several bugs here recently. A bigger constant would have left the divergence for someone to hit from another angle. The growth policy now lives once, in onnx-genai-kv: kv_capacity_bucket plus a KvCapacityGrowthBackend seam and a shared ensure_kv_capacity driver, consumed by both backends. The seam carries only primitives that genuinely differ -- allocate, copy the prefix, build the mask, invalidate capture, commit -- so a new backend implements primitives rather than re-deriving policy. Growth and CUDA graph capture coexist, as they already did on the ORT side: grow to the next bucket, device-to-device copy the prefix, rebuild the mask, release the capture, re-capture at the new size. The comment claiming capture required a frozen physical capacity described the fixed-capacity design rather than a real constraint, and is corrected. There is no derived VRAM ceiling. An earlier revision predicted a safe maximum from free memory times a fraction; that is guesswork dressed as policy, since the snapshot is stale as soon as another process allocates and the fraction is the kind of constant RULES.md rule 2 exists to prevent. The only remaining ceilings are facts: the model's declared max_sequence_length, and an explicit user cap. Unset means grow until the model's own limit or until the device refuses. That puts the weight on failing well. Growth is transactional: a failure at any fallible step leaves capacity, logical length, buffers and capture state exactly as they were, and the session stays usable. The commit phase is infallible assignment after all fallible work, which is why it has no failure case to cover. Failures are translated rather than passed through, naming the old and new bucket, the transient peak, bytes per token, device free and total, and the recovery levers. The transient peak is worth stating because it is unintuitive: growth holds the old and new buffers at once, so the requirement is roughly (old + new) * bytes_per_token. A grow can fail with what looks like ample free memory. Verified on an RTX 4060: forced 4 to 8 to 16 growth with real reallocation and re-capture, output stayed coherent -- a corrupted prefix copy produces wrong tokens rather than an error, so coherence is the assertion that matters. A real allocator OOM cannot be forced on a WDDM device, which virtualises and pages GPU memory; that is documented so it is not rediscovered. Coverage for the failure path is instead model-free injected failure at the shared seam -- allocation, prefix copy, mask allocation and capture invalidation each fail in turn, and each asserts no commit and an unchanged session, then proves a retry succeeds. Those run in CI on every PR rather than only on machines with a GPU. Also makes the Qwen CUDA smoke fixture env-overridable via ONNX_GENAI_QWEN_CUDA_SMOKE_MODEL; it previously skipped on a hardcoded Linux path. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 023d5a3f-c551-4aa8-afbe-a329af15fbdd --- .squad/agents/leon/history.md | 12 + crates/onnx-genai-engine/src/engine/load.rs | 15 +- .../src/native_decode/cuda.rs | 257 +++++++++--------- .../src/native_decode/tests.rs | 74 +++-- crates/onnx-genai-kv/src/lib.rs | 199 ++++++++++++++ crates/onnx-genai-ort/src/decode/dynamic.rs | 124 +++++++-- docs/native-ort-kv-capacity.md | 38 ++- 7 files changed, 519 insertions(+), 200 deletions(-) diff --git a/.squad/agents/leon/history.md b/.squad/agents/leon/history.md index 09d93c9abe..bf9c94106a 100644 --- a/.squad/agents/leon/history.md +++ b/.squad/agents/leon/history.md @@ -197,3 +197,15 @@ Authored the Phi-4-mini bit-exact native-CUDA-versus-ORT 64-token decode lock. F - Justin challenged the "metadata 131072, not clamped" DeepSeek result on the 8 GiB RTX 4060. The first clamp only checked final KV size (`0.85 * free / bytes_per_token`), but growth allocates the new bucket before freeing the old bucket. - Tightened the default hard maximum to the largest capacity whose worst old+new bucket transition fits inside the CUDA free-memory headroom budget. For the observed DeepSeek numbers (`free=5925502976`, `bytes/token=28680`), that clamps metadata 131072 to 110080 instead of allowing the 65536โ†’131072 transition to exceed the 85% growth budget. + +## 2026-07-28T10:53:59-07:00 โ€” VRAM ceiling removed; growth fails gracefully + +- Justin rejected the derived free-memory ceiling as guesswork. Removed the 85% CUDA-memory budget and the transient-growth hard-max arithmetic; only explicit user caps and `model.max_sequence_length` remain as factual ceilings. +- Kept the transient peak insight in docs/diagnostics: growth holds old and new buckets simultaneously, so failure can happen at roughly `(old + new) * bytes_per_token`. +- Native CUDA grow allocation/copy/mask failures now translate to an actionable, transactional error that says the old/new bucket, approximate new allocation, approximate transient peak, current CUDA free/total memory, bytes/token, recovery levers, and that session state was left unchanged. + +## 2026-07-28T11:43:45-07:00 โ€” Injected KV grow-failure coverage + +- Documented the Windows/WDDM finding: exhausting reported free memory does not reliably produce CUDA allocator failure because the GPU memory manager can page/virtualize. +- Added model-free injected-failure coverage at the shared `KvCapacityGrowthBackend` seam. Allocation failure now asserts the translated actionable error and proves the fake session's buffer id, logical length, capture generation, and capture-valid state remain unchanged before a successful retry. +- Added mid-sequence injected failures for prefix-copy, mask allocation, and capture invalidation; all prove no commit and unchanged session state. Native/ORT commit phases remain effectively infallible assignments after all fallible work has completed. diff --git a/crates/onnx-genai-engine/src/engine/load.rs b/crates/onnx-genai-engine/src/engine/load.rs index 5b04068f37..e803647880 100644 --- a/crates/onnx-genai-engine/src/engine/load.rs +++ b/crates/onnx-genai-engine/src/engine/load.rs @@ -348,16 +348,13 @@ fn resolve_metadata_and_decode_path( anyhow::bail!("Unsupported capabilities: {unsupported:?}"); } - // Optional cap on the runtime-owned fixed-capacity KV buffer. Foundry / + // Optional factual cap on runtime-owned KV growth. Foundry / // onnxruntime-genai `genai_config.json` models advertise the model's full - // `context_length` (e.g. 32k-131k) as their max sequence length, and the - // shared-buffer decode path pre-allocates a KV buffer of exactly that many - // tokens up front โ€” regardless of how many tokens a request will actually - // generate. On memory-constrained devices that over-allocation exhausts - // VRAM (spilling to shared system memory over PCIe) even for short runs. - // `ONNX_GENAI_KV_MAX_LEN` caps that capacity to the caller's real - // generation budget (prompt + max_new_tokens), mirroring the native - // path's `ONNX_GENAI_CUDA_KV_MAX_LEN`. Unset = unchanged (full context). + // `context_length` (e.g. 32k-131k) as their max sequence length. Shared KV + // now grows by buckets instead of pre-allocating that full length, but callers + // can still set `ONNX_GENAI_KV_MAX_LEN` as an explicit generation-budget cap, + // mirroring the native path's `ONNX_GENAI_CUDA_KV_MAX_LEN`. Unset = model + // metadata is the factual ceiling. let kv_shared_buffer_cap = shared_buffer_cap_from_env(); let metadata_max_context = metadata .model diff --git a/crates/onnx-genai-engine/src/native_decode/cuda.rs b/crates/onnx-genai-engine/src/native_decode/cuda.rs index e935bfc30c..0382acf4a8 100644 --- a/crates/onnx-genai-engine/src/native_decode/cuda.rs +++ b/crates/onnx-genai-engine/src/native_decode/cuda.rs @@ -48,9 +48,6 @@ pub(crate) fn resolve_graph_capture_enabled( structural.is_capture_safe() } -const CUDA_KV_MEMORY_BUDGET_NUMERATOR: usize = 85; -const CUDA_KV_MEMORY_BUDGET_DENOMINATOR: usize = 100; - #[derive(Clone, Debug, Eq, PartialEq)] pub struct CudaKvDebugStats { pub logical_len: usize, @@ -529,6 +526,33 @@ impl DecodeCudaState { cuda_kv_capacity_exceeded_message(requested, &self.capacity) } + fn growth_failed_error( + &self, + old_capacity: usize, + new_capacity: usize, + error: anyhow::Error, + memory_at_failure: Result, + ) -> anyhow::Error { + let memory = memory_at_failure + .map(|memory| { + format!( + "CUDA free={} bytes, total={} bytes", + memory.free_bytes, memory.total_bytes + ) + }) + .unwrap_or_else(|err| format!("CUDA free-memory query failed: {err}")); + let new_bytes = new_capacity.saturating_mul(self.capacity.bytes_per_token); + let transient_bytes = old_capacity + .saturating_add(new_capacity) + .saturating_mul(self.capacity.bytes_per_token); + anyhow::anyhow!( + "CUDA KV capacity growth failed while growing from {old_capacity} to {new_capacity} tokens: {error}. \ + The attempted new KV allocation is approximately {new_bytes} bytes and the transient peak is approximately {transient_bytes} bytes because growth keeps the old bucket live until the new bucket and valid-prefix copy are complete. \ + {memory}; KV bytes/token: {}. The session state was left unchanged; reset or retry with a shorter prompt/max_new_tokens, set ONNX_GENAI_CUDA_KV_MAX_LEN/load_with_cuda_kv_max_len to an explicit smaller cap, or free VRAM used by other processes.", + self.capacity.bytes_per_token + ) + } + fn ensure_capacity( &mut self, session: &mut InferenceSession, @@ -1208,49 +1232,58 @@ impl onnx_genai_kv::KvCapacityGrowthBackend for NativeCudaCapacityBackend<'_> { new_capacity: usize, valid_len: usize, ) -> anyhow::Result { - native_cuda_device_barrier(self.session)?; - let mut replacements = Vec::with_capacity(self.state.kv_binding_range.len()); - for index in self.state.kv_binding_range.clone() { - let old = &self.state.bindings[index]; - let mut physical_shape = old.physical_shape().to_vec(); - physical_shape[2] = new_capacity; - let mut logical_shape = physical_shape.clone(); - logical_shape[2] = valid_len; - let new_binding = self.session.allocate_device_binding( - old.input_name().to_owned(), - old.output_name().map(str::to_owned), - old.dtype, - physical_shape.clone(), - logical_shape, - )?; - let dst = new_binding.device_ptr() as usize; - let src = old.device_ptr() as usize; - let total_bytes = - checked_shape_bytes(&physical_shape, old.dtype).with_context(|| { - format!( - "CUDA KV grow allocation size overflow for '{}' shape {:?}", - old.input_name(), - physical_shape - ) - })?; - native_cuda_memset_zero(dst, total_bytes)?; - copy_kv_prefix_device_to_device( - dst, - src, - old.physical_shape(), - &physical_shape, - 2, - valid_len, - old.dtype.checked_storage_bytes(1).with_context(|| { - format!( - "CUDA KV grow copy element size overflow for '{}'", - old.input_name() - ) - })?, - )?; - replacements.push((index, new_binding)); - } - Ok(NativeCudaGrownBuffers { replacements }) + let old_capacity = self.state.max_len; + (|| { + native_cuda_device_barrier(self.session)?; + let mut replacements = Vec::with_capacity(self.state.kv_binding_range.len()); + for index in self.state.kv_binding_range.clone() { + let old = &self.state.bindings[index]; + let mut physical_shape = old.physical_shape().to_vec(); + physical_shape[2] = new_capacity; + let mut logical_shape = physical_shape.clone(); + logical_shape[2] = valid_len; + let new_binding = self.session.allocate_device_binding( + old.input_name().to_owned(), + old.output_name().map(str::to_owned), + old.dtype, + physical_shape.clone(), + logical_shape, + )?; + let dst = new_binding.device_ptr() as usize; + let src = old.device_ptr() as usize; + let total_bytes = + checked_shape_bytes(&physical_shape, old.dtype).with_context(|| { + format!( + "CUDA KV grow allocation size overflow for '{}' shape {:?}", + old.input_name(), + physical_shape + ) + })?; + native_cuda_memset_zero(dst, total_bytes)?; + copy_kv_prefix_device_to_device( + dst, + src, + old.physical_shape(), + &physical_shape, + 2, + valid_len, + old.dtype.checked_storage_bytes(1).with_context(|| { + format!( + "CUDA KV grow copy element size overflow for '{}'", + old.input_name() + ) + })?, + )?; + replacements.push((index, new_binding)); + } + Ok(NativeCudaGrownBuffers { replacements }) + })() + .map_err(|error| { + let memory = cuda_device_memory_snapshot(self.session.device_id().index as i32) + .map_err(|error| error.to_string()); + self.state + .growth_failed_error(old_capacity, new_capacity, error, memory) + }) } fn build_grown_mask( @@ -1258,19 +1291,28 @@ impl onnx_genai_kv::KvCapacityGrowthBackend for NativeCudaCapacityBackend<'_> { new_capacity: usize, valid_len: usize, ) -> anyhow::Result> { - let mut new_mask = self.session.allocate_device_binding( - self.state.bindings[0].input_name().to_owned(), - self.state.bindings[0].output_name().map(str::to_owned), - DataType::Int64, - vec![1, new_capacity], - vec![1, valid_len], - )?; - native_cuda_memset_zero( - new_mask.device_ptr() as usize, - new_capacity * std::mem::size_of::(), - )?; - new_mask.set_logical_shape(vec![1, valid_len])?; - Ok(Some(new_mask)) + let old_capacity = self.state.max_len; + (|| { + let mut new_mask = self.session.allocate_device_binding( + self.state.bindings[0].input_name().to_owned(), + self.state.bindings[0].output_name().map(str::to_owned), + DataType::Int64, + vec![1, new_capacity], + vec![1, valid_len], + )?; + native_cuda_memset_zero( + new_mask.device_ptr() as usize, + new_capacity * std::mem::size_of::(), + )?; + new_mask.set_logical_shape(vec![1, valid_len])?; + Ok(Some(new_mask)) + })() + .map_err(|error| { + let memory = cuda_device_memory_snapshot(self.session.device_id().index as i32) + .map_err(|error| error.to_string()); + self.state + .growth_failed_error(old_capacity, new_capacity, error, memory) + }) } fn invalidate_capture(&mut self) -> anyhow::Result<()> { @@ -1423,47 +1465,39 @@ pub(crate) fn resolve_cuda_kv_capacity( if bytes_per_token == 0 { bail!("CUDA KV bytes per token must be greater than zero"); } - let memory_budget_tokens = device_memory.and_then(|memory| { - let whole = (memory.free_bytes / CUDA_KV_MEMORY_BUDGET_DENOMINATOR) - .checked_mul(CUDA_KV_MEMORY_BUDGET_NUMERATOR)?; - let remainder = (memory.free_bytes % CUDA_KV_MEMORY_BUDGET_DENOMINATOR) - .checked_mul(CUDA_KV_MEMORY_BUDGET_NUMERATOR)? - / CUDA_KV_MEMORY_BUDGET_DENOMINATOR; - whole - .checked_add(remainder) - .and_then(|budget| budget.checked_div(bytes_per_token)) - }); - let memory_budget_len = memory_budget_tokens.map(max_kv_capacity_for_growth_budget); let (configured, source) = if let Some(max_len) = programmatic_max_len { - (max_len, "programmatic load_with_cuda_kv_max_len".to_owned()) - } else if let Some(max_len) = env_max_len { - (max_len, "ONNX_GENAI_CUDA_KV_MAX_LEN".to_owned()) - } else if let Some(max_len) = metadata_max_len { - let metadata_len = max_len; - let max_len = memory_budget_len.map_or(metadata_len, |budget| metadata_len.min(budget)); - ( - max_len, - match (device_memory, memory_budget_len) { - (Some(memory), Some(budget)) if budget < metadata_len => format!( - "model.max_sequence_length clamped by CUDA free-memory growth budget (free={} bytes, total={} bytes, budget={} tokens at {} bytes/token)", - memory.free_bytes, memory.total_bytes, budget, bytes_per_token + if let Some(metadata_len) = metadata_max_len + && metadata_len < max_len + { + ( + metadata_len, + format!( + "programmatic load_with_cuda_kv_max_len clamped by model.max_sequence_length ({metadata_len})" ), - (Some(memory), Some(budget)) => format!( - "model.max_sequence_length within CUDA free-memory growth budget (free={} bytes, total={} bytes, budget={} tokens at {} bytes/token)", - memory.free_bytes, memory.total_bytes, budget, bytes_per_token + ) + } else { + (max_len, "programmatic load_with_cuda_kv_max_len".to_owned()) + } + } else if let Some(max_len) = env_max_len { + if let Some(metadata_len) = metadata_max_len + && metadata_len < max_len + { + ( + metadata_len, + format!( + "ONNX_GENAI_CUDA_KV_MAX_LEN clamped by model.max_sequence_length ({metadata_len})" ), - _ => "model.max_sequence_length (CUDA free-memory query unavailable)".to_owned(), - }, - ) - } else if let Some(budget) = memory_budget_len { + ) + } else { + (max_len, "ONNX_GENAI_CUDA_KV_MAX_LEN".to_owned()) + } + } else if let Some(max_len) = metadata_max_len { + (max_len, "model.max_sequence_length".to_owned()) + } else { ( - budget, - "CUDA free-memory budget (metadata model.max_sequence_length unavailable)".to_owned(), + usize::MAX, + "unbounded (model.max_sequence_length unavailable)".to_owned(), ) - } else { - bail!( - "cannot derive native CUDA KV capacity: no explicit max length, model.max_sequence_length is unavailable, and CUDA free-memory could not be queried; set ONNX_GENAI_CUDA_KV_MAX_LEN or use load_with_cuda_kv_max_len" - ); }; if configured == 0 { bail!( @@ -1479,39 +1513,6 @@ pub(crate) fn resolve_cuda_kv_capacity( }) } -fn max_kv_capacity_for_growth_budget(budget_tokens: usize) -> usize { - if budget_tokens == 0 { - return 0; - } - let initial = onnx_genai_kv::kv_capacity_bucket(0, budget_tokens); - if initial >= budget_tokens { - return budget_tokens; - } - let mut previous = initial; - loop { - let Some(next) = previous.checked_mul(2) else { - return budget_tokens; - }; - if next >= budget_tokens { - return budget_tokens - .checked_sub(previous) - .filter(|candidate| *candidate > previous) - .unwrap_or(previous); - } - if previous - .checked_add(next) - .is_some_and(|transient| transient <= budget_tokens) - { - previous = next; - continue; - } - return budget_tokens - .checked_sub(previous) - .filter(|candidate| *candidate > previous) - .unwrap_or(previous); - } -} - pub(crate) fn cuda_kv_capacity_exceeded_message( requested: usize, capacity: &CudaKvCapacity, diff --git a/crates/onnx-genai-engine/src/native_decode/tests.rs b/crates/onnx-genai-engine/src/native_decode/tests.rs index 83f138755b..03caa54306 100644 --- a/crates/onnx-genai-engine/src/native_decode/tests.rs +++ b/crates/onnx-genai-engine/src/native_decode/tests.rs @@ -4,6 +4,20 @@ use onnx_runtime_ir::{Attribute, Graph, Node, NodeId, Shape, SymbolId, TensorDat use prost::Message; use std::collections::BTreeMap; +#[cfg(feature = "cuda")] +fn qwen_cuda_smoke_model_dir() -> Option { + let model_dir = std::env::var_os("ONNX_GENAI_QWEN_CUDA_SMOKE_MODEL") + .map(std::path::PathBuf::from) + .unwrap_or_else(|| std::path::PathBuf::from("/home/justinchu/qwen2.5-0.5b-int4-onnx")); + if !model_dir.join("model.onnx").is_file() { + eprintln!( + "skipping CUDA smoke; target model is not installed (set ONNX_GENAI_QWEN_CUDA_SMOKE_MODEL to its directory)" + ); + return None; + } + Some(model_dir) +} + #[test] fn tensor_argmax_reads_only_the_final_logits_row_and_keeps_first_tie() { let tensor = Tensor::from_f32(&[1, 2, 4], &[100.0, 0.0, 0.0, 0.0, 1.0, 7.0, 7.0, 2.0]).unwrap(); @@ -39,7 +53,7 @@ fn graph_capture_auto_enables_for_owned_cuda_kv() { } #[test] -fn cuda_kv_capacity_uses_metadata_with_memory_clamp() { +fn cuda_kv_capacity_uses_metadata_as_default_fact() { let capacity = resolve_cuda_kv_capacity( None, None, @@ -51,20 +65,16 @@ fn cuda_kv_capacity_uses_metadata_with_memory_clamp() { }), ) .unwrap(); - assert_eq!(capacity.max_len, 1024); - assert!( - capacity - .source - .contains("model.max_sequence_length clamped by CUDA free-memory growth budget") - ); + assert_eq!(capacity.max_len, 4096); + assert_eq!(capacity.source, "model.max_sequence_length"); } #[test] -fn cuda_kv_capacity_clamp_accounts_for_allocate_before_free_growth() { +fn cuda_kv_capacity_without_metadata_is_unbounded_until_growth_fails() { let capacity = resolve_cuda_kv_capacity( None, None, - Some(131_072), + None, 28_680, Some(CudaDeviceMemorySnapshot { free_bytes: 5_925_502_976, @@ -72,11 +82,10 @@ fn cuda_kv_capacity_clamp_accounts_for_allocate_before_free_growth() { }), ) .unwrap(); - assert_eq!(capacity.max_len, 110_080); - assert!( - capacity - .source - .contains("model.max_sequence_length clamped by CUDA free-memory growth budget") + assert_eq!(capacity.max_len, usize::MAX); + assert_eq!( + capacity.source, + "unbounded (model.max_sequence_length unavailable)" ); } @@ -85,7 +94,7 @@ fn cuda_kv_capacity_honors_env_before_metadata() { let capacity = resolve_cuda_kv_capacity( None, Some(8192), - Some(4096), + Some(16_384), 10, Some(CudaDeviceMemorySnapshot { free_bytes: 20_000, @@ -97,11 +106,24 @@ fn cuda_kv_capacity_honors_env_before_metadata() { assert_eq!(capacity.source, "ONNX_GENAI_CUDA_KV_MAX_LEN"); } +#[test] +fn cuda_kv_capacity_metadata_caps_oversized_explicit_override() { + let capacity = resolve_cuda_kv_capacity(None, Some(8192), Some(4096), 10, None).unwrap(); + assert_eq!(capacity.max_len, 4096); + assert!( + capacity + .source + .contains("ONNX_GENAI_CUDA_KV_MAX_LEN clamped by model.max_sequence_length"), + "{}", + capacity.source + ); +} + #[test] fn cuda_kv_capacity_error_explains_source_and_device_memory() { let capacity = CudaKvCapacity { - max_len: 1024, - source: "model.max_sequence_length clamped by CUDA free-memory growth budget".to_owned(), + max_len: 4096, + source: "model.max_sequence_length".to_owned(), metadata_max_len: Some(4096), device_memory: Some(CudaDeviceMemorySnapshot { free_bytes: 20_000, @@ -109,12 +131,12 @@ fn cuda_kv_capacity_error_explains_source_and_device_memory() { }), bytes_per_token: 10, }; - let message = cuda_kv_capacity_exceeded_message(1025, &capacity); + let message = cuda_kv_capacity_exceeded_message(4097, &capacity); assert!( - message.contains("requested context length 1025"), + message.contains("requested context length 4097"), "{message}" ); - assert!(message.contains("configured max_len 1024"), "{message}"); + assert!(message.contains("configured max_len 4096"), "{message}"); assert!( message.contains("source: model.max_sequence_length"), "{message}" @@ -1477,11 +1499,9 @@ fn native_cuda_qwen_decode_matches_cpu_tokens() -> anyhow::Result<()> { return Ok(()); } - let model_dir = Path::new("/home/justinchu/qwen2.5-0.5b-int4-onnx"); - if !model_dir.join("model.onnx").is_file() { - eprintln!("skipping CUDA smoke; target model is not installed"); + let Some(model_dir) = qwen_cuda_smoke_model_dir() else { return Ok(()); - } + }; let tokenizer = Tokenizer::from_file(model_dir.join("tokenizer.json"))?; let prompt = tokenizer.encode("Hello")?; const HORIZON: usize = 64; @@ -1607,11 +1627,9 @@ fn native_cuda_verify_rewind_no_kv_corruption() -> anyhow::Result<()> { eprintln!("skipping CUDA smoke; set ONNX_GENAI_RUN_CUDA_SMOKE=1 to run"); return Ok(()); } - let model_dir = Path::new("/home/justinchu/qwen2.5-0.5b-int4-onnx"); - if !model_dir.join("model.onnx").is_file() { - eprintln!("skipping CUDA smoke; target model is not installed"); + let Some(model_dir) = qwen_cuda_smoke_model_dir() else { return Ok(()); - } + }; let tokenizer = Tokenizer::from_file(model_dir.join("tokenizer.json"))?; let prompt = tokenizer.encode("The quick brown fox")?; diff --git a/crates/onnx-genai-kv/src/lib.rs b/crates/onnx-genai-kv/src/lib.rs index 3bffcad3a3..990441d28b 100644 --- a/crates/onnx-genai-kv/src/lib.rs +++ b/crates/onnx-genai-kv/src/lib.rs @@ -372,4 +372,203 @@ mod tests { assert_eq!(backend.events, ["buffers", "mask"]); assert_eq!(backend.current, 256); } + + #[derive(Clone, Copy, Debug, PartialEq, Eq)] + enum InjectedFailure { + Allocate, + PrefixCopy, + Mask, + InvalidateCapture, + } + + struct TransactionalFakeBackend { + current: usize, + hard: usize, + valid: usize, + bytes_per_token: usize, + buffer_id: usize, + capture_generation: usize, + capture_valid: bool, + events: Vec<&'static str>, + fail_at: Option, + } + + impl TransactionalFakeBackend { + fn new(fail_at: Option) -> Self { + Self { + current: 256, + hard: 4096, + valid: 123, + bytes_per_token: 32, + buffer_id: 7, + capture_generation: 11, + capture_valid: true, + events: Vec::new(), + fail_at, + } + } + + fn translated_growth_error(&self, old: usize, new: usize, raw: &str) -> String { + let new_bytes = new * self.bytes_per_token; + let transient_peak = (old + new) * self.bytes_per_token; + format!( + "injected KV capacity growth failed while growing from {old} to {new} tokens: {raw}. \ + The attempted new KV allocation is approximately {new_bytes} bytes and the transient peak is approximately {transient_peak} bytes. \ + KV bytes/token: {}. The session state was left unchanged; reset or retry with a shorter prompt/max_new_tokens, set an explicit KV max length cap, or free VRAM used by other processes.", + self.bytes_per_token + ) + } + + fn assert_unchanged(&self) { + assert_eq!(self.current, 256); + assert_eq!(self.valid, 123); + assert_eq!(self.buffer_id, 7); + assert_eq!(self.capture_generation, 11); + assert!(self.capture_valid); + } + } + + impl KvCapacityGrowthBackend for TransactionalFakeBackend { + type Error = String; + type GrownBuffers = usize; + type GrownMask = usize; + + fn current_capacity(&self) -> usize { + self.current + } + + fn hard_max_capacity(&self) -> usize { + self.hard + } + + fn valid_len(&self) -> usize { + self.valid + } + + fn capacity_exceeded(&self, required: usize) -> Self::Error { + format!("capacity exceeded at {required}") + } + + fn build_grown_buffers( + &mut self, + new_capacity: usize, + _valid_len: usize, + ) -> Result { + self.events.push("allocate"); + if self.fail_at == Some(InjectedFailure::Allocate) { + return Err(self.translated_growth_error(self.current, new_capacity, "raw OOM")); + } + self.events.push("prefix-copy"); + if self.fail_at == Some(InjectedFailure::PrefixCopy) { + return Err(self.translated_growth_error( + self.current, + new_capacity, + "raw copy failure", + )); + } + Ok(new_capacity + 1000) + } + + fn build_grown_mask( + &mut self, + new_capacity: usize, + _valid_len: usize, + ) -> Result, Self::Error> { + self.events.push("mask"); + if self.fail_at == Some(InjectedFailure::Mask) { + return Err(self.translated_growth_error( + self.current, + new_capacity, + "raw mask allocation failure", + )); + } + Ok(Some(new_capacity + 2000)) + } + + fn invalidate_capture(&mut self) -> Result<(), Self::Error> { + self.events.push("invalidate"); + if self.fail_at == Some(InjectedFailure::InvalidateCapture) { + return Err(self.translated_growth_error( + self.current, + self.current * 2, + "raw capture release failure", + )); + } + self.capture_valid = false; + Ok(()) + } + + fn commit_grown_capacity( + &mut self, + new_capacity: usize, + grown_buffers: Self::GrownBuffers, + grown_mask: Option, + ) -> Result<(), Self::Error> { + self.events.push("commit"); + assert_eq!(grown_buffers, new_capacity + 1000); + assert_eq!(grown_mask, Some(new_capacity + 2000)); + self.current = new_capacity; + self.buffer_id = grown_buffers; + self.capture_generation += 1; + self.capture_valid = true; + Ok(()) + } + } + + #[test] + fn injected_allocate_failure_is_actionable_transactional_and_retryable() { + let mut backend = TransactionalFakeBackend::new(Some(InjectedFailure::Allocate)); + let error = ensure_kv_capacity(&mut backend, 257).unwrap_err(); + + assert!(error.contains("growing from 256 to 512 tokens"), "{error}"); + assert!( + error.contains("transient peak is approximately 24576 bytes"), + "{error}" + ); + assert!(error.contains("KV bytes/token: 32"), "{error}"); + assert!( + error.contains("session state was left unchanged"), + "{error}" + ); + assert!(error.contains("shorter prompt/max_new_tokens"), "{error}"); + assert!(error.contains("free VRAM"), "{error}"); + assert_eq!(backend.events, ["allocate"]); + backend.assert_unchanged(); + + backend.fail_at = None; + let growth = ensure_kv_capacity(&mut backend, 257).unwrap(); + assert_eq!( + growth, + KvCapacityGrowth::Grew { + old_capacity: 256, + new_capacity: 512, + valid_len: 123, + } + ); + assert_eq!(backend.current, 512); + assert_eq!(backend.buffer_id, 1512); + assert_eq!(backend.capture_generation, 12); + assert!(backend.capture_valid); + } + + #[test] + fn injected_mid_sequence_failures_do_not_commit_or_invalidate_capture() { + for failure in [ + InjectedFailure::PrefixCopy, + InjectedFailure::Mask, + InjectedFailure::InvalidateCapture, + ] { + let mut backend = TransactionalFakeBackend::new(Some(failure)); + let error = ensure_kv_capacity(&mut backend, 257).unwrap_err(); + assert!( + error.contains("session state was left unchanged"), + "{error}" + ); + backend.assert_unchanged(); + assert!(!backend.events.contains(&"commit")); + if failure != InjectedFailure::InvalidateCapture { + assert!(!backend.events.contains(&"invalidate")); + } + } + } } diff --git a/crates/onnx-genai-ort/src/decode/dynamic.rs b/crates/onnx-genai-ort/src/decode/dynamic.rs index c50a2f4b64..1c92b21b4f 100644 --- a/crates/onnx-genai-ort/src/decode/dynamic.rs +++ b/crates/onnx-genai-ort/src/decode/dynamic.rs @@ -1291,28 +1291,32 @@ impl onnx_genai_kv::KvCapacityGrowthBackend for DecodeSession<'_> { // Build the replacement buffers first; only swap them in once every KV // tensor has been copied successfully, so a mid-way failure leaves the // session's existing state intact. - let mut grown = HashMap::with_capacity(self.kv_pairs.len()); - for pair in &self.kv_pairs { - let old = self.current_kv.get(&pair.past).ok_or_else(|| { - OrtError::InvalidArgument(format!( - "cannot grow KV: missing shared buffer for '{}'", - pair.past - )) - })?; - let mut new_shape = old.shape().to_vec(); - new_shape[pair.seq_axis] = i64::try_from(new_capacity) - .map_err(|_| OrtError::InvalidArgument("KV capacity exceeds i64".into()))?; - let new_value = grow_kv_value( - old, - &new_shape, - pair.seq_axis, - valid_len, - device, - self.kv_allocator.as_ref(), - )?; - grown.insert(pair.past.clone(), Arc::new(new_value)); - } - Ok(grown) + let old_capacity = self.kv_capacity; + (|| { + let mut grown = HashMap::with_capacity(self.kv_pairs.len()); + for pair in &self.kv_pairs { + let old = self.current_kv.get(&pair.past).ok_or_else(|| { + OrtError::InvalidArgument(format!( + "cannot grow KV: missing shared buffer for '{}'", + pair.past + )) + })?; + let mut new_shape = old.shape().to_vec(); + new_shape[pair.seq_axis] = i64::try_from(new_capacity) + .map_err(|_| OrtError::InvalidArgument("KV capacity exceeds i64".into()))?; + let new_value = grow_kv_value( + old, + &new_shape, + pair.seq_axis, + valid_len, + device, + self.kv_allocator.as_ref(), + )?; + grown.insert(pair.past.clone(), Arc::new(new_value)); + } + Ok(grown) + })() + .map_err(|error| self.kv_growth_failed_error(old_capacity, new_capacity, error)) } fn build_grown_mask( @@ -1333,13 +1337,15 @@ impl onnx_genai_kv::KvCapacityGrowthBackend for DecodeSession<'_> { { let mask_len = i64::try_from(new_capacity) .map_err(|_| OrtError::InvalidArgument("KV capacity exceeds i64".into()))?; - Ok(Some(self.build_grown_capture_mask( + self.build_grown_capture_mask( new_capacity, mask_len, valid_ones, on_device, mask_device_id, - )?)) + ) + .map(Some) + .map_err(|error| self.kv_growth_failed_error(self.kv_capacity, new_capacity, error)) } else { Ok(None) } @@ -1370,6 +1376,76 @@ impl onnx_genai_kv::KvCapacityGrowthBackend for DecodeSession<'_> { } impl DecodeSession<'_> { + fn kv_growth_bytes_per_token(&self) -> Option { + let mut bytes = if self.capture.is_some() { + std::mem::size_of::() + } else { + 0 + }; + for pair in &self.kv_pairs { + let value = self.current_kv.get(&pair.past)?; + let shape = value.shape(); + let mut elements = 1usize; + for (axis, dim) in shape.iter().copied().enumerate() { + if axis == pair.seq_axis { + continue; + } + let dim = usize::try_from(dim).ok()?; + elements = elements.checked_mul(dim)?; + } + bytes = bytes.checked_add(elements.checked_mul(value.dtype().size_of())?)?; + } + Some(bytes) + } + + fn kv_growth_failed_error( + &self, + old_capacity: usize, + new_capacity: usize, + error: OrtError, + ) -> OrtError { + let bytes_per_token = self.kv_growth_bytes_per_token(); + let byte_summary = bytes_per_token + .map(|bytes_per_token| { + let new_bytes = new_capacity.saturating_mul(bytes_per_token); + let transient_bytes = old_capacity + .saturating_add(new_capacity) + .saturating_mul(bytes_per_token); + format!( + "The attempted new KV allocation is approximately {new_bytes} bytes and the transient peak is approximately {transient_bytes} bytes because growth keeps the old bucket live until the new bucket and valid-prefix copy are complete. KV bytes/token: {bytes_per_token}." + ) + }) + .unwrap_or_else(|| { + "KV bytes/token could not be derived from the current shared buffers.".to_owned() + }); + let memory = self + .session + .cuda_device_id() + .and_then(|device_id| { + #[cfg(feature = "cuda")] + { + crate::cuda_rt::device_memory_info(device_id) + .ok() + .map(|memory| { + format!( + "CUDA free={} bytes, total={} bytes", + memory.free_bytes, memory.total_bytes + ) + }) + } + #[cfg(not(feature = "cuda"))] + { + let _ = device_id; + None + } + }) + .unwrap_or_else(|| "device free-memory query unavailable".to_owned()); + OrtError::InvalidArgument(format!( + "shared-buffer KV capacity growth failed while growing from {old_capacity} to {new_capacity} tokens: {error}. \ + {byte_summary} {memory}. The session state was left unchanged; reset or retry with a shorter prompt/max_new_tokens, set an explicit KV max length cap, or free VRAM used by other processes." + )) + } + /// Build a replacement captured attention mask of `capacity` elements /// (shape `[1, mask_len]`) with the leading `valid_ones` set to 1 and the /// rest zero, preserving the mask's residency: device-resident when diff --git a/docs/native-ort-kv-capacity.md b/docs/native-ort-kv-capacity.md index c1ab054a4d..4437718cc8 100644 --- a/docs/native-ort-kv-capacity.md +++ b/docs/native-ort-kv-capacity.md @@ -3,20 +3,17 @@ The runtime uses one KV capacity policy for ORT shared-buffer KV and native CUDA KV: `onnx_genai_kv::kv_capacity_bucket(required, hard_max)`. The policy rounds demand to a power-of-two bucket, applies the `ONNX_GENAI_KV_MIN_BUCKET` floor, -and clamps to a hard maximum derived from model metadata and device capacity. +and clamps only to factual limits: model metadata or an explicit user cap. ## Hard maximum -The hard maximum is not a model-family constant. It comes from explicit caller -configuration first (`load_with_cuda_kv_max_len`, `ONNX_GENAI_CUDA_KV_MAX_LEN`), -otherwise from `model.max_sequence_length`, clamped by a queried CUDA -free-memory growth budget using the graph's actual KV bytes per token. The -growth budget accounts for the allocation-before-free sequence used by both -backends: at a bucket boundary the old bucket remains live while the larger -bucket and mask are allocated and filled, so the default ceiling is the largest -capacity whose worst bucket transition fits inside the headroom budget. If -neither metadata nor a CUDA memory query can provide a limit, native CUDA fails -during load with an actionable configuration error instead of guessing. +The hard maximum is not a VRAM prediction. It comes from explicit caller +configuration first (`load_with_cuda_kv_max_len`, `ONNX_GENAI_CUDA_KV_MAX_LEN`, +or the ORT shared-buffer cap), otherwise from `model.max_sequence_length`. Those +are facts: the user asked for a cap, or the model cannot attend beyond its +declared context. Native CUDA does not derive a ceiling from a free-memory +snapshot; if metadata is unavailable and the user did not set a cap, it grows +until allocation fails and reports that failure transactionally. ## Growth @@ -31,6 +28,15 @@ capture, then commit the new capacity. ORT implements the primitive seam with rewrites. Adding another backend should implement those primitives rather than re-derive the growth policy. +Growth deliberately asks the device to allocate instead of predicting a safe +VRAM ceiling. A failed grow is expected to be graceful and transactional: the old +bucket, logical length, and capture state remain live, and the error reports the +target bucket, approximate new allocation, approximate transient peak, current +device free/total memory when available, bytes per token, and user levers. The +transient peak matters because growth keeps the old bucket alive while allocating +and filling the new bucket; peak KV+mask memory is roughly `(old + new) * +bytes_per_token`, not just `new * bytes_per_token`. + ## CUDA graph capture Growth and CUDA graph capture coexist by making growth a graph boundary. The @@ -46,3 +52,13 @@ path was exercised manually on an RTX 4060 Laptop GPU by forcing early buckets with `ONNX_GENAI_KV_MIN_BUCKET=4`, observing native CUDA growth from 4โ†’8โ†’16 with graph recapture at each bucket and coherent continued generation. Automated CUDA coverage still depends on a GPU runner and model fixtures being available. + +On Windows/WDDM, forcing a real CUDA allocation failure is not reliable: the GPU +memory manager virtualizes device memory and may page instead of making +`cudaMalloc` fail, even after a separate process reserves enough VRAM for +`cudaMemGetInfo` to report little or no free memory. The graceful-failure path is +therefore covered by injected model-free tests at the shared +`KvCapacityGrowthBackend` seam: allocation, prefix-copy, mask-allocation, and +capture-invalidation failures assert actionable errors and unchanged session +state. Linux/TCC GPU runners should still exercise real allocator failures when +available. From 3531331c21866b39ac0f5c30cee3ebb8b765ded1 Mon Sep 17 00:00:00 2001 From: justinchuby <223556219+Copilot@users.noreply.github.com> Date: Tue, 28 Jul 2026 13:26:48 -0700 Subject: [PATCH 3/4] docs(native-decode): correct stale KV capacity doc on load_with_cuda_kv_max_len MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The old comment said None "derives a default from model metadata clamped by queried CUDA free memory". That mechanism was removed by this PR: free memory no longer supplies a ceiling. The new resolution order is: explicit override โ†’ ONNX_GENAI_CUDA_KV_MAX_LEN โ†’ model.max_sequence_length โ†’ unbounded growth until allocation fails Update the doc comment to match. Sweep of all changed files found nothing else stale: no free-memory-derived ceiling, 85%-fraction, hardcoded-4096 default, or pre-allocated-KV-buffer language in any changed source or doc file. engine/load.rs mentions "pre-allocating" in the past tense as context for why callers may still want an env cap; that phrasing is accurate. docs/native-ort-kv-capacity.md already correctly describes the current design. Failure messages reviewed against RULES.md ยง1: - cuda_kv_capacity_exceeded_message: names requested length, configured max_len with its source, and user actions. Device free/total is diagnostic context, not the origin of the cap. โœ“ - growth_failed_error: names the growth range, transient peak, device memory at failure, and retry levers. Does not reference a removed ceiling. โœ“ Gates: cargo fmt --check โœ“ ยท clippy -D warnings โœ“ ยท cargo test -p onnx-genai-kv (101+1 passed) โœ“ Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- crates/onnx-genai-engine/src/native_decode/load.rs | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/crates/onnx-genai-engine/src/native_decode/load.rs b/crates/onnx-genai-engine/src/native_decode/load.rs index 542021b0f7..786cbaeaa8 100644 --- a/crates/onnx-genai-engine/src/native_decode/load.rs +++ b/crates/onnx-genai-engine/src/native_decode/load.rs @@ -73,9 +73,11 @@ impl NativeDecodeSession { ) } - /// Load with an explicit CUDA KV capacity. `None` uses - /// `ONNX_GENAI_CUDA_KV_MAX_LEN`, then derives a default from model metadata - /// clamped by queried CUDA free memory. + /// Load with an explicit CUDA KV capacity. `None` resolves in order: + /// `ONNX_GENAI_CUDA_KV_MAX_LEN`, then `model.max_sequence_length` from + /// inference metadata, then unbounded growth until the device refuses + /// allocation. No free-memory ceiling is derived; the hard maximum comes + /// only from an explicit override or the model's declared context limit. pub fn load_with_cuda_kv_max_len( path: impl AsRef, device: NativeDecodeDevice, From e10a9531ee041263a5501b5ab4a81eac80a24090 Mon Sep 17 00:00:00 2001 From: justinchuby <223556219+Copilot@users.noreply.github.com> Date: Tue, 28 Jul 2026 15:25:39 -0700 Subject: [PATCH 4/4] docs(kv): tighten ORT shared-buffer KV description to name initial bucket The comment in engine/load.rs and the prose in docs/native-ort-kv-capacity.md both said the shared-buffer path "grows by buckets instead of pre-allocating that full length", which is true but leaves a reader unsure whether anything is allocated at session start. Verified against source before writing: - Initial allocation: kv_capacity_bucket(0, max_length) in dynamic.rs:397 which evaluates to the minimum bucket (256 by default), not max_length. - Default floor: const MIN_BUCKET_DEFAULT: usize = 256 in onnx-genai-kv/src/lib.rs:64. - Env var: ONNX_GENAI_KV_MIN_BUCKET in lib.rs:68; the existing kv_capacity_bucket docstring already says it can "raise or lower" the floor, so "overridden by" is used here rather than "raised by". The update says both things explicitly: "pre-allocates at an initial power-of-two bucket (256 tokens by default, overridden by ONNX_GENAI_KV_MIN_BUCKET) and grows on demand up to the model's declared max_length; it does not pre-allocate the full context." Gates (verbatim from .github/workflows/ci.yml): cargo fmt --all -- --check exit 0 cargo clippy --locked --all-targets -p onnx-genai-kv \ -p onnx-genai-metadata -p onnx-genai-runtime-config \ -p onnx-genai-scheduler -p onnx-runtime-ir -p onnx-runtime-ep-api \ -p onnx-runtime-session -p onnx-runtime-tracer -p onnx-std \ -- -D warnings exit 0 cargo test -p onnx-genai-kv (101 unit + 1 integration, 0 failed) exit 0 cargo test -p onnx-genai-engine --lib (267 passed, 0 failed) exit 0 Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- crates/onnx-genai-engine/src/engine/load.rs | 14 ++++++++------ docs/native-ort-kv-capacity.md | 6 ++++-- 2 files changed, 12 insertions(+), 8 deletions(-) diff --git a/crates/onnx-genai-engine/src/engine/load.rs b/crates/onnx-genai-engine/src/engine/load.rs index a7e4e8e65d..6e3cf05bf2 100644 --- a/crates/onnx-genai-engine/src/engine/load.rs +++ b/crates/onnx-genai-engine/src/engine/load.rs @@ -379,13 +379,15 @@ fn resolve_metadata_and_decode_path( anyhow::bail!("Unsupported capabilities: {unsupported:?}"); } - // Optional factual cap on runtime-owned KV growth. Foundry / + // Optional explicit cap on runtime-owned KV growth. Foundry / // onnxruntime-genai `genai_config.json` models advertise the model's full - // `context_length` (e.g. 32k-131k) as their max sequence length. Shared KV - // now grows by buckets instead of pre-allocating that full length, but callers - // can still set `ONNX_GENAI_KV_MAX_LEN` as an explicit generation-budget cap, - // mirroring the native path's `ONNX_GENAI_CUDA_KV_MAX_LEN`. Unset = model - // metadata is the factual ceiling. + // `context_length` (e.g. 32k-131k) as their max sequence length. The + // shared-buffer decode path pre-allocates at an initial power-of-two bucket + // (256 tokens by default, overridden by `ONNX_GENAI_KV_MIN_BUCKET`) and grows + // on demand up to the model's declared `max_length`; it does not pre-allocate + // the full context. `ONNX_GENAI_KV_MAX_LEN` caps growth below the model + // maximum, mirroring the native path's `ONNX_GENAI_CUDA_KV_MAX_LEN`. + // Unset = model metadata is the factual ceiling. let kv_shared_buffer_cap = shared_buffer_cap_from_env(); let metadata_max_context = metadata .model diff --git a/docs/native-ort-kv-capacity.md b/docs/native-ort-kv-capacity.md index 4437718cc8..7eff79f8b4 100644 --- a/docs/native-ort-kv-capacity.md +++ b/docs/native-ort-kv-capacity.md @@ -17,8 +17,10 @@ until allocation fails and reports that failure transactionally. ## Growth -Both ORT shared-buffer KV and native CUDA grow on demand by the shared bucket -policy instead of pre-allocating the full context. The orchestration lives in +Both ORT shared-buffer KV and native CUDA pre-allocate at the minimum bucket +(256 tokens by default, overridden by `ONNX_GENAI_KV_MIN_BUCKET`) and grow on +demand by the shared bucket policy up to the hard maximum, never pre-allocating +the full declared context length. The orchestration lives in `onnx_genai_kv::ensure_kv_capacity` over the `KvCapacityGrowthBackend` trait: reject above the hard maximum, compute the next `kv_capacity_bucket`, build all fallible replacement state while the old state remains live, invalidate graph