Function dense_gemm_ba_gates_prefill 

pub fn dense_gemm_ba_gates_prefill(
    gpu: &dyn GpuBackend,
    kernel: 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<()>

Fused BA GEMM + GDN gates for prefill (token-parallel).

Replaces dense_gemm(normed, ba_weight) + compute_gdn_gates in the prefill path. Uses vectorized uint4 loads and warp-shuffle reduction per token, adding a token dimension via blockIdx.y. Skips the intermediate ba_out buffer entirely.

Output layout (shared gate_out buffer): 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)