Implementation:NVIDIA TransformerEngine Custom GEMM
| Field | Value |
|---|---|
| Sources | TransformerEngine |
| Domains | Deep_Learning, PyTorch, Quantization |
| Last Updated | 2026-02-07 14:00 GMT |
Overview
Dispatches GEMM operations to custom quantizer implementations, enabling user-defined quantized matrix multiplication logic for custom quantization recipes.
Description
The custom_gemm function provides a dispatch layer for custom quantization recipes. It routes GEMM calls to the qgemm method on the quantizer attached to the input QuantizedTensorStorage objects. It determines the GEMM type (FPROP, DGRAD, WGRAD) based on the grad flag and layout string, constructs MMParams with output dtype and split accumulator settings, and extracts the appropriate data/scale/transpose combinations for each GEMM type. The function supports 3D input reshaping for FPROP operations.
Usage
Used internally when a custom quantization recipe is active. The custom recipe's quantizer must implement the qgemm method.
Code Reference
Source Location
- Repository
NVIDIA/TransformerEngine- File
transformer_engine/pytorch/custom_recipes/gemm.py- Lines
- 1--137
Signature
def custom_gemm(
A: QuantizedTensorStorage,
B: QuantizedTensorStorage,
workspace: torch.Tensor,
out_dtype: Optional[torch.dtype] = None,
quantization_params: Optional[Quantizer] = None,
gelu: bool = False,
gelu_in: torch.Tensor = None,
accumulate: bool = False,
layout: str = "TN",
out: Optional[torch.Tensor] = None,
bias: Optional[torch.Tensor] = None,
use_split_accumulator: bool = False,
grad: bool = False,
) -> Iterable[Optional[torch.Tensor]]: ...
Import
from transformer_engine.pytorch.custom_recipes.gemm import custom_gemm
I/O Contract
Inputs
| Name | Type | Required | Description |
|---|---|---|---|
| A | QuantizedTensorStorage | Yes | First quantized tensor (must be custom) |
| B | QuantizedTensorStorage | Yes | Second quantized tensor (must be custom) |
| workspace | torch.Tensor | Yes | GEMM workspace buffer |
| layout | str | No | GEMM layout: "TN", "NN", or "NT"
|
| grad | bool | No | Whether this is a gradient GEMM |
| bias | Optional[torch.Tensor] | No | Optional bias for FPROP |
Outputs
| Name | Type | Description |
|---|---|---|
| result | torch.Tensor | GEMM result |
| (3 Nones) | None | Padding to match general_gemm return format
|
Usage Examples
from transformer_engine.pytorch.custom_recipes.gemm import custom_gemm
# Internal usage with custom quantized tensors:
result, _, _, _ = custom_gemm(
A=quantized_activation,
B=quantized_weight,
workspace=workspace,
out_dtype=torch.bfloat16,
layout="TN",
grad=False,
)