LLM 推論系統
RMM Prunes the Matrix Multiplication Contraction Axis Based on Input, Accelerating Llama 3.1 8B Long-Context Inference by Up to 1.40×
Reduced Matrix Multiplication leaves weights unchanged and, based on the column norms of the current activation, computes only the highest-scoring shared dimensions in a matrix multiplication. A100 tests show that gains increase with sequence length, but for short contexts, selection and kernel-launch overhead almost entirely offset the savings.
The newly released [Reduced Matrix Multiplication (RMM) paper](https://arxiv.org/abs/2608.13426) shifts the unit of approximation in Transformer inference—from tokens, neurons, or fixed subnetworks—down to the “contraction axis” of each matrix multiplication. For `Y=AB`, the system computes the L2 norm of each column in the current activation `A`, retains the Top-K dimensions specified by retention ratio `ρ`, and then executes the reduced multiplication `A[:,I]B[I,:]`. The same rule can be applied to the head dimension in QKᵀ, the token dimension when multiplying attention weights by V, and linear projections. Selections vary with the input, layer, head, and decoding step, requiring neither retraining nor permanent weight removal.
The authors evaluated the method on language models ranging from 1B to 70B parameters, as well as Qwen2.5-VL-7B. For Llama 3.1 8B on an A100 with batch size 1 and `ρ=0.8`, end-to-end latency at 1,024 tokens falls only from 109.39 ms to 103.91 ms, a 1.05× speedup. At 2,048 and 4,096 tokens, the speedups reach 1.27× and 1.40×, respectively. For the 70B model at 2,048 tokens, the speedup is 1.41×. At 4,096 tokens, the dense version runs out of memory while RMM completes inference, though this does not mean the model weights have been compressed.
Quality is highly sensitive to where pruning is applied. For the 8B model at `ρ=0.8`, CNN/DailyMail ROUGE scores remain nearly unchanged. At `ρ=0.5`, however, average accuracy across five QA tasks falls from 69.8 for the full model to 59.8. In ARC-Easy experiments at `ρ=0.7`, pruning only the attention components reduces accuracy by 3.52 points, whereas pruning the entire MLP causes an 18.78-point drop, supporting the use of component-specific retention ratios. The deployment results should still be treated cautiously: latency is averaged over only ten runs; the dense baseline uses Hugging Face SDPA, while RMM uses custom Triton kernels; and although the [OpenReview version](https://openreview.net/forum?id=2uxuiykvA4) documents an earlier form of the method, the GitHub repository listed in the paper was still returning a 404 at publication time. The key question is whether RMM can be integrated with vLLM, quantized kernels, continuous batching, and high-concurrency serving—not merely whether it works for long prefills at batch size 1.