The Final Solder Joint: Updating the C-ABI Layer for CUDA Graph Capture-Safety

[assistant] [edit] /home/theuser/glm-kimi-sm120-rtx6000bw/kdtree-engine/src/kernels/capi.cu
Edit applied successfully.

At first glance, this message from an opencode coding session appears almost trivial: a single file edit that succeeded. No fanfare, no benchmark numbers, no dramatic speedup claims. Yet this brief confirmation represents the completion of a critical architectural transition — the moment when a custom CUDA kernel for speculative decoding attention was made compatible with CUDA graph capture, unlocking production-grade performance for the Kimi K2.6 model on RTX PRO 6000 Blackwell GPUs. To understand why this three-word edit matters, one must trace the chain of reasoning that led to it.

The Mandate: Three Priorities

The story begins with the user's directive in [msg 12298]: "Fix cudagraphs, then do defrag and optimise marshaling." This was a clear prioritization. The assistant's custom DDTree verify attention kernel was already live and delivering a validated 2× speedup on long-context decoding (≥23k tokens) compared to the Triton baseline ([msg 12297]). But it came with a severe limitation: it required --disable-cuda-graph, which disabled CUDA graphs globally. On short contexts (5.7k tokens), this penalty pushed performance below parity at 0.93×. The production system needed CUDA graphs enabled — and that meant the verify kernel had to become capture-safe.

The Capture-Safety Constraint

CUDA graphs work by recording a sequence of GPU operations (kernel launches, memcpys, etc.) into a graph object that can be replayed with minimal CPU overhead. For SGLang's inference engine, this means pre-allocating static metadata buffers at fixed batch sizes and recording operations against them. On replay, SGLang copies new data into those same buffers and replays the graph, which executes the identical operations on the updated data.

The assistant's original kernel broke this contract in several ways ([msg 12299]). It performed host synchronizations via .item() calls to extract scalar values from GPU tensors. It used raw cudaMalloc in a workspace allocation function (ensure_ws) that bypassed PyTorch's caching allocator — a direct violation of capture constraints. Most critically, it rebuilt combined KV indices and visibility masks in Python loops, creating temporary tensors with different memory addresses than SGLang's static buffers. During graph replay, those stale pointers would cause the kernel to read old data.

The Investigation: Ground Truth Through Source Inspection

Before any code could be written, the assistant needed to know the exact data types of SGLang's static buffers. In [msg 12299], it ran a series of grep commands on the remote server to extract type information from SGLang's source code. The results were precise: kv_indices, qo_indptr, mask_indptr, and out_cache_loc were all int64; kv_indptr was int32; and custom_mask was uint8 during capture mode (versus bool in eager mode, both occupying one byte). This seemingly mundane step was crucial — the kernel would consume these buffers directly, and any type mismatch would cause silent data corruption.

The Three-Edits Sequence

The capture-safety refactoring unfolded across three coordinated edits, each targeting a different layer of the codebase:

  1. The kernel itself (verify_attn_flash_paged.cu, [msg 12300]): The CUDA kernel was rewritten to accept SGLang's native buffer pointers directly. Instead of receiving pre-combined KV indices, it reads prefix tokens from kv_indices and newly generated tokens from out_cache_loc, computing the distinction in-kernel from the kv_indptr offsets. Instead of a pre-gathered visibility mask, it indexes into custom_mask using mask_indptr offsets. The grid size was fixed with a constant NSPLIT=64, which adapts at runtime: for a 200k-token context, each split processes ~3125 tokens; for a 21-token context, only 21 of 64 splits are active, and the rest early-out cheaply.
  2. The header declaration (verify_attn.cuh, [msg 12301]): The C++ header was updated to declare the new function signature, replacing the old paged verification entry point with the capture-safe cooperative-groups-based interface.
  3. The C-ABI bridge (capi.cu, [msg 12302]): The C Application Binary Interface layer — the thin wrapper that translates between Python ctypes calls and CUDA kernel launches — was updated to match the new signature. This is the message we are examining.

Why the C-ABI Edit Matters

The capi.cu file is the critical glue in the system. Python cannot call CUDA kernels directly; it uses ctypes to invoke C functions that in turn launch CUDA kernels. The C-ABI functions defined in capi.cu must exactly match what the Python side expects and what the kernel side provides. If the kernel signature changes (new parameters, removed parameters, different pointer types), the C-ABI wrapper must be updated in lockstep. An inconsistency here means the Python side passes arguments that the C side interprets incorrectly — a recipe for crashes, silent memory corruption, or wrong results.

In this edit, the C-ABI was updated to pass SGLang's native buffer pointers (kv_indices, out_cache_loc, qo_indptr, kv_indptr, custom_mask, mask_indptr) directly to the kernel, along with a torch-allocated workspace pointer and scalar parameters (bs, q_len, H, Dl, Dr, scale, NSPLIT). The old parameters that required host-side preprocessing — combined indices, pre-gathered masks, workspace size — were removed. The function now derives everything it needs from the native buffers and scalar constants, with no host synchronization required.

Assumptions and Design Decisions

Several assumptions underpin this work. The assistant assumed that NSPLIT=64 fixed at capture time would not impose prohibitive overhead on short contexts — the early-out path for empty splits was deemed cheap enough. It assumed that PyTorch's caching allocator would handle workspace allocations correctly during capture (since torch.empty uses the allocator, not raw cudaMalloc). It assumed that SGLang's static buffer pointers remain valid across replay — a correct assumption given SGLang's design, but one that must hold for the entire approach to work.

The assistant also made a subtle but important architectural decision: rather than overriding SGLang's metadata replay function, it chose to consume the native buffers directly. This meant the kernel had to compute prefix lengths and slot indices in-kernel from kv_indptr and out_cache_loc rather than receiving them as pre-computed parameters. This increased kernel complexity but eliminated the need for any Python-side tensor manipulation during capture, keeping the entire data path on the GPU.

The Knowledge Chain

To produce this edit, the assistant needed deep knowledge spanning multiple domains: CUDA graph capture semantics (what operations are and aren't capturable), SGLang's internal buffer management (static pre-allocation, replay mechanics), the DDTree verify attention algorithm (split-K flash decode, online softmax, per-head warp-per-key), the Kimi K2.6 model architecture (MLA with absorbed projections, 61-layer transformer), and the sm_120 instruction set architecture of the RTX PRO 6000 Blackwell GPU. Each of these knowledge domains was accumulated over the preceding segments of the session, from initial environment setup through kernel development, debugging, and integration.

The Outcome

The chunk summary ([chunk 66.1]) confirms the result: "CUDA graph capture succeeded in ~1.5 seconds, and decode now replays the graph correctly with generations matching the Triton baseline." The capture-safe kernel, combined with subsequent tuning (increasing NSPLIT to 64 and adding 128-bit vectorized bf16 KV loads), delivered a dramatic 3–6× end-to-end decode speedup over Triton with CUDA graphs across all context lengths from 4k to 65k tokens. The C-ABI edit was the final solder joint in this chain — without it, the new kernel would never have been callable from Python, and the 3–6× speedup would have remained a simulation rather than a production reality.