MoE 提速新招:不用校准数据,直接挑通道
大模型里的专家混合(MoE)结构,推理时要把所有专家权重都读进内存,这是速度瓶颈。以前的做法是设一个阈值,只读权重高的通道,但阈值要拿离线数据校准,而且实际读多少并不受控——论文里校准成 18% 的读取预算,实际平均只读了 12.9%。这篇直接改成 top-k:每个 token 只读得分最高的 k 个通道,不需要任何校准数据,在 43 组对比里 40 组不输阈值版,还避免了校准阈值在 Gemma 4 上引发的重复循环崩溃。另外他们还用模型自身的权重构造了一个读取器,用三元草图筛通道、低秩映射补回跳过的部分,在三个模型上只读 17.8% 的专家权重就保住了 98.1% 的准确率,推理速度比优化过的 vLLM 快 1.6 到 1.9 倍。这不是你明天能直接用的东西,但它说明 MoE 提速可以甩掉校准数据这个包袱,靠更简单的规则和模型自身结构就能做到。
📄 原文摘要(英文)
Mixture-of-Experts (MoE) models are bottlenecked by loading routed-expert weights from memory during interactive decoding at small batch sizes. Because most intermediate channels in a routed expert contribute little to any single token, channel selectors predict which channels matter and read only their weights, typically keeping channels above a threshold calibrated on offline data. We show that channel selection needs no calibration data. First, a calibrated threshold does not fix how much a selector reads, because the number of channels passing it varies across inputs; Prox with thresholds calibrated for an 18% read budget actually reads only 12.9% on average across our benchmarks. Second, top-k keeps each token's k highest-scoring channels and needs no calibration data. At the calibrated budget, a selector's top-k version matches or beats its threshold version on the same score in 40 of 43 comparisons at 13% and 18% budgets, including all 23 on generation tasks. The top-k versions also avoid the repetition-loop collapse that calibrated thresholds cause on Gemma 4. Calibrating a threshold thus provides neither of the two things selectors use it for: a known read budget and a consistent accuracy gain at that budget. We furthermore introduce a reader built entirely from the model's weights, which screens channels with a ternary sketch and recovers the skipped ones with a low-rank map. Across three post-trained MoE models, our reader retains 98.1% of full-model accuracy at 17.8% of routed-expert reads, with Triton kernels running MoE blocks 1.60 to 1.90x faster than optimized vLLM at batch one and up to 2.69x at larger batches on the same three models.