feat(decider): bounded request-local CUDA batching

.github/workflows/ci.yml

@@ -42,6 +42,8 @@ jobs:
       - run: python -m pip install -r benchmarks/requirements.txt
       - run: python benchmarks/bench.py validate benchmarks/smoke.jsonl
       - run: python -m unittest discover -s tests/benchmarks -p 'test_*.py' -v
+      - name: Decider validation helpers
+        run: python -m unittest discover -s tests/decider -p 'test_*.py' -v
 
   cuda-contract:
     # The CUDA backend contract and its checker. No CUDA toolkit and no GPU: both

src/models/decider/native/Cargo.toml

@@ -34,6 +34,10 @@ path = "../../../../tests/decider/contract.rs"
 name = "config"
 path = "../../../../tests/decider/config.rs"
 
+[[test]]
+name = "batching"
+path = "../../../../tests/decider/batching.rs"
+
 [[example]]
 name = "decider-run"
 path = "../../../../tests/decider/runner.rs"

src/models/decider/native/src/batching.rs

@@ -0,0 +1,82 @@
+//! Bounded packing of independent complete rows within a single admitted request.
+use anyhow::{Context, Result, ensure};
+use std::ops::Range;
+
+#[derive(Clone, Copy, Debug)]
+pub struct BatchLimits {
+    max_rows: usize,
+    max_tokens: usize,
+}
+impl Default for BatchLimits {
+    fn default() -> Self {
+        Self {
+            max_rows: 1,
+            max_tokens: 4096,
+        }
+    }
+}
+impl BatchLimits {
+    pub fn new(max_rows: usize, max_tokens: usize) -> Result<Self> {
+        ensure!(
+            (1..=4).contains(&max_rows),
+            "batch rows must be within 1..=4"
+        );
+        ensure!(
+            (1..=4096).contains(&max_tokens),
+            "batch tokens must be within 1..=4096"
+        );
+        Ok(Self {
+            max_rows,
+            max_tokens,
+        })
+    }
+    pub fn from_values(rows: Option<&str>, tokens: Option<&str>) -> Result<Self> {
+        let rows = rows
+            .map_or(Ok(1), str::parse)
+            .context("DECIDER_BATCH_MAX_ROWS must be an integer")?;
+        let tokens = tokens
+            .map_or(Ok(4096), str::parse)
+            .context("DECIDER_BATCH_MAX_TOKENS must be an integer")?;
+        Self::new(rows, tokens)
+    }
+    pub fn from_env() -> Result<Self> {
+        let get = |name| match std::env::var(name) {
+            Ok(value) => Ok(Some(value)),
+            Err(std::env::VarError::NotPresent) => Ok(None),
+            Err(e) => Err(e),
+        };
+        let rows = get("DECIDER_BATCH_MAX_ROWS")?;
+        let tokens = get("DECIDER_BATCH_MAX_TOKENS")?;
+        Self::from_values(rows.as_deref(), tokens.as_deref())
+    }
+    pub fn max_rows(&self) -> usize {
+        self.max_rows
+    }
+    pub fn max_tokens(&self) -> usize {
+        self.max_tokens
+    }
+    /// Retain original row order; a row above the packing token budget runs alone.
+    /// Total request/row admission remains the processor/executor's separate limit.
+    pub fn ranges(&self, lengths: &[usize]) -> Vec<Range<usize>> {
+        let mut ranges = Vec::new();
+        let mut start = 0;
+        let mut tokens = 0usize;
+        for (index, &length) in lengths.iter().enumerate() {
+            if index > start
+                && (index - start == self.max_rows
+                    || tokens
+                        .checked_add(length)
+                        .is_none_or(|total| total > self.max_tokens))
+            {
+                ranges.push(start..index);
+                start = index;
+                tokens = 0;
+            }
+            tokens = tokens.saturating_add(length);
+        }
+        if start < lengths.len() {
+            ranges.push(start..lengths.len());
+        }
+        ranges
+    }
+}

src/models/decider/native/src/engine.rs

@@ -1,5 +1,8 @@
 //! Assemble verified artifacts, model-owned processing, and serial eager execution.
-use crate::{Config, Kind, Limits, Processor, checkpoint::Checkpoint, executor::Executor};
+use crate::{
+    Config, Kind, Limits, Processor, batching::BatchLimits, checkpoint::Checkpoint,
+    executor::Executor,
+};
 use anyhow::Result;
 use omni_runtime::SerialScheduler;
 use serde_json::{Value, json};
@@ -13,8 +16,15 @@ pub struct Engine {
 }
 impl Engine {
     pub async fn load(dir: &Path, library: &Path) -> Result<Self> {
+        Self::load_with_batch(dir, library, BatchLimits::from_env()?).await
+    }
+    pub async fn load_with_batch(
+        dir: &Path,
+        library: &Path,
+        batching: BatchLimits,
+    ) -> Result<Self> {
         let directory = dir.to_owned();
-        let (processor,checkpoint,metadata) = tokio::task::spawn_blocking(move || -> Result<_> {
+        let (processor,checkpoint,mut metadata) = tokio::task::spawn_blocking(move || -> Result<_> {
             let processor = Processor::load(&directory,Limits::default())?;
             let labels: Vec<u32> = processor.labels().iter().map(|label| label.id).collect();
             let checkpoint = Checkpoint::load(&directory,&labels)?;
@@ -23,7 +33,9 @@ impl Engine {
             Ok((processor,checkpoint,metadata))
         }).await??;
         let labels = processor.labels().iter().map(|label| label.id).collect();
-        let executor = Executor::load(dir, library, checkpoint, labels).await?;
+        let executor = Executor::load(dir, library, checkpoint, labels, batching).await?;
+        metadata["batch_limits"] =
+            json!({"rows":batching.max_rows(),"tokens":batching.max_tokens()});
         let engine = Self {
             processor: Arc::new(processor),
             executor,

src/models/decider/native/src/executor.rs

@@ -1,6 +1,7 @@
 //! Eager complete-row prefills and a selected tied-embedding CUDA readout.
 use crate::{
     Limits, RowInput,
+    batching::BatchLimits,
     checkpoint::{Checkpoint, HIDDEN, LABELS, PADDED_LABELS, VOCAB},
 };
 use anyhow::{Context, Result, ensure};
@@ -45,6 +46,7 @@ struct Head {
     weights: DeviceBuffer,
     input: DeviceBuffer,
     output: DeviceBuffer,
+    capacity: usize,
 }
 impl Drop for Head {
     fn drop(&mut self) {
@@ -53,7 +55,8 @@ impl Drop for Head {
     }
 }
 impl Head {
-    fn new(weights: &[u8]) -> Result<Self> {
+    fn new(weights: &[u8], capacity: usize) -> Result<Self> {
+        ensure!((1..=4).contains(&capacity), "invalid head batch capacity");
         ensure!(
             weights.len() == PADDED_LABELS * HIDDEN * 2,
             "invalid selected head buffer"
@@ -66,8 +69,9 @@ impl Head {
             stream,
             gemm,
             weights: DeviceBuffer::new(weights.len())?,
-            input: DeviceBuffer::new(HIDDEN * 2)?,
-            output: DeviceBuffer::new(PADDED_LABELS * 2)?,
+            input: DeviceBuffer::new(capacity * HIDDEN * 2)?,
+            output: DeviceBuffer::new(capacity * PADDED_LABELS * 2)?,
+            capacity,
         };
         // SAFETY: the device allocation has exactly weights.len() bytes.
         unsafe {
@@ -76,18 +80,27 @@ impl Head {
         Ok(head)
     }
     fn project(&self, hidden: &[f32], count: usize) -> Result<Vec<f32>> {
+        Ok(self.project_batch(&[hidden], &[count])?.pop().unwrap())
+    }
+    fn project_batch(&self, hidden: &[&[f32]], counts: &[usize]) -> Result<Vec<Vec<f32>>> {
+        let rows = hidden.len();
         ensure!(
-            hidden.len() == HIDDEN
-                && hidden.iter().all(|x| x.is_finite())
-                && (2..=LABELS).contains(&count),
-            "invalid hidden state/readout"
+            rows > 0
+                && rows <= self.capacity
+                && rows == counts.len()
+                && hidden
+                    .iter()
+                    .all(|h| h.len() == HIDDEN && h.iter().all(|x| x.is_finite()))
+                && counts.iter().all(|n| (2..=LABELS).contains(n)),
+            "invalid hidden batch/readout"
         );
         let input: Vec<u8> = hidden
             .iter()
+            .flat_map(|h| h.iter())
             .flat_map(|&x| bf16::from_f32(x).to_le_bytes())
             .collect();
-        // SAFETY: [1,HIDDEN] * [256,HIDDEN]^T -> [1,256]; allocated buffers match.
-        // The last weight row is zero padding and is never returned or normalized.
+        // SAFETY: [rows,HIDDEN] * [256,HIDDEN]^T -> [rows,256]; rows <= capacity.
+        // The final zero weight row is padding and never participates in normalization.
         unsafe {
             cuda::upload(self.input.at(0), &input, self.stream.0)?;
             cuda::check(
@@ -96,7 +109,7 @@ impl Head {
                     self.input.at(0),
                     self.weights.at(0),
                     self.output.at(0),
-                    1,
+                    rows as i32,
                     PADDED_LABELS as i32,
                     HIDDEN as i32,
                     PADDED_LABELS as i32,
@@ -105,36 +118,60 @@ impl Head {
                 "Decider label projection",
             )?;
         }
-        let mut output = vec![0; PADDED_LABELS * 2];
-        // SAFETY: download waits for GEMM and copies exactly the output allocation.
+        let mut output = vec![0; rows * PADDED_LABELS * 2];
+        // SAFETY: every returned row is inside the allocation; download synchronizes GEMM.
         unsafe {
             cuda::download(&mut output, self.output.at(0), self.stream.0)?;
         }
-        let logits: Vec<f32> = output.as_chunks::<2>().0[..count]
+        output
+            .as_chunks::<{ PADDED_LABELS * 2 }>()
+            .0
             .iter()
-            .map(|&b| bf16::from_le_bytes(b).to_f32())
-            .collect();
-        ensure!(
-            logits.iter().all(|x| x.is_finite()),
-            "nonfinite candidate logits"
-        );
-        Ok(logits)
+            .zip(counts)
+            .map(|(row, &count)| {
+                let logits: Vec<f32> = row.as_chunks::<2>().0[..count]
+                    .iter()
+                    .map(|&b| bf16::from_le_bytes(b).to_f32())
+                    .collect();
+                ensure!(
+                    logits.iter().all(|x| x.is_finite()),
+                    "nonfinite candidate logits"
+                );
+                Ok(logits)
+            })
+            .collect()
     }
 }
 struct Loaded {
     model: Model,
     head: Head,
+    batching: BatchLimits,
 }
 impl Loaded {
     fn execute(&mut self, rows: &[RowInput]) -> Result<Vec<Vec<f32>>> {
         cuda::set_device(0)?;
         let result = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
-            rows.iter()
-                .map(|row| {
-                    let hidden = self.model.forward(&row.ids)?;
-                    self.head.project(&hidden, row.candidate_ids.len())
-                })
-                .collect::<Result<Vec<_>>>()
+            let lengths: Vec<usize> = rows.iter().map(|row| row.ids.len()).collect();
+            let mut logits = Vec::with_capacity(rows.len());
+            for range in self.batching.ranges(&lengths) {
+                let batch = &rows[range];
+                if batch.len() == 1 {
+                    let hidden = self.model.forward(&batch[0].ids)?;
+                    logits.push(self.head.project(&hidden, batch[0].candidate_ids.len())?);
+                } else {
+                    let inputs: Vec<&[u32]> = batch.iter().map(|row| row.ids.as_slice()).collect();
+                    let hidden = self.model.forward_batch(&inputs)?;
+                    ensure!(
+                        hidden.len() == batch.len(),
+                        "backbone batch output count mismatch"
+                    );
+                    let hidden: Vec<&[f32]> = hidden.iter().map(Vec::as_slice).collect();
+                    let counts: Vec<usize> =
+                        batch.iter().map(|row| row.candidate_ids.len()).collect();
+                    logits.extend(self.head.project_batch(&hidden, &counts)?);
+                }
+            }
+            Ok(logits)
         }))
         .map_err(|_| anyhow::anyhow!("Decider execution panicked"))
         .and_then(|result| result);
@@ -158,6 +195,7 @@ impl Executor {
         library: &Path,
         checkpoint: Checkpoint,
         labels: Vec<u32>,
+        batching: BatchLimits,
     ) -> Result<Self> {
         ensure!(
             std::env::var("CUA_S1_GRAPH").as_deref() != Ok("1"),
@@ -166,8 +204,12 @@ impl Executor {
         let (dir, library) = (dir.to_owned(), library.to_owned());
         let loaded = tokio::task::spawn_blocking(move || -> Result<Loaded> {
             let model = Model::load(&dir, &library)?;
-            let head = Head::new(&checkpoint.head)?;
-            Ok(Loaded { model, head })
+            let head = Head::new(&checkpoint.head, batching.max_rows())?;
+            Ok(Loaded {
+                model,
+                head,
+                batching,
+            })
         })
         .await??;
         Ok(Self {

src/models/decider/native/src/lib.rs

@@ -31,6 +31,7 @@ pub const MODEL_ID: &str = "decider-2b-v11";
 pub const RUNTIME_REVISION: &str = "50d0be0d7cb43d2066965ce5fa7f3fe4e489a60f";
 pub const CHECKPOINT_REVISION: &str = "533964dae8be954c5b5e19fa4948e48408094c1e";
 
+pub mod batching;
 mod checkpoint;
 pub mod engine;
 pub mod executor;

tests/decider/batching.rs

@@ -0,0 +1,37 @@
+use omni_decider_native::batching::BatchLimits;
+#[test]
+fn bounded_contiguous_ranges_preserve_order_and_long_rows() {
+    let limits = BatchLimits::new(3, 8).unwrap();
+    assert_eq!(
+        limits.ranges(&[2, 3, 3, 1, 9, 2, 2, 2, 2]),
+        vec![0..3, 3..4, 4..5, 5..8, 8..9]
+    );
+    assert!(limits.ranges(&[]).is_empty());
+    assert_eq!(limits.ranges(&[9, 9]), vec![0..1, 1..2]);
+}
+#[test]
+fn row_limit_and_exact_token_boundary_are_independent() {
+    assert_eq!(
+        BatchLimits::new(2, 4096).unwrap().ranges(&[1, 1, 1]),
+        vec![0..2, 2..3]
+    );
+    assert_eq!(
+        BatchLimits::new(4, 4096).unwrap().ranges(&[2048, 2048, 1]),
+        vec![0..2, 2..3]
+    );
+    assert_eq!(
+        BatchLimits::default().ranges(&[1, 2, 3]),
+        vec![0..1, 1..2, 2..3]
+    );
+}
+#[test]
+fn startup_limits_reject_unbounded_or_malformed_settings() {
+    for (rows, tokens) in [(0, 1), (5, 1), (8, 1), (16, 1), (1, 0), (1, 4097)] {
+        assert!(BatchLimits::new(rows, tokens).is_err());
+    }
+    assert!(BatchLimits::from_values(Some("wrong"), None).is_err());
+    assert!(BatchLimits::from_values(None, Some("-1")).is_err());
+    let limits = BatchLimits::from_values(Some("4"), Some("4096")).unwrap();
+    assert_eq!(limits.max_rows(), 4);
+    assert_eq!(limits.max_tokens(), 4096);
+}