From 43ccbbcff99a94d5c6f744183b82e37853ed873c Mon Sep 17 00:00:00 2001 From: psv666 <150513104+psv666@users.noreply.github.com> Date: Wed, 28 Oct 2026 10:02:12 +0800 Subject: [PATCH 1/3] Optimize Open-Jev-9B prefill while preserving native outputs Snapshot of measured local head 95c0e70f2cdf2a72df763d0247ba7f20f8a9327d. Exact source tree: 044617e2b2370f44d8d7448b6feba1c86768c447. Includes accepted request-local prefix, batching, attention, GEMM and norm changes. Multi-executor diagnostic prototypes and local experiment artifacts are excluded. --- docs/architecture.md | 40 +- recipe/open_jev/native.md | 91 +- src/backends/cuda/qwen3_5/attention.cu | 85 +- src/backends/cuda/qwen3_5/gdn_prefill.cu | 19 +- src/backends/cuda/qwen3_5/gemm.cu | 83 +- src/backends/cuda/qwen3_5/norm.cu | 47 + src/backends/cuda/qwen3_5/ops.h | 25 +- src/models/open_jev/README.md | 27 + src/models/open_jev/native/src/batching.rs | 31 +- src/models/open_jev/native/src/contract.rs | 14 +- src/models/open_jev/native/src/engine.rs | 37 +- src/models/open_jev/native/src/executor.rs | 72 +- .../open_jev/native/src/executor/batched.rs | 212 +++++ src/models/qwen3_5/native/src/cuda.rs | 100 +- src/models/qwen3_5/native/src/model.rs | 890 ++++++++++++------ .../qwen3_5/native/src/prefix_compaction.rs | 339 +++++++ src/models/qwen3_5/native/src/prefix_tree.rs | 479 ++++++++++ src/runtime/Cargo.toml | 4 + src/runtime/README.md | 50 +- src/runtime/src/batching.rs | 220 +++++ src/runtime/src/lib.rs | 1 + tests/open_jev/batching.rs | 30 + tests/open_jev/contract.rs | 6 +- tests/open_jev/prefill_batch.rs | 284 ++++++ tests/open_jev/prefix.rs | 20 +- tests/qwen3_5/actual_row_gemm.rs | 106 +++ tests/qwen3_5/batched_attention.rs | 266 ++++++ tests/qwen3_5/batched_gdn.rs | 117 ++- tests/qwen3_5/capture_only.rs | 76 ++ tests/qwen3_5/equivalent_tiles.rs | 347 +++++++ tests/qwen3_5/kernels.rs | 94 +- tests/qwen3_5/multimodal_graph.rs | 2 +- tests/qwen3_5/prefix_tree.rs | 94 ++ tests/qwen3_5/tree_graph.rs | 165 ++++ tests/runtime/batching.rs | 314 ++++++ 35 files changed, 4427 insertions(+), 360 deletions(-) create mode 100644 src/models/open_jev/native/src/executor/batched.rs create mode 100644 src/models/qwen3_5/native/src/prefix_compaction.rs create mode 100644 src/models/qwen3_5/native/src/prefix_tree.rs create mode 100644 src/runtime/src/batching.rs create mode 100644 tests/qwen3_5/actual_row_gemm.rs create mode 100644 tests/qwen3_5/batched_attention.rs create mode 100644 tests/qwen3_5/capture_only.rs create mode 100644 tests/qwen3_5/equivalent_tiles.rs create mode 100644 tests/qwen3_5/prefix_tree.rs create mode 100644 tests/qwen3_5/tree_graph.rs create mode 100644 tests/runtime/batching.rs diff --git a/docs/architecture.md b/docs/architecture.md index 380e08d5..8da118de 100644 --- a/docs/architecture.md +++ b/docs/architecture.md @@ -12,12 +12,19 @@ to separately running workers. Cua-S1, Open-Jev, Decider, and StartLux Instant [Qwen3.5/3.8 executor](../src/models/qwen3_5/native/), which accepts 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 +it does not pack prompts or cache prefixes. Cua-S1 uses single-prompt calls. Open-Jev-9B and Open-Jev-27B-v1.1 +pack candidates within one request for 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 +For compatible 9B short prefixes, pointwise rows can be compacted while mixers +retain the complete candidate layout. The model owns expansion/gather maps and +stream lifetime; the executor retains request and candidate ordering. Exact9B +prefix capture explicitly requests state-only output; after final-layer K/V is +captured, unconsumed prefix hidden outputs are omitted. Hidden-row readouts +continue to request the complete window computation. +On 27B (and 9B with `OPEN_JEV_EXACT_PREFIX=1`), 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. @@ -33,8 +40,10 @@ coordinate independent processors and executors through [native runtime](../src/runtime/README.md) for FIFO admission and blocking dispatch per loaded executor. The [Decider native worker](../recipe/decider/README.md) also uses this runtime and can opt into outstanding-request and processed-token -budgets plus HTTP request deadlines. Shared processing orchestration, -cross-request batch assembly, compatibility grouping and dynamic batching are planned. +budgets plus HTTP request deadlines. Shared processing orchestration +and cross-worker coordination remain planned. Open-Jev-9B default execution uses +the runtime's weighted FIFO `BatchScheduler` to combine original candidate batches; +exact-prefix requests remain exclusive execution items. The pplx-decider reference worker runs the checkpoint's own loader instead of this shared Qwen executor: the checkpoint's full-attention layers are noncausal, and the language kernels in `src/backends/cuda/qwen3_5/attention.cu` @@ -116,7 +125,8 @@ its prefix rounded up to 1,024 positions, and the GEMM handle a 32 MiB workspace ## Native worker boundaries The native workers separate `processing.rs` from `executor.rs`; their worker -assembly owns a `SerialScheduler` per loaded executor, and the HTTP handler +assembly owns a scheduler per loaded executor (`BatchScheduler` for 9B default +execution, `SerialScheduler` for the other paths), and the HTTP handler coordinates the three stages. Preparation validates the entire request before any forward call and returns executor inputs plus a response context. The context retains question and candidate identity, usage, and response metadata outside the executor. @@ -131,13 +141,16 @@ and candidate identity, usage, and response metadata outside the executor. | StartLux Instant 4B | One unpadded token-ID prompt per question, or per knockout group and then a final prompt, in request order. | FP32 letter logits from a CPU dot of the tied embedding rows at the last position. | Per-type temperature, none-check split, wide-choice residual, and ordered `choice`, `noul`, and `score` answers with token usage. | | Clef | The raw `systemone` body. Preparation stays in the pinned Python encoder. | One joint-schema forward on the executor process main thread. The Rust closure only exchanges a frame. | The checkpoint's `systemone_answer` object, forwarded with its HTTP status unchanged. | -Cua-S1 and Open-Jev-9B input collections are serial work. For Open-Jev-27B-v1.1, +Cua-S1 input collections are serial work. For both Open-Jev checkpoints, Open-Jev's model-specific batch adapter packs up to 16 independent candidates and 4096 tokens per group; longer prompts -execute alone. It restores question/candidate grouping before normalization. -Laya batches questions within one request. Shared runtime +execute alone. For 9B default numerics, the runtime can combine these original +batches across requests up to 16 candidates and 8192 tokens. Input/gate-up GEMMs +retain original batch row counts; output/down GEMMs retain candidate lengths. +Request-local prefixes run as exclusive items. It restores question/candidate +grouping before normalization. Laya batches questions within one request. Shared runtime admission precedes blocking dispatch: Cua-S1 admits one question forward at a -time; Open-Jev, Decider, Laya, StartLux Instant, and Clef admit one complete request. Clef's admitted work is the frame exchange with its executor process. StartLux runs each prompt inside that admission one at a time. Cua-S1's CPU letter projection +time; serial Open-Jev modes, Decider, Laya, StartLux Instant, and Clef admit one complete request. Clef's admitted work is the frame exchange with its executor process. StartLux runs each prompt inside that admission one at a time. Cua-S1's CPU letter projection stays outside admission; Open-Jev's scalar heads and Laya's GPU heads and synchronized readback remain inside their request unit. Decider includes every backbone row and CUDA head projection, synchronizes both streams and retires @@ -224,6 +237,13 @@ cost; shared runtime primitives own capacity accounting, cancellation/deadline checks and FIFO execution admission. These request budgets do not add cross-request batching or replace the scheduler's single execution permit. +The implemented weighted batch scheduler similarly owns one FIFO execution queue +per loaded executor. It combines already queued atomic items without waiting for +more arrivals, enforces model-declared costs and exclusive items, and routes outputs +back to the requesting futures. Queued cancellation skips a request before its +first dispatch; afterward its complete work and resources are retained. A callback +must retire device work on failure or panic before releasing that ownership. + - Batch only work accepted by the same loaded executor, with compatible checkpoint/adapter identity, device, dtype, and input layout. Models declare additional constraints and limits; the scheduler owns queue and admission diff --git a/recipe/open_jev/native.md b/recipe/open_jev/native.md index 5e85a52e..57c39f47 100644 --- a/recipe/open_jev/native.md +++ b/recipe/open_jev/native.md @@ -69,7 +69,7 @@ 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 +Both checkpoints also require the additive packed-GEMM and batched-GDN symbols from this checkout; rebuild the worker and library together. ```sh @@ -159,7 +159,7 @@ uses the general kernel. MLP SiLU uses both BF16 rounding points; other layouts use the scalar path. 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 +both Open-Jev checkpoints, 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 combine consecutive candidates @@ -175,9 +175,10 @@ records latency, exact output checks and frozen controls. 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. +Open-Jev-9B packing keeps prefix reuse disabled by default. Explicit +`OPEN_JEV_PREFIX` modes retain their existing separate execution paths. -For one question with 2–16 candidates all longer than 2048 tokens, the worker can +On 27B, 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 @@ -242,7 +243,7 @@ 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 -27B packed text path retains at most 256 graphs; other modes retain 64. Keys are +packed Open-Jev 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 @@ -358,3 +359,83 @@ OPEN_JEV_MODEL=$PWD/weights/open-jev-9b-merged \ OPEN_JEV_CUDA_LIB=$PWD/target/release/libqwen3_5_cuda.so \ cargo test --release --locked -p omni-open-jev-native --test prefix -- --ignored ``` + +### Exact request-local prefix reuse on 9B + +Set `OPEN_JEV_EXACT_PREFIX=1` with `OPEN_JEV_PREFIX=0` to opt into the +retained-algorithm prefix path already used by 27B. This computes an identical +long prefix once within a single question, then evaluates independent candidate +suffixes. It preserves the complete prompts' GEMM algorithms and falls back to +ordinary packing when shapes are incompatible. All numeric prefix state is +recomputed for every request; no historical request or answer is cached. + +The current eligibility rule requires 2–16 candidates, each longer than 2048 +tokens, and a common prefix of at least 512 tokens rounded down to a 64-token +boundary. Input token usage still counts the complete candidate prompts. +`metadata.prefix_cache` reports actual reuse and saved computation. The switch +cannot be combined with fixed/shared/auto modes, which select a different +algorithm policy. `OPEN_JEV_EXACT_PREFIX=0` explicitly disables this exact path; +unset retains each checkpoint's default (9B off, 27B on). + +The explicit `Engine::load_with_prefix` API retains checkpoint defaults and does +not read this additional environment override. With the 9B export and the new +switch enabled, run the `text_prefix_buffers_grow_and_shrink_without_reusing_prior_contents` +GPU test in `prefill_batch` for exact output equality and prefix-state lifetime. + +Both pinned Open-Jev tokenizers use bounded batch encoding for independent +candidates (up to 16 at a time). This avoids computing unused character offsets +and allows the tokenizer's worker pool to encode candidates in parallel. Padded +tokenizers retain serial encoding, and any batch error is retried in original +order so validation reports the same first error. The exact-token regression +includes Unicode, special tokens, batch boundaries and oversized prompts for 9B. + +On SM90 with cuBLASLt 13.1.1, the 9B worker can combine 512–4096-row +projections using a tested engine-66 profile (stage 35, no split-K or reduction, +custom options 0/1/2/4/5/6/7, clusters 2×1×1, 4×1×1 or 1×2×1, undefined inner shape). It +requires the listed configuration; other options retain the original comparison. +Arithmetic equivalence between these options is empirically tested, not guaranteed +by the cuBLAS API. +This avoids redundant full-prefix computation for compatible long candidates. +The optional CUDA capability defaults off for other consumers; older libraries, +other hardware/library versions and unsupported shapes retain the conservative +algorithm comparison. It neither changes ordinary GEMM algorithm selection nor +shares numeric state between requests. Full-model output parity is required in +addition to standalone tile-equivalence checks. + +For 9B exact-prefix mode, candidates of at least 512 tokens can also reuse a +prefix when the original packed input/gate-up GEMMs are compatible with the +candidate projections. The check includes each original batch's total row count. +Any mismatch declines the entire sharing attempt before computation and retains +the original candidate packing. Other prefix modes and the 27B policy are unchanged. + +The 9B default execution path batches already queued requests through the shared +runtime, up to 16 candidates and 8192 tokens per GPU dispatch, without an added +batching timer. Original candidate batches remain atomic: input/gate-up projection +algorithms are selected at their original batch row counts, while output/down +projections retain per-candidate row counts. Request-local prefix work runs alone. +Question calibration waits for the complete request and retains candidate order. +Fixed/shared/auto prefix modes and 27B retain serial scheduling. CUDA Graph keys +include both original batch groups and candidate lengths. Queued cancellation, +post-dispatch resource retention and errors follow the [runtime contract](../../src/runtime/README.md). + +With 9B exact-prefix mode, a bounded single-question batch can compact an exact +16–511-token prefix for pointwise projections, residuals and the MLP. Before each +mixer, projected rows are expanded into the original complete candidate layout; +attention and GDN retain their original lengths, positions and sequence boundaries. +Their outputs are gathered back into the compact layout. All projection algorithms +are checked against the original packed input sizes and individual output sizes +before execution; unsupported layouts keep the complete forward. Short retained +checks use tested 64–511-row profiles, including custom option 3 and cluster 4×2×1, +under the same SM90/cuBLASLt scope. Ordinary packed GEMM grouping is unchanged. + +This path stores only reusable row-index buffers and recomputes every request. +It does not count projection-only reuse as full-prefix `saved_tokens`; reported +input usage still includes every original token. The original full-prefix path +remains available for longer prefixes. + +A 9B exact-prefix capture requests state only. Once the final full-attention +layer has stored its K/V rows, it omits that prefix's attention output, output +projection, MLP and final normalization: no caller reads those hidden rows. +The suffix starts from fresh embeddings and consumes only the captured state. +Other window/readout modes retain complete hidden outputs. State capture is +synchronized before it becomes available to a suffix. diff --git a/src/backends/cuda/qwen3_5/attention.cu b/src/backends/cuda/qwen3_5/attention.cu index 5193e05f..7deee166 100644 --- a/src/backends/cuda/qwen3_5/attention.cu +++ b/src/backends/cuda/qwen3_5/attention.cu @@ -11,6 +11,7 @@ // rounded to bfloat16 for the P*V product, as in flash attention; the running sums // stay float32. #include +#include #include "common.cuh" #include "mma.cuh" #include "ops.h" @@ -86,17 +87,17 @@ 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 +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, const bf16* __restrict__ prefix_k, - const bf16* __restrict__ prefix_v) { + const bf16* __restrict__ prefix_v, int query_blocks = 0) { extern __shared__ __align__(16) unsigned char smem[]; bf16* qs = reinterpret_cast(smem); bf16* ks = qs + BM * LDS; bf16* vs = ks + BN * LDS; const int h = blockIdx.y, hk = h / (Hq / Hk); - const int q0 = q_base + (gridDim.x - 1 - blockIdx.x) * BM; // the longest blocks first + const int q0 = q_base + ((Batched ? static_cast(query_blocks) : gridDim.x) - 1 - blockIdx.x) * BM; // the longest blocks first const int tid = threadIdx.x, warp = tid / 32, lane = tid % 32; const int g = lane / 4, t = lane % 4; const int row0 = q0 + warp * 16; // this warp's first query @@ -274,6 +275,28 @@ flash_shared_kernel(const bf16* __restrict__ q, const bf16* __restrict__ k, cons flash_body(q, k, v, ldv, gate, out, T, Hq, Hk, scale_log2, prefix, prefix_k, prefix_v); } +// Host descriptors are copied into the launch, including when a graph records it. +// Every z slice keeps its own original query tiling and online-softmax order. +struct SharedBatch { + int lengths[16]; + int offsets[16]; +}; +template +__global__ void __launch_bounds__(THREADS) +flash_batch_kernel(const bf16* q, const bf16* k, const bf16* v, int ldv, + const bf16* gate, bf16* out, SharedBatch batch, int prefix, + int Hq, int Hk, float scale_log2, const bf16* prefix_k, + const bf16* prefix_v) { + const int suffix = batch.lengths[blockIdx.z]; + const int blocks = (suffix + BM - 1) / BM; + if (blockIdx.x >= blocks) return; + const size_t offset = batch.offsets[blockIdx.z]; + flash_body(q + offset * Hq * D, k + offset * Hk * D, + v + offset * ldv, ldv, gate + offset * Hq * D, + out + offset * Hq * D, prefix + suffix, Hq, Hk, + scale_log2, prefix, prefix_k, prefix_v, blocks); +} + 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) { @@ -364,3 +387,59 @@ extern "C" int cs1_attention_gated_shared(const void* q, const void* k, const vo prefix, static_cast(prefix_k), static_cast(prefix_v)); return cudaGetLastError(); } + +template +static int attention_batch_launch(const void* q, const void* k, const void* v, + int ldv, const void* gate, void* out, + const void* prefix_k, const void* prefix_v, + const size_t* lengths, size_t count, int prefix, + int Hq, int Hk, int Dh, float scale, void* stream) { + if (!q || !k || !v || !gate || !out || !lengths || + count == 0 || count > 16 || prefix < 0 || prefix > INT_MAX - flash::BM || + (Shared && (!prefix_k || !prefix_v || prefix == 0 || prefix % flash::BM)) || + Dh != flash::D || Hq <= 0 || Hk <= 0 || Hk > INT_MAX / Dh || Hq % Hk || + ldv % 8 || ldv < Hk * Dh) return cudaErrorInvalidValue; + flash::SharedBatch batch{}; + size_t rows = 0; + int blocks = 0; + for (size_t i = 0; i < count; ++i) { + if (!lengths[i] || lengths[i] > static_cast(INT_MAX - flash::BM - prefix) || + lengths[i] > static_cast(INT_MAX - flash::BM) - rows) + return cudaErrorInvalidValue; + const int suffix = static_cast(lengths[i]); + batch.lengths[i] = suffix; + batch.offsets[i] = static_cast(rows); + rows += lengths[i]; + blocks = std::max(blocks, (suffix + flash::BM - 1) / flash::BM); + } + static const cudaError_t configured = cudaFuncSetAttribute( + flash::flash_batch_kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, + flash::SMEM_BYTES); + if (configured != cudaSuccess) return configured; + constexpr float LOG2E = 1.4426950408889634f; + flash::flash_batch_kernel<<(stream)>>>( + static_cast(q), static_cast(k), static_cast(v), + ldv, static_cast(gate), static_cast(out), batch, prefix, + Hq, Hk, scale * LOG2E, static_cast(prefix_k), static_cast(prefix_v)); + return cudaGetLastError(); +} + +// Same arithmetic as separate shared-suffix calls; candidate rows are disjoint. +extern "C" int cs1_attention_gated_shared_batch(const void* q, const void* k, const void* v, + int ldv, const void* gate, void* out, + const void* prefix_k, const void* prefix_v, + const size_t* lengths, size_t count, int prefix, + int Hq, int Hk, int Dh, float scale, void* stream) { + return attention_batch_launch(q, k, v, ldv, gate, out, prefix_k, prefix_v, + lengths, count, prefix, Hq, Hk, Dh, scale, stream); +} + +// Independent full sequences have their own key window and positions starting at zero. +extern "C" int cs1_attention_gated_batch(const void* q, const void* k, const void* v, int ldv, + const void* gate, void* out, const size_t* lengths, + size_t count, int Hq, int Hk, int Dh, float scale, + void* stream) { + return attention_batch_launch(q, k, v, ldv, gate, out, nullptr, nullptr, + lengths, count, 0, Hq, Hk, Dh, scale, stream); +} diff --git a/src/backends/cuda/qwen3_5/gdn_prefill.cu b/src/backends/cuda/qwen3_5/gdn_prefill.cu index 7201307f..6d4e649e 100644 --- a/src/backends/cuda/qwen3_5/gdn_prefill.cu +++ b/src/backends/cuda/qwen3_5/gdn_prefill.cu @@ -661,9 +661,9 @@ __global__ void __launch_bounds__(THREADS) gdn_batch_prep( 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) { +__global__ void __launch_bounds__(ST_THREADS) gdn_batch_state(Batch batch, const float* state) { const auto sequence = batch.sequences[blockIdx.z]; - gdn_chunk_state_body(sequence.ws, (sequence.tokens + C - 1) / C, nullptr, nullptr); + gdn_chunk_state_body(sequence.ws, (sequence.tokens + C - 1) / C, state, nullptr); } __global__ void __launch_bounds__(OUT_THREADS) gdn_batch_out(Batch batch, bf16* o, int H) { const auto sequence = batch.sequences[blockIdx.z]; @@ -786,9 +786,10 @@ size_t cs1_gdn_batch_workspace_floats(int tokens, int H, int count) { 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, +int cs1_gdn_prefill_batch_shared(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) { + const size_t* lengths, size_t count, int H, int HK, float scale, + const void* state, 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{}; @@ -817,9 +818,17 @@ int cs1_gdn_prefill_batch(const void* q, const void* k, const void* v, const flo 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_state<<>>(batch, static_cast(state)); gdn_batch_out<<>>(batch, static_cast(o), H); return cudaGetLastError(); } +// Legacy independent sequences use the same arithmetic with a zero initial state. +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) { + return cs1_gdn_prefill_batch_shared(q, k, v, g, beta, o, workspace, workspace_bytes, + lengths, count, H, HK, scale, nullptr, stream); +} + } // extern "C" diff --git a/src/backends/cuda/qwen3_5/gemm.cu b/src/backends/cuda/qwen3_5/gemm.cu index 92ebb54b..62a60dae 100644 --- a/src/backends/cuda/qwen3_5/gemm.cu +++ b/src/backends/cuda/qwen3_5/gemm.cu @@ -62,6 +62,7 @@ struct Gemm { cublasLtHandle_t handle = nullptr; void* workspace = nullptr; size_t workspace_bytes = 0; + bool equivalent_tiles = false; bool reference = false; bool matched = false; int reference_m = 0; // > 0: one algorithm per weight shape, chosen at this M @@ -73,6 +74,47 @@ struct Gemm { std::list packed_order; }; +template +bool config_value(const cublasLtMatmulAlgo_t& algo, cublasLtMatmulAlgoConfigAttributes_t attr, T& value) { + size_t written = 0; + return cublasLtMatmulAlgoConfigGetAttribute(&algo, attr, &value, sizeof(value), &written) + == CUBLAS_STATUS_SUCCESS && written == sizeof(value); +} + +// Empirically validated SM90/cuBLASLt 13.1.1 long-prefill profile. cuBLAS +// does not promise arithmetic equivalence between custom options; the optional +// capability is version/device scoped and covered by full-row/packed/graph and +// model precision tests. Keep split-K and all unknown configurations separate. +bool validated_long_profile(const Plan& p, bool short_reuse = false) { + for (auto pair : {std::make_pair(CUBLASLT_ALGO_CONFIG_ID, 66u), + std::make_pair(CUBLASLT_ALGO_CONFIG_SPLITK_NUM, 1u), + std::make_pair(CUBLASLT_ALGO_CONFIG_REDUCTION_SCHEME, 0u), + std::make_pair(CUBLASLT_ALGO_CONFIG_CTA_SWIZZLING, 0u), + std::make_pair(CUBLASLT_ALGO_CONFIG_STAGES_ID, 35u)}) { + uint32_t value = 0; + if (!config_value(p.algo, pair.first, value) || value != pair.second) return false; + } + uint32_t option = 0; + uint16_t inner = 0, cluster = 0; + if (!config_value(p.algo, CUBLASLT_ALGO_CONFIG_CUSTOM_OPTION, option) || + !config_value(p.algo, CUBLASLT_ALGO_CONFIG_INNER_SHAPE_ID, inner) || + !config_value(p.algo, CUBLASLT_ALGO_CONFIG_CLUSTER_SHAPE_ID, cluster)) return false; + return (option == 0 || option == 1 || option == 2 || option == 4 || option == 5 || option == 6 || option == 7 || (short_reuse && option == 3)) && inner == 0 && + (cluster == CUBLASLT_CLUSTER_SHAPE_2x1x1 || cluster == CUBLASLT_CLUSTER_SHAPE_4x1x1 || + cluster == CUBLASLT_CLUSTER_SHAPE_1x2x1 || + (short_reuse && cluster == CUBLASLT_CLUSTER_SHAPE_4x2x1)); +} + +bool compatible_algorithms(const Gemm& g, const Plan& a, const Plan& b, size_t ma, size_t mb, bool short_reuse = false) { + if (std::memcmp(&a.algo, &b.algo, sizeof(a.algo)) == 0) return true; + if (!g.equivalent_tiles || ma > 4096 || mb > 4096) return false; + if (ma >= 512 && mb >= 512) return validated_long_profile(a) && validated_long_profile(b); + // Only retained prefix/suffix projections can use the tested short profiles. + // Keep ordinary packed GEMM grouping and algorithm selection unchanged. + return short_reuse && ma >= 64 && mb >= 64 && + validated_long_profile(a, true) && validated_long_profile(b, true); +} + int status(cublasStatus_t s) { return s == CUBLAS_STATUS_SUCCESS ? 0 : 1000 + (int)s; } void destroy(Plan& p) { @@ -187,8 +229,9 @@ 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). +// Optional shape changes retain the original arithmetic profile. The original +// algorithm must still support the target shape; only opted-in compatible profiles +// may use the target-row heuristic. Return -1 keeps the complete-forward fallback. 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); @@ -212,6 +255,21 @@ int reused_plan(Gemm& g, int M, int N, int K, int original_M, Plan*& out) { destroy(p); return -1; } + // Prefix sharing changes the executed row count, sometimes by a large factor. + // Keep the original eligibility check above, then choose the actual-row plan + // only inside the empirically validated SM90/Lt profile. Other handles and + // split-K/unknown algorithms retain the exact original choice. + if (g.equivalent_tiles && M != original_M) { + Plan* actual = nullptr; + if (plan_for(g, M, N, K, N, actual) == 0 && + compatible_algorithms(g, *original, *actual, original_M, M, true)) { + cublasLtMatmulHeuristicResult_t candidate{}; + const auto valid = cublasLtMatmulAlgoCheck( + g.handle, p.op, p.a, p.b, p.c, p.c, &actual->algo, &candidate); + if (valid == CUBLAS_STATUS_SUCCESS && candidate.state == CUBLAS_STATUS_SUCCESS && + candidate.workspaceSize <= g.workspace_bytes) p.algo = actual->algo; + } + } if (g.reused_plans.size() == MAX_REUSED_PLANS) { const auto oldest = g.reused_plans.find(g.reused_order.front()); destroy(oldest->second.plan); @@ -271,7 +329,7 @@ extern "C" int cs1_gemm(void* gemm, const void* x, const void* w, void* y, int M // 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. +// shapes with compatible 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) { @@ -300,8 +358,7 @@ extern "C" int cs1_gemm_packed(void* gemm, const void* x, const void* w, void* y // 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) { + compatible_algorithms(*g, *originals[begin], *originals[end], lengths[begin], lengths[end])) { rows += lengths[end++]; } PackedSegment segment; @@ -381,7 +438,7 @@ extern "C" int cs1_gemm_reuse_check(void* gemm, const size_t* original_rows, siz 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; + if (!compatible_algorithms(*g, *original, *other, original_rows[0], original_rows[i], true)) return -1; } for (size_t i = 0; i < targets; ++i) { Plan* target = nullptr; @@ -455,3 +512,17 @@ extern "C" int cs1_gemm_compatible(void* ordinary, void* matched, int M, int N, if (plan_for(*static_cast(ordinary), M, N, K, N, a)) return 0; return std::memcmp(&a->algo, &b->algo, sizeof(a->algo)) == 0; } + +// Additive ABI 9 capability. Select before creating plans/capturing any graph. +// 1 = enabled, 0 = unsupported profile (unchanged fallback), negative = error. +extern "C" int cs1_gemm_enable_equivalent_tiles(void* gemm) { + auto* g = static_cast(gemm); + if (!g || g->reference || g->matched || g->reference_m || !g->plans.empty()) return -1; + int device = 0; + cudaDeviceProp prop{}; + auto error = cudaGetDevice(&device); + if (error == cudaSuccess) error = cudaGetDeviceProperties(&prop, device); + if (error != cudaSuccess) return -static_cast(error); + g->equivalent_tiles = prop.major == 9 && prop.minor == 0 && cublasLtGetVersion() == 130101; + return g->equivalent_tiles ? 1 : 0; +} diff --git a/src/backends/cuda/qwen3_5/norm.cu b/src/backends/cuda/qwen3_5/norm.cu index 5e1c210d..17347be6 100644 --- a/src/backends/cuda/qwen3_5/norm.cu +++ b/src/backends/cuda/qwen3_5/norm.cu @@ -72,6 +72,46 @@ __global__ void __launch_bounds__(NORM_THREADS) } } +// Vectorize memory traffic for 4096-wide rows while retaining the scalar +// kernel's FP32 reduction order. The shared transpose gives each thread its +// original strided elements; rounded BF16 values remain in registers for output. +__global__ void __launch_bounds__(NORM_THREADS) + add_rms_norm_vector_4096_kernel(Pack8* __restrict__ residual, const Pack8* __restrict__ delta, + const Pack8* __restrict__ w, Pack8* __restrict__ out, float eps) { + constexpr int D = 4096, PACKS = D / (8 * NORM_THREADS); + __shared__ float scratch[32]; + __shared__ __align__(16) bf16 shared_values[D]; + const size_t base = (size_t)blockIdx.x * (D / 8); + Pack8 values[PACKS]; +#pragma unroll + for (int i = 0; i < PACKS; ++i) { + const int col = threadIdx.x + i * NORM_THREADS; + const Pack8 a = residual[base + col], b = delta[base + col]; +#pragma unroll + for (int e = 0; e < 8; ++e) values[i].v[e] = to_bf16(f32(a.v[e]) + f32(b.v[e])); + residual[base + col] = values[i]; + reinterpret_cast(shared_values)[col] = values[i]; + } + __syncthreads(); + float ss = 0.f; +#pragma unroll + for (int i = 0; i < D / NORM_THREADS; ++i) { + const float v = f32(shared_values[threadIdx.x + i * NORM_THREADS]); + ss += v * v; + } + const float inv = rsqrtf(block_sum(ss, scratch) / D + eps); +#pragma unroll + for (int i = 0; i < PACKS; ++i) { + const int col = threadIdx.x + i * NORM_THREADS; + const Pack8 weight = w[col]; + Pack8 result; +#pragma unroll + for (int e = 0; e < 8; ++e) + result.v[e] = to_bf16(f32(values[i].v[e]) * inv * (1.f + f32(weight.v[e]))); + out[base + col] = result; + } +} + // Qwen3_5RMSNormGated with D = 128, one warp per row: the normalized value is // rounded to bfloat16, multiplied by w in bfloat16, then by silu(z) in float32. __global__ void gated_rms_norm_kernel(const bf16* __restrict__ x, const bf16* __restrict__ z, int ldz, @@ -113,6 +153,13 @@ extern "C" int cs1_rms_norm(const void* x, const void* w, void* out, int rows, i extern "C" int cs1_add_rms_norm(void* residual, const void* delta, const void* w, void* out, int rows, int D, float eps, void* stream) { if (rows <= 0) return cudaSuccess; + if (D == 4096 && ((reinterpret_cast(residual) | reinterpret_cast(delta) | + reinterpret_cast(w) | reinterpret_cast(out)) & 15) == 0) { + add_rms_norm_vector_4096_kernel<<(stream)>>>( + static_cast(residual), static_cast(delta), static_cast(w), + static_cast(out), eps); + return cudaGetLastError(); + } switch (D) { #define CACHED_NORM(width) \ case width: \ diff --git a/src/backends/cuda/qwen3_5/ops.h b/src/backends/cuda/qwen3_5/ops.h index 9ad9c739..d853715e 100644 --- a/src/backends/cuda/qwen3_5/ops.h +++ b/src/backends/cuda/qwen3_5/ops.h @@ -80,6 +80,19 @@ int cs1_attention_gated_shared(const void* q, const void* k, const void* v, int const void* prefix_v, int suffix, int prefix, int Hq, int Hk, int Dh, float scale, void* stream); +// Optional batch of 1..16 contiguous suffix sequences sharing immutable prefix KV. +// Lengths are host values consumed before return; launch descriptors are by value. +int cs1_attention_gated_shared_batch(const void* q, const void* k, const void* v, int ldv, + const void* gate, void* out, const void* prefix_k, + const void* prefix_v, const size_t* lengths, size_t count, + int prefix, int Hq, int Hk, int Dh, float scale, void* stream); + +// Optional independent full-sequence causal attention, same concatenated layouts. +// Candidate lengths are host values copied into by-value launch descriptors. +int cs1_attention_gated_batch(const void* q, const void* k, const void* v, int ldv, + const void* gate, void* out, const size_t* lengths, + size_t count, 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. int cs1_gdn_gates(const void* b, const void* a, int ld, const void* A_log, const void* dt_bias, @@ -92,6 +105,10 @@ 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); +// Optional: every suffix reads the same immutable FP32 [H,128,128] prefix state. +int cs1_gdn_prefill_batch_shared(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, const void* state, 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. @@ -171,6 +188,11 @@ void* cs1_gemm_create_fixed(size_t workspace_bytes, int reference_m); void* cs1_gemm_create_matched(size_t workspace_bytes, int reference_m); int cs1_gemm_compatible(void* ordinary, void* matched, int M, int N, int K); void cs1_gemm_destroy(void* gemm); +// Optional prefill tile equivalence; 1 enabled, 0 unsupported, negative error. +// Select on an ordinary handle before its first plan. Ordinary packing accepts +// tested512..4096-row profiles; retained reuse additionally accepts64..511 rows. +// Outside those profiles keep opaque equality; split-K is never relaxed. +int cs1_gemm_enable_equivalent_tiles(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); @@ -178,7 +200,8 @@ int cs1_gemm(void* gemm, const void* x, const void* w, void* y, int M, int N, in 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 +// Optional retained-arithmetic projections. Opted-in compatible profiles may select +// the actual-row algorithm after validating the original. 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); diff --git a/src/models/open_jev/README.md b/src/models/open_jev/README.md index df6c32f7..70b2dae9 100644 --- a/src/models/open_jev/README.md +++ b/src/models/open_jev/README.md @@ -69,3 +69,30 @@ every candidate with the same fixed GEMM algorithms as a control; see The engine owns a [shared serial scheduler](../../runtime/README.md) that admits the complete request before blocking dispatch. Cross-request batching and shared queue budgets remain planned. + + +The experimental 9B layerwise prefix tree can batch the independent suffix +attention queries into one CUDA launch per layer. Each candidate keeps its +original query blocks, causal key order, softmax accumulation and BF16 rounding; +only independent candidates execute together. Prefix K/V is immutable for that +layer and every suffix writes a disjoint output range. The optional +`cs1_attention_gated_shared_batch` entry consumes host lengths before returning +and passes descriptors by value, including under graph capture. Libraries +without the entry retain the separate candidate launches. The model continues +to own scratch and stream lifetimes; this adds no request cache or device state. + +The same optional batching covers ordinary and pointwise-compacted 9B causal +sequences: all candidates first prepare their separate Q/K/gate rows, then one +`cs1_attention_gated_batch` launch computes independent attention windows. +Sequence boundaries, original projection plans and output ordering remain +unchanged. Single sequences, reference numerics, bidirectional attention and +unvalidated adapter modes keep their previous entry points. Older CUDA +libraries without the optional symbol also retain separate launches. + +For opted-in SM90 / cuBLASLt 13.1.1 prefix projections, the original algorithm +must still support the actual row count. Within the already validated compatible +arithmetic profiles, the runtime can then select the heuristic for that actual +count. Unknown or split-K profiles and other hardware retain the original +algorithm. Candidate grouping and fallback eligibility stay unchanged; full-row +bitwise comparisons, changed-input graph replay and model output checks cover +the choice. cuBLASLt does not guarantee arithmetic equivalence across algorithms. diff --git a/src/models/open_jev/native/src/batching.rs b/src/models/open_jev/native/src/batching.rs index 58317242..f81fc7f1 100644 --- a/src/models/open_jev/native/src/batching.rs +++ b/src/models/open_jev/native/src/batching.rs @@ -27,8 +27,16 @@ pub(crate) fn ranges(inputs: &[&[u32]]) -> Vec> { /// 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) + shared_prefix_with_minimum(inputs, MAX_TOKENS / 2 + 1) +} + +/// Packed-row compatibility is checked by the executor before any prefix runs. +pub(crate) fn packed_shared_prefix(inputs: &[&[u32]]) -> usize { + shared_prefix_with_minimum(inputs, 512) +} + +fn shared_prefix_with_minimum(inputs: &[&[u32]], minimum: usize) -> usize { + if !(2..=MAX_SEQUENCES).contains(&inputs.len()) || inputs.iter().any(|ids| ids.len() < minimum) { return 0; } @@ -40,6 +48,25 @@ pub(crate) fn shared_prefix(inputs: &[&[u32]]) -> usize { if aligned >= 512 { aligned } else { 0 } } +/// Pointwise projection compaction keeps original mixers and one bounded batch. +pub(crate) fn projection_prefix(inputs: &[&[u32]]) -> usize { + if !(2..=MAX_SEQUENCES).contains(&inputs.len()) + || inputs.iter().any(|ids| ids.is_empty()) + || inputs.iter().map(|ids| ids.len()).sum::() > MAX_TOKENS + { + 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(); + if (16..512).contains(&common) { + common + } 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 17926ea9..b424fed4 100644 --- a/src/models/open_jev/native/src/contract.rs +++ b/src/models/open_jev/native/src/contract.rs @@ -16,10 +16,12 @@ pub struct Checkpoint { pub alias: &'static str, /// Hidden, intermediate, layers, attention, KV, linear key and linear value heads. pub backbone: (usize, usize, usize, usize, usize, usize, usize), - /// 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. + /// Pack a request's independent candidates for shared GEMMs. pub pack_candidates: bool, - /// Parallel preparation has exact-token coverage for the pinned 27B tokenizer. + /// Preserve the 27B default path's long-prefix reuse independently of packing. + /// 9B keeps prefix reuse opt-in through OPEN_JEV_PREFIX. + pub share_packed_prefix: bool, + /// Parallel preparation has exact-token coverage for both pinned tokenizers. pub parallel_tokenization: bool, } @@ -32,6 +34,7 @@ pub static CHECKPOINTS: [Checkpoint; 2] = [ alias: "open-jev-27b-v1.1", backbone: (5120, 17408, 64, 24, 4, 16, 48), pack_candidates: true, + share_packed_prefix: true, parallel_tokenization: true, }, Checkpoint { @@ -41,8 +44,9 @@ pub static CHECKPOINTS: [Checkpoint; 2] = [ checkpoint_revision: "47e966881e489511c0c7f5633a9e1960a676a551", alias: "open-jev-9b", backbone: (4096, 12288, 32, 16, 4, 16, 32), - pack_candidates: false, - parallel_tokenization: false, + pack_candidates: true, + share_packed_prefix: false, + parallel_tokenization: true, }, ]; diff --git a/src/models/open_jev/native/src/engine.rs b/src/models/open_jev/native/src/engine.rs index 007d727e..fc8eec93 100644 --- a/src/models/open_jev/native/src/engine.rs +++ b/src/models/open_jev/native/src/engine.rs @@ -4,17 +4,16 @@ use std::path::Path; use anyhow::{Context, Result, bail, ensure}; use omni_qwen3_5_native::prefix::PrefixMode; -use omni_runtime::SerialScheduler; use serde_json::Value; use crate::contract::Checkpoint; -use crate::executor::{DecisionHead, Executor}; +use crate::executor::{DecisionHead, Executor, Scheduler}; use crate::processing::Processor; pub struct Engine { pub checkpoint: &'static Checkpoint, pub processor: Processor, - pub scheduler: SerialScheduler, + pub scheduler: Scheduler, pub executor: Executor, } @@ -61,9 +60,36 @@ pub fn prefix_env() -> Result { ) } +/// Exact request-local reuse keeps each complete prompt's GEMM algorithms. +/// It is separate from the fixed-algorithm prefix modes. +pub fn exact_prefix_value(value: Option<&str>, prefix: PrefixMode) -> Result> { + let enabled = match value { + None => None, + Some("0") => Some(false), + Some("1") => Some(true), + Some(_) => bail!("OPEN_JEV_EXACT_PREFIX must be 0 or 1"), + }; + ensure!( + enabled != Some(true) || prefix == PrefixMode::Off, + "OPEN_JEV_EXACT_PREFIX cannot be combined with fixed/shared/auto prefix modes" + ); + Ok(enabled) +} + impl Engine { pub async fn load(dir: &Path, library: &Path) -> Result { - Self::load_with_prefix(dir, library, prefix_env()?).await + let prefix = prefix_env()?; + let raw = match std::env::var("OPEN_JEV_EXACT_PREFIX") { + Ok(value) => Some(value), + Err(std::env::VarError::NotPresent) => None, + Err(error) => return Err(error).context("OPEN_JEV_EXACT_PREFIX"), + }; + let exact = exact_prefix_value(raw.as_deref(), prefix)?; + let mut engine = Self::load_with_prefix(dir, library, prefix).await?; + if let Some(enabled) = exact { + engine.executor.set_exact_prefix(enabled); + } + Ok(engine) } pub async fn load_with_prefix( @@ -100,10 +126,11 @@ impl Engine { let head = DecisionHead::load(dir, &manifest, checkpoint)?; let processor = Processor::load(dir, checkpoint, prefix, suffix, temperature, max_length)?; let executor = Executor::load(dir, library, head, checkpoint, prefix_mode, graph).await?; + let scheduler = executor.scheduler(checkpoint)?; Ok(Self { checkpoint, processor, - scheduler: SerialScheduler::default(), + scheduler, executor, }) } diff --git a/src/models/open_jev/native/src/executor.rs b/src/models/open_jev/native/src/executor.rs index 2bed6d1c..b0c98810 100644 --- a/src/models/open_jev/native/src/executor.rs +++ b/src/models/open_jev/native/src/executor.rs @@ -9,6 +9,8 @@ use omni_qwen3_5_native::prefix::{PrefixMode, PrefixPlan, PrefixRow}; use omni_runtime::SerialScheduler; use serde_json::Value; +mod batched; + use crate::batching; use crate::contract::Checkpoint; @@ -104,10 +106,20 @@ pub struct PrefixUse { pub saved_tokens: usize, } +/// Each loaded model owns one execution queue; only 9B default numerics batch. +pub struct Scheduler(SchedulerKind); + +enum SchedulerKind { + Serial(SerialScheduler), + Batched(omni_runtime::batching::BatchScheduler), +} + pub struct Executor { model: Arc>, head: DecisionHead, pack: bool, + share_packed_prefix: bool, + share_packed_inputs: bool, prefix: PrefixMode, } @@ -123,12 +135,16 @@ impl Executor { let (d, lib) = (dir.to_path_buf(), library.to_path_buf()); let sharing = matches!(prefix, PrefixMode::Shared | PrefixMode::Auto); let batch_gdn = checkpoint.pack_candidates; + let equivalent_tiles = checkpoint.alias == "open-jev-9b"; let model = tokio::task::spawn_blocking(move || { let mut model = Model::load_with_graph(&d, &lib, graph)?; if batch_gdn { model.enable_batched_gdn()?; model.enable_packed_projections()?; } + if equivalent_tiles && model.enable_equivalent_projection_tiles()? { + eprintln!("Open-Jev-9B: validated SM90 long-prefill projection tiles enabled"); + } // Shared prefixes reproduce the fixed rows only where rows do not depend on M. if sharing { model @@ -142,20 +158,39 @@ impl Executor { model: Arc::new(Mutex::new(model)), head, pack: checkpoint.pack_candidates, + share_packed_prefix: checkpoint.share_packed_prefix, + share_packed_inputs: checkpoint.alias == "open-jev-9b", prefix, }) } + pub(crate) fn scheduler(&self, checkpoint: &Checkpoint) -> Result { + Ok(Scheduler( + if checkpoint.alias == "open-jev-9b" && self.prefix == PrefixMode::Off { + SchedulerKind::Batched(batched::new(self.model.clone(), self.head.clone())?) + } else { + SchedulerKind::Serial(SerialScheduler::default()) + }, + )) + } + + /// Configure exact within-question reuse before the engine admits requests. + pub(crate) fn set_exact_prefix(&mut self, enabled: bool) { + self.share_packed_prefix = enabled; + } + /// Inputs and outputs are grouped by question, then candidate, in prepared order. pub async fn execute( &self, - scheduler: &SerialScheduler, + scheduler: &Scheduler, ids: Vec>>, ) -> Result>> { Ok(self.execute_with_prefix(scheduler, ids).await?.0) } - /// Admit one whole request. With prefix reuse selected (and, for `auto`, worth it), + /// 9B default submits original batches or one exclusive prefix item to the + /// shared batch scheduler. Other modes admit a whole request as described below. + /// With prefix reuse selected (and, for `auto`, worth it), /// shared prefixes run once through `forward_shared`; with fixed or forced shared /// execution the candidates otherwise run one at a time with `forward_fixed`. /// Otherwise the default path runs: where the checkpoint packs candidates, a single @@ -164,12 +199,20 @@ impl Executor { /// one at a time. pub async fn execute_with_prefix( &self, - scheduler: &SerialScheduler, + scheduler: &Scheduler, ids: Vec>>, ) -> Result<(Vec>, PrefixUse)> { + let scheduler = match &scheduler.0 { + SchedulerKind::Batched(scheduler) => { + return batched::execute(scheduler, ids, self.share_packed_prefix).await; + } + SchedulerKind::Serial(scheduler) => scheduler, + }; let model = self.model.clone(); let head = self.head.clone(); let pack = self.pack; + let share_packed_prefix = self.share_packed_prefix; + let share_packed_inputs = self.share_packed_inputs; let mode = self.prefix; scheduler .run(move || { @@ -203,13 +246,30 @@ impl Executor { } } else { // The default path, which is also `auto`'s fallback. - let prefix_len = if pack && ids.len() == 1 { - batching::shared_prefix(&inputs) + let prefix_len = if pack && share_packed_prefix && ids.len() == 1 { + if share_packed_inputs { + batching::packed_shared_prefix(&inputs) + } else { + batching::shared_prefix(&inputs) + } } else { 0 }; let shared = if prefix_len > 0 { - model.try_forward_shared_prefix(&inputs, prefix_len)? + if share_packed_inputs && inputs.iter().any(|ids| ids.len() <= 2048) { + let mut packed_rows = vec![0; inputs.len()]; + for range in batching::ranges(&inputs) { + let total = inputs[range.clone()].iter().map(|ids| ids.len()).sum(); + packed_rows[range].fill(total); + } + model.try_forward_shared_packed_prefix( + &inputs, + prefix_len, + &packed_rows, + )? + } else { + model.try_forward_shared_prefix(&inputs, prefix_len)? + } } else { None }; diff --git a/src/models/open_jev/native/src/executor/batched.rs b/src/models/open_jev/native/src/executor/batched.rs new file mode 100644 index 00000000..6dfe045f --- /dev/null +++ b/src/models/open_jev/native/src/executor/batched.rs @@ -0,0 +1,212 @@ +//! Model execution units retain original numerical layouts across FIFO batching. +use std::panic::{AssertUnwindSafe, catch_unwind, resume_unwind}; +use std::sync::{Arc, Mutex}; + +use anyhow::{Result, ensure}; +use omni_qwen3_5_native::model::Model; +use omni_runtime::batching::{BatchCost, BatchLimits, BatchScheduler}; + +use super::{DecisionHead, PrefixUse}; +use crate::batching; + +pub(super) struct WorkItem { + inputs: Vec>, + prefix: usize, + compaction: usize, +} + +pub(super) struct Output { + scores: Vec, + used: PrefixUse, +} + +pub(super) fn new( + model: Arc>, + head: DecisionHead, +) -> Result> { + BatchScheduler::new( + BatchLimits { + items: 16, + cost: 8192, + }, + |item: &WorkItem| BatchCost { + items: item.inputs.len(), + cost: item.inputs.iter().map(Vec::len).sum(), + exclusive: item.prefix > 0 || item.compaction > 0, + }, + move |items: &[WorkItem]| { + let mut model = model + .lock() + .map_err(|_| anyhow::anyhow!("poisoned model"))?; + // Keep the model guard alive while retiring any partially submitted + // work, even on a panic. The runtime closes the worker after a panic. + let result = catch_unwind(AssertUnwindSafe(|| run(&mut model, &head, items))); + match result { + Ok(Ok(outputs)) => Ok(outputs), + Ok(Err(error)) => { + model.synchronize()?; + Err(error) + } + Err(payload) => { + let _ = model.synchronize(); + resume_unwind(payload) + } + } + }, + ) +} + +fn run(model: &mut Model, head: &DecisionHead, items: &[WorkItem]) -> Result> { + if items.len() == 1 { + let item = &items[0]; + let prefix = item.prefix.max(item.compaction / 64 * 64); + if prefix >= 64 { + let inputs: Vec<&[u32]> = item.inputs.iter().map(Vec::as_slice).collect(); + let mut original_rows = vec![0; inputs.len()]; + for range in batching::ranges(&inputs) { + let rows = inputs[range.clone()].iter().map(|ids| ids.len()).sum(); + original_rows[range].fill(rows); + } + if let Some(output) = model.try_forward_prefix_tree(&inputs, prefix, &original_rows)? { + return Ok(vec![Output { + scores: output + .hidden + .into_iter() + .map(|last| head.score(last)) + .collect::>()?, + used: PrefixUse { + shared: true, + saved_tokens: output.saved_tokens, + ..PrefixUse::default() + }, + }]); + } + } + } + if items.len() == 1 && items[0].compaction > 0 { + let item = &items[0]; + let inputs: Vec<&[u32]> = item.inputs.iter().map(Vec::as_slice).collect(); + let hidden = match model.try_forward_prefix_compaction(&inputs, item.compaction)? { + Some(hidden) => hidden, + None => model.forward_batch(&inputs)?, + }; + // Mixers still process every token. Do not report full-prefix token savings. + return Ok(vec![Output { + scores: hidden + .into_iter() + .map(|last| head.score(last)) + .collect::>()?, + used: PrefixUse::default(), + }]); + } + if items.len() == 1 && items[0].prefix > 0 { + let item = &items[0]; + let inputs: Vec<&[u32]> = item.inputs.iter().map(Vec::as_slice).collect(); + let shared = if inputs.iter().any(|ids| ids.len() <= 2048) { + let mut rows = vec![0; inputs.len()]; + for range in batching::ranges(&inputs) { + let total = inputs[range.clone()].iter().map(|ids| ids.len()).sum(); + rows[range].fill(total); + } + model.try_forward_shared_packed_prefix(&inputs, item.prefix, &rows)? + } else { + model.try_forward_shared_prefix(&inputs, item.prefix)? + }; + let mut scores = Vec::with_capacity(inputs.len()); + let mut used = PrefixUse::default(); + if let Some(output) = shared { + used.shared = true; + used.saved_tokens = output.saved_tokens; + for last in output.hidden { + scores.push(head.score(last)?); + } + } else { + for range in batching::ranges(&inputs) { + for last in model.forward_batch(&inputs[range])? { + scores.push(head.score(last)?); + } + } + } + return Ok(vec![Output { scores, used }]); + } + ensure!( + items + .iter() + .all(|item| item.prefix == 0 && item.compaction == 0), + "exclusive prefix item was batched" + ); + let inputs: Vec<&[u32]> = items + .iter() + .flat_map(|item| item.inputs.iter().map(Vec::as_slice)) + .collect(); + let groups: Vec = items.iter().map(|item| item.inputs.len()).collect(); + let hidden = model.forward_batch_preserving_groups(&inputs, &groups)?; + let mut scores = hidden.into_iter().map(|last| head.score(last)); + groups + .into_iter() + .map(|count| { + Ok(Output { + scores: scores.by_ref().take(count).collect::>>()?, + used: PrefixUse::default(), + }) + }) + .collect() +} + +pub(super) async fn execute( + scheduler: &BatchScheduler, + ids: Vec>>, + share_prefix: bool, +) -> Result<(Vec>, PrefixUse)> { + let counts: Vec = ids.iter().map(Vec::len).collect(); + let candidates: Vec> = ids.into_iter().flatten().collect(); + let inputs: Vec<&[u32]> = candidates.iter().map(Vec::as_slice).collect(); + let prefix = if share_prefix && counts.len() == 1 { + batching::packed_shared_prefix(&inputs) + } else { + 0 + }; + let compaction = if share_prefix && counts.len() == 1 && prefix == 0 { + batching::projection_prefix(&inputs) + } else { + 0 + }; + let items = if prefix > 0 || compaction > 0 { + vec![WorkItem { + inputs: candidates, + prefix, + compaction, + }] + } else { + let ranges = batching::ranges(&inputs); + let mut candidates = candidates.into_iter(); + ranges + .into_iter() + .map(|range| WorkItem { + inputs: candidates.by_ref().take(range.len()).collect(), + prefix: 0, + compaction: 0, + }) + .collect() + }; + let outputs = scheduler.run(items).await?; + let mut used = PrefixUse::default(); + let mut scores = Vec::new(); + for output in outputs { + scores.extend(output.scores); + used.shared |= output.used.shared; + used.saved_tokens += output.used.saved_tokens; + } + ensure!( + scores.len() == counts.iter().sum::(), + "candidate score count mismatch" + ); + let mut scores = scores.into_iter(); + Ok(( + counts + .into_iter() + .map(|count| scores.by_ref().take(count).collect()) + .collect(), + used, + )) +} diff --git a/src/models/qwen3_5/native/src/cuda.rs b/src/models/qwen3_5/native/src/cuda.rs index 612d49da..4a8c55e9 100644 --- a/src/models/qwen3_5/native/src/cuda.rs +++ b/src/models/qwen3_5/native/src/cuda.rs @@ -89,7 +89,7 @@ pub struct SharedSuffix { ) -> c_int, } -/// Optional shape changes using a validated original-shape GEMM algorithm. +/// Optional shape changes preserving a validated original-shape arithmetic profile. #[derive(Clone, Copy)] pub struct ReuseGemm { pub check: unsafe extern "C" fn( @@ -127,6 +127,62 @@ pub type PackedGemm = unsafe extern "C" fn( Stream, ) -> c_int; +/// Optional causal attention over independent full sequences. +pub type BatchedAttention = unsafe extern "C" fn( + *const c_void, + *const c_void, + *const c_void, + c_int, + *const c_void, + *mut c_void, + *const usize, + usize, + c_int, + c_int, + c_int, + f32, + Stream, +) -> c_int; + +/// Optional attention over independent suffixes sharing immutable prefix KV. +pub type SharedBatchedAttention = 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, + *const usize, + usize, + c_int, + c_int, + c_int, + c_int, + f32, + Stream, +) -> c_int; + +/// Optional batched continuations of one immutable FP32 prefix state. +pub type SharedBatchedGdn = 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, + *const c_void, + Stream, +) -> c_int; + /// Optional GDN prefill over up to sixteen independent token sequences. #[derive(Clone, Copy)] pub struct BatchedGdn { @@ -443,6 +499,18 @@ impl Api { } } + pub fn equivalent_tiles(&self) -> Option c_int> { + // SAFETY: optional additive signature matches ops.h; Api owns the library. + unsafe { + self._lib + .get:: c_int>( + b"cs1_gemm_enable_equivalent_tiles\0", + ) + .ok() + .map(|f| *f) + } + } + pub fn wide_gdn(&self) -> Option { // SAFETY: the additive signature matches ops.h; API retains the library. unsafe { @@ -494,6 +562,36 @@ impl Api { .context("packed projections require the optional cs1_gemm_packed CUDA symbol") } + pub fn batched_attention(&self) -> Option { + // SAFETY: optional signature matches ops.h; this API owns the library. + unsafe { + self._lib + .get(b"cs1_attention_gated_batch\0") + .ok() + .map(|f| *f) + } + } + + pub fn shared_batched_attention(&self) -> Option { + // SAFETY: optional signature matches ops.h; this API owns the library. + unsafe { + self._lib + .get(b"cs1_attention_gated_shared_batch\0") + .ok() + .map(|f| *f) + } + } + + pub fn shared_batched_gdn(&self) -> Option { + // SAFETY: optional signature matches ops.h; the API owns the library. + unsafe { + self._lib + .get(b"cs1_gdn_prefill_batch_shared\0") + .ok() + .map(|f| *f) + } + } + pub fn batched_gdn(&self) -> Result { // SAFETY: optional signatures match ops.h; the static API owns the library. unsafe { diff --git a/src/models/qwen3_5/native/src/model.rs b/src/models/qwen3_5/native/src/model.rs index f268e546..3f542a66 100644 --- a/src/models/qwen3_5/native/src/model.rs +++ b/src/models/qwen3_5/native/src/model.rs @@ -705,6 +705,19 @@ pub struct PrefixState { bytes: usize, } +#[path = "prefix_compaction.rs"] +mod prefix_compaction; +use prefix_compaction::{Compaction, Projection}; +#[path = "prefix_tree.rs"] +mod prefix_tree; +use prefix_tree::Tree; + +#[derive(Clone, Copy, PartialEq, Eq)] +enum WindowReadout { + Hidden, + StateOnly, +} + /// A packed suffix batch uses one freshly captured prefix and original algorithm. struct SuffixBatch<'a> { prefix: &'a PrefixState, @@ -862,6 +875,10 @@ 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)>, + /// Reusable row-index storage; contains no request values. + compaction_maps: Option<(usize, DeviceBuffer)>, + /// One layer's FP32 recurrent state and prefix V, overwritten on every use. + tree_workspace: Option<(usize, DeviceBuffer)>, /// Opt-in replay keyed by sequence lengths; bounded separately for each mode. graph_enabled: bool, graph_stats: GraphStats, @@ -1035,6 +1052,8 @@ impl Model { scratch: None, shared: [None, None], vocab_buffers: None, + compaction_maps: None, + tree_workspace: None, graph_enabled: graph, graph_stats: GraphStats { requested: graph, @@ -1066,6 +1085,27 @@ impl Model { Ok(()) } + /// Opt into the backend's validated long-prefill output-tile equivalence. + /// Unsupported devices/libraries keep conservative algorithm matching. + pub fn enable_equivalent_projection_tiles(&mut self) -> Result { + ensure!( + self.scratch.is_none() + && self.reference.is_none() + && self.lora.is_none() + && self.fp32_gdn.is_none(), + "select projection tiles on a merged model before the first forward" + ); + let Some(enable) = cuda::api().equivalent_tiles() else { + return Ok(false); + }; + // SAFETY: the ordinary GEMM handle belongs to this unstarted model. + match unsafe { enable(self.gemm) } { + 0 => Ok(false), + 1 => Ok(true), + code => anyhow::bail!("projection tile setup failed: {code}"), + } + } + /// Opt into independent GDN batching before allocating any request buffers. pub fn enable_batched_gdn(&mut self) -> Result<()> { ensure!( @@ -1392,6 +1432,53 @@ impl Model { /// Pack independent text prompts for input and gate/up GEMMs. /// Mixers reset at each boundary; final-position hidden states retain input order. pub fn forward_batch(&mut self, inputs: &[&[u32]]) -> Result>> { + self.forward_batch_layout(inputs, None) + } + + /// Combine original candidate batches without changing their input/gate-up + /// GEMM algorithms. `groups` counts candidates in each original batch; mixer + /// boundaries and output/down projections still use each candidate's length. + pub fn forward_batch_preserving_groups( + &mut self, + inputs: &[&[u32]], + groups: &[usize], + ) -> Result>> { + ensure!( + !groups.is_empty() + && groups.iter().all(|&n| n > 0) + && groups.iter().try_fold(0usize, |a, &b| a.checked_add(b)) == Some(inputs.len()), + "invalid original candidate batch groups" + ); + if groups.len() == 1 { + return self.forward_batch(inputs); + } + ensure!( + self.packed_gemm.is_some() + && self.lora.is_none() + && self.reference.is_none() + && self.fp32_gdn.is_none(), + "grouped batches require merged BF16 packed projections" + ); + let mut start = 0; + let original_rows: Result> = groups + .iter() + .map(|&count| { + let rows = inputs[start..start + count] + .iter() + .try_fold(0usize, |a, ids| a.checked_add(ids.len())) + .context("original batch row overflow"); + start += count; + rows + }) + .collect(); + self.forward_batch_layout(inputs, Some(&original_rows?)) + } + + fn forward_batch_layout( + &mut self, + inputs: &[&[u32]], + original_rows: Option<&[usize]>, + ) -> Result>> { ensure!(!inputs.is_empty(), "empty batch"); ensure!( inputs @@ -1420,7 +1507,18 @@ impl Model { let ids: Vec = inputs.iter().flat_map(|ids| ids.iter().copied()).collect(); self.embed_tokens(s, &ids)?; if self.graph_enabled { - if let Some((_, graph)) = self.graphs.iter().find(|(shape, _)| *shape == lengths) { + let graph_key = if let Some(rows) = original_rows { + // Zero cannot occur in a valid candidate length vector. Include both + // layouts: identical candidate lengths can have different input GEMMs. + let mut key = vec![0]; + key.extend_from_slice(rows); + key.push(0); + key.extend_from_slice(&lengths); + key + } else { + lengths.clone() + }; + if let Some((_, graph)) = self.graphs.iter().find(|(shape, _)| *shape == graph_key) { graph.launch(self.stream)?; self.graph_stats.replays += 1; if std::env::var("CUA_S1_GRAPH_TRACE").as_deref() == Ok("1") { @@ -1429,14 +1527,23 @@ impl Model { } else { // Warm plans and keep the eager result: run() advances the residual // in place, so replaying on this cache miss would advance it twice. - self.run(s, &lengths, false, false)?; - cuda::synchronize(self.stream)?; - let captured = - cuda::Graph::capture(self.stream, || self.run(s, &lengths, false, false)); - self.cache_graph(&lengths, false, captured)?; + self.run_with_layout(s, &lengths, false, false, None, original_rows, None, None)?; + // Pinned9B plans are warm after host submission. Capture records + // subsequent operations without executing them, so recording can + // overlap eager GPU work. Readback still retires that eager work. + if (self.cfg.hidden, self.cfg.intermediate, self.layers.len()) != (4096, 12288, 32) + || self.reference.is_some() + || self.lora.is_some() + { + cuda::synchronize(self.stream)?; + } + let captured = cuda::Graph::capture(self.stream, || { + self.run_with_layout(s, &lengths, false, false, None, original_rows, None, None) + }); + self.cache_graph(&graph_key, false, t, captured)?; } } else { - self.run(s, &lengths, false, false)?; + self.run_with_layout(s, &lengths, false, false, None, original_rows, None, None)?; } let s = self.scratch.as_ref().unwrap(); let mut end = 0; @@ -1457,6 +1564,37 @@ impl Model { &mut self, inputs: &[&[u32]], prefix_len: usize, + ) -> Result> { + self.shared_prefix_with_original_rows(inputs, prefix_len, None) + } + + /// Share a prefix only when it also preserves the input/gate-up projections + /// of an existing packed forward. `packed_rows[i]` is the total row count of + /// the original batch containing candidate i. Output/down projections retain + /// each candidate's original row count. Incompatibility declines the whole + /// operation before computation, allowing the caller to run its original batches. + pub fn try_forward_shared_packed_prefix( + &mut self, + inputs: &[&[u32]], + prefix_len: usize, + packed_rows: &[usize], + ) -> Result> { + ensure!( + packed_rows.len() == inputs.len() + && packed_rows + .iter() + .zip(inputs) + .all(|(&rows, ids)| { rows >= ids.len() && rows <= i32::MAX as usize }), + "invalid original packed projection rows" + ); + self.shared_prefix_with_original_rows(inputs, prefix_len, Some(packed_rows)) + } + + fn shared_prefix_with_original_rows( + &mut self, + inputs: &[&[u32]], + prefix_len: usize, + packed_rows: Option<&[usize]>, ) -> Result> { ensure!( self.lora.is_none(), @@ -1501,23 +1639,32 @@ impl Model { 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])); + for (weight, packed) in [ + (input, true), + (output, false), + (&layer.gate_up, true), + (&layer.down, false), + ] { + shapes.insert((weight.shape[0], weight.shape[1], packed)); } } let gemm = self.gemm; - let Some(groups) = compatible_prefix_groups(inputs.len(), |indices| { + let compatible = |indices: &[usize]| { 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 { + for &(n, k, packed) in &shapes { + let mut projection_rows = originals.clone(); + if packed && let Some(rows) = packed_rows { + projection_rows.extend(indices.iter().map(|&i| rows[i])); + } // 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(), + projection_rows.as_ptr(), + projection_rows.len(), windows.as_ptr(), windows.len(), n as i32, @@ -1530,9 +1677,20 @@ impl Model { check(rc, "prefix projection compatibility")?; } Ok(true) - })? - else { - return Ok(None); + }; + let all: Vec = (0..inputs.len()).collect(); + let groups = if packed_rows.is_some() { + // Partial groups would change the input GEMM size of fallback + // candidates. Keep their entire original packing plan on rejection. + if !compatible(&all)? { + return Ok(None); + } + vec![all] + } else { + let Some(groups) = compatible_prefix_groups(inputs.len(), compatible)? else { + return Ok(None); + }; + groups }; let (capacity, mut state) = match self.text_prefix_workspace.take() { Some((capacity, state)) if capacity >= prefix_len => (capacity, state), @@ -1562,7 +1720,7 @@ impl Model { let target = [rows]; let mut packed = ops.is_some() && indices.len() > 1; if packed { - for &(n, k) in &shapes { + for &(n, k, _) in &shapes { // SAFETY: synchronous shape-only check, before prefix capture. let rc = unsafe { (reuse.check)( @@ -1603,6 +1761,13 @@ impl Model { None, self.gemm, projection, + if (self.cfg.hidden, self.cfg.intermediate, self.layers.len()) + == (4096, 12288, 32) + { + WindowReadout::StateOnly + } else { + WindowReadout::Hidden + }, )?; self.synchronize()?; state.initialized = true; @@ -1647,6 +1812,7 @@ impl Model { Some(&state), self.gemm, projection, + WindowReadout::Hidden, )?; hidden[index] = self.last_hidden(s, ids.len())?; } @@ -1703,7 +1869,7 @@ impl Model { self.run(s, &[t], true, false)?; cuda::synchronize(self.stream)?; let captured = cuda::Graph::capture(self.stream, || self.run(s, &[t], true, false)); - self.cache_graph(&[t], true, captured)?; + self.cache_graph(&[t], true, t, captured)?; } } else { self.run(s, &[t], true, false)?; @@ -1716,6 +1882,7 @@ impl Model { &mut self, lengths: &[usize], multimodal: bool, + tokens: usize, captured: Result, ) -> Result<()> { let mode = if multimodal { "multimodal" } else { "text" }; @@ -1741,8 +1908,7 @@ impl Model { } cache.push_back((lengths.to_vec(), graph)); if std::env::var("CUA_S1_GRAPH_TRACE").as_deref() == Ok("1") { - let t: usize = lengths.iter().sum(); - eprintln!("CUA_S1_GRAPH mode={mode} event=capture tokens={t}"); + eprintln!("CUA_S1_GRAPH mode={mode} event=capture tokens={tokens}"); } } Err(error) => { @@ -1865,6 +2031,7 @@ impl Model { prefix, gemm, None, + WindowReadout::Hidden, )?; if let Some(cap) = capture { cap.initialized = true; @@ -2139,7 +2306,17 @@ 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, None)?; + self.run_window( + s, + 0, + t, + true, + Some(state), + None, + self.gemm, + None, + WindowReadout::Hidden, + )?; self.synchronize()?; state.initialized = true; Ok(()) @@ -2190,6 +2367,7 @@ impl Model { Some(prefix), self.gemm, None, + WindowReadout::Hidden, )?; self.last_hidden(s, tend) } @@ -2220,7 +2398,17 @@ 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, None)?; + self.run_window( + s, + 0, + t, + false, + None, + None, + gemm, + None, + WindowReadout::Hidden, + )?; self.last_hidden(s, t) } @@ -2292,7 +2480,17 @@ 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, None)?; + self.run_window( + s, + 0, + p, + false, + Some(&mut *state), + None, + gemm, + None, + WindowReadout::Hidden, + )?; Some(&*state) } None => None, @@ -2307,7 +2505,17 @@ 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, None)?; + self.run_window( + s, + p, + start, + false, + Some(&mut *state), + request, + gemm, + None, + WindowReadout::Hidden, + )?; Some(&*state) }; let rows = run @@ -2316,7 +2524,17 @@ 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, None)?; + self.run_window( + s, + start, + end, + false, + None, + from, + gemm, + None, + WindowReadout::Hidden, + )?; self.last_hidden(s, end) }) .collect::>>()?; @@ -2682,6 +2900,30 @@ impl Model { custom_positions: bool, bidirectional: bool, suffix: Option<&SuffixBatch<'_>>, + ) -> Result<()> { + self.run_with_layout( + s, + lengths, + custom_positions, + bidirectional, + suffix, + None, + None, + None, + ) + } + + #[allow(clippy::too_many_arguments)] // Original numerical layouts share one canonical layer loop. + fn run_with_layout( + &self, + s: &Scratch, + lengths: &[usize], + custom_positions: bool, + bidirectional: bool, + suffix: Option<&SuffixBatch<'_>>, + original_rows: Option<&[usize]>, + compaction: Option<&Compaction>, + tree: Option<&Tree>, ) -> Result<()> { let bidir = if bidirectional { ensure!( @@ -2707,6 +2949,7 @@ impl Model { let (ti, hi, eps) = (t 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); + let point_rows = compaction.map_or(t, |c| c.rows) as i32; let w = Widths::of(cfg); let p = |off: usize| s.at(off); let (cos, sin) = if custom_positions { @@ -2715,31 +2958,65 @@ impl Model { 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", - ) + let project = + |x: usize, weight: &Tensor, y: usize, group: Option, role: Projection| { + if let Some(tree) = tree { + return tree.project(self, s, x, weight, y, role); } - } 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) - } + if let Some(compact) = compaction { + return compact.project(self, s, x, weight, y, role); + } + 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 { + if let Some(rows) = original_rows { + self.gemm_sequences(s, x, weight, y, rows) + } else { + self.gemm_group(s, x, weight, y, t, lengths, group) + } + } else { + self.gemm_sequences(s, x, weight, y, lengths) + } + }; + // Batch only independent default 9B causal sequences. Other numerical + // modes and old libraries retain their original per-sequence entry points. + let attention_batch = if ( + cfg.hidden, + cfg.intermediate, + self.layers.len(), + cfg.heads, + cfg.kv_heads, + cfg.head_dim, + ) == (4096, 12288, 32, 16, 4, 256) + && self.packed_gemm.is_some() + && self.reference.is_none() + && self.lora.is_none() + && self.fp32_gdn.is_none() + && suffix.is_none() + && tree.is_none() + && !bidirectional + && (2..=16).contains(&lengths.len()) + { + cuda::api().batched_attention() + } else { + None }; let mut linear_index = 0; let mut attention_index = 0; @@ -2753,7 +3030,7 @@ impl Model { p(s.res), self.layers[0].input_norm.ptr, p(s.x), - ti, + point_rows, hi, eps, st, @@ -2772,110 +3049,74 @@ impl Model { for (i, layer) in self.layers.iter().enumerate() { match &layer.mixer { Mixer::Linear(la) => { - project(s.x, &la.in_proj, s.gdn_in, Some(Group::LinearInput))?; + project( + s.x, + &la.in_proj, + s.gdn_in, + Some(Group::LinearInput), + Projection::MixerInput, + )?; 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; - let b = z + vd * BF16; - let a = b + hv * BF16; - unsafe { - for &(offset, length) in &sequences { - check( - 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", - )?; - } - check( - (cuda::api().cs1_gdn_gates)( - p(b), - p(a), - ld, - la.a_log.ptr, - la.dt_bias.ptr, - p(s.beta), - p(s.g).cast(), - ti, - hv as i32, - st, - ), - "gdn gates", - )?; - if let Some(suffix) = suffix { + if let Some(tree) = tree { + tree.linear(self, s, la, lengths)?; + } else { + let ld = w.gdn_in as i32; + let z = s.gdn_in + w.conv * BF16; + let b = z + vd * BF16; + let a = b + hv * BF16; + unsafe { 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), - 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), - suffix.prefix.gdn_state[linear_index].at(0), - std::ptr::null_mut(), - st, - ), - "suffix GDN prefill", + 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", )?; } - } 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(), + (cuda::api().cs1_gdn_gates)( + p(b), + p(a), + ld, + la.a_log.ptr, + la.dt_bias.ptr, p(s.beta), - p(s.lo), - p(s.workspace).cast(), - s.workspace_bytes, - lengths.as_ptr(), - lengths.len(), + p(s.g).cast(), + ti, hv as i32, - cfg.lin_k_heads as i32, - (cfg.lin_k_dim as f32).powf(-0.5), st, ), - "batched GDN prefill", + "gdn gates", )?; - } else { - 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), @@ -2887,151 +3128,232 @@ 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 { - 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)( + p(s.lo), + p(z), + ld, + la.norm.ptr, + p(s.ln), + ti, + hv as i32, + cfg.lin_v_dim as i32, + eps, + st, + ), + "gated norm", + )?; } - check( - (cuda::api().cs1_gated_rms_norm)( - p(s.lo), - p(z), - ld, - la.norm.ptr, - p(s.ln), - ti, - hv as i32, - cfg.lin_v_dim as i32, - eps, - st, - ), - "gated norm", - )?; } - project(s.ln, &la.out, s.delta, None)?; + project(s.ln, &la.out, s.delta, None, Projection::MixerOutput)?; self.apply_lora(i, Group::LinearOutput, s, s.ln, s.delta, lengths)?; linear_index += 1; } Mixer::Full(fa) => { - project(s.x, &fa.qkv, s.attn_in, Some(Group::AttentionInput))?; + project( + s.x, + &fa.qkv, + s.attn_in, + Some(Group::AttentionInput), + Projection::MixerInput, + )?; 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; - let v = k + cfg.kv_heads * cfg.head_dim * BF16; - unsafe { - for &(offset, length) in &sequences { - check( - (self - .reference - .map_or(cuda::api().cs1_attn_prep, |r| r.attn_prep))( - p(s.attn_in + offset * w.attn_in * BF16), - p(k + offset * w.attn_in * BF16), - ld, - fa.q_norm.ptr, - fa.k_norm.ptr, - p(cos), - p(sin), - p(s.aq + offset * cfg.heads * cfg.head_dim * BF16), - p(s.agate + offset * cfg.heads * cfg.head_dim * BF16), - p(s.ak + offset * cfg.kv_heads * cfg.head_dim * BF16), - length, - hq, - hk, - hd, - cfg.rotary_half as i32, - eps, - st, - ), - "attention prep", - )?; - let scale = (cfg.head_dim as f32).powf(-0.5); - check( - 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), + if let Some(tree) = tree { + tree.attention(self, s, fa, lengths)?; + } else { + let ld = w.attn_in as i32; + let k = s.attn_in + w.attn_q * BF16; + let v = k + cfg.kv_heads * cfg.head_dim * BF16; + unsafe { + for &(offset, length) in &sequences { + check( + (self + .reference + .map_or(cuda::api().cs1_attn_prep, |r| r.attn_prep))( + p(s.attn_in + offset * w.attn_in * BF16), + p(k + 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)( + fa.q_norm.ptr, + fa.k_norm.ptr, + p(cos), + p(sin), 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), - length, - hq, - hk, - hd, - scale, - st, - ) - } else if let Some(reference) = self.reference { - (reference.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), length, hq, hk, - scale, - p(s.workspace), + hd, + cfg.rotary_half as i32, + eps, st, - ) - } else { - (cuda::api().cs1_attention_gated)( - 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), + ), + "attention prep", + )?; + // Every preparation writes disjoint rows. A single + // attention launch below waits for all of them. + if attention_batch.is_some() { + continue; + } + let scale = (cfg.head_dim as f32).powf(-0.5); + check( + 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), + 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), + length, + hq, + hk, + hd, + scale, + st, + ) + } else if let Some(reference) = self.reference { + (reference.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), + length, + hq, + hk, + scale, + p(s.workspace), + st, + ) + } else { + (cuda::api().cs1_attention_gated)( + 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), + length, + hq, + hk, + hd, + scale, + st, + ) + }, + "gated attention", + )?; + } + if let Some(batch) = attention_batch { + check( + batch( + p(s.aq), + p(s.ak), + p(v), ld, - p(s.agate + offset * cfg.heads * cfg.head_dim * BF16), - p(s.ao + offset * cfg.heads * cfg.head_dim * BF16), - length, + p(s.agate), + p(s.ao), + lengths.as_ptr(), + lengths.len(), hq, hk, hd, - scale, + (cfg.head_dim as f32).powf(-0.5), st, - ) - }, - "gated attention", - )?; + ), + "batched causal attention", + )?; + } } } - project(s.ao, &fa.o, s.delta, None)?; + project(s.ao, &fa.o, s.delta, None, Projection::MixerOutput)?; self.apply_lora(i, Group::AttentionOutput, s, s.ao, s.delta, lengths)?; attention_index += 1; } @@ -3050,7 +3372,7 @@ impl Model { p(s.delta), layer.post_norm.ptr, p(s.x), - ti, + point_rows, hi, eps, st, @@ -3058,7 +3380,13 @@ impl Model { "post-attention norm", )?; } - project(s.x, &layer.gate_up, s.gate_up, Some(Group::GateUp))?; + project( + s.x, + &layer.gate_up, + s.gate_up, + Some(Group::GateUp), + Projection::GateUp, + )?; self.apply_lora(i, Group::GateUp, s, s.x, s.gate_up, lengths)?; unsafe { check( @@ -3066,14 +3394,14 @@ impl Model { p(s.gate_up), (2 * cfg.intermediate) as i32, p(s.act), - ti, + point_rows, cfg.intermediate as i32, st, ), "silu mul", )?; } - project(s.act, &layer.down, s.delta, None)?; + project(s.act, &layer.down, s.delta, None, Projection::Down)?; self.apply_lora(i, Group::Down, s, s.act, s.delta, lengths)?; let next = self .layers @@ -3088,7 +3416,7 @@ impl Model { p(s.delta), next.ptr, p(s.x), - ti, + point_rows, hi, eps, st, @@ -3125,6 +3453,7 @@ impl Model { prefix: Option<&PrefixState>, gemm: *mut c_void, reuse: Option<(cuda::ReuseGemm, usize)>, + readout: WindowReadout, ) -> Result<()> { ensure!( self.fp32_gdn.is_none(), @@ -3136,6 +3465,11 @@ impl Model { && capture.as_deref().is_none_or(|cap| cap.len == tend), "prefix window shape" ); + ensure!( + readout != WindowReadout::StateOnly + || (capture.is_some() && prefix.is_none() && !custom_positions && reuse.is_some()), + "state-only execution requires an exact text-prefix capture" + ); let cfg = &self.cfg; let st = self.stream; let rows = tend - qb; @@ -3387,6 +3721,11 @@ impl Model { // K must be post-prep (normed + rotary); V is the raw column. cuda::copy_dd(cap.attn_kv[fa_i].0.at(0), p(s.ak), tend * kvb, st)?; cuda::copy2d(cap.attn_kv[fa_i].1.at(0), kvb, p(v), ab, kvb, tend, st)?; + if readout == WindowReadout::StateOnly && i + 1 == self.layers.len() { + // The suffix consumes K/V, not final prefix hidden rows. + // The caller synchronizes these copies before publishing state. + return Ok(()); + } } check( (cuda::api().cs1_attention_gated_prefix)( @@ -3419,6 +3758,9 @@ impl Model { fa_i += 1; } } + if readout == WindowReadout::StateOnly && i + 1 == self.layers.len() { + return Ok(()); + } unsafe { check( (self @@ -3517,3 +3859,15 @@ mod prefix_group_tests; #[cfg(test)] #[path = "../../../../../tests/qwen3_5/last_mlp.rs"] mod last_mlp_tests; + +#[cfg(test)] +#[path = "../../../../../tests/qwen3_5/capture_only.rs"] +mod capture_only_tests; + +#[cfg(test)] +#[path = "../../../../../tests/qwen3_5/prefix_tree.rs"] +mod prefix_tree_tests; + +#[cfg(test)] +#[path = "../../../../../tests/qwen3_5/tree_graph.rs"] +mod tree_graph_tests; diff --git a/src/models/qwen3_5/native/src/prefix_compaction.rs b/src/models/qwen3_5/native/src/prefix_compaction.rs new file mode 100644 index 00000000..887f3841 --- /dev/null +++ b/src/models/qwen3_5/native/src/prefix_compaction.rs @@ -0,0 +1,339 @@ +//! Share pointwise prefix rows while retaining each complete sequence's mixers. +use super::*; + +#[derive(Clone, Copy)] +pub(super) enum Projection { + MixerInput, + MixerOutput, + GateUp, + Down, +} + +pub(super) struct Compaction { + pub rows: usize, + full_rows: usize, + original_output_rows: usize, + expand: *const i32, + gather: *const i32, + gemm: cuda::ReuseGemm, +} + +struct RowMaps { + expand: Vec, + gather: Vec, + ends: Vec, +} + +fn row_maps(lengths: &[usize], prefix: usize) -> RowMaps { + let mut expand = Vec::new(); + let mut gather: Vec = (0..prefix as i32).collect(); + let mut ends = Vec::new(); + let mut full_start = 0; + let mut compact_start = prefix; + for &length in lengths { + expand.extend(0..prefix as i32); + expand.extend((compact_start..compact_start + length - prefix).map(|n| n as i32)); + gather.extend((full_start + prefix..full_start + length).map(|n| n as i32)); + compact_start += length - prefix; + full_start += length; + ends.push(compact_start); + } + RowMaps { + expand, + gather, + ends, + } +} + +impl Compaction { + /// Scratch has full sequence capacity; gate_up is temporary until the MLP. + pub(super) fn project( + &self, + model: &Model, + s: &Scratch, + x: usize, + weight: &Tensor, + y: usize, + role: Projection, + ) -> Result<()> { + let (n, k) = (weight.shape[0], weight.shape[1]); + let (source, target, original_rows) = match role { + Projection::MixerInput => (x, s.gate_up, self.full_rows), + Projection::MixerOutput => (s.x, y, self.original_output_rows), + Projection::GateUp => (x, y, self.full_rows), + Projection::Down => (x, y, self.original_output_rows), + }; + // SAFETY: checked maps index the full mixer rows or compact pointwise rows. + // All buffers are distinct scratch regions; the stream orders their reuse. + unsafe { + if matches!(role, Projection::MixerOutput) { + check( + (cuda::api().cs1_embed)( + self.gather, + s.at(x), + s.at(s.x), + self.rows as i32, + k as i32, + model.stream, + ), + "compact mixer output", + )?; + } + check( + (self.gemm.run)( + model.gemm, + s.at(source), + weight.ptr, + s.at(target), + self.rows as i32, + n as i32, + k as i32, + original_rows as i32, + model.stream, + ), + "compact retained projection", + )?; + if matches!(role, Projection::MixerInput) { + check( + (cuda::api().cs1_embed)( + self.expand, + s.at(target), + s.at(y), + self.full_rows as i32, + n as i32, + model.stream, + ), + "expand mixer input", + )?; + } + } + Ok(()) + } +} + +impl Model { + /// Capture only the pointwise-compacted layout. Model-owned maps and scratch + /// remain live; every request replaces values before replay. Prefix trees use + /// their eager path and never reference these graph keys. + fn run_compaction( + &mut self, + lengths: &[usize], + prefix: usize, + layout: &Compaction, + ) -> Result<()> { + let run = |model: &Self| { + model.run_with_layout( + model.scratch.as_ref().unwrap(), + lengths, + false, + false, + None, + None, + Some(layout), + None, + ) + }; + if !self.graph_enabled { + return run(self); + } + // Normal text keys contain positive lengths or zero-delimited original + // groups. The sentinel separates compact layouts even at equal row counts. + let mut key = vec![usize::MAX, 0, prefix]; + key.extend_from_slice(lengths); + if let Some((_, graph)) = self.graphs.iter().find(|(shape, _)| *shape == key) { + graph.launch(self.stream)?; + self.graph_stats.replays += 1; + return Ok(()); + } + run(self)?; + // Record while eager kernels execute; first capture must keep the eager + // result, since replay here would advance the residual twice. + let captured = cuda::Graph::capture(self.stream, || run(self)); + self.cache_graph(&key, false, layout.full_rows, captured) + } + + /// Share pointwise prefix rows, keeping original full-length causal mixers. + /// Unsupported algorithms decline before any request values are computed. + /// The pinned9B widths guarantee mixer inputs fit the gate/up temporary and + /// gathered mixer outputs fit the compact hidden buffer. + pub fn try_forward_prefix_compaction( + &mut self, + inputs: &[&[u32]], + prefix: usize, + ) -> Result>>> { + if self.cfg.hidden != 4096 + || self.cfg.intermediate != 12288 + || self.layers.len() != 32 + || ( + self.cfg.heads, + self.cfg.kv_heads, + self.cfg.head_dim, + self.cfg.lin_k_heads, + self.cfg.lin_v_heads, + self.cfg.lin_k_dim, + self.cfg.lin_v_dim, + ) != (16, 4, 256, 16, 32, 128, 128) + || self.packed_gemm.is_none() + || self.reference.is_some() + || self.lora.is_some() + || self.fp32_gdn.is_some() + { + return Ok(None); + } + ensure!( + (2..=16).contains(&inputs.len()) && (16..512).contains(&prefix), + "invalid projection prefix" + ); + ensure!( + inputs + .iter() + .all(|ids| ids.len() > prefix && ids.len() <= self.cfg.max_positions) + && inputs + .iter() + .all(|ids| ids[..prefix] == inputs[0][..prefix]), + "projection prefix differs or leaves an empty suffix" + ); + ensure!( + inputs + .iter() + .flat_map(|ids| ids.iter()) + .all(|&id| (id as usize) < self.embed.shape[0]), + "token id outside the vocabulary" + ); + let lengths: Vec = inputs.iter().map(|ids| ids.len()).collect(); + let full_rows = lengths + .iter() + .try_fold(0usize, |n, &m| n.checked_add(m)) + .context("projection row overflow")?; + if full_rows > 4096 { + return Ok(None); + } + let rows = full_rows - (inputs.len() - 1) * prefix; + let Some(gemm) = cuda::api().reuse_gemm() else { + return Ok(None); + }; + cuda::set_device(0)?; + 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, packed) in [ + (input, true), + (output, false), + (&layer.gate_up, true), + (&layer.down, false), + ] { + shapes.insert((weight.shape[0], weight.shape[1], packed)); + } + } + for (n, k, packed) in shapes { + let originals: &[usize] = if packed { + std::slice::from_ref(&full_rows) + } else { + &lengths + }; + // SAFETY: shape-only check reads live arrays; all projections are checked + // before embedding or modifying residuals, including original split-K. + let rc = unsafe { + (gemm.check)( + self.gemm, + originals.as_ptr(), + originals.len(), + &rows, + 1, + n as i32, + k as i32, + ) + }; + if rc == -1 { + return Ok(None); + } + check(rc, "projection compaction compatibility")?; + } + let maps = row_maps(&lengths, prefix); + let capacity = if self.graph_enabled { + 8192 + } else { + full_rows + rows + }; + if self + .compaction_maps + .as_ref() + .is_none_or(|(n, _)| *n < capacity) + { + // First allocation cannot invalidate any captured map address. Only + // replacement of an existing map allocation requires graph retirement. + if self.compaction_maps.is_some() { + let retired = self.synchronize(); + if !self.graphs.is_empty() { + self.graph_stats.invalidations += 1; + } + self.graphs.clear(); + retired?; + } + self.compaction_maps = Some((capacity, DeviceBuffer::new(capacity * 4)?)); + } + self.prepare_scratch(full_rows)?; + let map_start = self.compaction_maps.as_ref().unwrap().1.at(0); + let map_gather = self.compaction_maps.as_ref().unwrap().1.at(full_rows * 4); + let bytes: Vec = maps + .expand + .iter() + .chain(&maps.gather) + .flat_map(|n| n.to_le_bytes()) + .collect(); + let layout = Compaction { + rows, + full_rows, + original_output_rows: lengths[0], + expand: map_start.cast(), + gather: map_gather.cast(), + gemm, + }; + let ids: Vec = inputs[0][..prefix] + .iter() + .copied() + .chain(inputs.iter().flat_map(|ids| ids[prefix..].iter().copied())) + .collect(); + let result = (|| { + // SAFETY: persistent model-owned buffer holds exactly these live row maps. + unsafe { + cuda::upload(map_start, &bytes, self.stream)?; + } + let s = self.scratch.as_ref().unwrap(); + self.embed_tokens(s, &ids)?; + self.run_compaction(&lengths, prefix, &layout)?; + let s = self.scratch.as_ref().unwrap(); + maps.ends + .iter() + .map(|&end| self.last_hidden(s, end)) + .collect::>>() + })(); + // No caller may reuse maps/scratch while a failed operation still owns them. + self.synchronize()?; + result.map(Some) + } +} + +#[cfg(test)] +mod tests { + use super::*; + #[test] + fn maps_preserve_every_suffix_and_restore_original_row_order() { + let maps = row_maps(&[5, 6, 4], 2); + assert_eq!(maps.expand, [0, 1, 2, 3, 4, 0, 1, 5, 6, 7, 8, 0, 1, 9, 10]); + assert_eq!(maps.gather, [0, 1, 2, 3, 4, 7, 8, 9, 10, 13, 14]); + assert_eq!(maps.ends, [5, 9, 11]); + let full = [10, 11, 12, 13, 14, 10, 11, 22, 23, 24, 25, 10, 11, 32, 33]; + let compact: Vec<_> = maps.gather.iter().map(|&i| full[i as usize]).collect(); + assert_eq!( + maps.expand + .iter() + .map(|&i| compact[i as usize]) + .collect::>(), + full + ); + } +} diff --git a/src/models/qwen3_5/native/src/prefix_tree.rs b/src/models/qwen3_5/native/src/prefix_tree.rs new file mode 100644 index 00000000..5343c87f --- /dev/null +++ b/src/models/qwen3_5/native/src/prefix_tree.rs @@ -0,0 +1,479 @@ +//! A request-local prefix tree shares all layer work, retaining original GEMM plans. +use super::*; + +pub(super) struct Tree { + rows: usize, + prefix: usize, + sequences: Vec<(usize, usize)>, + input_rows: usize, + output_rows: usize, + gemm: cuda::ReuseGemm, + ops: cuda::SharedSuffix, + gdn: cuda::SharedBatchedGdn, + attention_batch: Option, + state: *mut c_void, + values: *mut c_void, +} + +impl Tree { + pub(super) fn project( + &self, + model: &Model, + s: &Scratch, + x: usize, + weight: &Tensor, + y: usize, + role: Projection, + ) -> Result<()> { + let original = match role { + Projection::MixerInput | Projection::GateUp => self.input_rows, + Projection::MixerOutput | Projection::Down => self.output_rows, + }; + // SAFETY: every weight shape and original group was checked before embedding. + // The compact rows and output dimensions fit distinct scratch regions. + unsafe { + check( + (self.gemm.run)( + model.gemm, + s.at(x), + weight.ptr, + s.at(y), + self.rows as i32, + weight.shape[0] as i32, + weight.shape[1] as i32, + original as i32, + model.stream, + ), + "tree retained projection", + ) + } + } + + pub(super) fn linear( + &self, + model: &Model, + s: &Scratch, + la: &LinearAttention, + lengths: &[usize], + ) -> Result<()> { + let cfg = &model.cfg; + let w = Widths::of(cfg); + let (kd, vd, hv) = (cfg.key_dim(), cfg.value_dim(), cfg.lin_v_heads); + let (ld, st) = (w.gdn_in as i32, model.stream); + let p = |off| s.at(off); + let z = s.gdn_in + w.conv * BF16; + let b = z + vd * BF16; + let a = b + hv * BF16; + // SAFETY: scratch holds one prefix then disjoint suffixes. Convolution + // history reads exactly three prefix rows; the model-owned FP32 state is + // written before every layer's suffixes and reused only after they finish. + unsafe { + for &(offset, length) in &self.sequences { + check( + if offset == 0 { + (model.convolution())( + p(s.gdn_in), + ld, + la.conv.ptr, + p(s.lq), + p(s.lk), + p(s.lv), + length as i32, + kd as i32, + vd as i32, + st, + ) + } else { + (self.ops.conv)( + p(s.gdn_in + offset * w.gdn_in * BF16), + p(s.gdn_in + (self.prefix - 3) * 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 as i32, + kd as i32, + vd as i32, + st, + ) + }, + "tree convolution", + )?; + } + check( + (cuda::api().cs1_gdn_gates)( + p(b), + p(a), + ld, + la.a_log.ptr, + la.dt_bias.ptr, + p(s.beta), + p(s.g).cast(), + self.rows as i32, + hv as i32, + st, + ), + "tree gates", + )?; + let scale = (cfg.lin_k_dim as f32).powf(-0.5); + check( + (cuda::api().cs1_gdn_prefill_x)( + p(s.lq), + p(s.lk), + p(s.lv), + p(s.g).cast(), + p(s.beta), + p(s.lo), + p(s.workspace).cast(), + self.prefix as i32, + hv as i32, + cfg.lin_k_heads as i32, + scale, + std::ptr::null(), + self.state, + st, + ), + "tree prefix state", + )?; + let offset = self.prefix; + check( + (self.gdn)( + 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(), + s.workspace_bytes, + lengths[1..].as_ptr(), + lengths.len() - 1, + hv as i32, + cfg.lin_k_heads as i32, + scale, + self.state, + st, + ), + "tree suffix states", + )?; + check( + (cuda::api().cs1_gated_rms_norm)( + p(s.lo), + p(z), + ld, + la.norm.ptr, + p(s.ln), + self.rows as i32, + hv as i32, + cfg.lin_v_dim as i32, + cfg.eps, + st, + ), + "tree gated norm", + ) + } + } + + pub(super) fn attention( + &self, + model: &Model, + s: &Scratch, + fa: &FullAttention, + lengths: &[usize], + ) -> Result<()> { + let cfg = &model.cfg; + let w = Widths::of(cfg); + let (hq, hk, hd) = (cfg.heads as i32, cfg.kv_heads as i32, cfg.head_dim as i32); + let (ld, st) = (w.attn_in as i32, model.stream); + let qb = cfg.heads * cfg.head_dim * BF16; + let kb = cfg.kv_heads * cfg.head_dim * BF16; + let p = |off| s.at(off); + let k = s.attn_in + w.attn_q * BF16; + let v = k + kb; + let scale = (cfg.head_dim as f32).powf(-0.5); + // SAFETY: custom rotary rows encode each branch's original position. + // Prefix K is contiguous in scratch; V is copied out of strided QKV once. + // Every suffix sees only that prefix and its own K/V, never another branch. + unsafe { + check( + (cuda::api().cs1_attn_prep)( + p(s.attn_in), + p(k), + ld, + fa.q_norm.ptr, + fa.k_norm.ptr, + p(s.custom_cos), + p(s.custom_sin), + p(s.aq), + p(s.agate), + p(s.ak), + self.rows as i32, + hq, + hk, + hd, + cfg.rotary_half as i32, + cfg.eps, + st, + ), + "tree attention prep", + )?; + cuda::copy2d(self.values, kb, p(v), w.attn_in * BF16, kb, self.prefix, st)?; + for &(offset, length) in &self.sequences { + if offset != 0 && self.attention_batch.is_some() { + break; + } + check( + if offset == 0 { + (cuda::api().cs1_attention_gated)( + p(s.aq), + p(s.ak), + p(v), + ld, + p(s.agate), + p(s.ao), + length as i32, + hq, + hk, + hd, + scale, + st, + ) + } else { + (self.ops.attention)( + p(s.aq + offset * qb), + p(s.ak + offset * kb), + p(v + offset * w.attn_in * BF16), + ld, + p(s.agate + offset * qb), + p(s.ao + offset * qb), + p(s.ak), + self.values, + length as i32, + self.prefix as i32, + hq, + hk, + hd, + scale, + st, + ) + }, + "tree attention", + )?; + } + if let Some(batch) = self.attention_batch { + let offset = self.prefix; + check( + batch( + p(s.aq + offset * qb), + p(s.ak + offset * kb), + p(v + offset * w.attn_in * BF16), + ld, + p(s.agate + offset * qb), + p(s.ao + offset * qb), + p(s.ak), + self.values, + lengths[1..].as_ptr(), + lengths.len() - 1, + self.prefix as i32, + hq, + hk, + hd, + scale, + st, + ), + "batched tree suffix attention", + )?; + } + } + Ok(()) + } +} + +impl Model { + /// Compute a 64-aligned common prefix once in each layer, together with its + /// candidate suffixes. Original input groups and per-candidate output plans + /// are checked separately. Unsupported shapes decline before request work. + pub fn try_forward_prefix_tree( + &mut self, + inputs: &[&[u32]], + prefix: usize, + input_rows: &[usize], + ) -> Result> { + if ( + self.cfg.hidden, + self.cfg.intermediate, + self.layers.len(), + self.cfg.heads, + self.cfg.kv_heads, + self.cfg.head_dim, + self.cfg.lin_k_heads, + self.cfg.lin_v_heads, + self.cfg.lin_k_dim, + self.cfg.lin_v_dim, + ) != (4096, 12288, 32, 16, 4, 256, 16, 32, 128, 128) + || self.packed_gemm.is_none() + || self.batched_gdn.is_none() + || self.reference.is_some() + || self.lora.is_some() + || self.fp32_gdn.is_some() + { + return Ok(None); + } + ensure!( + (2..=16).contains(&inputs.len()) + && input_rows.len() == inputs.len() + && prefix >= GDN_CHUNK + && prefix.is_multiple_of(GDN_CHUNK), + "invalid prefix tree layout" + ); + ensure!( + inputs.iter().all(|ids| ids.len() > prefix + && ids.len() <= self.cfg.max_positions + && ids[..prefix] == inputs[0][..prefix]), + "prefix tree input mismatch" + ); + ensure!( + inputs + .iter() + .flat_map(|ids| ids.iter()) + .all(|&id| (id as usize) < self.embed.shape[0]), + "token outside vocabulary" + ); + ensure!( + inputs + .iter() + .zip(input_rows) + .all(|(ids, &rows)| rows >= ids.len() && rows <= i32::MAX as usize), + "invalid original input projection rows" + ); + let output_rows: Vec<_> = inputs.iter().map(|ids| ids.len()).collect(); + let lengths: Vec<_> = std::iter::once(prefix) + .chain(inputs.iter().map(|ids| ids.len() - prefix)) + .collect(); + let rows: usize = lengths.iter().sum(); + // Bound scratch independently of candidate count; oversized requests keep + // the existing segmented shared-prefix path. + if rows > 8192 { + return Ok(None); + } + let (Some(gemm), Some(ops), Some(gdn)) = ( + cuda::api().reuse_gemm(), + cuda::api().shared_suffix(), + cuda::api().shared_batched_gdn(), + ) else { + return Ok(None); + }; + cuda::set_device(0)?; + 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, packed) in [ + (input, true), + (output, false), + (&layer.gate_up, true), + (&layer.down, false), + ] { + shapes.insert((weight.shape[0], weight.shape[1], packed)); + } + } + for (n, k, packed) in shapes { + let originals = if packed { input_rows } else { &output_rows }; + // SAFETY: shape-only checks consume live host arrays; no request data + // is executed unless every changed projection is compatible. + let rc = unsafe { + (gemm.check)( + self.gemm, + originals.as_ptr(), + originals.len(), + &rows, + 1, + n as i32, + k as i32, + ) + }; + if rc == -1 { + return Ok(None); + } + check(rc, "prefix tree projection compatibility")?; + } + let state_bytes = self.cfg.lin_v_heads * self.cfg.lin_k_dim * self.cfg.lin_v_dim * F32; + let bytes = state_bytes + prefix * self.cfg.kv_heads * self.cfg.head_dim * BF16; + if self + .tree_workspace + .as_ref() + .is_none_or(|(cap, _)| *cap < bytes) + { + self.synchronize()?; + self.tree_workspace = Some((bytes, DeviceBuffer::new(bytes)?)); + } + self.prepare_scratch(rows)?; + let buffer = &self.tree_workspace.as_ref().unwrap().1; + let mut offset = 0; + let sequences: Vec<_> = lengths + .iter() + .map(|&length| { + let start = offset; + offset += length; + (start, length) + }) + .collect(); + let tree = Tree { + rows, + prefix, + sequences, + input_rows: input_rows[0], + output_rows: output_rows[0], + gemm, + ops, + gdn, + attention_batch: cuda::api().shared_batched_attention(), + state: buffer.at(0), + values: buffer.at(state_bytes), + }; + let ids: Vec<_> = inputs[0][..prefix] + .iter() + .copied() + .chain(inputs.iter().flat_map(|ids| ids[prefix..].iter().copied())) + .collect(); + let result = (|| { + let s = self.scratch.as_ref().unwrap(); + let rb = self.cfg.rotary_half * BF16; + for &(offset, length) in &tree.sequences { + let position = if offset == 0 { 0 } else { prefix }; + // SAFETY: static tables cover every original position. Compact + // destination rows are disjoint and their total is scratch capacity. + unsafe { + cuda::copy_dd( + s.at(s.custom_cos + offset * rb), + s.at(s.cos + position * rb), + length * rb, + self.stream, + )?; + cuda::copy_dd( + s.at(s.custom_sin + offset * rb), + s.at(s.sin + position * rb), + length * rb, + self.stream, + )?; + } + } + self.embed_tokens(s, &ids)?; + self.run_with_layout(s, &lengths, true, false, None, None, None, Some(&tree))?; + let hidden = tree.sequences[1..] + .iter() + .map(|&(offset, length)| self.last_hidden(s, offset + length)) + .collect::>()?; + Ok(SharedPrefixOutput { + hidden, + saved_tokens: (inputs.len() - 1) * prefix, + }) + })(); + // Error paths must retire all consumers before scratch/workspace is reused. + self.synchronize()?; + result.map(Some) + } +} diff --git a/src/runtime/Cargo.toml b/src/runtime/Cargo.toml index 6fa06401..3b29d92f 100644 --- a/src/runtime/Cargo.toml +++ b/src/runtime/Cargo.toml @@ -27,3 +27,7 @@ path = "../../tests/runtime/admission.rs" [[test]] name = "control" path = "../../tests/runtime/control.rs" + +[[test]] +name = "batching" +path = "../../tests/runtime/batching.rs" diff --git a/src/runtime/README.md b/src/runtime/README.md index 00e13d50..d1a51415 100644 --- a/src/runtime/README.md +++ b/src/runtime/README.md @@ -1,13 +1,13 @@ # Native execution runtime -`omni-runtime` provides `SerialScheduler` for the Cua-S1, Open-Jev, Laya, -Decider, and Clef workers. Cua-S1, Open-Jev, Laya, and Decider assemble a +`omni-runtime` provides `SerialScheduler` for native workers and `BatchScheduler` +for Open-Jev-9B default execution. Cua-S1, Open-Jev, Laya, and Decider assemble a processor, one scheduler for the loaded executor, and the executor. HTTP handlers still coordinate `prepare` → `execute` → `finish`. Clef uses the same scheduler. Its admitted closure is the frame RPC to its Python executor and does not call CUDA. -The scheduler owns FIFO admission and blocking dispatch. A caller waits +The serial scheduler owns FIFO admission and blocking dispatch. A caller waits asynchronously for one execution permit before entering Tokio's blocking pool. Cloning a scheduler shares that admission queue; separate loaded executors have separate queues. The model mutex remains a guard for Rust's mutable executor @@ -16,14 +16,15 @@ state. Waiting requests no longer occupy blocking threads waiting for that mutex | Worker | Admitted unit | Work outside admission | | --- | --- | --- | | Cua-S1 | One question's unpadded Qwen forward. Release before admitting its next question. | Request preparation, CPU letter projection and response finishing. | -| Open-Jev | One complete request's independent candidate forwards and CPU scalar heads. | Request preparation and calibrated response finishing. | +| Open-Jev | 9B default: original candidate batches or one exclusive exact-prefix request through `BatchScheduler`; other modes: one complete request through `SerialScheduler`. | Request preparation and calibrated response finishing. | | Laya | One complete padded request, including GPU scorer/action head and copyback. | Request preparation/padding and calibrated response finishing. | | Decider | One complete request's independent rows, including within-request packing and LM-head projection. | Request preparation and calibrated response finishing. An optional request reservation covers both; token reservation begins after preparation. | | Clef | One complete text request. The closure writes one frame and waits for that same executor process. | HTTP body reads and the cached health response. CUDA stays on the executor process main thread, including the warmup that runs before the port opens. | Models retain weights, heads, device state, scratch buffers and graph caches. -The scheduler takes an owned closure and returns its result; it does not know -model identifiers, question/candidate mappings, tensor layouts or backend APIs. +Serial scheduling takes an owned closure; batch scheduling takes owned items and +one executor callback. Neither scheduler knows model identifiers, +question/candidate mappings, tensor layouts or backend APIs. The existing numerical, ordering, whole-request failure and readiness contracts remain in the processors, executors and worker handlers. @@ -74,14 +75,35 @@ See the [Decider recipe](../../recipe/decider/README.md) for configuration, health fields and overload/deadline responses. Other workers keep their current admission behavior unless they explicitly opt in. -Shared processing orchestration, compatibility grouping, cross-request dynamic -batching and a shared worker drain lifecycle remain planned. Models may pack -supported inputs within one request; this scheduler does not assemble GPU batches. -The Qwen executor supports single prompts, bounded packed prefill (Open-Jev -candidates or Decider complete rows), and request-local shared prefixes -(Decider opt-in). Graph, fixed and shared execution retain the same complete -request admission unit and context lifetime; the scheduler does not construct -GPU batches or coordinate separately running workers. +## Cross-request batching for Open-Jev-9B + +`BatchScheduler` owns a FIFO queue and one execution thread per loaded executor. +An idle worker runs immediately. It combines work already in the queue without a +batching delay. Each item declares an item count, a cost and whether it must run +alone; the generic scheduler enforces the configured budgets. An oversized atomic +item runs alone, without truncation. Queue admission itself is not capacity-limited. + +The 9B adapter preserves its original batches of at most 16 candidates and 4096 +tokens as atomic items. A GPU dispatch may combine these items up to 16 candidates +and 8192 tokens. Input/gate-up projections retain the original batch row counts; +output/down projections retain individual candidate lengths. Compatible GEMMs can +share a call, and incompatible ones keep separate calls. Prefix reuse is an +exclusive item using the existing request-local state path. Other Open-Jev modes +and checkpoints keep serial execution. Complete request outputs are reconstructed +before per-question calibration. + +Cancellation before the first item dispatch discards that queued request. Once +any item starts, all remaining items and resources stay alive through completion, +even when the caller disconnects. Errors fail affected complete requests; unrelated +queued requests remain ordered. The model callback synchronizes partially submitted +GPU work on errors or panics, and a panic closes the execution queue. The callback +and its model resources remain owned by the worker until it finishes. Dropping all +scheduler handles closes admission and joins the execution thread, waiting for +dispatched work and callback resources to retire before the owner exits. + +Shared processing orchestration and cross-worker coordination remain planned. +Batch layout keys include both candidate lengths and original input projection +groups, so CUDA Graph replay cannot conflate different numerical layouts. CPU tests cover FIFO serialization, queued cancellation, retention of admission and resources after dispatched cancellation, errors/panics, independent loaded diff --git a/src/runtime/src/batching.rs b/src/runtime/src/batching.rs new file mode 100644 index 00000000..bf617c2c --- /dev/null +++ b/src/runtime/src/batching.rs @@ -0,0 +1,220 @@ +//! FIFO batching of independent items with request-scoped output routing. + +use std::collections::VecDeque; +use std::panic::{AssertUnwindSafe, catch_unwind}; +use std::sync::{Arc, mpsc}; + +use anyhow::{Context, Result, anyhow, ensure}; +use tokio::sync::oneshot; + +/// A single oversized item runs alone; items are never truncated or split. +#[derive(Clone, Copy)] +pub struct BatchLimits { + pub items: usize, + pub cost: usize, +} + +/// Model-provided resource weight of one atomic execution item. +/// Exclusive items cannot share a dispatch, even when their costs fit. +#[derive(Clone, Copy)] +pub struct BatchCost { + pub items: usize, + pub cost: usize, + pub exclusive: bool, +} + +impl From for BatchCost { + fn from(cost: usize) -> Self { + Self { + items: 1, + cost, + exclusive: false, + } + } +} + +struct Request { + pending: VecDeque, + outputs: Vec, + response: oneshot::Sender>>, + started: bool, +} + +/// One owning execution thread per loaded executor. Already queued requests can +/// share a batch. An idle worker dispatches immediately, without a batching timer. +/// Cancellation before dispatch discards the request; after its first item is +/// dispatched, the complete request retains its inputs until execution finishes. +pub struct BatchScheduler { + worker: Arc>, +} + +struct Worker { + sender: Option>>, + thread: Option>, +} + +impl Drop for Worker { + fn drop(&mut self) { + // Close admission before joining, so an idle receiver wakes and any + // dispatched work and callback-owned device resources retire before exit. + self.sender.take(); + if let Some(thread) = self.thread.take() + && thread.thread().id() != std::thread::current().id() + { + let _ = thread.join(); + } + } +} + +impl Clone for BatchScheduler { + fn clone(&self) -> Self { + Self { + worker: self.worker.clone(), + } + } +} + +impl BatchScheduler { + /// The callback owns device execution and must synchronize before returning, + /// including errors. Its outputs must correspond to its inputs in order. + /// Dropping the last scheduler closes admission and waits for the worker and + /// callback resources to retire. Callbacks must not depend on that dropping + /// thread to unblock their execution; async callers can drop off-thread. + pub fn new( + limits: BatchLimits, + cost: impl Fn(&I) -> BatchCost + Send + 'static, + execute: impl FnMut(&[I]) -> Result> + Send + 'static, + ) -> Result { + ensure!(limits.items > 0 && limits.cost > 0, "empty batch budget"); + let (sender, receiver) = mpsc::channel(); + let thread = std::thread::Builder::new() + .name("omni-batch".into()) + .spawn(move || dispatch(receiver, limits, cost, execute)) + .context("start batch execution thread")?; + Ok(Self { + worker: Arc::new(Worker { + sender: Some(sender), + thread: Some(thread), + }), + }) + } + + pub async fn run(&self, items: Vec) -> Result> { + if items.is_empty() { + return Ok(Vec::new()); + } + let (response, receiver) = oneshot::channel(); + self.worker + .sender + .as_ref() + .expect("live scheduler owns its sender") + .send(Request { + outputs: Vec::with_capacity(items.len()), + pending: items.into(), + response, + started: false, + }) + .map_err(|_| anyhow!("batch executor is unavailable"))?; + receiver.await.context("batch execution stopped")? + } +} + +fn dispatch( + receiver: mpsc::Receiver>, + limits: BatchLimits, + cost: impl Fn(&I) -> BatchCost, + mut execute: impl FnMut(&[I]) -> Result>, +) { + let mut pending = None; + loop { + let mut inputs = Vec::with_capacity(limits.items); + let mut parts = Vec::new(); + let mut total = 0usize; + let mut item_total = 0usize; + let mut exclusive = false; + loop { + let next = pending.take().or_else(|| { + if inputs.is_empty() { + receiver.recv().ok() + } else { + receiver.try_recv().ok() + } + }); + let Some(mut request) = next else { break }; + if !request.started && request.response.is_closed() { + continue; + } + let start = inputs.len(); + while let Some(item) = request.pending.front() { + let weight = cost(item); + if !inputs.is_empty() + && (inputs.len() == limits.items + || weight.exclusive + || exclusive + || item_total >= limits.items + || weight.items > limits.items.saturating_sub(item_total) + || total > limits.cost + || weight.cost > limits.cost.saturating_sub(total)) + { + break; + } + total = total.saturating_add(weight.cost); + item_total = item_total.saturating_add(weight.items); + exclusive = weight.exclusive; + inputs.push(request.pending.pop_front().expect("checked pending item")); + request.started = true; + } + let count = inputs.len() - start; + if count == 0 { + pending = Some(request); + break; + } + let incomplete = !request.pending.is_empty(); + parts.push((request, count)); + if incomplete + || exclusive + || inputs.len() == limits.items + || item_total >= limits.items + || total >= limits.cost + { + break; + } + } + if inputs.is_empty() { + return; + } + let result = catch_unwind(AssertUnwindSafe(|| execute(&inputs))); + let panicked = result.is_err(); + let result = result + .unwrap_or_else(|_| Err(anyhow!("batch executor panicked"))) + .and_then(|outputs| { + ensure!(outputs.len() == inputs.len(), "batch output count mismatch"); + Ok(outputs) + }); + match result { + Ok(outputs) => { + let mut outputs = outputs.into_iter(); + for (mut request, count) in parts { + request.outputs.extend(outputs.by_ref().take(count)); + if request.pending.is_empty() { + let _ = request.response.send(Ok(request.outputs)); + } else { + debug_assert!(pending.is_none()); + pending = Some(request); + } + } + } + Err(error) => { + let message = format!("{error:#}"); + for (request, _) in parts { + let _ = request.response.send(Err(anyhow!(message.clone()))); + } + } + } + // A panicked device executor may have invalid state. Close its queue; + // callers receive errors and no new work runs on that executor. + if panicked { + return; + } + } +} diff --git a/src/runtime/src/lib.rs b/src/runtime/src/lib.rs index 1f0e7d2b..aac9a741 100644 --- a/src/runtime/src/lib.rs +++ b/src/runtime/src/lib.rs @@ -1,6 +1,7 @@ //! Shared admission, dispatch and the worker-side engine contract for a loaded executor. pub mod admission; +pub mod batching; pub mod control; pub mod engine; diff --git a/tests/open_jev/batching.rs b/tests/open_jev/batching.rs index 4de2af2c..cda1e2ac 100644 --- a/tests/open_jev/batching.rs +++ b/tests/open_jev/batching.rs @@ -29,3 +29,33 @@ fn shared_prefix_is_exact_aligned_and_leaves_a_suffix() { assert_eq!(shared_prefix(&[&a, &a[..2048]]), 0); assert_eq!(shared_prefix(&vec![a.as_slice(); 17]), 0); } + +#[test] +fn packed_prefix_admits_midlength_without_changing_legacy_policy() { + let a = vec![42u32; 801]; + let mut b = vec![42u32; 832]; + b[790] = 7; + assert_eq!(shared_prefix(&[&a, &b]), 0); + assert_eq!(super::packed_shared_prefix(&[&a, &b]), 768); + assert_eq!(super::packed_shared_prefix(&[&a, &b[..511]]), 0); + assert_eq!(super::packed_shared_prefix(&[&a, &b[..512]]), 0); + assert_eq!(super::packed_shared_prefix(&[&a]), 0); + assert_eq!(super::packed_shared_prefix(&vec![a.as_slice(); 17]), 0); +} + +#[test] +fn projection_prefix_keeps_unaligned_rows_and_original_batch_bounds() { + let a = vec![42; 128]; + let mut b = a.clone(); + b[43] = 7; + assert_eq!(super::projection_prefix(&[&a, &b]), 43); + assert_eq!(super::packed_shared_prefix(&[&a, &b]), 0); + assert_eq!(super::projection_prefix(&[&a, &a]), 127); + assert_eq!(super::projection_prefix(&[&a]), 0); + assert_eq!(super::projection_prefix(&vec![a.as_slice(); 17]), 0); + assert_eq!(super::projection_prefix(&[&a, &[]]), 0); + let big = vec![42; 3000]; + assert_eq!(super::projection_prefix(&[&big, &big]), 0); + b[15] = 8; + assert_eq!(super::projection_prefix(&[&a, &b]), 0); +} diff --git a/tests/open_jev/contract.rs b/tests/open_jev/contract.rs index 7a6d2ab0..877d8c59 100644 --- a/tests/open_jev/contract.rs +++ b/tests/open_jev/contract.rs @@ -131,9 +131,11 @@ fn exports_select_their_pinned_checkpoint() { } #[test] -fn only_validated_checkpoints_pack_candidates() { +fn packing_and_legacy_prefix_reuse_are_independent() { assert!(OPEN_JEV_27B.pack_candidates); - assert!(!OPEN_JEV_9B.pack_candidates); + assert!(OPEN_JEV_9B.pack_candidates); + assert!(OPEN_JEV_27B.share_packed_prefix); + assert!(!OPEN_JEV_9B.share_packed_prefix); } #[test] diff --git a/tests/open_jev/prefill_batch.rs b/tests/open_jev/prefill_batch.rs index c9d5386d..4f13fa9b 100644 --- a/tests/open_jev/prefill_batch.rs +++ b/tests/open_jev/prefill_batch.rs @@ -142,3 +142,287 @@ async fn text_prefix_buffers_grow_and_shrink_without_reusing_prior_contents() { } } } + +#[tokio::test(flavor = "current_thread")] +#[ignore = "needs the pinned Open-Jev-9B export, CUDA library and a GPU reservation"] +async fn nine_b_packing_keeps_prefix_off_and_independent_candidates() { + use omni_qwen3_5_native::prefix::PrefixMode; + + 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_with_prefix(&model, &library, PrefixMode::Off) + .await + .unwrap(); + assert_eq!(engine.checkpoint.alias, "open-jev-9b"); + let candidates: Vec> = (0..3) + .map(|index| { + let mut ids = vec![42; 2048]; + ids.extend(vec![7 + index; 64]); + ids + }) + .collect(); + let (packed, used) = engine + .executor + .execute_with_prefix(&engine.scheduler, vec![candidates.clone()]) + .await + .unwrap(); + assert!(!used.shared, "9B batching must not turn on prefix reuse"); + assert_eq!(used.saved_tokens, 0); + for (index, candidate) in candidates.into_iter().enumerate() { + let independent = engine + .executor + .execute(&engine.scheduler, vec![vec![candidate]]) + .await + .unwrap(); + assert_eq!(packed[0][index], independent[0][0]); + } +} + +#[tokio::test(flavor = "current_thread")] +#[ignore = "needs pinned9B export, OPEN_JEV_EXACT_PREFIX=1, CUDA and a GPU reservation"] +async fn packed_prefix_matches_original_batches_across_budget_boundaries() { + use omni_qwen3_5_native::prefix::PrefixMode; + 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 scenarios = [ + (512, vec![64, 65, 70], 42), + (768, vec![64, 65, 70, 72], 100), + (1024, vec![32, 64, 65, 70], 42), + (1984, vec![64, 65, 66], 100), + (512, vec![1; 8], 42), + (2048, vec![64; 3], 100), + (512, vec![64, 65, 70], 42), + ]; + let cases: Vec>> = scenarios + .iter() + .map(|(prefix, tails, token)| { + tails + .iter() + .enumerate() + .map(|(i, tail)| { + let mut ids = vec![*token; *prefix]; + ids.extend(vec![7 + i as u32; *tail]); + ids + }) + .collect() + }) + .collect(); + let baseline = Engine::load_with_prefix(&model, &library, PrefixMode::Off) + .await + .unwrap(); + let mut expected = Vec::new(); + for inputs in &cases { + expected.push( + baseline + .executor + .execute(&baseline.scheduler, vec![inputs.clone()]) + .await + .unwrap(), + ); + } + drop(baseline); + let candidate = Engine::load(&model, &library).await.unwrap(); + let mut reused = 0; + for ((inputs, expected), (prefix, _, _)) in cases.into_iter().zip(expected).zip(scenarios) { + let count = inputs.len(); + let (actual, used) = candidate + .executor + .execute_with_prefix(&candidate.scheduler, vec![inputs]) + .await + .unwrap(); + assert_eq!(actual, expected, "original packed projection arithmetic"); + if used.shared { + assert_eq!(used.saved_tokens, (count - 1) * prefix); + reused += 1; + } + } + assert!( + reused >= 3, + "exercise actual midlength prefix reuse, not only fallback" + ); +} + +#[test] +#[ignore = "needs pinned9B export, CUDA library and a GPU reservation"] +fn grouped_forward_preserves_original_projection_sizes_and_graph_keys() { + use omni_qwen3_5_native::model::Model; + let dir = PathBuf::from(std::env::var_os("OPEN_JEV_MODEL").unwrap()); + let lib = PathBuf::from(std::env::var_os("OPEN_JEV_CUDA_LIB").unwrap()); + let graph = std::env::var("CUA_S1_GRAPH").as_deref() == Ok("1"); + let mut model = Model::load_with_graph(&dir, &lib, graph).unwrap(); + model.enable_batched_gdn().unwrap(); + model.enable_packed_projections().unwrap(); + model.enable_equivalent_projection_tiles().unwrap(); + let lengths = [79, 81, 83, 255, 257, 513]; + for token in [42, 100, 42] { + let ids: Vec> = lengths + .iter() + .enumerate() + .map(|(i, &n)| vec![token + i as u32; n]) + .collect(); + let inputs: Vec<&[u32]> = ids.iter().map(Vec::as_slice).collect(); + for groups in [vec![3, 3], vec![2, 4], vec![1, 2, 3], vec![6], vec![3, 3]] { + let mut expected = Vec::new(); + let mut start = 0; + for &count in &groups { + expected.extend(model.forward_batch(&inputs[start..start + count]).unwrap()); + start += count; + } + let actual = model + .forward_batch_preserving_groups(&inputs, &groups) + .unwrap(); + assert_eq!( + actual, expected, + "original groups={groups:?}, token={token}" + ); + } + } + assert!( + model + .forward_batch_preserving_groups(&[&[42u32]], &[0, 1]) + .is_err() + ); + assert!( + model + .forward_batch_preserving_groups(&[&[42u32]], &[2]) + .is_err() + ); + if graph { + assert!(model.graph_stats().replays > 0); + } +} + +#[tokio::test(flavor = "current_thread")] +#[ignore = "needs pinned9B export, exact prefix enabled, CUDA and a GPU reservation"] +async fn concurrent_packed_and_prefix_requests_preserve_routing_and_state() { + let dir = PathBuf::from(std::env::var_os("OPEN_JEV_MODEL").unwrap()); + let lib = PathBuf::from(std::env::var_os("OPEN_JEV_CUDA_LIB").unwrap()); + let engine = Engine::load(&dir, &lib).await.unwrap(); + let shared = |prefix: usize, tails: &[usize], token: u32| -> Vec> { + tails + .iter() + .enumerate() + .map(|(i, &tail)| { + let mut ids = vec![token; prefix]; + ids.extend(vec![7 + i as u32; tail]); + ids + }) + .collect() + }; + let requests = [ + vec![shared(37, &[90, 91, 89, 91], 42)], + vec![vec![vec![45; 70]], vec![vec![46; 129], vec![47; 130]]], + vec![shared(768, &[64, 65, 66, 67], 100)], + vec![shared(1024, &[32, 32, 32, 32, 32], 42)], + vec![vec![vec![48; 4096]]], + ]; + let mut expected = Vec::new(); + for request in &requests { + let (scores, used) = engine + .executor + .execute_with_prefix(&engine.scheduler, request.clone()) + .await + .unwrap(); + expected.push((scores, used.shared, used.saved_tokens)); + } + assert!( + expected + .iter() + .any(|(_, shared, saved)| *shared && *saved > 0) + ); + for order in [[0, 1, 2, 3, 4], [4, 3, 2, 1, 0], [0, 1, 2, 3, 4]] { + let (a, b, c, d, e) = tokio::join!( + engine + .executor + .execute_with_prefix(&engine.scheduler, requests[order[0]].clone()), + engine + .executor + .execute_with_prefix(&engine.scheduler, requests[order[1]].clone()), + engine + .executor + .execute_with_prefix(&engine.scheduler, requests[order[2]].clone()), + engine + .executor + .execute_with_prefix(&engine.scheduler, requests[order[3]].clone()), + engine + .executor + .execute_with_prefix(&engine.scheduler, requests[order[4]].clone()), + ); + for (actual, index) in [a, b, c, d, e].into_iter().zip(order) { + let (scores, used) = actual.unwrap(); + assert_eq!((scores, used.shared, used.saved_tokens), expected[index]); + } + } +} + +#[test] +#[ignore = "needs pinned9B export, retained projection CUDA symbols and a GPU reservation"] +fn pointwise_prefix_compaction_preserves_full_mixers_and_changed_inputs() { + use omni_qwen3_5_native::model::Model; + let dir = PathBuf::from(std::env::var_os("OPEN_JEV_MODEL").unwrap()); + let lib = PathBuf::from(std::env::var_os("OPEN_JEV_CUDA_LIB").unwrap()); + let mut original = Model::load(&dir, &lib).unwrap(); + let mut candidate = Model::load(&dir, &lib).unwrap(); + for model in [&mut original, &mut candidate] { + model.enable_packed_projections().unwrap(); + model.enable_batched_gdn().unwrap(); + assert!(model.enable_equivalent_projection_tiles().unwrap()); + } + let mut used = 0; + for (lengths, prefix, seed) in [ + (vec![128, 126, 127, 128], 37, 42), + (vec![128; 8], 63, 43), + (vec![508, 510, 512], 433, 44), + (vec![512; 8], 480, 45), + (vec![128, 126, 127, 128], 37, 46), + (vec![128, 126, 127, 128], 37, 42), + ] { + let ids: Vec> = lengths + .iter() + .enumerate() + .map(|(i, &length)| { + (0..length) + .map(|n| { + if n < prefix { + seed + n as u32 % 11 + } else { + 1000 + (i * 7919 + n * 37 + seed as usize) as u32 % 30000 + } + }) + .collect() + }) + .collect(); + let inputs: Vec<&[u32]> = ids.iter().map(Vec::as_slice).collect(); + let expected = original.forward_batch(&inputs).unwrap(); + let actual = candidate + .try_forward_prefix_compaction(&inputs, prefix) + .unwrap(); + if let Some(actual) = actual { + assert_eq!( + actual, expected, + "prefix={prefix} lengths={lengths:?} seed={seed}" + ); + used += 1; + } + // A normal forward after compaction must see fresh embeddings and all rows. + assert_eq!(candidate.forward_batch(&inputs).unwrap(), expected); + } + assert!( + used >= 5, + "exercise compact projections instead of only fallback" + ); + // Different original split-K reductions must decline before modifying state. + let a = vec![42; 80]; + let mut b = vec![42; 83]; + b[20] = 43; + assert!( + candidate + .try_forward_prefix_compaction(&[&a, &b], 20) + .unwrap() + .is_none() + ); + assert_eq!( + candidate.forward_batch(&[&a, &b]).unwrap(), + original.forward_batch(&[&a, &b]).unwrap() + ); +} diff --git a/tests/open_jev/prefix.rs b/tests/open_jev/prefix.rs index 5d66e2d7..ee44b405 100644 --- a/tests/open_jev/prefix.rs +++ b/tests/open_jev/prefix.rs @@ -3,7 +3,7 @@ use std::path::PathBuf; use omni_open_jev_native::{ - engine::{Engine, check_graph, prefix_values}, + engine::{Engine, check_graph, exact_prefix_value, prefix_values}, executor::PrefixUse, }; use omni_qwen3_5_native::{model::Model, prefix::PrefixMode}; @@ -134,3 +134,21 @@ async fn shared_prefixes_match_fixed_rows_and_auto_keeps_small_requests_on_the_d "auto keeps small requests on the default path" ); } + +#[test] +fn exact_prefix_is_explicit_and_does_not_change_fixed_algorithms() { + assert_eq!(exact_prefix_value(None, PrefixMode::Off).unwrap(), None); + assert_eq!( + exact_prefix_value(Some("0"), PrefixMode::Off).unwrap(), + Some(false) + ); + assert_eq!( + exact_prefix_value(Some("1"), PrefixMode::Off).unwrap(), + Some(true) + ); + assert!(exact_prefix_value(Some("auto"), PrefixMode::Off).is_err()); + for mode in [PrefixMode::Fixed, PrefixMode::Shared, PrefixMode::Auto] { + assert!(exact_prefix_value(Some("1"), mode).is_err()); + assert_eq!(exact_prefix_value(Some("0"), mode).unwrap(), Some(false)); + } +} diff --git a/tests/qwen3_5/actual_row_gemm.rs b/tests/qwen3_5/actual_row_gemm.rs new file mode 100644 index 00000000..91485e19 --- /dev/null +++ b/tests/qwen3_5/actual_row_gemm.rs @@ -0,0 +1,106 @@ +//! Opted-in target-row heuristics preserve retained projection arithmetic. +use super::*; + +#[test] +#[ignore = "needs SM90, cuBLASLt 13.1.1 and a reserved GPU"] +fn actual_row_heuristics_match_retained_algorithms_and_graph_replay() { + let st = setup(); + let reuse = api().reuse_gemm().unwrap(); + // SAFETY: independent handles and workspaces remain live through every readback. + let original = unsafe { (api().cs1_gemm_create)(32 << 20) }; + let optimized = unsafe { (api().cs1_gemm_create)(32 << 20) }; + assert!(!original.is_null() && !optimized.is_null()); + assert_eq!(unsafe { api().equivalent_tiles().unwrap()(optimized) }, 1); + let mut checked = 0; + for (n, k) in [ + (4096usize, 4096usize), + (4096, 12288), + (12352, 4096), + (10240, 4096), + (24576, 4096), + ] { + let weights = to_device(&random(n * k, 521, 0.25), st); + for (full, rows) in [ + (31usize, 64usize), + (128, 65), + (128, 193), + (256, 127), + (512, 129), + (1024, 255), + (2048, 513), + (4096, 2049), + (2048, 4095), + ] { + let capability = |handle| unsafe { + // SAFETY: shape checks consume only these live dimension arrays. + (reuse.check)(handle, &full, 1, &rows, 1, n as i32, k as i32) + }; + let old = capability(original); + assert_eq!(old, capability(optimized), "eligibility changed"); + if old == -1 { + continue; + } + check(old, "retained shape").unwrap(); + let input = DeviceBuffer::new(rows * k * 2).unwrap(); + let sentinel = vec![bf16::from_f32(-7.0); rows * n + 16]; + let expected = to_device(&sentinel, st); + let actual = to_device(&sentinel, st); + let run = |handle, out: &DeviceBuffer| unsafe { + // SAFETY: complete target rows and weights fit; guards surround output. + check( + (reuse.run)( + handle, + input.at(0), + weights.at(0), + out.at(16), + rows as i32, + n as i32, + k as i32, + full as i32, + st, + ), + "actual-row retained projection", + ) + }; + let mut graph = None; + for (seed, scale) in [(522, 0.5), (523, 32.0)] { + let bytes: Vec<_> = random(rows * k, seed, scale) + .iter() + .flat_map(|x| x.to_le_bytes()) + .collect(); + // SAFETY: replace all live inputs before eager execution or replay. + unsafe { cuda::upload(input.at(0), &bytes, st).unwrap() }; + run(original, &expected).unwrap(); + if let Some(g) = &graph { + cuda::Graph::launch(g, st).unwrap(); + } else { + run(optimized, &actual).unwrap(); + cuda::synchronize(st).unwrap(); + graph = Some(cuda::Graph::capture(st, || run(optimized, &actual)).unwrap()); + } + let mut a = vec![0u8; sentinel.len() * 2]; + let mut b = vec![0u8; sentinel.len() * 2]; + unsafe { + // SAFETY: read complete output buffers, including both guard regions. + cuda::download(&mut a, expected.at(0), st).unwrap(); + cuda::download(&mut b, actual.at(0), st).unwrap(); + } + assert_eq!( + a, b, + "N={n} K={k} original={full} actual={rows} seed={seed}" + ); + let guard: Vec<_> = sentinel[..8].iter().flat_map(|x| x.to_le_bytes()).collect(); + assert_eq!(&b[..16], &guard); + assert_eq!(&b[b.len() - 16..], &guard); + } + checked += 1; + } + } + assert!(checked >= 35, "exercise actual shapes, not only fallbacks"); + unsafe { + // SAFETY: graph owners dropped in their iterations; readbacks retired all work. + (api().cs1_gemm_destroy)(optimized); + (api().cs1_gemm_destroy)(original); + (api().cs1_stream_destroy)(st); + } +} diff --git a/tests/qwen3_5/batched_attention.rs b/tests/qwen3_5/batched_attention.rs new file mode 100644 index 00000000..21cacbab --- /dev/null +++ b/tests/qwen3_5/batched_attention.rs @@ -0,0 +1,266 @@ +//! Independent suffix queries retain their exact per-sequence arithmetic. +use super::*; + +#[test] +#[ignore = "needs CUDA, a GPU reservation, and the batched shared-attention symbol"] +fn batched_shared_attention_matches_serial_bits_and_live_graph_inputs() { + let st = setup(); + let serial = api().shared_suffix().unwrap().attention; + let batch = api().shared_batched_attention().unwrap(); + for (hq, hk) in [(8usize, 2usize), (16, 4), (24, 4)] { + let d = 256; + let ld = hk * d + 8; + for (prefix, lengths) in [ + (64usize, vec![1usize]), + (64, vec![1, 65]), + (512, vec![31, 63, 64, 65, 127, 128, 129, 3]), + (2048, vec![17, 129, 33]), + (4096, (0..16).map(|i| 1 + i * 9).collect()), + ] { + let rows: usize = lengths.iter().sum(); + let q = to_device(&random(rows * hq * d, 451, 0.5), st); + let k = to_device(&random(rows * hk * d, 452, 0.5), st); + let v = to_device(&random(rows * ld, 453, 0.5), st); + let gate = to_device(&random(rows * hq * d, 454, 0.5), st); + let pk = to_device(&random(prefix * hk * d, 455, 0.5), st); + let pv = to_device(&random(prefix * hk * d, 456, 0.5), st); + let n = rows * hq * d; + let sentinel = vec![bf16::from_f32(-7.0); n + 16]; + let expected = to_device(&sentinel, st); + let actual = to_device(&sentinel, st); + let run = || unsafe { + // SAFETY: offsets cover disjoint suffix rows; both KV allocations + // cover prefix rows. The host descriptor is consumed by value. + check( + batch( + q.at(0), + k.at(0), + v.at(0), + ld as i32, + gate.at(0), + actual.at(16), + pk.at(0), + pv.at(0), + lengths.as_ptr(), + lengths.len(), + prefix as i32, + hq as i32, + hk as i32, + d as i32, + (d as f32).powf(-0.5), + st, + ), + "batched shared attention", + ) + }; + run().unwrap(); + cuda::synchronize(st).unwrap(); + let graph = cuda::Graph::capture(st, run).unwrap(); + for seed in [457, 458, 459] { + let changed = random(prefix * hk * d, seed, 0.5); + let bytes: Vec<_> = changed.iter().flat_map(|x| x.to_le_bytes()).collect(); + unsafe { + // SAFETY: overwrite live prefix values before replay; keep + // all inputs alive until downloads retire the stream. + cuda::upload(pk.at(0), &bytes, st).unwrap(); + let mut offset = 0; + for &length in &lengths { + check( + serial( + q.at(offset * hq * d * 2), + k.at(offset * hk * d * 2), + v.at(offset * ld * 2), + ld as i32, + gate.at(offset * hq * d * 2), + expected.at(16 + offset * hq * d * 2), + pk.at(0), + pv.at(0), + length as i32, + prefix as i32, + hq as i32, + hk as i32, + d as i32, + (d as f32).powf(-0.5), + st, + ), + "serial shared attention", + ) + .unwrap(); + offset += length; + } + } + if seed == 457 { + run().unwrap(); + } else { + graph.launch(st).unwrap(); + } + let mut a = vec![0; (n + 16) * 2]; + let mut b = a.clone(); + unsafe { + // SAFETY: downloads cover allocated outputs and both guards. + cuda::download(&mut a, expected.at(0), st).unwrap(); + cuda::download(&mut b, actual.at(0), st).unwrap(); + } + assert_eq!( + a, b, + "hq={hq} prefix={prefix} lengths={lengths:?} seed={seed}" + ); + let guard: Vec<_> = sentinel[..8].iter().flat_map(|x| x.to_le_bytes()).collect(); + assert_eq!(&b[..16], guard); + assert_eq!(&b[b.len() - 16..], guard); + } + for invalid in [vec![], vec![0], vec![usize::MAX], vec![1; 17]] { + let rc = unsafe { + // SAFETY: invalid lengths must be rejected before launching. + batch( + q.at(0), + k.at(0), + v.at(0), + ld as i32, + gate.at(0), + actual.at(16), + pk.at(0), + pv.at(0), + invalid.as_ptr(), + invalid.len(), + prefix as i32, + hq as i32, + hk as i32, + d as i32, + (d as f32).powf(-0.5), + st, + ) + }; + assert_ne!(rc, 0); + } + } + } + unsafe { + (api().cs1_stream_destroy)(st); + } +} + +#[test] +#[ignore = "needs CUDA, a GPU reservation, and the batched causal-attention symbol"] +fn batched_causal_attention_matches_serial_bits_and_live_graph_inputs() { + let st = setup(); + let batch = api().batched_attention().unwrap(); + for (hq, hk) in [(8usize, 2usize), (16, 4), (24, 4)] { + let d = 256; + let ld = hk * d + 8; + for lengths in [ + vec![1usize], + vec![1, 65], + vec![31, 63, 64, 65, 127, 128, 129, 3], + vec![512, 1024, 2048], + vec![64; 16], + vec![4096, 65], + ] { + let rows: usize = lengths.iter().sum(); + let q = to_device(&random(rows * hq * d, 481, 0.5), st); + let k = to_device(&random(rows * hk * d, 482, 0.5), st); + let v = to_device(&random(rows * ld, 483, 0.5), st); + let gate = to_device(&random(rows * hq * d, 484, 0.5), st); + let n = rows * hq * d; + let sentinel = vec![bf16::from_f32(-7.0); n + 16]; + let expected = to_device(&sentinel, st); + let actual = to_device(&sentinel, st); + let run = || unsafe { + // SAFETY: concatenated input/output rows have independent boundaries; + // the host lengths are copied into the launch before return. + check( + batch( + q.at(0), + k.at(0), + v.at(0), + ld as i32, + gate.at(0), + actual.at(16), + lengths.as_ptr(), + lengths.len(), + hq as i32, + hk as i32, + d as i32, + (d as f32).powf(-0.5), + st, + ), + "batched causal attention", + ) + }; + run().unwrap(); + cuda::synchronize(st).unwrap(); + let graph = cuda::Graph::capture(st, run).unwrap(); + for seed in [485, 486, 487] { + let changed = random(rows * hk * d, seed, 0.5); + let bytes: Vec<_> = changed.iter().flat_map(|x| x.to_le_bytes()).collect(); + unsafe { + // SAFETY: replace all sequence keys; serial calls use exactly the + // corresponding independent ranges, including uneven final blocks. + cuda::upload(k.at(0), &bytes, st).unwrap(); + let mut offset = 0; + for &length in &lengths { + check( + (api().cs1_attention_gated)( + q.at(offset * hq * d * 2), + k.at(offset * hk * d * 2), + v.at(offset * ld * 2), + ld as i32, + gate.at(offset * hq * d * 2), + expected.at(16 + offset * hq * d * 2), + length as i32, + hq as i32, + hk as i32, + d as i32, + (d as f32).powf(-0.5), + st, + ), + "serial causal attention", + ) + .unwrap(); + offset += length; + } + } + if seed == 485 { + run().unwrap(); + } else { + graph.launch(st).unwrap(); + } + let mut a = vec![0; (n + 16) * 2]; + let mut b = a.clone(); + unsafe { + // SAFETY: downloads include both allocated output guards. + cuda::download(&mut a, expected.at(0), st).unwrap(); + cuda::download(&mut b, actual.at(0), st).unwrap(); + } + assert_eq!(a, b, "hq={hq} lengths={lengths:?} seed={seed}"); + let guard: Vec<_> = sentinel[..8].iter().flat_map(|x| x.to_le_bytes()).collect(); + assert_eq!(&b[..16], guard); + assert_eq!(&b[b.len() - 16..], guard); + } + for invalid in [vec![], vec![0], vec![usize::MAX], vec![1; 17]] { + let rc = unsafe { + // SAFETY: invalid dimensions must be rejected before any launch. + batch( + q.at(0), + k.at(0), + v.at(0), + ld as i32, + gate.at(0), + actual.at(16), + invalid.as_ptr(), + invalid.len(), + hq as i32, + hk as i32, + d as i32, + (d as f32).powf(-0.5), + st, + ) + }; + assert_ne!(rc, 0); + } + } + } + unsafe { + (api().cs1_stream_destroy)(st); + } +} diff --git a/tests/qwen3_5/batched_gdn.rs b/tests/qwen3_5/batched_gdn.rs index 59080d5a..c6c6c006 100644 --- a/tests/qwen3_5/batched_gdn.rs +++ b/tests/qwen3_5/batched_gdn.rs @@ -4,9 +4,24 @@ 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() { + check_batched(false); +} + +#[test] +#[ignore = "needs a GPU reservation and shared-state batched GDN CUDA symbol"] +fn batched_gdn_shared_state_matches_serial_and_changed_state_replay() { + check_batched(true); +} + +fn check_batched(shared: bool) { let stream = setup(); let batched = api().batched_gdn().unwrap(); - for (heads, key_heads) in [(4usize, 2usize), (48, 16)] { + let dimensions = if shared { + [(4usize, 2usize), (32, 16)] + } else { + [(4, 2), (48, 16)] + }; + for (heads, key_heads) in dimensions { for lengths in [ vec![1usize, 63, 64, 65, 129], vec![1; 16], @@ -46,10 +61,29 @@ fn batched_gdn_matches_serial_with_partial_chunks_and_graph_replay() { ) }; let serial_scratch = DeviceBuffer::new(serial_size * 4).unwrap(); + let state = f32_to_device(&vec![0.0; heads * 128 * 128], stream); let scale = (128f32).powf(-0.5); let batch = || unsafe { // SAFETY: packed buffers cover every length; scratch uses the API's bound. - check( + let rc = if shared { + (api().shared_batched_gdn().unwrap())( + 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, + state.at(0), + stream, + ) + } else { (batched.run)( q.at(0), k.at(0), @@ -65,9 +99,9 @@ fn batched_gdn_matches_serial_with_partial_chunks_and_graph_replay() { key_heads as i32, scale, stream, - ), - "batch GDN", - ) + ) + }; + check(rc, "batch GDN") }; batch().unwrap(); cuda::synchronize(stream).unwrap(); @@ -75,29 +109,58 @@ fn batched_gdn_matches_serial_with_partial_chunks_and_graph_replay() { 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(); + let state_bytes: Vec = + random(heads * 128 * 128, if replay { 89 } else { 88 }, 0.01) + .iter() + .flat_map(|v| v.to_f32().to_le_bytes()) + .collect(); unsafe { // SAFETY: replace all value rows, then compare independent sequence runs. cuda::upload(v.at(0), &bytes, stream).unwrap(); + cuda::upload(state.at(0), &state_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(); + if shared { + check( + (api().cs1_gdn_prefill_x)( + 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, + state.at(0), + std::ptr::null_mut(), + stream, + ), + "serial GDN", + ) + .unwrap(); + } else { + 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; } } @@ -106,6 +169,12 @@ fn batched_gdn_matches_serial_with_partial_chunks_and_graph_replay() { } else { batch().unwrap(); } + let mut after = vec![0u8; state_bytes.len()]; + // SAFETY: the full allocation is downloaded after all consumers. + unsafe { + cuda::download(&mut after, state.at(0), stream).unwrap(); + } + assert_eq!(after, state_bytes, "shared input state was modified"); let a = from_device(&expected, values, stream); let b = from_device(&actual, values, stream); let changes = a @@ -115,7 +184,7 @@ fn batched_gdn_matches_serial_with_partial_chunks_and_graph_replay() { .count(); assert_eq!( changes, 0, - "heads={heads} lengths={lengths:?} graph={replay}" + "heads={heads} lengths={lengths:?} graph={replay} shared={shared}" ); } unsafe { diff --git a/tests/qwen3_5/capture_only.rs b/tests/qwen3_5/capture_only.rs new file mode 100644 index 00000000..07ad2f36 --- /dev/null +++ b/tests/qwen3_5/capture_only.rs @@ -0,0 +1,76 @@ +//! A state-only capture must retain every byte consumed by later suffixes. +use super::*; +use std::path::PathBuf; + +#[test] +#[ignore = "needs pinned9B export, CUDA retained GEMMs and a GPU reservation"] +fn state_only_capture_matches_complete_prefix_state_bytes() { + let dir = PathBuf::from(std::env::var_os("OPEN_JEV_MODEL").unwrap()); + let lib = PathBuf::from(std::env::var_os("OPEN_JEV_CUDA_LIB").unwrap()); + let mut model = Model::load(&dir, &lib).unwrap(); + model.enable_packed_projections().unwrap(); + model.enable_batched_gdn().unwrap(); + assert_eq!(model.cfg.hidden, 4096); + assert!(matches!(model.layers.last().unwrap().mixer, Mixer::Full(_))); + let reuse = cuda::api().reuse_gemm().unwrap(); + for (prefix, seed) in [(512, 42), (2048, 43), (4096, 44), (512, 45)] { + let ids: Vec = (0..prefix) + .map(|n| 256 + (n * 37 + seed) as u32 % 30000) + .collect(); + let mut complete = model.alloc_prefix(prefix).unwrap(); + let mut state_only = model.alloc_prefix(prefix).unwrap(); + model.prepare_scratch(prefix).unwrap(); + let s = model.scratch.as_ref().unwrap(); + for (state, readout) in [ + (&mut complete, WindowReadout::Hidden), + (&mut state_only, WindowReadout::StateOnly), + ] { + model.embed_tokens(s, &ids).unwrap(); + model + .run_window( + s, + 0, + prefix, + false, + Some(state), + None, + model.gemm, + Some((reuse, prefix)), + readout, + ) + .unwrap(); + model.synchronize().unwrap(); + } + let compare = |a: &DeviceBuffer, b: &DeviceBuffer, bytes: usize, label: &str| { + let mut want = vec![0u8; bytes]; + let mut got = vec![0u8; bytes]; + // SAFETY: sizes match the prefix-state allocations for the pinned model. + unsafe { + cuda::download(&mut want, a.at(0), model.stream).unwrap(); + cuda::download(&mut got, b.at(0), model.stream).unwrap(); + } + assert!(want == got, "prefix={prefix} seed={seed} buffer={label}"); + }; + for ((ka, va), (kb, vb)) in complete.attn_kv.iter().zip(&state_only.attn_kv) { + let bytes = prefix * model.cfg.kv_heads * model.cfg.head_dim * BF16; + compare(ka, kb, bytes, "K"); + compare(va, vb, bytes, "V"); + } + for (a, b) in complete.gdn_state.iter().zip(&state_only.gdn_state) { + compare( + a, + b, + model.cfg.lin_v_heads * model.cfg.lin_k_dim * model.cfg.lin_v_dim * F32, + "GDN state", + ); + } + for (a, b) in complete.conv_tail.iter().zip(&state_only.conv_tail) { + compare( + a, + b, + 3 * Widths::of(&model.cfg).gdn_in * BF16, + "convolution tail", + ); + } + } +} diff --git a/tests/qwen3_5/equivalent_tiles.rs b/tests/qwen3_5/equivalent_tiles.rs new file mode 100644 index 00000000..3fc2d85d --- /dev/null +++ b/tests/qwen3_5/equivalent_tiles.rs @@ -0,0 +1,347 @@ +//! Opt-in output-tile compatibility, exact arithmetic and unchanged fallback. +use super::*; + +#[test] +#[ignore = "needs SM90, cuBLASLt 13.1.1, the optional tile capability and a GPU reservation"] +fn equivalent_tiles_preserve_full_rows_packing_and_graph_replay() { + let st = setup(); + let enable = api().equivalent_tiles().expect("new optional tile symbol"); + let reuse = api().reuse_gemm().unwrap(); + let packed = api().packed_gemm().unwrap(); + // SAFETY: handles own independent plans/workspaces and are retired below. + let original = unsafe { (api().cs1_gemm_create)(32 << 20) }; + let optimized = unsafe { (api().cs1_gemm_create)(32 << 20) }; + assert!(!original.is_null() && !optimized.is_null()); + // SAFETY: no plans or graphs have been created on the ordinary handle. + assert_eq!( + unsafe { enable(optimized) }, + 1, + "exercise supported profile" + ); + let compatible = |handle, rows: &[usize], n, k| unsafe { + // SAFETY: capability check only reads the live row arrays. + (reuse.check)( + handle, + rows.as_ptr(), + rows.len(), + [512usize].as_ptr(), + 1, + n, + k, + ) + }; + assert_eq!(compatible(original, &[2048, 3782], 4096, 4096), -1); + assert_eq!(compatible(optimized, &[2048, 3782], 4096, 4096), 0); + for (n, k) in [ + (4096, 4096), + (4096, 12288), + (12352, 4096), + (10240, 4096), + (24576, 4096), + ] { + assert_eq!(compatible(optimized, &[3776, 3779, 3780, 3782], n, k), 0); + } + // Different split-K reductions must not become compatible. + assert_eq!(compatible(optimized, &[80, 90], 4096, 12288), -1); + // A capability change after plans exist would invalidate cached grouping. + assert!(unsafe { enable(optimized) } < 0); + + let rows = 4096usize; + let mut checked = 0; + for (n, k) in [ + (4096usize, 4096usize), + (4096, 12288), + (12352, 4096), + (10240, 4096), + (24576, 4096), + ] { + let weights = to_device(&random(n * k, 71, 0.25), st); + let input = to_device(&random(rows * k, 72, 0.5), st); + let expected = DeviceBuffer::new(rows * n * 2).unwrap(); + let actual = DeviceBuffer::new(rows * n * 2).unwrap(); + let prefix = || unsafe { + // SAFETY: input/output own at least 512 rows; weights are [n,k]. + check( + (reuse.run)( + optimized, + input.at(0), + weights.at(0), + actual.at(0), + 512, + n as i32, + k as i32, + 2048, + st, + ), + "equivalent prefix", + ) + }; + prefix().unwrap(); + cuda::synchronize(st).unwrap(); + let graph = cuda::Graph::capture(st, prefix).unwrap(); + for seed in [73, 74, 75] { + let bytes: Vec = random(rows * k, seed, if seed == 75 { 32.0 } else { 0.5 }) + .iter() + .flat_map(|x| x.to_le_bytes()) + .collect(); + // SAFETY: input owns exactly rows*k BF16 elements. + unsafe { cuda::upload(input.at(0), &bytes, st).unwrap() }; + for m in [ + 512usize, 513, 576, 577, 582, 768, 832, 833, 838, 840, 1024, 1056, 1088, 1089, + 1094, 1673, 1735, 2048, 2112, 2113, 2240, 2241, 2299, 2300, 2367, 2496, 3233, 3343, + 3419, 3483, 3591, 3776, 3779, 3780, 3782, 4096, + ] { + let supported = compatible(optimized, &[2048, m], n as i32, k as i32); + if supported == -1 { + continue; + } + check(supported, "tile compatibility").unwrap(); + // SAFETY: complete original rows fit both allocations. + unsafe { + check( + (api().cs1_gemm)( + original, + input.at(0), + weights.at(0), + expected.at(0), + m as i32, + n as i32, + k as i32, + n as i32, + st, + ), + "complete original projection", + ) + .unwrap() + }; + if seed == 74 { + graph.launch(st).unwrap(); + } else { + prefix().unwrap(); + } + assert_eq!( + from_device(&expected, 512 * n, st), + from_device(&actual, 512 * n, st), + "M={m} N={n} K={k} seed={seed}" + ); + checked += 1; + } + let lengths = [512usize, 1024, 2048]; + let mut offset = 0; + for m in lengths { + // SAFETY: each original submatrix is within the packed row buffers. + unsafe { + check( + (api().cs1_gemm)( + original, + input.at(offset * k * 2), + weights.at(0), + expected.at(offset * n * 2), + m as i32, + n as i32, + k as i32, + n as i32, + st, + ), + "independent packed reference", + ) + .unwrap() + }; + offset += m; + } + // SAFETY: all packed rows and dimensions fit the same allocations. + unsafe { + check( + packed( + optimized, + input.at(0), + weights.at(0), + actual.at(0), + lengths.as_ptr(), + lengths.len(), + n as i32, + k as i32, + st, + ), + "equivalent packed projection", + ) + .unwrap() + }; + assert_eq!( + from_device(&expected, offset * n, st), + from_device(&actual, offset * n, st) + ); + } + } + assert!( + checked >= 120, + "exercise actual equivalence across projection shapes" + ); + cuda::synchronize(st).unwrap(); + // SAFETY: downloads and synchronization retired all work and captured graphs. + unsafe { + (api().cs1_gemm_destroy)(original); + (api().cs1_gemm_destroy)(optimized); + } +} + +#[test] +#[ignore = "needs SM90, cuBLASLt13.1.1, expanded profiles and a GPU reservation"] +fn short_reuse_profiles_preserve_full_rows_and_ordinary_packing() { + let st = setup(); + let enable = api().equivalent_tiles().unwrap(); + let reuse = api().reuse_gemm().unwrap(); + let packed = api().packed_gemm().unwrap(); + // SAFETY: independent handles own all plans and are retired after synchronization. + let original = unsafe { (api().cs1_gemm_create)(32 << 20) }; + let optimized = unsafe { (api().cs1_gemm_create)(32 << 20) }; + assert!(!original.is_null() && !optimized.is_null()); + assert_eq!(unsafe { enable(optimized) }, 1); + let mut checked = 0; + for (n, k) in [ + (4096usize, 4096usize), + (4096, 12288), + (12352, 4096), + (10240, 4096), + (24576, 4096), + ] { + let weights = to_device(&random(n * k, 88, 0.25), st); + let input = DeviceBuffer::new(4096 * k * 2).unwrap(); + let expected = DeviceBuffer::new(4096 * n * 2).unwrap(); + let actual = DeviceBuffer::new(4096 * n * 2).unwrap(); + let lengths = [ + 63usize, 64, 79, 80, 83, 127, 128, 160, 192, 254, 256, 332, 380, 506, + ]; + let run_packed = || unsafe { + // SAFETY: summed ragged lengths fit the 4096-row input/output allocations. + check( + packed( + optimized, + input.at(0), + weights.at(0), + actual.at(0), + lengths.as_ptr(), + lengths.len(), + n as i32, + k as i32, + st, + ), + "short packed projection", + ) + }; + let mut graph: Option = None; + for (seed, scale) in [(89, 0.5), (90, 32.0)] { + let bytes: Vec = random(4096 * k, seed, scale) + .iter() + .flat_map(|x| x.to_le_bytes()) + .collect(); + // SAFETY: input owns exactly these 4096*k BF16 elements. + unsafe { + cuda::upload(input.at(0), &bytes, st).unwrap(); + } + for m in [ + 64usize, 79, 80, 83, 127, 128, 160, 192, 254, 256, 332, 380, 506, 512, + ] { + // SAFETY: shape-only check reads live original and target arrays. + let supported = unsafe { + (reuse.check)( + optimized, + [2048usize, m].as_ptr(), + 2, + [64usize].as_ptr(), + 1, + n as i32, + k as i32, + ) + }; + if supported == -1 { + continue; + } // Retain split-K/unsupported profiles. + check(supported, "short profile compatibility").unwrap(); + unsafe { + // SAFETY: original complete rows and selected 64 rows fit allocations. + check( + (api().cs1_gemm)( + original, + input.at(0), + weights.at(0), + expected.at(0), + m as i32, + n as i32, + k as i32, + n as i32, + st, + ), + "short complete original", + ) + .unwrap(); + check( + (reuse.run)( + optimized, + input.at(0), + weights.at(0), + actual.at(0), + 64, + n as i32, + k as i32, + 2048, + st, + ), + "short equivalent rows", + ) + .unwrap(); + } + assert_eq!( + from_device(&expected, 64 * n, st), + from_device(&actual, 64 * n, st), + "M={m} N={n} K={k} seed={seed}" + ); + checked += 1; + } + let mut offset = 0; + for m in lengths { + unsafe { + // SAFETY: each independent reference occupies its original slice. + check( + (api().cs1_gemm)( + original, + input.at(offset * k * 2), + weights.at(0), + expected.at(offset * n * 2), + m as i32, + n as i32, + k as i32, + n as i32, + st, + ), + "ragged original", + ) + .unwrap(); + } + offset += m; + } + if let Some(graph) = &graph { + graph.launch(st).unwrap(); + } else { + run_packed().unwrap(); + cuda::synchronize(st).unwrap(); + graph = Some(cuda::Graph::capture(st, run_packed).unwrap()); + } + assert_eq!( + from_device(&expected, offset * n, st), + from_device(&actual, offset * n, st), + "ragged N={n} K={k} seed={seed}" + ); + } + } + assert!( + checked >= 90, + "exercise short profiles rather than only falling back" + ); + cuda::synchronize(st).unwrap(); + // SAFETY: all graphs and pending device work have retired. + unsafe { + (api().cs1_gemm_destroy)(original); + (api().cs1_gemm_destroy)(optimized); + } +} diff --git a/tests/qwen3_5/kernels.rs b/tests/qwen3_5/kernels.rs index 5173a122..bf4f270a 100644 --- a/tests/qwen3_5/kernels.rs +++ b/tests/qwen3_5/kernels.rs @@ -68,9 +68,11 @@ fn residual_rms_norm_matches_rounded_reference() { let st = setup(); let eps = 1e-6f32; // Cua-S1 and Open-Jev widths, scalar fallbacks for other widths, and a BF16 - // element offset to check that cached loads need no packed alignment. + // element offset to exercise the scalar fallback for packed loads. for (d, offset) in [ (2560usize, 0usize), + (4096, 0), + (4096, 1), (5120, 0), (8192, 0), (257, 0), @@ -143,6 +145,89 @@ fn residual_rms_norm_matches_rounded_reference() { } } +#[test] +#[ignore = "needs a GPU and CUA_S1_CUDA_LIB"] +fn residual_rms_norm_vector_matches_each_unaligned_fallback() { + let st = setup(); + let d = 4096usize; + // Each pointer independently forces the scalar fallback on identical values. + // Both kernels must preserve BF16 rounding and the FP32 reduction order. + for rows in [1usize, 3, 107, 256] { + for amp in [0.5, 32.0] { + let n = rows * d; + let residual = random(n, 61, amp); + let delta = random(n, 62, amp / 8.0); + let weights = random(d, 63, 0.25); + let mut reference = None; + for offsets in [ + [8usize; 4], + [1, 8, 8, 8], + [8, 1, 8, 8], + [8, 8, 1, 8], + [8, 8, 8, 1], + ] { + let guard = bf16::from_f32(-7.0); + let mut host = [ + vec![guard; n + offsets[0] + 8], + vec![guard; n + offsets[1] + 8], + vec![guard; d + offsets[2] + 8], + vec![guard; n + offsets[3] + 8], + ]; + host[0][offsets[0]..offsets[0] + n].copy_from_slice(&residual); + host[1][offsets[1]..offsets[1] + n].copy_from_slice(&delta); + host[2][offsets[2]..offsets[2] + d].copy_from_slice(&weights); + let buffers = host.each_ref().map(|v| to_device(v, st)); + // SAFETY: distinct allocations, complete rows and weights after + // each BF16 offset, with eight guard elements after each region. + unsafe { + check( + (api().cs1_add_rms_norm)( + buffers[0].at(offsets[0] * 2), + buffers[1].at(offsets[1] * 2), + buffers[2].at(offsets[2] * 2), + buffers[3].at(offsets[3] * 2), + rows as i32, + d as i32, + 1e-6, + st, + ), + "vector residual norm", + ) + .unwrap(); + } + let mut actual = Vec::new(); + for i in 0..4 { + let got = from_device(&buffers[i], host[i].len(), st); + let len = if i == 2 { d } else { n }; + assert!(got[..offsets[i]].iter().all(|v| *v == guard.to_f32())); + assert!(got[offsets[i] + len..].iter().all(|v| *v == guard.to_f32())); + let bits: Vec<_> = got[offsets[i]..offsets[i] + len] + .iter() + .map(|v| bf16::from_f32(*v).to_bits()) + .collect(); + if i == 1 || i == 2 { + let original: Vec<_> = host[i][offsets[i]..offsets[i] + len] + .iter() + .map(|v| v.to_bits()) + .collect(); + assert_eq!(bits, original, "read-only buffer changed"); + } else { + actual.push(bits); + } + } + if let Some(expected) = &reference { + assert_eq!( + &actual, expected, + "rows={rows}, amp={amp}, offsets={offsets:?}" + ); + } else { + reference = Some(actual); + } + } + } + } +} + #[test] #[ignore = "needs a GPU and CUA_S1_CUDA_LIB"] fn silu_mul_matches_rounded_reference() { @@ -1592,3 +1677,10 @@ mod shared_suffix; mod vector_conv; mod wide_gdn; + +mod equivalent_tiles; + +mod batched_attention; + +#[path = "actual_row_gemm.rs"] +mod actual_row_gemm; diff --git a/tests/qwen3_5/multimodal_graph.rs b/tests/qwen3_5/multimodal_graph.rs index 7b9f7771..cfe47a9b 100644 --- a/tests/qwen3_5/multimodal_graph.rs +++ b/tests/qwen3_5/multimodal_graph.rs @@ -332,7 +332,7 @@ fn graph_recording_failure_keeps_eager_result_and_disables_both_modes() -> Resul .contains("injected record failure"), "must exercise the actual recording-error path" ); - model.cache_graph(&[text.len()], multimodal, captured)?; + model.cache_graph(&[text.len()], multimodal, text.len(), captured)?; assert_hidden_eq( &model.last_hidden(model.scratch.as_ref().unwrap(), text.len())?, &prior_eager, diff --git a/tests/qwen3_5/prefix_tree.rs b/tests/qwen3_5/prefix_tree.rs new file mode 100644 index 00000000..ae7a2bfb --- /dev/null +++ b/tests/qwen3_5/prefix_tree.rs @@ -0,0 +1,94 @@ +//! Whole hidden-vector equality across branching, growth and changed live inputs. +use super::*; +use std::path::PathBuf; + +#[test] +#[ignore = "needs pinned9B export, shared-state CUDA and GPU reservation"] +fn prefix_tree_preserves_original_groups_and_every_hidden_component() { + let dir = PathBuf::from(std::env::var_os("OPEN_JEV_MODEL").unwrap()); + let lib = PathBuf::from(std::env::var_os("OPEN_JEV_CUDA_LIB").unwrap()); + let mut model = Model::load(&dir, &lib).unwrap(); + model.enable_packed_projections().unwrap(); + model.enable_batched_gdn().unwrap(); + model.enable_equivalent_projection_tiles().unwrap(); + let mut executed = 0; + for (prefix, tails, group, seed) in [ + (64, vec![64; 3], true, 40), + (64, vec![64, 65, 63], true, 41), + (448, vec![64; 8], true, 42), + (512, vec![64; 3], true, 43), + (1984, vec![64; 4], false, 44), + (4032, vec![64; 3], false, 45), + (4160, vec![64; 3], false, 46), + (64, vec![1, 3, 65, 127], true, 47), + (512, vec![64; 16], false, 48), + (4032, vec![64; 16], false, 49), + ] { + let common: Vec = (0..prefix) + .map(|i| 256 + ((i * 37 + seed) % 30000) as u32) + .collect(); + let ids: Vec> = tails + .iter() + .enumerate() + .map(|(j, &length)| { + common + .iter() + .copied() + .chain((0..length).map(|i| 256 + ((i * 31 + seed + j * 701) % 30000) as u32)) + .collect() + }) + .collect(); + let inputs: Vec<&[u32]> = ids.iter().map(Vec::as_slice).collect(); + let original: Vec<_> = if group { + vec![ids.iter().map(Vec::len).sum(); ids.len()] + } else { + ids.iter().map(Vec::len).collect() + }; + let expected = if group { + model.forward_batch(&inputs).unwrap() + } else { + inputs + .iter() + .map(|ids| model.forward_batch(&[*ids]).unwrap().remove(0)) + .collect() + }; + let Some(output) = model + .try_forward_prefix_tree(&inputs, prefix, &original) + .unwrap() + else { + // Mixed original sizes can select incompatible algorithms. Declining + // must leave the accepted full-forward result unchanged. + assert!( + tails.windows(2).any(|pair| pair[0] != pair[1]) || tails.len() == 16, + "unexpected equal-size fallback prefix={prefix}" + ); + let repeated = if group { + model.forward_batch(&inputs).unwrap() + } else { + inputs + .iter() + .map(|ids| model.forward_batch(&[*ids]).unwrap().remove(0)) + .collect() + }; + assert_eq!(repeated, expected); + eprintln!("safe fallback prefix={prefix} tails={tails:?}"); + continue; + }; + executed += 1; + eprintln!("tree executed prefix={prefix} tails={tails:?}"); + assert_eq!(output.saved_tokens, (inputs.len() - 1) * prefix); + for (branch, (want, got)) in expected.iter().zip(output.hidden).enumerate() { + let changed = want.iter().zip(&got).filter(|(a, b)| a != b).count(); + let max = want + .iter() + .zip(&got) + .map(|(a, b)| (a - b).abs()) + .fold(0.0f32, f32::max); + assert_eq!( + changed, 0, + "prefix={prefix}, group={group}, branch={branch}, max={max}" + ); + } + } + assert!(executed >= 6, "insufficient executed tree layouts"); +} diff --git a/tests/qwen3_5/tree_graph.rs b/tests/qwen3_5/tree_graph.rs new file mode 100644 index 00000000..7dd6cfcb --- /dev/null +++ b/tests/qwen3_5/tree_graph.rs @@ -0,0 +1,165 @@ +//! Ordinary graph capture/replay remains exact around eager prefix-tree execution. +use super::*; +use std::path::PathBuf; + +fn inputs(length: usize, prefix: usize, count: usize, seed: usize) -> Vec> { + (0..count) + .map(|branch| { + (0..length) + .map(|i| { + 256 + ((i * 37 + seed * 17 + if i < prefix { 0 } else { branch * 701 }) % 30000) + as u32 + }) + .collect() + }) + .collect() +} + +#[test] +#[ignore = "needs pinned9B export allowing4224tokens, CUDA and GPU reservation"] +fn ordinary_graphs_preserve_live_inputs_around_eager_tree_and_growth() { + let dir = PathBuf::from(std::env::var_os("OPEN_JEV_MODEL").unwrap()); + let lib = PathBuf::from(std::env::var_os("OPEN_JEV_CUDA_LIB").unwrap()); + let load = |enabled| { + let mut model = Model::load_with_graph(&dir, &lib, enabled).unwrap(); + model.enable_packed_projections().unwrap(); + model.enable_batched_gdn().unwrap(); + assert!(model.enable_equivalent_projection_tiles().unwrap()); + model + }; + let (mut eager, mut graph) = (load(false), load(true)); + for (index, prefix) in [64, 512, 4032, 4160, 64].into_iter().enumerate() { + for seed in [index * 5 + 1, index * 5 + 2] { + let ids = inputs(128, 0, 1, seed); + let refs: Vec<_> = ids.iter().map(Vec::as_slice).collect(); + assert_eq!( + eager.forward_batch(&refs).unwrap(), + graph.forward_batch(&refs).unwrap() + ); + } + let before = graph.graph_stats(); + assert!(before.captures > 0 && before.replays > 0); + let ids = inputs(prefix + 64, prefix, 3, index * 5 + 3); + let refs: Vec<_> = ids.iter().map(Vec::as_slice).collect(); + let original = if ids.iter().map(Vec::len).sum::() <= 4096 { + vec![(prefix + 64) * 3; 3] + } else { + vec![prefix + 64; 3] + }; + let want = eager + .try_forward_prefix_tree(&refs, prefix, &original) + .unwrap() + .unwrap(); + let got = graph + .try_forward_prefix_tree(&refs, prefix, &original) + .unwrap() + .unwrap(); + assert_eq!(want.hidden, got.hidden); + assert_eq!(want.saved_tokens, got.saved_tokens); + let after = graph.graph_stats(); + assert_eq!(before.captures, after.captures, "tree must stay eager"); + assert_eq!( + before.replays, after.replays, + "tree must not replay ordinary graphs" + ); + } + assert!( + graph.graph_stats().invalidations > 0, + "exercise replacement of captured scratch" + ); + assert_eq!(graph.graph_stats().fallbacks, 0); + // Different candidate order and fresh IDs must retain independent causal rows. + for seed in [101, 102] { + let ids = inputs(128, 0, 2, seed); + for order in [[0, 1], [1, 0]] { + let refs: Vec<_> = order.iter().map(|&i| ids[i].as_slice()).collect(); + assert_eq!( + eager.forward_batch(&refs).unwrap(), + graph.forward_batch(&refs).unwrap() + ); + } + } +} + +#[test] +#[ignore = "needs pinned9B export allowing4224tokens, CUDA and GPU reservation"] +fn compact_graphs_keep_maps_keys_and_fresh_inputs_after_tree_growth() { + let dir = PathBuf::from(std::env::var_os("OPEN_JEV_MODEL").unwrap()); + let lib = PathBuf::from(std::env::var_os("OPEN_JEV_CUDA_LIB").unwrap()); + let load = |enabled| { + let mut m = Model::load_with_graph(&dir, &lib, enabled).unwrap(); + m.enable_packed_projections().unwrap(); + m.enable_batched_gdn().unwrap(); + assert!(m.enable_equivalent_projection_tiles().unwrap()); + m + }; + let (mut eager, mut graph) = (load(false), load(true)); + // An ordinary layout with the same candidate lengths must not alias compact keys. + let ids = inputs(128, 37, 4, 1); + let refs: Vec<_> = ids.iter().map(Vec::as_slice).collect(); + assert_eq!( + eager.forward_batch(&refs).unwrap(), + graph.forward_batch(&refs).unwrap() + ); + let initial_invalidations = graph.graph_stats().invalidations; + for (length, prefix, count) in [(128, 37, 4), (128, 63, 4), (512, 480, 8), (128, 37, 4)] { + let before = graph.graph_stats().replays; + for seed in [2, 3, 4] { + let ids = inputs(length, prefix, count, seed); + let refs: Vec<_> = ids.iter().map(Vec::as_slice).collect(); + assert_eq!( + eager + .try_forward_prefix_compaction(&refs, prefix) + .unwrap() + .unwrap(), + graph + .try_forward_prefix_compaction(&refs, prefix) + .unwrap() + .unwrap() + ); + } + assert!( + graph.graph_stats().replays >= before + 2, + "exercise live-map graph replay" + ); + } + assert_eq!( + graph.graph_stats().invalidations, + initial_invalidations, + "first map allocation must preserve ordinary graphs" + ); + let before = graph.graph_stats().invalidations; + let ids = inputs(4224, 4160, 3, 5); + let refs: Vec<_> = ids.iter().map(Vec::as_slice).collect(); + assert_eq!( + eager + .try_forward_prefix_tree(&refs, 4160, &[4224; 3]) + .unwrap() + .unwrap() + .hidden, + graph + .try_forward_prefix_tree(&refs, 4160, &[4224; 3]) + .unwrap() + .unwrap() + .hidden + ); + assert!(graph.graph_stats().invalidations > before); + let before = graph.graph_stats(); + for seed in [6, 7] { + let ids = inputs(128, 37, 4, seed); + let refs: Vec<_> = ids.iter().map(Vec::as_slice).collect(); + assert_eq!( + eager + .try_forward_prefix_compaction(&refs, 37) + .unwrap() + .unwrap(), + graph + .try_forward_prefix_compaction(&refs, 37) + .unwrap() + .unwrap() + ); + } + assert_eq!(graph.graph_stats().captures, before.captures + 1); + assert_eq!(graph.graph_stats().replays, before.replays + 1); + assert_eq!(graph.graph_stats().fallbacks, 0); +} diff --git a/tests/runtime/batching.rs b/tests/runtime/batching.rs new file mode 100644 index 00000000..b146065f --- /dev/null +++ b/tests/runtime/batching.rs @@ -0,0 +1,314 @@ +use std::future::{Future, poll_fn}; +use std::pin::Pin; +use std::sync::{Arc, Mutex, mpsc}; +use std::task::Poll; +use std::time::Duration; + +use anyhow::{Result, bail}; +use omni_runtime::batching::{BatchCost, BatchLimits, BatchScheduler}; +use tokio::sync::oneshot; + +async fn done(future: impl Future) -> T { + tokio::time::timeout(Duration::from_secs(5), future) + .await + .expect("batch scheduler stopped making progress") +} + +async fn queued(mut future: Pin<&mut impl Future>) { + poll_fn(|cx| { + assert!(future.as_mut().poll(cx).is_pending()); + Poll::Ready(()) + }) + .await; +} + +#[tokio::test] +async fn batches_ready_requests_and_restores_each_request_in_fifo_order() { + let calls = Arc::new(Mutex::new(Vec::new())); + let observed = calls.clone(); + let (started, ready) = oneshot::channel(); + let mut started = Some(started); + let (release, wait) = mpsc::channel(); + let scheduler = BatchScheduler::new( + BatchLimits { items: 4, cost: 6 }, + |&item| BatchCost::from(if item == 10 { 99 } else { 1 }), + move |items: &[usize]| { + observed.lock().unwrap().push(items.to_vec()); + if items == [0] { + started.take().unwrap().send(()).unwrap(); + wait.recv().unwrap(); + } + Ok(items.iter().map(|x| x * 100).collect()) + }, + ) + .unwrap(); + let mut first = Box::pin(scheduler.run(vec![0])); + queued(first.as_mut()).await; + done(ready).await.unwrap(); + let mut a = Box::pin(scheduler.run(vec![1, 2, 3])); + let mut b = Box::pin(scheduler.run(vec![4, 5])); + let mut large = Box::pin(scheduler.run(vec![10])); + let mut last = Box::pin(scheduler.run(vec![11, 12])); + queued(a.as_mut()).await; + queued(b.as_mut()).await; + queued(large.as_mut()).await; + queued(last.as_mut()).await; + release.send(()).unwrap(); + assert_eq!(done(first).await.unwrap(), [0]); + assert_eq!(done(a).await.unwrap(), [100, 200, 300]); + assert_eq!(done(b).await.unwrap(), [400, 500]); + assert_eq!(done(large).await.unwrap(), [1000]); + assert_eq!(done(last).await.unwrap(), [1100, 1200]); + assert_eq!( + *calls.lock().unwrap(), + vec![vec![0], vec![1, 2, 3, 4], vec![5], vec![10], vec![11, 12]] + ); +} + +#[tokio::test] +async fn cancellation_before_dispatch_skips_the_request() { + let calls = Arc::new(Mutex::new(Vec::new())); + let observed = calls.clone(); + let (started, ready) = oneshot::channel(); + let mut started = Some(started); + let (release, wait) = mpsc::channel(); + let scheduler = BatchScheduler::new( + BatchLimits { items: 4, cost: 4 }, + |_| BatchCost::from(1), + move |items: &[usize]| { + observed.lock().unwrap().extend_from_slice(items); + if items == [0] { + started.take().unwrap().send(()).unwrap(); + wait.recv().unwrap(); + } + Ok(items.to_vec()) + }, + ) + .unwrap(); + let mut first = Box::pin(scheduler.run(vec![0])); + queued(first.as_mut()).await; + done(ready).await.unwrap(); + let mut cancelled = Box::pin(scheduler.run(vec![1, 2])); + queued(cancelled.as_mut()).await; + drop(cancelled); + let mut next = Box::pin(scheduler.run(vec![3])); + queued(next.as_mut()).await; + release.send(()).unwrap(); + done(first).await.unwrap(); + assert_eq!(done(next).await.unwrap(), [3]); + assert_eq!(*calls.lock().unwrap(), [0, 3]); +} + +#[tokio::test] +async fn cancellation_after_dispatch_finishes_all_parts_before_later_requests() { + let calls = Arc::new(Mutex::new(Vec::new())); + let observed = calls.clone(); + let (started, ready) = oneshot::channel(); + let mut started = Some(started); + let (release, wait) = mpsc::channel(); + let scheduler = BatchScheduler::new( + BatchLimits { items: 1, cost: 1 }, + |_| BatchCost::from(1), + move |items: &[usize]| { + observed.lock().unwrap().extend_from_slice(items); + if items == [1] { + started.take().unwrap().send(()).unwrap(); + wait.recv().unwrap(); + } + Ok(items.to_vec()) + }, + ) + .unwrap(); + let mut cancelled = Box::pin(scheduler.run(vec![1, 2])); + queued(cancelled.as_mut()).await; + done(ready).await.unwrap(); + drop(cancelled); + let mut next = Box::pin(scheduler.run(vec![3])); + queued(next.as_mut()).await; + release.send(()).unwrap(); + assert_eq!(done(next).await.unwrap(), [3]); + assert_eq!(*calls.lock().unwrap(), [1, 2, 3]); +} + +#[tokio::test] +async fn errors_discard_partial_results_and_leave_later_requests_usable() { + for malformed in [false, true] { + let calls = Arc::new(Mutex::new(Vec::new())); + let observed = calls.clone(); + let scheduler = BatchScheduler::new( + BatchLimits { items: 1, cost: 1 }, + |_| BatchCost::from(1), + move |items: &[usize]| { + observed.lock().unwrap().extend_from_slice(items); + if items == [2] { + if malformed { + return Ok(Vec::new()); + } + bail!("device error"); + } + Ok(items.to_vec()) + }, + ) + .unwrap(); + assert!(done(scheduler.run(vec![1, 2, 3])).await.is_err()); + assert_eq!(done(scheduler.run(vec![4])).await.unwrap(), [4]); + assert_eq!(*calls.lock().unwrap(), [1, 2, 4]); + } +} + +#[tokio::test] +async fn a_panicked_executor_closes_admission_instead_of_reusing_bad_state() { + let scheduler: BatchScheduler = BatchScheduler::new( + BatchLimits { items: 1, cost: 1 }, + |_| BatchCost::from(1), + |_| -> Result> { panic!("device executor failed") }, + ) + .unwrap(); + assert!( + done(scheduler.run(vec![1])) + .await + .unwrap_err() + .to_string() + .contains("panicked") + ); + assert!(done(scheduler.run(vec![2])).await.is_err()); +} + +#[tokio::test] +async fn oversized_cost_cannot_overflow_the_batch_budget() { + let scheduler = BatchScheduler::new( + BatchLimits { + items: 3, + cost: usize::MAX - 1, + }, + |&item| BatchCost::from(item), + |items: &[usize]| { + assert!(items.len() == 1 || items.iter().sum::() < usize::MAX); + Ok(items.to_vec()) + }, + ) + .unwrap(); + assert_eq!( + done(scheduler.run(vec![usize::MAX, 1, 2])).await.unwrap(), + [usize::MAX, 1, 2] + ); + assert!(done(scheduler.run(Vec::new())).await.unwrap().is_empty()); +} + +#[tokio::test] +async fn dropping_callers_and_schedulers_retains_executing_resources() { + struct Resource(Option>); + impl Drop for Resource { + fn drop(&mut self) { + let _ = self.0.take().unwrap().send(()); + } + } + let (dropped, mut retired) = oneshot::channel(); + let resource = Resource(Some(dropped)); + let (started, ready) = oneshot::channel(); + let mut started = Some(started); + let (release, wait) = mpsc::channel(); + let scheduler = BatchScheduler::new( + BatchLimits { items: 1, cost: 1 }, + |_| BatchCost::from(1), + move |items: &[usize]| { + let _keep = &resource; + started.take().unwrap().send(()).unwrap(); + wait.recv().unwrap(); + Ok(items.to_vec()) + }, + ) + .unwrap(); + let mut request = Box::pin(scheduler.run(vec![1])); + queued(request.as_mut()).await; + done(ready).await.unwrap(); + drop(request); + // Dropping the last owner now waits for shutdown; do it off the async + // thread which supplies this test callback's release signal. + let shutdown = tokio::task::spawn_blocking(move || drop(scheduler)); + assert!(matches!( + retired.try_recv(), + Err(oneshot::error::TryRecvError::Empty) + )); + release.send(()).unwrap(); + done(retired).await.unwrap(); + done(shutdown).await.unwrap(); +} + +#[tokio::test] +async fn weighted_atomic_items_and_exclusive_work_keep_fifo_boundaries() { + let calls = Arc::new(Mutex::new(Vec::new())); + let seen = calls.clone(); + let (started, ready) = oneshot::channel(); + let mut started = Some(started); + let (release, wait) = mpsc::channel(); + let scheduler = BatchScheduler::new( + BatchLimits { + items: 4, + cost: 100, + }, + |&id| BatchCost { + items: match id { + 1 => 3, + 2 | 3 => 2, + _ => 1, + }, + cost: 1, + exclusive: id == 9, + }, + move |ids: &[usize]| { + seen.lock().unwrap().push(ids.to_vec()); + if ids == [0] { + started.take().unwrap().send(()).unwrap(); + wait.recv().unwrap(); + } + Ok(ids.iter().map(|id| id * 10).collect()) + }, + ) + .unwrap(); + let mut first = Box::pin(scheduler.run(vec![0])); + queued(first.as_mut()).await; + done(ready).await.unwrap(); + let mut a = Box::pin(scheduler.run(vec![1, 2, 3])); + let mut b = Box::pin(scheduler.run(vec![9, 4, 5])); + queued(a.as_mut()).await; + queued(b.as_mut()).await; + release.send(()).unwrap(); + assert_eq!(done(first).await.unwrap(), [0]); + assert_eq!(done(a).await.unwrap(), [10, 20, 30]); + assert_eq!(done(b).await.unwrap(), [90, 40, 50]); + assert_eq!( + *calls.lock().unwrap(), + [vec![0], vec![1], vec![2, 3], vec![9], vec![4, 5]] + ); +} + +#[test] +fn last_scheduler_drop_joins_callback_resource_retirement() { + use std::sync::atomic::{AtomicBool, Ordering}; + struct Retired(Arc); + impl Drop for Retired { + fn drop(&mut self) { + self.0.store(true, Ordering::SeqCst); + } + } + let retired = Arc::new(AtomicBool::new(false)); + let resource = Retired(retired.clone()); + let scheduler = BatchScheduler::new( + BatchLimits { items: 1, cost: 1 }, + |_| BatchCost::from(1), + move |items: &[usize]| { + let _keep_alive = &resource; + Ok(items.to_vec()) + }, + ) + .unwrap(); + let last = scheduler.clone(); + drop(scheduler); + assert!(!retired.load(Ordering::SeqCst)); + drop(last); + assert!( + retired.load(Ordering::SeqCst), + "callback resources must retire before owner exit" + ); +} From ecea1a31e287890fee5a595504203b53734ffe2a Mon Sep 17 00:00:00 2001 From: psv666 <150513104+psv666@users.noreply.github.com> Date: Thu, 29 Oct 2026 10:37:12 +0800 Subject: [PATCH 2/3] test(open-jev): compare mixed packing with frozen original native --- src/models/open_jev/native/Cargo.toml | 4 + tests/open_jev/data/packing_cases.json | 588 ++++++++++ tests/open_jev/data/packing_original.json | 1201 +++++++++++++++++++++ tests/open_jev/packing_parity.md | 95 ++ tests/open_jev/packing_parity.rs | 146 +++ tests/open_jev/prefill_batch.rs | 2 +- 6 files changed, 2035 insertions(+), 1 deletion(-) create mode 100644 tests/open_jev/data/packing_cases.json create mode 100644 tests/open_jev/data/packing_original.json create mode 100644 tests/open_jev/packing_parity.md create mode 100644 tests/open_jev/packing_parity.rs diff --git a/src/models/open_jev/native/Cargo.toml b/src/models/open_jev/native/Cargo.toml index dba85c30..b73ea0b9 100644 --- a/src/models/open_jev/native/Cargo.toml +++ b/src/models/open_jev/native/Cargo.toml @@ -25,3 +25,7 @@ path = "../../../../tests/open_jev/prefill_batch.rs" [[test]] name = "prefix" path = "../../../../tests/open_jev/prefix.rs" + +[[test]] +name = "packing_parity" +path = "../../../../tests/open_jev/packing_parity.rs" diff --git a/tests/open_jev/data/packing_cases.json b/tests/open_jev/data/packing_cases.json new file mode 100644 index 00000000..02dc2611 --- /dev/null +++ b/tests/open_jev/data/packing_cases.json @@ -0,0 +1,588 @@ +{ + "version": 1, + "token_rule": "token[i] = 1000 + (seed * 7919 + i * 37) % 30000", + "cases": [ + { + "id": "mixed-a", + "candidates": [ + [ + 128, + 1 + ], + [ + 257, + 2 + ], + [ + 511, + 3 + ] + ], + "groups": [ + 3 + ], + "questions": [ + 3 + ] + }, + { + "id": "mixed-reverse", + "candidates": [ + [ + 511, + 3 + ], + [ + 257, + 2 + ], + [ + 128, + 1 + ] + ], + "groups": [ + 3 + ], + "questions": [ + 3 + ] + }, + { + "id": "mixed-permute", + "candidates": [ + [ + 257, + 2 + ], + [ + 128, + 1 + ], + [ + 511, + 3 + ] + ], + "groups": [ + 3 + ], + "questions": [ + 3 + ] + }, + { + "id": "mixed-changed", + "candidates": [ + [ + 128, + 21 + ], + [ + 257, + 22 + ], + [ + 511, + 23 + ] + ], + "groups": [ + 3 + ], + "questions": [ + 3 + ] + }, + { + "id": "mixed-return-a", + "candidates": [ + [ + 128, + 1 + ], + [ + 257, + 2 + ], + [ + 511, + 3 + ] + ], + "groups": [ + 3 + ], + "questions": [ + 3 + ] + }, + { + "id": "singleton-after-pack", + "candidates": [ + [ + 128, + 1 + ] + ], + "groups": [ + 1 + ], + "questions": [ + 1 + ] + }, + { + "id": "unaligned", + "candidates": [ + [ + 37, + 1 + ], + [ + 79, + 2 + ], + [ + 81, + 3 + ], + [ + 129, + 4 + ], + [ + 255, + 5 + ], + [ + 513, + 6 + ] + ], + "groups": [ + 6 + ], + "questions": [ + 6 + ] + }, + { + "id": "unaligned-changed", + "candidates": [ + [ + 37, + 31 + ], + [ + 79, + 32 + ], + [ + 81, + 33 + ], + [ + 129, + 34 + ], + [ + 255, + 35 + ], + [ + 513, + 36 + ] + ], + "groups": [ + 6 + ], + "questions": [ + 6 + ] + }, + { + "id": "unaligned-return", + "candidates": [ + [ + 37, + 1 + ], + [ + 79, + 2 + ], + [ + 81, + 3 + ], + [ + 129, + 4 + ], + [ + 255, + 5 + ], + [ + 513, + 6 + ] + ], + "groups": [ + 6 + ], + "questions": [ + 6 + ] + }, + { + "id": "budget-4095", + "candidates": [ + [ + 511, + 1 + ], + [ + 1024, + 2 + ], + [ + 2560, + 3 + ] + ], + "groups": [ + 3 + ], + "questions": [ + 3 + ] + }, + { + "id": "budget-4096", + "candidates": [ + [ + 512, + 1 + ], + [ + 1024, + 2 + ], + [ + 2560, + 3 + ] + ], + "groups": [ + 3 + ], + "questions": [ + 3 + ] + }, + { + "id": "budget-4097", + "candidates": [ + [ + 513, + 1 + ], + [ + 1024, + 2 + ], + [ + 2560, + 3 + ] + ], + "groups": [ + 2, + 1 + ], + "questions": [ + 3 + ] + }, + { + "id": "count-16", + "candidates": [ + [ + 63, + 1 + ], + [ + 64, + 2 + ], + [ + 65, + 3 + ], + [ + 66, + 4 + ], + [ + 67, + 5 + ], + [ + 68, + 6 + ], + [ + 69, + 7 + ], + [ + 70, + 8 + ], + [ + 71, + 9 + ], + [ + 72, + 10 + ], + [ + 73, + 11 + ], + [ + 74, + 12 + ], + [ + 75, + 13 + ], + [ + 76, + 14 + ], + [ + 77, + 15 + ], + [ + 78, + 16 + ] + ], + "groups": [ + 16 + ], + "questions": [ + 16 + ] + }, + { + "id": "count-17", + "candidates": [ + [ + 63, + 1 + ], + [ + 64, + 2 + ], + [ + 65, + 3 + ], + [ + 66, + 4 + ], + [ + 67, + 5 + ], + [ + 68, + 6 + ], + [ + 69, + 7 + ], + [ + 70, + 8 + ], + [ + 71, + 9 + ], + [ + 72, + 10 + ], + [ + 73, + 11 + ], + [ + 74, + 12 + ], + [ + 75, + 13 + ], + [ + 76, + 14 + ], + [ + 77, + 15 + ], + [ + 78, + 16 + ], + [ + 79, + 17 + ] + ], + "groups": [ + 16, + 1 + ], + "questions": [ + 17 + ] + }, + { + "id": "questions-mixed", + "candidates": [ + [ + 79, + 1 + ], + [ + 128, + 2 + ], + [ + 257, + 3 + ], + [ + 511, + 4 + ], + [ + 129, + 5 + ], + [ + 83, + 6 + ] + ], + "groups": [ + 6 + ], + "questions": [ + 1, + 3, + 2 + ] + }, + { + "id": "questions-reverse", + "candidates": [ + [ + 83, + 6 + ], + [ + 129, + 5 + ], + [ + 511, + 4 + ], + [ + 257, + 3 + ], + [ + 128, + 2 + ], + [ + 79, + 1 + ] + ], + "groups": [ + 6 + ], + "questions": [ + 2, + 3, + 1 + ] + }, + { + "id": "long-split", + "candidates": [ + [ + 2112, + 1 + ], + [ + 2112, + 2 + ], + [ + 2112, + 3 + ] + ], + "groups": [ + 1, + 1, + 1 + ], + "questions": [ + 3 + ] + }, + { + "id": "mixed-return-after-long", + "candidates": [ + [ + 128, + 1 + ], + [ + 257, + 2 + ], + [ + 511, + 3 + ] + ], + "groups": [ + 3 + ], + "questions": [ + 3 + ] + } + ] +} diff --git a/tests/open_jev/data/packing_original.json b/tests/open_jev/data/packing_original.json new file mode 100644 index 00000000..c425c15d --- /dev/null +++ b/tests/open_jev/data/packing_original.json @@ -0,0 +1,1201 @@ +{ + "original_revision": "22b7717c9710de5ac5c6ccba51a6d6686dae90d9", + "manifest": { + "version": 1, + "token_rule": "token[i] = 1000 + (seed * 7919 + i * 37) % 30000", + "cases": [ + { + "id": "mixed-a", + "candidates": [ + [ + 128, + 1 + ], + [ + 257, + 2 + ], + [ + 511, + 3 + ] + ], + "groups": [ + 3 + ], + "questions": [ + 3 + ] + }, + { + "id": "mixed-reverse", + "candidates": [ + [ + 511, + 3 + ], + [ + 257, + 2 + ], + [ + 128, + 1 + ] + ], + "groups": [ + 3 + ], + "questions": [ + 3 + ] + }, + { + "id": "mixed-permute", + "candidates": [ + [ + 257, + 2 + ], + [ + 128, + 1 + ], + [ + 511, + 3 + ] + ], + "groups": [ + 3 + ], + "questions": [ + 3 + ] + }, + { + "id": "mixed-changed", + "candidates": [ + [ + 128, + 21 + ], + [ + 257, + 22 + ], + [ + 511, + 23 + ] + ], + "groups": [ + 3 + ], + "questions": [ + 3 + ] + }, + { + "id": "mixed-return-a", + "candidates": [ + [ + 128, + 1 + ], + [ + 257, + 2 + ], + [ + 511, + 3 + ] + ], + "groups": [ + 3 + ], + "questions": [ + 3 + ] + }, + { + "id": "singleton-after-pack", + "candidates": [ + [ + 128, + 1 + ] + ], + "groups": [ + 1 + ], + "questions": [ + 1 + ] + }, + { + "id": "unaligned", + "candidates": [ + [ + 37, + 1 + ], + [ + 79, + 2 + ], + [ + 81, + 3 + ], + [ + 129, + 4 + ], + [ + 255, + 5 + ], + [ + 513, + 6 + ] + ], + "groups": [ + 6 + ], + "questions": [ + 6 + ] + }, + { + "id": "unaligned-changed", + "candidates": [ + [ + 37, + 31 + ], + [ + 79, + 32 + ], + [ + 81, + 33 + ], + [ + 129, + 34 + ], + [ + 255, + 35 + ], + [ + 513, + 36 + ] + ], + "groups": [ + 6 + ], + "questions": [ + 6 + ] + }, + { + "id": "unaligned-return", + "candidates": [ + [ + 37, + 1 + ], + [ + 79, + 2 + ], + [ + 81, + 3 + ], + [ + 129, + 4 + ], + [ + 255, + 5 + ], + [ + 513, + 6 + ] + ], + "groups": [ + 6 + ], + "questions": [ + 6 + ] + }, + { + "id": "budget-4095", + "candidates": [ + [ + 511, + 1 + ], + [ + 1024, + 2 + ], + [ + 2560, + 3 + ] + ], + "groups": [ + 3 + ], + "questions": [ + 3 + ] + }, + { + "id": "budget-4096", + "candidates": [ + [ + 512, + 1 + ], + [ + 1024, + 2 + ], + [ + 2560, + 3 + ] + ], + "groups": [ + 3 + ], + "questions": [ + 3 + ] + }, + { + "id": "budget-4097", + "candidates": [ + [ + 513, + 1 + ], + [ + 1024, + 2 + ], + [ + 2560, + 3 + ] + ], + "groups": [ + 2, + 1 + ], + "questions": [ + 3 + ] + }, + { + "id": "count-16", + "candidates": [ + [ + 63, + 1 + ], + [ + 64, + 2 + ], + [ + 65, + 3 + ], + [ + 66, + 4 + ], + [ + 67, + 5 + ], + [ + 68, + 6 + ], + [ + 69, + 7 + ], + [ + 70, + 8 + ], + [ + 71, + 9 + ], + [ + 72, + 10 + ], + [ + 73, + 11 + ], + [ + 74, + 12 + ], + [ + 75, + 13 + ], + [ + 76, + 14 + ], + [ + 77, + 15 + ], + [ + 78, + 16 + ] + ], + "groups": [ + 16 + ], + "questions": [ + 16 + ] + }, + { + "id": "count-17", + "candidates": [ + [ + 63, + 1 + ], + [ + 64, + 2 + ], + [ + 65, + 3 + ], + [ + 66, + 4 + ], + [ + 67, + 5 + ], + [ + 68, + 6 + ], + [ + 69, + 7 + ], + [ + 70, + 8 + ], + [ + 71, + 9 + ], + [ + 72, + 10 + ], + [ + 73, + 11 + ], + [ + 74, + 12 + ], + [ + 75, + 13 + ], + [ + 76, + 14 + ], + [ + 77, + 15 + ], + [ + 78, + 16 + ], + [ + 79, + 17 + ] + ], + "groups": [ + 16, + 1 + ], + "questions": [ + 17 + ] + }, + { + "id": "questions-mixed", + "candidates": [ + [ + 79, + 1 + ], + [ + 128, + 2 + ], + [ + 257, + 3 + ], + [ + 511, + 4 + ], + [ + 129, + 5 + ], + [ + 83, + 6 + ] + ], + "groups": [ + 6 + ], + "questions": [ + 1, + 3, + 2 + ] + }, + { + "id": "questions-reverse", + "candidates": [ + [ + 83, + 6 + ], + [ + 129, + 5 + ], + [ + 511, + 4 + ], + [ + 257, + 3 + ], + [ + 128, + 2 + ], + [ + 79, + 1 + ] + ], + "groups": [ + 6 + ], + "questions": [ + 2, + 3, + 1 + ] + }, + { + "id": "long-split", + "candidates": [ + [ + 2112, + 1 + ], + [ + 2112, + 2 + ], + [ + 2112, + 3 + ] + ], + "groups": [ + 1, + 1, + 1 + ], + "questions": [ + 3 + ] + }, + { + "id": "mixed-return-after-long", + "candidates": [ + [ + 128, + 1 + ], + [ + 257, + 2 + ], + [ + 511, + 3 + ] + ], + "groups": [ + 3 + ], + "questions": [ + 3 + ] + } + ] + }, + "records": [ + { + "id": "mixed-a", + "lengths": [ + 128, + 257, + 511 + ], + "policy_groups": [ + 3 + ], + "packing_enabled": false, + "scores": [ + [ + -0.8836134076118469, + -2.377385139465332, + -2.1972479820251465 + ] + ], + "score_bits": [ + [ + 3210884221, + 3222808340, + 3222052790 + ] + ] + }, + { + "id": "mixed-reverse", + "lengths": [ + 511, + 257, + 128 + ], + "policy_groups": [ + 3 + ], + "packing_enabled": false, + "scores": [ + [ + -2.1972479820251465, + -2.377385139465332, + -0.8836134076118469 + ] + ], + "score_bits": [ + [ + 3222052790, + 3222808340, + 3210884221 + ] + ] + }, + { + "id": "mixed-permute", + "lengths": [ + 257, + 128, + 511 + ], + "policy_groups": [ + 3 + ], + "packing_enabled": false, + "scores": [ + [ + -2.377385139465332, + -0.8836134076118469, + -2.1972479820251465 + ] + ], + "score_bits": [ + [ + 3222808340, + 3210884221, + 3222052790 + ] + ] + }, + { + "id": "mixed-changed", + "lengths": [ + 128, + 257, + 511 + ], + "policy_groups": [ + 3 + ], + "packing_enabled": false, + "scores": [ + [ + -2.6016812324523926, + -2.0423574447631836, + -1.6241869926452637 + ] + ], + "score_bits": [ + [ + 3223749106, + 3221403132, + 3218072924 + ] + ] + }, + { + "id": "mixed-return-a", + "lengths": [ + 128, + 257, + 511 + ], + "policy_groups": [ + 3 + ], + "packing_enabled": false, + "scores": [ + [ + -0.8836134076118469, + -2.377385139465332, + -2.1972479820251465 + ] + ], + "score_bits": [ + [ + 3210884221, + 3222808340, + 3222052790 + ] + ] + }, + { + "id": "singleton-after-pack", + "lengths": [ + 128 + ], + "policy_groups": [ + 1 + ], + "packing_enabled": false, + "scores": [ + [ + -0.8836134076118469 + ] + ], + "score_bits": [ + [ + 3210884221 + ] + ] + }, + { + "id": "unaligned", + "lengths": [ + 37, + 79, + 81, + 129, + 255, + 513 + ], + "policy_groups": [ + 6 + ], + "packing_enabled": false, + "scores": [ + [ + -1.2803736925125122, + -1.9052419662475586, + 0.1501929759979248, + -0.8065001964569092, + -1.9822355508804321, + -2.403242826461792 + ] + ], + "score_bits": [ + [ + 3215188809, + 3220430584, + 1041878064, + 3209590476, + 3221076453, + 3222916795 + ] + ] + }, + { + "id": "unaligned-changed", + "lengths": [ + 37, + 79, + 81, + 129, + 255, + 513 + ], + "policy_groups": [ + 6 + ], + "packing_enabled": false, + "scores": [ + [ + -1.6164536476135254, + -1.3435107469558716, + -2.7300872802734375, + -1.6214776039123535, + -1.488254427909851, + -2.3121819496154785 + ] + ], + "score_bits": [ + [ + 3218008052, + 3215718441, + 3224287680, + 3218050196, + 3216932639, + 3222534858 + ] + ] + }, + { + "id": "unaligned-return", + "lengths": [ + 37, + 79, + 81, + 129, + 255, + 513 + ], + "policy_groups": [ + 6 + ], + "packing_enabled": false, + "scores": [ + [ + -1.2803736925125122, + -1.9052419662475586, + 0.1501929759979248, + -0.8065001964569092, + -1.9822355508804321, + -2.403242826461792 + ] + ], + "score_bits": [ + [ + 3215188809, + 3220430584, + 1041878064, + 3209590476, + 3221076453, + 3222916795 + ] + ] + }, + { + "id": "budget-4095", + "lengths": [ + 511, + 1024, + 2560 + ], + "policy_groups": [ + 3 + ], + "packing_enabled": false, + "scores": [ + [ + -2.5573623180389404, + -2.3214974403381348, + -2.9056589603424072 + ] + ], + "score_bits": [ + [ + 3223563219, + 3222573930, + 3225024081 + ] + ] + }, + { + "id": "budget-4096", + "lengths": [ + 512, + 1024, + 2560 + ], + "policy_groups": [ + 3 + ], + "packing_enabled": false, + "scores": [ + [ + -2.5780467987060547, + -2.3214974403381348, + -2.9056589603424072 + ] + ], + "score_bits": [ + [ + 3223649976, + 3222573930, + 3225024081 + ] + ] + }, + { + "id": "budget-4097", + "lengths": [ + 513, + 1024, + 2560 + ], + "policy_groups": [ + 2, + 1 + ], + "packing_enabled": false, + "scores": [ + [ + -1.235496163368225, + -2.3214974403381348, + -2.9056589603424072 + ] + ], + "score_bits": [ + [ + 3214812349, + 3222573930, + 3225024081 + ] + ] + }, + { + "id": "count-16", + "lengths": [ + 63, + 64, + 65, + 66, + 67, + 68, + 69, + 70, + 71, + 72, + 73, + 74, + 75, + 76, + 77, + 78 + ], + "policy_groups": [ + 16 + ], + "packing_enabled": false, + "scores": [ + [ + -1.8805922269821167, + -2.1084887981414795, + -3.2593448162078857, + -1.5717893838882446, + -1.755454659461975, + -3.6777944564819336, + -1.5548018217086792, + -0.7498553991317749, + -1.8789061307907104, + -1.4451743364334106, + -2.4365954399108887, + -1.0026410818099976, + -2.6995840072631836, + -2.279020309448242, + -1.955730676651001, + -2.530183792114258 + ] + ], + "score_bits": [ + [ + 3220223807, + 3221680507, + 3226507547, + 3217633381, + 3219174077, + 3228262652, + 3217490879, + 3208640134, + 3220209663, + 3216571257, + 3223056686, + 3212859019, + 3224159740, + 3222395768, + 3220854114, + 3223449224 + ] + ] + }, + { + "id": "count-17", + "lengths": [ + 63, + 64, + 65, + 66, + 67, + 68, + 69, + 70, + 71, + 72, + 73, + 74, + 75, + 76, + 77, + 78, + 79 + ], + "policy_groups": [ + 16, + 1 + ], + "packing_enabled": false, + "scores": [ + [ + -1.8805922269821167, + -2.1084887981414795, + -3.2593448162078857, + -1.5717893838882446, + -1.755454659461975, + -3.6777944564819336, + -1.5548018217086792, + -0.7498553991317749, + -1.8789061307907104, + -1.4451743364334106, + -2.4365954399108887, + -1.0026410818099976, + -2.6995840072631836, + -2.279020309448242, + -1.955730676651001, + -2.530183792114258, + -1.6548871994018555 + ] + ], + "score_bits": [ + [ + 3220223807, + 3221680507, + 3226507547, + 3217633381, + 3219174077, + 3228262652, + 3217490879, + 3208640134, + 3220209663, + 3216571257, + 3223056686, + 3212859019, + 3224159740, + 3222395768, + 3220854114, + 3223449224, + 3218330456 + ] + ] + }, + { + "id": "questions-mixed", + "lengths": [ + 79, + 128, + 257, + 511, + 129, + 83 + ], + "policy_groups": [ + 6 + ], + "packing_enabled": false, + "scores": [ + [ + -2.032999277114868 + ], + [ + -2.4287283420562744, + -3.1898374557495117, + -1.9095982313156128 + ], + [ + -2.477084159851074, + -3.2781989574432373 + ] + ], + "score_bits": [ + [ + 3221363881 + ], + [ + 3223023689, + 3226216012, + 3220467127 + ], + [ + 3223226508, + 3226586627 + ] + ] + }, + { + "id": "questions-reverse", + "lengths": [ + 83, + 129, + 511, + 257, + 128, + 79 + ], + "policy_groups": [ + 6 + ], + "packing_enabled": false, + "scores": [ + [ + -3.2781989574432373, + -2.477084159851074 + ], + [ + -1.9095982313156128, + -3.1898374557495117, + -2.4287283420562744 + ], + [ + -2.032999277114868 + ] + ], + "score_bits": [ + [ + 3226586627, + 3223226508 + ], + [ + 3220467127, + 3226216012, + 3223023689 + ], + [ + 3221363881 + ] + ] + }, + { + "id": "long-split", + "lengths": [ + 2112, + 2112, + 2112 + ], + "policy_groups": [ + 1, + 1, + 1 + ], + "packing_enabled": false, + "scores": [ + [ + -2.385683536529541, + -2.3607938289642334, + -2.01570725440979 + ] + ], + "score_bits": [ + [ + 3222843146, + 3222738751, + 3221291353 + ] + ] + }, + { + "id": "mixed-return-after-long", + "lengths": [ + 128, + 257, + 511 + ], + "policy_groups": [ + 3 + ], + "packing_enabled": false, + "scores": [ + [ + -0.8836134076118469, + -2.377385139465332, + -2.1972479820251465 + ] + ], + "score_bits": [ + [ + 3210884221, + 3222808340, + 3222052790 + ] + ] + } + ] +} \ No newline at end of file diff --git a/tests/open_jev/packing_parity.md b/tests/open_jev/packing_parity.md new file mode 100644 index 00000000..94d9d0a6 --- /dev/null +++ b/tests/open_jev/packing_parity.md @@ -0,0 +1,95 @@ +# Open-Jev-9B mixed packing parity + +`packing_parity.rs` replays 18 deterministic token cases (94 candidate logits) +against an oracle generated by original native commit +`22b7717c9710de5ac5c6ccba51a6d6686dae90d9`. That revision runs 9B candidates +individually. The verification checkout packs them. This is an execution +regression test, not a task-accuracy or performance benchmark. + +The fixture covers mixed lengths 128/257/511, permutations, changed tokens and +returning inputs, unaligned lengths, total lengths 4095/4096/4097, 16/17 +candidates, multiple questions, and a long singleton-group control. Each case +declares the expected original packing groups; the CPU check calls the actual +production grouping function. The GPU check asserts that 9B packing is enabled, +prefix reuse is disabled, and the returned FP32 learned-head logits match the +frozen original **bit for bit**, before response normalization. It also checks +question/candidate counts. New-code singleton output is not the oracle. + +The reference is scoped to the pinned BF16 9B export from +[the native recipe](../../recipe/open_jev/native.md), NVIDIA L20X / measured +SM90, CUDA 13.0. Bitwise agreement on other architectures/toolchains is +unverified. Arbitrary synthetic token IDs do not establish semantic accuracy. +The existing typed HTTP benchmarks separately check Noul/Choice/Score outputs. + +## Verify the committed reference + +Build the current Rust worker and matching CUDA library as in the recipe. Set +`OPEN_JEV_MODEL` to the pinned export and `OPEN_JEV_CUDA_LIB` to the current +library. On the measured host, build the CUDA library with architecture `90`. + +```bash +cargo test --locked -p omni-open-jev-native --test packing_parity + +# Run the following inside your host's GPU reservation, e.g. gpu run -- ... . +for graph in 0 1; do + OPEN_JEV_PREFIX=0 OPEN_JEV_EXACT_PREFIX=0 CUA_S1_GRAPH="$graph" \ + cargo test --release --locked -p omni-open-jev-native \ + --test packing_parity mixed_packing_matches_frozen_original_native \ + -- --ignored --exact --nocapture --test-threads=1 +done +``` + +The explicit `record_original_native_reference` test is intentionally ignored +and refuses to run with packing enabled. Do not run all ignored tests in this +target: recording and verification require different source revisions. + +## Independently regenerate the original oracle + +Use a fresh worktree and **separate Cargo target directories**. Sharing a target +directory across these checkouts can reuse a test binary from the wrong source. +The packing-enabled assertion detects that mistake, but isolation prevents it. + +From the PR checkout, with absolute `REPLAY_DIR` and `OPEN_JEV_MODEL` paths: + +```bash +PR_CHECKOUT="$PWD" +ORIGINAL=22b7717c9710de5ac5c6ccba51a6d6686dae90d9 +mkdir -p "$REPLAY_DIR" +git worktree add --detach "$REPLAY_DIR/original" "$ORIGINAL" +cp tests/open_jev/packing_parity.rs "$REPLAY_DIR/original/tests/open_jev/" +cp tests/open_jev/data/packing_{cases,original}.json \ + "$REPLAY_DIR/original/tests/open_jev/data/" +cat >> "$REPLAY_DIR/original/src/models/open_jev/native/Cargo.toml" <<'EOF' + +[[test]] +name = "packing_parity" +path = "../../../../tests/open_jev/packing_parity.rs" +EOF + +cd "$REPLAY_DIR/original" +test "$(git rev-parse HEAD)" = "$ORIGINAL" +# Configure NVCC/CPATH/LIBRARY_PATH for your CUDA installation as needed. +src/backends/cuda/qwen3_5/build.sh "$REPLAY_DIR/cuda-original" 90 +CARGO_TARGET_DIR="$REPLAY_DIR/target-original" \ + cargo test --release --locked -p omni-open-jev-native \ + --test packing_parity --no-run + +# Inside a GPU reservation; original production source is unchanged. +OPEN_JEV_CUDA_LIB="$REPLAY_DIR/cuda-original/libqwen3_5_cuda.so" \ +OPEN_JEV_ORACLE_OUTPUT="$REPLAY_DIR/original-reference.json" \ +OPEN_JEV_PREFIX=0 OPEN_JEV_EXACT_PREFIX=0 CUA_S1_GRAPH=0 \ +CARGO_TARGET_DIR="$REPLAY_DIR/target-original" \ + cargo test --release --locked -p omni-open-jev-native \ + --test packing_parity record_original_native_reference \ + -- --ignored --exact --nocapture --test-threads=1 + +cmp "$REPLAY_DIR/original-reference.json" \ + "$PR_CHECKOUT/tests/open_jev/data/packing_original.json" +cd "$PR_CHECKOUT" +``` + +Then run verification above with the PR's CUDA library and a distinct candidate +target directory. Keep stdout/stderr, source revisions, library/executable +SHA256 hashes, GPU identity, export-file hashes, and both successful and failed +attempts. A regenerated oracle mismatch is a failed check; do not overwrite the +committed oracle to make verification pass. diff --git a/tests/open_jev/packing_parity.rs b/tests/open_jev/packing_parity.rs new file mode 100644 index 00000000..47f26c21 --- /dev/null +++ b/tests/open_jev/packing_parity.rs @@ -0,0 +1,146 @@ +//! Replay the same token manifest on the original native checkout and this one. +//! See packing_parity.md. Synthetic token cases test execution, not task accuracy. + +use std::path::PathBuf; + +use omni_open_jev_native::engine::Engine; +use omni_qwen3_5_native::prefix::PrefixMode; +use serde_json::{Value, json}; + +// Exercise the real grouping policy without making a production-private API public. +#[allow(dead_code)] +#[path = "../../src/models/open_jev/native/src/batching.rs"] +mod packing; + +const MANIFEST: &str = include_str!("data/packing_cases.json"); +const ORIGINAL: &str = "22b7717c9710de5ac5c6ccba51a6d6686dae90d9"; + +fn inputs(case: &Value) -> Vec>> { + let candidates: Vec> = case["candidates"] + .as_array() + .unwrap() + .iter() + .map(|spec| { + let length = spec[0].as_u64().unwrap() as usize; + let seed = spec[1].as_u64().unwrap() as u32; + // No shared prefix: each candidate starts with its own seed. + (0..length) + .map(|i| 1000 + (seed * 7919 + i as u32 * 37) % 30000) + .collect() + }) + .collect(); + let mut candidates = candidates.into_iter(); + let result = case["questions"] + .as_array() + .unwrap() + .iter() + .map(|n| { + candidates + .by_ref() + .take(n.as_u64().unwrap() as usize) + .collect() + }) + .collect(); + assert!(candidates.next().is_none()); + result +} + +fn checked_groups(case: &Value, ids: &[Vec>]) -> Vec { + let flat: Vec<&[u32]> = ids.iter().flatten().map(Vec::as_slice).collect(); + let actual: Vec = packing::ranges(&flat).iter().map(|r| r.len()).collect(); + assert_eq!(json!(actual), case["groups"], "{}", case["id"]); + assert_eq!(flat.len(), case["candidates"].as_array().unwrap().len()); + actual +} + +#[test] +fn manifest_exercises_real_mixed_packing_and_both_limits() { + let manifest: Value = serde_json::from_str(MANIFEST).unwrap(); + let cases = manifest["cases"].as_array().unwrap(); + let mut mixed = 0; + for case in cases { + let ids = inputs(case); + let groups = checked_groups(case, &ids); + let lengths: Vec = ids.iter().flatten().map(Vec::len).collect(); + if groups.len() == 1 && groups[0] > 1 && lengths.windows(2).any(|v| v[0] != v[1]) { + mixed += 1; + } + } + assert!(mixed >= 6, "must exercise genuinely packed mixed lengths"); + assert!(cases.iter().any(|c| c["groups"] == json!([16, 1]))); + assert!(cases.iter().any(|c| c["id"] == "budget-4096")); + assert!(cases.iter().any(|c| c["id"] == "budget-4097")); +} + +async fn replay(original: bool) -> Value { + 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_with_prefix(&model, &library, PrefixMode::Off) + .await + .unwrap(); + assert_eq!(engine.checkpoint.alias, "open-jev-9b"); + assert_eq!(engine.checkpoint.pack_candidates, !original); + let manifest: Value = serde_json::from_str(MANIFEST).unwrap(); + let mut records = Vec::new(); + for case in manifest["cases"].as_array().unwrap() { + let ids = inputs(case); + let groups = checked_groups(case, &ids); + let lengths: Vec = ids.iter().flatten().map(Vec::len).collect(); + let (scores, used) = engine + .executor + .execute_with_prefix(&engine.scheduler, ids) + .await + .unwrap(); + assert!(!used.shared); + assert_eq!(used.saved_tokens, 0); + assert_eq!(scores.len(), case["questions"].as_array().unwrap().len()); + for (row, count) in scores.iter().zip(case["questions"].as_array().unwrap()) { + assert_eq!(row.len(), count.as_u64().unwrap() as usize); + assert!(row.iter().all(|x| x.is_finite())); + } + let bits: Vec> = scores + .iter() + .map(|r| r.iter().map(|x| x.to_bits()).collect()) + .collect(); + let record = json!({"id":case["id"], "lengths":lengths, + "policy_groups":groups, "packing_enabled":engine.checkpoint.pack_candidates, + "scores":scores, "score_bits":bits}); + println!("PACKING_RECORD {record}"); + records.push(record); + } + json!({"original_revision":ORIGINAL, "manifest":manifest, "records":records}) +} + +#[tokio::test(flavor = "current_thread")] +#[ignore = "run only on the frozen original checkout, pinned 9B and a reserved GPU"] +async fn record_original_native_reference() { + // Deliberately separate from verification: ordinary test runs cannot rewrite the oracle. + assert_ne!(std::env::var("CUA_S1_GRAPH").as_deref(), Ok("1")); + let output = std::env::var_os("OPEN_JEV_ORACLE_OUTPUT").expect("explicit output required"); + let records = replay(true).await; + std::fs::write(output, serde_json::to_vec_pretty(&records).unwrap()).unwrap(); +} + +#[tokio::test(flavor = "current_thread")] +#[ignore = "needs pinned 9B, matching CUDA library and a reserved GPU; run eager and graph"] +async fn mixed_packing_matches_frozen_original_native() { + let actual = replay(false).await; + let reference: Value = + serde_json::from_str(include_str!("data/packing_original.json")).unwrap(); + assert_eq!(reference["original_revision"], ORIGINAL); + assert_eq!(actual["manifest"], reference["manifest"]); + let expected = reference["records"].as_array().unwrap(); + let actual = actual["records"].as_array().unwrap(); + assert_eq!(actual.len(), expected.len()); + for (actual, expected) in actual.iter().zip(expected) { + assert_eq!(actual["id"], expected["id"]); + assert_eq!(actual["policy_groups"], expected["policy_groups"]); + // Stronger than the public returned-probability tolerance: compare the + // complete learned-head FP32 logits before response normalization/rounding. + assert_eq!( + actual["score_bits"], expected["score_bits"], + "{}", + actual["id"] + ); + } +} diff --git a/tests/open_jev/prefill_batch.rs b/tests/open_jev/prefill_batch.rs index 4f13fa9b..fc04dcfd 100644 --- a/tests/open_jev/prefill_batch.rs +++ b/tests/open_jev/prefill_batch.rs @@ -145,7 +145,7 @@ async fn text_prefix_buffers_grow_and_shrink_without_reusing_prior_contents() { #[tokio::test(flavor = "current_thread")] #[ignore = "needs the pinned Open-Jev-9B export, CUDA library and a GPU reservation"] -async fn nine_b_packing_keeps_prefix_off_and_independent_candidates() { +async fn nine_b_long_candidates_keep_prefix_off_and_independent() { use omni_qwen3_5_native::prefix::PrefixMode; let model = PathBuf::from(std::env::var_os("OPEN_JEV_MODEL").unwrap()); From 9f782c6737f80b05491283699e1e1102f1679f32 Mon Sep 17 00:00:00 2001 From: psv666 <150513104+psv666@users.noreply.github.com> Date: Sat, 31 Oct 2026 09:24:49 +0800 Subject: [PATCH 3/3] fix(open-jev): guard device profiles, add serial rollback and report failed workers Restrict equivalent projection tiles to the measured NVIDIA L20X / SM90 / 132-SM / cuBLASLt 130101 profile. Add OPEN_JEV_PACK=0 to restore eager per-candidate 9B execution and serial preparation/admission. Report permanent executor failure on /health. CPU: 410 workspace tests passed, 107 ignored; 2 pinned-tokenizer tests and 27 CUDA backend CPU tests passed. Format, Clippy, release build, CUDA compilation/contract and strict docs passed. New-head GPU execution is pending availability; prior GPU/performance evidence remains historical. Source tree matches locally validated commit f7eb675134f14afd00421cec03197499cf676b83. --- recipe/open_jev/native.md | 52 ++++++++++++++- src/backends/cuda/qwen3_5/gemm.cu | 4 +- src/backends/cuda/qwen3_5/gemm_profile.h | 15 +++++ src/models/open_jev/README.md | 21 ++++-- src/models/open_jev/native/src/engine.rs | 54 ++++++++++++++-- src/models/open_jev/native/src/executor.rs | 33 ++++++++-- src/models/open_jev/native/src/main.rs | 23 +++++-- src/models/open_jev/native/src/processing.rs | 8 ++- src/runtime/README.md | 9 +++ src/runtime/src/batching.rs | 32 ++++++++- tests/cuda/equivalent_tiles_profile.cpp | 13 ++++ tests/cuda/test_gemm_profile.py | 18 ++++++ tests/open_jev/health.rs | 37 +++++++++++ tests/open_jev/prefix.rs | 16 ++++- tests/open_jev/processing.rs | 10 +++ tests/open_jev/rollback.rs | 68 ++++++++++++++++++++ tests/runtime/batching.rs | 11 ++++ 17 files changed, 394 insertions(+), 30 deletions(-) create mode 100644 src/backends/cuda/qwen3_5/gemm_profile.h create mode 100644 tests/cuda/equivalent_tiles_profile.cpp create mode 100644 tests/cuda/test_gemm_profile.py create mode 100644 tests/open_jev/health.rs create mode 100644 tests/open_jev/rollback.rs diff --git a/recipe/open_jev/native.md b/recipe/open_jev/native.md index 57c39f47..114ef60e 100644 --- a/recipe/open_jev/native.md +++ b/recipe/open_jev/native.md @@ -128,7 +128,7 @@ and CUDA kernel tests are opt-in; the latter require a GPU reservation: ```sh # Tokenizer cases against either merged export, on CPU: OPEN_JEV_MODEL=weights/open-jev-9b-merged \ - cargo test --locked -p omni-open-jev-native --lib -- --ignored + cargo test --locked -p omni-open-jev-native --lib tokenization_tests -- --ignored # Inside a GPU reservation, after building the library: CUA_S1_CUDA_LIB=$PWD/target/release/libqwen3_5_cuda.so \ cargo test --release --locked -p omni-qwen3-5-native --test kernels -- --ignored @@ -389,7 +389,8 @@ tokenizers retain serial encoding, and any batch error is retried in original order so validation reports the same first error. The exact-token regression includes Unicode, special tokens, batch boundaries and oversized prompts for 9B. -On SM90 with cuBLASLt 13.1.1, the 9B worker can combine 512–4096-row +On the measured device profile (CUDA-reported `NVIDIA L20X`, SM 9.0, +132 multiprocessors, `cublasLtGetVersion() == 130101`), the 9B worker can combine 512–4096-row projections using a tested engine-66 profile (stage 35, no split-K or reduction, custom options 0/1/2/4/5/6/7, clusters 2×1×1, 4×1×1 or 1×2×1, undefined inner shape). It requires the listed configuration; other options retain the original comparison. @@ -408,6 +409,51 @@ candidate projections. The check includes each original batch's total row count. Any mismatch declines the entire sharing attempt before computation and retains the original candidate packing. Other prefix modes and the 27B policy are unchanged. +### 9B packing rollback + +`OPEN_JEV_PACK=0` is a startup escape hatch for Open-Jev-9B. It restores serial +tokenization and request scheduling, then runs each candidate through the eager +`Model::forward` path. It disables candidate packing, cross-request batching, +prefix reuse, packed GDN/projection capabilities and equivalent projection tiles. +It forces CUDA Graph off even if `CUA_S1_GRAPH=1` is inherited. Shared kernel +implementations still come from the loaded CUDA library; this is not a promise +of an identical historical binary or historical performance. + +```sh +OPEN_JEV_PACK=0 OPEN_JEV_PREFIX=0 OPEN_JEV_FIXED=0 OPEN_JEV_EXACT_PREFIX=0 \ +OPEN_JEV_MODEL=weights/open-jev-9b-merged \ +OPEN_JEV_CUDA_LIB=$PWD/target/release/libqwen3_5_cuda.so \ + target/release/omni-open-jev-native +``` + +The default is `OPEN_JEV_PACK=1`. Values other than `0`/`1`, a 27B rollback, +or combining rollback with fixed/shared/auto or enabled exact-prefix reuse fail +at startup. Prefix options may also be unset. `Engine::load` reads this switch; +embedded callers can use `Engine::load_with_packing(..., PrefixMode::Off, false)`. +The existing explicit `load_with_prefix` API keeps packing enabled. + +The frozen original-native oracle also covers the rollback route. Inside a GPU +reservation, with the pinned 9B export (8192-token capacity) and matching library, +run the following for both `CUA_S1_GRAPH=0` and `1`; both must take the eager route: + +```sh +OPEN_JEV_PACK=0 OPEN_JEV_PREFIX=0 OPEN_JEV_FIXED=0 OPEN_JEV_EXACT_PREFIX=0 \ +OPEN_JEV_MODEL=/absolute/path/to/open-jev-9b-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 --lib \ + rollback_uses_serial_candidates_and_matches_original_logits \ + -- --ignored --test-threads=1 +``` + +CPU tests check configuration rejection, tokenization equivalence, device-profile +gating and worker failure health responses. They do not substitute for this GPU +oracle check. An unavailable batch executor or poisoned model mutex makes +`/health` return HTTP 503 with `{"status":"failed","model":"..."}`; ordinary +request errors and a busy executor retain HTTP 200 / `ready`. Restart the worker +after permanent failure. This endpoint does not probe CUDA on each request. + +### Cross-request batching + The 9B default execution path batches already queued requests through the shared runtime, up to 16 candidates and 8192 tokens per GPU dispatch, without an added batching timer. Original candidate batches remain atomic: input/gate-up projection @@ -426,7 +472,7 @@ Their outputs are gathered back into the compact layout. All projection algorith are checked against the original packed input sizes and individual output sizes before execution; unsupported layouts keep the complete forward. Short retained checks use tested 64–511-row profiles, including custom option 3 and cluster 4×2×1, -under the same SM90/cuBLASLt scope. Ordinary packed GEMM grouping is unchanged. +under the same device-profile/cuBLASLt scope. Ordinary packed GEMM grouping is unchanged. This path stores only reusable row-index buffers and recomputes every request. It does not count projection-only reuse as full-prefix `saved_tokens`; reported diff --git a/src/backends/cuda/qwen3_5/gemm.cu b/src/backends/cuda/qwen3_5/gemm.cu index 62a60dae..b06ebd4f 100644 --- a/src/backends/cuda/qwen3_5/gemm.cu +++ b/src/backends/cuda/qwen3_5/gemm.cu @@ -25,6 +25,7 @@ #include #include "ops.h" +#include "gemm_profile.h" namespace { @@ -523,6 +524,7 @@ extern "C" int cs1_gemm_enable_equivalent_tiles(void* gemm) { auto error = cudaGetDevice(&device); if (error == cudaSuccess) error = cudaGetDeviceProperties(&prop, device); if (error != cudaSuccess) return -static_cast(error); - g->equivalent_tiles = prop.major == 9 && prop.minor == 0 && cublasLtGetVersion() == 130101; + g->equivalent_tiles = qwen3_5::equivalent_tiles_device( + prop.name, prop.major, prop.minor, prop.multiProcessorCount, cublasLtGetVersion()); return g->equivalent_tiles ? 1 : 0; } diff --git a/src/backends/cuda/qwen3_5/gemm_profile.h b/src/backends/cuda/qwen3_5/gemm_profile.h new file mode 100644 index 00000000..49f567e5 --- /dev/null +++ b/src/backends/cuda/qwen3_5/gemm_profile.h @@ -0,0 +1,15 @@ +#pragma once + +#include +#include + +namespace qwen3_5 { +// Empirical numerical evidence covers this reported device profile only. +// Matching SM90 alone is insufficient: do not enable on H100/H200 by inference. +// Unknown names, partition sizes and library versions keep exact algo matching. +constexpr bool equivalent_tiles_device(std::string_view name, int major, int minor, + int multiprocessors, std::size_t lt_version) { + return name == "NVIDIA L20X" && major == 9 && minor == 0 && + multiprocessors == 132 && lt_version == 130101; +} +} // namespace qwen3_5 diff --git a/src/models/open_jev/README.md b/src/models/open_jev/README.md index 70b2dae9..d92dc9c3 100644 --- a/src/models/open_jev/README.md +++ b/src/models/open_jev/README.md @@ -25,7 +25,7 @@ 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. On 27B, preparation tokenizes up to 16 candidates +calibration, usage, and metadata. By default, 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 @@ -61,14 +61,22 @@ 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. +Open-Jev-9B packs candidates and batches queued requests by default. Set +`OPEN_JEV_PACK=0` to restore serial tokenization, serial request scheduling and +eager per-candidate forwards, without prefix reuse or packed backend capabilities. +This is an execution-strategy fallback, not a restoration of every historical +kernel implementation. See the [rollback recipe](../../../recipe/open_jev/native.md#9b-packing-rollback) +for configuration constraints and validation commands. Optionally, the executor plans the candidates' shared token prefixes with the shared Qwen planner and runs each prefix once per request through `forward_shared`, or runs every candidate with the same fixed GEMM algorithms as a control; see [optional request-local shared prefixes](../../../recipe/open_jev/native.md#optional-request-local-shared-prefixes). -The engine owns a -[shared serial scheduler](../../runtime/README.md) that admits the complete request -before blocking dispatch. Cross-request batching and shared queue budgets remain planned. +The engine owns a [shared runtime scheduler](../../runtime/README.md): default 9B +uses cross-request batching; rollback, explicit prefix modes and 27B admit complete +requests serially. Shared queue capacity limits remain planned. `/health` reports +HTTP 503 with `status: failed` after a batch worker exits or the model mutex is +poisoned. A busy worker and ordinary request errors do not mark it failed; recovery +from permanent failure requires restarting the process. The experimental 9B layerwise prefix tree can batch the independent suffix @@ -89,7 +97,8 @@ unchanged. Single sequences, reference numerics, bidirectional attention and unvalidated adapter modes keep their previous entry points. Older CUDA libraries without the optional symbol also retain separate launches. -For opted-in SM90 / cuBLASLt 13.1.1 prefix projections, the original algorithm +For opted-in prefix projections on the measured device profile (CUDA-reported +`NVIDIA L20X`, SM 9.0, 132 multiprocessors, `cublasLtGetVersion() == 130101`), the original algorithm must still support the actual row count. Within the already validated compatible arithmetic profiles, the runtime can then select the heuristic for that actual count. Unknown or split-K profiles and other hardware retain the original diff --git a/src/models/open_jev/native/src/engine.rs b/src/models/open_jev/native/src/engine.rs index fc8eec93..9c384424 100644 --- a/src/models/open_jev/native/src/engine.rs +++ b/src/models/open_jev/native/src/engine.rs @@ -76,6 +76,20 @@ pub fn exact_prefix_value(value: Option<&str>, prefix: PrefixMode) -> Result, prefix: PrefixMode, exact: Option) -> Result { + let enabled = match value { + None | Some("1") => true, + Some("0") => false, + Some(_) => bail!("OPEN_JEV_PACK must be 0 or 1"), + }; + ensure!( + enabled || (prefix == PrefixMode::Off && exact != Some(true)), + "OPEN_JEV_PACK=0 requires OPEN_JEV_PREFIX=0, OPEN_JEV_FIXED=0 and OPEN_JEV_EXACT_PREFIX=0 (or unset)" + ); + Ok(enabled) +} + impl Engine { pub async fn load(dir: &Path, library: &Path) -> Result { let prefix = prefix_env()?; @@ -85,7 +99,12 @@ impl Engine { Err(error) => return Err(error).context("OPEN_JEV_EXACT_PREFIX"), }; let exact = exact_prefix_value(raw.as_deref(), prefix)?; - let mut engine = Self::load_with_prefix(dir, library, prefix).await?; + let pack = match std::env::var("OPEN_JEV_PACK") { + Ok(value) => packing_value(Some(&value), prefix, exact)?, + Err(std::env::VarError::NotPresent) => packing_value(None, prefix, exact)?, + Err(error) => return Err(error).context("OPEN_JEV_PACK"), + }; + let mut engine = Self::load_with_packing(dir, library, prefix, pack).await?; if let Some(enabled) = exact { engine.executor.set_exact_prefix(enabled); } @@ -97,14 +116,30 @@ impl Engine { library: &Path, prefix_mode: PrefixMode, ) -> Result { - // `Model::load` reads the switch the same way. - let graph = std::env::var("CUA_S1_GRAPH").as_deref() == Ok("1"); + Self::load_with_packing(dir, library, prefix_mode, true).await + } + + /// Explicit per-engine packing choice; does not read OPEN_JEV_PACK/PREFIX. + /// `pack=false` restores the eager, per-candidate 9B execution path. + pub async fn load_with_packing( + dir: &Path, + library: &Path, + prefix_mode: PrefixMode, + pack: bool, + ) -> Result { + packing_value(Some(if pack { "1" } else { "0" }), prefix_mode, None)?; + // Rollback always uses eager execution, even with an inherited graph switch. + let graph = pack && std::env::var("CUA_S1_GRAPH").as_deref() == Ok("1"); check_graph(prefix_mode, graph)?; let manifest: Value = serde_json::from_slice( &std::fs::read(dir.join("open_jev_export.json")) .context("export the merged checkpoint; see recipe/open_jev/native.md")?, )?; let checkpoint = Checkpoint::from_export(&manifest)?; + ensure!( + pack || checkpoint.alias == "open-jev-9b", + "OPEN_JEV_PACK=0 is supported only for Open-Jev-9B" + ); let temperature = manifest["temperature"].as_f64().context("temperature")?; ensure!( temperature.is_finite() && temperature > 0.0, @@ -124,8 +159,13 @@ impl Engine { .context("chat_suffix")? .to_owned(); let head = DecisionHead::load(dir, &manifest, checkpoint)?; - let processor = Processor::load(dir, checkpoint, prefix, suffix, temperature, max_length)?; - let executor = Executor::load(dir, library, head, checkpoint, prefix_mode, graph).await?; + let mut processor = + Processor::load(dir, checkpoint, prefix, suffix, temperature, max_length)?; + if !pack { + processor.disable_parallel_tokenization(); + } + let executor = + Executor::load(dir, library, head, checkpoint, prefix_mode, graph, pack).await?; let scheduler = executor.scheduler(checkpoint)?; Ok(Self { checkpoint, @@ -134,4 +174,8 @@ impl Engine { executor, }) } + + pub fn is_available(&self) -> bool { + self.scheduler.is_available() && self.executor.is_available() + } } diff --git a/src/models/open_jev/native/src/executor.rs b/src/models/open_jev/native/src/executor.rs index b0c98810..17d28e7d 100644 --- a/src/models/open_jev/native/src/executor.rs +++ b/src/models/open_jev/native/src/executor.rs @@ -114,6 +114,15 @@ enum SchedulerKind { Batched(omni_runtime::batching::BatchScheduler), } +impl Scheduler { + pub fn is_available(&self) -> bool { + match &self.0 { + SchedulerKind::Serial(_) => true, + SchedulerKind::Batched(scheduler) => scheduler.is_available(), + } + } +} + pub struct Executor { model: Arc>, head: DecisionHead, @@ -131,11 +140,13 @@ impl Executor { checkpoint: &Checkpoint, prefix: PrefixMode, graph: bool, + pack: bool, ) -> Result { let (d, lib) = (dir.to_path_buf(), library.to_path_buf()); let sharing = matches!(prefix, PrefixMode::Shared | PrefixMode::Auto); - let batch_gdn = checkpoint.pack_candidates; - let equivalent_tiles = checkpoint.alias == "open-jev-9b"; + let pack = pack && checkpoint.pack_candidates; + let batch_gdn = pack; + let equivalent_tiles = pack && checkpoint.alias == "open-jev-9b"; let model = tokio::task::spawn_blocking(move || { let mut model = Model::load_with_graph(&d, &lib, graph)?; if batch_gdn { @@ -143,7 +154,7 @@ impl Executor { model.enable_packed_projections()?; } if equivalent_tiles && model.enable_equivalent_projection_tiles()? { - eprintln!("Open-Jev-9B: validated SM90 long-prefill projection tiles enabled"); + eprintln!("Open-Jev-9B: validated device profile projection tiles enabled"); } // Shared prefixes reproduce the fixed rows only where rows do not depend on M. if sharing { @@ -157,16 +168,16 @@ impl Executor { Ok(Self { model: Arc::new(Mutex::new(model)), head, - pack: checkpoint.pack_candidates, - share_packed_prefix: checkpoint.share_packed_prefix, - share_packed_inputs: checkpoint.alias == "open-jev-9b", + pack, + share_packed_prefix: pack && checkpoint.share_packed_prefix, + share_packed_inputs: pack && checkpoint.alias == "open-jev-9b", prefix, }) } pub(crate) fn scheduler(&self, checkpoint: &Checkpoint) -> Result { Ok(Scheduler( - if checkpoint.alias == "open-jev-9b" && self.prefix == PrefixMode::Off { + if self.pack && checkpoint.alias == "open-jev-9b" && self.prefix == PrefixMode::Off { SchedulerKind::Batched(batched::new(self.model.clone(), self.head.clone())?) } else { SchedulerKind::Serial(SerialScheduler::default()) @@ -174,6 +185,10 @@ impl Executor { )) } + pub fn is_available(&self) -> bool { + !self.model.is_poisoned() + } + /// Configure exact within-question reuse before the engine admits requests. pub(crate) fn set_exact_prefix(&mut self, enabled: bool) { self.share_packed_prefix = enabled; @@ -309,3 +324,7 @@ impl Executor { #[cfg(test)] #[path = "../../../../../tests/open_jev/head.rs"] mod tests; + +#[cfg(test)] +#[path = "../../../../../tests/open_jev/rollback.rs"] +mod rollback_tests; diff --git a/src/models/open_jev/native/src/main.rs b/src/models/open_jev/native/src/main.rs index c637d29e..2649334c 100644 --- a/src/models/open_jev/native/src/main.rs +++ b/src/models/open_jev/native/src/main.rs @@ -21,6 +21,19 @@ fn error(status: StatusCode, message: impl ToString) -> Response { (status, Json(json!({"error": message.to_string()}))).into_response() } +fn health_response(model: &str, available: bool) -> Response { + let (status, state) = if available { + (StatusCode::OK, "ready") + } else { + (StatusCode::SERVICE_UNAVAILABLE, "failed") + }; + (status, Json(json!({"status": state, "model": model}))).into_response() +} + +async fn health(State(engine): State>) -> Response { + health_response(engine.checkpoint.model_id, engine.is_available()) +} + async fn decide(engine: &Engine, raw: &[u8]) -> Response { let prepared = match engine.processor.prepare(raw) { Ok(prepared) => prepared, @@ -83,12 +96,8 @@ async fn main() -> Result<()> { let port: u16 = std::env::var("OPEN_JEV_PORT") .map_or(Ok(8000), |v| v.parse()) .context("OPEN_JEV_PORT")?; - let model = engine.checkpoint.model_id; let app = Router::new() - .route( - "/health", - get(move || async move { Json(json!({"status": "ready", "model": model})) }), - ) + .route("/health", get(health)) .route("/v1/systemone", post(systemone)) .layer(DefaultBodyLimit::max(4 << 20)) .with_state(engine); @@ -97,3 +106,7 @@ async fn main() -> Result<()> { axum::serve(listener, app).await?; Ok(()) } + +#[cfg(test)] +#[path = "../../../../../tests/open_jev/health.rs"] +mod tests; diff --git a/src/models/open_jev/native/src/processing.rs b/src/models/open_jev/native/src/processing.rs index bf01b082..3923faac 100644 --- a/src/models/open_jev/native/src/processing.rs +++ b/src/models/open_jev/native/src/processing.rs @@ -13,6 +13,7 @@ use omni_qwen3_5_native::prefix::PrefixMode; pub struct Processor { checkpoint: &'static Checkpoint, + parallel_tokenization: bool, tokenizer: Tokenizer, prefix: String, suffix: String, @@ -49,6 +50,7 @@ impl Processor { Tokenizer::from_file(dir.join("tokenizer.json")).map_err(anyhow::Error::msg)?; Ok(Self { checkpoint, + parallel_tokenization: checkpoint.parallel_tokenization, tokenizer, prefix, suffix, @@ -57,6 +59,10 @@ impl Processor { }) } + pub(crate) fn disable_parallel_tokenization(&mut self) { + self.parallel_tokenization = false; + } + /// Validate every candidate length before inference; never truncate. pub fn prepare(&self, raw: &[u8]) -> Result { let questions = contract::compile(raw, self.checkpoint)?; @@ -67,7 +73,7 @@ impl Processor { &self.suffix, self.max_length, &questions, - self.checkpoint.parallel_tokenization, + self.parallel_tokenization, )?; let input_tokens = inputs.iter().flatten().map(Vec::len).sum(); let candidates = inputs.iter().map(Vec::len).sum(); diff --git a/src/runtime/README.md b/src/runtime/README.md index d1a51415..0d860e6a 100644 --- a/src/runtime/README.md +++ b/src/runtime/README.md @@ -101,6 +101,15 @@ and its model resources remain owned by the worker until it finishes. Dropping a scheduler handles closes admission and joins the execution thread, waiting for dispatched work and callback resources to retire before the owner exits. +`BatchScheduler::is_available()` reports worker liveness shared by all clones. +The worker marks itself unavailable before returning an executor-panic error and +also on thread exit. Subsequent requests fail, including empty requests. Ordinary +callback errors fail their requests but leave the worker available. Availability +does not measure queue capacity or probe GPU health. Open-Jev combines this state +with model-mutex poisoning in `/health`: HTTP 200 / `ready` while available and +HTTP 503 / `failed` after permanent failure. It requires a process restart; it +does not rebuild possibly corrupted model state automatically. + Shared processing orchestration and cross-worker coordination remain planned. Batch layout keys include both candidate lengths and original input projection groups, so CUDA Graph replay cannot conflate different numerical layouts. diff --git a/src/runtime/src/batching.rs b/src/runtime/src/batching.rs index bf617c2c..eb5abb8b 100644 --- a/src/runtime/src/batching.rs +++ b/src/runtime/src/batching.rs @@ -2,6 +2,7 @@ use std::collections::VecDeque; use std::panic::{AssertUnwindSafe, catch_unwind}; +use std::sync::atomic::{AtomicBool, Ordering}; use std::sync::{Arc, mpsc}; use anyhow::{Context, Result, anyhow, ensure}; @@ -51,6 +52,16 @@ pub struct BatchScheduler { struct Worker { sender: Option>>, thread: Option>, + available: Arc, +} + +// Also marks unavailable on an unexpected unwind outside the execute callback. +struct WorkerExit(Arc); + +impl Drop for WorkerExit { + fn drop(&mut self) { + self.0.store(false, Ordering::Release); + } } impl Drop for Worker { @@ -87,19 +98,32 @@ impl BatchScheduler { ) -> Result { ensure!(limits.items > 0 && limits.cost > 0, "empty batch budget"); let (sender, receiver) = mpsc::channel(); + let available = Arc::new(AtomicBool::new(true)); + let worker_available = available.clone(); let thread = std::thread::Builder::new() .name("omni-batch".into()) - .spawn(move || dispatch(receiver, limits, cost, execute)) + .spawn(move || { + let _exit = WorkerExit(worker_available.clone()); + dispatch(receiver, limits, cost, execute, &worker_available); + }) .context("start batch execution thread")?; Ok(Self { worker: Arc::new(Worker { sender: Some(sender), thread: Some(thread), + available, }), }) } + /// Whether the worker can still execute requests; shared by all clones. + /// A busy worker remains available. A retired worker requires replacement. + pub fn is_available(&self) -> bool { + self.worker.available.load(Ordering::Acquire) + } + pub async fn run(&self, items: Vec) -> Result> { + ensure!(self.is_available(), "batch executor is unavailable"); if items.is_empty() { return Ok(Vec::new()); } @@ -124,6 +148,7 @@ fn dispatch( limits: BatchLimits, cost: impl Fn(&I) -> BatchCost, mut execute: impl FnMut(&[I]) -> Result>, + available: &AtomicBool, ) { let mut pending = None; loop { @@ -185,6 +210,11 @@ fn dispatch( } let result = catch_unwind(AssertUnwindSafe(|| execute(&inputs))); let panicked = result.is_err(); + if panicked { + // Publish failure before notifying affected callers, so their next + // health probe cannot observe a ready but permanently dead worker. + available.store(false, Ordering::Release); + } let result = result .unwrap_or_else(|_| Err(anyhow!("batch executor panicked"))) .and_then(|outputs| { diff --git a/tests/cuda/equivalent_tiles_profile.cpp b/tests/cuda/equivalent_tiles_profile.cpp new file mode 100644 index 00000000..690d6a6c --- /dev/null +++ b/tests/cuda/equivalent_tiles_profile.cpp @@ -0,0 +1,13 @@ +#include "src/backends/cuda/qwen3_5/gemm_profile.h" + +using qwen3_5::equivalent_tiles_device; +static_assert(equivalent_tiles_device("NVIDIA L20X", 9, 0, 132, 130101)); +static_assert(!equivalent_tiles_device("NVIDIA H100 80GB HBM3", 9, 0, 132, 130101)); +static_assert(!equivalent_tiles_device("NVIDIA H200", 9, 0, 132, 130101)); +static_assert(!equivalent_tiles_device("", 9, 0, 132, 130101)); +static_assert(!equivalent_tiles_device("NVIDIA L20X", 8, 9, 132, 130101)); +static_assert(!equivalent_tiles_device("NVIDIA L20X", 9, 1, 132, 130101)); +static_assert(!equivalent_tiles_device("NVIDIA L20X", 9, 0, 66, 130101)); +static_assert(!equivalent_tiles_device("NVIDIA L20X", 9, 0, 132, 130100)); +static_assert(!equivalent_tiles_device("NVIDIA L20X", 9, 0, 132, 130102)); +int main() { return 0; } diff --git a/tests/cuda/test_gemm_profile.py b/tests/cuda/test_gemm_profile.py new file mode 100644 index 00000000..01296673 --- /dev/null +++ b/tests/cuda/test_gemm_profile.py @@ -0,0 +1,18 @@ +"""Compile the production device gate on CPU; no CUDA toolkit or GPU required.""" +from pathlib import Path +import subprocess +import tempfile +import unittest + +ROOT = Path(__file__).resolve().parents[2] + + +class GemmDeviceProfileTest(unittest.TestCase): + def test_only_measured_profile_is_allowed(self): + with tempfile.TemporaryDirectory(prefix="gemm-profile-") as tmp: + binary = Path(tmp) / "profile" + subprocess.run([ + "c++", "-std=c++17", "-Wall", "-Wextra", "-Werror", "-I", str(ROOT), + str(ROOT / "tests/cuda/equivalent_tiles_profile.cpp"), "-o", str(binary), + ], check=True, capture_output=True, text=True) + subprocess.run([str(binary)], check=True) diff --git a/tests/open_jev/health.rs b/tests/open_jev/health.rs new file mode 100644 index 00000000..87fc4a77 --- /dev/null +++ b/tests/open_jev/health.rs @@ -0,0 +1,37 @@ +//! HTTP health must follow the worker's actual lifecycle, without needing a GPU. +use super::health_response; +use axum::{body::to_bytes, http::StatusCode}; +use omni_runtime::batching::{BatchCost, BatchLimits, BatchScheduler}; +use serde_json::Value; + +#[tokio::test] +async fn health_fails_after_panic_but_not_after_an_ordinary_error() { + let scheduler: BatchScheduler = BatchScheduler::new( + BatchLimits { items: 1, cost: 1 }, + |_| BatchCost::from(1), + |items: &[usize]| { + if items == [1] { + anyhow::bail!("recoverable request failure"); + } + if items == [2] { + panic!("executor state lost"); + } + Ok(items.to_vec()) + }, + ) + .unwrap(); + assert!(scheduler.run(vec![1]).await.is_err()); + assert_eq!( + health_response("open-jev-9b", scheduler.is_available()).status(), + StatusCode::OK + ); + assert_eq!(scheduler.run(vec![3]).await.unwrap(), [3]); + assert!(scheduler.run(vec![2]).await.is_err()); + let response = health_response("open-jev-9b", scheduler.is_available()); + assert_eq!(response.status(), StatusCode::SERVICE_UNAVAILABLE); + let body = to_bytes(response.into_body(), 1024).await.unwrap(); + let json: Value = serde_json::from_slice(&body).unwrap(); + assert_eq!(json["status"], "failed"); + assert_eq!(json["model"], "open-jev-9b"); + assert!(scheduler.run(vec![3]).await.is_err()); +} diff --git a/tests/open_jev/prefix.rs b/tests/open_jev/prefix.rs index ee44b405..44038d5e 100644 --- a/tests/open_jev/prefix.rs +++ b/tests/open_jev/prefix.rs @@ -3,12 +3,26 @@ use std::path::PathBuf; use omni_open_jev_native::{ - engine::{Engine, check_graph, exact_prefix_value, prefix_values}, + engine::{Engine, check_graph, exact_prefix_value, packing_value, prefix_values}, executor::PrefixUse, }; use omni_qwen3_5_native::{model::Model, prefix::PrefixMode}; use serde_json::json; +#[test] +fn packing_rollback_rejects_conflicting_modes_and_bad_values() { + assert!(packing_value(None, PrefixMode::Off, None).unwrap()); + assert!(packing_value(Some("1"), PrefixMode::Off, Some(true)).unwrap()); + for exact in [None, Some(false)] { + assert!(!packing_value(Some("0"), PrefixMode::Off, exact).unwrap()); + } + for mode in [PrefixMode::Fixed, PrefixMode::Shared, PrefixMode::Auto] { + assert!(packing_value(Some("0"), mode, None).is_err()); + } + assert!(packing_value(Some("0"), PrefixMode::Off, Some(true)).is_err()); + assert!(packing_value(Some("auto"), PrefixMode::Off, None).is_err()); +} + #[test] fn switches_follow_decider() { assert_eq!(prefix_values(None, None).unwrap(), PrefixMode::Off); diff --git a/tests/open_jev/processing.rs b/tests/open_jev/processing.rs index daf917ae..23b28f93 100644 --- a/tests/open_jev/processing.rs +++ b/tests/open_jev/processing.rs @@ -56,6 +56,7 @@ fn processor(max_length: usize) -> Processor { )); Processor { checkpoint: &CHECKPOINTS[0], + parallel_tokenization: true, tokenizer, prefix: "".into(), suffix: "".into(), @@ -68,6 +69,15 @@ const REQUEST: &[u8] = br#"{"state":"x","questions":{ "z":{"type":"choice","instructions":"Route","criteria":{"technical":null,"billing":null}}, "a":{"type":"noul","instructions":"Verify"}}}"#; +#[test] +fn serial_preparation_preserves_tokens_and_question_mapping() { + let mut processor = processor(4096); + let expected = processor.prepare(REQUEST).unwrap().inputs; + processor.disable_parallel_tokenization(); + assert!(!processor.parallel_tokenization); + assert_eq!(processor.prepare(REQUEST).unwrap().inputs, expected); +} + #[test] fn preparation_applies_the_saved_chat_prefix_and_suffix() { let mut processor = processor(4096); diff --git a/tests/open_jev/rollback.rs b/tests/open_jev/rollback.rs new file mode 100644 index 00000000..f9c9097e --- /dev/null +++ b/tests/open_jev/rollback.rs @@ -0,0 +1,68 @@ +//! Run with OPEN_JEV_PACK=0; this checks actual routing and the frozen original oracle. +use super::SchedulerKind; +use crate::engine::Engine; +use serde_json::Value; +use std::path::PathBuf; + +#[tokio::test(flavor = "current_thread")] +#[ignore = "needs pinned 9B, OPEN_JEV_PACK=0, matching CUDA library and a reserved GPU"] +async fn rollback_uses_serial_candidates_and_matches_original_logits() { + assert_eq!(std::env::var("OPEN_JEV_PACK").as_deref(), Ok("0")); + 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(); + assert!(matches!(engine.scheduler.0, SchedulerKind::Serial(_))); + assert!(!engine.executor.pack); + assert!(!engine.executor.share_packed_prefix); + assert!(!engine.executor.share_packed_inputs); + assert!(!engine.executor.model.lock().unwrap().graph_stats().enabled); + assert!(engine.is_available()); + let oracle: Value = serde_json::from_str(include_str!("data/packing_original.json")).unwrap(); + for (case, expected) in oracle["manifest"]["cases"] + .as_array() + .unwrap() + .iter() + .zip(oracle["records"].as_array().unwrap()) + { + let candidates: Vec> = case["candidates"] + .as_array() + .unwrap() + .iter() + .map(|c| { + let length = c[0].as_u64().unwrap() as usize; + let seed = c[1].as_u64().unwrap() as u32; + (0..length) + .map(|i| 1000 + (seed * 7919 + i as u32 * 37) % 30000) + .collect() + }) + .collect(); + let mut candidates = candidates.into_iter(); + let ids = case["questions"] + .as_array() + .unwrap() + .iter() + .map(|n| { + candidates + .by_ref() + .take(n.as_u64().unwrap() as usize) + .collect() + }) + .collect(); + let (scores, used) = engine + .executor + .execute_with_prefix(&engine.scheduler, ids) + .await + .unwrap(); + assert!(!used.shared); + let bits: Vec> = scores + .iter() + .map(|r| r.iter().map(|x| x.to_bits()).collect()) + .collect(); + assert_eq!( + serde_json::json!(bits), + expected["score_bits"], + "{}", + case["id"] + ); + } +} diff --git a/tests/runtime/batching.rs b/tests/runtime/batching.rs index b146065f..5f42eefd 100644 --- a/tests/runtime/batching.rs +++ b/tests/runtime/batching.rs @@ -151,6 +151,10 @@ async fn errors_discard_partial_results_and_leave_later_requests_usable() { ) .unwrap(); assert!(done(scheduler.run(vec![1, 2, 3])).await.is_err()); + assert!( + scheduler.is_available(), + "ordinary errors do not retire the worker" + ); assert_eq!(done(scheduler.run(vec![4])).await.unwrap(), [4]); assert_eq!(*calls.lock().unwrap(), [1, 2, 4]); } @@ -164,6 +168,8 @@ async fn a_panicked_executor_closes_admission_instead_of_reusing_bad_state() { |_| -> Result> { panic!("device executor failed") }, ) .unwrap(); + let observer = scheduler.clone(); + assert!(observer.is_available()); assert!( done(scheduler.run(vec![1])) .await @@ -171,6 +177,11 @@ async fn a_panicked_executor_closes_admission_instead_of_reusing_bad_state() { .to_string() .contains("panicked") ); + assert!( + !observer.is_available(), + "failure is visible before the caller resumes" + ); + assert!(done(observer.run(Vec::new())).await.is_err()); assert!(done(scheduler.run(vec![2])).await.is_err()); }