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}