Designing a Custom sm_120 Verify Attention Kernel: Inside the AI's Design Reasoning

Introduction

In the course of a high-stakes machine learning deployment session, a critical moment arrives. The assistant, tasked with optimizing speculative decoding throughput on NVIDIA's new Blackwell RTX PRO 6000 GPUs (compute capability sm_120), pauses from execution to engage in an extended period of architectural reasoning. This message—index 12225 in the conversation—is not a tool call, not a code edit, not a command execution. It is a pure thinking trace, spanning hundreds of lines of internal deliberation, as the assistant designs a custom CUDA kernel from first principles.

The message reveals something rare: the moment-by-moment evolution of an AI's design process, complete with dead ends, trade-off analyses, constraint discoveries, and eventual convergence on a solution. The assistant is designing a "flash MLA verify attention kernel"—the computational heart of speculative decoding with DeepSeek-style Multi-head Latent Attention (MLA)—for an architecture (sm_120) that no existing optimized library supports. This article examines that message in depth: the reasoning that drove it, the decisions it made, the assumptions it relied on, and the knowledge it both consumed and produced.

The Strategic Context: Why This Message Exists

To understand message 12225, we must first understand the project it serves. The broader session (Segment 66 of the conversation) is focused on deploying the Kimi K2.6 model with DFlash speculative decoding on Blackwell hardware. The team has already achieved a working deployment with Triton-based attention kernels, but performance at long context lengths is severely bottlenecked. At 185k tokens, decode throughput drops to 0.7 tokens per second, with GPU tensor core utilization at only ~3% despite 99.8% SM occupancy reported by nvidia-smi.

The root cause was identified in the previous segment: the DDTree verify attention step—the phase where the target model evaluates draft tokens proposed by the speculative drafter—is locked to Triton MLA with page_size=1, causing scattered KV access at approximately 14 GB/s effective bandwidth, roughly 130× below the 1.8 TB/s peak of the RTX PRO 6000. The Triton-generated kernels cannot exploit the hardware's full memory bandwidth because they use a naive implementation that reads the shared MLA latent representation once per head, resulting in a 64× redundant read pattern.

The user's directive, established in the messages preceding 12225, is clear: build a custom sm_120 verify attention kernel from scratch, integrate it into the live SGLang serving stack, and implement KV cache defragmentation. The user explicitly approved the "owned kernel" approach, noting the smaller 100KB shared memory constraint of sm_120 (compared to 228KB on Hopper architectures). Message 12225 is the assistant's response to this directive—but rather than jumping immediately into code, the assistant first engages in a deep design exploration.

The message begins with the assistant setting up its execution plan: "I'm laying out the execution plan: first I'll create the spec document at plans/0002-sm120-verify-kernel-defrag.md with detailed phase breakdowns, then move into Phase 0/1." But within a few sentences, the assistant is pulled into a vortex of architectural reasoning, unable to proceed without first resolving fundamental design questions about thread mapping, shared memory allocation, and data flow.

The Core Problem: MLA's 64× Latent Read Bottleneck

The message's central intellectual challenge is the MLA verify attention computation. In DeepSeek's MLA formulation, the KV cache stores a per-token compressed latent representation (kv_c, dimension 512) and a decoupled RoPE key (k_pe, dimension 64). This latent is shared across all 64 heads—a design choice that dramatically reduces KV cache memory at the cost of requiring each head to read from the same shared representation.

The naive attention kernel, as implemented in the existing verify_attn.cu, processes one head at a time. For each head, it loads the full KV latent from global memory, computes dot products with the query, applies the softmax, and accumulates the output. This means the same 512-element KV latent vector is loaded from global memory 64 times—once per head. For a prefix of 100k tokens, this 64× redundancy translates to reading 64 × 100k × 512 × 2 bytes ≈ 6.4 GB of data that could have been read once and reused.

The assistant identifies this as the primary bottleneck: "The docs specifically call out that re-reading the shared MLA latent per head (64×) is the biggest bottleneck." The entire design effort in message 12225 is driven by the goal of eliminating this redundancy.

The Design Space: Exploring Kernel Architectures

The assistant's reasoning in message 12225 is structured as an iterative exploration of the design space, with each approach evaluated against the constraints of sm_120 hardware and the MLA computation pattern. Let us trace the evolution.

Approach 1: One Block Per (Batch, Query, Head) — The Naive Baseline

The assistant initially considers the simplest possible flash attention design: one CUDA block per (batch position, query position, head). Each block loads a tile of KV tokens into shared memory, computes scores for its single head, and accumulates the output. This approach is straightforward to implement and maps naturally onto CUDA's thread hierarchy.

However, the assistant immediately recognizes the problem: "But storing 64×512=32,768 float accumulators exceeds shared memory." If each block handles only one head, the KV latent must still be loaded 64 times across the 64 blocks that process different heads for the same query position. The latent read redundancy is not eliminated—it is merely shifted from global memory reads within a single kernel to separate kernel launches or separate blocks.

The assistant rejects this approach because it fails to address the core bottleneck.

Approach 2: One Block Per (Batch, Query) with Full Register Tiling

The assistant then pivots to a more ambitious design: one block processes all 64 heads for a single query position, with the output accumulator distributed across thread registers rather than stored in shared memory. The idea is that each thread holds a slice of the output matrix, and the KV latent is loaded into shared memory once per block and reused across all heads.

The assistant works through the thread mapping: "64 heads × 4 threads per head = 256 threads total, with each thread holding 128 accumulator floats for its assigned head and dimension range." For scoring, the four threads sharing a head cooperate to compute dot products between the query and each key in the KV tile, using warp shuffles to reduce partial dot products across the thread group.

This approach achieves the ideal—a single global memory read of the KV latent per query position—but at the cost of significant implementation complexity. The assistant wrestles with the details: "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."

The assistant eventually steps back from this approach, acknowledging the complexity: "But making this fully flexible across all test configs is getting messy—I think the cleaner approach is to accept a compromise."

Approach 3: Grouped Heads with Shared Memory Accumulators

The compromise is a grouped-heads design. Instead of processing all 64 heads in one block or one head per block, the assistant proposes processing a tile of HT heads per block, with HT being a tunable constant. The accumulator for the head group fits in shared memory alongside the KV tile, and the KV latent is loaded once per block and reused across the HT heads.

The assistant calculates the memory budget: "HT=16 with Tk=16 gives me about 72KB total (KV tile, accumulators, scores), which fits comfortably under the 100KB limit while keeping latent reads to a reasonable 4×." With HT=16, the 64 heads are processed in 4 groups of 16, reducing the latent read redundancy from 64× to 4×—a 16× improvement.

But then the assistant hits a wall: the shared memory calculation reveals that the full design, including query cache and score scratch space, pushes beyond the 100KB limit. "For HT=16 and Tk=16, the allocations add up to over 100KB, which exceeds the typical 96KB or 100KB limit on older architectures."

The assistant iterates: "I'm considering reducing HT to 8 and increasing Tk to 32 to balance the memory usage, though that still pushes close to the limit." With HT=8, the latent read redundancy is 8× (64/8), still an 8× improvement over the naive approach.

The Final Configuration

The assistant settles on HT=8 and Tk=32, estimating approximately 91KB of shared memory usage. This configuration:

Shared Memory Constraints and the sm_120 Reality

A recurring theme in message 12225 is the assistant's struggle with the sm_120 shared memory limit. The Blackwell RTX PRO 6000 (sm_120) has approximately 100KB of shared memory per block—significantly less than the 228KB available on Hopper (sm_90) architectures. This constraint shapes every design decision.

The assistant repeatedly recalculates the shared memory budget:

Precision Strategy: FP32 First, BF16 Later

Another key decision in message 12225 concerns numerical precision. The existing oracle and reference implementations use FP32 (float), while the live SGLang service operates in BF16 (bfloat16) for the KV cache. The assistant must decide which precision to target in the initial kernel.

The reasoning is pragmatic: "For now, I'll keep Phase 1 focused on validating the core algorithmic changes—streaming softmax and single-latent-read restructuring—using FP32 to match the existing oracle, then add BF16 I/O support in Phase 2."

This decision reflects a sound engineering principle: separate correctness validation from optimization. By keeping FP32 in Phase 1, the assistant can verify that the streaming softmax and grouped-heads architecture produce token-exact results compared to the naive oracle. Once correctness is established, BF16 I/O can be added in Phase 2, with the shared memory pressure naturally easing since BF16 values are half the size of FP32.

The assistant also notes: "For Phase 2 when I integrate bf16 inputs, the shared memory pressure naturally eases since the data is already in lower precision, so I can afford larger HT values then." This observation shows forward thinking—the Phase 1 design constraints (HT=8, Tk=32) are temporary; the BF16 Phase 2 can use more aggressive tiling.

Assumptions Made by the Assistant

Message 12225 rests on several assumptions, some explicit and some implicit:

  1. sm_120 shared memory limit is ~100KB: The assistant repeatedly calculates against this limit. This assumption is based on the user's earlier statement about "the smaller 100KB shared memory constraint of sm_120 (compared to 228KB on Hopper architectures)." If the actual limit differs, the tiling parameters would need adjustment.
  2. The 64× latent read is the primary bottleneck: The assistant assumes that eliminating this redundancy will yield proportional speedups. This is a reasonable assumption given the measured 14 GB/s effective bandwidth (130× below peak), but the actual speedup depends on whether other factors (instruction throughput, memory latency hiding, occupancy) become limiting once the bandwidth bottleneck is addressed.
  3. Online softmax is numerically stable for long prefixes: The assistant plans to use the standard online softmax algorithm (Griffiths & Dowling, 2019), which rescales accumulators as new tiles arrive. This algorithm is well-established but requires careful handling of floating-point precision for very long sequences (100k+ tokens).
  4. The grouped-heads approach with HT=8 provides sufficient occupancy: With HT=8 heads per block and 256 threads per block, the assistant assumes this configuration will achieve adequate SM occupancy. The actual occupancy depends on register usage, shared memory allocation, and the scheduler's ability to hide latency.
  5. The mask can be applied efficiently on the final tile: The assistant plans to apply the DDTree visibility mask only on the final KV tile (the tail region), processing the prefix unmasked. This assumes that the prefix is always fully visible, which is correct for the DDTree speculative decoding pattern.
  6. The C-ABI interface from Phase 0 will be reusable: The assistant assumes that the existing C-ABI shim (capi.cu) can be extended to expose the new flash kernel without major restructuring.

Mistakes and Incorrect Assumptions

While the assistant's reasoning is generally sound, several points merit critical examination:

  1. The shared memory calculation may be optimistic: The assistant estimates ~91KB for HT=8, Tk=32, but this calculation appears to omit several overheads: the score scratch space (HT × Tk = 256 floats = 1KB), the online softmax state (HT × 2 = 16 floats = negligible), and alignment padding. The actual footprint may be closer to 95-100KB, leaving little margin. If the kernel fails to launch due to shared memory limits, the assistant will need to reduce Tk or HT further.
  2. The register-tiled all-heads approach was abandoned too quickly: The assistant dismissed the full register-tiling approach (one block per query, all 64 heads) as "getting messy" and "complex." However, this approach would have achieved the ideal 1× latent read (vs. 8× in the compromise). Given that the 64×→8× improvement still leaves a non-trivial 8× redundancy, the full register-tiling approach might have been worth the extra implementation effort, especially since the kernel is a long-lived component of the production stack.
  3. The assumption that CPU orchestration is the bottleneck was disproven later: The chunk summary reveals that the assistant later discovered (through profiling) that CPU orchestration (tree-build at 1.8ms, mask-build at 0.18ms) was negligible—the real bottleneck was the verify attention itself. This suggests that the assistant's initial focus on the verify kernel was correct, but the reasoning in message 12225 does not yet have this empirical confirmation.
  4. The assumption that BF16 will "naturally ease" shared memory pressure is partially correct: BF16 halves the KV cache footprint, but the query vectors and accumulators remain in FP32 (since accumulation requires higher precision). The assistant correctly notes this but may underestimate the complexity of mixed-precision handling in the kernel.

Input Knowledge Required to Understand This Message

Message 12225 draws on a substantial body of technical knowledge:

  1. CUDA programming model: Thread blocks, shared memory, warp shuffles, grid/block hierarchy, register allocation. The assistant reasons about thread-to-data mappings, block-level parallelism, and shared memory budgets throughout.
  2. Flash Attention algorithm: The online softmax technique (tiling the KV sequence, maintaining running max and sum, rescaling accumulators). The assistant's entire kernel design is built on this foundation.
  3. Multi-head Latent Attention (MLA): DeepSeek's MLA formulation with absorbed KV, shared latent representation, and decoupled RoPE. The assistant must understand the data layout (kv_c of dimension 512, k_pe of dimension 64, per-head queries) to design the kernel.
  4. Speculative decoding with DDTree: The verify phase, tree masks, prefix/tail split, commit length. The assistant's mask handling strategy (unmasked prefix, masked tail) depends on understanding the DDTree algorithm.
  5. NVIDIA GPU architecture: sm_120 compute capability, shared memory limits, memory bandwidth, tensor core availability. The assistant references the 100KB shared memory limit and the 1.8 TB/s peak bandwidth.
  6. SGLang serving stack: Attention backend interface, CUDA graph capture, TP (tensor parallelism) rank handling. The Phase 2 integration plan depends on understanding SGLang's architecture.
  7. The existing codebase: The verify_attn.cu naive kernel, the C-ABI shim in capi.cu, the test harness in test_verify_attn.cu, the benchmark in bench_kernels.cu, the reference generator in gen_verify_attn_refs.py, and the build scripts.

Output Knowledge Created by This Message

Message 12225 produces several forms of output knowledge:

  1. A concrete kernel architecture: The grouped-heads flash MLA verify kernel with HT=8, Tk=32, online softmax, and single-latent-read per block. This architecture is documented in the assistant's reasoning and will be implemented in verify_attn_flash.cu.
  2. A phased execution plan: The assistant establishes a clear sequence: Phase 0 (scaffolding and parity harness), Phase 1 (flash kernel implementation and validation), Phase 2 (SGLang backend integration), Phase 3 (live validation), Phase 4 (KV defragmentation). Each phase has gates that must be passed before the next begins.
  3. Design trade-off documentation: The assistant's reasoning captures the exploration of alternative designs (one-head-per-block, all-heads register-tiled, grouped-heads) and the rationale for the chosen approach. This serves as implicit documentation for future developers who might ask "why was it done this way?"
  4. Shared memory budget analysis: The assistant calculates the memory footprint of each component (KV tile, query cache, accumulators, scores, softmax state), providing a reference for future tuning.
  5. Precision strategy: The decision to use FP32 in Phase 1 and add BF16 in Phase 2 establishes a clear correctness-first methodology.
  6. A todo list and tracking system: The assistant creates a structured todo list with priorities and status, enabling systematic execution across multiple phases.
  7. Risk identification: The assistant identifies several risks: cuda-graph capture-safety, 100KB smem tile pressure, exact MLA-absorb/bf16 numeric parity with triton, TP8 per-rank correctness, Tier-1 defrag vs radix sharing, and GPU contention with the live service.

The Thinking Process: A Window into AI Design Reasoning

Perhaps the most remarkable aspect of message 12225 is what it reveals about the assistant's thinking process. Unlike a human engineer who might sketch on a whiteboard, type notes in a document, or think silently before coding, the assistant externalizes its entire reasoning chain in natural language. This creates an unprecedented window into AI design cognition.

The reasoning exhibits several characteristic patterns:

Iterative refinement: The assistant does not arrive at the final design in one leap. It starts with a simple idea (one block per head), identifies the flaw (latent still read 64×), proposes a more complex alternative (all heads per block with register tiling), encounters implementation complexity, and settles on a compromise (grouped heads with shared memory). Each iteration is triggered by a constraint or a discovered limitation.

Constraint propagation: The assistant constantly checks design decisions against hardware constraints. The shared memory limit is the dominant constraint, but the assistant also considers thread count limits (1024 threads per block), warp size (32 threads), and the need for thread-count-to-work-item divisibility.

Mental simulation: The assistant simulates the kernel's execution in its "mind's eye," tracing the flow of data from global memory to shared memory to registers, through the score computation, softmax rescaling, and output accumulation. This mental simulation allows the assistant to identify issues (e.g., "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") before writing any code.

Self-correction: The assistant frequently catches its own mistakes. For example, after proposing a complex register-tiling scheme, it steps back: "Actually, wait—the docs specifically call out that re-reading the shared MLA latent per head (64×) is the biggest bottleneck. So I can't ignore that; I need the all-heads-per-block design to amortize those reads." This self-correction is a hallmark of careful reasoning.

Pragmatic compromise: Despite the ideal of a single-pass, all-heads kernel, the assistant repeatedly chooses simpler, more implementable designs over theoretically optimal ones. "The first kernel with 8× latency reduction and proper coalescing is solid and correct; I can tune the hyperparameters once I see the microbench results." This reflects an understanding that engineering progress requires shipping something that works, even if it's not perfect.

Meta-cognitive awareness: The assistant is aware of its own design process: "I need to stop designing and start executing: write a spec document, implement the verification kernel, wire up the C API, add comparison tests and benchmarks, then build and run everything on the CT200." This ability to recognize when analysis has reached diminishing returns and execution should begin is a sophisticated cognitive skill.

Conclusion

Message 12225 is a microcosm of the entire engineering process: problem identification, design space exploration, constraint analysis, trade-off evaluation, and execution planning—all compressed into a single reasoning trace. It captures the moment when an AI assistant transitions from understanding a problem to designing a solution, and it does so with a level of detail and transparency that is rare even in human engineering documentation.

The message's significance extends beyond the specific kernel it designs. It demonstrates that AI systems can engage in genuine design reasoning—not just retrieving and applying known patterns, but exploring novel configurations, evaluating trade-offs, and making informed decisions under constraints. The assistant does not simply implement a flash attention kernel; it designs one, adapting general principles (online softmax, tiling, shared memory reuse) to the specific constraints of sm_120 hardware and MLA computation.

For the project, this message sets the technical direction that will ultimately yield a 3–6× decode speedup over the Triton baseline, as the chunk summary reveals. The kernel designed in this message, once implemented, tuned, and integrated, becomes a key component of the production SGLang stack. The defragmentation work planned alongside it addresses the remaining bottleneck after attention is optimized.

For the observer, message 12225 offers a rare glimpse into the cognitive process of an AI system engaged in creative technical work. It shows that AI reasoning, at its best, is not a black box that produces answers from unknown internal processes, but a transparent, structured, and ultimately human-comprehensible chain of thought—one that can be studied, critiqued, and learned from.