Skip to main content

Module decode_flash

Module decode_flash 

Source
Expand description

Decode flash-arm emission — the safe, gated Op::Branch that makes the CUDA flash_decoding binding (commit f693da2c) reachable in decode.

§Why this module exists

Fuel dispatch is fail-fast with a key-only, shape-blind executor binding lookup: a registered kernel that returns Err fails the whole realize, and a wrapper cannot soft-decline back to the decomposed base map (the finding recorded in f693da2c). The CUDA (FlashAttn, [f16|bf16;4], Cuda) capacity-K bindings are live-validated, but only reachable safely if a CUDA FlashAttn node is only ever created for shapes the kernel supports: seq_q == 1, head_dim ∈ [1, 128], f16/bf16, no window / ALiBi / softcap. The plan-time infeasible-cost gate (cost::cost_flash_decoding_cuda) is defense in depth for unpinned placement; the guaranteed clean gate is graph-construction-time arm emission — only build the CUDA arm for supported decode shapes. That is this module.

§Constitutional posture (04-optimization / 06-runtime)

  • The optimizer emits/prunes arms; backends never decide. All the strategy — the shape/dtype/config gate, the capability gate, the CUDA pin, the Op::Branch construction — lives here, in the dispatch (optimizer) layer, not in the model. The model layer only supplies the region’s tensor handles and the live k_len it alone knows (the attended prefix = cached_len + seq), which is data, not a strategic choice.
  • The decomposed base map stays arm 0 / the correctness oracle. The Op::Branch’s arm 0 is the existing decomposed region the decode graph already builds (matmul → scale → mask → softmax → matmul over the capacity KV). finalize_branches enforces that the merge reads arm 0, so a finalized-but-unpicked graph — and every CPU/Vulkan build, where no CUDA flash arm is ever emitted — realizes on the decomposed route, byte-identical to today.
  • CUDA-only. The capacity-K prefix (k_len != sk) is a CUDA-only capability: vulkan_dispatch::flash_attn hard-rejects k_len != sk (fail-fast → fails realize) and the CPU flash wrapper would silently change results vs. the decomposed arm. So arm 1 is pinned to BackendId::Cuda and only emitted when a CUDA flash kernel is bound AND a CUDA device is in the current topology.

§What this module does NOT do (documented remaining slice)

Emitting the arm makes the flash kernel offered. Whether it is selected at dispatch is the runtime route picker’s job — and the shipped picker (ranker::route_picker) resolves placement choices (it ranks arms by their (backend, device) keyed on arm 0’s op), not a same-device kernel-variant choice (flash vs. decomposed, both on CUDA). Per 04-optimization, kernel-variant choice is “largely baked at optimize time”; selecting the flash arm on a same-device basis needs either a per-arm-op cost leg in the picker or an optimize-time route bake. That, plus wiring this emitter into fuel-core’s decode builder (with the persistent-decode attended-length symbol), is the remaining integration — see the module’s tests + the session report.

Structs§

DecodeFlashSpec
The decode-attention region the model hands to the optimizer as a flash-arm candidate. Arm 0 — the decomposed region — is already in the graph; this describes how to build arm 1 (the flash kernel) and where to splice the Op::Branch.
FlashArmCapability
The capability + topology gate for the CUDA flash decode arm: a CUDA FlashAttn kernel must be bound and a CUDA device must be present in the current topology. Split out (and injectable) so the gate logic is unit-testable without a CUDA build, while production reads the real process-global registry + topology via Self::production.

Functions§

flash_decode_admissible
Is a decode-attention region admissible for the CUDA flash arm?
offer_decode_flash_arm
Offer the CUDA flash decode arm for spec on graph.