flashinfer.quantization.block_scale_interleave

flashinfer.quantization.block_scale_interleave(unswizzled_sf: Tensor, *, out: Tensor | None = None) → Tensor

Swizzle a block-scale tensor for FP4 layouts.

Reorders an unswizzled FP4 block-scale tensor to optimize memory access patterns for FP4 GEMM/MoE kernels. The output is padded in the m dimension to a multiple of 128.

Parameters:

  • unswizzled_sf (torch.Tensor) – Input scale-factor tensor with dtype uint8 or bfloat16.

  • out (Optional[torch.Tensor]) – Preallocated contiguous, flat swizzled output on the input device, with the same dtype and padded size. Reusing it preserves its address across CUDA Graph captures and replay.

Returns:

1D flattened swizzled scale-factor buffer of shape (num_experts * expert_out_size,) where num_experts is unswizzled_sf.shape[0] for 3D inputs (and 1 otherwise) and expert_out_size is the padded swizzled size returned by _compute_swizzled_layout_sf_size. Note that this is not the same logical shape as unswizzled_sf; downstream FP4 GEMM/MoE kernels consume the flat buffer directly.

Return type:

torch.Tensor

Raises:

  • AssertionError – If the input dtype is not uint8 or bfloat16.

  • ValueError – If out has an incompatible device, dtype, size, or layout.