Keep Laya weights and workspace on CUDA

Cargo.toml

@@ -1,3 +1,3 @@
 [workspace]
-members = ["src/frontend", "src/models/cua_s1/native", "src/models/laya"]
+members = ["src/frontend", "src/models/cua_s1/native", "src/models/laya", "src/backends/cuda"]
 resolver = "3"

src/backends/cuda/Cargo.toml

@@ -0,0 +1,16 @@
+[package]
+name = "omni-cuda"
+version = "0.1.0"
+edition = "2024"
+publish = false
+
+[dependencies]
+anyhow = "1"
+libloading = "0.8"
+
+[dev-dependencies]
+tempfile = "3"
+
+[[test]]
+name = "runtime"
+path = "../../../tests/backends/cuda/runtime.rs"

src/backends/cuda/README.md

@@ -1,7 +1,56 @@
 # 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.
+[`qwen3_5/`](qwen3_5/) provides the prefill-only Qwen3.5 operations used by the
+Cua-S1 native worker, measured on sm_89.
 
-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.
+## Laya resources
 
-Status: [`qwen3_5/`](qwen3_5/) has the operations of a prefill-only Qwen3.5 forward pass, used by the Cua-S1 native worker and measured on sm_89. Other models are planned.
+`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.
+
+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.
+
+### 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.

src/backends/cuda/kernels/runtime.cu

@@ -0,0 +1,65 @@
+#include <cuda_runtime.h>
+#include <stdint.h>
+
+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<cudaError_t>(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<cudaStream_t>(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<cudaStream_t>(stream));
+}
+
+int laya_sync(void* stream) {
+  if (!stream) return invalid_argument;
+  return cudaStreamSynchronize(static_cast<cudaStream_t>(stream));
+}
+
+int laya_stream_free(void* stream) {
+  if (!stream) return invalid_argument;
+  return cudaStreamDestroy(static_cast<cudaStream_t>(stream));
+}
+}

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<Context>,
+}
+
+impl Cuda {
+    /// # Safety
+    /// `path` must name a trusted library implementing the complete runtime ABI.
+    pub unsafe fn load(path: &Path, device: i32) -> Result<Self> {
+        let library = unsafe { Library::new(path) }?;
+        let version =
+            unsafe { library.get::<unsafe extern "C" fn() -> 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<Buffer> {
+        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<Buffer> {
+        let buffer = self.alloc(bytes.len())?;
+        buffer.write(bytes)?;
+        Ok(buffer)
+    }
+
+    pub fn sync(&self) -> Result<()> {
+        self.ctx.sync()
+    }
+}
+
+pub struct Buffer {
+    ctx: Rc<Context>,
+    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<Vec<u8>> {
+        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);
+            }
+        }
+    }
+}

src/models/laya/Cargo.toml

@@ -8,11 +8,13 @@ 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"