This document is relevant for: Trn2, Trn3

nki.isa.nc_matmul_mx#

nki.isa.nc_matmul_mx(dst: NkiTensor, stationary: NkiTensor, moving: NkiTensor, stationary_scale: NkiTensor, moving_scale: NkiTensor, tile_position=None, tile_size=None, accumulate=None, name=None)[source]#

Compute matrix multiplication of MXFP8/MXFP4 quantized matrices with integrated dequantization using Tensor Engine.

Note

Available only on NeuronCore-v4 and newer.

The NeuronCore-v4 Tensor Engine supports matrix multiplication of MXFP8/MXFP4 quantized matrices as defined in the OCP Microscaling standard. This instruction performs matrix multiplication between quantized stationary and moving matrices while applying dequantization scales during computation. The micro-scaling group size is 32 elements in groups of 8 partitions × 4 elements per partition of both stationary and moving tensors. See Trainium3 arch guide for more detailed discussion.

Tiling Mode.

NeuronCore Tensor Engine is built upon a systolic array with 128 rows and 128 columns of processing elements (PEs). For nc_matmul_mx, Tensor Engine supports only row tiling mode, which allows multiple nc_matmul_mx instructions with a stationary partition dimension size smaller than 128 to run in parallel to improve hardware utilization. Row tiling mode slices the 128 PE rows into 2x 64 row tiles or 4x 32 row tiles.

tile_size is a tuple (row_size, column_size) that sets the row tile size. row_size must be 128, 64, or 32, and column_size must be 128. The stationary partition dimension size must not exceed row_size.

A given nc_matmul_mx can pick the exact row tile within the 128x128 systolic array by specifying the starting row in tile_position as a tuple (start_row, start_column), where start_column must be 0. The start_row must be a multiple of row_size, and start_row + row_size must not exceed 128. Place stationary, moving, and their scales at SBUF partition start_row.

Note

Known issue: NKI currently does not use row_size to set the row tile size. When tiling mode is enabled, NKI derives the row tile size from the stationary partition dimension size. Set row_size to that value so it matches the tile size actually used. A row tile smaller than 128 also currently requires a stationary free dimension size greater than 64. NKI also does not yet check that row_size is 128, 64, or 32, or that start_row + row_size does not exceed 128. Whether such a configuration compiles depends on the stationary partition dimension size, which sets the actual row tile size.

For example, setting tile_position to (64, 0) and tile_size to (64, 128) means using the bottom half of the systolic array.

Note, tile_position and tile_size must both be set to enable tiling mode. If they are not set, the default is to use the full systolic array, which is equivalent to tile_position=(0, 0) and tile_size=(128, 128). The values in tile_position and tile_size tuples can be integers or affine expressions.

Memory types.

The nc_matmul_mx instruction must read inputs from SBUF and write outputs to PSUM. Therefore, the stationary, moving, stationary_scale, and moving_scale must be SBUF tiles, and dst tile must be a PSUM tile.

Data types.

The input stationary and moving tiles must be float8_e5m2_x4, float8_e4m3fn_x4, or float4_e2m1fn_x4 (4-packed quantized data types). The stationary_scale and moving_scale tiles must be float8_e8m0fnu or uint8 (prefer float8_e8m0fnu: OCP MX standard; uint8 accepted for backward compatibility). The dst tile can be float32 or bfloat16.

Layout.

The contraction dimension of the matrix multiplication is along the partition dimension of stationary and moving tensors and also the x4 dimension within each packed data type element (float8_e5m2_x4, float8_e4m3fn_x4, or float4_e2m1fn_x4).

The free dimension of the stationary tile matches the partition dimension of the output dst tile in size, while the free dimension of the moving tile matches the free dimension of the dst tile in size.

The scale tensors follow a special layout requirement. See more details in nisa.quantize_mx API doc.

Tile size

  • The partition dimension size of stationary and moving must be identical and be either 32, 64, or 128.

  • The free dimension size of stationary must be even and not exceed 128.

  • The free dimension size of moving must not exceed 4096 when dst is in float32 or 8192 when dst is in bfloat16, matching the size of 8x PSUM banks (the entire PSUM).

  • The scale tensors have partition dimensions that depend on whether the data tensors span multiple quadrants. See more details in nisa.quantize_mx API doc.

  • dst can have at most three free dimensions (maximum rank 4). Each of stationary, moving, stationary_scale, and moving_scale must be contiguous in memory (multiple free dimensions are allowed if they collapse to a single contiguous run). These ranks include the partition dimension.

Profiler view syntax.

nc_matmul_mx uses the same profiler output format as nisa.nc_matmul, except the source access pattern is interpreted as an MX-quantized tensor: src=<dtype>@$MX[<data_addr>,<scale_addr>,<start_scale_partition>]@[<step_elem>][<num_elem>].

Tensor indirection.

dst supports tensor indirection (gather/scatter) independently by passing a view created with .indirect(index). The (stationary, stationary_scale) and (moving, moving_scale) pairs must agree on tensor indirection: either both members of a pair use .indirect() with the same index, or neither does — because in MX format the data bytes and scale exponents form a single logical tensor. Tensor indirection is incompatible with hardware tiling, so tile_position must be (0, 0) (or unset) when it is used. (nc_matmul_mx itself is only available on NeuronCore-v4 and later.)

When operands are manually allocated, their base partitions must satisfy:

  • the index of every .indirect() view starts on a quadrant boundary (a multiple of 32);

  • if stationary or moving uses .indirect(), the operand starts on the same partition as its index;

  • if dst uses .indirect(), dst’s index starts on the same partition as stationary.

Parameters:

  • dst – the matrix multiplication output (PSUM tile)

  • stationary – the stationary quantized matrix (SBUF tile)

  • moving – the moving quantized matrix (SBUF tile)

  • stationary_scale – the dequantization scales for stationary matrix (SBUF tile)

  • moving_scale – the dequantization scales for moving matrix (SBUF tile)

  • tile_position – (optional) a 2D tuple (start_row, start_column) to control the starting row in Tensor Engine tiling mode; start_column must be 0

  • tile_size – (optional) a 2D tuple (row_size, column_size) setting the row tile size in Tensor Engine tiling mode; column_size must be 128.

  • accumulate – (optional) if True, accumulate the matmul result into the existing dst PSUM tile content; if False, overwrite the existing content; if None (default), auto-detect based on whether this PSUM location was previously written

This document is relevant for: Trn2, Trn3