METRO — Balancing Activated Experts (not Tokens) for Memory-Bound MoE Serving #

1. TL;DR #

In memory-bound MoE decode, GPU runtime is set by the number of *activated

expert replicas* (weight loads), not tokens. Token-balancing load balancers

(EPLB) inadvertently inflate activated experts and hurt decode. METRO instead

minimizes max activated experts/GPU via a lock-guarded greedy kernel + an

all-gather dispatch for global top-k, cutting decode latency 11–22% and lifting

throughput 3–21%.

2. Q1 / Q2 / Q3 #

Q1 — 痛点 (the problem) #

All existing expert-parallel (EP) load balancers — EPLB and its many variants —

share one objective: balance the number of tokens each GPU processes. This

implicitly assumes runtime scales linearly with token count, which is only true

in the compute-bound regime (prefill, training). But MoE **decode is

memory-bound**: with one token per request, small batches, rising MoE sparsity,

and GPUs whose FLOPs/byte keeps climbing, the arithmetic intensity of a decode

FFN sits ~2 orders of magnitude below the hardware roofline (batch < 64) and

still 47%–3.0× below it at batch 1024. In this regime the MoE-layer runtime is

governed by how many distinct expert replicas get loaded from HBM — the

activation traffic is < 0.6% of the expert-weight traffic even at a 1K decode

batch. Token-balancing spreads each expert's tokens across all its replicas to

equalize counts, which activates more replicas per GPU and therefore raises

memory traffic. Measured on Qwen3-30B/vLLM, EPLB at 1.5× replication inflates

activated experts ~30%, raising decode latency 14% and dropping throughput 10%.

Q2 — 方法 (the approach) #

METRO (Minimum Expert Token ROuting) is a token-routing algorithm only — it

keeps EPLB's expert placement/replication untouched (so prefill is unaffected)

and swaps only the decode-phase routing objective: **minimize the maximum number

of activated expert replicas across GPUs**, formalized as an ILP called

MIN-EXP-ROUTING. A key lemma shows every feasible solution can route all of an

expert's tokens to a single replica without hurting the objective, collapsing

the problem to "pick one replica per active expert to minimize per-GPU replica

count." The exact optimum reduces to restricted-assignment makespan minimization

(binary search + bipartite matching / max-flow) but costs 31–104% of an FFN's

runtime — prohibitive. METRO instead uses a GPU-native greedy ($O(|A|)$):

each active expert is assigned to the candidate GPU with the fewest currently

activated experts, using per-GPU load counters guarded by locks acquired in GPU-ID

order (deadlock-free), all inside a single SM with counters/locks in shared

memory. To feed the greedy algorithm the global top-k (T[1..N]) that

all-to-all would keep local, METRO replaces the EP dispatch with an **all-gather

before top-k**: every GPU sees all tokens, computes global top-k, routes, runs

FFN, then does the usual all-to-all combine.

核心技术壁垒: The single hardest-to-replicate insight is the *objective reframing itself* — proving (Lemma 1) that minimizing memory traffic ≡ minimizing max activated replicas ≡ a single-replica assignment problem, and then realizing that a seemingly worse all-gather dispatch is effectively free in the memory-bound small-batch regime because NCCL launch cost (~100µs) dwarfs the transfer delta (~3µs vs ~400ns). Everything downstream (greedy kernel, single-SM placement) is engineering; the reframing is the load-bearing idea.

Q3 — 结果 (the results) #

Versus EPLB routing (identical placement/replication) on vLLM/8×A100 (real) and

a proprietary B200 simulator (8–16 GPUs): decode latency ↓ 1.9–21.8%, total

throughput ↑ 0.7–21.0%, gains growing with replication ratio. At a fixed decode

SLO (throughput-latency Pareto), METRO delivers 1.98×–4.11× higher decode

throughput by spending its latency headroom on larger batches. The greedy

router is within 10.9% of the optimal's activated-expert count and up to 42.3%

below EPLB's.

3. 架构 / 方法图 #

METRO lives entirely inside the EP MoE layer of a serving stack (vLLM). The

baseline EP workflow — top-k → token routing → all-to-all dispatch → expert FFN

→ all-to-all combine — is shown first:

*Paper's Figure 2, verbatim (caption: "Expert-parallel MoE inference workflow

with expert placement and replication, as well as token routing – the algorithm

to dynamically route tokens to expert replicas.").*

Each GPU computes top-k locally, the router chooses a physical replica ("which

replica?"), an all-to-all dispatch ships tokens to the hosting GPU, FFNs run,

and an all-to-all combine returns outputs. The router is the only stage METRO

changes; placement/replication (Expert1, Expert4 replicated here) stay EPLB's.

METRO's dataflow reorders the dispatch so global top-k is available before

routing:

*Paper's Figure 7, verbatim (caption: "METRO replaces the conventional all-to-all

with all-gather dispatch before top-k for every GPU to obtain the global top-k

knowledge as input to Algorithm 1, with minimal overhead.").*

Notice the ordering flip vs Fig. 2: AllGather Dispatch → Topk → FFN, then

routing decisions, then only an all-to-all combine. The all-gather is what

gives every GPU the global token-per-expert vector T[1..N] needed to minimize

activated experts globally rather than locally.

Request lifecycle & scheduler placement. METRO is not a request scheduler —

it delegates queueing/batching to vLLM's continuous-batching scheduler and only

intervenes at the per-layer routing decision. The lifecycle:

The memory manager is EPLB's, not METRO's: expert replicas are placed/sized

by EPLB's two-step algorithm (replica count ∝ prior-window token load, placed to

balance expected tokens); METRO never reallocates HBM. Cross-GPU transport is

NVLink (600 GB/s A100 / 900 GB/s B200), single domain; collectives are

all-gather (dispatch) + all-to-all (combine).

4. 作者证明 #

There is a formal model here — the MIN-EXP-ROUTING ILP plus a reduction to

restricted makespan minimization — not just an empirical throughput fit.

Notation table:

The ILP: minimize $\lambda$ subject to (for all $g,i$)

$$\sum_{i=1}^{N} y_{i,g} \leq \lambda \tag{1}$$

$$\sum_{g=1}^{G} x_{i,g} = T[i] \tag{2}$$

$$x_{i,g} = y_{i,g} = 0 \ \text{ if } A_{i,g}=0 \tag{3}$$

$$x_{i,g} \leq T[i]\cdot y_{i,g} \tag{4}$$

Once $y_{i,g}$ is chosen, Lemma 1 lets the rest be read off directly: $x_{i,g}=T[i]$ if $y_{i,g}=1$ else $0$, and $\lambda = \max_{g}\sum_i y_{i,g}$.

6 minimum checks:

1. 物理意义 of the objective (max, not sum). $\lambda$ is a max over GPUs because end-to-end latency is set by the slowest (most-loaded) GPU; equalizing/minimizing the peak is what shrinks the critical path. Summing would optimize total work, which is irrelevant when GPUs run in parallel.
2. Why (1) uses ≤ λ. Constraint (1) bounds every GPU's activated-expert count by the single scalar $\lambda$; minimizing $\lambda$ therefore squeezes the worst GPU. This is the memory-traffic proxy: activated replicas × per-replica weight bytes dominates HBM traffic (activations <0.6%), so bounding replicas bounds runtime.
3. Why (4) couples $x$ and $y$. $x_{i,g}\le T[i]\,y_{i,g}$ forces $y_{i,g}=1$ whenever any token flows to $(i,g)$ — you cannot route to a replica "for free"; every touched replica counts toward the memory-traffic objective. This is exactly what token-balancers ignore.
4. Lemma 1 (single-replica optimality). Any feasible solution either routes each expert's tokens to one replica, or maps to one that does without raising $\lambda$. Proof: given a solution splitting expert $i$ across replicas, collapse them onto one GPU; (2),(3),(4) still hold and no GPU's activated count increases, so (1) holds with $\lambda\le\lambda_0$. This is the load-bearing simplification: routing becomes a discrete replica-selection problem, not a token-fractioning LP.
5. Optimum is at a boundary / combinatorial. With Lemma 1 the feasible region is the integral assignment polytope; the optimum is a vertex (an integral single-replica assignment), reachable by binary search on $\lambda$ + bipartite matching where each GPU node matches ≤ $\lambda_0$ times. No interior optimum exists — hence a scheduling (makespan) reduction, not gradient methods.
6. Complexity accounting = why the optimum is rejected. Binary search factor $O(\log\lceil|A|G\rceil)$ (since $\lambda\le\lceil|A|G\rceil$); each feasibility test is a max-flow $O((N+G)^2(\lceil|A|G\rceil+N+G))$. Measured, this is 116–129µs (CPU) / 290–292µs (GPU) = 31–104% of one Qwen3-30B FFN — rejected. METRO's greedy is $O(|A|)$ instead of $O((N+G)^2(\lceil|A|/G\rceil+N+G)\log\lceil|A|/G\rceil)$.

First-order mapping (sanity of the memory-bound claim). Plug the case study

into the traffic argument: 32 decode tokens/GPU, fp16, 8 GPUs. All-to-all sends

256KB/GPU (~400ns at 600 GB/s), all-gather 2MB/GPU (~3µs) — both far below the

~100µs NCCL launch fixed cost, so the "expensive" all-gather is free in practice.

On the compute side, activation traffic <0.6% of expert-weight traffic at 1K

batch ⇒ runtime ∝ activated replicas, matching the strong measured correlation

between activated experts and decode latency (Fig. 5b vs 5d). Greedy quality:

within 10.9% of optimal, up to 42.3% below EPLB (Fig. 8) — the model's numbers

come out without sweep-fitting.

5. 实验与数据 #

Motivation — decode is memory-bound (roofline).



*Paper's Figure 3, verbatim (caption: "DeepSeek-V3 and Qwen3-30B attainable

operational intensities VS. FLOPs/byte ratio of H100 and B200. The former is two

orders of magnitude lower than the latter with batch size smaller than 64 tokens,

and 47% - 3.0× lower with a batch size of 1024 tokens.").*

The model FFN curves sit far below the hardware FLOPs/byte lines at all realistic

decode batch sizes — the empirical grounding for treating decode as memory-bound

and thus for counting replicas, not tokens.

The headline reversal — token-balancing backfires.



*Paper's Figure 5, verbatim (caption: "The performance impact of EPLB on prefill

latency (a), decode latency (b), overall token throughput (c), and maximum number

of activated experts across GPUs per decode batch (d) for Qwen3-30B on vLLM …

EPLB reduces prefill latency by 17% with batch size 32, but inflates the number

of activated experts by 30% with 1.5x replication. As a result, the decode

latency increases by 14% and the overall token throughput decreases by 10% with

1.5x replication.").*

Panels (b) and (d) move together: as replication rises, EPLB's activated experts

climb ~30% and decode latency climbs 14% — direct evidence that replica count,

not token count, drives decode. Panel (a) shows the opposite for compute-bound

prefill, exposing the co-deployment tension METRO resolves.

Routing quality — greedy vs optimal vs EPLB.



*Paper's Figure 8, verbatim (caption: "The maximum number of activated experts

per GPU per decode batch (32 tokens) for EPLB routing, the optimal algorithms,

and METRO … METRO is within 10.9% higher than the optimal algorithm and is lower

than EPLB by up to 42.3%.").*

METRO tracks the optimal bar closely (≤10.9% gap) while EPLB's bars are far

taller — the greedy approximation loses almost nothing yet costs $O(|A|)$.

End-to-end (simulated, larger models).



*Paper's Figure 10, verbatim (caption: "Simulated total token throughput and

decode latency … METRO's consistent benefit on decode latency (up to 21.8% at

1.5× replication) improves total token throughput to outperform EPLB routing for

every replication ratio (up to 21.0% at 1.5× replication).").*

Decode-heavy Humaneval (10a/10c) is where METRO shines (+13.5% / +21.0%);

prefill-heavy GSM8K (10b/10d) gains less (+3.0% / +4.2%) — the honest

regime-dependence, quantified in §6 below.

Fixed-SLO Pareto — the 4.11× headline.



*Paper's Figure 12, verbatim (caption: "Pareto curves of the decode phase …

For a fixed TPOT (representing a specific SLO), METRO delivers remarkably higher

decode throughput of 1.98× – 4.11× across models and datasets. Higher replication

ratio improves METRO's performance gain because it exacerbates EPLB's inflation

on activated experts.").*

DeepSeek-V3/Humaneval (12c) at 1.5× replication reaches 4.11× because METRO's

lower TPOT lets it run a 4× larger batch under the same SLO — the gain exceeds

§6-B's fixed-config numbers precisely because the SLO framing unlocks batch-size

freedom.

6. 论证链 #

#	Step	Evidence (paper-internal)
1	Decode is memory-bound: model operational intensity is 2 orders below GPU FLOPs/byte at batch<64, still 47%–3.0× below at 1024.	Fig. 3; §III-A
2	In the memory-bound regime, MoE-layer runtime ∝ memory traffic, dominated by expert-weight loads; activation traffic is <0.6% of weight traffic at 1K batch.	§III-B analytical model [13]
3	Therefore runtime is set by the number of activated expert replicas, confirmed by strong correlation between activated experts and measured decode latency.	Fig. 5b vs Fig. 5d
4	Token-balancing (EPLB) spreads each expert's tokens across replicas, inflating activated experts ~30% at 1.5×, raising decode latency 14% and cutting throughput 10%.	Fig. 5b/5c/5d; Fig. 4 toy example
5	Minimizing max activated replicas is an ILP (MIN-EXP-ROUTING); Lemma 1 collapses it to single-replica assignment; the exact optimum reduces to makespan minimization.	Eqs (1)–(4), Lemma 1; §IV-A/B
6	The exact optimum costs 31–104% of an FFN — prohibitive — so a greedy $O(\	A\	)$ single-SM kernel + all-gather dispatch approximates it within 10.9% of optimal.	Fig. 6; Fig. 8; Algorithm 1
7	Net result: decode latency ↓11–22%, throughput ↑3–21%, up to 4.11× decode throughput at fixed SLO, with all three overheads dominated by FFN savings.	Fig. 9/10/12; Fig. 11 breakdown

5b. Scheduling & resource management #

• Routing granularity: per-expert (each active expert → exactly one replica), decided per decode batch per layer. Not per-token, not per-request.
• Preemption / admission / fairness: none — METRO is a within-batch routing decision, not a request scheduler; queueing, preemption, and admission are delegated to vLLM's continuous-batching scheduler.
• Memory management: delegated to EPLB. Allocation unit is an expert replica; replica count ∝ prior-window token load; no eviction/swap/prefix-cache logic is introduced by METRO.
• Synchronization: per-GPU load-counter lock (test-and-set, shared-memory), acquired in GPU-ID total order to avoid deadlock; concurrency is bounded below 64 by locking, which is why a single SM suffices.

6b. Workload characterization — where it wins / loses #

Workload regime	METRO	EPLB routing	Why
Decode-heavy (Humaneval, InstructCoder, NuminaMath)	↑ up to 21.0% throughput, ↓ up to 21.8% TPOT	degrades with more replication (activated experts inflate)	decode dominates runtime; replica count is the bottleneck METRO minimizes
Prefill-heavy (GSM8K)	modest ↑ up to 4.2%	already benefits from replication (compute-bound prefill)	decode is a small fraction of runtime; less headroom to recover
Fixed strict SLO, larger batches allowed	1.98×–4.11× decode throughput	lower	METRO's TPOT headroom enables up to 4× larger batches under same SLO
Extreme low TPOT (1/TPOT>0.9, batch ≤64)	gains diminish; converges to full TP	same	batch leaves memory-bound regime; network latency dominates, full TP removes imbalance, so no EP balancing (incl. METRO) helps
Disaggregated decode	↑ 4.3% (1.125×) / 5.0% (1.5×) over no-replication	no-replication preferred	decoupled from prefill; smaller but positive gains

7. 实现 cross-reference #

• Greedy router (Algorithm 1): implemented as a single-SM CUDA kernel; parallelism bounded by expert count (128 Qwen / 256 DeepSeek), locking reduces concurrency below 64, so one A100 SM suffices; load counters L and locks l kept in SM-local shared memory for fast access, using a test-and-set lock [1]. (§V — [实现未公开] beyond this description; no repo link in source.)
• vLLM integration: METRO folded into vLLM's decode-phase CUDA Graphs via vLLM's compilation framework [2] to avoid extra kernel-launch overhead; graphs precompiled for power-of-two batch sizes up to 32 tok/GPU, non-powers padded up. (§V)
• All-gather dispatch: replaces the standard all-to-all EP dispatch so every GPU sees all tokens and computes global top-k before routing. (§IV-C, Fig. 7)
• API / usability: no user-facing API change surfaced — METRO is a drop-in swap of the decode routing algorithm under the existing EPLB placement, applied only to the memory-bound decode phase (prefill keeps EPLB routing). Migration for a vLLM shop is a router-kernel swap + all-gather dispatch patch, not a config flag or library replacement.
• Adoption: no upstream-merge or production-deployment claim stated; evaluated on vLLM/A100 (real) and a proprietary B200 simulator only.

关键实现细节 (easy-to-miss tricks):

1. Total-order lock acquisition on GPU-ID is what makes the parallel greedy deadlock-free while keeping it inside one SM — the "single SM" choice is not a limitation but a deliberate consequence of the ≤64 effective concurrency, letting counters/locks live in shared memory.
2. *The all-gather is inserted before top-k, not after* (contrast Fig. 2 vs Fig. 7). This reordering — plus padding to power-of-two batch sizes for CUDA-graph reuse — is why the "more expensive" collective adds no measurable time: the delta (~3µs vs ~400ns) is buried under NCCL's ~100µs launch cost in the small-batch decode regime.

核心技术壁垒 (dedicated paragraph). The replication-hardest part is the

objective reframing validated by Lemma 1: recognizing that in memory-bound

decode the loss function is "max activated replicas per GPU," proving this

collapses to a single-replica assignment problem, and — critically — that the

communication scheme required to compute it globally (all-gather) is free in

exactly the regime where the objective matters (small batches where NCCL launch

cost dominates). A competitor can copy Algorithm 1 in an afternoon; reproducing

the insight that token-balance is the wrong target and that the "expensive"

dispatch is actually cheap is the durable moat. Everything else (single-SM

kernel, CUDA-graph folding) is standard systems engineering.