This document is relevant for: Trn2, Trn3

RMSNorm Router Top-K TKG Kernel API Reference#

Fused RMSNorm (+ optional MX quantize) + Router TopK.

Background#

The rmsnorm_router_topk_tkg kernel fuses RMSNorm, optional MX (fp8) quantization, and router top-K expert selection for MoE token generation. It returns the (optionally MX-packed) normalized hidden states, top-K expert indices, and masked top-K affinities. Requires LNC=2 sharding.

API Reference#

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

rmsnorm_router_topk_tkg#

nkilib.experimental.moe_block.rmsnorm_router_topk_tkg(hidden_states: nl.ndarray, gamma: nl.ndarray, router_weights: nl.ndarray, router_bias: Optional[nl.ndarray] = None, eps: float = 1e-06, top_k: int = 1, hidden_actual: Optional[int] = None, quantization_type: QuantizationType = QuantizationType.NONE, router_mm_dtype=nl.bfloat16, router_act_fn: RouterActFnType = RouterActFnType.SIGMOID)#

Fused RMSNorm (+ optional MX quantize) + Router TopK.

Parameters:

  • hidden_states (nl.ndarray) – [B, S, H], Input tensor on HBM.

  • gamma (nl.ndarray) – [1, H], RMSNorm weights on HBM.

  • router_weights (nl.ndarray) – [H, E], Router weights on HBM.

  • router_bias (Optional[nl.ndarray]) – [1, E], Optional router bias on HBM.

  • eps (float) – Epsilon for RMSNorm. Default 1e-6.

  • top_k (int) – Number of top experts per token. Default 1.

  • hidden_actual (Optional[int]) – Actual hidden dim for padded inputs.

  • quantization_type (QuantizationType) – NONE or MX. Default NONE.

  • router_mm_dtype – Dtype for router matmul. Default nl.bfloat16.

  • router_act_fn (RouterActFnType) – SOFTMAX or SIGMOID. Default SIGMOID.

Returns:

[T, H] (NONE) or [T, H + H/4] FP8 packed quant‖scales (MX).

Return type:

nl.ndarray

Returns:

[T, K] int32 top-K indices.

Return type:

nl.ndarray

Returns:

[T, E] bfloat16 masked top-K affinities (zero elsewhere).

Return type:

nl.ndarray

Notes:

  • Requires LNC=2 sharding.

  • NONE: H must be divisible by 128; T must be a multiple of 256 (DLoC tiling).

  • MX: H must be divisible by 512 (MX block size).

This document is relevant for: Trn2, Trn3