Bridging Reference and Kernel: The Test-Data Generator That Validates a CUDA DDTree Engine

In the high-stakes world of custom inference engine development, the gap between a correct algorithm and a correct implementation is where most bugs hide. Message [msg 11858] captures the precise moment when an AI assistant bridges that gap: the creation of a test-data generator that will serve as the oracle for validating three custom CUDA kernels against their numpy reference implementations. This message, embedded within a larger effort to build a native C/C++/CUDA Draft-Driven Tree (DDTree) inference engine for the Kimi K2.6 language model, represents the critical infrastructure layer that transforms abstract algorithmic correctness into verifiable, deterministic test cases.

The Broader Mission: A Native DDTree Engine

To understand why this message matters, one must first understand the project it serves. The assistant is building a complete native DDTree inference engine from scratch — a specialized speculative decoding system that uses a lightweight "drafter" model to propose multiple candidate tokens in parallel, then verifies them against the full target model. The engine lives in a new repository called kdtree-engine/, organized into phases: Phase 0 establishes the build system and test infrastructure, Phase 1 implements three custom CUDA kernels (tree builder, tree-verify attention, and tree-accept), and Phase 2 assembles them into a working MVP transformer engine.

By message [msg 11858], the assistant has already laid the groundwork. The repository is scaffolded ([msg 11851]). The implementation plan is saved ([msg 11852]). A custom binary container format called KDTR has been defined with both a Python writer (kdtr_io.py in [msg 11855]) and a C++ reader (kdtr_io.h in [msg 11856]). Most critically, a faithful numpy reference implementation of the DDTree tree-building algorithm has been written (ddtree_ref.py in [msg 11857]), ported directly from SGLang's build_ddtree_tree_from_topk function.

What remains — and what message [msg 11858] provides — is the glue that connects these pieces: a script that generates realistic test data, runs it through the reference to produce expected outputs, and packages everything into the KDTR format that the C++ test harness can consume.

The Subject Message: Generating Test Data for Kernel Validation

The message itself is concise but dense with purpose. The assistant writes:

Now the test-data generator that writes input+expected bundles: [write] /home/theuser/glm-kimi-sm120-rtx6000bw/kdtree-engine/python/gen_tree_build_refs.py Wrote file successfully.

The reasoning block preceding the action reveals the key design constraint:

The key constraint is that top_log_probs need to be sorted in descending order for each depth row to match what the parallel head function produces, so I'll generate random values and sort each row accordingly, with token_ids being randomly selected and distinct per row.

This constraint is not arbitrary. In the actual DDTree inference pipeline, the "parallel head" of the drafter model produces log probabilities for the top-k tokens at each depth position, and these log probabilities are naturally sorted in descending order — the most likely token first, the least likely among the top-k last. The test-data generator must respect this invariant to produce inputs that mirror real-world conditions. If the generator produced unsorted logprobs, the reference implementation might still produce a valid tree, but the resulting test cases would not reflect the actual data distribution the CUDA kernels will encounter in production, potentially masking subtle ordering-dependent behaviors.

The generator creates multiple configuration files, each targeting different regions of the parameter space. The subsequent message ([msg 11859]) shows the results:

wrote .../tests/refs/tree_build_default.kdtr  B=6 D=7 K=4 budget=8 q_len=9 actual=[9,9]
wrote .../tests/refs/tree_build_mid.kdtr  B=4 D=7 K=4 budget=16 q_len=17 actual=[17,17]
wrote .../tests/refs/tree_build_wide.kdtr  B=4 D=7 K=8 budget=32 q_len=33 actual=[33,33]

Each file contains batches of test cases with varying budget, depth (D), top-k (K), and batch size (B). The actual values confirm that the reference produced the expected number of nodes — matching q_len (budget + 1 for the root) in the normal cases. Edge cases like underfull (where the heap exhausts before reaching budget) and chain (topk=1 producing a linear chain) are also generated, ensuring the test suite covers boundary conditions that often expose kernel bugs.

Why This Message Matters: The Validation Pipeline

The test-data generator occupies a pivotal position in the validation pipeline. The full flow works as follows:

1. Python generates random inputs (top_log_probs, top_token_ids) with the correct structure.
2. The reference implementation (pure numpy, ported from SGLang) processes these inputs to produce expected outputs: node tokens, depths, parent indices, visibility matrices, and retrieval pointers.
3. Both inputs and expected outputs are serialized to a .kdtr binary file using the custom container format.
4. The C++ test harness loads the .kdtr file, launches the CUDA kernel with the same inputs, and compares the kernel's outputs against the expected values. This pipeline embodies a crucial principle in systems programming: never trust that a kernel is correct just because it compiles and runs. The reference implementation — written in pure Python/numpy with no CUDA dependencies — serves as the ground truth. By generating test cases in Python and validating against Python, the assistant ensures that any discrepancy between the CUDA kernel and the reference is a bug in the kernel, not a misunderstanding of the algorithm.

Design Decisions and Assumptions

The generator embodies several deliberate design decisions, each carrying implicit assumptions.

Random generation with sorted rows. The assistant assumes that randomly generated log probabilities, sorted within each depth row, sufficiently approximate the distribution of real model outputs. This is a reasonable assumption for a correctness test: the tree-building algorithm is purely combinatorial, depending only on the relative ordering of log probabilities, not their absolute values. Any set of sorted values will exercise the same heap operations.

Distinct token IDs per row. By ensuring that token IDs are distinct within each depth row, the generator avoids ambiguity in the tree structure. If two nodes at the same depth had the same token ID, the tree would still be valid, but debugging would be harder. This simplification makes test failures easier to diagnose.

Multiple configurations for coverage. The generator produces separate files for default, mid, wide, underfull, and chain configurations. This reflects an assumption that correctness must be verified across the full operational envelope, not just at typical settings. The underfull case (heap empties before budget is reached) and chain case (topk=1) are particularly important because they exercise edge conditions in the heap management logic that might not be triggered by standard configurations.

The binary container format is stable. The assistant assumes that the KDTR format, once defined, will not change in ways that break compatibility between the Python writer and C++ reader. This is a safe assumption within a single project where both sides are developed together, but it highlights the importance of the format specification being correct from the start.

Input Knowledge Required

To fully grasp this message, one needs several pieces of context:

• Speculative decoding with draft trees. The DDTree algorithm uses a lightweight drafter model to propose a tree of candidate token sequences, which the target model then verifies in parallel. The tree is built using a best-first expansion from the drafter's top-k log probabilities at each depth.
• The KDTR binary container format. Defined in kdtr_io.py, this format uses a magic header, version number, and array metadata (name, dtype, shape) followed by raw binary data. It is designed to be trivially readable from both Python and C++ without external dependencies.
• The reference implementation. ddtree_ref.py contains a faithful numpy port of SGLang's build_ddtree_tree_from_topk, including the visibility matrix computation and retrieval pointer construction. The test-data generator calls this reference to produce expected outputs.
• CUDA kernel testing patterns. The standard approach is to generate test data in Python (where prototyping is fast), run a reference implementation, and compare against GPU kernel outputs in C++ (where the kernel lives). This message implements the Python side of that pattern.

Output Knowledge Created

Message [msg 11858] produces a reusable test-data generation script that:

1. Creates deterministic, reproducible test cases for the tree-builder kernel.
2. Covers normal operation (default, mid, wide configurations) and edge cases (underfull, chain).
3. Respects the structural invariants of real model outputs (sorted logprobs, distinct token IDs per row).
4. Packages inputs and expected outputs in a format directly consumable by the C++ test harness.
5. Can be re-run to regenerate test data if the reference implementation changes. The immediate output — the .kdtr files — is visible in the next message ([msg 11859]), where the assistant verifies the round-trip by loading a file and inspecting its contents. The verification confirms that the data structure is correct: keys include meta, actual, node_token_ids, node_depths, parents, next_token, next_sibling, and visibility, all with the expected shapes and values.

The Thinking Process: Precision in Infrastructure

The assistant's reasoning reveals a methodical approach to infrastructure. Rather than diving directly into kernel implementation, the assistant invests in the testing pipeline first. This is a classic strategy in systems programming: build the oracle before building the system under test, so that every implementation step can be validated immediately.

The attention to detail is notable. The assistant considers the sorting order of log probabilities, the distinctness of token IDs, the coverage of edge cases, and the stability of the binary format. These are not afterthoughts — they are first-order design decisions that determine whether the test suite will catch bugs or silently pass over them.

The assistant also demonstrates awareness of the limits of random testing. Random log probabilities with distinct values will almost never produce exact ties, but the real model might. The assistant addresses this implicitly by testing multiple configurations and explicitly by noting in earlier reasoning that tie-breaking behavior is documented as a minor divergence. The test-data generator, by using random values, will naturally produce cases without ties — which is fine for initial validation, but suggests that future work might need targeted tie-breaking tests.

Conclusion

Message [msg 11858] is a masterclass in building test infrastructure for systems-level GPU programming. It is not flashy — it does not contain a novel algorithm or a performance optimization — but it is essential. Without this test-data generator, the CUDA kernels that follow would be validated against hand-crafted test cases at best, or against nothing at all. With it, every kernel can be tested against hundreds of randomly generated cases spanning the full parameter space, with ground truth provided by a faithful reference implementation.

In the broader narrative of the DDTree engine development, this message represents the moment when the project transitions from planning and infrastructure into the actual implementation loop. The test-data generator is the first piece of the validation pipeline that will accompany every kernel, every optimization, and every bug fix throughout the project's lifetime. It is the silent partner that ensures the assistant can write "27 kernel tests passed bit-exact against the references" — a claim that would be meaningless without the rigorous test infrastructure this message establishes.