Validation at the Gate: Running Custom CUDA Kernels on Production Blackwell Hardware

In the lifecycle of any high-performance computing project, there comes a moment of truth: the transition from a development environment to the actual target hardware. Message [msg 11985] captures precisely such a moment in the development of a native C/C++/CUDA speculative decoding engine for the Kimi K2.6 language model. After building three custom CUDA kernels—a GPU best-first tree builder, a tree-verify MLA-absorb attention kernel, and a greedy tree-accept kernel—the assistant deploys the code to an 8× RTX PRO 6000 Blackwell machine (codename CT200) and runs the full unit-test suite. The message is deceptively short, but it represents the culmination of dozens of preceding messages: the validation of hand-written CUDA kernels on the exact hardware architecture they were designed for.

The Message in Full

The assistant's reasoning reveals the intent clearly:

Now I'm running the unit tests against the PRO 6000 hardware to validate correctness on the actual target device, using the virtual environment's Python for generating reference outputs.

The action is a single SSH command that executes the test runner on the remote machine:

timeout 120 ssh -o StrictHostKeyChecking=no root@10.1.2.200 'cd /root/kdtree-engine \
  && PYTHON=/root/venv_sglang211/bin/python bash scripts/run_tests_nvcc.sh 2>&1 | tail -45'

The output shows a clean sweep of passing tests:

PASS  tree_build_singlereq.kdtr     B=1 D=7 K=4 budget=12
PASS  tree_build_small_budget.kdtr  B=5 D=7 K=4 budget=4
PASS  tree_build_tiny.kdtr          B=8 D=3 K=3 budget=5
PASS  tree_build_underfull.kdtr     B=5 D=2 K=2 budget=8
PASS  tree_build_wide.kdtr          B=4 D=7 K=8 budget=32

== verify_attn ==
PASS  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  verify_attn_h64_p128_b32.kdtr  B=1 H=64 q_len=33 kv_len=161  max_abs=2.05e-08 max_rel...

Every test passes. The numerical precision is excellent—maximum absolute errors in the 10⁻⁸ range, with relative errors around 10⁻³ for the larger configuration. The kernels are correct.

The Road to This Moment

To appreciate what this message represents, one must understand the journey that preceded it. The assistant had spent the better part of a session (Segment 65, Chunk 0) building an entire inference engine from scratch. This was not a trivial modification of existing code but a ground-up implementation of a custom speculative decoding pipeline.

The engine comprises three novel CUDA kernels:

  1. tree_build: A GPU best-first tree builder that replaces SGLang's per-request CPU heapq-based tree construction. Given a set of candidate tokens with scores, it builds a draft tree directly on the GPU, avoiding the host-device round-trip that SGLang's implementation requires.
  2. verify_attn: A tree-verify MLA-absorb attention kernel. This is the heart of the speculative decoding verification step, implementing Multi-head Latent Attention (MLA) with a visibility mask that ensures each position in the draft tree only attends to its prefix. The kernel absorbs the latent projection into the attention computation, avoiding an explicit materialization step.
  3. tree_accept: A greedy tree-accept kernel that walks the draft tree to determine how many tokens from the speculative draft can be accepted given the target model's output. These kernels had been validated on the local development machine (a GeForce RTX 5070 Ti, also sm_120 Blackwell), but that validation revealed two issues. First, the tree_accept benchmark crashed due to cyclic random test data—the kernel's path-walk algorithm could loop indefinitely on malformed input. The assistant fixed this with a safety bound that caps the walk length at the tree size, which is a correct invariant for any valid tree. Second, the verify_attn kernel was embarrassingly slow on the 5070 Ti—over 6 milliseconds for moderate configurations—because the naive implementation had each of the 64 MLA heads independently re-reading the shared latent KV cache from HBM, wasting bandwidth by a factor of 64. The assistant recognized, however, that this performance issue was not a flaw in the kernel design but a confirmation of the production architecture. The plan had always been to reuse an optimized MLA kernel like FlashMLA for the long prefix attention and reserve the custom verify_attn kernel only for the small tree tail (typically 9–33 query tokens). The benchmark numbers on the 5070 Ti validated this design decision: the custom kernel is correct but bandwidth-bound, which is acceptable for the small-tree regime.

Why This Validation Matters

The transition from the 5070 Ti development machine to the RTX PRO 6000 Blackwell production hardware is non-trivial for several reasons.

Architecture differences: Although both GPUs are Blackwell (sm_120), the PRO 6000 has significantly different memory hierarchy characteristics—larger L2 cache, higher HBM bandwidth, and different warp scheduling behavior. A kernel that works correctly on one Blackwell GPU might still exhibit numerical divergence on another due to different instruction scheduling or floating-point accumulation order. The unit tests verify bit-exact (or near-bit-exact) agreement against numpy reference implementations, which are computed on the CPU with deterministic FP64 arithmetic. Passing these tests on the PRO 6000 confirms that the GPU kernels produce the same results regardless of which Blackwell GPU executes them.

CUDA toolkit version: The local machine used CUDA 13.2, while CT200 had CUDA 13.0 installed. Different CUDA versions can produce different code generation, especially for newer architectures like sm_120 where the compiler is still maturing. The successful test run confirms that the kernels compile and execute correctly under CUDA 13.0.

Environment consistency: The test runner explicitly specifies a Python virtual environment (/root/venv_sglang211/bin/python) for generating reference data. The KDTR binary container format—a custom format designed to share test data between Python (numpy reference generation) and C++ (kernel execution)—must be interpreted identically on both machines. The passing tests confirm end-to-end compatibility of the data pipeline.

Assumptions and Input Knowledge

The message rests on several layers of prior knowledge and assumptions.

The KDTR format: The test data files (.kdtr) contain serialized tensors representing tree structures, attention states, and verification targets. The assistant designed this format specifically for cross-language data sharing. Understanding the message requires knowing that each KDTR bundle encodes a complete test case: input parameters (batch size, tree depth, top-k, budget), input tensors (scores, hidden states, KV cache), and expected output tensors generated by the numpy reference.

The test infrastructure: The run_tests_nvcc.sh script iterates over KDTR bundles, launches each test binary, and compares kernel outputs against the reference using absolute and relative error thresholds. The assistant wrote this script as a cmake-free alternative for CT200, which lacks cmake. The tail -45 in the command suggests the output is long, with only the final results being of interest.

The remote environment: The assistant assumes that the SSH connection to root@10.1.2.200 is available, that the rsync in the previous message ([msg 11984]) successfully transferred the repository, that the nvcc build completed without errors, and that the CUDA 13.0 toolkit is correctly configured. The timeout 120 wrapper indicates an expectation that the test suite should complete within two minutes.

The Thinking Process

The assistant's reasoning in this message is concise but reveals a methodical deployment philosophy. The phrase "validate correctness on the actual target device" underscores a key principle: development GPUs and production GPUs are not interchangeable for correctness testing. The assistant could have assumed that passing tests on the 5070 Ti was sufficient, but it deliberately chose to re-run the full suite on the PRO 6000.

This is particularly important given the earlier crash in tree_accept. The crash was caused by random test data producing cyclic trees, which is impossible in production (the tree builder always produces acyclic trees). But the fix—adding a safety bound—needs to be verified on the target hardware because it changes the kernel's control flow. A bound that is correct on one GPU might interact differently with the memory subsystem on another. The PRO 6000 validation confirms that the safety bound works correctly and does not introduce spurious failures on valid inputs.

The choice of Python environment is also telling. By specifying the full path to the virtual environment's Python interpreter, the assistant ensures that the numpy reference computations use the same library versions as during development. Numerical libraries can differ in their default accumulation order or transcendental function implementations, and pinning the environment eliminates one source of variability.

Output Knowledge Created

This message produces several concrete pieces of knowledge:

  1. Correctness confirmation: All five tree_build test configurations pass on the PRO 6000, covering single-request, small-budget, tiny, underfull, and wide tree shapes. The parameter space (B=1–8, D=2–7, K=2–8, budget=4–32) spans the realistic operating range for the Kimi K2.6 speculative decoding workload.
  2. Numerical precision characterization: The verify_attn tests quantify the kernel's numerical accuracy. For the h16_p2048_b16 configuration (2 batches, 16 heads, 17 query tokens, 2065 KV tokens), the maximum absolute error is 2.28×10⁻⁸ and the maximum relative error is 2.33×10⁻³. For the h64_p128_b32 configuration (1 batch, 64 heads, 33 query tokens, 161 KV tokens), the maximum absolute error is 2.05×10⁻⁸. These numbers are well within acceptable bounds for FP32 inference, where the model's own weights and activations have much larger inherent uncertainty.
  3. End-to-end pipeline validation: The successful test run confirms that the entire toolchain—numpy reference generation, KDTR serialization, CUDA compilation, kernel execution, and result comparison—works correctly on the production machine. This is a prerequisite for the subsequent performance benchmarking that follows in [msg 11986].

The Broader Significance

In the context of the overall project—deploying Kimi K2.6 with DFlash speculative decoding on Blackwell hardware—this message represents a critical quality gate. The assistant is systematically working through a deployment checklist: build locally, fix crashes, deploy to target, build on target, validate correctness, benchmark performance, then integrate into the live service. Each step must pass before the next begins.

The passing tests also validate the assistant's design decisions. The GPU tree builder, which replaces SGLang's CPU-based implementation, produces correct trees. The verify_attn kernel, despite its naive memory access pattern, computes correct attention outputs. The tree_accept kernel, now protected by the safety bound, correctly identifies accepted tokens. All three kernels are ready for performance characterization.

What makes this message noteworthy is what it doesn't contain: there are no bugs to fix, no crashes to debug, no regressions to investigate. After the intense problem-solving of the preceding messages—the cyclic-tree crash, the slow verify_attn kernel, the CUDA ABI mismatch—this message is a moment of calm. The kernels work. The pipeline is sound. The assistant can now proceed to the next phase: measuring how fast these kernels run on the production hardware, and ultimately whether they deliver the promised throughput improvement over the autoregressive baseline.

This is the quiet satisfaction of a well-executed engineering plan: the moment when everything compiles, links, and passes, and the only remaining question is how fast it will go.