The Bridge: Writing a C ABI for GPU Kernel Integration

[assistant] [write] /home/theuser/glm-kimi-sm120-rtx6000bw/kdtree-engine/src/kernels/capi.cu
Wrote file successfully.

At first glance, message [msg 11907] appears unremarkable: a single tool call that writes a file, confirmed with "Wrote file successfully." There is no reasoning block, no elaborate explanation, no visible decision-making. Yet this message represents a critical architectural pivot — the moment when a collection of validated but isolated CUDA kernels was transformed into a usable software component. The file being written, capi.cu, is the C ABI (Application Binary Interface) shim layer that would bridge the gap between the GPU kernel implementations and the Python-based SGLang inference server they were designed to enhance.

The Context That Made This Message Necessary

To understand why this message was written, one must trace the arc of the session that preceded it. The assistant had just completed Phase 1 of building a native DDTree (Draft-Tree) inference engine for the Kimi K2.6 language model, running on 8× RTX PRO 6000 Blackwell GPUs. Three custom CUDA kernels had been implemented and rigorously validated: a GPU best-first tree builder (tree_build), a tree-verify MLA-absorb attention kernel (verify_attn), and a greedy tree-accept kernel (tree_accept). All 27 tests passed, including a critical on-device composition test that chained tree_build → tree_accept without any host round-trip, proving the device-buffer contract the engine would rely on.

But validated kernels are not yet a usable system. The kernels were written as C++ template functions and CUDA device launches, compiled into a static library (libkdtree_kernels.a). They accepted raw device pointers and dimension parameters — the natural interface for CUDA C++ but inaccessible from Python, the language in which SGLang's inference worker is written. The assistant faced a classic systems integration problem: how to make high-performance CUDA kernels callable from a Python inference framework without sacrificing performance or introducing unnecessary dependencies.

The Reasoning Behind the C ABI Approach

The assistant's reasoning, visible in the preceding message [msg 11906], reveals a deliberate architectural choice. The kernels already accepted raw device pointers — int* and float* arguments representing GPU memory buffers. The natural bridge was a thin C ABI layer: extern "C" wrapper functions that could be compiled into a shared library (.so), loaded via Python's ctypes module, and invoked by passing the data pointers of PyTorch CUDA tensors.

This approach carried several advantages. First, it imposed zero overhead at the call site — the C ABI is the standard calling convention for shared libraries on Linux, and ctypes can invoke it directly without any intermediate binding framework. Second, it maintained full compatibility with PyTorch's CUDA stream semantics: the wrapper could accept a CUDA stream handle and launch kernels on the caller's stream, ensuring correct synchronization with the surrounding computation. Third, it kept the integration surface minimal and testable: the shared library could be verified independently of PyTorch or SGLang, and the same symbols would work identically whether called from a C test harness or from Python.

The assistant explicitly noted the limitation: "I can build and verify the symbols locally, though the full torch integration testing would happen on CT200." This was a pragmatic scoping decision — the local development machine (an RTX 5070 Ti with the same sm_120 architecture) could compile and export the symbols, but the actual end-to-end integration with the 1-trillion-parameter model would require the target hardware.

Input Knowledge Required

Understanding this message requires familiarity with several layers of the software stack. The reader must know what a C ABI is and why extern "C" matters for shared library interop. They must understand the CUDA compilation model — that .cu files are compiled by nvcc and can produce both static archives (.a) and shared objects (.so). They must know Python's ctypes module and its ability to call C functions from shared libraries by declaring argument and return types. They must understand PyTorch's tensor memory model — that a CUDA tensor's data_ptr() returns the raw device pointer, and that torch.cuda.current_stream() provides the CUDA stream handle needed for correct kernel launch ordering. And they must understand the broader context: that SGLang's inference worker runs Python code that constructs draft trees on the CPU using heapq, and the goal is to replace that CPU path with GPU kernel calls.

Output Knowledge Created

This message produced a single file: src/kernels/capi.cu. The subsequent build step ([msg 11909]) confirmed that it compiled into libkdtree_kernels_c.so with four exported C symbols: kdtree_build_ddtree, kdtree_max_qlen, kdtree_tree_accept, and kdtree_verify_attn. The Python ctypes wrapper (python/kdtree_kernels.py, written in [msg 11910]) loaded the library and verified it worked by reading MAX_QLEN=65. The integration documentation (docs/sglang_integration.md, written in [msg 11912]) specified exactly how to wire these calls into the SGLang worker: Swap 1 replaces the CPU heapq tree builder with a single GPU kernel call, Swap 2 replaces the host-side greedy accept with the GPU tree_accept kernel.

Assumptions and Potential Pitfalls

The approach made several assumptions. It assumed that the SGLang worker could load a custom shared library and call into it without modifying the SGLang package itself — this required either monkey-patching or a custom fork. It assumed that the CUDA runtime version on CT200 was compatible with the compiled library (both using CUDA 13 for sm_120). It assumed that passing raw device pointers from PyTorch tensors was safe — that PyTorch would not move or free the tensors during kernel execution, which is true if the CUDA stream is properly synchronized. And it assumed that the performance gain from GPU tree construction would outweigh the overhead of the ctypes call and the kernel launch, which is a reasonable assumption given that the CPU heapq approach is O(budget × depth) per request and serializes across the batch.

The Broader Significance

Message [msg 11907] is the bridge between two worlds: the world of pure CUDA kernel development, where correctness is measured in bit-exact comparisons against numpy references, and the world of production inference serving, where a kernel is only valuable if it can be integrated into a running system. Writing capi.cu was the act of making the kernels usable — not just correct, but callable. It transformed a research prototype into a deployable component. In the architecture of the kdtree-engine repository, this file is the seam where the GPU-optimized compute meets the Python orchestration layer, and its existence is what makes the subsequent CT200 integration session possible.