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/README.md

@@ -56,3 +56,93 @@ 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
+(`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.

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));
+}
+}