Procedural weight generation boosts sparse expert model inference speed
FORGE: Procedural Weight Generation for Sparse MoE Inference
Hardware ArchitectureDistributed, Parallel, and Cluster ComputingMachine Learning
Summary
Large AI models with many specialized parts, called experts, normally store huge amounts of information that makes them slow to run. The authors show how to recreate this expert information on-the-fly using very compact math settings, saving a lot of memory and speeding up the process. This approach trades a bit of accuracy for much faster and more memory-efficient performance. The system they built, called the Rotor Generation Unit, is a small chip addition that creates these weights quickly without slowing down the main computing units.
What this means in practice
- •For ai infrastructure engineers: Increase throughput and reduce latency for sparse mixture-of-experts model inference on GPUs by regenerating expert weights on-chip instead of loading them from memory.
- •For distributed systems engineers: Reduce memory footprint and network traffic for multi-node expert parallelism in large AI models by compressing expert weight storage into small procedural seeds.
- •For graphics processor designers: Integrate a compact Rotor Generation Unit in GPU streaming multiprocessors to efficiently generate MoE expert weights without impacting tensor core utilization.$Commercial implications: Makes possible specialized GPU chips optimized for faster and more memory-efficient sparse expert model inference, improving AI workload performance.
Authors
Chanyoung Moon, Travis Pence, Daisuke Yamada, Jiaqi Mo, Yanxin Zhang, Vikas Singh, Karthikeyan Sankaralingam
Abstract
We introduce FORGE (Fast On-chip Rotor Generation for Experts), a procedural weight generation system for Sparse Mixture-of-Experts (MoE) inference. Rather than storing and loading every expert weight tensor from HBM, we show how a compact accelerator inside each Streaming Multiprocessor can regenerate them on-the-fly from $O(\log^2 d)$ rotor parameters per projection. Per replaced layer, gigabytes of stored expert weights are approximated with under a megabyte of seeds. This addresses the dual memory bottleneck of autoregressive Sparse MoE decode, where per-expert FFN operations are HBM-bandwidth-bound at small routed-token counts and the aggregate expert footprint forces multi-node Expert Parallelism (EP) with all-to-all traffic. We realize the approach as the Rotor Generation Unit (RGU), a fixed-function datapath that generates FFN weight tiles without competing with Tensor Core execution. With 30% of the MoE FFN layers replaced across five Sparse MoE models, we estimate 1.3-2.8x higher system throughput at matched GPU count and 1.10-1.15x lower per-user decode latency at short context. In exchange, downstream accuracy on three model families falls below that of the same model fine-tuned without rotor replacement, and the replacement ratio controls this tradeoff. FORGE also expands KV cache capacity by 1.2-2.6x. It is agnostic to the bit width at which weights are stored, so it can compose with quantization. The RGU is fully implemented/verified in SystemVerilog at ~0.9% SM area overhead (0.032 mm^2/SM, 7nm).