From 043f82c5f3a3da1c06829c28436a58b1fbaa36a2 Mon Sep 17 00:00:00 2001 From: levius <2114377220@qq.com> Date: Sun, 1 Nov 2026 09:43:02 +0800 Subject: [PATCH] cuda: add eager full-row prefix-bidirectional Qwen execution --- docs/architecture.md | 7 + src/backends/cuda/qwen3_5/README.md | 30 ++++ src/backends/cuda/qwen3_5/attention.cu | 14 ++ src/backends/cuda/qwen3_5/ops.h | 8 + src/models/qwen3_5/native/Cargo.toml | 4 + src/models/qwen3_5/native/src/cuda.rs | 28 +++ src/models/qwen3_5/native/src/model.rs | 82 ++++++++- tests/qwen3_5/README.md | 50 ++++++ tests/qwen3_5/graph_loading.rs | 27 +++ tests/qwen3_5/kernels.rs | 222 ++++++++++++++++++++++++ tests/qwen3_5/prefix_bidir.rs | 70 ++++++++ tests/qwen3_5/prefix_bidir_reference.py | 89 ++++++++++ 12 files changed, 630 insertions(+), 1 deletion(-) create mode 100644 tests/qwen3_5/prefix_bidir.rs create mode 100644 tests/qwen3_5/prefix_bidir_reference.py diff --git a/docs/architecture.md b/docs/architecture.md index 276143e3..4dc650d5 100644 --- a/docs/architecture.md +++ b/docs/architecture.md @@ -93,6 +93,13 @@ process runs the checkpoint's joint schema head on its main thread, which is the thread that initialized CUDA. The joint head is still that Python module. Clef does not use the shared Qwen CUDA executor. +The shared Qwen executor also exposes a single-row eager +`forward_prompt_prefix_bidirectional` prerequisite: state queries read state +keys, question queries read their complete row, and GDN remains causal. The +model supplies an exact split, with no chunk rounding. This is not state-once +continuation or a D3 native worker; it rejects Graph, LoRA and reference +numerics. See the [backend capability](../src/backends/cuda/qwen3_5/README.md#full-row-prefix-bidirectional-baseline). + The Qwen executor's `forward_shared` takes a request prefix, group prefixes and branches, ends each prefix at its last multiple of 64 tokens, and returns one final hidden state per branch in order diff --git a/src/backends/cuda/qwen3_5/README.md b/src/backends/cuda/qwen3_5/README.md index 090e54fb..79704e27 100644 --- a/src/backends/cuda/qwen3_5/README.md +++ b/src/backends/cuda/qwen3_5/README.md @@ -26,6 +26,36 @@ missing. Rebuild this library before starting that worker. This tree does not claim that the causal specialization's SASS is bit-identical to the previous build; GPU revalidation of Cua-S1 and Open-Jev was not run with this change. +## Full-row prefix-bidirectional baseline + +`Model::forward_prompt_prefix_bidirectional(ids, state_tokens)` executes one +complete, unpadded text row with `0 < state_tokens < ids.len()`. The state +queries attend only to the complete state; the question queries attend to the +state and their own complete question. GDN remains causal. The caller supplies +the exact semantic split: no 64-token rounding and no prefix continuation. + +The optional additive ABI-9 symbol +`cs1_attention_gated_prefix_bidirectional` composes existing bidirectional +launches on the caller's stream: a complete-row pass followed by a state-only +pass that overwrites the state outputs before the next layer. It does extra +state work and makes no acceleration claim. Existing attention templates and +causal entry points are unchanged. Inputs and output must be disjoint, with +the complete BF16 row/stride geometry in `ops.h`. + +The Rust entry is eager and merged; it rejects Graph, LoRA and reference +numerics. Call `require_prompt_prefix_bidirectional` during adapter load before +readiness. An old ABI-9 library still loads for existing callers and returns +an actionable rebuild error when the new capability is requested. Rebuild the +CUDA library to opt in. This primitive does not add a D3 worker, FP32 answer +head, media path, state-once reuse or batching. Full-model reference parity +requires a separate model-specific acceptance campaign. + +[Root regression instructions](../../../../tests/qwen3_5/README.md#prefix-bidirectional-full-row-baseline) +cover the independent dense FP64 mask, two-layer state isolation and a tiny +Transformers hybrid model with two full-attention layers and causal GDN. + +## Existing numerical profiles + HF Qwen3.5 folders store `linear_attn.A_log` and `linear_attn.norm.weight` as FP32. The shared loader rounds those two suffixes to BF16 at upload because the kernels consume BF16. Other non-BF16 language tensors still fail. That rounding diff --git a/src/backends/cuda/qwen3_5/attention.cu b/src/backends/cuda/qwen3_5/attention.cu index 5193e05f..90c0d194 100644 --- a/src/backends/cuda/qwen3_5/attention.cu +++ b/src/backends/cuda/qwen3_5/attention.cu @@ -337,6 +337,20 @@ extern "C" int cs1_attention_gated_bidirectional(const void* q, const void* k, c return flash::launch(q, k, v, ldv, gate, out, T, Hq, Hk, Dh, scale, 0, stream); } +// A correctness-first composition of the existing kernel, without modifying its +// causal or whole-prompt specializations. Question outputs from the first pass +// already read exactly their own row; overwrite the state outputs before any +// downstream layer can observe them. The split is semantic, not tile-aligned. +extern "C" int cs1_attention_gated_prefix_bidirectional(const void* q, const void* k, const void* v, int ldv, + const void* gate, void* out, int T, int state, + int Hq, int Hk, int Dh, float scale, void* stream) { + if (state <= 0 || state >= T || q == nullptr || k == nullptr || v == nullptr + || gate == nullptr || out == nullptr) return cudaErrorInvalidValue; + const int status = flash::launch(q, k, v, ldv, gate, out, T, Hq, Hk, Dh, scale, 0, stream); + if (status != cudaSuccess) return status; + return flash::launch(q, k, v, ldv, gate, out, state, Hq, Hk, Dh, scale, 0, stream); +} + // Same layout contract as the full pass; only the covered query rows differ: // q_base is a multiple of 64 and rows below it keep whatever was already in k/v. extern "C" int cs1_attention_gated_prefix(const void* q, const void* k, const void* v, int ldv, const void* gate, diff --git a/src/backends/cuda/qwen3_5/ops.h b/src/backends/cuda/qwen3_5/ops.h index 9ad9c739..2303c466 100644 --- a/src/backends/cuda/qwen3_5/ops.h +++ b/src/backends/cuda/qwen3_5/ops.h @@ -143,6 +143,14 @@ int cs1_attention_gated_bidirectional(const void* q, const void* k, const void* const void* gate, void* out, int T, int Hq, int Hk, int Dh, float scale, void* stream); +// Optional additive ABI 9 symbol. One unpadded state+question row, 0 < state < T. +// State queries see only state keys; question queries see all keys in this row. +// Eager correctness baseline: full bidirectional pass, then state-only overwrite. +// Buffers must be disjoint inputs/output; launches use the caller's stream. +int cs1_attention_gated_prefix_bidirectional(const void* q, const void* k, const void* v, int ldv, + const void* gate, void* out, int T, int state, + int Hq, int Hk, int Dh, float scale, void* stream); + // Windowed variant for the cached prefix: k and v cover rows [0, T) (the cached // prefix must already be in place); q, gate and out cover rows [q_base, T) only. // q_base must be a multiple of 64; q_base = 0 reduces to cs1_attention_gated. diff --git a/src/models/qwen3_5/native/Cargo.toml b/src/models/qwen3_5/native/Cargo.toml index 7db6a25f..62d38c05 100644 --- a/src/models/qwen3_5/native/Cargo.toml +++ b/src/models/qwen3_5/native/Cargo.toml @@ -47,3 +47,7 @@ path = "../../../../tests/qwen3_5/graph_loading.rs" [[test]] name = "device_features" path = "../../../../tests/qwen3_5/device_features.rs" + +[[test]] +name = "prefix_bidir" +path = "../../../../tests/qwen3_5/prefix_bidir.rs" diff --git a/src/models/qwen3_5/native/src/cuda.rs b/src/models/qwen3_5/native/src/cuda.rs index 612d49da..65fda61d 100644 --- a/src/models/qwen3_5/native/src/cuda.rs +++ b/src/models/qwen3_5/native/src/cuda.rs @@ -419,7 +419,35 @@ pub type GatedAttention = unsafe extern "C" fn( Stream, ) -> c_int; +/// Full-row prefix-bidirectional attention; `state` follows the row length. +pub type PrefixGatedAttention = unsafe extern "C" fn( + *const c_void, + *const c_void, + *const c_void, + c_int, + *const c_void, + *mut c_void, + c_int, + c_int, + c_int, + c_int, + c_int, + f32, + Stream, +) -> c_int; + impl Api { + /// Optional ABI-9 capability; older libraries continue serving causal users. + pub fn attention_gated_prefix_bidirectional(&self) -> Result { + // SAFETY: ops.h declares this signature; Api retains the loaded library. + unsafe { + self._lib + .get::(b"cs1_attention_gated_prefix_bidirectional\0") + .map(|symbol| *symbol) + .context("libqwen3_5_cuda.so has no cs1_attention_gated_prefix_bidirectional; rebuild src/backends/cuda/qwen3_5 (optional ABI 9 capability)") + } + } + pub fn matched_create(&self) -> Result *mut c_void> { unsafe { Ok(*self._lib.get(b"cs1_gemm_create_matched\0")?) } } diff --git a/src/models/qwen3_5/native/src/model.rs b/src/models/qwen3_5/native/src/model.rs index 05f8c6cb..1051e781 100644 --- a/src/models/qwen3_5/native/src/model.rs +++ b/src/models/qwen3_5/native/src/model.rs @@ -1389,6 +1389,44 @@ impl Model { Ok(()) } + /// One complete unpadded state+question row with a semantic, unrounded split. + /// Full-attention state rows see only state; question rows see this whole row. + /// GDN stays causal. No prefix continuation, Graph, LoRA or reference profile. + pub fn forward_prompt_prefix_bidirectional( + &mut self, + ids: &[u32], + state_tokens: usize, + ) -> Result> { + ensure!( + !ids.is_empty() && ids.len() <= self.cfg.max_positions, + "empty or oversized prompt" + ); + ensure!( + state_tokens > 0 && state_tokens < ids.len(), + "state boundary must be inside the prompt" + ); + ensure!( + ids.iter().all(|&id| (id as usize) < self.embed.shape[0]), + "token id outside the vocabulary" + ); + self.require_prompt_prefix_bidirectional()?; + self.prepare_scratch(ids.len())?; + let s = self.scratch.as_ref().unwrap(); + self.embed_tokens(s, ids)?; + self.run_with_attention(s, &[ids.len()], false, true, None, Some(state_tokens))?; + self.last_hidden(s, ids.len()) + } + + /// Check this optional capability during a model adapter's load/warmup. + pub fn require_prompt_prefix_bidirectional(&self) -> Result<()> { + ensure!( + !self.graph_enabled && self.reference.is_none() && self.lora.is_none(), + "prefix-bidirectional prefill is eager, merged and uses the default kernels" + ); + cuda::api().attention_gated_prefix_bidirectional()?; + Ok(()) + } + /// 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>> { @@ -2762,7 +2800,33 @@ impl Model { bidirectional: bool, suffix: Option<&SuffixBatch<'_>>, ) -> Result<()> { - let bidir = if bidirectional { + self.run_with_attention(s, lengths, custom_positions, bidirectional, suffix, None) + } + + fn run_with_attention( + &self, + s: &Scratch, + lengths: &[usize], + custom_positions: bool, + bidirectional: bool, + suffix: Option<&SuffixBatch<'_>>, + state_tokens: Option, + ) -> Result<()> { + let prefix_bidir = if let Some(state) = state_tokens { + self.require_prompt_prefix_bidirectional()?; + ensure!( + bidirectional && !custom_positions && suffix.is_none() && lengths.len() == 1, + "prefix-bidirectional prefill requires one complete text row" + ); + ensure!( + state > 0 && state < lengths[0], + "state boundary must be inside the prompt" + ); + Some(cuda::api().attention_gated_prefix_bidirectional()?) + } else { + None + }; + let bidir = if bidirectional && state_tokens.is_none() { ensure!( self.reference.is_none() && self.lora.is_none() && suffix.is_none(), "prompt-bidirectional prefill does not use reference numerics or an adapter" @@ -3060,6 +3124,22 @@ impl Model { (cfg.head_dim as f32).powf(-0.5), st, ) + } else if let Some(prefix_bidir) = prefix_bidir { + (prefix_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, + state_tokens.unwrap() as i32, + hq, + hk, + hd, + scale, + st, + ) } else if let Some(bidir) = bidir { (bidir)( p(s.aq + offset * cfg.heads * cfg.head_dim * BF16), diff --git a/tests/qwen3_5/README.md b/tests/qwen3_5/README.md index c8a8376f..cb004468 100644 --- a/tests/qwen3_5/README.md +++ b/tests/qwen3_5/README.md @@ -45,3 +45,53 @@ Coverage: These are executor regressions, not a screenshot-encoder or agent-accuracy benchmark. No latency, capture-cost or memory claim is made by these tests. + +## Prefix-bidirectional full-row baseline + +The optional `Model::forward_prompt_prefix_bidirectional` entry consumes one +complete state+question row. State rows read only state keys; question rows +read that state and their own question. The semantic split must be inside +the row and is not rounded. GDN remains causal; Graph, LoRA, reference numerics, +media, prefix continuation, branch batching and an HTTP worker are outside this +entry's supported scope. It is a correctness baseline, not a speed optimization. + +Build the backend for the actual reserved device (`80` for the A800 used in +qualification), then run the kernel checks: + +```sh +src/backends/cuda/qwen3_5/build.sh /absolute/backend-output 80 +export CUA_S1_CUDA_LIB=/absolute/backend-output/libqwen3_5_cuda.so +cargo test --release --locked -p omni-qwen3-5-native --test kernels \ + prefix_bidirectional -- --ignored --test-threads=1 --nocapture +``` + +The dense FP64 oracle covers unaligned state lengths, GQA, strided V and +nonzero gates. The fixed gate is the existing attention test's 0.016 maximum +row-normalized absolute error, including the existing BF16 gating policy. +A two-layer test ensures changed question contents do not affect state bits, +future state tokens influence earlier state outputs, and repeated rows agree. +Invalid dimensions use valid allocations; missing buffers are checked separately. + +For a small full-layer-loop oracle, use Python with `torch==2.12.1+cu130`, +`transformers==5.17.0` and their dependencies, on the same GPU. The generator +initializes a deterministic, all-BF16, three-layer Qwen model with full/linear/ +full attention and uses an independent dense boolean mask. It saves roughly +a few MB of generated weights and reference arrays outside the checkout; no +pretrained model download is needed. The native replay fixes the normalized +hidden-state maximum error gate at 0.03, repeats an earlier row and checks +that a subsequent causal call is bitwise unchanged. + +```sh +python tests/qwen3_5/prefix_bidir_reference.py /absolute/new-oracle-directory +export QWEN_PREFIX_BIDIR_FIXTURE=/absolute/new-oracle-directory/reference.json +export QWEN_PREFIX_BIDIR_OUTPUT=/absolute/native-replay.json +cargo test --release --locked -p omni-qwen3-5-native --test prefix_bidir \ + -- --ignored --test-threads=1 --nocapture +``` + +The output variable is optional and preserves raw hidden states when set. +The synthetic numerical gate is not a D3 decision/probability gate. An external +full-checkpoint replay needs its own frozen inputs, artifact identity and model +acceptance. Standard workspace checks compile these tests but leave GPU cases +ignored; `graph_loading` checks missing capability, unsupported graph/adapter +profiles and input rejection without a GPU. diff --git a/tests/qwen3_5/graph_loading.rs b/tests/qwen3_5/graph_loading.rs index 4b245d88..4fc2bb29 100644 --- a/tests/qwen3_5/graph_loading.rs +++ b/tests/qwen3_5/graph_loading.rs @@ -160,6 +160,26 @@ fn explicit_graph_mode_and_adapter_models_retire_the_stream_once_after_allocatio } for graph in [false, true] { let mut model = Model::load_with_graph(&directory, &library, graph).unwrap(); + // Older ABI-9 libraries must still load for causal callers. The new + // capability fails before scratch allocation, with an actionable error. + let error = model.require_prompt_prefix_bidirectional().unwrap_err(); + let detail = format!("{error:#}"); + assert!( + detail.contains(if graph { "eager" } else { "rebuild" }), + "{detail}" + ); + for (ids, state) in [ + (&[][..], 0), + (&[0, 1][..], 0), + (&[0, 1][..], 2), + (&[0, 4][..], 1), + ] { + assert!( + model + .forward_prompt_prefix_bidirectional(ids, state) + .is_err() + ); + } assert_eq!(model.graph_stats().requested, graph); assert_eq!(model.graph_stats().enabled, graph); // The fake has no arithmetic kernels. The fixed handle and scratch are @@ -300,6 +320,13 @@ fn explicit_graph_mode_and_adapter_models_retire_the_stream_once_after_allocatio ); } let mut model = Model::load_with_lora(&directory, &library, &adapter, 2.0).unwrap(); + assert!( + model + .require_prompt_prefix_bidirectional() + .unwrap_err() + .to_string() + .contains("merged") + ); assert_eq!( model.graph_stats().requested, std::env::var("CUA_S1_GRAPH").as_deref() == Ok("1") diff --git a/tests/qwen3_5/kernels.rs b/tests/qwen3_5/kernels.rs index 5173a122..e346023f 100644 --- a/tests/qwen3_5/kernels.rs +++ b/tests/qwen3_5/kernels.rs @@ -402,6 +402,228 @@ fn bidirectional_attention_reads_future_keys_after_shared_suffix_refactor() { check_flash_attention_reference(true); } +#[test] +#[ignore = "needs a GPU and CUA_S1_CUDA_LIB with prefix-bidirectional attention"] +fn prefix_bidirectional_attention_matches_dense_block_mask() { + let st = setup(); + let op = api().attention_gated_prefix_bidirectional().unwrap(); + for (t, state, hq, hk) in [ + (2, 1, 1, 1), + (65, 31, 4, 1), + (129, 63, 6, 2), + (137, 65, 24, 4), + (257, 129, 4, 2), + ] { + let dh = 256; + let ldv = hk * dh + 8; // strided, not a separately packed V matrix + let qh = random(t * hq * dh, 71, 0.5); + let kh = random(t * hk * dh, 72, 0.5); + let vh = random(t * ldv, 73, 1.0); + let gh = random(t * hq * dh, 74, 1.0); + let q = to_device(&qh, st); + let k = to_device(&kh, st); + let v = to_device(&vh, st); + let gate = to_device(&gh, st); + let out = DeviceBuffer::new(t * hq * dh * 2).unwrap(); + // SAFETY: all buffers contain complete unpadded rows, state is strictly + // inside the row, and the stream owns the complete sequence of launches. + unsafe { + check( + op( + q.at(0), + k.at(0), + v.at(0), + ldv as i32, + gate.at(0), + out.at(0), + t as i32, + state as i32, + hq as i32, + hk as i32, + dh as i32, + 0.0625, + st, + ), + "prefix bidirectional", + ) + .unwrap(); + } + let got = from_device(&out, t * hq * dh, st); + let mut worst = 0.0f64; + for row in 0..t { + let keys = if row < state { state } else { t }; + for head in 0..hq { + let group = head / (hq / hk); + let scores: Vec = (0..keys) + .map(|key| { + (0..dh) + .map(|d| { + qh[(row * hq + head) * dh + d].to_f64() + * kh[(key * hk + group) * dh + d].to_f64() + }) + .sum::() + * 0.0625 + }) + .collect(); + let max = scores.iter().copied().fold(f64::NEG_INFINITY, f64::max); + let weights: Vec = scores.iter().map(|x| (x - max).exp()).collect(); + let sum: f64 = weights.iter().sum(); + let want: Vec = (0..dh) + .map(|d| { + let a = (0..keys) + .map(|key| weights[key] * vh[key * ldv + group * dh + d].to_f64()) + .sum::() + / sum; + let gate = gh[(row * hq + head) * dh + d].to_f32(); + // Existing gated attention rounds attention and sigmoid to + // BF16 independently before the final BF16 multiplication. + bf16::from_f32( + bf16::from_f64(a).to_f32() + * bf16::from_f32(1.0 / (1.0 + (-gate).exp())).to_f32(), + ) + .to_f64() + }) + .collect(); + let scale = want.iter().map(|x| x.abs()).fold(1e-3f64, f64::max); + for (d, expected) in want.iter().enumerate() { + let actual = got[(row * hq + head) * dh + d] as f64; + assert!(actual.is_finite()); + worst = worst.max((actual - expected).abs() / scale); + } + } + } + eprintln!("prefix-bidir T={t}, state={state}, heads={hq}/{hk}: {worst:.6e}"); + assert!(worst < 1.6e-2, "dense block mask mismatch: {worst}"); + } + // Invalid geometry and semantic boundaries are rejected before dereference. + let valid = to_device(&vec![bf16::ZERO; 65 * 6 * 256], st); + let output = DeviceBuffer::new(65 * 6 * 256 * 2).unwrap(); + for (t, state, hq, hk, dh, ldv) in [ + (0, 0, 1, 1, 256, 256), + (65, 0, 1, 1, 256, 256), + (65, 65, 1, 1, 256, 256), + (65, -1, 1, 1, 256, 256), + (65, 66, 1, 1, 256, 256), + (65, 31, 3, 2, 256, 512), + (65, 31, 1, 1, 128, 256), + (65, 31, 1, 1, 256, 255), + (65, 31, 4, 2, 256, 256), + (65, 31, 0, 1, 256, 256), + (65, 31, 1, 0, 256, 256), + ] { + // SAFETY: every tuple is invalid and must return before accessing pointers. + let code = unsafe { + op( + valid.at(0), + valid.at(0), + valid.at(0), + ldv, + valid.at(0), + output.at(0), + t, + state, + hq, + hk, + dh, + 0.0625, + st, + ) + }; + assert_ne!(code, 0, "accepted invalid prefix-bidir geometry"); + } + for missing in 0..5 { + let mut pointers = [valid.at(0); 5]; + pointers[4] = output.at(0); + pointers[missing] = std::ptr::null_mut(); + // SAFETY: otherwise valid geometry must reject each missing buffer. + let code = unsafe { + op( + pointers[0], + pointers[1], + pointers[2], + 256, + pointers[3], + pointers[4], + 65, + 31, + 1, + 1, + 256, + 0.0625, + st, + ) + }; + assert_ne!(code, 0, "accepted null buffer {missing}"); + } + cuda::synchronize(st).unwrap(); +} + +#[test] +#[ignore = "needs a GPU and CUA_S1_CUDA_LIB with prefix-bidirectional attention"] +fn prefix_bidirectional_two_layers_isolate_state_from_question() { + let st = setup(); + let op = api().attention_gated_prefix_bidirectional().unwrap(); + for (t, state) in [(3, 1), (65, 63), (129, 65)] { + let n = t * 256; + let zero = to_device(&vec![bf16::ZERO; n], st); + let original: Vec = (0..n) + .map(|i| bf16::from_f32((i / 256) as f32 / 16.0)) + .collect(); + let run = |values: &[bf16]| { + let v = to_device(values, st); + let out = DeviceBuffer::new(n * 2).unwrap(); + let second = DeviceBuffer::new(n * 2).unwrap(); + // SAFETY: two sequential layers share only immutable zero Q/K/gates; + // each V input and output is a separate, correctly sized allocation. + unsafe { + for (input, output) in [(v.at(0), out.at(0)), (out.at(0), second.at(0))] { + check( + op( + zero.at(0), + zero.at(0), + input, + 256, + zero.at(0), + output, + t as i32, + state as i32, + 1, + 1, + 256, + 0.0625, + st, + ), + "two-layer prefix mask", + ) + .unwrap(); + } + } + from_device(&second, n, st) + }; + let ordinary = run(&original); + let mut question = original.clone(); + question[state * 256..].fill(bf16::from_f32(32.0)); + let changed_question = run(&question); + assert_eq!( + ordinary[..state * 256], + changed_question[..state * 256], + "question leaked into state" + ); + assert_ne!(ordinary[state * 256..], changed_question[state * 256..]); + let mut future_state = original.clone(); + future_state[(state - 1) * 256..state * 256].fill(bf16::from_f32(16.0)); + let changed_state = run(&future_state); + assert_ne!( + ordinary[..256], + changed_state[..256], + "state incorrectly used a causal mask" + ); + // A subsequent differently sized row must not consume an earlier row. + assert_eq!(ordinary, run(&original)); + } + cuda::synchronize(st).unwrap(); +} + fn check_flash_attention_reference(bidirectional: bool) { let st = setup(); let (hq, hk, dh) = (16usize, 4usize, 256usize); diff --git a/tests/qwen3_5/prefix_bidir.rs b/tests/qwen3_5/prefix_bidir.rs new file mode 100644 index 00000000..fa5c4aa8 --- /dev/null +++ b/tests/qwen3_5/prefix_bidir.rs @@ -0,0 +1,70 @@ +//! Frozen full-row reference replay, including two softmax layers and causal GDN. +use omni_qwen3_5_native::model::Model; +use serde::Deserialize; +use std::path::PathBuf; + +#[derive(Deserialize)] +struct Fixture { + schema: String, + rows: Vec, +} + +#[derive(Deserialize)] +struct Row { + id: String, + ids: Vec, + state_tokens: usize, + hidden: Vec, +} + +#[test] +#[ignore = "needs a GPU, CUA_S1_CUDA_LIB and QWEN_PREFIX_BIDIR_FIXTURE"] +fn multilayer_prefix_mask_matches_frozen_reference_and_restores_causal_execution() { + let file = PathBuf::from(std::env::var_os("QWEN_PREFIX_BIDIR_FIXTURE").expect("fixture path")); + let fixture: Fixture = serde_json::from_slice(&std::fs::read(&file).unwrap()).unwrap(); + assert_eq!(fixture.schema, "qwen-prefix-bidir-reference/1"); + assert!(!fixture.rows.is_empty()); + let checkpoint = file.parent().unwrap(); + let library = PathBuf::from(std::env::var_os("CUA_S1_CUDA_LIB").expect("CUDA library")); + let mut model = Model::load_with_graph(checkpoint, &library, false).unwrap(); + model.require_prompt_prefix_bidirectional().unwrap(); + assert!(model.cfg.full_attention.iter().filter(|&&x| x).count() >= 2); + assert!(model.cfg.full_attention.iter().any(|&x| !x)); + let first = &fixture.rows[0]; + let causal_before = model.forward(&first.ids).unwrap(); + let mut records = Vec::new(); + let mut worst = 0.0f32; + for row in fixture.rows.iter().chain([first]) { + let hidden = model + .forward_prompt_prefix_bidirectional(&row.ids, row.state_tokens) + .unwrap(); + assert_eq!(hidden.len(), row.hidden.len()); + assert!(hidden.iter().chain(&row.hidden).all(|x| x.is_finite())); + let scale = row.hidden.iter().map(|x| x.abs()).fold(1e-3f32, f32::max); + let error = hidden + .iter() + .zip(&row.hidden) + .map(|(a, b)| (a - b).abs()) + .fold(0.0f32, f32::max) + / scale; + eprintln!("{}: normalized hidden max error {error:.6e}", row.id); + worst = worst.max(error); + records.push(serde_json::json!({"id":row.id,"ids":row.ids,"state_tokens":row.state_tokens,"hidden":hidden,"normalized_max_error":error})); + } + assert_eq!( + records[0]["hidden"], + records.last().unwrap()["hidden"], + "earlier row changed after scratch reuse" + ); + assert_eq!( + causal_before, + model.forward(&first.ids).unwrap(), + "prefix mask leaked into a causal call" + ); + if let Some(output) = std::env::var_os("QWEN_PREFIX_BIDIR_OUTPUT") { + std::fs::write(output, serde_json::to_vec(&records).unwrap()).unwrap(); + } + // This synthetic BF16 kernel/backbone gate is deliberately separate from + // a model adapter's decision/probability gates. It is fixed before replay. + assert!(worst < 0.03, "full-row BF16 hidden gate failed: {worst}"); +} diff --git a/tests/qwen3_5/prefix_bidir_reference.py b/tests/qwen3_5/prefix_bidir_reference.py new file mode 100644 index 00000000..50086447 --- /dev/null +++ b/tests/qwen3_5/prefix_bidir_reference.py @@ -0,0 +1,89 @@ +"""Generate a small independent Transformers oracle; no pretrained weights. + +Run with torch 2.12.1+cu130 and transformers 5.17.0 on the qualification GPU: +python tests/qwen3_5/prefix_bidir_reference.py /absolute/output-directory +Generated weights/arrays are external evidence, not committed fixtures. +""" + +import argparse +import json +from pathlib import Path + +import torch +import transformers +from transformers import Qwen3_5TextConfig +from transformers.models.qwen3_5.modeling_qwen3_5 import Qwen3_5TextModel + + +@torch.inference_mode() +def generate(output): + if torch.__version__ != "2.12.1+cu130" or transformers.__version__ != "5.17.0": + raise RuntimeError("use the documented pinned oracle dependencies") + torch.manual_seed(20261101) + torch.backends.cuda.matmul.allow_tf32 = False + torch.backends.cudnn.allow_tf32 = False + config = Qwen3_5TextConfig( + vocab_size=128, + hidden_size=256, + intermediate_size=512, + num_hidden_layers=3, + num_attention_heads=1, + num_key_value_heads=1, + head_dim=256, + max_position_embeddings=512, + linear_num_key_heads=1, + linear_num_value_heads=2, + linear_key_head_dim=128, + linear_value_head_dim=128, + layer_types=["full_attention", "linear_attention", "full_attention"], + rope_parameters={ + "rope_type": "default", + "rope_theta": 10000.0, + "partial_rotary_factor": 0.25, + "mrope_section": [11, 11, 10], + "mrope_interleaved": True, + }, + ) + config._attn_implementation = "sdpa" + model = Qwen3_5TextModel(config).to(device="cuda", dtype=torch.bfloat16).eval() + output.mkdir(parents=True, exist_ok=False) + model.save_pretrained(output) + rows = [] + for length, state in [(3, 1), (65, 31), (129, 63), (137, 65), (257, 129)]: + ids = (torch.arange(length, device="cuda") * 17 + 7).remainder(128).view(1, -1) + # Dense boolean oracle, independent of the native two-pass composition. + visible = torch.ones((length, length), device="cuda", dtype=torch.bool) + visible[:state, state:] = False + mask = {"full_attention": visible[None, None], "linear_attention": None} + hidden = ( + model(input_ids=ids, attention_mask=mask, use_cache=False) + .last_hidden_state[0, -1] + .float() + ) + rows.append( + { + "id": f"t{length}-s{state}", + "ids": ids[0].tolist(), + "state_tokens": state, + "hidden": hidden.tolist(), + } + ) + result = { + "schema": "qwen-prefix-bidir-reference/1", + "rows": rows, + "metadata": { + "torch": torch.__version__, + "transformers": transformers.__version__, + "device": torch.cuda.get_device_name(), + "capability": torch.cuda.get_device_capability(), + "seed": 20261101, + "precision": "all BF16; dense block SDPA; causal GDN", + }, + } + (output / "reference.json").write_text(json.dumps(result, allow_nan=False) + "\n") + + +if __name__ == "__main__": + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("output", type=Path) + generate(parser.parse_args().output)