inferencelayer 0.2.4

//! Pocket TTS's codec on wgpu — the decode half (latent → 1920 samples) as ONE static dispatch
//! plan, replayed per frame into a single command encoder.
//!
//! Pocket's Mimi is Moshi's Mimi at different sizes: SEANet 6·5·4 instead of 4·5·6·8, a k32/s16
//! channel-wise upsample instead of k4/s2, two transformer layers instead of eight, and a plain
//! 1×1 convolution where the RVQ used to be. Every *kernel* is therefore already written and
//! parity-gated in [`crate::mimi_gpu`], whose `Builder` this module drives through the borrowed
//! views — so this file is a plan, not a second codec. The graph:
//!
//! ```text
//!   latent[32] ─1×1 conv─▶ [512,1] ─convtr k32/s16 (groups=512)─▶ [512,16] @200 Hz
//!            ─CT2TM─▶ 2-layer transformer (d512, 8H, ctx 250, LayerScale) ─TM2CT─▶
//!            ─SEANet: conv k7 → (ELU, convtr, resblock)×[6,5,4] → ELU, conv k3─▶ 1920 samples
//! ```
//!
//! All streaming state lives on the GPU exactly as in `mimi_gpu`: conv `prev` prefixes inside
//! persistent `xa` buffers shifted by a CARRY dispatch, transposed-conv bias-subtracted overlap
//! tails, and a 250-slot KV ring per layer with a position ring for window masking. Per frame the
//! CPU writes 32 floats and one uniform, submits once, and reads back 1920 samples.
//!
//! Measured on an M4 Max, same run: **4.98 ms/frame here vs 12.84 ms for the threaded CPU codec**
//! — 2.6x, which takes the whole pipeline from ~3.9x to ~6.4x real-time. (An extrapolation from
//! Moshi's Mimi predicted a wash; that was wrong, because Moshi's decode is dominated by its
//! EIGHT transformer layers where Pocket has two, so the dispatch count per frame is far lower
//! here.) On top of the latency win it buys what `mimi_gpu` was written for: a frame time that
//! does not track box load (the CPU path swings several-fold under it), the CPU cores back for
//! concurrent streams, and a path to WebGPU.
//!
//! Gate: `tests/pocket_tts_gpu_parity.rs` — cosine vs the CPU oracle frame by frame over a real
//! utterance, including past the ring wrap. Bitwise equality is NOT expected (GPU summation order,
//! WGSL transcendentals); the standard is the one `mimi_gpu` is held to.

use anyhow::Result;

use crate::GpuCtx;
use crate::mimi_gpu::{
    Builder, CT2TM, ConvTrView, ConvView, INVALID_POS, LayerView, MimiGpu, Step, TM2CT, TrShape,
};
use crate::pocket_tts::{Conv1dS, ConvTr1dS, GenOpts, Linear, MimiDec, PocketTts, TrLayer, Voice};

/// Transformer attention window, in encoder-rate steps (the reference's `context: 250`).
const CONTEXT: usize = 250;

/// A `Linear` is a 1×1 convolution; the packed form is a pure permutation, so `unpack` recovers
/// the `[out, in]` rows the conv kernel wants exactly.
fn linear_rows(l: &Linear) -> Vec<f32> {
    l.packed.unpack()
}

fn conv_view(c: &Conv1dS) -> ConvView<'_> {
    ConvView {
        w: &c.w,
        b: c.b.as_deref(),
        in_c: c.in_c,
        out_c: c.out_c,
        k: c.k,
        // No dilated convolution anywhere in this codec, so the effective kernel is the kernel.
        k_eff: c.k,
        stride: c.stride,
        replicate: c.replicate,
    }
}

fn convtr_view(c: &ConvTr1dS) -> ConvTrView<'_> {
    ConvTrView {
        w: &c.w,
        b: c.b.as_deref(),
        in_c: c.in_c,
        out_c: c.out_c,
        k: c.k,
        stride: c.stride,
        groups: c.groups,
    }
}

/// The per-layer weights, unpacked into owned rows that outlive the borrowed [`LayerView`]s.
struct LayerRows {
    in_proj: Vec<f32>,
    out_proj: Vec<f32>,
    lin1: Vec<f32>,
    lin2: Vec<f32>,
    ones: Vec<f32>,
}

impl LayerRows {
    fn new(l: &TrLayer, d_model: usize) -> Self {
        Self {
            in_proj: linear_rows(&l.attn.in_proj),
            out_proj: linear_rows(&l.attn.out_proj),
            lin1: linear_rows(&l.linear1),
            lin2: linear_rows(&l.linear2),
            // LayerScale is folded into the two output projections by the MV_LS kernel; a layer
            // without it (the FlowLM backbone) scales by ones.
            ones: vec![1.0; d_model],
        }
    }

    fn view<'a>(&'a self, l: &'a TrLayer) -> LayerView<'a> {
        let norm = |n: &'a crate::pocket_tts::LayerNorm| {
            (
                n.w.as_deref().expect("affine LayerNorm"),
                n.b.as_deref().expect("affine LayerNorm"),
            )
        };
        LayerView {
            norm1: norm(&l.norm1),
            norm2: norm(&l.norm2),
            in_proj: &self.in_proj,
            out_proj: &self.out_proj,
            lin1: &self.lin1,
            lin2: &self.lin2,
            ls1: l.ls1.as_deref().unwrap_or(&self.ones),
            ls2: l.ls2.as_deref().unwrap_or(&self.ones),
        }
    }
}

/// Pocket TTS's Mimi decoder, resident on the GPU.
pub struct PocketCodecGpu {
    steps: Vec<Step>,
    latent_in: wgpu::Buffer,
    pcm_out: wgpu::Buffer,
    uni: wgpu::Buffer,
    state: Vec<wgpu::Buffer>,
    pos_rings: Vec<wgpu::Buffer>,
    taps: std::collections::HashMap<String, (wgpu::Buffer, usize)>,
    /// Encoder-rate steps consumed so far (the transformer's rope/ring position).
    pos: u32,
    /// Encoder-rate steps per latent — how far `pos` advances per frame.
    steps_per_latent: u32,
    frame: usize,
    latent_dim: usize,
}

impl PocketCodecGpu {
    /// Build the decode plan from an already-loaded CPU model (which stays the parity oracle).
    pub fn new(ctx: &GpuCtx, tts: &PocketTts) -> Result<Self> {
        let cfg = tts.config();
        let m: &MimiDec = tts.mimi_dec();
        let (latent_dim, dim) = (cfg.latent_dim, cfg.mimi_dim);
        let steps_per_latent = cfg.steps_per_latent();

        // Unpack the transformer weights first: the views borrow from these.
        let rows: Vec<LayerRows> = m.tr.layers.iter().map(|l| LayerRows::new(l, dim)).collect();
        let views: Vec<LayerView<'_>> = rows
            .iter()
            .zip(&m.tr.layers)
            .map(|(r, l)| r.view(l))
            .collect();

        let quant_w = linear_rows(&m.quant_out);
        let quant = ConvView {
            w: &quant_w,
            b: m.quant_out.b.as_deref(),
            in_c: latent_dim,
            out_c: dim,
            k: 1,
            k_eff: 1,
            stride: 1,
            replicate: false,
        };

        let mut b = Builder::new(ctx);
        let latent_in = ctx.empty(latent_dim);
        let uni = crate::forward::uni(ctx, bytemuck::cast_slice(&[0u32, 0, 0, 0]));

        // 1×1 conv: the DummyQuantizer's output projection, 32 → 512.
        let (z, _) = b.conv(quant, &latent_in, false, 1, None, None);
        b.tap("quant", &z, dim);

        // Learnt channel-wise upsample 12.5 → 200 Hz.
        let (y_up, t_up) = b.convtr(convtr_view(&m.upsample), &z, 1);
        assert_eq!(
            t_up, steps_per_latent,
            "upsample must emit one latent's steps"
        );
        b.tap("up", &y_up, dim * t_up);

        // [C, T] → time-major for the transformer, and back.
        let xt = ctx.empty(t_up * dim);
        let p = b.u4(dim as u32, t_up as u32, 0, 0);
        let pl = b.pl("mimi_ct2tm", CT2TM);
        b.step(
            &pl,
            &[&y_up, &xt, &p],
            ((dim * t_up) as u32).div_ceil(256),
            1,
        );

        // The ring must outlive one frame's worth of pushes: all `t_up` positions land before
        // attention runs, so at capacity == window the earliest query would lose keys it is still
        // entitled to. CONTEXT + t_up slots, masked to the CONTEXT-wide causal window.
        b.transformer_shaped(
            &views,
            &xt,
            t_up,
            &uni,
            TrShape {
                dim,
                heads: 8,
                ffn: 2048,
                ring: CONTEXT + t_up,
                window: CONTEXT,
                ln_eps: 1e-5,
            },
        );
        b.tap("tr", &xt, t_up * dim);

        let y_ct = ctx.empty(dim * t_up);
        let p = b.u4(dim as u32, t_up as u32, 0, 0);
        let pl = b.pl("mimi_tm2ct", TM2CT);
        b.step(
            &pl,
            &[&xt, &y_ct, &p],
            ((dim * t_up) as u32).div_ceil(256),
            1,
        );

        // SEANet decoder.
        let (mut y, mut t) = b.conv(conv_view(&m.dec.first), &y_ct, false, t_up, None, None);
        let mut c = m.dec.first.out_c;
        for (up, blocks) in &m.dec.stages {
            b.elu_ip(&y, c * t);
            let (uy, ut) = b.convtr(convtr_view(up), &y, t);
            c = up.out_c;
            t = ut;
            y = uy;
            for rb in blocks {
                y = b.resblock(conv_view(&rb.c1), conv_view(&rb.c2), &y, c, t);
            }
        }
        b.elu_ip(&y, c * t);
        let (pcm_out, t_pcm) = b.conv(conv_view(&m.dec.last), &y, false, t, None, None);
        anyhow::ensure!(
            t_pcm == cfg.frame_size(),
            "decode emitted {t_pcm} samples, expected {}",
            cfg.frame_size()
        );

        let (steps, state, pos_rings, taps) = b.finish();
        // Rings start INVALID (empty window) — a zeroed ring would alias absolute position 0 onto
        // every slot and let garbage rows into the first softmax.
        for r in &pos_rings {
            let inval = vec![INVALID_POS; (r.size() / 4) as usize];
            ctx.queue.write_buffer(r, 0, bytemuck::cast_slice(&inval));
        }
        Ok(Self {
            steps,
            latent_in,
            pcm_out,
            uni,
            state,
            pos_rings,
            taps,
            pos: 0,
            steps_per_latent: steps_per_latent as u32,
            frame: t_pcm,
            latent_dim,
        })
    }

    /// One DE-NORMALIZED latent (`latent · emb_std + emb_mean`, the same input the CPU
    /// `decode_latents` takes) → one frame of 24 kHz samples.
    pub fn decode_frame(&mut self, ctx: &GpuCtx, latent: &[f32]) -> Result<Vec<f32>> {
        anyhow::ensure!(
            latent.len() == self.latent_dim,
            "latent is {} wide, expected {}",
            latent.len(),
            self.latent_dim
        );
        ctx.queue
            .write_buffer(&self.latent_in, 0, bytemuck::cast_slice(latent));
        ctx.queue
            .write_buffer(&self.uni, 0, bytemuck::cast_slice(&[self.pos, 0u32, 0, 0]));
        MimiGpu::run(ctx, &self.steps);
        self.pos += self.steps_per_latent;
        ctx.read(&self.pcm_out, self.frame)
    }

    /// Clear every streaming buffer — conv carries, transposed-conv tails, KV and position rings.
    pub fn reset(&mut self, ctx: &GpuCtx) {
        for b in &self.state {
            let n = (b.size() / 4) as usize;
            ctx.queue
                .write_buffer(b, 0, bytemuck::cast_slice(&vec![0f32; n]));
        }
        for r in &self.pos_rings {
            let inval = vec![INVALID_POS; (r.size() / 4) as usize];
            ctx.queue.write_buffer(r, 0, bytemuck::cast_slice(&inval));
        }
        self.pos = 0;
    }

    /// Read a named intermediate (`quant`, `up`, `tr`) — stage probes for the parity gate.
    pub fn debug_read(&self, ctx: &GpuCtx, name: &str) -> Result<Vec<f32>> {
        let (buf, len) = self
            .taps
            .get(name)
            .ok_or_else(|| anyhow::anyhow!("no tap {name}"))?;
        ctx.read(buf, *len)
    }
}

// =================================================================================================
// the FlowLM backbone on wgpu
// =================================================================================================

/// The autoregressive backbone (6 layers, d1024, 16 heads, unbounded causal attention) as a
/// one-position dispatch plan, plus the final norm and the EOS projection.
///
/// Two differences from the codec's stack, both handled by [`TrShape`]: there is **no LayerScale**
/// (the residual branches scale by ones — folded into the same `MV_LS` kernel so the plan shape is
/// unchanged) and there is **no attention window**, so `window == ring` and the ring is sized to
/// the longest utterance rather than to a 250-step context.
///
/// The prompt is fed one token at a time through this same plan rather than as a batched pass:
/// `backbone_is_chunk_invariant` proves the two are bit-identical on the CPU, and it keeps a
/// single static plan instead of one per prompt length.
pub struct PocketBackboneGpu {
    steps: Vec<Step>,
    x_in: wgpu::Buffer,
    hidden: wgpu::Buffer,
    eos: wgpu::Buffer,
    uni: wgpu::Buffer,
    /// `(k, v)` ring per layer, in plan order.
    kv: Vec<(wgpu::Buffer, wgpu::Buffer)>,
    pos_rings: Vec<wgpu::Buffer>,
    ring: usize,
    pos: u32,
    d_model: usize,
    /// `out_eos` has a bias the matvec kernel does not apply; added on read-back.
    eos_bias: f32,
}

impl PocketBackboneGpu {
    /// `max_positions` bounds one utterance: voice prompt + text tokens + generated frames.
    pub fn new(ctx: &GpuCtx, tts: &PocketTts, max_positions: usize) -> Result<Self> {
        let cfg = tts.config();
        let lm = tts.flow_lm();
        let (dim, heads) = (cfg.d_model, cfg.num_heads);

        let rows: Vec<LayerRows> = lm
            .tr
            .layers
            .iter()
            .map(|l| LayerRows::new(l, dim))
            .collect();
        let views: Vec<LayerView<'_>> = rows
            .iter()
            .zip(&lm.tr.layers)
            .map(|(r, l)| r.view(l))
            .collect();

        let mut b = Builder::new(ctx);
        let x_in = ctx.empty(dim);
        let uni = crate::forward::uni(ctx, bytemuck::cast_slice(&[0u32, 0, 0, 0]));

        let h = b.transformer_shaped(
            &views,
            &x_in,
            1,
            &uni,
            TrShape {
                dim,
                heads,
                ffn: cfg.ffn_dim,
                ring: max_positions,
                // Unbounded: the FlowLM attends its whole history, so nothing is masked out.
                window: max_positions,
                ln_eps: 1e-5,
            },
        );

        // out_norm, then the EOS logit (bias applied on the host).
        let nw = ctx.storage(lm.out_norm.w.as_deref().expect("affine out_norm"));
        let nb = ctx.storage(lm.out_norm.b.as_deref().expect("affine out_norm"));
        let ln = b.pl(&format!("ln_{dim}"), &crate::mimi_gpu::ln_src(dim, 1e-5));
        b.step(&ln, &[&x_in, &nw, &nb, &h], 1, 1);

        let eos_w = ctx.storage(&linear_rows(&lm.out_eos));
        let eos = ctx.empty(1);
        let p = b.u4(1, dim as u32, 0, 0);
        let mv = b.pl("mimi_matvec", crate::mimi_gpu::MATVEC);
        b.step(&mv, &[&eos_w, &h, &eos, &p], 1, 1);

        let (steps, state, pos_rings, _) = b.finish();
        anyhow::ensure!(
            state.len() == 2 * views.len(),
            "expected one (k, v) ring per layer, got {} buffers",
            state.len()
        );
        let kv: Vec<(wgpu::Buffer, wgpu::Buffer)> = state
            .chunks_exact(2)
            .map(|c| (c[0].clone(), c[1].clone()))
            .collect();
        for r in &pos_rings {
            let inval = vec![INVALID_POS; (r.size() / 4) as usize];
            ctx.queue.write_buffer(r, 0, bytemuck::cast_slice(&inval));
        }
        Ok(Self {
            steps,
            x_in,
            hidden: h,
            eos,
            uni,
            kv,
            pos_rings,
            ring: max_positions,
            pos: 0,
            d_model: dim,
            eos_bias: lm.out_eos.b.as_ref().map_or(0.0, |b| b[0]),
        })
    }

    /// Adopt a voice: upload its per-layer K/V straight into the rings. The CPU cache is
    /// `[positions][heads][head_dim]`, which is exactly the ring's slot-major layout, and the
    /// stored keys already carry rope — so this is a copy, not a replay.
    pub fn load_voice(&mut self, ctx: &GpuCtx, voice: &Voice) -> Result<()> {
        let layers = voice.state.layers();
        anyhow::ensure!(
            layers.len() == self.kv.len(),
            "voice has {} layers, plan has {}",
            layers.len(),
            self.kv.len()
        );
        let mut offset = 0usize;
        for (kv, (kb, vb)) in layers.iter().zip(&self.kv) {
            let (k, v, off) = kv.rows();
            anyhow::ensure!(
                off <= self.ring,
                "voice needs {off} positions, ring holds {}",
                self.ring
            );
            ctx.queue
                .write_buffer(kb, 0, bytemuck::cast_slice(&k[..off * self.d_model]));
            ctx.queue
                .write_buffer(vb, 0, bytemuck::cast_slice(&v[..off * self.d_model]));
            offset = off;
        }
        let pos: Vec<u32> = (0..offset as u32).collect();
        for r in &self.pos_rings {
            ctx.queue.write_buffer(r, 0, bytemuck::cast_slice(&pos));
        }
        self.pos = offset as u32;
        Ok(())
    }

    /// One position through the stack. `x` is the input embedding — a text row, or
    /// `input_linear · latent` for a generation step. Returns `(normalized hidden, EOS logit)`.
    pub fn step(&mut self, ctx: &GpuCtx, x: &[f32]) -> Result<(Vec<f32>, f32)> {
        anyhow::ensure!(x.len() == self.d_model, "input is {} wide", x.len());
        anyhow::ensure!(
            (self.pos as usize) < self.ring,
            "utterance exceeded the {}-position ring; rebuild with a larger max_positions",
            self.ring
        );
        ctx.queue
            .write_buffer(&self.x_in, 0, bytemuck::cast_slice(x));
        ctx.queue
            .write_buffer(&self.uni, 0, bytemuck::cast_slice(&[self.pos, 0u32, 0, 0]));
        MimiGpu::run(ctx, &self.steps);
        self.pos += 1;
        let hidden = ctx.read(&self.hidden, self.d_model)?;
        let eos = ctx.read(&self.eos, 1)?[0] + self.eos_bias;
        Ok((hidden, eos))
    }

    /// Absolute position of the next step.
    pub fn position(&self) -> u32 {
        self.pos
    }

    /// The normalized hidden state, on the GPU — bind this as the flow head's conditioning to
    /// keep it device-side.
    pub fn hidden_buffer(&self) -> &wgpu::Buffer {
        &self.hidden
    }

    /// Like [`Self::step`] but WITHOUT reading the 1024-wide hidden state back: it stays in
    /// `hidden_buffer` for a GPU flow head. Only the EOS scalar crosses the bus.
    pub fn step_device(&mut self, ctx: &GpuCtx, x: &[f32]) -> Result<f32> {
        anyhow::ensure!(x.len() == self.d_model, "input is {} wide", x.len());
        anyhow::ensure!(
            (self.pos as usize) < self.ring,
            "utterance exceeded the {}-position ring",
            self.ring
        );
        ctx.queue
            .write_buffer(&self.x_in, 0, bytemuck::cast_slice(x));
        ctx.queue
            .write_buffer(&self.uni, 0, bytemuck::cast_slice(&[self.pos, 0u32, 0, 0]));
        MimiGpu::run(ctx, &self.steps);
        self.pos += 1;
        Ok(ctx.read(&self.eos, 1)?[0] + self.eos_bias)
    }

    pub fn reset(&mut self, ctx: &GpuCtx) {
        for (k, v) in &self.kv {
            for b in [k, v] {
                let n = (b.size() / 4) as usize;
                ctx.queue
                    .write_buffer(b, 0, bytemuck::cast_slice(&vec![0f32; n]));
            }
        }
        for r in &self.pos_rings {
            let inval = vec![INVALID_POS; (r.size() / 4) as usize];
            ctx.queue.write_buffer(r, 0, bytemuck::cast_slice(&inval));
        }
        self.pos = 0;
    }
}

// =================================================================================================
// the whole pipeline
// =================================================================================================

/// The codec on the GPU, and optionally the backbone too.
///
/// **Which configuration is faster depends on the machine, and the answer here was not the
/// obvious one.** Measured on an M4 Max over one utterance, same run:
///
/// | | ms/frame | real-time |
/// |---|---|---|
/// | all-CPU | 12.9 | 6.2x |
/// | **CPU backbone + GPU codec** ([`Self::codec_only`]) | **~9** | **~8x** |
/// | all-GPU ([`Self::new`]) | 12.9 | 5.8x |
///
/// The codec alone is 1.8x faster on the GPU, but moving the BACKBONE there gives it all back:
/// generation is m=1 autoregressive, so each step is a submit plus a blocking read of the hidden
/// state, and that round trip costs about what the 16 CPU cores take to do the arithmetic. This is
/// the same asymmetry `whisper_gpu` documents for its decoder.
///
/// So: prefer [`Self::codec_only`] on a machine with a real CPU. The all-GPU path exists for
/// WebGPU/wasm, where there are no threads to fall back on and the comparison inverts.
///
/// In the all-GPU configuration the flow head runs device-side too, reading the backbone's hidden
/// state in place — so a generation step reads back only the EOS scalar and the 32-wide latent,
/// never the 1024-wide hidden. All that remains host-side are the embedding lookup and
/// `input_linear · latent`, which are a memcpy and a 32×1024 GEMV.
pub struct PocketGpu {
    backbone: Option<PocketBackboneGpu>,
    /// Present exactly when `backbone` is: with both device-side, a generation step reads back
    /// only the EOS scalar and the 32-wide latent instead of the 1024-wide hidden state.
    flow: Option<PocketFlowGpu>,
    codec: PocketCodecGpu,
    /// The CPU backbone's state, in the hybrid configuration.
    cpu_session: Option<crate::pocket_tts::Session>,
}

impl PocketGpu {
    /// Everything on the GPU — the browser configuration.
    pub fn new(ctx: &GpuCtx, tts: &PocketTts, max_positions: usize) -> Result<Self> {
        let backbone = PocketBackboneGpu::new(ctx, tts, max_positions)?;
        let flow = PocketFlowGpu::new(ctx, tts, backbone.hidden_buffer())?;
        Ok(Self {
            backbone: Some(backbone),
            flow: Some(flow),
            codec: PocketCodecGpu::new(ctx, tts)?,
            cpu_session: None,
        })
    }

    /// Codec on the GPU, backbone on the CPU — the fastest configuration on a desktop, and the
    /// one that also keeps the codec's frame time immune to box load.
    pub fn codec_only(ctx: &GpuCtx, tts: &PocketTts) -> Result<Self> {
        Ok(Self {
            backbone: None,
            flow: None,
            codec: PocketCodecGpu::new(ctx, tts)?,
            cpu_session: None,
        })
    }

    /// Text → 24 kHz samples. The loop itself lives in [`PocketTts::generate_with`]; this type
    /// only supplies frames, so the GPU path cannot drift from the CPU one on chunking, the EOS
    /// rule or the seed stream.
    pub fn generate(
        &mut self,
        ctx: &GpuCtx,
        tts: &PocketTts,
        voice: &Voice,
        text: &str,
        opts: &GenOpts,
    ) -> Result<Vec<f32>> {
        let mut pcm = Vec::new();
        // The frames borrow `ctx` for the duration of the call.
        let mut src = GpuFrames { gpu: self, ctx };
        tts.generate_with(&mut src, voice, text, opts, |f| pcm.extend_from_slice(f))?;
        Ok(pcm)
    }
}

/// [`PocketGpu`] bound to a device for one `generate` call.
struct GpuFrames<'a> {
    gpu: &'a mut PocketGpu,
    ctx: &'a GpuCtx,
}

impl crate::pocket_tts::FrameSource for GpuFrames<'_> {
    fn start(&mut self, tts: &PocketTts, voice: &Voice, ids: &[u32]) -> Result<()> {
        self.gpu.codec.reset(self.ctx);
        match self.gpu.backbone.as_mut() {
            Some(bb) => {
                bb.reset(self.ctx);
                bb.load_voice(self.ctx, voice)?;
                for id in ids {
                    bb.step_device(self.ctx, tts.text_embedding(*id))?;
                }
            }
            None => {
                let mut s = tts.session(voice, 0);
                tts.prompt_tokens(&mut s, ids);
                self.gpu.cpu_session = Some(s);
            }
        }
        Ok(())
    }

    fn next(
        &mut self,
        tts: &PocketTts,
        prev: Option<&[f32]>,
        noise: &[f32],
        opts: &GenOpts,
    ) -> Result<crate::pocket_tts::Step> {
        let (latent, eos_logit) = match (self.gpu.backbone.as_mut(), self.gpu.flow.as_mut()) {
            (Some(bb), Some(fl)) => {
                let eos = bb.step_device(self.ctx, &tts.latent_input(prev))?;
                (fl.step(self.ctx, noise)?, eos)
            }
            _ => {
                let session = self
                    .gpu
                    .cpu_session
                    .as_mut()
                    .ok_or_else(|| anyhow::anyhow!("start() was not called"))?;
                let (hidden, eos) = tts.backbone_step(session, prev);
                (
                    tts.sample_latent_with_noise(&hidden, noise, opts.lsd_steps),
                    eos,
                )
            }
        };
        let pcm = self
            .gpu
            .codec
            .decode_frame(self.ctx, &tts.denormalize(&latent))?;
        Ok(crate::pocket_tts::Step {
            latent,
            eos_logit,
            pcm,
        })
    }
}

// =================================================================================================
// the flow head on wgpu
// =================================================================================================

/// `y[n] = b[n] + Σ w[n,k]·x[k]`, optionally through SiLU. `p = (n, k, act, 0)`.
const MV_BIAS: &str = r#"
@group(0) @binding(0) var<storage, read>       w: array<f32>;
@group(0) @binding(1) var<storage, read>       x: array<f32>;
@group(0) @binding(2) var<storage, read>       b: array<f32>;
@group(0) @binding(3) var<storage, read_write> y: array<f32>;
@group(0) @binding(4) var<uniform> p: vec4<u32>;
@compute @workgroup_size(64)
fn main(@builtin(global_invocation_id) g: vec3<u32>) {
    let n = g.x;
    if (n >= p.x) { return; }
    var acc = b[n];
    let wo = n * p.y;
    for (var k = 0u; k < p.y; k++) { acc += w[wo + k] * x[k]; }
    if (p.z == 1u) { acc = acc / (1.0 + exp(-acc)); }
    y[n] = acc;
}
"#;

/// `y[i] = silu(a[i] + c[i])` — the adaLN conditioning vector: a constant timestep embedding plus
/// the projected backbone hidden state, then the SiLU every modulation head starts with.
const ADD_SILU: &str = r#"
@group(0) @binding(0) var<storage, read>       a: array<f32>;
@group(0) @binding(1) var<storage, read>       c: array<f32>;
@group(0) @binding(2) var<storage, read_write> y: array<f32>;
@group(0) @binding(3) var<uniform> p: vec4<u32>;
@compute @workgroup_size(64)
fn main(@builtin(global_invocation_id) g: vec3<u32>) {
    if (g.x >= p.x) { return; }
    let v = a[g.x] + c[g.x];
    y[g.x] = v / (1.0 + exp(-v));
}
"#;

/// `y[i] = x[i]·(1 + m[scale_off + i]) + m[shift_off + i]` — DiT's `modulate`. `p = (n, shift,
/// scale, 0)` as OFFSETS into the fused adaLN projection.
const MODULATE: &str = r#"
@group(0) @binding(0) var<storage, read>       x: array<f32>;
@group(0) @binding(1) var<storage, read>       m: array<f32>;
@group(0) @binding(2) var<storage, read_write> y: array<f32>;
@group(0) @binding(3) var<uniform> p: vec4<u32>;
@compute @workgroup_size(64)
fn main(@builtin(global_invocation_id) g: vec3<u32>) {
    if (g.x >= p.x) { return; }
    y[g.x] = x[g.x] * (1.0 + m[p.z + g.x]) + m[p.y + g.x];
}
"#;

/// `h[i] += m[gate_off + i] · u[i]` — the adaLN-Zero residual. `p = (n, gate_off, 0, 0)`.
const GATED_ADD: &str = r#"
@group(0) @binding(0) var<storage, read>       u: array<f32>;
@group(0) @binding(1) var<storage, read>       m: array<f32>;
@group(0) @binding(2) var<storage, read_write> h: array<f32>;
@group(0) @binding(3) var<uniform> p: vec4<u32>;
@compute @workgroup_size(64)
fn main(@builtin(global_invocation_id) g: vec3<u32>) {
    if (g.x >= p.x) { return; }
    h[g.x] += m[p.y + g.x] * u[g.x];
}
"#;

/// `x1[i] = x0[i] + dir[i] / steps` — the LSD update. `p = (n, 0, 0, 0)`, `q.x = 1/steps`.
const LSD_STEP: &str = r#"
@group(0) @binding(0) var<storage, read>       dir: array<f32>;
@group(0) @binding(1) var<storage, read_write> x:   array<f32>;
@group(0) @binding(2) var<uniform> p: vec4<u32>;
@group(0) @binding(3) var<uniform> q: vec4<f32>;
@compute @workgroup_size(64)
fn main(@builtin(global_invocation_id) g: vec3<u32>) {
    if (g.x >= p.x) { return; }
    x[g.x] += dir[g.x] * q.x;
}
"#;

/// `SimpleMLPAdaLN` as a dispatch plan, reading the backbone's hidden state **in place** on the
/// GPU — which is the point: with the flow head here, a generation step no longer reads 1024
/// floats back to the host, only the 32-wide latent and the EOS scalar.
///
/// Only `lsd_steps == 1` is planned (the shipped default). More steps would need the ODE loop
/// unrolled into the plan; the CPU path handles them.
///
/// With `s`/`t` fixed by the step count, both TimestepEmbedders are CONSTANT — they are evaluated
/// once at build time and uploaded, so neither their MLPs nor their unbiased-variance RMSNorm ever
/// run per frame.
pub struct PocketFlowGpu {
    steps: Vec<Step>,
    noise_in: wgpu::Buffer,
    latent_out: wgpu::Buffer,
    latent_dim: usize,
}

impl PocketFlowGpu {
    /// `cond` is the backbone's normalized hidden buffer — bound, not copied.
    pub fn new(ctx: &GpuCtx, tts: &PocketTts, cond: &wgpu::Buffer) -> Result<Self> {
        let cfg = tts.config();
        let f = &tts.flow_lm().flow;
        let (fd, ld, dm) = (f.dim, cfg.latent_dim, cfg.d_model);
        let mut b = Builder::new(ctx);

        let mv = b.pl("pt_mv_bias", MV_BIAS);
        let ln = b.pl(&format!("ln_{fd}_1e-6"), &crate::mimi_gpu::ln_src(fd, 1e-6));
        let modu = b.pl("pt_modulate", MODULATE);
        let gate = b.pl("pt_gated_add", GATED_ADD);

        let noise_in = ctx.empty(ld);
        // x = input_proj(noise); the LSD update lands back on the noise buffer at the end.
        let x = ctx.empty(fd);
        // Weights + bias as GPU buffers (the matvec kernel always reads a bias; biasless layers
        // would get zeros, though every layer in this head has one).
        let lin = |l: &Linear| {
            (
                ctx.storage(&linear_rows(l)),
                ctx.storage(&l.b.clone().unwrap_or_else(|| vec![0f32; l.out])),
            )
        };
        let (ipw, ipb) = lin(&f.input_proj);
        let p = b.u4(fd as u32, ld as u32, 0, 0);
        b.step(
            &mv,
            &[&ipw, &noise_in, &ipb, &x, &p],
            (fd as u32).div_ceil(64),
            1,
        );

        // y = silu(time_const + cond_embed(hidden)) — the conditioning every block modulates by.
        let (cw, cb) = lin(&f.cond);
        let cond_p = ctx.empty(fd);
        let p = b.u4(fd as u32, dm as u32, 0, 0);
        b.step(
            &mv,
            &[&cw, cond, &cb, &cond_p, &p],
            (fd as u32).div_ceil(64),
            1,
        );
        let tconst = ctx.storage(&f.time_constant(0.0, 1.0));
        let y = ctx.empty(fd);
        let p = b.u4(fd as u32, 0, 0, 0);
        let addsilu = b.pl("pt_add_silu", ADD_SILU);
        b.step(
            &addsilu,
            &[&tconst, &cond_p, &y, &p],
            (fd as u32).div_ceil(64),
            1,
        );

        let ones = ctx.storage(&vec![1f32; fd]);
        let zeros = ctx.storage(&vec![0f32; fd]);
        let normed = ctx.empty(fd);
        let modulated = ctx.empty(fd);
        let inner = ctx.empty(fd);
        let upd = ctx.empty(fd);
        for blk in &f.blocks {
            let (aw, ab) = lin(&blk.ada);
            let m = ctx.empty(3 * fd);
            let p = b.u4(3 * fd as u32, fd as u32, 0, 0);
            b.step(
                &mv,
                &[&aw, &y, &ab, &m, &p],
                (3 * fd as u32).div_ceil(64),
                1,
            );

            let nw = ctx.storage(blk.in_ln.w.as_deref().expect("affine in_ln"));
            let nb = ctx.storage(blk.in_ln.b.as_deref().expect("affine in_ln"));
            b.step(&ln, &[&x, &nw, &nb, &normed], 1, 1);
            // fused adaLN projection is [shift | scale | gate]
            let p = b.u4(fd as u32, 0, fd as u32, 0);
            b.step(
                &modu,
                &[&normed, &m, &modulated, &p],
                (fd as u32).div_ceil(64),
                1,
            );

            let (w0, b0) = lin(&blk.m0);
            let p = b.u4(fd as u32, fd as u32, 1, 0); // act = SiLU
            b.step(
                &mv,
                &[&w0, &modulated, &b0, &inner, &p],
                (fd as u32).div_ceil(64),
                1,
            );
            let (w2, b2) = lin(&blk.m2);
            let p = b.u4(fd as u32, fd as u32, 0, 0);
            b.step(
                &mv,
                &[&w2, &inner, &b2, &upd, &p],
                (fd as u32).div_ceil(64),
                1,
            );
            let p = b.u4(fd as u32, 2 * fd as u32, 0, 0); // gate is the third chunk
            b.step(&gate, &[&upd, &m, &x, &p], (fd as u32).div_ceil(64), 1);
        }

        // final layer: [shift | scale] over a NON-affine LayerNorm, then down to the latent
        let (faw, fab) = lin(&f.final_layer.ada);
        let fm = ctx.empty(2 * fd);
        let p = b.u4(2 * fd as u32, fd as u32, 0, 0);
        b.step(
            &mv,
            &[&faw, &y, &fab, &fm, &p],
            (2 * fd as u32).div_ceil(64),
            1,
        );
        b.step(&ln, &[&x, &ones, &zeros, &normed], 1, 1);
        let p = b.u4(fd as u32, 0, fd as u32, 0);
        b.step(
            &modu,
            &[&normed, &fm, &modulated, &p],
            (fd as u32).div_ceil(64),
            1,
        );
        let (flw, flb) = lin(&f.final_layer.lin);
        let dir = ctx.empty(ld);
        let p = b.u4(ld as u32, fd as u32, 0, 0);
        b.step(
            &mv,
            &[&flw, &modulated, &flb, &dir, &p],
            (ld as u32).div_ceil(64),
            1,
        );

        // x1 = x0 + dir / steps, in place on the noise buffer (steps == 1).
        let p = b.u4(ld as u32, 0, 0, 0);
        let q = crate::forward::uni(ctx, bytemuck::cast_slice(&[1f32, 0.0, 0.0, 0.0]));
        let lsd = b.pl("pt_lsd_step", LSD_STEP);
        b.step(
            &lsd,
            &[&dir, &noise_in, &p, &q],
            (ld as u32).div_ceil(64),
            1,
        );

        let (steps, _, _, _) = b.finish();
        Ok(Self {
            steps,
            latent_out: noise_in.clone(),
            noise_in,
            latent_dim: ld,
        })
    }

    /// One LSD step from `noise` (drawn host-side so the seed stream matches the CPU path
    /// exactly), conditioned on whatever the bound hidden buffer currently holds.
    pub fn step(&mut self, ctx: &GpuCtx, noise: &[f32]) -> Result<Vec<f32>> {
        anyhow::ensure!(noise.len() == self.latent_dim, "noise width");
        ctx.queue
            .write_buffer(&self.noise_in, 0, bytemuck::cast_slice(noise));
        MimiGpu::run(ctx, &self.steps);
        ctx.read(&self.latent_out, self.latent_dim)
    }
}