Skip to main content

offer_decode_flash_arm

fuel_dispatch::decode_flash

Function offer_decode_flash_arm 

Source 
pub fn offer_decode_flash_arm(
    graph: &mut Graph,
    spec: &DecodeFlashSpec,
    cap: FlashArmCapability,
) -> Result<Option<NodeId>>
Expand description

Offer the CUDA flash decode arm for spec on graph.

If flash_decode_admissible holds for spec’s q shape/dtype + config + cap, this builds arm 1 (the flash Op::Fused(FLASH_ATTN, …) node, CUDA-pinned) and records a 2-arm Op::Branch:

  • arm 0 = spec.decomposed_out (the decomposed oracle — the route an unpicked graph, and every non-CUDA build, realizes on);
  • arm 1 = the flash node.

Returns Ok(Some(branch)) when the arm was emitted, Ok(None) when the gate declined (the graph is left with the decomposed route only — byte-identical to today). Surfaces a build-time Err only on a malformed spec that fails finalize_branches’ validation (never a panic).

The flash node reads [q, k, v, (alibi)] and carries k_len in its FusedOpParams::FlashAttn — the flash_attn_dyn → SymEnv → OpParams path the pipelined lowering already resolves (pipelined.rs ~3660).