swizzle.h#

Functions to convert scaling factors into format expected by GEMM.

Functions

void nvte_swizzle_scaling_factors(
const NVTETensor input,
NVTETensor output,
cudaStream_t stream,
)#

Swizzling scaling factors into the required interleaved layout for GEMM.

Requirements:

  • scale_inv is stored in row-major.

  • scale_inv size is padded to 128x4 for row-scale and 4x128 for col-scale.

  • Input and output scale_inv buffer pointers used by this operation must be 16-byte aligned.

  • data is quantized along K-dimension, i.e. 1D-scaling block lies along the K-dimension.

Parameters:

  • input – [in] Input tensor with non-swizzled scale_inv.

  • output – [inout] Output tensor which hosts swizzled scale_inv.

  • stream – [in] CUDA stream used for the operation.

void nvte_multi_tensor_swizzle_scaling_factors(
const NVTETensor *inputs,
NVTETensor *outputs,
const size_t num_tensors,
cudaStream_t stream,
)#

Swizzling scaling factors into the required interleaved layout for GEMM.

Requirements:

  • scale_inv is stored in row-major.

  • scale_inv size is padded to 128x4 for row-scale and 4x128 for col-scale.

  • Input and output scale_inv buffer pointers used by this operation must be 16-byte aligned.

  • data is quantized along K-dimension, i.e. 1D-scaling block lies along the K-dimension.

Parameters:

  • inputs – [in] Input tensors with non-swizzled scale_inv.

  • outputs – [inout] Output tensors which hosts swizzled scale_inv.

  • num_tensors – [in] Number of input and output tensors.

  • stream – [in] CUDA stream used for the operation.

void nvte_multi_tensor_swizzle_scaling_factors_unchecked(
const NVTETensor *inputs,
NVTETensor *outputs,
const size_t num_tensors,
cudaStream_t stream,
)#

Same as nvte_multi_tensor_swizzle_scaling_factors, but skips scale_inv shape/padding validation.

Use this variant when the data and scale_inv tensors intentionally have different shapes, e.g. when scale_invs have been transposed for attention. The 16-byte alignment requirement for scale_inv buffer pointers used by this operation still applies.

void nvte_unswizzle_scaling_factors(
const NVTETensor input,
NVTETensor output,
cudaStream_t stream,
)#

Unswizzling scaling factors from the interleaved layout used by GEMM back to row-major.

Requirements:

  • scale_inv is stored in row-major in output.

  • scale_inv size is padded to 128x4 for row-scale and 4x128 for col-scale.

  • data is quantized along K-dimension, i.e. 1D-scaling block lies along the K-dimension.

Parameters:

  • input – [in] Input tensor with swizzled scale_inv.

  • output – [inout] Output tensor which hosts non-swizzled scale_inv.

  • stream – [in] CUDA stream used for the operation.

void nvte_multi_tensor_unswizzle_scaling_factors(
const NVTETensor *inputs,
NVTETensor *outputs,
const size_t num_tensors,
cudaStream_t stream,
)#

Unswizzling scaling factors from the interleaved layout used by GEMM back to row-major.

Requirements:

  • scale_inv is stored in row-major in output.

  • scale_inv size is padded to 128x4 for row-scale and 4x128 for col-scale.

  • data is quantized along K-dimension, i.e. 1D-scaling block lies along the K-dimension.

Parameters:

  • inputs – [in] Input tensors with swizzled scale_inv.

  • outputs – [inout] Output tensors which hosts non-swizzled scale_inv.

  • num_tensors – [in] Number of input and output tensors.

  • stream – [in] CUDA stream used for the operation.

void nvte_swizzle_block_scaling_to_mxfp8_scaling_factors(
const NVTETensor input,
NVTETensor output,
cudaStream_t stream,
)#

Swizzling FP8 block scaling scaling factors into mxfp8 interleaved layout for GEMM.

This function is used for emulating the FP8 block scaling recipe on Blackwell and newer as it not natively supported by cublasLt on architectures other than Hopper.

Requirements:

  • input is an FP8 block scaling tensor

  • input has rowwise usage

  • output is an MXFP8 tensor

  • output has rowwise usage

  • output.scale_inv has appropriate shape

Parameters:

  • input – [in] Input FP8 block-scaled tensor.

  • output – [inout] Output mxfp8 tensor which hosts swizzled scale_inv.

  • stream – [in] CUDA stream used for the operation.

void nvte_multi_tensor_swizzle_block_scaling_to_mxfp8_scaling_factors(
const NVTETensor *inputs,
NVTETensor *outputs,
const size_t num_tensors,
cudaStream_t stream,
)#

Swizzling FP8 block scaling scaling factors of multiple tensors into mxfp8 interleaved layout for GEMM.

Multi-tensor variant of nvte_swizzle_block_scaling_to_mxfp8_scaling_factors, with the same requirements for each pair of tensors.

Parameters:

  • inputs – [in] Input FP8 block-scaled tensors.

  • outputs – [inout] Output mxfp8 tensors which host the swizzled scale_inv.

  • num_tensors – [in] Number of input and output tensors.

  • stream – [in] CUDA stream used for the operation.

void nvte_swizzle_grouped_block_scaling_to_mxfp8_scaling_factors(
const NVTEGroupedTensor input,
NVTEGroupedTensor output,
cudaStream_t stream,
)#

Swizzle the FP8 block-scaling scaling factors of a grouped tensor into the MXFP8 interleaved layout for GEMM (grouped variant of nvte_swizzle_block_scaling_to_mxfp8_scaling_factors).

Used to emulate the FP8 block scaling recipe with MXFP8 grouped GEMM on Blackwell and newer. Per-tensor dimensions are read on the device, so the operation is CUDA-graph safe when they change between replays.

Requirements:

  • input is an FP8 block scaling (1D or 2D) grouped tensor with rowwise data and FP32 rowwise scale_inv in the compact per-tensor layout; per-tensor first dims or last dims (int64) may vary, but not both

  • output is an MXFP8 grouped tensor with the same rowwise data pointer, dims and logical shape, with_gemm_swizzled_scales set, and E8M0 rowwise scale_inv. Tensor t with rowwise data [f, l] is written at the cumulative offset of roundup(f, 128) * ceil(l / 128) * 4 bytes.

  • output scale_inv holds at least nvte_get_grouped_block_scaling_to_mxfp8_scale_inv_size(input) bytes

  • 24 * n + 8 bytes (per-tensor offset tables) fit in the device’s shared memory per block, where n is the number of tensors

Parameters:

  • input – [in] Input FP8 block-scaled grouped tensor.

  • output – [inout] Output MXFP8 grouped tensor which hosts the swizzled scale_inv.

  • stream – [in] CUDA stream used for the operation.

size_t nvte_get_grouped_block_scaling_to_mxfp8_scale_inv_size(
const NVTEGroupedTensor input,
)#

Size in bytes of the output scale_inv required by nvte_swizzle_grouped_block_scaling_to_mxfp8_scaling_factors.

Depends only on the logical shape and on which per-tensor dims vary, so the output can be allocated without reading per-tensor dims from the device.

Parameters:

input – [in] Input FP8 block-scaled grouped tensor.

void nvte_swizzle_grouped_scaling_factors(
const NVTEGroupedTensor input,
NVTEGroupedTensor output,
cudaStream_t stream,
)#

Swizzling scaling factors into the required interleaved layout for GEMM (grouped tensor).

Requirements(for now, more features will be added later):

  • scaling mode must be MXFP8 1D scaling.

  • scale_inv is stored in row-major per group.

  • scale_inv size is padded to 128x4 for row-scale and 4x128 for col-scale.

  • data is quantitized along K-dimension, i.e. 1D-scaling block lies along the K-dimension.

  • all tensors in the grouped tensor must have the same shape.

Parameters:

  • input – [in] Input grouped tensor with non-swizzled scale_inv.

  • output – [inout] Output grouped tensor which hosts swizzled scale_inv.

  • stream – [in] CUDA stream used for the operation.

void nvte_unswizzle_grouped_scaling_factors(
const NVTEGroupedTensor input,
NVTEGroupedTensor output,
cudaStream_t stream,
)#

Unswizzling scaling factors from the interleaved GEMM layout back to row-major (grouped).

Requirements:

  • scaling mode must be MXFP8 1D scaling.

  • scale_inv is stored in row-major in output.

  • scale_inv size is padded to 128x4 for row-scale and 4x128 for col-scale.

  • data is quantized along K-dimension, i.e. 1D-scaling block lies along the K-dimension.

  • all tensors in the grouped tensor must have the same shape.

Parameters:

  • input – [in] Input grouped tensor with swizzled scale_inv.

  • output – [inout] Output grouped tensor which hosts non-swizzled scale_inv.

  • stream – [in] CUDA stream used for the operation.