Router Latency

The gating network (or router) is a small neural network, typically a linear layer, that must process every input token. Its forward pass is the first and most fundamental source of latency.

• Operation: For a hidden state of dimension d_model, the router computes logits = x @ W_gate, where W_gate is a matrix of shape (d_model, num_experts). This is a dense matrix multiplication.
• Bottleneck: While small, this operation is synchronous for all tokens in a batch and adds fixed overhead before any expert computation can begin. Its cost scales with the number of experts.

After computing expert logits, the router must select the top k experts for each token. This involves a sorting or partitioning operation across the expert dimension.

• Algorithmic Cost: For num_experts=8 and k=2, this is cheap. For models with hundreds of experts, the O(n log n) cost of sorting all expert scores per token becomes significant.
• Implementation: Optimized kernels use vectorized top-k algorithms (e.g., bitonic sort) on GPU. Poorly implemented sorting on the CPU can become a major bottleneck.
• Noise Addition: In Noise Top-k Gating, adding tunable Gaussian noise to logits before selection adds a minor but non-zero computational step.

Once tokens are assigned to experts, they must be regrouped from their original sequence order into expert-centric batches. This data movement is a major latency source.

• On-Device: On a single GPU, this involves a permutation/scatter-gather operation in memory. Efficient kernels fuse this with subsequent matrix multiplies.
• Expert Parallelism: When experts are sharded across multiple devices (GPUs), this becomes a network-bound All-to-All collective communication. Tokens are scattered from all devices to their target expert devices, and results are gathered back. This cross-device communication often dominates latency in distributed MoE inference.

To prevent overloading a single expert, the routing logic must enforce an expert capacity limit, adding decision and data management overhead.

• Capacity Factor: A buffer multiplier (e.g., 1.25) on the theoretical per-expert load. The system must calculate capacity as (batch_size * seq_len * k / num_experts) * capacity_factor.
• Token Dropping Logic: If an expert is at capacity, a fallback policy (e.g., routing to the next-best expert, or passing the token through unchanged) must be executed. This conditional logic and potential re-routing adds non-deterministic latency.
• Auxiliary Loss Calculation: During training, the load balancing loss is computed to ensure uniform expert utilization. While not part of inference, its design influences router behavior and stability.

The final computational step is launching the actual expert network computations. The sparsity of activation makes efficient kernel execution challenging.

• Sparse Matrix Multiplication: The system must perform matrix multiplications using only the weight blocks for the activated k experts. Naively launching many small, separate matrix multiplies incurs severe kernel launch overhead.
• Fused MoE Kernels: High-performance systems use fused kernels (e.g., in NVIDIA's FasterTransformer or vLLM) that combine routing, permutation, and the batched expert GEMMs into a single GPU kernel. The absence of such kernels leads to significant latency from many small operations.
• Kernel Efficiency: The irregular and varying batch size per expert can lead to under-utilized GPU warps, reducing computational efficiency compared to dense models.

Router latency is magnified by its interaction with broader inference serving optimizations like continuous batching and KV cache management.

• Continuous Batching: Dynamic batching groups requests with different sequence lengths. The router must handle a jagged 3D tensor [num_requests, variable_seq_len, hidden_dim], complicating permutation and capacity logic.
• KV Cache per Expert: In some MoE transformer designs, the Key-Value (KV) cache may be managed per expert. Routing a token to a different expert in the next layer may require retrieving its KV states from a different memory location, adding latency.
• Scheduling Overhead: In a multi-tenant serving system, the scheduler must account for the variable compute cost of MoE layers, making predictable low-latency execution more complex than with dense models.

The gating network is the small, trainable neural network component responsible for the routing decision. It computes a set of scores (logits) for each input token to determine which expert(s) should process it. Its architecture and computational cost are the primary determinants of router latency.

• Function: Transforms a token's hidden state into expert probabilities.
• Impact: A simple, shallow network minimizes latency but may reduce routing quality. Complex routers increase accuracy at the cost of compute time.

Top-k gating is the predominant routing strategy where the router selects the k experts with the highest scores for each token. This enforces sparsity and defines the computational graph.

• k=1 (Switch Routing): Minimal routing overhead but no expert blending.
• k=2 (Common in models like Mixtral): Balances model capacity with conditional compute. The router must sort scores to find the top k, adding a small but non-negligible latency component compared to a simple forward pass.

In expert parallelism, where experts are distributed across multiple devices (GPUs), all-to-all communication is the collective operation that implements routing. After the local router decides assignments, tokens are scattered to the devices hosting their assigned experts, and outputs are gathered back.

• Latency Dominant: This network communication often becomes the largest contributor to total router latency in distributed MoE systems, especially at scale.
• Optimization Target: Overlapping this communication with computation is a key systems challenge.

Fused MoE kernels are custom, highly optimized GPU kernels that combine the routing logic, token permutation, and the sparse matrix multiplications for multiple experts into a single operation.

• Latency Reduction: By fusing these steps, they eliminate the overhead of launching multiple small kernels and reduce intermediate data writes to memory.
• Implementation Critical: The use of fused kernels (e.g., in frameworks like vLLM or NVIDIA's FasterTransformer) is essential for achieving low router latency and high throughput in production.

Expert capacity is a pre-defined limit on the number of tokens a single expert can process in a given forward pass. It's a buffer to enable efficient batched computation on experts.

• Latency Trade-off: A low capacity risks dropped tokens, which must be handled (e.g., skipped), potentially degrading quality. A high capacity ensures no drops but forces experts to pad their computations, increasing latency and wasting FLOPs on zero-padded operations.
• Tuning Parameter: Setting optimal expert capacity is crucial for balancing latency, throughput, and model quality.

Sparse activation is the fundamental property enabled by the router, where only a small subset of the model's total parameters (the chosen experts) are computed for a given input.

• Theoretical Benefit: This enables gigantic models (e.g., trillion parameters) with a manageable computational cost per token.
• Latency Reality: The router latency is the price paid to achieve this sparsity. The system overhead of making the routing decision must be less than the compute saved by activating fewer parameters, otherwise the MoE architecture provides no speed advantage.