mla_batched_gemv

Function mla_batched_gemv 

Source
pub fn mla_batched_gemv(
    gpu: &dyn GpuBackend,
    kernel: KernelHandle,
    input: DevicePtr,
    weight: DevicePtr,
    output: DevicePtr,
    n_out: u32,
    k: u32,
    num_heads: u32,
    input_stride: u32,
    output_stride: u32,
    stream: u64,
) -> Result<()>
Expand description

2026-09-25: MLA batched GEMV over all heads in one launch: output[head, n] = sum_k(weight[head, n, k] * input[head, k]), BF16, with the per-head input and output at input_stride / output_stride and the weight contiguous per head.