This document is relevant for: Inf2, Trn1, Trn2

nki.isa.activation#

nki.isa.activation(op, data, *, bias=None, scale=1.0, mask=None, dtype=None, **kwargs)[source]#

Apply an activation function on every element of the input tile using Scalar Engine. The activation function is specified in the op input field (see Supported Activation Functions for NKI ISA for a list of supported activation functions).

The activation instruction can optionally multiply the input data by a scalar or vector scale and then add another vector bias before the activation function is applied, at no additional performance cost:

\[output = f_{act}(data * scale + bias)\]

When the scale is a scalar, it must be a compile-time constant. In this case, the scale is broadcasted to all the elements in the input data tile. When the scale/bias is a vector, it must have the same partition axis size as the input data tile and only one element per partition. In this case, the element of scale/bias within each partition is broadcasted to elements of the input data tile in the same partition.

Note, the Scalar Engine always performs the math operations in float32 precision. Therefore, the engine automatically casts the input data tile to float32 before performing multiply/add/activate specified in the activation instruction. The engine is also capable of casting the float32 math results into another output data type specified by the dtype field at no additional performance cost. If dtype field is not specified, Neuron Compiler will set output data type of the instruction to be the same as input data type of data. On the other hand, the scale parameter must have a float32 data type, while the bias parameter can be float32/float16/bfloat16.

The input data tile can be an SBUF or PSUM tile. Similarly, the instruction can write the output tile into either SBUF or PSUM, which is specified using the buffer field. If not specified, nki.language.sbuf is selected by default.

Estimated instruction cost:

max(MIN_II, N) Scalar Engine cycles, where

  • N is the number of elements per partition in data.

  • MIN_II is the minimum instruction initiation interval for small input tiles. MIN_II is roughly 64 engine cycles.

Parameters:

  • op – an activation function (see Supported Activation Functions for NKI ISA for supported functions)

  • data – the input tile; layout: (partition axis <= 128, free axis)

  • bias – a vector with the same partition axis size as data for broadcast add (after broadcast multiply with scale)

  • scale – a scalar or a vector with the same partition axis size as data for broadcast multiply

  • dtype – (optional) data type to cast the output type to (see Supported Data Types for more information); if not specified, it will default to be the same as the data type of the input tile.

  • mask – (optional) a compile-time constant predicate that controls whether/how this instruction is executed (see NKI API Masking for details)

Returns:

output tile of the activation instruction; layout: same as input data tile

Example:

import neuronxcc.nki.language as nl
import neuronxcc.nki.isa as nisa

##################################################################
# Example 1: perform exponential function on matrix a of shape (128, 1024)
##################################################################
a = nl.load(a_tensor)
activated_a = nisa.activation(op=nl.exp, data=a)
nl.store(a_act_tensor, activated_a)

##################################################################
# Example 2: perform the following operations to matrix b of shape (128, 512)
# using a single activation instruction: np.square(b * 2.0) + c
# 1) compute `np.square(b * 2.0 + c)`
# 2) cast 1) results into bfloat16
##################################################################
b = nl.load(b_tensor)
c = nl.load(c_tensor)
activated_b = nisa.activation(op=np.square, data=b, bias=c, scale=2.0,
                              dtype=nl.bfloat16)
nl.store(b_act_tensor, activated_b)

This document is relevant for: Inf2, Trn1, Trn2