Function moe_sort_by_expert 

pub fn moe_sort_by_expert(
    gpu: &dyn GpuBackend,
    kernel: KernelHandle,
    topk_ids: DevicePtr,
    sorted_token_ids: DevicePtr,
    sorted_expert_ids: DevicePtr,
    expert_offsets: DevicePtr,
    token_to_perm: DevicePtr,
    total_expanded: u32,
    num_experts: u32,
    topk: u32,
    stream: u64,
) -> Result<()>

2026-09-25: Counting sort of the total_expanded (token, slot) routes by expert, in one block: writes sorted_token_ids and sorted_expert_ids in expert order, expert_offsets (the [num_experts + 1] prefix sum) and token_to_perm (each route’s sorted position).