Bridging the Gap: Reading SGLang's Triton Backend to Integrate a Custom CUDA Verify Kernel

In the high-stakes world of deploying large language models on novel hardware, the difference between a working prototype and a production-ready system often comes down to a single, deceptively simple act: reading someone else's code. Message [msg 12252] captures this exact moment in an ongoing effort to deploy a custom CUDA attention kernel for speculative decoding on NVIDIA RTX PRO 6000 Blackwell GPUs (sm_120 architecture). After weeks of building, tuning, and validating a custom flash-decode verify attention kernel that achieves 3–6× speedup over the Triton baseline, the assistant faces the final, most delicate challenge: integrating that kernel into the live SGLang inference service without breaking anything. This message is the first step of that integration — a targeted source-code reconnaissance mission into SGLang's attention backend.

The Context: Phase 1 Complete, Phase 2 Beckons

To understand why this message was written, we must trace the arc of the project. The assistant had just completed Phase 1 of a multi-phase plan (documented in plans/0002-sm120-verify-kernel-defrag.md): a custom CUDA kernel (verify_attn_flash.cu) implementing split-K, warp-per-key, online softmax attention for the DDTree speculative decoding drafter. The kernel was token-exact against a naive oracle baseline, delivered 1.7–4.2× speedup on microbenchmarks, and could handle sequence lengths of 65,536 tokens that the naive approach physically could not fit in shared memory ([msg 12247]). The float4 vectorization optimization had pushed achieved bandwidth from 403 GB/s to approximately 440 GB/s ([msg 12244]). Phase 1 was committed to git with the message "Phase 1: owned sm_120 flash-decode verify kernel" ([msg 12249]).

But a kernel that performs well in isolation is not yet a kernel that accelerates a production service. Phase 2 required integrating this custom kernel into SGLang's attention mechanism — specifically, subclassing TritonAttnBackend and overriding the DDTree verify forward pass to call the custom CUDA kernel via ctypes. The assistant recognized this would be the riskiest part of the project: "Phases 2-4 involve integrating into the live SGLang service on production, which is risky and complex" ([msg 12250]). The live service serves a 1-trillion-parameter model across 8 GPUs; a mistake could crash inference for all users.

The Message: A Targeted Source-Code Probe

Message [msg 12252] consists of two parts: an Agent Reasoning block that explicitly states the investigative goal, and a bash command that executes the investigation. The reasoning is direct and unambiguous:

I need to examine the forward_extend method in triton_backend.py to see what tensors it works with and how it invokes the extend_attention_fwd function with the KV pool and custom mask—this is the key method I'll need to override for the DDTree verification.

The bash command uses sed to extract lines 874–985 from the triton_backend.py file on the remote CT200 server. This is a surgical read — not the entire file, not a grep for keywords, but a precise line-range extraction targeting the forward_extend method. The choice of sed -n '874,985p' reveals that the assistant already knows (from the previous reconnaissance in [msg 12251]) approximately where this method lives in the file. The previous message had grepped for function definitions and found forward_extend at line 48 of the class definition, but the actual method body starts later. The assistant is now zeroing in.

The output returned shows the method signature and the beginning of its body:

def forward_extend(
    self,
    q: torch.Tensor,
    k: torch.Tensor,
    v: torch.Tensor,
    layer: RadixAttention,
    forward_batch: ForwardBatch,
    save_kv_cache=True,
    sinks=None,
):
    # TODO: reuse the buffer across layers
    if layer.qk_head_dim != layer.v_head_dim:
        o = q.new_empty((q.shape[0], layer.tp_q_head_num * layer.v_head_dim))
    else:
        o = torch.empty_like(q)

    if k is None and v is None...

This is the raw data the assistant needs to plan the integration. The method receives query (q), key (k), value (v) tensors, a RadixAttention layer object, and a ForwardBatch containing metadata including the custom mask and KV indices. The output tensor o is allocated based on the head dimension configuration.

The Reasoning Process: What the Assistant Is Thinking

The assistant's reasoning reveals several layers of strategic thinking. First, there is the recognition that the integration surface is the forward_extend method — this is where the DDTree verify pass occurs, using the extend_attention_fwd Triton function with a custom mask. The assistant needs to understand exactly what data flows into this method to determine whether the custom CUDA kernel can be substituted directly, or whether data transformations (like gathering paged KV into contiguous buffers) are needed.

Second, there is an implicit assumption that the integration will follow a subclass-and-override pattern. The assistant earlier stated: "I need to subclass TritonAttnBackend and override just the DDTree verify forward pass to call our custom kernel through ctypes" ([msg 12250]). This is a reasonable architectural approach — Python's inheritance model makes it natural to override a single method while keeping the rest of the backend intact. However, this assumption carries risk: if forward_extend is called from C++ code or if the Triton backend is instantiated in ways that make subclassing difficult, the approach may need revision.

Third, the assistant is aware of a fundamental data format mismatch. SGLang stores the KV cache as a paged latent pool where kv_indices point to scattered memory slots — a design optimized for memory efficiency with prefix sharing. The custom kernel, in its current form, expects contiguous tensors. The assistant had already identified this tension in [msg 12250]: "Gathering 200k tokens of KV into contiguous memory at each step would be expensive and defeat the purpose of paging. The better approach is to adapt the kernel itself to read paged KV directly via kv_indices." This means the integration may require modifying the CUDA kernel to accept a block table and perform indirect loads — a significant engineering effort that the assistant is mentally preparing for.

Input Knowledge Required

To fully understand this message, one needs knowledge spanning several domains. On the CUDA side, one must understand what a flash-decode verify attention kernel does — it computes attention scores for a small set of query tokens against a large KV cache, using online softmax and split-K reduction to avoid storing intermediate results. The sm_120 architecture refers to the compute capability of the RTX PRO 6000 Blackwell consumer GPU, which lacks the tensor core instructions (wgmma, TMA, tcgen05) available on Hopper and Blackwell-DC architectures — hence the need for a custom kernel rather than using off-the-shelf solutions like FlashMLA.

On the SGLang side, one must understand the TritonAttnBackend class, the forward_extend method's role in speculative decoding verification, and the paged KV cache design. The RadixAttention layer and ForwardBatch are SGLang abstractions for managing prefix caching and batched inference. The custom mask is a tree-structured attention mask used in DDTree speculative decoding, where multiple candidate continuations are verified in parallel.

On the engineering strategy side, one must appreciate the risk calculus: the assistant is deliberately choosing to investigate the codebase before writing any integration code, a classic "measure twice, cut once" approach. The alternative — guessing the interface and writing the integration blind — would risk breaking the production service.

Output Knowledge Created

This message produces concrete, actionable knowledge. The assistant now knows the exact signature of forward_extend, the parameter types, and the output allocation strategy. It can see that the method checks qk_head_dim != v_head_dim to decide the output shape — a detail that may affect how the custom kernel's output is handled. The truncation at if k is None and v is None... means the assistant needs to read further to understand the full data flow, but the critical entry point has been identified.

More subtly, the message creates knowledge about the process of integration. The assistant has established a pattern: read the source, understand the data flow, then design the integration. This pattern will be repeated for each subsequent phase. The message also implicitly documents the location of the integration surface (lines 874–985 of triton_backend.py), which will be useful for future reference.

Assumptions and Potential Pitfalls

The assistant makes several assumptions worth examining. It assumes that subclassing TritonAttnBackend and overriding forward_extend is the correct integration strategy. This is reasonable but may encounter practical obstacles: the Triton backend may be instantiated via factory methods that don't easily accommodate subclasses, or the override may need to intercept at a lower level (inside the Triton kernel launch itself) rather than at the Python method level.

The assistant also assumes that the custom kernel's output format is compatible with what SGLang expects from the attention layer. The kernel produces attention output tensors; if SGLang expects additional metadata (like attention probabilities for debugging, or specific memory layouts for the subsequent layers), the integration may need adjustments beyond simply calling the kernel.

There is also an assumption about the KV cache format. While the assistant correctly identifies the paged nature of SGLang's KV cache, it has not yet confirmed whether the forward_extend method receives already-gathered contiguous KV tensors (the k and v parameters might be gathered by an earlier step) or whether it receives indices into the paged pool. The output snippet shows if k is None and v is None... — suggesting that sometimes k and v are indeed passed as tensors, which would mean the gathering happens upstream. This is a critical detail that will determine whether the kernel needs paged support or can work with the already-contiguous tensors.

Conclusion

Message [msg 12252] is a quiet but pivotal moment in a complex engineering effort. It represents the transition from isolated kernel development to production integration — a phase where correctness and performance must be weighed against operational risk. The assistant's methodical approach — reading the source, understanding the data flow, and reasoning about the integration strategy before writing code — exemplifies disciplined engineering in the face of complexity. The output, while seemingly just a code snippet, is the foundation upon which the entire Phase 2 integration will be built. In the broader narrative of deploying speculative decoding on novel hardware, this message is where the rubber begins to meet the road.