Gather + GEMV Strategy
Matrix multiplication optimization for MoE inference replacing standard grouped GEMM operations.
The gather + GEMV (matrix-vector multiply) strategy is a kernel-level optimization for MoE (mixture-of-experts) inference that replaces the standard grouped general matrix-vector multiplication (grouped GEMM) approach with a two-step gather-then-multiply pattern. In standard MoE inference, each token is routed to a subset of experts, and the active parameters for that token are gathered from a larger weight matrix — typically implemented as a single grouped GEMM that performs all the expert computations in one optimized batched operation. The gather + GEMV strategy separates this into an explicit gather step (collecting the relevant weight shards for each token) followed by individual GEMV operations.
The motivation for this separation is latency and efficiency at the specific shapes encountered in bidirectional serving. Standard grouped GEMM implementations are optimized for the common case of large batch sizes and relatively uniform expert selection. In an interaction model's bidirectional serving scenario, the batch sizes are smaller and more variable (due to the streaming session design), and the attention to latency per token is higher. The gather + GEMV strategy allows each token's expert computation to be launched independently, with lower overhead per launch — at the cost of less efficient utilization of GPU compute resources when batch sizes are large.
The approach was identified through prior work from PyTorch and Cursor on MoE kernel optimization. In the interaction model context, it enables the per-200ms-chunk latency requirements to be met despite the overhead of frequent prefill and decode operations. The key insight is that the optimal kernel strategy depends on the serving configuration: what is optimal for large offline batches is not optimal for small real-time streams, and vice versa. Interaction model serving requires a different optimization profile than most existing LLM serving workloads.
The tradeoff is kernel complexity and portability. Gather + GEMV requires explicit management of gather operations, which introduces additional synchronization points and memory movement compared to the fused grouped GEMM. The optimization is also hardware-specific — it leverages particular memory bandwidth and launch latency characteristics of the target GPU architecture. Porting to different hardware requires re-tuning these tradeoffs.