From d717e497fac143cae37f9aa8c18bab00e3497c8b Mon Sep 17 00:00:00 2001 From: Tianyao Wu Date: Tue, 6 Oct 2026 13:48:45 +0800 Subject: [PATCH 1/4] qwen3_5: add a GEMM handle with one algorithm per weight shape cuBLASLt's heuristic picks an algorithm per M, so a row's result can change with the number of rows in the call, and a prefix and its branch run separately round differently from one pass over the same tokens. Add cs1_gemm_create_fixed: one algorithm per weight shape (N, K, ldy), the heuristic's first choice at a reference M among algorithms without split-K, for every M. The default handle is unchanged. Bump the ABI to 7 and add a GPU test that each projection's rows match across M. --- docs/supported-models.md | 2 +- recipe/cua_s1/native.md | 2 +- recipe/jev_vl/README.md | 4 +- recipe/open_jev/native.md | 2 +- src/backends/cuda/contract.md | 4 +- src/backends/cuda/qwen3_5/README.md | 19 ++++- src/backends/cuda/qwen3_5/gemm.cu | 67 +++++++++++++-- src/backends/cuda/qwen3_5/ops.h | 8 +- src/models/qwen3_5/native/src/cuda.rs | 3 +- tests/qwen3_5/kernels.rs | 117 ++++++++++++++++++++++++++ 10 files changed, 210 insertions(+), 18 deletions(-) diff --git a/docs/supported-models.md b/docs/supported-models.md index 34d7b465..55249138 100644 --- a/docs/supported-models.md +++ b/docs/supported-models.md @@ -17,7 +17,7 @@ Models that are being added are also tracked in issues labeled [new model](https | Cua-S1 4B 0.2, `multimodal` adapter | Reference worker on Transformers and PEFT, [`src/frontend/cua_s1.py`](../src/frontend/cua_s1.py); no recipe yet | Not supported | Validated ([#17](https://github.com/ThinkFlowLab/system1-omni/pull/17), [#18](https://github.com/ThinkFlowLab/system1-omni/pull/18)) | Not supported | The state is one PNG or JPEG image; upstream's `weights.lock.json` next to the base weights | | Open-Jev-27B-v1.1 | [Native Rust/CUDA worker](../recipe/open_jev/native.md) on the shared Qwen3.5/3.8 executor | Not supported | Validated on H200 (sm_90) for the [74 single-candidate workload](../recipe/open_jev/validation.md) | Not supported | Compute capability 8.0 or newer, CUDA toolkit to build, exported merged weights and trained head | | Open-Jev-9B | The same [native Rust/CUDA worker](../recipe/open_jev/native.md) | Not supported | Validated on compute capability 8.9 against the reference for [253 requests](../recipe/open_jev/validation-9b.md) | Not supported | Compute capability 8.0 or newer, CUDA toolkit to build, exported merged weights and trained head | -| autotrust/JEV-27B-VL | [Experimental native worker](../recipe/jev_vl/README.md), single-question API | Contract tests only; no CPU inference | H800 frozen-corpus check; [experimental scope and limits](../recipe/jev_vl/validation.md) | Not supported | Pinned merged checkpoint, rebuilt ABI 6 CUDA library; image inputs require offline Transformers preencoding | +| autotrust/JEV-27B-VL | [Experimental native worker](../recipe/jev_vl/README.md), single-question API | Contract tests only; no CPU inference | H800 frozen-corpus check; [experimental scope and limits](../recipe/jev_vl/validation.md) | Not supported | Pinned merged checkpoint, rebuilt ABI 7 CUDA library; image inputs require offline Transformers preencoding | | CLM-v0.1-8B | [External `clm-serve` recipe](../recipe/clm/README.md) with a CPU stub embeddings server | **Stub-encoder contract checks only** ([#23](https://github.com/ThinkFlowLab/system1-omni/pull/23)); not real Qwen3-8B decisions | Real encoder unverified by the merged recipe | Unverified | Python, upstream CLM and head checkpoint; a real encoder requires a separate embeddings server | - **Validated:** covered by the recipe on `main` or by the checks in the linked merged pull request. diff --git a/recipe/cua_s1/native.md b/recipe/cua_s1/native.md index 83322b25..55e43929 100644 --- a/recipe/cua_s1/native.md +++ b/recipe/cua_s1/native.md @@ -28,7 +28,7 @@ each exact prompt length warms the GEMM plans and captures the forward pass; later requests replay it with freshly uploaded token ids. At most eight lengths are cached. Growing the scratch allocation clears the captures before freeing their buffers. Capture adds first-use latency; leave the variable unset to use -the eager control. Rebuild both the worker and CUDA library together (ABI 6). +the eager control. Rebuild both the worker and CUDA library together (ABI 7). If capture fails, the worker returns the completed eager result and disables Graph capture/replay for its remaining lifetime, logging the failure to stderr. diff --git a/recipe/jev_vl/README.md b/recipe/jev_vl/README.md index 6304c364..03ddfe1e 100644 --- a/recipe/jev_vl/README.md +++ b/recipe/jev_vl/README.md @@ -39,7 +39,7 @@ The author measured the worker on one H800 80 GB (`sm_90`). A maintainer also on an earlier revision. These runs do not validate later code changes or maximum context lengths. Use an allocated GPU on scheduled hosts. Build the CUDA library and Rust workers from the same revision; -the integrated backend uses **ABI 6** and older libraries must be rebuilt. +the integrated backend uses **ABI 7** and older libraries must be rebuilt. ```sh src/backends/cuda/qwen3_5/build.sh target/release 90 @@ -141,7 +141,7 @@ The image-prefix tokenizer check requires the exported checkpoint. First [restore the frozen corpus](validation.md#restore-the-frozen-corpus). It uses the corpus's fixed `[1, 60, 60]` grid and synthetic embedding rows, so it checks token/position splitting rather than the vision encoder. -CUDA kernel tests require an allocated GPU and rebuilt ABI 6 library: +CUDA kernel tests require an allocated GPU and rebuilt ABI 7 library: ```sh JEV_VL_EXPORT=$PWD/weights/jev-vl-merged \ diff --git a/recipe/open_jev/native.md b/recipe/open_jev/native.md index 761c4f28..331f0d1e 100644 --- a/recipe/open_jev/native.md +++ b/recipe/open_jev/native.md @@ -67,7 +67,7 @@ or an incomplete export. The saved limit defaults to 4096 tokens per candidate; The CUDA kernels require compute capability 8.0 or newer. The current build target below is Ada (`89`); pass your GPU's compute capability explicitly. The CUDA shared library and Rust workers must be rebuilt together for ABI -version 6, which includes vision, CUDA Graph and prefix-continuation entry points. +version 7, which includes vision, CUDA Graph, prefix-continuation and fixed-GEMM entry points. ```sh src/backends/cuda/qwen3_5/build.sh target/release 89 diff --git a/src/backends/cuda/contract.md b/src/backends/cuda/contract.md index 41301af1..e8ccade6 100644 --- a/src/backends/cuda/contract.md +++ b/src/backends/cuda/contract.md @@ -49,7 +49,7 @@ change. Schema: ```json { "name": "qwen3_5", - "abi_version": 4, + "abi_version": 7, "status": "validated", "sources": ["common.cuh", "mma.cuh", "ops.h", "norm.cu", "elementwise.cu", "attention.cu", "gdn_prefill.cu", "gemm.cu", "runtime.cu"], @@ -136,7 +136,7 @@ reference for the shape. The contract fixes four points: whose value it does not know. The manifest's `abi_version` is **that number**, not a version of this document, and the checker reads the macro out of the declared sources and requires the two to agree. `#19`'s `ops.h` says - `CS1_ABI_VERSION 4` today, so a `qwen3_5` manifest declares 4. + `CS1_ABI_VERSION 7` today, so a `qwen3_5` manifest declares 7. Two backends may declare different values. They are independent libraries, and the repository layout says CUDA and Metal implementations need not share diff --git a/src/backends/cuda/qwen3_5/README.md b/src/backends/cuda/qwen3_5/README.md index 63825533..e1745051 100644 --- a/src/backends/cuda/qwen3_5/README.md +++ b/src/backends/cuda/qwen3_5/README.md @@ -11,8 +11,20 @@ The norm, elementwise and q/k preparation kernels round to bfloat16 where Transf `cs1_attention_gated` fuses the sigmoid gate into the attention epilogue, preserving the BF16 rounding of both attention and sigmoid before multiplication. The native workers use this entry point; the separate operations remain available for kernel -comparisons. Rebuild the library and workers together for ABI version 6, which -includes vision, CUDA Graph, gated attention and prefix-continuation entry points. +comparisons. Rebuild the library and workers together for ABI version 7, which +includes vision, CUDA Graph, gated attention, prefix-continuation entry points and +the fixed-algorithm GEMM handle below. + +cuBLASLt's heuristic picks a GEMM algorithm per M, so a row's result can change with +the number of rows in the call: a prefix and its branch, run separately, round +differently from the same tokens in one pass. `cs1_gemm_create_fixed` returns a handle +that keeps one algorithm per weight shape for every M, the heuristic's first choice at a +reference M among algorithms without split-K, so each row's result is the same whatever +M is and wherever the row sits. cuBLASLt doesn't document this property; +`fixed_gemm_rows_do_not_depend_on_m` checks it, so run it on a new GPU before relying +on it (it passed on sm_89). It is slower for some shapes and faster for others: on an +RTX 6000 Ada, the down and output projections take up to about 4 times as long at small +M without split-K. `cs1_gemm_create` remains the default. Gated DeltaNet preparation stores converted TF32 operands in three-byte component planes, preserves the original four-term TF32 accumulation, and writes U/W fragments directly as bfloat16. Dynamic shared memory is 72 KiB per block. The [H200 comparison](../../../../benchmarks/gdn/README.md) records complete GDN call latency, numerical checks, and the small end-to-end change measured with the Open-Jev worker from PR #55. @@ -24,6 +36,7 @@ entry points. Open-Jev uses this path within requests; Cua-S1 keeps single-promp Prefix continuation retains full-attention KV, three convolution input rows, and FP32 Gated DeltaNet state at 64-token chunk boundaries. It is used by the experimental JEV-VL worker; existing full-prefill callers keep their own paths. -ABI 6 adds mandatory copy and prefix entry points after ABI 5 native vision. +ABI 6 adds mandatory copy and prefix entry points after ABI 5 native vision, +and ABI 7 adds the fixed-algorithm GEMM handle. Rebuild the library and **all** Qwen worker binaries together; older libraries are rejected. Current-branch GPU regression is required before release. diff --git a/src/backends/cuda/qwen3_5/gemm.cu b/src/backends/cuda/qwen3_5/gemm.cu index 154ad698..f0febca4 100644 --- a/src/backends/cuda/qwen3_5/gemm.cu +++ b/src/backends/cuda/qwen3_5/gemm.cu @@ -6,7 +6,14 @@ // // Each shape uses cuBLASLt's first heuristic choice, excluding split-K reductions that // accumulate into the output in place, since their order, and so the rounding, is not -// fixed. +// fixed (vision's FP32 and biased GEMMs exclude all split-K). That choice depends on M, +// so a row's result can change with the number of rows in the call. A handle from +// cs1_gemm_create_fixed instead keeps one algorithm per weight shape (N, K, ldy, and for +// vision's GEMMs the data type and bias) for every M: the heuristic's first choice at a +// reference M among algorithms without split-K. Each output row then takes the same +// path whatever the other rows, so its result does not depend on M or on its row index. +// cuBLASLt does not document this; tests/qwen3_5/kernels.rs checks it on the GPU it +// runs on. #include #include @@ -28,7 +35,9 @@ struct Gemm { cublasLtHandle_t handle = nullptr; void* workspace = nullptr; size_t workspace_bytes = 0; + int reference_m = 0; // > 0: one algorithm per weight shape, chosen at this M std::map plans; + std::map, cublasLtMatmulAlgo_t> fixed; // N, K, ldy, FP32, bias }; int status(cublasStatus_t s) { return s == CUBLAS_STATUS_SUCCESS ? 0 : 1000 + (int)s; } @@ -58,15 +67,13 @@ int describe(int M, int N, int K, int ldy, Plan& p, bool fp32, bool bias) { return 0; } -// The heuristic's first choice. Vision disables all split-K to avoid BF16 -// intermediate reductions; existing language GEMMs exclude only in-place reductions. -int first_choice(Gemm& g, Plan& p, bool vision) { +// The heuristic's first choice among the given reduction schemes. +int first_choice(Gemm& g, Plan& p, uint32_t schemes) { cublasLtMatmulPreference_t pref; cublasStatus_t s = cublasLtMatmulPreferenceCreate(&pref); if (s != CUBLAS_STATUS_SUCCESS) return status(s); cublasLtMatmulPreferenceSetAttribute(pref, CUBLASLT_MATMUL_PREF_MAX_WORKSPACE_BYTES, &g.workspace_bytes, sizeof(g.workspace_bytes)); - const uint32_t schemes = vision ? CUBLASLT_REDUCTION_SCHEME_NONE : (CUBLASLT_REDUCTION_SCHEME_MASK & ~CUBLASLT_REDUCTION_SCHEME_INPLACE); cublasLtMatmulPreferenceSetAttribute(pref, CUBLASLT_MATMUL_PREF_REDUCTION_SCHEME_MASK, &schemes, sizeof(schemes)); cublasLtMatmulHeuristicResult_t r{}; @@ -79,6 +86,40 @@ int first_choice(Gemm& g, Plan& p, bool vision) { return 0; } +// The weight shape's algorithm, chosen at the reference M on first use. An M it cannot +// serve is an error, not a reason to switch algorithms. +int fixed_choice(Gemm& g, int N, int K, int ldy, bool fp32, bool bias, Plan& p) { + const std::tuple key{N, K, ldy, fp32, bias}; + auto it = g.fixed.find(key); + if (it == g.fixed.end()) { + Plan r; + int rc = describe(g.reference_m, N, K, ldy, r, fp32, bias); + if (rc == 0) rc = first_choice(g, r, CUBLASLT_REDUCTION_SCHEME_NONE); + if (rc == 0) { + // Check that the choice really has no split-K reduction. + int32_t splits = 0; + uint32_t scheme = 0; + size_t written = 0; + const bool read = + cublasLtMatmulAlgoConfigGetAttribute(&r.algo, CUBLASLT_ALGO_CONFIG_SPLITK_NUM, &splits, + sizeof(splits), &written) == CUBLAS_STATUS_SUCCESS && + cublasLtMatmulAlgoConfigGetAttribute(&r.algo, CUBLASLT_ALGO_CONFIG_REDUCTION_SCHEME, &scheme, + sizeof(scheme), &written) == CUBLAS_STATUS_SUCCESS; + if (!read || splits != 1 || scheme != CUBLASLT_REDUCTION_SCHEME_NONE) + rc = status(CUBLAS_STATUS_NOT_SUPPORTED); + } + if (rc == 0) it = g.fixed.emplace(key, r.algo).first; + destroy(r); + if (rc != 0) return rc; + } + p.algo = it->second; + cublasLtMatmulHeuristicResult_t check{}; + const cublasStatus_t s = cublasLtMatmulAlgoCheck(g.handle, p.op, p.a, p.b, p.c, p.c, &p.algo, &check); + if (s != CUBLAS_STATUS_SUCCESS) return status(s); + if (check.workspaceSize > g.workspace_bytes) return status(CUBLAS_STATUS_NOT_SUPPORTED); + return 0; +} + // The plan for a shape, created on first use. int plan_for(Gemm& g, int M, int N, int K, int ldy, Plan*& out, bool fp32 = false, bool bias = false) { const Key key{M, N, K, ldy, fp32, bias}; @@ -89,7 +130,14 @@ int plan_for(Gemm& g, int M, int N, int K, int ldy, Plan*& out, bool fp32 = fals } Plan p; int rc = describe(M, N, K, ldy, p, fp32, bias); - if (rc == 0) rc = first_choice(g, p, fp32 || bias); + if (rc == 0 && g.reference_m > 0) { + rc = fixed_choice(g, N, K, ldy, fp32, bias, p); + } else if (rc == 0) { + // Vision's FP32 and biased GEMMs exclude all split-K to avoid BF16 intermediate + // reductions; the language GEMMs exclude only in-place reductions. + rc = first_choice(g, p, fp32 || bias ? CUBLASLT_REDUCTION_SCHEME_NONE + : CUBLASLT_REDUCTION_SCHEME_MASK & ~CUBLASLT_REDUCTION_SCHEME_INPLACE); + } if (rc != 0) { destroy(p); return rc; @@ -100,6 +148,13 @@ int plan_for(Gemm& g, int M, int N, int K, int ldy, Plan*& out, bool fp32 = fals } // namespace +extern "C" void* cs1_gemm_create_fixed(size_t workspace_bytes, int reference_m) { + if (reference_m <= 0) return nullptr; + Gemm* g = static_cast(cs1_gemm_create(workspace_bytes)); + if (g) g->reference_m = reference_m; + return g; +} + extern "C" void* cs1_gemm_create(size_t workspace_bytes) { Gemm* g = new Gemm(); if (cublasLtCreate(&g->handle) != CUBLAS_STATUS_SUCCESS || diff --git a/src/backends/cuda/qwen3_5/ops.h b/src/backends/cuda/qwen3_5/ops.h index 50d3618a..5a4d03cd 100644 --- a/src/backends/cuda/qwen3_5/ops.h +++ b/src/backends/cuda/qwen3_5/ops.h @@ -13,7 +13,7 @@ #include // Bumped whenever the required interface below changes. -#define CS1_ABI_VERSION 6 +#define CS1_ABI_VERSION 7 #ifdef __cplusplus extern "C" { @@ -118,6 +118,12 @@ int cs1_silu_mul(const void* gate_up, int ld, void* out, int T, int I, void* str // y [M, N] (rows of ldy) = x [M, K] * w [N, K]^T through cuBLASLt, float32 accumulation, // with cuBLASLt's first heuristic choice for each shape (see gemm.cu). void* cs1_gemm_create(size_t workspace_bytes); +// A handle that keeps one algorithm per weight shape (N, K, ldy) for every M, chosen at +// reference_m rows among algorithms without split-K, so that a row's result depends +// neither on M nor on its row index (checked by tests/qwen3_5/kernels.rs on the GPU it +// runs on). Null if reference_m <= 0 or setup fails. cs1_gemm with this handle returns a +// cuBLAS status for an M the algorithm cannot serve, rather than switching algorithms. +void* cs1_gemm_create_fixed(size_t workspace_bytes, int reference_m); void cs1_gemm_destroy(void* gemm); int cs1_gemm(void* gemm, const void* x, const void* w, void* y, int M, int N, int K, int ldy, void* stream); diff --git a/src/models/qwen3_5/native/src/cuda.rs b/src/models/qwen3_5/native/src/cuda.rs index 5e8141de..b9194e8d 100644 --- a/src/models/qwen3_5/native/src/cuda.rs +++ b/src/models/qwen3_5/native/src/cuda.rs @@ -9,7 +9,7 @@ use std::sync::OnceLock; use anyhow::{Context, Result, bail, ensure}; /// `CS1_ABI_VERSION` in ops.h. -const ABI_VERSION: u32 = 6; +const ABI_VERSION: u32 = 7; pub const LIBRARY: &str = "libqwen3_5_cuda.so"; /// A `cudaStream_t`. @@ -125,6 +125,7 @@ api! { cs1_sigmoid_gate(x: *mut c_void, gate: *const c_void, n: usize, stream: Stream) -> c_int; cs1_silu_mul(gate_up: *const c_void, ld: c_int, out: *mut c_void, t: c_int, i: c_int, stream: Stream) -> c_int; cs1_gemm_create(workspace_bytes: usize) -> *mut c_void; + cs1_gemm_create_fixed(workspace_bytes: usize, reference_m: c_int) -> *mut c_void; cs1_gemm_destroy(gemm: *mut c_void); cs1_gemm( gemm: *mut c_void, x: *const c_void, w: *const c_void, y: *mut c_void, m: c_int, n: c_int, k: c_int, diff --git a/tests/qwen3_5/kernels.rs b/tests/qwen3_5/kernels.rs index 7a217c81..cce51549 100644 --- a/tests/qwen3_5/kernels.rs +++ b/tests/qwen3_5/kernels.rs @@ -1086,3 +1086,120 @@ fn gdn_prefill_two_phase_matches_one_shot_bit_for_bit() { ); } } + +#[test] +#[ignore = "needs a GPU and CUA_S1_CUDA_LIB"] +fn fixed_gemm_rows_do_not_depend_on_m() { + let st = setup(); + // Every projection of the 4B, 9B and 27B backbones as (N, K): the GDN input and + // output, the attention input and output, and the MLP's gate|up and down. + let shapes = [ + (12352usize, 2560usize), + (2560, 4096), + (10240, 2560), + (18432, 2560), + (2560, 9216), + (12352, 4096), + (4096, 4096), + (10240, 4096), + (24576, 4096), + (4096, 12288), + (16480, 5120), + (5120, 6144), + (14336, 5120), + (34816, 5120), + (5120, 17408), + ]; + let max_m = 4096usize; + // SAFETY: plain handle creation; checked for null below. + let fixed = unsafe { (api().cs1_gemm_create_fixed)(32 << 20, 64) }; + assert!(!fixed.is_null()); + // SAFETY: a non-positive reference M is refused before anything is allocated. + assert!(unsafe { (api().cs1_gemm_create_fixed)(32 << 20, 0) }.is_null()); + let nan = vec![0xffu8; max_m * 34816 * 2]; + for (i, (n, k)) in shapes.into_iter().enumerate() { + let x = random(max_m * k, 70 + i as u64, 2.0); + let w = random(n * k, 90 + i as u64, 0.05); + let (xd, wd) = (to_device(&x, st), to_device(&w, st)); + let y = DeviceBuffer::new(max_m * n * 2).unwrap(); + // the bfloat16 output of m rows starting at input row `at`, over a NaN-filled buffer + let rows = |at: usize, m: usize| { + let mut out = vec![0u8; m * n * 2]; + // SAFETY: x holds max_m rows of k, w is [n, k], y has room for max_m rows of n, + // and at + m <= max_m. + unsafe { + cuda::upload(y.at(0), &nan[..max_m * n * 2], st).unwrap(); + check( + (api().cs1_gemm)( + fixed, + xd.at(at * k * 2), + wd.at(0), + y.at(0), + m as i32, + n as i32, + k as i32, + n as i32, + st, + ), + "fixed gemm", + ) + .unwrap(); + cuda::download(&mut out, y.at(0), st).unwrap(); + } + out + }; + let all = rows(0, max_m); + // every output was written (the comparisons below are byte for byte against this) + let (values, _) = all.as_chunks::<2>(); + assert!(values.iter().all(|b| !bf16::from_le_bytes(*b).is_nan())); + let row = |r: usize| &all[r * n * 2..(r + 1) * n * 2]; + // fewer rows, and rows that sit lower in the full call, as a branch's tokens do + for (at, m) in [ + (0usize, 1usize), + (0, 2), + (0, 3), + (0, 8), + (0, 63), + (0, 64), + (0, 65), + (0, 256), + (0, 1000), + (0, 1024), + (0, 3109), + (1, 2), + (37, 64), + (64, 65), + (1000, 200), + (3000, 1096), + ] { + assert!( + rows(at, m) == all[at * n * 2..(at + m) * n * 2], + "rows {at}..{} changed at {n} x {k}", + at + m + ); + } + // the first and last rows against float64, so the product itself is right + for r in [0, max_m - 1] { + let want: Vec = (0..n) + .map(|j| { + (0..k) + .map(|c| x[r * k + c].to_f64() * w[j * k + c].to_f64()) + .sum() + }) + .collect(); + let (got, _) = row(r).as_chunks::<2>(); + let scale = want.iter().fold(0f64, |m, v| m.max(v.abs())); + let worst = got + .iter() + .zip(&want) + .map(|(b, v)| (bf16::from_le_bytes(*b).to_f64() - v).abs()) + .fold(0f64, f64::max); + assert!( + worst <= 1e-2 * scale, + "{n} x {k}, row {r}: {worst} vs {scale}" + ); + } + } + // SAFETY: created above and not destroyed before. + unsafe { (api().cs1_gemm_destroy)(fixed) }; +} From 2ff1396ab7c4a52c2f89d591269d9f2a403c7e3f Mon Sep 17 00:00:00 2001 From: Tianyao Wu Date: Fri, 9 Oct 2026 21:59:47 +0800 Subject: [PATCH 2/4] qwen3_5: run prompts that share prefixes once per prefix forward_shared captures each prefix with the prefix continuation from #96 and continues every branch from it, with the fixed-algorithm GEMM handle; forward_fixed is the same path without a prefix. run_window takes a GEMM handle, can capture after a continuation, and indexes the text rotary tables by position. --- docs/architecture.md | 20 +- src/backends/cuda/qwen3_5/README.md | 5 +- src/models/qwen3_5/native/Cargo.toml | 4 + src/models/qwen3_5/native/src/model.rs | 339 ++++++++++++++++++++++--- src/runtime/README.md | 6 +- tests/qwen3_5/model.rs | 42 +++ tests/qwen3_5/prefix.rs | 205 +++++++++++++++ 7 files changed, 584 insertions(+), 37 deletions(-) create mode 100644 tests/qwen3_5/model.rs create mode 100644 tests/qwen3_5/prefix.rs diff --git a/docs/architecture.md b/docs/architecture.md index 9fb69e11..c96891b3 100644 --- a/docs/architecture.md +++ b/docs/architecture.md @@ -12,7 +12,8 @@ that share the [Qwen3.5/3.8 executor](../src/models/qwen3_5/native/), which acce single prompts and bounded packed prefill. Cua-S1 and Open-Jev-9B use single-prompt calls; Open-Jev-27B-v1.1 packs candidates within one request for input and gate/up GEMMs while preserving per-sequence mixers and output/down GEMM -shapes. +shapes. The executor can also run prompts that share token prefixes, each prefix +once (below); no worker uses that yet. [Laya's native worker](../src/models/laya/README.md) uses a separate Hopper CUDA backend for one complete padded request. All three coordinate independent processors and executors through @@ -27,6 +28,23 @@ Graph captures Encoder/Decision; gather, scorer/action head and synchronized readback remain outside capture. Native Metal remains planned; Python workers retain their documented reference/serving roles. +The Qwen executor's `forward_shared` takes a request prefix, group prefixes and +branches, ends each prefix at its last multiple of 64 tokens, and returns one final +hidden state per branch in order +([#85](https://github.com/ThinkFlowLab/system1-omni/issues/85)). It runs on the +prefix continuation that JEV-VL's cache uses: each prefix is captured once (per +linear-attention layer the float32 recurrent state and the last three conv inputs, +per full-attention layer the prefix keys and values), and each branch then runs as +its own pass from that state. The two prefix states persist across calls and grow only for a +longer prefix; their contents are valid only within one call, which runs eagerly and +synchronizes its stream before returning, also on error. Its GEMMs keep one +algorithm per weight shape, so on GPUs where that keeps rows independent of M +(checked by `fixed_gemm_rows_do_not_depend_on_m`, so far on sm_89) each result equals +`forward_fixed` on the full prompt bit for bit; it differs from `forward` by +rounding. Each prefix state takes about 50 MiB on Qwen3.5-4B and 9B, plus 32 KiB per +position of its prefix rounded up to 1,024 positions, and the GEMM handle a 32 MiB +workspace. + ## Native worker boundaries The native workers separate `processing.rs` from `executor.rs`; their worker diff --git a/src/backends/cuda/qwen3_5/README.md b/src/backends/cuda/qwen3_5/README.md index e1745051..cc238001 100644 --- a/src/backends/cuda/qwen3_5/README.md +++ b/src/backends/cuda/qwen3_5/README.md @@ -35,7 +35,10 @@ entry points. Open-Jev uses this path within requests; Cua-S1 keeps single-promp Prefix continuation retains full-attention KV, three convolution input rows, and FP32 Gated DeltaNet state at 64-token chunk boundaries. It is used by the -experimental JEV-VL worker; existing full-prefill callers keep their own paths. +experimental JEV-VL worker and by the shared executor's `forward_shared`, which +runs prompts that share prefixes once per prefix with the fixed-algorithm GEMM +handle ([#85](https://github.com/ThinkFlowLab/system1-omni/issues/85)); existing +full-prefill callers keep their own paths. ABI 6 adds mandatory copy and prefix entry points after ABI 5 native vision, and ABI 7 adds the fixed-algorithm GEMM handle. Rebuild the library and **all** Qwen worker binaries together; older libraries diff --git a/src/models/qwen3_5/native/Cargo.toml b/src/models/qwen3_5/native/Cargo.toml index 57c2428f..25837a54 100644 --- a/src/models/qwen3_5/native/Cargo.toml +++ b/src/models/qwen3_5/native/Cargo.toml @@ -21,3 +21,7 @@ path = "../../../../tests/qwen3_5/config.rs" [[test]] name = "kernels" path = "../../../../tests/qwen3_5/kernels.rs" + +[[test]] +name = "prefix" +path = "../../../../tests/qwen3_5/prefix.rs" diff --git a/src/models/qwen3_5/native/src/model.rs b/src/models/qwen3_5/native/src/model.rs index 304c3083..07c691fe 100644 --- a/src/models/qwen3_5/native/src/model.rs +++ b/src/models/qwen3_5/native/src/model.rs @@ -8,6 +8,9 @@ //! their kernels). Text prompts use one position per //! token. The explicit multimodal boundary inserts adapted image rows and supplies //! the interleaved temporal/height/width rotary positions to the same layer loop. +//! +//! `Model::forward_shared` runs prompts that share token prefixes, running each prefix +//! once and every branch from the state after it (see `SharedPrompts`). use std::collections::{HashMap, VecDeque}; use std::ffi::c_void; @@ -24,6 +27,9 @@ const ALIGN: usize = 256; const BF16: usize = 2; const F32: usize = 4; const GEMM_WORKSPACE: usize = 32 << 20; +/// Rows at which `forward_shared` fixes each weight shape's GEMM algorithm (see +/// cs1_gemm_create_fixed); its branches run tens of tokens. +const FIXED_GEMM_M: i32 = 64; #[derive(Debug, Clone)] pub struct Config { @@ -565,6 +571,8 @@ pub struct PrefixState { /// Tokens this prefix covers; always a multiple of 64 (the GDN chunk length), /// which also aligns the flash-attention key tiles of the one-shot pass. len: usize, + /// Tokens the buffers hold; `forward_shared` reuses a state for shorter prefixes. + cap: usize, /// Post-prep K [len, Hk*Dh] and raw V [len, Hk*Dh] rows per full-attention layer. attn_kv: Vec<(DeviceBuffer, DeviceBuffer)>, /// Float32 [H, K, V] recurrent states, one per Gated DeltaNet layer. @@ -587,6 +595,62 @@ impl PrefixState { } } +/// Gated DeltaNet chunk length. `forward_shared` ends prefixes at its multiples, where +/// a continuation repeats the one-shot arithmetic (see cs1_gdn_prefill_x). +const GDN_CHUNK: usize = 64; + +/// Prompts that share token prefixes, for `Model::forward_shared`. Branch `b` of group `g` +/// is the prompt `prefix`, then `groups[g].prefix`, then `groups[g].branches[b]`. How the +/// prompts are grouped does not change any result where the fixed GEMM algorithms keep +/// rows independent of M (see `forward_shared`). +pub struct SharedPrompts<'a> { + pub prefix: &'a [u32], + pub groups: Vec>, +} + +pub struct PromptGroup<'a> { + pub prefix: &'a [u32], + pub branches: Vec<&'a [u32]>, +} + +/// The token runs `forward_shared` executes: the request prefix, then per group its own +/// prefix and branches. Each prefix ends at the last multiple of `GDN_CHUNK` within it; +/// its remaining tokens run with every group or branch after it. +struct Runs { + prefix: Vec, + groups: Vec, +} + +struct GroupRuns { + prefix: Vec, + branches: Vec>, +} + +fn runs(prompts: &SharedPrompts) -> Runs { + let end = |n: usize| n - n % GDN_CHUNK; + let p = end(prompts.prefix.len()); + let groups = prompts + .groups + .iter() + .map(|g| { + let shared = [prompts.prefix, g.prefix].concat(); + let q = end(shared.len()); + GroupRuns { + prefix: shared[p..q].to_vec(), + branches: g + .branches + .iter() + .map(|b| [&shared[q..], b].concat()) + .collect(), + } + }) + .collect(); + Runs { + prefix: prompts.prefix[..p].to_vec(), + groups, + } +} + pub struct Model { pub cfg: Config, prefix_owner: Arc<()>, @@ -596,22 +660,31 @@ pub struct Model { layers: Vec, stream: Stream, gemm: *mut c_void, + /// One GEMM algorithm per weight shape for `forward_shared`; created on first use. + fixed_gemm: *mut c_void, /// Buffers for the largest packed token count so far; grows as needed. scratch: Option, + /// The request and group prefix states of `forward_shared`; grow as needed. + shared: [Option; 2], /// Opt-in replay with at most 64 captures keyed by ordered sequence lengths. graph_enabled: bool, graphs: VecDeque<(Vec, cuda::Graph)>, } -// SAFETY: the raw pointers are device addresses and a cuBLASLt handle owned by the +// SAFETY: the raw pointers are device addresses and cuBLASLt handles owned by the // model; the engine runs one forward pass at a time behind a mutex. unsafe impl Send for Model {} impl Drop for Model { fn drop(&mut self) { self.graphs.clear(); - // SAFETY: created by cs1_gemm_create and not destroyed before. - unsafe { (cuda::api().cs1_gemm_destroy)(self.gemm) }; + // SAFETY: created by cs1_gemm_create(_fixed) and not destroyed before. + unsafe { + (cuda::api().cs1_gemm_destroy)(self.gemm); + if !self.fixed_gemm.is_null() { + (cuda::api().cs1_gemm_destroy)(self.fixed_gemm); + } + } } } @@ -689,7 +762,9 @@ impl Model { layers, stream, gemm, + fixed_gemm: std::ptr::null_mut(), scratch: None, + shared: [None, None], graph_enabled: std::env::var("CUA_S1_GRAPH").as_deref() == Ok("1"), graphs: VecDeque::new(), }; @@ -697,12 +772,25 @@ impl Model { } fn gemm(&self, s: &Scratch, x: usize, w: &Tensor, y: usize, m: usize) -> Result<()> { + self.gemm_with(self.gemm, s, x, w, y, m) + } + + /// `gemm` with the cuBLASLt handle `handle`. + fn gemm_with( + &self, + handle: *mut c_void, + s: &Scratch, + x: usize, + w: &Tensor, + y: usize, + m: usize, + ) -> Result<()> { let (n, k) = (w.shape[0] as i32, w.shape[1] as i32); // SAFETY: x and y are scratch buffers sized for m rows of w's shape. check( unsafe { (cuda::api().cs1_gemm)( - self.gemm, + handle, s.at(x), w.ptr, s.at(y), @@ -938,6 +1026,7 @@ impl Model { owner: Arc::clone(&self.prefix_owner), initialized: false, len, + cap: len, attn_kv, gdn_state, conv_tail, @@ -979,7 +1068,7 @@ impl Model { self.upload_positions(s, input.position_ids)?; self.embed_tokens(s, input.token_ids)?; self.overwrite_image_rows(s, input, 0)?; - self.run_window(s, 0, t, true, Some(state), None)?; + self.run_window(s, 0, t, true, Some(state), None, self.gemm)?; self.synchronize()?; state.initialized = true; Ok(()) @@ -1021,10 +1110,174 @@ impl Model { self.upload_positions(s, input.position_ids)?; self.embed_tokens_at(s, input.token_ids, prefix.len)?; self.overwrite_image_rows(s, input, prefix.len)?; - self.run_window(s, prefix.len, tend, true, None, Some(prefix))?; + self.run_window(s, prefix.len, tend, true, None, Some(prefix), self.gemm)?; self.last_hidden(s, tend) } + /// `forward` with the GEMM algorithms of `forward_shared`, one per weight shape, run + /// eagerly: the result that `forward_shared` reproduces for this prompt. + pub fn forward_fixed(&mut self, ids: &[u32]) -> Result> { + let result = self.run_fixed(ids); + if result.is_err() { + // Work already queued may still use the buffers. + let _ = cuda::synchronize(self.stream); + } + result + } + + fn run_fixed(&mut self, ids: &[u32]) -> Result> { + let t = ids.len(); + ensure!( + t > 0 && t <= self.cfg.max_positions, + "empty or oversized prompt" + ); + ensure!( + ids.iter().all(|&i| (i as usize) < self.embed.shape[0]), + "token id outside the vocabulary" + ); + self.prepare_scratch(t)?; + let gemm = self.fixed_gemm()?; + let s = self.scratch.as_ref().unwrap(); + self.embed_tokens(s, ids)?; + self.run_window(s, 0, t, false, None, None, gemm)?; + self.last_hidden(s, t) + } + + /// The final-norm hidden state at the last position of every branch, grouped and + /// ordered as in `prompts`. Each prefix runs once, up to its last multiple of 64 + /// tokens, capturing its state as `forward_multimodal_capture` does, and each branch + /// continues from the state after its prefixes, with its positions and causal mask + /// continuing theirs. A prefix's remaining tokens run with every group or branch + /// after it. + /// + /// On a GPU where the fixed GEMM algorithms keep rows independent of M (see + /// cs1_gemm_create_fixed), each result equals `forward_fixed` on the full prompt bit + /// for bit, however the prompts are grouped. It differs from `forward`, which picks + /// GEMM algorithms per row count, by rounding. + /// + /// Runs eagerly. The prefix states are overwritten by the next call, and the stream + /// is synchronized before this returns, also on error. + pub fn forward_shared(&mut self, prompts: &SharedPrompts) -> Result>>> { + let result = self.run_shared(prompts); + if result.is_err() { + // Work already queued may still use the buffers. + let _ = cuda::synchronize(self.stream); + } + result + } + + fn run_shared(&mut self, prompts: &SharedPrompts) -> Result>>> { + let vocab = self.embed.shape[0]; + ensure!(!prompts.groups.is_empty(), "no prompts"); + let mut longest = 0; + for g in &prompts.groups { + ensure!(!g.branches.is_empty(), "a prompt group without branches"); + for b in &g.branches { + ensure!(!b.is_empty(), "an empty branch"); + longest = longest.max(prompts.prefix.len() + g.prefix.len() + b.len()); + } + } + ensure!(longest <= self.cfg.max_positions, "oversized prompt"); + let all = prompts + .groups + .iter() + .flat_map(|g| g.prefix.iter().chain(g.branches.iter().copied().flatten())); + ensure!( + prompts + .prefix + .iter() + .chain(all) + .all(|&i| (i as usize) < vocab), + "token id outside the vocabulary" + ); + let runs = runs(prompts); + let p = runs.prefix.len(); + // the longest group prefix that runs on its own, if any + let q = runs + .groups + .iter() + .filter(|g| !g.prefix.is_empty()) + .map(|g| p + g.prefix.len()) + .max() + .unwrap_or(0); + self.prepare_scratch(longest)?; + let gemm = self.fixed_gemm()?; + let [mut request, mut group] = std::mem::take(&mut self.shared); + let result = (|| { + let request = self.shared_state(&mut request, p)?; + let group = self.shared_state(&mut group, q)?; + let s = self.scratch.as_ref().unwrap(); + let request = match request { + Some(state) => { + self.embed_tokens(s, &runs.prefix)?; + self.run_window(s, 0, p, false, Some(&mut *state), None, gemm)?; + Some(&*state) + } + None => None, + }; + let mut group = group; + let mut hidden = Vec::with_capacity(runs.groups.len()); + for run in &runs.groups { + let start = p + run.prefix.len(); + let from = if run.prefix.is_empty() { + request + } else { + let state = group.as_deref_mut().unwrap(); + state.len = start; + self.embed_tokens_at(s, &run.prefix, p)?; + self.run_window(s, p, start, false, Some(&mut *state), request, gemm)?; + Some(&*state) + }; + let rows = run + .branches + .iter() + .map(|b| { + let end = start + b.len(); + self.embed_tokens_at(s, b, start)?; + self.run_window(s, start, end, false, None, from, gemm)?; + self.last_hidden(s, end) + }) + .collect::>>()?; + hidden.push(rows); + } + Ok(hidden) + })(); + self.shared = [request, group]; + result + } + + /// `slot`'s state resized to hold `len` tokens, none for an empty prefix. + fn shared_state<'s>( + &self, + slot: &'s mut Option, + len: usize, + ) -> Result> { + if len == 0 { + return Ok(None); + } + if slot.as_ref().is_none_or(|state| state.cap < len) { + *slot = None; + let cap = len + .next_multiple_of(1024) + .min(self.cfg.max_positions - self.cfg.max_positions % GDN_CHUNK); + *slot = Some(self.alloc_prefix(cap.max(len))?); + } + let state = slot.as_mut().unwrap(); + state.len = len; + Ok(Some(state)) + } + + /// The GEMM handle with one algorithm per weight shape, created on first use. + fn fixed_gemm(&mut self) -> Result<*mut c_void> { + if self.fixed_gemm.is_null() { + // SAFETY: plain allocation; checked for null below. + self.fixed_gemm = + unsafe { (cuda::api().cs1_gemm_create_fixed)(GEMM_WORKSPACE, FIXED_GEMM_M) }; + ensure!(!self.fixed_gemm.is_null(), "cuBLASLt setup failed"); + } + Ok(self.fixed_gemm) + } + fn embed_tokens(&self, s: &Scratch, ids: &[u32]) -> Result<()> { self.embed_tokens_at(s, ids, 0) } @@ -1293,15 +1546,19 @@ impl Model { /// The R2d prefix-cache pass: run rows `[qb, tend)` of one prompt with /// explicit positions, touching per-layer buffers only for those rows. /// - /// `capture` (qb must be 0): collect post-prep K and pre-GEMM V columns of each + /// `capture` (`tend` rows): collect post-prep K and pre-GEMM V columns of each /// full-attention layer plus the float32 GDN state and the three pre-conv - /// projection columns of each Gated DeltaNet layer. `prefix` (qb = prefix.len): + /// projection columns of each Gated DeltaNet layer, for rows `[0, tend)`; with a + /// `prefix` too, the state after it continued. `prefix` (qb = prefix.len): /// seed the window from a captured entry — the conv tail columns are copied /// back into place before the convolution, and the prefix K/V columns before /// the attention. Everything else is `run()` with shifted row bases, so a /// captured+continued pass repeats the one-shot arithmetic token for token /// (qb is a multiple of the 64-row GDN chunk; the only remaining divergence - /// is cuBLASLt's M-shaped plan choice, an f32-accumulation-order ulp). + /// is cuBLASLt's M-shaped plan choice, an f32-accumulation-order ulp, which the + /// fixed-algorithm `gemm` handle removes). Custom positions hold the window's + /// rows from index 0; the text tables are indexed by position. + #[allow(clippy::too_many_arguments)] // One window: rows, positions, states and GEMM handle. fn run_window( &self, s: &Scratch, @@ -1310,9 +1567,12 @@ impl Model { custom_positions: bool, capture: Option<&mut PrefixState>, prefix: Option<&PrefixState>, + gemm: *mut c_void, ) -> Result<()> { ensure!( - qb <= tend && (qb == 0) == prefix.is_none() && (qb == 0 || capture.is_none()), + qb < tend + && prefix.map_or(qb == 0, |pre| pre.len == qb) + && capture.as_deref().is_none_or(|cap| cap.len == tend), "prefix window shape" ); let cfg = &self.cfg; @@ -1326,7 +1586,8 @@ impl Model { let (cos, sin) = if custom_positions { (s.custom_cos, s.custom_sin) } else { - (s.cos, s.sin) + let rope = qb * cfg.rotary_half * BF16; + (s.cos + rope, s.sin + rope) }; let hb = cfg.hidden * BF16; let ob = cfg.heads * cfg.head_dim * BF16; // post-prep q/gate and ao row bytes @@ -1341,7 +1602,8 @@ impl Model { let mut la_i = 0usize; // SAFETY (every kernel call below): pointers are weights in the arena or // scratch buffers laid out for tend tokens at the widths used here; prefix - // buffers hold exactly the shapes allocated by alloc_prefix(qb). + // buffers hold at least `len` rows (a capture's `tend`) of the shapes + // allocated by alloc_prefix. unsafe { check( (cuda::api().cs1_rms_norm)( @@ -1359,33 +1621,38 @@ impl Model { for (i, layer) in self.layers.iter().enumerate() { match &layer.mixer { Mixer::Linear(la) => { - self.gemm(s, s.x + qb * hb, &la.in_proj, s.gdn_in + qb * ldbb, rows)?; + self.gemm_with( + gemm, + s, + s.x + qb * hb, + &la.in_proj, + s.gdn_in + qb * ldbb, + rows, + )?; let ld = ldb as i32; let z = s.gdn_in + w.conv * BF16; let b = z + vd * BF16; let a = b + hv * BF16; - let (s_in, s_out): (*const c_void, *mut c_void) = - match (&prefix, capture.as_deref()) { - (Some(pre), _) => { - (pre.gdn_state[la_i].at(0).cast_const(), std::ptr::null_mut()) - } - (None, Some(cap)) => (std::ptr::null(), cap.gdn_state[la_i].at(0)), - (None, None) => (std::ptr::null(), std::ptr::null_mut()), - }; + let s_in: *const c_void = + prefix.map_or(std::ptr::null(), |pre| pre.gdn_state[la_i].at(0)); + let s_out: *mut c_void = capture + .as_deref() + .map_or(std::ptr::null_mut(), |cap| cap.gdn_state[la_i].at(0)); unsafe { - if let Some(cap) = capture.as_deref() { - // conv window tail: the last three pre-conv columns. + if qb > 0 { cuda::copy_dd( - cap.conv_tail[la_i].at(0), - p(s.gdn_in + (rows - 3) * ldbb), + p(s.gdn_in + (qb - 3) * ldbb), + prefix.unwrap().conv_tail[la_i].at(0), 3 * ldbb, st, )?; } - if qb > 0 { + if let Some(cap) = capture.as_deref() { + // conv window tail: the last three pre-conv columns, which + // include restored prefix columns for a window under three rows. cuda::copy_dd( - p(s.gdn_in + (qb - 3) * ldbb), - prefix.unwrap().conv_tail[la_i].at(0), + cap.conv_tail[la_i].at(0), + p(s.gdn_in + (tend - 3) * ldbb), 3 * ldbb, st, )?; @@ -1459,11 +1726,11 @@ impl Model { "gated norm", )?; } - self.gemm(s, s.ln + qb * vb, &la.out, s.delta + qb * hb, rows)?; + self.gemm_with(gemm, s, s.ln + qb * vb, &la.out, s.delta + qb * hb, rows)?; la_i += 1; } Mixer::Full(fa) => { - self.gemm(s, s.x + qb * hb, &fa.qkv, s.attn_in + qb * ab, rows)?; + self.gemm_with(gemm, s, s.x + qb * hb, &fa.qkv, s.attn_in + qb * ab, rows)?; let k = s.attn_in + w.attn_q * BF16; let v = k + kvb; let (q_base, t_flash) = (qb as i32, tend as i32); @@ -1521,7 +1788,7 @@ impl Model { "gated attention", )?; } - self.gemm(s, s.ao + qb * ob, &fa.o, s.delta + qb * hb, rows)?; + self.gemm_with(gemm, s, s.ao + qb * ob, &fa.o, s.delta + qb * hb, rows)?; fa_i += 1; } } @@ -1540,7 +1807,8 @@ impl Model { "post-attention norm", )?; } - self.gemm( + self.gemm_with( + gemm, s, s.x + qb * hb, &layer.gate_up, @@ -1560,7 +1828,8 @@ impl Model { "silu mul", )?; } - self.gemm( + self.gemm_with( + gemm, s, s.act + qb * cfg.intermediate * BF16, &layer.down, @@ -1590,3 +1859,7 @@ impl Model { Ok(()) } } + +#[cfg(test)] +#[path = "../../../../../tests/qwen3_5/model.rs"] +mod tests; diff --git a/src/runtime/README.md b/src/runtime/README.md index bd571b03..fcb2ec0f 100644 --- a/src/runtime/README.md +++ b/src/runtime/README.md @@ -34,8 +34,10 @@ existing failure behavior. This implements serial admission per executor within a worker process. Admitted concurrency is one. Queue-length limits, token budgets, compatibility grouping, shared processing orchestration and GPU batching remain planned. The Qwen -executor still accepts one prompt per forward; this scheduler does not pack -inputs or coordinate separately running workers. +executor runs single prompts, bounded packed prefill (Open-Jev-27B-v1.1's +candidates within a request) or prompts that share prefixes in one call, which no +worker uses yet; this scheduler does not pack inputs or coordinate separately +running workers. CPU tests cover FIFO serialization, queued cancellation, retention of admission and resources after dispatched cancellation, errors/panics, and independent diff --git a/tests/qwen3_5/model.rs b/tests/qwen3_5/model.rs new file mode 100644 index 00000000..da64d1a5 --- /dev/null +++ b/tests/qwen3_5/model.rs @@ -0,0 +1,42 @@ +use super::*; + +fn seq(n: usize, base: u32) -> Vec { + (0..n as u32).map(|i| base + i).collect() +} + +#[test] +fn runs_keep_every_prompt_and_end_prefixes_on_chunks() { + for p in [0usize, 1, 63, 64, 65, 130] { + let prefix = seq(p, 0); + let groups = [ + (seq(0, 1000), vec![seq(1, 2000), seq(70, 3000)]), + (seq(1, 4000), vec![seq(2, 5000)]), + (seq(64, 6000), vec![seq(5, 7000), seq(1, 8000)]), + ]; + let prompts = SharedPrompts { + prefix: &prefix, + groups: groups + .iter() + .map(|(prefix, branches)| PromptGroup { + prefix, + branches: branches.iter().map(Vec::as_slice).collect(), + }) + .collect(), + }; + let r = runs(&prompts); + let chunk = |n: usize| n - n % GDN_CHUNK; + assert_eq!(r.prefix.len(), chunk(p), "p = {p}"); + for ((group, branches), run) in groups.iter().zip(&r.groups) { + assert_eq!(r.prefix.len() + run.prefix.len(), chunk(p + group.len())); + assert_eq!(branches.len(), run.branches.len()); + for (branch, run_branch) in branches.iter().zip(&run.branches) { + assert!(!run_branch.is_empty()); + assert_eq!( + [&r.prefix[..], &run.prefix, run_branch].concat(), + [&prefix[..], group, branch].concat(), + "p = {p}" + ); + } + } + } +} diff --git a/tests/qwen3_5/prefix.rs b/tests/qwen3_5/prefix.rs new file mode 100644 index 00000000..1ca07061 --- /dev/null +++ b/tests/qwen3_5/prefix.rs @@ -0,0 +1,205 @@ +//! GPU checks of `Model::forward_shared` against `forward_fixed` on each full prompt. They +//! need a GPU, CUA_S1_CUDA_LIB pointing at libqwen3_5_cuda.so and QWEN3_5_MODEL pointing at +//! a merged Qwen3.5/3.8 checkpoint, such as a Cua-S1 or Open-Jev export: +//! +//! CUA_S1_CUDA_LIB=$PWD/target/release/libqwen3_5_cuda.so QWEN3_5_MODEL= \ +//! cargo test --release -p omni-qwen3-5-native --test prefix -- --ignored --test-threads=1 +//! +//! With CUA_S1_GRAPH=1, repeated plain `forward` calls of one length replay a CUDA Graph, +//! also between shared runs. + +use std::path::PathBuf; + +use omni_qwen3_5_native::model::{Model, PromptGroup, SharedPrompts}; + +fn model() -> Model { + let var = |name: &str| { + std::env::var_os(name) + .map(PathBuf::from) + .unwrap_or_else(|| panic!("{name} must be set")) + }; + Model::load(&var("QWEN3_5_MODEL"), &var("CUA_S1_CUDA_LIB")).unwrap() +} + +/// Token ids in [1000, 100000) from a fixed seed, inside every Qwen3.5 vocabulary. +fn tokens(n: usize, seed: u64) -> Vec { + let mut x = seed.wrapping_mul(0x9e37_79b9_7f4a_7c15) | 1; + (0..n) + .map(|_| { + x ^= x << 13; + x ^= x >> 7; + x ^= x << 17; + 1000 + (x >> 33) as u32 % 99_000 + }) + .collect() +} + +struct Case { + prefix: Vec, + groups: Vec<(Vec, Vec>)>, +} + +impl Case { + /// A request prefix of `p` tokens and groups of (prefix length, branch lengths). + fn new(p: usize, groups: &[(usize, &[usize])], seed: u64) -> Self { + let mut seed = seed * 1000; + let mut next = |n| { + seed += 1; + tokens(n, seed) + }; + Self { + prefix: next(p), + groups: groups + .iter() + .map(|&(q, branches)| (next(q), branches.iter().map(|&b| next(b)).collect())) + .collect(), + } + } + + fn prompts(&self) -> SharedPrompts<'_> { + SharedPrompts { + prefix: &self.prefix, + groups: self + .groups + .iter() + .map(|(prefix, branches)| PromptGroup { + prefix, + branches: branches.iter().map(Vec::as_slice).collect(), + }) + .collect(), + } + } + + fn full(&self, g: usize, b: usize) -> Vec { + let (prefix, branches) = &self.groups[g]; + [&self.prefix[..], prefix, &branches[b]].concat() + } +} + +fn rel_l2(a: &[f32], b: &[f32]) -> f64 { + let diff: f64 = a + .iter() + .zip(b) + .map(|(x, y)| (*x as f64 - *y as f64).powi(2)) + .sum(); + let norm: f64 = b.iter().map(|y| (*y as f64).powi(2)).sum(); + (diff / norm).sqrt() +} + +fn bits(rows: &[Vec>]) -> Vec { + rows.iter() + .flatten() + .flatten() + .map(|v| v.to_bits()) + .collect() +} + +#[test] +#[ignore = "needs a GPU, CUA_S1_CUDA_LIB and QWEN3_5_MODEL"] +fn shared_prefixes_match_full_prompts() { + let mut model = model(); + // Request prefixes around and inside 64-token chunks; groups without a prefix of their + // own, with one token, 70 and 128; branches of 1 to 200 tokens. + let groups: &[(usize, &[usize])] = + &[(0, &[1, 64]), (1, &[2, 65]), (70, &[1, 200]), (128, &[3])]; + for (i, p) in [0usize, 1, 63, 64, 65, 300, 1000].into_iter().enumerate() { + let case = Case::new(p, groups, i as u64 + 1); + let shared = model.forward_shared(&case.prompts()).unwrap(); + let mut to_forward = 0f64; + for (g, rows) in shared.iter().enumerate() { + for (b, row) in rows.iter().enumerate() { + let full = case.full(g, b); + let want = model.forward_fixed(&full).unwrap(); + assert!(row.iter().all(|v| v.is_finite())); + assert!( + row.iter() + .map(|v| v.to_bits()) + .eq(want.iter().map(|v| v.to_bits())), + "prefix {p}, group {g}, branch {b}" + ); + // the GEMM algorithms differ from `forward`'s, not the math + let plain = model.forward(&full).unwrap(); + to_forward = to_forward.max(rel_l2(row, &plain)); + } + } + eprintln!( + "prefix {p}: bit-identical to forward_fixed; largest relative L2 difference from forward {to_forward:.2e}" + ); + // A sanity bound for the change of GEMM algorithms, not a tolerance for reuse. + assert!(to_forward <= 3e-2, "prefix {p}: {to_forward}"); + } +} + +#[test] +#[ignore = "needs a GPU, CUA_S1_CUDA_LIB and QWEN3_5_MODEL"] +fn shared_prefix_state_stays_inside_one_call() { + let mut model = model(); + let a = Case::new(200, &[(10, &[3, 40, 1]), (0, &[64]), (65, &[2, 7])], 11); + // longer than any prompt before it, so the buffers grow in between + let b = Case::new(1100, &[(30, &[9, 100])], 12); + let plain = model.forward(&a.full(0, 1)).unwrap(); + // with CUA_S1_GRAPH=1, the first call captured a graph and this one replays it + assert_eq!(model.forward(&a.full(0, 1)).unwrap(), plain); + let first = model.forward_shared(&a.prompts()).unwrap(); + // the same request again, and after a longer one + let again = model.forward_shared(&a.prompts()).unwrap(); + assert_eq!(bits(&first), bits(&again)); + // the plain forward between shared runs, replayed with CUA_S1_GRAPH=1 + assert_eq!(model.forward(&a.full(0, 1)).unwrap(), plain); + model.forward_shared(&b.prompts()).unwrap(); + let after = model.forward_shared(&a.prompts()).unwrap(); + assert_eq!(bits(&first), bits(&after)); + // groups and branches in reverse order: every branch keeps its result + let reversed = Case { + prefix: a.prefix.clone(), + groups: a + .groups + .iter() + .rev() + .map(|(prefix, branches)| (prefix.clone(), branches.iter().rev().cloned().collect())) + .collect(), + }; + let back = model.forward_shared(&reversed.prompts()).unwrap(); + let mut back: Vec>> = back.into_iter().rev().collect(); + back.iter_mut().for_each(|branches| branches.reverse()); + assert_eq!(bits(&first), bits(&back)); + // the plain forward is unaffected (eager here: the longer request dropped the graphs) + assert_eq!(model.forward(&a.full(0, 1)).unwrap(), plain); + // Invalid requests fail without breaking the model. All of these are refused before + // any GPU work is queued. + let empty: &[u32] = &[]; + let outside: &[u32] = &[5, u32::MAX]; + let group = |prefix, branches| PromptGroup { prefix, branches }; + let bad = [ + SharedPrompts { + prefix: &a.prefix, + groups: vec![], + }, + SharedPrompts { + prefix: &a.prefix, + groups: vec![group(empty, vec![])], + }, + SharedPrompts { + prefix: &a.prefix, + groups: vec![group(empty, vec![empty])], + }, + SharedPrompts { + prefix: &a.prefix, + groups: vec![group(empty, vec![outside])], + }, + SharedPrompts { + prefix: &a.prefix, + groups: vec![group(outside, vec![&a.prefix])], + }, + SharedPrompts { + prefix: outside, + groups: vec![group(empty, vec![&a.prefix])], + }, + ]; + for prompts in &bad { + assert!(model.forward_shared(prompts).is_err()); + } + assert!(model.forward_fixed(outside).is_err()); + let last = model.forward_shared(&a.prompts()).unwrap(); + assert_eq!(bits(&first), bits(&last)); +} From d8cd04fe77bf4ab282de3ac3ee66555e13d6c9c6 Mon Sep 17 00:00:00 2001 From: Tianyao Wu Date: Sun, 11 Oct 2026 22:16:39 +0800 Subject: [PATCH 3/4] qwen3_5: declare ABI 7 in the backend manifest --- src/backends/cuda/qwen3_5/qwen3_5.backend.json | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/backends/cuda/qwen3_5/qwen3_5.backend.json b/src/backends/cuda/qwen3_5/qwen3_5.backend.json index 37484cef..cf06c19d 100644 --- a/src/backends/cuda/qwen3_5/qwen3_5.backend.json +++ b/src/backends/cuda/qwen3_5/qwen3_5.backend.json @@ -1,6 +1,6 @@ { "name": "qwen3_5", - "abi_version": 6, + "abi_version": 7, "status": "experimental", "sources": [ "ops.h", From 866bed555ac8b9a7b10046bfffd4362d17a1ab9e Mon Sep 17 00:00:00 2001 From: Tianyao Wu Date: Wed, 21 Oct 2026 23:31:12 +0800 Subject: [PATCH 4/4] Refuse other execution profiles in forward_fixed and forward_shared Both cover merged weights with the default numerics. A model with an unmerged LoRA adapter, reference numerics or FP32 GDN is now refused before any allocation or CUDA call. A CPU test loads a tiny checkpoint with an adapter on a fake CUDA library that counts calls, and checks that both entry points fail without a call or a buffer for each profile. --- docs/architecture.md | 7 +- src/models/qwen3_5/native/src/model.rs | 12 ++ tests/qwen3_5/model.rs | 226 +++++++++++++++++++++++++ tests/qwen3_5/profile_fake.c | 20 +++ 4 files changed, 262 insertions(+), 3 deletions(-) create mode 100644 tests/qwen3_5/profile_fake.c diff --git a/docs/architecture.md b/docs/architecture.md index 7f76d309..c4aa209e 100644 --- a/docs/architecture.md +++ b/docs/architecture.md @@ -46,9 +46,10 @@ synchronizes its stream before returning, also on error. Its GEMMs keep one algorithm per weight shape, so on GPUs where that keeps rows independent of M (checked by `fixed_gemm_rows_do_not_depend_on_m`, so far on sm_89) each result equals `forward_fixed` on the full prompt bit for bit; it differs from `forward` by -rounding. Each prefix state takes about 50 MiB on Qwen3.5-4B and 9B, plus 32 KiB per -position of its prefix rounded up to 1,024 positions, and the GEMM handle a 32 MiB -workspace. +rounding. Both take merged weights with the default numerics only, and refuse an +unmerged adapter, reference numerics or FP32 GDN before any device work. Each +prefix state takes about 50 MiB on Qwen3.5-4B and 9B, plus 32 KiB per position of +its prefix rounded up to 1,024 positions, and the GEMM handle a 32 MiB workspace. ## Native worker boundaries diff --git a/src/models/qwen3_5/native/src/model.rs b/src/models/qwen3_5/native/src/model.rs index b604e174..5aa60e8c 100644 --- a/src/models/qwen3_5/native/src/model.rs +++ b/src/models/qwen3_5/native/src/model.rs @@ -1451,6 +1451,7 @@ impl Model { /// `forward` with the GEMM algorithms of `forward_shared`, one per weight shape, run /// eagerly: the result that `forward_shared` reproduces for this prompt. pub fn forward_fixed(&mut self, ids: &[u32]) -> Result> { + self.ensure_fixed_profile()?; let result = self.run_fixed(ids); if result.is_err() { // Work already queued may still use the buffers. @@ -1492,6 +1493,7 @@ impl Model { /// Runs eagerly. The prefix states are overwritten by the next call, and the stream /// is synchronized before this returns, also on error. pub fn forward_shared(&mut self, prompts: &SharedPrompts) -> Result>>> { + self.ensure_fixed_profile()?; let result = self.run_shared(prompts); if result.is_err() { // Work already queued may still use the buffers. @@ -1601,6 +1603,16 @@ impl Model { Ok(Some(state)) } + /// `forward_fixed` and `forward_shared` cover merged weights with the default + /// numerics. Refuse the other profiles before any allocation or CUDA call. + fn ensure_fixed_profile(&self) -> Result<()> { + ensure!( + self.lora.is_none() && self.reference.is_none() && self.fp32_gdn.is_none(), + "fixed-algorithm GEMMs and shared prefixes support only merged weights with the default numerics" + ); + Ok(()) + } + /// The GEMM handle with one algorithm per weight shape, created on first use. fn fixed_gemm(&mut self) -> Result<*mut c_void> { if self.fixed_gemm.is_null() { diff --git a/tests/qwen3_5/model.rs b/tests/qwen3_5/model.rs index da64d1a5..6e66fd1c 100644 --- a/tests/qwen3_5/model.rs +++ b/tests/qwen3_5/model.rs @@ -40,3 +40,229 @@ fn runs_keep_every_prompt_and_end_prefixes_on_chunks() { } } } + +/// A two-layer checkpoint with hidden size 128, the smallest whose stacked weights stay +/// contiguous in the arena, and an FP32 adapter on layer 0's down projection. +#[cfg(unix)] +fn write_tiny_checkpoint(dir: &Path) { + use safetensors::{ + Dtype, + tensor::{TensorView, serialize}, + }; + let config = serde_json::json!({ + "text_config": { + "hidden_size": 128, "intermediate_size": 8, "num_hidden_layers": 2, + "layer_types": ["linear_attention", "full_attention"], + "num_attention_heads": 1, "num_key_value_heads": 1, "head_dim": 256, + "linear_num_key_heads": 1, "linear_num_value_heads": 1, + "linear_key_head_dim": 128, "linear_value_head_dim": 128, + "linear_conv_kernel_dim": 4, "hidden_act": "silu", "rms_norm_eps": 1e-6, + "max_position_embeddings": 4096, + "rope_parameters": { + "rope_type": "default", "rope_theta": 1e7, "partial_rotary_factor": 0.25, + "mrope_section": [11, 11, 10], "mrope_interleaved": true + } + } + }); + std::fs::write(dir.join("config.json"), config.to_string()).unwrap(); + let write = |file: &str, dtype: Dtype, tensors: &[(String, Vec)]| { + let size = if dtype == Dtype::F32 { 4 } else { 2 }; + let data: Vec> = tensors + .iter() + .map(|(_, shape)| vec![0; shape.iter().product::() * size]) + .collect(); + let views: Vec<_> = tensors + .iter() + .zip(&data) + .map(|((name, shape), bytes)| { + ( + name.clone(), + TensorView::new(dtype, shape.clone(), bytes).unwrap(), + ) + }) + .collect(); + std::fs::write(dir.join(file), serialize(views, None).unwrap()).unwrap(); + }; + let layer = |index: usize, mixer: &[(&str, Vec)]| { + [ + ("input_layernorm.weight", vec![128]), + ("post_attention_layernorm.weight", vec![128]), + ("mlp.gate_proj.weight", vec![8, 128]), + ("mlp.up_proj.weight", vec![8, 128]), + ("mlp.down_proj.weight", vec![128, 8]), + ] + .into_iter() + .chain(mixer.iter().cloned()) + .map(|(name, shape)| (format!("layers.{index}.{name}"), shape)) + .collect::>() + }; + let weights: Vec<_> = [ + ("embed_tokens.weight".to_owned(), vec![16, 128]), + ("norm.weight".to_owned(), vec![128]), + ] + .into_iter() + .chain(layer( + 0, + &[ + ("linear_attn.in_proj_qkv.weight", vec![384, 128]), + ("linear_attn.in_proj_z.weight", vec![128, 128]), + ("linear_attn.in_proj_b.weight", vec![1, 128]), + ("linear_attn.in_proj_a.weight", vec![1, 128]), + ("linear_attn.conv1d.weight", vec![384, 1, 4]), + ("linear_attn.A_log", vec![1]), + ("linear_attn.dt_bias", vec![1]), + ("linear_attn.norm.weight", vec![128]), + ("linear_attn.out_proj.weight", vec![128, 128]), + ], + )) + .chain(layer( + 1, + &[ + ("self_attn.q_proj.weight", vec![512, 128]), + ("self_attn.k_proj.weight", vec![256, 128]), + ("self_attn.v_proj.weight", vec![256, 128]), + ("self_attn.o_proj.weight", vec![128, 256]), + ("self_attn.q_norm.weight", vec![256]), + ("self_attn.k_norm.weight", vec![256]), + ], + )) + .map(|(name, shape)| (format!("model.language_model.{name}"), shape)) + .collect(); + write("model.safetensors", Dtype::BF16, &weights); + let target = "base_model.model.model.language_model.layers.0.mlp.down_proj"; + write( + "adapter.safetensors", + Dtype::F32, + &[ + (format!("{target}.lora_A.weight"), vec![1, 8]), + (format!("{target}.lora_B.weight"), vec![128, 1]), + ], + ); +} + +#[cfg(unix)] +#[test] +fn cpu_fixed_and_shared_refuse_other_profiles_before_device_work() { + // The library loads once per process, so this runs in a child process. + if std::env::var_os("QWEN_TEST_PROFILE_CHILD").is_none() { + let name = format!( + "{}::cpu_fixed_and_shared_refuse_other_profiles_before_device_work", + module_path!().split_once("::").unwrap().1 + ); + let output = std::process::Command::new(std::env::current_exe().unwrap()) + .args(["--exact", &name, "--nocapture"]) + .env("QWEN_TEST_PROFILE_CHILD", "1") + .output() + .unwrap(); + assert!( + output.status.success(), + "{}\n{}", + String::from_utf8_lossy(&output.stdout), + String::from_utf8_lossy(&output.stderr) + ); + return; + } + let dir = std::env::temp_dir().join(format!("qwen-profile-{}", std::process::id())); + std::fs::create_dir_all(&dir).unwrap(); + write_tiny_checkpoint(&dir); + let library = dir.join(if cfg!(target_os = "macos") { + "fake.dylib" + } else { + "fake.so" + }); + let cuda_source = include_str!("../../src/models/qwen3_5/native/src/cuda.rs"); + let reference_source = include_str!("../../src/models/qwen3_5/native/src/cuda/reference.rs"); + let abi = cuda_source + .lines() + .find_map(|line| line.strip_prefix("const ABI_VERSION: u32 =")) + .unwrap() + .trim() + .trim_end_matches(';'); + // Every other required and optional symbol counts the call and fails. + let mut code = include_str!("profile_fake.c").to_owned(); + let required = cuda_source + .split("api! {") + .nth(1) + .unwrap() + .split("\n}\n") + .next() + .unwrap() + .lines() + .filter_map(|line| line.trim().strip_prefix("cs1_")); + let optional = [cuda_source, reference_source] + .into_iter() + .flat_map(|source| source.split("b\"cs1_").skip(1)); + let mut names: Vec<_> = required + .chain(optional) + .map(|rest| format!("cs1_{}", rest.split(['(', '\\']).next().unwrap())) + .filter(|name| !code.contains(&format!(" {name}("))) + .collect(); + names.sort(); + names.dedup(); + for name in names { + code.push_str(&format!("int {name}(void) {{ ++calls; return 77; }}\n")); + } + let source = dir.join("fake.c"); + std::fs::write(&source, code).unwrap(); + let output = std::process::Command::new("cc") + .args(["-shared", "-fPIC"]) + .arg(format!("-DPROFILE_FAKE_ABI={abi}")) + .arg(&source) + .arg("-o") + .arg(&library) + .output() + .unwrap(); + assert!( + output.status.success(), + "{}", + String::from_utf8_lossy(&output.stderr) + ); + let mut model = + Model::load_with_lora(&dir, &library, &dir.join("adapter.safetensors"), 2.).unwrap(); + // SAFETY: the test library defines this counter and stays loaded in the process. + let metrics = unsafe { libloading::Library::new(&library).unwrap() }; + let calls = unsafe { + *metrics + .get:: i32>(b"profile_fake_calls\0") + .unwrap() + }; + let prompts = SharedPrompts { + prefix: &[1, 2], + groups: vec![PromptGroup { + prefix: &[3], + branches: vec![&[4], &[5, 6]], + }], + }; + let refuse = |model: &mut Model, profile: &str| { + let before = unsafe { calls() }; + for error in [ + model.forward_fixed(&[1, 2, 3]).unwrap_err(), + model.forward_shared(&prompts).unwrap_err(), + ] { + assert!( + error.to_string().contains("only merged weights"), + "{profile}: {error:#}" + ); + } + assert_eq!(unsafe { calls() }, before, "{profile}: CUDA calls"); + assert!( + model.scratch.is_none() + && model.fixed_gemm.is_null() + && model.shared.iter().all(Option::is_none), + "{profile}: buffers" + ); + }; + refuse(&mut model, "LoRA"); + // Each of the other profiles alone is refused too; without any, the check passes. + let adapter = model.lora.take(); + model.ensure_fixed_profile().unwrap(); + model.fp32_gdn = Some(cuda::api().fp32_gdn().unwrap()); + refuse(&mut model, "FP32 GDN"); + model.fp32_gdn = None; + model.reference = Some(cuda::api().language_reference().unwrap()); + refuse(&mut model, "reference numerics"); + model.reference = None; + model.lora = adapter; + drop(model); + let _ = std::fs::remove_dir_all(&dir); +} diff --git a/tests/qwen3_5/profile_fake.c b/tests/qwen3_5/profile_fake.c new file mode 100644 index 00000000..add89a17 --- /dev/null +++ b/tests/qwen3_5/profile_fake.c @@ -0,0 +1,20 @@ +// CPU stand-in for the CUDA library that counts every call, for the profile test. +#include +#include +#include +static int calls; +int profile_fake_calls(void) { return calls; } +uint32_t cs1_abi_version(void) { ++calls; return PROFILE_FAKE_ABI; } +const char* cs1_error_string(int code) { (void)code; ++calls; return "unexpected fake CUDA call"; } +int cs1_get_device(int* device) { ++calls; if (!device) return 1; *device=0; return 0; } +int cs1_set_device(int device) { (void)device; ++calls; return 0; } +int cs1_malloc(void** p, size_t n) { ++calls; *p=malloc(n ? n : 1); return *p ? 0 : 1; } +int cs1_free(void* p) { ++calls; free(p); return 0; } +int cs1_stream_create(void** stream) { ++calls; *stream=(void*)1; return 0; } +int cs1_stream_sync(void* stream) { (void)stream; ++calls; return 0; } +int cs1_stream_destroy(void* stream) { (void)stream; ++calls; return 0; } +int cs1_upload(void* dst, const void* src, size_t n, void* stream) { (void)stream; ++calls; memcpy(dst,src,n); return 0; } +void* cs1_gemm_create(size_t n) { (void)n; ++calls; return &calls; } +void cs1_gemm_destroy(void* gemm) { (void)gemm; ++calls; } +void* cs1_gdn_fp32_create(size_t n) { (void)n; ++calls; return &calls; } +void cs1_gdn_fp32_destroy(void* handle) { (void)handle; ++calls; }