This document is relevant for: Trn2, Trn3

MLP Kernel API Reference#

Implements Multi-Layer Perceptron with optional normalization fusion and quantization support.

The kernel supports:

  • Both context encoding (CTE) and token generation (TKG) modes

  • Optional normalization fusion (RMSNorm, LayerNorm)

  • Various activation functions

  • Residual connections via fused addition

  • Flexible tensor layouts and column tiling optimizations

  • Bias addition for all projections and normalization

  • FP8 quantization (static and row-wise, TKG mode only)

  • MXFP4/MXFP8 quantization (TKG mode)

  • Gate and up projection result clamping

  • Optional gate projection skipping

  • SBUF output for kernel fusion

Background#

The MLP kernel is a critical component in transformer architectures, responsible for processing token representations after the attention mechanism. This kernel optimizes the MLP computation by fusing it with optional normalization and supporting various optimizations for both context encoding and token generation scenarios.

Note

This kernel automatically selects between TKG (Token Generation) and CTE (Context Encoding) implementations based on the batch size × sequence length, ensuring optimal performance across different use cases.

API Reference#

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

mlp#

nkilib.core.mlp.mlp(hidden_tensor: nl.NkiTensor, gate_proj_weights_tensor: nl.NkiTensor, up_proj_weights_tensor: nl.NkiTensor, down_proj_weights_tensor: nl.NkiTensor, normalization_weights_tensor: Optional[nl.NkiTensor] = None, gate_proj_bias_tensor: Optional[nl.NkiTensor] = None, up_proj_bias_tensor: Optional[nl.NkiTensor] = None, down_proj_bias_tensor: Optional[nl.NkiTensor] = None, normalization_bias_tensor: Optional[nl.NkiTensor] = None, fused_add_tensor: Optional[nl.NkiTensor] = None, store_fused_add_result: bool = False, activation_fn: ActFnType = ActFnType.SiLU, normalization_type: NormType = NormType.NO_NORM, quantization_type: QuantizationType = QuantizationType.NONE, gate_w_scale: Optional[nl.NkiTensor] = None, up_w_scale: Optional[nl.NkiTensor] = None, down_w_scale: Optional[nl.NkiTensor] = None, gate_up_in_scale: Optional[nl.NkiTensor] = None, down_in_scale: Optional[nl.NkiTensor] = None, quant_clipping_bound: float = 0.0, output_dtype=None, store_output_in_sbuf: bool = False, eps: float = 1e-06, skip_gate_proj: bool = False, use_tkg_gate_up_proj_column_tiling: bool = True, use_tkg_down_proj_column_tiling: bool = True, use_tkg_down_proj_optimized_layout: bool = False, use_contiguous_x4_gate_up: bool = False, gate_clamp_upper_limit: Optional[float] = None, gate_clamp_lower_limit: Optional[float] = None, up_clamp_upper_limit: Optional[float] = None, up_clamp_lower_limit: Optional[float] = None, force_cte_mode: bool = False, mode: ComputationMode = ComputationMode.AUTO, sbm: Optional[BufferManager] = None, mx_dummy_scale_hbm: Optional[nl.NkiTensor] = None, transposed_in: bool = False, transposed_out: bool = False, dtype_mode: DtypeMode = DtypeMode.NON_OCP, gate_up_w_layout: MLPGateUpWeightLayout = MLPGateUpWeightLayout.CONTIGUOUS, use_folded_mx_scales: bool = False) → list[nl.NkiTensor]#

MLP (Multi-Layer Perceptron) Kernel implementation.

Performs the standard MLP computation with support for both context encoding (CTE) and token generation (TKG) modes. Automatically selects the appropriate implementation based on input dimensions and supports various optimizations.

Parameters:

  • hidden_tensor (nl.NkiTensor) – Input hidden states tensor with shape [B, S, H], SBUF layout, or transposed HBM layout [H0, n_prgs, H1_shard, BxS] when transposed_in=True.

  • gate_proj_weights_tensor (nl.NkiTensor) – Gate projection weight matrix with shape [H, I].

  • up_proj_weights_tensor (nl.NkiTensor) – Up projection weight matrix with shape [H, I].

  • down_proj_weights_tensor (nl.NkiTensor) – Down projection weight matrix with shape [I, H].

  • normalization_weights_tensor (Optional[nl.NkiTensor]) – Normalization weights with shape [1, H].

  • gate_proj_bias_tensor (Optional[nl.NkiTensor]) – Bias tensor for gate projection with shape [1, I].

  • up_proj_bias_tensor (Optional[nl.NkiTensor]) – Bias tensor for up projection with shape [1, I].

  • down_proj_bias_tensor (Optional[nl.NkiTensor]) – Bias tensor for down projection with shape [1, H].

  • normalization_bias_tensor (Optional[nl.NkiTensor]) – Bias tensor for normalization with shape [1, H]. Only applicable for layer normalization.

  • fused_add_tensor (Optional[nl.NkiTensor]) – tensor to fuse for the residual connection..

  • store_fused_add_result (bool) – If True, stores the fused_add output to HBM, and the kernel returns both the fused_add output and the MLP output. (default: False)

  • activation_fn (ActFnType) – Activation function type.

  • normalization_type (NormType) – Type of normalization.

  • quantization_type (QuantizationType) – Quantization type to use (default: QuantizationType.NONE). Supported values: - QuantizationType.NONE: No quantization - QuantizationType.STATIC: FP8 tensor-wise quantization with 2x perf mode (CTE and TKG) - QuantizationType.STATIC_MX: FP8 tensor-wise quantization with 4x perf mode (CTE) - If using STATIC_MX in CTE mode, the down weights should be swizzled as follows: I, H = w.shape w.reshape((ceil(I / 512), 128, 4, H)).transpose(1, 0, 3, 2).to_x4_dtype() - It expects up and gate weights to be swizzled as well. See gate_up_w_layout for details. - QuantizationType.ROW: FP8 row-wise quantization (CTE and TKG) - QuantizationType.MX: MXFP quantization (MXFP4/MXFP8, TKG only)

  • gate_w_scale (Optional[nl.NkiTensor]) – Dequantization scales for gate weights. - FP8: Shape [128, I] for row-wise, [128, 1] for tensor-wise quantization - MXFP (TKG only): Scale factors for MXFP quantized weights

  • up_w_scale (Optional[nl.NkiTensor]) – Dequantization scales for up weights. - FP8: Shape [128, I] for row-wise, [128, 1] for tensor-wise quantization - MXFP (TKG only): Scale factors for MXFP quantized weights

  • down_w_scale (Optional[nl.NkiTensor]) – Dequantization scales for down weights. - FP8: Shape [128, I] for row-wise, [128, 1] for tensor-wise quantization - MXFP (TKG only): Scale factors for MXFP quantized weights

  • gate_up_in_scale (Optional[nl.NkiTensor]) – FP8 dequantization scales for gate and up input. Used for tensor-wise quantization with shape [128, 1]. Defaults to None.

  • down_in_scale (Optional[nl.NkiTensor]) – FP8 dequantization scales for down input. Used for tensor-wise quantization with shape [128, 1]. Defaults to None.

  • quant_clipping_bound (float) – Clipping boundary for context encoding FP8 row quantization (default: 0.0)

  • output_dtype – Output tensor data type. Defaults to None; if None, the hidden tensor’s dtype is used.

  • store_output_in_sbuf (bool) – If True, stores the output in SBUF instead of HBM, allowing the next layer to read it directly without an additional load operation. This option is only available in TKG mode where output tensor is small enough to fit in SBUF. (default: False)

  • eps (float) – Epsilon value for numerical stability.

  • skip_gate_proj (bool) – Skip gate projection

  • use_tkg_gate_up_proj_column_tiling (bool) – If True, uses column tiling for the gate and up projection in TKG mode. (default: True)

  • use_tkg_down_proj_column_tiling (bool) – If True, uses column tiling for the down projection in TKG mode. (default: True)

  • use_tkg_down_proj_optimized_layout (bool) – If True, the standard down_weight tensor (shape [I, H]) is reinterpreted as [I, lnc, 128, H // (128 * lnc)], then transposed to [I, lnc, H // (128 * lnc), 128]. This layout provides unit-stride weight loading, reducing the matrix multiplication initiation interval. Only applied when use_tkg_down_proj_column_tiling is False. (default: False)

  • gate_clamp_upper_limit (Optional[float]) – upper bound value to clamp on gate projection results, does not perform clamping if the value is set to None

  • gate_clamp_lower_limit (Optional[float]) – lower bound value to clamp on gate projection results, does not perform clamping if the value is set to None

  • up_clamp_upper_limit (Optional[float]) – upper bound value to clamp on up projection results, does not perform clamping if the value is set to None

  • up_clamp_lower_limit (Optional[float]) – lower bound value to clamp on up projection results, does not perform clamping if the value is set to None

  • force_cte_mode (bool) – If True, forces the use of CTE mode. (default: False)

  • mode (ComputationMode) – Computation mode to use (default = AUTO) Supported values: - AUTO (default): automatically choose the more performant computation mode. Does not support fused normalization and quantization - PREFILL: use prefill computation mode, does not support fused quantization, but support quantized input. - DECDOE: use decode computation mode, not performant for large T, supported fused normalization and quantization.

  • sbm (Optional[BufferManager]) – Optional BufferManager for HBM tensor allocation with consistent naming.

  • mx_dummy_scale_hbm (Optional[nl.NkiTensor]) – Pre-filled HBM tensor of all 127 values used as dummy MX scale tensors for MX matmul operations. When provided, scales are DMA’d from HBM instead of memset, avoiding per-layer memset overhead. Default: None.

  • transposed_in (bool) – When True, input is in transposed HBM layout [H0, n_prgs, H1_shard, BxS] instead of [B, S, H]. Enables contiguous per-NC DMA loads. Only supported in TKG mode with NONE, STATIC, or ROW quantization. Not compatible with fused add or SBUF input. (default: False)

  • transposed_out (bool) – When True, output is in transposed HBM layout [H0, n_prgs, H1_shard, BxS] instead of [B, S, H]. Each NC DMAs its shard directly without transpose_store. Only supported in TKG mode. Not compatible with down_proj column tiling or SBUF output. (default: False)

  • dtype_mode (DtypeMode) – Explicit FP8 E4M3 dtype selection for TKG STATIC/ROW quantization weight tiles. - DtypeMode.NON_OCP (default): nl.float8_e4m3 (max=240). - DtypeMode.OCP: nl.float8_e4m3fn (max=448). TRN3 only. - DtypeMode.AUTO: nl.float8_e4m3fn on TRN3, nl.float8_e4m3 elsewhere. (default: DtypeMode.NON_OCP)

  • gate_up_w_layout (MLPGateUpWeightLayout) – The layout of gate_proj_weights_tensor and up_proj_weights_tensor (default: CONTIGUOUS) This parameter is required if quantization_type is either MX, ROW_MX, or STATIC_MX. Supported values: - CONTIGUOUS: Use if quantization_type is not one of MX, ROW_MX, or STATIC_MX - H_X4_INNERMOST: One of two alternative MX weight layouts. More optimal if hidden_tensor is already quantized. - H_X4_MIDDLE: One of two alternative MX weight layouts. More optimal if hidden_tensor is not quantized.

  • use_folded_mx_scales (bool) – CTE MX only. When True, gate_w_scale, up_w_scale, and down_w_scale are supplied pre-folded into the physical SBUF layout the kernel consumes, allowing the scales to be loaded with a single large DMA per operand instead of many small quadrant-fold copies: - gate/up: uint8 [128, ceil((H/512)/4), I]. Four H/512 tiles are folded into one 128-partition buffer; tile k occupies partitions (q*32 + (k%4)*4 : +4) for each quadrant q, at buffer index k//4. - down: uint8 [128, I/512, H]. The 16 scale rows of each I/512 tile are scattered to partitions q*32 : q*32+4 (q in 0..3); the remaining partitions are zero. Producers must emit all three scales folded. (default: False)

Returns:

The MLP output tensor(s): - HBM output: Tensor with shape [B, S, H]. - HBM transposed output (transposed_out=True): Tensor with shape [H0, n_prgs, H1_shard, BxS]. - SBUF output: Shape depends on the mode setting. - CTE : Not applicable - TKG when use_tkg_down_proj_column_tiling is True = [BxS, H] - TKG when use_tkg_down_proj_column_tiling is False = [128(p_max), H/128, BxS] - If store_fused_add_result is True, returns a list containing both the output and the stored fused output.

Return type:

list[nl.NkiTensor]

Notes:

  • Automatically dispatches to either CTE or TKG implementation based on batch size and sequence length.

  • Token generation mode (TKG) is used for small batch/sequence dimensions (batch_size × sequence_length ≤ 96), while context encoding (CTE) handles larger inputs.

  • Column tiling and tensor layout optimization (use_tkg_down_proj_optimized_layout) are valid only in TKG mode.

  • FP8 quantization support is available only in TKG mode.

  • Supported input data types: nl.bfloat16, nl.float16, nl.float32

Implementation Details#

The kernel implementation includes several key optimizations:

  1. Dual Implementation Strategy: Automatically selects between CTE (Context Encoding) and TKG (Token Generation) implementations based on batch_size × sequence_length.

  2. Normalization Fusion: Optionally fuses RMSNorm or LayerNorm operations with the MLP computation for improved performance.

  3. FP8 Quantization: Supports FP8 quantization with both static and row-wise dequantization scales. Available only in TKG mode for weights and activations.

  4. Flexible Tensor Layouts: Supports column tiling optimizations and tensor layout optimizations in TKG mode to improve memory access patterns.

  5. Activation Function Options: Supports multiple activation functions, including SiLU (Swish), GELU, and ReLU.

  6. Result Clamping: Provides optional clamping of gate and up projection results with configurable upper and lower bounds.

  7. Gate Projection Skipping: Allows skipping the gate projection computation when skip_gate_proj is enabled.

  8. Residual Connection Fusion: Can incorporate residual connections through fused_add_tensor for improved performance.

  9. SBUF Output Option: Provides the option to keep output in SBUF for fusion with subsequent operations (TKG mode only).

  10. Bias Addition: Supports optional bias addition for gate, up, and down projections, as well as for normalization.

  11. Optimized Weight Loading: In TKG mode, use_tkg_down_proj_optimized_layout enables unit-stride weight loading to reduce matrix multiplication initiation interval.

  12. Multi-Precision Support: Supports bfloat16, float16, and float32 input data types for flexible precision requirements.

See Also#

This document is relevant for: Trn2, Trn3