The Flash Attention Kernel That Had to Be Built: Designing sm_120 MLA Verify for Blackwell Consumer GPUs

Introduction

In the high-stakes world of large language model inference, every microsecond counts. When deploying a 200k-token-context speculative decoding system across eight RTX PRO 6000 Blackwell GPUs, the team behind the Kimi K2.6 model hit a wall: the existing ecosystem of optimized attention kernels simply didn't support their hardware. The NVIDIA sm_120 architecture powering these consumer Blackwell GPUs lacked support from FlashMLA, cutlass-MLA, and flashinfer-MLA—all compiled only for sm_90a (Hopper), sm_100a, or sm_103a (Blackwell datacenter). This left a single message in the conversation as the turning point: the moment when the assistant stopped surveying the landscape and started building a custom flash attention kernel from scratch.

Message [msg 12227] captures the critical transition from planning to execution. In this message, the assistant writes verify_attn_flash.cu—the first implementation of a custom sm_120 flash attention kernel for Multi-head Latent Attention (MLA) verification in the DFlash speculative decoding pipeline. But more importantly, the message's extensive reasoning section reveals the full design journey: a multi-iteration exploration of kernel architectures, shared memory budgets, thread occupancy, and the fundamental tension between algorithmic elegance and hardware constraints. This article examines that reasoning in depth, unpacking the decisions, assumptions, and trade-offs that shaped a kernel destined to become the backbone of a production inference system.

The Strategic Context: Why This Message Matters

To understand why [msg 12227] is a pivotal moment, we must trace the arc of the broader session. The team had been deploying Kimi K2.6 with DFlash speculative decoding—a technique where a lightweight "drafter" model proposes candidate tokens and the main model verifies them in parallel using a tree-structured attention mask. The verification step is the throughput bottleneck: for each decode step, the model must compute attention scores between the current query and every cached key-value token, then apply the tree mask to determine which candidates survive.

Earlier in the session ([msg 12225]), the assistant had conducted an exhaustive survey of available MLA kernels and discovered a stark reality: none supported sm_120. The RTX PRO 6000 Blackwell GPUs, despite being cutting-edge hardware, use a consumer-grade architecture that lacks the Tensor Memory Accelerator (TMA) and warp-group matrix multiply (wgmma) instructions found in datacenter Blackwell (sm_100a/sm_103a). The assistant had documented this in a comprehensive plan (plans/0002-sm120-verify-kernel-defrag.md) that laid out four phases: Phase 0 (baseline validation), Phase 1 (custom flash kernel), Phase 2 (SGLang backend integration), and Phase 3 (KV defragmentation).

Message [msg 12227] executes the transition from Phase 0 to Phase 1. The spec document is written; now the kernel must be built. The reasoning opens with a tactical decision about workflow efficiency:

"I need to execute the baseline phase first by syncing my local changes to the deployed repository at CT200 and running the existing verification tests to establish a baseline, then move on to writing the flash kernel in phase one. Actually, let me be more efficient: I'll write the flash kernel implementation locally along with the necessary API changes and test code, then sync everything to CT200 in one batch and build and test together."

This seemingly minor pivot—from sequential (baseline first, then kernel) to batched (write everything, then test)—reveals a key aspect of the assistant's operating style: it values minimizing round-trips to the remote machine. Each sync-and-test cycle on CT200 involves network transfer, build time, and potential debugging iterations. By writing the kernel, API changes, tests, and benchmarks locally before syncing, the assistant collapses multiple potential cycles into one. It's a bet that the kernel design is mature enough to get right in a single pass.

The Design Space: Iterating Toward a Viable Kernel Architecture

The bulk of the reasoning in [msg 12227] is a detailed exploration of kernel design alternatives. This is not a straightforward implementation—it's a genuine search through a multi-dimensional trade-off space. Let us trace the iterations.

Iteration 1: The Register-Tiled All-Heads Design

The assistant begins with the most ambitious approach: one block per (batch, query position) processing all 64 heads simultaneously, with the full output accumulator distributed across thread registers. The motivation is clear—the KV cache latent vector (kv_c[j], dimension 512) is reused across all heads, so reading it once and amortizing across 64 heads gives a 64× reduction in global memory reads. The accumulator (64 heads × 512 dimensions = 32,768 floats) doesn't fit in shared memory, so it must live in registers: 128 floats per thread across 256 threads.

But this design introduces a severe complexity: when a new KV tile updates the running softmax maximum for a head, all 512 floats of that head's accumulator (spread across multiple threads) must be rescaled. This requires careful thread-to-(head, dimension) mapping so each thread knows which rescale factor to apply. The assistant identifies this challenge and pivots:

"The challenge is that when a new KV tile updates the running max for a head, all 512 floats of that head's accumulator spread across threads need rescaling, which requires careful thread-to-(head, dimension) mapping so each thread knows which rescale factor to apply."

Iteration 2: Grouped Heads with Shared Memory Accumulators

The pivot is to a grouped-heads approach: process heads in smaller batches (e.g., 8 heads per block group) so the accumulator fits in shared memory alongside the KV tile and online softmax state. For TK=32 (32-token KV tile) and HT=8 (8 heads per tile), the shared memory footprint is approximately 90KB—tight but workable under the 100KB limit. The win is still substantial: redundant KV cache reads drop from 64× down to 8× (one read per head group instead of per head).

But the assistant is dissatisfied with "substantial but not ideal":

"The real bottleneck is global memory reads of the KV cache, so the ultimate goal is to read it only once per query."

Iteration 3: Full Register Tiling (v3)

This drives a return to the register-tiled approach, now with a refined thread layout: 64 heads × 4 threads per head = 256 threads, each holding 128 accumulator floats. The scoring computation uses warp shuffles to reduce partial dot products across the 4-thread group. The KV cache is read from global memory only once per query position block, achieving the full 64× reduction.

But the assistant hesitates:

"The v3 design is complex but follows a standard pattern, so I should implement it carefully. I'm reconsidering the thread mapping though—the QK computation (score per head and token) wants threads split over the reduction dimension, while the PV computation (accumulating output) wants threads split over the output dimension."

This tension between the QK phase (which wants threads distributed over the reduction dimension for dot products) and the PV phase (which wants threads distributed over the output dimension for accumulation) is a fundamental challenge in flash attention kernel design. The assistant considers generalizing beyond H=64 but finds the flexibility requirements messy:

"But making this fully flexible across all test configs is getting messy—I think the cleaner approach is to accept a compromise: use one block per (batch, position, head) for simplicity and correctness, even if it means more latent register reads."

The Final Architecture: HT=8, TK=16, Shared Memory

The assistant ultimately settles on a pragmatic middle ground. The key insight comes from revisiting the shared memory budget:

"I'll use a 2D grid where each block processes a (b, i) query position and a tile of HT heads, with the head-tile dimension handled by grid.z. This keeps the accumulator bounded to HT×Dl floats and lets me reuse the KV tile across all heads in the block, cutting latent reads from 64× down to ceil(H/HT)—so 4× for H=64 with HT=16."

Wait—HT=16? That contradicts the earlier conclusion that HT=16 with TK=16 exceeds shared memory. Let us trace the final numbers carefully.

The assistant recalculates: HT=16 with TK=16 gives approximately 72KB total (KV tile, accumulators, scores), fitting comfortably under 100KB. But then the reasoning shifts to HT=8 with TK=16 at 72KB. There's an inconsistency here—the assistant seems to be exploring multiple configurations without fully committing to one.

The resolution comes later:

"For now, I'll go with HT=8 and Tk=32 as my target, which should give me roughly 91KB of shared memory usage."

But then the actual implementation uses TK=16 with HT=8:

"Writing the kernel implementation... TK=16 gives me 72KB of shared memory, which fits safely under the 100KB limit with good margin. I'm settling on TK=16 with HT=8 and a block size of 256 threads."

The final choice is HT=8, TK=16, BLOCK=256, with q_nope cached in shared memory. The assistant calculates the performance improvement: the naive approach reads the KV cache twice per request (once for scoring, once for PV) across all H×q_len blocks, each touching the full kv_len×Dl. Flash attention tiles across (H/HT)×q_len blocks, each reading KV cache once and reusing it for both phases. Total global reads drop to naive/(2×HT) = 16× reduction with HT=8.

The Assumptions Underlying the Design

Every kernel design rests on assumptions about hardware behavior and workload characteristics. The reasoning in [msg 12227] reveals several critical assumptions:

Shared memory budget: The assistant assumes a 100KB shared memory limit per block on sm_120. This was established earlier in the session through documentation and experimentation. The assumption drives all tile-size decisions: HT=8 with TK=16 fits at ~72KB, leaving margin for the query cache and score scratch space.

Memory-bound regime: The kernel is designed for the memory-bound regime typical of long-context decode. The assistant explicitly notes that "the memory-bound nature of this kernel means extra tile iterations won't hurt performance as long as loads stay coalesced." This assumption justifies the relatively small TK=16 tile size—more tile iterations are acceptable because the bottleneck is global memory bandwidth, not compute.

Coalesced access patterns: The assistant assumes that float4 loads from the KV cache will be coalesced across threads. This is a standard CUDA optimization but depends on the memory layout of the KV cache. The MLA cache stores per-token latent vectors contiguously, so a tile of TK tokens can be loaded with coalesced float4 reads.

Online softmax numerical stability: The flash attention algorithm uses online softmax (running max and sum) to avoid storing full score rows. The assistant assumes this is numerically stable for the FP32 precision used in Phase 1, with the plan to validate against the oracle baseline.

BF16 migration in Phase 2: The assistant assumes that switching to BF16 I/O in Phase 2 will halve the KV cache footprint, making HT=16 feasible for a 16× reduction. This is a reasonable projection but depends on the BF16 hardware support on sm_120 and the numerical accuracy requirements of the verification step.

What Input Knowledge Is Required to Understand This Message

To fully grasp the reasoning in [msg 12227], one needs familiarity with several domains:

Multi-head Latent Attention (MLA): The attention mechanism used by DeepSeek models (DeepSeekV3, Kimi K2.6). MLA compresses the key-value cache into a latent representation (kv_c, dimension 512) plus a decoupled RoPE key (k_pe, dimension 64). The query is similarly decomposed into absorbed (q_nope) and positional (q_pe) components. The verification kernel computes attention scores as dot products between query components and their corresponding key components, then applies a tree-structured mask.

Flash Attention Algorithm: The technique of tiling the attention computation across the key-value sequence, using online softmax to avoid materializing the full attention matrix. The assistant's design follows the standard flash attention pattern: load a KV tile into shared memory, compute partial scores, update running softmax state (max and sum), accumulate weighted values, then normalize at the end.

CUDA Shared Memory Architecture: The assistant navigates the sm_120 shared memory hierarchy, balancing the 100KB limit against the needs of KV tiles (TK × (Dl + Dr) × sizeof(float)), query cache (HT × Dl × sizeof(float)), accumulators (HT × Dl × sizeof(float)), and score scratch space (HT × TK × sizeof(float)).

Speculative Decoding with DFlash: The verification step in DFlash uses a tree-structured attention mask where each query position attends to a subset of the KV cache determined by the drafter's candidate tree. The kernel must apply this mask during scoring, setting masked positions to -inf before softmax.

The Output Knowledge Created

Message [msg 12227] creates both tangible and intangible knowledge:

Tangible: The file verify_attn_flash.cu—a CUDA kernel implementing flash attention for MLA verification on sm_120. The kernel uses a 2D grid (batch × query position × head tile), shared memory for KV tiles and query vectors, online softmax for arbitrary prefix lengths, and tree-mask support for speculative decoding.

Intangible: A documented design rationale that future developers can reference. The reasoning section captures the trade-offs considered and rejected, the performance projections (16× reduction in KV cache reads), and the migration path to BF16 in Phase 2. This is valuable because kernel design decisions are often lost after implementation—the reasoning in [msg 12227] preserves the "why" behind the code.

Process knowledge: The assistant's workflow pattern—write locally, sync in batch, test once—is itself a piece of output knowledge. It demonstrates a strategy for minimizing remote iteration cycles when the design is mature enough to commit to in a single pass.

Mistakes and Incorrect Assumptions

The reasoning in [msg 12227] contains several points where the assistant corrects itself or where assumptions later proved incorrect:

The HT=16 shared memory miscalculation: The assistant initially believes HT=16 with TK=16 fits in ~72KB, then later calculates HT=16 with TK=16 at over 100KB. This inconsistency suggests the assistant was refining its shared memory budget model during the reasoning process. The final settled configuration (HT=8, TK=16) is conservative and safe.

The register-tiled v3 overreach: The assistant spent significant design effort on a fully register-tiled kernel (all 64 heads in one block, accumulators in registers) before concluding it was too complex for Phase 1. This exploration was not wasted—it informed the final design and provides a roadmap for future optimization—but it represents an incorrect initial assessment of what was achievable within the time budget.

The assumption that thread occupancy matters: The assistant worries that with HT=8 and TK=16, only 128 threads participate in the score step while 256 are available. This concern about thread occupancy is valid for compute-bound kernels but may be irrelevant for a memory-bound kernel where the bottleneck is global memory bandwidth, not thread utilization. The assistant implicitly acknowledges this by proceeding with the design anyway.

The 16× reduction projection: The assistant calculates a 16× reduction in KV cache reads (naive/(2×HT) with HT=8). This is a theoretical upper bound that assumes perfect coalescing and no overhead from tile boundary handling. In practice, the achieved speedup will be lower due to shared memory bandwidth limits, synchronization overhead, and the cost of the reduction kernel (which the assistant later discovers is a significant bottleneck in microbenchmarks).

The Broader Significance

Message [msg 12227] exemplifies a pattern that recurs throughout the session: the assistant operating at the intersection of deep algorithmic knowledge and practical hardware constraints. The design exploration is not abstract—it is grounded in the specific capabilities of the sm_120 architecture, the MLA attention mechanism of Kimi K2.6, and the real-time performance requirements of a production inference service.

The message also reveals the assistant's disciplined approach to decision-making. When faced with multiple viable designs, it does not default to the simplest or the most complex—it systematically evaluates trade-offs, projects performance, and selects a configuration that balances correctness (FP32 oracle validation), performance (16× read reduction), and implementability (shared memory accumulators over register tiling).

This discipline pays off in the subsequent messages. The kernel passes token-exact correctness tests against the oracle, and after tuning (increasing NSPLIT from 16 to 64, adding 128-bit vectorized loads), it delivers a 3–6× end-to-end decode speedup over the Triton baseline across all context lengths from 4k to 65k tokens. The design decisions made in [msg 12227]—the tile sizes, the shared memory layout, the head-grouping strategy—proved sound.

But perhaps the most important lesson from this message is the value of preserving design reasoning. The assistant's extensive thinking section, captured as part of the message, serves as a design document that explains not just what was built but why it was built that way. In a field where kernel optimization often relies on tribal knowledge and trial-and-error, this kind of documented reasoning is rare and valuable. It turns a single code-writing action into a teachable moment for anyone who wants to understand how to build custom CUDA kernels for modern attention mechanisms on emerging hardware architectures.