Skip to main content

grouped_gemm_f32

Function grouped_gemm_f32 

Source
pub fn grouped_gemm_f32(
    input: &[f32],
    weights: &[f32],
    m_sizes: &[i32],
    gather_indices: &[i32],
    total_tokens: usize,
    columns: usize,
    reduction: usize,
    input_row_stride: usize,
    weight_expert_stride: usize,
    weight_row_stride: usize,
    output_row_stride: usize,
    permute_input: bool,
    permute_output: bool,
    top_k: usize,
) -> Vec<f32>
Expand description

Grouped expert GEMM with optional input/output permutation metadata.

m_sizes gives the token count for each expert. For each expert group, rows are multiplied by that expert’s [columns, reduction] weight tile. When permute_input or permute_output is set, gather_indices maps sorted expert rows back to the original token/top-k routing order.