This document is relevant for: Trn2, Trn3

Matmul MXFP8 Backward Kernel API Reference#

Backward pass for matrix multiplication with MXFP8 quantization.

Computes both input gradients (dX) and weight gradients (dW) for a linear layer. Forward pass convention: Y = X @ W^T, where X is [M, K], W is [N, K], Y is [M, N] Backward pass (two separate matmuls with different dimensions): dX = dY @ W (shape [M, K]): M_logical=M, K_contraction=N, N_logical=K dW = dY^T @ X (shape [N, K]): M_logical=N, K_contraction=M, N_logical=K

Background#

The matmul_mxfp8_backward kernel computes the backward pass for matrix multiplication with MXFP8 quantization.

API Reference#

Source code for this kernel API can be found at: matmul_mxfp8_generic_backward_kernel.py

matmul_mxfp8_backward#

nkilib.experimental.matmul_mxfp8.matmul_mxfp8_backward(output_grad, weights, input_activation, input_grad_config: MatmulMxfp8KernelConfig = None, weight_grad_config: MatmulMxfp8KernelConfig = None, tile_loop_order: str = 'mnk', float8_dtype: str = 'float8_e5m2', output_dtype=nl.bfloat16, run_with_lnc2: bool = True, lnc_2_shard_rhs: bool = True, output_grad_scales=None, weight_scales=None, input_scales=None, use_scale_packing: bool = False, spill_reload: bool = False, output_grad_is_swizzled: bool = False, weights_is_swizzled: bool = False, input_is_swizzled: bool = False) → tuple#

Backward pass for matrix multiplication with MXFP8 quantization.

Parameters:

  • output_grad – Output gradient (dY), shape [M, N] in BF16.

  • weights – Weight matrix (W), shape [N, K] in BF16.

  • input_activation – Input activation (X), shape [M, K] in BF16.

  • input_grad_config (MatmulMxfp8KernelConfig) – MatmulMxfp8KernelConfig for the dX phase (auto-resolved if None).

  • weight_grad_config (MatmulMxfp8KernelConfig) – MatmulMxfp8KernelConfig for the dW phase (auto-resolved if None).

  • tile_loop_order (str) – Tile processing order within blocks, default ‘mnk’.

  • float8_dtype (str) – FP8 dtype for quantization, default “float8_e5m2”.

  • output_dtype – Output data type, default nl.bfloat16.

  • run_with_lnc2 (bool) – Enable LNC2 parallelization, default True.

  • lnc_2_shard_rhs (bool) – Shard on N dimension (RHS), default True.

  • output_grad_scales – Optional pre-computed scales for output gradient.

  • weight_scales – Optional pre-computed scales for weights.

  • input_scales – Optional pre-computed scales for input activation.

  • use_scale_packing (bool) – Assert packed scales for pre-quantized inputs.

  • spill_reload (bool) – Spill quantized blocks to HBM for reuse.

  • output_grad_is_swizzled (bool) – Whether output gradient is pre-swizzled.

  • weights_is_swizzled (bool) – Whether weights are pre-swizzled.

  • input_is_swizzled (bool) – Whether input activation is pre-swizzled.

Returns:

(input_grad, weight_grad) where: - input_grad: Shape [M, K], gradient with respect to input - weight_grad: Shape [N, K], gradient with respect to weights

Return type:

nl.ndarray

This document is relevant for: Trn2, Trn3