SwiGLU fusion, FlexAttention and a GPU optimization cheat-sheet

Bringing together hands-on PyTorch kernel tuning, PyTorch FlexAttention, and a comprehensive optimization cheat-sheet wraps up the practical workflow for modern GPU programming. Below is a deep dive into all three areas.

1. Tuning a Specific PyTorch Operation: SwiGLU Activation Fusion

In modern LLMs (such as Llama, Mistral, and DeepSeek), the Feed-Forward Network (FFN) uses SwiGLU:

SwiGLU(x)=SiLU(xW_(gate))\odot (xW_(up))

In eager PyTorch, computing the element-wise activation step F.silu(gate) * up launches multiple kernels and writes intermediate tensors to HBM. Here is how to benchmark and tune it across three optimization tiers:

import torch
import torch.nn.functional as F
import triton
import triton.language as tl

# Tier 1: Eager PyTorch Baseline (Multiple HBM read/write passes)
def eager_swiglu(gate: torch.Tensor, up: torch.Tensor) -> torch.Tensor:
    return F.silu(gate) * up

# Tier 2: Automated Compiler Fusion via torch.compile
compiled_swiglu = torch.compile(eager_swiglu, mode="max-autotune")

# Tier 3: Custom Autotuned Triton Kernel (Single HBM pass + Block Tuning)
@triton.autotune(
    configs=[
        triton.Config({'BLOCK_SIZE': 1024}, num_warps=4),
        triton.Config({'BLOCK_SIZE': 2048}, num_warps=8),
        triton.Config({'BLOCK_SIZE': 4096}, num_warps=8),
    ],
    key=['n_elements'],
)
@triton.jit
def swiglu_triton_kernel(gate_ptr, up_ptr, out_ptr, n_elements, BLOCK_SIZE: tl.constexpr):
    pid = tl.program_id(0)
    offsets = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
    mask = offsets < n_elements

    # Load both inputs once into registers in FP32 for numerical accuracy
    gate = tl.load(gate_ptr + offsets, mask=mask, other=0.0).to(tl.float32)
    up = tl.load(up_ptr + offsets, mask=mask, other=0.0).to(tl.float32)

    # SiLU(gate) = gate * sigmoid(gate), computed purely in registers
    silu_gate = gate * tl.sigmoid(gate)
    out = (silu_gate * up).to(tl.float16)

    tl.store(out_ptr + offsets, out, mask=mask)

def triton_swiglu(gate: torch.Tensor, up: torch.Tensor) -> torch.Tensor:
    out = torch.empty_like(gate)
    n_elements = gate.numel()
    grid = lambda meta: (triton.cdiv(n_elements, meta['BLOCK_SIZE']),)
    swiglu_triton_kernel[grid](gate, up, out, n_elements)
    return out

• Tuning Takeaway: By using @triton.autotune, Triton benchmarks different BLOCK_SIZE and num_warps combinations on your exact GPU hardware, picking the configuration that saturates HBM bandwidth (~85%–92% of physical peak) and cutting activation latency by 2.5× to 3.5× over eager PyTorch.

2. PyTorch FlexAttention: Programmable FlashAttention via Triton

Historically, researchers faced a painful dilemma: use FlashAttention for fast speed (which only supported rigid masks like causal or full attention) or write custom slow PyTorch attention masks (which materialized the quadratic N × N attention matrix in VRAM).

PyTorch FlexAttention solves this by letting you write standard Python functions to modify attention scores on the fly, which torch.compile lowers directly into a fused FlashAttention Triton kernel.

import torch
from torch.nn.attention.flex_attention import flex_attention, create_block_mask

# Example 1: Custom Score Modification (ALiBi Positional Bias + Soft-Capping)
def alibi_softcap_score_mod(score, b, h, q_idx, kv_idx):
    # Apply Gemma-style logit soft-capping and distance penalty inside SRAM
    score = 30.0 * torch.tanh(score / 30.0)
    distance = q_idx - kv_idx
    return score - (0.125 * (h + 1) * distance)

# Example 2: Sparse Block Mask (Sliding Window + Document Packing Isolation)
document_id = torch.tensor([0, 0, 0, 1, 1, 1, 1, 1], device="cuda")

def doc_sliding_window_mask(b, h, q_idx, kv_idx):
    causal = q_idx >= kv_idx
    window = (q_idx - kv_idx) <= 512
    same_doc = document_id[q_idx] == document_id[kv_idx]
    return causal & window & same_doc

# Pre-compile sparse block mask so completely masked 128x128 tiles are skipped entirely
block_mask = create_block_mask(doc_sliding_window_mask, B=1, H=8, Q_LEN=8, KV_LEN=8)
compiled_flex_attention = torch.compile(flex_attention)

# Executes at FlashAttention speeds with zero N x N VRAM materialization
# out = compiled_flex_attention(q, k, v, score_mod=alibi_softcap_score_mod, block_mask=block_mask)

• Why FlexAttention Is a Breakthrough:

• score_mod: Injects custom element-wise math (such as ALiBi slopes, relative positional biases, or attention logit soft-capping) directly into the FlashAttention inner loop inside fast SRAM.
• block_mask: Pre-computes which 128 × 128 tiles of the attention matrix are completely zeroed out (for example, cross-document boundaries during sequence packing) and skips computing those tiles altogether, yielding massive speedups over dense causal attention.

3. Master Cheat-Sheet: CUDA, Triton, and LLM Optimization Rules

Use this quick-reference guide when diagnosing bottlenecks or architecting AI training and inference pipelines: