This document is relevant for: Trn2, Trn3

Blockwise MM Shard-on-Block Kernel API Reference#

Blockwise matrix multiplication kernel for context-encoding MoE layers.

This kernel implements blockwise matrix multiplication for mixture-of-experts (MoE) layers, processing tokens through expert-specific gate, up, and down projections. The computation combines static optimization benefits with dynamic early-exit capabilities by using a hybrid loop structure. Optimized for block-level sharding with PING_PONG strategy and supports FP8 quantization, multiple expert affinity scaling modes, and TopK > 1 accumulation patterns. Optimized for block sizes 128-512 tokens, 8-64 experts, and sequence lengths up to 32K tokens. Best performance when I_TP >= 512 and batch size * sequence length <= 4096.

Background#

The bwmm_shard_on_block kernel implements blockwise matrix multiplication for context-encoding MoE layers, processing tokens through expert-specific gate, up, and down projections with block-level sharding. A hybrid static/dynamic loop structure combines static-scheduling benefits with dynamic early-exit, and the bwmm_shard_on_block_hybrid entry point exposes the hybrid path directly.

API Reference#

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

bwmm_shard_on_block#

nkilib.experimental.moe.moe_cte.bwmm_shard_on_block(hidden_states: nl.ndarray, expert_affinities_masked: nl.ndarray, gate_up_proj_weight: nl.ndarray, down_proj_weight: nl.ndarray, block_size: int, token_position_to_id: nl.ndarray, block_to_expert: nl.ndarray, gate_and_up_proj_bias: Optional[nl.ndarray] = None, down_proj_bias: Optional[nl.ndarray] = None, gate_up_proj_scale: Optional[nl.ndarray] = None, down_proj_scale: Optional[nl.ndarray] = None, down_activations: Optional[nl.ndarray] = None, activation_function: common_types.ActFnType = common_types.ActFnType.SiLU, skip_dma: SkipMode = SkipMode(False, False), compute_dtype: Any = nl.bfloat16, is_tensor_update_accumulating: bool = True, expert_affinities_scaling_mode: common_types.ExpertAffinityScaleMode = common_types.ExpertAffinityScaleMode.POST_SCALE, n_block_per_iter: int = 1, 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, block_sharding_strategy: BlockShardStrategy = BlockShardStrategy.PING_PONG, sbm: Optional[SbufManager] = None, num_static_block: Optional[int] = None, total_n_blocks: Optional[int] = None, down_bias_tp_degree: Optional[int] = None, down_bias_tp_rank: Optional[int] = None, non_overlapping_shards: bool = False)#

Blockwise matrix multiplication kernel for context-encoding MoE layers.

Parameters:

  • hidden_states (nl.ndarray) – [T, H], Input token embeddings in HBM

  • expert_affinities_masked (nl.ndarray) – [(T+1)*E, 1], Expert routing weights for token assignments in HBM

  • gate_up_proj_weight (nl.ndarray) – [E, H, 2, I_TP], Combined gate and up projection weights in HBM

  • down_proj_weight (nl.ndarray) – [E, I_TP, H], Down projection weights in HBM

  • block_size (int) – Number of tokens processed per block

  • token_position_to_id (nl.ndarray) – [N*B], Mapping from block positions to token IDs in HBM

  • block_to_expert (nl.ndarray) – [N, 1], Expert assignment for each block in HBM

  • gate_and_up_proj_bias (Optional[nl.ndarray]) – [E, 2, I_TP], Bias terms for gate/up projections in HBM

  • down_proj_bias (Optional[nl.ndarray]) – [E, 1, H], Bias terms for down projection in HBM

  • gate_up_proj_scale (Optional[nl.ndarray]) – [E, 1, 2*I_TP], Dequantization scales for gate/up weights in HBM

  • down_proj_scale (Optional[nl.ndarray]) – [E, 1, H], Dequantization scales for down weights in HBM

  • down_activations (Optional[nl.ndarray]) – [N, B, H], Storage for intermediate activations in HBM

  • activation_function (common_types.ActFnType) – Activation function type (SiLU, GELU, etc.)

  • skip_dma (SkipMode) – DMA skip configuration for memory optimization

  • compute_dtype (Any) – Data type for internal computations (default: bfloat16)

  • is_tensor_update_accumulating (bool) – Enable accumulation for TopK > 1 scenarios

  • expert_affinities_scaling_mode (common_types.ExpertAffinityScaleMode) – Expert affinity application mode

  • n_block_per_iter (int) – Number of blocks processed per iteration

  • gate_clamp_upper_limit (Optional[float]) – Upper clamp limit for gate projections

  • gate_clamp_lower_limit (Optional[float]) – Lower clamp limit for gate projections

  • up_clamp_upper_limit (Optional[float]) – Upper clamp limit for up projections

  • up_clamp_lower_limit (Optional[float]) – Lower clamp limit for up projections

  • block_sharding_strategy (BlockShardStrategy) – Block distribution strategy across cores

  • sbm (Optional[SbufManager]) – SBUF memory manager. If None, one is created internally.

  • num_static_block (Optional[int]) – Number of blocks for static loop. Defaults to N.

  • total_n_blocks (Optional[int]) – Total block count for shard partitioning. Defaults to num_static_block.

  • down_bias_tp_degree (Optional[int]) – TP degree for down projection bias sharding.

  • down_bias_tp_rank (Optional[int]) – TP rank for down projection bias sharding.

  • non_overlapping_shards (bool) – When True, shards write to the same output slot (slot 0) and skip zero-init and cross-shard reduce. Requires non-overlapping token routing across shards (e.g., HI_LO strategy with sequence-level sharding). Default: False.

Returns:

Expert-processed token representations in HBM. Shape depends on accumulation mode: - Single expert (is_tensor_update_accumulating=False): [T, H] - Multiple experts (is_tensor_update_accumulating=True): [T, 2, H+E] for cross-core accumulation (the trailing E columns hold expert affinities; the hidden payload occupies columns 0:H)

Return type:

nl.ndarray

Notes:

  • Supports the PING_PONG and HI_LO block sharding strategies (selected via block_sharding_strategy; default PING_PONG)

  • Static loop processes N-E blocks with compile-time optimizations

  • Dynamic loop handles remaining blocks with early-exit capability

  • Supports FP8 quantization with dequantization scales

  • Expert affinity scaling modes: PRE_SCALE, POST_SCALE, PRE_SCALE_DELAYED

  • Multi-shard execution requires num_shards == 2 for accumulation

Dimensions:

  • T: Total number of input tokens

  • H: Hidden dimension size

  • B: Block size (tokens per block)

  • E: Number of experts

  • N: Total number of blocks (T / B)

  • I_TP: Intermediate size divided by the tensor-parallelism degree

bwmm_shard_on_block_hybrid#

nkilib.experimental.moe.moe_cte.bwmm_shard_on_block_hybrid(conditions: nl.ndarray, hidden_states: nl.ndarray, expert_affinities_masked: nl.ndarray, gate_up_proj_weight: nl.ndarray, down_proj_weight: nl.ndarray, block_size: int, token_position_to_id: nl.ndarray, block_to_expert: nl.ndarray, gate_and_up_proj_bias: Optional[nl.ndarray] = None, down_proj_bias: Optional[nl.ndarray] = None, gate_up_proj_scale: Optional[nl.ndarray] = None, down_proj_scale: Optional[nl.ndarray] = None, down_activations: Optional[nl.ndarray] = None, activation_function: common_types.ActFnType = common_types.ActFnType.SiLU, skip_dma: SkipMode = SkipMode(False, False), compute_dtype: Any = nl.bfloat16, is_tensor_update_accumulating: bool = True, expert_affinities_scaling_mode: common_types.ExpertAffinityScaleMode = common_types.ExpertAffinityScaleMode.POST_SCALE, n_block_per_iter: int = 1, 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, block_sharding_strategy: BlockShardStrategy = BlockShardStrategy.PING_PONG, down_bias_tp_degree: Optional[int] = None, down_bias_tp_rank: Optional[int] = None, non_overlapping_shards: bool = False)#

Hybrid static/dynamic shard-on-block kernel.

Parameters:

conditions (nl.ndarray) – [ceil(N/num_shards)+1] per-shard condition vector. 1=active, 0=padded. Last entry must be 0 for loop termination. All other args: same as bwmm_shard_on_block.

compute_same_weights_block_parallel_hbm#

nkilib.experimental.moe.moe_cte.compute_same_weights_block_parallel_hbm(N: int, block_to_expert: nl.ndarray, num_shards: int, shard_id: int, shard_strat: BlockShardStrategy, sbm: Optional[SbufManager] = None) → nl.ndarray#

Compute weight reuse mask for block-parallel execution.

Parameters:

  • N (int) – Total number of blocks

  • block_to_expert (nl.ndarray) – Expert assignment for each block

  • num_shards (int) – Number of shards for parallel execution

  • shard_id (int) – Current shard identifier

  • shard_strat (BlockShardStrategy) – Block distribution strategy

load_down_proj_weight#

nkilib.experimental.moe.moe_cte.load_down_proj_weight(down_proj_weight: nl.ndarray, block_expert: nl.ndarray, compute_dtype, skip_dma: SkipMode = SkipMode(), load_dst: Optional[list] = None, sbm: Optional[SbufManager] = None) → list#

Load down projection weights.

Parameters:

  • down_proj_weight (nl.ndarray) – Weight tensor with shape [E, I_TP, H]

  • block_expert (nl.ndarray) – Expert index tensor with shape (1, 1) in SBUF

  • compute_dtype – Compute data type

  • skip_dma (SkipMode) – DMA skip configuration

  • load_dst (Optional[list]) – Optional pre-allocated destination list

Notes:

  • Assumes I_TP is divisible by 16 for vector operations

  • Partial tiles are zero-padded

  • Uses scalar_offset for dynamic expert indexing

load_gate_up_proj_weights#

nkilib.experimental.moe.moe_cte.load_gate_up_proj_weights(gate_up_proj_weight: nl.ndarray, block_expert: nl.ndarray, compute_dtype, skip_dma: SkipMode = SkipMode(), load_dst: Optional[list] = None, sbm: Optional[SbufManager] = None) → list#

Load gate and up projection weights.

Parameters:

  • gate_up_proj_weight (nl.ndarray) – Weight tensor with shape [E, H, 2, I_TP]

  • block_expert (nl.ndarray) – Expert index tensor with shape (1, 1) in SBUF

  • compute_dtype – Compute data type

  • skip_dma (SkipMode) – DMA skip configuration

  • load_dst (Optional[list]) – Optional pre-allocated destination list

Notes:

  • Gate and up projections are interleaved in dimension 2

  • Partial tiles are zero-padded

  • Uses scalar_offset for dynamic expert indexing

compute_block_output#

nkilib.experimental.moe.moe_cte.compute_block_output(intermediate_states, dp_weights, expert_affinity, block_old, down_activations, block_idx, H, I_TP, NUM_TILES, output_dtype, is_tensor_update_accumulating, down_bias_broadcasted=None, down_bias_raw=None, down_scale=None, sbm: Optional[SbufManager] = None, down_proj_weight_hbm=None, block_expert=None, skip_dma: SkipMode = SkipMode(), block_new_lst_pre=None, i_tp_offset=0, down_bias_h_offset: int = 0, down_bias_h_size: Optional[int] = None)#

Compute block output with down projection and expert affinity scaling.

Parameters:

  • intermediate_states – Intermediate activation states [gup_tile_count][TILE_SIZE, B]

  • dp_weights – Down projection weights [gup_tile_count][TILE_SIZE, H]

  • expert_affinity – Expert affinities [NUM_TILES][TILE_SIZE, 1]

  • block_old – Previous block outputs for accumulation [NUM_TILES][TILE_SIZE, H]

  • down_activations – Storage for intermediate activations

  • block_idx – Current block index

  • H – Hidden dimension size

  • I_TP – Intermediate dimension size

  • NUM_TILES – Number of tiles per block

  • output_dtype – Output data type

  • is_tensor_update_accumulating – Enable accumulation mode

  • down_bias_broadcasted – Broadcasted bias [TILE_SIZE, H]

  • down_scale – Dequantization scales

Returns:

Block output tensors [NUM_TILES][TILE_SIZE, H]

Return type:

list

Notes:

  • Supports FP8 dequantization with down_scale

  • Accumulation mode for TopK > 1 scenarios

  • Optional bias addition before affinity scaling

reduce_outputs#

nkilib.experimental.moe.moe_cte.reduce_outputs(output: nl.ndarray, num_tiles: int, reduce_tile_size: int, offset: int, dim_hidden: int, sbm: Optional[SbufManager] = None)#

Synchronize across axis=0 in output by performing FMA reduce and store.

Parameters:

  • output (nl.ndarray) – Output tensor, size [T, 2, H+E] (hidden payload in columns 0:H; trailing E columns hold expert affinities)

  • num_tiles (int) – Number of tiles (iterations)

  • reduce_tile_size (int) – Size of tile size on partition dimension

  • offset (int) – Output read/write offset on row

  • dim_hidden (int) – Hidden dimension

  • sbm (Optional[SbufManager]) – Optional SBUF manager for allocation.

load_and_transpose_gup_bias#

nkilib.experimental.moe.moe_cte.load_and_transpose_gup_bias(inps: InputTensors, dims: DimensionSizes, cfg: Configs, block_expert, skip_dma, sbm: Optional[SbufManager] = None)#

Load and transpose gate/up projection bias for current expert.

Parameters:

  • inps (InputTensors) – Input tensor container

  • dims (DimensionSizes) – Dimension configuration

  • cfg (Configs) – Kernel configuration

  • block_expert – Expert index for current block [1, 1]

  • skip_dma – DMA skip configuration

  • sbm (Optional[SbufManager]) – Optional SBUF manager for allocation.

shard_strat2blk_idx#

nkilib.experimental.moe.moe_cte.shard_strat2blk_idx(shard_strat: BlockShardStrategy, outer_block_iter: int, inner_block_iter: int) → int#

Convert shard strategy indices to global block index.

Parameters:

  • shard_strat (BlockShardStrategy) – Sharding strategy (HI_LO or PING_PONG)

  • outer_block_iter (int) – Outer block iteration index

  • inner_block_iter (int) – Inner block iteration index (0 to BLOCK_PARALLEL_FACTOR-1)

Returns:

Global block index

Return type:

int

shard_strat2new_blk_idx_offset#

nkilib.experimental.moe.moe_cte.shard_strat2new_blk_idx_offset(shard_id: int, shard_strat: BlockShardStrategy, n_blocks_per_shard: int) → int#

Calculate block index offset based on shard ID and strategy.

Parameters:

  • shard_id (int) – Current shard identifier (0 or 1)

  • shard_strat (BlockShardStrategy) – Sharding strategy

  • n_blocks_per_shard (int) – Number of blocks per shard

Returns:

Block index offset for the current shard

Return type:

int

load_and_broadcast_down_bias#

nkilib.experimental.moe.moe_cte.load_and_broadcast_down_bias(inps: InputTensors, dims: DimensionSizes, cfg: Configs, block_expert, skip_dma, sbm: Optional[SbufManager] = None, use_pe_broadcast: bool = False)#

Load and broadcast down projection bias for the current block.

Parameters:

  • inps (InputTensors) – Input tensor container

  • dims (DimensionSizes) – Dimension configuration

  • cfg (Configs) – Kernel configuration

  • block_expert – Expert index for current block

  • skip_dma – DMA skip configuration

  • sbm (Optional[SbufManager]) – Optional SBUF manager for allocation.

  • use_pe_broadcast (bool) – Use PE matmul broadcast instead of DVE StreamShuffle.

load_down_bias_raw#

nkilib.experimental.moe.moe_cte.load_down_bias_raw(inps, dims, cfg, block_expert, skip_dma, sbm=None, bias_h_size=None)#

Load raw (1, bias_h_size) down bias without broadcasting. bias_h_size defaults to H.

bwmm_output_initialization#

nkilib.experimental.moe.moe_cte.bwmm_output_initialization(output, shard_id=None, sbm: Optional[SbufManager] = None, expert_affinities_masked=None, E=0, H=0, skip_zero_init=False)#

Zero initialize buffer at output and optionally copy expert affinities.

Parameters:

  • output – External memory, shape (T, H) or (T, 2, H+E).

  • shard_id – Optionally provide shard ID.

  • sbm (Optional[SbufManager]) – Optional SBUF manager for allocation.

  • expert_affinities_masked – Expert affinities [(T+1)*E, 1], or None.

  • E – Number of experts.

  • H – Hidden dimension (excluding affinity columns).

  • skip_zero_init – Skip zero initialization of output[:, shard_id, :H]. Used with non_overlapping_shards where zero-init is unnecessary.

bwmm_load_old_block#

nkilib.experimental.moe.moe_cte.bwmm_load_old_block(output, token_indices, NUM_TILES, dtype, skip_dma: SkipMode = SkipMode(), shard_id=None, token_indices_offset=0, sbm: Optional[SbufManager] = None)#

Loads the partially computed output hidden states for the current block’s token indices.

This document is relevant for: Trn2, Trn3