The Last Piece of Phase 1: Extending the Benchmark to Validate a Custom Flash Attention Kernel

In the midst of an intensive session optimizing speculative decoding for the Kimi K2.6 language model on RTX PRO 6000 Blackwell GPUs, the assistant issued a short but consequential message. At first glance, it appears mundane: an edit to a benchmark file. But this message—message index 12233 in the conversation—represents the final, deliberate step in a carefully orchestrated Phase 1 of a multi-phase plan to build a custom CUDA kernel from scratch. Understanding why this message was written, and what it accomplished, requires unpacking the complex chain of reasoning, discovery, and architectural decisions that led to this moment.

The Message Itself

The assistant wrote:

## Agent Reasoning Now I'm updating the benchmark kernel to test both the naive and flash attention implementations, and I'll expand the prefix sizes to include much larger values up to 16384 to better stress-test the performance characteristics. [edit] /home/theuser/glm-kimi-sm120-rtx6000bw/kdtree-engine/tests/bench_kernels.cu Edit applied successfully.

This is the entirety of the message. It consists of a reasoning block (the agent's internal monologue) followed by a tool call that performed a file edit. The edit itself was applied to the benchmark harness bench_kernels.cu, which is the microbenchmarking tool for measuring kernel latency across various configurations.

The Crisis That Precipitated This Work

To understand why this message exists, we must go back to the root cause that triggered the entire kernel development effort. Earlier in the session, the assistant had deployed the Kimi K2.6 model with a 200k-token context length on the CT200 server. A comprehensive long-context benchmark revealed a catastrophic decode slowdown: at 185k tokens, throughput dropped to 0.7 tokens per second, with GPU tensor core utilization at only ~3%. This was despite nvidia-smi reporting 99.8% SM occupancy—a classic symptom of a memory-bound kernel that is spending nearly all its time waiting on global memory loads.

The assistant diagnosed the root cause: the DDTree verify attention kernel was locked to Triton's MLA implementation with page_size=1, which caused scattered, non-coalesced KV cache accesses. The effective bandwidth was approximately 14 GB/s, which is 130× below the GPU's theoretical peak of 1.8 TB/s. In essence, the GPU was sitting almost entirely idle, waiting for data to trickle in from memory.

The user then directed the assistant to investigate two paths: integrating an optimized kernel and implementing K/V defragmentation. The assistant discovered that all existing optimized MLA kernels—FlashMLA, cutlass-MLA, flashinfer-MLA—were compiled only for sm_90a (Hopper), sm_100a, or sm_103a architectures. None supported sm_120, the compute capability of the RTX PRO 6000 Blackwell consumer GPU. These kernels rely on Hopper/Blackwell-DC instructions like wgmma, TMA, and tcgen05 that simply do not exist on sm_120's Ada-like instruction set architecture. There was no off-the-shelf solution.

The Decision to Build an Owned Kernel

With no existing optimized kernel available for their hardware, the assistant and user made a pivotal decision: build a custom sm_120 verify attention kernel from scratch. The user approved this approach, noting the smaller 100KB shared memory constraint on sm_120 compared to the data-center Blackwell variants.

The assistant responded by creating a detailed per-phase plan document (plans/0002-sm120-verify-kernel-defrag.md) and began executing Phase 1: implementing a KV-split flash-decode MLA verify kernel. The core idea was to use a flash-attention-style tiled approach where the KV cache is loaded into shared memory once per block and reused across multiple heads, rather than having each head independently fetch the same KV data from global memory. This promised to reduce global memory reads by a factor proportional to the number of heads processed per tile—potentially 8× to 16×.

The assistant wrote the kernel (verify_attn_flash.cu), wired it into the C API header and implementation, updated the build script to compile the new source file, extended the correctness test suite to compare the flash kernel against both the naive kernel and a Python golden reference, and improved the test output to include kernel names in PASS/FAIL messages. Each of these steps was a separate message in the conversation.

Why This Message Matters

Message 12233 is the capstone of Phase 1 implementation. The assistant had already written the kernel, integrated it into the build and test infrastructure, and added correctness validation. What remained was ensuring that the benchmark harness—the tool that would measure whether the new kernel actually improved performance—could test both the naive and flash implementations side by side, and across a sufficiently wide range of prefix lengths to expose the kernel's strengths and weaknesses.

The decision to expand prefix sizes up to 16384 tokens was not arbitrary. The flash kernel's design centers on tiling the KV cache: it processes the prefix in chunks (tiles) of Tk tokens at a time, using online softmax to avoid storing full attention score matrices. At short prefix lengths (e.g., 256 tokens), the overhead of tile iteration and shared memory management can outweigh the benefits of coalesced access, making the flash kernel potentially slower than the naive approach. At long prefixes (thousands of tokens), the amortization of global memory reads across heads and tiles should dominate, yielding substantial speedups. The benchmark needed to capture this crossover point to validate the design.

Furthermore, testing both kernels in the same benchmark run was essential for fair comparison. The naive kernel served as the baseline—it was the existing Triton-based implementation that the team was trying to replace. By instrumenting the benchmark to call both kernels with identical inputs and measure their latencies, the assistant could generate the data needed to confirm whether the flash kernel actually delivered the promised memory-bandwidth reduction.

The Thinking Process Visible in the Reasoning

The assistant's reasoning reveals a clear understanding of what the benchmark must demonstrate. The phrase "to test both the naive and flash attention implementations" shows that the assistant recognized the need for a controlled comparison—not just measuring the new kernel in isolation, but benchmarking it against the existing implementation under identical conditions. This is a sign of rigorous experimental design.

The expansion "to include much larger values up to 16384" indicates the assistant was thinking about the kernel's performance profile across the full operating range. The flash kernel's tiled approach is expected to perform differently at different context lengths: at short prefixes, the overhead of loading tiles into shared memory and managing online softmax state may dominate; at longer prefixes, the reduction in global memory traffic should become the dominant factor. By sweeping prefix sizes from small to large, the benchmark would reveal the crossover point and confirm the kernel's asymptotic behavior.

The assistant also implicitly assumed that the benchmark harness already had the infrastructure to call multiple kernel variants. This assumption was reasonable given that the benchmark file was designed as a microbenchmark for custom kernels, and the assistant had already modified the test suite to support A/B comparison. The benchmark extension was the natural next step.

Input Knowledge Required

To understand this message, one needs to know several things:

  1. The architecture context: The RTX PRO 6000 Blackwell GPU uses sm_120 compute capability, which lacks the specialized tensor memory accelerator (TMA) and warp-group matrix multiply-accumulate (wgmma) instructions found in Hopper (sm_90) and data-center Blackwell (sm_100+) GPUs. This forced the team to build a custom kernel rather than using existing optimized libraries.
  2. The flash attention algorithm: The kernel uses a tiled approach where the KV cache is streamed into shared memory in chunks, and online softmax is used to compute attention outputs incrementally without storing full score matrices. This is a well-known technique from the FlashAttention literature, adapted here for the MLA (Multi-head Latent Attention) architecture used by DeepSeek and Kimi models.
  3. The benchmark harness structure: The file bench_kernels.cu is a standalone microbenchmark that measures kernel latency across sweeps of parameters like number of streams, query length, and prefix length. It uses random inputs (correctness is tested separately) and reports per-call latency in microseconds.
  4. The project structure: The kdtree-engine project implements custom CUDA kernels for speculative decoding with DFlash and DDTree. It has a C ABI layer for integration with SGLang, a Python reference implementation for generating golden test data, and a test suite that validates correctness against the reference.
  5. The Phase 1 goal: The immediate objective was to produce a flash attention kernel that passes token-exact correctness tests against the naive oracle and shows measurable performance improvement in microbenchmarks, particularly at long prefix lengths where memory bandwidth is the bottleneck.

Output Knowledge Created

This message produced a modified benchmark file that could run both the naive and flash attention kernels across prefix sizes up to 16384 tokens. The output knowledge created includes:

  1. A validated measurement framework: The benchmark could now generate comparative latency data for both kernel implementations under identical conditions, enabling the team to quantify the speedup (or slowdown) at each prefix length.
  2. A stress test for long prefixes: By extending the maximum prefix size to 16384, the benchmark could probe the kernel's behavior in the regime where the flash attention design should provide the greatest benefit—long contexts where global memory bandwidth dominates.
  3. A diagnostic tool for the crossover point: The sweep across prefix sizes would reveal the prefix length at which the flash kernel's tiling overhead is amortized enough to outperform the naive kernel, providing guidance for when to use each implementation.

Assumptions and Potential Mistakes

The assistant made several assumptions in this message:

  1. That the naive kernel implementation was already registered in the benchmark: The assistant assumed that the benchmark harness could already call the naive kernel and that only the flash kernel variant needed to be added. This was likely correct given that the benchmark existed before the flash kernel was created.
  2. That 16384 tokens is a sufficient upper bound: While 16384 is a reasonable stress-test value, the production deployment already supported 200k context lengths. The assistant may have chosen 16384 as a pragmatic limit for microbenchmarking—longer prefixes would increase benchmark runtime and may not fit in GPU memory for all configurations. However, this means the benchmark would not directly measure performance at the full 200k context length that the production service supports.
  3. That the benchmark should test both kernels in the same binary: The assistant edited the existing benchmark file rather than creating a separate benchmark for the flash kernel. This assumes that the two kernels can coexist in the same compilation unit and share the same test harness, which is reasonable given that they share the same interface.
  4. That the naive kernel is the correct baseline: The "naive" kernel here is the Triton-based MLA implementation that was already in use. However, the assistant had previously identified that the Triton implementation was using page_size=1, causing scattered memory access. The benchmark would therefore compare the new flash kernel against a known-bad baseline, which might overstate the improvement. A more rigorous comparison would also include an optimized version of the Triton kernel with better memory access patterns, but that was not available for sm_120.

The Broader Context: What Came Before and After

Before this message, the assistant had already written the flash kernel, integrated it into the build system, and added correctness tests. The test suite confirmed that the flash kernel produced token-exact results matching the naive oracle—a critical validation that the algorithmic changes (tiling, online softmax, shared memory reuse) did not introduce numerical errors.

After this message, the assistant would build and run the benchmark on the CT200 server. The initial microbenchmark results (reported in the next chunk's summary) showed that the flash kernel was slower than the naive kernel at short prefixes (0.4–0.7×) and only modestly faster at longer ones (1.7–2.0×). This was below the expected 8–16× improvement from reduced global memory reads, indicating that the reduction kernel and tile sizing needed further optimization for sm_120's architecture. The work would continue into subsequent phases, with the assistant eventually achieving a 3–6× end-to-end decode speedup over the Triton baseline after further tuning.

Conclusion

Message 12233 appears to be a routine edit to a benchmark file, but it represents a critical moment in the development of a custom CUDA kernel. It is the point where the assistant transitions from implementation to evaluation, ensuring that the new kernel can be measured fairly against the existing baseline across a meaningful range of operating conditions. The expansion to 16384-token prefixes reflects an understanding that the flash attention design's benefits are context-length-dependent, and that the benchmark must capture the full performance profile to guide further optimization. This message, though brief, embodies the disciplined engineering approach that characterizes the entire session: build, test, measure, and iterate based on real data rather than assumptions.