From d249c5e244de4e42544d6fa3837aef40700cd7dc Mon Sep 17 00:00:00 2001 From: levius <2114377220@qq.com> Date: Sun, 4 Oct 2026 12:49:32 +0800 Subject: [PATCH 1/2] feat(laya): cache encoder graphs within shape and memory limits --- src/models/laya/src/cache.rs | 118 +++++++++++++++++++++++++++++++ src/models/laya/src/lib.rs | 1 + src/models/laya/src/workspace.rs | 75 +++++++++++++------- 3 files changed, 169 insertions(+), 25 deletions(-) create mode 100644 src/models/laya/src/cache.rs diff --git a/src/models/laya/src/cache.rs b/src/models/laya/src/cache.rs new file mode 100644 index 00000000..3704a7db --- /dev/null +++ b/src/models/laya/src/cache.rs @@ -0,0 +1,118 @@ +//! Opt-in, bounded LRU storage for encoder graphs of supported shapes. +use crate::{ + encoder::{Encoder, validate_inputs}, + graph::ShapeEncoder, + workspace::Workspace, +}; +use anyhow::Result; +use std::collections::VecDeque; + +/// Limits retained graph workspaces. Either zero limit disables graph caching. +#[derive(Clone, Copy, Debug)] +pub struct CacheConfig { + pub max_shapes: usize, + /// GPU workspace bytes only: excludes weights, graph/driver overhead, host + /// staging, transient eager workspaces and buffers cloned by callers. + pub max_bytes: usize, +} + +impl Default for CacheConfig { + fn default() -> Self { + Self { + max_shapes: 4, + max_bytes: 512 * 1024 * 1024, + } + } +} + +/// Each cache borrows one encoder; shapes cannot reuse another encoder's graph. +/// Returned workspaces remain valid until the next mutable access to the cache. +pub struct EncoderCache<'a> { + encoder: &'a Encoder, + config: CacheConfig, + entries: VecDeque>, + cached_bytes: usize, + eager: Option, +} + +impl<'a> EncoderCache<'a> { + pub fn new(encoder: &'a Encoder, config: CacheConfig) -> Self { + Self { + encoder, + config, + entries: VecDeque::new(), + cached_bytes: 0, + eager: None, + } + } + + /// Validates before changing cache state, then uploads the current inputs. + /// Misses evict least recently used entries before allocation. Disabled or + /// oversized shapes reuse one eager workspace while the shape stays the same. + /// CUDA failures are always returned; failed requests do not promise outputs. + pub fn run( + &mut self, + batch: usize, + sequence: usize, + ids: &[i64], + lengths: &[i32], + types: &[i64], + ) -> Result<&Workspace> { + let bytes = Workspace::required_bytes(batch, sequence)?; + validate_inputs(ids, lengths, types, batch, sequence)?; + if self.config.max_shapes == 0 || bytes > self.config.max_bytes { + if !self.eager.as_ref().is_some_and(|workspace| { + workspace.batch() == batch && workspace.sequence() == sequence + }) { + self.eager = None; + self.eager = Some(Workspace::new(&self.encoder.cuda, batch, sequence)?); + } + let workspace = self.eager.as_ref().unwrap(); + self.encoder.run(ids, lengths, types, workspace)?; + return Ok(workspace); + } + self.eager = None; + if let Some(index) = self.entries.iter().position(|entry| { + entry.workspace().batch() == batch && entry.workspace().sequence() == sequence + }) { + self.entries[index].run(ids, lengths, types)?; + let entry = self.entries.remove(index).unwrap(); + self.entries.push_back(entry); + } else { + while self.entries.len() >= self.config.max_shapes + || self.cached_bytes > self.config.max_bytes - bytes + { + let evicted = self.entries.pop_front().unwrap(); + self.cached_bytes -= evicted.workspace().bytes(); + drop(evicted); + } + let workspace = Workspace::new(&self.encoder.cuda, batch, sequence)?; + let entry = ShapeEncoder::capture(self.encoder, workspace, ids, lengths, types)?; + entry.run(ids, lengths, types)?; + self.entries.push_back(entry); + self.cached_bytes += bytes; + } + Ok(self.entries.back().unwrap().workspace()) + } + + /// Releases all cached graphs and the reusable eager workspace. Caller-held + /// buffer clones remain alive; the encoder and cache limits are unchanged. + pub fn clear(&mut self) { + self.entries.clear(); + self.cached_bytes = 0; + self.eager = None; + } + + pub fn len(&self) -> usize { + self.entries.len() + } + + pub fn is_empty(&self) -> bool { + self.entries.is_empty() + } + + /// Total bytes owned by retained graph workspaces, subject to `max_bytes`. + pub fn cached_bytes(&self) -> usize { + self.cached_bytes + } +} diff --git a/src/models/laya/src/lib.rs b/src/models/laya/src/lib.rs index 4bc77cef..d8096215 100644 --- a/src/models/laya/src/lib.rs +++ b/src/models/laya/src/lib.rs @@ -1,4 +1,5 @@ mod artifacts; +pub mod cache; pub mod config; pub mod encoder; pub mod graph; diff --git a/src/models/laya/src/workspace.rs b/src/models/laya/src/workspace.rs index 1f86c4cf..974a3bb9 100644 --- a/src/models/laya/src/workspace.rs +++ b/src/models/laya/src/workspace.rs @@ -37,7 +37,7 @@ pub struct WorkspaceBuffers { } impl Workspace { - pub fn new(cuda: &Cuda, batch: usize, sequence: usize) -> Result { + fn buffer_sizes(batch: usize, sequence: usize) -> Result<[usize; 17]> { ensure!( batch.is_power_of_two() && batch <= 16, "workspace batch must be 1, 2, 4, 8 or 16" @@ -47,37 +47,62 @@ impl Workspace { "workspace sequence must be a multiple of 16 in 16..=512" ); let tokens = batch * sequence; - let mut bytes = 0; - let mut alloc = |size| { - let buffer = cuda.alloc(size)?; - bytes += size; - Ok::<_, anyhow::Error>(buffer) - }; + Ok([ + tokens * 8, + batch * 4, + batch * 8, + tokens * D * 4, + tokens * D * 2, + tokens * D * 6, + tokens * D * 2, + tokens * 2624 * 2, + tokens * 4096 * 2, + MAX_MARKERS * 4, + (batch + 1) * 4, + MAX_MARKERS * D * 2, + MAX_MARKERS * D * 2, + MAX_MARKERS * 2, + batch * 1028 * 2, + batch * 256 * 2, + batch * 2 * 2, + ]) + } + + /// Validates the shape and returns its GPU scratch size before allocating. + pub fn required_bytes(batch: usize, sequence: usize) -> Result { + Ok(Self::buffer_sizes(batch, sequence)?.iter().sum()) + } + + pub fn new(cuda: &Cuda, batch: usize, sequence: usize) -> Result { + let sizes = Self::buffer_sizes(batch, sequence)?; + let bytes = sizes.iter().sum(); + let mut sizes = sizes.into_iter(); + let mut alloc = || cuda.alloc(sizes.next().expect("workspace buffer layout")); let buffers = WorkspaceBuffers { - ids: alloc(tokens * 8)?, - lengths: alloc(batch * 4)?, - types: alloc(batch * 8)?, - residual: alloc(tokens * D * 4)?, - hidden: alloc(tokens * D * 2)?, - qkv: alloc(tokens * D * 6)?, - attention: alloc(tokens * D * 2)?, - gated: alloc(tokens * 2624 * 2)?, - feed_forward: alloc(tokens * 4096 * 2)?, - indices: alloc(MAX_MARKERS * 4)?, - offsets: alloc((batch + 1) * 4)?, - markers: alloc(MAX_MARKERS * D * 2)?, - scored: alloc(MAX_MARKERS * D * 2)?, - logits: alloc(MAX_MARKERS * 2)?, - features: alloc(batch * 1028 * 2)?, - action_hidden: alloc(batch * 256 * 2)?, - actions: alloc(batch * 2 * 2)?, + ids: alloc()?, + lengths: alloc()?, + types: alloc()?, + residual: alloc()?, + hidden: alloc()?, + qkv: alloc()?, + attention: alloc()?, + gated: alloc()?, + feed_forward: alloc()?, + indices: alloc()?, + offsets: alloc()?, + markers: alloc()?, + scored: alloc()?, + logits: alloc()?, + features: alloc()?, + action_hidden: alloc()?, + actions: alloc()?, }; Ok(Self { batch, sequence, bytes, buffers, - staging: RefCell::new(vec![0; tokens * 8 + batch * 12]), + staging: RefCell::new(vec![0; batch * sequence * 8 + batch * 12]), }) } From ae8e34fcd292d4c8f21c35d4d143033a8f56a796 Mon Sep 17 00:00:00 2001 From: levius <2114377220@qq.com> Date: Mon, 5 Oct 2026 23:59:37 +0800 Subject: [PATCH 2/2] ci: validate retargeted Laya host changes against main