🤖 machine learning

Lynx: Enabling Efficient MoE Inference through Dynamic Batch-Aware Expert Selection

Lynx is a workload-agnostic system that enables efficient Mixture-of-Experts (MoE) inference by exploiting training-induced activation skews and using a novel AffinityBinning technique to dynamically remap token-to-expert assignments, thereby reducing the number of invoked experts and improving throughput by up to 1.30x without significantly compromising accuracy.

Original authors: Vima Gupta, Jae Hyung Ju, Kartik Sinha, Ada Gavrilovska, Anand Padmanabha Iyer

Published 2026-05-20

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Vima Gupta, Jae Hyung Ju, Kartik Sinha, Ada Gavrilovska, Anand Padmanabha Iyer

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you run a massive, high-end restaurant called "The MoE Kitchen."

In this kitchen, instead of having one giant chef who tries to cook every single dish, you have a team of 64 specialized sous-chefs (the Experts). There's a head waiter (the Router) who looks at every order (a Token) and decides which 8 specific chefs should cook that dish.

This setup is brilliant because it's efficient: you don't need to pay for all 64 chefs to work on every single order. You only pay for the 8 who are needed. This is how modern AI models like Qwen or Llama work—they are huge, but they only "wake up" a small part of their brain for each word they generate.

The Problem: The Rush Hour Traffic Jam

The paper explains a major problem that happens when the restaurant gets busy (which is called Batching).

When you have just one customer, the head waiter sends their order to 8 specific chefs. Easy.
But when you have a batch of 16 customers, the head waiter looks at all 16 orders. Because every customer is different, the waiter ends up needing to call upon almost all 64 chefs to handle the group.

The Bottleneck:
The kitchen is organized so that the ingredients (the chefs' knowledge) are stored in a giant, high-speed pantry (the GPU Memory). To cook, the chefs have to run to the pantry to grab their specific ingredients.

The Issue: When 16 customers arrive, the pantry gets flooded. The chefs spend more time running back and forth grabbing ingredients than actually cooking. The kitchen slows down because the "running" (memory bandwidth) is the bottleneck, not the "cooking" (computation).
The Result: Even though the AI is supposed to be fast and efficient, it slows down to a crawl when handling multiple requests at once because it's stuck fetching data.

The Solution: LYNX (The Smart Waiter)

The authors created a new system called LYNX. Think of LYNX as a super-smart, dynamic manager who steps in only when the kitchen is busy (during the "decode" phase, which is like the restaurant serving food one bite at a time).

LYNX doesn't fire any chefs or change the menu. Instead, it uses a clever trick called AffinityBinning (a fancy way of saying "grouping by confidence").

Here is how LYNX works, step-by-step:

The "Confidence" Check:
Sometimes, the head waiter is 100% sure Chef A is the best for a dish. Other times, the waiter is unsure and picks Chef B just because the rules say "you must pick 8 different people." The paper found that these "uncertain" picks are often redundant.
- Analogy: If you ask a friend for a movie recommendation and they say, "I'm not sure, but maybe Movie X or Movie Y," they aren't really committed to either. If you ask them again, they might just pick Movie X again.
The "Binning" Strategy:
LYNX looks at the waiter's confidence scores. It groups the orders into "buckets."
- High Confidence: "This order must go to Chef A." (LYNX leaves this alone).
- Low Confidence: "This order is just sort of for Chef B." (LYNX says, "Actually, let's send this to Chef A instead, because Chef A is already in the kitchen for the other orders.")
The Great Remap:
LYNX takes those "low confidence" orders and redirects them to the chefs who are already being used by the batch.
- The Magic: Instead of running to the pantry to get ingredients for 64 different chefs, the kitchen now only needs to run to the pantry for, say, 30 chefs.
- Result: The chefs spend less time running and more time cooking. The kitchen moves much faster.

Why This is Special

The paper highlights three key reasons why LYNX is a big deal:

It's "Workload-Agnostic": LYNX doesn't need to be trained on a specific type of customer or menu. It figures out the pattern on the fly, every single time. It's like a manager who learns the crowd's habits instantly without needing a manual.
It Doesn't Break Anything: LYNX doesn't fire chefs or change the recipe. It just rearranges who does what for that specific moment. The paper shows that the food (the AI's answers) tastes just as good, and sometimes even better, because it stops forcing chefs to cook dishes they aren't confident about.
It Plays Well with Others: LYNX can be added on top of other speed-up tricks (like shrinking the chefs' aprons or sending them to a different building) to make them even faster.

The Results

When the authors tested this on real-world AI models (like Qwen, Mixtral, and Llama) across tasks like coding, math, and reasoning:

Speed: The system became 1.3 times faster (a 30% speed boost) in the worst cases.
Accuracy: The answers were just as accurate, or slightly better. In the worst case, accuracy dropped by less than 1%, which is barely noticeable.
Efficiency: It allowed the system to handle more customers at once without slowing down.

Summary

LYNX is like a traffic cop for a busy AI kitchen. It notices that when a group of orders comes in, the kitchen is wasting time running to the pantry for too many different chefs. So, it cleverly redirects the "maybe" orders to the chefs who are already working, reducing the traffic jam and getting the food out faster, all without changing the menu or firing anyone.

Technical Summary: LYNX – Enabling Efficient MoE Inference Through Dynamic Batch-Aware Expert Selection

1. Problem Statement

Modern foundational models increasingly utilize Mixture-of-Experts (MoE) architectures to decouple parameter count from computational cost. While MoEs activate only a small subset of parameters per token (e.g., 14 out of 57 billion in Qwen2-57B), this efficiency promise is broken in production serving due to the tension between latency and throughput.

To achieve high throughput, inference systems batch requests. However, because each token in a batch independently selects its top- $k$ experts, the set of experts that must be loaded from High-Bandwidth Memory (HBM) grows with the batch size and composition. At moderate batch sizes required to meet latency Service Level Objectives (SLOs), this dynamic, data-dependent memory access pattern renders MoE decoding memory-bandwidth-bound. Consequently, the latency scales with the total number of distinct experts activated across the entire batch, often saturating the entire expert pool (e.g., a batch of just 8 diverse requests can activate all 64 experts in Qwen2-57B). This negates the sparsity benefits of MoEs, making them perform similarly to dense models of equivalent capacity but with the added overhead of dynamic expert movement.

Existing solutions, such as static pruning, quantization, or clustering, rely on offline calibration and permanently alter the model, making them inflexible to workload changes. Dynamic per-token approaches fail to address the bottleneck because they do not reduce the aggregate number of experts fetched at the batch level.

2. Methodology: LYNX

LYNX is a workload-agnostic system designed to alleviate the memory bandwidth bottleneck in MoE inference without modifying the model architecture or requiring offline calibration. It leverages a key insight derived from MoE training dynamics: load-balancing losses force routers to distribute tokens broadly across experts, even when confidence in secondary selections is low. This creates batch-level redundancy where the router selects experts due to training constraints rather than genuine token-expert affinity.

LYNX exploits this redundancy through a novel AffinityBinning technique and a three-component pipeline that operates entirely at runtime:

2.1 Core Principles

The design is grounded in three empirical observations:

Phase Sensitivity: MoEs are compute-bound during the prefill phase but memory-bandwidth-bound during the decode phase. Expert remapping is beneficial only during decode. Furthermore, the prefill phase is highly sensitive to expert fidelity, while the decode phase is resilient to reassignment.
Token Confidence: Router confidence scores (derived from the gap between top-ranked expert logits) reliably identify which token-to-expert mappings are critical. Low-confidence tokens can be remapped with minimal accuracy degradation.
Expert Rank Hierarchy: The top-ranked expert (rank-0) dominates output quality. Lower-ranked experts among the top- $k$ exhibit high redundancy. Sensitivity to reassignment also varies across layers.

2.2 System Components

LYNX integrates with the serving engine (e.g., vLLM) via a phase-aware optimizer and three lightweight CUDA kernel components:

Phase-Aware Optimizer: Determines if the current iteration is memory-bound (decode) or compute-bound (prefill). It bypasses LYNX logic during prefill or chunked-prefill phases to avoid unnecessary overhead.
Confidence Analyzer: Intercepts router logits at every layer. It discretizes each token's router confidence into bins using AffinityBinning. The bin width and count are determined solely by the model's sparsity ratio ( $k/N$ ), making the system self-calibrating and independent of the workload. It identifies tokens with low affinity to their assigned experts.
Adaptive Expert Scorer: Computes a batch-wide score for each expert. It uses an exponential weighting scheme based on the bin values assigned by the Confidence Analyzer. This ensures that experts preferred by high-conviction tokens receive disproportionately higher scores. It determines the minimal set of "critical" experts required for the batch.
Expert Remapper: Redirects low-affinity token-to-expert assignments to the surviving critical expert set. Crucially, LYNX preserves the top- $k$ activation semantics; every token still activates exactly $k$ experts, but the batch collectively activates a smaller, consolidated set of unique experts.

2.3 Implementation

LYNX is implemented as a plug-and-play module on top of vLLM. Its core logic is realized using four fused Triton/CUDA kernels that intercept the router output. This design eliminates intermediate tensor data movement and reduces kernel launch overhead, ensuring the system adds negligible latency (<4% overhead) while operating in the critical path.

3. Key Contributions

First Workload-Agnostic System: LYNX is the first system to alleviate the memory bandwidth bottleneck in MoE inference without requiring calibration data or permanent model changes.
AffinityBinning Technique: A novel method to discretize router confidence based on architectural sparsity ratios, enabling automatic adaptation to any MoE architecture.
Batch-Level Optimization: Unlike prior per-token dynamic approaches, LYNX operates at the batch level to reduce the total number of unique experts fetched from HBM, directly addressing the memory bandwidth bottleneck.
Complementarity: LYNX is designed to work in conjunction with existing optimizations, including offloading and quantization.

4. Experimental Results

LYNX was evaluated on state-of-the-art models from four families (Qwen, DeepSeek, Mixtral, Llama) across eight benchmarks (code generation, mathematics, reasoning, and vision).

Latency Reduction: LYNX achieves up to 1.30× lower latency (Time Per Output Token) compared to baselines. In disaggregated serving scenarios (where prefill and decode are separated), reductions reach up to 1.35×.
Accuracy Impact: The system incurs less than 1% accuracy loss in the worst case. Notably, LYNX frequently improves accuracy on average (e.g., +1.9% on HumanEval for DeepSeek-Coder-V2), suggesting that training-time load balancing forces suboptimal expert selections that LYNX corrects.
Throughput: Under fixed SLO budgets (e.g., 25ms or 50ms), LYNX improves system throughput by up to 2.1×.
Compatibility:
- Offloading: When combined with offloading techniques (e.g., Fiddler), LYNX boosts performance by 31%.
- Quantization: When applied atop INT4 quantization, LYNX provides an additional 10% speedup.
- Parallelism: LYNX provides additive throughput gains on top of Expert Parallelism (EP) and Tensor Parallelism (TP).
Robustness: Performance gains are consistent across varying batch sizes (saturating at moderate sizes) and real-world traces with bursty arrival rates.

5. Significance

The paper posits that as GPU compute continues to outpace memory bandwidth growth, the arithmetic intensity gap for MoE decoding will widen, making memory-bound bottlenecks increasingly critical. LYNX addresses this by fundamentally changing how experts are selected at the batch level. By exploiting the redundancy introduced by training-time load balancing, LYNX enables MoE architectures to realize their efficiency promise on future hardware without requiring model retraining or offline profiling. The system demonstrates that dynamic, runtime-aware expert management can significantly decouple throughput from the memory bandwidth constraints inherent in current MoE serving paradigms.