Run Laya encoder and decision layers 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"

recipe/laya/native/export_encoder.py

@@ -0,0 +1,109 @@
+"""Export official Laya intermediates for the Rust encoder integration test."""
+import argparse
+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())
+    args.output.mkdir(parents=True, exist_ok=True)
+    inputs, rows = [], {}
+    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()
+
+        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())
+
+        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(), "origins": origins})
+        print("REFERENCE", name, b, l, flush=True)
+    (args.output / "cases.json").write_text(json.dumps(records))
+
+
+if __name__ == "__main__":
+    main()

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,67 @@
 # 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.
+
+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
+
+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.
+
+## 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.

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/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<String, Launch>,
+}
+
+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<Self> {
+        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::<Launch>(symbol.as_bytes())? };
+            functions.insert((*name).to_owned(), launch);
+        }
+        let init = unsafe { library.get::<unsafe extern "C" fn() -> 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();
+    }
+}