diff --git a/README.md b/README.md index 6e35c3ed..c3768df1 100644 --- a/README.md +++ b/README.md @@ -90,8 +90,9 @@ processors and executors. Native workers use shared FIFO admission and blocking dispatch per loaded executor. Processing orchestration, batch budgets, compatibility grouping and dynamic batching remain planned. -Open-Jev already packs candidates within one request for selected prefill GEMMs, -with bounded groups and independent sequence state. See the +Open-Jev-27B packs candidates within one request for compatible prefill GEMMs +and GDN launches, and can share a long prefix within one question. Each sequence +retains independent state. See the [native recipe](recipe/open_jev/native.md) and [matched H200 measurements](benchmarks/prefill_batching/README.md). @@ -182,7 +183,7 @@ checkpoint. Public231 concurrency-1 latency for heads `1ba07d1`, `f968e576`, | Cua-S1 4B 0.2 (`multimodal` adapter) | [Python CUDA reference](src/frontend/cua_s1.py) and [native Rust/CUDA screenshot worker](recipe/cua_s1/native_screenshot.md); one PNG/JPEG image, `choice`; scoped RTX 4090 validation | | JEMM (Qwen3.8-27B) | [Native Rust/CUDA worker](recipe/jemm/README.md); [text and 1–4-image frozen-corpus parity on A800](src/models/jemm/README.md#validation) | | Valen-Preview-0923 | [Python reference worker](recipe/valen/reference.md); text state or one inline PNG/JPEG image, `choice`/`noul`/`score`; author-run on RTX 5060 Laptop GPU (sm_120, WSL2), owner: community | -| Open-Jev-27B-v1.1 | [Native Rust/CUDA worker](recipe/open_jev/native.md); eager independent text candidates; [H200 validation](recipe/open_jev/validation.md) | +| Open-Jev-27B-v1.1 | [Native Rust/CUDA worker](recipe/open_jev/native.md); eager text decisions with within-request candidate optimizations; [historical H200 validation](recipe/open_jev/validation.md) | | Open-Jev-9B | The same [native Rust/CUDA worker](recipe/open_jev/native.md); eager independent text candidates; [reference comparison on sm_89](recipe/open_jev/validation-9b.md) | | StartLux-Decision v1.1 4B Instant | [Native Rust/CUDA text worker](recipe/startlux/README.md); `choice` / `noul` / `score`; weights CC BY-NC 4.0; Instant only (no AutoThink, no 27B); GPU parity unverified; [L20X latency recipe](recipe/startlux/latency.md) | | autotrust/JEV-27B-VL | [Experimental Rust/CUDA worker](recipe/jev_vl/README.md); single-question text, online inline PNG/JPEG and offline-preencoded image inputs; [validation scope and limits](recipe/jev_vl/validation.md) | diff --git a/docs/architecture.md b/docs/architecture.md index 197af0b6..5cfcaf14 100644 --- a/docs/architecture.md +++ b/docs/architecture.md @@ -14,8 +14,14 @@ single prompts and bounded packed prefill. StartLux Instant uses a bidirectional full-attention call on each prompt; its GDN stays causal, and it does not pack prompts or cache prefixes. 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. Decider optionally packs up to four complete rows and 4096 tokens within +input and gate/up GEMMs. Its opt-in output/down packing retains each original +matrix algorithm and checks the combined shape; incompatible groups keep separate +calls. Its batched linear-attention kernels preserve independent sequence state. +Long candidates from one question can share an exact prefix when every retained +projection algorithm supports the prefix and suffix shapes. This Open-Jev path +preserves the original complete-prompt arithmetic choices; it is separate from +the fixed-algorithm `forward_shared` interface described below. +Decider optionally packs up to four complete rows and 4096 tokens within one request, or selects request-local shared prefixes through the same canonical continuation path used by JEV-VL (below). Both optimizations default off. [Laya's native worker](../src/models/laya/README.md) diff --git a/docs/supported-models.md b/docs/supported-models.md index c5dbc4c2..767d71ae 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 | | Cua-S1 4B 0.2, `multimodal` adapter | [Native Rust/CUDA screenshot worker](../recipe/cua_s1/native_screenshot.md) | Preparation checks only; no CPU inference | RTX 4090 / sm_89, native vision and language Graph regressions; [scoped evidence and controls](../recipe/cua_s1/native_vision_graph.md) | Not supported | Verified BF16 language export, original vision checkpoint and FP32 vision adapter, rebuilt ABI 9 CUDA library; one PNG/JPEG and `choice` | | Valen-Preview-0923 | [Python reference worker](../recipe/valen/reference.md); text state or one inline PNG/JPEG image; `choice`, `noul`, and `score` questions | Not supported | Author-run on RTX 5060 Laptop (sm_120, WSL2): one text and one 2×2 PNG request; input and readout parity ≤1e-6 against reference (model outputs not compared); `noul`/`score` parity ≤1e-6 with exact legend and token agreement on the recipe's text fixtures. Owner: community. | Not supported | Python 3.12, pinned Valen source/checkpoint, Qwen3.5-2B base, BF16 CUDA | -| 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-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; rebuild worker and library together for the 27B candidate optimizations | | 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 | | StartLux-Decision v1.1 4B Instant | [Native Rust/CUDA worker](../recipe/startlux/README.md) | Contract checks only; no CPU inference | Unverified. Build target is L20X / Ada sm_89. The [latency recipe](../recipe/startlux/latency.md) is not a result | Not supported | Pinned CC BY-NC 4.0 Instant 4B checkpoint; rebuilt ABI 9 library with `cs1_attention_gated_bidirectional`. Instant only: AutoThink and 27B are refused | | Decider-2B v11 | [Native Rust/CUDA worker](../recipe/decider/README.md) | Contract checks only; no CPU inference | RTX 4090 (`sm_89`) synthetic parity corpus; [limits and provenance](../recipe/decider/validation.md) | Not supported | Pinned original BF16 checkpoint; ABI 9 library and consumers rebuilt together | diff --git a/recipe/README.md b/recipe/README.md index 041c091d..d087e6d9 100644 --- a/recipe/README.md +++ b/recipe/README.md @@ -25,7 +25,8 @@ For a first real decision, follow the [complete CPU walkthrough](../docs/getting multimodal language checkpoint and serve one-image decisions with Rust/CUDA. - [Open-Jev native text worker](open_jev/native.md): export the merged text backbone and trained decision head of Open-Jev-27B-v1.1 or Open-Jev-9B, then - serve with Rust and CUDA. [Open-Jev-9B validation](open_jev/validation-9b.md) + serve with Rust and CUDA; the 27B path supports compatible candidate batching + and within-question prefix sharing. [Open-Jev-9B validation](open_jev/validation-9b.md) compares the 9B worker with the reference. - [StartLux Instant 4B native worker](startlux/README.md): build the Ada sm_89 CUDA library, serve the pinned CC BY-NC Instant checkpoint, and run the diff --git a/recipe/open_jev/native.md b/recipe/open_jev/native.md index d78315d3..ccbe0687 100644 --- a/recipe/open_jev/native.md +++ b/recipe/open_jev/native.md @@ -69,11 +69,13 @@ 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 9, which includes shared vision, CUDA Graph, canonical prefix continuation, fixed-GEMM selection and device-feature validation. +The 27B worker also requires the additive packed-GEMM and batched-GDN +symbols from this checkout; rebuild the worker and library together. ```sh src/backends/cuda/qwen3_5/build.sh target/release 89 cargo build --release --locked -p omni-open-jev-native -p omni-jev -OPEN_JEV_MODEL=weights/open-jev-27b-merged \ +OPEN_JEV_MODEL=weights/open-jev-27b-merged RAYON_NUM_THREADS=4 \ target/release/omni-open-jev-native ``` @@ -81,6 +83,10 @@ For 9B, set `OPEN_JEV_MODEL=weights/open-jev-9b-merged`. `OPEN_JEV_HOST` and `OPEN_JEV_PORT` default to `127.0.0.1` and `8000`. `OPEN_JEV_CUDA_LIB` overrides the default library next to the executable. +For 27B, `RAYON_NUM_THREADS=4` bounds the tokenizer's thread pool in the example; +select a thread count appropriate to the worker's CPU allocation. Tokenization +processes at most 16 candidates from one question at a time. Single candidates, +configured padding and the 9B checkpoint keep per-candidate encoding. The worker loads all text weights onto visible CUDA device 0, performs a real warmup inference, then exposes `/health` and `/v1/systemone`. Use a reservation before any GPU command on hosts with a GPU scheduler. @@ -127,7 +133,8 @@ OPEN_JEV_MODEL=weights/open-jev-9b-merged \ CUA_S1_CUDA_LIB=$PWD/target/release/libqwen3_5_cuda.so \ cargo test --release --locked -p omni-qwen3-5-native --test kernels -- --ignored -# Full-checkpoint packing, ordering and graph-shape checks, in the reservation: +# Full-checkpoint packing, ordering, graph replay and prefix-buffer reuse checks: +# Run inside the reservation; repeat with CUA_S1_GRAPH=0 for eager execution. OPEN_JEV_MODEL=$PWD/weights/open-jev-27b-merged \ OPEN_JEV_CUDA_LIB=$PWD/target/release/libqwen3_5_cuda.so CUA_S1_GRAPH=1 \ cargo test --release --locked -p omni-open-jev-native --test prefill_batch \ @@ -155,24 +162,97 @@ This recipe leaves `CUA_S1_GRAPH` unset and uses eager prefill. For Open-Jev-27B-v1.1, candidates within a request are packed in prepared order, up to 16 sequences and 4096 total tokens per group; longer prompts execute alone without truncation. Input and gate/up -projections share GEMMs. Output/down projections preserve their per-prompt shapes -and reduction order, and each sequence retains independent attention, positions, -convolution and GDN state. Calibration still uses every candidate in its question. +projections share GEMMs. Output/down projections combine consecutive candidates +only when their original cuBLASLt algorithms are identical and the combined shape +passes its workspace check. Unsupported shapes retain separate calls. This keeps +the original algorithm and reduction configuration; numerical validation is still +required for each supported hardware/toolchain combination. Batched GDN launches +keep independent state and chunk-rounded workspace for every sequence. Attention, +positions and convolution also remain independent. Calibration still uses every +candidate in its question. The [H200 packing comparison](../../benchmarks/prefill_batching/README.md) records latency, exact output checks and frozen controls. -Packing validation covers H200 (sm_90); other CUDA architectures remain unverified. +That historical comparison covers the earlier packing implementation on H200 +(sm_90). The newer projection/GDN paths have exact-output regression coverage on +an SM90 device; their execution on other CUDA architectures remains unverified. Open-Jev-9B runs one forward pass per candidate: packing is not validated for it. +For one question with 2–16 candidates all longer than 2048 tokens, the worker can +share an exact token prefix within compatible candidate groups. The prefix must +be at least 512 tokens, aligned to 64 tokens, with a nonempty suffix for every +candidate. Compatibility checks retain the complete prompt's algorithms for every +prefix/suffix projection. A declined group uses complete forwards. Group outputs +are restored to the original candidate order before normalization. The model +retains allocation capacity for one prefix, up to the loaded maximum length; +every request invalidates and rewrites its numeric state. It does not cache request +contents across requests. The response reports executed reuse through +`metadata.prefix_cache`: + +| Field | Meaning in this worker | +| --- | --- | +| `policy` | `default`: no explicit prefix-policy override; compatible long candidates may still share automatically. | +| `enabled` | Whether this request actually reused a prefix after compatibility checks. | +| `scope` | `within_question` when sharing ran, otherwise `none`. | +| `saved_tokens` | Prefix tokens whose repeated computation was avoided, summed over executed compatible groups. | + +Each compatible group of `n` candidates with a prefix of `p` tokens saves +`(n - 1) * p` token computations. Singletons save zero; separate groups each +compute their own prefix. For three candidates with a 3072-token prefix, all +sharing saves 6144 tokens, while two sharing and one falling back saves only +3072. Splitting or packing the suffixes does not change this prefix count. +These four fields are the complete prefix metadata contract; `mode` and +`shared_tokens` are not emitted. `usage.input_tokens` still counts every complete +candidate's input. No sharing returns `enabled=false`, `scope=none` and +`saved_tokens=0`. This worker does not implement the explicit fixed/shared/auto +policy switches or cross-question reuse proposed in +[PR #179](https://github.com/InferMatrixLab/system1-omni/pull/179). + +The default 27B optimization is validated against the original default +complete-prompt execution with the same weights, inputs and temperature: +maximum absolute candidate-probability and ordinal-score error at most 0.001, +zero Choice changes, zero Noul flips at the `>= 0.5` threshold, and zero request +failures. These are comparison-test gates, not a second baseline inference on +every served request. Retaining the original projection algorithms and falling +back when incompatible remains the execution policy. Equality to an alternative +fixed-algorithm path alone does not establish these default-path gates. The +accepted historical latency tradeoff below does not relax numerical validation. + +Compatible suffixes can also share compact storage and projection calls, bounded +to 16 sequences and 4096 suffix tokens per group. Attention reads immutable prefix +keys/values separately; convolution reads the prefix's three retained rows; each +linear-attention sequence starts from its own copy of the group state. If the +optional compact-suffix symbols or an algorithm check are unavailable, suffixes +use the original continuation path. + +The 27B opt-in also uses vector convolution for aligned layouts, wider work +assignment for Gated DeltaNet preparation, and independent-column output tiling. +These preserve each accumulator's arithmetic and BF16 rounding boundaries. For +one ordinary text execution unit of 512–16384 tokens, only the final position +needs the last layer's feed-forward network and final norms. The model skips +unused rows only when both projections accept one row with their original +complete-prompt algorithms. Prefix/suffix, multimodal, bidirectional, unmerged, +reference and incompatible paths retain their complete computation. + +Packed projection plans and retained prefix/suffix plans each have a separate +512-entry host LRU cache. These hold shape/algorithm descriptions, not token values. +GPU tests cover replay of an older graph after plan eviction, changed request +contents and prefix-buffer growth/shrinkage. Batched GDN adds approximately +96 MiB of bounded workspace for this 27B configuration; prefix storage grows with +the largest shared prefix encountered. Include both allocations in memory +measurements and distinguish first traversal from warmed latency. + Set `CUA_S1_GRAPH=1` on the worker to enable CUDA Graph replay. The -shared backend retains at most 64 graphs, keyed by ordered sequence token lengths; -growing the scratch buffer clears them. A new shape first runs an eager forward +27B packed text path retains at most 256 graphs; other modes retain 64. Keys are +ordered sequence token lengths. Graph-enabled packed text initially reserves +scratch for at least 4096 rows, rounded up in 1024-row increments. Longer supported +inputs can still grow the buffer; growth synchronizes work and clears graphs +whose addresses would become invalid. A new shape first runs an eager forward to initialize its plans and captures the layer loop for later replay. This adds cost for new lengths, so graph mode remains opt-in. The earlier single-prompt graph comparison with Open-Jev-27B-v1.1 on the 74-case H200 workload measured mean HTTP latency 2.03% below eager execution after all lengths were -warmed. Combined -packing and graph performance remains unmeasured. Tokenization, transfers -and the CPU scalar head remain outside the graph. Prefix sharing, GEMM autotuning, +warmed. These historical numbers describe that earlier implementation. Tokenization, transfers +and the CPU scalar head remain outside the graph. Cross-request prefix-content caching, GEMM autotuning, quantization and multimodal inference are not implemented. The [H200 validation](validation.md) reports Open-Jev-27B-v1.1 results for 74 single-candidate requests, including probability differences and timing @@ -180,3 +260,46 @@ variability. It does not establish general accuracy parity or a speedup over OpenJev-Fast; the author's B300 results use different hardware and workloads. The [Open-Jev-9B validation](validation-9b.md) compares the 9B worker with the reference on one RTX 6000 Ada; it makes no claim about 27B. + +### Cumulative 27B experiment and accepted tradeoff + +The October 22–24, 2026 experiment compared frozen baseline +`d577bb824a5dbf0f5e385cf829bb6b9025dabf8d` with cumulative candidate P19 +(patch SHA-256 `5bb08cafffcf1825c1467de36705ae78baf07e0f7b24662c67ed7237b6504799`). +P19 predates integration with the current shared executor; its measurements are +historical evidence, not measurements of every later checkout. + +The workload was 231 fixed public JevBench requests (139 Choice, 74 Noul, +18 Score), single-request concurrency through the Rust frontend, with complete +response receipt, parsing and validation inside the client timer. Twelve balanced +blocks compared baseline, parent P18 and P19; each service start had one separate +feasibility traversal and eight measured traversals. Baseline used eager execution; +both candidates enabled Graph replay. Across all three arms, 66,528 responses +were timed and 8,316 were feasibility responses. No observations were excluded. + +| Historical metric | Baseline | P19 | Change | +| --- | ---: | ---: | ---: | +| Mean full-request latency | 216.31 ms | 123.02 ms | -43.13% | +| Median full-request latency | 61.04 ms | 44.67 ms | -26.82% | +| 95th-percentile full-request latency | 1081.56 ms | 471.11 ms | -56.44% | +| Noul, 513–2048 tokens, 95th percentile | 116.32 ms | 119.77 ms | **+2.97%** | +| Sampled whole-invocation peak device memory | 50,949 MiB | 52,035 MiB | +1086 MiB | + +Device: driver-reported **NVIDIA L20X**, measured compute capability **9.0**, +132 streaming multiprocessors, about 150.1 GB visible memory; the conflicting +name/properties do not establish a commercial model. CUDA 13.0.88, driver +570.133.20, BF16 merged weights and the unchanged FP32 scalar head were used. +Memory was sampled approximately every 0.5 seconds, including loading and warmup; +brief peaks may be missed. Loading is outside the request timer. + +All 74,844 historical responses had zero observed probability/Score drift, +decision flips and request failures. The six-request Noul group above regressed +in all 12 blocks; its paired approximate 95% interval was +2.32% to +3.70%. +The contributor accepts this specific tradeoff for submission. The original +2% all-group guard did not pass: four additional tail checks were inconclusive +(baseline: Score 129–512; parent: Noul 1–128, Score 1–128 and Score 129–512). +They remain inconclusive, and the conditional final synthetic timing stage was +not run. Earlier candidate synthetic passes do not validate P19's missing stage. +Multi-request concurrency, other architectures and other complete checkpoints +remain unverified for these optimizations. Source reconstruction, current-head +checks and exact reproduction/evidence records belong to the accompanying PR. diff --git a/src/backends/cuda/qwen3_5/README.md b/src/backends/cuda/qwen3_5/README.md index ee2dbeba..3becc3b8 100644 --- a/src/backends/cuda/qwen3_5/README.md +++ b/src/backends/cuda/qwen3_5/README.md @@ -51,11 +51,30 @@ 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. The shared Rust model can pack independent sequences for input and gate/up GEMMs. -Output/down GEMMs retain each prompt's original shape and reduction order; -attention, convolution and GDN calls remain sequence-local. Packing adds no CUDA -entry points. Open-Jev uses this path within requests; Cua-S1 keeps single-prompt calls. Decider optionally packs complete +Output/down GEMMs default to each prompt's original shape and reduction order; +attention, convolution and GDN calls remain sequence-local in that default path. +Cua-S1 keeps single-prompt calls. Decider optionally packs complete question rows within one admitted request. +Open-Jev-27B also opts into additive `cs1_gemm_packed` and +`cs1_gdn_prefill_batch` entry points. Packed projections combine only consecutive +shapes with byte-identical original cuBLASLt algorithms that pass `AlgoCheck` +within the existing workspace. Unsupported groups split back into smaller groups +or independent calls. Batched GDN uses a by-value descriptor for up to 16 sequences, +separate chunk-rounded workspace, and zero initial state for each sequence. The +original single-sequence entry points remain available to other consumers. + +Optional `cs1_gemm_reuse_check` / `cs1_gemm_reuse` entry points retain a complete +prompt's algorithm for a prefix or suffix window. A compatibility result of `-1` +means the caller must fall back before executing the shared prefix. The model +adapter owns candidate grouping and prefix state; the backend owns only the GEMM +plans and execution. Packed plans and window plans each use a 512-entry host LRU, +with graph replay tests after eviction. These additions keep ABI 9 and do not +require other model consumers to enable the new paths. Rebuild the library and +27B worker from the same checkout. New batched-kernel execution is validated on +SM90; SM80/SM89 execution remains unverified. + + 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 and by the shared executor's `forward_shared`, which diff --git a/src/backends/cuda/qwen3_5/attention.cu b/src/backends/cuda/qwen3_5/attention.cu index 24a03909..5193e05f 100644 --- a/src/backends/cuda/qwen3_5/attention.cu +++ b/src/backends/cuda/qwen3_5/attention.cu @@ -10,6 +10,7 @@ // 32, keeping the output and the online softmax in registers. The probabilities are // rounded to bfloat16 for the P*V product, as in flash attention; the running sums // stay float32. +#include #include "common.cuh" #include "mma.cuh" #include "ops.h" @@ -85,11 +86,11 @@ constexpr int SMEM_BYTES = (BM + 2 * BN) * LDS * 2; // q_base offsets row 0 of the q/gate/out buffers: this launch covers the window // [q_base, T_total) while k/v cover [0, T_total). Used by the cached-prefix // continuation; q_base = 0 reduces to the plain full pass. -template -__global__ void __launch_bounds__(THREADS) - flash_kernel(const bf16* __restrict__ q, const bf16* __restrict__ k, const bf16* __restrict__ v, int ldv, +template +__device__ __forceinline__ void flash_body(const bf16* __restrict__ q, const bf16* __restrict__ k, const bf16* __restrict__ v, int ldv, const bf16* __restrict__ gate, bf16* __restrict__ out, int T, int Hq, int Hk, - float scale_log2, int q_base) { + float scale_log2, int q_base, const bf16* __restrict__ prefix_k, + const bf16* __restrict__ prefix_v) { extern __shared__ __align__(16) unsigned char smem[]; bf16* qs = reinterpret_cast(smem); bf16* ks = qs + BM * LDS; @@ -102,7 +103,8 @@ __global__ void __launch_bounds__(THREADS) for (int c = tid; c < BM * (D / 8); c += THREADS) { const int r = c / (D / 8), col = (c % (D / 8)) * 8, row = q0 + r; - cp_async16(qs + r * LDS + col, q + ((size_t)min(row, T - 1) * Hq + h) * D + col, row < T); + const int source_row = min(row, T - 1) - (Shared ? q_base : 0); + cp_async16(qs + r * LDS + col, q + ((size_t)source_row * Hq + h) * D + col, row < T); } cp_async_commit(); @@ -116,12 +118,26 @@ __global__ void __launch_bounds__(THREADS) for (int k0 = 0; k0 < kv_end; k0 += BN) { for (int c = tid; c < BN * (D / 8); c += THREADS) { const int r = c / (D / 8), col = (c % (D / 8)) * 8, s = k0 + r; - cp_async16(ks + r * LDS + col, k + ((size_t)min(s, T - 1) * Hk + hk) * D + col, s < T); + if constexpr (Shared) { + const int row = min(s, T - 1); + const bf16* source = row < q_base ? prefix_k + ((size_t)row * Hk + hk) * D + : k + ((size_t)(row - q_base) * Hk + hk) * D; + cp_async16(ks + r * LDS + col, source + col, s < T); + } else { + cp_async16(ks + r * LDS + col, k + ((size_t)min(s, T - 1) * Hk + hk) * D + col, s < T); + } } cp_async_commit(); for (int c = tid; c < BN * (D / 8); c += THREADS) { const int r = c / (D / 8), col = (c % (D / 8)) * 8, s = k0 + r; - cp_async16(vs + r * LDS + col, v + (size_t)min(s, T - 1) * ldv + (size_t)hk * D + col, s < T); + if constexpr (Shared) { + const int row = min(s, T - 1); + const bf16* source = row < q_base ? prefix_v + (size_t)row * Hk * D + : v + (size_t)(row - q_base) * ldv; + cp_async16(vs + r * LDS + col, source + (size_t)hk * D + col, s < T); + } else { + cp_async16(vs + r * LDS + col, v + (size_t)min(s, T - 1) * ldv + (size_t)hk * D + col, s < T); + } } cp_async_commit(); cp_async_wait<1>(); // Q and K @@ -223,12 +239,13 @@ __global__ void __launch_bounds__(THREADS) for (int r = 0; r < 2; r++) { const int row = row0 + g + r * 8; if (row >= T) continue; - bf16* dst = out + ((size_t)row * Hq + h) * D + 2 * t; + const int output_row = row - (Shared ? q_base : 0); + bf16* dst = out + ((size_t)output_row * Hq + h) * D + 2 * t; #pragma unroll for (int n = 0; n < D / 8; n++) { float a = o[n][2 * r] * inv[r], b = o[n][2 * r + 1] * inv[r]; if constexpr (Gated) { - const size_t idx = ((size_t)row * Hq + h) * D + 2 * t + n * 8; + const size_t idx = ((size_t)output_row * Hq + h) * D + 2 * t + n * 8; // Match attention -> BF16 store -> BF16 sigmoid -> multiply. // Rounding before the multiply is required even without that store. a = round_bf16(a) * round_bf16(sigmoid(f32(gate[idx]))); @@ -239,6 +256,24 @@ __global__ void __launch_bounds__(THREADS) } } +// Keep the legacy entry points and parameter layouts. The shared specialization +// reads immutable prefix KV separately, while Q/gate/output hold suffix rows only. +template +__global__ void __launch_bounds__(THREADS) +flash_kernel(const bf16* __restrict__ q, const bf16* __restrict__ k, const bf16* __restrict__ v, int ldv, + const bf16* __restrict__ gate, bf16* __restrict__ out, int T, int Hq, int Hk, + float scale_log2, int q_base) { + flash_body(q, k, v, ldv, gate, out, T, Hq, Hk, scale_log2, q_base, nullptr, nullptr); +} + +__global__ void __launch_bounds__(THREADS) +flash_shared_kernel(const bf16* __restrict__ q, const bf16* __restrict__ k, const bf16* __restrict__ v, int ldv, + const bf16* __restrict__ gate, bf16* __restrict__ out, int T, int Hq, int Hk, + float scale_log2, int prefix, const bf16* __restrict__ prefix_k, + const bf16* __restrict__ prefix_v) { + flash_body(q, k, v, ldv, gate, out, T, Hq, Hk, scale_log2, prefix, prefix_k, prefix_v); +} + template int launch(const void* q, const void* k, const void* v, int ldv, const void* gate, void* out, int T, int Hq, int Hk, int Dh, float scale, int q_base, void* stream) { @@ -309,3 +344,23 @@ extern "C" int cs1_attention_gated_prefix(const void* q, const void* k, const vo void* stream) { return flash::launch(q, k, v, ldv, gate, out, T, Hq, Hk, Dh, scale, q_base, stream); } + +// Compact suffix tensors with one immutable, contiguous prefix KV allocation. +extern "C" int cs1_attention_gated_shared(const void* q, const void* k, const void* v, int ldv, + const void* gate, void* out, const void* prefix_k, + const void* prefix_v, int suffix, int prefix, int Hq, + int Hk, int Dh, float scale, void* stream) { + if (!q || !k || !v || !gate || !out || !prefix_k || !prefix_v || suffix <= 0 || prefix <= 0 || + prefix % flash::BM || prefix > INT_MAX - suffix || Dh != flash::D || Hq <= 0 || Hk <= 0 || + Hq % Hk || ldv % 8 || ldv < Hk * Dh) return cudaErrorInvalidValue; + static const cudaError_t configured = cudaFuncSetAttribute( + flash::flash_shared_kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, flash::SMEM_BYTES); + if (configured != cudaSuccess) return configured; + constexpr float LOG2E = 1.4426950408889634f; + flash::flash_shared_kernel<<(stream)>>>( + static_cast(q), static_cast(k), static_cast(v), ldv, + static_cast(gate), static_cast(out), prefix + suffix, Hq, Hk, scale * LOG2E, + prefix, static_cast(prefix_k), static_cast(prefix_v)); + return cudaGetLastError(); +} diff --git a/src/backends/cuda/qwen3_5/elementwise.cu b/src/backends/cuda/qwen3_5/elementwise.cu index 0e4c0526..ea1b9d3e 100644 --- a/src/backends/cuda/qwen3_5/elementwise.cu +++ b/src/backends/cuda/qwen3_5/elementwise.cu @@ -39,6 +39,64 @@ __global__ void gdn_conv_kernel(const bf16* __restrict__ qkv, int ld, const bf16 v[(size_t)t * value_dim + c - 2 * key_dim] = y; } +// Identical four-tap accumulation with the preceding three projection rows. +__global__ void gdn_conv_prefix_kernel(const bf16* __restrict__ qkv, const bf16* __restrict__ history, int ld, const bf16* __restrict__ w, + bf16* __restrict__ q, bf16* __restrict__ k, bf16* __restrict__ v, int T, + int key_dim, int value_dim) { + const int channels = 2 * key_dim + value_dim; + const size_t idx = (size_t)blockIdx.x * blockDim.x + threadIdx.x; + if (idx >= (size_t)T * channels) return; + const int t = idx / channels, c = idx % channels; + float acc = 0.f; +#pragma unroll + for (int j = 0; j < 4; j++) { + const int s = t - 3 + j; + const bf16* source = s < 0 ? history + (size_t)(s + 3) * ld : qkv + (size_t)s * ld; + acc = fmaf(f32(w[c * 4 + j]), f32(source[c]), acc); + } + const bf16 y = to_bf16(silu(round_bf16(acc))); + if (c < key_dim) + q[(size_t)t * key_dim + c] = y; + else if (c < 2 * key_dim) + k[(size_t)t * key_dim + c - key_dim] = y; + else + v[(size_t)t * value_dim + c - 2 * key_dim] = y; +} + +// Eight adjacent channels retain the original four FMAs and both BF16 roundings. +// Separate token/channel grid dimensions avoid per-element integer division. +template +__global__ void gdn_conv_vector_kernel(const Pack8* __restrict__ qkv, + const Pack8* __restrict__ history, int ld, + const Pack8* __restrict__ w, Pack8* __restrict__ q, + Pack8* __restrict__ k, Pack8* __restrict__ v, + int key_packs, int value_packs) { + const int t = blockIdx.y, c = blockIdx.x * blockDim.x + threadIdx.x; + if (c >= 2 * key_packs + value_packs) return; + const Pack8 weights[4] = {w[c * 4], w[c * 4 + 1], w[c * 4 + 2], w[c * 4 + 3]}; + float acc[8] = {}; +#pragma unroll + for (int j = 0; j < 4; j++) { + const int row = t - 3 + j; + if (Prefix || row >= 0) { + const Pack8 value = row < 0 ? history[(size_t)(row + 3) * ld + c] + : qkv[(size_t)row * ld + c]; +#pragma unroll + for (int e = 0; e < 8; e++) + acc[e] = fmaf(f32(weights[e / 2].v[(e % 2) * 4 + j]), f32(value.v[e]), acc[e]); + } + } + Pack8 out; +#pragma unroll + for (int e = 0; e < 8; e++) out.v[e] = to_bf16(silu(round_bf16(acc[e]))); + if (c < key_packs) + q[(size_t)t * key_packs + c] = out; + else if (c < 2 * key_packs) + k[(size_t)t * key_packs + c - key_packs] = out; + else + v[(size_t)t * value_packs + c - 2 * key_packs] = out; +} + // beta = sigmoid(b) in bfloat16; g = -exp(A_log) * softplus(a + dt_bias) in float32 // (F.softplus with threshold 20). __global__ void gdn_gates_kernel(const bf16* __restrict__ b, const bf16* __restrict__ a, int ld, @@ -160,3 +218,44 @@ extern "C" int cs1_bf16_finite(const void* x, size_t n, uint32_t* flag, void* st static_cast(x), n, flag); return cudaGetLastError(); } + +extern "C" int cs1_gdn_conv_prefix(const void* qkv, const void* history, int ld, const void* w, + void* q, void* k, void* v, int T, int key_dim, int value_dim, + void* stream) { + if (!qkv || !history || !w || !q || !k || !v || T <= 0 || key_dim <= 0 || value_dim <= 0 || + ld < 2 * key_dim + value_dim) return cudaErrorInvalidValue; + const size_t n = (size_t)T * (2 * key_dim + value_dim); + const uintptr_t alignment = reinterpret_cast(qkv) | reinterpret_cast(history) | + reinterpret_cast(w) | reinterpret_cast(q) | reinterpret_cast(k) | + reinterpret_cast(v); + if (T >= 64 && T <= 65535 && key_dim % 8 == 0 && value_dim % 8 == 0 && ld % 8 == 0 && !(alignment & 15)) { + const int packs = (2 * key_dim + value_dim) / 8; + gdn_conv_vector_kernel<<(stream)>>>( + static_cast(qkv), static_cast(history), ld / 8, + static_cast(w), static_cast(q), static_cast(k), + static_cast(v), key_dim / 8, value_dim / 8); + } else { + gdn_conv_prefix_kernel<<(stream)>>>( + static_cast(qkv), static_cast(history), ld, static_cast(w), + static_cast(q), static_cast(k), static_cast(v), T, key_dim, value_dim); + } + return cudaGetLastError(); +} + +// Additive opt-in entry point; unsupported layouts preserve the scalar path. +extern "C" int cs1_gdn_conv_vector(const void* qkv, int ld, const void* w, void* q, + void* k, void* v, int T, int key_dim, int value_dim, + void* stream) { + const uintptr_t alignment = reinterpret_cast(qkv) | reinterpret_cast(w) | + reinterpret_cast(q) | reinterpret_cast(k) | reinterpret_cast(v); + if (T < 64 || T > 65535 || key_dim <= 0 || value_dim <= 0 || + key_dim % 8 || value_dim % 8 || ld % 8 || ld < 2 * key_dim + value_dim || (alignment & 15)) + return cs1_gdn_conv(qkv, ld, w, q, k, v, T, key_dim, value_dim, stream); + const int packs = (2 * key_dim + value_dim) / 8; + gdn_conv_vector_kernel<<(stream)>>>( + static_cast(qkv), nullptr, ld / 8, static_cast(w), + static_cast(q), static_cast(k), static_cast(v), key_dim / 8, value_dim / 8); + return cudaGetLastError(); +} diff --git a/src/backends/cuda/qwen3_5/gdn_prefill.cu b/src/backends/cuda/qwen3_5/gdn_prefill.cu index 8355a771..7201307f 100644 --- a/src/backends/cuda/qwen3_5/gdn_prefill.cu +++ b/src/backends/cuda/qwen3_5/gdn_prefill.cu @@ -17,6 +17,8 @@ #include #include #include +#include +#include #include "mma.cuh" #include "ops.h" @@ -107,10 +109,11 @@ static_assert(C * HB * 2 <= C * CP * 4, "k exp(cum) as bf16 fits over A"); static_assert(C * sizeof(Tf32Row) <= R3 - R2, "packed qn fits in R2"); static_assert(C * HB * 2 <= R2 - R1, "v as bf16 fits over kn"); -__global__ void __launch_bounds__(THREADS) gdn_chunk_prep( +template +__device__ __forceinline__ void gdn_chunk_prep_body( const __nv_bfloat16* __restrict__ q, const __nv_bfloat16* __restrict__ k, const __nv_bfloat16* __restrict__ v, const float* __restrict__ g, - const __nv_bfloat16* __restrict__ beta, Work ws, int T, int H, int HK, float scale) { + const __nv_bfloat16* __restrict__ beta, Work ws, int T, int H, int HK, float scale, int NC) { extern __shared__ __align__(128) unsigned char sm[]; Tf32Row* kn = reinterpret_cast(sm + R1); Tf32Row* qn = reinterpret_cast(sm + R2); @@ -126,7 +129,7 @@ __global__ void __launch_bounds__(THREADS) gdn_chunk_prep( __nv_bfloat16* tb = reinterpret_cast<__nv_bfloat16*>(sm + R2 + TBF); __nv_bfloat16* kw = reinterpret_cast<__nv_bfloat16*>(sm + R3); - const int c = blockIdx.x, h = blockIdx.y, NC = gridDim.x, NCC = NC * C; + const int c = blockIdx.x, h = blockIdx.y, NCC = NC * C; const int hk = h / (H / HK); const int tid = threadIdx.x, lane = tid & 31, warp = tid >> 5; @@ -167,7 +170,7 @@ __global__ void __launch_bounds__(THREADS) gdn_chunk_prep( // 2. L2-normalize q/k in float32 and form decayed q/k before rounding to // bfloat16. Convert the pair-product operands to TF32 once per row. - for (int r = warp; r < C; r += THREADS / 32) { + for (int r = warp; r < C; r += BLOCK_THREADS / 32) { const int t = c * C + r; float kv[4], qv[4], ks = 0.f, qs = 0.f; #pragma unroll @@ -207,7 +210,7 @@ __global__ void __launch_bounds__(THREADS) gdn_chunk_prep( // 3. pair products on tensor cores: the 10 lower 16x16 tiles of k.k (into A) and of // q.k (into P, in global memory), then the masks and decays elementwise float* pout = ws.p + ((size_t)h * NC + c) * C * C; - for (int e = warp; e < 20; e += THREADS / 32) { + for (int e = warp; e < 20; e += BLOCK_THREADS / 32) { const int tile = e % 10; const int it = tile < 1 ? 0 : tile < 3 ? 1 : tile < 6 ? 2 : 3; const int jt = tile - it * (it + 1) / 2; @@ -239,7 +242,7 @@ __global__ void __launch_bounds__(THREADS) gdn_chunk_prep( } __syncthreads(); __nv_bfloat16* pb = ws.pb + ((size_t)h * NC + c) * C * C; - for (int x = tid; x < C * C; x += THREADS) { + for (int x = tid; x < C * C; x += BLOCK_THREADS) { const int i = x / C, j = x % C; const float dec = j <= i ? expf(cum[i] - cum[j]) : 0.f; am[i * CP + j] = j < i ? bet[i] * am[i * CP + j] * dec : 0.f; @@ -248,7 +251,7 @@ __global__ void __launch_bounds__(THREADS) gdn_chunk_prep( __syncthreads(); // 4. T = (I + A)^-1 in 16x16 blocks - for (int x = tid; x < C * CP; x += THREADS) tm[x] = 0.f; + for (int x = tid; x < C * CP; x += BLOCK_THREADS) tm[x] = 0.f; __syncthreads(); if (tid < C) { const int base = 16 * (tid / 16), col = tid % 16; @@ -266,7 +269,7 @@ __global__ void __launch_bounds__(THREADS) gdn_chunk_prep( __syncthreads(); for (int lv = 1; lv < 4; lv++) { const int nb = 4 - lv; - for (int x = tid; x < nb * 256; x += THREADS) { + for (int x = tid; x < nb * 256; x += BLOCK_THREADS) { const int bi = lv + x / 256, bj = bi - lv, r = (x % 256) / 16, cc = x % 16; float acc = 0.f; for (int kb = bj; kb < bi; kb++) @@ -276,7 +279,7 @@ __global__ void __launch_bounds__(THREADS) gdn_chunk_prep( sc[x] = acc; } __syncthreads(); - for (int x = tid; x < nb * 256; x += THREADS) { + for (int x = tid; x < nb * 256; x += BLOCK_THREADS) { const int bi = lv + x / 256, bj = bi - lv, r = (x % 256) / 16, cc = x % 16; const float* s0 = sc + (x / 256) * 256 + cc; float acc = 0.f; @@ -288,24 +291,24 @@ __global__ void __launch_bounds__(THREADS) gdn_chunk_prep( // 5. bfloat16 operands: k exp(cum) over A, T beta after T, then v over kn; // u = T (beta v) and w = T (beta exp(cum) k) on tensor cores, stored as bfloat16 - for (int x = tid; x < C * K; x += THREADS) { + for (int x = tid; x < C * K; x += BLOCK_THREADS) { const int j = x / K, d = x % K; const int t = c * C + j; const float kx = t < T ? __bfloat162float(k[((size_t)t * HK + hk) * K + d]) : 0.f; kw[j * HB + d] = __float2bfloat16((kx * kinvs[j]) * ecum[j]); } - for (int x = tid; x < C * C; x += THREADS) { + for (int x = tid; x < C * C; x += BLOCK_THREADS) { const int i = x / C, j = x % C; tb[i * TB + j] = __float2bfloat16(tm[i * CP + j] * bet[j]); } __syncthreads(); - for (int x = tid; x < C * V; x += THREADS) { + for (int x = tid; x < C * V; x += BLOCK_THREADS) { const int j = x / V, d = x % V; const int t = c * C + j; vb[j * HB + d] = t < T ? v[((size_t)t * H + h) * V + d] : __float2bfloat16(0.f); } __syncthreads(); - for (int e = warp; e < 64; e += THREADS / 32) { + for (int e = warp; e < 64; e += BLOCK_THREADS / 32) { const bool is_u = e < 32; const int it = (e % 32) / 8, dt = e % 8; const __nv_bfloat16* bsrc = is_u ? vb : kw; @@ -337,7 +340,7 @@ constexpr int SS_LD = BVS + 8; // bfloat16 row stride of the S copy and v_n constexpr int STAGE = C * WS_LD; // elements of one staged w or kd constexpr size_t SMEM2_BYTES = (4 * STAGE + K * SS_LD + C * SS_LD) * 2; -__global__ void __launch_bounds__(ST_THREADS) gdn_chunk_state(Work ws, int NC, const float* __restrict__ S_IN, +__device__ __forceinline__ void gdn_chunk_state_body(Work ws, int NC, const float* __restrict__ S_IN, float* __restrict__ S_OUT) { extern __shared__ __align__(128) unsigned char sm[]; bf16* wbuf = reinterpret_cast(sm); // [2][C][WS_LD] @@ -478,13 +481,13 @@ constexpr int P_LD = C + 8; // bfloat16 row stride of P constexpr size_t SMEM3_BYTES = (C * O_LD + K * O_LD + C * P_LD + C * O_LD) * 2; static_assert(K == V, "S, qd and v_new share a row stride"); -__global__ void __launch_bounds__(OUT_THREADS) gdn_chunk_out(Work ws, bf16* __restrict__ o, int T, int H) { +__device__ __forceinline__ void gdn_chunk_out_body(Work ws, bf16* __restrict__ o, int T, int H, int NC) { extern __shared__ __align__(128) unsigned char sm[]; bf16* qs = reinterpret_cast(sm); // [C][O_LD] bf16* ss = qs + C * O_LD; // [K][O_LD] bf16* ps = ss + K * O_LD; // [C][P_LD] bf16* vs = ps + C * P_LD; // [C][O_LD] - const int c = blockIdx.x, h = blockIdx.y, NC = gridDim.x, NCC = NC * C; + const int c = blockIdx.x, h = blockIdx.y, NCC = NC * C; const int tid = threadIdx.x, warp = tid / 32, lane = tid % 32, g = lane / 4, t = lane % 4; const size_t rows = (size_t)h * NCC + (size_t)c * C; for (int x = tid; x < C * K / 8; x += OUT_THREADS) { @@ -543,6 +546,132 @@ __global__ void __launch_bounds__(OUT_THREADS) gdn_chunk_out(Work ws, bf16* __re } } +// Both dispatch paths execute the same arithmetic body. A batched launch has +// separate state/workspace for each sequence; no operation crosses a boundary. +__global__ void __launch_bounds__(THREADS) gdn_chunk_prep( + const bf16* __restrict__ q, const bf16* __restrict__ k, const bf16* __restrict__ v, + const float* __restrict__ g, const bf16* __restrict__ beta, + Work ws, int T, int H, int HK, float scale) { + gdn_chunk_prep_body(q, k, v, g, beta, ws, T, H, HK, scale, gridDim.x); +} +// More warps share independent rows/tiles; each reduction and FMA order stays fixed. +template +__global__ void __launch_bounds__(N) gdn_chunk_prep_wide( + const bf16* q, const bf16* k, const bf16* v, const float* g, const bf16* beta, + Work ws, int T, int H, int HK, float scale) { + gdn_chunk_prep_body(q, k, v, g, beta, ws, T, H, HK, scale, gridDim.x); +} +__global__ void __launch_bounds__(ST_THREADS) gdn_chunk_state( + Work ws, int NC, const float* __restrict__ si, float* __restrict__ so) { + gdn_chunk_state_body(ws, NC, si, so); +} +__global__ void __launch_bounds__(OUT_THREADS) gdn_chunk_out(Work ws, bf16* __restrict__ o, int T, int H) { + gdn_chunk_out_body(ws, o, T, H, gridDim.x); +} + +// Split only independent value columns; preserve each accumulator's MMA order. +// The legacy output body stays separate for existing batch and stateful entries. +template +__device__ __forceinline__ void gdn_chunk_out_tiled_body(Work ws, bf16* __restrict__ o, int T, int H, int NC) { + static_assert(N == 256 || N == 512, "disjoint value-column partitions"); + extern __shared__ __align__(128) unsigned char sm[]; + bf16* qs = reinterpret_cast(sm); // [C][O_LD] + bf16* ss = qs + C * O_LD; // [K][O_LD] + bf16* ps = ss + K * O_LD; // [C][P_LD] + bf16* vs = ps + C * P_LD; // [C][O_LD] + const int c = blockIdx.x, h = blockIdx.y, NCC = NC * C; + const int tid = threadIdx.x, warp = tid / 32, lane = tid % 32, g = lane / 4, t = lane % 4; + constexpr int PART = V / (N / 128); + const int row_warp = warp % 4, vb0 = (warp / 4) * PART; + const size_t rows = (size_t)h * NCC + (size_t)c * C; + for (int x = tid; x < C * K / 8; x += N) { + const int r = x / (K / 8), col = (x % (K / 8)) * 8; + cs1::cp_async16(qs + r * O_LD + col, ws.qd + (rows + r) * K + col); + cs1::cp_async16(vs + r * O_LD + col, ws.vn + (rows + r) * V + col); + } + for (int x = tid; x < K * V / 8; x += N) { + const int r = x / (V / 8), col = (x % (V / 8)) * 8; + cs1::cp_async16(ss + r * O_LD + col, ws.s + ((size_t)h * NC + c) * K * V + (size_t)r * V + col); + } + for (int x = tid; x < C * C / 8; x += N) { + const int r = x / (C / 8), col = (x % (C / 8)) * 8; + cs1::cp_async16(ps + r * P_LD + col, ws.pb + ((size_t)h * NC + c) * C * C + r * C + col); + } + cs1::cp_async_commit(); + cs1::cp_async_wait<0>(); + __syncthreads(); + + float acc[PART / 8][4]; +#pragma unroll + for (int nt = 0; nt < PART / 8; nt++) acc[nt][0] = acc[nt][1] = acc[nt][2] = acc[nt][3] = 0.f; + // qd S +#pragma unroll 2 + for (int kk = 0; kk < K; kk += 16) { + uint32_t a[4]; + cs1::load_a(a, qs, O_LD, row_warp * 16, kk, lane); +#pragma unroll + for (int nt = 0; nt < PART / 8; nt += 2) { + uint32_t b[4]; + cs1::load_b_kn(b, ss, O_LD, kk, vb0 + nt * 8, lane); + cs1::mma16816(acc[nt], a, b[0], b[1]); + cs1::mma16816(acc[nt + 1], a, b[2], b[3]); + } + } + // P v_new: each row warp needs only its lower-triangular prefix. + for (int kk = 0; kk < (row_warp + 1) * 16; kk += 16) { + uint32_t a[4]; + cs1::load_a(a, ps, P_LD, row_warp * 16, kk, lane); +#pragma unroll + for (int nt = 0; nt < PART / 8; nt += 2) { + uint32_t b[4]; + cs1::load_b_kn(b, vs, O_LD, kk, vb0 + nt * 8, lane); + cs1::mma16816(acc[nt], a, b[0], b[1]); + cs1::mma16816(acc[nt + 1], a, b[2], b[3]); + } + } +#pragma unroll + for (int r = 0; r < 2; r++) { + const int tok = c * C + row_warp * 16 + g + r * 8; + if (tok >= T) continue; + bf16* dst = o + ((size_t)tok * H + h) * V + vb0 + 2 * t; +#pragma unroll + for (int nt = 0; nt < PART / 8; nt++) + *reinterpret_cast(dst + nt * 8) = cs1::pack_bf16(acc[nt][2 * r], acc[nt][2 * r + 1]); + } +} + + +template +__global__ void __launch_bounds__(N) gdn_chunk_out_tiled(Work ws, bf16* o, int T, int H) { + gdn_chunk_out_tiled_body(ws, o, T, H, gridDim.x); +} + +constexpr int MAX_SEQUENCES = 16; +struct Sequence { Work ws; int tokens; int offset; }; +struct Batch { Sequence sequences[MAX_SEQUENCES]; }; +static_assert(sizeof(Batch) < 2048, "batched descriptors fit pre-Hopper kernel parameters"); + +__global__ void __launch_bounds__(THREADS) gdn_batch_prep( + const bf16* q, const bf16* k, const bf16* v, const float* g, const bf16* beta, + Batch batch, int H, int HK, float scale) { + const auto sequence = batch.sequences[blockIdx.z]; + if (blockIdx.x >= (sequence.tokens + C - 1) / C) return; + const size_t offset = sequence.offset; + gdn_chunk_prep_body(q + offset * HK * K, k + offset * HK * K, + v + offset * H * V, g + offset * H, beta + offset * H, + sequence.ws, sequence.tokens, H, HK, scale, (sequence.tokens + C - 1) / C); +} +__global__ void __launch_bounds__(ST_THREADS) gdn_batch_state(Batch batch) { + const auto sequence = batch.sequences[blockIdx.z]; + gdn_chunk_state_body(sequence.ws, (sequence.tokens + C - 1) / C, nullptr, nullptr); +} +__global__ void __launch_bounds__(OUT_THREADS) gdn_batch_out(Batch batch, bf16* o, int H) { + const auto sequence = batch.sequences[blockIdx.z]; + if (blockIdx.x >= (sequence.tokens + C - 1) / C) return; + gdn_chunk_out_body(sequence.ws, o + static_cast(sequence.offset) * H * V, + sequence.tokens, H, (sequence.tokens + C - 1) / C); +} + Work split(float* workspace, int T, int H) { const Layout l(T, H); unsigned char* b = reinterpret_cast(workspace); @@ -567,7 +696,7 @@ size_t cs1_gdn_workspace_floats(int T, int H) { return (Layout(T, H).total + 3) static int gdn_prefill_run(const void* q, const void* k, const void* v, const float* g, const void* beta, void* o, float* workspace, int T, int H, int HK, - float scale, const void* s_in, void* s_out, void* stream) { + float scale, const void* s_in, void* s_out, void* stream, bool wide = false) { if (T < 0 || HK <= 0 || H % HK != 0) return cudaErrorInvalidValue; const float* si = static_cast(s_in); float* so = static_cast(s_out); @@ -595,11 +724,39 @@ static int gdn_prefill_run(const void* q, const void* k, const void* v, const fl const int NC = (T + C - 1) / C; const Work ws = split(workspace, T, H); cudaStream_t st = static_cast(stream); - gdn_chunk_prep<<>>( - static_cast(q), static_cast(k), - static_cast(v), g, static_cast(beta), ws, T, H, HK, scale); + if (wide) { + static const cudaError_t wide_configured = [] { + auto e = cudaFuncSetAttribute(gdn_chunk_prep_wide<512>, + cudaFuncAttributeMaxDynamicSharedMemorySize, (int)SMEM1_BYTES); + if (e == cudaSuccess) e = cudaFuncSetAttribute(gdn_chunk_prep_wide<768>, + cudaFuncAttributeMaxDynamicSharedMemorySize, (int)SMEM1_BYTES); + if (e == cudaSuccess) e = cudaFuncSetAttribute(gdn_chunk_out_tiled<256>, + cudaFuncAttributeMaxDynamicSharedMemorySize, (int)SMEM3_BYTES); + if (e == cudaSuccess) e = cudaFuncSetAttribute(gdn_chunk_out_tiled<512>, + cudaFuncAttributeMaxDynamicSharedMemorySize, (int)SMEM3_BYTES); + return e; + }(); + if (wide_configured != cudaSuccess) return wide_configured; + if (T <= 128) + gdn_chunk_prep_wide<768><<>>( + static_cast(q), static_cast(k), static_cast(v), + g, static_cast(beta), ws, T, H, HK, scale); + else + gdn_chunk_prep_wide<512><<>>( + static_cast(q), static_cast(k), static_cast(v), + g, static_cast(beta), ws, T, H, HK, scale); + } else { + gdn_chunk_prep<<>>( + static_cast(q), static_cast(k), static_cast(v), + g, static_cast(beta), ws, T, H, HK, scale); + } gdn_chunk_state<<>>(ws, NC, si, so); - gdn_chunk_out<<>>(ws, static_cast(o), T, H); + if (wide && T > 128 && T <= 2048) + gdn_chunk_out_tiled<512><<>>(ws, static_cast(o), T, H); + else if (wide) + gdn_chunk_out_tiled<256><<>>(ws, static_cast(o), T, H); + else + gdn_chunk_out<<>>(ws, static_cast(o), T, H); return cudaGetLastError(); } @@ -608,10 +765,61 @@ int cs1_gdn_prefill(const void* q, const void* k, const void* v, const float* g, return gdn_prefill_run(q, k, v, g, beta, o, workspace, T, H, HK, scale, nullptr, nullptr, stream); } +// Opt-in 27B head layout. Keep the legacy entry and longer-input capacity intact. +int cs1_gdn_prefill_wide(const void* q, const void* k, const void* v, const float* g, const void* beta, + void* o, float* workspace, int T, int H, int HK, float scale, void* stream) { + const bool wide = H == 48 && HK == 16 && T > 0 && T <= 16384; + return gdn_prefill_run(q, k, v, g, beta, o, workspace, T, H, HK, scale, nullptr, nullptr, stream, wide); +} + int cs1_gdn_prefill_x(const void* q, const void* k, const void* v, const float* g, const void* beta, void* o, float* workspace, int T, int H, int HK, float scale, const void* s_in, void* s_out, void* stream) { return gdn_prefill_run(q, k, v, g, beta, o, workspace, T, H, HK, scale, s_in, s_out, stream); } +// Bound the sum of per-sequence layouts: up to (count-1)*63 extra padded rows, +// plus one 256-byte alignment gap per layout part per sequence. +size_t cs1_gdn_batch_workspace_floats(int tokens, int H, int count) { + if (tokens <= 0 || H <= 0 || count <= 0 || count > MAX_SEQUENCES || + tokens > std::numeric_limits::max() - count * C) return 0; + return (Layout(tokens + (count - 1) * (C - 1), H).total + count * 9 * 256 + 3) / 4; +} + +int cs1_gdn_prefill_batch(const void* q, const void* k, const void* v, const float* g, + const void* beta, void* o, float* workspace, size_t workspace_bytes, + const size_t* lengths, size_t count, int H, int HK, float scale, void* stream) { + if (!q || !k || !v || !g || !beta || !o || !workspace || !lengths || !count || + count > MAX_SEQUENCES || H <= 0 || HK <= 0 || H % HK != 0) return cudaErrorInvalidValue; + Batch batch{}; + size_t rows = 0, bytes = 0; + int chunks = 0; + for (size_t i = 0; i < count; ++i) { + if (!lengths[i] || lengths[i] > static_cast(std::numeric_limits::max() - C) - rows) + return cudaErrorInvalidValue; + const int length = static_cast(lengths[i]); + const size_t needed = Layout(length, H).total; + if (bytes > workspace_bytes || needed > workspace_bytes - bytes) return cudaErrorInvalidValue; + batch.sequences[i] = {split(reinterpret_cast(reinterpret_cast(workspace) + bytes), length, H), + length, static_cast(rows)}; + rows += lengths[i]; + bytes += needed; + chunks = std::max(chunks, (length + C - 1) / C); + } + static const cudaError_t configured = [] { + cudaError_t e = cudaFuncSetAttribute(gdn_batch_prep, cudaFuncAttributeMaxDynamicSharedMemorySize, (int)SMEM1_BYTES); + if (e == cudaSuccess) e = cudaFuncSetAttribute(gdn_batch_state, cudaFuncAttributeMaxDynamicSharedMemorySize, (int)SMEM2_BYTES); + if (e == cudaSuccess) e = cudaFuncSetAttribute(gdn_batch_out, cudaFuncAttributeMaxDynamicSharedMemorySize, (int)SMEM3_BYTES); + return e; + }(); + if (configured != cudaSuccess) return configured; + const auto st = static_cast(stream); + gdn_batch_prep<<>>( + static_cast(q), static_cast(k), static_cast(v), g, + static_cast(beta), batch, H, HK, scale); + gdn_batch_state<<>>(batch); + gdn_batch_out<<>>(batch, static_cast(o), H); + return cudaGetLastError(); +} + } // extern "C" diff --git a/src/backends/cuda/qwen3_5/gemm.cu b/src/backends/cuda/qwen3_5/gemm.cu index d761f862..0f55faee 100644 --- a/src/backends/cuda/qwen3_5/gemm.cu +++ b/src/backends/cuda/qwen3_5/gemm.cu @@ -17,7 +17,12 @@ #include #include +#include +#include #include +#include +#include +#include #include "ops.h" @@ -30,6 +35,28 @@ struct Plan { }; using Key = std::tuple; // M, N, K, ldy, FP32, bias +using PackedKey = std::tuple, int, int>; + +struct PackedSegment { + Plan plan; + size_t offset = 0; + int rows = 0; + bool enabled = false; +}; + +using PackedSegments = std::vector; +constexpr size_t MAX_PACKED_PLANS = 512; +struct PackedPlan { + PackedSegments segments; + std::list::iterator age; +}; + +using ReuseKey = std::tuple; +constexpr size_t MAX_REUSED_PLANS = 512; +struct ReusedPlan { + Plan plan; + std::list::iterator age; +}; struct Gemm { cublasLtHandle_t handle = nullptr; @@ -39,6 +66,10 @@ struct Gemm { 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 + std::map reused_plans; + std::list reused_order; + std::map packed_plans; + std::list packed_order; }; int status(cublasStatus_t s) { return s == CUBLAS_STATUS_SUCCESS ? 0 : 1000 + (int)s; } @@ -150,6 +181,42 @@ int plan_for(Gemm& g, int M, int N, int K, int ldy, Plan*& out, bool fp32 = fals return 0; } +// Optional shape changes retain the exact original cuBLASLt algorithm. A return +// of -1 means unsupported (the caller must use its unchanged complete forward). +int reused_plan(Gemm& g, int M, int N, int K, int original_M, Plan*& out) { + const auto key = std::make_tuple(M, N, K, original_M); + auto it = g.reused_plans.find(key); + if (it != g.reused_plans.end()) { + if (std::next(it->second.age) != g.reused_order.end()) + g.reused_order.splice(g.reused_order.end(), g.reused_order, it->second.age); + out = &it->second.plan; + return 0; + } + Plan* original = nullptr; + int rc = plan_for(g, original_M, N, K, N, original); + if (rc) return rc; + Plan p; + rc = describe(M, N, K, N, p, false, false); + if (rc) { destroy(p); return rc; } + p.algo = original->algo; + cublasLtMatmulHeuristicResult_t checked{}; + const auto result = cublasLtMatmulAlgoCheck(g.handle, p.op, p.a, p.b, p.c, p.c, &p.algo, &checked); + if (result != CUBLAS_STATUS_SUCCESS || checked.state != CUBLAS_STATUS_SUCCESS || + checked.workspaceSize > g.workspace_bytes) { + destroy(p); + return -1; + } + if (g.reused_plans.size() == MAX_REUSED_PLANS) { + const auto oldest = g.reused_plans.find(g.reused_order.front()); + destroy(oldest->second.plan); + g.reused_plans.erase(oldest); + g.reused_order.pop_front(); + } + g.reused_order.push_back(key); + out = &g.reused_plans.emplace(key, ReusedPlan{p, std::prev(g.reused_order.end())}).first->second.plan; + return 0; +} + } // namespace extern "C" void* cs1_gemm_create_fixed(size_t workspace_bytes, int reference_m) { @@ -175,6 +242,9 @@ extern "C" void cs1_gemm_destroy(void* gemm) { Gemm* g = static_cast(gemm); if (!g) return; for (auto& kv : g->plans) destroy(kv.second); + for (auto& kv : g->reused_plans) destroy(kv.second.plan); + for (auto& kv : g->packed_plans) + for (auto& segment : kv.second.segments) destroy(segment.plan); if (g->workspace) cudaFree(g->workspace); cublasLtDestroy(g->handle); delete g; @@ -193,6 +263,141 @@ extern "C" int cs1_gemm(void* gemm, const void* x, const void* w, void* y, int M g->workspace, g->workspace_bytes, static_cast(stream))); } +// Additive opt-in for merged language models. Keep each original shape's +// algorithm, including split-K and its reduction scheme. Pack only consecutive +// shapes with identical algorithms, and split groups to fit the existing workspace. +extern "C" int cs1_gemm_packed(void* gemm, const void* x, const void* w, void* y, + const size_t* lengths, size_t count, int N, int K, + void* stream) { + auto* g = static_cast(gemm); + if (!g || !lengths || count == 0 || N <= 0 || K <= 0) return cudaErrorInvalidValue; + size_t total = 0; + for (size_t i = 0; i < count; ++i) { + if (lengths[i] == 0 || lengths[i] > static_cast(std::numeric_limits::max()) - total) + return cudaErrorInvalidValue; + total += lengths[i]; + } + if (count == 1) return cs1_gemm(g, x, w, y, static_cast(total), N, K, N, stream); + const PackedKey key{std::vector(lengths, lengths + count), N, K}; + auto it = g->packed_plans.find(key); + if (it == g->packed_plans.end()) { + std::vector originals(count); + for (size_t i = 0; i < count; ++i) { + const int rc = plan_for(*g, static_cast(lengths[i]), N, K, N, originals[i]); + if (rc) return rc; + } + PackedSegments packed; + size_t offset = 0; + for (size_t begin = 0; begin < count;) { + size_t end = begin + 1; + size_t rows = lengths[begin]; + // Equality of the opaque value is conservative: differences in + // reserved bytes only decline packing. Original map pointers are stable. + while (!g->reference && end < count && + std::memcmp(&originals[begin]->algo, &originals[end]->algo, + sizeof(cublasLtMatmulAlgo_t)) == 0) { + rows += lengths[end++]; + } + PackedSegment segment; + segment.offset = offset; + // Try the largest compatible group first. A smaller group may fit + // when split-K workspace or an algorithm shape limit rejects it. + while (end > begin + 1) { + const int rc = describe(static_cast(rows), N, K, N, segment.plan, false, false); + if (rc) { + destroy(segment.plan); + for (auto& previous : packed) destroy(previous.plan); + return rc; + } + auto& p = segment.plan; + p.algo = originals[begin]->algo; + cublasLtMatmulHeuristicResult_t checked{}; + const auto result = cublasLtMatmulAlgoCheck(g->handle, p.op, p.a, p.b, p.c, p.c, + &p.algo, &checked); + segment.enabled = result == CUBLAS_STATUS_SUCCESS && checked.state == CUBLAS_STATUS_SUCCESS + && checked.workspaceSize <= g->workspace_bytes; + if (segment.enabled) break; + destroy(segment.plan); + rows -= lengths[--end]; + } + segment.rows = static_cast(rows); + packed.push_back(segment); + offset += rows; + begin = end; + } + if (g->packed_plans.size() == MAX_PACKED_PLANS) { + const auto oldest = g->packed_plans.find(g->packed_order.front()); + // Descriptors are host-side cuBLAS inputs; captured graphs own their + // launch arguments. The eviction/replay regression covers this lifetime. + for (auto& segment : oldest->second.segments) destroy(segment.plan); + g->packed_plans.erase(oldest); + g->packed_order.pop_front(); + } + g->packed_order.push_back(key); + it = g->packed_plans.emplace(key, PackedPlan{std::move(packed), std::prev(g->packed_order.end())}).first; + } + if (std::next(it->second.age) != g->packed_order.end()) + g->packed_order.splice(g->packed_order.end(), g->packed_order, it->second.age); + const float alpha = 1.f, beta = 0.f; + for (auto& segment : it->second.segments) { + const auto* input = static_cast(x) + segment.offset * K * 2; + auto* output = static_cast(y) + segment.offset * N * 2; + int rc; + if (segment.enabled) { + auto& p = segment.plan; + rc = status(cublasLtMatmul(g->handle, p.op, &alpha, w, p.a, input, p.b, &beta, output, p.c, + output, p.c, &p.algo, g->workspace, g->workspace_bytes, + static_cast(stream))); + } else { + rc = cs1_gemm(g, input, w, output, segment.rows, N, K, N, stream); + } + if (rc) return rc; + } + return 0; +} + +extern "C" int cs1_gemm_reuse_check(void* gemm, const size_t* original_rows, size_t originals, + const size_t* target_rows, size_t targets, int N, int K) { + auto* g = static_cast(gemm); + if (!g || !original_rows || !target_rows || originals == 0 || targets == 0 || N <= 0 || K <= 0) + return cudaErrorInvalidValue; + if (g->reference) return -1; + for (size_t i = 0; i < originals; ++i) + if (!original_rows[i] || original_rows[i] > static_cast(std::numeric_limits::max())) + return cudaErrorInvalidValue; + for (size_t i = 0; i < targets; ++i) + if (!target_rows[i] || target_rows[i] > static_cast(std::numeric_limits::max())) + return cudaErrorInvalidValue; + Plan* original = nullptr; + int rc = plan_for(*g, static_cast(original_rows[0]), N, K, N, original); + if (rc) return rc; + for (size_t i = 1; i < originals; ++i) { + Plan* other = nullptr; + rc = plan_for(*g, static_cast(original_rows[i]), N, K, N, other); + if (rc) return rc; + if (std::memcmp(&original->algo, &other->algo, sizeof(original->algo))) return -1; + } + for (size_t i = 0; i < targets; ++i) { + Plan* target = nullptr; + rc = reused_plan(*g, static_cast(target_rows[i]), N, K, static_cast(original_rows[0]), target); + if (rc) return rc; + } + return 0; +} + +extern "C" int cs1_gemm_reuse(void* gemm, const void* x, const void* w, void* y, + int M, int N, int K, int original_M, void* stream) { + auto* g = static_cast(gemm); + if (!g || M <= 0 || N <= 0 || K <= 0 || original_M <= 0) return cudaErrorInvalidValue; + if (g->reference) return -1; + Plan* p = nullptr; + const int rc = reused_plan(*g, M, N, K, original_M, p); + if (rc) return rc; + const float alpha = 1.f, beta = 0.f; + return status(cublasLtMatmul(g->handle, p->op, &alpha, w, p->a, x, p->b, &beta, y, p->c, y, p->c, + &p->algo, g->workspace, g->workspace_bytes, static_cast(stream))); +} + // Vision's biased BF16 linears round only after adding bias. The patch // convolution passes zero bias here and applies its bias after BF16 rounding. extern "C" int cs1_vision_linear(void* gemm, const void* x, const void* w, const void* bias, diff --git a/src/backends/cuda/qwen3_5/ops.h b/src/backends/cuda/qwen3_5/ops.h index 73224df5..77b8bd1e 100644 --- a/src/backends/cuda/qwen3_5/ops.h +++ b/src/backends/cuda/qwen3_5/ops.h @@ -66,6 +66,19 @@ int cs1_gated_rms_norm(const void* x, const void* z, int ldz, const void* w, voi // q [T, key_dim], k [T, key_dim] and v [T, value_dim]. w is [key_dim*2 + value_dim, 4]. int cs1_gdn_conv(const void* qkv, int ld, const void* w, void* q, void* k, void* v, int T, int key_dim, int value_dim, void* stream); +// Optional vectorized convolution with exact scalar-layout fallback. +int cs1_gdn_conv_vector(const void* qkv, int ld, const void* w, void* q, void* k, void* v, int T, int key_dim, + int value_dim, void* stream); + +// Optional compact suffix convolution seeded by three immutable pre-conv rows [3,ld]. +int cs1_gdn_conv_prefix(const void* qkv, const void* history, int ld, const void* w, + void* q, void* k, void* v, int T, int key_dim, int value_dim, void* stream); +// Optional compact suffix attention. Prefix K/V are contiguous [prefix,Hk,Dh]; +// suffix Q/gate/output start at row zero and use absolute query position prefix. +int cs1_attention_gated_shared(const void* q, const void* k, const void* v, int ldv, + const void* gate, void* out, const void* prefix_k, + const void* prefix_v, int suffix, int prefix, int Hq, + int Hk, int Dh, float scale, void* stream); // beta = sigmoid(b) (bfloat16) and g = -exp(A_log) * softplus(a + dt_bias) (float32), [T, H]; // b and a are [T, H] in rows of ld. @@ -74,6 +87,11 @@ int cs1_gdn_gates(const void* b, const void* a, int ld, const void* A_log, const // Chunked gated delta rule, q and k L2-normalized inside, q scaled by `scale`. // q, k [T, HK, 128], v [T, H, 128], g float [T, H], beta [T, H], o [T, H, 128]. +// Optional independent sequence batching, with separate chunk padding/workspaces. +size_t cs1_gdn_batch_workspace_floats(int tokens, int H, int count); +int cs1_gdn_prefill_batch(const void* q, const void* k, const void* v, const float* g, + const void* beta, void* o, float* workspace, size_t workspace_bytes, + const size_t* lengths, size_t count, int H, int HK, float scale, void* stream); size_t cs1_gdn_workspace_floats(int T, int H); // Optional no-FLA text precision mode: FP32 normalized Q/K and recurrent state, // BF16 output. Context and scratch remain model-owned; prefix continuation is unsupported. @@ -86,6 +104,11 @@ int cs1_gdn_prefill_fp32(void* context,const void* q, const void* k, const void* int cs1_gdn_prefill(const void* q, const void* k, const void* v, const float* g, const void* beta, void* o, float* workspace, int T, int H, int HK, float scale, void* stream); +// Optional wider preparation/output launch for the opt-in 27B text path. Same layout, +// workspace, arithmetic and output as cs1_gdn_prefill; other shapes fall back. +int cs1_gdn_prefill_wide(const void* q, const void* k, const void* v, const float* g, const void* beta, + void* o, float* workspace, int T, int H, int HK, float scale, void* stream); + // Same scan with an explicit state: s_in is the [H, K, V] float32 state before // this window (null = zero), s_out receives the [H, K, V] float32 state after // it (null = skip). The float32 state makes a continuation starting exactly at @@ -146,6 +169,17 @@ 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); +// Optional packed projections retaining original per-sequence algorithms. +int cs1_gemm_packed(void* gemm, const void* x, const void* w, void* y, + const size_t* lengths, size_t count, int N, int K, void* stream); + +// Optional retained-algorithm projections. Check returns 0 if compatible, -1 if +// unsupported; other values are errors. Check all windows before computing any. +int cs1_gemm_reuse_check(void* gemm, const size_t* original_rows, size_t originals, + const size_t* target_rows, size_t targets, int N, int K); +int cs1_gemm_reuse(void* gemm, const void* x, const void* w, void* y, + int M, int N, int K, int original_M, void* stream); + // ---- vision: single image, 1024 hidden, 16 heads of 64; all BF16 except explicit float pointers ---- int cs1_vision_linear(void* gemm, const void* x, const void* w, const void* bias, void* y, int M, int N, int K, void* stream); int cs1_gemm_f32(void* gemm, const float* x, const float* w, float* y, int M, int N, int K, void* stream); diff --git a/src/models/open_jev/README.md b/src/models/open_jev/README.md index 48fdbe56..2d690c39 100644 --- a/src/models/open_jev/README.md +++ b/src/models/open_jev/README.md @@ -25,17 +25,42 @@ The native worker follows the [architecture contracts](../../../docs/architectur with independent [processing](native/src/processing.rs) and [executor](native/src/executor.rs) modules. Preparation returns token IDs grouped by question and candidate, plus the response context for identity/order, -calibration, usage, and metadata. The executor owns the trained head and returns -one FP32 scalar per candidate. Response finishing adds the `noul` false baseline +calibration, usage, and metadata. On 27B, preparation tokenizes up to 16 candidates +from one question in parallel when the tokenizer has no padding; singleton and +padded inputs retain per-candidate encoding. Token IDs, candidate order and the +first validation error are preserved. `RAYON_NUM_THREADS` controls tokenizer +parallelism. Batched encoding skips unused character offsets. The executor owns the trained head and returns one FP32 scalar per +candidate plus whether prefix sharing actually ran. Response finishing adds the `noul` false baseline and applies calibrated normalization across each complete question. The worker retains request-wide model locking and independent sequence semantics, with the scalar head on the CPU after CUDA prefill. For Open-Jev-27B-v1.1, its [batch adapter](native/src/batching.rs) packs at most 16 candidates and 4096 tokens per group, in prepared order; longer individual prompts execute alone without truncation. Input and gate/up -projections share packed GEMMs. Output/down projections retain their original -per-prompt GEMM shapes; attention, positions, convolution and GDN state reset -at each sequence boundary. Results are regrouped before question normalization. +projections share packed GEMMs. Output/down projections can combine consecutive +candidates when their original cuBLASLt algorithms match and the combined shape +fits the existing workspace; otherwise they retain separate calls. Batched GDN +launches process at most 16 sequences with separate state and workspace. Attention, +positions and convolution also reset at each sequence boundary. + +A single question with 2–16 candidates, each longer than 2048 tokens, can share +an exact token prefix of at least 512 tokens, rounded down to a multiple of 64. +Every candidate keeps a nonempty suffix. The executor groups candidates only +when all projection algorithms from their complete forwards remain compatible +with the prefix and suffix shapes. Each group captures its own prefix state; +singletons use complete forwards, and all results return to their original indices +before question-wide normalization. The model retains one bounded prefix-buffer +capacity, invalidates its numeric contents after each request, and captures new +contents for every group. This is within-question reuse, with no reuse of request +contents across requests. Response `metadata.prefix_cache` separates configured +`policy` (currently `default`) from actual `enabled`/`scope` and `saved_tokens`. +Savings count avoided prefix computations after compatibility grouping, with +zero for singleton fallbacks; suffix packing adds no prefix savings. These four +fields are the complete contract; `mode` and `shared_tokens` are not emitted. +Input usage still counts complete prompts. See the +[native recipe](../../../recipe/open_jev/native.md#validation-and-optimization-scope) +for field semantics and the unchanged default-path numerical validation gates. +Results are regrouped before question normalization. Open-Jev-9B runs candidates one at a time until packing is validated for it. The engine owns a [shared serial scheduler](../../runtime/README.md) that admits the complete request diff --git a/src/models/open_jev/native/src/batching.rs b/src/models/open_jev/native/src/batching.rs index 13f95c40..58317242 100644 --- a/src/models/open_jev/native/src/batching.rs +++ b/src/models/open_jev/native/src/batching.rs @@ -24,6 +24,22 @@ pub(crate) fn ranges(inputs: &[&[u32]]) -> Vec> { batches } +/// Only long candidates that the original packing policy executes separately. +/// Sharing is restricted to one question by the caller, with a nonempty suffix. +pub(crate) fn shared_prefix(inputs: &[&[u32]]) -> usize { + if !(2..=MAX_SEQUENCES).contains(&inputs.len()) + || inputs.iter().any(|ids| ids.len() <= MAX_TOKENS / 2) + { + return 0; + } + let shortest = inputs.iter().map(|ids| ids.len()).min().unwrap(); + let common = (0..shortest - 1) + .take_while(|&i| inputs.iter().all(|ids| ids[i] == inputs[0][i])) + .count(); + let aligned = common / 64 * 64; + if aligned >= 512 { aligned } else { 0 } +} + #[cfg(test)] #[path = "../../../../../tests/open_jev/batching.rs"] mod tests; diff --git a/src/models/open_jev/native/src/contract.rs b/src/models/open_jev/native/src/contract.rs index 009aa75a..17926ea9 100644 --- a/src/models/open_jev/native/src/contract.rs +++ b/src/models/open_jev/native/src/contract.rs @@ -19,6 +19,8 @@ pub struct Checkpoint { /// Pack a request's candidates for shared GEMMs. Packing was validated against /// per-candidate execution on 27B only, so 9B runs candidates one at a time. pub pack_candidates: bool, + /// Parallel preparation has exact-token coverage for the pinned 27B tokenizer. + pub parallel_tokenization: bool, } pub static CHECKPOINTS: [Checkpoint; 2] = [ @@ -30,6 +32,7 @@ pub static CHECKPOINTS: [Checkpoint; 2] = [ alias: "open-jev-27b-v1.1", backbone: (5120, 17408, 64, 24, 4, 16, 48), pack_candidates: true, + parallel_tokenization: true, }, Checkpoint { name: "Open-Jev-9B", @@ -39,6 +42,7 @@ pub static CHECKPOINTS: [Checkpoint; 2] = [ alias: "open-jev-9b", backbone: (4096, 12288, 32, 16, 4, 16, 32), pack_candidates: false, + parallel_tokenization: false, }, ]; diff --git a/src/models/open_jev/native/src/executor.rs b/src/models/open_jev/native/src/executor.rs index b06f085f..c93f16bb 100644 --- a/src/models/open_jev/native/src/executor.rs +++ b/src/models/open_jev/native/src/executor.rs @@ -68,6 +68,12 @@ impl DecisionHead { } } +/// Model results and execution facts for response metadata. +pub struct ExecutedRequest { + pub scores: Vec>, + pub saved_prefix_tokens: usize, +} + pub struct Executor { model: Arc>, head: DecisionHead, @@ -82,7 +88,16 @@ impl Executor { checkpoint: &Checkpoint, ) -> Result { let (d, lib) = (dir.to_path_buf(), library.to_path_buf()); - let model = tokio::task::spawn_blocking(move || Model::load(&d, &lib)).await??; + let batch_gdn = checkpoint.pack_candidates; + let model = tokio::task::spawn_blocking(move || { + let mut model = Model::load(&d, &lib)?; + if batch_gdn { + model.enable_batched_gdn()?; + model.enable_packed_projections()?; + } + Ok::<_, anyhow::Error>(model) + }) + .await??; Ok(Self { model: Arc::new(Mutex::new(model)), head, @@ -98,6 +113,15 @@ impl Executor { scheduler: &SerialScheduler, ids: Vec>>, ) -> Result>> { + Ok(self.execute_detailed(scheduler, ids).await?.scores) + } + + /// Include whether a prefix was actually reused after device compatibility checks. + pub async fn execute_detailed( + &self, + scheduler: &SerialScheduler, + ids: Vec>>, + ) -> Result { let model = self.model.clone(); let head = self.head.clone(); let pack = self.pack; @@ -108,7 +132,22 @@ impl Executor { .map_err(|_| anyhow::anyhow!("poisoned model"))?; let inputs: Vec<&[u32]> = ids.iter().flatten().map(Vec::as_slice).collect(); let mut scores = Vec::with_capacity(inputs.len()); - if pack { + let prefix_len = if pack && ids.len() == 1 { + batching::shared_prefix(&inputs) + } else { + 0 + }; + let shared = if prefix_len > 0 { + model.try_forward_shared_prefix(&inputs, prefix_len)? + } else { + None + }; + let saved_prefix_tokens = shared.as_ref().map_or(0, |output| output.saved_tokens); + if let Some(output) = shared { + for last in output.hidden { + scores.push(head.score(last)?); + } + } else if pack { for range in batching::ranges(&inputs) { for last in model.forward_batch(&inputs[range])? { scores.push(head.score(last)?); @@ -120,10 +159,13 @@ impl Executor { } } let mut scores = scores.into_iter(); - Ok(ids - .iter() - .map(|candidates| scores.by_ref().take(candidates.len()).collect()) - .collect()) + Ok(ExecutedRequest { + scores: ids + .iter() + .map(|candidates| scores.by_ref().take(candidates.len()).collect()) + .collect(), + saved_prefix_tokens, + }) }) .await } diff --git a/src/models/open_jev/native/src/main.rs b/src/models/open_jev/native/src/main.rs index 5f712f16..ed1108a0 100644 --- a/src/models/open_jev/native/src/main.rs +++ b/src/models/open_jev/native/src/main.rs @@ -27,11 +27,13 @@ async fn decide(engine: &Engine, raw: &[u8]) -> Response { Err(e) => return error(StatusCode::UNPROCESSABLE_ENTITY, e), }; let result = async { - let rows = engine + let executed = engine .executor - .execute(&engine.scheduler, prepared.inputs) + .execute_detailed(&engine.scheduler, prepared.inputs) .await?; - prepared.context.finish(rows) + prepared + .context + .finish_with_prefix(executed.scores, executed.saved_prefix_tokens) } .await; match result { diff --git a/src/models/open_jev/native/src/processing.rs b/src/models/open_jev/native/src/processing.rs index 8d77f701..f1d0bf1d 100644 --- a/src/models/open_jev/native/src/processing.rs +++ b/src/models/open_jev/native/src/processing.rs @@ -65,6 +65,7 @@ impl Processor { &self.suffix, self.max_length, &questions, + self.checkpoint.parallel_tokenization, )?; let input_tokens = inputs.iter().flatten().map(Vec::len).sum(); let candidates = inputs.iter().map(Vec::len).sum(); @@ -86,6 +87,15 @@ impl Processor { impl ResponseContext { /// Restore the Noul baseline and normalize across each complete question. pub fn finish(self, rows: Vec>) -> Result { + self.finish_with_prefix(rows, 0) + } + + /// Report executed prefix reuse, rather than inferring it from token lengths. + pub fn finish_with_prefix( + self, + rows: Vec>, + saved_prefix_tokens: usize, + ) -> Result { ensure!(rows.len() == self.questions.len(), "invalid logit rows"); let mut answers = Map::new(); for (q, mut logits) in self.questions.iter().zip(rows) { @@ -98,13 +108,19 @@ impl ResponseContext { contract::answer(q, &logits, self.temperature)?, ); } + let prefix_cache = if saved_prefix_tokens > 0 { + json!({"policy": "default", "enabled": true, "scope": "within_question", + "saved_tokens": saved_prefix_tokens}) + } else { + json!({"policy": "default", "enabled": false, "scope": "none", "saved_tokens": 0}) + }; Ok( json!({"model": self.checkpoint.model_id, "answers": answers, "usage": {"input_tokens": self.input_tokens, "output_tokens": 0}, "metadata": {"method": "native_merged_lora_decision_head", "temperature": self.temperature, "candidate_sequences": self.candidates, "inference_seconds": self.start.elapsed().as_secs_f64(), "base_revision": self.checkpoint.base_revision, "max_length": self.max_length, - "prefix_cache": {"enabled": false, "mode": "independent_candidates"}}}), + "prefix_cache": prefix_cache}}), ) } } @@ -115,15 +131,25 @@ fn encode_questions( suffix: &str, max_length: usize, questions: &[Question], + parallel: bool, ) -> Result>>> { questions .iter() .map(|q| { - q.prompts - .iter() - .map(|prompt| { - let chat = format!("{prefix}{}{suffix}", prompt.trim()); - let enc = tokenizer.encode(chat, true).map_err(anyhow::Error::msg)?; + // BatchLongest padding would change each candidate's IDs and length. + // Keep the existing per-candidate semantics for all padded tokenizers. + let chunk_size = if parallel && tokenizer.get_padding().is_none() { + 16 + } else { + 1 + }; + let mut inputs = Vec::with_capacity(q.prompts.len()); + for chunk in q.prompts.chunks(chunk_size) { + let chats: Vec<_> = chunk + .iter() + .map(|prompt| format!("{prefix}{}{suffix}", prompt.trim())) + .collect(); + let validate = |enc: tokenizers::Encoding| -> Result> { ensure!( enc.len() <= max_length, "question {}: {} tokens exceeds max_length={max_length}; no silent truncation", @@ -131,8 +157,26 @@ fn encode_questions( enc.len() ); Ok(enc.get_ids().to_vec()) - }) - .collect() + }; + if chats.len() > 1 + && let Ok(encodings) = tokenizer.encode_batch_fast(chats.iter().map(String::as_str).collect(), true) + { + // Only token IDs are consumed; skip unused character offsets. + // Batch encoding preserves input order; validate in that order too. + for enc in encodings { + inputs.push(validate(enc)?); + } + } else { + // A later tokenizer failure must not hide an earlier length error. + // Retrying serially also preserves the original first-error message. + for chat in chats { + inputs.push(validate( + tokenizer.encode(chat, true).map_err(anyhow::Error::msg)?, + )?); + } + } + } + Ok(inputs) }) .collect() } diff --git a/src/models/qwen3_5/native/src/cuda.rs b/src/models/qwen3_5/native/src/cuda.rs index feac0e8f..74cd16e1 100644 --- a/src/models/qwen3_5/native/src/cuda.rs +++ b/src/models/qwen3_5/native/src/cuda.rs @@ -24,6 +24,131 @@ pub struct Stream(*mut c_void); // from one thread at a time. unsafe impl Send for Stream {} +/// Optional GDN prefill with more threads and unchanged arithmetic. +pub type GdnPrefill = unsafe extern "C" fn( + *const c_void, + *const c_void, + *const c_void, + *const f32, + *const c_void, + *mut c_void, + *mut f32, + c_int, + c_int, + c_int, + f32, + Stream, +) -> c_int; + +/// Optional vectorized convolution with the original layout and rounding. +pub type GdnConv = unsafe extern "C" fn( + *const c_void, + c_int, + *const c_void, + *mut c_void, + *mut c_void, + *mut c_void, + c_int, + c_int, + c_int, + Stream, +) -> c_int; + +/// Optional compact suffix operators reading immutable per-layer prefix state. +#[derive(Clone, Copy)] +pub struct SharedSuffix { + pub conv: unsafe extern "C" fn( + *const c_void, + *const c_void, + c_int, + *const c_void, + *mut c_void, + *mut c_void, + *mut c_void, + c_int, + c_int, + c_int, + Stream, + ) -> c_int, + pub attention: unsafe extern "C" fn( + *const c_void, + *const c_void, + *const c_void, + c_int, + *const c_void, + *mut c_void, + *const c_void, + *const c_void, + c_int, + c_int, + c_int, + c_int, + c_int, + f32, + Stream, + ) -> c_int, +} + +/// Optional shape changes using a validated original-shape GEMM algorithm. +#[derive(Clone, Copy)] +pub struct ReuseGemm { + pub check: unsafe extern "C" fn( + *mut c_void, + *const usize, + usize, + *const usize, + usize, + c_int, + c_int, + ) -> c_int, + pub run: unsafe extern "C" fn( + *mut c_void, + *const c_void, + *const c_void, + *mut c_void, + c_int, + c_int, + c_int, + c_int, + Stream, + ) -> c_int, +} + +/// Optional packed BF16 projections retaining the original per-sequence algorithms. +pub type PackedGemm = unsafe extern "C" fn( + *mut c_void, + *const c_void, + *const c_void, + *mut c_void, + *const usize, + usize, + c_int, + c_int, + Stream, +) -> c_int; + +/// Optional GDN prefill over up to sixteen independent token sequences. +#[derive(Clone, Copy)] +pub struct BatchedGdn { + pub workspace_floats: unsafe extern "C" fn(c_int, c_int, c_int) -> usize, + pub run: unsafe extern "C" fn( + *const c_void, + *const c_void, + *const c_void, + *const f32, + *const c_void, + *mut c_void, + *mut f32, + usize, + *const usize, + usize, + c_int, + c_int, + f32, + Stream, + ) -> c_int, +} + /// Additive dimension-aware vision entry points. Libraries at the required ABI may omit these optional symbols. pub struct VisionV2 { pub position: unsafe extern "C" fn( @@ -309,6 +434,67 @@ impl Api { } } + pub fn wide_gdn(&self) -> Option { + // SAFETY: the additive signature matches ops.h; API retains the library. + unsafe { + self._lib + .get::(b"cs1_gdn_prefill_wide\0") + .ok() + .map(|f| *f) + } + } + + pub fn vector_conv(&self) -> Option { + // SAFETY: the additive signature matches ops.h; API retains the library. + unsafe { + self._lib + .get::(b"cs1_gdn_conv_vector\0") + .ok() + .map(|f| *f) + } + } + + pub fn shared_suffix(&self) -> Option { + // SAFETY: additive signatures match ops.h; the API retains the library. + unsafe { + Some(SharedSuffix { + conv: *self._lib.get(b"cs1_gdn_conv_prefix\0").ok()?, + attention: *self._lib.get(b"cs1_attention_gated_shared\0").ok()?, + }) + } + } + + pub fn reuse_gemm(&self) -> Option { + // SAFETY: signatures match optional ops.h declarations; API owns the library. + unsafe { + Some(ReuseGemm { + check: *self._lib.get(b"cs1_gemm_reuse_check\0").ok()?, + run: *self._lib.get(b"cs1_gemm_reuse\0").ok()?, + }) + } + } + + pub fn packed_gemm(&self) -> Result { + // SAFETY: the additive signature matches cs1_gemm_packed in ops.h; + // this API retains the library for the process lifetime. + unsafe { + self._lib + .get::(b"cs1_gemm_packed\0") + .map(|f| *f) + } + .context("packed projections require the optional cs1_gemm_packed CUDA symbol") + } + + pub fn batched_gdn(&self) -> Result { + // SAFETY: optional signatures match ops.h; the static API owns the library. + unsafe { + Ok(BatchedGdn { + workspace_floats: *self._lib.get(b"cs1_gdn_batch_workspace_floats\0")?, + run: *self._lib.get(b"cs1_gdn_prefill_batch\0")?, + }) + } + } + pub fn fp32_gdn(&self) -> Result { // SAFETY: these optional signatures match ops.h; the library lives in static API. let (create, destroy, run, size) = unsafe { diff --git a/src/models/qwen3_5/native/src/model.rs b/src/models/qwen3_5/native/src/model.rs index b47f4489..01cefd5d 100644 --- a/src/models/qwen3_5/native/src/model.rs +++ b/src/models/qwen3_5/native/src/model.rs @@ -32,6 +32,42 @@ const GEMM_WORKSPACE: usize = 32 << 20; /// cs1_gemm_create_fixed); its branches run tens of tokens. const FIXED_GEMM_M: i32 = 64; +/// Partition by backend compatibility while retaining original output indices. +/// The all-compatible path needs one check; singleton groups use complete forwards. +fn compatible_prefix_groups( + count: usize, + mut compatible: impl FnMut(&[usize]) -> Result, +) -> Result>>> { + let all: Vec = (0..count).collect(); + if compatible(&all)? { + return Ok(Some(vec![all])); + } + let mut groups: Vec> = Vec::new(); + for index in 0..count { + let mut placed = false; + for group in &mut groups { + group.push(index); + if compatible(group)? { + placed = true; + break; + } + group.pop(); + } + if !placed { + groups.push(vec![index]); + } + } + Ok(groups.iter().any(|g| g.len() > 1).then_some(groups)) +} + +/// Each nonempty compatible group computes its prefix once. Singletons save nothing. +fn shared_prefix_savings(groups: &[Vec], prefix_len: usize) -> usize { + groups + .iter() + .map(|group| (group.len() - 1) * prefix_len) + .sum() +} + #[derive(Debug, Clone)] pub struct Config { pub hidden: usize, @@ -468,6 +504,7 @@ struct Scratch { lo: usize, ln: usize, workspace: usize, + workspace_bytes: usize, attn_in: usize, aq: usize, agate: usize, @@ -490,6 +527,7 @@ impl Scratch { lora: Option<&lora::Lora>, reference: Option<&cuda::Reference>, fp32_gdn: Option<&cuda::Fp32Gdn>, + batched_gdn: Option<&cuda::BatchedGdn>, ) -> Result { let (h, kd, vd, hv) = (cfg.hidden, cfg.key_dim(), cfg.value_dim(), cfg.lin_v_heads); let (hq, hk, hd) = (cfg.heads, cfg.kv_heads, cfg.head_dim); @@ -518,6 +556,14 @@ impl Scratch { } else { ws_floats }; + let ws_floats = if let Some(gdn) = batched_gdn { + // SAFETY: validated model dimensions and a bounded sequence count. + let batched = unsafe { (gdn.workspace_floats)(cap as i32, hv as i32, 16) }; + ensure!(batched > 0, "cannot size batched GDN workspace"); + ws_floats.max(batched) + } else { + ws_floats + }; let offsets = [ take(cap * 4), take(cap * h * BF16), @@ -599,6 +645,7 @@ impl Scratch { lo, ln, workspace, + workspace_bytes: ws_floats * F32, attn_in, aq, agate, @@ -656,6 +703,34 @@ pub struct PrefixState { bytes: usize, } +/// A packed suffix batch uses one freshly captured prefix and original algorithm. +struct SuffixBatch<'a> { + prefix: &'a PrefixState, + ops: cuda::SharedSuffix, + gemm: cuda::ReuseGemm, + original_rows: usize, +} + +/// Bound only the suffix work; oversized suffixes retain individual windows. +fn suffix_segments(indices: &[usize], lengths: &[usize], prefix: usize) -> Vec> { + let mut groups = Vec::new(); + let mut group = Vec::new(); + let mut tokens = 0; + for &index in indices { + let suffix = lengths[index] - prefix; + if !group.is_empty() && (group.len() == 16 || tokens + suffix > 4096) { + groups.push(std::mem::take(&mut group)); + tokens = 0; + } + group.push(index); + tokens += suffix; + } + if !group.is_empty() { + groups.push(group); + } + groups +} + impl PrefixState { /// Number of tokens covered by this nonempty prefix. pub fn token_count(&self) -> usize { @@ -713,6 +788,13 @@ pub struct PromptGroup<'a> { pub branches: Vec<&'a [u32]>, } +/// A successful compatible-prefix execution, with results in original input order. +pub struct SharedPrefixOutput { + pub hidden: Vec>, + /// Prefix tokens not recomputed, accounting for separate groups and singletons. + pub saved_tokens: usize, +} + /// 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. @@ -754,6 +836,8 @@ fn runs(prompts: &SharedPrompts) -> Runs { pub struct Model { pub cfg: Config, prefix_owner: Arc<()>, + /// Allocation capacity only: state is invalidated and fully recaptured per request. + text_prefix_workspace: Option<(usize, PrefixState)>, _weights: Weights, embed: Tensor, final_norm: Tensor, @@ -765,6 +849,8 @@ pub struct Model { fixed_gemm: *mut c_void, reference: Option<&'static cuda::Reference>, fp32_gdn: Option, + batched_gdn: Option, + packed_gemm: Option, /// 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. @@ -772,7 +858,7 @@ pub struct Model { /// Hidden rows and logits for `forward_multimodal_readout`'s vocabulary pass, and /// the host copy of the logits; allocated on first use. vocab_buffers: Option<(DeviceBuffer, DeviceBuffer, Vec)>, - /// Opt-in replay with at most 64 captures per mode, keyed by sequence lengths. + /// Opt-in replay keyed by sequence lengths; bounded separately for each mode. graph_enabled: bool, graph_stats: GraphStats, graphs: VecDeque<(Vec, cuda::Graph)>, @@ -924,6 +1010,7 @@ impl Model { let model = Self { cfg, prefix_owner: Arc::new(()), + text_prefix_workspace: None, _weights: weights, embed, final_norm, @@ -934,6 +1021,8 @@ impl Model { fixed_gemm: std::ptr::null_mut(), reference: None, fp32_gdn: None, + batched_gdn: None, + packed_gemm: None, scratch: None, shared: [None, None], vocab_buffers: None, @@ -957,6 +1046,30 @@ impl Model { } } + /// Opt into packing output/down projections when the original algorithms agree. + /// Select before the first forward; other model consumers retain separate GEMMs. + pub fn enable_packed_projections(&mut self) -> Result<()> { + ensure!( + self.scratch.is_none() && self.reference.is_none() && self.lora.is_none(), + "select packed projections on a merged model before the first forward" + ); + self.packed_gemm = Some(cuda::api().packed_gemm()?); + Ok(()) + } + + /// Opt into independent GDN batching before allocating any request buffers. + pub fn enable_batched_gdn(&mut self) -> Result<()> { + ensure!( + self.scratch.is_none() + && self.reference.is_none() + && self.fp32_gdn.is_none() + && self.lora.is_none(), + "select batched GDN on a merged model before the first forward" + ); + self.batched_gdn = Some(cuda::api().batched_gdn()?); + Ok(()) + } + /// Cua-S1's unmerged text path: framework norm/GEMM/attention plus FP32 GDN. pub fn enable_reference_text_numerics(&mut self) -> Result<()> { ensure!( @@ -1129,7 +1242,7 @@ impl Model { } } - /// Preserve each prompt's output/down GEMM shape and split-K reduction order. + /// Keep per-prompt output/down projections, or opt into packing with their algorithms. fn gemm_sequences( &self, s: &Scratch, @@ -1139,6 +1252,30 @@ impl Model { lengths: &[usize], ) -> Result<()> { let (n, k) = (w.shape[0], w.shape[1]); + if let Some(packed) = self.packed_gemm + && self.reference.is_none() + && self.lora.is_none() + { + // SAFETY: x/y contain all validated packed rows, w is [n,k], and + // lengths is read synchronously by the host planner. Device work + // remains on this serialized model's stream and workspace. + return unsafe { + check( + packed( + self.gemm, + s.at(x), + w.ptr, + s.at(y), + lengths.as_ptr(), + lengths.len(), + n as i32, + k as i32, + self.stream, + ), + "packed projection", + ) + }; + } let mut offset = 0; for &length in lengths { self.gemm(s, x + offset * k * BF16, w, y + offset * n * BF16, length)?; @@ -1147,6 +1284,26 @@ impl Model { Ok(()) } + fn gdn_prefill(&self) -> cuda::GdnPrefill { + let api = cuda::api(); + if let Some(reference) = self.reference { + reference.gdn_prefill + } else if self.packed_gemm.is_some() { + api.wide_gdn().unwrap_or(api.cs1_gdn_prefill) + } else { + api.cs1_gdn_prefill + } + } + + fn convolution(&self) -> cuda::GdnConv { + let api = cuda::api(); + if self.packed_gemm.is_some() { + api.vector_conv().unwrap_or(api.cs1_gdn_conv) + } else { + api.cs1_gdn_conv + } + } + /// Finish queued work before releasing external execution admission. pub fn synchronize(&self) -> Result<()> { cuda::set_device(0)?; @@ -1155,7 +1312,16 @@ impl Model { fn prepare_scratch(&mut self, t: usize) -> Result<()> { cuda::set_device(0)?; - if self.scratch.as_ref().is_none_or(|s| t > s.cap) { + // Graphs retain scratch addresses. Reserve the bounded packing window on + // the first opt-in text execution so later shapes within it do not discard + // already-captured graphs. Larger supported inputs still grow normally. + let capacity = if self.graph_enabled && self.packed_gemm.is_some() { + t.max(4096) + } else { + t + } + .next_multiple_of(1024); + if self.scratch.as_ref().is_none_or(|s| capacity > s.cap) { cuda::synchronize(self.stream)?; if !self.graphs.is_empty() || !self.mm_graphs.is_empty() { self.graph_stats.invalidations += 1; @@ -1165,11 +1331,12 @@ impl Model { self.scratch = None; self.scratch = Some(Scratch::new( &self.cfg, - t.next_multiple_of(1024), + capacity, self.stream, self.lora.as_ref(), self.reference, self.fp32_gdn.as_ref(), + self.batched_gdn.as_ref(), )?); } Ok(()) @@ -1273,6 +1440,223 @@ impl Model { .collect() } + /// Evaluate candidates sharing an exact, chunk-aligned token prefix once. + /// Retain each complete forward's algorithms, sharing only within compatible + /// candidate groups. Return None before computation when no group can share. + /// Every request recaptures all numeric state; only allocation capacity persists. + pub fn try_forward_shared_prefix( + &mut self, + inputs: &[&[u32]], + prefix_len: usize, + ) -> Result> { + ensure!( + self.lora.is_none(), + "shared text prefix requires merged weights" + ); + ensure!( + self.reference.is_none() && self.fp32_gdn.is_none(), + "shared text prefix requires standard BF16 numerics" + ); + ensure!( + prefix_len >= 64 && prefix_len.is_multiple_of(64), + "shared prefix must be a positive chunk multiple" + ); + ensure!(inputs.len() >= 2, "shared prefix needs multiple candidates"); + ensure!( + inputs + .iter() + .all(|ids| ids.len() > prefix_len && ids.len() <= self.cfg.max_positions), + "empty suffix or oversized prompt" + ); + ensure!( + inputs + .iter() + .flat_map(|ids| ids.iter()) + .all(|&id| (id as usize) < self.embed.shape[0]), + "token id outside the vocabulary" + ); + ensure!( + inputs + .iter() + .all(|ids| ids[..prefix_len] == inputs[0][..prefix_len]), + "candidate prefixes differ" + ); + cuda::set_device(0)?; + let Some(reuse) = cuda::api().reuse_gemm() else { + return Ok(None); + }; + let original_rows: Vec = inputs.iter().map(|ids| ids.len()).collect(); + let mut shapes = std::collections::BTreeSet::new(); + for layer in &self.layers { + let (input, output) = match &layer.mixer { + Mixer::Linear(la) => (&la.in_proj, &la.out), + Mixer::Full(fa) => (&fa.qkv, &fa.o), + }; + for weight in [input, output, &layer.gate_up, &layer.down] { + shapes.insert((weight.shape[0], weight.shape[1])); + } + } + let gemm = self.gemm; + let Some(groups) = compatible_prefix_groups(inputs.len(), |indices| { + let originals: Vec = indices.iter().map(|&i| original_rows[i]).collect(); + let mut windows: Vec = originals.iter().map(|&n| n - prefix_len).collect(); + windows.push(prefix_len); + for &(n, k) in &shapes { + // SAFETY: host arrays stay alive for the synchronous capability + // check. No device output or prefix state is produced here. + let rc = unsafe { + (reuse.check)( + gemm, + originals.as_ptr(), + originals.len(), + windows.as_ptr(), + windows.len(), + n as i32, + k as i32, + ) + }; + if rc == -1 { + return Ok(false); + } + check(rc, "prefix projection compatibility")?; + } + Ok(true) + })? + else { + return Ok(None); + }; + let (capacity, mut state) = match self.text_prefix_workspace.take() { + Some((capacity, state)) if capacity >= prefix_len => (capacity, state), + old => { + // Release the smaller allocation before growth, bounding peak memory. + drop(old); + (prefix_len, self.alloc_prefix(prefix_len)?) + } + }; + state.len = prefix_len; + state.initialized = false; + let result = (|| { + self.prepare_scratch(inputs.iter().map(|ids| ids.len()).max().unwrap())?; + let mut hidden = vec![Vec::new(); inputs.len()]; + let saved_tokens = shared_prefix_savings(&groups, prefix_len); + for group in groups { + if group.len() == 1 { + hidden[group[0]] = self.forward(inputs[group[0]])?; + continue; + } + let ops = cuda::api().shared_suffix(); + let mut segments = Vec::new(); + let mut required = inputs.iter().map(|ids| ids.len()).max().unwrap(); + for indices in suffix_segments(&group, &original_rows, prefix_len) { + let rows: usize = indices.iter().map(|&i| original_rows[i] - prefix_len).sum(); + let original = [original_rows[indices[0]]]; + let target = [rows]; + let mut packed = ops.is_some() && indices.len() > 1; + if packed { + for &(n, k) in &shapes { + // SAFETY: synchronous shape-only check, before prefix capture. + let rc = unsafe { + (reuse.check)( + gemm, + original.as_ptr(), + 1, + target.as_ptr(), + 1, + n as i32, + k as i32, + ) + }; + if rc == -1 { + packed = false; + break; + } + check(rc, "packed suffix compatibility")?; + } + } + if packed { + required = required.max(rows); + } + segments.push((indices, packed)); + } + self.prepare_scratch(required)?; + // Different groups can have different original algorithms. Recompute + // every captured value; buffers never make a later request's state valid. + state.initialized = false; + let projection = Some((reuse, original_rows[group[0]])); + let s = self.scratch.as_ref().unwrap(); + self.embed_tokens(s, &inputs[group[0]][..prefix_len])?; + self.run_window( + s, + 0, + prefix_len, + false, + Some(&mut state), + None, + self.gemm, + projection, + )?; + self.synchronize()?; + state.initialized = true; + for (indices, packed) in segments { + if packed { + let lengths: Vec = indices + .iter() + .map(|&i| inputs[i].len() - prefix_len) + .collect(); + let ids: Vec = indices + .iter() + .flat_map(|&i| inputs[i][prefix_len..].iter().copied()) + .collect(); + self.embed_tokens(s, &ids)?; + self.run_with_suffix( + s, + &lengths, + false, + false, + Some(&SuffixBatch { + prefix: &state, + ops: ops.unwrap(), + gemm: reuse, + original_rows: original_rows[indices[0]], + }), + )?; + let mut end = 0; + for (index, length) in indices.into_iter().zip(lengths) { + end += length; + hidden[index] = self.last_hidden(s, end)?; + } + } else { + for index in indices { + let ids = inputs[index]; + self.embed_tokens_at(s, &ids[prefix_len..], prefix_len)?; + self.run_window( + s, + prefix_len, + ids.len(), + false, + None, + Some(&state), + self.gemm, + projection, + )?; + hidden[index] = self.last_hidden(s, ids.len())?; + } + } + } + } + Ok(SharedPrefixOutput { + hidden, + saved_tokens, + }) + })(); + // A failed forward may have queued work. Finish it before retaining buffers; + // on a synchronization error, drop them instead of recycling their state. + self.synchronize()?; + state.initialized = false; + self.text_prefix_workspace = Some((capacity, state)); + result.map(Some) + } + /// Prefill one unpadded prompt with already-adapted BF16 image embeddings and /// explicit `[3, 1, sequence]` T/H/W positions. No vision tower runs here. /// Load a checkpoint with the matching multimodal language adapter merged. @@ -1329,12 +1713,20 @@ impl Model { match captured { Ok(graph) => { self.graph_stats.captures += 1; + // Candidate packing produces many length vectors. Keep a bounded + // larger text cache for opt-in packed consumers; other modes retain + // their existing budget. Captures store operations, never input data. + let capacity = if !multimodal && self.packed_gemm.is_some() { + 256 + } else { + 64 + }; let cache = if multimodal { &mut self.mm_graphs } else { &mut self.graphs }; - if cache.len() == 64 { + if cache.len() == capacity { cuda::synchronize(self.stream)?; cache.pop_front(); } @@ -1455,7 +1847,16 @@ impl Model { cap.initialized = false; cap.len = tend; } - self.run_window(s, qb, tend, true, capture.as_deref_mut(), prefix, gemm)?; + self.run_window( + s, + qb, + tend, + true, + capture.as_deref_mut(), + prefix, + gemm, + None, + )?; if let Some(cap) = capture { cap.initialized = true; } @@ -1729,7 +2130,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.gemm)?; + self.run_window(s, 0, t, true, Some(state), None, self.gemm, None)?; self.synchronize()?; state.initialized = true; Ok(()) @@ -1771,7 +2172,16 @@ 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.gemm)?; + self.run_window( + s, + prefix.len, + tend, + true, + None, + Some(prefix), + self.gemm, + None, + )?; self.last_hidden(s, tend) } @@ -1801,7 +2211,7 @@ impl Model { 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.run_window(s, 0, t, false, None, None, gemm, None)?; self.last_hidden(s, t) } @@ -1873,7 +2283,7 @@ impl Model { let request = match request { Some(state) => { self.embed_tokens(s, &runs.prefix)?; - self.run_window(s, 0, p, false, Some(&mut *state), None, gemm)?; + self.run_window(s, 0, p, false, Some(&mut *state), None, gemm, None)?; Some(&*state) } None => None, @@ -1888,7 +2298,7 @@ impl Model { 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)?; + self.run_window(s, p, start, false, Some(&mut *state), request, gemm, None)?; Some(&*state) }; let rows = run @@ -1897,7 +2307,7 @@ impl Model { .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.run_window(s, start, end, false, None, from, gemm, None)?; self.last_hidden(s, end) }) .collect::>>()?; @@ -1999,19 +2409,157 @@ impl Model { .collect()) } + /// Keep the complete prompt's algorithms when its final MLP can use one row. + fn last_mlp_reuse(&self, t: usize) -> Result> { + if self.packed_gemm.is_none() + || self.reference.is_some() + || self.lora.is_some() + || self.fp32_gdn.is_some() + || self.cfg.hidden != 5120 + || self.cfg.intermediate != 17408 + || self.layers.len() != 64 + || !(512..=16384).contains(&t) + { + return Ok(None); + } + let Some(reuse) = cuda::api().reuse_gemm() else { + return Ok(None); + }; + let layer = self.layers.last().expect("validated nonempty model"); + for weight in [&layer.gate_up, &layer.down] { + // SAFETY: both host row arrays remain live for this synchronous check. + // Check both projections before any final-layer output is changed. + let rc = unsafe { + (reuse.check)( + self.gemm, + &t, + 1, + &1usize, + 1, + weight.shape[0] as i32, + weight.shape[1] as i32, + ) + }; + if rc == -1 { + return Ok(None); + } + check(rc, "last-row MLP compatibility")?; + } + Ok(Some(reuse)) + } + + /// Only last_hidden() consumes this single prompt's final layer outputs. + /// Earlier rows need no final MLP; the next call embeds a fresh residual. + fn run_last_mlp( + &self, + s: &Scratch, + layer: &Layer, + t: usize, + reuse: cuda::ReuseGemm, + ) -> Result<()> { + let cfg = &self.cfg; + let row = t - 1; + let hidden_offset = row * cfg.hidden * BF16; + let gate_offset = row * 2 * cfg.intermediate * BF16; + let act_offset = row * cfg.intermediate * BF16; + let p = |offset| s.at(offset); + // SAFETY: the complete prompt already ran its mixer, scratch holds t rows, + // and the retained-algorithm check validated both one-row projections. + // These row-local operations keep the original BF16 rounding boundaries. + unsafe { + check( + (cuda::api().cs1_add_rms_norm)( + p(s.res + hidden_offset), + p(s.delta + hidden_offset), + layer.post_norm.ptr, + p(s.x + hidden_offset), + 1, + cfg.hidden as i32, + cfg.eps, + self.stream, + ), + "last-row post-attention norm", + )?; + check( + (reuse.run)( + self.gemm, + p(s.x + hidden_offset), + layer.gate_up.ptr, + p(s.gate_up + gate_offset), + 1, + (2 * cfg.intermediate) as i32, + cfg.hidden as i32, + t as i32, + self.stream, + ), + "last-row gate/up projection", + )?; + check( + (cuda::api().cs1_silu_mul)( + p(s.gate_up + gate_offset), + (2 * cfg.intermediate) as i32, + p(s.act + act_offset), + 1, + cfg.intermediate as i32, + self.stream, + ), + "last-row silu mul", + )?; + check( + (reuse.run)( + self.gemm, + p(s.act + act_offset), + layer.down.ptr, + p(s.delta + hidden_offset), + 1, + cfg.hidden as i32, + cfg.intermediate as i32, + t as i32, + self.stream, + ), + "last-row down projection", + )?; + check( + (cuda::api().cs1_add_rms_norm)( + p(s.res + hidden_offset), + p(s.delta + hidden_offset), + self.final_norm.ptr, + p(s.x + hidden_offset), + 1, + cfg.hidden as i32, + cfg.eps, + self.stream, + ), + "last-row final norm", + )?; + } + Ok(()) + } + /// Queue language layers over packed embeddings, resetting sequence positions. - /// Final-norm hidden states end up in `s.x`. `bidirectional` changes only the - /// full-attention mask; Gated DeltaNet stays causal and sequence-local. + /// Final-norm hidden states end up in `s.x`. Bidirectional attention + /// changes only the full-attention mask; Gated DeltaNet stays causal. fn run( &self, s: &Scratch, lengths: &[usize], custom_positions: bool, bidirectional: bool, + ) -> Result<()> { + self.run_with_suffix(s, lengths, custom_positions, bidirectional, None) + } + + fn run_with_suffix( + &self, + s: &Scratch, + lengths: &[usize], + custom_positions: bool, + bidirectional: bool, + suffix: Option<&SuffixBatch<'_>>, ) -> Result<()> { let bidir = if bidirectional { ensure!( - self.reference.is_none() && self.lora.is_none(), + self.reference.is_none() && self.lora.is_none() && suffix.is_none(), "prompt-bidirectional prefill does not use reference numerics or an adapter" ); Some(cuda::api().attention_gated_bidirectional()?) @@ -2038,8 +2586,37 @@ impl Model { let (cos, sin) = if custom_positions { (s.custom_cos, s.custom_sin) } else { - (s.cos, s.sin) + let offset = suffix.map_or(0, |x| x.prefix.len * cfg.rotary_half * BF16); + (s.cos + offset, s.sin + offset) + }; + let project = |x: usize, weight: &Tensor, y: usize, group: Option| { + if let Some(suffix) = suffix { + // SAFETY: combined rows were checked against the original algorithm; + // scratch owns the compact batch and all work uses the model stream. + unsafe { + check( + (suffix.gemm.run)( + self.gemm, + p(x), + weight.ptr, + p(y), + t as i32, + weight.shape[0] as i32, + weight.shape[1] as i32, + suffix.original_rows as i32, + st, + ), + "packed suffix projection", + ) + } + } else if let Some(group) = group { + self.gemm_group(s, x, weight, y, t, lengths, group) + } else { + self.gemm_sequences(s, x, weight, y, lengths) + } }; + let mut linear_index = 0; + let mut attention_index = 0; // SAFETY (every kernel call below): the pointers are weights in the arena or // scratch buffers laid out for at least t tokens with the widths used here. unsafe { @@ -2058,18 +2635,18 @@ impl Model { "input norm", )?; } + let last_mlp = + if !bidirectional && !custom_positions && suffix.is_none() && lengths.len() == 1 { + self.last_mlp_reuse(t)? + } else { + None + }; + let convolution = self.convolution(); + let gdn_prefill = self.gdn_prefill(); for (i, layer) in self.layers.iter().enumerate() { match &layer.mixer { Mixer::Linear(la) => { - self.gemm_group( - s, - s.x, - &la.in_proj, - s.gdn_in, - t, - lengths, - Group::LinearInput, - )?; + project(s.x, &la.in_proj, s.gdn_in, Some(Group::LinearInput))?; self.apply_lora(i, Group::LinearInput, s, s.x, s.gdn_in, lengths)?; let ld = w.gdn_in as i32; let z = s.gdn_in + w.conv * BF16; @@ -2078,18 +2655,34 @@ impl Model { unsafe { for &(offset, length) in &sequences { check( - (cuda::api().cs1_gdn_conv)( - p(s.gdn_in + offset * w.gdn_in * BF16), - ld, - la.conv.ptr, - p(s.lq + offset * kd * BF16), - p(s.lk + offset * kd * BF16), - p(s.lv + offset * vd * BF16), - length, - kd as i32, - vd as i32, - st, - ), + if let Some(suffix) = suffix { + (suffix.ops.conv)( + p(s.gdn_in + offset * w.gdn_in * BF16), + suffix.prefix.conv_tail[linear_index].at(0), + ld, + la.conv.ptr, + p(s.lq + offset * kd * BF16), + p(s.lk + offset * kd * BF16), + p(s.lv + offset * vd * BF16), + length, + kd as i32, + vd as i32, + st, + ) + } else { + (convolution)( + p(s.gdn_in + offset * w.gdn_in * BF16), + ld, + la.conv.ptr, + p(s.lq + offset * kd * BF16), + p(s.lk + offset * kd * BF16), + p(s.lv + offset * vd * BF16), + length, + kd as i32, + vd as i32, + st, + ) + }, "gdn conv", )?; } @@ -2108,10 +2701,10 @@ impl Model { ), "gdn gates", )?; - for &(offset, length) in &sequences { - check( - if let Some(gdn) = &self.fp32_gdn { - gdn.run( + if let Some(suffix) = suffix { + for &(offset, length) in &sequences { + check( + (cuda::api().cs1_gdn_prefill_x)( p(s.lq + offset * kd * BF16), p(s.lk + offset * kd * BF16), p(s.lv + offset * vd * BF16), @@ -2123,28 +2716,72 @@ impl Model { hv as i32, cfg.lin_k_heads as i32, (cfg.lin_k_dim as f32).powf(-0.5), + suffix.prefix.gdn_state[linear_index].at(0), + std::ptr::null_mut(), st, - ) - } else { - (self - .reference - .map_or(cuda::api().cs1_gdn_prefill, |r| r.gdn_prefill))( - p(s.lq + offset * kd * BF16), - p(s.lk + offset * kd * BF16), - p(s.lv + offset * vd * BF16), - p(s.g + offset * hv * F32).cast(), - p(s.beta + offset * hv * BF16), - p(s.lo + offset * vd * BF16), - p(s.workspace).cast(), - length, - hv as i32, - cfg.lin_k_heads as i32, - (cfg.lin_k_dim as f32).powf(-0.5), - st, - ) - }, - "gdn prefill", + ), + "suffix GDN prefill", + )?; + } + } else if let Some(gdn) = &self.batched_gdn + && (2..=16).contains(&lengths.len()) + { + check( + (gdn.run)( + p(s.lq), + p(s.lk), + p(s.lv), + p(s.g).cast(), + p(s.beta), + p(s.lo), + p(s.workspace).cast(), + s.workspace_bytes, + lengths.as_ptr(), + lengths.len(), + hv as i32, + cfg.lin_k_heads as i32, + (cfg.lin_k_dim as f32).powf(-0.5), + st, + ), + "batched GDN prefill", )?; + } else { + for &(offset, length) in &sequences { + check( + if let Some(gdn) = &self.fp32_gdn { + gdn.run( + p(s.lq + offset * kd * BF16), + p(s.lk + offset * kd * BF16), + p(s.lv + offset * vd * BF16), + p(s.g + offset * hv * F32).cast(), + p(s.beta + offset * hv * BF16), + p(s.lo + offset * vd * BF16), + p(s.workspace).cast(), + length, + hv as i32, + cfg.lin_k_heads as i32, + (cfg.lin_k_dim as f32).powf(-0.5), + st, + ) + } else { + gdn_prefill( + p(s.lq + offset * kd * BF16), + p(s.lk + offset * kd * BF16), + p(s.lv + offset * vd * BF16), + p(s.g + offset * hv * F32).cast(), + p(s.beta + offset * hv * BF16), + p(s.lo + offset * vd * BF16), + p(s.workspace).cast(), + length, + hv as i32, + cfg.lin_k_heads as i32, + (cfg.lin_k_dim as f32).powf(-0.5), + st, + ) + }, + "gdn prefill", + )?; + } } check( (cuda::api().cs1_gated_rms_norm)( @@ -2162,19 +2799,12 @@ impl Model { "gated norm", )?; } - self.gemm_sequences(s, s.ln, &la.out, s.delta, lengths)?; + project(s.ln, &la.out, s.delta, None)?; self.apply_lora(i, Group::LinearOutput, s, s.ln, s.delta, lengths)?; + linear_index += 1; } Mixer::Full(fa) => { - self.gemm_group( - s, - s.x, - &fa.qkv, - s.attn_in, - t, - lengths, - Group::AttentionInput, - )?; + project(s.x, &fa.qkv, s.attn_in, Some(Group::AttentionInput))?; self.apply_lora(i, Group::AttentionInput, s, s.x, s.attn_in, lengths)?; let ld = w.attn_in as i32; let k = s.attn_in + w.attn_q * BF16; @@ -2207,7 +2837,25 @@ impl Model { )?; let scale = (cfg.head_dim as f32).powf(-0.5); check( - if let Some(bidir) = bidir { + if let Some(suffix) = suffix { + (suffix.ops.attention)( + p(s.aq + offset * cfg.heads * cfg.head_dim * BF16), + p(s.ak + offset * cfg.kv_heads * cfg.head_dim * BF16), + p(v + offset * w.attn_in * BF16), + ld, + p(s.agate + offset * cfg.heads * cfg.head_dim * BF16), + p(s.ao + offset * cfg.heads * cfg.head_dim * BF16), + suffix.prefix.attn_kv[attention_index].0.at(0), + suffix.prefix.attn_kv[attention_index].1.at(0), + length, + suffix.prefix.len as i32, + hq, + hk, + hd, + (cfg.head_dim as f32).powf(-0.5), + st, + ) + } else if let Some(bidir) = bidir { (bidir)( p(s.aq + offset * cfg.heads * cfg.head_dim * BF16), p(s.ak + offset * cfg.kv_heads * cfg.head_dim * BF16), @@ -2257,10 +2905,16 @@ impl Model { )?; } } - self.gemm_sequences(s, s.ao, &fa.o, s.delta, lengths)?; + project(s.ao, &fa.o, s.delta, None)?; self.apply_lora(i, Group::AttentionOutput, s, s.ao, s.delta, lengths)?; + attention_index += 1; } } + if i + 1 == self.layers.len() + && let Some(reuse) = last_mlp + { + return self.run_last_mlp(s, layer, t, reuse); + } unsafe { check( (self @@ -2278,7 +2932,7 @@ impl Model { "post-attention norm", )?; } - self.gemm_group(s, s.x, &layer.gate_up, s.gate_up, t, lengths, Group::GateUp)?; + project(s.x, &layer.gate_up, s.gate_up, Some(Group::GateUp))?; self.apply_lora(i, Group::GateUp, s, s.x, s.gate_up, lengths)?; unsafe { check( @@ -2293,7 +2947,7 @@ impl Model { "silu mul", )?; } - self.gemm_sequences(s, s.act, &layer.down, s.delta, lengths)?; + project(s.act, &layer.down, s.delta, None)?; self.apply_lora(i, Group::Down, s, s.act, s.delta, lengths)?; let next = self .layers @@ -2330,12 +2984,11 @@ impl Model { /// 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, 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. + /// captured+continued pass retains the sequence-local kernel arithmetic. + /// Fixed-profile callers use their fixed-algorithm handle. Open-Jev text + /// reuse instead validates and retains each complete prompt's algorithm. + /// Custom position tables contain window rows; text tables use absolute rows. + #[allow(clippy::too_many_arguments)] // Window, prefix states, GEMM handle and optional original shape. fn run_window( &self, s: &Scratch, @@ -2345,6 +2998,7 @@ impl Model { capture: Option<&mut PrefixState>, prefix: Option<&PrefixState>, gemm: *mut c_void, + reuse: Option<(cuda::ReuseGemm, usize)>, ) -> Result<()> { ensure!( self.fp32_gdn.is_none(), @@ -2359,6 +3013,32 @@ impl Model { let cfg = &self.cfg; let st = self.stream; let rows = tend - qb; + let project = |x, weight: &Tensor, y, group| -> Result<()> { + if let Some((api, original_rows)) = reuse { + // SAFETY: compatibility was checked before prefix capture. Scratch + // owns rows at the offsets passed by the unchanged layer layout. + unsafe { + check( + (api.run)( + gemm, + s.at(x), + weight.ptr, + s.at(y), + rows as i32, + weight.shape[0] as i32, + weight.shape[1] as i32, + original_rows as i32, + self.stream, + ), + "retained prefix projection", + ) + } + } else if let Some(group) = group { + self.gemm_group_with(gemm, s, x, weight, y, rows, &[rows], group) + } else { + self.gemm_with(gemm, s, x, weight, y, rows) + } + }; let (ti, hi, eps) = (rows as i32, cfg.hidden as i32, cfg.eps); let (kd, vd, hv) = (cfg.key_dim(), cfg.value_dim(), cfg.lin_v_heads); let (hq, hk, hd) = (cfg.heads as i32, cfg.kv_heads as i32, cfg.head_dim as i32); @@ -2401,18 +3081,15 @@ impl Model { "input norm", )?; } + let convolution = self.convolution(); for (i, layer) in self.layers.iter().enumerate() { match &layer.mixer { Mixer::Linear(la) => { - self.gemm_group_with( - gemm, - s, + project( s.x + qb * hb, &la.in_proj, s.gdn_in + qb * ldbb, - rows, - &[rows], - Group::LinearInput, + Some(Group::LinearInput), )?; self.apply_lora( i, @@ -2455,7 +3132,7 @@ impl Model { // those three outputs so continuation rows see all taps. let conv_start = qb.saturating_sub(3); check( - (cuda::api().cs1_gdn_conv)( + (convolution)( p(s.gdn_in + conv_start * ldbb), ld, la.conv.ptr, @@ -2519,7 +3196,7 @@ impl Model { "gated norm", )?; } - self.gemm_with(gemm, s, s.ln + qb * vb, &la.out, s.delta + qb * hb, rows)?; + project(s.ln + qb * vb, &la.out, s.delta + qb * hb, None)?; self.apply_lora( i, Group::LinearOutput, @@ -2531,15 +3208,11 @@ impl Model { la_i += 1; } Mixer::Full(fa) => { - self.gemm_group_with( - gemm, - s, + project( s.x + qb * hb, &fa.qkv, s.attn_in + qb * ab, - rows, - &[rows], - Group::AttentionInput, + Some(Group::AttentionInput), )?; self.apply_lora( i, @@ -2608,7 +3281,7 @@ impl Model { "gated attention", )?; } - self.gemm_with(gemm, s, s.ao + qb * ob, &fa.o, s.delta + qb * hb, rows)?; + project(s.ao + qb * ob, &fa.o, s.delta + qb * hb, None)?; self.apply_lora( i, Group::AttentionOutput, @@ -2637,15 +3310,11 @@ impl Model { "post-attention norm", )?; } - self.gemm_group_with( - gemm, - s, + project( s.x + qb * hb, &layer.gate_up, s.gate_up + qb * 2 * cfg.intermediate * BF16, - rows, - &[rows], - Group::GateUp, + Some(Group::GateUp), )?; self.apply_lora( i, @@ -2668,13 +3337,11 @@ impl Model { "silu mul", )?; } - self.gemm_with( - gemm, - s, + project( s.act + qb * cfg.intermediate * BF16, &layer.down, s.delta + qb * hb, - rows, + None, )?; self.apply_lora( i, @@ -2716,3 +3383,11 @@ mod graph_tests; #[cfg(test)] #[path = "../../../../../tests/qwen3_5/model.rs"] mod tests; + +#[cfg(test)] +#[path = "../../../../../tests/qwen3_5/prefix_groups.rs"] +mod prefix_group_tests; + +#[cfg(test)] +#[path = "../../../../../tests/qwen3_5/last_mlp.rs"] +mod last_mlp_tests; diff --git a/tests/open_jev/batching.rs b/tests/open_jev/batching.rs index b9632330..4de2af2c 100644 --- a/tests/open_jev/batching.rs +++ b/tests/open_jev/batching.rs @@ -1,4 +1,4 @@ -use super::ranges; +use super::{ranges, shared_prefix}; #[test] fn preserves_all_candidates_across_token_and_sequence_limits() { @@ -15,3 +15,17 @@ fn preserves_all_candidates_across_token_and_sequence_limits() { (0..inputs.len()).collect::>() ); } + +#[test] +fn shared_prefix_is_exact_aligned_and_leaves_a_suffix() { + let a = vec![1u32; 2100]; + let mut b = a.clone(); + b[2000] = 2; + assert_eq!(shared_prefix(&[&a, &b]), 1984); + b[100] = 3; + assert_eq!(shared_prefix(&[&a, &b]), 0); + assert_eq!(shared_prefix(&[&a, &a]), 2048); + assert_eq!(shared_prefix(&[&a]), 0); + assert_eq!(shared_prefix(&[&a, &a[..2048]]), 0); + assert_eq!(shared_prefix(&vec![a.as_slice(); 17]), 0); +} diff --git a/tests/open_jev/prefill_batch.rs b/tests/open_jev/prefill_batch.rs index 4fbfad15..74f5846e 100644 --- a/tests/open_jev/prefill_batch.rs +++ b/tests/open_jev/prefill_batch.rs @@ -96,3 +96,49 @@ async fn packed_candidates_preserve_isolation_order_and_graph_shapes() { assert!((row[0] - expected).abs() <= 0.01); } } + +#[tokio::test(flavor = "current_thread")] +#[ignore = "needs the pinned Open-Jev export, CUDA library and a GPU reservation"] +async fn text_prefix_buffers_grow_and_shrink_without_reusing_prior_contents() { + let model = PathBuf::from(std::env::var_os("OPEN_JEV_MODEL").unwrap()); + let library = PathBuf::from(std::env::var_os("OPEN_JEV_CUDA_LIB").unwrap()); + let engine = Engine::load(&model, &library).await.unwrap(); + let mut first = None; + for (prefix, token) in [(2048, 42), (3072, 43), (2048, 44), (2048, 42)] { + let inputs: Vec> = (0..3) + .map(|index| { + let mut ids = vec![token; prefix]; + ids.extend(vec![7 + index; 256]); + ids + }) + .collect(); + let mut expected = Vec::new(); + for ids in &inputs { + expected.push( + engine + .executor + .execute(&engine.scheduler, vec![vec![ids.clone()]]) + .await + .unwrap()[0][0], + ); + } + let actual = engine + .executor + .execute_detailed(&engine.scheduler, vec![inputs]) + .await + .unwrap(); + assert_eq!( + actual.saved_prefix_tokens, + 2 * prefix, + "exercise actual prefix reuse" + ); + assert_eq!(actual.scores, vec![expected.clone()]); + if token == 42 { + if let Some(first) = &first { + assert_eq!(&actual.scores, first, "returning input after changed state"); + } else { + first = Some(actual.scores); + } + } + } +} diff --git a/tests/open_jev/processing.rs b/tests/open_jev/processing.rs index 95742530..36f0bd3f 100644 --- a/tests/open_jev/processing.rs +++ b/tests/open_jev/processing.rs @@ -131,12 +131,49 @@ fn prepared_candidates_have_exact_ids_and_finished_responses_keep_mapping() { ); assert_eq!( body["metadata"]["prefix_cache"], - json!({"enabled":false,"mode":"independent_candidates"}) + json!({"policy":"default","enabled":false,"scope":"none","saved_tokens":0}) ); assert!(body["metadata"]["inference_seconds"].as_f64().unwrap() >= 0.0); assert_eq!(body["model"], CHECKPOINTS[0].model_id); } +#[test] +fn prefix_metadata_uses_executed_savings_without_changing_answers_or_usage() { + let request = serde_json::to_vec(&json!({ + "state": "x ".repeat(768), + "questions": {"q": {"type": "choice", "instructions": "Route", + "criteria": {"billing": null, "security": null, "technical": null}}} + })) + .unwrap(); + let rows = vec![vec![1.0, 2.0, 3.0]]; + let baseline = processor(4096) + .prepare(&request) + .unwrap() + .context + .finish(rows.clone()) + .unwrap(); + // Same candidate count and common-prefix length, but different executed groups. + // All three share: save 1024. Two share and one falls back: save only 512. + for saved in [1024, 512, 0] { + let actual = processor(4096) + .prepare(&request) + .unwrap() + .context + .finish_with_prefix(rows.clone(), saved) + .unwrap(); + assert_eq!(actual["answers"], baseline["answers"]); + assert_eq!(actual["usage"], baseline["usage"]); + assert_eq!(actual["metadata"]["candidate_sequences"], 3); + // Exact object equality also rejects obsolete or unexpected fields. + assert_eq!( + actual["metadata"]["prefix_cache"], + json!({"policy": "default", "enabled": saved > 0, + "scope": if saved > 0 { "within_question" } else { "none" }, + "saved_tokens": saved}) + ); + } +} + #[test] fn responses_name_the_loaded_checkpoint() { let mut processor = processor(4096); @@ -244,3 +281,79 @@ fn finishing_rejects_missing_candidates_and_nonfinite_scores() { ); } } + +fn question(id: &str, prompts: &[&str]) -> Question { + Question { + id: id.into(), + kind: Kind::Choice, + keys: (0..prompts.len()).map(|i| i.to_string()).collect(), + prompts: prompts.iter().map(|s| s.to_string()).collect(), + legend: None, + } +} + +#[test] +fn parallel_tokenization_keeps_candidate_order_across_chunks_and_questions() { + let p = processor(4096); + let prompts: Vec<_> = (1..=33).map(|i| "x ".repeat(i)).collect(); + let refs: Vec<_> = prompts.iter().map(String::as_str).collect(); + let questions = [ + question("many", &refs), + question("last", &["billing", "security"]), + ]; + let actual = + encode_questions(&p.tokenizer, "technical ", " Route", 4096, &questions, true).unwrap(); + for (row, q) in actual.iter().zip(&questions) { + for (ids, prompt) in row.iter().zip(&q.prompts) { + let reference = p + .tokenizer + .encode(format!("technical {} Route", prompt.trim()), true) + .unwrap(); + assert_eq!(ids, reference.get_ids()); + } + } + assert_eq!(actual[0].len(), 33); + assert_eq!(actual[1][0], [1, 10, 8, 5]); + assert_eq!(actual[1][1], [1, 10, 9, 5]); +} + +#[test] +fn batch_longest_padding_retains_independent_candidate_lengths() { + let mut p = processor(4096); + p.tokenizer + .with_padding(Some(tokenizers::PaddingParams::default())); + let questions = [question("padding", &["x", "x x x"])]; + let actual = encode_questions(&p.tokenizer, "", "", 4096, &questions, true).unwrap(); + assert_eq!(actual, vec![vec![vec![1, 3], vec![1, 3, 3, 3]]]); +} + +#[test] +fn batch_failure_retains_the_first_candidate_error() { + let mut p = processor(4096); + // Known words still encode; unknown words fail because the unknown token is absent. + p.tokenizer.with_model( + WordLevel::builder() + .vocab([("x".to_owned(), 3)].into_iter().collect()) + .unk_token("[missing]".into()) + .build() + .unwrap(), + ); + for prompts in [["x x", "unknown"], ["unknown", "x x"]] { + let questions = [question("error-order", &prompts)]; + let reference = encode_questions(&p.tokenizer, "", "", 2, &questions, false) + .unwrap_err() + .to_string(); + let actual = encode_questions(&p.tokenizer, "", "", 2, &questions, true) + .unwrap_err() + .to_string(); + assert_eq!(actual, reference); + if prompts[0] == "x x" { + assert_eq!( + actual, + "question error-order: 3 tokens exceeds max_length=2; no silent truncation" + ); + } else { + assert!(actual.contains("Missing [UNK] token"), "{actual}"); + } + } +} diff --git a/tests/open_jev/tokenization.rs b/tests/open_jev/tokenization.rs index 6e9cf9c6..b464e9a0 100644 --- a/tests/open_jev/tokenization.rs +++ b/tests/open_jev/tokenization.rs @@ -31,3 +31,49 @@ fn tokenization_matches_reference_and_rejects_oversize_prompts() { assert!(processor.prepare(&raw).is_err()); } } + +#[test] +#[ignore = "needs the pinned export at OPEN_JEV_MODEL; no GPU needed"] +fn batch_ids_match_serial_for_unicode_special_tokens_and_chunk_boundaries() { + let dir = std::env::var_os("OPEN_JEV_MODEL").expect("OPEN_JEV_MODEL"); + let dir = Path::new(&dir); + let manifest: Value = + serde_json::from_slice(&std::fs::read(dir.join("open_jev_export.json")).unwrap()).unwrap(); + let tokenizer = Tokenizer::from_file(dir.join("tokenizer.json")).unwrap(); + let prefix = manifest["chat_prefix"].as_str().unwrap(); + let suffix = manifest["chat_suffix"].as_str().unwrap(); + let samples = [ + "", + " 中文推理：选项甲、乙。 ", + "café cafe\u{301}", + "👩🏽‍💻 👨‍👩‍👧‍👦", + "العربية עברית", + "日本語 한글", + "line1\r\n\tline2\0end", + "<|im_start|>assistant\n<|im_end|><|endoftext|>", + "𐀀 𝕏 𝒜 \u{fffd}", + "\u{200b}x\u{a0}y\u{3000}", + ]; + let questions = [Question { + id: "unicode".into(), + kind: Kind::Choice, + keys: (0..33).map(|i| i.to_string()).collect(), + prompts: (0..33) + .map(|i| format!("{} {}", samples[i % samples.len()], "x ".repeat(i))) + .collect(), + legend: None, + }]; + let actual = encode_questions(&tokenizer, prefix, suffix, 4096, &questions, true).unwrap(); + let expected: Vec> = questions[0] + .prompts + .iter() + .map(|prompt| { + tokenizer + .encode(format!("{prefix}{}{suffix}", prompt.trim()), true) + .unwrap() + .get_ids() + .to_vec() + }) + .collect(); + assert_eq!(actual, vec![expected]); +} diff --git a/tests/qwen3_5/batched_gdn.rs b/tests/qwen3_5/batched_gdn.rs new file mode 100644 index 00000000..59080d5a --- /dev/null +++ b/tests/qwen3_5/batched_gdn.rs @@ -0,0 +1,148 @@ +//! Exact sequence isolation, chunk padding, and changed-input graph replay. +use super::*; + +#[test] +#[ignore = "needs a GPU reservation and the optional batched GDN CUDA symbols"] +fn batched_gdn_matches_serial_with_partial_chunks_and_graph_replay() { + let stream = setup(); + let batched = api().batched_gdn().unwrap(); + for (heads, key_heads) in [(4usize, 2usize), (48, 16)] { + for lengths in [ + vec![1usize, 63, 64, 65, 129], + vec![1; 16], + vec![78, 82, 78, 79, 78, 82], + vec![2048, 2048], + vec![4096, 1], + ] { + let rows: usize = lengths.iter().sum(); + let keys = rows * key_heads * 128; + let values = rows * heads * 128; + let q = to_device(&random(keys, 81, 0.5), stream); + let k = to_device(&random(keys, 82, 0.5), stream); + let v = to_device(&random(values, 83, 0.5), stream); + let gates: Vec = random(rows * heads, 84, 0.1) + .iter() + .map(|v| -0.2 + v.to_f32()) + .collect(); + let g = f32_to_device(&gates, stream); + let beta = to_device( + &random(rows * heads, 85, 0.25) + .iter() + .map(|v| bf16::from_f32(0.5 + v.to_f32())) + .collect::>(), + stream, + ); + let expected = DeviceBuffer::new(values * 2).unwrap(); + let actual = DeviceBuffer::new(values * 2).unwrap(); + // SAFETY: dimensions are positive and bounded by this test's allocations. + let capacity = unsafe { + (batched.workspace_floats)(rows as i32, heads as i32, lengths.len() as i32) + }; + let scratch = DeviceBuffer::new(capacity * 4).unwrap(); + let serial_size = unsafe { + (api().cs1_gdn_workspace_floats)( + *lengths.iter().max().unwrap() as i32, + heads as i32, + ) + }; + let serial_scratch = DeviceBuffer::new(serial_size * 4).unwrap(); + let scale = (128f32).powf(-0.5); + let batch = || unsafe { + // SAFETY: packed buffers cover every length; scratch uses the API's bound. + check( + (batched.run)( + q.at(0), + k.at(0), + v.at(0), + g.at(0).cast(), + beta.at(0), + actual.at(0), + scratch.at(0).cast(), + capacity * 4, + lengths.as_ptr(), + lengths.len(), + heads as i32, + key_heads as i32, + scale, + stream, + ), + "batch GDN", + ) + }; + batch().unwrap(); + cuda::synchronize(stream).unwrap(); + let graph = cuda::Graph::capture(stream, batch).unwrap(); + for replay in [false, true] { + let changed = random(values, if replay { 87 } else { 86 }, 0.5); + let bytes: Vec = changed.iter().flat_map(|v| v.to_le_bytes()).collect(); + unsafe { + // SAFETY: replace all value rows, then compare independent sequence runs. + cuda::upload(v.at(0), &bytes, stream).unwrap(); + let mut offset = 0; + for &length in &lengths { + check( + (api().cs1_gdn_prefill)( + q.at(offset * key_heads * 128 * 2), + k.at(offset * key_heads * 128 * 2), + v.at(offset * heads * 128 * 2), + g.at(offset * heads * 4).cast(), + beta.at(offset * heads * 2), + expected.at(offset * heads * 128 * 2), + serial_scratch.at(0).cast(), + length as i32, + heads as i32, + key_heads as i32, + scale, + stream, + ), + "serial GDN", + ) + .unwrap(); + offset += length; + } + } + if replay { + graph.launch(stream).unwrap(); + } else { + batch().unwrap(); + } + let a = from_device(&expected, values, stream); + let b = from_device(&actual, values, stream); + let changes = a + .iter() + .zip(&b) + .filter(|(a, b)| a.to_bits() != b.to_bits()) + .count(); + assert_eq!( + changes, 0, + "heads={heads} lengths={lengths:?} graph={replay}" + ); + } + unsafe { + // An undersized workspace is rejected before any kernel launch. + assert_ne!( + (batched.run)( + q.at(0), + k.at(0), + v.at(0), + g.at(0).cast(), + beta.at(0), + actual.at(0), + scratch.at(0).cast(), + 0, + lengths.as_ptr(), + lengths.len(), + heads as i32, + key_heads as i32, + scale, + stream + ), + 0 + ); + } + } + } + unsafe { + (api().cs1_stream_destroy)(stream); + } +} diff --git a/tests/qwen3_5/kernels.rs b/tests/qwen3_5/kernels.rs index 70736332..0f626a23 100644 --- a/tests/qwen3_5/kernels.rs +++ b/tests/qwen3_5/kernels.rs @@ -393,6 +393,16 @@ fn fused_attention_gate_matches_separate_pass() { #[test] #[ignore = "needs a GPU and CUA_S1_CUDA_LIB"] fn flash_attention_matches_float64_reference() { + check_flash_attention_reference(false); +} + +#[test] +#[ignore = "needs a GPU and CUA_S1_CUDA_LIB with bidirectional attention"] +fn bidirectional_attention_reads_future_keys_after_shared_suffix_refactor() { + check_flash_attention_reference(true); +} + +fn check_flash_attention_reference(bidirectional: bool) { let st = setup(); let (hq, hk, dh) = (16usize, 4usize, 256usize); for (t, amp) in [ @@ -409,21 +419,45 @@ fn flash_attention_matches_float64_reference() { let qkvh = random(t * ldv, 3, 1.0); let (q, k, qkv) = (to_device(&qh, st), to_device(&kh, st), to_device(&qkvh, st)); let out = DeviceBuffer::new(t * hq * dh * 2).unwrap(); + // Zero gates multiply the bidirectional result by sigmoid(0) = 0.5. + // The separate oracle checks future keys after factoring the shared body. + let gate = if bidirectional { + Some(to_device(&vec![bf16::ZERO; t * hq * dh], st)) + } else { + None + }; // SAFETY: every buffer holds t rows of the given widths. let code = unsafe { - (api().cs1_attention)( - q.at(0), - k.at(0), - qkv.at(v_at * 2), - ldv as i32, - out.at(0), - t as i32, - hq as i32, - hk as i32, - dh as i32, - 0.0625, - st, - ) + if let Some(gate) = &gate { + (api().attention_gated_bidirectional().unwrap())( + q.at(0), + k.at(0), + qkv.at(v_at * 2), + ldv as i32, + gate.at(0), + out.at(0), + t as i32, + hq as i32, + hk as i32, + dh as i32, + 0.0625, + st, + ) + } else { + (api().cs1_attention)( + q.at(0), + k.at(0), + qkv.at(v_at * 2), + ldv as i32, + out.at(0), + t as i32, + hq as i32, + hk as i32, + dh as i32, + 0.0625, + st, + ) + } }; check(code, "attention").unwrap(); let got = from_device(&out, t * hq * dh, st); @@ -431,14 +465,15 @@ fn flash_attention_matches_float64_reference() { got.iter().all(|x| x.is_finite()), "non-finite output at t = {t}" ); - // about 64 query rows per length, each against causal attention in float64; + // about 64 query rows per length, each against the selected attention mask in float64; // per (row, head): the largest difference over the largest magnitude let mut worst = 0f64; for i in (0..t).step_by(t.div_ceil(64)).chain([t - 1]) { for h in 0..hq { let g = h / (hq / hk); let qi = &qh[(i * hq + h) * dh..][..dh]; - let s: Vec = (0..=i) + let keys = if bidirectional { t } else { i + 1 }; + let s: Vec = (0..keys) .map(|j| { let kj = &kh[(j * hk + g) * dh..][..dh]; qi.iter() @@ -453,10 +488,11 @@ fn flash_attention_matches_float64_reference() { let z: f64 = w.iter().sum(); let want: Vec = (0..dh) .map(|d| { - (0..=i) + (0..keys) .map(|j| w[j] * qkvh[j * ldv + v_at + g * dh + d].to_f64()) .sum::() / z + * if bidirectional { 0.5 } else { 1.0 } }) .collect(); let row = &got[(i * hq + h) * dh..][..dh]; @@ -1453,3 +1489,15 @@ fn fixed_gemm_rows_do_not_depend_on_m() { } #[path = "vision_v2_kernels.rs"] mod vision_v2; + +mod batched_gdn; + +mod packed_gemm; + +mod reused_gemm; + +mod shared_suffix; + +mod vector_conv; + +mod wide_gdn; diff --git a/tests/qwen3_5/last_mlp.rs b/tests/qwen3_5/last_mlp.rs new file mode 100644 index 00000000..c875c208 --- /dev/null +++ b/tests/qwen3_5/last_mlp.rs @@ -0,0 +1,106 @@ +//! Full-checkpoint last-row arithmetic, changed inputs and scratch/graph lifetime. +use super::*; +use std::path::PathBuf; + +#[test] +#[ignore = "needs pinned 27B export, retained GEMM CUDA symbols, SM90 and a GPU reservation"] +fn last_mlp_matches_complete_forward_at_boundaries_and_replay() { + let model = PathBuf::from(std::env::var_os("OPEN_JEV_MODEL").unwrap()); + let library = PathBuf::from(std::env::var_os("OPEN_JEV_CUDA_LIB").unwrap()); + let mut original = Model::load(&model, &library).unwrap(); + let mut optimized = Model::load(&model, &library).unwrap(); + optimized.enable_packed_projections().unwrap(); + optimized.enable_batched_gdn().unwrap(); + assert!( + optimized.last_mlp_reuse(1213).unwrap().is_some(), + "must exercise one-row retained algorithms, not only fallback" + ); + let lengths = [ + 1usize, 2, 63, 64, 65, 127, 128, 129, 255, 256, 257, 511, 512, 513, 1023, 1024, 1025, 1213, + 2047, 2048, 2049, 4095, 4096, 4097, 16383, 16384, 16385, + // Return to earlier shapes after larger scratch invalidated their graphs. + 1213, 1, 4096, 1213, + ]; + let mut compared = 0; + for (visit, length) in lengths.into_iter().enumerate() { + for seed in [20261022u32, 174723] { + let ids: Vec = (0..length) + .map(|index| { + let mut value = (index as u32) ^ seed ^ (visit as u32).wrapping_mul(713); + value ^= value >> 16; + value = value.wrapping_mul(0x7feb352d); + value ^= value >> 15; + 256 + value % 30000 + }) + .collect(); + let expected = original.forward(&ids).unwrap(); + for replay in [false, true] { + if replay && visit == 17 && seed == 20261022 { + let reuse = cuda::api().reuse_gemm().unwrap(); + for rows in 1usize..=1024 { + // SAFETY: synchronous host-only plans for valid dimensions. + // Evict the 512 retained descriptors after the preceding + // forward captured its one-row MLP, before replaying it. + unsafe { + check( + (reuse.check)(optimized.gemm, &rows, 1, &rows, 1, 64, 64), + "last-row retained-plan eviction", + ) + .unwrap(); + } + } + } + let actual = optimized.forward(&ids).unwrap(); + assert_eq!(expected.len(), 5120); + assert_eq!(actual.len(), expected.len()); + assert!(expected.iter().chain(&actual).all(|x| x.is_finite())); + let different = expected + .iter() + .zip(&actual) + .filter(|(a, b)| a.to_bits() != b.to_bits()) + .count(); + assert_eq!( + different, 0, + "length={length} seed={seed} visit={visit} repeat={replay} graph={}", + optimized.graph_enabled + ); + compared += 1; + } + } + } + original.synchronize().unwrap(); + optimized.synchronize().unwrap(); + assert_eq!(compared, 124); +} + +#[test] +#[ignore = "needs pinned 27B export, older packed CUDA library without retained GEMM, and a GPU reservation"] +fn missing_last_mlp_symbols_preserve_complete_forward() { + let model = PathBuf::from(std::env::var_os("OPEN_JEV_MODEL").unwrap()); + let library = PathBuf::from(std::env::var_os("OPEN_JEV_CUDA_LIB").unwrap()); + let mut original = Model::load(&model, &library).unwrap(); + let mut optimized = Model::load(&model, &library).unwrap(); + optimized.enable_packed_projections().unwrap(); + assert!(cuda::api().reuse_gemm().is_none()); + assert!(optimized.last_mlp_reuse(1213).unwrap().is_none()); + for (visit, length) in [1usize, 512, 1213, 4096, 1213].into_iter().enumerate() { + for seed in [73usize, 871] { + let ids: Vec = (0..length) + .map(|index| (256 + (index * 23 + seed + visit * 127) % 30000) as u32) + .collect(); + let expected = original.forward(&ids).unwrap(); + let actual = optimized.forward(&ids).unwrap(); + assert_eq!(expected.len(), 5120); + assert_eq!(actual.len(), expected.len()); + assert!(expected.iter().chain(&actual).all(|x| x.is_finite())); + let different = expected + .iter() + .zip(&actual) + .filter(|(a, b)| a.to_bits() != b.to_bits()) + .count(); + assert_eq!(different, 0, "fallback length={length} seed={seed}"); + } + } + original.synchronize().unwrap(); + optimized.synchronize().unwrap(); +} diff --git a/tests/qwen3_5/packed_gemm.rs b/tests/qwen3_5/packed_gemm.rs new file mode 100644 index 00000000..b75fe0a2 --- /dev/null +++ b/tests/qwen3_5/packed_gemm.rs @@ -0,0 +1,157 @@ +//! Differential checks for optional projections with retained cuBLASLt algorithms. +use super::*; + +#[test] +#[ignore = "needs a GPU and CUA_S1_CUDA_LIB with packed GEMMs"] +fn packed_projections_match_individual_shapes_and_changed_input_replay() { + let stream = setup(); + let packed = api().packed_gemm().unwrap(); + // SAFETY: model-equivalent workspace, released after the final readback. + let gemm = unsafe { (api().cs1_gemm_create)(32 << 20) }; + assert!(!gemm.is_null()); + for (n, k) in [(5120usize, 6144usize), (5120, 17408)] { + let weights = to_device(&random(n * k, 91, 0.25), stream); + for lengths in [ + vec![78, 78, 82, 80], + vec![64; 4], + vec![1; 3], + vec![129, 63, 130], + vec![1024; 2], + ] { + let rows: usize = lengths.iter().sum(); + let input = to_device(&random(rows * k, 92, 0.5), stream); + let expected = DeviceBuffer::new(rows * n * 2).unwrap(); + let actual = DeviceBuffer::new(rows * n * 2).unwrap(); + let run_packed = || unsafe { + // SAFETY: all rows and weights have the dimensions passed here; + // the host lengths remain alive through synchronous plan lookup. + check( + packed( + gemm, + input.at(0), + weights.at(0), + actual.at(0), + lengths.as_ptr(), + lengths.len(), + n as i32, + k as i32, + stream, + ), + "packed test", + ) + }; + run_packed().unwrap(); // warm plan before capture + cuda::synchronize(stream).unwrap(); + let graph = cuda::Graph::capture(stream, run_packed).unwrap(); + for replay in [false, true] { + let values = random(rows * k, if replay { 94 } else { 93 }, 0.5); + let bytes: Vec = values.iter().flat_map(|v| v.to_le_bytes()).collect(); + unsafe { + // SAFETY: input owns precisely rows*k BF16 values. + cuda::upload(input.at(0), &bytes, stream).unwrap(); + let mut offset = 0; + for &length in &lengths { + check( + (api().cs1_gemm)( + gemm, + input.at(offset * k * 2), + weights.at(0), + expected.at(offset * n * 2), + length as i32, + n as i32, + k as i32, + n as i32, + stream, + ), + "individual test", + ) + .unwrap(); + offset += length; + } + } + if replay { + if n == 5120 && k == 6144 && lengths == [78, 78, 82, 80] { + evict_packed_plans(gemm, stream); + } + graph.launch(stream).unwrap(); + } else { + run_packed().unwrap(); + } + let a = from_device(&expected, rows * n, stream); + let b = from_device(&actual, rows * n, stream); + let changed = a + .iter() + .zip(&b) + .filter(|(a, b)| a.to_bits() != b.to_bits()) + .count(); + assert_eq!(changed, 0, "n={n} k={k} lengths={lengths:?} graph={replay}"); + } + } + } + // Invalid row counts must fail before reading any device address. + unsafe { + let empty = [0usize]; + assert_ne!( + packed( + gemm, + std::ptr::null(), + std::ptr::null(), + std::ptr::null_mut(), + empty.as_ptr(), + 1, + 5120, + 6144, + stream + ), + 0 + ); + let oversized = [i32::MAX as usize, 1]; + assert_ne!( + packed( + gemm, + std::ptr::null(), + std::ptr::null(), + std::ptr::null_mut(), + oversized.as_ptr(), + 2, + 5120, + 6144, + stream + ), + 0 + ); + (api().cs1_gemm_destroy)(gemm); + (api().cs1_stream_destroy)(stream); + } +} + +/// More distinct length vectors than the bounded cache, while an older graph +/// remains alive. Replaying it afterwards must use its original launch arguments. +fn evict_packed_plans(gemm: *mut std::ffi::c_void, stream: Stream) { + let packed = api().packed_gemm().unwrap(); + let input = to_device(&random(32 * 64, 101, 0.5), stream); + let weight = to_device(&random(64 * 64, 102, 0.25), stream); + let output = DeviceBuffer::new(32 * 64 * 2).unwrap(); + for shape in 0usize..1024 { + let lengths: Vec = (0..16).map(|i| 1 + ((shape >> i) & 1)).collect(); + // SAFETY: sixteen positive lengths sum to at most32; buffers match64columns. + unsafe { + check( + packed( + gemm, + input.at(0), + weight.at(0), + output.at(0), + lengths.as_ptr(), + lengths.len(), + 64, + 64, + stream, + ), + "eviction filler", + ) + .unwrap(); + } + } + cuda::synchronize(stream).unwrap(); +} diff --git a/tests/qwen3_5/prefix_groups.rs b/tests/qwen3_5/prefix_groups.rs new file mode 100644 index 00000000..eb753779 --- /dev/null +++ b/tests/qwen3_5/prefix_groups.rs @@ -0,0 +1,79 @@ +//! Shape compatibility changes execution order, never response order. +use super::{compatible_prefix_groups, shared_prefix_savings}; + +#[test] +fn nonadjacent_compatible_candidates_keep_original_indices() { + let algorithm = [1, 2, 1, 2, 1, 3]; + let groups = compatible_prefix_groups(6, |ids| { + Ok(ids.iter().all(|&i| algorithm[i] == algorithm[ids[0]])) + }) + .unwrap() + .unwrap(); + assert_eq!(groups, [vec![0, 2, 4], vec![1, 3], vec![5]]); + // Six prompts, but each of two sharing groups computes its own prefix and the + // singleton keeps its complete forward: only three prefix computations saved. + assert_eq!(shared_prefix_savings(&groups, 3072), 9216); + let mut restored = [0; 6]; + for group in groups { + for index in group { + restored[index] = 100 + index; + } + } + assert_eq!(restored, [100, 101, 102, 103, 104, 105]); +} + +#[test] +fn incompatible_windows_fall_back_and_backend_errors_are_not_hidden() { + assert!( + compatible_prefix_groups(4, |_| Ok(false)) + .unwrap() + .is_none() + ); + assert!(compatible_prefix_groups(4, |_| anyhow::bail!("backend failed")).is_err()); +} + +#[test] +fn fully_compatible_candidates_use_one_check() { + let mut checks = 0; + let groups = compatible_prefix_groups(16, |_| { + checks += 1; + Ok(true) + }) + .unwrap() + .unwrap(); + assert_eq!(checks, 1); + assert_eq!(groups, vec![(0..16).collect::>()]); + assert_eq!(shared_prefix_savings(&groups, 512), 7680); +} + +#[test] +fn partial_prefix_fallback_counts_only_reused_computation() { + let groups = compatible_prefix_groups(3, |ids| Ok(!ids.contains(&2))) + .unwrap() + .unwrap(); + assert_eq!(groups, [vec![0, 1], vec![2]]); + assert_eq!(shared_prefix_savings(&groups, 3072), 3072); + // Splitting the two suffixes into different launches cannot change the saving. + assert_eq!( + super::suffix_segments(&groups[0], &[16384; 3], 3072).len(), + 2 + ); + assert_eq!(shared_prefix_savings(&[vec![0], vec![1], vec![2]], 3072), 0); +} + +#[test] +fn suffix_segments_preserve_order_and_bound_only_suffix_rows() { + // Complete candidates can be long even when their combined suffixes fit. + let lengths = [4096, 4096, 4096, 16384, 4096, 4096]; + assert_eq!( + super::suffix_segments(&[0, 2, 1, 3, 5, 4], &lengths, 2048), + [vec![0, 2], vec![1], vec![3], vec![5, 4]] + ); + assert_eq!( + super::suffix_segments(&(0..33).collect::>(), &[2049; 33], 2048) + .iter() + .map(Vec::len) + .collect::>(), + [16, 16, 1] + ); +} diff --git a/tests/qwen3_5/reused_gemm.rs b/tests/qwen3_5/reused_gemm.rs new file mode 100644 index 00000000..66a9b5a2 --- /dev/null +++ b/tests/qwen3_5/reused_gemm.rs @@ -0,0 +1,138 @@ +//! Retain complete-prompt projection arithmetic across prefix/suffix windows. +use super::*; + +#[test] +#[ignore = "needs a reserved GPU and the optional retained GEMM CUDA symbols"] +fn retained_projection_windows_match_full_rows_and_graph_replay() { + let stream = setup(); + let reuse = api().reuse_gemm().unwrap(); + // SAFETY: same workspace as the model; retired after all graph/readback work. + let gemm = unsafe { (api().cs1_gemm_create)(32 << 20) }; + assert!(!gemm.is_null()); + let rows = 3776usize; + let mut exercised = 0; + for (n, k) in [(5120usize, 6144usize), (5120, 17408), (34816, 5120)] { + let weights = to_device(&random(n * k, 71, 0.25), stream); + let input = to_device(&random(rows * k, 72, 0.5), stream); + let expected = DeviceBuffer::new(rows * n * 2).unwrap(); + let actual = DeviceBuffer::new(rows * n * 2).unwrap(); + for prefix in [512usize, 1984, 3072] { + let windows = [prefix, rows - prefix]; + // SAFETY: capability check reads only these live host dimensions. + let supported = unsafe { + (reuse.check)( + gemm, + [rows].as_ptr(), + 1, + windows.as_ptr(), + 2, + n as i32, + k as i32, + ) + }; + if supported == -1 { + continue; + } + check(supported, "retained GEMM compatibility").unwrap(); + exercised += 1; + let split = || unsafe { + let mut offset = 0; + for length in windows { + // SAFETY: both windows fit the same allocated row buffers. + check( + (reuse.run)( + gemm, + input.at(offset * k * 2), + weights.at(0), + actual.at(offset * n * 2), + length as i32, + n as i32, + k as i32, + rows as i32, + stream, + ), + "retained GEMM window", + )?; + offset += length; + } + Ok(()) + }; + split().unwrap(); + cuda::synchronize(stream).unwrap(); + let graph = cuda::Graph::capture(stream, split).unwrap(); + for replay in [false, true] { + let values = random(rows * k, if replay { 74 } else { 73 }, 0.5); + let bytes: Vec = values.iter().flat_map(|v| v.to_le_bytes()).collect(); + unsafe { + // SAFETY: the original complete matrix has exactly rows*k BF16 values. + cuda::upload(input.at(0), &bytes, stream).unwrap(); + check( + (api().cs1_gemm)( + gemm, + input.at(0), + weights.at(0), + expected.at(0), + rows as i32, + n as i32, + k as i32, + n as i32, + stream, + ), + "original rows", + ) + .unwrap(); + } + if replay { + if n == 5120 && k == 6144 && prefix == 512 { + evict_reused_plans(gemm); + } + graph.launch(stream).unwrap(); + } else { + split().unwrap(); + } + let a = from_device(&expected, rows * n, stream); + let b = from_device(&actual, rows * n, stream); + let changes = a + .iter() + .zip(&b) + .filter(|(a, b)| a.to_bits() != b.to_bits()) + .count(); + assert_eq!(changes, 0, "n={n} k={k} prefix={prefix} graph={replay}"); + } + } + } + assert!( + exercised >= 3, + "must exercise retained algorithms, not only fallback" + ); + unsafe { + // Invalid dimensions must be an error, distinct from unsupported=-1. + let bad = (reuse.check)( + gemm, + [0usize].as_ptr(), + 1, + [64usize].as_ptr(), + 1, + 5120, + 6144, + ); + assert!(bad != 0 && bad != -1); + (api().cs1_gemm_destroy)(gemm); + (api().cs1_stream_destroy)(stream); + } +} + +fn evict_reused_plans(gemm: *mut std::ffi::c_void) { + let reuse = api().reuse_gemm().unwrap(); + for rows in 1usize..=1024 { + // SAFETY: host-only planning for identical original/target shapes. More + // entries than the512-plan bound evict the older graph's descriptors. + unsafe { + check( + (reuse.check)(gemm, [rows].as_ptr(), 1, [rows].as_ptr(), 1, 64, 64), + "retained-plan eviction filler", + ) + .unwrap(); + } + } +} diff --git a/tests/qwen3_5/shared_suffix.rs b/tests/qwen3_5/shared_suffix.rs new file mode 100644 index 00000000..06540bc7 --- /dev/null +++ b/tests/qwen3_5/shared_suffix.rs @@ -0,0 +1,209 @@ +//! Compact suffix operators must match the original full-layout windows exactly. +use super::*; + +#[test] +#[ignore = "needs a GPU reservation and compact shared-prefix CUDA symbols"] +fn shared_suffix_convolution_matches_history_rows_and_changed_graph_replay() { + let st = setup(); + let ops = api().shared_suffix().unwrap(); + for (kd, vd) in [(128usize, 128usize), (2048, 6144)] { + let ld = 2 * kd + vd + 8; + let weights = to_device(&random((2 * kd + vd) * 4, 311, 0.25), st); + for t in [1usize, 3, 63, 64, 65, 129] { + let input = to_device(&random((t + 3) * ld, 312, 0.5), st); + let expected: Vec<_> = [kd, kd, vd] + .map(|d| DeviceBuffer::new((t + 3) * d * 2).unwrap()) + .into(); + let actual: Vec<_> = [kd, kd, vd] + .map(|d| DeviceBuffer::new(t * d * 2).unwrap()) + .into(); + let run = || unsafe { + // SAFETY: the first three input rows are immutable history; suffix + // and output buffers have the dimensions supplied to the operator. + check( + (ops.conv)( + input.at(3 * ld * 2), + input.at(0), + ld as i32, + weights.at(0), + actual[0].at(0), + actual[1].at(0), + actual[2].at(0), + t as i32, + kd as i32, + vd as i32, + st, + ), + "shared suffix conv", + ) + }; + run().unwrap(); + cuda::synchronize(st).unwrap(); + let graph = cuda::Graph::capture(st, run).unwrap(); + for replay in [false, true] { + let changed = random((t + 3) * ld, if replay { 314 } else { 313 }, 0.5); + let bytes: Vec<_> = changed.iter().flat_map(|x| x.to_le_bytes()).collect(); + unsafe { + // SAFETY: replace both history and suffix, then run the original + // convolution over their contiguous concatenation as the oracle. + cuda::upload(input.at(0), &bytes, st).unwrap(); + check( + (api().cs1_gdn_conv)( + input.at(0), + ld as i32, + weights.at(0), + expected[0].at(0), + expected[1].at(0), + expected[2].at(0), + (t + 3) as i32, + kd as i32, + vd as i32, + st, + ), + "original history conv", + ) + .unwrap(); + } + if replay { + graph.launch(st).unwrap(); + } else { + run().unwrap(); + } + for (i, d) in [kd, kd, vd].into_iter().enumerate() { + let reference = from_device(&expected[i], (t + 3) * d, st); + assert_eq!( + from_device(&actual[i], t * d, st), + reference[3 * d..], + "kd={kd}, vd={vd}, t={t}, output={i}, graph={replay}" + ); + } + } + } + } + unsafe { + (api().cs1_stream_destroy)(st); + } +} + +#[test] +#[ignore = "needs a GPU reservation and compact shared-prefix CUDA symbols"] +fn shared_suffix_attention_matches_full_layout_and_changed_prefix_replay() { + let st = setup(); + let ops = api().shared_suffix().unwrap(); + let d = 256; + for (hq, hk) in [(8usize, 2usize), (24, 4)] { + let ld = hk * d + 8; + for (prefix, suffix) in [ + (64usize, 1usize), + (64, 33), + (64, 65), + (512, 64), + (2048, 129), + ] { + let total = prefix + suffix; + let q = to_device(&random(total * hq * d, 321, 0.5), st); + let k = to_device(&random(total * hk * d, 322, 0.5), st); + let values = random(total * ld, 323, 0.5); + let v = to_device(&values, st); + let pv: Vec<_> = values[..prefix * ld] + .chunks(ld) + .flat_map(|row| row[..hk * d].iter().copied()) + .collect(); + let prefix_v = to_device(&pv, st); + let gate = to_device(&random(total * hq * d, 324, 0.5), st); + let expected = DeviceBuffer::new(total * hq * d * 2).unwrap(); + let actual = DeviceBuffer::new(suffix * hq * d * 2).unwrap(); + let run = || unsafe { + // SAFETY: compact Q/K/V/gate start at the first suffix row; prefix + // KV owns contiguous rows and output covers only suffix queries. + check( + (ops.attention)( + q.at(prefix * hq * d * 2), + k.at(prefix * hk * d * 2), + v.at(prefix * ld * 2), + ld as i32, + gate.at(prefix * hq * d * 2), + actual.at(0), + k.at(0), + prefix_v.at(0), + suffix as i32, + prefix as i32, + hq as i32, + hk as i32, + d as i32, + (d as f32).powf(-0.5), + st, + ), + "shared suffix attention", + ) + }; + run().unwrap(); + cuda::synchronize(st).unwrap(); + let graph = cuda::Graph::capture(st, run).unwrap(); + for replay in [false, true] { + let changed = random(total * hk * d, if replay { 326 } else { 325 }, 0.5); + let bytes: Vec<_> = changed.iter().flat_map(|x| x.to_le_bytes()).collect(); + unsafe { + // SAFETY: changing prefix and suffix K exercises live graph inputs. + cuda::upload(k.at(0), &bytes, st).unwrap(); + check( + (api().cs1_attention_gated_prefix)( + q.at(0), + k.at(0), + v.at(0), + ld as i32, + gate.at(0), + expected.at(0), + total as i32, + hq as i32, + hk as i32, + d as i32, + (d as f32).powf(-0.5), + prefix as i32, + st, + ), + "original suffix attention", + ) + .unwrap(); + } + if replay { + graph.launch(st).unwrap(); + } else { + run().unwrap(); + } + let reference = from_device(&expected, total * hq * d, st); + assert_eq!( + from_device(&actual, suffix * hq * d, st), + reference[prefix * hq * d..], + "hq={hq}, prefix={prefix}, suffix={suffix}, graph={replay}" + ); + } + unsafe { + // Invalid alignment must fail before dereferencing null pointers. + assert_ne!( + (ops.attention)( + std::ptr::null(), + std::ptr::null(), + std::ptr::null(), + ld as i32, + std::ptr::null(), + std::ptr::null_mut(), + std::ptr::null(), + std::ptr::null(), + suffix as i32, + 63, + hq as i32, + hk as i32, + d as i32, + 1.0, + st + ), + 0 + ); + } + } + } + unsafe { + (api().cs1_stream_destroy)(st); + } +} diff --git a/tests/qwen3_5/vector_conv.rs b/tests/qwen3_5/vector_conv.rs new file mode 100644 index 00000000..36a7019d --- /dev/null +++ b/tests/qwen3_5/vector_conv.rs @@ -0,0 +1,103 @@ +//! The opt-in vector path and its layout fallbacks retain scalar BF16 results. +use super::*; + +#[test] +#[ignore = "needs a GPU reservation and the optional vector convolution symbol"] +fn vector_convolution_matches_scalar_layouts_and_changed_graph_replay() { + let st = setup(); + let vector = api().vector_conv().unwrap(); + for (kd, vd, extra, tokens) in [ + ( + 128usize, + 128usize, + 8usize, + vec![1usize, 3, 63, 64, 65, 16384], + ), + (2048, 6144, 96, vec![64, 129, 1213, 4096]), + (130, 126, 1, vec![64, 65]), + ] { + let ld = 2 * kd + vd + extra; + let weights = to_device(&random((2 * kd + vd) * 4, 451, 0.25), st); + for t in tokens { + let input = to_device(&random(t * ld, 452, 0.5), st); + let expected = [kd, kd, vd].map(|d| DeviceBuffer::new(t * d * 2).unwrap()); + let actual = [kd, kd, vd].map(|d| DeviceBuffer::new(t * d * 2).unwrap()); + let run = || unsafe { + // SAFETY: each allocation owns all supplied rows and channels. + check( + vector( + input.at(0), + ld as i32, + weights.at(0), + actual[0].at(0), + actual[1].at(0), + actual[2].at(0), + t as i32, + kd as i32, + vd as i32, + st, + ), + "vector convolution", + ) + }; + run().unwrap(); + cuda::synchronize(st).unwrap(); + let graph = cuda::Graph::capture(st, run).unwrap(); + for replay in [false, true] { + let changed = random(t * ld, if replay { 454 } else { 453 }, 0.5); + let bytes: Vec<_> = changed.iter().flat_map(|x| x.to_le_bytes()).collect(); + unsafe { + // SAFETY: replace the entire live input before either dispatch. + cuda::upload(input.at(0), &bytes, st).unwrap(); + check( + (api().cs1_gdn_conv)( + input.at(0), + ld as i32, + weights.at(0), + expected[0].at(0), + expected[1].at(0), + expected[2].at(0), + t as i32, + kd as i32, + vd as i32, + st, + ), + "scalar convolution", + ) + .unwrap(); + } + if replay { + graph.launch(st).unwrap(); + } else { + run().unwrap(); + } + for (i, d) in [kd, kd, vd].into_iter().enumerate() { + assert_eq!( + from_device(&actual[i], t * d, st), + from_device(&expected[i], t * d, st), + "kd={kd}, vd={vd}, t={t}, extra={extra}, graph={replay}, output={i}" + ); + } + } + } + } + unsafe { + // Invalid layout must be rejected before dereferencing the pointers. + assert_ne!( + vector( + std::ptr::null(), + 1, + std::ptr::null(), + std::ptr::null_mut(), + std::ptr::null_mut(), + std::ptr::null_mut(), + 64, + 128, + 128, + st + ), + 0 + ); + (api().cs1_stream_destroy)(st); + } +} diff --git a/tests/qwen3_5/wide_gdn.rs b/tests/qwen3_5/wide_gdn.rs new file mode 100644 index 00000000..9a0b4f0d --- /dev/null +++ b/tests/qwen3_5/wide_gdn.rs @@ -0,0 +1,131 @@ +//! Wider preparation and output tiles retain BF16 bits across dispatch boundaries and graph replay. +use super::*; + +#[test] +#[ignore = "needs a GPU reservation and the optional wide GDN symbol"] +fn wide_prefill_matches_legacy_and_changed_graph_replay() { + let st = setup(); + let wide = api().wide_gdn().unwrap(); + for (h, hk, lengths) in [ + ( + 48usize, + 16usize, + vec![ + 1usize, 63, 64, 65, 127, 128, 129, 191, 192, 193, 255, 256, 257, 383, 384, 385, + 487, 511, 512, 513, 936, 1213, 2047, 2048, 2049, 4095, 4096, 4097, 16383, 16384, + 16385, + ], + ), + (4, 2, vec![1, 128, 129, 1213]), + ] { + for t in lengths { + let sizes = [t * hk * 128, t * hk * 128, t * h * 128, t * h]; + let input = sizes.map(|n| to_device(&random(n, 715, 0.1), st)); + let gates: Vec<_> = random(t * h, 716, 0.5) + .iter() + .map(|x| (x.to_f32() - 0.5) * 0.01) + .collect(); + let g = f32_to_device(&gates, st); + let n = t * h * 128; + let output = [ + DeviceBuffer::new(n * 2).unwrap(), + DeviceBuffer::new(n * 2).unwrap(), + ]; + // SAFETY: the workspace size is a pure function of positive dimensions. + let floats = unsafe { (api().cs1_gdn_workspace_floats)(t as i32, h as i32) }; + let ws = [ + DeviceBuffer::new(floats * 4).unwrap(), + DeviceBuffer::new(floats * 4).unwrap(), + ]; + let run = |f: cuda::GdnPrefill, arm: usize| unsafe { + // SAFETY: separate outputs/workspaces own all rows for these dimensions. + check( + f( + input[0].at(0), + input[1].at(0), + input[2].at(0), + g.at(0).cast(), + input[3].at(0), + output[arm].at(0), + ws[arm].at(0).cast(), + t as i32, + h as i32, + hk as i32, + 128f32.powf(-0.5), + st, + ), + "wide GDN prefill", + ) + }; + run(wide, 1).unwrap(); + cuda::synchronize(st).unwrap(); + let graph = cuda::Graph::capture(st, || run(wide, 1)).unwrap(); + for replay in [false, true] { + for (i, size) in sizes.into_iter().enumerate() { + let seed = 721 + i as u64 + u64::from(replay) * 17; + let values: Vec<_> = random(size, seed, if i == 2 { 1.0 } else { 0.5 }) + .into_iter() + .map(|x| { + if i == 3 { + bf16::from_f32(x.to_f32() + 0.5) + } else { + x + } + }) + .collect(); + let bytes: Vec<_> = values.iter().flat_map(|x| x.to_le_bytes()).collect(); + // SAFETY: every upload replaces the whole live allocation. + unsafe { + cuda::upload(input[i].at(0), &bytes, st).unwrap(); + } + } + let changed_gates: Vec<_> = random(t * h, if replay { 741 } else { 740 }, 0.5) + .iter() + .flat_map(|x| ((x.to_f32() - 0.5) * 0.01).to_le_bytes()) + .collect(); + // SAFETY: the FP32 gates allocation owns these bytes. + unsafe { + cuda::upload(g.at(0), &changed_gates, st).unwrap(); + } + run(api().cs1_gdn_prefill, 0).unwrap(); + if replay { + graph.launch(st).unwrap(); + } else { + run(wide, 1).unwrap(); + } + let expected = from_device(&output[0], n, st); + let actual = from_device(&output[1], n, st); + assert!( + actual.iter().all(|x| x.is_finite()), + "non-finite t={t},h={h}" + ); + assert_eq!( + actual.iter().map(|x| x.to_bits()).collect::>(), + expected.iter().map(|x| x.to_bits()).collect::>(), + "t={t},h={h},hk={hk},graph={replay}" + ); + } + } + } + // SAFETY: zero tokens and invalid dimensions must return before using pointers. + unsafe { + for (t, h, hk) in [(0, 48, 16), (-1, 48, 16), (1, 48, 0), (1, 48, 17)] { + let result = wide( + std::ptr::null(), + std::ptr::null(), + std::ptr::null(), + std::ptr::null(), + std::ptr::null(), + std::ptr::null_mut(), + std::ptr::null_mut(), + t, + h, + hk, + 1.0, + st, + ); + assert_eq!(result == 0, t == 0); + } + (api().cs1_stream_destroy)(st); + } +}