Function dense_gemm_ba_gates_prefill 

pub fn dense_gemm_ba_gates_prefill(
    gpu: &dyn GpuBackend,
    kernel: KernelHandle,
    twin: KernelHandle,
    input: DevicePtr,
    ba_weight: &DenseWeight,
    a_log: DevicePtr,
    dt_bias: DevicePtr,
    gate_out: DevicePtr,
    m: u32,
    n: u32,
    k: u32,
    k_stride: u32,
    gate_stride: u32,
    nv: u32,
    vheads_per_group: u32,
    stream: u64,
) -> Result<()>

2026-09-25: BA GEMM and the GDN gate transforms for m tokens in one kernel, with no intermediate BA buffer.

input is [m, k_stride] BF16 activations and ba_weight [n, k] BF16, row-major. gate_out is one FP32 buffer, gate_stride floats per token: gate_out[token * gate_stride + vh] = gate (alpha→exp transform) gate_out[token * gate_stride + nv + vh] = beta (sigmoid)

Kernel: dense_gemm_ba_gates_prefill(A, B, A_log, dt_bias, gate_out, M, N, K, K_stride, gate_stride, nv, vpg) Grid: (ceil(N/4), M_tokens, 1) Block: (256, 1, 1)

twin is the one-CTA-per-token dense_gemm_ba_gates_prefill_hopper (module ssm_ba_gates_hopper), or KernelHandle(0) where the target lacks it. ba_gates_pick chooses between the two here, once, so the call sites do not branch. Both take the same arguments; the native_ssm_ba_gates_hopper_microtest example compares their outputs byte for byte.