Enum GraphAttn 

pub enum GraphAttn<'a> {
    Full {
        wq: GraphW<'a>,
        wk: GraphW<'a>,
        wv: GraphW<'a>,
        wo: GraphW<'a>,
        q_norm: Option<&'a [f32]>,
        k_norm: Option<&'a [f32]>,
        late_qk_norm: bool,
        bias: Option<(&'a [f32], &'a [f32], &'a [f32])>,
        output_gate: bool,
        cpu_k: &'a [Vec<f32>],
        cpu_v: &'a [Vec<f32>],
        geom: Option<GraphAttnGeom<'a>>,
    },
    Gdn {

        qkv: GraphW<'a>,
        z: GraphW<'a>,
        a: GraphW<'a>,
        b: GraphW<'a>,
        out: GraphW<'a>,
        conv1d: &'a [f32],
        a_log: &'a [f32],
        dt_bias: &'a [f32],
        norm: &'a [f32],
        nv: usize,
        nk: usize,
        dk: usize,
        dv: usize,
        kk: usize,
        cpu_state: &'a [f32],
    
},
    ShortConv {
        inp: GraphW<'a>,
        out: GraphW<'a>,
        taps: &'a [f32],
        kernel: usize,
        cpu_state: &'a [f32],
    },
}

A layer’s token-mixing op: standard attention or a GDN (linear-attention) block. The surrounding norms + SwiGLU FFN are common to both.

Variants

Full

Fields

wq: GraphW<'a>

wk: GraphW<'a>

wv: GraphW<'a>

wo: GraphW<'a>

q_norm: Option<&'a [f32]>

k_norm: Option<&'a [f32]>

late_qk_norm: bool

HunYuan dense: q/k norm after RoPE (rope-kernel flag bit 32).

bias: Option<(&'a [f32], &'a [f32], &'a [f32])>

(bq, bk, bv) attention biases (Qwen2). None ⇒ no bias.

output_gate: bool

Qwen3.5 gated attention: wq emits 2·nh·hd (q||gate per head), the attention output is scaled by sigmoid(gate) before the O projection.

cpu_k: &'a [Vec<f32>]

cpu_v: &'a [Vec<f32>]

geom: Option<GraphAttnGeom<'a>>

This layer’s own attention geometry, when the model’s layers do not share one (MiMo-V2: 4/8 KV heads, 128-wide V under 192-wide heads, sliding windows with learned sinks, two RoPE tables). None = the call-wide (nkv, hd, rd, invf), V as wide as K, full context and a plain softmax — the historical contract, whose kernels and dispatch are untouched.

Gdn

Fields

qkv: GraphW<'a>

z: GraphW<'a>

a: GraphW<'a>

b: GraphW<'a>

out: GraphW<'a>

conv1d: &'a [f32]

a_log: &'a [f32]

dt_bias: &'a [f32]

norm: &'a [f32]

nv: usize

nk: usize

dk: usize

dv: usize

kk: usize

cpu_state: &'a [f32]

CPU recurrent state [ring (kk-1)·cdim | S nv·dk·dv] — seeds the device mirror when prefill ran on the host (o1 collection, CPU fallback): a zero-initialized device state at decode is exactly the “coherent but contextless” garble.

ShortConv

LFM2 gated short convolution: a fused (B, C, x) projection, a depthwise causal conv over a (kernel−1)-deep per-channel ring, C-gating, and an output projection. This mixer is what most of an LFM2 stack is (22 of the 2.6B’s 30 layers), and before it had a graph arm the whole model fell to the per-op path — ~100 submits a token, 22 tok/s on an A100 for a 1.4 GB file.

Fields

inp: GraphW<'a>

[3·hidden, hidden] fused input projection.

out: GraphW<'a>

[hidden, hidden] output projection.

taps: &'a [f32]

[hidden · kernel] depthwise taps, [channel][tap], tap kernel−1 multiplying the current position.

kernel: usize

cpu_state: &'a [f32]

CPU conv ring [channel][kernel−1], slot 0 newest — seeds the device mirror when prefill ran on the host, which for this mixer is always (the batch graph declines it).