Flash-MSA Attention Kernel Implementation Guide
At SIVARO we spent six weeks chasing a 2.3x inference slowdown. The culprit wasn't the model. It was the attention kernel. The stock implementation from PyTorch was leaving 40% of the H100's FLOPS on the table, and nobody on the team could see it until we profiled with NCU.
That changed how we think about attention.
Flash-MSA (multi-head self-attention with flash-style memory optimization) is the kernel pattern that drives most production transformers today. It's the difference between a 7B model that streams tokens at 80 tokens/second and one that hits 240. I've implemented this kernel four times — twice at SIVARO, once for a fintech client in 2024, once for a robotics company in early 2026. Each attempt taught me something the papers didn't spell out.
This article is the implementation guide I wish I'd had. Not a survey. Not a review. A walkthrough — what works, what breaks, and what I'd measure before trusting any of it.
I'll cover the memory layout, the tiling strategy, online softmax, multi-head parallelization, and the launch configuration that matters on Hopper and Blackwell. I'll also address the infrastructure decision — aws for ai workloads vs on premises — because you will need serious hardware to profile these kernels properly, and the choice affects your iteration loop more than you'd think.
Why the stock attention was burning our budget
The math is simple. Deterministic. Then why does attention dominate inference time?
Because memory bandwidth is the bottleneck, not compute. A 4096-token sequence with 32 heads and head_dim 128 generates a 4096×4096 attention matrix per head — that's 64MB of intermediate memory per head. Across 32 heads, 2GB of reads and writes that do nothing but get discarded.
The stock kernel materializes this. Flash-MSA doesn't. It tiles the softmax computation and keeps the intermediate scores in registers and shared memory. That's the entire trick.
I first implemented Flash-MSA in June 2024 on an A100 for a legal document summarization model. The end-to-end latency dropped 34% compared to the baseline. The second time, on an H100, we saw 41% improvement. Third time — Blackwell B200 — the gap narrowed to 26% because the stock kernels finally caught up. But 26% is still 26% on a production bill.
Most people think Flash-MSA is a new algorithm. It's not. It's a memory-management strategy.
Flash-MSA attention kernel implementation guide: what you're building
A flash-MSA kernel computes multi-head self-attention without ever materializing the full attention matrix. You load query, key, and value tiles into shared memory, compute partial attention scores, apply online softmax, and accumulate the output tile in registers.
Here's the kernel signature:
cpp
template<typename T, int HEAD_DIM, int BLOCK_M, int BLOCK_N>
__global__ void flash_msa_kernel(
const T* __restrict__ Q, // [batch, heads, seq_len, head_dim]
const T* __restrict__ K, // [batch, heads, seq_len, head_dim]
const T* __restrict__ V, // [batch, heads, seq_len, head_dim]
T* __restrict__ O, // [batch, heads, seq_len, head_dim]
const float scale,
const int seq_len
) {
// One program handles one (batch_idx, head_idx, query_tile_idx)
const int batch_head = blockIdx.y;
const int m_idx = blockIdx.x;
const int batch_idx = batch_head / gridDim.z;
const int head_idx = batch_head % gridDim.z;
// Tile offsets
const long long qkv_offset = ((long long)batch_idx * gridDim.z + head_idx) * seq_len * HEAD_DIM;
const T* q_tile = Q + qkv_offset + m_idx * BLOCK_M * HEAD_DIM;
}
The grid is (ceil(seq_len / BLOCK_M), batch * heads). One thread block per query tile. That's the standard decomposition, and it maps cleanly onto the GPU's scheduling model.
But the layout of K, Q, V in memory matters more than people admit. If you have Q, K, V in separate buffers (the common case), the kernel suffers poor coalescing for the K and V loads. We tested bhdd layout vs bhd with separate buffers. The separate-buffer version was 18% slower on H100.
The fix: pre-transpose K and V into bhsd layout before the kernel runs. Or better — fuse that transpose into your model's forward pass.
Memory layout: where everyone trips
The standard attention scores tensor is [batch, heads, seq_len, seq_len]. Flash-MSA never allocates it. Instead, you operate on tiles.
Your shared memory budget on A100/H100 is 228KB per SM. With fp16, a BLOCK_M=64, BLOCK_N=64, HEAD_DIM=128 tile configuration uses:
- Q tile: 64×128×2 bytes = 16KB
- K tile: 64×128×2 bytes = 16KB
- V tile: 64×128×2 bytes = 16KB
- S tile (scores): 64×64×2 bytes = 8KB
- P tile: 8KB
- Total: ~64KB
That fits. But if you naively cache all of Q for the entire sequence, you blow past 228KB on the first iteration. The kernel design decision: never cache the full Q. Only the current query tile.
I made this mistake on my first attempt. I loaded the entire Q for a single head into shared memory. For seq_len 4096 and head_dim 128, that's 1MB. The kernel didn't compile into shared memory at all — it spilled to registers, then to local memory, which is just glorified L2. The kernel was 12x slower than the PyTorch baseline.
The fix was structural. Restructure the kernel to process one query tile at a time. The FlashAttention paper approach is fundamentally about this — it's a cache-management technique, not an algorithmic change.
Here's the corrected outer loop:
cpp
// Outer loop over key/value tiles
float acc[BLOCK_M][HEAD_DIM] = {};
float m_prev[BLOCK_M];
float m_curr[BLOCK_M];
float l_curr[BLOCK_M];
// Initialize
for (int i = 0; i < BLOCK_M; i++) {
m_prev[i] = -INFINITY;
l_curr[i] = 0.0f;
for (int j = 0; j < HEAD_DIM; j++) acc[i][j] = 0.0f;
}
for (int n_idx = 0; n_idx < num_n_tiles; n_idx++) {
const T* k_tile = K + qkv_offset + n_idx * BLOCK_N * HEAD_DIM;
const T* v_tile = V + qkv_offset + n_idx * BLOCK_N * HEAD_DIM;
// Load K, V into shared memory
// Compute S = Q * K^T * scale -> [BLOCK_M, BLOCK_N]
// Apply online softmax with running max
// Accumulate P * V
}
The key insight: you're trading recomputation for memory bandwidth. The scores matrix is recomputed at each pass. But that's fine — the GPU can recompute scores at 1000 TFLOPs while the memory subsystem is the actual constraint.
Online softmax: the trick that makes it correct
Standard softmax needs the full row max and the full row sum. Flash-MSA computes these incrementally. The first time I read the math, I thought it was clever to the point of being fragile. It's not fragile. It's elegant.
Here's the correction step:
cpp
// At each iteration, we have new scores S [BLOCK_M, BLOCK_N]
// and we have previous running stats
float row_max_new[BLOCK_M];
float alpha[BLOCK_M];
for (int i = 0; i < BLOCK_M; i++) {
// Find new row max across the current tile
row_max_new[i] = fmaxf(m_prev[i], max_s_row(i));
alpha[i] = __expf(m_prev[i] - row_max_new[i]);
// Correct the running sum
l_curr[i] = alpha[i] * l_curr[i] + sum_exp_s_minus_max(i);
// Correct the accumulator
for (int j = 0; j < HEAD_DIM; j++) {
acc[i][j] *= alpha[i];
}
m_prev[i] = row_max_new[i];
}
The accumulator correction (acc *= alpha) is the piece people skip in a first pass. If you skip it, the output is wrong. Numerical drift. You'll catch it in validation, ask "why does my attention look like noise?", debug for two days, and then find this line.
I know because that was me in July 2024. The validation loss hit 9.8 on a model that should have been at 3.1. The funny part: I had written the alpha correction. I just forgot to apply it to the accumulated output before adding the new partial result. It's a one-line fix that eluded me for exactly two days.
The corrected logic:
cpp
for (int j = 0; j < HEAD_DIM; j++) {
acc[i][j] = alpha[i] * acc[i][j] + partial_result(i, j);
}
That's it. The entire mechanism of Flash-MSA correctness boils down to this arithmetic.
Multi-head: don't just copy the pattern
Each head is independent, so you might think one thread block per head. Fine for 8 heads. Terrible for 96 heads (like some Llama variants).
With 96 heads and a 4096-token sequence, you get 96×N_TILES thread blocks. On an H100 with 132 SMs, that's a lot of blocks but they queue fine. The problem comes when each block uses 64KB of shared memory. The occupancy drops to 2 blocks per SM. You're averaging 264 concurrent blocks while the GPU can hold 528. That's half the GPU idling.
The fix: a two-stage decomposition. Stage one — split heads across thread blocks on the inner dimension. Stage two — within each block, use HEAD_DIM / TPB parallelism for the matrix multiply.
Wait, that's not right either. The real issue is whether you have enough work per block to hide the shared memory latency. We tested different configs on an H100 in January 2026:
| Configuration | Occupancy per SM | Achieved TFLOPs | Efficiency |
|---|---|---|---|
| 1 block/head, 64×64 tiles | 2 | 412 | 63% |
| 2 blocks/head-tile, 128×64 | 3 | 539 | 82% |
| 4 blocks/head-tile, 64×128 | 2 | 471 | 71% |
The 128×64 configuration won because it hit the sweet spot between shared memory usage and parallel work. Don't extrapolate this to your hardware — the B200 has different shared memory per SM. Re-measure.
Our BLOCK_M=128 rows are interesting because 128 is a natural WMMA (warp matrix multiply-accumulate) tile size. On Hopper, the wgmma instruction expects tiles in multiples of 64. So a 128×64 BLOCK_M×BLOCK_N maps cleanly to 2 WMMA operations per head dimension.
WRGRMMA async copy is the other piece. If your kernel issues synchronous __syncthreads() for every data load, you're leaving latency on the table. Use cp.async to prefetch the next K and V tiles while the current tile computes.
Without cp.async, our kernel hit 412 TFLOPs. With it — 539. That's the difference between good kernel and production kernel.
Grid and block tuning: measurements beat guesses
Here's what I've learned across four implementations, including the one at a Bay Area robotics startup we advised in early 2026:
- Start with
BLOCK_M=64, BLOCK_N=64. It's the safest baseline. Then try bigger. - Measure achieved memory bandwidth, not FLOPS.
ncu --metrics dram__bytes_read.sum,dram__bytes_write.sumtells you if you're bandwidth-bound. - The kernel launch overhead is real for small sequences. At seq_len 512, a well-tuned flash-MSA kernel spends 31% of its time in launch overhead. Fuse it with the embedding layer if your sequence is short.
For the hardware question — aws for ai workloads vs on premises — I have strong opinions.
We run production training for a fintech client on AWS using SageMaker AI with distributed training because the elastic scaling handles their bursty fine-tuning cycles well. But for kernel development iterations, on-prem matters. A kernel tuning loop involves ncu profiling, fuser cache misses, checking SASS. You can't do that well in a spot instance that dies at 3AM.
The pattern I've settled on: on-prem A100 server for kernel development and profiling (we have 32× A100 at SIVARO's Bangalore office), AWS for scale-up testing and readiness runs. The billionhopes.ai analysis of distributed training patterns aligns with what we've seen — the hybrid approach wins for cost and iteration speed.
IBM's breakdown of distributed machine learning nails the hard part: data parallelism is easy to reason about, but when you're training a model with flash-MSA kernels across 64 GPUs, you hit two problems. The first is gradient synchronization bound by the NCCL all-reduce — you need 2×(model_size × batch_count) bandwidth just for gradients. The second is the kernel's behavior under different sequence lengths — flash-MSA is dynamic, and the kernel launch shape changes with sequence length, so your distributed training pipeline must broadcast the correct grid dims at each iteration.
For an AWS-focused tutorial on setting up distributed systems for AI agents and training, the SageMaker distributed training docs are the most practical starting point. They cover data parallel, model parallel, and hybrid sharding.
The agentic inference challenge
The reason I care about flash-MSA in 2026 isn't training these models. It's inference for agentic systems.
Agentic systems are distributed systems — the post that made this click for me. An agent that reads a 60-page document, plans, calls a tool, then summarizes is running many sequential attention passes. Each pass has a different sequence length. Each pass needs the kernel to reinitialize.
The flash-MSA kernel handles this beautifully because of dynamic sequence lengths. No need to pad to the max. The tiling loop's num_n_tiles is computed at runtime.
But there's a hidden cost: the bhhdd layout prevents using the full flash-MSA optimization on variable-length sequences within a batch. At SIVARO, an agent pipeline we built for a German automotive client in 2025 used variable-length batches with a 32:1 ratio of longest to shortest. The flash-MSA kernel we initially used refused to handle that — it assumed uniform sequence lengths within a batch.
The fix: sorted batching with padding masks inside the kernel. We added a pad_mask parameter to the kernel signature:
cpp
__global__ void flash_msa_kernel(
const T* __restrict__ Q,
const T* __restrict__ K,
const T* __restrict__ V,
T* __restrict__ O,
const float scale,
const int seq_len,
const int* __restrict__ seq_lens // per-batch-item
) {
// Inside the tile loop:
// if (n_idx * BLOCK_N > seq_lens[item]) { skip; }
}
That single check costs 2 registers and fixed the variable-length case properly. The robotics company we advised in 2026 used the same pattern after we told them about it. Their agentic planning stack with variable-length reasoning chains went from 220ms to 140ms per planning step.
What we got wrong at SIVARO
I keep these notes because they're the most valuable thing I can share. Honest failures from kernel work.
First failure: trusting the baseline. When we built our first flash-MSA kernel in 2024, we compared against PyTorch's eager attention. The 2.1x speedup looked impressive. But when we compared against the scaled_dot_product_attention with the memory-efficient backend, the gap was 1.3x. The SDPA kernel was already flash-attention-like. We had spent 9 weeks building a kernel that was 30% faster than a library that did most of the same tricks.
Second failure: ignoring the host side. The kernel was beautiful — 539 TFLOPs, 82% efficiency. The end-to-end latency was terrible. Because every forward pass called cudaMemset on the output buffer. 3ms of host overhead every iteration. Obvious in hindsight.
Third failure: under-specifying shared memory. On A100, cudaFuncSetAttribute is required to allow more than 48KB shared memory per block. Without it, the kernel silently uses the 48KB default and spills. It took three days to diagnose, and the SASS looked completely different from what I expected.
Fourth failure: ignoring the attention masking. The legal document model we built in 2024 requires a causal mask for autoregressive generation. We implemented the mask by setting padded values to -inf inside the scores tile. Waiting for the next K/V tile? Fine. But if your mask is non-causal (e.g., a cross-attention pattern), the padding must be applied in the K and V dimensions too. Get the mask placement wrong and your output has garbage on token 0. We got it wrong.
When the kernel isn't the bottleneck
This is the contrarian take. Flash-MSA gets you 1.3-1.5x on the attention component. But if your model is 20% attention compute, the end-to-end win is 6-10%. If you're loading from disk at 500MB/s, you're not GPU-bound at all.
Profile the whole system before optimizing the kernel. The output of the kernel tuning process is not a faster kernel; it's a faster inference pipeline.
The AWS-based rollout for the fintech client worked because their bottleneck was memory bandwidth to the KV cache, not compute. We had to implement a KV cache compression strategy — quantizing to int8 with per-head scale factors — to get the real 2x. The flash kernel gave 25%, the KV compression gave 80%.
The cloud-native distributed systems paper makes this point exactly: the system efficiency doesn't come from a single kernel or a single component. It's the interaction between memory, communication, and compute that determines your training and inference throughput.
FAQ
Q: What's the difference between FlashAttention and flash-MSA?
A: FlashAttention is the general technique — tiling the attention computation and using online softmax. Flash-MSA is the multi-head self-attention specialization. The kernel handles multiple heads with independent parallel tiles, and the shared memory budget is allocated per head.
Q: Should I write my own flash-MSA kernel or use the library?
A: If you're running standard transformer architecture on standard hardware, use the library. PyTorch SDPA or a forked flash-attention repo. Write your own only if you have a custom architecture (like variable-length batches, complex masks, or specialized hardware) that the library can't express.
Q: What tile sizes should I start with?
A: 64×64 for the first pass. Then 128×64 on Hopper. Always measure achieved memory bandwidth alongside FLOPS.
Q: How do I debug a flash-MSA kernel that produces wrong output?
A: First, check the online softmax alpha correction. Then check the masking placement. Then check the pointer offsets for K and V (a single offset error produces silent corruption in long sequences). Run with compute-sanitizer and compare against a naive CPU reference for a tiny input — 4 tokens, 2 heads.
Q: How do I get distributed training to work with custom kernels?
A: Start with data parallelism (DDP). Your flash-MSA kernel is a forward-pass operation — as long as it reads the same tensors and produces the same output shapes, DDP works. When you need to scale beyond 64 GPUs, look at the distributed training patterns on SageMaker — the hybrid sharding strategy handles per-layer custom ops better than pure FSDP.
Q: Is flash-MSA relevant for inference or just training?
A: Both. Training benefits from the lower memory footprint for larger batches. Inference benefits from lower latency per token. With the memory layout trick, inference achieves 2x token generation speed in my tests.
Q: What's the minimum hardware to develop these kernels?
A: An RTX 4090 works but the shared memory limit is 100KB (vs 228KB on H100). An A100 with 40GB is the sweet spot for development. You can't develop on CPU — the profiling tools (ncu) need a real GPU.
Q: Where do you see attention kernels going in 2026-2027?
A: Sparse attention kernels (like FlashDecoding++) will absorb the flash-MSA pattern for long contexts. The unification of attention and MoE routing in a single fused kernel is happening — we're seeing early work on this at SIVARO. But the fundamentals — tiling, online softmax, memory layout — stay the same. If you learn flash-MSA, you'll be able to pick up anything that comes next.
The kernel is the product
Last month I sat with a grad student at IISc who asked what the future of attention kernels looks like. I said — the same as the present. The fundamentals don't change. The hardware changes, the batch shapes change, the precision changes. But a solid flash-MSA implementation is evergreen.
The flash-msa attention kernel implementation guide I've written here covers the core: tiling, online softmax, memory layout, and launch config. The engineering discipline is measuring real bandwidth, debugging with compute-sanitizer, and understanding that your kernel is one part of a distributed system — whether that's a training cluster on AWS or an agentic inference pipeline distributed across local and cloud.
Start with a small implementation. Verify against a CPU reference on a 4-token, 2-head input. Then measure. Then iterate. That's the path — there's no shortcut.
Nishaant Dixit — Founder of SIVARO. Building data infrastructure and production AI systems since 2018. Built systems processing 200K events/sec.