Continuous Batching for MoE

In a continuous batch, incoming requests contain varied token sequences, leading to highly uneven loads across the expert networks. Without intervention, this causes GPU underutilization as some experts are overloaded while others sit idle.

Solutions:

• Dynamic Capacity Factor Adjustment: The system monitors token routing in real-time and temporarily increases the capacity factor for overloaded experts to prevent token dropping.
• Expert-Aware Batch Scheduling: The scheduler groups requests with complementary token-expert affinities to create a more balanced load across all experts in the running batch.

Under expert parallelism, tokens must be scattered to the GPUs hosting their assigned experts and gathered back after processing. This all-to-all communication becomes a severe bottleneck in a continuous batch where tokens are constantly arriving and departing.

Solutions:

• Fused MoE Kernels: Kernels like FlashFFN or Triton MoE combine routing, permutation, and the expert computation into a single operation, drastically reducing the need for separate, costly communication collectives.
• Overlap Communication and Computation: Advanced pipelines hide communication latency by prefetching expert weights and streaming token data while other parts of the model (e.g., attention) are computing.

The core MoE operation is a sparse matrix multiplication where only the weight blocks for activated experts are used. In a dynamic batch, this sparsity pattern changes constantly, making it difficult to leverage dense tensor cores on modern GPUs efficiently.

Solutions:

• Blocked Sparse Layouts: Expert weights are stored in a blocked format aligned with GPU memory and compute units (e.g., 16x16 or 32x32 blocks), allowing the sparse computation to still use highly optimized tensor core instructions.
• Predictive Kernel Selection: The serving system profiles different kernel implementations (fully sparse, dense with masking, blocked) and selects the optimal one based on the real-time expert activation sparsity of the current batch.

In transformer-based MoE models, each expert may maintain a separate Key-Value (KV) cache for the sequences it processes. In continuous batching, as sequences finish and new ones start, this leads to highly fragmented cache memory across experts, increasing memory overhead and reducing usable batch size.

Solutions:

• PagedAttention for MoE: Extending mechanisms like vLLM's PagedAttention to manage the KV cache at a per-expert level. Memory is allocated in fixed-size blocks that can be non-contiguously assigned to different sequences, eliminating fragmentation.
• Expert-Shared Cache Proposals: Research into architectures where the attention mechanism is kept dense and shared, while only the FFN experts are sparse, simplifying cache management.

The gating network (router) must execute for every token in the dynamic batch, adding a sequential, non-parallelizable step. This router latency can become a dominant factor in total per-token latency, especially with large vocabulary sizes or complex routing functions.

Solutions:

• Lightweight Router Architectures: Employing extremely efficient router designs, such as linear layers with top-k gating, that minimize FLOPs.
• Router Computation Overlap: The router for the next set of tokens in the sequence is computed concurrently with the expert network processing of the current tokens, effectively hiding its latency within the expert computation time.

An inference optimization technique where incoming requests are dynamically added to and removed from a running batch as they finish generation, rather than waiting for an entire static batch to complete. This maximizes GPU utilization by eliminating idle time.

• Core Mechanism: Maintains a persistent batch of active sequences, scheduling new tokens for computation as previous ones are completed.
• Contrast with Static Batching: Static batching processes fixed-size batches, leading to padding and underutilization when sequences finish at different times.
• Key Benefit for MoE: Essential for handling the variable computational load per token introduced by sparse expert activation, allowing the system to efficiently pack tokens with similar expert assignments.

A model parallelism strategy designed for Mixture of Experts where different expert networks are placed on different devices (e.g., GPUs). Tokens are routed and communicated between devices based on the gating network's decisions.

• Communication Pattern: Relies heavily on All-to-All communication operations. After routing, tokens are scattered from all devices to the specific devices hosting their assigned experts; processed outputs are then gathered back.
• System Challenge: The efficiency of this data exchange is a major bottleneck. Continuous batching must coordinate with this pattern to minimize communication overhead across the dynamic batch.
• Integration: Optimized systems co-design continuous batching schedules with expert placement to reduce cross-device traffic for tokens routed to the same experts.

A property of conditional computation models, like Mixture of Experts, where only a small, dynamically chosen subset of the model's total parameters is activated and computed for a given input.

• Contrast with Dense Models: In dense models (e.g., standard transformers), all parameters are used for every token. Sparse activation provides computational savings proportional to the sparsity.
• Routing Determines Sparsity: The gating network (router) selects which experts to activate per token, typically using a Top-k Gating policy (e.g., top-2).
• Batching Implication: This creates a variable computational graph per token. Continuous batching must group tokens not just by sequence length, but by their activated expert sets to enable efficient Sparse Matrix Multiplication.

In MoE, a set of techniques designed to prevent the router from consistently favoring a small subset of experts, ensuring uniform computational load and parameter utilization.

• Problem: Without balancing, a few 'popular' experts become bottlenecks, while others are underutilized, wasting capacity and harming model performance.
• Mechanisms: Often enforced via an Auxiliary Loss (Load Loss) during training and operational constraints like the Capacity Factor during inference.
• Interaction with Batching: Continuous batching must respect expert capacity limits. If an expert's Expert Capacity is full, excess tokens may be Dropped or handled by a fallback, impacting throughput and quality. The batch scheduler must be aware of these limits.

Highly optimized GPU kernels that combine the routing logic, token permutation, and the sparse matrix multiplications of multiple experts into a single, efficient operation.

• Purpose: To minimize the latency overhead of launching many small, separate kernels and to reduce costly memory movements between routing and computation phases.
• Key Optimization: Fuses the Top-k selection, token sorting/expert assignment, and the batched matrix multiplications for the activated expert weights.
• System Impact: Enables continuous batching to achieve high throughput by making the per-token routing and sparse computation extremely fast. The kernel efficiently handles the irregular batch structure created by dynamic routing.