fix: default to host KV on CUDA EP to avoid IoBinding segfault

docs/ARCHITECTURE.md

@@ -291,7 +291,7 @@ Sessions are configured through `RuntimeOptions`: execution provider, intra-op t
 
 The default execution provider is CPU. CUDA is available when the crate is built with `--features cuda` and a CUDA-enabled ONNX Runtime is present. Selecting CUDA without that feature fails at session creation with a clear error; with the feature enabled, CUDA EP registration fails hard if the provider cannot initialize (no silent CPU fallback).
 
-On the CUDA path, autoregressive decode keeps KV past/present tensors on device via ORT IoBinding. Token sampling still runs on the host from final-position logits. The host `KvCache` type and CPU decode path remain the default when CUDA is not selected; CUDA helpers are compile-gated behind the `cuda` feature so CPU builds are unchanged.
+On the CUDA path, autoregressive decode still uses the host `KvCache` and `Session::run` by default (graphs may run on GPU via the CUDA EP; KV traffic crosses the host). An experimental IoBinding device-KV path is available behind `FAST_LIGHTONOCR_CUDA_DEVICE_KV=1`; empty past tensors start on the host and present outputs are promoted to device after the first step. Token sampling always runs on the host from final-position logits. CUDA helpers are compile-gated behind the `cuda` feature so CPU builds are unchanged.
 
 The runtime provides lightweight, strongly typed wrappers around the exported ONNX models, exposing Rust domain types such as `ImageTensor`, `ImageFeatures`, `InputEmbeddings`, `AttentionMask`, `Logits`, and `KvCache` rather than raw ONNX tensors.
 

docs/ROADMAP.md

@@ -95,7 +95,7 @@ changing model contracts.
 ### Completed
 
 - ✅ CUDA EP registration and session wiring (`--features cuda`)
-- ✅ Device-resident decoder KV via IoBinding (CUDA generate path; host path unchanged)
+- ✅ CUDA EP generate path (default: host KV + `Session::run`; opt-in IoBinding device KV via `FAST_LIGHTONOCR_CUDA_DEVICE_KV`)
 - ✅ Expose EP / thread options through Python `runtime_kwargs`
 - ✅ Packaging notes for GPU / accelerator runtimes (wheels stay CPU-default)
 

src/model/decoder/decoder.rs

@@ -181,20 +181,21 @@ impl Decoder {
             return Ok(GenerationOutput::new(Vec::new(), FinishReason::Length));
         }
 
+        // Optional IoBinding device-KV path. Default CUDA EP uses the host KV
+        // loop below (Session::run): zero-length CUDA past tensors have been
+        // observed to segfault with ORT IoBinding on some stacks (e.g. Colab).
+        // Enable with FAST_LIGHTONOCR_CUDA_DEVICE_KV=1 for continued debugging.
         #[cfg(feature = "cuda")]
-        if let ExecutionProvider::Cuda { device_id } = self.execution_provider {
+        if let ExecutionProvider::Cuda { device_id } = self.execution_provider
+            && std::env::var_os("FAST_LIGHTONOCR_CUDA_DEVICE_KV").is_some()
+        {
             let device_id =
                 i32::try_from(device_id).map_err(|_| Error::OnnxRuntimeCompatibility {
                     reason: format!("CUDA device_id {device_id} is out of range for i32"),
                 })?;
 
             let mut generated = Vec::with_capacity(self.generation_config.max_new_tokens);
-            let mut kv_state = CudaKvState::empty(
-                &self.session,
-                &self.config,
-                decoder_input.batch_size(),
-                device_id,
-            )?;
+            let mut kv_state = CudaKvState::empty(&self.config, decoder_input.batch_size())?;
             attention_mask.reserve(self.generation_config.max_new_tokens);
 
             let hidden_size = decoder_input.hidden_size();

src/model/decoder/kv_cache.rs

@@ -7,9 +7,7 @@
 use super::DecoderConfig;
 
 #[cfg(feature = "cuda")]
-use ort::memory::{AllocationDevice, Allocator, AllocatorType, MemoryInfo, MemoryType};
-#[cfg(feature = "cuda")]
-use ort::session::Session;
+use ort::memory::{AllocationDevice, AllocatorType, MemoryInfo, MemoryType};
 #[cfg(feature = "cuda")]
 use ort::value::{DynValue, Tensor, TensorElementType, ValueType};
 
@@ -241,26 +239,19 @@ pub(crate) struct CudaKvState {
 
 #[cfg(feature = "cuda")]
 impl CudaKvState {
-    /// Allocates empty `(batch, kv_heads, 0, head_dim)` past tensors on CUDA.
-    pub(crate) fn empty(
-        session: &Session,
-        config: &DecoderConfig,
-        batch_size: usize,
-        device_id: i32,
-    ) -> Result<Self> {
+    /// Creates empty `(batch, kv_heads, 0, head_dim)` past tensors on the **host**.
+    ///
+    /// Zero-length past tensors are allocated on CPU on purpose: CUDA allocations
+    /// with a zero sequence dimension are unreliable with ORT IoBinding and have
+    /// been observed to segfault. After the first decode step, [`Self::promote_present`]
+    /// replaces these with device-resident present tensors.
+    pub(crate) fn empty(config: &DecoderConfig, batch_size: usize) -> Result<Self> {
         if batch_size == 0 {
             return Err(Error::InvalidKvCache {
                 reason: "batch size must be greater than zero".to_owned(),
             });
         }
 
-        let allocator =
-            Allocator::new(session, cuda_memory_info(device_id)?).map_err(|source| {
-                Error::OnnxRuntimeCompatibility {
-                    reason: format!("failed to create CUDA allocator: {source}"),
-                }
-            })?;
-
         let shape = [
             batch_size as i64,
             config.num_key_value_heads as i64,
@@ -271,16 +262,18 @@ impl CudaKvState {
         let mut past_keys = Vec::with_capacity(config.num_hidden_layers);
         let mut past_values = Vec::with_capacity(config.num_hidden_layers);
         for _ in 0..config.num_hidden_layers {
-            let key = Tensor::<f32>::new(&allocator, shape).map_err(|source| {
-                Error::OnnxRuntimeCompatibility {
-                    reason: format!("failed to allocate empty CUDA past key: {source}"),
-                }
-            })?;
-            let value = Tensor::<f32>::new(&allocator, shape).map_err(|source| {
+            // Empty f32 buffer with a zero-length sequence axis.
+            let key = Tensor::<f32>::from_array((shape, Vec::<f32>::new())).map_err(|source| {
                 Error::OnnxRuntimeCompatibility {
-                    reason: format!("failed to allocate empty CUDA past value: {source}"),
+                    reason: format!("failed to create empty host past key: {source}"),
                 }
             })?;
+            let value =
+                Tensor::<f32>::from_array((shape, Vec::<f32>::new())).map_err(|source| {
+                    Error::OnnxRuntimeCompatibility {
+                        reason: format!("failed to create empty host past value: {source}"),
+                    }
+                })?;
             past_keys.push(key.into_dyn());
             past_values.push(value.into_dyn());
         }