3.4 A deep dive into the Transformer feed-forward network (FFN)

The feed-forward network (FFN) is the largest parameter block in a Transformer, about 2/3 of a single Block’s params. It applies an independent nonlinear transform to each token and is the main carrier of the model’s “stored knowledge” and “deeper reasoning”. This article starts from FFN structure, then covers activation-function history, parameter counting, the SwiGLU gate, and AI Infra practice such as tensor-parallel splits, MoE expert parallelism, and CUDA kernel fusion.

📑 Table of contents

• 1. Where FFN sits in the Transformer
• 2. The “expand then compress” FFN structure
• 3. Activation functions in detail
• 4. SwiGLU: the mainstream choice in today’s LLMs
• 5. Parameter counts and comparisons
• 6. Choosing the FFN inner width: why 4x and 8/3x
• 7. Tensor-parallel splits of the FFN
• 8. MoE: splitting the FFN into experts
• 9. CUDA kernel fusion in the FFN
• Self-check list
• References

---

1. Where FFN sits in the Transformer

To see what the FFN does, first recall a Transformer Block. Each Block has two core submodules, Self-Attention and FFN, chained with residual connections and LayerNorm.

Self-Attention’s job is information mixing: every token in the sequence can “see” other tokens and gather related info. If a Transformer Block is a factory line, Self-Attention is the “materials sorting station”: it picks related materials (Value) from warehouse shelves (Key) according to the order (Query) and gathers what is needed.

The FFN is the “processing shop” right after that. Once materials are sorted, each batch gets its own deep processing (cut, melt, shape) that turns rough info into structured output. The key point: each line in the shop runs independently. Each token walks its own line; tokens no longer interact.

In more technical terms:

• Self-Attention is between-token compute: the output depends on combinations across the sequence
• FFN is within-token compute: the same nonlinear transform is applied independently to each token’s representation

That “mix first, then process” design matters. Pure linear aggregation (Attention is a weighted sum, so it is linear) limits expressiveness hard. The FFN’s nonlinear transform gives the model the power to approximate complex functions. By the Universal Approximation Theorem, a two-layer feed-forward net with a nonlinear activation can approximate any continuous function.

Recent work also finds that the FFN acts as a “knowledge store” in large models. A lot of factual knowledge (“Paris is the capital of France”, “water is H2O”) lives in FFN params, while Attention mostly organizes and retrieves that knowledge. That also explains why the FFN holds most of the model’s parameters: it needs enough “storage capacity” for a huge amount of knowledge.

---

2. The “expand then compress” FFN structure

2.1 Basic structure

The standard FFN formula is short:

$$\text{FFN}(x) = \sigma(xW_1 + b_1)W_2 + b_2$$

Where:

• $x$ is the input, shape $(N, d_{model})$
• $W_1$ has shape $(d_{model}, d_{ff})$ and raises the dim from $d_{model}$ to $d_{ff}$
• $\sigma$ is a nonlinear activation
• $W_2$ has shape $(d_{ff}, d_{model})$ and compresses back from $d_{ff}$ to $d_{model}$
• $b_1, b_2$ are biases (modern LLMs usually drop them)

With $d_{model} = 4096$ and $d_{ff} = 4 \times d_{model} = 16384$, the data dims move as:

2.2 Why expand first, then compress

The “expand then compress” pattern is not arbitrary. There is a clear information-processing logic behind it.

✅ Linear separability is better in high dimensions. Classic statistical learning. Red and blue balls mixed on a line (1D) are hard to split with a point. Toss them into 3D and a plane separates them more easily. FFN’s first layer $W_1$ does that: project into a higher-dimensional space so feature patterns are easier to separate and manipulate.

✅ Activations work better in high dimensions. A nonlinear activation selectively keeps or kills channels. In a high-dim space each dim acts like a feature detector; the activation decides which detectors fire. More dims mean more patterns can be checked at once, so expressiveness rises.

✅ Projection down is information distillation. The second layer $W_2$ compresses the activation-filtered high-dim info back to the original dim. That forces the model to keep only what matters, like reducing a soup: water out, essence left.

✅ Contrast with an Autoencoder. If you know Autoencoders, the FFN is the reverse. An Autoencoder compresses then reconstructs (encode-decode) to learn a low-dim code. The FFN expands then compresses (expand-distill) to run a complex nonlinear transform in high dimensions and pull useful info out.

2.3 PyTorch: standard FFN

A commented standard FFN:

import torch
import torch.nn as nn
import torch.nn.functional as F

class StandardFFN(nn.Module):
    """Standard feed-forward net: expand -> activate -> project down"""

    def __init__(self, d_model: int, d_ff: int, activation: str = "relu"):
        """
        Args:
            d_model: model hidden dim, e.g. 4096
            d_ff: FFN inner dim, usually 4 * d_model, e.g. 16384
            activation: activation type, "relu" or "gelu"
        """
        super().__init__()

        # Up projection: (d_model, d_ff), map from the original space to high dim
        self.w_up = nn.Linear(d_model, d_ff, bias=False)

        # Down projection: (d_ff, d_model), compress processed high-dim info back
        self.w_down = nn.Linear(d_ff, d_model, bias=False)

        # Pick the activation
        if activation == "relu":
            self.activation = nn.ReLU()
        elif activation == "gelu":
            self.activation = nn.GELU()
        else:
            raise ValueError(f"Unsupported activation: {activation}")

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        Forward pass
        Args:
            x: input tensor, shape (batch_size, seq_len, d_model)
        Returns:
            output tensor, shape (batch_size, seq_len, d_model)
        """
        # Step 1: expand (batch_size, seq_len, d_model) -> (batch_size, seq_len, d_ff)
        h = self.w_up(x)

        # Step 2: nonlinear activation, feature selection in high dim
        h = self.activation(h)

        # Step 3: project down (batch_size, seq_len, d_ff) -> (batch_size, seq_len, d_model)
        output = self.w_down(h)

        return output

# Usage example
d_model = 4096
d_ff = 4 * d_model  # 16384

ffn = StandardFFN(d_model, d_ff, activation="gelu")
x = torch.randn(2, 128, d_model)  # batch=2, seq_len=128
out = ffn(x)  # output shape: (2, 128, 4096)
print(f"Input shape: {x.shape}, output shape: {out.shape}")
# Params: 4096*16384 + 16384*4096 = 134,217,728 (~134M)
print(f"Params: {sum(p.numel() for p in ffn.parameters()):,}")

---

3. Activation functions in detail

The activation is a critical FFN piece. Without it, two stacked linear maps stay linear ($xW_1 W_2$ equals one linear map $x(W_1 W_2)$), and the FFN has no point. The nonlinearity from the activation is where expressiveness comes from.

3.1 ReLU (Rectified Linear Unit)

✴️ Formula

$$\text{ReLU}(x) = \max(0, x)$$

💬 Intuition

ReLU is the simplest “switch”: positive inputs pass through, negative inputs shut off. Think of a one-way valve: signal only flows in the positive direction; the negative side is fully cut.

Output range and properties

• Output range: $[0, +\infty)$
• Gradient on the positive half is always 1, so no vanishing-gradient issue there
• Gradient on the negative half is always 0; output is always 0
• Extremely cheap: one compare

✅ Pros

• Simple and hardware friendly: one max on the GPU
• Positive-half gradient of 1 eases vanishing gradients in deep nets (vs sigmoid/tanh)
• Sparseness: negatives become zero, so activated vectors often have many zeros, which helps sparse feature codes

❌ Cons: dying ReLU

This is ReLU’s worst flaw and deserves detail.

If a neuron’s input lands in the negative region during training, ReLU outputs 0 and the gradient is 0. Zero gradient means that neuron’s weights get no update on the backward pass. With no update, the next forward pass likely stays negative, so the neuron “dies” permanently and stops contributing.

Large learning rates make this worse. One big update can push weights into a bad region where the neuron is negative on every training sample and never recovers. In practice you sometimes see more than 40% of ReLU neurons dead, which shrinks effective capacity a lot.

LeakyReLU softens this with a tiny negative slope (for example 0.01x), but that is a patch. It does not fix the non-smooth, non-differentiable kink at zero.

3.2 GELU (Gaussian Error Linear Unit)

✴️ Formula

$$\text{GELU}(x) = x \cdot \Phi(x) = x \cdot \frac{1}{2}\left[1 + \text{erf}\left(\frac{x}{\sqrt{2}}\right)\right]$$

$\Phi(x)$ is the standard normal CDF; $\text{erf}$ is the error function.

In practice people often use the approximation:

$$\text{GELU}(x) \approx 0.5x\left[1 + \tanh\left(\sqrt{\frac{2}{\pi}}(x + 0.044715x^3)\right)\right]$$

💬 Intuition: probabilistic gating

GELU treats activation as “probabilistic gating”. For input $x$, GELU is the expected value of “keep $x$ with probability $\Phi(x)$, zero it with probability $1-\Phi(x)$”.

Why a normal CDF as the keep probability? The intuition: if neural-net inputs are roughly normal (by the central limit theorem, sums of random variables tend toward normality), a larger input (farther right of the mean) is more likely a meaningful signal than noise, so it should be kept more often. $\Phi(x)$ gives a keep probability that rises monotonically with the input.

Unlike ReLU’s hard switch (0 or full pass), GELU is a soft gate: near-zero inputs are not fully cut; they are kept with a probability tied to their magnitude.

Output range and properties

• Output range: about $[-0.17, +\infty)$; a small negative pocket on the negative half
• Smooth and differentiable everywhere; no ReLU kink at zero
• Negative region output is tiny but nonzero, and gradients stay nonzero, so no dying neurons
• Close to ReLU but smoother: for large positive $x$, GELU(x) ≈ x; for large negative $x$, ≈ 0

✅ Pros:

• Smooth and differentiable, more stable training
• Probabilistic gating keeps information more finely
• Shown better than ReLU in BERT, GPT-2, and similar models

❌ Cons:

• Costlier than ReLU (erf or a tanh approx)
• On modern GPUs the gap is almost noise next to GEMM

3.3 Swish

✴️ Formula

$$\text{Swish}(x) = x \cdot \sigma(\beta x) = \frac{x}{1 + e^{-\beta x}}$$

$\sigma$ is the sigmoid; $\beta$ is learnable or fixed. When $\beta=1$, people usually write $\text{SiLU}(x) = x \cdot \sigma(x)$.

💬 Intuition

Swish is a kind of adaptive interpolation between ReLU and a linear map. Sigmoid $\sigma(x)$ outputs a gate in $[0, 1]$; multiply by $x$ and you get “the input gates how much of itself passes”. That is self-gating.

For large positive $x$, $\sigma(x) \to 1$, so $\text{Swish}(x) \approx x$ (near linear). For large negative $x$, $\sigma(x) \to 0$, so $\text{Swish}(x) \approx 0$ (near cutoff). Near a small negative $x$, Swish dips slightly below zero: a small “bowl”, minimum about $-0.278$ near $x \approx -1.278$.

Output range and properties

• Output range: about $[-0.278, +\infty)$
• Smooth and differentiable everywhere
• Non-monotonic: a small dip on the negative half
• As $\beta \to \infty$, Swish collapses to ReLU; at $\beta = 0$ it collapses to the linear map $x/2$

✅ Smoothness benefits

Swish smoothness helps training in practice. ReLU’s gradient jumps at $x=0$ (left derivative 0, right derivative 1), so the loss landscape has “corners”. Gradient descent can oscillate near those non-smooth points.

Everywhere-differentiable Swish makes the loss surface smoother; gradients change continuously over the domain, so optimization is more stable. Google Brain work found Swish consistently beats ReLU in deep nets, especially past about 40 layers.

3.4 Comparison summary

---

4. SwiGLU: the mainstream choice in today’s LLMs

4.1 The GLU gating idea

Before SwiGLU, understand GLU (Gated Linear Unit).

A standard FFN applies one “uniform” activation: every dim gets the same nonlinear map. GLU takes a different path: let the model learn which information passes and which is suppressed.

✴️ Core GLU formula:

$$\text{GLU}(x) = (xW_{up}) \otimes \sigma(xW_{gate})$$

$\otimes$ is elementwise multiply; $\sigma$ is sigmoid.

That introduces two independent linear projections:

• $xW_{up}$: candidate content ready to pass
• $\sigma(xW_{gate})$: gate signal, a pass probability per dim in $(0, 1)$

Elementwise multiply keeps dims with large gates and suppresses dims with small gates. That is more flexible than one shared activation: the model can learn different gating policies for different inputs.

4.2 Full SwiGLU definition

SwiGLU is a GLU variant from Noam Shazeer’s 2020 paper “GLU Variants Improve Transformer”. It replaces the sigmoid gate with Swish:

$$\text{SwiGLU}(x) = (\text{Swish}(xW_{gate})) \otimes (xW_{up})$$

A full SwiGLU FFN has three weight matrices:

$$\text{FFN}_{SwiGLU}(x) = \left[\text{Swish}(xW_\text{gate}) \otimes (xW_\text{up})\right] W_\text{down}$$

Where:

• $W_{gate}$: shape $(d_{model}, d_{ff})$, gate projection; Swish of its output is the gate signal
• $W_{up}$: shape $(d_{model}, d_{ff})$, content projection; candidate information to pass
• $W_{down}$: shape $(d_{ff}, d_{model})$, down projection; compresses back to the original dim

4.3 Why gating works

A few angles on why gating helps.

✅ Finer information selection. In a standard FFN, one global nonlinear rule hits every expanded dim. SwiGLU gating lets the model make a different keep/kill decision per dim. That supports more conditional compute: “when the input has feature A, keep these dims; when it has feature B, keep those.”

✅ Richer gradient signal. Gate and content paths are two independent linear maps, so backprop can send gradients down both, which gives richer update signals. Think of it as an implicit “gradient highway”, in the same spirit as residual connections.

✅ Empirics. Shazeer’s paper compared several GLU variants at matched param count; SwiGLU won consistently on many downstream tasks. Later LLaMA, Mistral, Qwen, DeepSeek and other mainstream LLMs adopted SwiGLU, which further validated it in production.

4.4 PyTorch: SwiGLU FFN

import torch
import torch.nn as nn
import torch.nn.functional as F

class SwiGLUFFN(nn.Module):
    """SwiGLU feed-forward: gate + content projections -> elementwise multiply -> down"""

    def __init__(self, d_model: int, d_ff: int):
        """
        Args:
            d_model: model hidden dim, e.g. 4096
            d_ff: FFN inner dim; SwiGLU usually uses (8/3) * d_model, e.g. 11008
        """
        super().__init__()

        # Gate projection: (d_model, d_ff)
        # After Swish, becomes the gate that controls how much information passes
        self.w_gate = nn.Linear(d_model, d_ff, bias=False)

        # Content projection: (d_model, d_ff)
        # Candidate content ready to pass
        self.w_up = nn.Linear(d_model, d_ff, bias=False)

        # Down projection: (d_ff, d_model)
        # Compress gated high-dim info back to the original dim
        self.w_down = nn.Linear(d_ff, d_model, bias=False)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        Forward pass
        Args:
            x: input tensor, shape (batch_size, seq_len, d_model)
        Returns:
            output tensor, shape (batch_size, seq_len, d_model)
        """
        # Step 1: gate path: project to high dim, then Swish
        # (batch_size, seq_len, d_model) -> (batch_size, seq_len, d_ff)
        gate = F.silu(self.w_gate(x))  # silu is Swish(x) = x * sigmoid(x)

        # Step 2: content path: project to high dim (no activation)
        # (batch_size, seq_len, d_model) -> (batch_size, seq_len, d_ff)
        up = self.w_up(x)

        # Step 3: gated multiply: elementwise gate * content
        # Large gate dims keep info; small gate dims suppress it
        # (batch_size, seq_len, d_ff)
        hidden = gate * up

        # Step 4: down project: compress back to the original dim
        # (batch_size, seq_len, d_ff) -> (batch_size, seq_len, d_model)
        output = self.w_down(hidden)

        return output

# Example: LLaMA-2-7B config
d_model = 4096
d_ff = 11008  # (8/3) * 4096, rounded to a multiple of 256

swiglu_ffn = SwiGLUFFN(d_model, d_ff)
x = torch.randn(2, 128, d_model)  # batch=2, seq_len=128
out = swiglu_ffn(x)  # output shape: (2, 128, 4096)
print(f"Input shape: {x.shape}, output shape: {out.shape}")
# Params: 4096*11008 * 2 (gate + up) + 11008*4096 (down) = 135,266,304 (~135M)
print(f"Params: {sum(p.numel() for p in swiglu_ffn.parameters()):,}")

---

5. Parameter counts and comparisons

FFN params dominate a Transformer Block. Accurate counting is the base for memory planning and parallel strategy design.

5.1 Standard FFN vs SwiGLU FFN

With $d_{model} = 4096$, compare both structures.

Standard FFN ($d_{ff} = 4 \times d_{model} = 16384$)

SwiGLU FFN ($d_{ff} = \frac{8}{3} \times d_{model} \approx 10922$, actually 11008)

5.2 Side-by-side

Key finding: total params are almost the same ($8 \times d_{model}^2$), but SwiGLU gets better quality through gating. The cost is 3 forward GEMMs instead of 2, so compute rises a bit.

5.3 Vs the Attention module

FFN params are about 2× Attention, roughly 2/3 of a Block. Any FFN optimization (parallel split, quantization, MoE) moves overall efficiency a lot.

---

6. Choosing the FFN inner width: why 4x and 8/3x

6.1 Why standard FFN uses 4× expansion

$d_{ff} = 4 \times d_{model}$ comes from the original Transformer paper “Attention Is All You Need”. Vaswani et al. set $d_{model}=512, d_{ff}=2048$, a 4× ratio that later became the de facto standard.

Why 4× rather than 2× or 8×? Mostly empirical. Intuition:

• Too small (for example 2×): the high-dim space is too tight, nonlinear expressiveness suffers, quality drops
• Too large (for example 8×): params explode while gains diminish, so param efficiency falls
• 4× is a balance: a solid tradeoff between quality and param efficiency

Later work (Kaplan et al. on Scaling Laws) also suggests that, at fixed total params, $d_{ff} / d_{model}$ near 4 is roughly optimal. Much larger or smaller ratios hurt at the same param budget.

6.2 Why SwiGLU uses 8/3× expansion

SwiGLU turns the standard FFN’s 2 matrices into 3. Keep $d_{ff} = 4 \times d_{model}$ and total params become:

$$3 \times d_{model} \times 4d_{model} = 12 \times d_{model}^2$$

That is 50% more than the standard FFN’s $8 \times d_{model}^2$. To add gating without growing the param budget, shrink the inner dim:

$$3 \times d_{model} \times d_{ff} = 8 \times d_{model}^2$$ $$d_{ff} = \frac{8}{3} \times d_{model} \approx 2.667 \times d_{model}$$

At $d_{model} = 4096$: $d_{ff} = \frac{8}{3} \times 4096 = 10922.67$.

In engineering, $d_{ff}$ is usually rounded to a multiple of 128 or 256 so matrix dims line up with GPU Tensor Cores. Tensor Cores want certain alignments for matmul (for example multiples of 8 in FP16; aligning to 128 or 256 often raises utilization in practice).

LLaMA-2-7B example: $10922.67$ rounds to $11008$ ($= 43 \times 256$). That is where the LLaMA config $d_{ff} = 11008$ comes from.

---

7. Tensor-parallel splits of the FFN

When a model no longer fits on one GPU, you shard params across GPUs: Tensor Parallelism (TP). The FFN is 2/3 of Block params, so how you split it largely sets overall TP efficiency.

7.1 Megatron-LM column + row split

Megatron-LM’s scheme is neat: split $W_{up}$ (and $W_{gate}$) by columns, split $W_{down}$ by rows, so intermediates need no communication and only the final output needs one AllReduce.

On 2 GPUs, a standard FFN splits like this:

Step 1: column-split $W_{up}$

Shard $W_{up}$‘s columns (output dim) evenly:

GPU 0: W_up[:, :d_ff/2]    shape (d_model, d_ff/2)
GPU 1: W_up[:, d_ff/2:]    shape (d_model, d_ff/2)

Each GPU computes its own expand result:

GPU 0: h_0 = activation(x @ W_up_0)    shape (N, d_ff/2)
GPU 1: h_1 = activation(x @ W_up_1)    shape (N, d_ff/2)

Note: each GPU needs a full copy of $x$ (from a prior AllReduce or broadcast). The column-split win: the activation can run independently after the split. Activations are elementwise and do not depend on other dims, so activating the full vector equals activating each shard.

Step 2: row-split $W_{down}$

Shard $W_{down}$‘s rows (input dim) evenly:

GPU 0: W_down[:d_ff/2, :]    shape (d_ff/2, d_model)
GPU 1: W_down[d_ff/2:, :]    shape (d_ff/2, d_model)

Each GPU multiplies its intermediate by its $W_{down}$ shard:

GPU 0: out_0 = h_0 @ W_down_0    shape (N, d_model)
GPU 1: out_1 = h_1 @ W_down_1    shape (N, d_model)

Step 3: AllReduce sum

The final output is the sum of both shards:

output = AllReduce(out_0, out_1) = out_0 + out_1    shape (N, d_model)

Because (row-vector view: $h$ split on columns, $W_{down}$ split on rows):

$$h W_{down} = \begin{bmatrix} h_0 & h_1 \end{bmatrix} \begin{bmatrix} W_{down,0} \\ W_{down,1} \end{bmatrix} = h_0 W_{down,0} + h_1 W_{down,1}$$

7.2 Splitting SwiGLU

For SwiGLU, column-split both $W_{gate}$ and $W_{up}$, row-split $W_{down}$. Same logic:

GPU 0: gate_0 = Swish(x @ W_gate_0)    # (N, d_ff/2)
        up_0  = x @ W_up_0              # (N, d_ff/2)
        mid_0 = gate_0 * up_0           # elementwise, (N, d_ff/2)
        out_0 = mid_0 @ W_down_0        # (N, d_model)

GPU 1: gate_1 = Swish(x @ W_gate_1)    # (N, d_ff/2)
        up_1  = x @ W_up_1              # (N, d_ff/2)
        mid_1 = gate_1 * up_1           # (N, d_ff/2)
        out_1 = mid_1 @ W_down_1        # (N, d_model)

output = AllReduce(out_0, out_1)        # (N, d_model)

The whole FFN forward needs one AllReduce. The three GEMMs (gate/up/down) and elementwise ops on each GPU stay local. TP communication on the FFN is cheap: one AllReduce per Block’s FFN (Attention adds another, so two AllReduces per Block).

7.3 Engineering notes

In deployment, $d_{ff}$ must be divisible by the TP GPU count. For example $d_{ff} = 11008$ splits to 5504 per GPU at TP=2, 2752 at TP=4, 1376 at TP=8: all divide cleanly. That is one reason $d_{ff}$ favors powers of two or multiples of 128/256.

---

8. MoE: splitting the FFN into experts

8.1 Core idea

Mixture of Experts (MoE) grows model capacity without growing compute in proportion. The FFN is the main thing it rewrites.

The idea is simple: instead of one huge FFN, split it into several smaller “experts”, each an independent FFN. For each input token, activate only a few experts; the rest stay idle.

Example: an MoE layer may hold 64 expert FFNs, but each token only picks Top-2 for the forward pass. Then:

• Total model params = sum of all 64 experts (large)
• Compute per token = compute of 2 experts (small)

That is “large params, small compute”: enough capacity to store knowledge, while train/infer compute stays manageable.

8.2 Router

The key MoE piece is the Router / Gate, which decides which experts each token goes to.

The common pattern is Top-K routing:

$$g(x) = \text{softmax}(x W_r)$$ $$\text{TopK}(g(x)) \to \text{pick the K experts with highest probability}$$

$W_r$ is the router weight matrix, shape (d_model, n_experts). It maps each token to an n_experts-dim distribution; the top K experts are chosen.

The final output is a weighted sum of the chosen experts:

$$\text{MoE}(x) = \sum_{i \in \text{TopK}} g_i(x) \cdot \text{Expert}_i(x)$$

Weights $g_i(x)$ are re-normalized to sum to 1 so the output scale matches a single FFN.

8.3 Load balancing

A hard MoE training issue is load imbalance: the router may send most tokens to a few “star experts”, so others get little training, degrade, and imbalance gets worse.

Common mitigations:

✅ Auxiliary loss (Load Balancing Loss)

Add a loss term that punishes uneven load. A common form:

$$L_{aux} = \alpha \cdot n_{experts} \cdot \sum_{i=1}^{n_{experts}} f_i \cdot p_i$$

$f_i$ is the fraction of tokens expert $i$ actually got, $p_i$ is the average router probability for expert $i$, $\alpha$ is a balance coefficient. The loss is minimized when load is even.

✅ Capacity factor

Cap how many tokens each expert may take. Overflow tokens are dropped (sent along the residual path) or re-routed. That hard constraint stops one expert from overloading.

✅ Expert Choice routing

Unlike token-choose-expert, Expert Choice lets each expert pick its own top-k tokens. That naturally gives every expert the same token count and attacks imbalance at the root.

8.4 AI Infra link: Expert Parallelism

MoE adds a new parallel axis: Expert Parallelism (EP). Because each expert is an independent FFN, different experts can live on different GPUs.

That brings a distinct communication pattern: classic TP uses AllReduce (every rank joins), while MoE EP uses All-to-All: each GPU sends tokens meant for remote experts and receives tokens meant for local experts. That pattern stresses network topology and the comm library differently from AllReduce, and is a core MoE systems challenge.

Successful deployments like Mixtral 8x7B and DeepSeek-V2 need carefully designed EP scheduling and communication optimization.

---

9. CUDA kernel fusion in the FFN

9.1 Why fuse kernels

The GPU’s basic compute unit is a kernel: the CPU schedules one kernel at a time onto the GPU. Each launch has fixed overhead (usually a few microseconds), and each standalone kernel reads inputs from HBM and writes outputs back.

An FFN forward has several steps: GEMM, activation, elementwise multiply, and so on. If each step is its own kernel, intermediates bounce between HBM and registers/SRAM over and over. For SwiGLU, the elementwise work in one forward includes:

1. Swish: $\text{gate} = x\_{gate} \cdot \sigma(x\_{gate})$ (sigmoid plus elementwise multiply)
2. Gated multiply: $\text{hidden} = \text{gate} \otimes \text{up}$

Run separately, the gate vector is written to HBM then read back for the multiply. Fused, both ops finish in one kernel: the gate intermediate stays in registers, multiplies with up, and never touches HBM.

9.2 Typical FFN fusion cases

🔶 Case 1: fuse Swish + gated multiply

Fuse Swish(gate) * up into one kernel. Two $(N, d_{ff})$ inputs become one $(N, d_{ff})$ output; HBM accesses drop from 5 to 3: unfused reads gate and up, writes then rereads the Swish intermediate, then writes hidden; fused only reads two inputs and writes one output, skipping the intermediate write/read.

🔶 Case 2: GEMM + activation fusion (Epilogue Fusion)

Embed the activation in the GEMM kernel’s epilogue. For the $W_{gate}$ GEMM, the naive path is:

1. kernel A: $x_{\text{gate}} = x W_{gate}$, write HBM
2. kernel B: $\text{gate} = \text{Swish}(x_{\text{gate}})$, read $x_{\text{gate}}$ from HBM and write gate back

Fused: in the last GEMM step (before writing HBM), apply Swish in registers per output element, then write once. One kernel, one HBM write.

cuBLAS, CUTLASS, and similar GEMM libraries expose epilogue fusion so you can append custom elementwise ops after the matmul.

🔶 Case 3: merge $W_{gate}$ and $W_{up}$ GEMMs

Both $W_{gate}$ and $W_{up}$ take the same $x$. Concatenate them along columns into one $(d_{model}, 2 \times d_{ff})$ matrix, run one big GEMM, then split the output into gate and up. That turns two GEMMs into one, cuts launch overhead, and a larger GEMM usually uses the GPU better than two smaller ones.

# Separate (2 GEMM launches):
gate = x @ W_gate    # (N, d_model) x (d_model, d_ff) = (N, d_ff)
up   = x @ W_up      # (N, d_model) x (d_model, d_ff) = (N, d_ff)

# Fused (1 GEMM launch):
W_fused = concat(W_gate, W_up, dim=1)  # (d_model, 2*d_ff)
fused_out = x @ W_fused                 # (N, 2*d_ff)
gate, up = split(fused_out, d_ff, dim=1)

🔶 Case 4: fuse bias add + residual add + LayerNorm

FFN output usually goes through a residual add and a following LayerNorm. Fusing those elementwise ops into one kernel lets you read FFN output, residual input, and LayerNorm params once, finish add and normalize on-chip, then write back, avoiding multiple HBM round trips.

9.3 Performance impact

In real LLM inference, FFN kernel fusion gains vary by phase:

• Decode (one token at a time): GEMM becomes GEMV, compute is small, launch overhead and HBM bandwidth dominate, so fusion helps a lot (often 20-40% speedup)
• Prefill (full prompt): GEMM compute dominates, elementwise HBM cost is a smaller share, so fusion helps less (often 5-15%)

---

🎯 Self-check list

After this article, check whether you really understand the FFN:

• Explain the FFN’s role in a Transformer Block: why Attention still needs an FFN after it, and what the FFN owns
• Draw standard FFN and SwiGLU FFN on a whiteboard, with each matrix’s dims and data-flow direction
• Hand-compute param counts for both, and explain why they nearly match
• Explain why the inner width is 4× (standard) and 8/3× (SwiGLU), and where LLaMA’s 11008 comes from
• State ReLU, GELU, and Swish formulas, intuition, and tradeoffs
• Explain why dying ReLU happens, and how GELU and Swish avoid it
• Explain SwiGLU gating: what $W_{gate}$ and $W_{up}$ each do, and why gating beats one shared activation
• Describe Megatron-LM’s TP split for the FFN: column-split $W_{up}$/$W_{gate}$, row-split $W_{down}$, and why that needs only one AllReduce
• Sketch MoE: multiple expert FFNs, routing, load balancing, and Expert Parallelism’s communication pattern
• Name at least two concrete CUDA kernel fusion cases inside the FFN

📚 References

Papers

• Attention Is All You Need (Vaswani et al., 2017): https://arxiv.org/abs/1706.03762 — original Transformer paper; source of the FFN structure
• GLU Variants Improve Transformer (Shazeer, 2020): https://arxiv.org/abs/2002.05202 — SwiGLU and other gated activations, with experiments
• Gaussian Error Linear Units (GELUs) (Hendrycks & Gimpel, 2016): https://arxiv.org/abs/1606.08415 — GELU activation
• Searching for Activation Functions (Ramachandran et al., 2017): https://arxiv.org/abs/1710.05941 — discovery of Swish
• Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism (Shoeybi et al., 2019): https://arxiv.org/abs/1909.08053 — tensor-parallel split strategy
• Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity (Fedus et al., 2022): https://arxiv.org/abs/2101.03961 — simplified MoE at large scale
• LLaMA: Open and Efficient Foundation Language Models (Touvron et al., 2023): https://arxiv.org/abs/2302.13971 — SwiGLU in LLaMA
• Mixtral of Experts (Jiang et al., 2024): https://arxiv.org/abs/2401.04088 — sparse MoE architecture

Tutorials and blogs

• The Illustrated Transformer (Jay Alammar): https://jalammar.github.io/illustrated-transformer/ — illustrated Transformer intro
• Andrej Karpathy: Let’s build GPT from scratch: https://www.youtube.com/watch?v=kCc8FmEb1nY — GPT from scratch
• The Annotated Transformer (Harvard NLP): https://nlp.seas.harvard.edu/annotated-transformer/ — paper lined up with a PyTorch implementation