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.NkiTensor, gamma: nl.NkiTensor, router_weights: nl.NkiTensor, router_bias: Optional[nl.NkiTensor] = 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, store_eager_affi_only: bool = False)#

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

Parameters:

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

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

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

  • router_bias (Optional[nl.NkiTensor]) – [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.

  • store_eager_affi_only (bool) – If True, skip building/storing the sparse [T, E] expert affinities (no [T, E] HBM allocation, no indirect-DMA scatter) and instead return the dense [T, K] eager affinities. Requires T <= 128. Downstream consumers must read affinities in the dense [T, K] form (co-indexed with expert_index) rather than the sparse [T, E] form.

Returns:

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

Returns:

[T, K] int32 top-K indices.

Returns:

when store_eager_affi_only=False, the sparse [T, E] bfloat16 masked top-K affinities (zero elsewhere); when store_eager_affi_only=True, the dense [T, K] bf16 eager affinities co-indexed with expert_index.

Return type:

nl.NkiTensor

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