Pack Instant V into contiguous scratch before bidirectional attention

recipe/startlux/latency.md

@@ -187,8 +187,11 @@ Device scratch grows to the next multiple of 1024 tokens on the first prompt
 that exceeds the current cap, with a stream sync inside the request. Process
 warmup is one short choice, and `bench.py` warms five manifest rows, so a
 later longer Public231 prompt can allocate during a measured traversal.
-A static count of the Instant 4B scratch layout (the byte offsets in
-`Scratch::new`, including `cs1_gdn_workspace_floats`) is about 240 MiB at
-1024 tokens and about 15 GiB at 65536 tokens, before the weight arena.
+A static count of the Instant 4B scratch layout before the contiguous V
+buffer (the other byte offsets in `Scratch::new`, including
+`cs1_gdn_workspace_floats`) is about 240 MiB at 1024 tokens and about 15 GiB
+at 65536 tokens, before the weight arena. The V buffer is
+`kv_heads * head_dim` bfloat16 values per token (2048 bytes on Instant 4B):
+2 MiB at 1024 tokens and 128 MiB at 65536 tokens.
 That 65536-token buffer is not reserved. L20X free memory was not measured.
 Hypothesis: the in-request realloc can move p95. Measured delta: none.

src/backends/cuda/qwen3_5/README.md

@@ -18,7 +18,11 @@ the BF16 rounding of both attention and sigmoid before multiplication. The nativ
 workers use this entry point; the separate operations remain available for kernel
 comparisons. `cs1_attention_gated_bidirectional` is an optional ABI 9 symbol with
 the same argument list. It is the Instant prompt mask: every query reads every
-key in that prompt, and a non-zero `q_base` is rejected. Causal
+key in that prompt, and a non-zero `q_base` is rejected. The Instant caller
+packs V into contiguous `[T, Hk, Dh]` rows and passes `ldv = Hk * Dh` (1024 on
+Instant 4B). The kernel still accepts a larger `ldv`, which is how causal
+attention reads V from the QKV row. The pack is a Rust `cs1_copy2d`; it does
+not add a symbol. Causal
 `cs1_attention_gated` and prefix continuation stay causal. The second template
 parameter defaults to false, so existing causal launches keep that mask. Causal
 callers do not resolve the new symbol. StartLux Instant fails closed when it is

src/backends/cuda/qwen3_5/ops.h

@@ -114,8 +114,11 @@ int cs1_attention_gated(const void* q, const void* k, const void* v, int ldv, co
                         void* out, int T, int Hq, int Hk, int Dh, float scale, void* stream);
 
 // Additive, optional at ABI 9. Same tensors as cs1_attention_gated, but every query
-// reads every key in the prompt. Missing on libraries built before this entry point;
-// causal callers do not resolve it. q_base-style prefix continuation is not supported.
+// reads every key in the prompt. `ldv` is the V row stride in elements. The Instant
+// caller packs V into contiguous [T, Hk, Dh] and passes ldv = Hk*Dh; a larger stride
+// (V still inside the QKV row) remains valid. Missing on libraries built before this
+// entry point; causal callers do not resolve it. q_base-style prefix continuation is
+// not supported.
 int cs1_attention_gated_bidirectional(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, void* stream);

src/models/qwen3_5/native/Cargo.toml

@@ -20,6 +20,10 @@ serde_json = { version = "1.0.149", features = ["float_roundtrip", "preserve_ord
 name = "config"
 path = "../../../../tests/qwen3_5/config.rs"
 
+[[test]]
+name = "v_pack"
+path = "../../../../tests/qwen3_5/v_pack.rs"
+
 [[test]]
 name = "kernels"
 path = "../../../../tests/qwen3_5/kernels.rs"

src/models/qwen3_5/native/src/model.rs

@@ -450,6 +450,57 @@ impl Widths {
     }
 }
 
+/// V layout for one prompt-bidirectional full-attention call.
+///
+/// Q and K leave `cs1_attn_prep` in contiguous scratch. V is produced in the
+/// QKV row. Instant copies those V rows into contiguous `[T, Hk, Dh]` scratch
+/// and passes `ldv`. Causal attention does not pack V; it keeps `source_stride`
+/// as `ldv`.
+#[derive(Clone, Copy, Debug, PartialEq, Eq)]
+pub struct PromptBidirectionalVLayout {
+    /// Elements per packed V row, and the `ldv` argument (`kv_heads * head_dim`).
+    pub ldv: usize,
+    /// QKV row stride in elements (`attn_in`).
+    pub source_stride: usize,
+    /// Element index of V inside that row, after Q/gate and K.
+    pub source_column: usize,
+}
+
+impl PromptBidirectionalVLayout {
+    /// Bytes in one packed bfloat16 V row. This is also the contiguous destination pitch.
+    pub fn row_bytes(self) -> usize {
+        self.ldv * BF16
+    }
+
+    /// Bytes between QKV rows. This is the source pitch of the pack.
+    pub fn source_pitch_bytes(self) -> usize {
+        self.source_stride * BF16
+    }
+
+    /// Byte offset of packed V for `token` inside the contiguous scratch.
+    pub fn packed_offset(self, token: usize) -> usize {
+        token * self.row_bytes()
+    }
+}
+
+/// Layout of the V pack used by prompt-bidirectional attention.
+///
+/// `heads`, `kv_heads`, and `head_dim` are the text-config counts. Instant 4B
+/// is 16, 4, and 256, so `ldv` is 1024 and the QKV stride is 10240.
+pub fn prompt_bidirectional_v_layout(
+    heads: usize,
+    kv_heads: usize,
+    head_dim: usize,
+) -> PromptBidirectionalVLayout {
+    let attn_q = heads * head_dim * 2;
+    let packed = kv_heads * head_dim;
+    PromptBidirectionalVLayout {
+        ldv: packed,
+        source_stride: attn_q + 2 * packed,
+        source_column: attn_q + packed,
+    }
+}
+
 /// Per-request buffers for up to `cap` tokens, as byte offsets into one allocation,
 /// plus the rotary tables for positions below `cap`.
 struct Scratch {
@@ -472,6 +523,9 @@ struct Scratch {
     aq: usize,
     agate: usize,
     ak: usize,
+    /// Contiguous V `[cap, Hk, Dh]` for prompt-bidirectional attention.
+    /// Causal attention does not read it.
+    av: usize,
     ao: usize,
     gate_up: usize,
     act: usize,
@@ -536,6 +590,7 @@ impl Scratch {
             take(cap * hq * hd * BF16),
             take(cap * hq * hd * BF16),
             take(cap * hk * hd * BF16),
+            take(cap * hk * hd * BF16),
             take(cap * hq * hd * BF16),
             take(cap * 2 * cfg.intermediate * BF16),
             take(cap * cfg.intermediate * BF16),
@@ -563,6 +618,7 @@ impl Scratch {
             aq,
             agate,
             ak,
+            av,
             ao,
             gate_up,
             act,
@@ -603,6 +659,7 @@ impl Scratch {
             aq,
             agate,
             ak,
+            av,
             ao,
             gate_up,
             act,
@@ -1181,8 +1238,10 @@ impl Model {
     }
 
     /// One prompt whose full-attention layers read the whole prompt in both
-    /// directions. Linear attention stays causal. CUDA graphs, LoRA and the
-    /// reference numerics are refused: those paths are causal-only.
+    /// directions. Linear attention stays causal. V is packed into contiguous
+    /// `[T, Hk, Dh]` scratch and attention is called with `ldv = Hk * Dh`.
+    /// CUDA graphs, LoRA and the reference numerics are refused: those paths
+    /// are causal-only.
     pub fn forward_prompt_bidirectional(&mut self, ids: &[u32]) -> Result<Vec<f32>> {
         ensure!(
             !self.graph_enabled && self.reference.is_none() && self.lora.is_none(),
@@ -2000,8 +2059,10 @@ impl Model {
     }
 
     /// Queue language layers over packed embeddings, resetting sequence positions.
-    /// Final-norm hidden states end up in `s.x`. `bidirectional` changes only the
-    /// full-attention mask; Gated DeltaNet stays causal and sequence-local.
+    /// Final-norm hidden states end up in `s.x`. `bidirectional` changes the
+    /// full-attention mask and packs V into contiguous `[T, Hk, Dh]` scratch
+    /// (`ldv = Hk * Dh`). Causal full attention still reads V from the QKV row
+    /// (`ldv = attn_in`). Gated DeltaNet stays causal and sequence-local.
     fn run(
         &self,
         s: &Scratch,
@@ -2034,6 +2095,16 @@ impl Model {
         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 w = Widths::of(cfg);
+        let packed_v = if bidirectional {
+            let layout = prompt_bidirectional_v_layout(cfg.heads, cfg.kv_heads, cfg.head_dim);
+            ensure!(
+                layout.source_stride == w.attn_in && layout.source_column + layout.ldv == w.attn_in,
+                "bidirectional V pack does not match the QKV row"
+            );
+            Some(layout)
+        } else {
+            None
+        };
         let p = |off: usize| s.at(off);
         let (cos, sin) = if custom_positions {
             (s.custom_cos, s.custom_sin)
@@ -2179,6 +2250,26 @@ impl Model {
                     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;
+                    // Instant only: copy V out of the QKV row into contiguous
+                    // [T, Hk, Dh] scratch. Causal attention keeps the strided column.
+                    let (v_base, v_pitch, v_ld) = if let Some(layout) = packed_v {
+                        if t > 0 {
+                            unsafe {
+                                cuda::copy2d(
+                                    p(s.av),
+                                    layout.row_bytes(),
+                                    p(v),
+                                    layout.source_pitch_bytes(),
+                                    layout.row_bytes(),
+                                    t,
+                                    st,
+                                )?;
+                            }
+                        }
+                        (s.av, layout.row_bytes(), layout.ldv as i32)
+                    } else {
+                        (v, w.attn_in * BF16, ld)
+                    };
                     unsafe {
                         for &(offset, length) in &sequences {
                             check(
@@ -2206,13 +2297,14 @@ impl Model {
                                 "attention prep",
                             )?;
                             let scale = (cfg.head_dim as f32).powf(-0.5);
+                            let v_row = p(v_base + offset * v_pitch);
                             check(
                                 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,
+                                        v_row,
+                                        v_ld,
                                         p(s.agate + offset * cfg.heads * cfg.head_dim * BF16),
                                         p(s.ao + offset * cfg.heads * cfg.head_dim * BF16),
                                         length,
@@ -2226,8 +2318,8 @@ impl Model {
                                     (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,
+                                        v_row,
+                                        v_ld,
                                         p(s.agate + offset * cfg.heads * cfg.head_dim * BF16),
                                         p(s.ao + offset * cfg.heads * cfg.head_dim * BF16),
                                         length,
@@ -2241,8 +2333,8 @@ impl Model {
                                     (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,
+                                        v_row,
+                                        v_ld,
                                         p(s.agate + offset * cfg.heads * cfg.head_dim * BF16),
                                         p(s.ao + offset * cfg.heads * cfg.head_dim * BF16),
                                         length,

src/models/startlux/README.md

@@ -27,7 +27,11 @@ Instant is Qwen3.5-4B with bidirectional full attention on one of every four
 layers and a letter readout. It is not an Open-Jev LoRA plus a scalar head.
 The worker reuses the shared Qwen3.5 loader, GDN, MLP, norms, RoPE, and
 GEMMs. Full-attention layers call the additive
-`cs1_attention_gated_bidirectional` symbol. Gated DeltaNet stays causal.
+`cs1_attention_gated_bidirectional` symbol. After the QKV projection, V is
+packed into a contiguous `[tokens, kv_heads, head_dim]` scratch (the same
+width as post-prep K). The call passes `ldv = kv_heads * head_dim` (1024 on
+Instant 4B), not the QKV row stride 10240. Causal workers still read V from
+the QKV row. Gated DeltaNet stays causal.
 There is no prefix cache and no CUDA graph on this path.
 
 Processing (`processing.rs`) loads the tokenizer once and shares it.

tests/qwen3_5/kernels.rs

@@ -404,7 +404,9 @@ fn flash_attention_matches_float64_reference() {
         (2048, 0.5),
     ] {
         let (qh, kh) = (random(t * hq * dh, 1, amp), random(t * hk * dh, 2, amp));
-        // v is read in place from the q|k|v projection output, rows of 10240 as in the model
+        // Causal attention still reads V from the QKV row (stride 10240 on this
+        // 4B shape). Instant bidirectional packs those rows and passes ldv = Hk*Dh;
+        // this launch keeps the strided kernel contract.
         let (ldv, v_at) = (10240usize, (hq * 2 + hk) * dh);
         let qkvh = random(t * ldv, 3, 1.0);
         let (q, k, qkv) = (to_device(&qh, st), to_device(&kh, st), to_device(&qkvh, st));

tests/qwen3_5/v_pack.rs

@@ -0,0 +1,82 @@
+//! CPU contract for the Instant V pack. No CUDA library and no GPU.
+//!
+//! Prompt-bidirectional attention copies V out of the QKV row into contiguous
+//! `[T, Hk, Dh]` scratch. The device copy is `cs1_copy2d`: `tokens` rows,
+//! width = destination pitch = one packed bfloat16 row, source pitch = the
+//! QKV row. Each sequence then starts at `packed_offset(token)`.
+
+use omni_qwen3_5_native::model::{PromptBidirectionalVLayout, prompt_bidirectional_v_layout};
+
+fn pack_rows(src: &[u32], layout: PromptBidirectionalVLayout, tokens: usize) -> Vec<u32> {
+    assert_eq!(layout.source_column + layout.ldv, layout.source_stride);
+    assert!(layout.row_bytes() <= layout.source_pitch_bytes());
+    let mut dst = vec![0; tokens * layout.ldv];
+    for token in 0..tokens {
+        let from = token * layout.source_stride + layout.source_column;
+        let to = token * layout.ldv;
+        dst[to..to + layout.ldv].copy_from_slice(&src[from..from + layout.ldv]);
+    }
+    dst
+}
+
+#[test]
+fn instant_4b_v_ldv_is_hk_dh_not_the_qkv_stride() {
+    let (heads, kv, dh) = (16usize, 4usize, 256usize);
+    let layout = prompt_bidirectional_v_layout(heads, kv, dh);
+    assert_eq!(layout.ldv, kv * dh);
+    assert_eq!(layout.ldv, 1024);
+    assert_eq!(layout.source_stride, heads * dh * 2 + 2 * kv * dh);
+    assert_eq!(layout.source_stride, 10240);
+    assert_eq!(layout.source_column, heads * dh * 2 + kv * dh);
+    assert_eq!(layout.source_column, 9216);
+    assert_eq!(layout.source_column + layout.ldv, layout.source_stride);
+    assert_eq!(layout.row_bytes(), 2048);
+    assert_eq!(layout.source_pitch_bytes(), 20480);
+    assert_eq!(layout.ldv % 8, 0, "kernel requires ldv % 8 == 0");
+    assert!(layout.ldv < layout.source_stride);
+    assert_ne!(layout.ldv, layout.source_stride);
+}
+
+#[test]
+fn packed_rows_match_the_strided_v_column_and_skip_q_and_k() {
+    let layout = prompt_bidirectional_v_layout(16, 4, 256);
+    let tokens = 5usize;
+    let mut src = vec![0u32; tokens * layout.source_stride];
+    for token in 0..tokens {
+        for col in 0..layout.source_stride {
+            src[token * layout.source_stride + col] = (token * layout.source_stride + col) as u32;
+        }
+    }
+    let packed = pack_rows(&src, layout, tokens);
+    assert_eq!(packed.len(), tokens * layout.ldv);
+    assert_eq!(packed[0], layout.source_column as u32);
+    assert_ne!(packed[0], src[0], "Q must not occupy packed V");
+    for token in 0..tokens {
+        for i in 0..layout.ldv {
+            assert_eq!(
+                packed[token * layout.ldv + i],
+                (token * layout.source_stride + layout.source_column + i) as u32
+            );
+        }
+        assert_eq!(
+            layout.packed_offset(token) / 2,
+            token * layout.ldv,
+            "sequence base is a token count times the packed row"
+        );
+    }
+    let seq = 3usize;
+    assert_eq!(layout.packed_offset(seq), seq * layout.row_bytes());
+    assert_eq!(
+        packed[seq * layout.ldv],
+        src[seq * layout.source_stride + layout.source_column]
+    );
+}
+
+#[test]
+fn other_head_counts_keep_ldv_at_hk_dh() {
+    let layout = prompt_bidirectional_v_layout(24, 4, 256);
+    assert_eq!(layout.ldv, 1024);
+    assert_eq!(layout.source_stride, 14336);
+    assert_eq!(layout.source_column + layout.ldv, layout.source_stride);
+    assert_eq!(layout.row_bytes(), layout.ldv * 2);
+}

tests/startlux/contract.rs

@@ -440,6 +440,26 @@ fn rejections_match_the_instant_contract() {
     assert_eq!(named.status, 422);
 }
 
+#[test]
+fn instant_4b_packs_bidirectional_v_to_hk_dh() {
+    let config: Value =
+        serde_json::from_str(include_str!("fixtures/instant_4b_config.json")).unwrap();
+    let text = &config["text_config"];
+    let heads = text["num_attention_heads"].as_u64().unwrap() as usize;
+    let kv = text["num_key_value_heads"].as_u64().unwrap() as usize;
+    let dh = text["head_dim"].as_u64().unwrap() as usize;
+    assert_eq!((heads, kv, dh), (16, 4, 256));
+    let layout = omni_qwen3_5_native::model::prompt_bidirectional_v_layout(heads, kv, dh);
+    assert_eq!(layout.ldv, 1024);
+    assert_eq!(layout.source_stride, 10240);
+    assert_eq!(layout.source_column + layout.ldv, layout.source_stride);
+    assert_eq!(layout.row_bytes(), 2048);
+    assert_eq!(layout.source_pitch_bytes(), 20_480);
+    assert_eq!(layout.packed_offset(0), 0);
+    assert_eq!(layout.packed_offset(3), 3 * 2048);
+    assert_ne!(layout.ldv, layout.source_stride);
+}
+
 #[test]
 fn temperatures_and_confidence_follow_the_pinned_readout() {
     let even = contract::probabilities(&[0.0, 0.0], Kind::Choice.temperature());