fuel-graph 0.20.0

// SPDX-License-Identifier: MIT OR Apache-2.0
//! FlashAttn backward — produces dQ, dK, dV given (Q, K, V, dO, `[alibi]`).
//!
//! Three separate FusedOpId variants (FLASH_ATTN_BACKWARD_Q/K/V) share
//! one `FusedOpParams::FlashAttnBackward` payload. Each variant's
//! output shape matches the corresponding forward input's shape:
//!   Q-backward → Q-shape, K-backward → K-shape, V-backward → V-shape.
//!
//! v1 design: each backward variant recomputes the softmax state
//! independently. A single-launch / multi-output variant that shares
//! the recompute pass across all three gradients is a follow-up; the
//! current design wins on Op-IR-additivity (no multi-output enum
//! infrastructure) at the cost of 3× recompute on the math-definition
//! CPU reference. GPU kernels can fuse the recompute when they ship.
//!
//! Pattern recognition is `SubgraphPattern::None` — no autogenerated
//! lowering rule fires this op family from primitives; autograd emits
//! it directly when computing gradients through a FlashAttn forward.
//!
//! ## Increment C, C-T4 — Q/K recipe migration + the V bind-gap
//!
//! The Q and K gradients reference every fused-node input, so their
//! `decompose` bodies migrated from imperative `graph.push` calls onto a
//! declarative [`PatternNode`] recipe re-emitted through
//! [`decompose_via_recipe`] + `tag_to_op`/`emit` — reusing the C-T3
//! `super::flash_attn::recompute_probs_recipe` for the shared softmax-state
//! recompute (now returning `k_rep` + `softcap_tanh` too) and the C-T2
//! nested-fused carrier for the softmax **backward** (`Op::Fused(
//! SOFTMAX_LAST_DIM_BACKWARD)` rides as an `OpTag::Fused` node).
//!
//! The **V** gradient is `dV = Pᵀ·dO`, which does NOT depend on V — so a recipe
//! over the node's `[q, k, v, do, …]` inputs would leave bind 2 (`v`)
//! unreferenced. That is a `NonContiguousBinds` decline from the ONE shared
//! `validate_recipe` invariant (bind `i` ⇔ input `i`, contiguous `[0, n)` —
//! load-bearing for the runtime-fused ingestion identity, deliberately NOT
//! weakened for one op). So the V variant stays the always-correct **imperative
//! primitive form** (unchanged from the pre-C-T4 body): a documented
//! registry-layer expressibility gap, never a crash — the sibling of
//! `flash_attn`'s symbolic-`k_len` gap.

use crate::registry::{
    BackwardKind, FusedOpEntry, FusedOpFamily, FusedOpId, FusedOpParams, FusedOps, PatternMatch,
    SubgraphPattern, decompose_via_recipe,
};
use crate::{Graph, Node, NodeId, Op};
use fuel_ir::{DType, Shape};
use fuel_kernel_seam_types::{OpAttrs, OpTag, PatternNode};

/// Metadata-side registry entry for FlashAttnBackwardQ.
pub fn entry_q() -> FusedOpEntry {
    FusedOpEntry {
        destructive_input: None,
        id: FusedOps::FLASH_ATTN_BACKWARD_Q,
        name: "FlashAttnBackwardQ",
        family: FusedOpFamily::Attention,
        pattern: SubgraphPattern::Callable(no_pattern),
        decompose,
        backward: BackwardKind::NotDifferentiable,
        shape_rule: shape_rule_q,
        dtype_rule,
        output_views: None,
    }
}

/// Metadata-side registry entry for FlashAttnBackwardK.
pub fn entry_k() -> FusedOpEntry {
    FusedOpEntry {
        destructive_input: None,
        id: FusedOps::FLASH_ATTN_BACKWARD_K,
        name: "FlashAttnBackwardK",
        family: FusedOpFamily::Attention,
        pattern: SubgraphPattern::Callable(no_pattern),
        decompose,
        backward: BackwardKind::NotDifferentiable,
        shape_rule: shape_rule_k,
        dtype_rule,
        output_views: None,
    }
}

/// Metadata-side registry entry for FlashAttnBackwardV.
pub fn entry_v() -> FusedOpEntry {
    FusedOpEntry {
        destructive_input: None,
        id: FusedOps::FLASH_ATTN_BACKWARD_V,
        name: "FlashAttnBackwardV",
        family: FusedOpFamily::Attention,
        pattern: SubgraphPattern::Callable(no_pattern),
        decompose,
        backward: BackwardKind::NotDifferentiable,
        shape_rule: shape_rule_v,
        dtype_rule,
        output_views: None,
    }
}

/// Shape rule for dQ: matches input 0 (Q).
fn shape_rule_q(input_shapes: &[Shape], _params: &FusedOpParams) -> Shape {
    debug_assert!(
        input_shapes.len() == 4 || input_shapes.len() == 5,
        "FlashAttnBackwardQ takes 4-5 inputs (q, k, v, do, [alibi])",
    );
    input_shapes[0].clone()
}

/// Shape rule for dK: matches input 1 (K).
fn shape_rule_k(input_shapes: &[Shape], _params: &FusedOpParams) -> Shape {
    debug_assert!(
        input_shapes.len() == 4 || input_shapes.len() == 5,
        "FlashAttnBackwardK takes 4-5 inputs (q, k, v, do, [alibi])",
    );
    input_shapes[1].clone()
}

/// Shape rule for dV: matches input 2 (V).
fn shape_rule_v(input_shapes: &[Shape], _params: &FusedOpParams) -> Shape {
    debug_assert!(
        input_shapes.len() == 4 || input_shapes.len() == 5,
        "FlashAttnBackwardV takes 4-5 inputs (q, k, v, do, [alibi])",
    );
    input_shapes[2].clone()
}

/// Dtype rule: gradients preserve the forward dtype. All inputs share
/// dtype; output matches input 0 (Q).
fn dtype_rule(input_dtypes: &[DType], _params: &FusedOpParams) -> DType {
    debug_assert!(
        input_dtypes.len() == 4 || input_dtypes.len() == 5,
        "FlashAttnBackward* takes 4-5 inputs",
    );
    input_dtypes[0]
}

/// No pattern recognition — autograd emits these ops directly.
fn no_pattern(_g: &Graph, _root: NodeId) -> Option<PatternMatch> {
    None
}

// --- recipe builders (Increment C, C-T4) -------------------------------------
// Tiny constructors for the portable [`PatternNode`] recipe data. Every extent
// and diagonal is a concrete integer at decompose time, so shape-changer nodes
// carry a BAKED absolute `target_shape` and scalar-param ops bake their concrete
// value (the FSCE / flash_attn posture) — the Q/K recipe mirrors the pre-C-T4
// imperative body node-for-node. Reuses the C-T2 nested-fused carrier for the
// softmax backward (mechanism 2a).

fn r_op(op: OpTag, attrs: OpAttrs, operands: Vec<PatternNode>) -> PatternNode {
    PatternNode::Op {
        op,
        attrs,
        operands,
    }
}
fn r_bind(i: u8) -> PatternNode {
    PatternNode::Bind { index: i }
}
/// A baked absolute `target_shape` attr for a shape-changer (Reshape).
fn r_shape(dims: &[usize]) -> OpAttrs {
    OpAttrs {
        target_shape: dims.iter().map(|&d| d as i64).collect(),
        ..OpAttrs::default()
    }
}
/// A baked scalar for a scalar-param op (MulScalar/AddScalar) — NOT an open slot.
fn r_scalar(v: f64) -> OpAttrs {
    OpAttrs {
        scalars: vec![v],
        ..OpAttrs::default()
    }
}
/// An absolute permutation attr.
fn r_perm(p: Vec<u8>) -> OpAttrs {
    OpAttrs {
        perm: p,
        ..OpAttrs::default()
    }
}
/// A single-axis attr (SumDim reduces `attrs.axis`).
fn r_axis(a: i64) -> OpAttrs {
    OpAttrs {
        axis: Some(a),
        ..OpAttrs::default()
    }
}
/// The nested-softmax-**backward** selector attr (mechanism 2a): names the
/// `SoftmaxLastDimBackward` registry entry the C-T2 carrier reconstructs to
/// (the param-less `fid -> params` map already covers it).
fn fused_softmax_backward_attr() -> OpAttrs {
    OpAttrs {
        fused_op: Some("SoftmaxLastDimBackward".to_string()),
        ..OpAttrs::default()
    }
}

/// GQA fold as recipe DATA (mirrors the imperative `fold` closure — the inverse
/// of the forward's head-repeat): sum the `g` repeated heads back from `Hq` to
/// `Hkv` via `Reshape([b, hkv, g, s, d]) → SumDim(2)` (`SumDim` REMOVES the
/// reduced axis, so the result is `[b, hkv, s, d]`; no keepdim). Identity when
/// `Hq == Hkv`. `SumDim`'s output shape is derived by `primitive_shape`; only the
/// `Reshape` carries a baked target.
fn fold_recipe(x: PatternNode, b: usize, hkv: usize, hq: usize, s: usize, d: usize) -> PatternNode {
    if hq == hkv {
        return x;
    }
    let g = hq / hkv;
    let r5 = r_op(OpTag::Reshape, r_shape(&[b, hkv, g, s, d]), vec![x]);
    r_op(OpTag::SumDim, r_axis(2), vec![r5])
}

/// The Q/K backward gradient as portable [`PatternNode`] DATA (Increment C,
/// C-T4) — the structure-preserving migration of the pre-C-T4 imperative
/// `decompose` body for the Q and K variants. `variant` selects dQ vs dK via
/// ordinary Rust control flow (mechanism 1); the shared softmax-state recompute
/// rides [`super::flash_attn::recompute_probs_recipe`], and the softmax backward
/// rides an `OpTag::Fused(SOFTMAX_LAST_DIM_BACKWARD)` carrier (mechanism 2a). All
/// extents / diagonals are baked per call — ZERO open scalar slots.
///
/// Bind space (the fused node's input order): `0 = q`, `1 = k`, `2 = v`,
/// `3 = dO`, `4 = alibi_slopes` (present iff `has_alibi`). Both Q and K reference
/// every bind — Q uses `k_rep` for dQ and `v_rep` for dP; K uses `q` for dK and
/// `v_rep` for dP — so the binds are contiguous `[0, n)` (unlike V, which drops
/// `v` and therefore stays imperative).
///
/// The math (backward is always the top-left `q_pos_offset = 0` / full-K form —
/// the backward never slices K/V):
///
/// ```text
///   P      = recompute softmax( scale·QKᵀ [+softcap][+alibi][+mask] )
///   dP     = dO · V_repᵀ
///   dScaled = softmax_bwd(P, dP)   [softcap: · (1 − tanh²(s/cap))]
///   dScores = scale · dScaled
///   dQ     = dScores · K_rep                              # Q variant
///   dK     = fold( dScoresᵀ · Q )   (Hq → Hkv sum-fold)   # K variant
/// ```
#[allow(clippy::too_many_arguments)]
fn recipe_qk(
    variant: FusedOpId,
    b: usize,
    hq: usize,
    sq: usize,
    d: usize,
    hkv: usize,
    sk: usize,
    scale: f32,
    causal: bool,
    window_l: Option<usize>,
    window_r: Option<usize>,
    softcap: Option<f32>,
    has_alibi: bool,
    dtype: DType,
) -> PatternNode {
    // Shared recompute of P (softmax probs), the repeated K/V, and (softcap) the
    // saved tanh — over q=bind0, k=bind1, v=bind2, alibi=bind4. The backward
    // always attends the full K bottom-left-aligned (`q_pos_offset = 0`).
    let alibi = if has_alibi { Some(r_bind(4)) } else { None };
    let rc = super::flash_attn::recompute_probs_recipe(
        r_bind(0),
        r_bind(1),
        r_bind(2),
        alibi,
        b,
        hq,
        sq,
        sk,
        d,
        hkv,
        scale,
        causal,
        window_l,
        window_r,
        softcap,
        /* q_pos_offset */ 0,
        dtype,
    );

    // dP = dO · V_repᵀ ; dScaled = softmax_bwd(P, dP).
    let vt = r_op(OpTag::Permute, r_perm(vec![0, 1, 3, 2]), vec![rc.v_rep]);
    let dp = r_op(OpTag::MatMul, OpAttrs::default(), vec![r_bind(3), vt]);
    let mut dscaled = r_op(
        OpTag::Fused,
        fused_softmax_backward_attr(),
        vec![rc.probs, dp],
    );

    // softcap backprop: d/ds[cap·tanh(s/cap)] = 1 − tanh²(s/cap). The saved
    // `tanh` (`t`) is the SAME node emit derives inside `probs` (identity-share).
    if let Some(t) = rc.softcap_tanh {
        let t2 = r_op(OpTag::Mul, OpAttrs::default(), vec![t.clone(), t]);
        let neg = r_op(OpTag::MulScalar, r_scalar(-1.0), vec![t2]);
        let omt2 = r_op(OpTag::AddScalar, r_scalar(1.0), vec![neg]);
        dscaled = r_op(OpTag::Mul, OpAttrs::default(), vec![dscaled, omt2]);
    }
    let dscores = r_op(OpTag::MulScalar, r_scalar(scale as f64), vec![dscaled]);

    if variant == FusedOps::FLASH_ATTN_BACKWARD_Q {
        // dQ = dScores · K_rep  →  [b, hq, sq, d]
        return r_op(OpTag::MatMul, OpAttrs::default(), vec![dscores, rc.k_rep]);
    }
    // K variant: dK = dScoresᵀ · Q  →  [b, hq, sk, d]  →  fold → [b, hkv, sk, d]
    let dst = r_op(OpTag::Permute, r_perm(vec![0, 1, 3, 2]), vec![dscores]);
    let dk_rep = r_op(OpTag::MatMul, OpAttrs::default(), vec![dst, r_bind(0)]);
    fold_recipe(dk_rep, b, hkv, hq, sk, d)
}

/// Decompose to the standard SDPA backward. Re-materializing the `[Sq, Sk]`
/// score matrix is exactly what `decompose` is allowed to do (the cost-guided
/// optimizer keeps the fused kernel when present; this is the always-correct
/// primitive form). The single `FusedOpId` (Q/K/V) selects which gradient to
/// return; the recompute of `P` (softmax probs) reuses the forward's shared
/// recompute so it is byte-identical.
///
/// With `P` (probs), `K_rep`, `V_rep` and `dO` (input 3):
///   dV = Pᵀ·dO ; dP = dO·V_repᵀ ; dScaled = softmax_bwd(P, dP)
///   [softcap: dScaled·(1−tanh²)] ; dScores = scale·dScaled
///   dQ = dScores·K_rep ; dK = dScoresᵀ·Q
/// GQA folds dK/dV back from `Hq` to `Hkv` heads (the inverse of the forward's
/// head-repeat). Masked positions need no special-casing: there `P = 0`, so
/// `softmax_bwd` already yields 0.
///
/// Since C-T4 the **Q and K** variants re-emit the [`recipe_qk`] data through the
/// [`decompose_via_recipe`] bridge (structure-preserving: the emitted base map is
/// node-for-node identical to the pre-C-T4 imperative body — see the parity test
/// in `tests`). The **V** variant keeps the imperative form: `dV = Pᵀ·dO` does
/// not reference V, so a recipe would leave bind 2 unreferenced (a
/// `NonContiguousBinds` decline from the shared bind-contiguity invariant). Per
/// G2 `decompose` is total + never-panic: a wrong-params payload or a malformed
/// node returns `id` (the driver's fixpoint signal) before any emission.
fn decompose(graph: &mut Graph, id: NodeId, params: &FusedOpParams) -> NodeId {
    // The FusedOpId distinguishes which gradient (Q/K/V); params are shared.
    let variant = if let Op::Fused(fid, _) = &graph.node(id).op {
        *fid
    } else {
        return id;
    };
    let (scale, causal, window_l, window_r, softcap) = match params {
        FusedOpParams::FlashAttnBackward {
            softmax_scale,
            causal,
            window_size_left,
            window_size_right,
            softcap,
        } => (
            *softmax_scale,
            *causal,
            *window_size_left,
            *window_size_right,
            *softcap,
        ),
        // Wrong params for this id — can't decompose; return self.
        _ => return id,
    };

    // Shape/arity guards (never panic): a well-formed backward node has 4 (q,k,v,
    // do) or 5 (+alibi) 4-D inputs. Malformed → fixpoint self-return.
    let (q_shape, k_shape, n_inputs, dtype) = {
        let n = graph.node(id);
        if n.inputs.len() != 4 && n.inputs.len() != 5 {
            return id;
        }
        (
            graph.node(n.inputs[0]).shape.clone(),
            graph.node(n.inputs[1]).shape.clone(),
            n.inputs.len(),
            n.dtype,
        )
    };
    let q_dims = q_shape.dims();
    let k_dims = k_shape.dims();
    if q_dims.len() != 4 || k_dims.len() != 4 {
        return id;
    }
    let (b, hq, sq, d) = (q_dims[0], q_dims[1], q_dims[2], q_dims[3]);
    let (hkv, sk) = (k_dims[1], k_dims[2]);
    let has_alibi = n_inputs == 5;

    // Q/K migrate to the declarative recipe (they reference every input → binds
    // are contiguous). The nested softmax-backward rides an `OpTag::Fused`
    // carrier; shared recompute + softcap backprop + the GQA fold are all data.
    if variant == FusedOps::FLASH_ATTN_BACKWARD_Q || variant == FusedOps::FLASH_ATTN_BACKWARD_K {
        let recipe_node = recipe_qk(
            variant, b, hq, sq, d, hkv, sk, scale, causal, window_l, window_r, softcap, has_alibi,
            dtype,
        );
        // No open scalar slots (scale / softcap consts / diagonals are baked).
        return decompose_via_recipe(graph, id, &recipe_node, Some(Vec::new()));
    }
    if variant != FusedOps::FLASH_ATTN_BACKWARD_V {
        // Wrong fid for this entry family — decline (fixpoint).
        return id;
    }

    // --- V variant: imperative primitive form (documented bind-gap) -----------
    // `dV = Pᵀ·dO` does NOT reference V, so a recipe over [q,k,v,do,…] would
    // leave bind 2 (`v`) unreferenced — a `NonContiguousBinds` decline from the
    // shared `validate_recipe` invariant (load-bearing for runtime-fused
    // ingestion identity, deliberately not weakened). This is the always-correct
    // primitive form, unchanged from the pre-C-T4 body.
    let g = hq / hkv;
    let (q_id, k_id, v_id, do_id, alibi_id) = {
        let n = graph.node(id);
        let alibi = if n.inputs.len() == 5 {
            Some(n.inputs[4])
        } else {
            None
        };
        (n.inputs[0], n.inputs[1], n.inputs[2], n.inputs[3], alibi)
    };
    let rc = super::flash_attn::recompute_probs(
        graph, q_id, k_id, v_id, alibi_id, b, hq, sq, sk, d, hkv, scale, causal, window_l,
        window_r, softcap, /* q_pos_offset */ 0, dtype,
    );

    // GQA fold: [b, hq, s, d] → [b, hkv, s, d] by summing the `g` repeated heads
    // (the inverse of `repeat_kv_heads`, hkv-major / g-minor).
    let fold = |graph: &mut Graph, x: NodeId, s: usize| -> NodeId {
        if hq == hkv {
            return x;
        }
        let r5 = graph.push(Node {
            op: Op::Reshape(Shape::from_dims(&[b, hkv, g, s, d])),
            inputs: vec![x],
            shape: Shape::from_dims(&[b, hkv, g, s, d]),
            dtype,
        });
        graph.push(Node {
            op: Op::SumDim(2),
            inputs: vec![r5],
            shape: Shape::from_dims(&[b, hkv, s, d]),
            dtype,
        })
    };

    // dV = Pᵀ · dO  →  [b, hq, sk, d]  →  fold → [b, hkv, sk, d]
    let pt = graph.push(Node {
        op: Op::Permute(vec![0, 1, 3, 2]),
        inputs: vec![rc.probs],
        shape: Shape::from_dims(&[b, hq, sk, sq]),
        dtype,
    });
    let dv_rep = graph.push(Node {
        op: Op::MatMul,
        inputs: vec![pt, do_id],
        shape: Shape::from_dims(&[b, hq, sk, d]),
        dtype,
    });
    fold(graph, dv_rep, sk)
}

#[cfg(test)]
mod tests {
    use super::*;

    /// FROZEN copy of the pre-C-T4 imperative `flash_attn_backward::decompose`
    /// body, verbatim (the explicit-`Shape` `graph.push` spelling + the shared
    /// imperative `recompute_probs` + `fold` closure), before C-T4 replaced the
    /// live Q/K bodies with the data recipe. The C-T4 structure-preservation
    /// oracle: the migrated Q/K recipe re-emit must produce a graph structurally
    /// identical to this. The V arm is UNCHANGED between frozen and live (the
    /// documented bind-gap), so V's parity is trivially the same imperative body.
    fn frozen_legacy_decompose(graph: &mut Graph, id: NodeId, params: &FusedOpParams) -> NodeId {
        let variant = if let Op::Fused(fid, _) = &graph.node(id).op {
            *fid
        } else {
            return id;
        };
        let (q_id, k_id, v_id, do_id, alibi_id, q_shape, k_shape, dtype) = {
            let n = graph.node(id);
            let q_shape = graph.node(n.inputs[0]).shape.clone();
            let k_shape = graph.node(n.inputs[1]).shape.clone();
            let alibi = if n.inputs.len() == 5 {
                Some(n.inputs[4])
            } else {
                None
            };
            (
                n.inputs[0],
                n.inputs[1],
                n.inputs[2],
                n.inputs[3],
                alibi,
                q_shape,
                k_shape,
                n.dtype,
            )
        };
        let (scale, causal, window_l, window_r, softcap) = match params {
            FusedOpParams::FlashAttnBackward {
                softmax_scale,
                causal,
                window_size_left,
                window_size_right,
                softcap,
            } => (
                *softmax_scale,
                *causal,
                *window_size_left,
                *window_size_right,
                *softcap,
            ),
            _ => return id,
        };

        let q_dims = q_shape.dims();
        let k_dims = k_shape.dims();
        let (b, hq, sq, d) = (q_dims[0], q_dims[1], q_dims[2], q_dims[3]);
        let (hkv, sk) = (k_dims[1], k_dims[2]);
        let g = hq / hkv;

        let rc = super::super::flash_attn::recompute_probs(
            graph, q_id, k_id, v_id, alibi_id, b, hq, sq, sk, d, hkv, scale, causal, window_l,
            window_r, softcap, /* q_pos_offset */ 0, dtype,
        );
        let scores_shape = rc.scores_shape.clone();

        let fold = |graph: &mut Graph, x: NodeId, s: usize| -> NodeId {
            if hq == hkv {
                return x;
            }
            let r5 = graph.push(Node {
                op: Op::Reshape(Shape::from_dims(&[b, hkv, g, s, d])),
                inputs: vec![x],
                shape: Shape::from_dims(&[b, hkv, g, s, d]),
                dtype,
            });
            graph.push(Node {
                op: Op::SumDim(2),
                inputs: vec![r5],
                shape: Shape::from_dims(&[b, hkv, s, d]),
                dtype,
            })
        };

        if variant == FusedOps::FLASH_ATTN_BACKWARD_V {
            let pt = graph.push(Node {
                op: Op::Permute(vec![0, 1, 3, 2]),
                inputs: vec![rc.probs],
                shape: Shape::from_dims(&[b, hq, sk, sq]),
                dtype,
            });
            let dv_rep = graph.push(Node {
                op: Op::MatMul,
                inputs: vec![pt, do_id],
                shape: Shape::from_dims(&[b, hq, sk, d]),
                dtype,
            });
            return fold(graph, dv_rep, sk);
        }

        let vt = graph.push(Node {
            op: Op::Permute(vec![0, 1, 3, 2]),
            inputs: vec![rc.v_rep],
            shape: Shape::from_dims(&[b, hq, d, sk]),
            dtype,
        });
        let dp = graph.push(Node {
            op: Op::MatMul,
            inputs: vec![do_id, vt],
            shape: scores_shape.clone(),
            dtype,
        });
        let mut dscaled = graph.push(Node {
            op: Op::Fused(
                FusedOps::SOFTMAX_LAST_DIM_BACKWARD,
                FusedOpParams::SoftmaxLastDimBackward,
            ),
            inputs: vec![rc.probs, dp],
            shape: scores_shape.clone(),
            dtype,
        });
        if let Some(t) = rc.softcap_tanh {
            let t2 = graph.push(Node {
                op: Op::Mul,
                inputs: vec![t, t],
                shape: scores_shape.clone(),
                dtype,
            });
            let neg = graph.push(Node {
                op: Op::MulScalar(-1.0),
                inputs: vec![t2],
                shape: scores_shape.clone(),
                dtype,
            });
            let omt2 = graph.push(Node {
                op: Op::AddScalar(1.0),
                inputs: vec![neg],
                shape: scores_shape.clone(),
                dtype,
            });
            dscaled = graph.push(Node {
                op: Op::Mul,
                inputs: vec![dscaled, omt2],
                shape: scores_shape.clone(),
                dtype,
            });
        }
        let dscores = graph.push(Node {
            op: Op::MulScalar(scale as f64),
            inputs: vec![dscaled],
            shape: scores_shape,
            dtype,
        });

        if variant == FusedOps::FLASH_ATTN_BACKWARD_Q {
            return graph.push(Node {
                op: Op::MatMul,
                inputs: vec![dscores, rc.k_rep],
                shape: q_shape,
                dtype,
            });
        }

        let dst = graph.push(Node {
            op: Op::Permute(vec![0, 1, 3, 2]),
            inputs: vec![dscores],
            shape: Shape::from_dims(&[b, hq, sk, sq]),
            dtype,
        });
        let dk_rep = graph.push(Node {
            op: Op::MatMul,
            inputs: vec![dst, q_id],
            shape: Shape::from_dims(&[b, hq, sk, d]),
            dtype,
        });
        fold(graph, dk_rep, sk)
    }

    /// Recursively assert two subgraphs are node-for-node identical (op, shape,
    /// dtype, arity, recursively-equal inputs). A shared leaf (same NodeId — a
    /// bound external input) matches by identity. Shape/dtype-sensitive at EVERY
    /// node, so it catches any structural drift the recipe migration introduces
    /// (a wrong band diagonal, a lost softcap node, a wrong fold axis, a lost
    /// share).
    fn assert_structural_eq(g: &Graph, a: NodeId, b: NodeId) {
        if a == b {
            return;
        }
        let na = g.node(a);
        let nb = g.node(b);
        assert_eq!(na.op, nb.op, "op mismatch: {:?} vs {:?}", na.op, nb.op);
        assert_eq!(
            na.shape, nb.shape,
            "shape mismatch at {:?}: {:?} vs {:?}",
            na.op, na.shape, nb.shape
        );
        assert_eq!(na.dtype, nb.dtype, "dtype mismatch at {:?}", na.op);
        assert_eq!(
            na.inputs.len(),
            nb.inputs.len(),
            "arity mismatch at {:?}",
            na.op
        );
        for (&ia, &ib) in na.inputs.iter().zip(nb.inputs.iter()) {
            assert_structural_eq(g, ia, ib);
        }
    }

    /// A representative FlashAttn-backward config for the parity matrix.
    #[derive(Clone, Copy)]
    struct Cfg {
        hq: usize,
        hkv: usize,
        causal: bool,
        window_l: Option<usize>,
        window_r: Option<usize>,
        softcap: Option<f32>,
        alibi: bool,
    }

    /// Build a fused FLASH_ATTN_BACKWARD_{Q,K,V} node over `q/dO [b,hq,sq,d]`,
    /// `k/v [b,hkv,sk,d]`, optional `alibi [hq]`, all F32 `Op::Const` leaves. The
    /// node's declared shape is the produced gradient's shape (Q→q, K/V→k/v).
    fn backward_node(
        g: &mut Graph,
        variant: FusedOpId,
        b: usize,
        sq: usize,
        sk: usize,
        d: usize,
        cfg: Cfg,
    ) -> NodeId {
        let q_shape = Shape::from_dims(&[b, cfg.hq, sq, d]);
        let kv_shape = Shape::from_dims(&[b, cfg.hkv, sk, d]);
        let q = g.push(Node {
            op: Op::Const,
            inputs: vec![],
            shape: q_shape.clone(),
            dtype: DType::F32,
        });
        let k = g.push(Node {
            op: Op::Const,
            inputs: vec![],
            shape: kv_shape.clone(),
            dtype: DType::F32,
        });
        let v = g.push(Node {
            op: Op::Const,
            inputs: vec![],
            shape: kv_shape.clone(),
            dtype: DType::F32,
        });
        // dO has the forward-output shape (= q's shape).
        let do_ = g.push(Node {
            op: Op::Const,
            inputs: vec![],
            shape: q_shape.clone(),
            dtype: DType::F32,
        });
        let mut inputs = vec![q, k, v, do_];
        if cfg.alibi {
            let al = g.push(Node {
                op: Op::Const,
                inputs: vec![],
                shape: Shape::from_dims(&[cfg.hq]),
                dtype: DType::F32,
            });
            inputs.push(al);
        }
        let out_shape = if variant == FusedOps::FLASH_ATTN_BACKWARD_Q {
            q_shape
        } else {
            kv_shape
        };
        g.push(Node {
            op: Op::Fused(
                variant,
                FusedOpParams::FlashAttnBackward {
                    softmax_scale: 0.5,
                    causal: cfg.causal,
                    window_size_left: cfg.window_l,
                    window_size_right: cfg.window_r,
                    softcap: cfg.softcap,
                },
            ),
            inputs,
            shape: out_shape,
            dtype: DType::F32,
        })
    }

    /// C-T4: across all three variants and a representative config matrix (plain
    /// / causal / windowed / GQA / softcap / alibi), each variant's live
    /// `decompose` is node-for-node identical to the FROZEN pre-C-T4 imperative
    /// body — the structure-preservation contract. Q/K exercise the recipe
    /// re-emit (variant branch + nested softmax-backward `OpTag::Fused` carrier +
    /// softcap `1 − tanh²` + GQA sum-fold); V exercises the retained imperative
    /// path (trivially identical to frozen — the documented bind-gap). Asserted
    /// at F32 (per the plan).
    #[test]
    fn flash_attn_backward_recipe_matches_frozen_legacy_across_variants_and_configs() {
        let (b, sq, sk, d) = (1usize, 3usize, 3usize, 2usize);
        let masks: [(bool, Option<usize>, Option<usize>); 5] = [
            (false, None, None),       // plain
            (true, None, None),        // causal
            (false, Some(2), Some(1)), // both windows
            (false, Some(2), None),    // left-only window
            (true, Some(2), Some(1)),  // causal + windows (kitchen sink)
        ];
        let variants = [
            FusedOps::FLASH_ATTN_BACKWARD_Q,
            FusedOps::FLASH_ATTN_BACKWARD_K,
            FusedOps::FLASH_ATTN_BACKWARD_V,
        ];
        let mut fired = 0usize;
        for variant in variants {
            for (causal, window_l, window_r) in masks {
                for softcap in [None, Some(30.0f32)] {
                    for alibi in [false, true] {
                        for (hq, hkv) in [(2usize, 2usize), (4usize, 2usize)] {
                            let cfg = Cfg {
                                hq,
                                hkv,
                                causal,
                                window_l,
                                window_r,
                                softcap,
                                alibi,
                            };
                            let mut g = Graph::new();
                            let fused = backward_node(&mut g, variant, b, sq, sk, d, cfg);
                            let out_shape = g.node(fused).shape.clone();
                            let params = match &g.node(fused).op {
                                Op::Fused(_, p) => p.clone(),
                                _ => unreachable!(),
                            };

                            let new_root = decompose(&mut g, fused, &params);
                            assert_ne!(
                                new_root, fused,
                                "decompose fires (variant={variant:?}, causal={causal}, \
                                 wl={window_l:?}, wr={window_r:?}, softcap={softcap:?}, \
                                 alibi={alibi}, hq={hq}, hkv={hkv})"
                            );
                            assert_eq!(
                                g.node(new_root).shape,
                                out_shape,
                                "gradient output = shape_rule (variant={variant:?})"
                            );
                            assert_eq!(g.node(new_root).dtype, DType::F32);

                            let legacy_root = frozen_legacy_decompose(&mut g, fused, &params);
                            assert_structural_eq(&g, new_root, legacy_root);
                            fired += 1;
                        }
                    }
                }
            }
        }
        assert_eq!(
            fired,
            3 * 5 * 2 * 2 * 2,
            "every (variant × config) was checked"
        );
    }

    /// The V-variant bind-gap, made explicit: `dV = Pᵀ·dO` does not reference V,
    /// so the V decompose stays the imperative primitive form. It still fires
    /// (produces a primitive graph, not a self-return) and matches the frozen
    /// legacy — this test documents WHY V is not recipe-migrated (a recipe would
    /// drop bind 2 → `NonContiguousBinds`). GQA config so the fold is exercised.
    #[test]
    fn flash_attn_backward_v_stays_imperative_bind_gap() {
        let cfg = Cfg {
            hq: 4,
            hkv: 2,
            causal: true,
            window_l: None,
            window_r: None,
            softcap: Some(30.0),
            alibi: true,
        };
        let mut g = Graph::new();
        let fused = backward_node(&mut g, FusedOps::FLASH_ATTN_BACKWARD_V, 1, 3, 3, 2, cfg);
        let params = match &g.node(fused).op {
            Op::Fused(_, p) => p.clone(),
            _ => unreachable!(),
        };
        let new_root = decompose(&mut g, fused, &params);
        assert_ne!(
            new_root, fused,
            "V decompose fires (imperative primitive form, not a self-return)"
        );
        // dV has V's shape [b, hkv, sk, d].
        assert_eq!(g.node(new_root).shape, Shape::from_dims(&[1, 2, 3, 2]));
        let legacy_root = frozen_legacy_decompose(&mut g, fused, &params);
        assert_structural_eq(&g, new_root, legacy_root);
    }

    /// Totality (G2): a wrong params payload declines to a fixpoint, never a
    /// crash, before any emission — for every variant.
    #[test]
    fn flash_attn_backward_wrong_params_is_a_fixpoint_not_a_crash() {
        for variant in [
            FusedOps::FLASH_ATTN_BACKWARD_Q,
            FusedOps::FLASH_ATTN_BACKWARD_K,
            FusedOps::FLASH_ATTN_BACKWARD_V,
        ] {
            let cfg = Cfg {
                hq: 2,
                hkv: 2,
                causal: false,
                window_l: None,
                window_r: None,
                softcap: None,
                alibi: false,
            };
            let mut g = Graph::new();
            let fused = backward_node(&mut g, variant, 1, 3, 3, 2, cfg);
            let before = g.len();
            let out = decompose(&mut g, fused, &FusedOpParams::Rope);
            assert_eq!(
                out, fused,
                "wrong params => typed decline => fixpoint ({variant:?})"
            );
            assert_eq!(
                g.len(),
                before,
                "declined before any emission ({variant:?})"
            );
        }
    }
}