From ebf4e8ac5fd26a17a4c8b2db6c3debc5abdf71e5 Mon Sep 17 00:00:00 2001 From: linear3735 Date: Wed, 30 Sep 2026 11:37:39 +0800 Subject: [PATCH 01/15] Add Laya CUDA resource ownership --- Cargo.lock | 63 +++++ Cargo.toml | 2 +- src/backends/cuda/Cargo.toml | 12 + src/backends/cuda/README.md | 50 +++- src/backends/cuda/kernels/runtime.cu | 65 ++++++ src/backends/cuda/src/lib.rs | 188 +++++++++++++++ src/backends/cuda/tests/fixtures/runtime.c | 113 +++++++++ src/backends/cuda/tests/runtime.rs | 254 +++++++++++++++++++++ 8 files changed, 743 insertions(+), 4 deletions(-) create mode 100644 src/backends/cuda/Cargo.toml create mode 100644 src/backends/cuda/kernels/runtime.cu create mode 100644 src/backends/cuda/src/lib.rs create mode 100644 src/backends/cuda/tests/fixtures/runtime.c create mode 100644 src/backends/cuda/tests/runtime.rs diff --git a/Cargo.lock b/Cargo.lock index be37964..92fc916 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2,6 +2,12 @@ # It is not intended for manual editing. version = 4 +[[package]] +name = "anyhow" +version = "1.0.104" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "330a5ed07fa54e4702c9d6c4174f74427fc0ef6e214bbd677ae50a5099946470" + [[package]] name = "atomic-waker" version = "1.1.2" @@ -153,6 +159,12 @@ dependencies = [ "windows-sys 0.61.2", ] +[[package]] +name = "fastrand" +version = "2.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "da7c62ceae207dd37ea5b845da6a0696c799f85e97da1ab5b7910be3c1c80223" + [[package]] name = "find-msvc-tools" version = "0.1.13" @@ -494,6 +506,22 @@ version = "0.2.189" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "3eaf3ede3fee6db1a4c2ee091bf8a8b4dccdc6d17f656fb07896ee72867612f2" +[[package]] +name = "libloading" +version = "0.8.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d7c4b02199fee7c5d21a5ae7d8cfa79a6ef5bb2fc834d6e9058e89c825efdc55" +dependencies = [ + "cfg-if", + "windows-link", +] + +[[package]] +name = "linux-raw-sys" +version = "0.12.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "32a66949e030da00e8c7d4434b251670a91556f4144941d37452769c25d58a53" + [[package]] name = "litemap" version = "0.8.3" @@ -541,6 +569,15 @@ dependencies = [ "windows-sys 0.61.2", ] +[[package]] +name = "omni-cuda" +version = "0.1.0" +dependencies = [ + "anyhow", + "libloading", + "tempfile", +] + [[package]] name = "omni-jev" version = "0.1.0" @@ -745,6 +782,19 @@ version = "2.1.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6b1e7f9a428571be2dc5bc0505c13fb6bf936822b894ec87abf8a08a4e51742d" +[[package]] +name = "rustix" +version = "1.1.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "891efababe418670775f199f0d233d84843c227a0949a883ce15b37c78d6629d" +dependencies = [ + "bitflags", + "errno", + "libc", + "linux-raw-sys", + "windows-sys 0.61.2", +] + [[package]] name = "rustls" version = "0.23.45" @@ -938,6 +988,19 @@ dependencies = [ "syn", ] +[[package]] +name = "tempfile" +version = "3.27.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "32497e9a4c7b38532efcdebeef879707aa9f794296a4f0244f6f69e9bc8574bd" +dependencies = [ + "fastrand", + "getrandom 0.4.3", + "once_cell", + "rustix", + "windows-sys 0.61.2", +] + [[package]] name = "thiserror" version = "2.0.21" diff --git a/Cargo.toml b/Cargo.toml index 8036e4b..1bf3a16 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -1,3 +1,3 @@ [workspace] -members = ["src/frontend"] +members = ["src/frontend", "src/backends/cuda"] resolver = "3" diff --git a/src/backends/cuda/Cargo.toml b/src/backends/cuda/Cargo.toml new file mode 100644 index 0000000..5d028b5 --- /dev/null +++ b/src/backends/cuda/Cargo.toml @@ -0,0 +1,12 @@ +[package] +name = "omni-cuda" +version = "0.1.0" +edition = "2024" +publish = false + +[dependencies] +anyhow = "1" +libloading = "0.8" + +[dev-dependencies] +tempfile = "3" diff --git a/src/backends/cuda/README.md b/src/backends/cuda/README.md index 1ba4c67..39161c3 100644 --- a/src/backends/cuda/README.md +++ b/src/backends/cuda/README.md @@ -1,7 +1,51 @@ # CUDA backend -Planned home for high-performance NVIDIA GPU operations and kernel integration. Implement the operations required by the first model, with hardware-specific optimizations where needed. +`omni-cuda` loads Laya's CUDA resource library at runtime. It owns one device and +stream per context, plus the buffers allocated through that context. Rust builds +and CPU tests need no CUDA toolkit. -Model orchestration, batching policy, state management, and kernel selection remain with the model engine. CUDA and Metal implementations do not need identical internal structures or a universal tensor abstraction. +This first slice covers allocation, copies, synchronization and cleanup. Model +initialization, weights, kernel calls, Graphs and hardware-specific optimizations +remain separate. It does not yet run Laya inference. -Status: planned; no CUDA implementation or validated hardware coverage yet. +## Build and check + +On a machine with the CUDA toolkit, build the resource library: + +```sh +nvcc -shared -Xcompiler=-fPIC -O2 src/backends/cuda/kernels/runtime.cu -o /tmp/liblaya_cuda.so +LAYA_CUDA_LIBRARY=/tmp/liblaya_cuda.so LAYA_CUDA_DEVICE=0 \ + cargo test --locked -p omni-cuda --test runtime -- --ignored +``` + +The device is an ordinal after `CUDA_VISIBLE_DEVICES` filtering. The library +contains no generated kernels and needs neither TileLang nor cuBLAS. This command +builds only the resource slice; the complete model bundle has a separate build. + +The normal CPU tests compile a small C fixture with `cc`. They check the dynamic +loader, errors, copy bounds and resource lifetime. They do not validate CUDA or +hardware support. The ignored test exercises real allocation and copy roundtrips. + +## Ownership and ABI + +Load only a trusted library with the matching ABI. `Cuda::load(path, device)` +checks `laya_abi_version() == 1` and all required symbols before creating a stream. +The old prototype's `laya_init` library has no version symbol and is rejected. + +`Cuda` and `Buffer` stay on their creating thread. A buffer keeps its stream and +library alive even after the caller drops `Cuda`. Operations select the owning +device before using its resources. Destruction attempts synchronization and +cleanup; call `sync()` explicitly when errors need to reach the caller. + +`write` and `read` check byte limits and synchronize before returning, so borrowed +host memory cannot outlive a queued copy. They are not Graph-capture operations. +Allocation of zero bytes is rejected; empty reads and writes are no-ops. + +The native resource entry points return zero on success and CUDA error codes on +failure; code 1000 means an invalid runtime argument. `laya_error_string` explains +the code. Upload and download take the caller's stream as their last argument and +do not synchronize internally. No Hopper requirement or model initialization is +hidden in stream creation. + +These are Laya's resource entry points, not a new shared tensor interface. A common +runtime can be extracted when another model needs the same implementation. diff --git a/src/backends/cuda/kernels/runtime.cu b/src/backends/cuda/kernels/runtime.cu new file mode 100644 index 0000000..da4d69c --- /dev/null +++ b/src/backends/cuda/kernels/runtime.cu @@ -0,0 +1,65 @@ +#include +#include + +namespace { +constexpr int invalid_argument = 1000; +} + +extern "C" { +uint32_t laya_abi_version() { return 1; } + +const char* laya_error_string(int code) { + return code == invalid_argument ? "invalid runtime argument" + : cudaGetErrorString(static_cast(code)); +} + +int laya_set_device(int device) { return cudaSetDevice(device); } + +int laya_stream_create(void** stream) { + if (!stream) return invalid_argument; + *stream = nullptr; + cudaStream_t created = nullptr; + cudaError_t status = cudaStreamCreateWithFlags(&created, cudaStreamNonBlocking); + if (status == cudaSuccess) *stream = created; + return status; +} + +int laya_alloc(void** p, size_t bytes) { + if (!p) return invalid_argument; + *p = nullptr; + if (!bytes) return invalid_argument; + void* allocated = nullptr; + cudaError_t status = cudaMalloc(&allocated, bytes); + if (status == cudaSuccess) *p = allocated; + return status; +} + +int laya_free(void* p) { + if (!p) return invalid_argument; + return cudaFree(p); +} + +int laya_upload(void* dst, const void* src, size_t bytes, void* stream) { + if (!bytes) return 0; + if (!dst || !src || !stream) return invalid_argument; + return cudaMemcpyAsync(dst, src, bytes, cudaMemcpyHostToDevice, + static_cast(stream)); +} + +int laya_download(void* dst, const void* src, size_t bytes, void* stream) { + if (!bytes) return 0; + if (!dst || !src || !stream) return invalid_argument; + return cudaMemcpyAsync(dst, src, bytes, cudaMemcpyDeviceToHost, + static_cast(stream)); +} + +int laya_sync(void* stream) { + if (!stream) return invalid_argument; + return cudaStreamSynchronize(static_cast(stream)); +} + +int laya_stream_free(void* stream) { + if (!stream) return invalid_argument; + return cudaStreamDestroy(static_cast(stream)); +} +} diff --git a/src/backends/cuda/src/lib.rs b/src/backends/cuda/src/lib.rs new file mode 100644 index 0000000..4338ade --- /dev/null +++ b/src/backends/cuda/src/lib.rs @@ -0,0 +1,188 @@ +//! CUDA resources confined to one thread. CPU builds do not link CUDA. +use anyhow::{Result, anyhow, ensure}; +use libloading::Library; +use std::{ + ffi::{CStr, c_char, c_void}, + path::Path, + rc::Rc, +}; + +type Ptr = *mut c_void; + +struct Functions { + error: unsafe extern "C" fn(i32) -> *const c_char, + set_device: unsafe extern "C" fn(i32) -> i32, + stream_create: unsafe extern "C" fn(*mut Ptr) -> i32, + alloc: unsafe extern "C" fn(*mut Ptr, usize) -> i32, + free: unsafe extern "C" fn(Ptr) -> i32, + upload: unsafe extern "C" fn(Ptr, *const u8, usize, Ptr) -> i32, + download: unsafe extern "C" fn(*mut u8, Ptr, usize, Ptr) -> i32, + sync: unsafe extern "C" fn(Ptr) -> i32, + stream_free: unsafe extern "C" fn(Ptr) -> i32, +} + +impl Functions { + fn check(&self, code: i32) -> Result<()> { + if code == 0 { + return Ok(()); + } + let message = unsafe { (self.error)(code) }; + if message.is_null() { + return Err(anyhow!("CUDA error {code}")); + } + Err(anyhow!("CUDA {code}: {}", unsafe { + CStr::from_ptr(message).to_string_lossy() + })) + } +} + +struct Context { + _library: Library, + functions: Functions, + device: i32, + stream: Ptr, +} + +impl Context { + fn activate(&self) -> Result<()> { + self.functions + .check(unsafe { (self.functions.set_device)(self.device) }) + } + + fn sync(&self) -> Result<()> { + self.activate()?; + self.functions + .check(unsafe { (self.functions.sync)(self.stream) }) + } +} + +impl Drop for Context { + fn drop(&mut self) { + if self.activate().is_ok() { + unsafe { + (self.functions.sync)(self.stream); + (self.functions.stream_free)(self.stream); + } + } + } +} + +#[derive(Clone)] +pub struct Cuda { + ctx: Rc, +} + +impl Cuda { + /// # Safety + /// `path` must name a trusted library implementing the complete runtime ABI. + pub unsafe fn load(path: &Path, device: i32) -> Result { + let library = unsafe { Library::new(path) }?; + let version = + unsafe { library.get:: u32>(b"laya_abi_version\0")?() }; + ensure!(version == 1, "unsupported CUDA runtime ABI {version}"); + let functions = unsafe { + Functions { + error: *library.get(b"laya_error_string\0")?, + set_device: *library.get(b"laya_set_device\0")?, + stream_create: *library.get(b"laya_stream_create\0")?, + alloc: *library.get(b"laya_alloc\0")?, + free: *library.get(b"laya_free\0")?, + upload: *library.get(b"laya_upload\0")?, + download: *library.get(b"laya_download\0")?, + sync: *library.get(b"laya_sync\0")?, + stream_free: *library.get(b"laya_stream_free\0")?, + } + }; + functions.check(unsafe { (functions.set_device)(device) })?; + let mut stream = std::ptr::null_mut(); + functions.check(unsafe { (functions.stream_create)(&mut stream) })?; + ensure!(!stream.is_null(), "CUDA runtime returned a null stream"); + Ok(Self { + ctx: Rc::new(Context { + _library: library, + functions, + device, + stream, + }), + }) + } + + pub fn alloc(&self, bytes: usize) -> Result { + ensure!(bytes > 0, "zero CUDA allocation"); + self.ctx.activate()?; + let mut ptr = std::ptr::null_mut(); + self.ctx + .functions + .check(unsafe { (self.ctx.functions.alloc)(&mut ptr, bytes) })?; + ensure!(!ptr.is_null(), "CUDA runtime returned a null allocation"); + Ok(Buffer { + ctx: self.ctx.clone(), + ptr, + bytes, + }) + } + + pub fn upload(&self, bytes: &[u8]) -> Result { + let buffer = self.alloc(bytes.len())?; + buffer.write(bytes)?; + Ok(buffer) + } + + pub fn sync(&self) -> Result<()> { + self.ctx.sync() + } +} + +pub struct Buffer { + ctx: Rc, + ptr: Ptr, + bytes: usize, +} + +impl Buffer { + pub fn bytes(&self) -> usize { + self.bytes + } + + pub fn write(&self, bytes: &[u8]) -> Result<()> { + ensure!(bytes.len() <= self.bytes, "upload exceeds allocation"); + if bytes.is_empty() { + return Ok(()); + } + self.ctx.activate()?; + let functions = &self.ctx.functions; + let copied = + unsafe { (functions.upload)(self.ptr, bytes.as_ptr(), bytes.len(), self.ctx.stream) }; + // Even a failed copy may have queued work using the borrowed host memory. + let synced = unsafe { (functions.sync)(self.ctx.stream) }; + functions.check(copied)?; + functions.check(synced) + } + + pub fn read(&self, bytes: usize) -> Result> { + ensure!(bytes <= self.bytes, "download exceeds allocation"); + let mut data = vec![0; bytes]; + if bytes == 0 { + return Ok(data); + } + self.ctx.activate()?; + let functions = &self.ctx.functions; + let copied = + unsafe { (functions.download)(data.as_mut_ptr(), self.ptr, bytes, self.ctx.stream) }; + let synced = unsafe { (functions.sync)(self.ctx.stream) }; + functions.check(copied)?; + functions.check(synced)?; + Ok(data) + } +} + +impl Drop for Buffer { + fn drop(&mut self) { + if self.ctx.activate().is_ok() { + unsafe { + (self.ctx.functions.sync)(self.ctx.stream); + (self.ctx.functions.free)(self.ptr); + } + } + } +} diff --git a/src/backends/cuda/tests/fixtures/runtime.c b/src/backends/cuda/tests/fixtures/runtime.c new file mode 100644 index 0000000..ff26bb4 --- /dev/null +++ b/src/backends/cuda/tests/fixtures/runtime.c @@ -0,0 +1,113 @@ +#include +#include +#include +#include +#include + +#ifndef LAYA_TEST_ABI +#define LAYA_TEST_ABI 1 +#endif + +enum { + invalid_argument = 1000, + mode_create_error = 1, + mode_copy_error = 2, + mode_sync_error = 3, + mode_copy_and_sync_error = 4, + mode_alloc_error = 5 +}; + +typedef struct { int device, live; } Stream; +typedef struct { int device; unsigned char data[]; } Allocation; +static Stream streams[16]; +static int device = -1, mode, next_stream, live; +static char trace[4096], error[64]; +static size_t trace_len, pending_bytes; +static void *pending_dst; +static const void *pending_src; + +static void record(const char *name) { + trace_len += (size_t)snprintf(trace + trace_len, sizeof(trace) - trace_len, + "%s%d ", name, device); +} +static Allocation *allocation(void *p) { + return (Allocation *)((unsigned char *)p - offsetof(Allocation, data)); +} +static int valid_stream(void *p) { + Stream *s = p; + return s && s->live && s->device == device; +} +void laya_test_mode(int value) { mode = value; } +const char *laya_test_trace(void) { return trace; } +int laya_test_live(void) { return live; } +uint32_t laya_abi_version(void) { return LAYA_TEST_ABI; } +const char *laya_error_string(int code) { + snprintf(error, sizeof(error), "fixture-error-%d", code); + return error; +} +int laya_set_device(int value) { + if (value < 0) return 11; + device = value; + record("device"); + return 0; +} +int laya_stream_create(void **out) { + record("create"); + if (mode == mode_create_error) return 23; + Stream *s = &streams[next_stream++]; + *s = (Stream){device, 1}; + *out = s; + live++; + return 0; +} +int laya_alloc(void **out, size_t bytes) { + record("alloc"); + if (mode == mode_alloc_error) return 31; + Allocation *a = calloc(1, sizeof(*a) + bytes); + if (!a) return 32; + a->device = device; + *out = a->data; + live++; + return 0; +} +int laya_free(void *p) { + if (!p) return invalid_argument; + record("free"); + if (allocation(p)->device != device || pending_bytes) return 91; + free(allocation(p)); + live--; + return 0; +} +static int copy(void *dst, const void *src, size_t bytes, void *stream) { + if (!valid_stream(stream)) return 91; + pending_dst = dst; + pending_src = src; + pending_bytes = bytes; + return mode == mode_copy_error || mode == mode_copy_and_sync_error ? 41 : 0; +} +int laya_upload(void *dst, const unsigned char *src, size_t bytes, void *stream) { + record("upload"); + if (allocation(dst)->device != device) return 91; + return copy(dst, src, bytes, stream); +} +#ifndef LAYA_TEST_NO_DOWNLOAD +int laya_download(unsigned char *dst, void *src, size_t bytes, void *stream) { + record("download"); + if (allocation(src)->device != device) return 91; + return copy(dst, src, bytes, stream); +} +#endif +int laya_sync(void *stream) { + record("sync"); + if (!valid_stream(stream)) return 91; + if (pending_bytes) memcpy(pending_dst, pending_src, pending_bytes); + pending_bytes = 0; + return mode == mode_sync_error || mode == mode_copy_and_sync_error ? 42 : 0; +} +int laya_stream_free(void *stream) { + record("destroy"); + if (!valid_stream(stream) || pending_bytes) return 91; + ((Stream *)stream)->live = 0; + live--; + return 0; +} diff --git a/src/backends/cuda/tests/runtime.rs b/src/backends/cuda/tests/runtime.rs new file mode 100644 index 0000000..1101021 --- /dev/null +++ b/src/backends/cuda/tests/runtime.rs @@ -0,0 +1,254 @@ +#![cfg(unix)] + +use libloading::Library; +use omni_cuda::Cuda; +use std::{ffi::CStr, path::PathBuf, process::Command}; +use tempfile::{TempDir, tempdir}; + +const MODE_NORMAL: i32 = 0; +const MODE_CREATE_ERROR: i32 = 1; +const MODE_COPY_ERROR: i32 = 2; +const MODE_SYNC_ERROR: i32 = 3; +const MODE_COPY_AND_SYNC_ERROR: i32 = 4; +const MODE_ALLOC_ERROR: i32 = 5; + +struct Fixture { + _dir: TempDir, + path: PathBuf, + library: Option, +} +impl Fixture { + fn new(defines: &[&str]) -> Self { + let dir = tempdir().unwrap(); + let path = dir + .path() + .join(format!("runtime{}", std::env::consts::DLL_SUFFIX)); + let mut command = Command::new("cc"); + command.arg(if cfg!(target_os = "macos") { + "-dynamiclib" + } else { + "-shared" + }); + let result = command + .args(["-fPIC", "-std=c11", "-Wall", "-Wextra", "-Werror"]) + .args(defines) + .arg(PathBuf::from(env!("CARGO_MANIFEST_DIR")).join("tests/fixtures/runtime.c")) + .arg("-o") + .arg(&path) + .output() + .unwrap(); + assert!( + result.status.success(), + "{}", + String::from_utf8_lossy(&result.stderr) + ); + let library = Some(unsafe { Library::new(&path) }.unwrap()); + Self { + _dir: dir, + path, + library, + } + } + fn mode(&self, mode: i32) { + unsafe { + self.library + .as_ref() + .unwrap() + .get::(b"laya_test_mode\0") + .unwrap()(mode); + } + } + fn trace(&self) -> String { + unsafe { + let f = self + .library + .as_ref() + .unwrap() + .get:: *const std::ffi::c_char>(b"laya_test_trace\0") + .unwrap(); + CStr::from_ptr(f()).to_string_lossy().into_owned() + } + } + fn live(&self) -> i32 { + unsafe { + self.library + .as_ref() + .unwrap() + .get:: i32>(b"laya_test_live\0") + .unwrap()() + } + } + fn load(&self, device: i32) -> anyhow::Result { + unsafe { Cuda::load(&self.path, device) } + } +} + +#[test] +fn loading_rejects_incompatible_libraries_before_creating_resources() { + for defines in [vec!["-DLAYA_TEST_ABI=2"], vec!["-DLAYA_TEST_NO_DOWNLOAD"]] { + let fixture = Fixture::new(&defines); + assert!(fixture.load(0).is_err()); + assert_eq!(fixture.live(), 0); + assert!(!fixture.trace().contains("create")); + } + let fixture = Fixture::new(&[]); + assert!(fixture.load(-1).is_err()); + fixture.mode(MODE_CREATE_ERROR); + let error = fixture.load(0).err().unwrap().to_string(); + assert!(error.contains("23"), "{error}"); + assert_eq!(fixture.live(), 0); + assert!(!fixture.trace().contains("destroy")); +} + +#[test] +fn copies_round_trip_and_bounds_are_checked() { + let fixture = Fixture::new(&[]); + let cuda = fixture.load(0).unwrap(); + assert!(cuda.alloc(0).is_err()); + let buffer = cuda.upload(&[1, 2, 3, 4]).unwrap(); + assert_eq!(buffer.bytes(), 4); + assert_eq!(buffer.read(4).unwrap(), [1, 2, 3, 4]); + buffer.write(&[9, 8]).unwrap(); + assert_eq!(buffer.read(4).unwrap(), [9, 8, 3, 4]); + buffer.write(&[]).unwrap(); + assert!(buffer.read(0).unwrap().is_empty()); + let before = fixture.trace(); + assert!(buffer.write(&[0; 5]).is_err()); + assert!(buffer.read(5).is_err()); + assert_eq!(fixture.trace(), before); + drop(buffer); + drop(cuda); + assert_eq!(fixture.live(), 0); +} + +#[test] +fn copy_errors_still_synchronize_and_async_errors_are_reported() { + let fixture = Fixture::new(&[]); + let cuda = fixture.load(0).unwrap(); + let buffer = cuda.alloc(3).unwrap(); + fixture.mode(MODE_COPY_ERROR); + let error = buffer.write(&[4, 5, 6]).unwrap_err().to_string(); + assert!(error.contains("fixture-error-41"), "{error}"); + assert!(fixture.trace().contains("upload0 sync0 ")); + fixture.mode(MODE_NORMAL); + assert_eq!(buffer.read(3).unwrap(), [4, 5, 6]); + fixture.mode(MODE_COPY_ERROR); + assert!(buffer.read(3).unwrap_err().to_string().contains("41")); + assert!(fixture.trace().contains("download0 sync0 ")); + fixture.mode(MODE_SYNC_ERROR); + assert!( + buffer + .write(&[7, 8, 9]) + .unwrap_err() + .to_string() + .contains("42") + ); + assert!(buffer.read(3).unwrap_err().to_string().contains("42")); + assert!(cuda.sync().unwrap_err().to_string().contains("42")); + fixture.mode(MODE_COPY_AND_SYNC_ERROR); + assert!( + buffer + .write(&[1, 2, 3]) + .unwrap_err() + .to_string() + .contains("41") + ); + fixture.mode(MODE_NORMAL); + assert_eq!(buffer.read(3).unwrap(), [1, 2, 3]); + drop(buffer); + drop(cuda); + assert_eq!(fixture.live(), 0); +} + +#[test] +fn failed_allocations_and_uploads_release_partial_resources() { + let fixture = Fixture::new(&[]); + let cuda = fixture.load(0).unwrap(); + fixture.mode(MODE_ALLOC_ERROR); + assert!(cuda.alloc(4).err().unwrap().to_string().contains("31")); + assert_eq!(fixture.live(), 1); + fixture.mode(MODE_COPY_ERROR); + assert!( + cuda.upload(&[1, 2]) + .err() + .unwrap() + .to_string() + .contains("41") + ); + assert_eq!(fixture.live(), 1); + fixture.mode(MODE_NORMAL); + drop(cuda); + assert_eq!(fixture.live(), 0); +} + +#[test] +fn buffers_keep_the_library_and_stream_alive_after_cuda_is_dropped() { + let mut fixture = Fixture::new(&[]); + let cuda = fixture.load(0).unwrap(); + let cloned = cuda.clone(); + let buffer = cuda.upload(&[7, 8, 9]).unwrap(); + // Remove the test's dlopen handle too: only the buffer may keep this library alive. + drop(fixture.library.take()); + drop(cuda); + drop(cloned); + assert_eq!(buffer.read(3).unwrap(), [7, 8, 9]); + fixture.library = Some(unsafe { Library::new(&fixture.path) }.unwrap()); + assert_eq!(fixture.live(), 2); + assert!(!fixture.trace().contains("destroy")); + drop(buffer); + assert_eq!(fixture.live(), 0); + let trace = fixture.trace(); + assert!(trace.rfind("sync0 ").unwrap() < trace.rfind("destroy0 ").unwrap()); + assert!(trace.rfind("free0 ").unwrap() < trace.rfind("destroy0 ").unwrap()); +} + +#[test] +fn operations_and_drops_restore_the_owning_device() { + let fixture = Fixture::new(&[]); + let cuda0 = fixture.load(0).unwrap(); + let buffer0 = cuda0.upload(&[1, 2]).unwrap(); + let cuda1 = fixture.load(1).unwrap(); + let buffer1 = cuda1.upload(&[3, 4]).unwrap(); + assert_eq!(buffer0.read(2).unwrap(), [1, 2]); + buffer1.write(&[5, 6]).unwrap(); + cuda0.sync().unwrap(); + drop(buffer1); + drop(cuda1); + drop(buffer0); + drop(cuda0); + assert_eq!(fixture.live(), 0); + let trace = fixture.trace(); + for event in [ + "device0 download0", + "device1 upload1", + "device0 sync0", + "device1 sync1 free1", + "device0 sync0 free0", + ] { + assert!(trace.contains(event), "missing {event}: {trace}"); + } +} + +#[test] +#[ignore = "requires an approved GPU and LAYA_CUDA_LIBRARY plus LAYA_CUDA_DEVICE"] +fn real_gpu_round_trip() { + let path = PathBuf::from(std::env::var_os("LAYA_CUDA_LIBRARY").expect("set LAYA_CUDA_LIBRARY")); + let device = std::env::var("LAYA_CUDA_DEVICE") + .expect("set LAYA_CUDA_DEVICE") + .parse() + .unwrap(); + { + let library = unsafe { Library::new(&path) }.unwrap(); + let free = unsafe { + library.get:: i32>(b"laya_free\0") + } + .unwrap(); + assert_eq!(unsafe { free(std::ptr::null_mut()) }, 1000); + } + let cuda = unsafe { Cuda::load(&path, device) }.unwrap(); + let expected: Vec = (0..4096).map(|i| (i % 251) as u8).collect(); + let buffer = cuda.upload(&expected).unwrap(); + assert_eq!(buffer.read(expected.len()).unwrap(), expected); + drop(cuda); + assert_eq!(buffer.read(expected.len()).unwrap(), expected); +} From ff8d0170f4fb06394fae35f4356eca212415bf31 Mon Sep 17 00:00:00 2001 From: linear3735 Date: Wed, 30 Sep 2026 11:42:11 +0800 Subject: [PATCH 02/15] Keep Laya checkpoint weights resident on CUDA --- Cargo.lock | 2 + src/models/laya/Cargo.toml | 2 + src/models/laya/README.md | 17 +++ src/models/laya/src/lib.rs | 1 + src/models/laya/src/resident.rs | 84 +++++++++++ src/models/laya/src/resident/tests.rs | 199 ++++++++++++++++++++++++++ 6 files changed, 305 insertions(+) create mode 100644 src/models/laya/src/resident.rs create mode 100644 src/models/laya/src/resident/tests.rs diff --git a/Cargo.lock b/Cargo.lock index be3de8d..ed4b2ca 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -668,7 +668,9 @@ version = "0.1.0" dependencies = [ "anyhow", "half", + "libloading", "memmap2", + "omni-cuda", "safetensors", "serde", "serde_json", diff --git a/src/models/laya/Cargo.toml b/src/models/laya/Cargo.toml index d9e73e8..69ce281 100644 --- a/src/models/laya/Cargo.toml +++ b/src/models/laya/Cargo.toml @@ -8,10 +8,12 @@ publish = false anyhow = "1" half = "2" memmap2 = "0.9" +omni-cuda = { path = "../../backends/cuda" } safetensors = "0.6" serde = { version = "1", features = ["derive"] } serde_json = "1" [dev-dependencies] +libloading = "0.8" sha2 = "0.10" tempfile = "3" diff --git a/src/models/laya/README.md b/src/models/laya/README.md index 858da19..147979d 100644 --- a/src/models/laya/README.md +++ b/src/models/laya/README.md @@ -22,3 +22,20 @@ cargo test --release --locked -p omni-laya --test weights -- --ignored ``` These two CPU tests check all 206 tensor names and shapes, 618 conversion hashes, and the legacy temperature buffer. The normal CI job skips them because it does not download the full checkpoint. + +## GPU weight residency + +`ResidentWeights::upload(&cuda, &weights)` validates the checkpoint inventory and +uploads the weights once. Embeddings use FP16; encoder norms, head norms and biases, +and the scorer input norm use FP32; other weights use BF16. Layouts stay unchanged. +The legacy `temperature` buffer is validated but not uploaded. + +`get(name)` returns the resident buffer. `bytes()` reports weight allocations only, +excluding CUDA context and allocator overhead. Buffers keep their CUDA context alive +after the caller drops the source mapping or `Cuda`. Failed loads release partial +allocations. Workspace, rotary tables and inference are separate modules. + +The ignored `real_checkpoint_residency_matches_torch` test requires +`LAYA_CUDA_LIBRARY`, `LAYA_CUDA_DEVICE`, `LAYA_CHECKPOINT` and `LAYA_WEIGHT_ORACLE`. +It uploads all 205 used tensors and compares readback hashes with the existing Torch +oracle; it does not test model outputs or latency. diff --git a/src/models/laya/src/lib.rs b/src/models/laya/src/lib.rs index fbc136f..315c26c 100644 --- a/src/models/laya/src/lib.rs +++ b/src/models/laya/src/lib.rs @@ -1,2 +1,3 @@ pub mod config; +pub mod resident; pub mod weights; diff --git a/src/models/laya/src/resident.rs b/src/models/laya/src/resident.rs new file mode 100644 index 0000000..b46b867 --- /dev/null +++ b/src/models/laya/src/resident.rs @@ -0,0 +1,84 @@ +use crate::weights::{TensorSpec, Weights, checkpoint_tensors}; +use anyhow::{Context, Result}; +use omni_cuda::{Buffer, Cuda}; +use std::collections::HashMap; + +/// Checkpoint weights owned by one CUDA context; no inference workspace or tables. +pub struct ResidentWeights { + buffers: HashMap, + bytes: usize, +} + +impl ResidentWeights { + pub fn upload(cuda: &Cuda, source: &Weights) -> Result { + let tensors = checkpoint_tensors(); + source.validate_names(tensors.iter().map(|t| t.name.as_str()))?; + // Account for the legacy buffer, but do not keep unused calibration on GPU. + source.f32("temperature", &[3])?; + Self::upload_tensors( + cuda, + source, + tensors.iter().filter(|t| t.name != "temperature"), + ) + } + + fn upload_tensors<'a>( + cuda: &Cuda, + source: &Weights, + tensors: impl IntoIterator, + ) -> Result { + let mut resident = Self { + buffers: HashMap::new(), + bytes: 0, + }; + for spec in tensors { + let data = packed(source, spec)?; + let buffer = cuda + .upload(&data) + .with_context(|| format!("upload {}", spec.name))?; + resident.bytes += buffer.bytes(); + resident.buffers.insert(spec.name.clone(), buffer); + } + Ok(resident) + } + + pub fn get(&self, name: &str) -> Result<&Buffer> { + self.buffers + .get(name) + .with_context(|| format!("no resident weight: {name}")) + } + + /// Weight allocation bytes, excluding CUDA context and allocator overhead. + pub fn bytes(&self) -> usize { + self.bytes + } +} + +fn packed(source: &Weights, spec: &TensorSpec) -> Result> { + let name = spec.name.as_str(); + let shape = &spec.shape; + if name == "encoder.embeddings.tok_embeddings.weight" { + Ok(source + .f16(name, shape)? + .into_iter() + .flat_map(u16::to_le_bytes) + .collect()) + } else if (shape.len() == 1 && (name.starts_with("encoder.") || name.starts_with("head."))) + || name.starts_with("scorer.0.") + { + Ok(source + .f32(name, shape)? + .into_iter() + .flat_map(f32::to_le_bytes) + .collect()) + } else { + Ok(source + .bf16(name, shape)? + .into_iter() + .flat_map(u16::to_le_bytes) + .collect()) + } +} + +#[cfg(test)] +mod tests; diff --git a/src/models/laya/src/resident/tests.rs b/src/models/laya/src/resident/tests.rs new file mode 100644 index 0000000..b77d42e --- /dev/null +++ b/src/models/laya/src/resident/tests.rs @@ -0,0 +1,199 @@ +use super::*; +use libloading::Library; +use safetensors::{Dtype, tensor::TensorView}; +use sha2::{Digest, Sha256}; +use std::{fs, path::PathBuf, process::Command}; +use tempfile::{TempDir, tempdir}; + +const EMBED: &str = "encoder.embeddings.tok_embeddings.weight"; + +struct Fixture { + dir: TempDir, + library: Library, + cuda: Cuda, +} + +impl Fixture { + fn new() -> Self { + let dir = tempdir().unwrap(); + let path = dir + .path() + .join(format!("runtime{}", std::env::consts::DLL_SUFFIX)); + let output = Command::new("cc") + .args([ + if cfg!(target_os = "macos") { + "-dynamiclib" + } else { + "-shared" + }, + "-fPIC", + ]) + .arg( + PathBuf::from(env!("CARGO_MANIFEST_DIR")) + .join("../../backends/cuda/tests/fixtures/runtime.c"), + ) + .arg("-o") + .arg(&path) + .output() + .unwrap(); + assert!( + output.status.success(), + "{}", + String::from_utf8_lossy(&output.stderr) + ); + let library = unsafe { Library::new(&path) }.unwrap(); + let cuda = unsafe { Cuda::load(&path, 0) }.unwrap(); + Self { dir, library, cuda } + } + + fn live(&self) -> i32 { + unsafe { + self.library + .get:: i32>(b"laya_test_live\0") + .unwrap()() + } + } + + fn source(&self, specs: &[TensorSpec]) -> Weights { + let values: Vec = [1.0f32, -2.0] + .into_iter() + .flat_map(f32::to_le_bytes) + .collect(); + let tensors: Vec<_> = specs + .iter() + .map(|s| { + ( + s.name.as_str(), + TensorView::new(Dtype::F32, s.shape.clone(), &values).unwrap(), + ) + }) + .collect(); + let path = self.dir.path().join("weights.safetensors"); + safetensors::tensor::serialize_to_file(tensors, None, &path).unwrap(); + Weights::open(&path).unwrap() + } +} + +fn spec(name: &str, shape: &[usize]) -> TensorSpec { + TensorSpec { + name: name.into(), + shape: shape.into(), + } +} + +#[test] +fn precision_and_residency_preserve_expected_bytes() { + let fixture = Fixture::new(); + let specs = [ + spec(EMBED, &[1, 2]), + spec("encoder.layers.0.mlp_norm.weight", &[2]), + spec("encoder.layers.0.attn.Wqkv.weight", &[1, 2]), + spec("head.layers.0.norm1.weight", &[2]), + spec("head.layers.0.self_attn.in_proj_bias", &[2]), + spec("head.layers.0.self_attn.in_proj_weight", &[1, 2]), + spec("scorer.0.bias", &[2]), + spec("scorer.1.bias", &[2]), + spec("act_head.0.bias", &[2]), + spec("type_emb.weight", &[1, 2]), + ]; + let source = fixture.source(&specs); + let weights = ResidentWeights::upload_tensors(&fixture.cuda, &source, &specs).unwrap(); + drop(source); + let f16 = vec![0x00, 0x3c, 0x00, 0xc0]; + let bf16 = vec![0x80, 0x3f, 0x00, 0xc0]; + let f32 = vec![0x00, 0x00, 0x80, 0x3f, 0x00, 0x00, 0x00, 0xc0]; + let expected = [ + &f16, &f32, &bf16, &f32, &f32, &bf16, &f32, &bf16, &bf16, &bf16, + ]; + for (s, bytes) in specs.iter().zip(expected) { + let buffer = weights.get(&s.name).unwrap(); + assert_eq!(&buffer.read(buffer.bytes()).unwrap(), bytes, "{}", s.name); + } + assert_eq!(weights.bytes(), 56); + assert!(weights.get("missing").is_err()); + assert_eq!(fixture.live(), 11); + drop(fixture.cuda); + assert_eq!(weights.get(EMBED).unwrap().read(4).unwrap(), f16); + drop(weights); + assert_eq!( + unsafe { + fixture + .library + .get:: i32>(b"laya_test_live\0") + .unwrap()() + }, + 0 + ); +} + +#[test] +fn invalid_checkpoint_and_partial_load_release_allocations() { + let fixture = Fixture::new(); + let mut specs = [ + spec(EMBED, &[1, 2]), + spec("head.layers.0.norm1.weight", &[2]), + ]; + let source = fixture.source(&specs); + assert!(ResidentWeights::upload(&fixture.cuda, &source).is_err()); + assert_eq!(fixture.live(), 1); + specs[1].shape = vec![3]; + let error = ResidentWeights::upload_tensors(&fixture.cuda, &source, &specs) + .err() + .unwrap(); + assert!(error.to_string().contains("head.layers.0.norm1.weight")); + assert_eq!(fixture.live(), 1); + specs[1].shape = vec![2]; + unsafe { + fixture + .library + .get::(b"laya_test_mode\0") + .unwrap()(2); + } + assert!(ResidentWeights::upload_tensors(&fixture.cuda, &source, &specs).is_err()); + assert_eq!(fixture.live(), 1); +} + +#[test] +#[ignore = "requires approved GPU, LAYA_CUDA_LIBRARY, LAYA_CUDA_DEVICE, LAYA_CHECKPOINT and LAYA_WEIGHT_ORACLE"] +fn real_checkpoint_residency_matches_torch() { + let library = PathBuf::from(std::env::var_os("LAYA_CUDA_LIBRARY").unwrap()); + let device = std::env::var("LAYA_CUDA_DEVICE").unwrap().parse().unwrap(); + let checkpoint = PathBuf::from(std::env::var_os("LAYA_CHECKPOINT").unwrap()); + let rows: Vec = + serde_json::from_slice(&fs::read(std::env::var_os("LAYA_WEIGHT_ORACLE").unwrap()).unwrap()) + .unwrap(); + let cuda = unsafe { Cuda::load(&library, device) }.unwrap(); + let source = Weights::open(&checkpoint.join("model.safetensors")).unwrap(); + let weights = ResidentWeights::upload(&cuda, &source).unwrap(); + drop(source); + drop(cuda); + assert_eq!(rows.len(), 206); + assert_eq!(weights.buffers.len(), 205); + assert!(weights.get("temperature").is_err()); + let mut names = std::collections::HashSet::new(); + let mut bytes = 0; + for row in rows { + let name = row["name"].as_str().unwrap(); + assert!(names.insert(name.to_owned())); + if name == "temperature" { + continue; + } + let buffer = weights.get(name).unwrap(); + let hash = format!("{:x}", Sha256::digest(buffer.read(buffer.bytes()).unwrap())); + // The oracle contains independent Torch conversions at all three precisions. + let shape = row["shape"].as_array().unwrap(); + let dtype = if name == EMBED { + "f16" + } else if name.starts_with("scorer.0.") + || (shape.len() == 1 && (name.starts_with("encoder.") || name.starts_with("head."))) + { + "f32" + } else { + "bf16" + }; + assert_eq!(hash, row[dtype].as_str().unwrap(), "{name} {dtype}"); + bytes += buffer.bytes(); + } + assert_eq!(weights.bytes(), bytes); + println!("verified_resident_tensors=205 weight_allocation_bytes={bytes}"); +} From 4b2c0d46e9ac2775d7578ce4fea786b44128de57 Mon Sep 17 00:00:00 2001 From: linear3735 Date: Wed, 30 Sep 2026 12:13:30 +0800 Subject: [PATCH 03/15] Allocate bounded Laya inference workspaces --- src/backends/cuda/tests/fixtures/runtime.c | 5 +- src/models/laya/README.md | 31 +++ src/models/laya/src/lib.rs | 1 + src/models/laya/src/workspace.rs | 100 ++++++++++ src/models/laya/src/workspace/tests.rs | 216 +++++++++++++++++++++ 5 files changed, 352 insertions(+), 1 deletion(-) create mode 100644 src/models/laya/src/workspace.rs create mode 100644 src/models/laya/src/workspace/tests.rs diff --git a/src/backends/cuda/tests/fixtures/runtime.c b/src/backends/cuda/tests/fixtures/runtime.c index ff26bb4..b37eb61 100644 --- a/src/backends/cuda/tests/fixtures/runtime.c +++ b/src/backends/cuda/tests/fixtures/runtime.c @@ -21,6 +21,7 @@ typedef struct { int device, live; } Stream; typedef struct { int device; unsigned char data[]; } Allocation; static Stream streams[16]; static int device = -1, mode, next_stream, live; +static int allocations_before_failure = -1; static char trace[4096], error[64]; static size_t trace_len, pending_bytes; static void *pending_dst; @@ -38,6 +39,7 @@ static int valid_stream(void *p) { return s && s->live && s->device == device; } void laya_test_mode(int value) { mode = value; } +void laya_test_fail_alloc_after(int count) { allocations_before_failure = count; } const char *laya_test_trace(void) { return trace; } int laya_test_live(void) { return live; } uint32_t laya_abi_version(void) { return LAYA_TEST_ABI; } @@ -62,7 +64,8 @@ int laya_stream_create(void **out) { } int laya_alloc(void **out, size_t bytes) { record("alloc"); - if (mode == mode_alloc_error) return 31; + if (mode == mode_alloc_error || allocations_before_failure == 0) return 31; + if (allocations_before_failure > 0) allocations_before_failure--; Allocation *a = calloc(1, sizeof(*a) + bytes); if (!a) return 32; a->device = device; diff --git a/src/models/laya/README.md b/src/models/laya/README.md index 147979d..65a410c 100644 --- a/src/models/laya/README.md +++ b/src/models/laya/README.md @@ -39,3 +39,34 @@ The ignored `real_checkpoint_residency_matches_torch` test requires `LAYA_CUDA_LIBRARY`, `LAYA_CUDA_DEVICE`, `LAYA_CHECKPOINT` and `LAYA_WEIGHT_ORACLE`. It uploads all 205 used tensors and compares readback hashes with the existing Torch oracle; it does not test model outputs or latency. + +## Inference workspace + +`Workspace::new(&cuda, batch, sequence)` allocates fixed scratch buffers for one +shape. Batch must be 1, 2, 4, 8 or 16; sequence must be a multiple of 16 in 16..=512. +Invalid shapes fail before any allocation; a failed allocation releases the partial +workspace. Contents are uninitialized and must be written before use. + +`buffers()` borrows the named buffers without allowing allocations to be replaced. +The workspace outlives the caller's `Cuda` handle. `bytes()` reports scratch +allocations only, excluding resident weights and CUDA overhead. + +Let `B` be batch, `L` sequence, `D=1024`, and `M=MAX_MARKERS=2048`. Layouts are: + +| Buffers | Shape and dtype | +| --- | --- | +| ids / lengths / types | `[B,L]` int64 / `[B]` int32 / `[B]` int64 | +| residual / hidden / attention | `[B,L,D]` FP32 / BF16 / BF16 | +| qkv / gated / feed_forward | `[B,L,3D]` / `[B,L,2624]` / `[B,L,4096]`, BF16 | +| indices / offsets | `[M]` / `[B+1]`, int32 | +| markers / scored / logits | `[M,D]` / `[M,D]` / `[M]`, BF16 | +| features / action_hidden / actions | `[B,1028]` / `[B,256]` / `[B,2]`, BF16 | + +At `(B,L)=(1,512)`, allocations total 22,628,896 bytes; at `(16,512)`, +236,048,836 bytes. The caller must enforce at most `MAX_MARKERS` scored positions. +No Graph cache, kernel launch, cuBLAS workspace or inference is included. + +The ignored `real_gpu_workspace_capacity_and_reuse` test requires +`LAYA_CUDA_LIBRARY` and `LAYA_CUDA_DEVICE`. It writes and reads all 17 buffers twice +at three shapes, including both capacity bounds, using deterministic byte patterns. +This tests allocation and transfer, not model numerics or latency. diff --git a/src/models/laya/src/lib.rs b/src/models/laya/src/lib.rs index 315c26c..a4be935 100644 --- a/src/models/laya/src/lib.rs +++ b/src/models/laya/src/lib.rs @@ -1,3 +1,4 @@ pub mod config; pub mod resident; pub mod weights; +pub mod workspace; diff --git a/src/models/laya/src/workspace.rs b/src/models/laya/src/workspace.rs new file mode 100644 index 0000000..1d2a1e9 --- /dev/null +++ b/src/models/laya/src/workspace.rs @@ -0,0 +1,100 @@ +use anyhow::{Result, ensure}; +use omni_cuda::{Buffer, Cuda}; + +pub const MAX_MARKERS: usize = 2048; +const D: usize = 1024; + +/// Fixed-shape scratch allocations. Contents are uninitialized until written. +pub struct Workspace { + batch: usize, + sequence: usize, + bytes: usize, + buffers: WorkspaceBuffers, +} + +/// Buffer layouts consumed by Laya's encoder, decision head and scorer. +/// Access through `Workspace::buffers` keeps allocations fixed for its lifetime. +pub struct WorkspaceBuffers { + pub ids: Buffer, + pub lengths: Buffer, + pub types: Buffer, + pub residual: Buffer, + pub hidden: Buffer, + pub qkv: Buffer, + pub attention: Buffer, + pub gated: Buffer, + pub feed_forward: Buffer, + pub indices: Buffer, + pub offsets: Buffer, + pub markers: Buffer, + pub scored: Buffer, + pub logits: Buffer, + pub features: Buffer, + pub action_hidden: Buffer, + pub actions: Buffer, +} + +impl Workspace { + pub fn new(cuda: &Cuda, batch: usize, sequence: usize) -> Result { + ensure!( + batch.is_power_of_two() && batch <= 16, + "workspace batch must be 1, 2, 4, 8 or 16" + ); + ensure!( + (16..=512).contains(&sequence) && sequence.is_multiple_of(16), + "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) + }; + 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)?, + }; + Ok(Self { + batch, + sequence, + bytes, + buffers, + }) + } + + pub fn batch(&self) -> usize { + self.batch + } + + pub fn sequence(&self) -> usize { + self.sequence + } + + /// Total scratch allocation bytes, excluding weights and CUDA overhead. + pub fn bytes(&self) -> usize { + self.bytes + } + + pub fn buffers(&self) -> &WorkspaceBuffers { + &self.buffers + } +} + +#[cfg(test)] +mod tests; diff --git a/src/models/laya/src/workspace/tests.rs b/src/models/laya/src/workspace/tests.rs new file mode 100644 index 0000000..2a9fd76 --- /dev/null +++ b/src/models/laya/src/workspace/tests.rs @@ -0,0 +1,216 @@ +use super::*; +use libloading::Library; +use std::{path::PathBuf, process::Command}; +use tempfile::{TempDir, tempdir}; + +struct Fixture { + _dir: TempDir, + library: Library, + cuda: Cuda, +} +impl Fixture { + fn new() -> Self { + let dir = tempdir().unwrap(); + let path = dir + .path() + .join(format!("runtime{}", std::env::consts::DLL_SUFFIX)); + let output = Command::new("cc") + .args([ + if cfg!(target_os = "macos") { + "-dynamiclib" + } else { + "-shared" + }, + "-fPIC", + ]) + .arg( + PathBuf::from(env!("CARGO_MANIFEST_DIR")) + .join("../../backends/cuda/tests/fixtures/runtime.c"), + ) + .arg("-o") + .arg(&path) + .output() + .unwrap(); + assert!( + output.status.success(), + "{}", + String::from_utf8_lossy(&output.stderr) + ); + let library = unsafe { Library::new(&path) }.unwrap(); + let cuda = unsafe { Cuda::load(&path, 0) }.unwrap(); + Self { + _dir: dir, + library, + cuda, + } + } + fn live(&self) -> i32 { + live(&self.library) + } + fn trace(&self) -> String { + unsafe { + let trace = self + .library + .get:: *const std::ffi::c_char>(b"laya_test_trace\0") + .unwrap(); + std::ffi::CStr::from_ptr(trace()) + .to_string_lossy() + .into_owned() + } + } +} +fn live(library: &Library) -> i32 { + unsafe { + library + .get:: i32>(b"laya_test_live\0") + .unwrap()() + } +} +fn buffers(s: &Workspace) -> [&Buffer; 17] { + let s = s.buffers(); + [ + &s.ids, + &s.lengths, + &s.types, + &s.residual, + &s.hidden, + &s.qkv, + &s.attention, + &s.gated, + &s.feed_forward, + &s.indices, + &s.offsets, + &s.markers, + &s.scored, + &s.logits, + &s.features, + &s.action_hidden, + &s.actions, + ] +} + +#[test] +fn capacities_match_consumer_layouts_at_both_bounds() { + let fixture = Fixture::new(); + for (batch, sequence, expected) in [ + ( + 1, + 16, + [ + 128, 4, 8, 65536, 32768, 98304, 32768, 83968, 131072, 8192, 8, 4194304, 4194304, + 4096, 2056, 512, 4, + ], + ), + ( + 16, + 512, + [ + 65536, 64, 128, 33554432, 16777216, 50331648, 16777216, 42991616, 67108864, 8192, + 68, 4194304, 4194304, 4096, 32896, 8192, 64, + ], + ), + ] { + let workspace = Workspace::new(&fixture.cuda, batch, sequence).unwrap(); + assert_eq!(workspace.batch(), batch); + assert_eq!(workspace.sequence(), sequence); + assert_eq!(buffers(&workspace).map(Buffer::bytes), expected); + assert_eq!(workspace.bytes(), expected.iter().sum::()); + assert_eq!(fixture.live(), 18); + drop(workspace); + assert_eq!(fixture.live(), 1); + } +} + +#[test] +fn invalid_shapes_allocate_nothing() { + let fixture = Fixture::new(); + let before = fixture.trace(); + for batch in [0, 3, 17, 32, usize::MAX] { + assert!(Workspace::new(&fixture.cuda, batch, 16).is_err()); + } + for sequence in [0, 1, 15, 17, 511, 513, usize::MAX] { + assert!(Workspace::new(&fixture.cuda, 1, sequence).is_err()); + } + assert_eq!(fixture.live(), 1); + assert_eq!(fixture.trace(), before); +} + +#[test] +fn failed_allocations_release_the_partial_workspace() { + for completed in [0, 1, 8, 16] { + let fixture = Fixture::new(); + unsafe { + fixture + .library + .get::(b"laya_test_fail_alloc_after\0") + .unwrap()(completed); + } + assert!(Workspace::new(&fixture.cuda, 1, 16).is_err()); + assert_eq!(fixture.live(), 1, "failed after {completed} allocations"); + } +} + +#[test] +fn workspaces_do_not_alias_and_outlive_the_cuda_handle() { + let fixture = Fixture::new(); + let first = Workspace::new(&fixture.cuda, 1, 16).unwrap(); + let second = Workspace::new(&fixture.cuda, 2, 32).unwrap(); + assert_eq!(fixture.live(), 35); + drop(fixture.cuda); + for (i, b) in buffers(&first) + .into_iter() + .chain(buffers(&second)) + .enumerate() + { + b.write(&vec![i as u8; b.bytes().min(32)]).unwrap(); + } + for (i, b) in buffers(&first) + .into_iter() + .chain(buffers(&second)) + .enumerate() + { + assert_eq!( + b.read(b.bytes().min(32)).unwrap(), + vec![i as u8; b.bytes().min(32)] + ); + } + drop(first); + assert_eq!(live(&fixture.library), 18); + drop(second); + assert_eq!(live(&fixture.library), 0); +} + +#[test] +#[ignore = "requires approved GPU, LAYA_CUDA_LIBRARY and LAYA_CUDA_DEVICE"] +fn real_gpu_workspace_capacity_and_reuse() { + let path = PathBuf::from(std::env::var_os("LAYA_CUDA_LIBRARY").unwrap()); + let device = std::env::var("LAYA_CUDA_DEVICE").unwrap().parse().unwrap(); + for (batch, sequence, expected_bytes) in + [(1, 16, 8848032), (1, 512, 22628896), (16, 512, 236048836)] + { + let cuda = unsafe { Cuda::load(&path, device) }.unwrap(); + let workspace = Workspace::new(&cuda, batch, sequence).unwrap(); + drop(cuda); + assert_eq!(workspace.bytes(), expected_bytes); + for pass in 0..2 { + let pattern = |size, index| { + (0..size) + .map(|offset| ((offset % 251 + index * 7 + pass * 89) % 256) as u8) + .collect::>() + }; + for (i, b) in buffers(&workspace).into_iter().enumerate() { + b.write(&pattern(b.bytes(), i)).unwrap(); + } + for (i, b) in buffers(&workspace).into_iter().enumerate() { + assert_eq!( + b.read(b.bytes()).unwrap(), + pattern(b.bytes(), i), + "buffer {i}, pass {pass}" + ); + } + } + println!( + "batch={batch} sequence={sequence} buffers=17 bytes={expected_bytes} verified_passes=2" + ); + } +} From ad1c7b800763b99392cba1f12e0a06f2f26b2c6a Mon Sep 17 00:00:00 2001 From: linear3735 Date: Wed, 30 Sep 2026 13:41:56 +0800 Subject: [PATCH 04/15] Run Laya encoder and decision layers through native CUDA --- recipe/laya/native/export_encoder.py | 90 ++++++++ src/backends/cuda/src/kernels.rs | 79 +++++++ src/backends/cuda/src/lib.rs | 2 + src/backends/cuda/tests/fixtures/runtime.c | 14 ++ src/backends/cuda/tests/runtime.rs | 39 ++++ src/models/laya/Cargo.toml | 2 +- src/models/laya/README.md | 38 ++++ src/models/laya/src/artifacts.rs | 77 +++++++ src/models/laya/src/artifacts/tests.rs | 178 +++++++++++++++ src/models/laya/src/encoder.rs | 251 +++++++++++++++++++++ src/models/laya/src/encoder/tests.rs | 118 ++++++++++ src/models/laya/src/lib.rs | 2 + 12 files changed, 889 insertions(+), 1 deletion(-) create mode 100644 recipe/laya/native/export_encoder.py create mode 100644 src/backends/cuda/src/kernels.rs create mode 100644 src/models/laya/src/artifacts.rs create mode 100644 src/models/laya/src/artifacts/tests.rs create mode 100644 src/models/laya/src/encoder.rs create mode 100644 src/models/laya/src/encoder/tests.rs diff --git a/recipe/laya/native/export_encoder.py b/recipe/laya/native/export_encoder.py new file mode 100644 index 0000000..eb3cbae --- /dev/null +++ b/recipe/laya/native/export_encoder.py @@ -0,0 +1,90 @@ +"""Export official Laya intermediates for the Rust encoder integration test.""" +import argparse +import copy +import importlib.metadata +import json +from pathlib import Path + +import torch +from laya import Agent +from laya.common import collate_items + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument("checkpoint") + parser.add_argument("requests", type=Path) + parser.add_argument("output", type=Path) + args = parser.parse_args() + assert importlib.metadata.version("laya") == "0.3.20" + agent = Agent(args.checkpoint, device="cuda", fast=False, compile=False) + assert agent.accelerate(use_graphs=False, strict=True) + fast = agent._fast + assert fast is not None and not fast.use_graphs + cases = json.loads(args.requests.read_text()) + # Exercise the largest supported batch with real, independently packed questions. + widest = copy.deepcopy(next(c for c in cases if c["name"] == "long_3")) + widest["name"] = "long_16" + question = next(iter(widest["request"]["questions"].values())) + widest["request"]["questions"] = {f"q{i}": copy.deepcopy(question) for i in range(16)} + cases.append(widest) + args.output.mkdir(parents=True, exist_ok=True) + records = [] + for case in cases: + name, request = case["name"], case["request"] + questions = request["questions"] + internal = {k: agent._to_internal(v) for k, v in questions.items()} + items = agent._encode_state(request["state"], list(questions), internal) + packed = collate_items([items], agent.tok.pad_token_id) + n, length = packed["input_ids"].shape + b, l = 1 << (n - 1).bit_length(), (length + 15) // 16 * 16 + assert 1 <= b <= 16 and 16 <= l <= 512 + ids = torch.zeros((b, l), dtype=torch.int64, device="cuda") + lens = torch.zeros(b, dtype=torch.int32, device="cuda") + types = torch.zeros(b, dtype=torch.int64, device="cuda") + ids[:n, :length] = packed["input_ids"].cuda() + lens[:n] = packed["attention_mask"].sum(-1).to("cuda", torch.int32) + types[:n] = packed["qtype"].cuda() + + def save(stage, value): + suffix = f"-{stage}" if stage else "" + (args.output / f"{name}{suffix}.f32").write_bytes(value.float().contiguous().cpu().numpy().tobytes()) + + original_ln = fast.k_addln + original_head_norm, original_ffn2 = fast.k_ln_b, fast.k_ffn2 + head_state, head_count = [None], [0] + + def head_norm(*values): + head_state[0] = values[0] + original_head_norm(*values) + + def ffn2(*values): + original_ffn2(*values) + save(f"head{head_count[0]}", head_state[0] + values[-1].float()) + head_count[0] += 1 + count = [0] + + def addln(*values): + original_ln(*values) + count[0] += 1 + if count[0] in (2, 4, 6, 56): + save(f"encoder{count[0] // 2 - 1}", values[0]) + + fast.k_addln = addln + fast.k_ln_b, fast.k_ffn2 = head_norm, ffn2 + with torch.no_grad(): + embedding = torch.nn.functional.embedding(ids, fast.emb_w).reshape(-1, 1024).float() + save("embedding", torch.nn.functional.layer_norm(embedding, (1024,), fast.emb_ln, None, fast.eps)) + hidden = fast._encode(ids, lens, types) + save("", hidden) + fast.k_addln = original_ln + fast.k_ln_b, fast.k_ffn2 = original_head_norm, original_ffn2 + records.append({"name": name, "batch": b, "sequence": l, + "ids": ids.flatten().cpu().tolist(), "lengths": lens.cpu().tolist(), + "types": types.cpu().tolist()}) + print("REFERENCE", name, b, l, flush=True) + (args.output / "cases.json").write_text(json.dumps(records)) + + +if __name__ == "__main__": + main() diff --git a/src/backends/cuda/src/kernels.rs b/src/backends/cuda/src/kernels.rs new file mode 100644 index 0000000..3e7c310 --- /dev/null +++ b/src/backends/cuda/src/kernels.rs @@ -0,0 +1,79 @@ +//! Calls into a separately built, trusted CUDA kernel library. +use super::*; +use std::collections::HashMap; + +type Launch = unsafe extern "C" fn(*mut Ptr, i32, i32, i32, Ptr) -> i32; + +pub struct Kernels { + cuda: Cuda, + _library: Library, + functions: HashMap, +} + +impl Kernels { + /// # Safety + /// The library must implement the named pointer-array launch ABI and + /// `laya_kernels_init`, and be compiled for this device's architecture. + pub unsafe fn load(cuda: &Cuda, path: &Path, names: &[&str]) -> Result { + cuda.ctx.activate()?; + let library = unsafe { Library::new(path) }?; + let mut functions = HashMap::new(); + for name in names { + let symbol = format!("laya_{name}\0"); + let launch = unsafe { *library.get::(symbol.as_bytes())? }; + functions.insert((*name).to_owned(), launch); + } + let init = unsafe { library.get:: i32>(b"laya_kernels_init\0")? }; + cuda.ctx.functions.check(unsafe { init() })?; + Ok(Self { + cuda: cuda.clone(), + _library: library, + functions, + }) + } + + /// # Safety + /// Argument count, sizes, contents, dtypes and aliasing must match the kernel. + /// This checks context identity and shape bounds, not tensor semantics. + pub unsafe fn launch( + &self, + name: &str, + args: &[&Buffer], + batch: usize, + sequence: usize, + ) -> Result<()> { + ensure!( + batch.is_power_of_two() + && batch <= 16 + && (16..=512).contains(&sequence) + && sequence.is_multiple_of(16), + "invalid kernel shape" + ); + ensure!( + args.iter().all(|b| Rc::ptr_eq(&b.ctx, &self.cuda.ctx)), + "kernel buffer belongs to another CUDA context" + ); + let launch = self + .functions + .get(name) + .ok_or_else(|| anyhow!("kernel not loaded: {name}"))?; + self.cuda.ctx.activate()?; + let mut pointers: Vec<_> = args.iter().map(|b| b.ptr).collect(); + self.cuda.ctx.functions.check(unsafe { + launch( + pointers.as_mut_ptr(), + batch as i32, + sequence as i32, + (batch * sequence) as i32, + self.cuda.ctx.stream, + ) + }) + } +} + +impl Drop for Kernels { + fn drop(&mut self) { + // Pending launches must finish before unloading their device code. + let _ = self.cuda.sync(); + } +} diff --git a/src/backends/cuda/src/lib.rs b/src/backends/cuda/src/lib.rs index 4338ade..1770d35 100644 --- a/src/backends/cuda/src/lib.rs +++ b/src/backends/cuda/src/lib.rs @@ -186,3 +186,5 @@ impl Drop for Buffer { } } } + +pub mod kernels; diff --git a/src/backends/cuda/tests/fixtures/runtime.c b/src/backends/cuda/tests/fixtures/runtime.c index b37eb61..4a6175b 100644 --- a/src/backends/cuda/tests/fixtures/runtime.c +++ b/src/backends/cuda/tests/fixtures/runtime.c @@ -114,3 +114,17 @@ int laya_stream_free(void *stream) { live--; return 0; } + +#ifdef LAYA_TEST_KERNELS +int laya_kernels_init(void) { + record("kernels_init"); + return mode == mode_create_error ? 23 : 0; +} +int laya_fill(void **args, int b, int l, int m, void *stream) { + static const unsigned char value = 73; + record("kernel"); + if (m != b * l || !valid_stream(stream)) return 91; + if (mode == mode_copy_error) return 41; + return copy(args[0], &value, 1, stream); +} +#endif diff --git a/src/backends/cuda/tests/runtime.rs b/src/backends/cuda/tests/runtime.rs index 1101021..6c4cd3f 100644 --- a/src/backends/cuda/tests/runtime.rs +++ b/src/backends/cuda/tests/runtime.rs @@ -252,3 +252,42 @@ fn real_gpu_round_trip() { drop(cuda); assert_eq!(buffer.read(expected.len()).unwrap(), expected); } + +#[test] +fn kernels_validate_context_and_shape_and_keep_the_library_alive() { + use omni_cuda::kernels::Kernels; + let f = Fixture::new(&["-DLAYA_TEST_KERNELS"]); + let cuda = unsafe { Cuda::load(&f.path, 0) }.unwrap(); + let buffer = cuda.alloc(1).unwrap(); + let kernels = unsafe { Kernels::load(&cuda, &f.path, &["fill"]) }.unwrap(); + assert!(unsafe { Kernels::load(&cuda, &f.path, &["absent"]) }.is_err()); + f.mode(MODE_CREATE_ERROR); + assert!(unsafe { Kernels::load(&cuda, &f.path, &["fill"]) }.is_err()); + f.mode(MODE_NORMAL); + let other = unsafe { Cuda::load(&f.path, 1) }.unwrap(); + let foreign = other.alloc(1).unwrap(); + let same_device = unsafe { Cuda::load(&f.path, 0) }.unwrap(); + let foreign_stream = same_device.alloc(1).unwrap(); + let before = f.trace(); + for args in [&foreign, &foreign_stream] { + assert!(unsafe { kernels.launch("fill", &[args], 1, 16) }.is_err()); + } + for (b, l) in [(0, 16), (3, 16), (32, 16), (1, 0), (1, 17), (1, 528)] { + assert!(unsafe { kernels.launch("fill", &[&buffer], b, l) }.is_err()); + } + assert!(unsafe { kernels.launch("absent", &[&buffer], 1, 16) }.is_err()); + assert_eq!( + before, + f.trace(), + "rejected calls must not enter the runtime" + ); + f.mode(MODE_COPY_ERROR); + assert!(unsafe { kernels.launch("fill", &[&buffer], 1, 16) }.is_err()); + f.mode(MODE_NORMAL); + drop(cuda); + unsafe { kernels.launch("fill", &[&buffer], 1, 16) }.unwrap(); + assert!(f.trace().ends_with("device0 kernel0 ")); + drop(kernels); + assert!(f.trace().ends_with("device0 sync0 ")); + assert_eq!(buffer.read(1).unwrap(), [73]); +} diff --git a/src/models/laya/Cargo.toml b/src/models/laya/Cargo.toml index 69ce281..a99881d 100644 --- a/src/models/laya/Cargo.toml +++ b/src/models/laya/Cargo.toml @@ -12,8 +12,8 @@ omni-cuda = { path = "../../backends/cuda" } safetensors = "0.6" serde = { version = "1", features = ["derive"] } serde_json = "1" +sha2 = "0.10" [dev-dependencies] libloading = "0.8" -sha2 = "0.10" tempfile = "3" diff --git a/src/models/laya/README.md b/src/models/laya/README.md index 65a410c..d8203aa 100644 --- a/src/models/laya/README.md +++ b/src/models/laya/README.md @@ -70,3 +70,41 @@ The ignored `real_gpu_workspace_capacity_and_reuse` test requires `LAYA_CUDA_LIBRARY` and `LAYA_CUDA_DEVICE`. It writes and reads all 17 buffers twice at three shapes, including both capacity bounds, using deterministic byte patterns. This tests allocation and transfer, not model numerics or latency. + +## Eager encoder + +`Encoder::load(&cuda, checkpoint, bundle)` loads the existing trusted native Laya +bundle and resident weights. `run(ids, lengths, types, &workspace)` executes all +28 encoder layers and both decision transformer layers, synchronizes, and leaves +FP32 hidden states in `workspace.buffers().residual`. Padded rows have length zero. +IDs must be within the vocabulary; lengths and types are checked before upload. + +This uses the existing original RoPE entry point and dynamic full/local attention. +No Scorer, output decoding, HTTP service, Graph capture or cache is included. +The resource library supplied to `Cuda::load` remains separate from the operator +bundle; CPU builds need neither library. The operator bundle currently targets +Hopper `sm_90a`. `Encoder::load` is unsafe because callers must trust the native +code and provide a compatible GPU; hashes bind artifacts, not code trust. + +Reuse the [existing CUDA build entry](https://github.com/linear3735/system1-omni/blob/5ff41a5/src/backends/cuda/build.sh) +and its model source/export tools. Export rotary tables with that version's +`tools/export_tables.py`. Keep `liblaya_cuda.so`, `build-manifest.json`, +`tables.json` and the four rotary table files in one bundle directory. These build +tools are a separate dependency, not duplicated by this encoder change. Startup +checks checkpoint and bundle hashes before loading native code. + +For GPU validation, supply real request fixtures to +`recipe/laya/native/export_encoder.py CHECKPOINT REQUESTS OUTPUT`. It uses Laya +0.3.20 and PyTorch on CUDA as the reference, independently of the Rust runtime. +Set `LAYA_CUDA_LIBRARY`, `LAYA_CHECKPOINT`, `LAYA_KERNEL_BUNDLE` and +`LAYA_ENCODER_ORACLE`, then run: + +```sh +cargo test --release --locked -p omni-laya --lib real_encoder_matches_official_hidden_states -- --ignored --nocapture +``` + +The test checks selected encoder intermediates, both head layers, final hidden +states, shape changes and repeated workspace reuse. It compares valid tokens; +empty-key padding differs intentionally from the original attention. Candidate +padding must still be finite. This is numerical parity, not model quality or a +performance benchmark. The intermediate capture hooks compile only in tests. diff --git a/src/models/laya/src/artifacts.rs b/src/models/laya/src/artifacts.rs new file mode 100644 index 0000000..128a02b --- /dev/null +++ b/src/models/laya/src/artifacts.rs @@ -0,0 +1,77 @@ +//! Bind a compiled bundle to its read-only checkpoint before CUDA startup. +use anyhow::{Context, Result, ensure}; +use sha2::{Digest, Sha256}; +use std::{fs, io::Read, path::Path}; + +pub const CHECKPOINT_ARTIFACTS: [&str; 5] = [ + "rl_agent_config.json", + "encoder/config.json", + "model.safetensors", + "tokenizer/tokenizer.json", + "tokenizer/tokenizer_config.json", +]; + +fn sha256_file(path: &Path) -> Result { + let mut file = fs::File::open(path).with_context(|| format!("open {}", path.display()))?; + let mut digest = Sha256::new(); + let mut buffer = [0u8; 64 * 1024]; + loop { + let count = file + .read(&mut buffer) + .with_context(|| format!("hash {}", path.display()))?; + if count == 0 { + break; + } + digest.update(&buffer[..count]); + } + Ok(format!("{:x}", digest.finalize())) +} + +pub fn validate_bundle(checkpoint: &Path, bundle: &Path) -> Result<()> { + let tables: serde_json::Value = serde_json::from_slice(&fs::read(bundle.join("tables.json"))?)?; + let build: serde_json::Value = + serde_json::from_slice(&fs::read(bundle.join("build-manifest.json"))?)?; + ensure!( + tables["abi"] == 1 + && tables["laya"] == "0.3.20" + && tables["hidden_size"] == 1024 + && tables["head_dim"] == 64 + && tables["max_len"] == 512 + && build["abi"] == 1 + && build["arch"] == "sm_90a", + "unsupported CUDA bundle" + ); + let check = |path: std::path::PathBuf, expected: Option<&str>| -> Result<()> { + let hash = sha256_file(&path)?; + ensure!( + expected == Some(hash.as_str()), + "bundle hash mismatch: {}", + path.display() + ); + Ok(()) + }; + for name in CHECKPOINT_ARTIFACTS { + let expected = tables["checkpoint_sha256"][name] + .as_str() + .with_context(|| { + format!("missing checkpoint hash for {name}; regenerate tables.json") + })?; + check(checkpoint.join(name), Some(expected))?; + } + for name in [ + "rope_full_cos.f32", + "rope_full_sin.f32", + "rope_local_cos.f32", + "rope_local_sin.f32", + ] { + check(bundle.join(name), tables["tables"][name].as_str())?; + } + check( + bundle.join("liblaya_cuda.so"), + build["library_sha256"].as_str(), + )?; + Ok(()) +} + +#[cfg(test)] +mod tests; diff --git a/src/models/laya/src/artifacts/tests.rs b/src/models/laya/src/artifacts/tests.rs new file mode 100644 index 0000000..84fa9cd --- /dev/null +++ b/src/models/laya/src/artifacts/tests.rs @@ -0,0 +1,178 @@ +use super::*; +use serde_json::json; +use std::{ + path::PathBuf, + sync::atomic::{AtomicU64, Ordering}, +}; + +static NEXT: AtomicU64 = AtomicU64::new(0); +struct Fixture { + root: PathBuf, + checkpoint: PathBuf, + bundle: PathBuf, +} +impl Fixture { + fn new() -> Self { + let root = std::env::temp_dir().join(format!( + "laya-artifacts-{}-{}", + std::process::id(), + NEXT.fetch_add(1, Ordering::Relaxed) + )); + fs::create_dir(&root).unwrap(); + let checkpoint = root.join("checkpoint"); + let bundle = root.join("bundle"); + fs::create_dir_all(checkpoint.join("encoder")).unwrap(); + fs::create_dir_all(checkpoint.join("tokenizer")).unwrap(); + fs::create_dir(&bundle).unwrap(); + let mut hashes = serde_json::Map::new(); + for name in CHECKPOINT_ARTIFACTS { + // The weight fixture crosses several bounded hash reads; no GPU or real model required. + let data = if name == "model.safetensors" { + vec![42; 192 * 1024 + 1] + } else { + name.as_bytes().to_vec() + }; + fs::write(checkpoint.join(name), &data).unwrap(); + hashes.insert(name.into(), json!(format!("{:x}", Sha256::digest(&data)))); + } + let mut tables = serde_json::Map::new(); + for name in [ + "rope_full_cos.f32", + "rope_full_sin.f32", + "rope_local_cos.f32", + "rope_local_sin.f32", + ] { + fs::write(bundle.join(name), b"table").unwrap(); + tables.insert( + name.into(), + json!(format!("{:x}", Sha256::digest(b"table"))), + ); + } + let table_manifest = json!({ + "abi": 1, + "laya": "0.3.20", + "hidden_size": 1024, + "head_dim": 64, + "max_len": 512, + "tables": tables, + "checkpoint_sha256": hashes, + }); + fs::write( + bundle.join("tables.json"), + serde_json::to_vec(&table_manifest).unwrap(), + ) + .unwrap(); + fs::write(bundle.join("liblaya_cuda.so"), b"not loaded in CPU test").unwrap(); + let build_manifest = json!({ + "abi": 1, + "arch": "sm_90a", + "library_sha256": format!("{:x}", Sha256::digest(b"not loaded in CPU test")), + }); + fs::write( + bundle.join("build-manifest.json"), + serde_json::to_vec(&build_manifest).unwrap(), + ) + .unwrap(); + Self { + root, + checkpoint, + bundle, + } + } + fn validate(&self) -> Result<()> { + validate_bundle(&self.checkpoint, &self.bundle) + } +} +impl Drop for Fixture { + fn drop(&mut self) { + let _ = fs::remove_dir_all(&self.root); + } +} + +#[test] +fn matching_artifacts_pass_without_loading_cuda() { + Fixture::new().validate().unwrap(); +} + +#[test] +fn each_checkpoint_artifact_is_bound_to_the_bundle() { + for name in CHECKPOINT_ARTIFACTS { + let f = Fixture::new(); + let path = f.checkpoint.join(name); + let mut data = fs::read(&path).unwrap(); + data[0] ^= 1; // Same size: shape/file-size checks alone would not catch substitution. + fs::write(path, data).unwrap(); + let error = f.validate().unwrap_err().to_string(); + assert!( + error.contains("hash mismatch") && error.contains(name), + "{error}" + ); + } +} + +#[test] +fn missing_checkpoint_hashes_fail_closed() { + for name in CHECKPOINT_ARTIFACTS { + let f = Fixture::new(); + let path = f.bundle.join("tables.json"); + let mut manifest: serde_json::Value = + serde_json::from_slice(&fs::read(&path).unwrap()).unwrap(); + manifest["checkpoint_sha256"] + .as_object_mut() + .unwrap() + .remove(name); + fs::write(path, serde_json::to_vec(&manifest).unwrap()).unwrap(); + let error = f.validate().unwrap_err().to_string(); + assert!( + error.contains("missing checkpoint hash") && error.contains(name), + "{error}" + ); + } + let f = Fixture::new(); + let path = f.bundle.join("tables.json"); + let mut manifest: serde_json::Value = + serde_json::from_slice(&fs::read(&path).unwrap()).unwrap(); + manifest + .as_object_mut() + .unwrap() + .remove("checkpoint_sha256"); + fs::write(path, serde_json::to_vec(&manifest).unwrap()).unwrap(); + assert!( + f.validate() + .unwrap_err() + .to_string() + .contains("missing checkpoint hash") + ); +} + +#[test] +fn missing_checkpoint_file_fails_closed() { + for name in CHECKPOINT_ARTIFACTS { + let f = Fixture::new(); + fs::remove_file(f.checkpoint.join(name)).unwrap(); + let error = f.validate().unwrap_err().to_string(); + assert!(error.contains(name), "{error}"); + } +} + +#[test] +fn changed_tables_and_native_library_fail_identity_checks() { + for name in [ + "rope_full_cos.f32", + "rope_full_sin.f32", + "rope_local_cos.f32", + "rope_local_sin.f32", + "liblaya_cuda.so", + ] { + let f = Fixture::new(); + let path = f.bundle.join(name); + let mut data = fs::read(&path).unwrap(); + data[0] ^= 1; + fs::write(path, data).unwrap(); + let error = f.validate().unwrap_err().to_string(); + assert!( + error.contains("hash mismatch") && error.contains(name), + "{error}" + ); + } +} diff --git a/src/models/laya/src/encoder.rs b/src/models/laya/src/encoder.rs new file mode 100644 index 0000000..02a1985 --- /dev/null +++ b/src/models/laya/src/encoder.rs @@ -0,0 +1,251 @@ +//! Eager execution of the fixed Laya encoder and decision transformer. +use crate::{ + artifacts, config::Config, resident::ResidentWeights, weights::Weights, workspace::Workspace, +}; +use anyhow::{Result, ensure}; +use omni_cuda::{Buffer, Cuda, kernels::Kernels}; +use std::{collections::HashMap, fs, path::Path}; + +const NAMES: &[&str] = &[ + "embed", + "qkv", + "rope_original", + "attn_full", + "attn_local", + "out", + "addln", + "geglu", + "down", + "type", + "ln_bias", + "head_in", + "head_out", + "addln_bias", + "ffn1", + "ffn2", + "residual", +]; + +pub struct Encoder { + cuda: Cuda, + kernels: Kernels, + weights: ResidentWeights, + tables: HashMap, + zeros: Buffer, + #[cfg(test)] + checkpoints: std::cell::RefCell)>>, +} + +impl Encoder { + /// # Safety + /// `bundle` must contain trusted Laya native code built for this device. + /// Hash validation binds files together; it does not establish code trust. + pub unsafe fn load(cuda: &Cuda, checkpoint: &Path, bundle: &Path) -> Result { + Config::load(checkpoint)?; + artifacts::validate_bundle(checkpoint, bundle)?; + let kernels = unsafe { Kernels::load(cuda, &bundle.join("liblaya_cuda.so"), NAMES) }?; + let source = Weights::open(&checkpoint.join("model.safetensors"))?; + let weights = ResidentWeights::upload(cuda, &source)?; + let mut tables = HashMap::new(); + for kind in ["full", "local"] { + for part in ["cos", "sin"] { + let name = format!("rope_{kind}_{part}"); + let data = fs::read(bundle.join(format!("{name}.f32")))?; + ensure!(data.len() == 512 * 32 * 4, "invalid rotary table size"); + tables.insert(name, cuda.upload(&data)?); + } + } + Ok(Self { + #[cfg(test)] + checkpoints: Default::default(), + cuda: cuda.clone(), + kernels, + weights, + tables, + zeros: cuda.upload(&vec![0; 3072 * 4])?, + }) + } + + /// Writes FP32 final hidden states into `workspace.buffers().residual`. + /// Inputs include padded rows; a zero length marks a dummy row. + pub fn run( + &self, + ids: &[i64], + lengths: &[i32], + types: &[i64], + workspace: &Workspace, + ) -> Result<()> { + validate_inputs(ids, lengths, types, workspace.batch(), workspace.sequence())?; + let s = workspace.buffers(); + s.ids + .write(&ids.iter().flat_map(|x| x.to_le_bytes()).collect::>())?; + s.lengths.write( + &lengths + .iter() + .flat_map(|x| x.to_le_bytes()) + .collect::>(), + )?; + s.types.write( + &types + .iter() + .flat_map(|x| x.to_le_bytes()) + .collect::>(), + )?; + self.execute(workspace)?; + self.cuda.sync() + } + + fn execute(&self, workspace: &Workspace) -> Result<()> { + let (b, l) = (workspace.batch(), workspace.sequence()); + let s = workspace.buffers(); + let w = |name: &str| self.weights.get(name); + // Layouts are fixed by ResidentWeights and Workspace; indices were checked on CPU. + let call = |name: &str, args: &[&Buffer]| unsafe { self.kernels.launch(name, args, b, l) }; + let z = &self.zeros; + call( + "embed", + &[ + &s.ids, + w("encoder.embeddings.tok_embeddings.weight")?, + w("encoder.embeddings.norm.weight")?, + &s.residual, + &s.hidden, + ], + )?; + #[cfg(test)] + self.record("embedding", &s.residual)?; + for i in 0..28 { + let prefix = format!("encoder.layers.{i}"); + let layer = |name: &str| w(&format!("{prefix}.{name}")); + call("qkv", &[&s.hidden, layer("attn.Wqkv.weight")?, z, &s.qkv])?; + let kind = if i % 3 == 0 { "full" } else { "local" }; + call( + "rope_original", + &[ + &s.qkv, + &self.tables[&format!("rope_{kind}_cos")], + &self.tables[&format!("rope_{kind}_sin")], + ], + )?; + call(&format!("attn_{kind}"), &[&s.qkv, &s.lengths, &s.attention])?; + call( + "out", + &[&s.attention, layer("attn.Wo.weight")?, z, &s.hidden], + )?; + call( + "addln", + &[ + &s.residual, + &s.hidden, + layer("mlp_norm.weight")?, + z, + &s.hidden, + ], + )?; + call("geglu", &[&s.hidden, layer("mlp.Wi.weight")?, &s.gated])?; + call("down", &[&s.gated, layer("mlp.Wo.weight")?, z, &s.hidden])?; + let norm = if i < 27 { + format!("encoder.layers.{}.attn_norm.weight", i + 1) + } else { + "encoder.final_norm.weight".into() + }; + call("addln", &[&s.residual, &s.hidden, w(&norm)?, z, &s.hidden])?; + #[cfg(test)] + if [0, 1, 2, 27].contains(&i) { + self.record(&format!("encoder{i}"), &s.residual)?; + } + } + call( + "type", + &[&s.hidden, w("type_emb.weight")?, &s.types, &s.residual], + )?; + for i in 0..2 { + let prefix = format!("head.layers.{i}"); + let layer = |name: &str| w(&format!("{prefix}.{name}")); + call( + "ln_bias", + &[ + &s.residual, + &s.hidden, + layer("norm1.weight")?, + layer("norm1.bias")?, + &s.hidden, + ], + )?; + call( + "head_in", + &[ + &s.hidden, + layer("self_attn.in_proj_weight")?, + layer("self_attn.in_proj_bias")?, + &s.qkv, + ], + )?; + call("attn_full", &[&s.qkv, &s.lengths, &s.attention])?; + call( + "head_out", + &[ + &s.attention, + layer("self_attn.out_proj.weight")?, + layer("self_attn.out_proj.bias")?, + &s.hidden, + ], + )?; + call( + "addln_bias", + &[ + &s.residual, + &s.hidden, + layer("norm2.weight")?, + layer("norm2.bias")?, + &s.hidden, + ], + )?; + call( + "ffn1", + &[ + &s.hidden, + layer("linear1.weight")?, + layer("linear1.bias")?, + &s.feed_forward, + ], + )?; + call( + "ffn2", + &[ + &s.feed_forward, + layer("linear2.weight")?, + layer("linear2.bias")?, + &s.hidden, + ], + )?; + call("residual", &[&s.residual, &s.hidden])?; + #[cfg(test)] + self.record(&format!("head{i}"), &s.residual)?; + } + Ok(()) + } +} + +fn validate_inputs(ids: &[i64], lengths: &[i32], types: &[i64], b: usize, l: usize) -> Result<()> { + ensure!( + ids.len() == b * l && lengths.len() == b && types.len() == b, + "encoder input shape mismatch" + ); + ensure!( + ids.iter().all(|id| (0..50368).contains(id)), + "token ID outside vocabulary" + ); + ensure!( + lengths.iter().all(|n| *n >= 0 && *n as usize <= l), + "invalid sequence length" + ); + ensure!( + types.iter().all(|t| (0..=2).contains(t)), + "invalid question type" + ); + Ok(()) +} + +#[cfg(test)] +mod tests; diff --git a/src/models/laya/src/encoder/tests.rs b/src/models/laya/src/encoder/tests.rs new file mode 100644 index 0000000..fae8f21 --- /dev/null +++ b/src/models/laya/src/encoder/tests.rs @@ -0,0 +1,118 @@ +use super::*; + +#[test] +fn input_bounds_cover_padding_and_vocabulary_edges() { + assert!(validate_inputs(&[0; 32], &[16, 0], &[0, 2], 2, 16).is_ok()); + assert!(validate_inputs(&[50367; 16], &[16], &[1], 1, 16).is_ok()); + for ids in [vec![-1; 16], vec![50368; 16], vec![0; 15]] { + assert!(validate_inputs(&ids, &[16], &[0], 1, 16).is_err()); + } + for lens in [vec![-1], vec![17], vec![]] { + assert!(validate_inputs(&[0; 16], &lens, &[0], 1, 16).is_err()); + } + for types in [vec![-1], vec![3], vec![]] { + assert!(validate_inputs(&[0; 16], &[16], &types, 1, 16).is_err()); + } +} + +#[test] +#[ignore = "requires approved GPU, checkpoint, trusted kernel bundle and encoder oracle"] +fn real_encoder_matches_official_hidden_states() -> Result<()> { + let path = |name| { + std::env::var_os(name) + .map(std::path::PathBuf::from) + .expect(name) + }; + let cuda = unsafe { Cuda::load(&path("LAYA_CUDA_LIBRARY"), 0) }?; + let encoder = + unsafe { Encoder::load(&cuda, &path("LAYA_CHECKPOINT"), &path("LAYA_KERNEL_BUNDLE")) }?; + let root = path("LAYA_ENCODER_ORACLE"); + let cases: serde_json::Value = serde_json::from_slice(&fs::read(root.join("cases.json"))?)?; + for case in cases + .as_array() + .unwrap() + .iter() + .chain(cases.as_array().unwrap().iter().rev()) + { + let name = case["name"].as_str().unwrap(); + let b = case["batch"].as_u64().unwrap() as usize; + let l = case["sequence"].as_u64().unwrap() as usize; + let ids: Vec = serde_json::from_value(case["ids"].clone())?; + let lengths: Vec = serde_json::from_value(case["lengths"].clone())?; + let types: Vec = serde_json::from_value(case["types"].clone())?; + let workspace = Workspace::new(&cuda, b, l)?; + // Kernels are asynchronous; readback must observe the completed final layer. + encoder.run(&ids, &lengths, &types, &workspace)?; + let got = workspace.buffers().residual.read(b * l * 1024 * 4)?; + check_output( + name, + &got, + &fs::read(root.join(format!("{name}.f32")))?, + &lengths, + l, + )?; + for (stage, output) in encoder.checkpoints.take() { + check_output( + &format!("{name}/{stage}"), + &output, + &fs::read(root.join(format!("{name}-{stage}.f32")))?, + &lengths, + l, + )?; + } + // Reusing the same allocations must not accumulate residuals across requests. + encoder.run(&ids, &lengths, &types, &workspace)?; + encoder.checkpoints.borrow_mut().clear(); + assert_eq!( + got, + workspace.buffers().residual.read(got.len())?, + "{name}: repeat changed output" + ); + } + Ok(()) +} + +impl Encoder { + pub(super) fn record(&self, stage: &str, buffer: &Buffer) -> Result<()> { + self.checkpoints + .borrow_mut() + .push((stage.into(), buffer.read(buffer.bytes())?)); + Ok(()) + } +} + +fn check_output(name: &str, got: &[u8], expected: &[u8], lengths: &[i32], l: usize) -> Result<()> { + assert_eq!(expected.len(), got.len(), "{name}"); + let mut err2 = 0f64; + let mut ref2 = 0f64; + let mut max_error = 0f32; + let mut max_ref = 0f32; + for (j, (g, r)) in got + .as_chunks::<4>() + .0 + .iter() + .zip(expected.as_chunks::<4>().0) + .enumerate() + { + let g = f32::from_le_bytes(*g); + let r = f32::from_le_bytes(*r); + assert!(g.is_finite(), "{name}: nonfinite output at {j}"); + // Padding queries have no observable output; patched empty-key attention differs there. + if (j / 1024) % l >= lengths[j / (l * 1024)] as usize { + continue; + } + assert!(r.is_finite()); + let error = g - r; + err2 += f64::from(error).powi(2); + ref2 += f64::from(r).powi(2); + max_error = max_error.max(error.abs()); + max_ref = max_ref.max(r.abs()); + } + let nrms = (err2 / ref2.max(1e-24)).sqrt(); + assert!( + nrms <= 0.005 && max_error <= 0.02 * max_ref.max(1.), + "{name}: nrms={nrms}, max_abs={max_error}, ref_max={max_ref}" + ); + println!("ENCODER {name} nrms={nrms} max_abs={max_error} ref_max={max_ref}"); + Ok(()) +} diff --git a/src/models/laya/src/lib.rs b/src/models/laya/src/lib.rs index a4be935..4b2f248 100644 --- a/src/models/laya/src/lib.rs +++ b/src/models/laya/src/lib.rs @@ -1,4 +1,6 @@ +mod artifacts; pub mod config; +pub mod encoder; pub mod resident; pub mod weights; pub mod workspace; From 5b5518bb87f9863db7e57a2e91249461de9eaf62 Mon Sep 17 00:00:00 2001 From: linear3735 Date: Wed, 30 Sep 2026 13:44:34 +0800 Subject: [PATCH 05/15] Check encoder reuse without intermediate synchronization --- recipe/laya/native/export_encoder.py | 2 +- src/models/laya/src/encoder.rs | 2 +- src/models/laya/src/encoder/tests.rs | 12 +++++++----- 3 files changed, 9 insertions(+), 7 deletions(-) diff --git a/recipe/laya/native/export_encoder.py b/recipe/laya/native/export_encoder.py index eb3cbae..b0c572f 100644 --- a/recipe/laya/native/export_encoder.py +++ b/recipe/laya/native/export_encoder.py @@ -23,7 +23,7 @@ def main(): assert fast is not None and not fast.use_graphs cases = json.loads(args.requests.read_text()) # Exercise the largest supported batch with real, independently packed questions. - widest = copy.deepcopy(next(c for c in cases if c["name"] == "long_3")) + widest = copy.deepcopy(next(c for c in cases if c["name"] == "truncated_3")) widest["name"] = "long_16" question = next(iter(widest["request"]["questions"].values())) widest["request"]["questions"] = {f"q{i}": copy.deepcopy(question) for i in range(16)} diff --git a/src/models/laya/src/encoder.rs b/src/models/laya/src/encoder.rs index 02a1985..5bc8a97 100644 --- a/src/models/laya/src/encoder.rs +++ b/src/models/laya/src/encoder.rs @@ -33,7 +33,7 @@ pub struct Encoder { tables: HashMap, zeros: Buffer, #[cfg(test)] - checkpoints: std::cell::RefCell)>>, + checkpoints: tests::Checkpoints, } impl Encoder { diff --git a/src/models/laya/src/encoder/tests.rs b/src/models/laya/src/encoder/tests.rs index fae8f21..4b9a033 100644 --- a/src/models/laya/src/encoder/tests.rs +++ b/src/models/laya/src/encoder/tests.rs @@ -1,5 +1,7 @@ use super::*; +pub(super) type Checkpoints = std::cell::RefCell)>>>; + #[test] fn input_bounds_cover_padding_and_vocabulary_edges() { assert!(validate_inputs(&[0; 32], &[16, 0], &[0, 2], 2, 16).is_ok()); @@ -42,6 +44,7 @@ fn real_encoder_matches_official_hidden_states() -> Result<()> { let types: Vec = serde_json::from_value(case["types"].clone())?; let workspace = Workspace::new(&cuda, b, l)?; // Kernels are asynchronous; readback must observe the completed final layer. + *encoder.checkpoints.borrow_mut() = Some(Vec::new()); encoder.run(&ids, &lengths, &types, &workspace)?; let got = workspace.buffers().residual.read(b * l * 1024 * 4)?; check_output( @@ -51,7 +54,7 @@ fn real_encoder_matches_official_hidden_states() -> Result<()> { &lengths, l, )?; - for (stage, output) in encoder.checkpoints.take() { + for (stage, output) in encoder.checkpoints.take().unwrap() { check_output( &format!("{name}/{stage}"), &output, @@ -62,7 +65,6 @@ fn real_encoder_matches_official_hidden_states() -> Result<()> { } // Reusing the same allocations must not accumulate residuals across requests. encoder.run(&ids, &lengths, &types, &workspace)?; - encoder.checkpoints.borrow_mut().clear(); assert_eq!( got, workspace.buffers().residual.read(got.len())?, @@ -74,9 +76,9 @@ fn real_encoder_matches_official_hidden_states() -> Result<()> { impl Encoder { pub(super) fn record(&self, stage: &str, buffer: &Buffer) -> Result<()> { - self.checkpoints - .borrow_mut() - .push((stage.into(), buffer.read(buffer.bytes())?)); + if let Some(stages) = self.checkpoints.borrow_mut().as_mut() { + stages.push((stage.into(), buffer.read(buffer.bytes())?)); + } Ok(()) } } From 0c50122cf1f351d988b13a8e43f2479e92ce44d9 Mon Sep 17 00:00:00 2001 From: linear3735 Date: Wed, 30 Sep 2026 13:48:25 +0800 Subject: [PATCH 06/15] Clarify eager encoder inference scope --- src/models/laya/README.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/models/laya/README.md b/src/models/laya/README.md index d8203aa..c4c59f3 100644 --- a/src/models/laya/README.md +++ b/src/models/laya/README.md @@ -6,7 +6,7 @@ GPU operations and kernel implementations belong in [`backends/cuda/`](../../bac The `omni-laya` crate currently reads and checks the English Laya 0.3.20 checkpoint. `Config::load` validates the architecture and temperatures; `Weights` checks tensor names and shapes and converts FP32, FP16 and BF16 values. `checkpoint_tensors()` lists the 206 expected tensors. Each backend chooses its own storage precision. -Keep checkpoint files unchanged while `Weights` holds a read-only memory mapping. This crate does not yet execute inference. +Keep checkpoint files unchanged while `Weights` holds a read-only memory mapping. The eager encoder is described below; complete request-to-result inference is not yet implemented. ## CPU checks From d37da99f8bc416e1761bc2ea9ee5a5af0d9abb65 Mon Sep 17 00:00:00 2001 From: linear3735 Date: Wed, 30 Sep 2026 13:49:45 +0800 Subject: [PATCH 07/15] Document separate CUDA operator loading --- src/backends/cuda/README.md | 17 ++++++++++++++--- 1 file changed, 14 insertions(+), 3 deletions(-) diff --git a/src/backends/cuda/README.md b/src/backends/cuda/README.md index 39161c3..bef609a 100644 --- a/src/backends/cuda/README.md +++ b/src/backends/cuda/README.md @@ -4,9 +4,9 @@ stream per context, plus the buffers allocated through that context. Rust builds and CPU tests need no CUDA toolkit. -This first slice covers allocation, copies, synchronization and cleanup. Model -initialization, weights, kernel calls, Graphs and hardware-specific optimizations -remain separate. It does not yet run Laya inference. +The resource library covers allocation, copies, synchronization and cleanup. +`Kernels` loads operator code separately; model execution order belongs to Laya. +Graphs and hardware-specific optimizations remain separate. ## Build and check @@ -49,3 +49,14 @@ hidden in stream creation. These are Laya's resource entry points, not a new shared tensor interface. A common runtime can be extracted when another model needs the same implementation. + +## Kernel library + +`Kernels::load(&cuda, path, names)` resolves the requested Laya pointer-array +entry points and initializes their launch attributes on the owning device. +`launch` checks the supported batch/sequence bounds and rejects buffers from +another context, including another stream on the same device. Tensor sizes, +dtypes, argument counts, contents and aliasing remain the unsafe caller's contract. +The library stays loaded until pending work has synchronized. Loading requires +trusted native code compiled for the selected GPU; the resource library itself +does not impose the operator bundle's architecture restrictions. From 684470ae01c3b786825abf064193b80baba79e2b Mon Sep 17 00:00:00 2001 From: linear3735 Date: Wed, 30 Sep 2026 14:00:15 +0800 Subject: [PATCH 08/15] Validate distinct mixed requests at maximum encoder batch --- recipe/laya/native/export_encoder.py | 37 +++++++++++++++++++++------- src/models/laya/src/encoder/tests.rs | 18 ++++++++++++++ 2 files changed, 46 insertions(+), 9 deletions(-) diff --git a/recipe/laya/native/export_encoder.py b/recipe/laya/native/export_encoder.py index b0c572f..198a0b8 100644 --- a/recipe/laya/native/export_encoder.py +++ b/recipe/laya/native/export_encoder.py @@ -1,6 +1,5 @@ """Export official Laya intermediates for the Rust encoder integration test.""" import argparse -import copy import importlib.metadata import json from pathlib import Path @@ -22,14 +21,8 @@ def main(): fast = agent._fast assert fast is not None and not fast.use_graphs cases = json.loads(args.requests.read_text()) - # Exercise the largest supported batch with real, independently packed questions. - widest = copy.deepcopy(next(c for c in cases if c["name"] == "truncated_3")) - widest["name"] = "long_16" - question = next(iter(widest["request"]["questions"].values())) - widest["request"]["questions"] = {f"q{i}": copy.deepcopy(question) for i in range(16)} - cases.append(widest) args.output.mkdir(parents=True, exist_ok=True) - records = [] + inputs, rows = [], {} for case in cases: name, request = case["name"], case["request"] questions = request["questions"] @@ -46,6 +39,32 @@ def main(): lens[:n] = packed["attention_mask"].sum(-1).to("cuda", torch.int32) types[:n] = packed["qtype"].cuda() + inputs.append((name, ids, lens, types, [])) + for index in range(n): + valid = int(lens[index]) + tokens = ids[index, :valid].cpu().tolist() + kind = int(types[index]) + rows.setdefault((tuple(tokens), kind), (name, index, tokens, kind)) + + # Use distinct real rows, including a maximum-length row, to expose batch indexing errors. + selected = list(rows.values()) + longest = max(selected, key=lambda row: len(row[2])) + selected = [longest] + [row for row in selected if row != longest][:15] + assert len(selected) == 16 and len(longest[2]) == 512 + assert {row[3] for row in selected} == {0, 1, 2} + assert len({len(row[2]) for row in selected}) > 1 + ids = torch.zeros((16, 512), dtype=torch.int64, device="cuda") + lens = torch.zeros(16, dtype=torch.int32, device="cuda") + types = torch.zeros(16, dtype=torch.int64, device="cuda") + for index, (_, _, tokens, kind) in enumerate(selected): + ids[index, :len(tokens)] = torch.tensor(tokens, dtype=torch.int64, device="cuda") + lens[index], types[index] = len(tokens), kind + origins = [{"case": row[0], "row": row[1]} for row in selected] + inputs.append(("mixed_16", ids, lens, types, origins)) + records = [] + for name, ids, lens, types, origins in inputs: + b, l = ids.shape + def save(stage, value): suffix = f"-{stage}" if stage else "" (args.output / f"{name}{suffix}.f32").write_bytes(value.float().contiguous().cpu().numpy().tobytes()) @@ -81,7 +100,7 @@ def addln(*values): fast.k_ln_b, fast.k_ffn2 = original_head_norm, original_ffn2 records.append({"name": name, "batch": b, "sequence": l, "ids": ids.flatten().cpu().tolist(), "lengths": lens.cpu().tolist(), - "types": types.cpu().tolist()}) + "types": types.cpu().tolist(), "origins": origins}) print("REFERENCE", name, b, l, flush=True) (args.output / "cases.json").write_text(json.dumps(records)) diff --git a/src/models/laya/src/encoder/tests.rs b/src/models/laya/src/encoder/tests.rs index 4b9a033..26799e4 100644 --- a/src/models/laya/src/encoder/tests.rs +++ b/src/models/laya/src/encoder/tests.rs @@ -42,6 +42,24 @@ fn real_encoder_matches_official_hidden_states() -> Result<()> { let ids: Vec = serde_json::from_value(case["ids"].clone())?; let lengths: Vec = serde_json::from_value(case["lengths"].clone())?; let types: Vec = serde_json::from_value(case["types"].clone())?; + if name == "mixed_16" { + assert_eq!((b, l), (16, 512)); + let rows: std::collections::HashSet<_> = lengths + .iter() + .enumerate() + .map(|(row, n)| (&ids[row * l..row * l + *n as usize], types[row])) + .collect(); + assert_eq!(rows.len(), 16, "maximum batch must have distinct real rows"); + assert_eq!( + types + .iter() + .copied() + .collect::>() + .len(), + 3 + ); + assert!(lengths.iter().any(|n| *n < 512) && lengths.contains(&512)); + } let workspace = Workspace::new(&cuda, b, l)?; // Kernels are asynchronous; readback must observe the completed final layer. *encoder.checkpoints.borrow_mut() = Some(Vec::new()); From 0a349f3398be73c5149f163e66a3b1361f2dd158 Mon Sep 17 00:00:00 2001 From: linear3735 Date: Wed, 30 Sep 2026 14:26:22 +0800 Subject: [PATCH 09/15] Validate rotary sizes before loading kernels and honor GPU selection --- src/models/laya/README.md | 2 +- src/models/laya/src/artifacts.rs | 4 ++++ src/models/laya/src/artifacts/tests.rs | 31 +++++++++++++++++++++----- src/models/laya/src/encoder/tests.rs | 3 ++- 4 files changed, 33 insertions(+), 7 deletions(-) diff --git a/src/models/laya/README.md b/src/models/laya/README.md index c4c59f3..c1fe735 100644 --- a/src/models/laya/README.md +++ b/src/models/laya/README.md @@ -96,7 +96,7 @@ checks checkpoint and bundle hashes before loading native code. For GPU validation, supply real request fixtures to `recipe/laya/native/export_encoder.py CHECKPOINT REQUESTS OUTPUT`. It uses Laya 0.3.20 and PyTorch on CUDA as the reference, independently of the Rust runtime. -Set `LAYA_CUDA_LIBRARY`, `LAYA_CHECKPOINT`, `LAYA_KERNEL_BUNDLE` and +Set `LAYA_CUDA_LIBRARY`, `LAYA_CUDA_DEVICE`, `LAYA_CHECKPOINT`, `LAYA_KERNEL_BUNDLE` and `LAYA_ENCODER_ORACLE`, then run: ```sh diff --git a/src/models/laya/src/artifacts.rs b/src/models/laya/src/artifacts.rs index 128a02b..1f52e68 100644 --- a/src/models/laya/src/artifacts.rs +++ b/src/models/laya/src/artifacts.rs @@ -64,6 +64,10 @@ pub fn validate_bundle(checkpoint: &Path, bundle: &Path) -> Result<()> { "rope_local_cos.f32", "rope_local_sin.f32", ] { + ensure!( + fs::metadata(bundle.join(name))?.len() == 512 * 32 * 4, + "invalid rotary table size: {name}" + ); check(bundle.join(name), tables["tables"][name].as_str())?; } check( diff --git a/src/models/laya/src/artifacts/tests.rs b/src/models/laya/src/artifacts/tests.rs index 84fa9cd..8da07f0 100644 --- a/src/models/laya/src/artifacts/tests.rs +++ b/src/models/laya/src/artifacts/tests.rs @@ -42,11 +42,9 @@ impl Fixture { "rope_local_cos.f32", "rope_local_sin.f32", ] { - fs::write(bundle.join(name), b"table").unwrap(); - tables.insert( - name.into(), - json!(format!("{:x}", Sha256::digest(b"table"))), - ); + let data = vec![0u8; 512 * 32 * 4]; + fs::write(bundle.join(name), &data).unwrap(); + tables.insert(name.into(), json!(format!("{:x}", Sha256::digest(&data)))); } let table_manifest = json!({ "abi": 1, @@ -94,6 +92,29 @@ fn matching_artifacts_pass_without_loading_cuda() { Fixture::new().validate().unwrap(); } +#[test] +fn wrong_sized_tables_fail_even_with_matching_hashes() { + for name in [ + "rope_full_cos.f32", + "rope_full_sin.f32", + "rope_local_cos.f32", + "rope_local_sin.f32", + ] { + for size in [0, 512 * 32 * 4 - 1, 512 * 32 * 4 + 1] { + let f = Fixture::new(); + let data = vec![0u8; size]; + fs::write(f.bundle.join(name), &data).unwrap(); + let path = f.bundle.join("tables.json"); + let mut manifest: serde_json::Value = + serde_json::from_slice(&fs::read(&path).unwrap()).unwrap(); + manifest["tables"][name] = json!(format!("{:x}", Sha256::digest(&data))); + fs::write(path, serde_json::to_vec(&manifest).unwrap()).unwrap(); + let error = f.validate().unwrap_err().to_string(); + assert!(error.contains("invalid rotary table size") && error.contains(name)); + } + } +} + #[test] fn each_checkpoint_artifact_is_bound_to_the_bundle() { for name in CHECKPOINT_ARTIFACTS { diff --git a/src/models/laya/src/encoder/tests.rs b/src/models/laya/src/encoder/tests.rs index 26799e4..345eb86 100644 --- a/src/models/laya/src/encoder/tests.rs +++ b/src/models/laya/src/encoder/tests.rs @@ -25,7 +25,8 @@ fn real_encoder_matches_official_hidden_states() -> Result<()> { .map(std::path::PathBuf::from) .expect(name) }; - let cuda = unsafe { Cuda::load(&path("LAYA_CUDA_LIBRARY"), 0) }?; + let device = std::env::var("LAYA_CUDA_DEVICE")?.parse()?; + let cuda = unsafe { Cuda::load(&path("LAYA_CUDA_LIBRARY"), device) }?; let encoder = unsafe { Encoder::load(&cuda, &path("LAYA_CHECKPOINT"), &path("LAYA_KERNEL_BUNDLE")) }?; let root = path("LAYA_ENCODER_ORACLE"); From a20b2da91805d247141c50559021442a004a0e36 Mon Sep 17 00:00:00 2001 From: linear3735 Date: Thu, 1 Oct 2026 17:54:59 +0800 Subject: [PATCH 10/15] Move CUDA resource tests to the root test directory --- src/backends/cuda/Cargo.toml | 4 ++++ .../cuda/tests => tests/backends/cuda}/fixtures/runtime.c | 0 {src/backends/cuda/tests => tests/backends/cuda}/runtime.rs | 5 ++++- 3 files changed, 8 insertions(+), 1 deletion(-) rename {src/backends/cuda/tests => tests/backends/cuda}/fixtures/runtime.c (100%) rename {src/backends/cuda/tests => tests/backends/cuda}/runtime.rs (97%) diff --git a/src/backends/cuda/Cargo.toml b/src/backends/cuda/Cargo.toml index 5d028b5..cdc7644 100644 --- a/src/backends/cuda/Cargo.toml +++ b/src/backends/cuda/Cargo.toml @@ -10,3 +10,7 @@ libloading = "0.8" [dev-dependencies] tempfile = "3" + +[[test]] +name = "runtime" +path = "../../../tests/backends/cuda/runtime.rs" diff --git a/src/backends/cuda/tests/fixtures/runtime.c b/tests/backends/cuda/fixtures/runtime.c similarity index 100% rename from src/backends/cuda/tests/fixtures/runtime.c rename to tests/backends/cuda/fixtures/runtime.c diff --git a/src/backends/cuda/tests/runtime.rs b/tests/backends/cuda/runtime.rs similarity index 97% rename from src/backends/cuda/tests/runtime.rs rename to tests/backends/cuda/runtime.rs index 1101021..ab241d4 100644 --- a/src/backends/cuda/tests/runtime.rs +++ b/tests/backends/cuda/runtime.rs @@ -32,7 +32,10 @@ impl Fixture { let result = command .args(["-fPIC", "-std=c11", "-Wall", "-Wextra", "-Werror"]) .args(defines) - .arg(PathBuf::from(env!("CARGO_MANIFEST_DIR")).join("tests/fixtures/runtime.c")) + .arg( + PathBuf::from(env!("CARGO_MANIFEST_DIR")) + .join("../../../tests/backends/cuda/fixtures/runtime.c"), + ) .arg("-o") .arg(&path) .output() From 2f5ff4be2a84729121d604a190111b2d4fa9fae9 Mon Sep 17 00:00:00 2001 From: linear3735 Date: Thu, 1 Oct 2026 17:55:38 +0800 Subject: [PATCH 11/15] Move Laya residency and workspace tests to the root test directory --- src/models/laya/src/resident.rs | 1 + src/models/laya/src/workspace.rs | 1 + .../laya/src/resident/tests.rs => tests/laya/unit/resident.rs | 2 +- .../laya/src/workspace/tests.rs => tests/laya/unit/workspace.rs | 2 +- 4 files changed, 4 insertions(+), 2 deletions(-) rename src/models/laya/src/resident/tests.rs => tests/laya/unit/resident.rs (98%) rename src/models/laya/src/workspace/tests.rs => tests/laya/unit/workspace.rs (98%) diff --git a/src/models/laya/src/resident.rs b/src/models/laya/src/resident.rs index b46b867..9c48854 100644 --- a/src/models/laya/src/resident.rs +++ b/src/models/laya/src/resident.rs @@ -81,4 +81,5 @@ fn packed(source: &Weights, spec: &TensorSpec) -> Result> { } #[cfg(test)] +#[path = "../../../../tests/laya/unit/resident.rs"] mod tests; diff --git a/src/models/laya/src/workspace.rs b/src/models/laya/src/workspace.rs index 1d2a1e9..da3c0f1 100644 --- a/src/models/laya/src/workspace.rs +++ b/src/models/laya/src/workspace.rs @@ -97,4 +97,5 @@ impl Workspace { } #[cfg(test)] +#[path = "../../../../tests/laya/unit/workspace.rs"] mod tests; diff --git a/src/models/laya/src/resident/tests.rs b/tests/laya/unit/resident.rs similarity index 98% rename from src/models/laya/src/resident/tests.rs rename to tests/laya/unit/resident.rs index b77d42e..92bf8c0 100644 --- a/src/models/laya/src/resident/tests.rs +++ b/tests/laya/unit/resident.rs @@ -30,7 +30,7 @@ impl Fixture { ]) .arg( PathBuf::from(env!("CARGO_MANIFEST_DIR")) - .join("../../backends/cuda/tests/fixtures/runtime.c"), + .join("../../../tests/backends/cuda/fixtures/runtime.c"), ) .arg("-o") .arg(&path) diff --git a/src/models/laya/src/workspace/tests.rs b/tests/laya/unit/workspace.rs similarity index 98% rename from src/models/laya/src/workspace/tests.rs rename to tests/laya/unit/workspace.rs index 2a9fd76..c4f8e18 100644 --- a/src/models/laya/src/workspace/tests.rs +++ b/tests/laya/unit/workspace.rs @@ -25,7 +25,7 @@ impl Fixture { ]) .arg( PathBuf::from(env!("CARGO_MANIFEST_DIR")) - .join("../../backends/cuda/tests/fixtures/runtime.c"), + .join("../../../tests/backends/cuda/fixtures/runtime.c"), ) .arg("-o") .arg(&path) From ff4d8fa13b2c8d52027aa2c565a0b97940d1c0ca Mon Sep 17 00:00:00 2001 From: linear3735 Date: Thu, 1 Oct 2026 17:56:00 +0800 Subject: [PATCH 12/15] Move Laya encoder tests to the root test directory --- src/models/laya/src/artifacts.rs | 1 + src/models/laya/src/encoder.rs | 1 + .../laya/src/artifacts/tests.rs => tests/laya/unit/artifacts.rs | 0 .../laya/src/encoder/tests.rs => tests/laya/unit/encoder.rs | 0 4 files changed, 2 insertions(+) rename src/models/laya/src/artifacts/tests.rs => tests/laya/unit/artifacts.rs (100%) rename src/models/laya/src/encoder/tests.rs => tests/laya/unit/encoder.rs (100%) diff --git a/src/models/laya/src/artifacts.rs b/src/models/laya/src/artifacts.rs index 1f52e68..6ed86cf 100644 --- a/src/models/laya/src/artifacts.rs +++ b/src/models/laya/src/artifacts.rs @@ -78,4 +78,5 @@ pub fn validate_bundle(checkpoint: &Path, bundle: &Path) -> Result<()> { } #[cfg(test)] +#[path = "../../../../tests/laya/unit/artifacts.rs"] mod tests; diff --git a/src/models/laya/src/encoder.rs b/src/models/laya/src/encoder.rs index 5bc8a97..26dc1d9 100644 --- a/src/models/laya/src/encoder.rs +++ b/src/models/laya/src/encoder.rs @@ -248,4 +248,5 @@ fn validate_inputs(ids: &[i64], lengths: &[i32], types: &[i64], b: usize, l: usi } #[cfg(test)] +#[path = "../../../../tests/laya/unit/encoder.rs"] mod tests; diff --git a/src/models/laya/src/artifacts/tests.rs b/tests/laya/unit/artifacts.rs similarity index 100% rename from src/models/laya/src/artifacts/tests.rs rename to tests/laya/unit/artifacts.rs diff --git a/src/models/laya/src/encoder/tests.rs b/tests/laya/unit/encoder.rs similarity index 100% rename from src/models/laya/src/encoder/tests.rs rename to tests/laya/unit/encoder.rs From df85135b846f054efae104e3afb7a6b21226d83b Mon Sep 17 00:00:00 2001 From: linear3735 Date: Sun, 4 Oct 2026 01:08:31 +0800 Subject: [PATCH 13/15] docs(laya): publish native encoder validation steps --- recipe/laya/README.md | 41 +++++++++++++++++++++++++++++++++++++++ src/models/laya/README.md | 19 +++--------------- 2 files changed, 44 insertions(+), 16 deletions(-) diff --git a/recipe/laya/README.md b/recipe/laya/README.md index 254c12c..84c9a86 100644 --- a/recipe/laya/README.md +++ b/recipe/laya/README.md @@ -56,3 +56,44 @@ if the worker requires a bearer token. See the [frontend documentation](../../src/frontend/README.md) for configuration and transport behavior. + +## Native encoder validation + +The Rust eager encoder runs on Hopper with the English Laya 0.3.20 checkpoint +(`convaiinnovations/laya`, revision `55cf4c4ebb4ebe31b2550e8bdf3bd21b99753851`). +It covers embedding, 28 encoder layers and two decision transformer layers. +Scorer, decoding, HTTP and CUDA Graphs are separate steps. + +Use a CUDA-enabled PyTorch environment with `laya==0.3.20` and its TileLang fast-path +dependencies for the reference. Build the operator bundle with the +[existing Hopper build entry](https://github.com/linear3735/system1-omni/blob/5ff41a5/src/backends/cuda/build.sh) +and export rotary tables with that version's `tools/export_tables.py`. Keep +`liblaya_cuda.so`, `build-manifest.json`, `tables.json` and the four rotary table +files in one directory. The operator bundle is separate from the resource library +built below. Load only trusted native libraries on a compatible Hopper GPU. + +Run GPU checks on an allocated device. Device 0 below means the first visible GPU. +`requests.json` is a list of `{ "name": "...", "request": { "state": ..., "questions": ... } }` +cases, each with 1–16 questions. The exporter constructs `mixed_16` from the +longest row and the first 15 other distinct token/type rows. That selection must +cover all three question types and include both a 512-token row and a shorter row. +Use a new output directory for each validation run. + +```sh +export LAYA_CHECKPOINT=/path/to/laya/snapshot +export LAYA_KERNEL_BUNDLE=/path/to/hopper-bundle +export LAYA_CUDA_DEVICE=0 +export LAYA_CUDA_LIBRARY=/tmp/liblaya_resources.so +export LAYA_ENCODER_ORACLE=/path/to/new-encoder-oracle +nvcc -shared -Xcompiler=-fPIC -O2 src/backends/cuda/kernels/runtime.cu -o "$LAYA_CUDA_LIBRARY" +python recipe/laya/native/export_encoder.py "$LAYA_CHECKPOINT" requests.json "$LAYA_ENCODER_ORACLE" +cargo test --release --locked -p omni-laya --lib real_encoder_matches_official_hidden_states -- --ignored --nocapture +``` + +The test compares selected intermediates and final hidden states, reverses request +order and repeats each workspace. It checks valid tokens and requires finite +candidate padding. This validates implementation parity, not model quality or latency. + +The [frozen H800 benchmark](https://gist.github.com/linear3735/c777b0dc449676ebf1c26f0990a90b16) +contains the five-case inputs, runner scripts, host/CUDA-event timing boundaries, +software versions and commands for the separate latency comparison. diff --git a/src/models/laya/README.md b/src/models/laya/README.md index f1d3a66..e40d9ef 100644 --- a/src/models/laya/README.md +++ b/src/models/laya/README.md @@ -86,22 +86,9 @@ bundle; CPU builds need neither library. The operator bundle currently targets Hopper `sm_90a`. `Encoder::load` is unsafe because callers must trust the native code and provide a compatible GPU; hashes bind artifacts, not code trust. -Reuse the [existing CUDA build entry](https://github.com/linear3735/system1-omni/blob/5ff41a5/src/backends/cuda/build.sh) -and its model source/export tools. Export rotary tables with that version's -`tools/export_tables.py`. Keep `liblaya_cuda.so`, `build-manifest.json`, -`tables.json` and the four rotary table files in one bundle directory. These build -tools are a separate dependency, not duplicated by this encoder change. Startup -checks checkpoint and bundle hashes before loading native code. - -For GPU validation, supply real request fixtures to -`recipe/laya/native/export_encoder.py CHECKPOINT REQUESTS OUTPUT`. It uses Laya -0.3.20 and PyTorch on CUDA as the reference, independently of the Rust runtime. -Set `LAYA_CUDA_LIBRARY`, `LAYA_CUDA_DEVICE`, `LAYA_CHECKPOINT`, `LAYA_KERNEL_BUNDLE` and -`LAYA_ENCODER_ORACLE`, then run: - -```sh -cargo test --release --locked -p omni-laya --lib real_encoder_matches_official_hidden_states -- --ignored --nocapture -``` +Build prerequisites, bundle preparation and the GPU test commands are in the +[native encoder validation recipe](../../../recipe/laya/README.md#native-encoder-validation). +Startup checks checkpoint and bundle hashes before loading operator code. The test checks selected encoder intermediates, both head layers, final hidden states, shape changes and repeated workspace reuse. It compares valid tokens; From 7574f19fc9aa3028ec6f12c236a2f580f57d8462 Mon Sep 17 00:00:00 2001 From: linear3735 Date: Sun, 4 Oct 2026 02:34:38 +0800 Subject: [PATCH 14/15] docs(laya): publish residency and workspace validation --- recipe/laya/README.md | 49 +++++++++++++++++++++++++++++++++++++++ src/models/laya/README.md | 17 +++++++------- 2 files changed, 58 insertions(+), 8 deletions(-) diff --git a/recipe/laya/README.md b/recipe/laya/README.md index 254c12c..8111a0f 100644 --- a/recipe/laya/README.md +++ b/recipe/laya/README.md @@ -56,3 +56,52 @@ if the worker requires a bearer token. See the [frontend documentation](../../src/frontend/README.md) for configuration and transport behavior. + +## Native residency and workspace validation + +These opt-in Rust checks test CUDA allocations and transfers. They do not need the +Python worker or frontend and do not test inference, model outputs or latency. +The normal CPU tests skip them. + +Use a Linux host with an approved CUDA GPU, a working NVIDIA driver, the CUDA +toolkit (`nvcc`) and Rust. Build the trusted resource library from this checkout; +it uses ABI version 1 and needs neither TileLang nor cuBLAS. `LAYA_CUDA_DEVICE` is +the approved device ordinal after `CUDA_VISIBLE_DEVICES` filtering. + +```sh +export LAYA_CUDA_LIBRARY=/tmp/liblaya-resources.so +export LAYA_CUDA_DEVICE=0 +nvcc -shared -Xcompiler=-fPIC -O2 src/backends/cuda/kernels/runtime.cu \ + -o "$LAYA_CUDA_LIBRARY" +``` + +The workspace check needs no checkpoint. It writes and reads all 17 buffers twice +at `(batch, sequence) = (1, 16), (1, 512), (16, 512)`, using deterministic byte +patterns: + +```sh +cargo test --release --locked -p omni-laya --lib \ + workspace::tests::real_gpu_workspace_capacity_and_reuse \ + -- --ignored --exact --nocapture +``` + +For the residency check, use an unchanged local snapshot of +`convaiinnovations/laya` revision `55cf4c4ebb4ebe31b2550e8bdf3bd21b99753851`, +including `model.safetensors`. Generate the oracle from that same snapshot in a +Python environment with PyTorch, safetensors and NumPy: + +```sh +export LAYA_CHECKPOINT=/path/to/laya/snapshot +export LAYA_WEIGHT_ORACLE=/tmp/laya-weight-oracle.json +python recipe/laya/native/export_weights.py "$LAYA_CHECKPOINT" "$LAYA_WEIGHT_ORACLE" +cargo test --release --locked -p omni-laya --lib \ + resident::tests::real_checkpoint_residency_matches_torch \ + -- --ignored --exact --nocapture +``` + +The residency check validates all 206 checkpoint tensors, uploads the 205 used +tensors and compares readback hashes with the Torch conversion oracle. The legacy +`temperature` buffer is validated but not uploaded. Reported allocation bytes +exclude CUDA context and library overhead. See the +[model contracts](https://github.com/linear3735/system1-omni/blob/codex/laya-workspace/src/models/laya/README.md) for storage precision, +workspace layouts and ownership. diff --git a/src/models/laya/README.md b/src/models/laya/README.md index dab0230..015d40f 100644 --- a/src/models/laya/README.md +++ b/src/models/laya/README.md @@ -35,10 +35,10 @@ excluding CUDA context and allocator overhead. Buffers keep their CUDA context a after the caller drops the source mapping or `Cuda`. Failed loads release partial allocations. Workspace, rotary tables and inference are separate modules. -The ignored `real_checkpoint_residency_matches_torch` test requires -`LAYA_CUDA_LIBRARY`, `LAYA_CUDA_DEVICE`, `LAYA_CHECKPOINT` and `LAYA_WEIGHT_ORACLE`. -It uploads all 205 used tensors and compares readback hashes with the existing Torch -oracle; it does not test model outputs or latency. +The ignored GPU check uploads all 205 used tensors and compares readback hashes +with the Torch conversion oracle; it does not test model outputs or latency. +Prerequisites and commands are in the +[native validation recipe](../../../recipe/laya/README.md#native-residency-and-workspace-validation). ## Inference workspace @@ -66,10 +66,11 @@ At `(B,L)=(1,512)`, allocations total 22,628,896 bytes; at `(16,512)`, 236,048,836 bytes. The caller must enforce at most `MAX_MARKERS` scored positions. No Graph cache, kernel launch, cuBLAS workspace or inference is included. -The ignored `real_gpu_workspace_capacity_and_reuse` test requires -`LAYA_CUDA_LIBRARY` and `LAYA_CUDA_DEVICE`. It writes and reads all 17 buffers twice -at three shapes, including both capacity bounds, using deterministic byte patterns. -This tests allocation and transfer, not model numerics or latency. +The ignored GPU check writes and reads all 17 buffers twice at three shapes, +including both capacity bounds, using deterministic byte patterns. This tests +allocation and transfer, not model numerics or latency. Prerequisites and commands +are in the +[native validation recipe](../../../recipe/laya/README.md#native-residency-and-workspace-validation). ## Python worker From 047d018c144588a31a1840606fcb0b0d85c29844 Mon Sep 17 00:00:00 2001 From: linear3735 Date: Sun, 4 Oct 2026 02:34:38 +0800 Subject: [PATCH 15/15] docs(laya): include parent resource validation steps --- recipe/laya/README.md | 49 +++++++++++++++++++++++++++++++++++++++ src/models/laya/README.md | 17 +++++++------- 2 files changed, 58 insertions(+), 8 deletions(-) diff --git a/recipe/laya/README.md b/recipe/laya/README.md index 84c9a86..dc9137a 100644 --- a/recipe/laya/README.md +++ b/recipe/laya/README.md @@ -57,6 +57,55 @@ if the worker requires a bearer token. See the [frontend documentation](../../src/frontend/README.md) for configuration and transport behavior. +## Native residency and workspace validation + +These opt-in Rust checks test CUDA allocations and transfers. They do not need the +Python worker or frontend and do not test inference, model outputs or latency. +The normal CPU tests skip them. + +Use a Linux host with an approved CUDA GPU, a working NVIDIA driver, the CUDA +toolkit (`nvcc`) and Rust. Build the trusted resource library from this checkout; +it uses ABI version 1 and needs neither TileLang nor cuBLAS. `LAYA_CUDA_DEVICE` is +the approved device ordinal after `CUDA_VISIBLE_DEVICES` filtering. + +```sh +export LAYA_CUDA_LIBRARY=/tmp/liblaya-resources.so +export LAYA_CUDA_DEVICE=0 +nvcc -shared -Xcompiler=-fPIC -O2 src/backends/cuda/kernels/runtime.cu \ + -o "$LAYA_CUDA_LIBRARY" +``` + +The workspace check needs no checkpoint. It writes and reads all 17 buffers twice +at `(batch, sequence) = (1, 16), (1, 512), (16, 512)`, using deterministic byte +patterns: + +```sh +cargo test --release --locked -p omni-laya --lib \ + workspace::tests::real_gpu_workspace_capacity_and_reuse \ + -- --ignored --exact --nocapture +``` + +For the residency check, use an unchanged local snapshot of +`convaiinnovations/laya` revision `55cf4c4ebb4ebe31b2550e8bdf3bd21b99753851`, +including `model.safetensors`. Generate the oracle from that same snapshot in a +Python environment with PyTorch, safetensors and NumPy: + +```sh +export LAYA_CHECKPOINT=/path/to/laya/snapshot +export LAYA_WEIGHT_ORACLE=/tmp/laya-weight-oracle.json +python recipe/laya/native/export_weights.py "$LAYA_CHECKPOINT" "$LAYA_WEIGHT_ORACLE" +cargo test --release --locked -p omni-laya --lib \ + resident::tests::real_checkpoint_residency_matches_torch \ + -- --ignored --exact --nocapture +``` + +The residency check validates all 206 checkpoint tensors, uploads the 205 used +tensors and compares readback hashes with the Torch conversion oracle. The legacy +`temperature` buffer is validated but not uploaded. Reported allocation bytes +exclude CUDA context and library overhead. See the +[model contracts](https://github.com/linear3735/system1-omni/blob/codex/laya-encoder/src/models/laya/README.md) for storage precision, +workspace layouts and ownership. + ## Native encoder validation The Rust eager encoder runs on Hopper with the English Laya 0.3.20 checkpoint diff --git a/src/models/laya/README.md b/src/models/laya/README.md index e40d9ef..008f071 100644 --- a/src/models/laya/README.md +++ b/src/models/laya/README.md @@ -35,10 +35,10 @@ excluding CUDA context and allocator overhead. Buffers keep their CUDA context a after the caller drops the source mapping or `Cuda`. Failed loads release partial allocations. Workspace, rotary tables and inference are separate modules. -The ignored `real_checkpoint_residency_matches_torch` test requires -`LAYA_CUDA_LIBRARY`, `LAYA_CUDA_DEVICE`, `LAYA_CHECKPOINT` and `LAYA_WEIGHT_ORACLE`. -It uploads all 205 used tensors and compares readback hashes with the existing Torch -oracle; it does not test model outputs or latency. +The ignored GPU check uploads all 205 used tensors and compares readback hashes +with the Torch conversion oracle; it does not test model outputs or latency. +Prerequisites and commands are in the +[native validation recipe](../../../recipe/laya/README.md#native-residency-and-workspace-validation). ## Inference workspace @@ -66,10 +66,11 @@ At `(B,L)=(1,512)`, allocations total 22,628,896 bytes; at `(16,512)`, 236,048,836 bytes. The caller must enforce at most `MAX_MARKERS` scored positions. No Graph cache, kernel launch, cuBLAS workspace or inference is included. -The ignored `real_gpu_workspace_capacity_and_reuse` test requires -`LAYA_CUDA_LIBRARY` and `LAYA_CUDA_DEVICE`. It writes and reads all 17 buffers twice -at three shapes, including both capacity bounds, using deterministic byte patterns. -This tests allocation and transfer, not model numerics or latency. +The ignored GPU check writes and reads all 17 buffers twice at three shapes, +including both capacity bounds, using deterministic byte patterns. This tests +allocation and transfer, not model numerics or latency. Prerequisites and commands +are in the +[native validation recipe](../../../recipe/laya/README.md#native-residency-and-workspace-validation). ## Eager encoder