moe_permute_tokens

Function moe_permute_tokens 

Source
pub fn moe_permute_tokens(
    gpu: &dyn GpuBackend,
    kernel: KernelHandle,
    hidden_states: DevicePtr,
    permuted: DevicePtr,
    sorted_token_ids: DevicePtr,
    hidden: u32,
    total_expanded: u32,
    stream: u64,
) -> Result<()>
Expand description

Gather token rows into expert-sorted order: permuted[i] = hidden[sorted_token_ids[i]]. permuted is [total_expanded, hidden]. One block per output row, threads stride over hidden. Used by the FP4 grouped gate_up path (the CUTLASS escape-hatch needs contiguous per-expert rows; the FP8 fused kernel gathers internally so it doesn’t need this).

Retained for the legacy FP4 escape-hatch + potential reuse; the live FP4 path now uses the fused kernel (in-kernel gather), so this is currently uncalled.