Grouped-query attention and mixture of experts

Grouped-query attention (GQA)

Module 2 computed that the key-value cache for 4,096 tokens takes 2 GiB in a 32-layer model with 32 heads. At 32,768 tokens that is 16 GiB per conversation, which limits how many users one GPU can serve. The observation behind GQA: the query heads need to be many, but the key and value heads can be shared by groups of query heads.

The code gives both effects (real output; model width 4,096 and 32 query heads):

A 4x smaller cache means roughly 4x more simultaneous long conversations per GPU, which feeds directly into your price per token (module 6). In published comparisons the quality cost is small, which is why many recent models use grouped-query attention. (I did not measure quality here; the code has random weights.)

Mixture of experts (MoE)

A dense model uses all its parameters for every token. A mixture of experts replaces the MLP in a block with many parallel MLPs (the experts) and a small router that picks a few of them for each token.

def moe_layer(x, router_w, experts, top_k=2):
    probs = softmax(x @ router_w)                   # how much each expert suits each token
    chosen = np.argsort(-probs, axis=1)[:, :top_k]  # keep the best top_k
    # mix the chosen experts' outputs by their renormalised router weights

Consequences, each tested:

• Each token goes to exactly top_k distinct experts, and the mixing weights sum to 1.
• With identical experts, the layer equals a single expert, which shows the mixing is correct.
• Total parameters are large (all experts), but active parameters per token are small (top_k of them). With 2 of 8 experts active, a quarter of the expert parameters are used per token. You get the capacity of a big model at closer to the speed of a small one.
• Load balance matters. If the router sends most tokens to a few experts, the rest are wasted and those few become a bottleneck. The demo prints each expert's share of tokens (real output with random weights and 8 experts): [0.1, 0.12, 0.18, 0.18, 0.1, 0.09, 0.09, 0.15]. A perfectly even split would be 0.125 each. Real training adds a penalty that pushes the router towards even use.
• Memory is still the full parameter count: you must hold all the experts even if each token uses two.