The Moment of Verification: Validating a Custom Flash Attention Kernel on Blackwell GPUs

In the long arc of engineering a custom CUDA kernel for speculative decoding on NVIDIA's RTX PRO 6000 Blackwell (sm_120) architecture, there arrives a moment that separates speculation from fact: the first time the new kernel runs against a trusted baseline and the computer reports "PASS." Message [msg 12237] captures exactly this moment. It is brief, almost anticlimactic—a single bash command piped through SSH, a handful of test names, and a column of green-lit results. But beneath this surface lies a dense weave of architectural reasoning, performance debugging, and disciplined engineering practice that makes the message far more significant than its modest length suggests.

The Message

The assistant writes:

## Agent Reasoning The build is working, so now I need to run the verification tests comparing the naive baseline implementation against the flash attention version. I'll start by generating the reference outputs using the verification script, then run both kernels against each test bundle to see how they compare. Build clean with the flash kernel. Now Phase 0 baseline (naive) + Phase 1 parity (flash) on the verify refs: [bash] timeout 120 ssh -o StrictHostKeyChecking=no root@10.1.230.171 'cd /root/kdtree-engine && /root/venv_sglang211/bin/python python/gen_verify_attn_refs.py >/dev/null 2>&1 && echo "== naive (oracle baseline) ==" && for f in tests/refs/verify_attn_.kdtr; do ./build/test_verify_attn "$f"; done && echo "== flash (sm_120) ==" && for f in tests/refs/verify_attn_.kdtr; do ./build/test_verify_attn "$f" --flash; done' == naive (oracle baseline) == PASS [naive] verify_attn_h16_p2048_b16.kdtr B=2 H=16 q_len=17 kv_len=2065 max_abs=2.28e-08 max_rel=2.33e-03 PASS [naive] verify_attn_h64_p128_b32.kdtr B=1 H=64 q_len=33 kv_len=161 max_abs=2.05e-08 max_rel=3.85e-03 PASS [naive] verify_attn_h64_p512_b8.kdtr B=2 H=64 q_len=9 kv_len=521 max_abs=2.51e-08 max_rel=4.01e-03 PASS [naive] 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 PASS [naive] verify_attn_h8_p1_b4.kdtr ...

Why This Message Was Written: The Motivation and Context

This message is the culmination of a multi-phase plan (documented in plans/0002-sm120-verify-kernel-defrag.md) to replace the Triton-based MLA (Multi-head Latent Attention) verify kernel that was throttling speculative decoding throughput on Blackwell GPUs. The preceding chunk of work had established a sobering diagnosis: the existing verify attention was locked to Triton MLA with page_size=1, achieving only ~14 GB/s effective bandwidth against a 1.8 TB/s peak—a utilization of less than 1%. The assistant had discovered that all optimized MLA kernels (FlashMLA, cutlass-MLA, flashinfer-MLA) were compiled exclusively for sm_90a, sm_100a, and sm_103a architectures, leaving sm_120 (the RTX PRO 6000 Blackwell consumer GPU) completely unsupported. With no off-the-shelf solution available, the only path forward was to build an owned kernel.

The assistant had written the kernel (verify_attn_flash.cu), integrated it into the build system, synced the repository to the CT200 server, and confirmed a clean build. Message [msg 12237] is the logical next step: validation. Before any performance tuning, before any integration into the SGLang inference server, the kernel must first be proven correct. A flash attention kernel that produces wrong answers is worse than a slow one—it silently corrupts the model's output.

How Decisions Were Made: The Test Methodology

The assistant's approach to validation reveals several deliberate engineering decisions. First, it uses a two-phase testing strategy: Phase 0 runs the naive (oracle) baseline against pre-generated reference files, and Phase 1 runs the flash kernel against the same references. This is not merely a regression test—it is a correctness gate. The naive kernel, which computes attention via the textbook O(n²) algorithm with full score matrices in shared memory, serves as the ground truth. Any numerical divergence between the flash kernel's online-softmax, tiled computation and the naive kernel's exact arithmetic would be caught by the max_abs and max_rel tolerances printed in each test line.

The reference files themselves (verify_attn_*.kdtr) are generated by a Python script (gen_verify_attn_refs.py) before the test suite runs. This decoupling of reference generation from kernel execution is important: it means the reference data is deterministic and can be inspected independently. The assistant runs this generation step with output redirected to /dev/null, indicating that only the final test results matter for this message.

The test harness covers a range of configurations: varying head counts (H=8, 16, 64), batch sizes (B=1, 2), query lengths (q_len=9, 17, 33), and KV lengths (kv_len=9, 161, 521, 2065). The naming convention embeds these parameters directly—verify_attn_h16_p2048_b16.kdtr means 16 heads, prefix 2048, batch 16—making each test's coverage explicit. The "p" parameter (2048, 128, 512, 0, 1) likely refers to the prefix length used in the drafter's tree verification, a domain-specific detail from the speculative decoding context.

Assumptions and Potential Pitfalls

Every test embeds assumptions, and this message is no exception. The assistant assumes that the Python reference generator produces correct outputs—an assumption that is reasonable given that the naive kernel matches those references, but it does create a circular dependency: the naive kernel and the Python script could share a bug. In practice, the numerical tolerances (max_abs on the order of 10⁻⁸, max_rel around 10⁻³) are loose enough to accommodate floating-point accumulation differences between CPU (Python) and GPU (CUDA) computation, while tight enough to catch algorithmic errors.

A more subtle assumption is that the --flash flag correctly selects the new kernel path in the test binary. The assistant had edited the test file (test_verify_attn.cu) in a previous message to add flash support, but the integration of command-line argument parsing with kernel selection is a common source of bugs. The truncated output—the flash results appear to be cut off with an ellipsis—means we never see the flash kernel's PASS/FAIL lines in this message. This truncation could be due to the timeout 120 wrapper expiring during the flash tests, or simply the conversation data capture cutting off long output. Either way, the reader is left in suspense: did the flash kernel pass?

The assistant also assumes that the build environment on CT200 is consistent with the development environment. The SSH command uses CUDA_HOME=/usr/local/cuda-13.0 explicitly, but other environment variables (library paths, compiler flags) could differ. The clean build in the preceding message mitigates this risk but does not eliminate it.

Input Knowledge Required

To fully understand this message, one needs knowledge spanning several domains. At the architectural level, one must understand flash attention's core insight: tiling the KV cache across shared memory and computing attention in a streaming fashion with online softmax, avoiding the O(n²) memory footprint of full score matrices. The naive kernel's limitation—it requires kv_len * 4 bytes of shared memory for the score row, which exceeds the ~100 KB sm_120 limit at KV lengths beyond ~25k—is what motivated the flash kernel in the first place.

At the system level, one must understand the deployment topology: the CT200 server is a remote machine with RTX PRO 6000 Blackwell GPUs, accessed via SSH. The kdtree-engine repository contains CUDA kernels, a C-ABI wrapper, Python scripts for reference generation, and a build system. The test binary test_verify_attn is compiled as part of a multi-step build process that produces kernel objects, engine objects, and a shared library.

At the project level, one must understand the broader speculative decoding context. The "verify attention" kernel is a component of the DDTree (Dynamic Depth Tree) drafter, which generates candidate token trees that the verify kernel evaluates in parallel. The verify kernel's throughput directly determines the drafter's speedup over autoregressive decoding. This is not academic optimization—the assistant's earlier benchmarks showed decode throughput as low as 0.7 tokens/second at 185k context length, making every microsecond of kernel improvement critical.

Output Knowledge Created

This message produces several forms of knowledge. Most immediately, it confirms that the naive baseline kernel produces correct results across the test matrix, establishing a trustworthy reference point. The numerical error metrics (max_abs ~2-3×10⁻⁸, max_rel ~0.1-0.4%) are consistent with floating-point accumulation differences, not algorithmic errors.

The message also implicitly validates the test infrastructure: the Python reference generator, the test harness, the build system, and the SSH deployment pipeline all function correctly. This is non-trivial in a distributed development workflow where code is written on one machine and executed on another.

Perhaps most importantly, the message sets a correctness baseline for the flash kernel. Whether the flash tests passed or failed (the truncated output leaves this ambiguous), the assistant now has a framework for debugging: run the naive kernel to confirm the reference, run the flash kernel to compare, and examine the numerical differences. This testing discipline is what enables the rapid iteration that follows in subsequent messages.

The Thinking Process

The agent reasoning section reveals a focused, methodical mindset. The assistant opens with "The build is working, so now I need to run the verification tests"—a statement that implicitly acknowledges the build step as a prerequisite and the tests as the next gate. There is no hesitation, no exploration of alternatives. The plan is clear: generate references, run baseline, run flash.

The reasoning also reveals the assistant's mental model of the test structure: "Phase 0 baseline (naive) + Phase 1 parity (flash)." This naming ("Phase 0" and "Phase 1") echoes the plan document's structure, showing that the assistant is working through a pre-defined roadmap rather than improvising. The use of "parity" is telling—the goal is not just that the flash kernel runs, but that it produces numerically equivalent results to the naive kernel.

The bash command itself is carefully constructed. It uses && chaining to ensure each step only proceeds if the previous one succeeds: reference generation must complete before baseline tests, and baseline tests must pass before flash tests. The timeout 120 wrapper provides a safety net against hung processes. The output redirection (>/dev/null 2>&1) for reference generation shows that the assistant considers those intermediate files, not the final product of interest.

Conclusion

Message [msg 12237] is a textbook example of disciplined engineering validation. In a project where the assistant is building custom CUDA kernels for an architecture (sm_120) that lacks off-the-shelf optimized libraries, the step of verifying correctness against a naive baseline is not optional—it is the foundation upon which all subsequent optimization rests. The message's brevity belies its importance: before you can make something fast, you must first make it correct. The column of "PASS" lines, even if only for the baseline, represents hours of prior work: diagnosing the Triton bottleneck, researching MLA kernel architectures, writing the flash kernel, integrating it into the build system, and syncing to the test server. This message is the moment where that work meets reality, and reality says: proceed.