The Benchmark That Bridges Two Worlds: Writing bench_tree_build_vs_sglang.py

"The cleanest meaningful benchmarks without destabilizing the live service: (a) GPU tree-build vs SGLang's actual CPU build_ddtree_tree_from_topk on identical inputs, (b) verify_attn at real K2.6 shapes on a PRO 6000, (c) the live DDTree baseline tok/s."

In the span of a single, deceptively short message, the assistant confirmed the creation of a file that encapsulates an entire strategic pivot:

[assistant] [write] /home/theuser/glm-kimi-sm120-rtx6000bw/kdtree-engine/python/bench_tree_build_vs_sglang.py
Wrote file successfully.

This is message [msg 11975] in the conversation — a terse success notification that a Python benchmark script was written to the repository. On its surface, it is barely a sentence. But to understand why this particular file exists, at this particular moment, requires unpacking a cascade of reasoning, risk assessment, and architectural judgment that unfolded over the preceding messages.

The Strategic Crossroads

The user's directive was clear: "proceed with next phases up to getting meaningful benchmarks on CT200" ([msg 11968]). CT200 is the 8× RTX PRO 6000 Blackwell machine running the live Kimi K2.6 DDTree speculative decoding service. The assistant had just delivered a fully validated native C/C++/CUDA DDTree inference engine ([msg 11967]) — a complete MLA+MoE transformer with three custom CUDA kernels (tree builder, tree-verify attention, tree acceptor), proven greedy-exact against numpy references across two model configurations. The engine was ready. The hardware was ready. The next phase was benchmarking on the real target.

But here the assistant faced a fork in the road. Option A was the grand vision: scale the native engine to the full 1-trillion-parameter K2.6 model with INT4 Marlin GEMMs, TP-8 NCCL sharding, and the real DFlash drafter. This would be the ultimate validation — a complete from-scratch inference stack running on Blackwell hardware. But as the assistant's reasoning in [msg 11969] makes clear, this path "would take months of work to fully reimplement from scratch" and was "too risky and not feasible autonomously in a reasonable timeframe."

Option B was the documented Phase-1 integration path: drop the validated kernels (tree_build, tree_accept, verify_attn) into the running SGLang DDTree worker via a ctypes bridge, restart the service, and measure the end-to-end throughput improvement on the actual 1T model. This was the pragmatic high-value route — real results on real hardware, leveraging what was already built.

Option C was the safest: build microbenchmarks of individual kernel components at K2.6 scale, producing meaningful numbers without touching the live service at all.

The Decision That Shaped the File

The critical insight came in [msg 11972]. The assistant realized that Option B — patching the live SGLang service — carried real risk. The service was actively serving requests across 8 GPUs, each consuming ~87GB of memory. Restarting it meant downtime. Patching it meant potential crashes. And there was a subtler concern: even if the GPU tree builder worked, isolating its contribution to end-to-end throughput would be confounded by other variables (CUDA graph fragmentation, drafter acceptance rates, KV cache behavior).

The assistant then arrived at a brilliant compromise: instead of modifying the live service, benchmark the CPU tree-build function in isolation by importing SGLang's build_ddtree_tree_from_topk directly and timing it against the GPU kernel for identical inputs. As the reasoning states: "This gives me a clean apples-to-apples comparison of the bottleneck removal without touching the service at all."

This decision is the direct parent of the target message. The file bench_tree_build_vs_sglang.py is the Python harness that makes this comparison possible. It uses ctypes to load the compiled CUDA kernel shared object, generates realistic top-k inputs matching K2.6's hidden-state dimensions, calls both the GPU kernel and SGLang's CPU implementation on identical data, and reports wall-clock time and throughput for budget sizes of 8, 16, and 32.

What the File Represents

The file is more than a benchmark script. It is a bridge between two worlds. On one side is the native C/C++/CUDA engine — a ground-up reimplementation of the DDTree speculative decoding loop, built with the discipline of bit-exact validation against numpy references. On the other side is SGLang — the production inference framework that actually serves the Kimi K2.6 model to users. The benchmark script is the first point of contact between these two codebases, the moment where the custom CUDA kernels meet the real inference stack.

The script's design reflects the assistant's deep understanding of the system. It tests at budget sizes of 8, 16, and 32 because the earlier DDTree findings report ([msg 11964]) had identified that larger budgets improved acceptance rates (from ~4.5 to ~5.3–6.4 tokens/step) but were bottlenecked by the CPU heapq implementation. The script tests with batch size 64 because that matches the DDTree worker's typical batch configuration. It tests on the actual PRO 6000 GPU because the Blackwell SM120 architecture has specific performance characteristics that cannot be simulated on other hardware.

Input Knowledge Required

To write this file, the assistant needed to know:

1. The SGLang DDTree worker internals: Specifically that build_ddtree_tree_from_topk is the CPU function performing per-request heapq-based tree construction, located in ddtree_utils.py within the SGLang speculative module.
2. The CUDA kernel interface: The exact function signatures, memory layouts, and calling conventions of the GPU tree builder kernel, which was built in the preceding session ([msg 11967]).
3. The ctypes bridge mechanism: How to load a .so shared object from Python, allocate GPU memory via CUDA runtime API calls, marshal data between numpy arrays and device pointers, and synchronize after kernel launches.
4. The K2.6 model's hidden-state dimensions: The top-k logit shape that determines the input size for tree construction, which varies with the model's vocabulary size and the DDTree budget parameter.
5. The CT200 environment: That the target machine has CUDA 13.0, no cmake (hence the nvcc-direct build), and a running SGLang service consuming most GPU memory, meaning the benchmark must be memory-efficient.

Output Knowledge Created

The file produces a quantitative comparison that answers a specific engineering question: How much faster is the GPU tree builder than the CPU heapq at realistic K2.6 shapes on Blackwell hardware? The answer to this question determines the entire next phase of the project. If the GPU kernel is 10× faster at budget 32, then the bottleneck is removed and the SGLang integration (Option B) becomes the obvious next step. If the GPU kernel is only marginally faster, then the bottleneck lies elsewhere — perhaps in the verify attention kernel or the drafter itself — and the optimization strategy must be rethought.

The file also creates a reproducible benchmark that can be re-run after any code change, serving as a regression test for the tree builder's performance. This is crucial for an iterative development process where kernel optimizations (shared memory tiling, warp-level reductions, etc.) are applied incrementally.

Assumptions and Their Risks

The benchmark script makes several assumptions that carry risk. It assumes that the CPU build_ddtree_tree_from_topk function can be imported and called in isolation without side effects from the running SGLang service — this is likely true but depends on SGLang's internal module structure. It assumes that the ctypes bridge to the GPU kernel works correctly on CT200's CUDA 13.0 toolchain — the kernel was validated on a local RTX 5070 Ti with the same SM120 architecture, but the PRO 6000 has different memory bandwidth and cache hierarchy. It assumes that the benchmark input distribution (random top-k indices) is representative of real model behavior — in practice, the drafter's top-k distribution may have structure (e.g., sparsity patterns) that affects tree-building performance differently.

The most subtle assumption is that tree-building time is the dominant bottleneck worth measuring in isolation. If the verify attention kernel or the drafter forward pass dominates the step time, then even a 100× improvement in tree building will yield negligible end-to-end gains. The assistant is aware of this — the benchmark package includes verify_attn microbenchmarks as well — but the tree-build comparison is the headline measurement because it directly validates the core claim of the native engine project: that GPU-parallel tree construction can replace the CPU heapq bottleneck and enable larger budgets.

The Thinking Process Visible in the Reasoning

The assistant's reasoning traces a clear arc from ambition to pragmatism. In [msg 11969], the initial instinct is toward the full native engine — "the pragmatic high-value path is to (1) build + unit-test the kernels on the actual target hardware, (2) microbenchmark them at K2.6-realistic shapes, and (3) do the Phase-1 SGLang integration." But by [msg 11972], after exploring CT200's environment and confirming the live service's constraints, the reasoning has sharpened: "Actually, I can get the same insight with much lower risk."

This is the hallmark of experienced engineering judgment — the ability to recognize when a simpler experiment answers the same question with less risk. The assistant resists the temptation to over-engineer the integration, instead choosing the measurement that provides the highest information value per unit of risk. The file bench_tree_build_vs_sglang.py is the concrete manifestation of that judgment.

Conclusion

Message [msg 11975] appears to be a minor administrative note — a file was written, success was confirmed. But in the context of the broader session, it represents a carefully calibrated engineering decision. The assistant chose to build a bridge between the native CUDA engine and the production SGLang stack, not by modifying either system, but by creating a measurement apparatus that can compare them on equal terms. The benchmark script is the first step toward answering whether the GPU tree builder can deliver on its promise of enabling larger DDTree budgets on Blackwell hardware — and that answer will determine the trajectory of the entire project.