The Test That Proves the Kernel: Validating a GPU Tree Builder Against a Golden Reference

Introduction

In the middle of building a native C/C++/CUDA inference engine for the Kimi K2.6 speculative decoding drafter, the assistant reached a critical juncture: the GPU kernel for best-first tree construction had been written, but it was untested. The message at [msg 11862] captures the moment when the assistant created the test harness that would validate the kernel against a set of golden reference outputs. This seemingly simple act — writing a C++/CUDA unit test file — represents a deliberate quality-assurance strategy that underpins the entire engine's correctness.

The message itself is brief, containing only the assistant's reasoning and a file-write action:

## Agent Reasoning Writing the test runner... Now the C++/CUDA unit test that validates the kernel against each reference bundle: [write] /home/theuser/glm-kimi-sm120-rtx6000bw/kdtree-engine/tests/test_tree_build.cu Wrote file successfully.

But behind this concise utterance lies a carefully constructed validation architecture, a sequence of prior decisions about test data generation, binary container formats, and reference implementations — all of which converge in this single file.

Why This Message Was Written: The Need for Kernel Validation

The assistant had just completed writing two GPU kernel files — a header (tree_build.cuh at [msg 11860]) and an implementation (tree_build.cu at [msg 11861]) — that implement a best-first tree builder for speculative decoding. This kernel replaces SGLang's per-request CPU-based heapq tree construction with a GPU-native implementation, which is essential for the throughput demands of production inference.

However, a GPU kernel written in CUDA is only as trustworthy as its test coverage. The tree-building algorithm is non-trivial: it performs a best-first expansion over candidate tokens using cumulative log-probabilities, constructs a visibility matrix for attention masking, and computes retrieval pointers for the verify pass. Any subtle bug in heap ordering, tie-breaking, or index arithmetic would silently corrupt the speculative decoding tree, leading to incorrect token acceptance or, worse, silently wrong model outputs.

The assistant's reasoning makes this motivation explicit: "Now the C++/CUDA unit test that validates the kernel against each reference bundle." The word "validates" is carefully chosen — this is not a smoke test or a sanity check, but a bit-exact comparison against a golden reference computed by a faithful numpy port of the original SGLang algorithm.

The Validation Architecture: A Multi-Layered Approach

To understand the significance of this test file, one must trace the validation chain that precedes it. The assistant had already established three critical pieces of infrastructure:

First, a binary container format (KDTR) was defined in python/kdtr_io.py ([msg 11855]) and its C++ counterpart src/common/kdtr_io.h ([msg 11856]). This format allows Python-generated reference data to be loaded directly by the C++/CUDA test harness, with no JSON parsing, no text serialization, and no format conversion. The KDTR format stores typed arrays with metadata in a compact binary layout, ensuring that the float32 log-probabilities and int32 token IDs arrive at the GPU test with zero precision loss.

Second, a faithful numpy reference implementation was written in python/ddtree_ref.py ([msg 11857]). This is not a simplified version of the algorithm — it is a line-for-line port of SGLang's build_ddtree_tree_from_topk function, reproducing the exact arithmetic operations, heap ordering, and tie-breaking behavior. The reference uses double-precision accumulation to match Python's float64, ensuring that the GPU kernel (which the assistant planned to run in float32) could be compared with known numerical behavior.

Third, a test-data generator (python/gen_tree_build_refs.py, [msg 11858]) produces random top-log-probabilities and token IDs, runs them through the reference, and saves both inputs and expected outputs as KDTR bundles. The generator covers multiple configurations: default parameters (budget=8, depth=7, topk=4), mid-range (budget=16), wide (topk=8), edge cases like underfull (where the tree runs out of candidates before reaching the budget), and chain (topk=1, producing a linear sequence). The generator was verified to produce correct outputs ([msg 11859]), with the underfull case yielding actual=7 (below the budget of 9) and the chain case yielding actual=8.

What the Test File Does

The test file test_tree_build.cu is the culmination of this infrastructure. It performs the following operations:

  1. Loads a KDTR bundle containing the input tensors (top_log_probs, top_token_ids, next_token, next_sibling) and expected output tensors (node_token_ids, node_depths, parents, visibility, retrieve_indices).
  2. Allocates GPU memory and copies the input data to the device.
  3. Launches the tree-building kernel with the same configuration parameters (budget, depth limit, topk) that were used to generate the reference.
  4. Copies the kernel output back to the host.
  5. Compares each output tensor element-by-element against the expected reference values, reporting any mismatches.
  6. Returns a pass/fail result that CMake's ctest framework can report. The test is structured so that each KDTR bundle becomes a separate ctest test case, identified by the bundle name (e.g., tree_build_default, tree_build_chain, tree_build_underfull). This design allows individual test cases to be run, debugged, and reported independently.

Assumptions Embedded in the Test Design

The test harness makes several important assumptions, each of which reflects a deliberate engineering choice:

Bit-exact determinism is assumed. The test compares GPU outputs against the reference using exact equality (==), not approximate comparison with tolerances. This is a strong assumption: it requires that the GPU kernel reproduces the exact same numerical operations in the exact same order as the Python reference. The assistant explicitly reasoned about this in the kernel design ([msg 11860]), noting that float tie-breaking is "measure-zero for continuous random floats" and that the kernel would accumulate log-probabilities in the same order as the reference because it replays the same best-first expansion sequence.

The KDTR format is trusted as a round-trip medium. The test assumes that writing reference data from Python to a KDTR file and reading it back in C++ introduces no data corruption or precision loss. This is a reasonable assumption for a custom binary format with explicit type handling, but it does place the format itself outside the test's scope — a bug in kdtr_io.h would produce false test failures.

Each reference bundle is self-contained and correct. The test does not re-run the reference algorithm; it trusts that the pre-computed expected outputs in the KDTR file are correct. This means the reference implementation and the test-data generator must be validated independently, which the assistant did by inspecting the generated outputs manually ([msg 11859]).

The kernel API is stable. The test file hard-codes the kernel launch parameters and output layout. If the kernel interface changes (e.g., adding a new output tensor or changing the padding scheme), the test must be updated in lockstep.

Mistakes and Corrections

The test file itself was not immune to issues. Immediately after writing it, the assistant discovered a missing #include <cstring> header ([msg 11863]), which was required for the memcpy and string functions used in the test. This was fixed with an edit command before the first build attempt.

More subtly, the kernel implementation contained a dead variable remaining_mask_count that was set but never used ([msg 11866]). The compiler emitted a warning, and the assistant removed it in a subsequent edit ([msg 11867]). While this did not affect correctness, it reflects the iterative nature of kernel development — the test file and kernel were refined together as build errors and warnings were resolved.

The Outcome: All Tests Pass

The test harness proved its worth immediately. After fixing the missing include and the dead variable, the assistant built and ran the tests ([msg 11868]):

Test project /home/theuser/glm-kimi-sm120-rtx6000bw/kdtree-engine/build
    Start 1: tree_build_chain .................   Passed    0.26 sec
    Start 2: tree_build_default ...............   Passed    0.24 sec
    Start 3: tree_build_mid ...................   Passed    0.24 sec

All eight initial test cases passed. The assistant then added two more configurations — a maximum-budget case (budget=64) and a high-batch case (batch=64) — and all eleven tests passed ([msg 11871]). This success validated the entire chain: the reference implementation, the test-data generator, the KDTR binary format, the kernel code, and the test harness itself.

Conclusion

The message at [msg 11862] is outwardly unremarkable — a single file write accompanied by a brief reasoning note. But it represents the moment when the assistant committed to a rigorous validation strategy for a complex GPU kernel. The test file it created is the keystone of an architecture that spans Python reference implementations, custom binary formats, and CUDA kernel code, all working together to ensure that the GPU tree builder produces exactly the same results as the trusted SGLang implementation.

In the broader context of the kdtree-engine project, this test harness was the first of several. The assistant would go on to write similar tests for the verify-attention kernel and the tree-accept kernel, culminating in 27 kernel tests and an end-to-end composition test. But it all started here, with a single file that asked the question: does the kernel produce the right answer?