diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml
index 594ef0f0..1b285855 100644
--- a/.github/workflows/ci.yml
+++ b/.github/workflows/ci.yml
@@ -23,13 +23,13 @@ jobs:
- name: Format
run: cargo fmt --all --check
- name: Clippy
- run: cargo clippy --workspace --locked --all-targets -- -D warnings
+ run: cargo clippy --workspace --locked --all-targets --features omni-laya/serve -- -D warnings
- name: Frontend API tests
run: cargo test -p omni-jev --test frontend --locked
- name: Test
- run: cargo test --workspace --locked
+ run: cargo test --workspace --locked --features omni-laya/serve
- name: Build
- run: cargo build --workspace --release --locked
+ run: cargo build --workspace --release --locked --features omni-laya/serve
benchmarks:
runs-on: ubuntu-latest
diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md
index 1aa9639a..ea0321f5 100644
--- a/CONTRIBUTING.md
+++ b/CONTRIBUTING.md
@@ -76,9 +76,9 @@ The Rust checks used by [CI](.github/workflows/ci.yml) are:
```sh
cargo fmt --all --check
-cargo clippy --workspace --locked --all-targets -- -D warnings
-cargo test --workspace --locked
-cargo build --workspace --release --locked
+cargo clippy --workspace --locked --all-targets --features omni-laya/serve -- -D warnings
+cargo test --workspace --locked --features omni-laya/serve
+cargo build --workspace --release --locked --features omni-laya/serve
```
For an agent-assisted self-review, use the repository's
@@ -96,29 +96,6 @@ remain responsible for understanding the changes and verifying the results.
Self-review helps maintainers focus on design and correctness; it does not
replace maintainer review.
-### Large code changes
-
-PRs with **more than 3,000 changed lines of authored code** need extra contributor
-attention before requesting review. Count additions plus deletions against the
-PR's merge base in source files, tests, and build or validation scripts. Report
-this count separately from the total diff size; exclude documentation, generated
-output, lockfiles, and static fixtures from the code count, while still reviewing
-those files for relevance and correctness.
-
-- Complete a full self-review of every affected component and its integration
- boundaries. A quick precheck alone is insufficient; keep the PR in draft until
- the contributor self-review is complete.
-- Consider splitting independent features, refactors, and cleanup into focused
- PRs. If the change needs to stay together, explain why in the PR description
- and provide a component map and suggested review order.
-- Include the code-line count and a validation summary for each affected area in
- the PR description: commands, results, and unverified behavior with reasons.
- Cover changed interfaces between components as well as individual components.
-
-Size signals the need for closer review; it is not itself a correctness finding.
-Choose checks based on the changed behavior and risk. Crossing this threshold
-alone does not require GPU benchmarks or other expensive experiments.
-
## Documentation site
The site at is built with MkDocs
diff --git a/Cargo.lock b/Cargo.lock
index c43ff226..71d54f9c 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -1165,6 +1165,7 @@ version = "0.1.0"
dependencies = [
"anyhow",
"libloading",
+ "tempfile",
]
[[package]]
diff --git a/src/backends/cuda/Cargo.toml b/src/backends/cuda/Cargo.toml
index 58885bb2..abc68858 100644
--- a/src/backends/cuda/Cargo.toml
+++ b/src/backends/cuda/Cargo.toml
@@ -7,3 +7,10 @@ publish = false
[dependencies]
anyhow = "1"
libloading = "0.8"
+
+[dev-dependencies]
+tempfile = "3"
+
+[[test]]
+name = "laya_runtime"
+path = "../../../tests/cuda/laya_runtime.rs"
diff --git a/src/backends/cuda/kernels/runtime.cu b/src/backends/cuda/kernels/runtime.cu
index 185de53d..91723c5b 100644
--- a/src/backends/cuda/kernels/runtime.cu
+++ b/src/backends/cuda/kernels/runtime.cu
@@ -62,14 +62,26 @@ int laya_capture_begin(void* stream) {
cudaStreamCaptureModeThreadLocal);
}
+// End capture and release temporary/partial graph ownership on every outcome.
int laya_capture_end(void* stream, void** executable) {
+ if (!executable) return -1;
+ *executable = nullptr;
+ if (!stream) return -1;
cudaGraph_t graph = nullptr;
- auto e = cudaStreamEndCapture(static_cast(stream), &graph);
- if (e != cudaSuccess)
- return e;
- e = cudaGraphInstantiate(reinterpret_cast(executable), graph, 0);
- cudaGraphDestroy(graph);
- return e;
+ cudaGraphExec_t created = nullptr;
+ auto status = cudaStreamEndCapture(static_cast(stream), &graph);
+ if (status == cudaSuccess && graph)
+ status = cudaGraphInstantiate(&created, graph, 0);
+ if (graph) {
+ auto destroyed = cudaGraphDestroy(graph);
+ if (status == cudaSuccess) status = destroyed;
+ }
+ if (status != cudaSuccess || !created) {
+ if (created) cudaGraphExecDestroy(created);
+ return status != cudaSuccess ? status : -1;
+ }
+ *executable = created;
+ return 0;
}
int laya_graph_run(void* executable, void* stream) {
diff --git a/src/backends/cuda/src/lib.rs b/src/backends/cuda/src/lib.rs
index 5343723c..a72f008e 100644
--- a/src/backends/cuda/src/lib.rs
+++ b/src/backends/cuda/src/lib.rs
@@ -2,17 +2,26 @@
use anyhow::{Result, anyhow, ensure};
use libloading::Library;
use std::{
+ cell::Cell,
ffi::{CStr, c_void},
path::Path,
rc::Rc,
};
pub type Ptr = *mut c_void;
-type Kernel = unsafe extern "C" fn(*mut Ptr, i32, i32, i32, Ptr) -> i32;
+type Launch = unsafe extern "C" fn(*mut Ptr, i32, i32, i32, Ptr) -> i32;
struct Context {
lib: Library,
stream: Ptr,
+ capturing: Cell,
}
impl Context {
+ fn ensure_not_capturing(&self) -> Result<()> {
+ ensure!(
+ !self.capturing.get(),
+ "operation prohibited during CUDA capture"
+ );
+ Ok(())
+ }
fn check(&self, code: i32) -> Result<()> {
if code == 0 {
return Ok(());
@@ -31,6 +40,7 @@ impl Context {
unsafe { Ok(*self.lib.get::(name)?) }
}
fn sync(&self) -> Result<()> {
+ self.ensure_not_capturing()?;
let f = self.symbol:: i32>(b"laya_sync\0")?;
self.check(unsafe { f(self.stream) })
}
@@ -61,11 +71,16 @@ impl Cuda {
let mut stream = std::ptr::null_mut();
let init = unsafe { lib.get:: i32>(b"laya_init\0") }?;
let code = unsafe { init(&mut stream) };
- let ctx = Rc::new(Context { lib, stream });
+ let ctx = Rc::new(Context {
+ lib,
+ stream,
+ capturing: Cell::new(false),
+ });
ctx.check(code)?;
Ok(Self { ctx })
}
pub fn alloc(&self, bytes: usize) -> Result {
+ self.ctx.ensure_not_capturing()?;
ensure!(bytes > 0, "zero CUDA allocation");
let f = self
.ctx
@@ -83,9 +98,59 @@ impl Cuda {
b.write(bytes)?;
Ok(b)
}
+ /// Complete a prevalidated group with one synchronization before host borrows end.
+ /// A failed submission is also drained; this does not overlap copies and compute.
+ pub fn write_many(&self, writes: &[(&Buffer, &[u8])]) -> Result<()> {
+ self.ctx.ensure_not_capturing()?;
+ for (buffer, bytes) in writes {
+ ensure!(
+ Rc::ptr_eq(&self.ctx, &buffer.ctx),
+ "upload buffer belongs to another CUDA context"
+ );
+ ensure!(bytes.len() <= buffer.bytes, "upload exceeds allocation");
+ }
+ if writes.iter().all(|(_, bytes)| bytes.is_empty()) {
+ return Ok(());
+ }
+ let upload = self
+ .ctx
+ .symbol:: i32>(b"laya_upload\0")?;
+ let sync = self
+ .ctx
+ .symbol:: i32>(b"laya_sync\0")?;
+ let mut copied = 0;
+ for (buffer, bytes) in writes {
+ if bytes.is_empty() {
+ continue;
+ }
+ copied = unsafe { upload(buffer.p, bytes.as_ptr(), bytes.len(), self.ctx.stream) };
+ if copied != 0 {
+ break;
+ }
+ }
+ // Resolve both entry points before submission and keep every source borrowed
+ // through the drain even when a copy reports an error after queuing work.
+ let synced = unsafe { sync(self.ctx.stream) };
+ match (self.ctx.check(copied), self.ctx.check(synced)) {
+ (Err(copy), Err(sync)) => Err(anyhow!(
+ "{copy}; stream synchronization also failed: {sync}"
+ )),
+ (Err(e), _) | (_, Err(e)) => Err(e),
+ (Ok(()), Ok(())) => Ok(()),
+ }
+ }
pub fn sync(&self) -> Result<()> {
self.ctx.sync()
}
+ /// Resolve a kernel once while retaining the owning runtime and native code.
+ pub fn resolve(&self, name: &str) -> Result {
+ Ok(Kernel {
+ ctx: self.ctx.clone(),
+ launch: self
+ .ctx
+ .symbol::(format!("laya_{name}\0").as_bytes())?,
+ })
+ }
/// # Safety
/// Tensor shape, dtype, layout, aliasing and allocation sizes must match the generated kernel.
/// Buffers must belong to this context and stay alive until synchronization or graph destruction.
@@ -96,7 +161,7 @@ impl Cuda {
);
let k = self
.ctx
- .symbol::(format!("laya_{name}\0").as_bytes())?;
+ .symbol::(format!("laya_{name}\0").as_bytes())?;
self.ctx.check(unsafe {
k(
args.as_ptr() as *mut Ptr,
@@ -108,36 +173,31 @@ impl Cuda {
})
}
/// # Safety
- /// Every allocation referenced by `work` must outlive the returned graph. No allocation or copy
- /// that may synchronize is permitted inside work. This context is confined to one OS thread.
+ /// Every allocation and kernel referenced by `work` must outlive the graph.
+ /// Work must only launch on this stream and cannot release resources or invoke
+ /// capture-incompatible CUDA operations, including through raw FFI symbols.
+ /// This context is confined to one OS thread. Capture ends on errors or unwind.
pub unsafe fn capture(&self, work: impl FnOnce() -> Result<()>) -> Result {
- self.sync()?;
+ self.ctx.ensure_not_capturing()?;
+ // Resolve the complete optional API before beginning, including cleanup.
+ let functions = GraphFunctions {
+ end: self.ctx.symbol(b"laya_capture_end\0")?,
+ run: self.ctx.symbol(b"laya_graph_run\0")?,
+ free: self.ctx.symbol(b"laya_graph_free\0")?,
+ };
let begin = self
.ctx
.symbol:: i32>(b"laya_capture_begin\0")?;
- let end = self
- .ctx
- .symbol:: i32>(b"laya_capture_end\0")?;
+ self.sync()?;
self.ctx.check(unsafe { begin(self.ctx.stream) })?;
- let result = work();
- let mut p = std::ptr::null_mut();
- let code = unsafe { end(self.ctx.stream, &mut p) };
- if result.is_err() || code != 0 {
- if !p.is_null() {
- let f = self
- .ctx
- .symbol:: i32>(b"laya_graph_free\0")?;
- unsafe {
- f(p);
- }
- }
- result?;
- self.ctx.check(code)?;
- }
- Ok(Graph {
- ctx: self.ctx.clone(),
- p,
- })
+ self.ctx.capturing.set(true);
+ let guard = CaptureGuard {
+ cuda: self,
+ functions,
+ active: true,
+ };
+ work()?;
+ guard.finish()
}
/// # Safety
/// Caller supplies the precise dimensions and allocation sizes required by this glue kernel.
@@ -151,7 +211,7 @@ impl Cuda {
) -> Result<()> {
let k = self
.ctx
- .symbol::(format!("laya_{name}\0").as_bytes())?;
+ .symbol::(format!("laya_{name}\0").as_bytes())?;
self.ctx.check(unsafe {
k(
args.as_ptr() as *mut Ptr,
@@ -187,15 +247,13 @@ impl Buffer {
self.p
}
pub fn write(&self, bytes: &[u8]) -> Result<()> {
- ensure!(bytes.len() <= self.bytes, "upload exceeds allocation");
- let f = self
- .ctx
- .symbol:: i32>(b"laya_upload\0")?;
- self.ctx
- .check(unsafe { f(self.p, bytes.as_ptr(), bytes.len(), self.ctx.stream) })?;
- self.ctx.sync()
+ Cuda {
+ ctx: self.ctx.clone(),
+ }
+ .write_many(&[(self, bytes)])
}
pub fn read(&self, bytes: usize) -> Result> {
+ self.ctx.ensure_not_capturing()?;
ensure!(bytes <= self.bytes, "download exceeds allocation");
let mut data = vec![0; bytes];
let f = self
@@ -220,28 +278,93 @@ impl Drop for Buffer {
}
}
}
+#[derive(Clone, Copy)]
+struct GraphFunctions {
+ end: unsafe extern "C" fn(Ptr, *mut Ptr) -> i32,
+ run: unsafe extern "C" fn(Ptr, Ptr) -> i32,
+ free: unsafe extern "C" fn(Ptr) -> i32,
+}
pub struct Graph {
ctx: Rc,
+ functions: GraphFunctions,
p: Ptr,
}
impl Graph {
pub fn replay(&self) -> Result<()> {
- let f = self
- .ctx
- .symbol:: i32>(b"laya_graph_run\0")?;
- self.ctx.check(unsafe { f(self.p, self.ctx.stream) })
+ self.ctx.ensure_not_capturing()?;
+ self.ctx
+ .check(unsafe { (self.functions.run)(self.p, self.ctx.stream) })
}
}
impl Drop for Graph {
fn drop(&mut self) {
let _ = self.ctx.sync();
- if let Ok(f) = self
- .ctx
- .symbol:: i32>(b"laya_graph_free\0")
- {
+ unsafe {
+ (self.functions.free)(self.p);
+ }
+ }
+}
+struct CaptureGuard<'a> {
+ cuda: &'a Cuda,
+ functions: GraphFunctions,
+ active: bool,
+}
+impl CaptureGuard<'_> {
+ fn finish(mut self) -> Result {
+ let mut p = std::ptr::null_mut();
+ let code = unsafe { (self.functions.end)(self.cuda.ctx.stream, &mut p) };
+ self.active = false;
+ self.cuda.ctx.capturing.set(false);
+ // Own partial handles before status validation, so failure also frees them.
+ let graph = (!p.is_null()).then(|| Graph {
+ ctx: self.cuda.ctx.clone(),
+ functions: self.functions,
+ p,
+ });
+ self.cuda.ctx.check(code)?;
+ ensure!(graph.is_some(), "CUDA runtime returned a null graph");
+ Ok(graph.unwrap())
+ }
+}
+impl Drop for CaptureGuard<'_> {
+ fn drop(&mut self) {
+ if self.active {
+ let mut p = std::ptr::null_mut();
unsafe {
- f(self.p);
+ (self.functions.end)(self.cuda.ctx.stream, &mut p);
+ if !p.is_null() {
+ (self.functions.free)(p);
+ }
}
+ self.cuda.ctx.capturing.set(false);
}
}
}
+
+/// A resolved entry point retaining its stream and native library.
+#[derive(Clone)]
+pub struct Kernel {
+ ctx: Rc,
+ launch: Launch,
+}
+impl Kernel {
+ /// # Safety
+ /// Shapes, dtype, layout, aliasing and pointer lifetimes must match this kernel.
+ /// Every pointer belongs to this context and stays alive through synchronization
+ /// or destruction of any graph that captures the launch.
+ pub unsafe fn launch(&self, args: &[Ptr], b: usize, l: usize) -> Result<()> {
+ ensure!(
+ b > 0 && b <= 16 && l > 0 && l <= 512 && l.is_multiple_of(16),
+ "invalid CUDA shape"
+ );
+ self.ctx.check(unsafe {
+ (self.launch)(
+ args.as_ptr() as *mut Ptr,
+ b as i32,
+ l as i32,
+ (b * l) as i32,
+ self.ctx.stream,
+ )
+ })
+ }
+}
diff --git a/src/models/laya/README.md b/src/models/laya/README.md
index 5131c1d7..d9363698 100644
--- a/src/models/laya/README.md
+++ b/src/models/laya/README.md
@@ -37,6 +37,10 @@ and [measured scope](../../../recipe/laya/native/VALIDATION.md).
## CPU checks
The normal workspace tests cover configuration errors, malformed tensors, inventory mismatches and conversion boundaries without downloading weights.
+Feature-enabled CI also runs [CPU host/ABI regressions](../../../tests/cuda/laya/README.md)
+for resolved dispatch, grouped uploads and Graph capture cleanup. They compile a small
+native fixture with `cc` on Linux/macOS and require no CUDA toolkit or weights;
+they do not establish GPU numerical parity.
To check the complete checkpoint, use `convaiinnovations/laya` revision `55cf4c4ebb4ebe31b2550e8bdf3bd21b99753851` and a Python environment with PyTorch, safetensors and NumPy:
diff --git a/src/models/laya/src/model.rs b/src/models/laya/src/model.rs
index ba25a769..80b3e05e 100644
--- a/src/models/laya/src/model.rs
+++ b/src/models/laya/src/model.rs
@@ -2,7 +2,7 @@
use crate::{config::Config, packing::Batch, weights::Weights};
use anyhow::{Context, Result, ensure};
use half::bf16;
-use omni_cuda::{Buffer, Cuda, Graph, Ptr};
+use omni_cuda::{Buffer, Cuda, Graph, Kernel, Ptr};
use std::{
collections::{HashMap, VecDeque},
fs,
@@ -32,9 +32,6 @@ fn bytes32(v: &[f32]) -> Vec {
fn i32bytes(v: &[i32]) -> Vec {
v.iter().flat_map(|x| x.to_le_bytes()).collect()
}
-fn i64bytes(v: &[i64]) -> Vec {
- v.iter().flat_map(|x| x.to_le_bytes()).collect()
-}
fn decode_bf16(v: &[u8]) -> Vec {
v.as_chunks::<2>()
.0
@@ -120,6 +117,7 @@ struct Workspace {
ids: Buffer,
lens: Buffer,
types: Buffer,
+ staging: Vec,
x: Buffer,
y: Buffer,
qkv: Buffer,
@@ -150,6 +148,7 @@ impl Workspace {
ids: alloc(m * 8)?,
lens: alloc(b * 4)?,
types: alloc(b * 8)?,
+ staging: vec![0; m * 8 + b * 12],
x: alloc(m * D * 4)?,
y: alloc(m * D * 2)?,
qkv: alloc(m * D * 6)?,
@@ -169,14 +168,314 @@ impl Workspace {
}
}
+#[derive(Clone, Copy)]
+enum Scratch {
+ Ids,
+ Lens,
+ Types,
+ X,
+ Y,
+ Qkv,
+ O,
+ G,
+ Ff,
+}
+impl Scratch {
+ fn buffer(self, s: &Workspace) -> &Buffer {
+ match self {
+ Self::Ids => &s.ids,
+ Self::Lens => &s.lens,
+ Self::Types => &s.types,
+ Self::X => &s.x,
+ Self::Y => &s.y,
+ Self::Qkv => &s.qkv,
+ Self::O => &s.o,
+ Self::G => &s.g,
+ Self::Ff => &s.ff,
+ }
+ }
+}
+#[derive(Clone, Copy)]
+enum Argument {
+ Weight(Ptr),
+ Scratch(Scratch),
+}
+impl Argument {
+ fn pointer(self, s: &Workspace) -> Ptr {
+ match self {
+ Self::Weight(p) => p,
+ Self::Scratch(slot) => slot.buffer(s).ptr(),
+ }
+ }
+}
+enum SelectedKernel {
+ Fixed(Kernel),
+ Attention([Option; 3]),
+}
+impl SelectedKernel {
+ fn resolve(cuda: &Cuda, name: &str) -> Result {
+ if name == "attn_full" || name == "attn_local" {
+ // Sparse trusted bundles may omit wrappers for shapes never requested.
+ // Preserve their eager usability and fail when a missing shape is used.
+ Ok(Self::Attention([
+ Some(cuda.resolve(name)?),
+ cuda.resolve(&format!("{name}_b1_l512")).ok(),
+ cuda.resolve(&format!("{name}_b4_l512")).ok(),
+ ]))
+ } else {
+ Ok(Self::Fixed(cuda.resolve(name)?))
+ }
+ }
+ fn select(&self, b: usize, l: usize) -> Result<&Kernel> {
+ match self {
+ Self::Fixed(k) => Ok(k),
+ Self::Attention(k) => k[if l == 512 && b == 1 {
+ 1
+ } else if l == 512 && b == 4 {
+ 2
+ } else {
+ 0
+ }]
+ .as_ref()
+ .ok_or_else(|| {
+ anyhow::anyhow!("missing specialized attention kernel for B={b}, L={l}")
+ }),
+ }
+ }
+}
+enum Step {
+ Launch {
+ kernel: SelectedKernel,
+ args: [Argument; 5],
+ count: usize,
+ },
+ Dump {
+ name: String,
+ slot: Scratch,
+ bf: bool,
+ },
+}
+
+fn prepare_encoder(
+ cuda: &Cuda,
+ weights: &HashMap,
+ original_rope: bool,
+) -> Result> {
+ let steps = std::cell::RefCell::new(Vec::new());
+ let weight = |name: &str| Argument::Weight(weights[name].ptr());
+ let call = |name: &str, args: &[Argument]| -> Result<()> {
+ ensure!(args.len() <= 5, "encoder argument count");
+ let mut resolved = [Argument::Weight(std::ptr::null_mut()); 5];
+ resolved[..args.len()].copy_from_slice(args);
+ steps.borrow_mut().push(Step::Launch {
+ kernel: SelectedKernel::resolve(cuda, name)?,
+ args: resolved,
+ count: args.len(),
+ });
+ Ok(())
+ };
+ let dump = |name: &str, slot: Scratch, bf: bool| {
+ steps.borrow_mut().push(Step::Dump {
+ name: name.to_owned(),
+ slot,
+ bf,
+ });
+ Ok::<_, anyhow::Error>(())
+ };
+ let z = weight("zeros.1024");
+ let attention = |label: &str| format!("attn_{label}");
+ call(
+ "embed",
+ &[
+ Argument::Scratch(Scratch::Ids),
+ weight("encoder.embeddings.tok_embeddings.weight"),
+ weight("encoder.embeddings.norm.weight"),
+ Argument::Scratch(Scratch::X),
+ Argument::Scratch(Scratch::Y),
+ ],
+ )?;
+ dump("embedding", Scratch::X, false)?;
+ for i in 0..28 {
+ let p = format!("encoder.layers.{i}");
+ let w = |n: &str| weight(&format!("{p}.{n}"));
+ call(
+ "qkv",
+ &[
+ Argument::Scratch(Scratch::Y),
+ w("attn.Wqkv.weight"),
+ weight("zeros.3072"),
+ Argument::Scratch(Scratch::Qkv),
+ ],
+ )?;
+ let kind = if i % 3 == 0 { "full" } else { "local" };
+ call(
+ if original_rope {
+ "rope_original"
+ } else {
+ "rope"
+ },
+ &[
+ Argument::Scratch(Scratch::Qkv),
+ weight(&format!("rope_{kind}_cos")),
+ weight(&format!("rope_{kind}_sin")),
+ ],
+ )?;
+ call(
+ &attention(if i % 3 == 0 { "full" } else { "local" }),
+ &[
+ Argument::Scratch(Scratch::Qkv),
+ Argument::Scratch(Scratch::Lens),
+ Argument::Scratch(Scratch::O),
+ ],
+ )?;
+ call(
+ "out",
+ &[
+ Argument::Scratch(Scratch::O),
+ w("attn.Wo.weight"),
+ z,
+ Argument::Scratch(Scratch::Y),
+ ],
+ )?;
+ call(
+ "addln",
+ &[
+ Argument::Scratch(Scratch::X),
+ Argument::Scratch(Scratch::Y),
+ w("mlp_norm.weight"),
+ z,
+ Argument::Scratch(Scratch::Y),
+ ],
+ )?;
+ call(
+ "geglu",
+ &[
+ Argument::Scratch(Scratch::Y),
+ w("mlp.Wi.weight"),
+ Argument::Scratch(Scratch::G),
+ ],
+ )?;
+ call(
+ "down",
+ &[
+ Argument::Scratch(Scratch::G),
+ w("mlp.Wo.weight"),
+ z,
+ Argument::Scratch(Scratch::Y),
+ ],
+ )?;
+ let next = if i < 27 {
+ weight(&format!("encoder.layers.{}.attn_norm.weight", i + 1))
+ } else {
+ weight("encoder.final_norm.weight")
+ };
+ call(
+ "addln",
+ &[
+ Argument::Scratch(Scratch::X),
+ Argument::Scratch(Scratch::Y),
+ next,
+ z,
+ Argument::Scratch(Scratch::Y),
+ ],
+ )?;
+ if [0, 1, 2, 27].contains(&i) {
+ dump(&format!("encoder{i}_residual"), Scratch::X, false)?;
+ dump(&format!("encoder{i}_normalized"), Scratch::Y, true)?;
+ }
+ }
+ call(
+ "type",
+ &[
+ Argument::Scratch(Scratch::Y),
+ weight("type_emb.weight"),
+ Argument::Scratch(Scratch::Types),
+ Argument::Scratch(Scratch::X),
+ ],
+ )?;
+ for i in 0..2 {
+ let p = format!("head.layers.{i}");
+ let w = |n: &str| weight(&format!("{p}.{n}"));
+ call(
+ "ln_bias",
+ &[
+ Argument::Scratch(Scratch::X),
+ Argument::Scratch(Scratch::Y),
+ w("norm1.weight"),
+ w("norm1.bias"),
+ Argument::Scratch(Scratch::Y),
+ ],
+ )?;
+ call(
+ "head_in",
+ &[
+ Argument::Scratch(Scratch::Y),
+ w("self_attn.in_proj_weight"),
+ w("self_attn.in_proj_bias"),
+ Argument::Scratch(Scratch::Qkv),
+ ],
+ )?;
+ call(
+ &attention("full"),
+ &[
+ Argument::Scratch(Scratch::Qkv),
+ Argument::Scratch(Scratch::Lens),
+ Argument::Scratch(Scratch::O),
+ ],
+ )?;
+ call(
+ "head_out",
+ &[
+ Argument::Scratch(Scratch::O),
+ w("self_attn.out_proj.weight"),
+ w("self_attn.out_proj.bias"),
+ Argument::Scratch(Scratch::Y),
+ ],
+ )?;
+ call(
+ "addln_bias",
+ &[
+ Argument::Scratch(Scratch::X),
+ Argument::Scratch(Scratch::Y),
+ w("norm2.weight"),
+ w("norm2.bias"),
+ Argument::Scratch(Scratch::Y),
+ ],
+ )?;
+ call(
+ "ffn1",
+ &[
+ Argument::Scratch(Scratch::Y),
+ w("linear1.weight"),
+ w("linear1.bias"),
+ Argument::Scratch(Scratch::Ff),
+ ],
+ )?;
+ call(
+ "ffn2",
+ &[
+ Argument::Scratch(Scratch::Ff),
+ w("linear2.weight"),
+ w("linear2.bias"),
+ Argument::Scratch(Scratch::Y),
+ ],
+ )?;
+ call(
+ "residual",
+ &[Argument::Scratch(Scratch::X), Argument::Scratch(Scratch::Y)],
+ )?;
+ }
+ Ok(steps.into_inner())
+}
+
pub struct Model {
pub config: Config,
cuda: Cuda,
blas: Blas,
weights: HashMap,
+ plan: Vec,
cache: VecDeque,
graphs: bool,
- original_rope: bool,
}
impl Drop for Model {
fn drop(&mut self) {
@@ -246,160 +545,38 @@ impl Model {
weights.len()
);
}
+ let plan = prepare_encoder(&cuda, &weights, original_rope)?;
Ok(Self {
config,
cuda,
blas,
weights,
+ plan,
cache: VecDeque::new(),
graphs,
- original_rope,
})
}
fn w(&self, n: &str) -> &Buffer {
&self.weights[n]
}
fn encode(&self, s: &Workspace) -> Result<()> {
- let (b, l) = (s.b, s.l);
- let z = self.w("zeros.1024").ptr();
- let attention = |label: &str| {
- if l == 512 && (b == 1 || b == 4) {
- format!("attn_{label}_b{b}_l512")
- } else {
- format!("attn_{label}")
- }
- };
- // All pointers refer to checked fixed-shape, resident allocations in this worker.
- let call = |name: &str, args: &[Ptr]| unsafe { self.cuda.launch(name, args, b, l) };
- call(
- "embed",
- &[
- s.ids.ptr(),
- self.w("encoder.embeddings.tok_embeddings.weight").ptr(),
- self.w("encoder.embeddings.norm.weight").ptr(),
- s.x.ptr(),
- s.y.ptr(),
- ],
- )?;
- self.dump("embedding", &s.x, false)?;
- for i in 0..28 {
- let p = format!("encoder.layers.{i}");
- let w = |n: &str| self.w(&format!("{p}.{n}")).ptr();
- call(
- "qkv",
- &[
- s.y.ptr(),
- w("attn.Wqkv.weight"),
- self.w("zeros.3072").ptr(),
- s.qkv.ptr(),
- ],
- )?;
- let kind = if i % 3 == 0 { "full" } else { "local" };
- call(
- if self.original_rope {
- "rope_original"
- } else {
- "rope"
- },
- &[
- s.qkv.ptr(),
- self.w(&format!("rope_{kind}_cos")).ptr(),
- self.w(&format!("rope_{kind}_sin")).ptr(),
- ],
- )?;
- call(
- &attention(if i % 3 == 0 { "full" } else { "local" }),
- &[s.qkv.ptr(), s.lens.ptr(), s.o.ptr()],
- )?;
- call("out", &[s.o.ptr(), w("attn.Wo.weight"), z, s.y.ptr()])?;
- call(
- "addln",
- &[s.x.ptr(), s.y.ptr(), w("mlp_norm.weight"), z, s.y.ptr()],
- )?;
- call("geglu", &[s.y.ptr(), w("mlp.Wi.weight"), s.g.ptr()])?;
- call("down", &[s.g.ptr(), w("mlp.Wo.weight"), z, s.y.ptr()])?;
- let next = if i < 27 {
- self.w(&format!("encoder.layers.{}.attn_norm.weight", i + 1))
- } else {
- self.w("encoder.final_norm.weight")
- };
- call("addln", &[s.x.ptr(), s.y.ptr(), next.ptr(), z, s.y.ptr()])?;
- if [0, 1, 2, 27].contains(&i) {
- self.dump(&format!("encoder{i}_residual"), &s.x, false)?;
- self.dump(&format!("encoder{i}_normalized"), &s.y, true)?;
+ for step in &self.plan {
+ match step {
+ Step::Launch {
+ kernel,
+ args,
+ count,
+ } => {
+ let pointers = args.map(|arg| arg.pointer(s));
+ unsafe {
+ kernel
+ .select(s.b, s.l)?
+ .launch(&pointers[..*count], s.b, s.l)
+ }?;
+ }
+ Step::Dump { name, slot, bf } => self.dump(name, slot.buffer(s), *bf)?,
}
}
- call(
- "type",
- &[
- s.y.ptr(),
- self.w("type_emb.weight").ptr(),
- s.types.ptr(),
- s.x.ptr(),
- ],
- )?;
- for i in 0..2 {
- let p = format!("head.layers.{i}");
- let w = |n: &str| self.w(&format!("{p}.{n}")).ptr();
- call(
- "ln_bias",
- &[
- s.x.ptr(),
- s.y.ptr(),
- w("norm1.weight"),
- w("norm1.bias"),
- s.y.ptr(),
- ],
- )?;
- call(
- "head_in",
- &[
- s.y.ptr(),
- w("self_attn.in_proj_weight"),
- w("self_attn.in_proj_bias"),
- s.qkv.ptr(),
- ],
- )?;
- call(&attention("full"), &[s.qkv.ptr(), s.lens.ptr(), s.o.ptr()])?;
- call(
- "head_out",
- &[
- s.o.ptr(),
- w("self_attn.out_proj.weight"),
- w("self_attn.out_proj.bias"),
- s.y.ptr(),
- ],
- )?;
- call(
- "addln_bias",
- &[
- s.x.ptr(),
- s.y.ptr(),
- w("norm2.weight"),
- w("norm2.bias"),
- s.y.ptr(),
- ],
- )?;
- call(
- "ffn1",
- &[
- s.y.ptr(),
- w("linear1.weight"),
- w("linear1.bias"),
- s.ff.ptr(),
- ],
- )?;
- call(
- "ffn2",
- &[
- s.ff.ptr(),
- w("linear2.weight"),
- w("linear2.bias"),
- s.y.ptr(),
- ],
- )?;
- call("residual", &[s.x.ptr(), s.y.ptr()])?;
- }
Ok(())
}
fn dump(&self, name: &str, buffer: &Buffer, bf: bool) -> Result<()> {
@@ -466,9 +643,19 @@ impl Model {
} else {
Workspace::new(&self.cuda, batch.b, batch.l)?
};
- s.ids.write(&i64bytes(&batch.input_ids))?;
- s.lens.write(&i32bytes(&batch.lens))?;
- s.types.write(&i64bytes(&batch.qtypes))?;
+ let (ids, rest) = s.staging.split_at_mut(batch.input_ids.len() * 8);
+ let (lens, types) = rest.split_at_mut(batch.lens.len() * 4);
+ for (bytes, value) in ids.as_chunks_mut::<8>().0.iter_mut().zip(&batch.input_ids) {
+ bytes.copy_from_slice(&value.to_le_bytes());
+ }
+ for (bytes, value) in lens.as_chunks_mut::<4>().0.iter_mut().zip(&batch.lens) {
+ bytes.copy_from_slice(&value.to_le_bytes());
+ }
+ for (bytes, value) in types.as_chunks_mut::<8>().0.iter_mut().zip(&batch.qtypes) {
+ bytes.copy_from_slice(&value.to_le_bytes());
+ }
+ self.cuda
+ .write_many(&[(&s.ids, ids), (&s.lens, lens), (&s.types, types)])?;
if self.graphs {
if s.graph.is_none() {
self.encode(&s)?;
diff --git a/tests/cuda/laya/README.md b/tests/cuda/laya/README.md
new file mode 100644
index 00000000..ace9064f
--- /dev/null
+++ b/tests/cuda/laya/README.md
@@ -0,0 +1,27 @@
+# Laya CPU host regressions
+
+`fixture.c` implements the Laya native ABI using host allocations, deterministic
+copy/failure counters, a pointer launch trace and a minimal captured fill
+operation. It performs no Laya tensor math and needs no CUDA toolkit, GPU,
+checkpoint, Python or network access. The trace oracle in
+`tests/laya/original_encode.rs` freezes the former encoder launch sequence from
+`7f39ac40902c374803992407bb26eeba29c8a588`; it checks dispatch pointers/order,
+not numerical outputs.
+
+Run from the repository root with stable Rust and a C compiler available as `cc`:
+
+```sh
+cargo test --locked -p omni-cuda --test laya_runtime
+cargo test --locked -p omni-laya --features serve --lib model::tests
+```
+
+The Rust helper embeds the maintained C source and compiles it into a unique
+`tempfile` directory for every test. It uses `cc -shared -fPIC` on Linux and
+`cc -dynamiclib` on macOS. Unique library paths isolate process-global counters
+while Cargo runs tests concurrently; test-owned temporary directories retain the
+library until their CUDA handles are dropped. These native fixture tests are
+Unix-only. CI enables `omni-laya/serve`, compiling the native model/worker and
+running its root-tree private tests on Linux.
+
+The registered tests at this stage cover resolved dispatch, grouped uploads and Graph capture cleanup.
+They validate host behavior without Laya tensor math or a CUDA driver.
diff --git a/tests/cuda/laya/capture.rs b/tests/cuda/laya/capture.rs
new file mode 100644
index 00000000..e2c29ff9
--- /dev/null
+++ b/tests/cuda/laya/capture.rs
@@ -0,0 +1,99 @@
+use super::fixture::Fixture;
+use omni_cuda::Cuda;
+
+#[test]
+fn capture_unwind_recovery_and_prohibited_operations() {
+ let fixture = Fixture::new(&[]);
+ let path = fixture.path();
+ let cuda = unsafe { Cuda::load(path) }.unwrap();
+ let buffer = cuda.alloc(64).unwrap();
+ let kernel = cuda.resolve("fill").unwrap();
+ let allocations = unsafe {
+ cuda.symbol:: i32>(b"test_allocations\0")
+ .unwrap()
+ };
+ let nested = unsafe {
+ cuda.capture(|| {
+ let before = allocations();
+ assert!(cuda.alloc(64).is_err());
+ assert_eq!(allocations(), before);
+ assert!(cuda.sync().is_err());
+ assert!(buffer.write(&[0; 64]).is_err());
+ assert!(buffer.read(64).is_err());
+ assert!(cuda.capture(|| Ok(())).is_err());
+ kernel.launch(&[buffer.ptr()], 1, 16)
+ })
+ }
+ .unwrap();
+ nested.replay().unwrap();
+ cuda.sync().unwrap();
+ assert_eq!(buffer.read(64).unwrap(), [73; 64]);
+ let error = unsafe { cuda.capture(|| Err(anyhow::anyhow!("injected closure error"))) };
+ assert!(error.is_err());
+ let panic = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| unsafe {
+ cuda.capture(|| panic!("injected panic"))
+ }));
+ assert!(panic.is_err());
+ let recovered = unsafe { cuda.capture(|| kernel.launch(&[buffer.ptr()], 1, 16)) }.unwrap();
+ recovered.replay().unwrap();
+ cuda.sync().unwrap();
+ let mode = unsafe {
+ cuda.symbol::(b"test_graph_mode\0")
+ .unwrap()
+ };
+ let freed = unsafe {
+ cuda.symbol:: i32>(b"test_freed\0")
+ .unwrap()
+ };
+ let before = unsafe { freed() };
+ unsafe { mode(1) };
+ assert!(unsafe { cuda.capture(|| Ok(())) }.is_err());
+ assert_eq!(unsafe { freed() }, before + 1);
+ unsafe { mode(2) };
+ assert!(unsafe { cuda.capture(|| Ok(())) }.is_err());
+ unsafe { mode(0) };
+ drop(cuda);
+ recovered.replay().unwrap();
+ drop(recovered);
+ drop(nested);
+ assert_eq!(buffer.read(64).unwrap(), [73; 64]);
+}
+
+#[test]
+fn eager_library_without_graph_cleanup_symbols_stays_usable() {
+ let fixture = Fixture::new(&["OMIT_GRAPH_FREE"]);
+ let path = fixture.path();
+ let cuda = unsafe { Cuda::load(path) }.unwrap();
+ let begins = unsafe {
+ cuda.symbol:: i32>(b"test_begins\0")
+ .unwrap()
+ };
+ let before = unsafe { begins() };
+ assert!(unsafe { cuda.capture(|| Ok(())) }.is_err());
+ assert_eq!(unsafe { begins() }, before);
+ let buffer = cuda.upload(&[9; 64]).unwrap();
+ assert_eq!(buffer.read(64).unwrap(), [9; 64]);
+}
+
+#[test]
+fn replay_is_rejected_inside_capture_before_entering_driver() {
+ let fixture = Fixture::new(&[]);
+ let path = fixture.path();
+ let cuda = unsafe { Cuda::load(path) }.unwrap();
+ let buffer = cuda.alloc(64).unwrap();
+ let kernel = cuda.resolve("fill").unwrap();
+ let first = unsafe { cuda.capture(|| kernel.launch(&[buffer.ptr()], 1, 16)) }.unwrap();
+ let second = unsafe {
+ cuda.capture(|| {
+ assert!(first.replay().is_err());
+ kernel.launch(&[buffer.ptr()], 1, 16)
+ })
+ }
+ .unwrap();
+ first.replay().unwrap();
+ cuda.sync().unwrap();
+ drop(first);
+ second.replay().unwrap();
+ cuda.sync().unwrap();
+ assert_eq!(buffer.read(64).unwrap(), [73; 64]);
+}
diff --git a/tests/cuda/laya/fixture.c b/tests/cuda/laya/fixture.c
new file mode 100644
index 00000000..896f4e3b
--- /dev/null
+++ b/tests/cuda/laya/fixture.c
@@ -0,0 +1,163 @@
+// CPU-only ABI fixture: record launches and ownership; do not perform model math.
+#include
+#include
+#include
+#include
+
+static int capture_active, graph_mode, begins, ends, freed, allocations;
+static int stream_value, stream_frees, copy_mode, sync_count;
+static void *captured;
+void test_graph_mode(int mode) { graph_mode = mode; }
+int test_begins(void) { return begins; }
+int test_ends(void) { return ends; }
+int test_freed(void) { return freed; }
+int test_allocations(void) { return allocations; }
+int test_stream_frees(void) { return stream_frees; }
+void test_copy_mode(int mode) { copy_mode = mode; }
+int test_syncs(void) { return sync_count; }
+
+typedef struct Allocation {
+ void *pointer;
+ size_t bytes;
+ struct Allocation *next;
+} Allocation;
+static Allocation *all_allocations;
+static size_t live_bytes, peak_bytes;
+size_t test_live_bytes(void) { return live_bytes; }
+size_t test_peak_bytes(void) { return peak_bytes; }
+void test_reset_peak(void) { peak_bytes = live_bytes; }
+
+int laya_init(void **stream) { *stream = &stream_value; return 0; }
+const char *laya_error(int code) { return "fixture CUDA error"; }
+int laya_alloc(void **pointer, size_t bytes) {
+ allocations++;
+ if (capture_active) return 90;
+ Allocation *allocation = malloc(sizeof(*allocation));
+ if (!allocation) return 1;
+ *pointer = calloc(1, bytes);
+ if (!*pointer) { free(allocation); return 1; }
+ *allocation = (Allocation){*pointer, bytes, all_allocations};
+ all_allocations = allocation;
+ live_bytes += bytes;
+ if (live_bytes > peak_bytes) peak_bytes = live_bytes;
+ return 0;
+}
+int laya_free(void *pointer) {
+ Allocation **link = &all_allocations;
+ while (*link && (*link)->pointer != pointer) link = &(*link)->next;
+ if (*link) {
+ Allocation *allocation = *link;
+ *link = allocation->next;
+ live_bytes -= allocation->bytes;
+ free(allocation);
+ }
+ free(pointer);
+ return 0;
+}
+int laya_upload(void *destination, const void *source, size_t bytes, void *stream) {
+ memcpy(destination, source, bytes);
+ // Submission may report failure after copying; Rust still must drain it.
+ return copy_mode & 1 ? 77 : 0;
+}
+int laya_download(void *destination, const void *source, size_t bytes, void *stream) {
+ memcpy(destination, source, bytes);
+ return 0;
+}
+int laya_sync(void *stream) {
+ sync_count++;
+ if (capture_active) return 90;
+ return copy_mode & 2 ? 78 : 0;
+}
+int laya_stream_free(void *stream) { stream_frees++; return 0; }
+int laya_fill(void **pointers, int batch, int length, int rows, void *stream) {
+ if (capture_active) captured = pointers[0];
+ else memset(pointers[0], 73, 64);
+ return 0;
+}
+
+static char trace[1048576];
+static size_t used;
+void test_reset(void) { used = 0; trace[0] = 0; }
+const char *test_trace(void) { return trace; }
+static int launch(const char *name, void **pointers, int batch, int length,
+ int rows, int count) {
+ // Reserve a complete record before snprintf so long cache tests cannot overflow.
+ if (sizeof(trace) - used < 512) return -1;
+ used += snprintf(trace + used, sizeof(trace) - used, "%s(%d,%d,%d)",
+ name, batch, length, rows);
+ for (int i = 0; i < count; i++)
+ used += snprintf(trace + used, sizeof(trace) - used, "/%p", pointers[i]);
+ used += snprintf(trace + used, sizeof(trace) - used, "\n");
+ return 0;
+}
+#define KERNEL(name, count) \
+ int laya_##name(void **p, int b, int l, int rows, void *stream) { \
+ return launch(#name, p, b, l, rows, count); \
+ }
+KERNEL(embed, 5)
+KERNEL(qkv, 4)
+KERNEL(rope, 3)
+KERNEL(rope_original, 3)
+KERNEL(attn_full, 3)
+KERNEL(attn_local, 3)
+KERNEL(out, 4)
+KERNEL(addln, 5)
+KERNEL(geglu, 3)
+KERNEL(down, 4)
+KERNEL(type, 4)
+KERNEL(ln_bias, 5)
+KERNEL(head_in, 4)
+KERNEL(head_out, 4)
+KERNEL(addln_bias, 5)
+KERNEL(ffn1, 4)
+KERNEL(ffn2, 4)
+KERNEL(residual, 2)
+KERNEL(gather, 5)
+KERNEL(features, 4)
+#ifndef OMIT_SPECIALIZED_ATTN
+KERNEL(attn_full_b1_l512, 3)
+KERNEL(attn_full_b4_l512, 3)
+KERNEL(attn_local_b1_l512, 3)
+KERNEL(attn_local_b4_l512, 3)
+#endif
+#undef KERNEL
+
+int laya_blas_create(void **pointer, void *stream) {
+ *pointer = malloc(1);
+ return *pointer ? 0 : 1;
+}
+int laya_blas_free(void *pointer) { free(pointer); return 0; }
+int laya_linear(void *handle, void *input, void *weight, void *bias, void *output,
+ int rows, int columns, int width, int gelu, void *stream) {
+ memset(output, 0, (size_t)rows * columns * 2);
+ return 0;
+}
+
+typedef struct { void *target; } TestGraph;
+int laya_capture_begin(void *stream) {
+ begins++;
+ if (capture_active) return 90;
+ capture_active = 1;
+ captured = 0;
+ return 0;
+}
+int laya_capture_end(void *stream, void **output) {
+ ends++;
+ capture_active = 0;
+ *output = 0;
+ if (graph_mode == 2) return 0; // Null-success injection.
+ TestGraph *graph = malloc(sizeof(*graph));
+ if (!graph) return 1;
+ graph->target = captured;
+ *output = graph;
+ return graph_mode == 1 ? 91 : 0; // Partial ownership with an error.
+}
+int laya_graph_run(void *pointer, void *stream) {
+ if (!pointer) return 90;
+ TestGraph *graph = pointer;
+ if (graph->target) memset(graph->target, 73, 64);
+ return 0;
+}
+#ifndef OMIT_GRAPH_FREE
+int laya_graph_free(void *pointer) { freed++; free(pointer); return 0; }
+#endif
diff --git a/tests/cuda/laya/fixture.rs b/tests/cuda/laya/fixture.rs
new file mode 100644
index 00000000..1b4fdbc5
--- /dev/null
+++ b/tests/cuda/laya/fixture.rs
@@ -0,0 +1,51 @@
+//! Build a distinct library per test so its global counters never leak between tests.
+use std::{
+ path::{Path, PathBuf},
+ process::Command,
+};
+use tempfile::{TempDir, tempdir};
+
+pub struct Fixture {
+ _directory: TempDir,
+ path: PathBuf,
+}
+impl Fixture {
+ pub fn new(defines: &[&str]) -> Self {
+ let directory = tempdir().expect("fixture directory");
+ let path = directory.path().join(if cfg!(target_os = "macos") {
+ "liblaya_fixture.dylib"
+ } else {
+ "liblaya_fixture.so"
+ });
+ let source = directory.path().join("fixture.c");
+ std::fs::write(&source, include_str!("fixture.c")).expect("write fixture source");
+ let mut compiler = Command::new("cc");
+ compiler.arg("-std=c11");
+ if cfg!(target_os = "macos") {
+ compiler.arg("-dynamiclib");
+ } else {
+ compiler.args(["-shared", "-fPIC"]);
+ }
+ for define in defines {
+ compiler.arg(format!("-D{define}"));
+ }
+ let output = compiler
+ .arg(&source)
+ .arg("-o")
+ .arg(&path)
+ .output()
+ .expect("a C compiler is required for the CPU ABI fixture");
+ assert!(
+ output.status.success(),
+ "fixture compilation failed: {}",
+ String::from_utf8_lossy(&output.stderr)
+ );
+ Self {
+ _directory: directory,
+ path,
+ }
+ }
+ pub fn path(&self) -> &Path {
+ &self.path
+ }
+}
diff --git a/tests/cuda/laya/grouped.rs b/tests/cuda/laya/grouped.rs
new file mode 100644
index 00000000..41219874
--- /dev/null
+++ b/tests/cuda/laya/grouped.rs
@@ -0,0 +1,60 @@
+use super::fixture::Fixture;
+use omni_cuda::Cuda;
+
+#[test]
+fn grouped_upload_prevalidates_and_drains_copy_failures() {
+ let fixture = Fixture::new(&[]);
+ let path = fixture.path();
+ let cuda = unsafe { Cuda::load(path) }.unwrap();
+ let other = unsafe { Cuda::load(path) }.unwrap();
+ let a = cuda.alloc(64).unwrap();
+ let b = cuda.alloc(64).unwrap();
+ let foreign = other.alloc(64).unwrap();
+ let syncs = unsafe {
+ cuda.symbol:: i32>(b"test_syncs\0")
+ .unwrap()
+ };
+ let before = unsafe { syncs() };
+ cuda.write_many(&[(&a, &[1; 64]), (&b, &[2; 64])]).unwrap();
+ assert_eq!(
+ unsafe { syncs() },
+ before + 1,
+ "one drain for the complete group"
+ );
+ assert_eq!(a.read(64).unwrap(), [1; 64]);
+ assert_eq!(b.read(64).unwrap(), [2; 64]);
+ assert!(
+ cuda.write_many(&[(&a, &[3; 64]), (&foreign, &[4; 64])])
+ .is_err()
+ );
+ assert_eq!(a.read(64).unwrap(), [1; 64]);
+ assert!(
+ cuda.write_many(&[(&a, &[3; 64][..]), (&b, &[4; 65][..])])
+ .is_err()
+ );
+ assert_eq!(a.read(64).unwrap(), [1; 64]);
+ let mode = unsafe {
+ cuda.symbol::(b"test_copy_mode\0")
+ .unwrap()
+ };
+ unsafe { mode(1) };
+ let before = unsafe { syncs() };
+ let error = cuda
+ .write_many(&[(&a, &[5; 64]), (&b, &[6; 64])])
+ .unwrap_err();
+ assert!(error.to_string().contains("CUDA"));
+ assert_eq!(unsafe { syncs() }, before + 1);
+ assert_eq!(
+ b.read(64).unwrap(),
+ [2; 64],
+ "stop submission after the failed member"
+ );
+ unsafe { mode(3) };
+ let error = cuda.write_many(&[(&a, &[5; 64])]).unwrap_err();
+ assert!(error.to_string().contains("also failed"));
+ unsafe { mode(0) };
+ let before = unsafe { syncs() };
+ cuda.write_many(&[]).unwrap();
+ cuda.write_many(&[(&a, &[])]).unwrap();
+ assert_eq!(unsafe { syncs() }, before);
+}
diff --git a/tests/cuda/laya/resolved.rs b/tests/cuda/laya/resolved.rs
new file mode 100644
index 00000000..6cdf456c
--- /dev/null
+++ b/tests/cuda/laya/resolved.rs
@@ -0,0 +1,30 @@
+use super::fixture::Fixture;
+use omni_cuda::Cuda;
+
+#[test]
+fn resolved_handle_retains_context_after_owner_drop() {
+ let fixture = Fixture::new(&[]);
+ let path = fixture.path();
+ let observer = unsafe { libloading::Library::new(path) }.unwrap();
+ let frees = unsafe {
+ *observer
+ .get:: i32>(b"test_stream_frees\0")
+ .unwrap()
+ };
+ let cuda = unsafe { Cuda::load(path) }.unwrap();
+ let buffer = cuda.upload(&[0; 64]).unwrap();
+ let kernel = cuda.resolve("fill").unwrap();
+ assert!(cuda.resolve("missing").is_err());
+ drop(cuda);
+ unsafe { kernel.launch(&[buffer.ptr()], 1, 16) }.unwrap();
+ assert_eq!(buffer.read(64).unwrap(), [73; 64]);
+ assert!(unsafe { kernel.launch(&[buffer.ptr()], 0, 16) }.is_err());
+ drop(buffer);
+ assert_eq!(
+ unsafe { frees() },
+ 0,
+ "Kernel is the sole remaining Context owner"
+ );
+ drop(kernel);
+ assert_eq!(unsafe { frees() }, 1);
+}
diff --git a/tests/cuda/laya_runtime.rs b/tests/cuda/laya_runtime.rs
new file mode 100644
index 00000000..f1b79717
--- /dev/null
+++ b/tests/cuda/laya_runtime.rs
@@ -0,0 +1,10 @@
+//! CPU tests of production Rust bindings against a deterministic native ABI.
+#![cfg(unix)]
+#[path = "laya/capture.rs"]
+mod capture;
+#[path = "laya/fixture.rs"]
+mod fixture;
+#[path = "laya/grouped.rs"]
+mod grouped;
+#[path = "laya/resolved.rs"]
+mod resolved;
diff --git a/tests/laya/model.rs b/tests/laya/model.rs
index 12dff82e..ab143385 100644
--- a/tests/laya/model.rs
+++ b/tests/laya/model.rs
@@ -1,4 +1,6 @@
use super::storage_dtype;
+#[cfg(unix)]
+use super::*;
#[test]
fn resident_precision_matches_cuda_consumers() {
@@ -27,3 +29,22 @@ fn resident_precision_matches_cuda_consumers() {
1
);
}
+
+// Private model regression helpers remain in the repository-level tests tree.
+#[cfg(unix)]
+#[path = "../cuda/laya/fixture.rs"]
+mod fixture;
+#[cfg(unix)]
+#[path = "model_fixture.rs"]
+mod model_fixture;
+#[cfg(unix)]
+use model_fixture::fixture_model_at;
+#[cfg(unix)]
+#[path = "model_dispatch.rs"]
+mod dispatch;
+#[cfg(unix)]
+#[path = "model_grouped.rs"]
+mod grouped;
+#[cfg(unix)]
+#[path = "original_encode.rs"]
+mod original_encode;
diff --git a/tests/laya/model_dispatch.rs b/tests/laya/model_dispatch.rs
new file mode 100644
index 00000000..d2feb506
--- /dev/null
+++ b/tests/laya/model_dispatch.rs
@@ -0,0 +1,58 @@
+use super::*;
+
+#[test]
+fn prepared_plan_preserves_original_launch_trace_for_all_shape_selectors() {
+ let fixture = fixture::Fixture::new(&[]);
+ let lib = unsafe { Cuda::load(fixture.path()) }.unwrap();
+ let reset = unsafe {
+ lib.symbol::(b"test_reset\0")
+ .unwrap()
+ };
+ let trace = unsafe {
+ lib.symbol:: *const i8>(b"test_trace\0")
+ .unwrap()
+ };
+ for original_rope in [false, true] {
+ let model = fixture_model_at(original_rope, fixture.path());
+ for b in [1, 2, 4, 8, 16] {
+ for l in [16, 512] {
+ let s = Workspace::new(&model.cuda, b, l).unwrap();
+ unsafe { reset() };
+ model.encode_original(&s, original_rope).unwrap();
+ let expected = unsafe { std::ffi::CStr::from_ptr(trace()) }
+ .to_bytes()
+ .to_vec();
+ unsafe { reset() };
+ model.encode(&s).unwrap();
+ let actual = unsafe { std::ffi::CStr::from_ptr(trace()) }
+ .to_bytes()
+ .to_vec();
+ assert_eq!(
+ expected, actual,
+ "B={b}, L={l}, original_rope={original_rope}"
+ );
+ assert_eq!(actual.split(|x| *x == b'\n').count() - 1, 242);
+ }
+ }
+ }
+}
+
+#[test]
+fn missing_specialized_attention_only_rejects_its_shape() {
+ let fixture = fixture::Fixture::new(&["OMIT_SPECIALIZED_ATTN"]);
+ let model = fixture_model_at(false, fixture.path());
+ let short = Workspace::new(&model.cuda, 1, 16).unwrap();
+ model.encode(&short).unwrap();
+ let generic_long = Workspace::new(&model.cuda, 2, 512).unwrap();
+ model.encode(&generic_long).unwrap();
+ for b in [1, 4] {
+ let specialized = Workspace::new(&model.cuda, b, 512).unwrap();
+ assert!(
+ model
+ .encode(&specialized)
+ .unwrap_err()
+ .to_string()
+ .contains("missing specialized attention")
+ );
+ }
+}
diff --git a/tests/laya/model_fixture.rs b/tests/laya/model_fixture.rs
new file mode 100644
index 00000000..0fd19b00
--- /dev/null
+++ b/tests/laya/model_fixture.rs
@@ -0,0 +1,70 @@
+use super::*;
+use serde_json::{Value, json};
+use tempfile::tempdir;
+fn configs() -> (Value, Value) {
+ let encoder = json!({
+ "model_type": "modernbert", "hidden_activation": "gelu",
+ "attention_bias": false, "mlp_bias": false, "norm_bias": false,
+ "hidden_size": 1024, "intermediate_size": 2624,
+ "num_attention_heads": 16, "num_hidden_layers": 28,
+ "vocab_size": 50368, "norm_eps": 0.00001, "local_attention": 128,
+ "layer_types": (["full_attention", "sliding_attention", "sliding_attention"]
+ .repeat(10)[..28]),
+ "rope_parameters": {
+ "full_attention": {"rope_type": "default", "rope_theta": 160000.0},
+ "sliding_attention": {"rope_type": "default", "rope_theta": 10000.0}
+ }
+ });
+ let agent = json!({
+ "max_len": 512, "head_max_len": 192, "head_layers": 2,
+ "temperature": [0.5, 1.0, 2.0], "temperature_by_options": {"4": 1.5}
+ });
+ (encoder, agent)
+}
+
+fn load_config(encoder: &Value, agent: &Value) -> anyhow::Result {
+ let dir = tempdir()?;
+ fs::create_dir(dir.path().join("encoder"))?;
+ fs::write(dir.path().join("encoder/config.json"), encoder.to_string())?;
+ fs::write(dir.path().join("rl_agent_config.json"), agent.to_string())?;
+ Config::load(dir.path())
+}
+
+pub(super) fn fixture_model_at(original_rope: bool, path: &Path) -> Model {
+ let cuda = unsafe { Cuda::load(path) }.unwrap();
+ let mut weights = HashMap::new();
+ for spec in crate::weights::checkpoint_tensors() {
+ let width = if storage_dtype(&spec.name) == "f32" {
+ 4
+ } else {
+ 2
+ };
+ weights.insert(
+ spec.name,
+ cuda.alloc(spec.shape.iter().product::() * width)
+ .unwrap(),
+ );
+ }
+ for n in [D, 3 * D, 4 * D] {
+ weights.insert(format!("zeros.{n}"), cuda.alloc(n * 4).unwrap());
+ }
+ for kind in ["full", "local"] {
+ for part in ["cos", "sin"] {
+ weights.insert(
+ format!("rope_{kind}_{part}"),
+ cuda.alloc(512 * 32 * 4).unwrap(),
+ );
+ }
+ }
+ let (encoder, agent) = configs();
+ let plan = prepare_encoder(&cuda, &weights, original_rope).unwrap();
+ Model {
+ config: load_config(&encoder, &agent).unwrap(),
+ blas: Blas::new(&cuda).unwrap(),
+ cuda,
+ weights,
+ plan,
+ cache: VecDeque::new(),
+ graphs: false,
+ }
+}
diff --git a/tests/laya/model_grouped.rs b/tests/laya/model_grouped.rs
new file mode 100644
index 00000000..36d08955
--- /dev/null
+++ b/tests/laya/model_grouped.rs
@@ -0,0 +1,58 @@
+use super::*;
+
+#[test]
+fn grouped_model_inputs_reuse_staging_and_update_all_values() {
+ let fixture = fixture::Fixture::new(&[]);
+ let mut model = fixture_model_at(false, fixture.path());
+ let mut batch = Batch {
+ b: 1,
+ l: 16,
+ input_ids: vec![0; 16],
+ lens: vec![16],
+ qtypes: vec![0],
+ markers: vec![vec![0]],
+ };
+ let syncs = unsafe {
+ model
+ .cuda
+ .symbol:: i32>(b"test_syncs\0")
+ .unwrap()
+ };
+ let before = unsafe { syncs() };
+ model.infer(&batch).unwrap();
+ assert_eq!(
+ unsafe { syncs() },
+ before + 5,
+ "one grouped input sync, two marker uploads, two synchronized readbacks"
+ );
+ let first = model.cache.back().unwrap().staging.as_ptr();
+ batch.input_ids.fill(7);
+ batch.lens[0] = 12;
+ batch.qtypes[0] = 2;
+ model.infer(&batch).unwrap();
+ let workspace = model.cache.back().unwrap();
+ assert_eq!(first, workspace.staging.as_ptr());
+ assert_eq!(
+ workspace.ids.read(16 * 8).unwrap(),
+ [7i64.to_le_bytes(); 16].concat()
+ );
+ assert_eq!(workspace.lens.read(4).unwrap(), 12i32.to_le_bytes());
+ assert_eq!(workspace.types.read(8).unwrap(), 2i64.to_le_bytes());
+ batch.input_ids.fill(0);
+ batch.lens[0] = 16;
+ batch.qtypes[0] = 0;
+ model.infer(&batch).unwrap();
+ assert_eq!(
+ model.cache.back().unwrap().ids.read(16 * 8).unwrap(),
+ [0; 16 * 8]
+ );
+ let before = unsafe { syncs() };
+ batch.qtypes[0] = 3;
+ assert!(model.infer(&batch).is_err());
+ assert_eq!(
+ unsafe { syncs() },
+ before,
+ "invalid request does not submit or synchronize"
+ );
+ assert_eq!(first, model.cache.back().unwrap().staging.as_ptr());
+}
diff --git a/tests/laya/original_encode.rs b/tests/laya/original_encode.rs
new file mode 100644
index 00000000..12f17e8a
--- /dev/null
+++ b/tests/laya/original_encode.rs
@@ -0,0 +1,148 @@
+// Frozen pre-prepared-plan encode oracle (merge base 7f39ac4), renamed for tests.
+use super::*;
+
+impl Model {
+ pub(super) fn encode_original(&self, s: &Workspace, original_rope: bool) -> Result<()> {
+ let (b, l) = (s.b, s.l);
+ let z = self.w("zeros.1024").ptr();
+ let attention = |label: &str| {
+ if l == 512 && (b == 1 || b == 4) {
+ format!("attn_{label}_b{b}_l512")
+ } else {
+ format!("attn_{label}")
+ }
+ };
+ // All pointers refer to checked fixed-shape, resident allocations in this worker.
+ let call = |name: &str, args: &[Ptr]| unsafe { self.cuda.launch(name, args, b, l) };
+ call(
+ "embed",
+ &[
+ s.ids.ptr(),
+ self.w("encoder.embeddings.tok_embeddings.weight").ptr(),
+ self.w("encoder.embeddings.norm.weight").ptr(),
+ s.x.ptr(),
+ s.y.ptr(),
+ ],
+ )?;
+ self.dump("embedding", &s.x, false)?;
+ for i in 0..28 {
+ let p = format!("encoder.layers.{i}");
+ let w = |n: &str| self.w(&format!("{p}.{n}")).ptr();
+ call(
+ "qkv",
+ &[
+ s.y.ptr(),
+ w("attn.Wqkv.weight"),
+ self.w("zeros.3072").ptr(),
+ s.qkv.ptr(),
+ ],
+ )?;
+ let kind = if i % 3 == 0 { "full" } else { "local" };
+ call(
+ if original_rope {
+ "rope_original"
+ } else {
+ "rope"
+ },
+ &[
+ s.qkv.ptr(),
+ self.w(&format!("rope_{kind}_cos")).ptr(),
+ self.w(&format!("rope_{kind}_sin")).ptr(),
+ ],
+ )?;
+ call(
+ &attention(if i % 3 == 0 { "full" } else { "local" }),
+ &[s.qkv.ptr(), s.lens.ptr(), s.o.ptr()],
+ )?;
+ call("out", &[s.o.ptr(), w("attn.Wo.weight"), z, s.y.ptr()])?;
+ call(
+ "addln",
+ &[s.x.ptr(), s.y.ptr(), w("mlp_norm.weight"), z, s.y.ptr()],
+ )?;
+ call("geglu", &[s.y.ptr(), w("mlp.Wi.weight"), s.g.ptr()])?;
+ call("down", &[s.g.ptr(), w("mlp.Wo.weight"), z, s.y.ptr()])?;
+ let next = if i < 27 {
+ self.w(&format!("encoder.layers.{}.attn_norm.weight", i + 1))
+ } else {
+ self.w("encoder.final_norm.weight")
+ };
+ call("addln", &[s.x.ptr(), s.y.ptr(), next.ptr(), z, s.y.ptr()])?;
+ if [0, 1, 2, 27].contains(&i) {
+ self.dump(&format!("encoder{i}_residual"), &s.x, false)?;
+ self.dump(&format!("encoder{i}_normalized"), &s.y, true)?;
+ }
+ }
+ call(
+ "type",
+ &[
+ s.y.ptr(),
+ self.w("type_emb.weight").ptr(),
+ s.types.ptr(),
+ s.x.ptr(),
+ ],
+ )?;
+ for i in 0..2 {
+ let p = format!("head.layers.{i}");
+ let w = |n: &str| self.w(&format!("{p}.{n}")).ptr();
+ call(
+ "ln_bias",
+ &[
+ s.x.ptr(),
+ s.y.ptr(),
+ w("norm1.weight"),
+ w("norm1.bias"),
+ s.y.ptr(),
+ ],
+ )?;
+ call(
+ "head_in",
+ &[
+ s.y.ptr(),
+ w("self_attn.in_proj_weight"),
+ w("self_attn.in_proj_bias"),
+ s.qkv.ptr(),
+ ],
+ )?;
+ call(&attention("full"), &[s.qkv.ptr(), s.lens.ptr(), s.o.ptr()])?;
+ call(
+ "head_out",
+ &[
+ s.o.ptr(),
+ w("self_attn.out_proj.weight"),
+ w("self_attn.out_proj.bias"),
+ s.y.ptr(),
+ ],
+ )?;
+ call(
+ "addln_bias",
+ &[
+ s.x.ptr(),
+ s.y.ptr(),
+ w("norm2.weight"),
+ w("norm2.bias"),
+ s.y.ptr(),
+ ],
+ )?;
+ call(
+ "ffn1",
+ &[
+ s.y.ptr(),
+ w("linear1.weight"),
+ w("linear1.bias"),
+ s.ff.ptr(),
+ ],
+ )?;
+ call(
+ "ffn2",
+ &[
+ s.ff.ptr(),
+ w("linear2.weight"),
+ w("linear2.bias"),
+ s.y.ptr(),
+ ],
+ )?;
+ call("residual", &[s.x.ptr(), s.y.ptr()])?;
+ }
+ Ok(())
+ }
+}