Sparse MQA + Fused MoE Mega Kernel + Hyper-connections: The Trifecta Behind Modern Frontier Models

MHA → GQA → MQA: The KV Head Reduction Story

Multi-Head Attention allocates separate K and V projection matrices per head. For sequence length S, batch B, H heads, and head dimension D, the KV-cache size is:

# MHA: H separate K/V heads
KV_size = 2 × B × S × H × D × dtype_bytes
# 70B model: B=1, S=1M, H=64, D=128, FP16 (2 bytes)
# KV_size ≈ 2 × 1 × 1M × 64 × 128 × 2 = 32 GB per layer
# × 80 layers = 2.56 TB total — far beyond H200's 141 GB HBM

# GQA (Grouped Query Attention): G groups, G << H
KV_size = 2 × B × S × G × D × dtype_bytes
# Llama-3.1 405B: G=8 heads → 8x KV reduction vs MHA

# MQA (Multi-Query): single shared K/V head
KV_size = 2 × B × S × 1 × D × dtype_bytes
# At H=64: 64x KV reduction vs MHA. Quality cost ~1–2% on benchmarks.

The reduction is substantial: GQA with 8 KV heads (as used in Llama-3 models) achieves 8× KV compression. Full MQA achieves 64× at H=64, enabling the same GPU memory to serve far longer contexts or larger batches. The quality penalty is modest for most tasks but measurable on long-document reasoning benchmarks.

Adding Sparsity: Local+Sink and A² Routing

Pure MQA still scales linearly with context length — it just scales with a smaller constant. For million-token contexts, even MQA's single KV head can overflow HBM. Attention sparsity cuts the growth from O(S) to constant or O(S·K/S) = O(K). Two dominant patterns:

Combined Sparse MQA — single KV head with local+sink sparsity pattern — is the mechanism by which models in the Gemini 1.5 class achieve million-token contexts without KV cache exceeding the GPU memory budget. A 2M-token context with this approach can fit in under 500 GB of KV memory vs. the multi-terabyte requirement of naive MHA.

Active vs Total Parameters: Why MoE Changes the Cost Function

Mixture-of-Experts replaces the FFN block in each transformer layer with N expert networks. A learned router selects top-k experts per token (typically k=2). Each token only passes through k experts, regardless of N.

The economics are stark: DeepSeek-V3 delivers 671B total parameters — with access to the full representational capacity of a frontier-class model — while activating only 37B parameters per token. Serving cost scales with active parameters, not total. Inference pricing should be quoted against active parameters.

Why Naive MoE Is Slow: The 7-Kernel Problem

A straightforward PyTorch loop launches seven separate CUDA kernels per layer and streams expert weights from HBM multiple times:

# Naive: 7 kernel launches, 7 HBM round-trips per layer
logits = x @ W_gate           # kernel 1: gate GEMM
topk_idx = topk(logits, k=2)  # kernel 2: sort + select
for i in range(k):
    x_i  = gather(x, topk_idx[i])   # kernel 3, 5: irregular gather
    y_i  = expert_ffn_i(x_i)        # kernel 4, 6: expert GEMM (HBM load)
x_out = scatter_add(y_i, weights)   # kernel 7: weighted combine

# Problems:
# - Each kernel launch: ~5–10µs overhead
# - Expert weights reloaded from HBM per kernel: no L2 residency
# - Dispatch imbalance causes tail latency: slow experts block fast ones
# - At B=256, S=2048: ~30% slower than equivalent dense layer

Mega Kernel Fusion: Three Key Tricks

A Mega Kernel fuses all seven operations into a single persistent CUDA kernel that launches once and handles all experts. Three tricks make it work:

Combined effect: 2.1× throughput vs naive MoE, and 1.3× vs dense at equal active parameter counts. MoE is no longer just cheaper than dense — it is faster. TRT-LLM and MegaBlocks implement production variants of this.

Why Standard Residuals Fail at Depth

Pre-LN transformers use residual connections of the form x_{l+1} = x_l + F_l(LN(x_l)). The gradient of the loss with respect to the input at layer l=0 is the product of Jacobians across all subsequent layers:

∂L/∂x_0 = ∏_{l=0}^{L-1} (I + ∂F_l/∂x_l)

# If ||∂F_l/∂x_l|| < 1 consistently: gradient → 0  (vanishing)
# If ||∂F_l/∂x_l|| > 1 consistently: gradient → ∞  (exploding)
# At L=120: even a mean factor of 0.99 per layer gives 0.99^120 = 0.30
# → 70% signal attenuation to layer 0. Training converges poorly.

Empirically, standard Pre-LN transformers above approximately 80 layers for models above 10B parameters show training instability, slower convergence, and diminishing returns from additional depth. This caps the practical depth of naive architectures well below what is theoretically desirable.

Hyper-connections: Learned Gradient Highways

Hyper-connections (DeepSeek AI, arXiv:2409.19606, September 2024) generalize the residual connection by introducing learned per-layer weight vectors that allow each layer to attend to multiple prior representations, not just the immediately preceding one. The core formulation:

# Standard Pre-LN residual
x_{l+1} = x_l + F_l(LN(x_l))

# Hyper-connection (simplified scalar form)
# alpha, beta are learned parameters — initialised near 1/L for stability
x_2 = alpha_{2,0} * x_0 + alpha_{2,1} * x_1 + F_2(LN(x_1))
x_3 = alpha_{3,0} * x_0 + alpha_{3,1} * x_1 + alpha_{3,2} * x_2 + F_3(LN(x_2))

# Full form: connection width n, expansion factor r
# Each layer l receives a weighted sum of n prior hidden states
# The weight matrix W_c (n×n) is learned; initial value = I/n
h_l = W_c[l] @ [x_{l-n}, ..., x_{l-1}]   # blend prior n states
x_l = h_l + F_l(LN(h_l))

The mechanism provides three guarantees:

• → Direct gradient path to all layers. Because layer 120's output depends directly on layer 0's output through the learned alpha terms, the gradient can reach layer 0 in a single step regardless of depth. Vanishing gradients across the full depth are structurally impossible.
• → Feature reuse across depth. Layer 80 can directly access token embeddings from layer 0 when the learned alpha weights place significant mass on early-layer representations. This is architecturally equivalent to the DenseNet insight applied to transformers.
• → Learned layer skipping. If a layer's output is destabilising (e.g. during training), the model can learn to down-weight that layer's contribution through its alpha parameters while maintaining gradient flow through other paths. This is automatic regularisation of depth.

The DeepSeek hyper-connections paper reports zero training divergence events at 100+ layers with hyper-connections vs approximately 40% divergence rate for equivalent Pre-LN architectures. The parameter overhead is small: for connection width n=4 and model width d, the additional parameters per layer are n × d — roughly 0.3% of total parameter count for typical models.

TMA and Async Copies: Hiding Dispatch Latency

The fundamental bottleneck in MoE is not arithmetic — it is the irregular memory access pattern of token dispatch. Token 5 routes to Expert 3, token 6 to Expert 17, token 7 back to Expert 3: the memory access pattern is unpredictable and cannot be coalesced into efficient sequential HBM reads. This creates latency rather than bandwidth pressure.

H100's Tensor Memory Accelerator (TMA) provides cp.async.bulk.tensor — an asynchronous bulk copy instruction that transfers tensor tiles from HBM to shared memory at up to 512 bytes/cycle (vs 128 bytes/cycle for standard cp.async). Critically, TMA operations proceed concurrently with tensor core execution. The Mega Kernel uses a software pipeline to overlap Expert N's compute with Expert N+1's weight load:

// CUDA C++ — Mega Kernel TMA pipeline (simplified)
__shared__ half smem_weights_A[EXPERT_TILE_SIZE];
__shared__ half smem_weights_B[EXPERT_TILE_SIZE];  // double buffer

cuda::pipeline pipe = cuda::make_pipeline();

// Stage 0: start loading expert[0] weights asynchronously
cuda::memcpy_async(smem_weights_A, &gmem_expert[0],
                   CUTE_TMA_LOAD, pipe.producer_acquire());
pipe.producer_commit();

for (int expert = 0; expert < n_active_experts; ++expert) {
  // While executing expert[expert-1]...
  // ...prefetch expert[expert] weights in background
  if (expert + 1 < n_active_experts) {
    auto& next_buf = (expert % 2 == 0) ? smem_weights_B : smem_weights_A;
    cuda::memcpy_async(next_buf, &gmem_expert[expert+1],
                       CUTE_TMA_LOAD, pipe.producer_acquire());
    pipe.producer_commit();
  }

  pipe.consumer_wait();          // wait for this expert's weights
  __syncthreads();               // SMEM coherence

  // Tensor core GEMM — expert weights already in SMEM
  compute_expert_gemm(smem_weights_current, token_tiles);
}
// Result: HBM load latency (~200 cycles) fully hidden behind 256-token GEMM

The net effect: expert weight load from HBM is completely hidden behind tensor core arithmetic. The kernel is now compute-bound, not memory-bandwidth-bound. This is the primary reason H100 MoE throughput is approximately 2× that of A100 for the same model — not clock frequency or FLOP count, but the TMA enabling this overlap.

FP8 E4M3 vs E5M2: Right Format for the Right Tensor

FP8 is not a single format — it is two distinct 8-bit floating-point formats with different precision/range tradeoffs. Using the wrong format for the wrong tensor is a silent quality failure.

H100 tensor cores support mixed FP8 formats natively: the input operands can be E4M3 and E5M2 independently, with FP32 accumulation. No explicit dequantization kernel is required.

FP4 Block-Scaling: The B200 Kernel Rewrite

Blackwell B200 adds FP4 precision (NVFP4): 4-bit values with one shared scale factor per 16-element block. This provides 2× weight bandwidth vs FP8, enabling more experts or larger experts for the same HBM bandwidth budget. But it requires a fundamentally different kernel memory layout.

• →Layout change. FP8 weights are stored as contiguous row-major tensors. FP4 weights must be stored as interleaved [blockscales | packed_fp4_values] structures. TMA descriptors must load scales and values as a single bulk tensor to avoid two separate HBM transactions per block.
• →Dequant in register. FP4 values are dequantized to BF16 in register during the GEMM epilogue — keeping the L2 footprint in FP4 but performing arithmetic in BF16. The index math overhead prevents keeping dequantized values in L2.
• →Tile size constraints. Grouped GEMM tile sizes must be multiples of 16 (one block = 16 elements). This sets a floor on minimum expert size and forces expert count down or parameters per expert up for B200-optimised models.

The architectural implication: B200 shifts the optimal MoE design toward more experts with fewer parameters each, since FP4 makes bandwidth-bound expert weight reads cheaper. Expect models designed for B200 to push toward 256+ expert counts rather than the 64–128 common on H100.

Why MoE Routing Is Inherently Non-Deterministic

MoE routing computes softmax(W_gate @ x) and selects top-k. Two sources of non-determinism make reproducibility difficult:

Practical Mitigation Strategies