PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Gunho Park, Kyoungho Jeun, Juntaek Oh, Byeongjun Shin, Baeseong Park, Minsoo Rhu

arXiv:2609.34117v1cs.LGcs.AI

TL;DR

Conventional expert pruning can sacrifice model quality during compute-bound prefill for little throughput benefit, while batched decoding makes expert-weight traffic a bottleneck. SlimWise uses full-model prefill and pruned decode with a shared KV cache; at 50% expert pruning, it improves decode throughput by up to 1.81× with minimal accuracy loss.

  • Problem

    Conventional pruning applies the same reduced expert pool to prefill and decode, sacrificing model quality during compute-bound prefill for little throughput benefit.

    Source: §1 · p. 2

  • Method

    SlimWise combines full-model prefill with a pruned decoder that directly reuses the prefill-generated KV cache, with low-cost distillation to further improve the decoder.

    Source: §1 · pp. 2–3

  • Results

    At 50% expert pruning, SlimWise achieves up to 1.81× decode throughput with minimal accuracy loss.

    Source: §1 · p. 3

  • Takeaways & Limitations

    The KV cache handoff recovers much of the accuracy lost to pruning in many evaluated settings without additional training.

    Source: §5 · p. 9

    End-to-end speedup depends on the fraction of serving time spent in decode, and the KV cache handoff can leave excessive generation unresolved in some settings.

    Source: §3.2 · p. 6, §5 · p. 9

Abstract

from arXiv · show

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.

1 INTRODUCTION

Batched MoE decoding can load weights for nearly the full expert pool, while conventional pruning also reduces experts during compute-bound prefill for little throughput benefit. SlimWise keeps full-model prefill and prunes decode, reusing the prefill KV cache; at 50% expert pruning it improves decode throughput by up to 1.81× with minimal accuracy loss.

1 INTRODUCTION

  • At 50% expert pruning, SlimWise improves decode throughput by up to 1.81× with minimal accuracy loss.The comparison is under PD-disaggregated serving and a 50 tokens/s per-user generation target.Source: Figure 1 · p. 2, §1 · p. 3
  • Benchmark accuracy can conceal pruning-induced changes in generation behavior, including shortened reasoning on some tasks and excessively long traces on others.The handoff alleviates some generation-length distortions but does not fully resolve excessive generation.Source: §1 · p. 2
  • Full-model prefill followed by pruned decode avoids pruning the compute-bound phase that gains little throughput from a smaller expert pool.The pruned decoder directly reuses the prefill-generated KV cache; expert pruning preserves the attention architecture and cache format, so no conversion is needed.Source: §1 · pp. 1–2
  • The KV cache handoff recovers a substantial portion of pruning-induced accuracy loss without additional training across multiple evaluated settings.The gains occur across three pruning criteria without changing the retained expert sets.Source: §1 · p. 2
  • Low-cost distillation trains the pruned decoder to continue from full-model KV caches, improving accuracy and bringing generation lengths closer to the full model.The distillation stage updates only a small subset of parameters.Source: §1 · pp. 2–3
  • SlimWise supports PD-disaggregated and PD-colocated serving, using phase-aware router masking for colocated execution.In PD disaggregation, the decode instance stores only retained expert weights, freeing memory for additional KV cache capacity and larger batches.Source: §1 · pp. 2–3

2 BACKGROUND

Batched decoding can access nearly the full expert pool, leaving decode dominated by expert-weight movement, while prefill has greater weight reuse and little throughput benefit from pruning. The background explains this phase asymmetry, the shared KV-cache interface, and the pruning criteria used.

2.1 BATCHED MOE INFERENCE

  • As batch size grows, the expected number of distinct experts accessed approaches the full pool, despite each token activating only k experts.Qwen3.6-35B-A3B activates 8 of 256 experts per token, versus 8 of 128 in Qwen3-30B-A3B, halving k/m despite comparable total parameter counts.Source: §2.1 · p. 3
  • With skewed routing, saturation is delayed but not prevented: batch size 64 accesses approximately 75% of the 60 GiB expert pool, and the largest evaluated batch accesses nearly all of it.At saturation, retaining m′ of m experts reduces accessed expert-weight volume to approximately m′/m of the original pool.Source: §2.1 · p. 3
  • Larger batches increase weight reuse, but KV-cache capacity and per-token latency constrain concurrency, leaving decode dominated by expert-weight movement.As the accessed set approaches the full pool, each expert processes fewer than ten tokens per step even at the largest evaluated batch size.Source: §2.1 · p. 3
  • Across evaluated pruning ratios, prefill throughput stays within a few percent of the unpruned model with no consistent improvement, while decode throughput improves by up to approximately 1.5×.These measurements motivate allocating expert capacity separately to prefill and decode.Source: §2.1 · p. 3, Figure 2 · p. 4

2.2 PREFILL–DECODE DISAGGREGATION

  • Prefill and decode can communicate through the same KV-cache interface because pruning preserves the attention architecture and cache format, allowing direct handoff without conversion.The cache carries keys and values for standard attention and recurrent states for linear attention in hybrid models such as Qwen3.6.Source: §2.2 · p. 4

2.3 REDUCING EXPERT TRAFFIC

  • The evaluated pruning criteria are routing mass, Expert Activation Norm, and Router-weighted Expert Activation Pruning; each selects a fixed retained set while keeping k unchanged.The criteria are computed from calibration rollouts and rank experts within each layer.Source: §2.3 · p. 4

3 METHODOLOGY

SlimWise separates expert capacity by inference phase: full-model prefill supplies a KV cache directly reused by a pruned decoder. The section evaluates this handoff’s effects on accuracy and generation length, then describes low-cost distillation to address residual loss and distortions.

3.1 SLIMWISE SERVING FRAMEWORK

  • SlimWise runs prefill with the full expert pool and decode with a ranked, retained subset, whose decoder directly reuses the prefill KV cache without conversion.A separate pruned checkpoint or router masking can implement decode; phase-aware masking supports both PD-disaggregated and PD-colocated serving.Source: §3.1 · pp. 4–5

3.2 BENEFITS AND LIMITATIONS OF SLIMWISE KV CACHE HANDOFF

  • Across three pruning criteria, full-model prefill substantially narrows the accuracy gap to the full model in several affected settings without changing retained experts.The gains use the same pruned decoder and require no additional training.Source: §3.2 · p. 5
  • With REAP pruning, conventional pruning cuts HumanEval+ median output length from 2.3k to 0.3k tokens and raises MATH-500 median length from 2.4k to 6.6k while accuracy stays close to baseline.SlimWise restores HumanEval+ median length to approximately 2.3k tokens, but MATH-500 remains approximately 5.8k, more than twice baseline; on MBPP+, handoff further increases the elevated median.Source: §3.2 · pp. 5–6
  • The KV-cache handoff alleviates but does not eliminate generation distortions, motivating distillation; under greedy decoding, generations can reach the token budget without naturally terminating.The reported non-termination occurs both with and without the handoff.Source: §3.2 · p. 6

3.3 RECOVERING THE REMAINING LOSS WITH A LOW-COST DISTILLATION

  • Distillation trains the pruned decoder to continue from teacher-generated prefix KV caches, using uniformly sampled sequence split points to expose varied prefix lengths.The student continues from the teacher prefix cache; conventional distillation is the special case where the split point is zero.Source: §3.3 · p. 6
  • Training combines cross-entropy and teacher-distribution distillation with α = 0.1, using the teacher’s top 64 tokens plus a residual probability bucket.The residual bucket aggregates the probability mass of all tokens outside the teacher’s top-64 set.Source: §3.3 · p. 6
  • Distillation updates only always-on components, totaling 0.42% of Qwen3.6-35B-A3B parameters and 2.12% of Gemma 4-26B-A4B parameters, while routed experts and attention remain frozen.Frozen tensors are shared between teacher and student, and router learning rate is one-tenth that of other trainable parameters.Source: §3.3 · p. 6

4 EVALUATION RESULTS

Across two model backbones and three pruning criteria, full-model prefill recovers accuracy in many settings. SlimWise also improves decode throughput under both PD-disaggregated and PD-colocated serving.

4.1 EXPERIMENTAL SETUP

  • Evaluation spans Qwen3.6 and Gemma 4 math, coding, and tool-use benchmarks, with mean accuracy over three sampled-decoding runs and a 32,768-token response budget.Distillation uses 64k sequences of at most 12k tokens and takes approximately six hours per run on four B200 GPUs.Source: §4.1 · p. 7

4.2 ACCURACY EVALUATION

  • Across two backbones and three pruning criteria, full-model prefill recovers accuracy in many settings, especially Gemma 4 tool use under REAP pruning.The comparison covers accuracy and median output length at 50% pruning.Source: §4.2 · pp. 7–8, Table 2 · p. 8
  • Under EAN pruning, KV-cache handoff leaves median output length at the generation budget on math and coding benchmarks; distillation improves accuracy and brings lengths closer to baseline.Without distillation, SlimWise accuracy on GSM8K and MATH-500 falls below conventional pruning in these settings.Source: §4.2 · p. 8
  • At m′=64, distillation improves accuracy on all four Qwen3.6 benchmarks and cuts MATH-500 p90 from 32.8k to 14.0k tokens, still above the baseline’s 6.1k.At m′=192, conventional pruning largely preserves accuracy, but SlimWise and distillation reduce MBPP+ p90 from 8.1k to 6.7k and 4.2k tokens, respectively.Source: §4.2 · p. 8, Table 2 · p. 8
  • At 75% pruning, inherited teacher caches improve conventionally distilled accuracy on all four benchmarks; cache-based distillation further improves three, while MBPP+ trades 0.929 accuracy for 3.5k-token p90, from 8.1k.The same-cache handoff also reduces p90 on GSM8K and MATH-500 compared with the conventionally distilled student.Source: §4.2 · p. 8, Table 2 · p. 8

4.3 THROUGHPUT EVALUATION

  • At 50% pruning, SlimWise delivers 1.40–1.81× full-model decode throughput under PD disaggregation and 1.36–1.72× under PD-colocation.The gains are larger at 75% pruning: 1.65–2.39× under disaggregation and 1.57–2.28× under colocation.Source: Figure 4 · p. 9, §4.3 · p. 9

5 CONCLUSION AND LIMITATIONS

SlimWise pairs full-model prefill with pruned decode, using KV-cache handoff and low-cost distillation to recover accuracy and mitigate generation-length distortions. Its throughput gains and memory effects depend on the serving configuration and decode share.

5 CONCLUSION AND LIMITATIONS

  • At 50% expert pruning, SlimWise improves decode throughput by 1.4–1.8× while leaving prefill throughput unchanged.The low-cost distillation stage further narrows the accuracy gap and mitigates generation-length distortions.Source: §5 · p. 9
  • End-to-end speedup depends on the fraction of serving time spent in decode; PD-colocated serving retains the full expert pool and does not reduce model-weight memory..Source: §5 · p. 9

A.1 EXPERIMENTAL DETAILS

The appendix details the expert-ranking, benchmark, and distillation protocols, alongside the greedy-decoding evaluation setup.

A.1 EXPERIMENTAL DETAILS

  • Expert statistics come from one pass over the REAP calibration mixture; routing-mass, EAN, and REAP criteria retain the top m′ experts per layer while keeping k unchanged.The mixture contains six components of 4,096 packed sequences, each 2,048 tokens long.Source: §A.1 · p. 12
  • Math, coding, and tool-use benchmarks use task-specific prompting and evaluation protocols, and each score is averaged over three runs.HumanEval+ and MBPP+ are zero-shot; GSM8K uses five-shot prompting on 500 problems, MATH-500 four-shot prompting, and BFCLv4 native function calling.Source: §A.1 · p. 12
  • Distillation trains on 64k sequences for 8,000 steps, updating selected expert or MLP parameters, routers, and normalization layers.Training totals approximately 282M tokens for Qwen3.6 and 250M for Gemma 4; each run takes about six hours on four B200 GPUs.Source: §A.1 · pp. 12–13
  • Figure 5 evaluates 50% pruning with greedy decoding and thinking enabled, comparing accuracy across three criteria and REAP output-length distributions with and without distillation.Its second panel reports p50, p90, and accuracy against full-model accuracy.Source: Figure 5 · p. 13

A.2 SLIMWISE WITH GREEDY DECODING

Under greedy decoding, the generation-length distortions persist and can become more pronounced; SlimWise and distillation affect accuracy and output length differently across tasks.

A.2 SLIMWISE WITH GREEDY DECODING

  • On HumanEval+ under REAP pruning, SlimWise restores accuracy from conventional pruning’s 0.746 to 0.915 and median output length from 0.3k to 2.4k tokens.On MBPP+ and MATH-500, p90 reaches the generation budget with or without cache handoff; distillation brings the distributions closer to baseline and below budget.Source: §A.2 · p. 13

A.3 EFFECTS OF EXPERT PRUNING IN PREFILL AND DECODE

The section compares four prefill–decode configurations under REAP pruning on Qwen3.6-35B-A3B. Which phase matters more varies by task: full-model prefill helps on HumanEval+, while full-model decode helps on math benchmarks.

A.3 EFFECTS OF EXPERT PRUNING IN PREFILL AND DECODE

  • At 50% pruning, Reverse brings MBPP+ accuracy and median output length close to baseline, showing that coding-task results also vary by phase..Source: §A.3 · p. 14
  • On GSM8K and MATH-500, Reverse stays near baseline accuracy at 50% pruning and outperforms training-free SlimWise on both at 75%.At 75% pruning on MATH-500, Reverse’s median output is 2.2k tokens versus SlimWise’s 4.7k; distillation cuts SlimWise’s 90th-percentile length to 14.0k, still above baseline’s 6.1k.Source: §A.3 · p. 14
  • SlimWise distills the pruned decoder because that targets decode’s expert-weight traffic, whereas pruning compute-bound prefill offers little throughput benefit.Under PD disaggregation, decoder distillation also complements increased KV cache capacity.Source: §A.3 · p. 14
  • 0.929 HumanEval+ accuracy and 2.3k median output length: SlimWise at 50% pruning, versus Reverse’s 0.890 accuracy and 0.3k median length.SlimWise restores HumanEval+ median length to the baseline’s 2.3k tokens; at 75% pruning it also exceeds Reverse on both coding benchmarks, though neither fully recovers baseline accuracy.Source: §A.3 · p. 13, Table 3 · p. 14

Related papers

Loading 2609.34117v1…