Mixtral (by Mistral AI)

Mixtral 8x7B uses a top-2 routing policy, meaning for each token in a sequence, only 2 out of 8 experts in each layer are activated. While the model has ~47B total parameters, only ~13B are used per forward pass. This conditional computation means its computational cost (FLOPs) is closer to a 13B dense model, not a 47B one, but it retains the knowledge capacity of the larger parameter count.

• Key Implication: Throughput is gated by the active parameters, but memory footprint is dominated by the full parameter set.

To fit Mixtral's large parameter count across multiple GPUs, Expert Parallelism is used. Each of the 8 experts in a layer is placed on a different device. The routing decision triggers an All-to-All communication collective:

1. Scatter: Tokens are sent from all GPUs to the specific GPUs hosting their assigned experts.
2. Compute: Each expert processes its assigned tokens locally.
3. Gather: The expert outputs are sent back to the original GPUs.

This communication overhead is a primary bottleneck and must be optimized with fused kernels and efficient networking.

A naive implementation of MoE—running the router, permuting tokens, and performing 8 separate matrix multiplications—is highly inefficient. Fused MoE kernels (e.g., in DeepSpeed or FlashAttention) combine these steps:

• They perform the top-k gating and token sorting in-register.
• Execute a batched, sparse matrix multiplication across only the activated expert weights.
• This minimizes kernel launch overhead and expensive memory movements between HBM and SRAM.

Minimizing router latency—the time for the small gating network to compute scores—is critical for low-latency inference.

Standard continuous batching groups requests with different sequence lengths into a single batch to maximize GPU utilization. For Mixtral, this is more complex because each request's computational graph varies based on router decisions.

• The system must dynamically batch tokens after routing, grouping tokens destined for the same expert across different requests.
• This requires sophisticated scheduling to balance expert capacity (the buffer per expert) and avoid dropped tokens while maintaining high GPU utilization.
• Serving engines like vLLM and TGI have extended their memory management and batching systems to handle MoE's sparse patterns.

Mixtral's memory demand has two major components:

• Parameter Memory: ~90GB in FP16 for the 47B parameters. Must be loaded, but only a fraction is active per token.
• KV Cache: The Key-Value cache for the attention mechanism. In MoE, a critical design choice is whether the KV cache is shared across experts or expert-specific.
  • Shared KV Cache: More memory-efficient, but may limit expert specialization.
  • Expert-Specific KV Cache: More expressive but multiplies the cache size by the number of experts (k), drastically increasing memory pressure.

Optimizing MoE KV cache eviction policies (like PagedAttention) is an active area of development.

To reduce the high memory footprint of the 47B parameters, aggressive model quantization is essential:

• Post-Training Quantization (PTQ): Converting weights to INT8 or FP8 can halve memory needs with minimal accuracy loss.
• GPTQ/AWQ: More advanced weight-only quantization methods that are highly effective for MoE models.
• Weight Offloading: For extreme memory savings, inactive expert parameters can be offloaded to CPU RAM or NVMe SSD, then fetched on-demand when routed. This trades off memory for increased latency due to PCIe/IO transfers.

These techniques make serving Mixtral on more affordable hardware (e.g., 2-4 GPUs) feasible.

A neural network architecture where the model is composed of multiple specialized sub-networks called experts. A gating network (router) dynamically selects a sparse subset of these experts to process each input. This allows for models with massive parameter counts (e.g., hundreds of billions) while maintaining a manageable conditional computational cost, as only a fraction of the total parameters are active per token.

The defining property of conditional computation models like MoE. Unlike dense models where all parameters are used for every input, sparse activation means only a small, dynamically chosen subset of the model's total parameters is activated and computed. In Mixtral, this is achieved via top-2 routing, where each token activates only 2 out of 8 experts in a layer, leading to significant compute savings.

The routing strategy at the heart of MoE models. A small, trainable gating network (router) computes scores for each expert based on the input token. The top-k experts with the highest scores are selected for processing that token.

• Mixtral uses top-2 gating.
• Noise Top-k Gating is a common variant that adds noise to scores during training to encourage balanced expert utilization.
• The router's decision is a primary source of router latency.

Critical training techniques to prevent router collapse, where the model overuses a few experts and ignores others.

• Load Balancing ensures computational load is spread evenly across experts.
• An Auxiliary Loss (load loss) is added to the main training objective. It penalizes the variance in token assignment across experts, incentivizing the router to utilize all experts effectively.

A model parallelism strategy designed for MoE, where different experts are placed on different hardware devices (e.g., GPUs). This is necessary to fit ultra-large MoE models into memory. It requires efficient all-to-all communication operations to:

1. Scatter tokens to the devices hosting their assigned experts.
2. Gather the processed outputs back. This communication overhead is a key bottleneck in distributed MoE inference.

The core computational kernel for MoE inference. The weight matrix for an MoE layer is a large, block-sparse matrix where each block corresponds to a different expert. During forward pass, the system performs a sparse matrix multiplication, calculating results only for the blocks belonging to the activated experts (e.g., 2 out of 8). Highly optimized fused MoE kernels combine routing and computation to maximize GPU efficiency.