def b3_sel_f

bendtorch-gpu@1.3.0.0 · main.bend:1107

import Base

1107# backward 3, dense: cells [0, nw) are dQ's, [nw, nw + nk) dK's, the next nk dV's (S8 as the causal lane's)
1108def b3_sel_f(a: Bool, b: Bool, +e: U32, +nw: U32, +nk: U32, +c: EAC, s: S8) -> S8:
1109  match a b s:
1110    case True{} _ S8{q, k, d, g, p, dq, dk, v}: b3_q(e, q, d, p, dq, dk, v, pv_val_f(e, c, g, k))
1111    case False{} True{} S8{q, k, d, g, p, dq, dk, v}: b3_k((e - nw : U32), k, d, p, dq, dk, v, kv_val_f((e - nw : U32), c, g, q))
1112    case False{} False{} S8{q, k, d, g, p, dq, dk, v}: b3_v(((e - nw : U32) - nk : U32), q, k, g, dq, dk, v, kv_val_f(((e - nw : U32) - nk : U32), c, p, d))