//! MTPLX model operations using its pinned custom and runtime 0.32.2 shaders.
//! Rust owns selection and dispatch. No runtime C/C++ library is linked.
use super::*;
use std::cell::RefCell;
#[cfg(test)]
use std::cell::{Cell, Ref};
use std::ffi::CString;
use std::rc::Rc;

pub(super) mod allocator;
mod ar;
mod array;
mod attention;
mod binary;
mod binding;
mod cache;
mod compile_graph;
mod compiled;
mod context_copy;
mod decoder;
mod dense;
mod encoder;
mod eval;
mod execution;
mod fused;
mod gdn;
mod graph_sync;
mod hyper;
mod indexing;
#[cfg(test)]
mod issue103_analysis;
#[cfg(test)]
mod issue104;
#[cfg(test)]
mod issue105;
#[cfg(test)]
mod issue106;
#[cfg(test)]
mod issue108;
#[cfg(test)]
mod issue109;
mod linear;
mod load;
mod mlp;
mod model;
mod moe;
mod mtp;
mod ngram;
mod ngram_sidecar;
mod ngram_stage;
mod normalization;
mod ops;
mod owner;
mod ple;
mod policy;
mod prefill;
mod qsa_attention;
mod qsa_compile;
mod qsa_core;
#[cfg(test)]
mod qsa_graph;
mod qsa_indexer;
mod qsa_prefill;
mod qsa_prepare;
mod qsa_select;
mod random;
mod reduce;
mod request;
mod residency;
mod rope;
mod sampling;
mod scheduler;
mod sdpa;
mod session_cache;
mod session_checkpoint;
mod session_codec;
mod sort;
mod sparse_attention;
mod speculative;
mod stream;
#[cfg(test)]
mod submission;
mod unary;
mod views;
mod vision;
mod vision_kernels;
mod weights;
#[cfg(test)]
pub(super) use encoder::Scope;
pub(super) use encoder::active_allocator;
pub(in crate::engine) use owner::QwenExecutor;
#[cfg(test)]
pub(super) use qsa_graph::{QsaCacheInput, QsaGraphBank};
pub(in crate::engine) use request::LiveSession as QwenResidentState;
#[cfg(test)]
pub(super) use stream::{Device as StreamDevice, Streams};

struct Binding<'a> {
    raw: MetalBinding,
    allocation: Option<BufferAllocationHold>,
    access: Option<encoder::Access>,
    _borrow: std::marker::PhantomData<&'a ()>,
}

fn tensor(index: u32, buffer: &Buffer) -> Binding<'_> {
    Binding {
        raw: MetalBinding {
            tensor: buffer.raw(),
            bytes: std::ptr::null(),
            size: 0,
            index,
            offset: 0,
        },
        allocation: Some(buffer.hold_allocation()),
        access: None,
        _borrow: std::marker::PhantomData,
    }
}

fn bytes<T: ?Sized>(index: u32, value: &T) -> Binding<'_> {
    Binding {
        raw: MetalBinding {
            tensor: std::ptr::null(),
            bytes: std::ptr::from_ref(value).cast(),
            size: u32::try_from(std::mem::size_of_val(value)).expect("Metal inline bytes size"),
            index,
            offset: 0,
        },
        allocation: None,
        access: None,
        _borrow: std::marker::PhantomData,
    }
}

fn scalar_buffer(value: &[u8]) -> Result<Buffer, String> {
    let buffer = Buffer::mtplx_bytes(value.len() as u64)?;
    buffer.write(0, value)?;
    Ok(buffer)
}

fn compile_custom(
    name: &str,
    source: impl FnOnce() -> Result<String, String>,
) -> Result<(), String> {
    let name = CString::new(name).map_err(|e| e.to_string())?;
    if unsafe { ds4_gpu_mtplx_compile(name.as_ptr(), std::ptr::null()) } == 0 {
        let source = CString::new(source()?).map_err(|e| e.to_string())?;
        check(
            unsafe { ds4_gpu_mtplx_compile(name.as_ptr(), source.as_ptr()) },
            "compiling original MTPLX custom kernel",
        )?;
    }
    Ok(())
}

fn qsa_output_args(section: &serde_json::Value) -> Result<Vec<String>, String> {
    section["outputs"]
        .as_array()
        .ok_or("invalid QSA output names")?
        .iter()
        .map(|name| {
            let name = name.as_str().ok_or("invalid QSA output name")?;
            let ty = if name.ends_with("ids") {
                "int"
            } else if name == "adjusted_scores" {
                "float"
            } else {
                "bool"
            };
            Ok(format!("device {ty}* {name}"))
        })
        .collect()
}

fn qsa_shader_source(header: &str, name: &str, args: &[String], body: &str) -> String {
    let args = args
        .iter()
        .enumerate()
        .map(|(i, a)| format!("{a} [[buffer({i})]]"))
        .collect::<Vec<_>>()
        .join(",\n");
    format!(
        "{header}kernel void {name}(\n{args},\nuint3 threadgroup_position_in_grid [[threadgroup_position_in_grid]],\nuint3 thread_position_in_threadgroup [[thread_position_in_threadgroup]]) {{{body}}}\n"
    )
}

fn dispatch(
    name: &str,
    bindings: &[Binding<'_>],
    groups: [u32; 3],
    threads: [u32; 3],
) -> Result<(), String> {
    dispatch_specialized(name, bindings, &[], groups, threads)
}

fn dispatch_specialized(
    name: &str,
    bindings: &[Binding<'_>],
    constants: &[MetalConstant],
    groups: [u32; 3],
    threads: [u32; 3],
) -> Result<(), String> {
    dispatch_geometry(name, bindings, constants, groups, threads, false)
}

fn dispatch_geometry(
    name: &str,
    bindings: &[Binding<'_>],
    constants: &[MetalConstant],
    grid: [u32; 3],
    threads: [u32; 3],
    individual_threads: bool,
) -> Result<(), String> {
    if let Some(result) = encoder::with_active(|encoder| {
        encoder.dispatch(name, bindings, constants, grid, threads, individual_threads)
    }) {
        return result;
    }
    dispatch_geometry_on(
        name,
        bindings,
        constants,
        grid,
        threads,
        individual_threads,
        None,
    )
}

fn dispatch_geometry_on(
    name: &str,
    bindings: &[Binding<'_>],
    constants: &[MetalConstant],
    grid: [u32; 3],
    threads: [u32; 3],
    individual_threads: bool,
    target: Option<(*mut c_void, *mut c_void, bool)>,
) -> Result<(), String> {
    let name = CString::new(name).map_err(|e| e.to_string())?;
    let raw = bindings.iter().map(|b| b.raw).collect::<Vec<_>>();
    let mut seen = std::collections::HashSet::new();
    let allocations = bindings
        .iter()
        .filter_map(|b| b.allocation.as_ref())
        .filter(|a| seen.insert(a.identity()))
        .cloned()
        .collect::<Vec<_>>();
    unsafe extern "C" fn release(context: *mut c_void) {
        // The bridge takes this Box exactly once at callback registration,
        // including when a submitted standalone command subsequently fails.
        drop(unsafe { Box::from_raw(context.cast::<Vec<BufferAllocationHold>>()) });
    }
    let mut owner = Box::into_raw(Box::new(allocations)).cast::<c_void>();
    let (command, encoder, barrier) =
        target.unwrap_or((std::ptr::null_mut(), std::ptr::null_mut(), false));
    let result = unsafe {
        ds4_gpu_mtplx_encode(
            name.as_ptr(),
            raw.as_ptr(),
            raw.len() as u32,
            constants.as_ptr(),
            constants.len() as u32,
            grid.as_ptr(),
            threads.as_ptr(),
            u32::from(individual_threads),
            &mut owner,
            Some(release),
            command,
            encoder,
            u32::from(barrier),
        )
    };
    if !owner.is_null() {
        // Validation/pipeline/encoder failure: no ownership was transferred.
        unsafe { release(owner) };
    }
    let result = check(result, &format!("dispatching MTPLX {name:?}"));
    #[cfg(test)]
    if result.is_ok() {
        super::qwen_mtplx_tests::record_dispatch(
            name.to_str().unwrap(),
            grid,
            threads,
            individual_threads,
        );
    }
    result
}

#[derive(Clone, Copy)]
#[cfg(test)]
pub(super) enum SortDtype {
    U32,
    Bf16,
    F32,
}

#[cfg(test)]
pub(super) struct RouterOutput {
    pub(super) probabilities: Buffer,
    /// Last ten columns of argsort, not a contiguous [rows,10] allocation.
    pub(super) indices: Buffer,
    pub(super) index_row_stride: u32,
    pub(super) scores: Buffer,
    _temporaries: Vec<Buffer>,
}

/// SparseMoeBlock's precise softmax -> argpartition -> take_along_axis ->
/// optional normalization. Consumes donatable logits; retains GPU scratch.
#[cfg(test)]
pub(super) fn router(
    logits: Buffer,
    [rows, experts]: [u32; 2],
    normalize: bool,
) -> Result<RouterOutput, String> {
    let count = u64::from(rows) * u64::from(experts);
    if rows == 0
        || experts < 10
        || count > i32::MAX as u64
        || unsafe { ds4_gpu_tensor_bytes(logits.raw()) } < count * 2
    {
        return Err("invalid MTPLX top-10 router shape or buffer".into());
    }
    let group = if experts <= 4096 {
        experts.div_ceil(4).div_ceil(32) * 32
    } else {
        1024 // Pinned looped-softmax pipeline maximum on M5 Max.
    };
    let threads = rows
        .checked_mul(group)
        .ok_or("MTPLX softmax grid overflow")?;
    dispatch_geometry(
        if experts <= 4096 {
            "block_softmax_precise_bfloat16"
        } else {
            "looped_softmax_precise_bfloat16"
        },
        &[
            tensor(0, &logits).input(count as usize),
            tensor(1, &logits).output(count as usize),
            bytes(2, &experts),
        ],
        &[],
        [threads, 1, 1],
        [group, 1, 1],
        true,
    )?;
    let order = Buffer::mtplx_bytes(count * 4)?;
    let mut temporary = argsort(&logits, &order, [rows, experts], SortDtype::Bf16)?;
    let offset = u64::from(experts - 10) * 4;
    let indices = order.view(offset, count * 4 - offset)?;
    let scores = Buffer::mtplx_bytes(u64::from(rows) * 10 * 2)?;
    let grid = [1, 10, rows];
    dispatch_geometry(
        if rows == 1 || experts == 10 {
            "kernel_qwen_mtplx_gather_axis_bf16_int_11"
        } else {
            "kernel_qwen_mtplx_gather_axis_bf16_int_10"
        },
        &[
            tensor(0, &logits).input(count as usize),
            tensor(1, &indices).input(((u64::from(rows) - 1) * u64::from(experts) + 10) as usize),
            tensor(2, &scores).output((u64::from(rows) * 10) as usize),
            bytes(3, &rows),
            bytes(4, &u64::from(experts)),
            bytes(5, &u64::from(experts)),
            bytes(6, &1_u64),
            bytes(7, &1_u32),
            bytes(8, &experts),
            bytes(9, &1_u64),
            bytes(10, &1_u64),
        ],
        &[],
        grid,
        block_dims(grid),
        true,
    )?;
    if normalize {
        let sums = Buffer::mtplx_bytes(u64::from(rows) * 2)?;
        dispatch_geometry(
            "row_reduce_small_1_reduce_sumbfloat16",
            &[
                tensor(0, &scores).input((u64::from(rows) * 10) as usize),
                tensor(1, &sums).output(rows as usize),
                bytes(2, &10_u64),
                bytes(3, &1_u64),
                bytes(4, &rows),
                bytes(5, &10_u64),
                bytes(6, &1_u32),
                bytes(7, &0_u32),
                bytes(8, &0_u64),
                bytes(9, &0_u32),
            ],
            &[],
            [rows, 1, 1],
            [rows.min(1024), 1, 1],
            true,
        )?;
        if rows == 1 {
            dispatch_geometry(
                "vs_Dividebfloat16",
                &[
                    tensor(0, &scores).input((u64::from(rows) * 10) as usize),
                    tensor(1, &sums).input(rows as usize),
                    tensor(2, &scores).output((u64::from(rows) * 10) as usize),
                    bytes(3, &10_u32),
                ],
                &[],
                [10, 1, 1],
                [10, 1, 1],
                true,
            )?;
        } else {
            let grid = [10, rows, 1];
            dispatch_geometry(
                "g2_Dividebfloat16",
                &[
                    tensor(0, &scores).input((u64::from(rows) * 10) as usize),
                    tensor(1, &sums).input(rows as usize),
                    tensor(2, &scores).output((u64::from(rows) * 10) as usize),
                    bytes(3, &[10_u64, 1]),
                    bytes(4, &[1_u64, 0]),
                ],
                &[],
                grid,
                block_dims(grid),
                true,
            )?;
        }
        temporary.push(sums);
    }
    temporary.push(order);
    Ok(RouterOutput {
        probabilities: logits,
        indices,
        index_row_stride: experts,
        scores,
        _temporaries: temporary,
    })
}

/// Contiguous last-axis argsort used by routing and _gather_sort. The reference
/// also implements argpartition by this complete sort, not a local top-k kernel.
#[cfg(test)]
pub(super) fn argsort(
    input: &Buffer,
    output: &Buffer,
    [rows, width]: [u32; 2],
    dtype: SortDtype,
) -> Result<Vec<Buffer>, String> {
    use array::{Array, Dtype};
    let dtype = match dtype {
        SortDtype::U32 => Dtype::U32,
        SortDtype::Bf16 => Dtype::BF16,
        SortDtype::F32 => Dtype::F32,
    };
    let count = u64::from(rows) * u64::from(width);
    let item_bytes = dtype.itemsize() as u64;
    if rows == 0
        || width == 0
        || count > i32::MAX as u64
        || unsafe { ds4_gpu_tensor_bytes(input.raw()) } < count * item_bytes
        || unsafe { ds4_gpu_tensor_bytes(output.raw()) } < count * 4
    {
        return Err("invalid MTPLX contiguous argsort shape or buffer".into());
    }
    let shape = if rows == 1 {
        vec![width as i32]
    } else {
        vec![rows as i32, width as i32]
    };
    let input = Array::new(&shape, dtype, input.view(0, count * item_bytes)?)?;
    let output = Array::new(&shape, Dtype::U32, output.view(0, count * 4)?)?;
    sort::sort_into(&input, &output, shape.len() - 1, true)
}

fn block_dims(grid: [u32; 3]) -> [u32; 3] {
    block_dims_pow2(grid, 10)
}

fn block_dims_pow2(grid: [u32; 3], pow2: u32) -> [u32; 3] {
    let mut group = [1, 1, 1];
    let mut used = 0;
    loop {
        let before = used;
        for axis in 0..3 {
            if grid[axis] >= group[axis] * 2 {
                group[axis] *= 2;
                used += 1;
            }
            if used == pow2 {
                break;
            }
        }
        if used == pow2 || used == before {
            break;
        }
    }
    group
}

/// Actual compiled nn.silu(gate) followed by the reference's separate * up.
/// Input is the fused [rows,2*intermediate] BF16 pack, with strided half views.
/// Keep returned SiLU storage alive until the enclosing GPU batch completes.
#[cfg(test)]
pub(super) fn gate_up_bf16(
    input: &Buffer,
    output: &Buffer,
    [rows, intermediate]: [u32; 2],
) -> Result<Buffer, String> {
    let count = u64::from(rows) * u64::from(intermediate);
    if rows == 0
        || intermediate == 0
        || !intermediate.is_multiple_of(2)
        || count > i32::MAX as u64
        || unsafe { ds4_gpu_tensor_bytes(input.raw()) } < count * 4
        || unsafe { ds4_gpu_tensor_bytes(output.raw()) } < count * 2
    {
        return Err("invalid MTPLX fused gate/up buffer or shape".into());
    }
    let activated = Buffer::mtplx_bytes(count * 2)?;
    let up = input.view(
        u64::from(intermediate) * 2,
        count * 4 - u64::from(intermediate) * 2,
    )?;
    let count = count as u32;
    if rows == 1 {
        let work = if count < 65536 { 1 } else { 4 };
        let threads = count.div_ceil(work);
        dispatch_geometry(
            if work == 1 {
                "kernel_qwen_mtplx_silu_bf16_contiguous"
            } else {
                "kernel_qwen_mtplx_silu_bf16_contiguous_n"
            },
            &[
                tensor(0, input).input((u64::from(count) * 2 - u64::from(intermediate)) as usize),
                tensor(1, &activated).output(count as usize),
                bytes(2, &count),
            ],
            &[],
            [threads, 1, 1],
            [threads.min(1024), 1, 1],
            true,
        )?;
        dispatch_geometry(
            if work == 1 {
                "vv_Multiplybfloat16"
            } else {
                "vvn_Multiplybfloat16"
            },
            &[
                tensor(0, &activated).input(count as usize),
                tensor(1, &up).input((u64::from(count) * 2 - u64::from(intermediate)) as usize),
                tensor(2, output).output(count as usize),
                bytes(3, &count),
            ],
            &[],
            [threads, 1, 1],
            [threads.min(1024), 1, 1],
            true,
        )?;
    } else {
        let shape = [rows, intermediate];
        let in_strides = [u64::from(intermediate) * 2, 1];
        let out_strides = [u64::from(intermediate), 1];
        let grid = [intermediate, rows, 1];
        let group = block_dims(grid);
        dispatch_geometry(
            "kernel_qwen_mtplx_silu_bf16_strided_2",
            &[
                tensor(0, input).input((u64::from(count) * 2 - u64::from(intermediate)) as usize),
                bytes(1, &in_strides),
                tensor(2, &activated).output(count as usize),
                bytes(3, &shape),
            ],
            &[],
            grid,
            group,
            true,
        )?;
        dispatch_geometry(
            "g2_Multiplybfloat16",
            &[
                tensor(0, &activated).input(count as usize),
                tensor(1, &up).input((u64::from(count) * 2 - u64::from(intermediate)) as usize),
                tensor(2, output).output(count as usize),
                bytes(3, &out_strides),
                bytes(4, &in_strides),
            ],
            &[],
            grid,
            group,
            true,
        )?;
    }
    Ok(activated)
}

#[cfg(test)]
fn rms_norm_bf16(
    input: &Buffer,
    weight: Option<&Buffer>,
    output: &Buffer,
    [rows, width]: [u32; 2],
    eps: f32,
) -> Result<(), String> {
    let count = u64::from(rows) * u64::from(width);
    if rows == 0
        || width == 0
        || count > i32::MAX as u64
        || !eps.is_finite()
        || eps < 0.0
        || [input, output]
            .iter()
            .any(|b| unsafe { ds4_gpu_tensor_bytes(b.raw()) } < count * 2)
        || weight.is_some_and(|w| unsafe { ds4_gpu_tensor_bytes(w.raw()) } < u64::from(width) * 2)
    {
        return Err("invalid MTPLX BF16 RMS norm".into());
    }
    let default_weight = weight
        .is_none()
        .then(|| scalar_buffer(&0x3f80_u16.to_le_bytes()))
        .transpose()?;
    let bound_weight = weight.or(default_weight.as_ref()).unwrap();
    normalization::dispatch(
        [
            (input, 0, count as usize),
            (
                bound_weight,
                0,
                if weight.is_some() { width as usize } else { 1 },
            ),
            (output, 0, count as usize),
        ],
        array::Dtype::BF16,
        width,
        u32::from(weight.is_some()),
        eps,
    )
}

/// The shared HC/PLE GroupedRMSNorm: unweighted per-group RMS followed by
/// separate full-feature weighting. Folding the weight into RMS changes rounding.
#[cfg(test)]
pub(super) fn grouped_rms_norm(
    input: &Buffer,
    weight: &Buffer,
    [rows, width, group_width]: [u32; 3],
    eps: f32,
) -> Result<Buffer, String> {
    if group_width == 0
        || !width.is_multiple_of(group_width)
        || unsafe { ds4_gpu_tensor_bytes(weight.raw()) } < u64::from(width) * 2
    {
        return Err("invalid MTPLX grouped RMS shape or weight".into());
    }
    let count = u64::from(rows) * u64::from(width);
    if count == 0 || count > i32::MAX as u64 {
        return Err("invalid MTPLX grouped RMS size".into());
    }
    let output = Buffer::mtplx_bytes(count * 2)?;
    rms_norm_bf16(
        input,
        None,
        &output,
        [count as u32 / group_width, group_width],
        eps,
    )?;
    if rows == 1 {
        let work = if width < 65536 { 1 } else { 4 };
        let threads = width.div_ceil(work);
        dispatch_geometry(
            if work == 1 {
                "vv_Multiplybfloat16"
            } else {
                "vvn_Multiplybfloat16"
            },
            &[
                tensor(0, &output).input(count as usize),
                tensor(1, weight).input(width as usize),
                tensor(2, &output).output(count as usize),
                bytes(3, &width),
            ],
            &[],
            [threads, 1, 1],
            [threads.min(1024), 1, 1],
            true,
        )?;
    } else {
        let grid = [width, rows, 1];
        dispatch_geometry(
            "g2_Multiplybfloat16",
            &[
                tensor(0, &output).input(count as usize),
                tensor(1, weight).input(width as usize),
                tensor(2, &output).output(count as usize),
                bytes(3, &[u64::from(width), 1]),
                bytes(4, &[0_u64, 1]),
            ],
            &[],
            grid,
            block_dims(grid),
            true,
        )?;
    }
    Ok(output)
}

#[cfg(test)]
fn cast_bf16_f32(input: &Buffer, output: &Buffer, count: u32, to_f32: bool) -> Result<(), String> {
    let (in_bytes, out_bytes) = if to_f32 { (2, 4) } else { (4, 2) };
    if count == 0
        || count > i32::MAX as u32
        || unsafe { ds4_gpu_tensor_bytes(input.raw()) } < u64::from(count) * in_bytes
        || unsafe { ds4_gpu_tensor_bytes(output.raw()) } < u64::from(count) * out_bytes
    {
        return Err("invalid MTPLX BF16/FP32 cast".into());
    }
    let work = if count < 65536 {
        1
    } else if to_f32 {
        2
    } else {
        4
    };
    let threads = count.div_ceil(work);
    let dtype = if to_f32 {
        "bfloat16float32"
    } else {
        "float32bfloat16"
    };
    dispatch_geometry(
        &format!("v{}_copy{dtype}", if work == 1 { "" } else { "n" }),
        &[
            tensor(0, input).input(count as usize),
            tensor(1, output).output(count as usize),
            bytes(2, &count),
        ],
        &[],
        [threads, 1, 1],
        [threads.min(1024), 1, 1],
        true,
    )
}

#[cfg(test)]
pub(super) struct GatedNormOutput {
    pub(super) hidden: Buffer,
    _temporaries: Vec<Buffer>,
}

/// SigmoidRMSNormGated, including the projected Z view's token stride.
/// RMS output rounds to BF16 before its FP32 sigmoid/multiply epilogue.
#[cfg(test)]
pub(super) fn sigmoid_rms_norm(
    input: &Buffer,
    weight: &Buffer,
    gate: Option<&Buffer>,
    [tokens, heads, width, gate_stride]: [u32; 4],
    eps: f32,
) -> Result<GatedNormOutput, String> {
    let row_width = heads
        .checked_mul(width)
        .ok_or("MTPLX gated RMS width overflow")?;
    let count = tokens
        .checked_mul(row_width)
        .filter(|n| *n > 0 && *n <= i32::MAX as u32)
        .ok_or("invalid MTPLX gated RMS count")?;
    let rows = tokens
        .checked_mul(heads)
        .ok_or("MTPLX gated RMS row overflow")?;
    let hidden = Buffer::mtplx_bytes(u64::from(count) * 2)?;
    rms_norm_bf16(input, Some(weight), &hidden, [rows, width], eps)?;
    let Some(gate) = gate else {
        return Ok(GatedNormOutput {
            hidden,
            _temporaries: Vec::new(),
        });
    };
    if gate_stride < row_width
        || unsafe { ds4_gpu_tensor_bytes(gate.raw()) }
            < (u64::from(tokens - 1) * u64::from(gate_stride) + u64::from(row_width)) * 2
    {
        return Err("invalid MTPLX gated RMS Z view".into());
    }
    let g = Buffer::mtplx_bytes(u64::from(count) * 4)?;
    if tokens == 1 || gate_stride == row_width {
        cast_bf16_f32(gate, &g, count, true)?;
    } else {
        let grid = [row_width, tokens, 1];
        dispatch_geometry(
            "g2_copybfloat16float32",
            &[
                tensor(0, gate).input(
                    (u64::from(tokens - 1) * u64::from(gate_stride) + u64::from(row_width))
                        as usize,
                ),
                tensor(1, &g).output(count as usize),
                bytes(3, &[u64::from(gate_stride), 1]),
            ],
            &[],
            grid,
            block_dims(grid),
            true,
        )?;
    }
    let work = if count < 65536 { 1 } else { 2 };
    let threads = count.div_ceil(work);
    dispatch_geometry(
        if work == 1 {
            "v_Sigmoidfloat32float32"
        } else {
            "vn_Sigmoidfloat32float32"
        },
        &[
            tensor(0, &g).input(count as usize),
            tensor(1, &g).output(count as usize),
            bytes(2, &count),
        ],
        &[],
        [threads, 1, 1],
        [threads.min(1024), 1, 1],
        true,
    )?;
    let x = Buffer::mtplx_bytes(u64::from(count) * 4)?;
    cast_bf16_f32(&hidden, &x, count, true)?;
    dispatch_geometry(
        if work == 1 {
            "vv_Multiplyfloat32"
        } else {
            "vvn_Multiplyfloat32"
        },
        &[
            tensor(0, &g).input(count as usize),
            tensor(1, &x).input(count as usize),
            tensor(2, &x).output(count as usize),
            bytes(3, &count),
        ],
        &[],
        [threads, 1, 1],
        [threads.min(1024), 1, 1],
        true,
    )?;
    cast_bf16_f32(&x, &hidden, count, false)?;
    Ok(GatedNormOutput {
        hidden,
        _temporaries: vec![g, x],
    })
}

/// Row gather used by MTPLX's _gather_sort and _scatter_unsort. Indices must
/// originate from the checked routing/sort path, not unchecked external data.
#[cfg(test)]
pub(super) fn gather_front(
    input: &Buffer,
    indices: &Buffer,
    output: &Buffer,
    [rows, width, routes]: [u32; 3],
    dtype: SortDtype,
) -> Result<(), String> {
    let (tag, item_bytes) = match dtype {
        SortDtype::U32 => ("u32", 4),
        SortDtype::Bf16 => ("bf16", 2),
        SortDtype::F32 => ("f32", 4),
    };
    let src_count = u64::from(rows) * u64::from(width);
    let out_count = u64::from(routes) * u64::from(width);
    if [rows, width, routes]
        .iter()
        .any(|v| *v == 0 || *v > i32::MAX as u32)
        || unsafe { ds4_gpu_tensor_bytes(input.raw()) } < src_count * item_bytes
        || unsafe { ds4_gpu_tensor_bytes(indices.raw()) } < u64::from(routes) * 4
        || unsafe { ds4_gpu_tensor_bytes(output.raw()) } < out_count * item_bytes
    {
        return Err("invalid MTPLX gather-front shape or buffer".into());
    }
    let work = if width > 8 && item_bytes < 4 { 2 } else { 1 };
    let grid = [width.div_ceil(work), routes, 1];
    let group = block_dims(grid);
    let loc = if src_count > i32::MAX as u64 || out_count > i32::MAX as u64 {
        "int64_t"
    } else {
        "int"
    };
    dispatch_geometry(
        &format!("kernel_qwen_mtplx_gather_front_{tag}_{loc}_{work}"),
        &[
            tensor(0, input).input(src_count as usize),
            tensor(1, indices).input(routes as usize),
            tensor(2, output).output(out_count as usize),
            bytes(3, &u64::from(width)),
            bytes(4, &rows),
        ],
        &[],
        grid,
        group,
        true,
    )
}

#[cfg(test)]
pub(super) struct GatherSorted {
    pub(super) x: Buffer,
    pub(super) indices: Buffer,
    pub(super) inverse: Buffer,
    _temporaries: Vec<Buffer>,
}

/// The connected reference _gather_sort path. Retain the result through GPU
/// completion; the caller applies the reference's >=64-route sort condition.
#[cfg(test)]
pub(super) fn gather_sort(
    input: &Buffer,
    indices: &Buffer,
    [tokens, top_k, width]: [u32; 3],
) -> Result<GatherSorted, String> {
    let routes = tokens
        .checked_mul(top_k)
        .filter(|v| *v > 0 && *v <= i32::MAX as u32)
        .ok_or("invalid MTPLX route count")?;
    if width == 0
        || width > i32::MAX as u32
        || unsafe { ds4_gpu_tensor_bytes(input.raw()) } < u64::from(tokens) * u64::from(width) * 2
    {
        return Err("invalid MTPLX gather-sort input".into());
    }
    let order = Buffer::mtplx_bytes(u64::from(routes) * 4)?;
    let inverse = Buffer::mtplx_bytes(u64::from(routes) * 4)?;
    let token_rows = Buffer::mtplx_bytes(u64::from(routes) * 4)?;
    let sorted_indices = Buffer::mtplx_bytes(u64::from(routes) * 4)?;
    let x = Buffer::mtplx_bytes(u64::from(routes) * u64::from(width) * 2)?;
    let mut temporary = argsort(indices, &order, [1, routes], SortDtype::U32)?;
    temporary.extend(argsort(&order, &inverse, [1, routes], SortDtype::U32)?);
    let divisor = scalar_buffer(&top_k.to_le_bytes())?;
    let work = if routes < 65536 { 1 } else { 2 };
    let threads = routes.div_ceil(work);
    dispatch_geometry(
        if work == 1 {
            "vs_Divideuint32"
        } else {
            "vsn_Divideuint32"
        },
        &[
            tensor(0, &order).input(routes as usize),
            tensor(1, &divisor).input(1),
            tensor(2, &token_rows).output(routes as usize),
            bytes(3, &routes),
        ],
        &[],
        [threads, 1, 1],
        [threads.min(1024), 1, 1],
        true,
    )?;
    gather_front(
        input,
        &token_rows,
        &x,
        [tokens, width, routes],
        SortDtype::Bf16,
    )?;
    gather_front(
        indices,
        &order,
        &sorted_indices,
        [routes, 1, routes],
        SortDtype::U32,
    )?;
    temporary.extend([order, token_rows]);
    Ok(GatherSorted {
        x,
        indices: sorted_indices,
        inverse,
        _temporaries: temporary,
    })
}

fn qmv_limit(k: u32, n: u32) -> u32 {
    if k <= 2048 && n <= 2048 {
        33
    } else if k <= 4096 && n <= 4096 {
        25
    } else {
        13
    }
}

#[derive(Debug, PartialEq)]
pub(super) struct AffinePlan {
    pub(super) name: String,
    pub(super) groups: [u32; 3],
    pub(super) threads: [u32; 3],
    pub(super) splits: u32,
    pub(super) parameters: Vec<u32>,
}

/// Contiguous, non-batched affine x @ W.T on the requested M5 Max.
/// Selection follows quantized.cpp at 1f8e74e, not a local speed heuristic.
pub(super) fn affine_plan(
    shape: [u32; 3],
    bits: u32,
    group: u32,
    nax: bool,
) -> Result<AffinePlan, String> {
    affine_plan_batched(shape, bits, group, nax, 1)
}

/// QuantizedMatmul uses M=x.shape(-2) and the original batch strides when
/// x is not row-contiguous. Only the B=1 matrix route may use Split-K.
fn affine_plan_batched(
    [m, n, k]: [u32; 3],
    bits: u32,
    group: u32,
    nax: bool,
    batches: u32,
) -> Result<AffinePlan, String> {
    if [m, n, k].iter().any(|d| *d == 0 || *d > i32::MAX as u32)
        || !matches!(bits, 4 | 8)
        || !matches!(group, 32 | 64)
        || !k.is_multiple_of(group)
        || batches == 0
        || u64::from(m)
            .checked_mul(u64::from(n))
            .and_then(|v| v.checked_mul(u64::from(batches)))
            .is_none_or(|v| v > i32::MAX as u64)
    {
        return Err("invalid MTPLX BF16 affine shape or installed-pack quantization".into());
    }
    let suffix = format!("bfloat16_t_gs_{group}_b_{bits}");
    let batch = u8::from(batches > 1);
    let limit = qmv_limit(k, n);
    let mut plan = AffinePlan {
        name: String::new(),
        groups: [1, 1, 1],
        threads: [32, 2, 1],
        splits: 1,
        parameters: vec![k, n],
    };
    if m < limit {
        if matches!(k, 64 | 128) {
            plan.name = format!("affine_qmv_quad_{suffix}_d_{k}_batch_{batch}");
            plan.groups = [m, n.div_ceil(64), batches];
            plan.threads = [32, 1, 1];
        } else if m >= 2 {
            let vectors = m.div_ceil(m.div_ceil(5));
            plan.name = format!("affine_qmv_wide_{suffix}_nv_{vectors}_kl_8_batch_{batch}");
            plan.groups = [m.div_ceil(vectors), n.div_ceil(8), batches];
            plan.parameters.push(m);
        } else {
            let fast = n.is_multiple_of(8) && k.is_multiple_of((32 / bits) * 2 * 32);
            plan.name = format!(
                "affine_qmv{}_{suffix}_batch_{batch}",
                if fast { "_fast" } else { "" }
            );
            plan.groups = [m, n.div_ceil(8), batches];
        }
        return Ok(plan);
    }
    let tiles = u64::from(n.div_ceil(32)) * u64::from(m.div_ceil(32));
    let alignment = group.max(32);
    let mut splits = if batches == 1 {
        ((512 / tiles) as u32).max(1).min(k / alignment)
    } else {
        1
    };
    while splits > 1 && !k.is_multiple_of(splits * alignment) {
        splits -= 1;
    }
    plan.splits = splits;
    plan.parameters.push(m);
    plan.threads = [32, 2, 2];
    if splits > 1 {
        plan.name = format!("affine_qmm_t_splitk_{suffix}_alN_{}", n.is_multiple_of(32));
        plan.groups = [n.div_ceil(32), m.div_ceil(32), splits];
        plan.parameters.extend([k / splits, m * n]);
    } else if nax && k.is_multiple_of(64) {
        let bm = if m <= 32 { 32 } else { 64 };
        plan.name = format!(
            "affine_qmm_t_nax_{suffix}_bm{bm}_bn64_bk64_wm2_wn2_alN_{}_batch_{batch}",
            n.is_multiple_of(64)
        );
        plan.groups = [n.div_ceil(64), m.div_ceil(bm), batches];
    } else {
        plan.name = format!(
            "affine_qmm_t_{suffix}_alN_{}_batch_{batch}",
            n.is_multiple_of(32)
        );
        plan.groups = [n.div_ceil(32), m.div_ceil(32), batches];
    }
    Ok(plan)
}

/// Normalized x[X,M,K], affine weights[E,N,K], index vectors[R], y[R,M,N].
/// For sorted=true x is already gathered in rhs order and lhs is identity,
/// matching _gather_sort and the reference's implicit-lhs sorted dispatch.
/// Indices originate from the checked router/sort path, not unchecked user data.
#[cfg(test)]
pub(super) fn gather_bf16(
    buffers: [&Buffer; 7],
    shape: [u32; 6],
    quant: [u32; 2],
    sorted: bool,
    nax: bool,
) -> Result<(), String> {
    gather_bf16_broadcast(buffers, shape, quant, sorted, nax, 1, [shape[5]; 2])
}

#[cfg(test)]
fn gather_bf16_broadcast(
    buffers: [&Buffer; 7],
    [x_batches, m, n, k, experts, routes]: [u32; 6],
    [bits, group]: [u32; 2],
    sorted: bool,
    nax: bool,
    lhs_repeat: u32,
    [rhs_width, rhs_stride]: [u32; 2],
) -> Result<(), String> {
    affine_plan([m, n, k], bits, group, nax)?;
    if [x_batches, experts, routes]
        .iter()
        .any(|d| *d == 0 || *d > i32::MAX as u32)
        || u64::from(m) * u64::from(routes) * u64::from(n) > i32::MAX as u64
        || (sorted && x_batches != routes)
        || lhs_repeat == 0
        || !routes.is_multiple_of(lhs_repeat)
        || (sorted && lhs_repeat != 1)
        || rhs_width == 0
        || rhs_stride < rhs_width
        || !routes.is_multiple_of(rhs_width)
        || (rhs_stride != rhs_width && lhs_repeat != 1 && lhs_repeat != rhs_width)
        || (sorted && rhs_stride != rhs_width)
    {
        return Err("invalid normalized MTPLX gather shape".into());
    }
    let [w, s, b, x, lhs, rhs, y] = buffers;
    let packed_cols = k / (32 / bits);
    let param_cols = k / group;
    let w_bytes = (u64::from(experts) * u64::from(n)).checked_mul(u64::from(packed_cols) * 4);
    let p_bytes = (u64::from(experts) * u64::from(n)).checked_mul(u64::from(param_cols) * 2);
    let expected = [
        w_bytes,
        p_bytes,
        p_bytes,
        (u64::from(x_batches) * u64::from(m)).checked_mul(u64::from(k) * 2),
        Some(u64::from(routes / lhs_repeat) * 4),
        Some(
            (u64::from(routes / rhs_width - 1) * u64::from(rhs_stride) + u64::from(rhs_width)) * 4,
        ),
        Some(u64::from(routes) * u64::from(m) * u64::from(n) * 2),
    ];
    for (buffer, size) in buffers.iter().zip(expected) {
        if size.is_none_or(|size| unsafe { ds4_gpu_tensor_bytes(buffer.raw()) } < size) {
            return Err("MTPLX gather buffer is smaller than its shape".into());
        }
    }
    let elements: [usize; 7] =
        std::array::from_fn(|i| (expected[i].unwrap() / [4, 2, 2, 2, 4, 4, 2][i]) as usize);
    let strided_rhs = rhs_stride != rhs_width && rhs_width < routes;
    let batch_dim = if strided_rhs || (lhs_repeat > 1 && lhs_repeat < routes) {
        2_u32
    } else {
        1
    };
    let batch_shape = if batch_dim == 2 {
        let width = if strided_rhs { rhs_width } else { lhs_repeat };
        [routes / width, width]
    } else {
        [routes, 1]
    };
    let lhs_strides = if strided_rhs && lhs_repeat == 1 {
        [u64::from(rhs_width), 1]
    } else if lhs_repeat > 1 && batch_dim == 1 {
        [0, 0]
    } else {
        [1_u64, 0]
    };
    let rhs_strides = if strided_rhs {
        [u64::from(rhs_stride), 1]
    } else if batch_dim == 1 {
        [1_u64, 0]
    } else {
        [u64::from(lhs_repeat), 1]
    };
    let x_shape = [x_batches, m, k];
    let x_strides = [u64::from(m) * u64::from(k), u64::from(k), 1];
    let w_shape = [experts, n, packed_cols];
    let w_strides = [
        u64::from(n) * u64::from(packed_cols),
        u64::from(packed_cols),
        1,
    ];
    let p_strides = [
        u64::from(n) * u64::from(param_cols),
        u64::from(param_cols),
        1,
    ];
    let batch_shape = batch_shape.map(|d| d as i32);
    let lhs_strides = lhs_strides.map(|d| d as i64);
    let rhs_strides = rhs_strides.map(|d| d as i64);

    gather_bf16_layout(
        [
            tensor(0, w).input(elements[0]),
            tensor(1, s).input(elements[1]),
            tensor(2, b).input(elements[2]),
            tensor(3, x).input(elements[3]),
            tensor(4, lhs).input(elements[4]),
            tensor(5, rhs).input(elements[5]),
            tensor(6, y).output(elements[6]),
        ],
        [m, n, k, experts, routes, routes],
        [bits, group],
        sorted,
        nax,
        [
            &x_shape.map(|d| d as i32),
            &w_shape.map(|d| d as i32),
            &batch_shape[..batch_dim as usize],
        ],
        [
            &x_strides.map(|d| d as i64),
            &w_strides.map(|d| d as i64),
            &p_strides.map(|d| d as i64),
            &p_strides.map(|d| d as i64),
            &lhs_strides[..batch_dim as usize],
            &rhs_strides[..batch_dim as usize],
        ],
    )
}

/// Shared original GatherQMM dispatch; array batch strides are passed verbatim,
/// while the two broadcast index layouts have jointly collapsed dimensions.
fn gather_bf16_layout(
    bindings: [Binding<'_>; 7],
    [m, n, k, experts, routes, sorted_rows]: [u32; 6],
    [bits, group]: [u32; 2],
    sorted: bool,
    nax: bool,
    [x_shape, w_shape, batch_shape]: [&[i32]; 3],
    [
        x_strides,
        w_strides,
        s_strides,
        b_strides,
        lhs_strides,
        rhs_strides,
    ]: [&[i64]; 6],
) -> Result<(), String> {
    let suffix = format!("bfloat16_t_gs_{group}_b_{bits}");
    if m == 1 && routes >= 16 && sorted && routes / experts >= 4 {
        let (bm, bn, bk, wm, wn) = if nax {
            (
                if sorted_rows / experts < 64 { 32 } else { 64 },
                64,
                64,
                2,
                2,
            )
        } else {
            (16, 32, 32, 1, 2)
        };
        let name = format!(
            "affine_gather_qmm_rhs{}_nt_{suffix}_bm_{bm}_bn_{bn}_bk_{bk}_wm_{wm}_wn_{wn}",
            if nax { "_nax" } else { "" }
        );
        let constants = [
            (200, sorted_rows.is_multiple_of(bm)),
            (201, n.is_multiple_of(bn)),
            (202, k.is_multiple_of(bk)),
        ]
        .map(|(index, value)| MetalConstant {
            index,
            value: u32::from(value),
            kind: 0,
        });

        let [w, s, b, x, _lhs, rhs, y] = bindings;
        let mut sorted_bindings = [x, w, s, b, rhs, y];
        for (index, binding) in sorted_bindings.iter_mut().enumerate() {
            binding.raw.index = index as u32;
        }
        let mut bindings = Vec::from(sorted_bindings);
        bindings.extend([bytes(6, &sorted_rows), bytes(7, &n), bytes(8, &k)]);
        return dispatch_specialized(
            &name,
            &bindings,
            &constants,
            [n.div_ceil(bn), sorted_rows.div_ceil(bm), 1],
            [32, wn, wm],
        );
    }
    let matrix = m >= qmv_limit(k, n);
    let (name, groups, threads) = if matrix {
        if nax && k.is_multiple_of(64) {
            (
                format!(
                    "affine_gather_qmm_t_nax_{suffix}_bm64_bn64_bk64_wm2_wn2_alN_{}",
                    n.is_multiple_of(64)
                ),
                [n.div_ceil(64), m.div_ceil(64), routes],
                [32, 2, 2],
            )
        } else {
            (
                format!("affine_gather_qmm_t_{suffix}_alN_{}", n.is_multiple_of(32)),
                [n.div_ceil(32), m.div_ceil(32), routes],
                [32, 2, 2],
            )
        }
    } else {
        let fast = n.is_multiple_of(8) && k.is_multiple_of((32 / bits) * 2 * 32);
        (
            format!(
                "affine_gather_qmv{}_{suffix}",
                if fast { "_fast" } else { "" }
            ),
            [m, n.div_ceil(8), routes],
            [32, 2, 1],
        )
    };

    let x_dim = (x_shape.len() - 2) as u32;
    let w_dim = (w_shape.len() - 2) as u32;
    let batch_dim = batch_shape.len() as u32;
    let mut bindings = Vec::from(bindings);
    bindings.extend([bytes(7, &k), bytes(8, &n)]);
    if matrix {
        bindings.push(bytes(9, &m));
    }
    let base = bindings.len() as u32;
    bindings.extend([
        bytes(base, &x_dim),
        bytes(base + 1, x_shape),
        bytes(base + 2, x_strides),
        bytes(base + 3, &w_dim),
        bytes(base + 4, w_shape),
        bytes(base + 5, w_strides),
        bytes(base + 6, s_strides),
        bytes(base + 7, b_strides),
        bytes(base + 8, &batch_dim),
        bytes(base + 9, batch_shape),
        bytes(base + 10, lhs_strides),
        bytes(base + 11, rhs_strides),
    ]);
    dispatch(&name, &bindings, groups, threads)
}

#[derive(Clone, Copy)]
#[cfg(test)]
pub(super) enum GateUpWeights<'a> {
    Fused {
        pack: [&'a Buffer; 3],
        quant: [u32; 2],
    },
    Separate {
        gate: [&'a Buffer; 3],
        up: [&'a Buffer; 3],
        gate_quant: [u32; 2],
        up_quant: [u32; 2],
    },
}

#[cfg(test)]
pub(super) struct SwitchWeights<'a> {
    pub(super) input: GateUpWeights<'a>,
    pub(super) down: [&'a Buffer; 3],
    pub(super) down_quant: [u32; 2],
}

/// Actual shapeless compiled swiglu for two contiguous projection outputs.
#[cfg(test)]
fn swiglu_bf16(gate: &Buffer, up: &Buffer, output: &Buffer, count: u32) -> Result<(), String> {
    if count == 0
        || count > i32::MAX as u32
        || [gate, up, output]
            .iter()
            .any(|buffer| unsafe { ds4_gpu_tensor_bytes(buffer.raw()) } < u64::from(count) * 2)
    {
        return Err("invalid MTPLX contiguous SwiGLU input".into());
    }
    let work = if count < 65536 { 1 } else { 4 };
    let threads = count.div_ceil(work);
    dispatch_geometry(
        if work == 1 {
            "kernel_qwen_mtplx_swiglu_bf16_contiguous"
        } else {
            "kernel_qwen_mtplx_swiglu_bf16_contiguous_n"
        },
        &[
            tensor(0, gate).input(count as usize),
            tensor(1, up).input(count as usize),
            tensor(2, output).output(count as usize),
            bytes(3, &count),
        ],
        &[],
        [threads, 1, 1],
        [threads.min(1024), 1, 1],
        true,
    )
}

#[cfg(test)]
pub(super) struct SwitchOutput {
    pub(super) hidden: Buffer,
    _temporaries: Vec<Buffer>,
    _sorted: Option<GatherSorted>,
}

#[cfg(test)]
fn flatten_indices(input: &Buffer, [rows, width, stride]: [u32; 3]) -> Result<Buffer, String> {
    let count = u64::from(rows) * u64::from(width);
    if rows == 0
        || width == 0
        || stride < width
        || count > i32::MAX as u64
        || unsafe { ds4_gpu_tensor_bytes(input.raw()) }
            < (u64::from(rows - 1) * u64::from(stride) + u64::from(width)) * 4
    {
        return Err("invalid MTPLX index flatten view".into());
    }
    if rows == 1 || stride == width {
        return input.view(0, count * 4);
    }
    let flat = Buffer::mtplx_bytes(count * 4)?;
    let grid = [width, rows, 1];
    dispatch_geometry(
        "g2_copyuint32uint32",
        &[
            tensor(0, input)
                .input((u64::from(rows - 1) * u64::from(stride) + u64::from(width)) as usize),
            tensor(1, &flat).output(count as usize),
            bytes(3, &[u64::from(stride), 1]),
        ],
        &[],
        grid,
        block_dims(grid),
        true,
    )?;
    Ok(flat)
}

/// Complete SwitchGLU or _FusedGateUpSwitchGLU with already-routed expert indices.
/// Own the result until GPU completion; router/shared-expert combination belongs
/// to SparseMoeBlock, not this reference submodule.
#[cfg(test)]
pub(super) fn switch_glu(
    weights: SwitchWeights<'_>,
    input: &Buffer,
    indices: &Buffer,
    [tokens, top_k, k, intermediate, experts]: [u32; 5],
) -> Result<SwitchOutput, String> {
    switch_glu_strided(
        weights,
        input,
        indices,
        [tokens, top_k, k, intermediate, experts],
        top_k,
    )
}

#[cfg(test)]
fn switch_glu_strided(
    weights: SwitchWeights<'_>,
    input: &Buffer,
    indices: &Buffer,
    [tokens, top_k, k, intermediate, experts]: [u32; 5],
    index_stride: u32,
) -> Result<SwitchOutput, String> {
    let routes = tokens
        .checked_mul(top_k)
        .filter(|v| *v > 0 && *v <= i32::MAX as u32)
        .ok_or("invalid MTPLX switch route count")?;
    if [k, intermediate, experts]
        .iter()
        .any(|v| *v == 0 || *v > i32::MAX as u32)
        || u64::from(routes) * u64::from(intermediate) * 2 > i32::MAX as u64
        || u64::from(routes) * u64::from(k) > i32::MAX as u64
        || index_stride < top_k
        || unsafe { ds4_gpu_tensor_bytes(indices.raw()) }
            < (u64::from(tokens - 1) * u64::from(index_stride) + u64::from(top_k)) * 4
    {
        return Err("invalid MTPLX switch dimensions".into());
    }
    let arange = |size: u32| -> Result<Buffer, String> {
        let result = Buffer::mtplx_bytes(u64::from(size) * 4)?;
        dispatch_geometry(
            "arangeuint32",
            &[
                bytes(0, &0_u32),
                bytes(1, &1_u32),
                tensor(2, &result).output(size as usize),
            ],
            &[],
            [size, 1, 1],
            [size.min(1024), 1, 1],
            true,
        )?;
        Ok(result)
    };
    // _gather_sort flattens the strided top-k view. Only that flatten needs
    // materialization; the unsorted GatherQMM path accepts original strides.
    let flat_indices = if routes >= 64 {
        Some(flatten_indices(indices, [tokens, top_k, index_stride])?)
    } else {
        None
    };
    let sorted = if routes >= 64 {
        Some(gather_sort(
            input,
            flat_indices.as_ref().unwrap_or(indices),
            [tokens, top_k, k],
        )?)
    } else {
        None
    };
    let route_range = arange(routes)?;
    let token_range = if sorted.is_none() {
        Some(arange(tokens)?)
    } else {
        None
    };
    let (x, rhs, lhs, batches, repeat) = if let Some(sorted) = &sorted {
        (&sorted.x, &sorted.indices, &route_range, routes, 1)
    } else {
        (input, indices, token_range.as_ref().unwrap(), tokens, top_k)
    };
    let rhs_layout = if sorted.is_some() {
        [routes; 2]
    } else {
        [top_k, index_stride]
    };
    let h = Buffer::mtplx_bytes(u64::from(routes) * u64::from(intermediate) * 2)?;
    let y = Buffer::mtplx_bytes(u64::from(routes) * u64::from(k) * 2)?;
    let activation_temporary = match weights.input {
        GateUpWeights::Fused {
            pack: [w, s, b],
            quant,
        } => {
            let gu = Buffer::mtplx_bytes(u64::from(routes) * u64::from(intermediate) * 4)?;
            gather_bf16_broadcast(
                [w, s, b, x, lhs, rhs, &gu],
                [batches, 1, intermediate * 2, k, experts, routes],
                quant,
                sorted.is_some(),
                true,
                repeat,
                rhs_layout,
            )?;
            let activated = gate_up_bf16(&gu, &h, [routes, intermediate])?;
            vec![gu, activated]
        }
        GateUpWeights::Separate {
            gate,
            up,
            gate_quant,
            up_quant,
        } => {
            let gate_out = Buffer::mtplx_bytes(u64::from(routes) * u64::from(intermediate) * 2)?;
            let up_out = Buffer::mtplx_bytes(u64::from(routes) * u64::from(intermediate) * 2)?;
            // SwitchGLU constructs Up before Gate; Qwen3NextMLP uses the reverse.
            for ([w, s, b], quant, output) in
                [(up, up_quant, &up_out), (gate, gate_quant, &gate_out)]
            {
                gather_bf16_broadcast(
                    [w, s, b, x, lhs, rhs, output],
                    [batches, 1, intermediate, k, experts, routes],
                    quant,
                    sorted.is_some(),
                    true,
                    repeat,
                    rhs_layout,
                )?;
            }
            swiglu_bf16(&gate_out, &up_out, &h, routes * intermediate)?;
            vec![gate_out, up_out]
        }
    };
    let [w, s, b] = weights.down;
    gather_bf16_broadcast(
        [w, s, b, &h, &route_range, rhs, &y],
        [routes, 1, k, intermediate, experts, routes],
        weights.down_quant,
        sorted.is_some(),
        true,
        1,
        rhs_layout,
    )?;
    let hidden = if let Some(sorted) = &sorted {
        let unsorted = Buffer::mtplx_bytes(u64::from(routes) * u64::from(k) * 2)?;
        gather_front(
            &y,
            &sorted.inverse,
            &unsorted,
            [routes, k, routes],
            SortDtype::Bf16,
        )?;
        unsorted
    } else {
        y.view(0, u64::from(routes) * u64::from(k) * 2)?
    };
    let mut temporary = vec![route_range, h, y];
    temporary.extend(activation_temporary);
    temporary.extend(token_range);
    temporary.extend(flat_indices);
    Ok(SwitchOutput {
        hidden,
        _temporaries: temporary,
        _sorted: sorted,
    })
}

#[cfg(test)]
pub(super) struct MoeWeights<'a> {
    pub(super) experts: SwitchWeights<'a>,
    pub(super) shared: SwitchWeights<'a>,
    pub(super) router: [&'a Buffer; 3],
    pub(super) router_quant: [u32; 2],
    pub(super) shared_gate: [&'a Buffer; 3],
    pub(super) shared_gate_quant: [u32; 2],
}

#[cfg(test)]
pub(super) struct MoeOutput {
    pub(super) hidden: Buffer,
    pub(super) used_fused: bool,
    _router: RouterOutput,
    _switch: Option<SwitchOutput>,
    _temporaries: Vec<Buffer>,
}

/// Connected SparseMoeBlock with separate or sanitize-fused expert/shared MLPs. The two flags
/// are MTPLX_FUSED_MOE_DECODE / VERIFY; both default off in the reference.
#[cfg(test)]
pub(super) fn sparse_moe(
    weights: MoeWeights<'_>,
    input: &Buffer,
    [tokens, experts, k, intermediate, shared_intermediate]: [u32; 5],
    normalize: bool,
    [fused_decode, fused_verify]: [bool; 2],
) -> Result<MoeOutput, String> {
    let count = u64::from(tokens) * u64::from(k);
    if tokens == 0
        || experts < 10
        || k == 0
        || count * 10 > i32::MAX as u64
        || u64::from(tokens) * u64::from(experts) > i32::MAX as u64
        || [intermediate, shared_intermediate]
            .iter()
            .any(|d| *d == 0 || u64::from(tokens) * u64::from(*d) * 20 > i32::MAX as u64)
    {
        return Err("invalid MTPLX sparse MoE shape".into());
    }
    let logits = Buffer::mtplx_bytes(u64::from(tokens) * u64::from(experts) * 2)?;
    let [w, s, b] = weights.router;
    let [bits, group] = weights.router_quant;
    let mut temporary = affine_bf16(
        [w, s, b, input, &logits],
        [tokens, experts, k],
        bits,
        group,
        true,
    )?;
    let routing = router(logits, [tokens, experts], normalize)?;
    let used_fused = ((tokens == 1 && fused_decode) || ((2..=4).contains(&tokens) && fused_verify))
        && matches!(
            weights.experts.input,
            GateUpWeights::Fused {
                quant: [4, 32 | 64],
                ..
            }
        )
        && weights.experts.down_quant[0] == 4
        && matches!(weights.experts.down_quant[1], 32 | 64);
    let hidden = Buffer::mtplx_bytes(count * 2)?;
    let switched = if used_fused {
        if k != 2560 || intermediate != 640 {
            return Err("MTPLX fused MoE requires the installed 2560/640 model shape".into());
        }
        temporary.extend(fused_experts(
            &weights.experts,
            input,
            &routing,
            &hidden,
            [tokens, experts],
        )?);
        None
    } else {
        let switched = switch_glu_strided(
            weights.experts,
            input,
            &routing.indices,
            [tokens, 10, k, intermediate, experts],
            routing.index_row_stride,
        )?;
        let grid = [k, tokens * 10, 1];
        dispatch_geometry(
            "g2_Multiplybfloat16",
            &[
                tensor(0, &switched.hidden).input((count * 10) as usize),
                tensor(1, &routing.scores).input((u64::from(tokens) * 10) as usize),
                tensor(2, &switched.hidden).output((count * 10) as usize),
                bytes(3, &[u64::from(k), 1]),
                bytes(4, &[1_u64, 0]),
            ],
            &[],
            grid,
            block_dims(grid),
            true,
        )?;
        let gx = k.div_ceil(4).min(32);
        dispatch(
            "col_reduce_small_1_reduce_sumbfloat16",
            &[
                tensor(0, &switched.hidden).input((count * 10) as usize),
                tensor(1, &hidden).output(count as usize),
                bytes(2, &10_u64),
                bytes(3, &u64::from(k)),
                bytes(4, &tokens),
                bytes(5, &(u64::from(k) * 10)),
                bytes(6, &1_u32),
                bytes(7, &10_u32),
                bytes(8, &u64::from(k)),
                bytes(9, &1_u32),
                bytes(10, &1_u64),
            ],
            [k.div_ceil(4).div_ceil(gx), tokens, 1],
            [gx, 8, 1],
        )?;
        Some(switched)
    };
    // The fused branch evaluates sigmoid(shared_gate(x)) before shared_expert(x);
    // the parent stock branch evaluates the MLP first. Keep both source orders.
    let gate = || {
        shared_gate_sigmoid(
            weights.shared_gate,
            weights.shared_gate_quant,
            input,
            [tokens, k],
        )
    };
    let early_gate = if used_fused { Some(gate()?) } else { None };
    let shared_h = Buffer::mtplx_bytes(u64::from(tokens) * u64::from(shared_intermediate) * 2)?;
    let shared = Buffer::mtplx_bytes(count * 2)?;
    match weights.shared.input {
        GateUpWeights::Fused {
            pack: [w, s, b],
            quant: [bits, group],
        } => {
            let gu = Buffer::mtplx_bytes(u64::from(tokens) * u64::from(shared_intermediate) * 4)?;
            temporary.extend(affine_bf16(
                [w, s, b, input, &gu],
                [tokens, shared_intermediate * 2, k],
                bits,
                group,
                true,
            )?);
            temporary.push(gate_up_bf16(&gu, &shared_h, [tokens, shared_intermediate])?);
            temporary.push(gu);
        }
        GateUpWeights::Separate {
            gate,
            up,
            gate_quant,
            up_quant,
        } => {
            let gate_out =
                Buffer::mtplx_bytes(u64::from(tokens) * u64::from(shared_intermediate) * 2)?;
            let up_out =
                Buffer::mtplx_bytes(u64::from(tokens) * u64::from(shared_intermediate) * 2)?;
            for ([w, s, b], [bits, group], output) in
                [(gate, gate_quant, &gate_out), (up, up_quant, &up_out)]
            {
                temporary.extend(affine_bf16(
                    [w, s, b, input, output],
                    [tokens, shared_intermediate, k],
                    bits,
                    group,
                    true,
                )?);
            }
            swiglu_bf16(&gate_out, &up_out, &shared_h, tokens * shared_intermediate)?;
            temporary.extend([gate_out, up_out]);
        }
    }
    let [w, s, b] = weights.shared.down;
    let [bits, group] = weights.shared.down_quant;
    temporary.extend(affine_bf16(
        [w, s, b, &shared_h, &shared],
        [tokens, k, shared_intermediate],
        bits,
        group,
        true,
    )?);
    let (shared_gate, gate_temporary) = match early_gate {
        Some(result) => result,
        None => gate()?,
    };
    temporary.extend(gate_temporary);
    if tokens == 1 {
        let work = if k < 65536 { 1 } else { 4 };
        let threads = k.div_ceil(work);
        dispatch_geometry(
            if work == 1 {
                "sv_Multiplybfloat16"
            } else {
                "svn_Multiplybfloat16"
            },
            &[
                tensor(0, &shared_gate).input(tokens as usize),
                tensor(1, &shared).input(count as usize),
                tensor(2, &shared).output(count as usize),
                bytes(3, &k),
            ],
            &[],
            [threads, 1, 1],
            [threads.min(1024), 1, 1],
            true,
        )?;
    } else {
        let grid = [k, tokens, 1];
        dispatch_geometry(
            "g2_Multiplybfloat16",
            &[
                tensor(0, &shared_gate).input(tokens as usize),
                tensor(1, &shared).input(count as usize),
                tensor(2, &shared).output(count as usize),
                bytes(3, &[1_u64, 0]),
                bytes(4, &[u64::from(k), 1]),
            ],
            &[],
            grid,
            block_dims(grid),
            true,
        )?;
    }
    let count = count as u32;
    let work = if count < 65536 { 1 } else { 4 };
    let threads = count.div_ceil(work);
    dispatch_geometry(
        if work == 1 {
            "vv_Addbfloat16"
        } else {
            "vvn_Addbfloat16"
        },
        &[
            tensor(0, &hidden).input(count as usize),
            tensor(1, &shared).input(count as usize),
            tensor(2, &hidden).output(count as usize),
            bytes(3, &count),
        ],
        &[],
        [threads, 1, 1],
        [threads.min(1024), 1, 1],
        true,
    )?;
    temporary.extend([shared_h, shared, shared_gate]);
    Ok(MoeOutput {
        hidden,
        used_fused,
        _router: routing,
        _switch: switched,
        _temporaries: temporary,
    })
}

#[cfg(test)]
fn shared_gate_sigmoid(
    [w, s, b]: [&Buffer; 3],
    [bits, group]: [u32; 2],
    input: &Buffer,
    [tokens, k]: [u32; 2],
) -> Result<(Buffer, Vec<Buffer>), String> {
    let gate = Buffer::mtplx_bytes(u64::from(tokens) * 2)?;
    let temporary = affine_bf16([w, s, b, input, &gate], [tokens, 1, k], bits, group, true)?;
    let work = if tokens < 65536 { 1 } else { 4 };
    let threads = tokens.div_ceil(work);
    dispatch_geometry(
        if work == 1 {
            "v_Sigmoidbfloat16bfloat16"
        } else {
            "vn_Sigmoidbfloat16bfloat16"
        },
        &[
            tensor(0, &gate).input(tokens as usize),
            tensor(1, &gate).output(tokens as usize),
            bytes(2, &tokens),
        ],
        &[],
        [threads, 1, 1],
        [threads.min(1024), 1, 1],
        true,
    )?;
    Ok((gate, temporary))
}

#[cfg(test)]
fn fused_experts(
    weights: &SwitchWeights<'_>,
    input: &Buffer,
    routing: &RouterOutput,
    output: &Buffer,
    [tokens, experts]: [u32; 2],
) -> Result<Vec<Buffer>, String> {
    let GateUpWeights::Fused {
        pack: gu_pack,
        quant: [4, gu_group],
    } = weights.input
    else {
        return Err("MTPLX fused MoE requires a fused Q4 gate/up pack".into());
    };
    // Entry conditions are checked by SparseMoeBlock, including the installed
    // dimensions for these template instantiations. Validate all six packs.
    for (pack, n, k, group) in [
        (gu_pack, 1280_u64, 2560_u64, gu_group),
        (weights.down, 2560, 640, weights.down_quant[1]),
    ] {
        let count = u64::from(experts) * n * k;
        for (buffer, size) in pack.into_iter().zip([
            count / 2,
            count / u64::from(group) * 2,
            count / u64::from(group) * 2,
        ]) {
            if unsafe { ds4_gpu_tensor_bytes(buffer.raw()) } < size {
                return Err("MTPLX fused MoE pack smaller than shape".into());
            }
        }
    }
    let indices = flatten_indices(&routing.indices, [tokens, 10, routing.index_row_stride])?;
    let routes = tokens * 10;
    let scores = Buffer::mtplx_bytes(u64::from(routes) * 4)?;
    dispatch_geometry(
        "v_copybfloat16float32",
        &[
            tensor(0, &routing.scores).input(routes as usize),
            tensor(1, &scores).output(routes as usize),
            bytes(2, &routes),
        ],
        &[],
        [routes, 1, 1],
        [routes, 1, 1],
        true,
    )?;
    let h = Buffer::mtplx_bytes(u64::from(routes) * 640 * 2)?;
    let suffix = |group| {
        if tokens == 1 {
            format!("g{group}_bf16")
        } else {
            format!("m_m{tokens}_g{group}_bf16")
        }
    };
    let [w, s, b] = gu_pack;
    dispatch(
        &format!("kernel_qwen_mtplx_moe_glu_h_{}", suffix(gu_group)),
        &[
            tensor(0, input).input(tokens as usize * 2560),
            tensor(1, w).input(experts as usize * 1280 * 2560 / 8),
            tensor(2, s).input(experts as usize * 1280 * 2560 / gu_group as usize),
            tensor(3, b).input(experts as usize * 1280 * 2560 / gu_group as usize),
            tensor(4, &indices).input(routes as usize),
            tensor(5, &h).output(routes as usize * 640),
        ],
        [(routes * 640).div_ceil(32), 1, 1],
        [1024, 1, 1],
    )?;
    let [w, s, b] = weights.down;
    dispatch(
        &format!(
            "kernel_qwen_mtplx_moe_down_y_{}",
            suffix(weights.down_quant[1])
        ),
        &[
            tensor(0, &h).input(routes as usize * 640),
            tensor(1, w).input(experts as usize * 2560 * 640 / 8),
            tensor(2, s).input(experts as usize * 2560 * 640 / weights.down_quant[1] as usize),
            tensor(3, b).input(experts as usize * 2560 * 640 / weights.down_quant[1] as usize),
            tensor(4, &indices).input(routes as usize),
            tensor(5, &scores).input(routes as usize),
            tensor(6, output).output(tokens as usize * 2560),
        ],
        [(tokens * 2560).div_ceil(32), 1, 1],
        [1024, 1, 1],
    )?;
    Ok(vec![indices, scores, h])
}

/// Keep returned intermediate buffers alive until the enclosing command batch completes.
#[cfg(test)]
pub(super) fn affine_bf16(
    buffers: [&Buffer; 5],
    shape: [u32; 3],
    bits: u32,
    group: u32,
    nax: bool,
) -> Result<Vec<Buffer>, String> {
    affine_bf16_layout(buffers, shape, bits, group, nax, None)
}

fn affine_bf16_layout(
    buffers: [&Buffer; 5],
    shape: [u32; 3],
    bits: u32,
    group: u32,
    nax: bool,
    input_layout: Option<[(u64, usize); 4]>,
) -> Result<Vec<Buffer>, String> {
    let [m, n, k] = shape;
    let plan = affine_plan(shape, bits, group, nax)?;
    let expected = [
        u64::from(n) * u64::from(k) / (8 / u64::from(bits)),
        u64::from(n) * u64::from(k / group) * 2,
        u64::from(n) * u64::from(k / group) * 2,
        u64::from(m) * u64::from(k) * 2,
        u64::from(m) * u64::from(n) * 2,
    ];
    let elements: [usize; 5] = std::array::from_fn(|i| (expected[i] / [4, 2, 2, 2, 2][i]) as usize);
    let input_layout = input_layout.unwrap_or_else(|| std::array::from_fn(|i| (0, elements[i])));
    for (i, (buffer, size)) in buffers.iter().zip(expected).enumerate() {
        let offset = input_layout.get(i).map_or(0, |layout| layout.0);
        if offset
            .checked_add(size)
            .is_none_or(|end| unsafe { ds4_gpu_tensor_bytes(buffer.raw()) } < end)
        {
            return Err("MTPLX affine binding is smaller than its shape".into());
        }
    }
    let mut temporary = Vec::new();
    if plan.splits > 1 {
        temporary.push(Buffer::mtplx_bytes(expected[4] * u64::from(plan.splits))?);
        encoder::add_temporaries(&temporary);
    }
    let output = temporary.first().unwrap_or(buffers[4]);
    let mut bindings = buffers[..4]
        .iter()
        .enumerate()
        .map(|(i, b)| {
            tensor(i as u32, b)
                .at_byte_offset(input_layout[i].0)
                .input(input_layout[i].1)
        })
        .collect::<Vec<_>>();
    bindings.push(tensor(4, output).output(elements[4] * plan.splits as usize));
    bindings.extend(
        plan.parameters
            .iter()
            .enumerate()
            .map(|(i, p)| bytes(i as u32 + 5, p)),
    );
    dispatch(&plan.name, &bindings, plan.groups, plan.threads)?;
    if plan.splits > 1
        && let Some(extra) = reduce_split_bf16(output, buffers[4], plan.splits, m * n)?
    {
        temporary.push(extra);
    }
    Ok(temporary)
}

fn reduce_split_bf16(
    input: &Buffer,
    output: &Buffer,
    splits: u32,
    width: u32,
) -> Result<Option<Buffer>, String> {
    // ColReduceArgs for contiguous [split, M, N], reducing axis 0.
    let size = u64::from(splits);
    let stride = u64::from(width);
    let zero = 0_u32;
    let zero_stride = 0_u64;
    let ndim = 1_u32;
    let non_col = 1_u64;
    let two_pass = splits > 256 && width / 32 < 1024;
    let intermediate = if two_pass {
        Some(Buffer::mtplx_bytes(stride * 32 * 2)?)
    } else {
        None
    };
    if let Some(buffer) = &intermediate {
        encoder::add_temporaries(std::slice::from_ref(buffer));
    }
    let mut bindings = vec![
        tensor(0, input).input((size * stride) as usize),
        tensor(1, intermediate.as_ref().unwrap_or(output))
            .output((stride * if two_pass { 32 } else { 1 }) as usize),
        bytes(2, &size),
        bytes(3, &stride),
        bytes(4, &zero),
        bytes(5, &zero_stride),
        bytes(6, &zero),
        bytes(7, &splits),
        bytes(8, &stride),
        bytes(9, &ndim),
        bytes(10, &non_col),
    ];
    let large = if size * stride > i32::MAX as u64 {
        "_large"
    } else {
        ""
    };
    if splits < 32 {
        let x = width.div_ceil(4).min(32);
        dispatch(
            &format!("col_reduce_small{large}_1_reduce_sumbfloat16"),
            &bindings,
            [width.div_ceil(4).div_ceil(x), 1, 1],
            [x, splits.min(8), 1],
        )?;
    } else if two_pass {
        bindings.push(bytes(11, &non_col));
        dispatch(
            &format!("col_reduce_2pass{large}_1_32_32_reduce_sumbfloat16"),
            &bindings,
            [width.div_ceil(32), 32, 1],
            [256, 1, 1],
        )?;
        reduce_split_bf16(intermediate.as_ref().unwrap(), output, 32, width)?;
    } else {
        dispatch(
            &format!("col_reduce_looped{large}_1_32_32_reduce_sumbfloat16"),
            &bindings,
            [width.div_ceil(32), 1, 1],
            [256, 1, 1],
        )?;
    }
    Ok(intermediate)
}

#[cfg(test)]
pub(super) struct ConvNormOutput {
    pub(super) q: Buffer,
    pub(super) k: Buffer,
    pub(super) v: Buffer,
    pub(super) conv: Buffer,
    /// Pre-convolution rows retained by the reference's verify capture.
    pub(super) qkv: Buffer,
    pub(super) used_fused: bool,
    _temporaries: Vec<Buffer>,
}

/// Stock GatedDeltaNet._l2norm, including its BF16 boundary before Q scaling.
#[cfg(test)]
fn gdn_l2_norm(input: &Buffer, output: &Buffer, rows: u32, q: bool) -> Result<Vec<Buffer>, String> {
    let count = rows * 2048;
    let heads = rows * 16;
    let xf = Buffer::mtplx_bytes(u64::from(count) * 4)?;
    let sq = Buffer::mtplx_bytes(u64::from(count) * 4)?;
    let sum = Buffer::mtplx_bytes(u64::from(heads) * 4)?;
    if rows == 1 {
        cast_bf16_f32(input, &xf, count, true)?;
    } else {
        let grid = [2048, rows, 1];
        dispatch_geometry(
            "g2_copybfloat16float32",
            &[
                tensor(0, input).input(((rows - 1) * 10240 + 2048) as usize),
                tensor(1, &xf).output(count as usize),
                bytes(3, &[10240_u64, 1]),
            ],
            &[],
            grid,
            block_dims(grid),
            true,
        )?;
    }
    let work = if count < 65536 { 1 } else { 2 };
    let threads = count.div_ceil(work);
    dispatch_geometry(
        if work == 1 {
            "vv_Multiplyfloat32"
        } else {
            "vvn_Multiplyfloat32"
        },
        &[
            tensor(0, &xf).input(count as usize),
            tensor(1, &xf).input(count as usize),
            tensor(2, &sq).output(count as usize),
            bytes(3, &count),
        ],
        &[],
        [threads, 1, 1],
        [threads.min(1024), 1, 1],
        true,
    )?;
    if heads >= 32 {
        dispatch_geometry(
            "row_reduce_simple_sumfloat32",
            &[
                tensor(0, &sq).input(count as usize),
                tensor(1, &sum).output(heads as usize),
                bytes(2, &128_u64),
                bytes(3, &u64::from(heads)),
            ],
            &[],
            [32, heads.div_ceil(4), 1],
            [32, 1, 1],
            true,
        )?;
    } else {
        dispatch_geometry(
            "row_reduce_looped_1_reduce_sumfloat32",
            &[
                tensor(0, &sq).input(count as usize),
                tensor(1, &sum).output(heads as usize),
                bytes(2, &128_u64),
                bytes(3, &1_u64),
                bytes(4, &heads),
                bytes(5, &128_u64),
                bytes(6, &1_u32),
                bytes(7, &0_u32),
                bytes(8, &0_u64),
                bytes(9, &0_u32),
            ],
            &[],
            [32, heads, 1],
            [32, 1, 1],
            true,
        )?;
    }
    let epsilon = scalar_buffer(&1e-6_f32.to_le_bytes())?;
    let work = if heads < 65536 { 1 } else { 2 };
    let threads = heads.div_ceil(work);
    dispatch_geometry(
        if work == 1 {
            "vs_Addfloat32"
        } else {
            "vsn_Addfloat32"
        },
        &[
            tensor(0, &sum).input(heads as usize),
            tensor(1, &epsilon).input(1),
            tensor(2, &sum).output(heads as usize),
            bytes(3, &heads),
        ],
        &[],
        [threads, 1, 1],
        [threads.min(1024), 1, 1],
        true,
    )?;
    dispatch_geometry(
        if work == 1 {
            "v_Rsqrtfloat32float32"
        } else {
            "vn_Rsqrtfloat32float32"
        },
        &[
            tensor(0, &sum).input(heads as usize),
            tensor(1, &sum).output(heads as usize),
            bytes(2, &heads),
        ],
        &[],
        [threads, 1, 1],
        [threads.min(1024), 1, 1],
        true,
    )?;
    let grid = [128, heads, 1];
    dispatch_geometry(
        "g2_Multiplyfloat32",
        &[
            tensor(0, &xf).input(count as usize),
            tensor(1, &sum).input(heads as usize),
            tensor(2, &xf).output(count as usize),
            bytes(3, &[128_u64, 1]),
            bytes(4, &[1_u64, 0]),
        ],
        &[],
        grid,
        block_dims(grid),
        true,
    )?;
    cast_bf16_f32(&xf, output, count, false)?;
    if q {
        // Weak Python scalar converts to BF16 BEFORE Multiplybfloat16.
        let bits = 128_f32.powf(-0.5).to_bits();
        let scale = ((bits + 0x7fff + ((bits >> 16) & 1)) >> 16) as u16;
        let scale = scalar_buffer(&scale.to_le_bytes())?;
        let work = if count < 65536 { 1 } else { 4 };
        let threads = count.div_ceil(work);
        dispatch_geometry(
            if work == 1 {
                "sv_Multiplybfloat16"
            } else {
                "svn_Multiplybfloat16"
            },
            &[
                tensor(0, &scale).input(1),
                tensor(1, output).input(count as usize),
                tensor(2, output).output(count as usize),
                bytes(3, &count),
            ],
            &[],
            [threads, 1, 1],
            [threads.min(1024), 1, 1],
            true,
        )?;
    }
    Ok(vec![xf, sq, sum])
}

/// GDN conv/state and Q/K/V preparation for B=1. `fused` contains the
/// model's two conv-norm flags; cache absence/lengths/mask suppress fusion.
/// `length` is the cache's optional one-element INT32 valid-length tensor.
#[cfg(test)]
pub(super) fn gdn_conv_norm(
    input: &Buffer,
    weight: &Buffer,
    [rows, stride]: [u32; 2],
    [state, length]: [Option<&Buffer>; 2],
    mask: Option<&Buffer>,
    fused: [bool; 2],
    has_cache: bool,
) -> Result<ConvNormOutput, String> {
    let count = u64::from(rows) * 10240;
    if rows == 0
        || stride < 10240
        || (u64::from(rows) + 3) * 10240 > i32::MAX as u64
        || unsafe { ds4_gpu_tensor_bytes(input.raw()) }
            < (u64::from(rows - 1) * u64::from(stride) + 10240) * 2
        || unsafe { ds4_gpu_tensor_bytes(weight.raw()) } < 10240 * 4 * 2
        || state.is_some_and(|x| unsafe { ds4_gpu_tensor_bytes(x.raw()) } < 3 * 10240 * 2)
        || mask.is_some_and(|x| unsafe { ds4_gpu_tensor_bytes(x.raw()) } < u64::from(rows))
        || length.is_some_and(|x| unsafe { ds4_gpu_tensor_bytes(x.raw()) } < 4)
    {
        return Err("invalid MTPLX GDN conv inputs".into());
    }
    let mut temporary = Vec::new();
    let input_span = ((rows - 1) as usize * stride as usize) + 10240;
    let qkv = if let Some(mask) = mask {
        let qkv = Buffer::mtplx_bytes(count * 2)?;
        let zero = scalar_buffer(&0_u16.to_le_bytes())?;
        let grid = [10240, rows, 1];
        dispatch_geometry(
            if rows == 1 {
                "g1_Selectbfloat16"
            } else {
                "g2_Selectbfloat16"
            },
            &[
                tensor(0, mask).input(rows as usize),
                tensor(1, input).input(input_span),
                tensor(2, &zero).input(1),
                tensor(3, &qkv).output(count as usize),
                bytes(4, &if rows == 1 { [0_u64, 0] } else { [1_u64, 0] }),
                bytes(
                    5,
                    &if rows == 1 {
                        [1_u64, 0]
                    } else {
                        [u64::from(stride), 1]
                    },
                ),
                bytes(6, &[0_u64, 0]),
            ],
            &[],
            grid,
            block_dims(grid),
            true,
        )?;
        qkv
    } else {
        input.view(0, (u64::from(rows - 1) * u64::from(stride) + 10240) * 2)?
    };
    let stride = if mask.is_some() { 10240 } else { stride };
    let qkv_span = (rows - 1) as usize * stride as usize + 10240;
    let q = Buffer::mtplx_bytes(u64::from(rows) * 2048 * 2)?;
    let k = Buffer::mtplx_bytes(u64::from(rows) * 2048 * 2)?;
    let v;
    let conv;
    let used_fused = has_cache
        && mask.is_none()
        && length.is_none()
        && ((rows == 1 && fused[0]) || ((2..=6).contains(&rows) && fused[1]));
    if used_fused {
        v = Buffer::mtplx_bytes(u64::from(rows) * 6144 * 2)?;
        conv = Buffer::mtplx_bytes(3 * 10240 * 2)?;
        let packed;
        let source = if rows == 1 || stride == 10240 {
            &qkv
        } else {
            packed = Buffer::mtplx_bytes(count * 2)?;
            let grid = [10240, rows, 1];
            dispatch_geometry(
                "g2_copybfloat16bfloat16",
                &[
                    tensor(0, &qkv).input(qkv_span),
                    tensor(1, &packed).output(count as usize),
                    bytes(3, &[u64::from(stride), 1]),
                ],
                &[],
                grid,
                block_dims(grid),
                true,
            )?;
            temporary.push(packed);
            temporary.last().unwrap()
        };
        let source = source.view(0, count * 2)?;
        let zero;
        let state = if let Some(state) = state {
            state
        } else {
            zero = Buffer::mtplx_bytes(3 * 10240 * 2)?;
            let scalar = scalar_buffer(&0_u16.to_le_bytes())?;
            // CustomKernel materializes the broadcast scalar with General copy.
            dispatch_geometry(
                "g1_copybfloat16bfloat16",
                &[
                    tensor(0, &scalar).input(1),
                    tensor(1, &zero).output(3 * 10240),
                    bytes(3, &0_u64),
                ],
                &[],
                [3 * 10240, 1, 1],
                [1024, 1, 1],
                true,
            )?;
            temporary.push(zero);
            temporary.last().unwrap()
        };
        let state = state.view(0, 3 * 10240 * 2)?;
        dispatch(
            &if rows == 1 {
                "kernel_qwen_mtplx_gdn_conv_norm_bf16".into()
            } else {
                format!("kernel_qwen_mtplx_gdn_conv_norm_rows_s{rows}_bf16")
            },
            &[
                tensor(0, &source).input(count as usize),
                tensor(1, &state).input(3 * 10240),
                tensor(2, weight).input(4 * 10240),
                tensor(3, &q).output(rows as usize * 2048),
                tensor(4, &k).output(rows as usize * 2048),
                tensor(5, &v).output(rows as usize * 6144),
                tensor(6, &conv).output(3 * 10240),
            ],
            [10, 1, 1],
            [1024, 1, 1],
        )?;
        encoder::add_temporaries(&temporary);
    } else {
        conv = Buffer::mtplx_bytes(3 * 10240 * 2)?;
        let window = Buffer::mtplx_bytes((count + 3 * 10240) * 2)?;
        let tail = window.view(3 * 10240 * 2, count * 2)?;
        // Original concatenate_gpu uses concurrent disjoint out_slice writers.
        let concurrent = encoder::concurrent_scope();
        let zero = state
            .is_none()
            .then(|| scalar_buffer(&0_u16.to_le_bytes()))
            .transpose()?;
        // Concatenate uses GeneralGeneral copies, even for contiguous slices.
        dispatch_geometry(
            "gg1_copybfloat16bfloat16",
            &[
                tensor(0, state.or(zero.as_ref()).unwrap()).input(if state.is_some() {
                    3 * 10240
                } else {
                    1
                }),
                tensor(1, &window).output(3 * 10240),
                bytes(3, &u64::from(state.is_some())),
                bytes(4, &1_u64),
            ],
            &[],
            [3 * 10240, 1, 1],
            [1024, 1, 1],
            true,
        )?;
        let contiguous = rows == 1 || stride == 10240;
        let grid = if contiguous {
            [count as u32, 1, 1]
        } else {
            [10240, rows, 1]
        };
        dispatch_geometry(
            if contiguous {
                "gg1_copybfloat16bfloat16"
            } else {
                "gg2_copybfloat16bfloat16"
            },
            &[
                tensor(0, &qkv).input(qkv_span),
                tensor(1, &tail).output(count as usize),
                bytes(
                    3,
                    &if contiguous {
                        [1_u64, 0]
                    } else {
                        [u64::from(stride), 1]
                    },
                ),
                bytes(
                    4,
                    &if contiguous {
                        [1_u64, 0]
                    } else {
                        [10240_u64, 1]
                    },
                ),
            ],
            &[],
            grid,
            block_dims(grid),
            true,
        )?;
        drop(concurrent);
        if let Some(length) = length {
            let zero = scalar_buffer(&0_i32.to_le_bytes())?;
            let upper = scalar_buffer(&(rows as i32).to_le_bytes())?;
            let end = Buffer::mtplx_bytes(4)?;
            let positions = Buffer::mtplx_bytes(12)?;
            let range = Buffer::mtplx_bytes(12)?;
            dispatch_geometry(
                "vs_Maximumint32",
                &[
                    tensor(0, length).input(1),
                    tensor(1, &zero).input(1),
                    tensor(2, &end).output(1),
                    bytes(3, &1_u32),
                ],
                &[],
                [1, 1, 1],
                [1, 1, 1],
                true,
            )?;
            dispatch_geometry(
                "vs_Minimumint32",
                &[
                    tensor(0, &end).input(1),
                    tensor(1, &upper).input(1),
                    tensor(2, &end).output(1),
                    bytes(3, &1_u32),
                ],
                &[],
                [1, 1, 1],
                [1, 1, 1],
                true,
            )?;
            dispatch_geometry(
                "arangeint32",
                &[
                    bytes(0, &0_i32),
                    bytes(1, &1_i32),
                    tensor(2, &range).output(3),
                ],
                &[],
                [3, 1, 1],
                [3, 1, 1],
                true,
            )?;
            dispatch_geometry(
                "sv_Addint32",
                &[
                    tensor(0, &end).input(1),
                    tensor(1, &range).input(3),
                    tensor(2, &positions).output(3),
                    bytes(3, &3_u32),
                ],
                &[],
                [3, 1, 1],
                [3, 1, 1],
                true,
            )?;
            let grid = [10240, 3, 1];
            dispatch_geometry(
                "kernel_qwen_mtplx_gather_axis_bf16_idxi32_int_10",
                &[
                    tensor(0, &window).input(count as usize + 3 * 10240),
                    tensor(1, &positions).input(3),
                    tensor(2, &conv).output(3 * 10240),
                    bytes(3, &10240_u32),
                    bytes(4, &1_u64),
                    bytes(5, &0_u64),
                    bytes(6, &1_u64),
                    bytes(7, &0_u32),
                    bytes(8, &(rows + 3)),
                    bytes(9, &10240_u64),
                    bytes(10, &1_u64),
                ],
                &[],
                grid,
                block_dims(grid),
                true,
            )?;
            temporary.extend([end, positions, range]);
        } else {
            let view = window.view(u64::from(rows) * 10240 * 2, 3 * 10240 * 2)?;
            dispatch_geometry(
                "g1_copybfloat16bfloat16",
                &[
                    tensor(0, &view).input(3 * 10240),
                    tensor(1, &conv).output(3 * 10240),
                    bytes(3, &1_u64),
                ],
                &[],
                [3 * 10240, 1, 1],
                [1024, 1, 1],
                true,
            )?;
        }
        let convolved = Buffer::mtplx_bytes(count * 2)?;
        let grid = [10240, rows, 1];
        dispatch_geometry(
            "depthwise_conv_1d_bfloat16",
            &[
                tensor(0, &window).input(count as usize + 3 * 10240),
                tensor(1, weight).input(4 * 10240),
                tensor(2, &convolved).output(count as usize),
                bytes(3, &[(rows + 3) * 10240, 10240, 1]),
                bytes(4, &4_u32),
            ],
            &[],
            grid,
            block_dims(grid),
            true,
        )?;
        let count = count as u32;
        let work = if count < 65536 { 1 } else { 4 };
        let threads = count.div_ceil(work);
        dispatch_geometry(
            if work == 1 {
                "kernel_qwen_mtplx_silu_bf16_contiguous"
            } else {
                "kernel_qwen_mtplx_silu_bf16_contiguous_n"
            },
            &[
                tensor(0, &convolved).input(count as usize),
                tensor(1, &convolved).output(count as usize),
                bytes(2, &count),
            ],
            &[],
            [threads, 1, 1],
            [threads.min(1024), 1, 1],
            true,
        )?;
        temporary.extend(gdn_l2_norm(&convolved, &q, rows, true)?);
        let key = convolved.view(2048 * 2, u64::from(count - 2048) * 2)?;
        temporary.extend(gdn_l2_norm(&key, &k, rows, false)?);
        let value = convolved.view(4096 * 2, u64::from(count - 4096) * 2)?;
        if rows == 1 {
            v = value;
        } else {
            v = Buffer::mtplx_bytes(u64::from(rows) * 6144 * 2)?;
            let grid = [6144, rows, 1];
            dispatch_geometry(
                "g2_copybfloat16bfloat16",
                &[
                    tensor(0, &value).input((rows - 1) as usize * 10240 + 6144),
                    tensor(1, &v).output(rows as usize * 6144),
                    bytes(3, &[10240_u64, 1]),
                ],
                &[],
                grid,
                block_dims(grid),
                true,
            )?;
        }
        temporary.extend([window, convolved]);
    }
    Ok(ConvNormOutput {
        q,
        k,
        v,
        conv,
        qkv,
        used_fused,
        _temporaries: temporary,
    })
}

#[cfg(test)]
pub(super) struct DeltaUpdateOutput {
    pub(super) hidden: Buffer,
    pub(super) delta: Buffer,
    pub(super) g: Buffer,
    pub(super) beta: Buffer,
    _initial_state: Option<Buffer>,
}

/// Actual gated_delta_update inference branch, including compute_g, sigmoid,
/// optional zero initial state and mask. Q/K/V are contiguous as required by
/// the reference custom-kernel dispatcher; A/B retain projection row strides.
#[cfg(test)]
pub(super) fn gated_delta_update(
    [q, k, v]: [&Buffer; 3],
    [a, b, a_log, dt_bias]: [&Buffer; 4],
    [rows, gate_stride]: [u32; 2],
    state: Option<&Buffer>,
    mask: Option<&Buffer>,
) -> Result<DeltaUpdateOutput, String> {
    let count = u64::from(rows) * 48;
    let state_count = 48 * 128 * 128_u32;
    if rows == 0 || gate_stride < 48 || u64::from(rows) * 6144 > i32::MAX as u64
        || [q, k].iter().any(|x| unsafe { ds4_gpu_tensor_bytes(x.raw()) } < u64::from(rows) * 2048 * 2)
        || unsafe { ds4_gpu_tensor_bytes(v.raw()) } < u64::from(rows) * 6144 * 2
        || [a, b].iter().any(|x| unsafe { ds4_gpu_tensor_bytes(x.raw()) }
            < (u64::from(rows - 1) * u64::from(gate_stride) + 48) * 2)
        || [a_log, dt_bias].iter().any(|x| unsafe { ds4_gpu_tensor_bytes(x.raw()) } < 48 * 2)
        || state.is_some_and(|x| unsafe { ds4_gpu_tensor_bytes(x.raw()) } < u64::from(state_count) * 4)
        || mask.is_some_and(|x| unsafe { ds4_gpu_tensor_bytes(x.raw()) } < u64::from(rows))
    { return Err("invalid MTPLX gated delta inputs".into()); }
    let gate_span = (rows - 1) as usize * gate_stride as usize + 48;
    let beta = Buffer::mtplx_bytes(count * 2)?;
    let g = Buffer::mtplx_bytes(count * 4)?;
    if rows == 1 || gate_stride == 48 {
        let count = count as u32;
        let work = if count < 65536 { 1 } else { 4 };
        let threads = count.div_ceil(work);
        dispatch_geometry(
            if work == 1 {
                "v_Sigmoidbfloat16bfloat16"
            } else {
                "vn_Sigmoidbfloat16bfloat16"
            },
            &[
                tensor(0, b).input(gate_span),
                tensor(1, &beta).output(count as usize),
                bytes(2, &count),
            ],
            &[],
            [threads, 1, 1],
            [threads.min(1024), 1, 1],
            true,
        )?;
    } else {
        let grid = [48, rows, 1];
        dispatch_geometry(
            "gn1_Sigmoidbfloat16bfloat16",
            &[
                tensor(0, b).input(gate_span),
                tensor(1, &beta).output(count as usize),
                bytes(2, &[rows, 48]),
                bytes(3, &[u64::from(gate_stride), 1]),
                bytes(4, &2_u32),
            ],
            &[],
            grid,
            block_dims(grid),
            true,
        )?;
    }
    if rows == 1 {
        dispatch_geometry(
            "kernel_qwen_mtplx_compute_g_bf16_contiguous",
            &[
                tensor(0, a_log).input(48),
                tensor(1, a).input(gate_span),
                tensor(2, dt_bias).input(48),
                tensor(3, &g).output(count as usize),
                bytes(4, &48_u32),
            ],
            &[],
            [48, 1, 1],
            [48, 1, 1],
            true,
        )?;
    } else {
        let grid = [48, rows, 1];
        dispatch_geometry(
            "kernel_qwen_mtplx_compute_g_bf16_strided_2",
            &[
                tensor(0, a_log).input(48),
                tensor(1, a).input(gate_span),
                tensor(2, dt_bias).input(48),
                bytes(3, &[0_u64, 1, u64::from(gate_stride), 1, 0, 1]),
                tensor(4, &g).output(count as usize),
                bytes(5, &[rows, 48]),
            ],
            &[],
            grid,
            block_dims(grid),
            true,
        )?;
    }
    let hidden = Buffer::mtplx_bytes(u64::from(rows) * 6144 * 2)?;
    let delta = Buffer::mtplx_bytes(u64::from(state_count) * 4)?;
    let initial_state = if state.is_none() {
        let zero = Buffer::mtplx_bytes(u64::from(state_count) * 4)?;
        let scalar = scalar_buffer(&0_f32.to_le_bytes())?;
        dispatch_geometry(
            "g1_copyfloat32float32",
            &[
                tensor(0, &scalar).input(1),
                tensor(1, &zero).output(state_count as usize),
                bytes(3, &0_u64),
            ],
            &[],
            [state_count, 1, 1],
            [1024, 1, 1],
            true,
        )?;
        Some(zero)
    } else {
        None
    };
    let state = state.or(initial_state.as_ref()).unwrap();
    let length = scalar_buffer(&(rows as i32).to_le_bytes())?;
    let mut bindings = vec![
        tensor(0, q).input(rows as usize * 2048),
        tensor(1, k).input(rows as usize * 2048),
        tensor(2, v).input(rows as usize * 6144),
        tensor(3, &g).input(count as usize),
        tensor(4, &beta).input(count as usize),
        tensor(5, state).input(state_count as usize),
        tensor(6, &length).input(1),
    ];
    if let Some(mask) = mask {
        bindings.push(tensor(7, mask).input(rows as usize));
    }
    let index = bindings.len() as u32;
    bindings.extend([
        tensor(index, &hidden).output(rows as usize * 6144),
        tensor(index + 1, &delta).output(state_count as usize),
    ]);
    dispatch(
        if mask.is_some() {
            "kernel_qwen_mtplx_gated_delta_masked_bf16"
        } else {
            "kernel_qwen_mtplx_gated_delta_bf16"
        },
        &bindings,
        [1, 32, 48],
        [32, 4, 1],
    )?;
    if let Some(initial_state) = &initial_state {
        encoder::add_temporaries(std::slice::from_ref(initial_state));
    }
    Ok(DeltaUpdateOutput {
        hidden,
        delta,
        g,
        beta,
        _initial_state: initial_state,
    })
}

#[cfg(test)]
fn qsa_copy_span(shape: [u32; 3], strides: [u64; 3]) -> Result<usize, String> {
    if shape.contains(&0) {
        return Ok(0);
    }
    shape
        .into_iter()
        .zip(strides)
        .try_fold(1_usize, |span, (dim, stride)| {
            usize::try_from(stride)
                .ok()
                .and_then(|stride| (dim as usize - 1).checked_mul(stride))
                .and_then(|n| span.checked_add(n))
                .ok_or_else(|| "QSA copy element span overflow".into())
        })
}

/// GeneralGeneral copy used by the reference's slice updates/concatenations.
/// Collapse only dimensions contiguous in BOTH source and destination.
#[cfg(test)]
fn qsa_copy(
    source: Option<&Buffer>,
    target: &Buffer,
    shape: [u32; 3],
    source_strides: [u64; 3],
    target_strides: [u64; 3],
    width: u32,
    output_elements: usize,
) -> Result<(), String> {
    qsa_copy_typed(
        source,
        target,
        shape,
        source_strides,
        target_strides,
        if width == 2 { "bfloat16" } else { "float32" },
        output_elements,
    )
}

#[cfg(test)]
fn qsa_copy_typed(
    source: Option<&Buffer>,
    target: &Buffer,
    shape: [u32; 3],
    source_strides: [u64; 3],
    target_strides: [u64; 3],
    dtype: &str,
    output_elements: usize,
) -> Result<(), String> {
    if shape.contains(&0) {
        return Ok(());
    }
    let itemsize = match dtype {
        "bool_" => 1,
        "bfloat16" => 2,
        "float32" | "int32" | "uint32" => 4,
        _ => return Err("unsupported QSA copy dtype".into()),
    };
    let zero = source
        .is_none()
        .then(|| scalar_buffer(&[0; 4][..itemsize]))
        .transpose()?;
    qsa_copy_offsets(
        [
            tensor(0, source.or(zero.as_ref()).unwrap()).input(if source.is_some() {
                qsa_copy_span(shape, source_strides)?
            } else {
                1
            }),
            tensor(1, target).output(output_elements),
        ],
        shape,
        source_strides,
        target_strides,
        dtype,
        [None; 2],
    )
}

#[cfg(test)]
fn qsa_copy_offsets(
    buffers: [Binding<'_>; 2],
    shape: [u32; 3],
    source_strides: [u64; 3],
    target_strides: [u64; 3],
    dtype: &str,
    offsets: [Option<&Buffer>; 2],
) -> Result<(), String> {
    if shape.contains(&0) {
        return Ok(());
    }
    let mut dims: Vec<(u32, u64, u64)> = Vec::new();
    for ((n, s), d) in shape.into_iter().zip(source_strides).zip(target_strides) {
        if n == 1 {
            continue;
        }
        if let Some(last) = dims.last_mut()
            && last.1 == u64::from(n) * s
            && last.2 == u64::from(n) * d
        {
            last.0 = last.0.checked_mul(n).ok_or("QSA copy shape overflow")?;
            last.1 = s;
            last.2 = d;
        } else {
            dims.push((n, s, d));
        }
    }
    if dims.is_empty() {
        dims.push((1, 0, 0));
    }
    let mut ss = [0_u64; 3];
    let mut ds = [0_u64; 3];
    let mut grid = [1_u32; 3];
    for (i, (n, s, d)) in dims.iter().enumerate() {
        ss[i] = *s;
        ds[i] = *d;
        grid[dims.len() - 1 - i] = *n;
    }
    let dynamic = offsets.iter().any(Option::is_some);
    let zero = 0_i64;
    let [source, target] = buffers;
    let bindings = [
        source,
        target,
        bytes(3, &ss),
        bytes(4, &ds),
        offsets[0].map_or_else(|| bytes(6, &zero), |b| tensor(6, b).input(1)),
        offsets[1].map_or_else(|| bytes(7, &zero), |b| tensor(7, b).input(1)),
    ];
    dispatch_geometry(
        &format!(
            "gg{}{}_copy{dtype}{dtype}",
            dims.len(),
            if dynamic { "_dynamic" } else { "" }
        ),
        &bindings[..if dynamic { 6 } else { 4 }],
        &[],
        grid,
        block_dims(grid),
        true,
    )
}

/// Original DynamicSlice/Update row-copy stage for B1 BF16 indexer leaves.
/// Start leaves are int32[1] graph values with validated in-bounds semantics;
/// they stay on-device, just as in the reference's private GPU primitives.
/// Returns offset temporaries so callers retain the entire asynchronous chain.
#[cfg(test)]
pub(super) fn qsa_dynamic_copy_rows(
    [source, target]: [&Buffer; 2],
    [source_stride, target_stride]: [u64; 2],
    [rows, width]: [u32; 2],
    starts: [Option<&Buffer>; 2],
    data_sizes: [usize; 2],
) -> Result<Vec<Buffer>, String> {
    if rows == 0
        || width == 0
        || rows.checked_mul(width).is_none_or(|n| n > i32::MAX as u32)
        || source_stride < u64::from(width)
        || target_stride < u64::from(width)
        || starts
            .iter()
            .flatten()
            .any(|b| unsafe { ds4_gpu_tensor_bytes(b.raw()) } < 4)
        || [source, target]
            .into_iter()
            .zip([source_stride, target_stride])
            .any(|(b, stride)| {
                u64::from(rows - 1)
                    .checked_mul(stride)
                    .and_then(|n| n.checked_add(u64::from(width)))
                    .and_then(|n| n.checked_mul(2))
                    .is_none_or(|n| unsafe { ds4_gpu_tensor_bytes(b.raw()) } < n)
            })
    {
        return Err("invalid QSA dynamic copy buffers or geometry".into());
    }
    for ((buffer, size), stride) in [source, target]
        .into_iter()
        .zip(data_sizes)
        .zip([source_stride, target_stride])
    {
        if size < qsa_copy_span([1, rows, width], [0, stride, 1])?
            || size
                .checked_mul(2)
                .is_none_or(|n| n as u64 > unsafe { ds4_gpu_tensor_bytes(buffer.raw()) })
        {
            return Err("invalid QSA dynamic copy array data_size".into());
        }
    }
    let mut offsets = [None, None];
    for (i, (start, stride)) in starts
        .into_iter()
        .zip([source_stride, target_stride])
        .enumerate()
    {
        if let Some(start) = start {
            let indices = array::Array::new(&[1], array::Dtype::I32, start.view(0, 4)?)?;
            let stride = i64::try_from(stride).map_err(|_| "QSA dynamic copy stride overflow")?;
            let offset = views::dynamic_offset(&indices, &[0, stride, 1], &[1])?;
            offsets[i] = Some(offset.buffer().view(0, 8)?);
        }
    }
    qsa_copy_offsets(
        [
            tensor(0, source).input(data_sizes[0]),
            tensor(1, target).output(data_sizes[1]),
        ],
        [1, rows, width],
        [0, source_stride, 1],
        [0, target_stride, 1],
        "bfloat16",
        [offsets[0].as_ref(), offsets[1].as_ref()],
    )?;
    Ok(offsets.into_iter().flatten().collect())
}

#[cfg(test)]
fn qsa_zero(count: u32, width: u32) -> Result<Buffer, String> {
    let output = Buffer::mtplx_bytes(u64::from(count) * u64::from(width))?;
    let zero = scalar_buffer(&[0; 4][..width as usize])?;
    let work = if count < 65536 { 1 } else { 8 / width };
    let threads = count.div_ceil(work);
    let dtype = if width == 2 { "bfloat16" } else { "float32" };
    dispatch_geometry(
        &format!("s{}_copy{dtype}{dtype}", if work == 1 { "" } else { "n" }),
        &[
            tensor(0, &zero).input(1),
            tensor(1, &output).output(count as usize),
            bytes(2, &count),
        ],
        &[],
        [threads, 1, 1],
        [threads.min(1024), 1, 1],
        true,
    )?;
    Ok(output)
}

#[cfg(test)]
fn qsa_grown_capacity(end: u32, current: u32) -> Result<u32, String> {
    end.checked_add(255)
        .map(|n| n / 256 * 256)
        .zip(current.checked_mul(2))
        .map(|(a, b)| a.max(b))
        .filter(|cap| *cap <= i32::MAX as u32 / 512)
        .ok_or_else(|| "QSA capacity overflow".into())
}

/// Four logical state leaves, sharing immutable backing until a later write.
/// Python array identity and its separately shared storage are different things:
/// __setitem__ overwrites the array descriptor, whereas slicing creates an array.
#[cfg(test)]
pub(super) struct QsaArray {
    storage: array::Array,
    pub(super) head_stride: Cell<u64>,
}

#[cfg(test)]
impl QsaArray {
    pub(super) fn new(buffer: Buffer, head_stride: u64) -> Result<Rc<Self>, String> {
        let elements = i32::try_from(unsafe { ds4_gpu_tensor_bytes(buffer.raw()) } / 2)
            .map_err(|_| "QSA backing dimension overflow")?;
        Ok(Rc::new(Self {
            storage: array::Array::new(&[elements], array::Dtype::BF16, buffer)?,
            head_stride: Cell::new(head_stride),
        }))
    }

    pub(super) fn buffer(&self) -> Ref<'_, Buffer> {
        self.storage.buffer()
    }

    fn slice(&self) -> Rc<Self> {
        Rc::new(Self {
            storage: self.storage.share_data(),
            head_stride: Cell::new(self.head_stride.get()),
        })
    }

    fn shared_storage(&self) -> bool {
        self.storage.shared_storage()
    }

    // Installed BF16 vector-copy contract: descriptor and Data must both be
    // exclusive, and at most 16 KiB of the root allocation may be unused.
    fn can_donate_bf16(array: &Rc<Self>, input: &Buffer, output_bytes: u64) -> bool {
        let storage = array.buffer();
        Rc::strong_count(array) == 1
            && array.storage.is_donatable()
            && storage.identity() == input.identity()
            && input.unique_data()
            && unsafe { ds4_gpu_tensor_bytes(input.raw()) } >= output_bytes
            && input.allocation_bytes().saturating_sub(output_bytes) <= 16384
    }
}

#[derive(Clone)]
#[cfg(test)]
pub(super) struct QsaSnapshot {
    pub(super) leaves: [Option<Rc<QsaArray>>; 4], // keys, values, raw, pooled
    pub(super) offset: u32,
    pub(super) pooled_len: u32,
}

/// Installed single-sequence QSACache + KVCache geometry (KV2x256, index128).
/// Array aliases follow MTPLX's state setter; distinct views use storage COW.
#[derive(Default)]
#[cfg(test)]
pub(super) struct QsaCache {
    pub(super) keys: Option<Rc<QsaArray>>,
    pub(super) values: Option<Rc<QsaArray>>,
    pub(super) raw: Option<Rc<QsaArray>>,
    pub(super) pooled: Option<Rc<QsaArray>>,
    pub(super) mirror: Option<Rc<Buffer>>,
    pub(super) offset: u32,
    pub(super) kv_cap: u32,
    pub(super) kv_head_stride: u64,
    pub(super) raw_cap: u32,
    pub(super) pooled_cap: u32,
    pub(super) pooled_len: u32,
    pub(super) mirror_cap: u32,
    pub(super) mirror_strides: [u64; 2], // D, blocks; initial transpose is a VIEW
    reserved_raw: u32,
    reserved_pooled: u32,
}

/// Original precompute_qsa_replay_capacity for the installed ratio4/step256.
/// The fixed ceil(window/4) staging window can exceed the logical frontier.
#[cfg(test)]
pub(super) fn qsa_replay_capacity(
    start: u32,
    window: u32,
    existing: [u32; 2],
) -> Result<[u32; 2], String> {
    let end = start
        .checked_add(window)
        .ok_or("QSA replay frontier overflow")?;
    let staging = window.div_ceil(4);
    let staging_tokens = staging.checked_mul(4).ok_or("QSA replay window overflow")?;
    let [raw, pooled] = [
        existing[0].max(end).max(staging_tokens),
        existing[1].max(end / 4).max(staging),
    ]
    .map(|required| {
        if required == 0 {
            return Ok(0);
        }
        required
            .max(256)
            .checked_next_power_of_two()
            .filter(|n| *n <= i32::MAX as u32 / 512)
            .ok_or_else(|| "QSA replay capacity overflow".to_owned())
    });
    Ok([raw?, pooled?])
}

#[cfg(test)]
impl QsaCache {
    fn grow_index(
        buffer: &mut Option<Rc<QsaArray>>,
        cap: &mut u32,
        target: u32,
    ) -> Result<(), String> {
        if target == 0 || target > i32::MAX as u32 / 512 {
            return Err("invalid QSA capacity".into());
        }
        let grown = qsa_zero(target * 128, 2)?;
        if let Some(old) = buffer.as_ref() {
            qsa_copy(
                Some(&old.buffer()),
                &grown,
                [1, *cap, 128],
                [0, 128, 1],
                [0, 128, 1],
                2,
                target as usize * 128,
            )?;
        }
        *buffer = Some(QsaArray::new(grown, 0)?);
        *cap = target;
        Ok(())
    }

    fn index_cow(buffer: &Option<Rc<QsaArray>>, cap: u32) -> Result<(), String> {
        if let Some(old) = buffer.as_ref()
            && old.shared_storage()
        {
            let copy = Buffer::mtplx_bytes(u64::from(cap) * 128 * 2)?;
            let count = cap * 128;
            let work = if count < 65536 { 1 } else { 4 };
            let threads = count.div_ceil(work);
            dispatch_geometry(
                if work == 1 {
                    "v_copybfloat16bfloat16"
                } else {
                    "vn_copybfloat16bfloat16"
                },
                &[
                    tensor(0, &old.buffer()).input(count as usize),
                    tensor(1, &copy).output(count as usize),
                    bytes(2, &count),
                ],
                &[],
                [threads, 1, 1],
                [threads.min(1024), 1, 1],
                true,
            )?;
            old.storage.replace(copy);
        }
        Ok(())
    }

    pub(super) fn reserve(&mut self, raw: u32, pooled: u32) -> Result<(), String> {
        let raw = raw.max(self.raw_cap).max(self.reserved_raw);
        let pooled = pooled.max(self.pooled_cap).max(self.reserved_pooled);
        if raw < self.offset
            || pooled < self.pooled_len
            || raw > i32::MAX as u32 / 512
            || pooled > i32::MAX as u32 / 512
        {
            return Err("invalid QSA reservation".into());
        }
        self.reserved_raw = raw;
        self.reserved_pooled = pooled;
        if self.raw.is_some() && raw > self.raw_cap {
            Self::grow_index(&mut self.raw, &mut self.raw_cap, raw)?;
        }
        if self.pooled.is_some() && pooled > self.pooled_cap {
            Self::grow_index(&mut self.pooled, &mut self.pooled_cap, pooled)?;
        }
        Ok(())
    }

    /// _ensure_compiled_backings includes pending reservations and materializes
    /// pristine leaves. It neither advances KV.offset nor the pooled frontier.
    pub(super) fn ensure_compiled_backings(&mut self, rows: u32) -> Result<(), String> {
        if rows == 0 {
            return Err("compiled QSA requires positive rows".into());
        }
        let [raw, pooled] = qsa_replay_capacity(
            self.offset,
            rows,
            [
                self.raw_cap.max(self.reserved_raw),
                self.pooled_cap.max(self.reserved_pooled),
            ],
        )?;
        self.reserve(raw, pooled.max(256))?;
        if self.raw.is_none() {
            Self::grow_index(&mut self.raw, &mut self.raw_cap, self.reserved_raw)?;
        }
        if self.pooled.is_none() {
            Self::grow_index(&mut self.pooled, &mut self.pooled_cap, self.reserved_pooled)?;
        }
        Ok(())
    }

    pub(super) fn write_raw(
        &mut self,
        keys: &Buffer,
        rows: u32,
        stride: u64,
    ) -> Result<(), String> {
        let end = self
            .offset
            .checked_add(rows)
            .filter(|n| *n <= i32::MAX as u32 / 512)
            .ok_or("QSA raw range overflow")?;
        let bytes_needed = u64::from(rows.saturating_sub(1))
            .checked_mul(stride)
            .and_then(|n| n.checked_add(128))
            .and_then(|n| n.checked_mul(2));
        if rows == 0
            || stride < 128
            || bytes_needed.is_none_or(|n| unsafe { ds4_gpu_tensor_bytes(keys.raw()) } < n)
        {
            return Err("invalid QSA raw input".into());
        }
        if self.raw.is_none() || end > self.raw_cap {
            let target = qsa_grown_capacity(end, self.raw_cap)?.max(self.reserved_raw);
            Self::grow_index(&mut self.raw, &mut self.raw_cap, target)?;
        } else {
            Self::index_cow(&self.raw, self.raw_cap)?;
        }
        let target = self
            .raw
            .as_ref()
            .unwrap()
            .buffer()
            .view(u64::from(self.offset) * 128 * 2, u64::from(rows) * 128 * 2)?;
        qsa_copy(
            Some(keys),
            &target,
            [1, rows, 128],
            [0, stride, 1],
            [0, 128, 1],
            2,
            self.raw_cap as usize * 128,
        )
    }

    pub(super) fn write_pooled(
        &mut self,
        blocks: &Buffer,
        start: u32,
        total: u32,
    ) -> Result<(), String> {
        if total <= start
            || total > i32::MAX as u32 / 512
            || unsafe { ds4_gpu_tensor_bytes(blocks.raw()) } < u64::from(total - start) * 128 * 2
        {
            return Err("invalid QSA pooled input".into());
        }
        if self.pooled.is_none() || total > self.pooled_cap {
            let target = qsa_grown_capacity(total, self.pooled_cap)?.max(self.reserved_pooled);
            Self::grow_index(&mut self.pooled, &mut self.pooled_cap, target)?;
        } else {
            Self::index_cow(&self.pooled, self.pooled_cap)?;
        }
        let target = self.pooled.as_ref().unwrap().buffer().view(
            u64::from(start) * 128 * 2,
            u64::from(total - start) * 128 * 2,
        )?;
        qsa_copy(
            Some(blocks),
            &target,
            [1, total - start, 128],
            [0, 128, 1],
            [0, 128, 1],
            2,
            self.pooled_cap as usize * 128,
        )?;
        if self.mirror.is_none() {
            self.rebuild_mirror()?;
        } else {
            if self.mirror_cap < self.pooled_cap {
                let grown = qsa_zero(self.pooled_cap * 128, 4)?;
                qsa_copy(
                    self.mirror.as_deref(),
                    &grown,
                    [1, 128, self.mirror_cap],
                    [0, self.mirror_strides[0], self.mirror_strides[1]],
                    [0, u64::from(self.pooled_cap), 1],
                    4,
                    self.pooled_cap as usize * 128,
                )?;
                self.mirror = Some(Rc::new(grown));
                self.mirror_cap = self.pooled_cap;
                self.mirror_strides = [u64::from(self.pooled_cap), 1];
            }
            let f32blocks = Buffer::mtplx_bytes(u64::from(total - start) * 128 * 4)?;
            cast_bf16_f32(blocks, &f32blocks, (total - start) * 128, true)?;
            let at = u64::from(start) * self.mirror_strides[1] * 4;
            let span = (127 * self.mirror_strides[0]
                + u64::from(total - start - 1) * self.mirror_strides[1]
                + 1)
                * 4;
            let target = self.mirror.as_ref().unwrap().view(at, span)?;
            qsa_copy(
                Some(&f32blocks),
                &target,
                [1, 128, total - start],
                [0, 1, 128],
                [0, self.mirror_strides[0], self.mirror_strides[1]],
                4,
                self.mirror_cap as usize * 128,
            )?;
        }
        self.pooled_len = total;
        Ok(())
    }

    fn rebuild_mirror(&mut self) -> Result<(), String> {
        let pooled = self
            .pooled
            .as_ref()
            .ok_or("QSA mirror needs pooled state")?;
        let mirror = Buffer::mtplx_bytes(u64::from(self.pooled_cap) * 128 * 4)?;
        cast_bf16_f32(&pooled.buffer(), &mirror, self.pooled_cap * 128, true)?;
        self.mirror = Some(Rc::new(mirror));
        self.mirror_cap = self.pooled_cap;
        self.mirror_strides = [1, 128];
        Ok(())
    }

    pub(super) fn pooled_f32_view(&mut self, nb: u32) -> Result<(&Buffer, [u64; 2]), String> {
        if nb == 0 || nb > self.pooled_cap {
            return Err("invalid QSA mirror frontier".into());
        }
        if self.mirror.is_none() || self.mirror_cap < nb {
            self.rebuild_mirror()?;
        }
        Ok((self.mirror.as_deref().unwrap(), self.mirror_strides))
    }

    pub(super) fn extend_pooled(
        &mut self,
        total: u32,
        norm: &Buffer,
        freq: &Buffer,
        eager: bool,
    ) -> Result<Vec<Buffer>, String> {
        let nb = total / 4;
        let old = self.pooled_len.min(nb);
        if total > self.raw_cap {
            return Err("QSA pool exceeds raw backing".into());
        }
        if nb > old {
            let fresh = self
                .raw
                .as_ref()
                .ok_or("QSA pool needs raw keys")?
                .buffer()
                .view(
                    u64::from(old) * 4 * 128 * 2,
                    u64::from(nb - old) * 4 * 128 * 2,
                )?;
            let (blocks, mut temporary) = if eager {
                qsa_pool_keys_eager([&fresh, norm, freq], nb - old, [128, 1], old as i32)?
            } else {
                (
                    qsa_pool_keys([&fresh, norm, freq], nb - old, [128, 1], old as i32)?,
                    Vec::new(),
                )
            };
            self.write_pooled(&blocks, old, nb)?;
            temporary.push(blocks);
            return Ok(temporary);
        }
        Ok(Vec::new())
    }

    pub(super) fn update_kv(
        &mut self,
        keys: &Buffer,
        values: &Buffer,
        rows: u32,
        strides: [u64; 2],
    ) -> Result<(), String> {
        let end = self
            .offset
            .checked_add(rows)
            .filter(|n| *n <= i32::MAX as u32 / 512)
            .ok_or("QSA KV range overflow")?;
        let bytes_needed = u64::from(rows.saturating_sub(1))
            .checked_mul(strides[1])
            .and_then(|n| n.checked_add(strides[0]))
            .and_then(|n| n.checked_add(256))
            .and_then(|n| n.checked_mul(2));
        if rows == 0
            || bytes_needed.is_none_or(|n| {
                [keys, values]
                    .iter()
                    .any(|b| unsafe { ds4_gpu_tensor_bytes(b.raw()) } < n)
            })
        {
            return Err("invalid QSA KV input".into());
        }
        if self.keys.is_none() || end > self.kv_cap {
            let keep = if !self.offset.is_multiple_of(256) {
                self.offset
            } else {
                self.kv_cap
            };
            let extra = rows.div_ceil(256) * 256;
            let cap = keep
                .checked_add(extra)
                .filter(|n| *n <= i32::MAX as u32 / 512)
                .ok_or("QSA KV allocation overflow")?;
            for buffer in [&mut self.keys, &mut self.values] {
                let grown = if buffer.is_none() {
                    qsa_zero(cap * 512, 2)?
                } else {
                    Buffer::mtplx_bytes(u64::from(cap) * 512 * 2)?
                };
                let concurrent = encoder::concurrent_scope();
                if let Some(old) = buffer.as_ref() {
                    qsa_copy(
                        Some(&old.buffer()),
                        &grown,
                        [2, keep, 256],
                        [self.kv_head_stride, 256, 1],
                        [u64::from(cap) * 256, 256, 1],
                        2,
                        keep as usize * 512,
                    )?;
                }
                if buffer.is_some() {
                    let tail = grown.view(
                        u64::from(keep) * 256 * 2,
                        (u64::from(cap) * 256 + u64::from(extra) * 256) * 2,
                    )?;
                    qsa_copy(
                        None,
                        &tail,
                        [2, extra, 256],
                        [0; 3],
                        [u64::from(cap) * 256, 256, 1],
                        2,
                        extra as usize * 512,
                    )?;
                }
                drop(concurrent);
                *buffer = Some(QsaArray::new(grown, u64::from(cap) * 256)?);
            }
            self.kv_cap = cap;
            self.kv_head_stride = u64::from(cap) * 256;
        } else {
            // Distinct slices can prevent donation; aliases of the same array
            // object observe its descriptor replacement instead.
            let need_copy = [&self.keys, &self.values]
                .iter()
                .any(|b| b.as_ref().unwrap().shared_storage())
                || self.kv_head_stride != u64::from(self.kv_cap) * 256;
            if need_copy {
                for buffer in [&mut self.keys, &mut self.values] {
                    let copy = Buffer::mtplx_bytes(u64::from(self.kv_cap) * 512 * 2)?;
                    if self.kv_head_stride == u64::from(self.kv_cap) * 256 {
                        let count = self.kv_cap * 512;
                        let work = if count < 65536 { 1 } else { 4 };
                        let threads = count.div_ceil(work);
                        dispatch_geometry(
                            if work == 1 {
                                "v_copybfloat16bfloat16"
                            } else {
                                "vn_copybfloat16bfloat16"
                            },
                            &[
                                tensor(0, &buffer.as_ref().unwrap().buffer()).input(
                                    (self.kv_head_stride + u64::from(self.kv_cap) * 256) as usize,
                                ),
                                tensor(1, &copy).output(self.kv_cap as usize * 512),
                                bytes(2, &count),
                            ],
                            &[],
                            [threads, 1, 1],
                            [threads.min(1024), 1, 1],
                            true,
                        )?;
                    } else {
                        let grid = [self.kv_cap * 256, 2, 1];
                        dispatch_geometry(
                            "g2_copybfloat16bfloat16",
                            &[
                                tensor(0, &buffer.as_ref().unwrap().buffer()).input(
                                    (self.kv_head_stride + u64::from(self.kv_cap) * 256) as usize,
                                ),
                                tensor(1, &copy).output(self.kv_cap as usize * 512),
                                bytes(3, &[self.kv_head_stride, 1]),
                            ],
                            &[],
                            grid,
                            block_dims(grid),
                            true,
                        )?;
                    }
                    let array = buffer.as_ref().unwrap();
                    array.storage.replace(copy);
                    array.head_stride.set(u64::from(self.kv_cap) * 256);
                }
                self.kv_head_stride = u64::from(self.kv_cap) * 256;
            }
        }
        for (source, target) in [keys, values].into_iter().zip([&self.keys, &self.values]) {
            let span = (self.kv_head_stride + u64::from(rows) * 256) * 2;
            let target = target
                .as_ref()
                .unwrap()
                .buffer()
                .view(u64::from(self.offset) * 256 * 2, span)?;
            qsa_copy(
                Some(source),
                &target,
                [2, rows, 256],
                [strides[0], strides[1], 1],
                [self.kv_head_stride, 256, 1],
                2,
                self.kv_cap as usize * 512,
            )?;
        }
        self.offset = end;
        Ok(())
    }

    pub(super) fn trim(&mut self, n: u32) -> u32 {
        let trimmed = n.min(self.offset);
        self.offset -= trimmed;
        self.pooled_len = self.pooled_len.min(self.offset / 4);
        trimmed
    }

    pub(super) fn nbytes(&self) -> u64 {
        // QSACache reports logical backing capacities, excluding its FP32 mirror.
        u64::from(if self.keys.is_some() { self.kv_cap } else { 0 }) * 2048
            + u64::from(self.raw_cap) * 256
            + u64::from(self.pooled_cap) * 256
    }

    pub(super) fn snapshot(&self) -> QsaSnapshot {
        let nb = self.pooled_len.min(self.offset / 4);
        let kv = |array: &Option<Rc<QsaArray>>| {
            array.as_ref().map(|a| {
                if self.offset == self.kv_cap {
                    a.clone()
                } else {
                    a.slice()
                }
            })
        };
        QsaSnapshot {
            leaves: [
                kv(&self.keys),
                kv(&self.values),
                self.raw.as_ref().map(|a| a.slice()),
                if nb == 0 {
                    None
                } else {
                    self.pooled.as_ref().map(|a| a.slice())
                },
            ],
            offset: self.offset,
            pooled_len: nb,
        }
    }

    pub(super) fn restore(&mut self, state: &QsaSnapshot) -> Result<(), String> {
        let head_stride = state.leaves[0].as_ref().map_or(0, |a| a.head_stride.get());
        let kv_span = head_stride
            .checked_add(u64::from(state.offset) * 256)
            .and_then(|n| n.checked_mul(2));
        if state.offset>i32::MAX as u32/512 || state.pooled_len>state.offset/4
            || state.leaves[0].is_some()!=state.leaves[1].is_some()
            || (state.offset>0 && state.leaves[0].is_none())
            || (state.pooled_len>0 && state.leaves[3].is_none())
            || state.leaves[..2].iter().flatten().any(|b|kv_span.is_none_or(|n|unsafe{ds4_gpu_tensor_bytes(b.buffer().raw())}<n))
            || state.leaves[2].as_ref().is_some_and(|b|unsafe{ds4_gpu_tensor_bytes(b.buffer().raw())}<u64::from(state.offset)*128*2)
            || state.leaves[3].as_ref().is_some_and(|b|unsafe{ds4_gpu_tensor_bytes(b.buffer().raw())}<u64::from(state.pooled_len)*128*2) {
            return Err("invalid QSA snapshot".into());
        }
        [self.keys, self.values, self.raw, self.pooled] = state.leaves.clone();
        self.offset = state.offset;
        self.kv_cap = state.offset;
        self.kv_head_stride = head_stride;
        self.raw_cap = if self.raw.is_some() { state.offset } else { 0 };
        self.pooled_cap = state.pooled_len;
        self.pooled_len = state.pooled_len;
        self.reserved_raw = self.raw_cap;
        self.reserved_pooled = self.pooled_cap;
        self.mirror = None;
        self.mirror_cap = 0;
        self.mirror_strides = [0; 2];
        Ok(())
    }
}

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub(super) enum QsaSelectMode {
    Blocks,
    DenseMask,
    RowTokens,
}

impl QsaSelectMode {
    fn name(self) -> &'static str {
        match self {
            Self::Blocks => "blocks",
            Self::DenseMask => "dense_mask",
            Self::RowTokens => "row_tokens",
        }
    }
}

/// Original selector specialization for the installed H4/D128/ratio4 model.
/// Logical frontiers remain device scalars, not shader specializations.
#[derive(Clone, Copy)]
#[cfg(test)]
pub(super) struct QsaSelectPlan {
    pub(super) rows: u32,
    pub(super) blocks: u32,
    pub(super) topk: u32,
    pub(super) pos: u32,
    pub(super) total: u32,
    pub(super) logical: u32,
    pub(super) dense_width: u32,
    pub(super) mode: QsaSelectMode,
    pub(super) tf32: bool,
    pub(super) input_f32: [bool; 2],
}

#[cfg(test)]
impl QsaSelectPlan {
    fn widths(self) -> Vec<(u32, u32)> {
        match self.mode {
            QsaSelectMode::Blocks => vec![(self.topk, 4), (self.topk, 1), (self.topk, 4)],
            QsaSelectMode::DenseMask => vec![(self.dense_width, 1)],
            QsaSelectMode::RowTokens => vec![(self.topk * 4 + 4, 4), (self.topk * 4 + 4, 1)],
        }
    }

    fn canonical(self) -> qsa_select::Kernel {
        use array::Dtype;
        qsa_select::Kernel {
            heads: 4,
            dim: 128,
            blocks: self.blocks as i32,
            topk: self.topk as i32,
            ratio: 4,
            dense_width: if self.mode == QsaSelectMode::DenseMask {
                self.dense_width as i32
            } else {
                0
            },
            mode: self.mode,
            tf32: self.tf32,
            dtypes: self
                .input_f32
                .map(|f32| if f32 { Dtype::F32 } else { Dtype::BF16 }),
        }
    }

    pub(super) fn source(self) -> Result<String, String> {
        self.canonical().source()
    }
}

/// One original fused selector dispatch, including its otherwise private scratch.
#[cfg(test)]
pub(super) struct QsaSelection {
    pub(super) outputs: Vec<Buffer>,
    pub(super) row_widths: Vec<(u32, u32)>, // elements, storage bytes
    pub(super) row_strides: Vec<u32>,
    _temporaries: Vec<Buffer>,
}

/// One original fused selector dispatch, including its otherwise private scratch.
#[cfg(test)]
pub(super) fn qsa_select(
    [q, pooled]: [&Buffer; 2],
    q_strides: [u64; 3],
    pooled_strides: [u64; 2],
    plan: QsaSelectPlan,
) -> Result<QsaSelection, String> {
    if plan.rows == 0
        || plan.rows > i32::MAX as u32 / 512
        || !(1..=512).contains(&plan.topk)
        || plan.total > i32::MAX as u32
        || plan
            .pos
            .checked_add(plan.rows)
            .is_none_or(|n| n > plan.total)
        || plan.logical > plan.blocks
        || plan.logical != plan.total / 4
        || (plan.mode == QsaSelectMode::DenseMask && plan.dense_width < plan.total)
    {
        return Err("invalid QSA selector geometry or frontier".into());
    }
    let positions =
        [plan.pos, plan.total, plan.logical].map(|value| scalar_buffer(&value.to_le_bytes()));
    let [pos, total, logical] = positions;
    let positions = [pos?, total?, logical?];
    let mut out = qsa_select_device(
        [q, pooled],
        q_strides,
        pooled_strides,
        plan,
        [&positions[0], &positions[1], &positions[2]],
    )?;
    out._temporaries.extend(positions);
    Ok(out)
}

/// Fixed selector geometry with explicit int32[1] graph frontier leaves.
/// plan.pos/total/logical are unused here; only the device leaves are consumed.
#[cfg(test)]
pub(super) fn qsa_select_device(
    [q, pooled]: [&Buffer; 2],
    q_strides: [u64; 3],
    pooled_strides: [u64; 2],
    plan: QsaSelectPlan,
    [pos, total, logical]: [&Buffer; 3],
) -> Result<QsaSelection, String> {
    if plan.rows == 0
        || plan.rows > i32::MAX as u32 / 512
        || plan.blocks > i32::MAX as u32
        || !(1..=512).contains(&plan.topk)
        || [pos, total, logical]
            .into_iter()
            .any(|b| unsafe { ds4_gpu_tensor_bytes(b.raw()) } < 4)
    {
        return Err("invalid QSA device selector geometry or leaves".into());
    }
    for (buffer, shape, strides, f32) in [
        (q, [plan.rows, 4, 128], q_strides, plan.input_f32[0]),
        (
            pooled,
            [1, plan.blocks, 128],
            [0, pooled_strides[0], pooled_strides[1]],
            plan.input_f32[1],
        ),
    ] {
        let span = shape
            .into_iter()
            .zip(strides)
            .try_fold(0_u64, |n, (dim, stride)| {
                u64::from(dim.saturating_sub(1))
                    .checked_mul(stride)
                    .and_then(|v| n.checked_add(v))
            })
            .and_then(|n| n.checked_add(1))
            .and_then(|n| n.checked_mul(if f32 { 4 } else { 2 }));
        if !shape.contains(&0)
            && span.is_none_or(|n| unsafe { ds4_gpu_tensor_bytes(buffer.raw()) } < n)
        {
            return Err("QSA selector input exceeds backing buffer".into());
        }
    }
    use array::{Array, Dtype};
    let kernel = plan.canonical();
    let q = buffer_array(
        q,
        &[1, plan.rows as i32, 4, 128],
        &[0, q_strides[0], q_strides[1], q_strides[2]],
        kernel.dtypes[0],
    )?;
    let pooled = buffer_array(
        pooled,
        &[1, plan.blocks as i32, 128],
        &[0, pooled_strides[0], pooled_strides[1]],
        kernel.dtypes[1],
    )?;
    let inputs = vec![
        q,
        pooled,
        Array::new(&[1], Dtype::I32, pos.view(0, 4)?)?,
        Array::new(&[1], Dtype::I32, total.view(0, 4)?)?,
        Array::new(&[1], Dtype::I32, logical.view(0, 4)?)?,
    ];
    let arrays = kernel
        .layouts(plan.rows as i32)?
        .into_iter()
        .map(|layout| {
            Array::new(
                layout.shape(),
                layout.dtype(),
                Buffer::mtplx_bytes((layout.nbytes() as u64).max(4))?,
            )
        })
        .collect::<Result<Vec<_>, String>>()?;
    qsa_select::encode(&inputs, &arrays, kernel)?;
    let mut buffers = arrays
        .iter()
        .map(|a| a.buffer().view(0, (a.layout().nbytes() as u64).max(4)))
        .collect::<Result<Vec<_>, _>>()?;
    let scratch = buffers.pop().ok_or("QSA scratch sibling missing")?;
    let widths = plan.widths();
    Ok(QsaSelection {
        outputs: buffers,
        row_strides: widths.iter().map(|(w, _)| *w).collect(),
        row_widths: widths,
        _temporaries: vec![scratch],
    })
}

/// Preserve an existing Buffer view when delegating to the canonical backend.
#[cfg(test)]
fn buffer_array(
    buffer: &Buffer,
    shape: &[i32],
    strides: &[u64],
    dtype: array::Dtype,
) -> Result<array::Array, String> {
    use array::{Array, Flags};
    let strides = strides
        .iter()
        .map(|&s| i64::try_from(s).map_err(|_| "QSA stride overflow"))
        .collect::<Result<Vec<_>, _>>()?;
    let size = if shape.contains(&0) {
        0
    } else {
        shape.iter().zip(&strides).try_fold(1_i64, |n, (&d, &s)| {
            n.checked_add(
                i64::from(d.saturating_sub(1))
                    .checked_mul(s)
                    .ok_or("QSA span overflow")?,
            )
            .ok_or("QSA span overflow")
        })? as usize
    };
    let a = Array::unallocated(shape, dtype)?;
    let (_, row, col) = a.layout().contiguity(&strides)?;
    a.set_strided_data(
        buffer.view(0, (size * dtype.itemsize()) as u64)?,
        size,
        &strides,
        Flags {
            contiguous: row || col,
            row_contiguous: row,
            col_contiguous: col,
        },
        0,
    )?;
    Ok(a)
}

/// QSAIndexer._select_fused: bound score scratch per dispatch, concatenate
/// original typed outputs, then return a zero-copy logical dense-mask slice.
#[cfg(test)]
pub(super) fn qsa_select_chunked(
    input: [&Buffer; 2],
    q_strides: [u64; 3],
    pooled_strides: [u64; 2],
    mut plan: QsaSelectPlan,
) -> Result<QsaSelection, String> {
    if plan.rows == 0
        || plan.rows > i32::MAX as u32 / 512
        || !(1..=512).contains(&plan.topk)
        || plan.total > i32::MAX as u32
        || plan
            .pos
            .checked_add(plan.rows)
            .is_none_or(|n| n > plan.total)
        || plan.logical > plan.blocks
        || plan.logical != plan.total / 4
    {
        return Err("invalid chunked QSA geometry".into());
    }
    if plan.mode == QsaSelectMode::DenseMask {
        plan.dense_width = plan
            .blocks
            .checked_add(1)
            .and_then(|n| n.checked_mul(4))
            .ok_or("QSA dense capacity overflow")?;
    }
    let chunk_rows = u64::from(plan.rows)
        .min(((32 * 1024 * 1024) / (u64::from(plan.blocks.max(1)) * 4)).max(1))
        as u32;
    let mut chunks = Vec::new();
    for start in (0..plan.rows).step_by(chunk_rows as usize) {
        let rows = chunk_rows.min(plan.rows - start);
        let at = u64::from(start)
            .checked_mul(q_strides[0])
            .and_then(|n| n.checked_mul(if plan.input_f32[0] { 4 } else { 2 }))
            .ok_or("QSA query slice overflow")?;
        let bytes = unsafe { ds4_gpu_tensor_bytes(input[0].raw()) }
            .checked_sub(at)
            .ok_or("QSA query slice outside backing")?;
        let query = input[0].view(at, bytes)?;
        let pos = plan
            .pos
            .checked_add(start)
            .ok_or("QSA query position overflow")?;
        chunks.push(qsa_select(
            [&query, input[1]],
            q_strides,
            pooled_strides,
            QsaSelectPlan { rows, pos, ..plan },
        )?);
    }
    let mut result = qsa_concat_selection(chunks, plan, chunk_rows)?;
    if plan.mode == QsaSelectMode::DenseMask {
        result.row_widths[0].0 = plan.total;
    }
    Ok(result)
}

#[cfg(test)]
fn qsa_concat_selection(
    mut chunks: Vec<QsaSelection>,
    plan: QsaSelectPlan,
    chunk_rows: u32,
) -> Result<QsaSelection, String> {
    Ok(if chunks.len() == 1 {
        chunks.pop().unwrap()
    } else {
        let widths = plan.widths();
        let outputs = widths
            .iter()
            .map(|(width, bytes)| {
                Buffer::mtplx_bytes(u64::from(plan.rows) * u64::from(*width) * u64::from(*bytes))
            })
            .collect::<Result<Vec<_>, _>>()?;
        for (i, (target, (width, bytes))) in outputs.iter().zip(&widths).enumerate() {
            let dtype = if *bytes == 1 {
                "bool_"
            } else if i == 0 {
                "int32"
            } else {
                "float32"
            };
            let concurrent = encoder::concurrent_scope();
            for (chunk_index, chunk) in chunks.iter().enumerate() {
                let start = chunk_index as u32 * chunk_rows;
                let rows = chunk_rows.min(plan.rows - start);
                let view = target.view(
                    u64::from(start) * u64::from(*width) * u64::from(*bytes),
                    u64::from(rows) * u64::from(*width) * u64::from(*bytes),
                )?;
                qsa_copy_typed(
                    Some(&chunk.outputs[i]),
                    &view,
                    [1, rows, *width],
                    [0, u64::from(*width), 1],
                    [0, u64::from(*width), 1],
                    dtype,
                    (rows) as usize * (*width) as usize,
                )?;
            }
            drop(concurrent);
        }
        let mut temporaries = Vec::new();
        for chunk in chunks {
            temporaries.extend(chunk.outputs);
            temporaries.extend(chunk._temporaries);
        }
        QsaSelection {
            outputs,
            row_strides: widths.iter().map(|(w, _)| *w).collect(),
            row_widths: widths,
            _temporaries: temporaries,
        }
    })
}

/// The producer-aware 128MiB score-plane budget, including the original
/// alignment-down rule. The general FP32 producer needs H4+1 planes.
#[cfg(test)]
pub(super) fn qsa_prefill_chunk_rows(
    rows: u32,
    blocks: u32,
    budget: u64,
    alignment: u32,
    mpp: bool,
) -> Result<u32, String> {
    if rows == 0 || blocks == 0 || budget == 0 || alignment == 0 {
        return Err("invalid QSA prefill budget".into());
    }
    let per_row = u64::from(blocks) * 4 * if mpp { 1 } else { 5 };
    let mut chunk = u64::from(rows).min((budget / per_row).max(1)) as u32;
    if chunk >= alignment && chunk < rows {
        chunk = alignment.max(chunk / alignment * alignment);
    }
    Ok(chunk)
}

/// Original General copy (including cast), after runtime dimension collapse.
#[cfg(test)]
fn qsa_copy_f32(
    input: &Buffer,
    shape: [u32; 3],
    strides: [u64; 3],
    bf16: bool,
) -> Result<Buffer, String> {
    qsa_copy_general(
        input,
        shape,
        strides,
        if bf16 { "bfloat16" } else { "float32" },
        "float32",
        4,
    )
}

#[cfg(test)]
fn qsa_copy_general(
    input: &Buffer,
    shape: [u32; 3],
    strides: [u64; 3],
    source_type: &str,
    output_type: &str,
    output_bytes: u32,
) -> Result<Buffer, String> {
    let count = shape
        .into_iter()
        .try_fold(1_u32, |n, d| n.checked_mul(d))
        .filter(|n| *n > 0 && *n <= i32::MAX as u32)
        .ok_or("QSA General copy exceeds its index range")?;
    let output = Buffer::mtplx_bytes(u64::from(count) * u64::from(output_bytes))?;
    let mut dims: Vec<(u32, u64)> = Vec::new();
    for (d, s) in shape.into_iter().zip(strides).filter(|(d, _)| *d > 1) {
        if let Some(last) = dims.last_mut()
            && last.1 == u64::from(d) * s
        {
            last.0 *= d;
            last.1 = s;
        } else {
            dims.push((d, s));
        }
    }
    let mut grid = [1; 3];
    let mut ss = [0_u64; 3];
    for (i, &(d, s)) in dims.iter().enumerate() {
        grid[dims.len() - 1 - i] = d;
        ss[i] = s;
    }
    dispatch_geometry(
        &format!("g{}_copy{source_type}{output_type}", dims.len()),
        &[
            tensor(0, input).input(qsa_copy_span(shape, strides)?),
            tensor(1, &output).output(count as usize),
            bytes(3, &ss),
        ],
        &[],
        grid,
        block_dims(grid),
        true,
    )?;
    Ok(output)
}

// AsType Vector preserves physically contiguous layouts, including transposes
// and broadcasts; General materializes all logical entries.
#[cfg(test)]
fn qsa_cast_f32(input: &Buffer, shape: [u32; 3], strides: &mut [u64; 3]) -> Result<Buffer, String> {
    let count = shape.into_iter().try_fold(1_u32, |n, d| n.checked_mul(d));
    let span = shape
        .into_iter()
        .zip(*strides)
        .try_fold(1_u64, |n, (d, s)| {
            u64::from(d.saturating_sub(1))
                .checked_mul(s)
                .and_then(|v| n.checked_add(v))
        })
        .and_then(|n| n.checked_mul(2));
    if shape.contains(&0)
        || count.is_none_or(|n| n > i32::MAX as u32)
        || span.is_none_or(|n| unsafe { ds4_gpu_tensor_bytes(input.raw()) } < n)
    {
        return Err("invalid QSA cast input".into());
    }
    let mut dims = shape
        .into_iter()
        .zip(*strides)
        .filter(|(d, s)| *d > 1 && *s != 0)
        .collect::<Vec<_>>();
    dims.sort_unstable_by_key(|(_, s)| *s);
    let mut size = 1_u64;
    let contiguous = dims.into_iter().all(|(d, s)| {
        let fits = s == size;
        size *= u64::from(d);
        fits
    });
    if contiguous {
        let cast = Buffer::mtplx_bytes(size * 4)?;
        cast_bf16_f32(input, &cast, size as u32, true)?;
        Ok(cast)
    } else {
        let cast = qsa_copy_f32(input, shape, *strides, true)?;
        *strides = [
            u64::from(shape[1]) * u64::from(shape[2]),
            u64::from(shape[2]),
            1,
        ];
        Ok(cast)
    }
}

#[repr(C)]
struct QsaGemmParams {
    matrix: [i32; 8],
    batch_strides: [i64; 3],
    swizzle: i32,
    k_iterations: i32,
    batch_ndim: i32,
    padding: i32,
}

/// qsa_indexer_prefill_scores / _select_eager, installed H4/D128 geometry.
/// Preserve the pinned runtime's casts, layout checks, batch collapse and
/// M5 Max GEMV/Steel/NAX/Split-K selection, followed by per-head ReLU and sum.
/// The caller retains the returned temporaries through GPU completion.
#[cfg(test)]
pub(super) fn qsa_prefill_f32_scores(
    [q, pooled_t]: [&Buffer; 2],
    [rows, blocks]: [u32; 2],
    mut q_strides: [u64; 3],
    pooled_strides: [u64; 2],
    q_bf16: bool,
    tf32: bool,
) -> Result<(Buffer, Vec<Buffer>), String> {
    if rows == 0
        || blocks == 0
        || rows.checked_mul(512).is_none_or(|n| n > i32::MAX as u32)
        || rows
            .checked_mul(4)
            .and_then(|n| n.checked_mul(blocks))
            .is_none_or(|n| n > i32::MAX as u32)
        || blocks.checked_mul(128).is_none_or(|n| n > i32::MAX as u32)
    {
        return Err("invalid QSA FP32 score geometry".into());
    }
    for (b, shape, strides, width) in [
        (q, [rows, 4, 128], q_strides, if q_bf16 { 2 } else { 4 }),
        (
            pooled_t,
            [1, 128, blocks],
            [0, pooled_strides[0], pooled_strides[1]],
            4,
        ),
    ] {
        let span = shape
            .into_iter()
            .zip(strides)
            .try_fold(1_u64, |n, (d, s)| {
                u64::from(d - 1)
                    .checked_mul(s)
                    .and_then(|v| n.checked_add(v))
            })
            .and_then(|n| n.checked_mul(width));
        if strides.iter().any(|s| *s > i32::MAX as u64)
            || span.is_none_or(|n| unsafe { ds4_gpu_tensor_bytes(b.raw()) } < n)
        {
            return Err("QSA FP32 input outside backing or index range".into());
        }
    }
    let mut temporary = Vec::new();
    let cast = if q_bf16 {
        Some(qsa_cast_f32(q, [rows, 4, 128], &mut q_strides)?)
    } else {
        None
    };
    let a = cast.as_ref().unwrap_or(q);
    let per_head_size = rows as usize * 4 * blocks as usize;
    // Matmul allocates its output before check_transpose's input copies.
    let per_head = Buffer::mtplx_bytes(per_head_size as u64 * 4)?;
    let copied_a = if q_strides[2] != 1 && q_strides[1] != 1 {
        let copied = qsa_copy_f32(a, [rows, 4, 128], q_strides, false)?;
        q_strides = [512, 128, 1];
        Some(copied)
    } else {
        None
    };
    let a = copied_a.as_ref().unwrap_or(a);
    let ta = q_strides[2] != 1;
    let lda = q_strides[if ta { 2 } else { 1 }] as i32;
    let [bd, bn_stride] = pooled_strides;
    let b_normal = bn_stride == 1 && (blocks != 1 || bd == u64::from(blocks));
    let b_transposed = bd == 1 && (blocks != 1 || bn_stride == 128);
    // check_transpose sees the already broadcast B. A General copy therefore
    // duplicates B across rows and prevents collapsing the rows into M.
    let copied_b = if b_normal || b_transposed {
        None
    } else {
        Some(qsa_copy_f32(
            pooled_t,
            [rows, 128, blocks],
            [0, bd, bn_stride],
            false,
        )?)
    };
    let b = copied_b.as_ref().unwrap_or(pooled_t);
    let tb = copied_b.is_none() && !b_normal;
    let ldb = if copied_b.is_some() {
        blocks
    } else if tb {
        bn_stride as u32
    } else {
        bd as u32
    } as i32;
    let mut batch = rows;
    let mut m = 4_u32;
    let mut a_batch = q_strides[0] as i64;
    let mut b_batch = if copied_b.is_some() {
        i64::from(blocks) * 128
    } else {
        0
    };
    if rows > 1 && !ta && q_strides[1] == 128 && a_batch == 512 && b_batch == 0 {
        m *= rows;
        batch = 1;
        a_batch = 0;
        b_batch = 0;
    }
    let a_size = qsa_copy_span([rows, 4, 128], q_strides)?;
    let b_size = if copied_b.is_some() {
        rows as usize * 128 * blocks as usize
    } else {
        qsa_copy_span([1, 128, blocks], [0, bd, bn_stride])?
    };
    let trans = format!(
        "{}{}",
        if ta { 't' } else { 'n' },
        if tb { 't' } else { 'n' }
    );
    if blocks == 1 {
        let (bm, bn, sm, sn) = if ta {
            (
                1,
                if m >= 2048 {
                    16
                } else if m >= 512 {
                    4
                } else {
                    2
                },
                8,
                4,
            )
        } else if 128 >= 16 * m {
            (1, 8, 1, 32)
        } else {
            (if m >= 4096 { 8 } else { 4 }, 1, 1, 32)
        };
        let out_per_group = if ta { bn * sn * 4 } else { bm * sm * 4 };
        dispatch(
            &format!(
                "gemv{}float32_bm{bm}_bn{bn}_sm{sm}_sn{sn}_tm4_tn4_nc0_axpby0",
                if ta { "_t_" } else { "_" }
            ),
            &[
                tensor(0, a).input(a_size),
                tensor(1, b).input(b_size),
                tensor(3, &per_head).output(per_head_size),
                bytes(4, &128_i32),
                bytes(5, &m),
                bytes(6, &lda),
                bytes(9, &1_i32),
                bytes(10, &batch),
                bytes(11, &b_batch),
                bytes(12, &a_batch),
            ],
            [m.div_ceil(out_per_group), 1, batch],
            [32, bn, bm],
        )?;
    } else if batch == 1 && ((!tf32 && m.max(blocks) <= 128) || (tf32 && m.max(blocks) < 64)) {
        let (bm, bn, bk, partitions, partition_size) = if tf32 {
            (64, 64, 256, 2, 64)
        } else {
            let parts = (8 / (m.div_ceil(32) * blocks.div_ceil(32)))
                .next_power_of_two()
                .clamp(2, 32);
            (
                if m < 40 { 16 } else { 32 },
                if blocks < 40 { 16 } else { 32 },
                16,
                parts,
                128 / parts,
            )
        };
        let scratch_size = partitions as usize * m as usize * blocks as usize;
        let scratch = Buffer::mtplx_bytes(scratch_size as u64 * 4)?;
        let tn = blocks.div_ceil(bn);
        let tm = m.div_ceil(bm);
        let params = [
            m as i32,
            blocks as i32,
            128,
            lda,
            ldb,
            blocks as i32,
            tn as i32,
            tm as i32,
            partitions as i32,
            (m * blocks) as i32,
            partition_size as i32,
            0,
            (partition_size / bk) as i32,
        ];
        let constants = [
            (200, m.is_multiple_of(bm)),
            (201, blocks.is_multiple_of(bn)),
        ]
        .map(|(index, value)| MetalConstant {
            index,
            value: u32::from(value),
            kind: 0,
        });
        let name = if tf32 {
            format!("steel_gemm_splitk_nax_{trans}_float32_float32_bm64_bn64_bk256_wm2_wn2")
        } else {
            format!(
                "steel_gemm_splitk_{trans}_float32_float32_bm{bm}_bn{bn}_bk16_wm2_wn2_MN_{}aligned_K_taligned",
                if m.is_multiple_of(bm) && blocks.is_multiple_of(bn) {
                    't'
                } else {
                    'n'
                }
            )
        };
        dispatch_specialized(
            &name,
            &[
                tensor(0, a).input(a_size),
                tensor(1, b).input(b_size),
                tensor(2, &scratch).output(scratch_size),
                bytes(3, &params),
            ],
            if tf32 { &constants } else { &[] },
            if tf32 {
                [tn * tm * partitions, 1, 1]
            } else {
                [tn, tm, partitions]
            },
            [32, 2, 2],
        )?;
        let grid = [blocks, m, 1];
        dispatch_geometry(
            "steel_gemm_splitk_accum_float32_float32",
            &[
                tensor(0, &scratch).input(scratch_size),
                tensor(1, &per_head).output(per_head_size),
                bytes(2, &partitions),
                bytes(3, &(m * blocks)),
                bytes(4, &blocks),
            ],
            &[],
            grid,
            block_dims(grid),
            true,
        )?;
        encoder::add_temporaries(std::slice::from_ref(&scratch));
        temporary.push(scratch);
    } else {
        let (bn, bk, wn, swizzle) = if tf32 {
            (128, 256, 4, 2)
        } else {
            (64, 16, 2, 0)
        };
        let tn = blocks.div_ceil(bn);
        let tm = m.div_ceil(64);
        let params = QsaGemmParams {
            matrix: [
                m as i32,
                blocks as i32,
                128,
                lda,
                ldb,
                blocks as i32,
                tn as i32,
                tm as i32,
            ],
            batch_strides: [a_batch, b_batch, i64::from(m) * i64::from(blocks)],
            swizzle,
            k_iterations: 128 / bk,
            batch_ndim: 1,
            padding: 0,
        };
        let constants = [
            (10, false),
            (100, false),
            (110, false),
            (200, m.is_multiple_of(64)),
            (201, blocks.is_multiple_of(bn)),
            (202, !tf32),
        ]
        .map(|(index, value)| MetalConstant {
            index,
            value: u32::from(value),
            kind: 0,
        });
        dispatch_specialized(
            &format!(
                "steel_gemm_fused_{}{trans}_float32_float32_bm64_bn{bn}_bk{bk}_wm2_wn{wn}",
                if tf32 { "nax_" } else { "" }
            ),
            &[
                tensor(0, a).input(a_size),
                tensor(1, b).input(b_size),
                tensor(3, &per_head).output(per_head_size),
                bytes(4, &params),
            ],
            &constants,
            [tn * (1 << swizzle), tm.div_ceil(1 << swizzle), batch],
            [32, wn, 2],
        )?;
    }
    encoder::add_temporaries(copied_a.as_slice());
    encoder::add_temporaries(copied_b.as_slice());
    temporary.extend(cast);
    temporary.extend(copied_a);
    temporary.extend(copied_b);
    let zero = scalar_buffer(&0_f32.to_le_bytes())?;
    let count = rows * 4 * blocks;
    let work = if count < 65536 { 1 } else { 2 };
    dispatch_geometry(
        &format!("vs{}_Maximumfloat32", if work == 1 { "" } else { "n" }),
        &[
            tensor(0, &per_head).input(per_head_size),
            tensor(1, &zero).input(1),
            tensor(2, &per_head).output(per_head_size),
            bytes(3, &count),
        ],
        &[],
        [count.div_ceil(work), 1, 1],
        [count.div_ceil(work).min(1024), 1, 1],
        true,
    )?;
    let scores_size = rows as usize * blocks as usize;
    let scores = Buffer::mtplx_bytes(scores_size as u64 * 4)?;
    let ndim = u32::from(rows > 1);
    let outer_shape = if rows > 1 { rows } else { 0 };
    let outer_stride = if rows > 1 { u64::from(blocks) * 4 } else { 0 };
    if blocks == 1 {
        dispatch_geometry(
            "row_reduce_small_1_reduce_sumfloat32",
            &[
                tensor(0, &per_head).input(per_head_size),
                tensor(1, &scores).output(scores_size),
                bytes(2, &4_u64),
                bytes(3, &1_u64),
                bytes(4, &outer_shape),
                bytes(5, &outer_stride),
                bytes(6, &ndim),
                bytes(7, &0_i32),
                bytes(8, &0_u64),
                bytes(9, &0_i32),
            ],
            &[],
            [rows, 1, 1],
            [rows.min(1024), 1, 1],
            true,
        )?;
    } else {
        let gx = blocks.div_ceil(4).min(32);
        dispatch(
            "col_reduce_small_1_reduce_sumfloat32",
            &[
                tensor(0, &per_head).input(per_head_size),
                tensor(1, &scores).output(scores_size),
                bytes(2, &4_u64),
                bytes(3, &u64::from(blocks)),
                bytes(4, &outer_shape),
                bytes(5, &outer_stride),
                bytes(6, &ndim),
                bytes(7, &4_u32),
                bytes(8, &u64::from(blocks)),
                bytes(9, &1_i32),
                bytes(10, &1_u64),
            ],
            [blocks.div_ceil(4).div_ceil(gx), rows, 1],
            [gx, 4, 1],
        )?;
    }
    let divisor = scalar_buffer(&(128_f64.sqrt() as f32).to_le_bytes())?;
    let count = rows * blocks;
    let work = if count < 65536 { 1 } else { 2 };
    dispatch_geometry(
        &format!("vs{}_Dividefloat32", if work == 1 { "" } else { "n" }),
        &[
            tensor(0, &scores).input(scores_size),
            tensor(1, &divisor).input(1),
            tensor(2, &scores).output(scores_size),
            bytes(3, &count),
        ],
        &[],
        [count.div_ceil(work), 1, 1],
        [count.div_ceil(work).min(1024), 1, 1],
        true,
    )?;
    temporary.push(per_head);
    Ok((scores, temporary))
}

#[cfg(test)]
fn qsa_vector(
    prefix: &str,
    op: &str,
    bindings: &[Binding<'_>],
    count: u32,
    width: u32,
) -> Result<(), String> {
    let work = if count < 65536 { 1 } else { 8 / width };
    let suffix = if work == 1 || prefix == "ss" { "" } else { "n" };
    dispatch_geometry(
        &format!("{prefix}{suffix}_{op}"),
        bindings,
        &[],
        [count.div_ceil(work), 1, 1],
        [count.div_ceil(work).min(1024), 1, 1],
        true,
    )
}

#[cfg(test)]
pub(super) struct QsaEagerTopk {
    rows: u32,
    pos_start: u32,
    tiled: bool,
    /// Untiled: last K columns of ArgPartition, stride N. Tiled: concat, stride K.
    pub(super) indices: Buffer,
    pub(super) valid: Buffer,
    pub(super) masked_scores: Option<Buffer>,
    pub(super) positions: Buffer,
    pub(super) complete_blocks: Buffer,
    pub(super) width: u32,
    pub(super) row_stride: u32,
    pub(super) index_row_stride: u32,
    _temporaries: Vec<Buffer>,
}

/// _select_eager, including _tiled_topk's synchronous per-tile eval and concat.
/// ArgPartition in the pinned runtime is a complete argsort, not a replacement
/// custom selector. The pooled FP32 mirror and tie vector are shared by tiles.
#[cfg(test)]
pub(super) fn qsa_eager_topk(
    input: [&Buffer; 2],
    q_strides: [u64; 3],
    pooled_t_strides: [u64; 2],
    plan: QsaSelectPlan,
    tile_rows: u32,
    mut commands: Option<&mut Commands>,
) -> Result<QsaEagerTopk, String> {
    if plan
        .pos
        .checked_add(plan.rows)
        .is_none_or(|n| n > plan.total)
        || plan.total > i32::MAX as u32
        || plan.blocks != plan.total / 4
        || !plan.input_f32[1]
        || !(1..=512).contains(&plan.topk)
        || plan.rows == 0
        || plan.blocks == 0
        || plan
            .rows
            .checked_mul(plan.blocks)
            .is_none_or(|n| n > i32::MAX as u32)
    {
        return Err("invalid QSA eager frontier".into());
    }
    let rows = plan.rows;
    let blocks = plan.blocks;
    let positions = Buffer::mtplx_bytes(u64::from(rows) * 4)?;
    let complete = Buffer::mtplx_bytes(u64::from(rows) * 4)?;
    let blk = Buffer::mtplx_bytes(u64::from(blocks) * 4)?;
    let tie = Buffer::mtplx_bytes(u64::from(blocks) * 4)?;
    let one = scalar_buffer(&1_i32.to_le_bytes())?;
    let ratio = scalar_buffer(&4_i32.to_le_bytes())?;
    let tie_scale = scalar_buffer(&1e-12_f32.to_le_bytes())?;
    let neg = scalar_buffer(&f32::NEG_INFINITY.to_le_bytes())?;
    // Original vector dispatch threshold for int32/float32; SS entries have
    // no size argument, which the shared Metal binder may harmlessly leave set.
    let vector = |prefix: &str, op: &str, bindings: &[Binding<'_>], count: u32| {
        qsa_vector(prefix, op, bindings, count, 4)
    };
    for (out, count, start) in [(&positions, rows, plan.pos), (&blk, blocks, 0)] {
        dispatch_geometry(
            "arangeint32",
            &[
                bytes(0, &start),
                bytes(1, &1_i32),
                tensor(2, out).output(count as usize),
            ],
            &[],
            [count, 1, 1],
            [count.min(1024), 1, 1],
            true,
        )?;
    }
    let prefix = if rows == 1 { "ss" } else { "vs" };
    vector(
        prefix,
        "Addint32",
        &[
            tensor(0, &positions).input(rows as usize),
            tensor(1, &one).input(1),
            tensor(2, &complete).output(rows as usize),
            bytes(3, &rows),
        ],
        rows,
    )?;
    vector(
        prefix,
        "Divideint32",
        &[
            tensor(0, &complete).input(rows as usize),
            tensor(1, &ratio).input(1),
            tensor(2, &complete).output(rows as usize),
            bytes(3, &rows),
        ],
        rows,
    )?;
    let work = if blocks < 65536 { 1 } else { 2 };
    dispatch_geometry(
        &format!("v{}_copyint32float32", if work == 1 { "" } else { "n" }),
        &[
            tensor(0, &blk).input(blocks as usize),
            tensor(1, &tie).output(blocks as usize),
            bytes(2, &blocks),
        ],
        &[],
        [blocks.div_ceil(work), 1, 1],
        [blocks.div_ceil(work).min(1024), 1, 1],
        true,
    )?;
    vector(
        if blocks == 1 { "ss" } else { "vs" },
        "Multiplyfloat32",
        &[
            tensor(0, &tie).input(blocks as usize),
            tensor(1, &tie_scale).input(1),
            tensor(2, &tie).output(blocks as usize),
            bytes(3, &blocks),
        ],
        blocks,
    )?;
    let width = plan.topk.min(blocks);
    let tiled = rows > 1 && tile_rows > 0 && tile_rows < rows;
    let (indices, valid, masked_scores, mut temporary) = if tiled {
        let mut parts = Vec::new();
        for start in (0..rows).step_by(tile_rows as usize) {
            let len = tile_rows.min(rows - start);
            let complete_view = complete.view(u64::from(start) * 4, u64::from(len) * 4)?;
            let valid = qsa_eager_valid(&blk, &complete_view, len, blocks)?;
            let item_bytes = if plan.input_f32[0] { 4 } else { 2 };
            let q_len = u64::from(len - 1)
                .checked_mul(q_strides[0])
                .and_then(|n| n.checked_add(3_u64.checked_mul(q_strides[1])?))
                .and_then(|n| n.checked_add(127_u64.checked_mul(q_strides[2])?))
                .and_then(|n| n.checked_add(1))
                .and_then(|n| n.checked_mul(item_bytes))
                .ok_or("QSA eager tile query span overflow")?;
            let offset = u64::from(start)
                .checked_mul(q_strides[0])
                .and_then(|n| n.checked_mul(item_bytes))
                .ok_or("QSA eager tile query offset overflow")?;
            let q = input[0].view(offset, q_len)?;
            let (part, _scores, _scratch) = qsa_eager_rank(
                [&q, input[1]],
                q_strides,
                pooled_t_strides,
                QsaSelectPlan {
                    rows: len,
                    pos: plan.pos + start,
                    ..plan
                },
                &valid,
                &tie,
                &neg,
            )?;
            // mx.eval(top_t) completes before constructing the next tile.
            // Keep only the evaluated ArgPartition view/backing in parts;
            // score intermediates are not retained for the entire query chunk.
            if let Some(result) = encoder::with_active(|encoder| encoder.synchronize()) {
                result?;
            } else {
                commands
                    .as_deref_mut()
                    .ok_or("QSA eager eval requires an active command owner")?
                    .eval_boundary()?;
            }
            parts.push((part, len));
        }
        let indices = Buffer::mtplx_bytes(u64::from(rows) * u64::from(width) * 4)?;
        let mut start = 0;
        let mut temporary = Vec::new();
        let concurrent = encoder::concurrent_scope();
        for (part, len) in parts {
            let dst = indices.view(
                u64::from(start) * u64::from(width) * 4,
                u64::from(len) * u64::from(width) * 4,
            )?;
            qsa_copy_typed(
                Some(&part),
                &dst,
                [1, len, width],
                [0, u64::from(blocks), 1],
                [0, u64::from(width), 1],
                "uint32",
                (len) as usize * (width) as usize,
            )?;
            temporary.push(part);
            start += len;
        }
        drop(concurrent);
        // Global valid is a different lazy expression from valid_t; evaluate
        // it for the epilogue after the per-tile evals, as in the reference.
        let valid = qsa_eager_valid(&blk, &complete, rows, blocks)?;
        (indices, valid, None, temporary)
    } else {
        let valid = qsa_eager_valid(&blk, &complete, rows, blocks)?;
        let (indices, scores, temporary) =
            qsa_eager_rank(input, q_strides, pooled_t_strides, plan, &valid, &tie, &neg)?;
        (indices, valid, Some(scores), temporary)
    };
    temporary.extend([blk, tie, one, ratio, tie_scale, neg]);
    Ok(QsaEagerTopk {
        rows,
        pos_start: plan.pos,
        tiled,
        indices,
        valid,
        masked_scores,
        positions,
        complete_blocks: complete,
        width,
        row_stride: blocks,
        index_row_stride: if tiled { width } else { blocks },
        _temporaries: temporary,
    })
}

#[cfg(test)]
fn qsa_eager_valid(
    blk: &Buffer,
    complete: &Buffer,
    rows: u32,
    blocks: u32,
) -> Result<Buffer, String> {
    let count = rows * blocks;
    let valid = Buffer::mtplx_bytes(u64::from(count))?;
    if rows == 1 || blocks == 1 {
        qsa_vector(
            if count == 1 {
                "ss"
            } else if rows == 1 {
                "vs"
            } else {
                "sv"
            },
            "Lessint32",
            &[
                tensor(0, blk).input(blocks as usize),
                tensor(1, complete).input(rows as usize),
                tensor(2, &valid).output(count as usize),
                bytes(3, &count),
            ],
            count,
            4,
        )?;
    } else {
        let grid = [blocks, rows, 1];
        dispatch_geometry(
            "g2_Lessint32",
            &[
                tensor(0, blk).input(blocks as usize),
                tensor(1, complete).input(rows as usize),
                tensor(2, &valid).output(count as usize),
                bytes(3, &[0_u64, 1]),
                bytes(4, &[1_u64, 0]),
            ],
            &[],
            grid,
            block_dims(grid),
            true,
        )?;
    }
    Ok(valid)
}

#[cfg(test)]
fn qsa_eager_rank(
    input: [&Buffer; 2],
    q_strides: [u64; 3],
    pooled_t_strides: [u64; 2],
    plan: QsaSelectPlan,
    valid: &Buffer,
    tie: &Buffer,
    neg: &Buffer,
) -> Result<(Buffer, Buffer, Vec<Buffer>), String> {
    let rows = plan.rows;
    let blocks = plan.blocks;
    let count = rows * blocks;
    let (scores, mut temporary) = qsa_prefill_f32_scores(
        input,
        [rows, blocks],
        q_strides,
        pooled_t_strides,
        !plan.input_f32[0],
        plan.tf32,
    )?;
    qsa_vector(
        if count == 1 { "v" } else { "vs" },
        "Selectfloat32",
        &[
            tensor(0, valid).input(count as usize),
            tensor(1, &scores).input(count as usize),
            tensor(2, neg).input(1),
            tensor(3, &scores).output(count as usize),
            bytes(4, &count),
        ],
        count,
        4,
    )?;
    if rows == 1 || blocks == 1 {
        qsa_vector(
            if count == 1 {
                "ss"
            } else if blocks == 1 {
                "vs"
            } else {
                "vv"
            },
            "Subtractfloat32",
            &[
                tensor(0, &scores).input(count as usize),
                tensor(1, tie).input(blocks as usize),
                tensor(2, &scores).output(count as usize),
                bytes(3, &count),
            ],
            count,
            4,
        )?;
    } else {
        let grid = [blocks, rows, 1];
        dispatch_geometry(
            "g2_Subtractfloat32",
            &[
                tensor(0, &scores).input(count as usize),
                tensor(1, tie).input(blocks as usize),
                tensor(2, &scores).output(count as usize),
                bytes(3, &[u64::from(blocks), 1]),
                bytes(4, &[0_u64, 1]),
            ],
            &[],
            grid,
            block_dims(grid),
            true,
        )?;
    }
    let sorted = Buffer::mtplx_bytes(u64::from(count) * 4)?;
    temporary.extend(argsort(&scores, &sorted, [rows, blocks], SortDtype::F32)?);
    let width = plan.topk.min(blocks);
    let offset = u64::from(blocks - width) * 4;
    let indices = sorted.view(offset, u64::from(count) * 4 - offset)?;
    temporary.push(sorted);
    Ok((indices, scores, temporary))
}

/// The eager selector's output switches. `large_prefill` is the common
/// phase/rows/earliest-position eligibility decision, also used by the fused
/// and compiled selectors; it must not be inferred from rows > 1 alone.
pub(super) struct QsaEagerOptions {
    pub(super) large_prefill: bool,
    pub(super) flash: bool,
    pub(super) gather_decode: bool,
    pub(super) gather: bool,
    pub(super) gather_max_rows: u32,
    pub(super) gather_min_context: u32,
}

#[cfg(test)]
pub(super) enum QsaIndexerSelection {
    Dense(QsaSelection),
    PrefillBlocks(QsaSelection),
    FlashBlocks(QsaSelection, u32),
    DecodeTokens(QsaSelection),
    RowTokens(QsaSelection),
}

/// Complete _select_eager return-branch order. Inactive lazy expressions
/// (especially the dense scatter mask in flash/gather modes) are not encoded.
#[cfg(test)]
pub(super) fn qsa_eager_output(
    top: QsaEagerTopk,
    total: u32,
    options: QsaEagerOptions,
) -> Result<QsaIndexerSelection, String> {
    if top.pos_start.checked_add(top.rows) != Some(total)
        || top.row_stride != total / 4
        || top.row_stride <= 512
    {
        return Err("invalid active eager selector output frontier".into());
    }
    if top.rows > 1 && options.large_prefill {
        return qsa_eager_blocks(top).map(QsaIndexerSelection::PrefillBlocks);
    }
    if top.rows == 1 && options.flash {
        let (blocks, tail) = qsa_eager_decode(top, total, false)?;
        return Ok(QsaIndexerSelection::FlashBlocks(blocks, tail));
    }
    if top.rows == 1 && options.gather_decode {
        return qsa_eager_decode(top, total, true)
            .map(|(tokens, _)| QsaIndexerSelection::DecodeTokens(tokens));
    }
    if top.rows > 1
        && !top.tiled
        && options.gather
        && top.rows <= options.gather_max_rows
        && total >= options.gather_min_context
    {
        return qsa_eager_gather_rows(top).map(QsaIndexerSelection::RowTokens);
    }
    qsa_eager_dense(top, total).map(QsaIndexerSelection::Dense)
}

/// Eager flash_prefill epilogue: chronological int32 blocks, gathered validity,
/// then invalid slots replaced by zero. No custom selector substitution.
#[cfg(test)]
pub(super) fn qsa_eager_blocks(top: QsaEagerTopk) -> Result<QsaSelection, String> {
    let rows = top.rows;
    let width = top.width;
    let count = rows
        .checked_mul(width)
        .filter(|n| *n > 0 && *n <= i32::MAX as u32)
        .ok_or("invalid eager QSA block output shape")?;
    let (sorted, ids) = qsa_eager_sorted_ids(&top)?;
    let indices64 = Buffer::mtplx_bytes(u64::from(count) * 8)?;
    let valid = Buffer::mtplx_bytes(u64::from(count))?;
    let work = if count < 65536 { 1 } else { 2 };
    // int64 output has one element per thread even above the vector threshold.
    dispatch_geometry(
        "v_copyint32int64",
        &[
            tensor(0, &sorted).input(count as usize),
            tensor(1, &indices64).output(count as usize),
            bytes(2, &count),
        ],
        &[],
        [count, 1, 1],
        [count.min(1024), 1, 1],
        true,
    )?;
    let grid = [1, width, rows];
    dispatch_geometry(
        "kernel_qwen_mtplx_gather_axis_bool_idxi64_int_11",
        &[
            tensor(0, &top.valid).input(rows as usize * top.row_stride as usize),
            tensor(1, &indices64).input(count as usize),
            tensor(2, &valid).output(count as usize),
            bytes(3, &rows),
            bytes(4, &u64::from(top.row_stride)),
            bytes(5, &u64::from(width)),
            bytes(6, &1_u64),
            bytes(7, &1_u32),
            bytes(8, &top.row_stride),
            bytes(9, &1_u64),
            bytes(10, &1_u64),
        ],
        &[],
        grid,
        block_dims(grid),
        true,
    )?;
    let zero = scalar_buffer(&0_i32.to_le_bytes())?;
    dispatch_geometry(
        &format!(
            "{}{}_Selectint32",
            if count == 1 { "v" } else { "vs" },
            if work == 1 { "" } else { "n" }
        ),
        &[
            tensor(0, &valid).input(count as usize),
            tensor(1, &sorted).input(count as usize),
            tensor(2, &zero).input(1),
            tensor(3, &sorted).output(count as usize),
            bytes(4, &count),
        ],
        &[],
        [count.div_ceil(work), 1, 1],
        [count.div_ceil(work).min(1024), 1, 1],
        true,
    )?;
    let mut temporary = top._temporaries;
    temporary.extend(top.masked_scores);
    temporary.extend([
        ids,
        indices64,
        top.indices,
        top.valid,
        top.positions,
        top.complete_blocks,
    ]);
    Ok(QsaSelection {
        outputs: vec![sorted, valid],
        row_widths: vec![(width, 4), (width, 1)],
        row_strides: vec![width; 2],
        _temporaries: temporary,
    })
}

#[cfg(test)]
fn qsa_eager_sorted_ids(top: &QsaEagerTopk) -> Result<(Buffer, Buffer), String> {
    let width = top.width;
    let rows = top.rows;
    let ids = qsa_eager_cast_indices(top, "int32", 4)?;
    let sorted = Buffer::mtplx_bytes(u64::from(rows) * u64::from(width) * 4)?;
    let bn = width.div_ceil(4).next_power_of_two().clamp(32, 512);
    let segment = if rows == 1 { i32::MAX as u32 } else { width };
    dispatch(
        &format!("c_block_sort_int32_int32_bn{bn}_tn4"),
        &[
            tensor(0, &ids).input(rows as usize * width as usize),
            tensor(1, &sorted).output(rows as usize * width as usize),
            bytes(2, &width),
            bytes(3, &1_u32),
            bytes(4, &1_u32),
            bytes(5, &segment),
            bytes(6, &segment),
        ],
        [1, rows, 1],
        [bn, 1, 1],
    )?;
    Ok((sorted, ids))
}

/// Untiled eager decode epilogues. The dead selected-mask expression in the
/// Python source is not evaluated by these lazy output branches.
#[cfg(test)]
pub(super) fn qsa_eager_decode(
    top: QsaEagerTopk,
    total: u32,
    gather: bool,
) -> Result<(QsaSelection, u32), String> {
    if top.rows != 1
        || top.pos_start.checked_add(1) != Some(total)
        || total / 4 != top.row_stride
        || top.row_stride <= 512
    {
        return Err("invalid active QSA eager decode frontier".into());
    }
    let tail_start = total / 4 * 4;
    let k = top.width;
    let (sorted, cast) = qsa_eager_sorted_ids(&top)?;
    let mut temporary = top._temporaries;
    temporary.extend(top.masked_scores);
    temporary.extend([
        cast,
        top.indices,
        top.valid,
        top.positions,
        top.complete_blocks,
    ]);
    let (out, width) = if !gather {
        (sorted, k)
    } else {
        let ratio = scalar_buffer(&4_i32.to_le_bytes())?;
        qsa_vector(
            "vs",
            "Multiplyint32",
            &[
                tensor(0, &sorted).input(k as usize),
                tensor(1, &ratio).input(1),
                tensor(2, &sorted).output(k as usize),
                bytes(3, &k),
            ],
            k,
            4,
        )?;
        let ratio_range = Buffer::mtplx_bytes(16)?;
        dispatch_geometry(
            "arangeint32",
            &[
                bytes(0, &0_u32),
                bytes(1, &1_u32),
                tensor(2, &ratio_range).output(4),
            ],
            &[],
            [4, 1, 1],
            [4, 1, 1],
            true,
        )?;
        let block_tokens = Buffer::mtplx_bytes(u64::from(k) * 16)?;
        let grid = [4, k, 1];
        dispatch_geometry(
            "g2_Addint32",
            &[
                tensor(0, &sorted).input(k as usize),
                tensor(1, &ratio_range).input(4),
                tensor(2, &block_tokens).output(k as usize * 4),
                bytes(3, &[1_u64, 0]),
                bytes(4, &[0_u64, 1]),
            ],
            &[],
            grid,
            block_dims(grid),
            true,
        )?;
        let tail_len = total - tail_start;
        let width = k * 4 + tail_len;
        let tail = if tail_len > 0 {
            let tail = Buffer::mtplx_bytes(u64::from(tail_len) * 4)?;
            dispatch_geometry(
                "arangeint32",
                &[
                    bytes(0, &tail_start),
                    bytes(1, &1_u32),
                    tensor(2, &tail).output(tail_len as usize),
                ],
                &[],
                [tail_len, 1, 1],
                [tail_len, 1, 1],
                true,
            )?;
            Some(tail)
        } else {
            None
        };
        let out = Buffer::mtplx_bytes(u64::from(width) * 4)?;
        let concurrent = encoder::concurrent_scope();
        qsa_copy_typed(
            Some(&block_tokens),
            &out,
            [1, 1, k * 4],
            [0, 0, 1],
            [0, 0, 1],
            "int32",
            (k * 4) as usize,
        )?;
        if let Some(tail) = tail {
            let dst = out.view(u64::from(k) * 16, u64::from(tail_len) * 4)?;
            qsa_copy_typed(
                Some(&tail),
                &dst,
                [1, 1, tail_len],
                [0, 0, 1],
                [0, 0, 1],
                "int32",
                (tail_len) as usize,
            )?;
            temporary.push(tail);
        }
        drop(concurrent);
        temporary.extend([sorted, ratio, ratio_range, block_tokens]);
        (out, width)
    };
    Ok((
        QsaSelection {
            outputs: vec![out],
            row_widths: vec![(width, 4)],
            row_strides: vec![width],
            _temporaries: temporary,
        },
        tail_start,
    ))
}

#[cfg(test)]
fn qsa_eager_cast_indices(
    top: &QsaEagerTopk,
    dtype: &str,
    bytes_per_item: u32,
) -> Result<Buffer, String> {
    let count = top.rows * top.width;
    let input_size = (top.rows as usize - 1) * top.index_row_stride as usize + top.width as usize;
    let output = Buffer::mtplx_bytes(u64::from(count) * u64::from(bytes_per_item))?;
    if top.rows == 1 || top.width == top.index_row_stride {
        qsa_vector(
            "v",
            &format!("copyuint32{dtype}"),
            &[
                tensor(0, &top.indices).input(input_size),
                tensor(1, &output).output(count as usize),
                bytes(2, &count),
            ],
            count,
            bytes_per_item,
        )?;
    } else {
        let grid = [top.width, top.rows, 1];
        dispatch_geometry(
            &format!("g2_copyuint32{dtype}"),
            &[
                tensor(0, &top.indices).input(input_size),
                tensor(1, &output).output(count as usize),
                bytes(3, &[u64::from(top.index_row_stride), 1]),
            ],
            &[],
            grid,
            block_dims(grid),
            true,
        )?;
    }
    Ok(output)
}

/// _select_eager's rows-gather epilogue. Keeps argsort order (unlike the
/// chronological flash block contract), then appends each row's visible tail.
#[cfg(test)]
pub(super) fn qsa_eager_gather_rows(top: QsaEagerTopk) -> Result<QsaSelection, String> {
    let rows = top.rows;
    let k = top.width;
    let count = rows * k;
    let width = k * 4 + 4;
    if rows <= 1
        || top.row_stride <= 512
        || rows.checked_mul(width).is_none_or(|n| n > i32::MAX as u32)
    {
        return Err("invalid active QSA gather-rows geometry".into());
    }
    let idx64 = qsa_eager_cast_indices(&top, "int64", 8)?;
    let blk_ok = Buffer::mtplx_bytes(u64::from(count))?;
    let grid = [1, k, rows];
    dispatch_geometry(
        "kernel_qwen_mtplx_gather_axis_bool_idxi64_int_11",
        &[
            tensor(0, &top.valid).input(rows as usize * top.row_stride as usize),
            tensor(1, &idx64).input(count as usize),
            tensor(2, &blk_ok).output(count as usize),
            bytes(3, &rows),
            bytes(4, &u64::from(top.row_stride)),
            bytes(5, &u64::from(k)),
            bytes(6, &1_u64),
            bytes(7, &1_u32),
            bytes(8, &top.row_stride),
            bytes(9, &1_u64),
            bytes(10, &1_u64),
        ],
        &[],
        grid,
        block_dims(grid),
        true,
    )?;
    let idx32 = qsa_eager_cast_indices(&top, "int32", 4)?;
    let ratio = scalar_buffer(&4_i32.to_le_bytes())?;
    qsa_vector(
        "vs",
        "Multiplyint32",
        &[
            tensor(0, &idx32).input(count as usize),
            tensor(1, &ratio).input(1),
            tensor(2, &idx32).output(count as usize),
            bytes(3, &count),
        ],
        count,
        4,
    )?;
    let ratio_range = Buffer::mtplx_bytes(16)?;
    dispatch_geometry(
        "arangeint32",
        &[
            bytes(0, &0_u32),
            bytes(1, &1_u32),
            tensor(2, &ratio_range).output(4),
        ],
        &[],
        [4, 1, 1],
        [4, 1, 1],
        true,
    )?;
    let tokens = Buffer::mtplx_bytes(u64::from(count) * 16)?;
    let blocks_ok = Buffer::mtplx_bytes(u64::from(count) * 4)?;
    let grid = [4, count, 1];
    dispatch_geometry(
        "g2_Addint32",
        &[
            tensor(0, &idx32).input(count as usize),
            tensor(1, &ratio_range).input(4),
            tensor(2, &tokens).output(count as usize * 4),
            bytes(3, &[1_u64, 0]),
            bytes(4, &[0_u64, 1]),
        ],
        &[],
        grid,
        block_dims(grid),
        true,
    )?;
    dispatch_geometry(
        "g2_copybool_bool_",
        &[
            tensor(0, &blk_ok).input(count as usize),
            tensor(1, &blocks_ok).output(count as usize * 4),
            bytes(3, &[1_u64, 0]),
        ],
        &[],
        grid,
        block_dims(grid),
        true,
    )?;
    let tail_start = top.complete_blocks;
    let tail_ratio = scalar_buffer(&4_i32.to_le_bytes())?;
    qsa_vector(
        "vs",
        "Multiplyint32",
        &[
            tensor(0, &tail_start).input(rows as usize),
            tensor(1, &tail_ratio).input(1),
            tensor(2, &tail_start).output(rows as usize),
            bytes(3, &rows),
        ],
        rows,
        4,
    )?;
    let tail_range = Buffer::mtplx_bytes(16)?;
    dispatch_geometry(
        "arangeint32",
        &[
            bytes(0, &0_u32),
            bytes(1, &1_u32),
            tensor(2, &tail_range).output(4),
        ],
        &[],
        [4, 1, 1],
        [4, 1, 1],
        true,
    )?;
    let tail_tokens = Buffer::mtplx_bytes(u64::from(rows) * 16)?;
    let tail_ok = Buffer::mtplx_bytes(u64::from(rows) * 4)?;
    let grid = [4, rows, 1];
    dispatch_geometry(
        "g2_Addint32",
        &[
            tensor(0, &tail_start).input(rows as usize),
            tensor(1, &tail_range).input(4),
            tensor(2, &tail_tokens).output(rows as usize * 4),
            bytes(3, &[1_u64, 0]),
            bytes(4, &[0_u64, 1]),
        ],
        &[],
        grid,
        block_dims(grid),
        true,
    )?;
    dispatch_geometry(
        "g2_LessEqualint32",
        &[
            tensor(0, &tail_tokens).input(rows as usize * 4),
            tensor(1, &top.positions).input(rows as usize),
            tensor(2, &tail_ok).output(rows as usize * 4),
            bytes(3, &[4_u64, 1]),
            bytes(4, &[1_u64, 0]),
        ],
        &[],
        grid,
        block_dims(grid),
        true,
    )?;
    let ids = Buffer::mtplx_bytes(u64::from(rows) * u64::from(width) * 4)?;
    let valid = Buffer::mtplx_bytes(u64::from(rows) * u64::from(width))?;
    for (block, tail, out, dtype, b) in [
        (&tokens, &tail_tokens, &ids, "int32", 4_u64),
        (&blocks_ok, &tail_ok, &valid, "bool_", 1),
    ] {
        let _concurrent = encoder::concurrent_scope();
        qsa_copy_typed(
            Some(block),
            out,
            [1, rows, k * 4],
            [0, u64::from(k) * 4, 1],
            [0, u64::from(width), 1],
            dtype,
            (rows) as usize * (k * 4) as usize,
        )?;
        let offset = u64::from(k) * 4 * b;
        let dst = out.view(offset, u64::from(rows) * u64::from(width) * b - offset)?;
        qsa_copy_typed(
            Some(tail),
            &dst,
            [1, rows, 4],
            [0, 4, 1],
            [0, u64::from(width), 1],
            dtype,
            (rows) as usize * 4,
        )?;
    }
    let count = rows * width;
    let zero = scalar_buffer(&0_i32.to_le_bytes())?;
    qsa_vector(
        "vs",
        "Selectint32",
        &[
            tensor(0, &valid).input(count as usize),
            tensor(1, &ids).input(count as usize),
            tensor(2, &zero).input(1),
            tensor(3, &ids).output(count as usize),
            bytes(4, &count),
        ],
        count,
        4,
    )?;
    let mut temporary = top._temporaries;
    temporary.extend(top.masked_scores);
    temporary.extend([
        idx64,
        idx32,
        blk_ok,
        blocks_ok,
        ratio,
        ratio_range,
        tail_ratio,
        tail_range,
        tokens,
        tail_start,
        tail_tokens,
        tail_ok,
        top.indices,
        top.valid,
        top.positions,
    ]);
    Ok(QsaSelection {
        outputs: vec![ids, valid],
        row_widths: vec![(width, 4), (width, 1)],
        row_strides: vec![width; 2],
        _temporaries: temporary,
    })
}

/// Remaining dense-mask epilogue of _select_eager, after the active indexer
/// threshold (N>512). Uses original ScatterAxis/Repeat/Concatenate operations.
#[cfg(test)]
pub(super) fn qsa_eager_dense(top: QsaEagerTopk, total: u32) -> Result<QsaSelection, String> {
    let rows = top.rows;
    let blocks = top.row_stride;
    if blocks <= 512
        || total / 4 != blocks
        || top.pos_start.checked_add(rows) != Some(total)
        || total > i32::MAX as u32
        || rows.checked_mul(total).is_none_or(|n| n > i32::MAX as u32)
    {
        return Err("invalid active QSA dense-mask geometry".into());
    }
    let k = top.width;
    let index_count = rows * k;
    let count = rows * blocks;
    let indices64 = qsa_eager_cast_indices(&top, "int64", 8)?;
    let selected = Buffer::mtplx_bytes(u64::from(count))?;
    let zero = scalar_buffer(&[0_u8])?;
    let one = scalar_buffer(&[1_u8])?;
    qsa_vector(
        "s",
        "copybool_bool_",
        &[
            tensor(0, &zero).input(1),
            tensor(1, &selected).output(count as usize),
            bytes(2, &count),
        ],
        count,
        1,
    )?;
    let grid = [1, k, rows];
    dispatch_geometry(
        if index_count == 1 {
            "kernel_qwen_mtplx_scatter_axis_bool_idxi64_int_11"
        } else {
            "kernel_qwen_mtplx_scatter_axis_bool_idxi64_int_01"
        },
        &[
            tensor(0, &one).input(1),
            tensor(1, &indices64).input(index_count as usize),
            tensor(2, &selected).output(count as usize),
            bytes(3, &rows),
            bytes(4, &0_u64),
            bytes(5, &u64::from(k)),
            bytes(6, &1_u64),
            bytes(7, &1_u32),
            bytes(8, &blocks),
            bytes(9, &0_u64),
            bytes(10, &1_u64),
        ],
        &[],
        grid,
        block_dims(grid),
        true,
    )?;
    qsa_vector(
        "vv",
        "BitwiseAndbool_",
        &[
            tensor(0, &selected).input(count as usize),
            tensor(1, &top.valid).input(count as usize),
            tensor(2, &selected).output(count as usize),
            bytes(3, &count),
        ],
        count,
        1,
    )?;
    // repeat inserts a broadcast axis and its final reshape is a General copy.
    let repeated = Buffer::mtplx_bytes(u64::from(count) * 4)?;
    let grid = [4, count, 1];
    dispatch_geometry(
        "g2_copybool_bool_",
        &[
            tensor(0, &selected).input(count as usize),
            tensor(1, &repeated).output(count as usize * 4),
            bytes(3, &[1_u64, 0]),
        ],
        &[],
        grid,
        block_dims(grid),
        true,
    )?;
    let mut temporary = top._temporaries;
    temporary.extend(top.masked_scores);
    let mask = if total == blocks * 4 {
        repeated
    } else {
        let pad_count = rows * (total - blocks * 4);
        let pad = Buffer::mtplx_bytes(u64::from(pad_count))?;
        let pad_zero = scalar_buffer(&[0_u8])?;
        qsa_vector(
            "s",
            "copybool_bool_",
            &[
                tensor(0, &pad_zero).input(1),
                tensor(1, &pad).output(pad_count as usize),
                bytes(2, &pad_count),
            ],
            pad_count,
            1,
        )?;
        let padded = Buffer::mtplx_bytes(u64::from(rows) * u64::from(total))?;
        let concurrent = encoder::concurrent_scope();
        qsa_copy_typed(
            Some(&repeated),
            &padded,
            [1, rows, blocks * 4],
            [0, u64::from(blocks) * 4, 1],
            [0, u64::from(total), 1],
            "bool_",
            (rows) as usize * (blocks * 4) as usize,
        )?;
        let tail = padded.view(
            u64::from(blocks) * 4,
            u64::from(rows) * u64::from(total) - u64::from(blocks) * 4,
        )?;
        qsa_copy_typed(
            Some(&pad),
            &tail,
            [1, rows, total - blocks * 4],
            [0, u64::from(total - blocks * 4), 1],
            [0, u64::from(total), 1],
            "bool_",
            (rows) as usize * (total - blocks * 4) as usize,
        )?;
        drop(concurrent);
        temporary.extend([repeated, pad, pad_zero]);
        padded
    };
    let tpos = Buffer::mtplx_bytes(u64::from(total) * 4)?;
    dispatch_geometry(
        "arangeint32",
        &[
            bytes(0, &0_u32),
            bytes(1, &1_i32),
            tensor(2, &tpos).output(total as usize),
        ],
        &[],
        [total, 1, 1],
        [total.min(1024), 1, 1],
        true,
    )?;
    let tail_start = top.complete_blocks;
    let ratio = scalar_buffer(&4_i32.to_le_bytes())?;
    qsa_vector(
        if rows == 1 { "ss" } else { "vs" },
        "Multiplyint32",
        &[
            tensor(0, &tail_start).input(rows as usize),
            tensor(1, &ratio).input(1),
            tensor(2, &tail_start).output(rows as usize),
            bytes(3, &rows),
        ],
        rows,
        4,
    )?;
    let count = rows * total;
    let tail = Buffer::mtplx_bytes(u64::from(count))?;
    let causal = Buffer::mtplx_bytes(u64::from(count))?;
    for (op, rhs, out) in [
        ("GreaterEqualint32", &tail_start, &tail),
        ("LessEqualint32", &top.positions, &causal),
    ] {
        if rows == 1 {
            qsa_vector(
                "vs",
                op,
                &[
                    tensor(0, &tpos).input(total as usize),
                    tensor(1, rhs).input(rows as usize),
                    tensor(2, out).output(count as usize),
                    bytes(3, &count),
                ],
                count,
                4,
            )?;
        } else {
            let grid = [total, rows, 1];
            dispatch_geometry(
                &format!("g2_{op}"),
                &[
                    tensor(0, &tpos).input(total as usize),
                    tensor(1, rhs).input(rows as usize),
                    tensor(2, out).output(count as usize),
                    bytes(3, &[0_u64, 1]),
                    bytes(4, &[1_u64, 0]),
                ],
                &[],
                grid,
                block_dims(grid),
                true,
            )?;
        }
    }
    qsa_vector(
        "vv",
        "BitwiseOrbool_",
        &[
            tensor(0, &mask).input(count as usize),
            tensor(1, &tail).input(count as usize),
            tensor(2, &mask).output(count as usize),
            bytes(3, &count),
        ],
        count,
        1,
    )?;
    qsa_vector(
        "vv",
        "BitwiseAndbool_",
        &[
            tensor(0, &mask).input(count as usize),
            tensor(1, &causal).input(count as usize),
            tensor(2, &mask).output(count as usize),
            bytes(3, &count),
        ],
        count,
        1,
    )?;
    temporary.extend([
        selected,
        zero,
        one,
        ratio,
        indices64,
        tpos,
        tail_start,
        tail,
        causal,
        top.indices,
        top.valid,
        top.positions,
    ]);
    Ok(QsaSelection {
        outputs: vec![mask],
        row_widths: vec![(total, 1)],
        row_strides: vec![total],
        _temporaries: temporary,
    })
}

/// Original TensorOps score producer, not a replacement QMM implementation.
#[cfg(test)]
pub(super) fn qsa_prefill_mpp_scores(
    [q, pooled]: [&Buffer; 2],
    [rows, blocks]: [u32; 2],
    q_strides: [u64; 3],
    pooled_strides: [u64; 2],
) -> Result<Buffer, String> {
    let tiles = rows
        .div_ceil(16)
        .checked_mul(blocks.div_ceil(32))
        .and_then(|n| n.checked_mul(128));
    if rows == 0
        || blocks == 0
        || rows > i32::MAX as u32 / 512
        || blocks > i32::MAX as u32
        || tiles.is_none()
    {
        return Err("invalid QSA MPP score geometry".into());
    }
    for (b, shape, strides) in [
        (q, [rows, 4, 128], q_strides),
        (
            pooled,
            [1, blocks, 128],
            [0, pooled_strides[0], pooled_strides[1]],
        ),
    ] {
        let span = shape
            .into_iter()
            .zip(strides)
            .try_fold(1_u64, |n, (d, s)| {
                u64::from(d - 1)
                    .checked_mul(s)
                    .and_then(|v| n.checked_add(v))
            })
            .and_then(|n| n.checked_mul(2));
        if span.is_none_or(|n| unsafe { ds4_gpu_tensor_bytes(b.raw()) } < n) {
            return Err("QSA MPP input outside backing".into());
        }
    }
    use array::{Array, Dtype};
    let inputs = [
        buffer_array(
            q,
            &[1, rows as i32, 4, 128],
            &[0, q_strides[0], q_strides[1], q_strides[2]],
            Dtype::BF16,
        )?,
        buffer_array(
            pooled,
            &[1, blocks as i32, 128],
            &[0, pooled_strides[0], pooled_strides[1]],
            Dtype::BF16,
        )?,
    ];
    let scores = Array::unallocated(&[rows as i32, blocks as i32], Dtype::F32)?;
    qsa_prefill::mpp_evaluate(&inputs, &scores)?;
    let buffer = scores.buffer().view(0, scores.layout().nbytes() as u64)?;
    Ok(buffer)
}

#[cfg(test)]
impl QsaSelectPlan {
    fn prefill_plan(self) -> qsa_prefill::TopK {
        qsa_prefill::TopK {
            blocks: self.blocks,
            topk: self.topk,
            ratio: 4,
            dense_width: if self.mode == QsaSelectMode::DenseMask {
                self.dense_width
            } else {
                0
            },
            mode: self.mode,
        }
    }
    pub(super) fn prefill_source(self) -> Result<String, String> {
        self.prefill_plan().prefill_source()
    }
}

#[cfg(test)]
pub(super) fn qsa_prefill_topk(
    scores: &Buffer,
    strides: [u64; 2],
    plan: QsaSelectPlan,
) -> Result<QsaSelection, String> {
    if plan.rows == 0
        || plan.blocks == 0
        || plan.rows > i32::MAX as u32 / 512
        || !(1..=512).contains(&plan.topk)
        || plan.total > i32::MAX as u32
        || plan
            .pos
            .checked_add(plan.rows)
            .is_none_or(|n| n > plan.total)
        || plan.logical > plan.blocks
        || plan.logical != plan.total / 4
        || plan
            .rows
            .checked_mul(plan.blocks)
            .is_none_or(|n| n > i32::MAX as u32)
        || (plan.mode == QsaSelectMode::DenseMask && plan.dense_width < plan.total)
    {
        return Err("invalid QSA prefill selector geometry".into());
    }
    let positions = [
        scalar_buffer(&plan.pos.to_le_bytes())?,
        scalar_buffer(&plan.total.to_le_bytes())?,
        scalar_buffer(&plan.logical.to_le_bytes())?,
    ];
    let mut out = qsa_prefill_topk_device(
        scores,
        strides,
        plan,
        [&positions[0], &positions[1], &positions[2]],
    )?;
    out._temporaries.extend(positions);
    Ok(out)
}

/// Original prefill top-k with moving frontiers kept on the GPU.
#[cfg(test)]
pub(super) fn qsa_prefill_topk_device(
    scores: &Buffer,
    strides: [u64; 2],
    plan: QsaSelectPlan,
    [pos, total, logical]: [&Buffer; 3],
) -> Result<QsaSelection, String> {
    if plan.rows == 0
        || plan.blocks == 0
        || plan.rows > i32::MAX as u32 / 512
        || !(1..=512).contains(&plan.topk)
        || plan
            .rows
            .checked_mul(plan.blocks)
            .is_none_or(|n| n > i32::MAX as u32)
        || [pos, total, logical]
            .into_iter()
            .any(|b| unsafe { ds4_gpu_tensor_bytes(b.raw()) } < 4)
    {
        return Err("invalid QSA device prefill geometry or leaves".into());
    }
    let span = u64::from(plan.rows - 1)
        .checked_mul(strides[0])
        .and_then(|n| {
            u64::from(plan.blocks - 1)
                .checked_mul(strides[1])
                .and_then(|m| n.checked_add(m))
        })
        .and_then(|n| n.checked_add(1))
        .and_then(|n| n.checked_mul(4));
    if span.is_none_or(|n| unsafe { ds4_gpu_tensor_bytes(scores.raw()) } < n) {
        return Err("QSA score view outside backing".into());
    }
    use array::{Array, Dtype};
    let inputs = [
        buffer_array(
            scores,
            &[plan.rows as i32, plan.blocks as i32],
            &strides,
            Dtype::F32,
        )?,
        Array::new(&[1], Dtype::I32, pos.view(0, 4)?)?,
        Array::new(&[1], Dtype::I32, total.view(0, 4)?)?,
        Array::new(&[1], Dtype::I32, logical.view(0, 4)?)?,
    ];
    let kernel = plan.prefill_plan();
    let outputs = kernel
        .layouts(plan.rows as i32)?
        .into_iter()
        .map(|l| Array::unallocated(l.shape(), l.dtype()))
        .collect::<Result<Vec<_>, _>>()?;
    qsa_prefill::topk_evaluate(&inputs, &outputs, kernel)?;
    let outputs = outputs
        .iter()
        .map(|a| a.buffer().view(0, a.layout().nbytes() as u64))
        .collect::<Result<Vec<_>, _>>()?;
    let widths = plan.widths();
    Ok(QsaSelection {
        outputs,
        row_strides: widths.iter().map(|(w, _)| *w).collect(),
        row_widths: widths,
        _temporaries: Vec::new(),
    })
}

/// qsa_indexer_prefill_metal on M5: BF16 pairs select TensorOps, otherwise the
/// general FP32 expression. Dispatch failures never switch producers.
#[cfg(test)]
pub(super) fn qsa_prefill(
    input: [&Buffer; 2],
    q_strides: [u64; 3],
    pooled_strides: [u64; 2],
    plan: QsaSelectPlan,
    workspace: u64,
) -> Result<QsaSelection, String> {
    if plan.rows <= 1
        || plan.rows > i32::MAX as u32 / 512
        || plan
            .pos
            .checked_add(plan.rows)
            .is_none_or(|n| n > plan.total)
        || plan.logical > plan.blocks
        || plan.logical != plan.total / 4
        || !(1..=512).contains(&plan.topk)
    {
        return Err("invalid QSA prefill input".into());
    }
    let mpp = plan.input_f32 == [false; 2];
    let chunk_rows = qsa_prefill_chunk_rows(plan.rows, plan.blocks, workspace, 32, mpp)?;
    let mut pooled_s = [0, pooled_strides[0], pooled_strides[1]];
    let pooled_cast = if !mpp && !plan.input_f32[1] {
        Some(qsa_cast_f32(
            input[1],
            [1, plan.blocks, 128],
            &mut pooled_s,
        )?)
    } else {
        None
    };
    let pooled = pooled_cast.as_ref().unwrap_or(input[1]);
    let mut chunks = Vec::new();
    for start in (0..plan.rows).step_by(chunk_rows as usize) {
        let rows = chunk_rows.min(plan.rows - start);
        let offset = u64::from(start)
            .checked_mul(q_strides[0])
            .and_then(|n| n.checked_mul(if plan.input_f32[0] { 4 } else { 2 }))
            .ok_or("QSA prefill offset overflow")?;
        let bytes = unsafe { ds4_gpu_tensor_bytes(input[0].raw()) }
            .checked_sub(offset)
            .ok_or("QSA prefill slice outside backing")?;
        let q = input[0].view(offset, bytes)?;
        let (scores, scratch) = if mpp {
            (
                qsa_prefill_mpp_scores(
                    [&q, pooled],
                    [rows, plan.blocks],
                    q_strides,
                    pooled_strides,
                )?,
                Vec::new(),
            )
        } else {
            qsa_prefill_f32_scores(
                [&q, pooled],
                [rows, plan.blocks],
                q_strides,
                [pooled_s[2], pooled_s[1]],
                !plan.input_f32[0],
                plan.tf32,
            )?
        };
        let mut selected = qsa_prefill_topk(
            &scores,
            [u64::from(plan.blocks), 1],
            QsaSelectPlan {
                rows,
                pos: plan.pos + start,
                ..plan
            },
        )?;
        selected._temporaries.push(scores);
        selected._temporaries.extend(scratch);
        chunks.push(selected);
    }
    let mut selected = qsa_concat_selection(chunks, plan, chunk_rows)?;
    selected._temporaries.extend(pooled_cast);
    Ok(selected)
}

#[cfg(test)]
pub(super) fn qsa_prefill_cached(
    cache: &QsaCache,
    prepared_q: &Buffer,
    rows: u32,
    mode: QsaSelectMode,
    workspace: u64,
) -> Result<QsaSelection, String> {
    let total = cache
        .offset
        .checked_add(rows)
        .ok_or("QSA prefill cache position overflow")?;
    let logical = total / 4;
    if cache.pooled_len < logical {
        return Err("QSA prefill needs complete pooled keys".into());
    }
    let pooled = cache
        .pooled
        .as_ref()
        .ok_or("QSA prefill needs pooled backing")?;
    // The large-prefill branch receives _extend_pooled's logical slice, not
    // the reserved cache capacity used by the legacy decode selector.
    qsa_prefill(
        [prepared_q, &pooled.buffer()],
        [512, 128, 1],
        [128, 1],
        QsaSelectPlan {
            rows,
            blocks: logical,
            topk: 512,
            pos: cache.offset,
            total,
            logical,
            dense_width: total,
            mode,
            tf32: true,
            input_f32: [false; 2],
        },
        workspace,
    )
}

#[cfg(test)]
pub(super) fn qsa_select_cached(
    cache: &QsaCache,
    prepared_q: &Buffer,
    rows: u32,
    mode: QsaSelectMode,
    tf32: bool,
) -> Result<QsaSelection, String> {
    let total = cache
        .offset
        .checked_add(rows)
        .ok_or("QSA cache position overflow")?;
    if cache.pooled_len < total / 4 {
        return Err("QSA selector needs complete pooled keys".into());
    }
    let pooled = cache
        .pooled
        .as_ref()
        .ok_or("QSA selector needs pooled backing")?;
    qsa_select_chunked(
        [prepared_q, &pooled.buffer()],
        [512, 128, 1],
        [128, 1],
        QsaSelectPlan {
            rows,
            blocks: cache.pooled_cap,
            topk: 512,
            pos: cache.offset,
            total,
            logical: total / 4,
            dense_width: 0,
            mode,
            tf32,
            input_f32: [false; 2],
        },
    )
}

/// Original eager _rope_cos_sin + _apply_partial_rope, for normalized BF16
/// QSA queries (H4) or pooled keys (H1). Every cast/concat stays explicit.
#[cfg(test)]
fn qsa_eager_rope(
    normalized: Buffer,
    freq: &Buffer,
    [rows, heads]: [u32; 2],
    [position, position_scale]: [i32; 2],
    attention_scale: f32,
) -> Result<(Buffer, Vec<Buffer>), String> {
    let count = rows * heads * 64;
    let positions = Buffer::mtplx_bytes(u64::from(rows) * 4)?;
    dispatch_geometry(
        "arangeint32",
        &[
            bytes(0, &position),
            bytes(1, &1_i32),
            tensor(2, &positions).output(rows as usize),
        ],
        &[],
        [rows, 1, 1],
        [rows.min(1024), 1, 1],
        true,
    )?;
    if position_scale != 1 {
        let scale = scalar_buffer(&position_scale.to_le_bytes())?;
        qsa_vector(
            if rows == 1 { "ss" } else { "vs" },
            "Multiplyint32",
            &[
                tensor(0, &positions).input(rows as usize),
                tensor(1, &scale).input(1),
                tensor(2, &positions).output(rows as usize),
                bytes(3, &rows),
            ],
            rows,
            4,
        )?;
    }
    let fp_positions = Buffer::mtplx_bytes(u64::from(rows) * 4)?;
    qsa_vector(
        "v",
        "copyint32float32",
        &[
            tensor(0, &positions).input(rows as usize),
            tensor(1, &fp_positions).output(rows as usize),
            bytes(2, &rows),
        ],
        rows,
        4,
    )?;
    let angles = Buffer::mtplx_bytes(u64::from(rows) * 32 * 4)?;
    if rows == 1 {
        qsa_vector(
            "sv",
            "Multiplyfloat32",
            &[
                tensor(0, &fp_positions).input(rows as usize),
                tensor(1, freq).input(32),
                tensor(2, &angles).output(rows as usize * 32),
                bytes(3, &32_u32),
            ],
            32,
            4,
        )?;
    } else {
        let grid = [32, rows, 1];
        dispatch_geometry(
            "g2_Multiplyfloat32",
            &[
                tensor(0, &fp_positions).input(rows as usize),
                tensor(1, freq).input(32),
                tensor(2, &angles).output(rows as usize * 32),
                bytes(3, &[1_u64, 0]),
                bytes(4, &[0_u64, 1]),
            ],
            &[],
            grid,
            block_dims(grid),
            true,
        )?;
    }
    let emb = Buffer::mtplx_bytes(u64::from(rows) * 64 * 4)?;
    let concurrent = encoder::concurrent_scope();
    for offset in [0, 32 * 4] {
        let dst = emb.view(offset, (u64::from(rows) * 64 - 32) * 4)?;
        qsa_copy_typed(
            Some(&angles),
            &dst,
            [1, rows, 32],
            [0, 32, 1],
            [0, 64, 1],
            "float32",
            (rows) as usize * 32,
        )?;
    }
    drop(concurrent);
    let cosine = Buffer::mtplx_bytes(u64::from(rows) * 64 * 4)?;
    let sine = Buffer::mtplx_bytes(u64::from(rows) * 64 * 4)?;
    for (op, out) in [("Cosfloat32float32", &cosine), ("Sinfloat32float32", &sine)] {
        qsa_vector(
            "v",
            op,
            &[
                tensor(0, &emb).input(rows as usize * 64),
                tensor(1, out).output(rows as usize * 64),
                bytes(2, &(rows * 64)),
            ],
            rows * 64,
            4,
        )?;
        if attention_scale != 1.0 {
            let scale = scalar_buffer(&attention_scale.to_le_bytes())?;
            qsa_vector(
                "vs",
                "Multiplyfloat32",
                &[
                    tensor(0, out).input(rows as usize * 64),
                    tensor(1, &scale).input(1),
                    tensor(2, out).output(rows as usize * 64),
                    bytes(3, &(rows * 64)),
                ],
                rows * 64,
                4,
            )?;
        }
    }
    let norm_bytes = u64::from(rows) * u64::from(heads) * 128 * 2;
    let x2 = normalized.view(32 * 2, norm_bytes - 32 * 2)?;
    let negative = Buffer::mtplx_bytes(u64::from(rows) * u64::from(heads) * 32 * 2)?;
    if rows * heads == 1 {
        qsa_vector(
            "v",
            "Negativebfloat16bfloat16",
            &[
                tensor(0, &x2).input(((rows * heads - 1) * 128 + 32) as usize),
                tensor(1, &negative).output(rows as usize * heads as usize * 32),
                bytes(2, &32_u32),
            ],
            32,
            2,
        )?;
    } else {
        let grid = [32, rows * heads, 1];
        dispatch_geometry(
            "gn1_Negativebfloat16bfloat16",
            &[
                tensor(0, &x2).input(((rows * heads - 1) * 128 + 32) as usize),
                tensor(1, &negative).output(rows as usize * heads as usize * 32),
                bytes(2, &[rows * heads, 32]),
                bytes(3, &[128_u64, 1]),
                bytes(4, &2_u32),
            ],
            &[],
            grid,
            block_dims(grid),
            true,
        )?;
    }
    let rotated = Buffer::mtplx_bytes(u64::from(count) * 2)?;
    let concurrent = encoder::concurrent_scope();
    for (src, stride, offset) in [(&negative, 32_u64, 0), (&normalized, 128, 32 * 2)] {
        let dst = rotated.view(offset, u64::from(count) * 2 - offset)?;
        qsa_copy_typed(
            Some(src),
            &dst,
            [1, rows * heads, 32],
            [0, stride, 1],
            [0, 64, 1],
            "bfloat16",
            (rows * heads) as usize * 32,
        )?;
    }
    drop(concurrent);
    let mut norm_strides = [u64::from(heads) * 128, 128, 1];
    let x_float = qsa_cast_f32(&normalized, [rows, heads, 64], &mut norm_strides)?;
    let rotated_float = Buffer::mtplx_bytes(u64::from(count) * 4)?;
    cast_bf16_f32(&rotated, &rotated_float, count, true)?;
    for (x, table) in [(&x_float, &cosine), (&rotated_float, &sine)] {
        if heads == 1 {
            qsa_vector(
                "vv",
                "Multiplyfloat32",
                &[
                    tensor(0, x).input(count as usize),
                    tensor(1, table).input(rows as usize * 64),
                    tensor(2, x).output(count as usize),
                    bytes(3, &count),
                ],
                count,
                4,
            )?;
        } else if rows == 1 {
            let grid = [64, heads, 1];
            dispatch_geometry(
                "g2_Multiplyfloat32",
                &[
                    tensor(0, x).input(count as usize),
                    tensor(1, table).input(rows as usize * 64),
                    tensor(2, x).output(count as usize),
                    bytes(3, &[64_u64, 1]),
                    bytes(4, &[0_u64, 1]),
                ],
                &[],
                grid,
                block_dims(grid),
                true,
            )?;
        } else {
            let grid = [64, heads, rows];
            dispatch_geometry(
                "g3_Multiplyfloat32",
                &[
                    tensor(0, x).input(count as usize),
                    tensor(1, table).input(rows as usize * 64),
                    tensor(2, x).output(count as usize),
                    bytes(3, &[u64::from(heads) * 64, 64, 1]),
                    bytes(4, &[64_u64, 0, 1]),
                ],
                &[],
                grid,
                block_dims(grid),
                true,
            )?;
        }
    }
    qsa_vector(
        "vv",
        "Addfloat32",
        &[
            tensor(0, &x_float).input(count as usize),
            tensor(1, &rotated_float).input(count as usize),
            tensor(2, &x_float).output(count as usize),
            bytes(3, &count),
        ],
        count,
        4,
    )?;
    let rope = Buffer::mtplx_bytes(u64::from(count) * 2)?;
    cast_bf16_f32(&x_float, &rope, count, false)?;
    let output = Buffer::mtplx_bytes(norm_bytes)?;
    let pass = normalized.view(64 * 2, norm_bytes - 64 * 2)?;
    let concurrent = encoder::concurrent_scope();
    for (src, stride, offset) in [(&rope, 64_u64, 0), (&pass, 128, 64 * 2)] {
        let dst = output.view(offset, norm_bytes - offset)?;
        qsa_copy_typed(
            Some(src),
            &dst,
            [1, rows * heads, 64],
            [0, stride, 1],
            [0, 128, 1],
            "bfloat16",
            (rows * heads) as usize * 64,
        )?;
    }
    drop(concurrent);
    Ok((
        output,
        vec![
            normalized,
            positions,
            fp_positions,
            angles,
            emb,
            cosine,
            sine,
            negative,
            rotated,
            x_float,
            rotated_float,
            rope,
        ],
    ))
}

#[cfg(test)]
pub(super) fn qsa_prepare_queries_eager(
    [input, norm, freq]: [&Buffer; 3],
    rows: u32,
    strides: [u64; 3],
    position: i32,
) -> Result<(Buffer, Vec<Buffer>), String> {
    let span = u64::from(rows.saturating_sub(1))
        .checked_mul(strides[0])
        .and_then(|n| n.checked_add(3_u64.checked_mul(strides[1])?))
        .and_then(|n| n.checked_add(127_u64.checked_mul(strides[2])?))
        .and_then(|n| n.checked_add(1))
        .and_then(|n| n.checked_mul(2));
    if rows == 0
        || rows > i32::MAX as u32 / 512
        || position.checked_add(rows as i32 - 1).is_none()
        || span.is_none_or(|n| unsafe { ds4_gpu_tensor_bytes(input.raw()) } < n)
        || unsafe { ds4_gpu_tensor_bytes(freq.raw()) } < 128
    {
        return Err("invalid eager QSA query preparation".into());
    }
    // RMSNorm::eval_gpu materializes noncontiguous projection slices, then
    // donates that copy to its output. A contiguous caller input is borrowed.
    let copied = if strides[2] != 1 || strides[1] != 128 || (rows > 1 && strides[0] != 512) {
        Some(qsa_copy_general(
            input,
            [rows, 4, 128],
            strides,
            "bfloat16",
            "bfloat16",
            2,
        )?)
    } else {
        None
    };
    let normalized = if let Some(copy) = copied {
        rms_norm_bf16(&copy, Some(norm), &copy, [rows * 4, 128], 1e-6)?;
        copy
    } else {
        let out = Buffer::mtplx_bytes(u64::from(rows) * 512 * 2)?;
        rms_norm_bf16(input, Some(norm), &out, [rows * 4, 128], 1e-6)?;
        out
    };
    qsa_eager_rope(normalized, freq, [rows, 4], [position, 1], 1.0)
}

#[cfg(test)]
pub(super) fn qsa_pool_keys_eager(
    [input, norm, freq]: [&Buffer; 3],
    blocks: u32,
    [row_stride, dim_stride]: [u64; 2],
    start: i32,
) -> Result<(Buffer, Vec<Buffer>), String> {
    if blocks == 0
        || blocks > i32::MAX as u32 / 512
        || start
            .checked_add(blocks as i32 - 1)
            .and_then(|n| n.checked_mul(4))
            .is_none()
        || unsafe { ds4_gpu_tensor_bytes(freq.raw()) } < 128
    {
        return Err("invalid eager QSA pooling".into());
    }
    let mut strides = [
        row_stride
            .checked_mul(4)
            .ok_or("QSA pool stride overflow")?,
        row_stride,
        dim_stride,
    ];
    let fresh = qsa_cast_f32(input, [blocks, 4, 128], &mut strides)?;
    if strides != [512, 128, 1] {
        return Err("unsupported eager pool physical layout".into());
    }
    let sum = Buffer::mtplx_bytes(u64::from(blocks) * 128 * 4)?;
    let ndim = u32::from(blocks > 1);
    dispatch(
        "col_reduce_small_1_reduce_sumfloat32",
        &[
            tensor(0, &fresh).input(blocks as usize * 512),
            tensor(1, &sum).output(blocks as usize * 128),
            bytes(2, &4_u64),
            bytes(3, &128_u64),
            bytes(4, &if blocks > 1 { blocks } else { 0 }),
            bytes(5, &if blocks > 1 { 512_u64 } else { 0 }),
            bytes(6, &ndim),
            bytes(7, &4_u32),
            bytes(8, &128_u64),
            bytes(9, &1_u32),
            bytes(10, &1_u64),
        ],
        [1, blocks, 1],
        [32, 4, 1],
    )?;
    let count = blocks * 128;
    let mean_scale = scalar_buffer(&0.25_f32.to_le_bytes())?;
    qsa_vector(
        "vs",
        "Multiplyfloat32",
        &[
            tensor(0, &sum).input(count as usize),
            tensor(1, &mean_scale).input(1),
            tensor(2, &sum).output(count as usize),
            bytes(3, &count),
        ],
        count,
        4,
    )?;
    let pooled = Buffer::mtplx_bytes(u64::from(count) * 2)?;
    cast_bf16_f32(&sum, &pooled, count, false)?;
    rms_norm_bf16(&pooled, Some(norm), &pooled, [blocks, 128], 1e-6)?;
    let (out, mut temporary) = qsa_eager_rope(pooled, freq, [blocks, 1], [start, 4], 1.0)?;
    temporary.extend([fresh, sum]);
    Ok((out, temporary))
}

/// Original installed QSA indexer geometry: H4, D128, rotary64, ratio4,
/// epsilon1e-6 and attention scaling1. Strides are physical BF16 elements.
#[cfg(test)]
pub(super) fn qsa_prepare_queries(
    input: [&Buffer; 3],
    rows: u32,
    strides: [u64; 3],
    pos_start: i32,
) -> Result<Buffer, String> {
    if rows == 0 || rows > i32::MAX as u32 / 512 || pos_start.checked_add(rows as i32 - 1).is_none()
    {
        return Err("invalid QSA query position".into());
    }
    let pos = scalar_buffer(&pos_start.to_le_bytes())?;
    qsa_prepare_queries_bound(input, rows, strides, tensor(6, &pos).input(1))
}

#[cfg(test)]
pub(super) fn qsa_prepare_queries_device(
    input: [&Buffer; 3],
    rows: u32,
    strides: [u64; 3],
    pos: &Buffer,
) -> Result<Buffer, String> {
    if unsafe { ds4_gpu_tensor_bytes(pos.raw()) } < 4 {
        return Err("invalid QSA device query position".into());
    }
    qsa_prepare_queries_bound(input, rows, strides, tensor(6, pos).input(1))
}

#[cfg(test)]
fn qsa_prepare_queries_bound(
    [input, norm, inv_freq]: [&Buffer; 3],
    rows: u32,
    [row_stride, head_stride, dim_stride]: [u64; 3],
    pos_start: Binding<'_>,
) -> Result<Buffer, String> {
    let end = u64::from(rows.saturating_sub(1))
        .checked_mul(row_stride)
        .and_then(|n| n.checked_add(3_u64.checked_mul(head_stride)?))
        .and_then(|n| n.checked_add(127_u64.checked_mul(dim_stride)?))
        .and_then(|n| n.checked_add(1))
        .and_then(|n| n.checked_mul(2));
    if rows == 0
        || rows > i32::MAX as u32 / 512
        || dim_stride == 0
        || end.is_none_or(|n| unsafe { ds4_gpu_tensor_bytes(input.raw()) } < n)
        || unsafe { ds4_gpu_tensor_bytes(norm.raw()) } < 128 * 2
        || unsafe { ds4_gpu_tensor_bytes(inv_freq.raw()) } < 32 * 4
    {
        return Err("invalid MTPLX QSA query preparation buffers or shape".into());
    }
    let output = Buffer::mtplx_bytes(u64::from(rows) * 512 * 2)?;
    dispatch(
        "kernel_qwen_mtplx_qsa_prepare_q_bf16",
        &[
            tensor(0, input).input((end.unwrap() / 2) as usize),
            bytes(1, &[0_u64, row_stride, head_stride, dim_stride]),
            tensor(2, norm).input(128),
            bytes(3, &1_u64),
            tensor(4, inv_freq).input(32),
            bytes(5, &1_u64),
            pos_start,
            tensor(7, &output).output(rows as usize * 512),
        ],
        [rows * 4, 1, 1],
        [32, 1, 1],
    )?;
    Ok(output)
}

#[cfg(test)]
pub(super) fn qsa_pool_keys(
    inputs: [&Buffer; 3],
    blocks: u32,
    strides: [u64; 2],
    block_start: i32,
) -> Result<Buffer, String> {
    if blocks == 0
        || blocks > i32::MAX as u32 / 4
        || block_start
            .checked_add(blocks as i32 - 1)
            .and_then(|n| n.checked_mul(4))
            .is_none()
    {
        return Err("invalid MTPLX QSA pooled-key position".into());
    }
    let start = scalar_buffer(&block_start.to_le_bytes())?;
    qsa_pool_keys_bound(inputs, blocks, strides, tensor(6, &start).input(1))
}

#[cfg(test)]
fn qsa_pool_keys_bound(
    [input, norm, inv_freq]: [&Buffer; 3],
    blocks: u32,
    [row_stride, dim_stride]: [u64; 2],
    block_start: Binding<'_>,
) -> Result<Buffer, String> {
    let tokens = u64::from(blocks) * 4;
    let end = tokens
        .saturating_sub(1)
        .checked_mul(row_stride)
        .and_then(|n| n.checked_add(127_u64.checked_mul(dim_stride)?))
        .and_then(|n| n.checked_add(1))
        .and_then(|n| n.checked_mul(2));
    if blocks == 0
        || tokens > i32::MAX as u64
        || dim_stride == 0
        || end.is_none_or(|n| unsafe { ds4_gpu_tensor_bytes(input.raw()) } < n)
        || unsafe { ds4_gpu_tensor_bytes(norm.raw()) } < 128 * 2
        || unsafe { ds4_gpu_tensor_bytes(inv_freq.raw()) } < 32 * 4
    {
        return Err("invalid MTPLX QSA pooled-key preparation buffers or shape".into());
    }
    let output = Buffer::mtplx_bytes(u64::from(blocks) * 128 * 2)?;
    dispatch(
        "kernel_qwen_mtplx_qsa_pool_k_bf16",
        &[
            tensor(0, input).input((end.unwrap() / 2) as usize),
            bytes(1, &[0_u64, row_stride, dim_stride]),
            tensor(2, norm).input(128),
            bytes(3, &1_u64),
            tensor(4, inv_freq).input(32),
            bytes(5, &1_u64),
            block_start,
            tensor(7, &output).output(blocks as usize * 128),
        ],
        [blocks, 1, 1],
        [32, 1, 1],
    )?;
    Ok(output)
}

fn qsa_compiled_clamp_name(max_new: u32, max_start: u32) -> String {
    format!("kernel_qwen_mtplx_qsa_clamp_n{max_new}_p{max_start}")
}

/// Original compiler's three constant-CSE layouts, with only shape constants
/// and the exported symbol specialized. Dynamic frontier leaves stay inputs.
pub(super) fn qsa_compiled_clamp_source(max_new: u32, max_start: u32) -> Result<String, String> {
    if max_new == 0 || max_new > i32::MAX as u32 || max_start > i32::MAX as u32 {
        return Err("invalid compiled QSA clamp geometry".into());
    }
    let bank: serde_json::Value = serde_json::from_str(include_str!(
        "../../../metal/mtplx-qsa-compiled-scalars.json"
    ))
    .map_err(|e| e.to_string())?;
    let kind = if max_start == 0 {
        "zero"
    } else if max_start == max_new {
        "equal"
    } else {
        "distinct"
    };
    let variant = &bank["clamp"][kind];
    let body = variant["kernels"]
        .as_str()
        .ok_or("missing original compiled QSA clamp")?;
    let original_name = body
        .split_once("host_name(\"")
        .and_then(|(_, s)| s.split_once('"'))
        .ok_or("invalid compiled QSA symbol")?
        .0;
    let mut body = body.replace(original_name, &qsa_compiled_clamp_name(max_new, max_start));
    body = body.replace(
        &format!("tmp_D = static_cast<int32_t>({});", variant["max_new"]),
        &format!("tmp_D = static_cast<int32_t>({max_new});"),
    );
    if kind == "distinct" {
        body = body.replace(
            &format!("tmp_F = static_cast<int32_t>({});", variant["max_start"]),
            &format!("tmp_F = static_cast<int32_t>({max_start});"),
        );
    }
    Ok(body)
}

#[cfg(test)]
fn qsa_compiled_frontiers(
    [logical, pooled_len]: [&Buffer; 2],
    [max_new, max_start]: [u32; 2],
) -> Result<[Buffer; 2], String> {
    let block_start = qsa_compiled_block_start([logical, pooled_len], [max_new, max_start])?;
    let raw_start = qsa_compiled_raw_start(&block_start)?;
    Ok([block_start, raw_start])
}

#[cfg(test)]
fn qsa_compiled_block_start(
    [logical, pooled_len]: [&Buffer; 2],
    [max_new, max_start]: [u32; 2],
) -> Result<Buffer, String> {
    qsa_compile::Scalar::Clamp { max_new, max_start }.buffer(&[logical, pooled_len])
}

#[cfg(test)]
fn qsa_compiled_raw_start(block_start: &Buffer) -> Result<Buffer, String> {
    qsa_compile::Scalar::Multiply(4).buffer(&[block_start])
}

#[cfg(test)]
fn qsa_retained_bf16_copy(input: &Buffer, count: u32) -> Result<Buffer, String> {
    let output = Buffer::mtplx_bytes(u64::from(count) * 2)?;
    let work = if count < 65536 { 1 } else { 4 };
    let threads = count.div_ceil(work);
    dispatch_geometry(
        if work == 1 {
            "v_copybfloat16bfloat16"
        } else {
            "vn_copybfloat16bfloat16"
        },
        &[
            tensor(0, input).input(count as usize),
            tensor(1, &output).output(count as usize),
            bytes(2, &count),
        ],
        &[],
        [threads, 1, 1],
        [threads.min(1024), 1, 1],
        true,
    )?;
    Ok(output)
}

#[cfg(test)]
pub(super) struct QsaCompiledCacheWindow {
    pub(super) raw: Buffer,
    pub(super) pooled: Buffer,
    pub(super) block_start: Buffer,
    pub(super) raw_start: Buffer,
    _temporaries: Vec<Buffer>,
}

/// Connected cache-maintenance subgraph of QSACompiledIndexerCore.from_qk_rows.
/// This entry retains the old leaves, so the reference's non-donation path is
/// required. Graph-bank ownership/donation and BFS scheduling are separate.
/// Position/logical/frontier are validated int32[1] device leaves; no readback.
#[cfg(test)]
pub(super) fn qsa_compiled_cache_window(
    [qk, raw, pooled, norm, freq]: [&Buffer; 5],
    [pos, logical, pooled_len]: [&Buffer; 3],
    rows: u32,
    qk_stride: u64,
    [raw_cap, pooled_cap]: [u32; 2],
) -> Result<QsaCompiledCacheWindow, String> {
    let max_new = rows.div_ceil(4);
    let max_start = (raw_cap / 4).min(pooled_cap).checked_sub(max_new);
    if rows == 0
        || rows > raw_cap
        || qk_stride < 640
        || max_start.is_none()
        || [raw_cap, pooled_cap]
            .into_iter()
            .any(|n| n < 256 || !n.is_power_of_two() || n > i32::MAX as u32 / 512)
        || [pos, logical, pooled_len]
            .into_iter()
            .any(|b| unsafe { ds4_gpu_tensor_bytes(b.raw()) } < 4)
        || unsafe { ds4_gpu_tensor_bytes(raw.raw()) } < u64::from(raw_cap) * 256
        || unsafe { ds4_gpu_tensor_bytes(pooled.raw()) } < u64::from(pooled_cap) * 256
        || u64::from(rows - 1)
            .checked_mul(qk_stride)
            .and_then(|n| n.checked_add(640))
            .and_then(|n| n.checked_mul(2))
            .is_none_or(|n| unsafe { ds4_gpu_tensor_bytes(qk.raw()) } < n)
    {
        return Err("invalid compiled QSA cache-window inputs or geometry".into());
    }
    let raw_next = qsa_retained_bf16_copy(raw, raw_cap * 128)?;
    let k = qk.view(512 * 2, (u64::from(rows - 1) * qk_stride + 128) * 2)?;
    let mut temporary = qsa_dynamic_copy_rows(
        [&k, &raw_next],
        [qk_stride, 128],
        [rows, 128],
        [None, Some(pos)],
        [
            ((rows - 1) as u64 * qk_stride + 128) as usize,
            raw_cap as usize * 128,
        ],
    )?;
    let [block_start, raw_start] =
        qsa_compiled_frontiers([logical, pooled_len], [max_new, max_start.unwrap()])?;
    let window = Buffer::mtplx_bytes(u64::from(max_new) * 4 * 128 * 2)?;
    temporary.extend(qsa_dynamic_copy_rows(
        [&raw_next, &window],
        [128, 128],
        [max_new * 4, 128],
        [Some(&raw_start), None],
        [raw_cap as usize * 128, max_new as usize * 512],
    )?);
    let pool_window = qsa_pool_keys_bound(
        [&window, norm, freq],
        max_new,
        [128, 1],
        tensor(6, &block_start).input(1),
    )?;
    let pooled_next = qsa_retained_bf16_copy(pooled, pooled_cap * 128)?;
    temporary.extend(qsa_dynamic_copy_rows(
        [&pool_window, &pooled_next],
        [128, 128],
        [max_new, 128],
        [None, Some(&block_start)],
        [max_new as usize * 128, pooled_cap as usize * 128],
    )?);
    temporary.extend([k, window, pool_window]);
    Ok(QsaCompiledCacheWindow {
        raw: raw_next,
        pooled: pooled_next,
        block_start,
        raw_start,
        _temporaries: temporary,
    })
}

#[cfg(test)]
pub(super) struct QsaProjected {
    pub(super) q: Buffer,
    pub(super) raw_keys: Buffer,
    _temporary: Vec<Buffer>,
}

#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
pub(super) enum QsaCompiledMode {
    UpdateOnly,
    Blocks,
    PrefillBlocks,
    RowTokens,
    DenseMask,
}

pub(super) fn qsa_compiled_row_position_source(start: u32) -> Result<String, String> {
    if start > i32::MAX as u32 {
        return Err("invalid compiled QSA row offset".into());
    }
    let bank: serde_json::Value = serde_json::from_str(include_str!(
        "../../../tests/fixtures/mtplx-qsa-select-jit.json"
    ))
    .map_err(|e| e.to_string())?;
    let body = bank[3]["kernels"]
        .as_str()
        .ok_or("missing original compiled row offset")?;
    let symbol = body
        .split_once("host_name(\"")
        .and_then(|(_, s)| s.split_once('"'))
        .ok_or("invalid compiled row offset symbol")?
        .0;
    Ok(body
        .replace(symbol, &format!("kernel_qwen_mtplx_qsa_pos_add_{start}"))
        .replace(
            "tmp_C = static_cast<int32_t>(3);",
            &format!("tmp_C = static_cast<int32_t>({start});"),
        ))
}

#[cfg(test)]
fn qsa_compiled_row_position(pos: &Buffer, start: u32) -> Result<Buffer, String> {
    qsa_compile::Scalar::Add(start).buffer(&[pos])
}

#[cfg(test)]
pub(super) struct QsaIndexerWeights<'a> {
    pub(super) projection: [&'a Buffer; 3],
    pub(super) quant: [u32; 2],
    pub(super) q_norm: &'a Buffer,
    pub(super) k_norm: &'a Buffer,
    pub(super) inv_freq: &'a Buffer,
}

/// Resolved MTPLX switches for both indexer routes. Prefill's
/// device-dependent AUTO value must be resolved before entering this path.
pub(super) struct QsaIndexerOptions {
    pub(super) fused: bool,
    pub(super) compiled: bool,
    pub(super) prefill: bool,
    pub(super) phase_prefill: bool,
    pub(super) prefill_min_rows: u32,
    pub(super) prefill_compile_rows: u32,
    pub(super) prefill_min_context: u32,
    pub(super) prefill_workspace: u64,
    pub(super) tile_rows: u32,
    pub(super) flash: bool,
    pub(super) gather_decode: bool,
    pub(super) gather: bool,
    pub(super) gather_max_rows: u32,
    pub(super) gather_min_context: u32,
    #[cfg_attr(not(test), expect(dead_code, reason = "Reference diagnostic API"))]
    pub(super) tf32: bool,
}

impl QsaIndexerOptions {
    fn large_prefill(&self, rows: u32, pos: u32) -> bool {
        self.prefill
            && self.phase_prefill
            && rows >= self.prefill_min_rows
            && pos >= self.prefill_min_context
    }

    /// Original _compiled_mode; total is the post-update token frontier.
    pub(super) fn compiled_mode(
        &self,
        rows: u32,
        total: u32,
        ratio: u32,
        topk: u32,
    ) -> QsaCompiledMode {
        if total / ratio <= topk {
            QsaCompiledMode::UpdateOnly
        } else if rows == 1 && self.flash {
            QsaCompiledMode::Blocks
        } else if rows != 1 && self.large_prefill(rows, total.saturating_sub(rows)) {
            QsaCompiledMode::PrefillBlocks
        } else if rows != 1
            && self.gather
            && rows <= self.gather_max_rows
            && total >= self.gather_min_context
        {
            QsaCompiledMode::RowTokens
        } else {
            QsaCompiledMode::DenseMask
        }
    }

    /// Cache/phase portion of _compiled_route_supported, BEFORE capacity
    /// bucketing. QsaCache fixes B1/BF16/D128/ratio4. This does not validate a
    /// source tensor, projection/norm dtype or device; those remain entry-point
    /// obligations. In particular it is not a generic tensor eligibility test.
    #[cfg(test)]
    pub(super) fn compiled_cache_mode(
        &self,
        cache: &QsaCache,
        rows: u32,
        pos: u32,
    ) -> Option<QsaCompiledMode> {
        let total = pos.checked_add(rows)?;
        if !self.fused || !self.compiled || rows == 0 || pos != cache.offset {
            return None;
        }
        let mode = self.compiled_mode(rows, total, 4, 512);
        if rows != 1 {
            if mode == QsaCompiledMode::UpdateOnly && self.phase_prefill {
                return None;
            }
            if rows >= self.prefill_min_rows
                && (!matches!(
                    mode,
                    QsaCompiledMode::PrefillBlocks | QsaCompiledMode::UpdateOnly
                ) || rows != self.prefill_compile_rows)
            {
                return None;
            }
        }
        if pos > 0 && (cache.raw.is_none() || cache.raw_cap < pos) {
            return None;
        }
        if cache.pooled.is_none() {
            if cache.pooled_len != 0 {
                return None;
            }
        } else if cache.pooled_len > cache.pooled_cap {
            return None;
        }
        let logical = total / 4;
        if logical - cache.pooled_len.min(logical) > rows.div_ceil(4) {
            return None;
        }
        if rows == 1
            && !matches!(mode, QsaCompiledMode::UpdateOnly | QsaCompiledMode::Blocks)
            && self.gather_decode
        {
            return None;
        }
        Some(mode)
    }
}

#[cfg(test)]
pub(super) struct QsaIndexed {
    pub(super) selection: Option<QsaIndexerSelection>,
    pub(super) lane: &'static str,
    _temporaries: Vec<Buffer>,
}

/// Connected _call_rows host flow for the installed, typed B1/BF16 indexer.
/// Replays the parameter-bound indexer graph with its exclusive cache state.
/// Shared descriptors/Data retain the reference's copying path.
/// Eligibility rejection selects eager; dispatch errors are returned, not hidden.
#[allow(clippy::too_many_arguments)] // Explicit model inputs, cache and evaluator state.
#[cfg(test)]
pub(super) fn qsa_indexer_rows(
    weights: QsaIndexerWeights<'_>,
    hidden: &Buffer,
    supplied_qk: Option<(&Buffer, u64)>,
    rows: u32,
    cache: &mut QsaCache,
    options: &QsaIndexerOptions,
    commands: Option<&mut Commands>,
    graph: &mut QsaGraphBank,
) -> Result<QsaIndexed, String> {
    let total = qsa_validate_indexer_options(rows, cache.offset, options)?;
    let Some(mode) = options.compiled_cache_mode(cache, rows, cache.offset) else {
        return qsa_indexer_rows_uncompiled(
            weights,
            hidden,
            supplied_qk,
            rows,
            cache,
            options,
            commands,
        );
    };
    cache.ensure_compiled_backings(rows)?;
    let logical = total / 4;
    let frontiers = [cache.offset, total, logical, cache.pooled_len.min(logical)]
        .map(|value| scalar_buffer(&(value as i32).to_le_bytes()));
    let [pos, end, blocks, frontier] = frontiers;
    let frontiers = [pos?, end?, blocks?, frontier?];
    let budget = if mode == QsaCompiledMode::PrefillBlocks {
        options.prefill_workspace
    } else {
        32 * 1024 * 1024
    };
    let geometry = [rows, cache.raw_cap, cache.pooled_cap];
    let (updated, selected) = graph.run(
        weights,
        supplied_qk.map_or(hidden, |(q, _)| q),
        QsaCacheInput::State([cache.raw.as_mut().unwrap(), cache.pooled.as_mut().unwrap()]),
        frontiers.each_ref(),
        geometry,
        supplied_qk.map(|(_, stride)| stride),
        mode,
        budget,
        options.tf32,
    )?;
    // Commit leaves/frontier without waiting or reading GPU scalars. Attention
    // alone advances KV.offset. A former eager mirror is now invalid.
    cache.raw = Some(QsaArray::new(updated.raw, 0)?);
    cache.pooled = Some(QsaArray::new(updated.pooled, 0)?);
    cache.pooled_len = logical;
    cache.mirror = None;
    cache.mirror_cap = 0;
    cache.mirror_strides = [0; 2];
    let mut temporary = updated._temporaries;
    temporary.extend(frontiers);
    let selection = if mode == QsaCompiledMode::UpdateOnly {
        None
    } else {
        let mut out = selected.ok_or("compiled QSA omitted its selection")?;
        Some(match mode {
            QsaCompiledMode::Blocks => {
                out._temporaries.extend(out.outputs.drain(1..));
                out.row_widths.truncate(1);
                out.row_strides.truncate(1);
                QsaIndexerSelection::FlashBlocks(out, total / 4 * 4)
            }
            QsaCompiledMode::PrefillBlocks => {
                out._temporaries.extend(out.outputs.drain(2..));
                out.row_widths.truncate(2);
                out.row_strides.truncate(2);
                QsaIndexerSelection::PrefillBlocks(out)
            }
            QsaCompiledMode::RowTokens => QsaIndexerSelection::RowTokens(out),
            QsaCompiledMode::DenseMask => {
                out.row_widths[0] = (total, 1);
                QsaIndexerSelection::Dense(out)
            }
            QsaCompiledMode::UpdateOnly => unreachable!(),
        })
    };
    Ok(QsaIndexed {
        selection,
        lane: "compiled",
        _temporaries: temporary,
    })
}

#[cfg(test)]
fn qsa_validate_indexer_options(
    rows: u32,
    pos: u32,
    options: &QsaIndexerOptions,
) -> Result<u32, String> {
    let total = pos
        .checked_add(rows)
        .filter(|n| rows > 0 && *n <= i32::MAX as u32 / 512)
        .ok_or("invalid QSA indexer frontier")?;
    if options.prefill_min_rows < 2
        || options.prefill_compile_rows < 2
        || options.prefill_min_context < 2049
        || options.prefill_workspace < 1024 * 1024
        || options.gather_max_rows < 2
    {
        return Err("unresolved or invalid QSA indexer options".into());
    }
    Ok(total)
}

/// _call_rows after compiled-route rejection, for the installed B1 BF16 model.
/// Does not replace an eligible compiled route. Attention advances KV.offset
/// later; this call only writes the indexer's raw/pool state.
#[cfg(test)]
pub(super) fn qsa_indexer_rows_uncompiled(
    weights: QsaIndexerWeights<'_>,
    hidden: &Buffer,
    supplied_qk: Option<(&Buffer, u64)>,
    rows: u32,
    cache: &mut QsaCache,
    options: &QsaIndexerOptions,
    commands: Option<&mut Commands>,
) -> Result<QsaIndexed, String> {
    let total = qsa_validate_indexer_options(rows, cache.offset, options)?;
    let mut temporary = Vec::new();
    let mut projected = None;
    let (qk, stride) = if let Some((qk, stride)) = supplied_qk {
        let span = u64::from(rows - 1)
            .checked_mul(stride)
            .and_then(|n| n.checked_add(640))
            .and_then(|n| n.checked_mul(2));
        if stride < 640 || span.is_none_or(|n| unsafe { ds4_gpu_tensor_bytes(qk.raw()) } < n) {
            return Err("invalid supplied QSA projection view".into());
        }
        (qk, stride)
    } else {
        let output = Buffer::mtplx_bytes(u64::from(rows) * 640 * 2)?;
        let [w, s, b] = weights.projection;
        temporary.extend(affine_bf16(
            [w, s, b, hidden, &output],
            [rows, 640, 2560],
            weights.quant[0],
            weights.quant[1],
            true,
        )?);
        projected = Some(output);
        (projected.as_ref().unwrap(), 640)
    };
    // In the dense==sparse regime q preparation is a dead lazy expression.
    // Projection/raw writes and completed-key pooling still belong to cache.
    let query = if total / 4 > 512 {
        if options.fused {
            Some(qsa_prepare_queries(
                [qk, weights.q_norm, weights.inv_freq],
                rows,
                [stride, 128, 1],
                cache.offset as i32,
            )?)
        } else {
            let (q, scratch) = qsa_prepare_queries_eager(
                [qk, weights.q_norm, weights.inv_freq],
                rows,
                [stride, 128, 1],
                cache.offset as i32,
            )?;
            temporary.extend(scratch);
            Some(q)
        }
    } else {
        None
    };
    let raw_len = (u64::from(rows - 1) * stride + 128) * 2;
    let raw = qk.view(512 * 2, raw_len)?;
    cache.write_raw(&raw, rows, stride)?;
    temporary.extend(cache.extend_pooled(
        total,
        weights.k_norm,
        weights.inv_freq,
        !options.fused,
    )?);
    temporary.extend(projected);
    let Some(query) = query else {
        return Ok(QsaIndexed {
            selection: None,
            lane: "update_only",
            _temporaries: temporary,
        });
    };
    let decode = rows == 1;
    let large_prefill = options.large_prefill(rows, cache.offset);
    // All prepared B1/H4/D128 BF16 arrays meet the original M5 eligibility.
    let (selection, lane) = if !decode && large_prefill {
        let mut out = qsa_prefill_cached(
            cache,
            &query,
            rows,
            QsaSelectMode::Blocks,
            options.prefill_workspace,
        )?;
        out._temporaries.extend(out.outputs.drain(2..));
        out.row_widths.truncate(2);
        out.row_strides.truncate(2);
        (QsaIndexerSelection::PrefillBlocks(out), "metal_prefill")
    } else if options.fused
        && (decode || rows < options.prefill_min_rows)
        && !(decode && options.gather_decode)
    {
        if decode && options.flash {
            let mut out =
                qsa_select_cached(cache, &query, rows, QsaSelectMode::Blocks, options.tf32)?;
            out._temporaries.extend(out.outputs.drain(1..));
            out.row_widths.truncate(1);
            out.row_strides.truncate(1);
            (
                QsaIndexerSelection::FlashBlocks(out, total / 4 * 4),
                "legacy_fused",
            )
        } else if !decode
            && options.gather
            && rows <= options.gather_max_rows
            && total >= options.gather_min_context
        {
            (
                QsaIndexerSelection::RowTokens(qsa_select_cached(
                    cache,
                    &query,
                    rows,
                    QsaSelectMode::RowTokens,
                    options.tf32,
                )?),
                "legacy_fused",
            )
        } else {
            (
                QsaIndexerSelection::Dense(qsa_select_cached(
                    cache,
                    &query,
                    rows,
                    QsaSelectMode::DenseMask,
                    options.tf32,
                )?),
                "legacy_fused",
            )
        }
    } else {
        let pos = cache.offset;
        let (pooled, strides) = cache.pooled_f32_view(total / 4)?;
        let top = qsa_eager_topk(
            [&query, pooled],
            [512, 128, 1],
            strides,
            QsaSelectPlan {
                rows,
                blocks: total / 4,
                topk: 512,
                pos,
                total,
                logical: total / 4,
                dense_width: total,
                mode: QsaSelectMode::Blocks,
                tf32: options.tf32,
                input_f32: [false, true],
            },
            options.tile_rows,
            commands,
        )?;
        (
            qsa_eager_output(
                top,
                total,
                QsaEagerOptions {
                    large_prefill,
                    flash: options.flash,
                    gather_decode: options.gather_decode,
                    gather: options.gather,
                    gather_max_rows: options.gather_max_rows,
                    gather_min_context: options.gather_min_context,
                },
            )?,
            "eager",
        )
    };
    temporary.push(query);
    Ok(QsaIndexed {
        selection: Some(selection),
        lane,
        _temporaries: temporary,
    })
}

/// Shared-input indexer projection followed by the selected original fused or
/// eager preparation. Raw keys retain the projection stride for cache.write_raw.
#[cfg(test)]
pub(super) fn qsa_project_prepare(
    pack: [&Buffer; 3],
    [norm, inv_freq]: [&Buffer; 2],
    input: &Buffer,
    rows: u32,
    quant: [u32; 2],
    pos_start: i32,
    eager: bool,
) -> Result<QsaProjected, String> {
    if rows == 0 || rows > i32::MAX as u32 / 640 {
        return Err("invalid MTPLX QSA projection rows".into());
    }
    let projected = Buffer::mtplx_bytes(u64::from(rows) * 640 * 2)?;
    let [w, s, b] = pack;
    let mut temporary = affine_bf16(
        [w, s, b, input, &projected],
        [rows, 640, 2560],
        quant[0],
        quant[1],
        true,
    )?;
    let q = if eager {
        let (q, scratch) = qsa_prepare_queries_eager(
            [&projected, norm, inv_freq],
            rows,
            [640, 128, 1],
            pos_start,
        )?;
        temporary.extend(scratch);
        q
    } else {
        qsa_prepare_queries([&projected, norm, inv_freq], rows, [640, 128, 1], pos_start)?
    };
    let raw_keys = projected.view(512 * 2, (u64::from(rows) * 640 - 512) * 2)?;
    temporary.push(projected);
    Ok(QsaProjected {
        q,
        raw_keys,
        _temporary: temporary,
    })
}

#[cfg(test)]
pub(super) enum GdnInputWeights<'a> {
    Fused {
        pack: [&'a Buffer; 3],
        quant: [u32; 2],
    },
    /// Original qkv, z, b, a module order, each with its own quantization.
    Separate {
        packs: [[&'a Buffer; 3]; 4],
        quant: [[u32; 2]; 4],
    },
}

#[cfg(test)]
pub(super) struct GdnWeights<'a> {
    pub(super) input: GdnInputWeights<'a>,
    pub(super) output: [&'a Buffer; 3],
    pub(super) output_quant: [u32; 2], // bits, group
    pub(super) conv: &'a Buffer,
    pub(super) a_log: &'a Buffer,
    pub(super) dt_bias: &'a Buffer,
    pub(super) norm: &'a Buffer,
}

#[cfg(test)]
struct GdnProjected {
    qkv: Buffer,
    z: Buffer,
    b: Buffer,
    a: Buffer,
    strides: [u32; 3], // qkv, z, a/b
    temporary: Vec<Buffer>,
}

#[cfg(test)]
fn gdn_project(
    weights: &GdnInputWeights<'_>,
    input: &Buffer,
    rows: u32,
) -> Result<GdnProjected, String> {
    let count = u64::from(rows) * 16480;
    if rows == 0 || count > i32::MAX as u64 {
        return Err("invalid MTPLX GDN projection rows".into());
    }
    let mut temporary = Vec::new();
    match weights {
        GdnInputWeights::Fused {
            pack: [w, s, b],
            quant: [bits, group],
        } => {
            let projected = Buffer::mtplx_bytes(count * 2)?;
            temporary.extend(affine_bf16(
                [w, s, b, input, &projected],
                [rows, 16480, 2560],
                *bits,
                *group,
                true,
            )?);
            let qkv = projected.view(0, (count - 6240) * 2)?;
            let z = projected.view(10240 * 2, (count - 10336) * 2)?;
            let b = projected.view(16384 * 2, (count - 16432) * 2)?;
            let a = projected.view(16432 * 2, (count - 16432) * 2)?;
            temporary.push(projected);
            Ok(GdnProjected {
                qkv,
                z,
                b,
                a,
                strides: [16480; 3],
                temporary,
            })
        }
        GdnInputWeights::Separate { packs, quant } => {
            let mut project = |index: usize, width| {
                let [w, s, b] = packs[index];
                let [bits, group] = quant[index];
                let output = Buffer::mtplx_bytes(u64::from(rows) * u64::from(width) * 2)?;
                temporary.extend(affine_bf16(
                    [w, s, b, input, &output],
                    [rows, width, 2560],
                    bits,
                    group,
                    true,
                )?);
                Ok::<_, String>(output)
            };
            // Do not sort these by width or coalesce independently quantized modules.
            let qkv = project(0, 10240)?;
            let z = project(1, 6144)?;
            let b = project(2, 48)?;
            let a = project(3, 48)?;
            Ok(GdnProjected {
                qkv,
                z,
                b,
                a,
                strides: [10240, 6144, 48],
                temporary,
            })
        }
    }
}

#[cfg(test)]
pub(super) struct GdnDecodeOutput {
    pub(super) hidden: Buffer,
    pub(super) conv: Buffer,
    pub(super) delta: Buffer,
    // Like the reference encoder's temporaries, live through command completion.
    _temporaries: Vec<Buffer>,
}

#[cfg(test)]
pub(super) struct GdnStagedOutput {
    pub(super) hidden: Buffer,
    pub(super) prepared: ConvNormOutput,
    pub(super) recurrent: DeltaUpdateOutput,
    _norm: Option<GatedNormOutput>,
    _temporaries: Vec<Buffer>,
}

// Keep the fixed GPU-handle bundle inline, without an allocation per layer.
#[allow(clippy::large_enum_variant)]
#[cfg(test)]
pub(super) enum GdnForwardOutput {
    Fused(GdnDecodeOutput),
    Staged(GdnStagedOutput),
}

/// GatedDeltaNet's inference branch order for the installed B=1 geometry.
/// Flags: fused step, conv-single, conv-verify, fused output, verify capture.
/// State bindings: BF16 conv, FP32 delta, optional INT32 valid length.
#[cfg(test)]
pub(super) fn gdn_forward(
    weights: &GdnWeights<'_>,
    input: &Buffer,
    rows: u32,
    cache: [Option<&Buffer>; 3],
    mask: Option<&Buffer>,
    flags: [bool; 5],
    has_cache: bool,
) -> Result<GdnForwardOutput, String> {
    let [conv, delta, length] = if has_cache { cache } else { [None; 3] };
    if rows == 1
        && has_cache
        && flags[0]
        && !flags[4]
        && mask.is_none()
        && length.is_none()
        && let Some(delta) = delta
    {
        gdn_decode(weights, input, conv, delta).map(GdnForwardOutput::Fused)
    } else {
        gdn_staged(
            weights,
            input,
            rows,
            [conv, delta, length],
            mask,
            [flags[1], flags[2], flags[3]],
            has_cache,
        )
        .map(GdnForwardOutput::Staged)
    }
}

/// Connected cold/prefill/verify forward with either input projection layout.
/// Fused-step selection remains the caller's earlier branch; these are the
/// conv-single, conv-verify and output-fusion flags from the actual model.
#[cfg(test)]
pub(super) fn gdn_staged(
    weights: &GdnWeights<'_>,
    input: &Buffer,
    rows: u32,
    [conv, delta, length]: [Option<&Buffer>; 3],
    mask: Option<&Buffer>,
    flags: [bool; 3],
    has_cache: bool,
) -> Result<GdnStagedOutput, String> {
    let GdnProjected {
        qkv,
        z,
        b,
        a,
        strides,
        mut temporary,
    } = gdn_project(&weights.input, input, rows)?;
    let prepared = gdn_conv_norm(
        &qkv,
        weights.conv,
        [rows, strides[0]],
        [conv, length],
        mask,
        [flags[0], flags[1]],
        has_cache,
    )?;
    let recurrent = gated_delta_update(
        [&prepared.q, &prepared.k, &prepared.v],
        [&a, &b, weights.a_log, weights.dt_bias],
        [rows, strides[2]],
        delta,
        mask,
    )?;
    let hidden = Buffer::mtplx_bytes(u64::from(rows) * 2560 * 2)?;
    let norm = if rows == 1
        && flags[2]
        && weights.output_quant[0] == 4
        && [32, 64].contains(&weights.output_quant[1])
    {
        let [w, s, b] = weights.output;
        dispatch(
            &format!(
                "kernel_qwen_mtplx_gdn_out_fused_gs{}_bf16",
                weights.output_quant[1]
            ),
            &[
                tensor(0, &recurrent.hidden).input(6144),
                tensor(1, &z).input(6144),
                tensor(2, weights.norm).input(128),
                tensor(3, w).input(2560 * 6144 / 8),
                tensor(4, s).input(2560 * 6144 / weights.output_quant[1] as usize),
                tensor(5, b).input(2560 * 6144 / weights.output_quant[1] as usize),
                tensor(6, &hidden).output(2560),
            ],
            [2560_u32.div_ceil(32), 1, 1],
            [1024, 1, 1],
        )?;
        None
    } else {
        let norm = sigmoid_rms_norm(
            &recurrent.hidden,
            weights.norm,
            Some(&z),
            [rows, 48, 128, strides[1]],
            1e-6,
        )?;
        let [w, s, b] = weights.output;
        temporary.extend(affine_bf16(
            [w, s, b, &norm.hidden, &hidden],
            [rows, 2560, 6144],
            weights.output_quant[0],
            weights.output_quant[1],
            true,
        )?);
        Some(norm)
    };
    temporary.extend([qkv, z, b, a]);
    Ok(GdnStagedOutput {
        hidden,
        prepared,
        recurrent,
        _norm: norm,
        _temporaries: temporary,
    })
}

/// The reference's live-state, unmasked, single-row fused-step branch.
/// Caller checks the model branch gate; cold/prefill/verify use different graphs.
/// Keep this result alive until the enclosing command batch completes.
#[cfg(test)]
pub(super) fn gdn_decode(
    weights: &GdnWeights<'_>,
    input: &Buffer,
    conv: Option<&Buffer>,
    delta: &Buffer,
) -> Result<GdnDecodeOutput, String> {
    for (buffer, size) in [
        (delta, 48 * 128 * 128 * 4),
        (weights.conv, 10240 * 4 * 2),
        (weights.a_log, 48 * 2),
        (weights.dt_bias, 48 * 2),
        (weights.norm, 128 * 2),
    ] {
        if unsafe { ds4_gpu_tensor_bytes(buffer.raw()) } < size {
            return Err("MTPLX GDN state/weight binding is too small".into());
        }
    }
    if conv.is_some_and(|conv| unsafe { ds4_gpu_tensor_bytes(conv.raw()) } < 3 * 10240 * 2) {
        return Err("MTPLX GDN conv state binding is too small".into());
    }
    let GdnProjected {
        qkv,
        z,
        b,
        a,
        mut temporary,
        ..
    } = gdn_project(&weights.input, input, 1)?;
    let gated = Buffer::mtplx_bytes(6144 * 2)?;
    let next_conv = Buffer::mtplx_bytes(3 * 10240 * 2)?;
    let next_delta = Buffer::mtplx_floats(48 * 128 * 128)?;
    // CustomKernel allocates outputs before materializing a broadcast input.
    // Keep the cold conv copy after the input projection, as in the reference.
    let initial_conv = if conv.is_none() {
        let zero = Buffer::mtplx_bytes(3 * 10240 * 2)?;
        let scalar = scalar_buffer(&0_u16.to_le_bytes())?;
        dispatch_geometry(
            "g1_copybfloat16bfloat16",
            &[
                tensor(0, &scalar).input(1),
                tensor(1, &zero).output(3 * 10240),
                bytes(3, &0_u64),
            ],
            &[],
            [3 * 10240, 1, 1],
            [1024, 1, 1],
            true,
        )?;
        Some(zero)
    } else {
        None
    };
    let conv = conv.or(initial_conv.as_ref()).unwrap();
    gdn::dispatch_step(&[
        tensor(0, &qkv).input(10240),
        tensor(1, &z).input(6144),
        tensor(2, &a).input(48),
        tensor(3, &b).input(48),
        tensor(4, conv).input(3 * 10240),
        tensor(5, weights.conv).input(4 * 10240),
        tensor(6, weights.a_log).input(48),
        tensor(7, weights.dt_bias).input(48),
        tensor(8, delta).input(48 * 128 * 128),
        tensor(9, weights.norm).input(128),
        tensor(10, &gated).output(6144),
        tensor(11, &next_conv).output(3 * 10240),
        tensor(12, &next_delta).output(48 * 128 * 128),
    ])?;
    if let Some(initial_conv) = initial_conv {
        encoder::add_temporaries(std::slice::from_ref(&initial_conv));
        temporary.push(initial_conv);
    }
    let hidden = Buffer::mtplx_bytes(2560 * 2)?;
    let [w, s, bias] = weights.output;
    temporary.extend(affine_bf16(
        [w, s, bias, &gated, &hidden],
        [1, 2560, 6144],
        weights.output_quant[0],
        weights.output_quant[1],
        true,
    )?);
    temporary.extend([qkv, z, b, a, gated]);
    Ok(GdnDecodeOutput {
        hidden,
        conv: next_conv,
        delta: next_delta,
        _temporaries: temporary,
    })
}