BR-MoE: A Memory-Budget-Aware Co-Design Framework for Efficient MoE Inference on Resource-Constrained GPUs
The Mixture-of-Experts (MoE) architecture has become the dominant paradigm for scaling Large Language Models, yet its explosive parameter growth routinely overflows the GPU memory budget. Weight quantization offers a practical way to fit these models into limited GPU memory, but reducing precision below 4 bits causes severe accuracy loss. Existing approaches, mixed-precision quantization and low-rank compensation, have been treated as independent strategies, overlooking the fact that both draw from the same hardware memory budget. Consequently, optimizing either dimension in isolation systematically overshoots on some experts while starving others, leaving the achievable memory–accuracy frontier well below what a joint allocation could reach. We present BR-MoE, a memory-budget-aware co-design framework that treats per-expert bit-width and compensator rank as joint variables under a single memory constraint. Guided by a lightweight layer-wise sensitivity proxy, BR-MoE casts the global allocation task as a Multiple-Choice Knapsack Problem (MCKP) and solves it to optimality via Integer Linear Programming (ILP). To turn memory savings into real inference speedups, BR-MoE further integrates a fused grouped expert GEMM backend that consolidates fragmented per-expert computations into bandwidth-efficient kernels. Evaluated across multiple MoE models on NVIDIA A100 GPUs, BR-MoE establishes a superior memory–accuracy trade-off over state-of-the-art quantization methods, delivers substantial end-to-end inference speedups, and re-derives an optimal allocation for any new memory budget in seconds, making it a deploy-time tunable system rather than a one-shot compression recipe.