Mixture of Experts

Mixture of Experts (MoE) scales model capacity by routing each token to only K of E total experts (where K and E are model parameters defined below) rather than all of them. This sparse activation allows many parameters while keeping inference cost manageable. This example shows how to implement MoE on TCP hardware, focusing on two key challenges: replacing control-flow-based TopK routing with branchless matrix operations, and executing sparse expert computations blockwise.

Background: Basic FFN

Readers familiar with FFNs and MoE can skip to MoE Implementation on TCP.

A basic FFN consists of two linear projections (up and down) that together map T × D → T × F → T × D. The following describes this simplified FFN (up/down projection only, without gate projection) as the baseline that MoE replaces:

• Input
  • x_ffn_norm: T x D
• Weights
  • W_up: D x F (up projection)
  • W_down: F x D (down projection)
• Output
  • ffn_out: T x D
• Operations
  • Up projection:
    • up = einsum(x_ffn_norm, W_up)
    • (T x D), (D x F) -> T x F
  • Down projection:
    • ffn_out = einsum(up, W_down)
    • (T x F), (F x D) -> T x D

MoE Structure

MoE replaces a single FFN with E independent FFNs called experts. Each expert has its own weights:

• W_up[0], W_up[1], ..., W_up[E-1]
• W_down[0], W_down[1], ..., W_down[E-1]

Computing all experts would increase computation by E times. To avoid this, MoE uses a router to select only the Top-K most suitable experts per token, enabling sparse computation.

Model Parameters

The following arguments define an MoE layer:

• T: number of tokens
  • prefill: T = B * S_in
  • decode: T = B
• D: hidden dimension
• F: intermediate dimension of ffn up projection result
• E: number of total experts (typically 128)
• K: number of experts to apply ffn
  • llama4: 1, gpt-oss: 4, qwen3: 8

MoE Processing Steps

MoE processing consists of four steps: gating (computing per-expert scores), Top-K selection (choosing the K best experts per token), sparse expert computation (applying the selected experts), and combining (merging expert outputs with routing weights).

1. Gating (Router)

The router calculates a score for each expert for every token, determining which experts should process each token:

• Input
  • x_norm: T x D
• Weight
  • W_router: D x E (Gating network weights)
• Output
  • scores: T x E
• Operation
  • scores = einsum(x_norm, W_router)
  • (T x D), (D x E) -> T x E
  • Calculate the score (Logit) for E Experts per token

2. Top-K Selection

This step selects the Top-K Experts based on the router scores from Step 1 and calculates the weight for each selected Expert:

• Input
  • scores: T x E
• Outputs
  • topk_indices: T x K (selected Expert ID per token)
  • routing_weights: T x K (weight of selected Expert per token)
• Operations
  • Top-K Selection:
    • raw_weights, topk_indices = topk(scores, K)
    • Extract the K Expert indices and scores with the highest scores per token
  • Softmax Normalization:
    • routing_weights = softmax(raw_weights)
    • Convert the selected K scores to probability values (sum is 1 per token)
    • softmax(x)[i] = exp(x[i]) / sum(exp(x[j]) for j in 0..K)

The output for each token t consists of:

• topk_indices[t, :]: K Expert IDs (0 <= e < E)
• routing_weights[t, :]: weights of those Experts (sum is 1)

3. Sparse Expert Computation

Only selected Experts perform computation, making this stage sparse. A total of T * K Expert calls occur, but each Expert only computes for the tokens that selected it.

For each token t in [0, T-1] and selected Expert k in [0, K-1]:

• Selected Expert ID: e = topk_indices[t, k]
• Input
  • x_norm[t]: D (input of token t)
• Weights (weights of Expert e)
  • W_up[e]: D x F
  • W_down[e]: F x D
• Output
  • y[t, k]: D (k-th Expert output of token t)
• Operations
  • Up projection:
    • up = einsum(x_norm[t], W_up[e])
    • D, (D x F) -> F
  • Down projection:
    • y[t, k] = einsum(up, W_down[e])
    • F, (F x D) -> D

The results for all (t, k) pairs are collected into y_experts: T x K x D.

4. Weighted Sum (Combine)

The final step combines the K Expert outputs using the routing weights calculated earlier:

• Inputs
  • y_experts: T x K x D
  • routing_weights: T x K (weight of each Expert)
• Output
  • ffn_out: T x D
• Operations
  • ffn_out = einsum(y_experts, routing_weights)

The result is that each token receives the weighted average output of its selected K Experts.

MoE Implementation on TCP

The TCP implementation of MoE uses two hardware-specific techniques: branchless TopK via the Vector Engine’s bit manipulation and filter operations, and blockwise execution via static-shape scatter/gather patterns.

1. Overview and Design Philosophy

1.1. Bridging Logical and Physical Execution

The two fundamental challenges when implementing MoE on TCP are:

• Challenge 1: Conflict between control flow and parallel structure
  • Problem: General Top-K algorithms use branch statements where the execution path varies depending on data values. Such branch statements cause performance degradation in SIMT-based accelerators that process thousands of elements with a single instruction.
  • Solution: Completely removing control flow and using Branchless Top-K technique with matrix operations and bit manipulation is essential.
• Challenge 2: Gap between logical Routing and physical execution
  • Problem: Logically, MoE is a process where each token finds the Expert that suits it (Token-centric). However, if implemented as is, memory access becomes irregular and the number of tokens to process per Expert changes dynamically, reducing TCP compiler efficiency.
  • Solution: The perspective must be shifted to a method where the Expert becomes the subject and collects tokens (Expert-centric).

1.2. Core Techniques for TCP Implementation

Two core techniques address these challenges:

• Branchless TopK: Performs routing via matrix operations only, eliminating all control flow
• Blockwise execution: Processes only selected Experts with data packed in fixed-size Block units

The following sections describe each technique in detail.

2. Branchless TopK

Branchless TopK replaces control-flow-based sorting with pure matrix operations. This approach consists of three stages: bit packing to combine score and index, parallel ranking to determine order, and filtering to extract the top K results.

2.1. Bit Packing (Combining Score and Index)

Bit packing bundles score and index into a single value so the Expert ID is preserved when scores are reordered during sorting. The TCP Vector Engine executes a fixed instruction sequence across all 256 slices simultaneously, so any operation whose address or control path depends on runtime data values must be rewritten as a fixed sequence of matrix operations:

• Inputs
  • scores: T x E
  • Index_expert: E
    • Index_expert(e) = e where e = 0, 1, 2, ..., E - 1
• Output
  • Packed_Value: T x E
    • Tensor with (score, index) packed.
  • Packed_Value_cmp: T x E
    • Tensor with (score, index) packed, preprocessed to enable comparison of score magnitude using integer comparison.
• Operation
  • Packing
    • Place Expert Score (e.g., bf16) in the upper bits and Expert Index (e.g., int16) in the lower bits to create a single 32-bit integer (or float).
    • Packed_Value_unprocessed = (Score << 16) | Index
    • Processed in Vector Engine.
  • Comparison Trick
    • This preprocessing enables magnitude comparison of score values using simple integer comparison.
    • Bit Flipping preprocessing solves the problem of negative magnitude relationships being reversed when comparing float values as integer values. This enables accurate Top-K selection with only integer comparators.

    • Packed_Value_cmp = if Packed_Value >= 0 {
          Packed_Value
      } else {
          Packed_Value ^ 0x7fff0000
      }

2.2. Parallel Ranking (All-to-All Comparison)

Parallel ranking determines the order of all experts simultaneously instead of sequential sorting. This approach requires E x E comparisons, but TCP efficiency remains high because only matrix operations without control flow are used:

• Input
  • Packed_Value_cmp: T x E
    • 32-bit Packed Tensor with Comparison Trick applied.
• Output
  • Rank: T x E
    • Rank of each Expert (0-based rank). Higher scores are closer to 0.
• Operations
  • Broadcast & Compare
    • Replicate (Tile) Packed_Value_cmp along the E axis to expand to T x E x E shape. Compare magnitude relationships for all Expert pairs (i, j).
    • Compare[t, i, j] = 1 if Packed_Value_cmp[t, j] > Packed_Value_cmp[t, i] else 0
    • Meaning: “Is Expert j’s score higher than Expert i’s?”
  • Rank Calculation (ReduceSum)
    • Sum along the E (comparison target) axis to calculate rank.
    • Rank[t, i] = sum(Compare[t, i, j] for j in 0..E)
    • Meaning: “The total number of Experts with higher scores than me” becomes my rank.

2.3. Filtering & Unpacking

Filtering extracts the top K entries based on rank, then unpacking separates the packed scores and indices:

• Inputs
  • Rank: T x E
  • Packed_Value: T x E
    • Note: The original Packed Value before Comparison Trick was applied must be used to restore accurate Score/Index later.
• Outputs
  • TopK_Indices: T x K
  • TopK_Scores: T x K
  • routing_weights: T x K (weights for K selected experts per Token)
• Operations
  • Filtering (FilterCompaction)
    • Only elements satisfying the Top-K condition (Rank < K) are kept.
    • Mask[t, i] = 1 if Rank[t, i] < K else 0
    • Only Packed_Value at positions where Mask is True are collected and compressed to T x K size.
    • Result: Selected_Packed: T x K
    • Uses the filter function of Vector Engine.
  • Unpacking
    • Restore scores and indices through bit operations from the selected 32-bit values.
    • Score Extraction: TopK_Scores = Selected_Packed >> 16 (then reinterpreted as bf16 type)
    • Index Extraction: TopK_Indices = Selected_Packed & 0xffff
  • Softmax Normalization
    • Softmax is applied to the extracted Top-K Scores to calculate final weights. This is used in the Combine stage later.
    • routing_weights[t, k] = exp(TopK_Scores[t, k]) / sum(exp(TopK_Scores[t, j]) for j in 0..K)

3. Blockwise Execution

Blockwise execution physically rearranges data based on Top-K routing decisions while satisfying TCP’s static shape constraints.

3.1. Problem: Dynamic Shape & Memory Explosion

The core challenge is that the number of tokens L_e assigned per Expert varies dynamically depending on the input. In the worst case, if all tokens are concentrated on a specific Expert, L_e ~ T.

Two approaches address this challenge:

• Naive Solution: Allocating a buffer of maximum size T for all Experts requires memory of E x T x D size, most of which is wasted as padding.
• Blockwise Solution: Instead of variable length L_e, manage data in fixed-size Block (B) units to optimize memory usage to approximately T x K level.

3.2. Grid Size Calculation

Tokens for the same Expert are grouped into blocks of B tokens, enabling blockwise computation with a single expert loaded. Grid size (the total number of blocks across all experts) determines how many blocks are needed to process all tokens.

The total number of blocks needed (Grid Size, G) is calculated as the sum of blocks required per expert:

• Number of blocks allocated to Expert e
  • Number of tokens allocated to e: Count_e
  • Number of blocks: ceil(Count_e / B)
• G = sum(ceil(Count_e / B) for e in 0..E)

The compiler calculates the worst-case G value and allocates memory space. At runtime, sparse operations skip execution for empty Grids.

In the worst case where all Experts include a grid containing only one token, (T*K - E) / B + E Grids are required.

3.3. Index and Expert ID Generation (Cumsum-based Address Calculation)

Each token’s destination block address (Scatter_Idx) and the expert assignment per block (Expert_IDs) are computed in parallel using cumsum over the token-to-expert assignment masks. (Cumsum is implemented in the Vector Engine using branch logging; see Section 4 for the hardware implementation.) This approach avoids loops and enables efficient parallel execution:

• Inputs
  • TopK_Indices: T x K
  • Expert_Indices: E = [0, 1, ..., E-1]
  • Block_Range: G = [0, 1, ..., G-1] (sequence of maximum block count, e.g., 32)
• Outputs
  • Scatter_Idx: T x K (final 1D address where each token will move)
  • Expert_IDs: G (Expert number each Block is responsible for)
• Operations
  • Mask Generation (One-Hot)
    • Convert indices to computable mask form.
    • Expert_Mask: T x K x E = one_hot(TopK_Indices, depth=E)
  • Histogram
    • Sum the masks to count the number of tokens allocated per Expert.
    • Count: E = reduce_sum(Expert_Mask, axis: (T, K))
  • Block calculation
    • Calculate the number of Blocks needed for each Expert.
      • Num_Blocks: E = ceil(Count / B)
  • Global offset Calculation
    • Through Cumsum, obtain the Block Start Index where each Expert starts in the entire Grid (G).
    • Global_Offset: E = cumsum(Num_Blocks) - Num_Blocks
  • Local Offset Calculation
    • Using Mask and Cumsum, calculate what position each token is in the Expert’s queue.
    • Cumsum_Mask: T x K x E = cumsum(Expert_Mask, axis: (T, K))
    • Token_Rank: T x K = gather(Cumsum_Mask, index: TopK_Indices)
    • Local_Offset: T x K = Token_Rank - 1
  • Expert ID expansion
    • Diff: E x G = Num_Blocks - Block_Range
    • Grid: E x G

      • Grid(e, i) = if Diff(e, i) > 0 {
            Expert_Indices(e)
        } else {
            -1
        }

    • Expert_IDs: G = filter_compaction(Grid, condition=(Grid >= 0))
    • e.g.)
      • expert 0: 2 blocks, expert 1: 3 blocks, expert 3: 3 blocks
      • Diff[0] = [2, 1, 0, -1, -2, …], Diff[1] = [3, 2, 1, 0, -1, …]: has positive terms equal to the number of allocated blocks per expert.
      • Grid[0] = [0, 0, -1,-1, …], Grid[1] = [1, 1, 1, -1, -1, …]: has expert id equal to the number of allocated blocks per expert.
      • Expert_IDs = [0, 0, 1, 1, 1, 3, 3, 3]: Filter only values >= 0 (expert id) from Grid.
  • Address Synthesis
    • Scatter_Idx = (Global_Offset * B) + Local_Offset
    • Calculate which block and which position within the block each of the T tokens corresponds to. Scatter_Idx in [0, G * B)

3.4. Dispatch (Blockwise Scatter)

Dispatch physically rearranges tokens using the computed addresses, placing each token in its designated block position:

• Input
  • x_norm: T x D (Input after Attention and norm)
  • Scatter_Idx: T x K (Final 1D address where each token will move)
• Output
  • x_blocked: G x B x D (Rearranged Blocked Tensor)
• Operation
  • Scatter
    • Place tokens x_norm at Scatter_Idx positions.

3.5. Sparse Computation (Weight Gather)

Sparse computation applies Expert weights to the sorted Blocks. The key insight is that weights are gathered only for Experts that have assigned tokens:

• Inputs
  • x_blocked: G x B x D
  • Expert_IDs: G (Expert number each Block is responsible for)
• Output
  • y_blocked: G x B x D
• Operations
  • Weight Gather
    • Using Expert_IDs as indices, only the necessary weights are fetched.
    • W_gathered_up: G x D x F = gather(W_up, index: Expert_IDs)
    • W_gathered_down: G x F x D = gather(W_down, index: Expert_IDs)
  • Sparse MLP
    • Operations are performed only for valid Blocks (G).
    • up: G x B x F = einsum(x_blocked, W_gathered_up)
    • y_blocked: G x B x D = einsum(up, W_gathered_down)

3.6. Combine (Weighted Sum)

Combine restores results to original token order and applies Routing probabilities. This is the final step that produces the MoE layer output:

• Inputs
  • y_blocked: G x B x D
  • Scatter_Idx: T x K
  • routing_weights: T x K
• Output
  • moe_out: T x D (Final MoE layer output)
• Operations
  • Gather
    • Using Scatter_Idx in reverse, results are fetched from y_blocked in the original token order.
    • y_restored: T x K x D = gather(y_blocked, index: Scatter_Idx)
  • Weighted Sum
    • The final output is summed by multiplying with routing_weights obtained from the Top-K process.
    • y_weighted: T x K x D = einsum(y_restored, routing_weights)
    • moe_out: T x D = reduce_sum(y_weighted, axis: K)

4. Cumsum Implementation on TCP

On TCP, cumsum is implemented in the Vector Engine using branch logging:

1. Create a static branch logger: For the axis (of size n) over which the sum is computed,

   branch(i) = if i == 0 {
       0
   } else if i < n - 1 {
       1
   } else {
       2  // i == n - 1
   }

2. Configure the Vector Engine as follows:

   add %mainstream, OperandRead(branch = 1, 2)
   WriteOperand(branch = 0, 1)