flashinfer.quantization.expand_csf_scales

flashinfer.quantization.expand_csf_scales(csf: CsfScales, out: Tensor, token_selected_experts: Tensor | None = None, layout: SfLayout = SfLayout.layout_linear, swap_row_halves: bool = False, local_expert_offset: int = 0) → Tensor

Expand CSF block scales into a scale buffer, one matrix per expert.

Nothing is read back to the host, so the expansion can run under CUDA graph capture, with the routing changing between replays.

Parameters:

  • csf (CsfScales) – Scales from encode_csf_scales(), one matrix per expert. Expansion takes about 192 bytes of shared memory per column, so cols can be at most 864 on SM80, 1200 on SM90 and SM100, and 518 on SM86, SM89 and SM120.

  • out (torch.Tensor) –

    Contiguous destination in a one-byte dtype on the scales’ device, written in place:

    • layout_128x4: [num_experts, round_up(rows, 128), round_up(cols, 4)] in the 128x4 swizzled layout of the fused MoE kernels, the padding written as zeros.

    • layout_linear: [num_experts, rows, cols].

  • token_selected_experts (Optional[torch.Tensor]) – Contiguous int32 expert ids of any shape, such as the [num_tokens, top_k] routing of a fused MoE call. Only these experts are expanded and the rest of out keeps its bytes; ids outside [local_expert_offset, local_expert_offset + num_experts) are skipped. None expands every expert. An expert named more than once is expanded once: in each 128-row tile, every id is compared with all the ids before it, so N ids take O(N^2) comparisons per tile. None costs the same for any routing. Pass ids only for small routings; with 512 experts on SM80, for example, 640 ids expanded faster than None and 1280 slower.

  • layout (SfLayout) – SfLayout.layout_128x4 or SfLayout.layout_linear.

  • swap_row_halves (bool) – Swap the two halves of each matrix’s rows, e.g. to turn [w1; w3] checkpoint order into the [w3; w1] order of the fused MoE kernels. Needs rows % 32 == 0.

  • local_expert_offset (int) – The id in token_selected_experts of csf’s first expert. Under expert parallelism the routing holds global ids while csf holds this rank’s experts, so pass the global id of this rank’s first expert.

Returns:

out.

Return type:

torch.Tensor