Implementation:NVIDIA TransformerEngine Fp8Padding
| Field | Value |
|---|---|
| Sources | TransformerEngine |
| Domains | Deep_Learning, PyTorch, Quantization |
| Last Updated | 2026-02-07 14:00 GMT |
Overview
Applies row padding to input tensors to satisfy FP8/MXFP8/NVFP4 alignment requirements for Grouped GEMM operations.
Description
Fp8Padding is a torch.nn.Module that pads input tensor rows to meet alignment constraints required by quantized GEMM operations. The alignment size is automatically determined from the active FP8/FP4 recipe (32 for MXFP8/NVFP4, 16 for others) or can be specified manually. The implementation uses a custom autograd function (_Fp8Padding) that calls tex.fused_multi_row_padding for the forward pass and tex.fused_multi_row_unpadding for the backward pass, ensuring correct gradient flow through padded regions.
Usage
Use before Grouped GEMM operations when input tensor dimensions may not be aligned to the quantization block size. Typically paired with Fp8Unpadding after the GEMM.
Code Reference
Source Location
- Repository
NVIDIA/TransformerEngine- File
transformer_engine/pytorch/module/fp8_padding.py- Lines
- 1--144
Signature
class Fp8Padding(torch.nn.Module):
def __init__(self, num_gemms: int, align_size: Optional[int] = None) -> None: ...
def forward(self, inp: torch.Tensor, m_splits: List[int]) -> Tuple[torch.Tensor, List[int]]: ...
Import
from transformer_engine.pytorch.module.fp8_padding import Fp8Padding
I/O Contract
Inputs
| Name | Type | Required | Description |
|---|---|---|---|
| inp | torch.Tensor | Yes | Input tensor to pad |
| m_splits | List[int] | Yes | List of row counts per GEMM group (must match num_gemms)
|
Outputs
| Name | Type | Description |
|---|---|---|
| out | torch.Tensor | Padded tensor with aligned row dimensions |
| padded_m_splits | List[int] | Updated split sizes after padding |
Usage Examples
from transformer_engine.pytorch.module.fp8_padding import Fp8Padding
padding = Fp8Padding(num_gemms=4)
padded_input, padded_splits = padding(input_tensor, m_splits=[128, 65, 200, 33])
# padded_splits will be aligned to the recipe's block size