The Plan Document: A Pivot Point in Custom CUDA Kernel Development

Introduction

In the midst of a high-stakes optimization session for speculative decoding on NVIDIA RTX PRO 6000 Blackwell (sm_120) GPUs, the assistant issued a single, deceptively simple message: it wrote a file. The file was plans/0002-sm120-verify-kernel-defrag.md, a specification document that would chart the course for building a custom CUDA kernel to replace Triton-generated attention kernels in the SGLang inference serving stack. The message itself ([msg 12226]) contains only the confirmation of a file write operation, but it represents the crystallization of an extraordinarily detailed design process spanning dozens of architectural tradeoffs, hardware constraints, and algorithmic innovations. This article examines that pivotal moment—the reasoning that motivated it, the decisions it encoded, and the knowledge it created.

Context: The Bottleneck That Demanded a Custom Kernel

The assistant was deep in a multi-session effort to deploy the Kimi K2.6 large language model with DFlash speculative decoding on Blackwell GPUs. A preceding investigation ([msg 12225]) had revealed a severe performance bottleneck: the DDTree verify attention kernel—the component responsible for computing attention scores during the speculative decoding verification step—was locked to a Triton-based MLA (Multi-head Latent Attention) implementation that used page_size=1 KV cache paging. This caused scattered, non-coalesced memory accesses that achieved only approximately 14 GB/s of effective bandwidth, roughly 130× below the GPU's theoretical 1.8 TB/s peak. The result was catastrophic decode slowdown: at 185k context length, throughput collapsed to 0.7 tokens per second.

The user had directed the assistant to investigate two avenues: optimized kernel integration and K/V defragmentation. The assistant discovered a critical hardware limitation—all optimized MLA kernels (FlashMLA, cutlass-MLA, flashinfer-MLA) were compiled only for sm_90a (Hopper), sm_100a, and sm_103a architectures. None supported sm_120, the compute capability of the RTX PRO 6000 Blackwell consumer GPU. These kernels relied on Hopper/Blackwell-DC instructions like wgmma, TMA, and tcgen05 that simply do not exist on sm_120's Ada-like instruction set architecture. The only path forward was to build a custom kernel from scratch.

WHY This Message Was Written: The Need for a Roadmap

The plan document was written because the assistant recognized that building a custom CUDA attention kernel is not a single task but a multi-phase engineering project requiring careful sequencing. The reasoning in [msg 12225] reveals that the assistant had been iterating through kernel designs in its head—at least three distinct architectural approaches—and needed to commit to a concrete plan before beginning implementation.

The motivation was threefold. First, the assistant needed to align with the user on the overall strategy before investing significant implementation effort. Writing a plan document served as a communication artifact that the user could review, modify, or approve. Second, the project involved multiple interdependent phases—baseline validation, kernel implementation, SGLang backend integration, and KV defragmentation—each of which could block the others if ordered incorrectly. Third, the assistant needed to externalize its own reasoning to avoid losing track of design decisions during implementation. The plan document became the authoritative reference for "why was HT chosen as 8?" and "why was TK set to 16?"—questions that would be difficult to reconstruct from code alone.

The Design Reasoning: A Deep Dive into the Thinking Process

The reasoning in [msg 12225] reveals an extraordinarily thorough design exploration. The assistant considered at least three major architectural variants for the flash attention kernel before settling on the approach documented in the plan.

Variant 1: One block per (batch, position, head). This was the simplest approach, where each CUDA block handles a single query position for a single head. It would be easy to implement and correct, but it would read the KV cache from global memory once per head—64 redundant reads of the same 512-dimensional latent vector. The assistant correctly identified this as insufficient, noting that "re-reading the shared MLA latent per head (64×) is the biggest bottleneck."

Variant 2: One block per (batch, position) with all heads, register-tiled accumulators. This ambitious design would have a single block process all 64 heads simultaneously, distributing the 64×512 output accumulator across threads' registers. Each thread would hold 128 floats for its assigned head and dimension range. The KV cache would be read from global memory only once per query position, achieving a theoretical 64× reduction in global reads. However, the assistant identified a critical complexity: when a new KV tile updates the running softmax maximum for a head, all 512 floats of that head's accumulator spread across threads need rescaling, requiring careful thread-to-(head, dimension) mapping and warp-level coordination. The complexity of managing register-based accumulators across warp shuffles for both the scoring (QK) and accumulation (PV) phases was deemed too high for a first implementation.

Variant 3 (chosen): One block per (batch, position, head-tile) with shared-memory accumulators. This design splits heads into groups (HT=8 heads per block), stores accumulators in shared memory, and tiles the KV cache (TK=16 tokens per tile). The KV cache is read once per block and reused across the 8 heads in the group, reducing global reads from 64× to 8×. The assistant calculated the shared memory budget meticulously: KV tile (16 tokens × 512 latent dims × 4 bytes = 32KB), query vectors (8 heads × 512 dims × 4 bytes = 16KB), output accumulators (8 heads × 512 dims × 4 bytes = 16KB), and score scratch space (8 heads × 16 tokens × 4 bytes = 512 bytes), plus online softmax state—totaling approximately 72KB, safely under the 100KB shared memory limit of sm_120.

Assumptions and Their Implications

The assistant made several key assumptions in the design. First, it assumed that the dominant bottleneck was global memory bandwidth for KV cache reads rather than compute throughput or instruction latency. This was supported by the earlier profiling showing ~3% tensor core utilization during decode, but it was still an assumption that would need validation once the kernel was running. Second, the assistant assumed that shared memory capacity on sm_120 was 100KB—a figure derived from the RTX PRO 6000 Blackwell's specifications but not yet confirmed through direct query. Third, it assumed that the existing FP32 oracle and reference implementations would be sufficient for correctness testing, deferring BF16 I/O support to Phase 2.

The most significant assumption was that HT=8 and TK=16 represented the optimal balance between latency reduction and shared memory pressure. The assistant acknowledged this was a starting point: "I can tune the hyperparameters once I see the microbench results." This assumption would later prove partially incorrect—in the actual implementation phase ([msg 12227] onward), the assistant would discover that occupancy in the TP8 regime (only 8 heads per rank) required different tuning, eventually increasing NSPLIT from 16 to 64 and adding 128-bit vectorized loads to achieve a 3–6× speedup over Triton.

Input Knowledge Required

Understanding this message requires substantial background knowledge. The reader must understand the MLA (Multi-head Latent Attention) architecture used in DeepSeek V3 and Kimi K2.6 models, where a compressed latent representation (kv_c, dimension 512) is shared across all attention heads and combined with per-head queries. The reader must understand CUDA memory hierarchy—global memory, shared memory, registers—and the performance characteristics of each. Knowledge of online softmax (safe softmax with running maximum and sum) is essential, as is familiarity with tiled matrix computations and warp-level primitives like __shfl_xor_sync. The sm_120 architecture constraints (no TMA, no wgmma, 100KB shared memory limit) are critical context that explains why existing optimized kernels could not be used.

Output Knowledge Created

The plan document created a structured roadmap with four phases: Phase 0 (baseline validation and parity harness), Phase 1 (flash attention kernel with KV tiling and online softmax), Phase 2 (SGLang backend integration with BF16 support), and Phase 3 (KV defragmentation). It established the HT=8, TK=16, BLOCK=256 parameterization as the starting point for kernel development. It documented the shared memory layout and thread mapping strategy, creating a reference that would guide the implementation in subsequent messages. Most importantly, it committed the assistant to a concrete execution path, transforming an abstract design space into an actionable sequence of coding tasks.

Conclusion

The message that wrote plans/0002-sm120-verify-kernel-defrag.md appears trivial—a one-line confirmation of a file operation. But it represents the culmination of an intensive design process and the launch point for a complex multi-phase engineering effort. The reasoning that preceded it ([msg 12225]) reveals the assistant systematically exploring architectural alternatives, calculating shared memory budgets, weighing complexity against performance, and ultimately committing to a design that balanced correctness, performance, and implementability. This pattern—extensive internal reasoning followed by a concise external action—is characteristic of effective technical work, and the plan document it produced would prove essential in guiding the subsequent 3–6× decode throughput improvement that the team would achieve.