From cda46d559fef2eede1ca937c1bbb860aac86f1e3 Mon Sep 17 00:00:00 2001 From: Mathieu Poumeyrol Date: Mon, 29 Jun 2026 08:38:58 +0000 Subject: [PATCH] cuda: add a cuda-batched runtime that micro-batches concurrent run() calls Independent caller threads each fire tiny per-token predictor/joiner calls; with many concurrent streams that is one serialized CPU<->GPU round-trip per token per stream. A per-model background worker drains the in-flight queue, concatenates inputs along the symbolic batch axis, runs the inner CUDA runnable once, and scatters the row-slices back to each caller. WIP parked for later. --- Cargo.lock | 1 + cuda/Cargo.toml | 1 + cuda/examples/batch_smoke.rs | 91 ++++++++++ cuda/src/batch.rs | 311 +++++++++++++++++++++++++++++++++++ cuda/src/lib.rs | 2 + 5 files changed, 406 insertions(+) create mode 100644 cuda/examples/batch_smoke.rs create mode 100644 cuda/src/batch.rs diff --git a/Cargo.lock b/Cargo.lock index da0b29c07a..0d68f29133 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -4784,6 +4784,7 @@ dependencies = [ "rustfft", "tract-core", "tract-gpu", + "tract-nnef", "tract-pulse-opl", "tract-transformers", ] diff --git a/cuda/Cargo.toml b/cuda/Cargo.toml index 3a7a3c2607..818e46feae 100644 --- a/cuda/Cargo.toml +++ b/cuda/Cargo.toml @@ -41,6 +41,7 @@ proptest.workspace = true rand.workspace = true rustfft.workspace = true num-complex.workspace = true +tract-nnef.workspace = true [features] # Pick exactly one cuda-XXXXX to bind cudarc against. Higher minor versions diff --git a/cuda/examples/batch_smoke.rs b/cuda/examples/batch_smoke.rs new file mode 100644 index 0000000000..b7a4ce70e6 --- /dev/null +++ b/cuda/examples/batch_smoke.rs @@ -0,0 +1,91 @@ +//! Validates the `cuda-batched` runtime on a real transducer predictor. +//! +//! Usage: cargo run --release -p tract-cuda --example batch_smoke -- +//! +//! Checks (a) correctness — a batched call returns the same logits as plain +//! `cuda`; (b) that concurrent calls actually get batched (mean batch > 1); and +//! (c) throughput of `cuda-batched` vs `cuda` as concurrency rises. + +use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering}; +use std::sync::Arc; +use std::thread; +use std::time::{Duration, Instant}; + +use tract_core::internal::*; +use tract_core::runtime::runtime_for_name; + +const B: usize = 4; // beam width per call + +/// Build zero inputs matching the model's input facts, substituting the +/// symbolic (batch) axis with B. Works for predictor and joiner alike. +fn inputs_for(r: &dyn Runnable) -> TVec { + let mut out = tvec!(); + for ix in 0..r.input_count() { + let f = r.input_fact(ix).unwrap(); + let shape: Vec = f.shape.dims().iter().map(|d| d.as_i64().map(|v| v as usize).unwrap_or(B)).collect(); + out.push(Tensor::zero_dt(f.datum_type, &shape).unwrap().into()); + } + out +} + +fn throughput(runnable: &Arc>, n: usize, dur: Duration) -> f64 { + let total = Arc::new(AtomicUsize::new(0)); + let stop = Arc::new(AtomicBool::new(false)); + let start = Instant::now(); + let handles: Vec<_> = (0..n) + .map(|_| { + let r = Arc::clone(runnable); + let total = Arc::clone(&total); + let stop = Arc::clone(&stop); + thread::spawn(move || { + let mut st = r.spawn().unwrap(); + while !stop.load(Ordering::Relaxed) { + st.run(inputs_for(&**r)).unwrap(); + total.fetch_add(1, Ordering::Relaxed); + } + }) + }) + .collect(); + thread::sleep(dur); + stop.store(true, Ordering::Relaxed); + for h in handles { + h.join().unwrap(); + } + total.load(Ordering::Relaxed) as f64 / start.elapsed().as_secs_f64() +} + +fn main() -> TractResult<()> { + let path = std::env::args().nth(1).expect("usage: batch_smoke "); + let model = tract_nnef::nnef().model_for_path(&path)?.into_decluttered()?; + + let cuda = runtime_for_name("cuda")?.expect("cuda runtime not compiled in"); + let cuda_run: Arc> = Arc::new(cuda.prepare(model.clone())?); + + let batched = runtime_for_name("cuda-batched")?.expect("cuda-batched runtime not compiled in"); + let batched_run: Arc> = Arc::new(batched.prepare(model.clone())?); + + // (a) correctness + let reference = cuda_run.run(inputs_for(&**cuda_run))?; + let mut st = batched_run.spawn()?; + let got = st.run(inputs_for(&**batched_run))?; + for (i, (x, y)) in reference.iter().zip(got.iter()).enumerate() { + x.close_enough(y, true).map_err(|e| format_err!("output {i} mismatch: {e}"))?; + } + eprintln!( + "correctness OK: cuda-batched == cuda (output shapes {:?})", + reference.iter().map(|t| t.shape().to_vec()).collect::>() + ); + + // (b)+(c) throughput vs concurrency + let dur = Duration::from_secs(3); + println!("# threads | cuda calls/s | cuda-batched calls/s | speedup | mean_batch"); + for &n in &[1usize, 2, 4, 8, 16, 32, 64] { + let c = throughput(&cuda_run, n, dur); + let (b0, j0) = tract_cuda::batch_stats(); + let cb = throughput(&batched_run, n, dur); + let (b1, j1) = tract_cuda::batch_stats(); + let mean_batch = (j1 - j0) as f64 / (b1 - b0).max(1) as f64; + println!("{n:>9} | {c:>12.0} | {cb:>20.0} | {:>6.2}x | {mean_batch:>10.2}", cb / c.max(1.0)); + } + Ok(()) +} diff --git a/cuda/src/batch.rs b/cuda/src/batch.rs new file mode 100644 index 0000000000..7cd91294e7 --- /dev/null +++ b/cuda/src/batch.rs @@ -0,0 +1,311 @@ +//! Dynamic micro-batching runtime that wraps the CUDA runtime. +//! +//! Registered as `cuda-batched`. It is transparent and hot-swappable by runtime +//! name: a model prepared through it behaves exactly like one prepared on `cuda`, +//! except that concurrent `run()` calls from independent caller threads are +//! gathered and executed as ONE batched device call. +//! +//! Motivation: the transducer decode loop fires one tiny predictor/joiner call +//! per token, per stream. With N concurrent streams that is N tiny serialized +//! CPU↔GPU round-trips per step (measured: ~360k ctx-switches/s, GPU busy with +//! launch overhead rather than work). Each model here already carries a symbolic +//! batch axis (the beam), so the fix is to grow that axis across streams: one +//! background worker per prepared model drains the queue of in-flight calls, +//! concatenates their inputs along the batch axis, runs the inner CUDA runnable +//! once, and scatters the row-slices back to each caller. +//! +//! Contract for a batchable model: every input and output must carry the batch +//! dimension on exactly one axis (a free symbol), and all other dims must be +//! identical across concurrent callers. Inputs/outputs with no symbolic axis are +//! treated as shared (broadcast once / replicated back). This holds for the +//! predictor and (after its `encoder_outputs` is reshaped to `[B,1024]`) the +//! joiner; it does NOT hold for the encoder (variable sequence length) — that +//! needs padding and is out of scope here. +//! +//! Knobs (env): `TRACT_BATCH_LINGER_US` (default 0) waits this long after the +//! first queued call to let more accumulate — the throughput↔latency dial; +//! `TRACT_BATCH_MAX` (default 256) caps batch size. + +use std::sync::atomic::{AtomicU64, Ordering}; +use std::sync::mpsc::{channel, Receiver, Sender}; +use std::sync::Mutex; +use std::thread; +use std::time::Duration; + +use tract_core::internal::*; + +use crate::CudaRuntime; + +/// Observability: number of batched device calls and total jobs served, since +/// process start. `mean batch = jobs / batches`. +pub static BATCHES: AtomicU64 = AtomicU64::new(0); +pub static JOBS: AtomicU64 = AtomicU64::new(0); + +/// (batched device calls, total jobs served) — for `mean batch = jobs/batches`. +pub fn batch_stats() -> (u64, u64) { + (BATCHES.load(Ordering::Relaxed), JOBS.load(Ordering::Relaxed)) +} + +/// Reply channel: the worker sends each caller its row-slices (owned tensors). +type Reply = Sender>>; + +struct Job { + inputs: TVec, + reply: Reply, +} + +#[derive(Debug)] +struct BatchingCudaRuntime; + +impl Runtime for BatchingCudaRuntime { + fn name(&self) -> StaticName { + "cuda-batched".into() + } + + fn check(&self) -> TractResult<()> { + CudaRuntime.check() + } + + fn prepare_with_options( + &self, + model: TypedModel, + options: &RunOptions, + ) -> TractResult> { + let inner = CudaRuntime.prepare_with_options(model, options)?; + BatchingRunnable::wrap(inner) + } +} + +register_runtime!(BatchingCudaRuntime = BatchingCudaRuntime); + +struct Shared { + tx: Mutex>, + model: Option>, + plan: Option>, +} + +#[derive(Clone)] +struct BatchingRunnable { + shared: Arc, +} + +impl std::fmt::Debug for BatchingRunnable { + fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result { + write!(f, "BatchingRunnable(cuda-batched)") + } +} + +/// The batch axis of a fact: its single symbolic axis, or `None` for a +/// shared/broadcast tensor (no symbolic axis). +fn detect_axis(fact: &TypedFact) -> TractResult> { + let symbolic: Vec = fact + .shape + .dims() + .iter() + .enumerate() + .filter(|(_, d)| d.as_i64().is_none()) + .map(|(i, _)| i) + .collect(); + match symbolic.len() { + 0 => Ok(None), + 1 => Ok(Some(symbolic[0])), + _ => bail!("cuda-batched: model has {} symbolic axes in fact {:?}; need exactly one batch axis", symbolic.len(), fact), + } +} + +impl BatchingRunnable { + fn wrap(inner: Box) -> TractResult> { + let model = inner.typed_model().cloned(); + let plan = inner.typed_plan().cloned(); + + let mut batch_in = Vec::with_capacity(inner.input_count()); + for ix in 0..inner.input_count() { + batch_in.push(detect_axis(inner.input_fact(ix)?)?); + } + let mut batch_out = Vec::with_capacity(inner.output_count()); + for ix in 0..inner.output_count() { + batch_out.push(detect_axis(inner.output_fact(ix)?)?); + } + + let linger_us: u64 = std::env::var("TRACT_BATCH_LINGER_US").ok().and_then(|s| s.parse().ok()).unwrap_or(0); + let max_batch: usize = std::env::var("TRACT_BATCH_MAX").ok().and_then(|s| s.parse().ok()).unwrap_or(256); + + let (tx, rx) = channel::(); + thread::Builder::new() + .name("tract-batcher".into()) + .spawn(move || worker(inner, rx, batch_in, batch_out, linger_us, max_batch))?; + + Ok(Box::new(BatchingRunnable { + shared: Arc::new(Shared { tx: Mutex::new(tx), model, plan }), + })) + } +} + +impl Runnable for BatchingRunnable { + fn spawn(&self) -> TractResult> { + let tx = self.shared.tx.lock().map_err(|_| format_err!("cuda-batched sender poisoned"))?.clone(); + Ok(Box::new(BatchingState { tx, runnable: self.clone() })) + } + + fn typed_plan(&self) -> Option<&Arc> { + self.shared.plan.as_ref() + } + + fn typed_model(&self) -> Option<&Arc> { + self.shared.model.as_ref() + } +} + +struct BatchingState { + tx: Sender, + runnable: BatchingRunnable, +} + +impl std::fmt::Debug for BatchingState { + fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result { + write!(f, "BatchingState(cuda-batched)") + } +} + +impl State for BatchingState { + fn run(&mut self, inputs: TVec) -> TractResult> { + let inputs: TVec = inputs.into_iter().map(|v| v.into_tensor()).collect(); + let (rtx, rrx) = channel(); + self.tx.send(Job { inputs, reply: rtx }).map_err(|_| format_err!("cuda-batched worker is gone"))?; + let out = rrx.recv().map_err(|_| format_err!("cuda-batched worker dropped the reply"))??; + Ok(out.into_iter().map(|t| t.into()).collect()) + } + + fn runnable(&self) -> &dyn Runnable { + &self.runnable + } + + fn freeze(&self) -> Box { + Box::new(FrozenBatchingState { tx: self.tx.clone(), runnable: self.runnable.clone() }) + } +} + +#[derive(Clone)] +struct FrozenBatchingState { + tx: Sender, + runnable: BatchingRunnable, +} + +impl std::fmt::Debug for FrozenBatchingState { + fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result { + write!(f, "FrozenBatchingState(cuda-batched)") + } +} + +impl FrozenState for FrozenBatchingState { + fn unfreeze(&self) -> Box { + Box::new(BatchingState { tx: self.tx.clone(), runnable: self.runnable.clone() }) + } + fn input_count(&self) -> usize { + self.runnable.input_count() + } + fn output_count(&self) -> usize { + self.runnable.output_count() + } +} + +fn worker( + inner: Box, + rx: Receiver, + batch_in: Vec>, + batch_out: Vec>, + linger_us: u64, + max_batch: usize, +) { + loop { + let first = match rx.recv() { + Ok(j) => j, + Err(_) => return, // all senders dropped: the runnable was dropped + }; + let mut jobs = vec![first]; + if linger_us > 0 { + thread::sleep(Duration::from_micros(linger_us)); + } + while jobs.len() < max_batch { + match rx.try_recv() { + Ok(j) => jobs.push(j), + Err(_) => break, + } + } + BATCHES.fetch_add(1, Ordering::Relaxed); + JOBS.fetch_add(jobs.len() as u64, Ordering::Relaxed); + match run_batch(&*inner, &jobs, &batch_in, &batch_out) { + Ok(per_job) => { + for (job, out) in jobs.into_iter().zip(per_job.into_iter()) { + let _ = job.reply.send(Ok(out)); + } + } + Err(e) => { + let msg = format!("{e:#}"); + for job in jobs { + let _ = job.reply.send(Err(format_err!("cuda-batched run failed: {msg}"))); + } + } + } + } +} + +fn run_batch( + inner: &dyn Runnable, + jobs: &[Job], + batch_in: &[Option], + batch_out: &[Option], +) -> TractResult>> { + let n_in = jobs[0].inputs.len(); + + // Per-job row count, read off the first input that carries a batch axis. + let probe = batch_in.iter().position(|a| a.is_some()); + let Some(probe) = probe else { + // No batchable axis anywhere: run each job independently (no win, but correct). + let mut out = Vec::with_capacity(jobs.len()); + for job in jobs { + let inputs: TVec = job.inputs.iter().cloned().map(|t| t.into()).collect(); + out.push(inner.run(inputs)?.into_iter().map(|v| v.into_tensor()).collect()); + } + return Ok(out); + }; + let probe_axis = batch_in[probe].unwrap(); + let rows: Vec = jobs.iter().map(|j| j.inputs[probe].shape()[probe_axis]).collect(); + + // Concatenate inputs along their batch axis (shared inputs taken from job 0). + let mut batched: TVec = tvec!(); + for i in 0..n_in { + match batch_in[i] { + Some(ax) => { + let refs: Vec<&Tensor> = jobs.iter().map(|j| &j.inputs[i]).collect(); + batched.push(Tensor::stack_tensors(ax, &refs)?.into()); + } + None => batched.push(jobs[0].inputs[i].clone().into()), + } + } + + let out: TVec = inner.run(batched)?; + + // Scatter each output's row-slices back to the originating jobs. + let mut result: Vec> = (0..jobs.len()).map(|_| tvec!()).collect(); + for (o, out_value) in out.into_iter().enumerate() { + match batch_out.get(o).copied().flatten() { + Some(ax) => { + let t = out_value.into_tensor(); + let mut off = 0; + for (ji, &r) in rows.iter().enumerate() { + result[ji].push(t.slice(ax, off, off + r)?); + off += r; + } + ensure!(off == t.shape()[ax], "cuda-batched: output rows {} != batch {}", off, t.shape()[ax]); + } + None => { + let t = out_value.into_tensor(); + for slot in result.iter_mut() { + slot.push(t.clone()); + } + } + } + } + Ok(result) +} diff --git a/cuda/src/lib.rs b/cuda/src/lib.rs index 24ebbb06b7..2984ae0c64 100644 --- a/cuda/src/lib.rs +++ b/cuda/src/lib.rs @@ -1,3 +1,5 @@ +mod batch; +pub use batch::batch_stats; mod context; pub mod kernels; pub mod ops;