inferencelayer 0.2.4

//! Moshi's depth transformer ("depformer") on wgpu — the last CPU-resident half of the LM
//! step. The 8 codebook steps are SEQUENTIAL (each consumes the previous step's sampled
//! token), which forbids per-step readbacks: the whole cycle is a STATIC plan in one compute
//! pass with GPU-side token feedback — per-step argmax writes into a 9-slot token buffer
//! (slot 0 = the text token, copied from the temporal argmax; slots 1..=8 = the audio
//! codebooks) and the next step's embedding gather reads it. One submit for temporal +
//! text argmax + depformer; ONE 36-byte readback per 80 ms frame.
//!
//! Weights ride the Q8_0 band (the FFN-fragility result applies here too — the CPU i8
//! depformer was 81/81, so 8-bit is the proven width). Per-STEP weight sets (8× everything;
//! norms shared per layer), NO rope, KV lives only within the frame (slots 0..step are
//! rewritten each frame — never read past the current step, so no clearing).
//!
//! Sampling WITHOUT readbacks: greedy argmax collapses Moshi to silence in live conversation
//! (proven by the moshi-talk A/B — greedy says one opener then goes mute; sampled replies).
//! But the token feedback must stay on-GPU, so instead of top-k we use the Gumbel-max
//! identity `sample(softmax(l/T)) = argmax(l/T + g) = argmax(l + T·g)`: the CPU uploads
//! T-scaled Gumbel noise once per frame (~190 KB, no sync) and a tiny ADD_NOISE dispatch
//! folds it into each head's logits before the UNCHANGED argmax kernels. Zero noise (the
//! buffer's initial state) = greedy, bit-identical to the old path.

use std::path::Path;

use anyhow::{Context, Result};

use crate::GpuCtx;
use crate::forward::{Lfm2Gpu, make_bg, pipeline, uni};
use crate::weights::{LazySt, f32_to_f16_bytes, quantize_q8_0};

const DDIM: usize = 1024;
const DLAYERS: usize = 6;
const STEPS: usize = 8;
const DHID: usize = 2816; // 2/3 · 4.125 · 1024
const CARD: usize = 2048;

const RMSNORM1024: &str = r#"
@group(0) @binding(0) var<storage, read>       x:     array<f32>;
@group(0) @binding(1) var<storage, read>       alpha: array<f32>;
@group(0) @binding(2) var<storage, read_write> y:     array<f32>;
var<workgroup> partial: array<f32, 256>;
@compute @workgroup_size(256)
fn main(@builtin(local_invocation_id) lid: vec3<u32>) {
    let t = lid.x;
    var s = 0.0;
    for (var i = t; i < 1024u; i += 256u) { let v = x[i]; s += v * v; }
    partial[t] = s;
    workgroupBarrier();
    var stride = 128u;
    while (stride > 0u) {
        if (t < stride) { partial[t] += partial[t + stride]; }
        workgroupBarrier();
        stride /= 2u;
    }
    let inv = 1.0 / sqrt(partial[0] / 1024.0 + 1e-8);
    for (var i = t; i < 1024u; i += 256u) { y[i] = x[i] * alpha[i] * inv; }
}
"#;

const ADD_EMB: &str = r#"
@group(0) @binding(0) var<storage, read_write> y:     array<f32>;
@group(0) @binding(1) var<storage, read>       table: array<f32>;
@group(0) @binding(2) var<storage, read>       tok:   array<u32>;
@group(0) @binding(3) var<uniform>             d:     vec4<u32>;   // (slot, _, _, _)
@compute @workgroup_size(256)
fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
    let i = gid.x;
    if (i < 1024u) { y[i] = y[i] + table[tok[d.x] * 1024u + i]; }
}
"#;

const COPY_KV: &str = r#"
@group(0) @binding(0) var<storage, read>       qkv: array<f32>;
@group(0) @binding(1) var<storage, read_write> kc:  array<f32>;
@group(0) @binding(2) var<storage, read_write> vc:  array<f32>;
@group(0) @binding(3) var<uniform>             d:   vec4<u32>;   // (step, _, _, _)
@compute @workgroup_size(256)
fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
    let i = gid.x;
    if (i < 1024u) {
        kc[d.x * 1024u + i] = qkv[1024u + i];
        vc[d.x * 1024u + i] = qkv[2048u + i];
    }
}
"#;

const DEP_ATTN: &str = r#"
@group(0) @binding(0) var<storage, read>       qkv: array<f32>;
@group(0) @binding(1) var<storage, read>       kc:  array<f32>;
@group(0) @binding(2) var<storage, read>       vc:  array<f32>;
@group(0) @binding(3) var<storage, read_write> y:   array<f32>;
@group(0) @binding(4) var<uniform>             d:   vec4<u32>;   // (S = steps so far, _, _, _)
var<workgroup> part: array<f32, 64>;
var<workgroup> sc: array<f32, 8>;
@compute @workgroup_size(64)
fn main(@builtin(workgroup_id) wid: vec3<u32>, @builtin(local_invocation_id) lid: vec3<u32>) {
    let h = wid.x;
    let t = lid.x;
    let n = d.x;
    let ql = qkv[h * 64u + t];
    for (var s = 0u; s < n; s++) {
        part[t] = ql * kc[s * 1024u + h * 64u + t];
        workgroupBarrier();
        var stride = 32u;
        while (stride > 0u) {
            if (t < stride) { part[t] += part[t + stride]; }
            workgroupBarrier();
            stride /= 2u;
        }
        if (t == 0u) { sc[s] = part[0] * 0.125; }
        workgroupBarrier();
    }
    if (t == 0u) {
        var mx = sc[0];
        for (var s = 1u; s < n; s++) { mx = max(mx, sc[s]); }
        var den = 0.0;
        for (var s = 0u; s < n; s++) { sc[s] = exp(sc[s] - mx); den += sc[s]; }
        for (var s = 0u; s < n; s++) { sc[s] /= den; }
    }
    workgroupBarrier();
    var acc = 0.0;
    for (var s = 0u; s < n; s++) { acc += sc[s] * vc[s * 1024u + h * 64u + t]; }
    y[h * 64u + t] = acc;
}
"#;

const ARGMAX2048: &str = r#"
@group(0) @binding(0) var<storage, read>       logits: array<f32>;
@group(0) @binding(1) var<storage, read_write> tok:    array<u32>;
@group(0) @binding(2) var<uniform>             d:      vec4<u32>;   // (out slot, _, _, _)
var<workgroup> bv: array<f32, 256>;
var<workgroup> bi: array<u32, 256>;
@compute @workgroup_size(256)
fn main(@builtin(local_invocation_id) lid: vec3<u32>) {
    let t = lid.x;
    var best = -1e30;
    var idx = 0u;
    for (var i = t; i < 2048u; i += 256u) {
        if (logits[i] > best) { best = logits[i]; idx = i; }
    }
    bv[t] = best;
    bi[t] = idx;
    workgroupBarrier();
    var stride = 128u;
    while (stride > 0u) {
        if (t < stride && bv[t + stride] > bv[t]) {
            bv[t] = bv[t + stride];
            bi[t] = bi[t + stride];
        }
        workgroupBarrier();
        stride /= 2u;
    }
    if (t == 0u) { tok[d.x] = bi[0]; }
}
"#;

/// `logits[i] += noise[d.x + i]` for i < d.y — Gumbel-max sampling folded into the argmax
/// kernels that follow. The noise buffer starts (and in greedy mode stays) all-zero.
const ADD_NOISE: &str = r#"
@group(0) @binding(0) var<storage, read_write> logits: array<f32>;
@group(0) @binding(1) var<storage, read>       noise:  array<f32>;
@group(0) @binding(2) var<uniform>             d:      vec4<u32>;   // (noise offset, n, _, _)
@compute @workgroup_size(256)
fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
    if (gid.x < d.y) { logits[gid.x] += noise[d.x + gid.x]; }
}
"#;

/// Text-head top-k on GPU, stage 1: 256 workgroups scan vocab slices → per-block (max, idx).
/// p = (VOCAB, PER_BLOCK, 0, 0).
const TK_BLOCK: &str = r#"
@group(0) @binding(0) var<storage, read>       logits: array<f32>;
@group(0) @binding(1) var<storage, read_write> pv:     array<f32>;
@group(0) @binding(2) var<storage, read_write> pi:     array<u32>;
@group(0) @binding(3) var<uniform> p: vec4<u32>;
var<workgroup> bv: array<f32, 256>;
var<workgroup> bi: array<u32, 256>;
@compute @workgroup_size(256)
fn main(@builtin(workgroup_id) wid: vec3<u32>, @builtin(local_invocation_id) lid: vec3<u32>) {
    let start = wid.x * p.y;
    var best = -3.0e38;
    var idx = 0u;
    for (var k = lid.x; k < p.y; k += 256u) {
        let i = start + k;
        if (i < p.x && logits[i] > best) { best = logits[i]; idx = i; }
    }
    bv[lid.x] = best;
    bi[lid.x] = idx;
    workgroupBarrier();
    var s = 128u;
    while (s > 0u) {
        if (lid.x < s && bv[lid.x + s] > bv[lid.x]) { bv[lid.x] = bv[lid.x + s]; bi[lid.x] = bi[lid.x + s]; }
        workgroupBarrier();
        s /= 2u;
    }
    if (lid.x == 0u) { pv[wid.x] = bv[0]; pi[wid.x] = bi[0]; }
}
"#;

/// Top-k stage 2 (×K iterations): merge the 256 block winners → record (value, id) into slot
/// `p.x` of the top-k list, MASK the winner in the logits, and note its block for the rescan.
/// p = (ITER, PER_BLOCK, 0, 0).
const TK_MERGE: &str = r#"
@group(0) @binding(0) var<storage, read>       pv:     array<f32>;
@group(0) @binding(1) var<storage, read>       pi:     array<u32>;
@group(0) @binding(2) var<storage, read_write> logits: array<f32>;
@group(0) @binding(3) var<storage, read_write> tkv:    array<f32>;
@group(0) @binding(4) var<storage, read_write> tki:    array<u32>;
@group(0) @binding(5) var<storage, read_write> bhit:   array<u32>;
@group(0) @binding(6) var<uniform> p: vec4<u32>;
var<workgroup> bv: array<f32, 256>;
var<workgroup> bi: array<u32, 256>;
@compute @workgroup_size(256)
fn main(@builtin(local_invocation_id) lid: vec3<u32>) {
    bv[lid.x] = pv[lid.x];
    bi[lid.x] = pi[lid.x];
    workgroupBarrier();
    var s = 128u;
    while (s > 0u) {
        if (lid.x < s && bv[lid.x + s] > bv[lid.x]) { bv[lid.x] = bv[lid.x + s]; bi[lid.x] = bi[lid.x + s]; }
        workgroupBarrier();
        s /= 2u;
    }
    if (lid.x == 0u) {
        tkv[p.x] = bv[0];
        tki[p.x] = bi[0];
        logits[bi[0]] = -3.0e38;
        bhit[0] = bi[0] / p.y;
    }
}
"#;

/// Top-k stage 3 (×K−1): rescan ONLY the block that lost its winner (read from `bhit`).
/// p = (VOCAB, PER_BLOCK, 0, 0).
const TK_RESCAN: &str = r#"
@group(0) @binding(0) var<storage, read>       logits: array<f32>;
@group(0) @binding(1) var<storage, read_write> pv:     array<f32>;
@group(0) @binding(2) var<storage, read_write> pi:     array<u32>;
@group(0) @binding(3) var<storage, read>       bhit:   array<u32>;
@group(0) @binding(4) var<uniform> p: vec4<u32>;
var<workgroup> bv: array<f32, 256>;
var<workgroup> bi: array<u32, 256>;
@compute @workgroup_size(256)
fn main(@builtin(local_invocation_id) lid: vec3<u32>) {
    let b = bhit[0];
    let start = b * p.y;
    var best = -3.0e38;
    var idx = 0u;
    for (var k = lid.x; k < p.y; k += 256u) {
        let i = start + k;
        if (i < p.x && logits[i] > best) { best = logits[i]; idx = i; }
    }
    bv[lid.x] = best;
    bi[lid.x] = idx;
    workgroupBarrier();
    var s = 128u;
    while (s > 0u) {
        if (lid.x < s && bv[lid.x + s] > bv[lid.x]) { bv[lid.x] = bv[lid.x + s]; bi[lid.x] = bi[lid.x + s]; }
        workgroupBarrier();
        s /= 2u;
    }
    if (lid.x == 0u) { pv[b] = bv[0]; pi[b] = bi[0]; }
}
"#;

/// Top-k stage 4: Gumbel-max over the K winners — `argmax(v_j + noise_j)` where the noise is
/// pre-scaled by temp_text (exact top-k temperature sampling) — into `tok[0]`. p = (K,).
const TK_PICK: &str = r#"
@group(0) @binding(0) var<storage, read>       tkv:   array<f32>;
@group(0) @binding(1) var<storage, read>       tki:   array<u32>;
@group(0) @binding(2) var<storage, read>       noise: array<f32>;
@group(0) @binding(3) var<storage, read_write> tok:   array<u32>;
@group(0) @binding(4) var<uniform> p: vec4<u32>;
var<workgroup> bv: array<f32, 32>;
var<workgroup> bi: array<u32, 32>;
@compute @workgroup_size(32)
fn main(@builtin(local_invocation_id) lid: vec3<u32>) {
    var s = -3.0e38;
    if (lid.x < p.x) { s = tkv[lid.x] + noise[lid.x]; }
    bv[lid.x] = s;
    bi[lid.x] = lid.x;
    workgroupBarrier();
    var st = 16u;
    while (st > 0u) {
        if (lid.x < st && bv[lid.x + st] > bv[lid.x]) { bv[lid.x] = bv[lid.x + st]; bi[lid.x] = bi[lid.x + st]; }
        workgroupBarrier();
        st /= 2u;
    }
    if (lid.x == 0u) { tok[0] = tki[bi[0]]; }
}
"#;

struct Q8Buf {
    scales: wgpu::Buffer,
    quants: wgpu::Buffer,
}

struct PlanStep {
    pl: wgpu::ComputePipeline,
    bg: wgpu::BindGroup,
    gx: u32,
}

pub struct MoshiDepGpu {
    /// slot 0 = text token (copied from the temporal argmax); 1..=8 = the audio codebooks.
    tok: wgpu::Buffer,
    plan: Vec<PlanStep>,
    /// Gumbel noise: `[0..text_vocab)` for the temporal text head, then 8×2048 per codebook.
    noise: wgpu::Buffer,
    text_vocab: usize,
    /// Adds the text-head slice of `noise` onto the temporal logits (recorded between the
    /// temporal forward and its argmax).
    text_noise: PlanStep,
    /// On-GPU top-k(25) text sampler over the temporal logits → `tok[0]` (see
    /// [`moshi_full_step_topk`]): faithful `LMGen` text semantics with zero readbacks.
    topk_plan: Vec<PlanStep>,
    tk_noise: wgpu::Buffer,
    _keep: Vec<wgpu::Buffer>,
}

/// The reference text top-k width (`LMGen` top_k_text).
pub const TOPK_TEXT: usize = 25;

/// Build the iterative masked-argmax top-k plan over `logits` (destructive: winners get
/// masked to -3e38 — the buffer is regenerated every frame). Standalone so the unit gate can
/// drive it over a synthetic buffer. Returns (plan, noise25) — noise pre-scaled by the text
/// temperature turns the final argmax into an EXACT top-k temperature sample.
fn build_topk_plan(
    ctx: &GpuCtx,
    logits: &wgpu::Buffer,
    tok: &wgpu::Buffer,
    vocab: usize,
) -> (Vec<PlanStep>, wgpu::Buffer) {
    let per_block = vocab.div_ceil(256);
    let pv = ctx.empty(256);
    let pi = ctx.empty(256);
    let tkv = ctx.empty(TOPK_TEXT);
    let tki = ctx.empty(TOPK_TEXT);
    let bhit = ctx.empty(1);
    let tk_noise = ctx.empty(TOPK_TEXT);
    let blk = pipeline(ctx, "tk_block", TK_BLOCK);
    let mrg = pipeline(ctx, "tk_merge", TK_MERGE);
    let rsc = pipeline(ctx, "tk_rescan", TK_RESCAN);
    let pck = pipeline(ctx, "tk_pick", TK_PICK);
    let u = |a: u32, b: u32| uni(ctx, bytemuck::cast_slice(&[a, b, 0u32, 0]));
    let mut plan = Vec::with_capacity(2 + 2 * TOPK_TEXT);
    let p_scan = u(vocab as u32, per_block as u32);
    plan.push(PlanStep {
        pl: blk.clone(),
        bg: make_bg(ctx, &blk, &[logits, &pv, &pi], &p_scan),
        gx: 256,
    });
    for it in 0..TOPK_TEXT {
        let p = u(it as u32, per_block as u32);
        plan.push(PlanStep {
            pl: mrg.clone(),
            bg: make_bg(ctx, &mrg, &[&pv, &pi, logits, &tkv, &tki, &bhit], &p),
            gx: 1,
        });
        if it + 1 < TOPK_TEXT {
            let p = u(vocab as u32, per_block as u32);
            plan.push(PlanStep {
                pl: rsc.clone(),
                bg: make_bg(ctx, &rsc, &[logits, &pv, &pi, &bhit], &p),
                gx: 1,
            });
        }
    }
    let p = u(TOPK_TEXT as u32, 0);
    plan.push(PlanStep {
        pl: pck.clone(),
        bg: make_bg(ctx, &pck, &[&tkv, &tki, &tk_noise, tok], &p),
        gx: 1,
    });
    (plan, tk_noise)
}

impl MoshiDepGpu {
    pub fn load(ctx: &GpuCtx, dir: impl AsRef<Path>, gpu: &Lfm2Gpu) -> Result<Self> {
        // Subgroup Q8 kernels where available, portable workgroup-tree twins otherwise — the
        // depformer runs on any adapter (the custom kernels below are already subgroup-free).
        let sg = ctx.subgroups;
        let hnorm = gpu.hnorm_buf();
        let st = LazySt::open(dir.as_ref()).context("moshi-lm checkpoint")?;
        let q8 = |name: &str, rows: usize, k: usize| -> Result<Q8Buf> {
            let v = st.tensor_f32(name)?;
            anyhow::ensure!(v.len() == rows * k, "{name} len {} != {rows}x{k}", v.len());
            let (sc, qz) = quantize_q8_0(&v, rows, k);
            Ok(Q8Buf {
                scales: ctx.storage_bytes(&f32_to_f16_bytes(&sc)),
                quants: ctx.storage(bytemuck::cast_slice(&qz)),
            })
        };
        let fbuf = |name: &str| -> Result<wgpu::Buffer> { Ok(ctx.storage(&st.tensor_f32(name)?)) };

        // pipelines
        let rmsnorm = pipeline(ctx, "dep_rmsnorm", RMSNORM1024);
        let add_emb = pipeline(ctx, "dep_add_emb", ADD_EMB);
        let copy_kv = pipeline(ctx, "dep_copy_kv", COPY_KV);
        let attn = pipeline(ctx, "dep_attn", DEP_ATTN);
        let argmax = pipeline(ctx, "dep_argmax", ARGMAX2048);
        let add_noise = pipeline(ctx, "dep_add_noise", ADD_NOISE);
        let (gemv_src, mlp_src) = if sg {
            (crate::gemv_q8_k_sg_src(), crate::mlp_gate_q8_k_sg_src())
        } else {
            (crate::gemv_q8_k_src(), crate::mlp_gate_q8_k_src())
        };
        let gemv = pipeline(ctx, "dep_gemv_q8", &gemv_src);
        let mlp = pipeline(ctx, "dep_mlp_q8", &mlp_src);

        // scratch
        let x = ctx.empty(DDIM);
        let h = ctx.empty(DDIM);
        let qkv = ctx.empty(3 * DDIM);
        let attn_out = ctx.empty(DDIM);
        let gate_mid = ctx.empty(DHID);
        let logits = ctx.empty(CARD);
        let tok = ctx.empty(STEPS + 1);
        let text_vocab = gpu.text_vocab();
        let noise = ctx.empty(text_vocab + STEPS * CARD); // zero-init = greedy
        let kc: Vec<wgpu::Buffer> = (0..DLAYERS).map(|_| ctx.empty(STEPS * DDIM)).collect();
        let vc: Vec<wgpu::Buffer> = (0..DLAYERS).map(|_| ctx.empty(STEPS * DDIM)).collect();

        // shared uniforms (shapes are static)
        let u = |a: u32, b: u32, c: u32, dd: u32| uni(ctx, bytemuck::cast_slice(&[a, b, c, dd]));
        let d_depin = u(DDIM as u32, 4096, 0, 1);
        let d_qkv = u(3 * DDIM as u32, DDIM as u32, 0, 1);
        let d_out = u(DDIM as u32, DDIM as u32, 1, 1); // acc into x
        let d_mlp = u(DHID as u32, DDIM as u32, 0, 1);
        let d_down = u(DDIM as u32, DHID as u32, 1, 1); // acc into x
        let d_head = u(CARD as u32, DDIM as u32, 0, 1);
        let epsm = uni(ctx, bytemuck::cast_slice(&[1e-8f32, 0.0, 0.0, 0.0]));
        let slot_u: Vec<wgpu::Buffer> = (0..=STEPS as u32).map(|s| u(s, 0, 0, 0)).collect();

        // norms (shared across steps) + embedding tables
        let mut norm1 = Vec::with_capacity(DLAYERS);
        let mut norm2 = Vec::with_capacity(DLAYERS);
        for l in 0..DLAYERS {
            norm1.push(fbuf(&format!("depformer.layers.{l}.norm1.alpha"))?);
            norm2.push(fbuf(&format!("depformer.layers.{l}.norm2.alpha"))?);
        }
        let text_emb = fbuf("depformer_text_emb.weight")?;
        let dep_emb: Vec<wgpu::Buffer> = (0..STEPS - 1)
            .map(|k| fbuf(&format!("depformer_emb.{k}.weight")))
            .collect::<Result<_>>()?;

        // two-uniform bind group for the fused MLP (dims @7, epsm @8)
        let bg2 = |pl: &wgpu::ComputePipeline, bufs: &[&wgpu::Buffer], d1: &wgpu::Buffer, d2: &wgpu::Buffer| {
            let mut entries: Vec<wgpu::BindGroupEntry> = bufs
                .iter()
                .enumerate()
                .map(|(i, b)| wgpu::BindGroupEntry { binding: i as u32, resource: b.as_entire_binding() })
                .collect();
            entries.push(wgpu::BindGroupEntry { binding: bufs.len() as u32, resource: d1.as_entire_binding() });
            entries.push(wgpu::BindGroupEntry { binding: bufs.len() as u32 + 1, resource: d2.as_entire_binding() });
            ctx.device.create_bind_group(&wgpu::BindGroupDescriptor {
                label: None,
                layout: &pl.get_bind_group_layout(0),
                entries: &entries,
            })
        };

        // Rows per workgroup matches the selected Q8 GEMV kernel (subgroup 8, portable 16).
        let rpw = if sg { 8u32 } else { 16u32 };
        let gx8 = |m: usize| (m as u32).div_ceil(rpw);
        let mut plan: Vec<PlanStep> = Vec::with_capacity(STEPS * (2 + DLAYERS * 7 + 2));
        for s in 0..STEPS {
            // x = dep_in[s] · hnorm  (+ emb[prev token])
            let w = q8(&format!("depformer_in.{s}.weight"), DDIM, 4096)?;
            plan.push(PlanStep {
                pl: gemv.clone(),
                bg: make_bg(ctx, &gemv, &[&w.scales, &w.quants, hnorm, &x], &d_depin),
                gx: gx8(DDIM),
            });
            let table = if s == 0 { &text_emb } else { &dep_emb[s - 1] };
            plan.push(PlanStep {
                pl: add_emb.clone(),
                bg: make_bg(ctx, &add_emb, &[&x, table, &tok], &slot_u[s]),
                gx: 4,
            });
            for l in 0..DLAYERS {
                let p = format!("depformer.layers.{l}");
                plan.push(PlanStep {
                    pl: rmsnorm.clone(),
                    bg: make_bg(ctx, &rmsnorm, &[&x, &norm1[l]], &h),
                    gx: 1,
                });
                let wq = q8(&format!("{p}.self_attn.in_projs.{s}.weight"), 3 * DDIM, DDIM)?;
                plan.push(PlanStep {
                    pl: gemv.clone(),
                    bg: make_bg(ctx, &gemv, &[&wq.scales, &wq.quants, &h, &qkv], &d_qkv),
                    gx: gx8(3 * DDIM),
                });
                plan.push(PlanStep {
                    pl: copy_kv.clone(),
                    bg: make_bg(ctx, &copy_kv, &[&qkv, &kc[l], &vc[l]], &slot_u[s]),
                    gx: 4,
                });
                plan.push(PlanStep {
                    pl: attn.clone(),
                    bg: make_bg(ctx, &attn, &[&qkv, &kc[l], &vc[l], &attn_out], &slot_u[s + 1]),
                    gx: 16,
                });
                let wo = q8(&format!("{p}.self_attn.out_projs.{s}.weight"), DDIM, DDIM)?;
                plan.push(PlanStep {
                    pl: gemv.clone(),
                    bg: make_bg(ctx, &gemv, &[&wo.scales, &wo.quants, &attn_out, &x], &d_out),
                    gx: gx8(DDIM),
                });
                // fused norm2 + silu-gated MLP: linear_in halves are [gate | mult]
                let gu = st.tensor_f32(&format!("{p}.gating.{s}.linear_in.weight"))?;
                anyhow::ensure!(gu.len() == 2 * DHID * DDIM, "gating.{s} linear_in shape");
                let (s1, z1) = quantize_q8_0(&gu[..DHID * DDIM], DHID, DDIM);
                let (s3, z3) = quantize_q8_0(&gu[DHID * DDIM..], DHID, DDIM);
                let g1 = Q8Buf {
                    scales: ctx.storage_bytes(&f32_to_f16_bytes(&s1)),
                    quants: ctx.storage(bytemuck::cast_slice(&z1)),
                };
                let g3 = Q8Buf {
                    scales: ctx.storage_bytes(&f32_to_f16_bytes(&s3)),
                    quants: ctx.storage(bytemuck::cast_slice(&z3)),
                };
                plan.push(PlanStep {
                    pl: mlp.clone(),
                    bg: bg2(
                        &mlp,
                        &[&g1.scales, &g1.quants, &g3.scales, &g3.quants, &x, &norm2[l], &gate_mid],
                        &d_mlp,
                        &epsm,
                    ),
                    gx: gx8(DHID),
                });
                let wd = q8(&format!("{p}.gating.{s}.linear_out.weight"), DDIM, DHID)?;
                plan.push(PlanStep {
                    pl: gemv.clone(),
                    bg: make_bg(ctx, &gemv, &[&wd.scales, &wd.quants, &gate_mid, &x], &d_down),
                    gx: gx8(DDIM),
                });
            }
            // head + noise + argmax → tok[s+1] (depformer_norms are Identity: the head reads
            // raw x; the noise add turns the argmax into a Gumbel-max sample when non-zero)
            let wh = q8(&format!("linears.{s}.weight"), CARD, DDIM)?;
            plan.push(PlanStep {
                pl: gemv.clone(),
                bg: make_bg(ctx, &gemv, &[&wh.scales, &wh.quants, &x, &logits], &d_head),
                gx: gx8(CARD),
            });
            let noise_u = u((text_vocab + s * CARD) as u32, CARD as u32, 0, 0);
            plan.push(PlanStep {
                pl: add_noise.clone(),
                bg: make_bg(ctx, &add_noise, &[&logits, &noise], &noise_u),
                gx: (CARD as u32).div_ceil(256),
            });
            plan.push(PlanStep {
                pl: argmax.clone(),
                bg: make_bg(ctx, &argmax, &[&logits, &tok], &slot_u[s + 1]),
                gx: 1,
            });
            ctx.queue.submit(std::iter::empty());
            let _ = ctx.device.poll(wgpu::PollType::wait_indefinitely());
        }
        let text_noise_u = u(0, text_vocab as u32, 0, 0);
        let text_noise = PlanStep {
            pl: add_noise.clone(),
            bg: make_bg(ctx, &add_noise, &[gpu.logits_buf(), &noise], &text_noise_u),
            gx: (text_vocab as u32).div_ceil(256),
        };
        let (topk_plan, tk_noise) = build_topk_plan(ctx, gpu.logits_buf(), &tok, text_vocab);
        let keep = vec![x, h, qkv, attn_out, gate_mid, logits];
        Ok(Self { tok, plan, noise, text_vocab, text_noise, topk_plan, tk_noise, _keep: keep })
    }

    /// Fresh Gumbel noise for the 25 top-k text slots, pre-scaled by `temp_text`.
    pub fn upload_topk_noise(&self, ctx: &GpuCtx, rng: &mut u64, temp_text: f32) {
        let mut v = [0f32; TOPK_TEXT];
        let mut x = (*rng).max(1);
        for o in v.iter_mut() {
            x ^= x << 13;
            x ^= x >> 7;
            x ^= x << 17;
            let u = ((x >> 11) as f64 / (1u64 << 53) as f64).max(1e-300);
            *o = -((-u.ln()).ln()) as f32 * temp_text;
        }
        *rng = x;
        ctx.queue.write_buffer(&self.tk_noise, 0, bytemuck::cast_slice(&v));
    }

    fn record(&self, enc: &mut wgpu::CommandEncoder) {
        let mut p = enc.begin_compute_pass(&wgpu::ComputePassDescriptor::default());
        for s in &self.plan {
            p.set_pipeline(&s.pl);
            p.set_bind_group(0, &s.bg, &[]);
            p.dispatch_workgroups(s.gx, 1, 1);
        }
    }

    /// The 9-slot token buffer (slot 0 = text, 1..=8 = audio) — the Mimi-GPU decoder binds it
    /// directly so the depformer's outputs feed the codec without a CPU round-trip.
    pub fn tok_buf(&self) -> &wgpu::Buffer {
        &self.tok
    }

    /// Record the text-head noise add (between the temporal forward and its argmax).
    fn record_text_noise(&self, enc: &mut wgpu::CommandEncoder) {
        let mut p = enc.begin_compute_pass(&wgpu::ComputePassDescriptor::default());
        p.set_pipeline(&self.text_noise.pl);
        p.set_bind_group(0, &self.text_noise.bg, &[]);
        p.dispatch_workgroups(self.text_noise.gx, 1, 1);
    }

    /// Upload fresh T-scaled Gumbel noise for the next step: the on-GPU argmaxes become
    /// exact softmax samples at `temp_text` (text head) / `temp` (audio heads). Seeded
    /// xorshift64 — a session replays. Call once per frame BEFORE [`moshi_full_step`];
    /// never calling it leaves the buffer zero, i.e. greedy.
    pub fn upload_gumbel_noise(&self, ctx: &GpuCtx, rng: &mut u64, temp: f32, temp_text: f32) {
        let n = self.text_vocab + STEPS * CARD;
        let mut v = vec![0f32; n];
        let mut x = (*rng).max(1);
        for (i, o) in v.iter_mut().enumerate() {
            x ^= x << 13;
            x ^= x >> 7;
            x ^= x << 17;
            let u = ((x >> 11) as f64 / (1u64 << 53) as f64).max(1e-300);
            let g = -(-u.ln()).ln() as f32;
            *o = g * if i < self.text_vocab { temp_text } else { temp };
        }
        *rng = x;
        ctx.queue.write_buffer(&self.noise, 0, bytemuck::cast_slice(&v));
    }

    /// Audio-only Gumbel noise (the 8 depformer heads) — for the split step where the TEXT
    /// token is sampled faithfully on the CPU (top-k 25) instead of full-vocab Gumbel.
    pub fn upload_gumbel_noise_audio(&self, ctx: &GpuCtx, rng: &mut u64, temp: f32) {
        let mut v = vec![0f32; STEPS * CARD];
        let mut x = (*rng).max(1);
        for o in v.iter_mut() {
            x ^= x << 13;
            x ^= x >> 7;
            x ^= x << 17;
            let u = ((x >> 11) as f64 / (1u64 << 53) as f64).max(1e-300);
            *o = -((-u.ln()).ln()) as f32 * temp;
        }
        *rng = x;
        ctx.queue
            .write_buffer(&self.noise, (self.text_vocab * 4) as u64, bytemuck::cast_slice(&v));
    }
}

/// The full Moshi LM step on the GPU: temporal plan + text noise + text argmax + the 8-step
/// depformer cycle in ONE submit; the only readback is the 9-token buffer (36 bytes).
/// Greedy by default; sampled when [`MoshiDepGpu::upload_gumbel_noise`] ran this frame.
pub fn moshi_full_step(
    gpu: &Lfm2Gpu,
    dep: &MoshiDepGpu,
    ctx: &GpuCtx,
    emb: &[f32],
    pos: usize,
) -> Result<(u32, [u32; STEPS])> {
    moshi_full_step_after(gpu, dep, ctx, emb, pos, None)
}

/// The production sampled step: temporal forward → ON-GPU top-k(25) text sample (iterative
/// masked argmax + Gumbel-25 pick, ~51 tiny dispatches, zero readbacks) → Gumbel-sampled
/// depformer — ONE submit, one 36-byte sync. Faithful `LMGen` text semantics at fused speed.
/// Call [`MoshiDepGpu::upload_gumbel_noise_audio`] + [`MoshiDepGpu::upload_topk_noise`] first.
pub fn moshi_full_step_topk(
    gpu: &Lfm2Gpu,
    dep: &MoshiDepGpu,
    ctx: &GpuCtx,
    emb: &[f32],
    pos: usize,
    mut after_submit: Option<&mut dyn FnMut()>,
) -> Result<(u32, [u32; STEPS])> {
    let mut enc = ctx
        .device
        .create_command_encoder(&wgpu::CommandEncoderDescriptor::default());
    gpu.record_from_embeds(ctx, &mut enc, emb, pos)?;
    {
        let mut p = enc.begin_compute_pass(&wgpu::ComputePassDescriptor::default());
        for s in &dep.topk_plan {
            p.set_pipeline(&s.pl);
            p.set_bind_group(0, &s.bg, &[]);
            p.dispatch_workgroups(s.gx, 1, 1);
        }
    }
    dep.record(&mut enc);
    ctx.queue.submit([enc.finish()]);
    if let Some(f) = after_submit.as_deref_mut() {
        f();
    }
    let toks = ctx.read_u32(&dep.tok, STEPS + 1)?;
    let mut audio = [0u32; STEPS];
    audio.copy_from_slice(&toks[1..=STEPS]);
    Ok((toks[0], audio))
}

/// [`moshi_full_step`] with a piggyback point: queue more work that consumes the token buffer
/// (the Mimi-GPU decode bound via `bind_lm_tokens`) after the submit, before the one sync.
pub fn moshi_full_step_after(
    gpu: &Lfm2Gpu,
    dep: &MoshiDepGpu,
    ctx: &GpuCtx,
    emb: &[f32],
    pos: usize,
    mut after_submit: Option<&mut dyn FnMut()>,
) -> Result<(u32, [u32; STEPS])> {
    let mut enc = ctx
        .device
        .create_command_encoder(&wgpu::CommandEncoderDescriptor::default());
    gpu.record_from_embeds(ctx, &mut enc, emb, pos)?;
    dep.record_text_noise(&mut enc);
    let chosen = gpu.record_argmax(ctx, &mut enc);
    enc.copy_buffer_to_buffer(chosen, 0, &dep.tok, 0, 4);
    dep.record(&mut enc);
    ctx.queue.submit([enc.finish()]);
    if let Some(f) = after_submit.as_deref_mut() {
        f();
    }
    let toks = ctx.read_u32(&dep.tok, STEPS + 1)?;
    let mut audio = [0u32; STEPS];
    audio.copy_from_slice(&toks[1..=STEPS]);
    Ok((toks[0], audio))
}

/// The split full step: temporal submit → read the text logits → the caller samples the text
/// token on CPU (reference top-k 25 semantics — the inner monologue steers Moshi, and the
/// full-vocab Gumbel tail measurably degrades it: dropped letters, drifting topics) → the
/// token is written into the depformer chain and the audio heads run Gumbel-sampled on-GPU.
/// Costs one extra sync (~1-3 ms) over [`moshi_full_step`].
pub fn moshi_full_step_text_cpu(
    gpu: &Lfm2Gpu,
    dep: &MoshiDepGpu,
    ctx: &GpuCtx,
    emb: &[f32],
    pos: usize,
    pick_text: &mut dyn FnMut(&[f32]) -> u32,
    mut after_submit: Option<&mut dyn FnMut()>,
) -> Result<(u32, [u32; STEPS])> {
    let mut enc = ctx
        .device
        .create_command_encoder(&wgpu::CommandEncoderDescriptor::default());
    gpu.record_from_embeds(ctx, &mut enc, emb, pos)?;
    ctx.queue.submit([enc.finish()]);
    let logits = ctx.read(gpu.logits_buf(), dep.text_vocab)?;
    let text_token = pick_text(&logits);
    ctx.queue.write_buffer(&dep.tok, 0, bytemuck::bytes_of(&text_token));
    let mut enc = ctx
        .device
        .create_command_encoder(&wgpu::CommandEncoderDescriptor::default());
    dep.record(&mut enc);
    ctx.queue.submit([enc.finish()]);
    // Piggyback point: queue more work that CONSUMES the token buffer (the Mimi-GPU decode
    // bound via `bind_lm_tokens`) before the blocking token read — one sync covers both.
    if let Some(f) = after_submit.as_deref_mut() {
        f();
    }
    let toks = ctx.read_u32(&dep.tok, STEPS + 1)?;
    let mut audio = [0u32; STEPS];
    audio.copy_from_slice(&toks[1..=STEPS]);
    Ok((text_token, audio))
}

#[cfg(test)]
mod portable_q8_tests {
    //! The portable (non-subgroup) Q8 GEMV/MLP kernels are the fallback for adapters without
    //! guaranteed subgroups. A portable kernel compiles and runs on subgroup-capable hardware
    //! too, so we validate both the portable and the subgroup kernel against a CPU reference on
    //! the SAME device — the tolerance is FP reduction-order only (both dequantize identically).

    use crate::GpuCtx;
    use crate::forward::{make_bg, pipeline, uni};
    use crate::weights::{f32_to_f16_bytes, quantize_q8_0};

    fn rng(seed: &mut u64) -> f32 {
        *seed ^= *seed << 13;
        *seed ^= *seed >> 7;
        *seed ^= *seed << 17;
        (*seed >> 40) as f32 / (1u32 << 24) as f32 - 0.5
    }

    /// Dequantize a Q8_0 row-block back to f32 (the exact values the kernels see).
    fn deq(scales: &[f32], quants: &[u32], rows: usize, cols: usize) -> Vec<f32> {
        let nblk = cols / 32;
        let mut w = vec![0f32; rows * cols];
        for r in 0..rows {
            for b in 0..nblk {
                let d = scales[r * nblk + b];
                for wi in 0..8 {
                    let word = quants[(r * nblk + b) * 8 + wi];
                    for byte in 0..4 {
                        let q = (((word >> (8 * byte)) & 0xff) as u8) as i8 as f32;
                        w[r * cols + b * 32 + wi * 4 + byte] = q * d;
                    }
                }
            }
        }
        w
    }

    fn maxd(a: &[f32], b: &[f32]) -> f32 {
        a.iter().zip(b).map(|(x, y)| (x - y).abs()).fold(0f32, f32::max)
    }

    #[test]
    fn portable_q8_gemv_matches_cpu_and_subgroup() {
        let Ok(ctx) = GpuCtx::new() else {
            eprintln!("SKIP: no gpu");
            return;
        };
        let (rows, cols) = (64usize, 128usize);
        let mut seed = 0x1234_5678u64;
        let wf: Vec<f32> = (0..rows * cols).map(|_| rng(&mut seed) * 0.05).collect();
        let x: Vec<f32> = (0..cols).map(|_| rng(&mut seed)).collect();
        let (scales, quants) = quantize_q8_0(&wf, rows, cols);
        let wdeq = deq(&scales, &quants, rows, cols);
        let cpu: Vec<f32> = (0..rows)
            .map(|r| (0..cols).map(|c| wdeq[r * cols + c] * x[c]).sum())
            .collect();

        let sbuf = ctx.storage_bytes(&f32_to_f16_bytes(&scales));
        let qbuf = ctx.storage(bytemuck::cast_slice(&quants));
        let xbuf = ctx.storage(&x);
        let dims = uni(&ctx, bytemuck::cast_slice(&[rows as u32, cols as u32, 0u32, 1u32]));

        for (name, src, rpw) in [
            ("portable", crate::gemv_q8_k_src(), 16u32),
            ("subgroup", crate::gemv_q8_k_sg_src(), if ctx.subgroups { 8 } else { 16 }),
        ] {
            let ybuf = ctx.empty(rows);
            let pl = pipeline(&ctx, name, &src);
            let bg = make_bg(&ctx, &pl, &[&sbuf, &qbuf, &xbuf, &ybuf], &dims);
            let mut enc = ctx
                .device
                .create_command_encoder(&wgpu::CommandEncoderDescriptor::default());
            {
                let mut p = enc.begin_compute_pass(&wgpu::ComputePassDescriptor::default());
                p.set_pipeline(&pl);
                p.set_bind_group(0, &bg, &[]);
                p.dispatch_workgroups((rows as u32).div_ceil(rpw), 1, 1);
            }
            ctx.queue.submit([enc.finish()]);
            let got = ctx.read(&ybuf, rows).expect("read y");
            let d = maxd(&got, &cpu);
            println!("q8 gemv {name}: max|Δ| vs cpu {d:.3e}");
            assert!(d < 1e-3, "q8 gemv {name} max|Δ| {d}");
        }
    }

    #[test]
    fn portable_q8_mlp_matches_cpu_and_subgroup() {
        let Ok(ctx) = GpuCtx::new() else {
            eprintln!("SKIP: no gpu");
            return;
        };
        let (im, cols) = (64usize, 128usize); // gate/up rows, hidden
        let eps = 1e-8f32;
        let mut seed = 0x0bad_c0deu64;
        let w1: Vec<f32> = (0..im * cols).map(|_| rng(&mut seed) * 0.05).collect();
        let w3: Vec<f32> = (0..im * cols).map(|_| rng(&mut seed) * 0.05).collect();
        let x: Vec<f32> = (0..cols).map(|_| rng(&mut seed)).collect();
        let wn: Vec<f32> = (0..cols).map(|_| 0.9 + 0.2 * rng(&mut seed)).collect();
        let (s1, z1) = quantize_q8_0(&w1, im, cols);
        let (s3, z3) = quantize_q8_0(&w3, im, cols);
        let (w1d, w3d) = (deq(&s1, &z1, im, cols), deq(&s3, &z3, im, cols));

        // CPU reference: rmsnorm(x)·wn → gate/up → silu(gate)·up.
        let sq: f32 = x.iter().map(|v| v * v).sum();
        let inv = 1.0 / (sq / cols as f32 + eps).sqrt();
        let xn: Vec<f32> = (0..cols).map(|i| x[i] * inv * wn[i]).collect();
        let silu = |g: f32| g / (1.0 + (-g).exp());
        let cpu: Vec<f32> = (0..im)
            .map(|j| {
                let g: f32 = (0..cols).map(|i| w1d[j * cols + i] * xn[i]).sum();
                let u: f32 = (0..cols).map(|i| w3d[j * cols + i] * xn[i]).sum();
                silu(g) * u
            })
            .collect();

        let s1b = ctx.storage_bytes(&f32_to_f16_bytes(&s1));
        let q1b = ctx.storage(bytemuck::cast_slice(&z1));
        let s3b = ctx.storage_bytes(&f32_to_f16_bytes(&s3));
        let q3b = ctx.storage(bytemuck::cast_slice(&z3));
        let xb = ctx.storage(&x);
        let wnb = ctx.storage(&wn);
        let dims = uni(&ctx, bytemuck::cast_slice(&[im as u32, cols as u32, 0u32, 1u32]));
        let epsm = uni(&ctx, bytemuck::cast_slice(&[eps, 0.0f32, 0.0, 0.0]));

        for (name, src, rpw) in [
            ("portable", crate::mlp_gate_q8_k_src(), 16u32),
            ("subgroup", crate::mlp_gate_q8_k_sg_src(), if ctx.subgroups { 8 } else { 16 }),
        ] {
            let yb = ctx.empty(im);
            let pl = pipeline(&ctx, name, &src);
            // two trailing uniforms (dims @7, epsm @8) — make_bg only appends one.
            let bufs = [&s1b, &q1b, &s3b, &q3b, &xb, &wnb, &yb];
            let mut entries: Vec<wgpu::BindGroupEntry> = bufs
                .iter()
                .enumerate()
                .map(|(i, b)| wgpu::BindGroupEntry { binding: i as u32, resource: b.as_entire_binding() })
                .collect();
            entries.push(wgpu::BindGroupEntry { binding: 7, resource: dims.as_entire_binding() });
            entries.push(wgpu::BindGroupEntry { binding: 8, resource: epsm.as_entire_binding() });
            let bg = ctx.device.create_bind_group(&wgpu::BindGroupDescriptor {
                label: None,
                layout: &pl.get_bind_group_layout(0),
                entries: &entries,
            });
            let mut enc = ctx
                .device
                .create_command_encoder(&wgpu::CommandEncoderDescriptor::default());
            {
                let mut p = enc.begin_compute_pass(&wgpu::ComputePassDescriptor::default());
                p.set_pipeline(&pl);
                p.set_bind_group(0, &bg, &[]);
                p.dispatch_workgroups((im as u32).div_ceil(rpw), 1, 1);
            }
            ctx.queue.submit([enc.finish()]);
            let got = ctx.read(&yb, im).expect("read y");
            let d = maxd(&got, &cpu);
            println!("q8 mlp {name}: max|Δ| vs cpu {d:.3e}");
            assert!(d < 1e-3, "q8 mlp {name} max|Δ| {d}");
        }
    }
}

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

    fn xs(v: &mut u64) -> u64 {
        *v ^= *v << 13;
        *v ^= *v >> 7;
        *v ^= *v << 17;
        *v
    }

    /// EXACT gate for the on-GPU top-k text sampler: synthetic logits + known noise → the
    /// pick must equal the CPU reference (true top-25 by value, then argmax(v + noise)).
    #[test]
    fn topk_pick_matches_cpu_exactly() {
        let Ok(ctx) = GpuCtx::new() else {
            eprintln!("SKIP: no gpu");
            return;
        };
        const V: usize = 32001;
        let mut lrng = 0x1234_5678_9abc_def0u64;
        let mut nrng = 0x0fed_cba9_8765_4321u64;
        for round in 0..25 {
            let logits: Vec<f32> = (0..V)
                .map(|_| ((xs(&mut lrng) >> 11) as f64 / (1u64 << 53) as f64 * 16.0 - 8.0) as f32)
                .collect();
            let lbuf = ctx.storage(&logits);
            let tok = ctx.empty(1);
            let (plan, tk_noise) = build_topk_plan(&ctx, &lbuf, &tok, V);
            let mut noise = [0f32; TOPK_TEXT];
            for o in noise.iter_mut() {
                let u = ((xs(&mut nrng) >> 11) as f64 / (1u64 << 53) as f64).max(1e-300);
                *o = -((-u.ln()).ln()) as f32 * 0.7;
            }
            ctx.queue.write_buffer(&tk_noise, 0, bytemuck::cast_slice(&noise));
            let mut enc = ctx
                .device
                .create_command_encoder(&wgpu::CommandEncoderDescriptor::default());
            {
                let mut p = enc.begin_compute_pass(&wgpu::ComputePassDescriptor::default());
                for s in &plan {
                    p.set_pipeline(&s.pl);
                    p.set_bind_group(0, &s.bg, &[]);
                    p.dispatch_workgroups(s.gx, 1, 1);
                }
            }
            ctx.queue.submit([enc.finish()]);
            let got = ctx.read_u32(&tok, 1).unwrap()[0];

            // CPU reference (random f32s: ties have ~zero probability)
            let mut idx: Vec<(f32, u32)> =
                logits.iter().enumerate().map(|(i, &l)| (l, i as u32)).collect();
            idx.select_nth_unstable_by(TOPK_TEXT - 1, |a, b| b.0.total_cmp(&a.0));
            idx.truncate(TOPK_TEXT);
            idx.sort_by(|a, b| b.0.total_cmp(&a.0)); // GPU extraction order: descending value
            let mut best = (f32::NEG_INFINITY, 0u32);
            for (j, (v, id)) in idx.iter().enumerate() {
                let s = v + noise[j];
                if s > best.0 {
                    best = (s, *id);
                }
            }
            assert_eq!(got, best.1, "round {round}: gpu {got} vs cpu {}", best.1);
        }
        eprintln!("25/25 rounds: GPU top-k pick == CPU reference exactly");
    }
}