dense_gemm_ba_gates_prefill_hopper

Function dense_gemm_ba_gates_prefill_hopper 

Source
pub fn dense_gemm_ba_gates_prefill_hopper(
    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<()>
Expand description

Launch dense_gemm_ba_gates_prefill_hopper — one CTA per token.

Same ABI as the parent, deliberately: the two differ in grid and in nothing else a caller can see, which is what lets ba_gates_pick return a handle rather than a closure.