When Theory Meets Silicon: Validating a Custom CUDA Verify Attention Kernel on Blackwell GPUs

In the high-stakes world of speculative decoding for large language models, every microsecond counts. When a production inference service is delivering only 0.7 tokens per second at 185k context length — with GPU tensor cores idling at 3% utilization — the bottleneck demands surgical intervention. This is the story of a single message in an opencode coding session where an AI assistant took a carefully reasoned kernel redesign and subjected it to the unforgiving test of real hardware: syncing source code to a remote machine, rebuilding, validating numerical correctness, and benchmarking performance. Message [msg 12243] is the bridge between theory and measurement — the moment a sixth-generation custom CUDA kernel for MLA (Multi-Head Latent Attention) verify attention meets the RTX PRO 6000 Blackwell GPU.

The Message

The message is a single bash command executed by the assistant, followed by its output:

[assistant] [bash] cd /home/theuser/glm-kimi-sm120-rtx6000bw/kdtree-engine && rsync -az src/kernels/verify_attn_flash.cu root@10.1.230.171:/root/kdtree-engine/src/kernels/ 2>&1 && timeout 300 ssh -o StrictHostKeyChecking=no root@10.1.230.171 'cd /root/kdtree-engine && CUDA_HOME=/usr/local/cuda-13.0 bash scripts/build_nvcc.sh 2>&1 | tail -2 && echo "== parity ==" && for f in tests/refs/verify_attn_*.kdtr; do ./build/test_verify_attn "$f" --flash; done && echo "== bench ==" && ./build/bench_kernels 2>&1 | sed -n "/verify_attn/,/tree_build/p" | head -26'
[5/5] engine tests + demo (needs cublas)
done -> build/
== parity ==
PASS  [flash] verify_attn_h16_p2048_b16.kdtr  B=2 H=16 q_len=17 kv_len=2065  max_abs=1.77e-08 max_rel=2.11e-03
PASS  [flash] verify_attn_h64_p128_b32.kdtr  B=1 H=64 q_len=33 kv_len=161  max_abs=2.24e-08 max_rel=3.76e-03
PASS  [flash] verify_attn_h64_p512_b8.kdtr  B=2 H=64 q_len=9 kv_len=521  max_abs=2.79e-08 max_rel=3.09e-03
PASS  [flash] verify_attn_h8_p0_b8.kdtr   B=2 H=8 q_len=9 kv_len=9  max_abs=2.98e-08 max_rel=1.20e-03
PA...

On the surface, this is a routine deployment command: sync a file, build, test, benchmark. But in context, it represents the culmination of an intense optimization odyssey spanning multiple kernel iterations, each one driven by a deepening understanding of the Blackwell GPU's architecture and the unique demands of speculative decoding with DDTree (draft-tree-based speculative decoding).

Context and Motivation: The Occupancy Crisis

To understand why this message was written, we must look at what immediately preceded it. In [msg 12242], the assistant engaged in an extraordinarily detailed reasoning session — over 1,500 words of analysis — dissecting why the previous kernel design (v4/v5 with KV-splitting) was still failing to deliver the expected performance.

The numbers were stark. At a 16,384-token prefix with 9 query tokens, the "flash" kernel was taking 11.6 milliseconds — a mere 1.7× speedup over the naive baseline of 19.6ms. At shorter prefixes, the flash kernel was actually slower than naive (0.4–0.7×). The theoretical peak compute for the workload was 0.22ms (based on ~1e10 FMAs at 90 TFLOPs), meaning the kernel was achieving only about 2% of peak. Something fundamental was wrong.

The assistant's analysis identified the root cause: occupancy starvation. The previous design used 70KB of shared memory per block, which forced the GPU's scheduler to place at most one block per streaming multiprocessor (SM). With only 128 threads per block, that meant just 6% of the SM's 2048-thread capacity was utilized. In a latency-bound decode workload, those idle warps meant no ability to hide memory latency — every stall on a global memory load or shared memory access would freeze the SM.

The assistant's reasoning in [msg 12242] is a masterclass in GPU optimization thinking. It walks through the math of shared memory budgets, thread occupancy, and the tension between amortizing latent reads (which reduces memory traffic but increases shared memory footprint) and maximizing occupancy (which keeps SMs busy but may increase redundant reads). The key insight was that for decode — where batch size is 1 and the workload is fundamentally latency-bound — occupancy matters more than memory traffic. Reducing the number of times the KV cache is read per head is a secondary concern if the GPU is mostly idle anyway.

The v6 Design: Per-Head, Warp-Per-Key, Small-Smem

The design that emerged from this analysis was a radical departure from the previous approach. Instead of amortizing multiple heads per block (HT=8), the v6 kernel uses a per-head design (HT=1) where each block handles exactly one head. This dramatically reduces shared memory requirements from ~70KB to approximately 23KB, allowing the GPU to fit four blocks per SM instead of one. Combined with 256 threads per block (up from 128), this yields approximately 67% occupancy — a 10× improvement over the 6% of the previous design.

The threading strategy also changed. Instead of having one thread compute an entire 576-element dot product serially while others idle, the v6 kernel uses a warp-per-key approach: each warp of 32 threads cooperates on one key's dot product, with each thread handling 18 elements followed by a warp-level reduction. This keeps all 256 threads active during the score computation phase, eliminating the thread-idling bottleneck of earlier designs.

The KV-split mechanism from v4/v5 was retained, partitioning long prefixes across multiple blocks to generate enough parallelism to fill all 188 SMs on the RTX PRO 6000. But the critical change was making the per-block shared memory small enough that occupancy — not block count — was the primary mechanism for keeping SMs busy.

Assumptions Embedded in This Message

This deployment message carries several assumptions, some explicit and some implicit:

That the v6 design is correct. The parity tests validate this assumption — all four test bundles pass with maximum absolute errors around 2e-8 and relative errors around 0.2–0.4%, which is token-exact for bf16 inference. The assistant clearly expected this (the design was carefully reasoned), but correctness is never guaranteed until tested.

That the remote build environment matches expectations. The command uses CUDA_HOME=/usr/local/cuda-13.0 and builds with scripts/build_nvcc.sh. This assumes the remote machine (CT200, IP 10.1.230.171) has CUDA 13.0 installed at that path and that the build script is compatible. Earlier in the session, the assistant had resolved CUDA ABI mismatches and other build issues, so this assumption was well-grounded.

That the benchmark harness produces meaningful numbers. The bench_kernels binary tests specific configurations (H=64, Dl=512, Dr=64, various prefix lengths and query counts). The assistant assumes these configurations are representative of the production workload. This is a reasonable assumption for a microbenchmark, but the true test would come later in production integration.

That the naive kernel is the correct baseline. The benchmark compares against a "naive" implementation that the assistant had previously validated. The assumption is that any improvement over naive represents genuine progress, even if the absolute performance is still far from theoretical peaks.

What This Message Creates: Output Knowledge

This message produces several forms of output knowledge:

Numerical correctness validation. All four parity tests pass, confirming that the v6 kernel produces the same results as the oracle implementation across a range of configurations: H=16 with 2048-length prefix, H=64 with 128-length prefix, H=64 with 512-length prefix, and H=8 with a tiny 9-length prefix. The maximum absolute errors (1.77e-08 to 2.98e-08) are at the limits of bf16 precision, confirming the kernel is token-exact.

Build success. The project compiles cleanly with CUDA 13.0 on the remote machine, confirming that the kernel source is syntactically valid and compatible with the target toolchain.

A benchmark run. The output is truncated at "PA..." — we don't see the full benchmark results in this message. However, the next message ([msg 12244]) reveals that v6 achieved 1.3–1.5× speedup on short sequences and 2.7–3.3× at 16,384 tokens, with the ability to handle 65,536 tokens (which the naive kernel could not). This represents a significant improvement over the v4/v5 design that managed only 1.7× at 16k and was slower at short prefixes.

Confidence to proceed. With correctness validated and performance improved, the assistant could move forward to the next phase: adding float4 vectorized loads to further improve memory throughput, and eventually integrating the kernel into the production SGLang inference server.

The Thinking Process Visible in Reasoning

The assistant's reasoning in [msg 12242] reveals a disciplined, analytical approach to GPU kernel optimization. Several patterns stand out:

Quantitative reasoning throughout. The assistant doesn't hand-wave about bottlenecks — it computes actual numbers. It calculates that ~1e10 FMAs at 90 TFLOPs should take 0.22ms, then compares this to the actual 11.6ms to conclude the kernel is 50× off peak. It computes shared memory budgets (70KB vs 49KB vs 23KB) and maps them to occupancy (1 block/SM vs 2 blocks/SM vs 4 blocks/SM). It estimates memory bandwidth requirements (2.4GB of KV data at 403 GB/s achieved vs 1.8 TB/s peak). Every claim is backed by arithmetic.

Iterative hypothesis testing. The assistant considers multiple explanations for the performance gap — memory bandwidth, compute throughput, occupancy, thread utilization — and tests each against the numbers. It recognizes that the production problem (scattered gather at 14 GB/s with page_size=1) is different from the microbenchmark (contiguous KV memory), and correctly notes that the real win might be larger in production than in the microbench.

Awareness of architectural constraints. The assistant knows the RTX PRO 6000 Blackwell (sm_120) has specific characteristics: 188 SMs, 48 warps/SM max, ~90 TFLOPs fp32, 1.8 TB/s memory bandwidth, and crucially, a smaller 100KB shared memory limit per SM compared to data-center Blackwell. The design decisions are explicitly tied to these numbers.

Recognition of the fundamental tension. The core insight — that occupancy and memory traffic are in tension for MLA verify attention — is articulated clearly. The assistant correctly identifies that for decode (latency-bound, batch=1), occupancy wins. This is not obvious: a naive designer might focus on reducing memory reads (the conventional flash-attention optimization) and miss that the GPU is mostly idle anyway.

Disciplined decision-making under uncertainty. When the assistant isn't sure which configuration will perform best, it plans to test empirically: "Rather than overthinking this, let me just test a few configs empirically." The v6 design was chosen not because it was guaranteed optimal, but because it had the best theoretical properties and could be tested quickly.

The Broader Significance

Message [msg 12243] sits at a critical juncture in a larger narrative about deploying speculative decoding on novel GPU hardware. The RTX PRO 6000 Blackwell (sm_120) is a consumer-grade GPU that lacks the specialized tensor-core instructions (wgmma, TMA, tcgen05) available on data-center Blackwell (sm_100/sm_103). This means the assistant cannot use optimized MLA kernels like FlashMLA or cutlass-MLA — it must build its own. The verify attention kernel is the heart of the DDTree speculative decoding engine, responsible for computing attention scores and value aggregations for the draft tree's verification step. Every microsecond saved here directly translates to faster token generation.

The fact that the assistant is iterating through six kernel designs in a single session — each one informed by measurement and analysis — speaks to the complexity of GPU kernel optimization. There is no single "right answer"; the optimal design depends on the specific GPU architecture, workload characteristics (decode vs prefill, batch size, context length), and the constraints of the surrounding system (CUDA graph capture, memory fragmentation, tensor parallelism).

Conclusion

Message [msg 12243] appears, at first glance, to be a routine build-and-test command. But in the context of the broader session, it represents the critical moment when a carefully reasoned kernel redesign — born from deep analysis of occupancy, shared memory budgets, and thread utilization — is validated against real hardware. The parity tests pass, confirming correctness. The benchmarks (visible in the subsequent message) show 1.3–3.3× speedup over the naive baseline. The v6 kernel works.

This is the essence of GPU kernel development: a cycle of measurement, analysis, redesign, and validation. Each iteration reveals new bottlenecks and new insights. The assistant's disciplined approach — quantifying every claim, testing every assumption, and letting data drive decisions — is a model for how to optimize for exotic hardware where off-the-shelf solutions don't exist. The message is a testament to the fact that in high-performance computing, theory is only half the battle; the other half is watching the numbers come back from the hardware.