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/src/lib.rs b/src/backends/cuda/src/lib.rs index 5343723c..9d4235e8 100644 --- a/src/backends/cuda/src/lib.rs +++ b/src/backends/cuda/src/lib.rs @@ -7,7 +7,7 @@ use std::{ 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, @@ -86,6 +86,15 @@ impl Cuda { 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 +105,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, @@ -151,7 +160,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, @@ -245,3 +254,31 @@ impl Drop for Graph { } } } + +/// 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..b6645ec7 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-handle lifetimes, exact dispatch traces and sparse wrappers. 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..5b404963 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, @@ -169,14 +169,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,159 +546,37 @@ 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}") + 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)?, } - }; - // 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)?; - } - } - 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(()) } diff --git a/tests/cuda/laya/README.md b/tests/cuda/laya/README.md new file mode 100644 index 00000000..2f4963d2 --- /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-handle lifetimes, exact dispatch traces and sparse wrappers. +They validate host behavior without Laya tensor math or a CUDA driver. 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/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..84676bb2 --- /dev/null +++ b/tests/cuda/laya_runtime.rs @@ -0,0 +1,6 @@ +//! CPU tests of production Rust bindings against a deterministic native ABI. +#![cfg(unix)] +#[path = "laya/fixture.rs"] +mod fixture; +#[path = "laya/resolved.rs"] +mod resolved; diff --git a/tests/laya/model.rs b/tests/laya/model.rs index 12dff82e..36e6b204 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,19 @@ 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 = "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/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(()) + } +}