Skip to main content

fuel_dispatch/
decode_flash.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2//! Decode flash-arm emission — the safe, gated `Op::Branch` that makes the
3//! CUDA `flash_decoding` binding (commit `f693da2c`) reachable in decode.
4//!
5//! ## Why this module exists
6//!
7//! Fuel dispatch is **fail-fast** with a key-only, shape-blind executor
8//! binding lookup: a registered kernel that returns `Err` fails the whole
9//! `realize`, and a wrapper *cannot* soft-decline back to the decomposed
10//! base map (the finding recorded in `f693da2c`). The CUDA
11//! `(FlashAttn, [f16|bf16;4], Cuda)` capacity-K bindings are live-validated,
12//! but only reachable **safely** if a CUDA `FlashAttn` node is *only ever
13//! created for shapes the kernel supports*: `seq_q == 1`, `head_dim ∈
14//! [1, 128]`, `f16`/`bf16`, no window / ALiBi / softcap. The plan-time
15//! infeasible-cost gate (`cost::cost_flash_decoding_cuda`) is defense in
16//! depth for unpinned placement; the *guaranteed* clean gate is
17//! **graph-construction-time arm emission** — only build the CUDA arm for
18//! supported decode shapes. That is this module.
19//!
20//! ## Constitutional posture (04-optimization / 06-runtime)
21//!
22//! - **The optimizer emits/prunes arms; backends never decide.** All the
23//!   strategy — the shape/dtype/config gate, the capability gate, the
24//!   CUDA pin, the `Op::Branch` construction — lives here, in the dispatch
25//!   (optimizer) layer, not in the model. The model layer only supplies
26//!   the region's tensor handles and the live `k_len` it alone knows (the
27//!   attended prefix = `cached_len + seq`), which is **data**, not a
28//!   strategic choice.
29//! - **The decomposed base map stays arm 0 / the correctness oracle.** The
30//!   `Op::Branch`'s arm 0 is the *existing decomposed region* the decode
31//!   graph already builds (`matmul → scale → mask → softmax → matmul` over
32//!   the capacity KV). `finalize_branches` enforces that the merge reads
33//!   arm 0, so a finalized-but-unpicked graph — and every CPU/Vulkan build,
34//!   where no CUDA flash arm is ever emitted — realizes on the decomposed
35//!   route, byte-identical to today.
36//! - **CUDA-only.** The capacity-K prefix (`k_len != sk`) is a CUDA-only
37//!   capability: `vulkan_dispatch::flash_attn` hard-rejects `k_len != sk`
38//!   (fail-fast → fails realize) and the CPU flash wrapper would silently
39//!   change results vs. the decomposed arm. So arm 1 is pinned to
40//!   `BackendId::Cuda` and only emitted when a CUDA flash kernel is bound
41//!   AND a CUDA device is in the current topology.
42//!
43//! ## What this module does NOT do (documented remaining slice)
44//!
45//! Emitting the arm makes the flash kernel **offered**. Whether it is
46//! **selected** at dispatch is the runtime route picker's job — and the
47//! shipped picker (`ranker::route_picker`) resolves *placement* choices
48//! (it ranks arms by their `(backend, device)` keyed on arm 0's op), not a
49//! **same-device kernel-variant** choice (flash vs. decomposed, both on
50//! CUDA). Per 04-optimization, kernel-variant choice is "largely baked at
51//! optimize time"; selecting the flash arm on a same-device basis needs
52//! either a per-arm-op cost leg in the picker or an optimize-time route
53//! bake. That, plus wiring this emitter into `fuel-core`'s decode builder
54//! (with the persistent-decode attended-length symbol), is the remaining
55//! integration — see the module's tests + the session report.
56
57use fuel_graph::registry::{FusedOpParams, FusedOps};
58use fuel_graph::{Graph, Node, NodeId, Op};
59use fuel_ir::probe::BackendId;
60use fuel_ir::{DType, DynScalar, Result, Shape};
61
62use crate::runtime_fused_kernels::fused_kernel_available;
63use crate::topology::SystemTopology;
64
65/// The capability + topology gate for the CUDA flash decode arm: a CUDA
66/// `FlashAttn` kernel must be bound **and** a CUDA device must be present
67/// in the current topology. Split out (and injectable) so the gate logic is
68/// unit-testable without a CUDA build, while production reads the real
69/// process-global registry + topology via [`Self::production`].
70#[derive(Clone, Copy, Debug)]
71pub struct FlashArmCapability {
72    /// A `(FlashAttn, Cuda)` kernel is bound (static or runtime-adopted).
73    pub cuda_flash_kernel: bool,
74    /// A CUDA device is present in the current [`SystemTopology`].
75    pub cuda_in_topology: bool,
76}
77
78impl FlashArmCapability {
79    /// Read the real process-global capability.
80    pub fn production() -> Self {
81        Self {
82            cuda_flash_kernel: fused_kernel_available(FusedOps::FLASH_ATTN, BackendId::Cuda),
83            cuda_in_topology: SystemTopology::current()
84                .backends()
85                .contains(&BackendId::Cuda),
86        }
87    }
88
89    /// A capability that always admits the arm — for unit tests that drive
90    /// the shape/dtype/config gate without a CUDA build.
91    #[cfg(test)]
92    pub fn all_available() -> Self {
93        Self {
94            cuda_flash_kernel: true,
95            cuda_in_topology: true,
96        }
97    }
98
99    /// The flash arm may be offered only when both halves hold.
100    pub fn available(&self) -> bool {
101        self.cuda_flash_kernel && self.cuda_in_topology
102    }
103}
104
105/// The decode-attention region the model hands to the optimizer as a
106/// flash-arm **candidate**. Arm 0 — the decomposed region — is already in
107/// the graph; this describes how to build arm 1 (the flash kernel) and
108/// where to splice the `Op::Branch`.
109///
110/// `q` is both the flash node's first input and the branch's shared
111/// **diverge** point (the decomposed chain and the flash node both depart
112/// from `q`). `decomposed_out` is arm 0's exit (the region's materialized
113/// output, `[B, Hq, Sq, D]` — same shape/dtype as the flash output).
114/// `reconverge` is the sole consumer of `decomposed_out` (the merge).
115#[derive(Clone, Debug)]
116pub struct DecodeFlashSpec {
117    /// `q` — `[B, Hq, Sq, D]`. Also the branch diverge point.
118    pub q: NodeId,
119    /// `k` — the capacity KV buffer `[B, Hkv, capacity, D]`.
120    pub k: NodeId,
121    /// `v` — the capacity KV buffer `[B, Hkv, capacity, D]`.
122    pub v: NodeId,
123    /// Optional ALiBi slopes `[Hq]`. Its presence *disqualifies* the arm
124    /// (the capacity-K CUDA kernel has no ALiBi param) — carried so the
125    /// gate can reject it explicitly.
126    pub alibi: Option<NodeId>,
127    /// `1/sqrt(head_dim)` (or the model's scale).
128    pub softmax_scale: f32,
129    /// Causal masking. Accepted-and-ignored by the kernel (`seq_q == 1`
130    /// attends the full `[0, k_len)` prefix), so it does not disqualify.
131    pub causal: bool,
132    /// Left sliding-window bound — disqualifies (no kernel support).
133    pub window_size_left: Option<usize>,
134    /// Right sliding-window bound — disqualifies (no kernel support).
135    pub window_size_right: Option<usize>,
136    /// Logit softcap — disqualifies (no kernel support).
137    pub softcap: Option<f32>,
138    /// The live attended prefix length (`cached_len + seq`), resolved via
139    /// the per-pass `SymEnv` at realize. For a per-step (non-persistent)
140    /// decode graph this is `DynScalar::Concrete(cached_len + seq)`; for
141    /// persistent (plan-once) decode it is a `Sym` bound to the live
142    /// length each token.
143    pub k_len: DynScalar,
144    /// Arm 0's exit: the decomposed region's output node (the oracle).
145    pub decomposed_out: NodeId,
146    /// The sole consumer of `decomposed_out` — the branch merge point.
147    pub reconverge: NodeId,
148}
149
150/// Is a decode-attention region admissible for the CUDA flash arm?
151///
152/// Rejects (→ `false`, decomposed-only) anything the capacity-K CUDA
153/// binding cannot do, matching `cost::cost_flash_decoding_cuda`'s
154/// infeasible set and the kernel's `_can_implement` gate:
155///
156/// - capability unavailable (no CUDA flash kernel / no CUDA device);
157/// - dtype not `f16` / `bf16`;
158/// - `q` not rank-4, `seq_q != 1`, or `head_dim` `0` or `> 128`;
159/// - sliding window (`window_size_{left,right}`), softcap, or ALiBi set.
160///
161/// `causal` is intentionally NOT a disqualifier: with `seq_q == 1` the
162/// single query attends the whole `[0, k_len)` prefix, so the causal mask
163/// is a no-op the kernel omits.
164pub fn flash_decode_admissible(
165    q_shape: &Shape,
166    dtype: DType,
167    window_size_left: Option<usize>,
168    window_size_right: Option<usize>,
169    softcap: Option<f32>,
170    has_alibi: bool,
171    cap: FlashArmCapability,
172) -> bool {
173    // Capability: a CUDA FlashAttn kernel bound AND CUDA present.
174    if !cap.available() {
175        return false;
176    }
177    // Dtype: the CUDA flash_decoding binding is f16 / bf16 only.
178    if !matches!(dtype, DType::F16 | DType::BF16) {
179        return false;
180    }
181    // Shape: decode is seq_q == 1; head_dim in [1, 128].
182    let dims = q_shape.dims();
183    if dims.len() != 4 {
184        return false;
185    }
186    let sq = dims[2];
187    let d = dims[3];
188    if sq != 1 || d == 0 || d > 128 {
189        return false;
190    }
191    // Config the capacity-K kernel does not implement.
192    if window_size_left.is_some() || window_size_right.is_some() {
193        return false;
194    }
195    if softcap.is_some() {
196        return false;
197    }
198    if has_alibi {
199        return false;
200    }
201    true
202}
203
204/// Offer the CUDA flash decode arm for `spec` on `graph`.
205///
206/// If [`flash_decode_admissible`] holds for `spec`'s `q` shape/dtype +
207/// config + `cap`, this builds arm 1 (the flash `Op::Fused(FLASH_ATTN, …)`
208/// node, CUDA-pinned) and records a 2-arm `Op::Branch`:
209///
210/// - **arm 0** = `spec.decomposed_out` (the decomposed oracle — the route
211///   an unpicked graph, and every non-CUDA build, realizes on);
212/// - **arm 1** = the flash node.
213///
214/// Returns `Ok(Some(branch))` when the arm was emitted, `Ok(None)` when the
215/// gate declined (the graph is left with the decomposed route only —
216/// byte-identical to today). Surfaces a build-time `Err` only on a
217/// malformed spec that fails `finalize_branches`' validation (never a
218/// panic).
219///
220/// The flash node reads `[q, k, v, (alibi)]` and carries `k_len` in its
221/// `FusedOpParams::FlashAttn` — the `flash_attn_dyn → SymEnv → OpParams`
222/// path the pipelined lowering already resolves (`pipelined.rs` ~3660).
223pub fn offer_decode_flash_arm(
224    graph: &mut Graph,
225    spec: &DecodeFlashSpec,
226    cap: FlashArmCapability,
227) -> Result<Option<NodeId>> {
228    let (q_shape, q_dtype) = {
229        let n = graph.node(spec.q);
230        (n.shape.clone(), n.dtype)
231    };
232
233    if !flash_decode_admissible(
234        &q_shape,
235        q_dtype,
236        spec.window_size_left,
237        spec.window_size_right,
238        spec.softcap,
239        spec.alibi.is_some(),
240        cap,
241    ) {
242        return Ok(None);
243    }
244
245    // Arm 1: the flash Op::Fused(FLASH_ATTN, { k_len }) node. Shape + dtype
246    // MUST equal arm 0 (the decomposed output) — both are q's shape/dtype —
247    // or finalize_branches' cast-to-uniform rule rejects the branch.
248    let mut inputs = vec![spec.q, spec.k, spec.v];
249    if let Some(a) = spec.alibi {
250        inputs.push(a);
251    }
252    let flash = graph.push(Node {
253        op: Op::Fused(
254            FusedOps::FLASH_ATTN,
255            FusedOpParams::FlashAttn {
256                softmax_scale: spec.softmax_scale,
257                causal: spec.causal,
258                window_size_left: spec.window_size_left,
259                window_size_right: spec.window_size_right,
260                softcap: spec.softcap,
261                k_len: Some(spec.k_len),
262            },
263        ),
264        inputs,
265        shape: q_shape,
266        dtype: q_dtype,
267    });
268    // Pin arm 1 to CUDA — the only backend whose FlashAttn binding handles a
269    // capacity-K prefix (Vulkan hard-rejects k_len != sk; the CPU flash
270    // wrapper would silently change results vs. the decomposed arm).
271    graph.set_target_backend(flash, BackendId::Cuda);
272
273    // Emit the branch: arm 0 = decomposed oracle, arm 1 = flash. The diverge
274    // is `q` (both arms depart from it); the merge is `reconverge`, which
275    // reads arm 0 (the runnability invariant `finalize_branches` enforces).
276    let mut builder = graph.open_branch(spec.q);
277    builder.add_arm(spec.decomposed_out); // arm 0
278    builder.add_arm(flash); // arm 1
279    let branch = builder.finalize_branches(graph, spec.reconverge)?;
280    Ok(branch)
281}
282
283#[cfg(test)]
284mod tests {
285    use super::*;
286
287    /// Push a Const node of the given shape/dtype (a stand-in leaf tensor).
288    fn leaf(g: &mut Graph, dims: &[usize], dtype: DType) -> NodeId {
289        g.push(Node {
290            op: Op::Const,
291            inputs: vec![],
292            shape: Shape::from_dims(dims),
293            dtype,
294        })
295    }
296
297    /// Build a synthetic decode-attention region:
298    ///   q [1,H,1,D], k/v [1,H,SK,D]  (dtype)
299    ///   scores  = MatMul(q, Permute(k))           [1,H,1,SK]
300    ///   scaled  = MulScalar(scale)(scores)
301    ///   masked  = Add(scaled, mask[1,H,1,SK])
302    ///   probs   = Fused(SOFTMAX_LAST_DIM)(masked)
303    ///   attn_v  = MatMul(probs, v)                [1,H,1,D]   (arm-0 exit)
304    ///   merged  = Permute(attn_v)                 (reconverge / sole consumer)
305    /// Returns `(graph, spec-parts)`: (q, k, v, attn_v, merged).
306    fn decode_region(
307        h: usize,
308        d: usize,
309        sk: usize,
310        dtype: DType,
311    ) -> (Graph, NodeId, NodeId, NodeId, NodeId, NodeId) {
312        let mut g = Graph::new();
313        let q = leaf(&mut g, &[1, h, 1, d], dtype);
314        let k = leaf(&mut g, &[1, h, sk, d], dtype);
315        let v = leaf(&mut g, &[1, h, sk, d], dtype);
316        let mask = leaf(&mut g, &[1, h, 1, sk], dtype);
317
318        let kt = g.push(Node {
319            op: Op::Permute(vec![0, 1, 3, 2]),
320            inputs: vec![k],
321            shape: Shape::from_dims(&[1, h, d, sk]),
322            dtype,
323        });
324        let scores = g.push(Node {
325            op: Op::MatMul,
326            inputs: vec![q, kt],
327            shape: Shape::from_dims(&[1, h, 1, sk]),
328            dtype,
329        });
330        let scaled = g.push(Node {
331            op: Op::MulScalar(0.5),
332            inputs: vec![scores],
333            shape: Shape::from_dims(&[1, h, 1, sk]),
334            dtype,
335        });
336        let masked = g.push(Node {
337            op: Op::Add,
338            inputs: vec![scaled, mask],
339            shape: Shape::from_dims(&[1, h, 1, sk]),
340            dtype,
341        });
342        let probs = g.push(Node {
343            op: Op::Fused(FusedOps::SOFTMAX_LAST_DIM, FusedOpParams::SoftmaxLastDim),
344            inputs: vec![masked],
345            shape: Shape::from_dims(&[1, h, 1, sk]),
346            dtype,
347        });
348        let attn_v = g.push(Node {
349            op: Op::MatMul,
350            inputs: vec![probs, v],
351            shape: Shape::from_dims(&[1, h, 1, d]),
352            dtype,
353        });
354        // reconverge: the sole consumer of attn_v.
355        let merged = g.push(Node {
356            op: Op::Permute(vec![0, 2, 1, 3]),
357            inputs: vec![attn_v],
358            shape: Shape::from_dims(&[1, 1, h, d]),
359            dtype,
360        });
361        (g, q, k, v, attn_v, merged)
362    }
363
364    fn spec_for(
365        q: NodeId,
366        k: NodeId,
367        v: NodeId,
368        attn_v: NodeId,
369        merged: NodeId,
370        sk: usize,
371    ) -> DecodeFlashSpec {
372        DecodeFlashSpec {
373            q,
374            k,
375            v,
376            alibi: None,
377            softmax_scale: 0.5,
378            causal: true,
379            window_size_left: None,
380            window_size_right: None,
381            softcap: None,
382            k_len: DynScalar::Concrete(sk),
383            decomposed_out: attn_v,
384            reconverge: merged,
385        }
386    }
387
388    /// GREEN target: a supported decode region (f16, seq_q==1, head_dim<=64,
389    /// capacity KV) on a CUDA-capable topology gets the flash arm offered —
390    /// arm 0 stays the decomposed oracle, arm 1 is a CUDA-pinned
391    /// Fused(FLASH_ATTN, { k_len: Some }).
392    #[test]
393    fn flash_arm_offered_for_supported_decode_shape() {
394        let (mut g, q, k, v, attn_v, merged) = decode_region(4, 64, 37, DType::F16);
395        let spec = spec_for(q, k, v, attn_v, merged, 37);
396
397        let before = g.len();
398        let branch = offer_decode_flash_arm(&mut g, &spec, FlashArmCapability::all_available())
399            .expect("well-formed spec")
400            .expect("supported shape ⇒ arm offered");
401
402        // A Branch was recorded.
403        assert!(
404            matches!(g.node(branch).op, Op::Branch { .. }),
405            "branch node emitted"
406        );
407        let arms = g.node(branch).inputs.clone();
408        assert_eq!(arms.len(), 2, "2-arm branch (decomposed + flash)");
409        // arm 0 = the decomposed oracle.
410        assert_eq!(
411            arms[0], attn_v,
412            "arm 0 is the decomposed region output (the oracle)"
413        );
414        // arm 1 = a CUDA-pinned Fused(FLASH_ATTN, { k_len: Some }).
415        let flash = arms[1];
416        match &g.node(flash).op {
417            Op::Fused(fid, FusedOpParams::FlashAttn { k_len, .. }) => {
418                assert_eq!(*fid, FusedOps::FLASH_ATTN, "arm 1 is FLASH_ATTN");
419                assert_eq!(
420                    *k_len,
421                    Some(DynScalar::Concrete(37)),
422                    "arm 1 carries the live k_len"
423                );
424            }
425            other => panic!("arm 1 must be Fused(FLASH_ATTN, FlashAttn), got {other:?}"),
426        }
427        assert_eq!(g.node(flash).inputs, vec![q, k, v], "flash reads q, k, v");
428        assert_eq!(
429            g.target_backend(flash),
430            Some(BackendId::Cuda),
431            "arm 1 pinned to CUDA"
432        );
433        // Arm-0 runnability: the merge still reads the decomposed output.
434        assert!(
435            g.node(merged).inputs.contains(&attn_v),
436            "reconverge reads arm 0 (an unpicked graph realizes decomposed)",
437        );
438        assert!(g.len() > before, "the flash node + branch were appended");
439    }
440
441    /// GUARD: f32 dtype ⇒ no arm (the CUDA binding is f16/bf16 only) ⇒
442    /// decomposed-only, graph untouched.
443    #[test]
444    fn f32_dtype_gets_no_flash_arm() {
445        let (mut g, q, k, v, attn_v, merged) = decode_region(4, 64, 37, DType::F32);
446        let spec = spec_for(q, k, v, attn_v, merged, 37);
447        let before = g.len();
448        let branch = offer_decode_flash_arm(&mut g, &spec, FlashArmCapability::all_available())
449            .expect("well-formed spec");
450        assert!(branch.is_none(), "f32 ⇒ no flash arm");
451        assert_eq!(
452            g.len(),
453            before,
454            "graph untouched (no flash node, no branch)"
455        );
456    }
457
458    /// GUARD: head_dim > 128 ⇒ no arm (outside the kernel's `_can_implement`).
459    #[test]
460    fn head_dim_over_128_gets_no_flash_arm() {
461        let (mut g, q, k, v, attn_v, merged) = decode_region(2, 256, 8, DType::F16);
462        let spec = spec_for(q, k, v, attn_v, merged, 8);
463        let before = g.len();
464        let branch = offer_decode_flash_arm(&mut g, &spec, FlashArmCapability::all_available())
465            .expect("well-formed spec");
466        assert!(branch.is_none(), "head_dim 256 > 128 ⇒ no flash arm");
467        assert_eq!(g.len(), before, "graph untouched");
468    }
469
470    /// GUARD: seq_q != 1 (prefill-shaped) ⇒ no arm. The capacity-K decode
471    /// kernel is `seq_q == 1` only.
472    #[test]
473    fn multi_query_gets_no_flash_arm() {
474        // Build a region with seq_q = 8.
475        let (mut g, _q, _k, _v, _av, _m) = decode_region(4, 64, 37, DType::F16);
476        // Re-make q with seq_q = 8 and rewire a fresh minimal region.
477        let dtype = DType::F16;
478        let q = leaf(&mut g, &[1, 4, 8, 64], dtype);
479        let k = leaf(&mut g, &[1, 4, 37, 64], dtype);
480        let v = leaf(&mut g, &[1, 4, 37, 64], dtype);
481        let attn_v = g.push(Node {
482            op: Op::MatMul,
483            inputs: vec![q, v],
484            shape: Shape::from_dims(&[1, 4, 8, 64]),
485            dtype,
486        });
487        let merged = g.push(Node {
488            op: Op::Permute(vec![0, 2, 1, 3]),
489            inputs: vec![attn_v],
490            shape: Shape::from_dims(&[1, 8, 4, 64]),
491            dtype,
492        });
493        let spec = spec_for(q, k, v, attn_v, merged, 37);
494        let branch = offer_decode_flash_arm(&mut g, &spec, FlashArmCapability::all_available())
495            .expect("well-formed spec");
496        assert!(branch.is_none(), "seq_q = 8 != 1 ⇒ no flash arm");
497    }
498
499    /// GUARD: capability unavailable (no CUDA flash kernel / no CUDA device)
500    /// ⇒ no arm — a CPU/Vulkan build stays byte-identical to today.
501    #[test]
502    fn no_cuda_capability_gets_no_flash_arm() {
503        let (mut g, q, k, v, attn_v, merged) = decode_region(4, 64, 37, DType::F16);
504        let spec = spec_for(q, k, v, attn_v, merged, 37);
505        let before = g.len();
506
507        // No kernel bound.
508        let no_kernel = FlashArmCapability {
509            cuda_flash_kernel: false,
510            cuda_in_topology: true,
511        };
512        assert!(
513            offer_decode_flash_arm(&mut g, &spec, no_kernel)
514                .unwrap()
515                .is_none()
516        );
517        // No CUDA device present.
518        let no_device = FlashArmCapability {
519            cuda_flash_kernel: true,
520            cuda_in_topology: false,
521        };
522        assert!(
523            offer_decode_flash_arm(&mut g, &spec, no_device)
524                .unwrap()
525                .is_none()
526        );
527
528        assert_eq!(
529            g.len(),
530            before,
531            "no capability ⇒ decomposed-only, graph untouched"
532        );
533    }
534
535    /// GUARD: window / softcap / ALiBi each disqualify (no kernel support).
536    #[test]
537    fn window_softcap_alibi_get_no_flash_arm() {
538        for mutate in [
539            |s: &mut DecodeFlashSpec| s.window_size_left = Some(64),
540            |s: &mut DecodeFlashSpec| s.window_size_right = Some(64),
541            |s: &mut DecodeFlashSpec| s.softcap = Some(30.0),
542        ] {
543            let (mut g, q, k, v, attn_v, merged) = decode_region(4, 64, 37, DType::F16);
544            let mut spec = spec_for(q, k, v, attn_v, merged, 37);
545            mutate(&mut spec);
546            assert!(
547                offer_decode_flash_arm(&mut g, &spec, FlashArmCapability::all_available())
548                    .unwrap()
549                    .is_none(),
550                "window/softcap disqualifies the flash arm",
551            );
552        }
553        // ALiBi: add a slopes leaf and point the spec at it.
554        let (mut g, q, k, v, attn_v, merged) = decode_region(4, 64, 37, DType::F16);
555        let alibi = leaf(&mut g, &[4], DType::F16);
556        let mut spec = spec_for(q, k, v, attn_v, merged, 37);
557        spec.alibi = Some(alibi);
558        assert!(
559            offer_decode_flash_arm(&mut g, &spec, FlashArmCapability::all_available())
560                .unwrap()
561                .is_none(),
562            "ALiBi disqualifies the flash arm",
563        );
564    }
565
566    /// `causal` does NOT disqualify (seq_q == 1 attends the whole prefix, so
567    /// the causal mask is a no-op the kernel omits).
568    #[test]
569    fn causal_flag_does_not_disqualify() {
570        let (mut g, q, k, v, attn_v, merged) = decode_region(4, 64, 37, DType::F16);
571        let mut spec = spec_for(q, k, v, attn_v, merged, 37);
572        spec.causal = true;
573        assert!(
574            offer_decode_flash_arm(&mut g, &spec, FlashArmCapability::all_available())
575                .unwrap()
576                .is_some(),
577            "causal is accepted-and-ignored ⇒ arm still offered",
578        );
579    }
580
581    /// bf16 is supported alongside f16.
582    #[test]
583    fn bf16_dtype_gets_flash_arm() {
584        let (mut g, q, k, v, attn_v, merged) = decode_region(4, 64, 37, DType::BF16);
585        let spec = spec_for(q, k, v, attn_v, merged, 37);
586        assert!(
587            offer_decode_flash_arm(&mut g, &spec, FlashArmCapability::all_available())
588                .unwrap()
589                .is_some(),
590            "bf16 is a supported decode dtype",
591        );
592    }
593}