This document is relevant for: Trn2, Trn3

nki.isa.nc_matmul#

nki.isa.nc_matmul(dst: NkiTensor, stationary: NkiTensor, moving: NkiTensor, is_stationary_onezero=False, is_moving_onezero=False, is_transpose=False, accumulate=None, tile_position=(), tile_size=(), perf_mode=matmul_perf_mode.none, name=None)[source]#

Compute dst = stationary.T @ moving matrix multiplication using Tensor Engine.

The figure below illustrates how to map a matrix multiplication from a mathematical definition to nisa.nc_matmul on Tensor Engine. The stationary tensor is loaded into the systolic array first and stays in place, while the moving tensor streams through the array during computation. For more detailed discussion of Tensor Engine capabilities, see Trainium arch guide.

../../../_images/matmul.png — Fig. 86 MxKxN Matrix Multiplication Visualization.#

Performance mode.

Tensor Engine supports FP8 double performance mode, enabled by setting performance mode to double_row. See Trainium2 arch guide for more details. double_row performance mode cannot be combined with Tensor Engine column tiling mode (details below).

Tiling mode. NeuronCore Tensor Engine is built upon a systolic array with 128 rows and 128 columns of processing elements (PEs). Tensor Engine supports both row and column tiling modes, which allow multiple nc_matmul instructions with a stationary tile size smaller than [128, 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. Column tiling mode slices the 128 PE columns in the same fashion. tile_size is a tuple (row_size, column_size) that sets the tile size. Each of row_size and column_size must be 128, 64, or 32. The stationary partition dimension size must not exceed row_size, and its free dimension size must not exceed column_size.

In addition, a given nc_matmul can also pick the exact row and column tile within the 128x128 systolic array, by specifying the starting row and starting column in tile_position as a tuple (start_row, start_column). Each of start_row and start_column must be a multiple of the corresponding tile_size value, and start plus that value must not exceed 128. Place stationary and moving at SBUF partition start_row, and dst at PSUM partition start_column.

Note

Known issue: NKI currently does not use tile_size to set the tile size. When tiling mode is enabled, NKI derives the tile size from the stationary partition and free dimension sizes, each rounded up to 32, 64, or 128. Set tile_size to those rounded values so it matches the tile size actually used. NKI also does not yet check that each tile_size value is 128, 64, or 32, or that start plus that value does not exceed 128. Whether such a configuration compiles depends on the stationary dimension sizes, which set the actual 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.

Accumulation mode.

The accumulate parameter controls whether the matmul result overwrites or accumulates onto the dst PSUM tile:

accumulate=False: Overwrites the PSUM destination. The first matmul targeting a PSUM location must use this overwrite mode to start the accumulation.
accumulate=True: Accumulates the result to the existing PSUM content. Only valid after the PSUM location has been initialized by a matmul with accumulate=False. Using accumulate=True on an uninitialized PSUM location is undefined behavior on all targets.
accumulate=None (default): The compiler automatically infers the correct flag — first write overwrites, subsequent writes accumulate.

Hint: Use accumulate=(i > 0) in a loop to explicitly overwrite on the first iteration and accumulate on subsequent iterations, or simply omit the parameter to let the compiler handle it.

Note

On NeuronCore-v3, matmul accumulating into a PSUM location that was value-initialized by a non-matmul instruction (e.g., memset, tensor_copy) is not supported and is undefined behavior on hardware.

Transpose mode.

Tensor Engine can transpose a tile in SBUF by loading it as a stationary tile and using an identity matrix as the moving tile. Turning on transpose mode by setting is_transpose=True enables bit-accurate data transpose, which can transpose tensors with NaN/Inf values properly. See Trainium2 arch guide for more details.

Memory types.

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

Data types.

The input stationary and moving tiles can be one of these supported data types: float8_e4m3/float8_e5m2/bfloat16/float16/tfloat32/float32. Starting NeuronCore-v4, float8_e4m3fn (OCP FP8) is also supported; float8_e4m3 (legacy) and float8_e4m3fn cannot be mixed in the same matmul. The stationary and moving tiles can have different data types, with one exception: if one of the input tiles is tfloat32/float32, the other tile must also be tfloat32/float32. When performance mode is double_row, stationary and moving tiles must be one of float8_e4m3 or float8_e5m2 (plus float8_e4m3fn on NeuronCore-v4), but the two input tiles can have different float8 formats (except that legacy float8_e4m3 and OCP float8_e4m3fn cannot appear together).

The accumulation precision internal to Tensor Engine is float32. The dst tile must be a float32 tile in NeuronCore-v3. Starting NeuronCore-v4, dst can either be a float32 or bfloat16 tile.

Layout.

If performance mode is off, the contraction dimension of the matmul must be along the partition dimension in both stationary and moving tiles.

If performance mode is double_row, the contraction dimension of the matmul is split between the partition dimension and the first free dimension after the partition dimension in both stationary and moving tiles. The first free dimension must be 2. For example, to perform a matmul of [1, 256]@[256, 3]=[1, 3], the stationary tile is of shape [128, 2, 1], while the moving tile is of shape [128, 2, 3].

Regardless of performance mode, 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.

Tile size.

The partition dimension sizes of the stationary and moving tiles must be identical. They must not exceed 128 when tiling mode is off or row_size specified in tile_size when tiling mode is on. The free dimension size of stationary must not exceed 128 when tiling mode is off or column_size in tile_size when tiling mode is on.

On NeuronCore-v3, the free dimension size of moving tile must not exceed 512, matching the maximum number of float32 elements per PSUM bank. Starting NeuronCore-v4, the free dimension size of moving tile can go up to 4096 for float32 dst or 8192 for bfloat16 dst, matching the size of 8x PSUM banks (the entire PSUM).

The dst, stationary, and moving tiles can have at most three free dimensions (maximum rank 4, including the partition dimension).

Explicit tiling is required when the high-level matmul operation exceeds the tile size limits of nc_matmul.

Profiler view syntax.

Each nc_matmul call lowers to two ISA instructions in the profiler: a load instruction (to load the stationary operand into the Tensor Engine) followed by a multiply instruction. Both instructions will appear in profiler output for a single nc_matmul call.

The multiply instruction operands are displayed in a compact ISA syntax:

src=<dtype>@<address>[<strides>][<num_elem>]
dst=<dtype>@<address>[<strides>][<num_elem>]
<M>*<K> acc_flags=<flags> psum_zero=<val>

Where:

<dtype>: data type (e.g., bfloat16, fp8e4, fp8e5)
<address>: hex memory address in SBUF (for src) or PSUM (for dst)
[<strides>]: element strides per dimension (multi-dimensional)
[<num_elem>]: number of elements per dimension (multi-dimensional)
<M>*<K>: matmul dimensions (M rows × K contraction)
acc_flags: accumulator control flags (e.g., 2 = reset accumulator)
psum_zero: PSUM zero-initialization control value

Tensor indirection.

On NeuronCore-v4 and later, dst, stationary, and moving support tensor indirection (gather/scatter) by passing a view created with .indirect(index). When is_transpose=True, moving cannot use tensor indirection. Tensor indirection is incompatible with hardware tiling, so tile_position must be (0, 0) (or unset) when it is used. The indirect operands must be a Tensor-Engine float dtype (float16, bfloat16, float8_e4m3, float8_e5m2, float8_e4m3fn).

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.

Dynamic access patterns.

stationary, moving and dst may each carry a dynamic access pattern built with nki.language.NkiTensor.ap(), whose stride and count slots hold a nki.isa.VirtualRegister read at runtime.

The partition dimension (dim 0) is the exception, and both its stride and its count must be Python int values, so put the dynamic slots on a free dimension. stationary’s free dimension count must be static as well, because that count becomes the partition dimension of dst. A dynamic access pattern is also unsupported with is_transpose=True and with perf_mode=nki.isa.matmul_perf_mode.double_row.

Parameters:

dst – the matmul output
stationary – the stationary operand
moving – the moving operand
is_stationary_onezero – hints to the compiler whether the stationary operand is a tile with ones/zeros only; setting this field explicitly could lead to 2x better performance if stationary tile is in float32; the field has no impact for non-float32 stationary
is_moving_onezero – hints to the compiler whether the moving operand is a tile with ones/zeros only; setting this field explicitly could lead to 2x better performance if moving tile is in float32; the field has no impact for non-float32 moving
is_transpose – controls Tensor Engine transpose mode on/off
accumulate – 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. Not exposed for nc_transpose.
tile_position – a 2D tuple (start_row, start_column) to control starting row and column in Tensor Engine tiling mode
tile_size – a 2D tuple (row_size, column_size) setting the row and column tile sizes in Tensor Engine tiling mode
perf_mode – controls Tensor Engine FP8 double performance mode on/off: matmul_perf_mode.none (default) disables double FP8 mode; matmul_perf_mode.double_row enables double FP8 mode which achieves 2x matmul throughput by packing two FP8 weight/ifmap element pairs and computing two multiplications in parallel per cycle; cannot be combined with column tiling mode. See the Trainium2 arch guide for more information.

Examples:

@nki.jit
def nc_matmul_accumulate_kernel(lhsT, rhs):
    '''Tiled matmul with accumulate=(i > 0).

    Tiles K=512 into 4 tiles of 128, accumulating partial products in PSUM.
    The first iteration overwrites PSUM (accumulate=False via i > 0 == False),
    subsequent iterations accumulate (accumulate=True via i > 0 == True).
    '''
    K, M = lhsT.shape
    K_, N = rhs.shape
    TILE_K = 128

    result = nl.ndarray((M, N), dtype=nl.float32, buffer=nl.shared_hbm)
    res_psum = nl.ndarray((M, N), dtype=nl.float32, buffer=nl.psum)

    for i in nl.static_range(K // TILE_K):
        lhs_sb = nl.ndarray((TILE_K, M), dtype=lhsT.dtype, buffer=nl.sbuf)
        rhs_sb = nl.ndarray((TILE_K, N), dtype=rhs.dtype, buffer=nl.sbuf)
        nisa.dma_copy(dst=lhs_sb, src=lhsT[i * TILE_K : (i + 1) * TILE_K, :])
        nisa.dma_copy(dst=rhs_sb, src=rhs[i * TILE_K : (i + 1) * TILE_K, :])
        nisa.nc_matmul(
            dst=res_psum, stationary=lhs_sb, moving=rhs_sb, accumulate=(i > 0)
        )

    res_sb = nl.ndarray((M, N), dtype=nl.float32, buffer=nl.sbuf)
    nisa.tensor_copy(dst=res_sb, src=res_psum)
    nisa.dma_copy(dst=result, src=res_sb)
    return result

This document is relevant for: Trn2, Trn3

nki.isa.nc_matmul

Contents

nki.isa.nc_matmul#