SlimWise: Decoupling Expert Pruning Across Prefill and Decode for Efficient MoE Serving

AI summary

Overview

• Research area: Efficient inference and serving for sparse mixture-of-experts (MoE) language models.
• Technical level: Intermediate.
• Scope: The paper presents SlimWise, a serving approach that uses the full model to process prompts but a pruned model to generate responses, passing the prefill-generated KV cache directly between them.

What This Paper Is About

MoE models activate only a small set of experts for each token, but a batch of requests can collectively access much of the model’s expert pool. That makes loading expert weights a major cost during token-by-token decoding, even though pruning experts can hurt quality if it is applied to both prompt processing and generation.

SlimWise aims to reduce decoding costs without unnecessarily pruning the compute-bound prompt-processing phase. It also investigates whether transferring the full model’s KV cache to the pruned decoder can preserve quality, and whether a small amount of targeted distillation can address remaining accuracy and generation-length problems.

Key Contributions

1. Separates expert capacity by inference phase. SlimWise uses all experts during prefill and a selected subset during decode. Because pruning preserves the attention architecture and KV cache format, the pruned decoder can use the full model’s KV cache without conversion.

2. Supports two serving configurations. The framework supports prefill–decode disaggregation, where separate instances handle the phases, and colocated serving, where router masking applies the full expert pool during prefill and restricts decoding to retained experts.

3. Identifies generation-length distortions that accuracy scores can miss. The paper shows that pruning can make responses too short on some tasks and excessively long on others, even when benchmark accuracy remains close to the full model.

4. Adds distillation tailored to the KV cache handoff. The pruned decoder is trained to continue from teacher-generated KV caches. Training updates only a small subset of parameters, while the routed expert weights and attention parameters remain frozen.

Main Findings

• The KV cache handoff often improves accuracy without training. Across Qwen3.6-35B-A3B and Gemma 4-26B-A4B-it, and across routing-mass, EAN, and REAP pruning, SlimWise frequently recovers some of the accuracy lost by conventional pruning. The size of the recovery depends on the model, pruning criterion, and benchmark.

• Generation behavior can change substantially. With REAP pruning on Qwen3.6-35B-A3B, conventional pruning reduced HumanEval+ median output length from 2.3k to 0.3k tokens, while on MATH-500 it increased the median from 2.4k to 6.6k tokens. SlimWise brought the HumanEval+ median back to approximately 2.3k tokens, but the MATH-500 median remained approximately 5.8k tokens before distillation.

• Distillation can improve quality and response lengths. The method trains on the inherited full-model KV cache, rather than only the pruned model’s own prefill states. For example, at 75% REAP pruning on Qwen3.6-35B-A3B, the MATH-500 90th-percentile output length was 32.8k tokens with the KV cache handoff alone and 14.0k tokens after the handoff-aware distillation, compared with 6.1k tokens for the baseline. The distortion was reduced, but not eliminated.

• Decode throughput improves, while prefill throughput changes little. Under PD-disaggregated serving on Qwen3.6-35B-A3B, SlimWise achieved 1.40–1.81× the full model’s decode throughput at 50% pruning and 1.65–2.39× at 75% pruning across the evaluated per-user generation-speed targets. Under PD-colocated serving, the corresponding ranges were 1.36–1.72× and 1.57–2.28×. The paper reports that prefill throughput stays within a few percent of the unpruned model, with no consistent improvement.

• Disaggregation offers an additional memory benefit. A pruned decode instance stores fewer expert weights, leaving more memory available for KV cache and potentially larger batches. Colocated serving still reduces decode expert traffic, but keeps the full expert pool resident and therefore does not gain that memory capacity.

Methodology in Plain English

The researchers first rank experts within each layer using three criteria: routing mass, Expert Activation Norm (EAN), and Router-weighted Expert Activation Pruning (REAP). They then retain a chosen number of experts while keeping the number selected for each token unchanged.

In SlimWise, the full model processes the prompt and creates the KV cache. A decoder with the retained experts then continues generation from that cache. For colocated serving, the same behavior is implemented by masking the router during decode so it cannot select removed experts. For disaggregated serving, the decoder can instead load a checkpoint containing only retained experts.

To improve the decoder further, the researchers distill it using training sequences split at sampled positions. The full model processes each sequence prefix and supplies its KV cache; the pruned decoder learns to continue from that state. The training updates only selected components, such as routers and normalization layers, along with the model-specific shared expert or dense MLP; routed expert weights and attention parameters stay frozen. The evaluation covers Qwen3.6-35B-A3B and Gemma 4-26B-A4B-it, using math, coding, and tool-use benchmarks.

Why This Matters

For MoE serving, reducing the total expert pool can lower the amount of expert-weight data that must be accessed during decoding. SlimWise’s phase-specific design aims to get that benefit without applying the same pruning to prompt processing, where the paper finds little throughput advantage from pruning.

Potential applications include:

• Reasoning assistants: Serving models that produce long math or analytical responses, where both accuracy and generation length affect cost and user experience.
• Code-generation systems: Handling coding prompts while checking not only benchmark accuracy but also whether responses are prematurely shortened.
• Tool-using assistants: Serving models that generate function calls, an area evaluated in the paper with BFCLv4.
• High-throughput language-model APIs: Using decode-side pruning to improve throughput under per-user generation-speed targets.

For industry, the work is relevant because it addresses serving efficiency in vLLM and supports both disaggregated and colocated deployments. Its results also highlight a practical evaluation concern: accuracy alone may not reveal changes in output length, which can affect latency and compute use.

Future Directions

• Test more models and serving workloads. The evaluation covers two MoE backbones; further work could establish how broadly the accuracy, output-length, and throughput results generalize.
• Improve generation-length control. Distillation reduces some excessive or shortened generations, but some output-length percentiles remain above the baseline.
• Explore pruning choices jointly with phase-specific serving. The paper compares three pruning criteria; further research could investigate how expert selection and retained-expert counts should vary with task and workload.
• Measure end-to-end gains across different request mixes. Prefill cost remains unchanged, so total serving speedup depends on how much time a workload spends in decode. Evaluating diverse prompt and response lengths could clarify when the approach is most beneficial.

Target Audience

This paper is most useful to researchers and engineers working on MoE model compression, inference optimization, and LLM serving systems. It may also interest practitioners deploying reasoning, coding, or tool-using models who need to balance throughput, accuracy, memory use, and response length.

Authors’ abstract

Mixture-of-experts (MoE) models activate few experts per token, yet batched decoding can access nearly the entire expert pool, making expert-weight traffic a major bottleneck. Expert pruning reduces this traffic, but conventional approaches also prune compute-bound prefill, sacrificing model quality for little throughput benefit. We present SlimWise, a serving framework that tailors the expert pool to each inference phase. SlimWise performs prefill with the full model and decode with a pruned model that directly reuses the prefill-generated KV cache without conversion. Across two MoE backbones and three pruning criteria, this training-free KV cache handoff substantially narrows accuracy gaps relative to the full model in many settings. We also show that benchmark accuracy can conceal substantial pruning-induced changes in generation length. To address these distortions and residual accuracy loss, SlimWise introduces a low-cost distillation stage that trains the decoder to continue from full-model KV caches while updating only a small subset of parameters. Implemented in vLLM, SlimWise supports both prefill-decode (PD) disaggregation and PD-colocated serving. On Qwen3.6-35B-A3B, SlimWise improves decode throughput by up to 1.81x at 50% expert pruning with minimal accuracy loss.