The Reconnaissance Phase: Mapping the Integration Surface Between a Custom CUDA Kernel and SGLang's Triton MLA Backend

Introduction

In the lifecycle of any deep systems integration project, there comes a moment when the developer must transition from building an isolated component to understanding how it will interface with a complex production system. Message [msg 12251] captures precisely this moment in the ongoing effort to deploy a custom sm_120 flash-decode verify attention kernel into SGLang's production inference stack. After successfully designing, implementing, and optimizing a custom CUDA kernel that achieves 3–6× speedup over the Triton baseline in isolation, the assistant now faces the critical task of understanding the integration surface—the exact data formats, control flow, and API boundaries through which the new kernel must communicate with SGLang's attention backend.

This message represents a reconnaissance operation: a deliberate, focused investigation into SGLang's Triton-based Multi-head Latent Attention (MLA) backend to map out precisely how the verify path works, what tensors flow through it, and how features like paged KV cache and custom masking are handled. The assistant is not yet modifying code or deploying anything—it is gathering intelligence to inform the next phase of work.

The Message in Full

The message consists of two parts: an agent reasoning block that articulates the investigation's purpose, and a bash command that executes the reconnaissance on the production server at 10.1.230.171.

Agent Reasoning:

I'm diving into the Triton MLA backend to understand exactly how the verify path works—specifically what data flows through, how it processes paged KV cache, and applies custom masking. I need to trace through the forward_extend method and locate the MLA-specific implementation details, including how custom_mask and kv_indices interact.

Bash Command:

SRT=/root/venv_sglang211/lib/python3.12/site-packages/sglang/srt
timeout 30 ssh -o StrictHostKeyChecking=no root@10.1.230.171 "
echo '### triton backend: MLA / forward_extend / custom_mask / kv_indices ###'
grep -nE 'def forward_extend|def forward_decode|use_mla|def _forward|custom_mask|mla|q_rope|kv_lora|extend_attention|decode_attention' $SRT/layers/attention/triton_backend.py | head -40
echo '### is there a separate MLA forward in triton ops? ###'
ls $SRT/layers/attention/triton_ops/ 2>/dev/null
"

Output:

### triton backend: MLA / forward_extend / custom_mask / kv_indices ###
48:    custom_mask: torch.Tensor
67:        from sglang.srt.layers.attention.triton_ops.decode_attention import (
68:            decode_attention_fwd,
70:        from sglang.srt.layers.attention.triton_ops.extend_attention import (
72:            extend_attention_fwd,
73:            extend_attention_fwd_unified,
78:        self.decode_attention_fwd = torch.compiler.disable(decode_attention_fwd)
79:        self.extend_attenti...

Why This Message Was Written: The Context and Motivation

To understand the motivation behind this message, we must look at the broader arc of the project. The assistant has been working on deploying the Kimi K2.6 model with DFlash speculative decoding on RTX PRO 6000 Blackwell GPUs (sm_120 architecture). A critical bottleneck was identified: the DDTree verify attention kernel, which runs during speculative decoding to verify draft tokens, was using SGLang's Triton-based MLA implementation with page_size=1, causing scattered KV cache accesses that achieved only ~14 GB/s effective bandwidth—roughly 130× below the GPU's 1.8 TB/s peak.

The assistant built a custom CUDA kernel (verify_attn_flash.cu) that solved this problem in isolation, achieving 3–6× speedup over the Triton baseline across context lengths from 4k to 65k tokens. Phase 1 was successfully completed and committed to the repository. However, a custom CUDA kernel that works in a standalone test harness is not yet a production asset. To realize the speedup in the live SGLang service, the kernel must be integrated into SGLang's attention backend—replacing the Triton-based verify path with the custom implementation while preserving all the surrounding machinery for paged KV cache, custom masking, tensor parallelism, and CUDA graph capture.

This message is the opening move of Phase 2: integration. The assistant's stated goal is to "scope the integration surface precisely." Before writing any integration code, it must understand exactly what data structures and calling conventions the Triton backend uses, so it can design a compatible replacement.

How Decisions Were Made

The decision-making in this message is subtle but important. The assistant faces a fork in the road: it could either (a) dive into the SGLang source code locally and study it, (b) experiment by trial and error with live integration attempts, or (c) use targeted instrumentation to map the integration surface from the production server. It chooses option (c), which reflects several implicit judgments:

First, the assistant judges that understanding the exact production code path is critical. SGLang is under active development, and the version running on the server (sglang installed in a virtual environment at /root/venv_sglang211/) may differ from the public repository. By grepping the actual installed package, the assistant ensures it is studying the exact code that will be modified.

Second, the assistant chooses a lightweight, non-destructive reconnaissance method. Rather than attaching a debugger, modifying code, or running the service with special flags, it uses grep to extract relevant lines from the source files. This is safe—it cannot crash the production service—and fast, completing within a 30-second SSH timeout.

Third, the assistant prioritizes specific search terms that reflect its understanding of the architecture. The grep pattern targets forward_extend, forward_decode, use_mla, _forward, custom_mask, mla, q_rope, kv_lora, extend_attention, and decode_attention. These terms are not chosen randomly; they reflect the assistant's mental model of how SGLang's attention backend is structured. The presence of q_rope and kv_lora indicates awareness that MLA uses a latent representation with rotary position embeddings and low-rank projections—a key architectural detail of the Kimi K2.6 model.

Fourth, the assistant also checks whether there is a separate directory for Triton attention ops (triton_ops/), which would indicate a modular structure where different attention implementations live in separate files. This is a structural reconnaissance question: is the MLA verify logic inline in the backend class, or is it delegated to a separate Triton kernel file?

Assumptions Made

This message rests on several assumptions, some explicit and some implicit:

Assumption 1: The verify path goes through forward_extend or forward_decode. The assistant searches for these method names, assuming that the DDTree verify attention is dispatched through one of these two entry points. This is a reasonable assumption given SGLang's architecture, where forward_extend handles prefill and extended-context attention, while forward_decode handles token-by-token autoregressive decoding. However, the DDTree verify path is neither pure prefill nor pure decode—it verifies multiple candidate tokens against a prefix—so it might use a specialized path.

Assumption 2: The integration surface is primarily in triton_backend.py. The assistant focuses on a single file: /layers/attention/triton_backend.py. This assumes that the MLA verify logic is implemented in the Triton backend class, rather than being distributed across multiple files or handled by a different backend entirely (e.g., FlashInfer or a native CUDA backend).

Assumption 3: custom_mask and kv_indices are the key data structures to understand. The assistant specifically wants to understand how custom masking and paged KV indices interact. This reflects its understanding that the DDTree verify path uses a custom attention mask (to handle the tree structure of candidate tokens) and that the KV cache is paged (requiring kv_indices to map logical positions to physical slots). Both of these are correct assumptions based on the earlier diagnostic work.

Assumption 4: The SSH connection will succeed and the grep will complete within 30 seconds. This is an operational assumption that proves correct—the output is returned successfully.

Mistakes or Incorrect Assumptions

The message itself does not contain obvious mistakes—it is a reconnaissance command that returns useful information. However, we can identify a potential blind spot in the investigation strategy:

The investigation is narrow in scope. The assistant searches only triton_backend.py and checks for the existence of triton_ops/. It does not examine the actual Triton kernel implementations in decode_attention.py or extend_attention.py, nor does it look at how the backend class is instantiated or how the attention forward pass is dispatched from higher-level code. The grep output shows that decode_attention_fwd and extend_attention_fwd are imported from separate files, but the assistant does not immediately follow up by examining those files. This narrow focus means the assistant may need additional reconnaissance rounds before it has a complete picture.

Additionally, the assistant's search pattern may miss important details. It searches for use_mla but the output does not show any matches for this term in the first 40 lines. It searches for q_rope and kv_lora but these also do not appear in the truncated output. The assistant might need to broaden its search or examine the file more thoroughly to find the MLA-specific code paths.

Input Knowledge Required

To understand this message, the reader needs knowledge in several domains:

SGLang Architecture: The reader must understand that SGLang uses a modular attention backend system where different attention implementations (Triton, FlashInfer, native CUDA) can be selected at runtime. The TritonAttnBackend class in triton_backend.py is the Triton-based implementation, and it provides forward_extend and forward_decode methods that are called by the model's attention layers.

MLA (Multi-head Latent Attention): The Kimi K2.6 model uses MLA, a form of attention where the key and value are projected into a low-dimensional latent space before being cached. This means the KV cache stores latent representations rather than full key/value vectors, and the attention computation involves absorbing the query into the latent space. The assistant's search for q_rope and kv_lora reflects this architecture.

DDTree Speculative Decoding: The DFlash speculative decoding framework uses a draft model to propose candidate tokens in a tree structure, and a verify step that computes attention over all candidates simultaneously. This verify step is the target of the custom kernel.

Paged KV Cache: Production inference servers use paged KV cache allocation to manage memory efficiently. The KV cache is stored in fixed-size pages, and a mapping (kv_indices) translates logical token positions to physical page slots. The custom kernel must handle this paged format.

CUDA Graph Capture: SGLang uses CUDA graphs to reduce kernel launch overhead during decode. The custom kernel must be capture-safe—meaning it cannot perform host-device synchronization, dynamic memory allocation, or other operations that are incompatible with graph capture.

Output Knowledge Created

This message produces several pieces of knowledge:

1. Confirmation of the file structure. The output confirms that triton_backend.py exists at the expected path and contains the expected class structure. Line 48 shows custom_mask: torch.Tensor, confirming that the backend class has a custom_mask attribute. Lines 67–79 show the import structure for decode_attention_fwd and extend_attention_fwd, confirming that the Triton kernels are in separate files under triton_ops/.

2. The triton_ops/ directory exists. The ls command succeeds (no error output), confirming that there is a separate directory for Triton attention ops. This tells the assistant that the actual Triton kernel implementations are modular and can potentially be replaced individually.

3. torch.compiler.disable is used. Line 78 shows self.decode_attention_fwd = torch.compiler.disable(decode_attention_fwd). This is a significant detail: it tells the assistant that SGLang explicitly disables Torch's compiler for the decode attention forward function, presumably to prevent unwanted graph optimizations from interfering with the Triton kernel. The custom kernel integration may need to follow the same pattern.

4. The integration surface is scoped. The assistant now knows that it needs to understand two files: the backend class (triton_backend.py) and the Triton kernel files (decode_attention.py and extend_attention.py). The next step would be to examine those kernel files to understand the exact tensor signatures and calling conventions.

The Thinking Process Visible in the Reasoning

The agent reasoning block reveals a structured investigative methodology. The assistant states three specific things it needs to understand:

  1. "what data flows through" — the tensor formats and shapes
  2. "how it processes paged KV cache" — the mechanism for handling non-contiguous KV storage
  3. "how it applies custom masking" — the tree-structured attention mask used in DDTree verification The phrase "trace through the forward_extend method" indicates an intention to follow the control flow from the high-level API down to the low-level Triton kernel calls. This is classic reverse engineering: start at the public interface and follow the code path until you reach the primitive operations that need to be replaced. The mention of "MLA-specific implementation details" shows that the assistant recognizes MLA is not standard attention—it has unique features like latent absorption, rope, and LoRA projections that must be handled correctly. The custom kernel already handles MLA absorption (as noted in the Phase 1 commit message: "MLA absorb"), but it must match SGLang's exact convention. The final phrase, "including how custom_mask and kv_indices interact," identifies the most complex part of the integration. The DDTree verify path uses a custom attention mask that defines which tokens each candidate can attend to, and the paged KV cache means that the mask indices must be translated through the page table. Understanding this interaction is essential for correctly implementing the custom kernel's invocation.

Conclusion

Message [msg 12251] is a small but pivotal step in a complex systems integration project. It represents the transition from isolated kernel development to production integration, from "does my kernel work correctly?" to "how do I plug my kernel into the existing system?" The assistant's methodical approach—using targeted grep commands to map the integration surface without risking production stability—reflects a disciplined engineering mindset. The knowledge gained from this reconnaissance will inform the design of the Python backend wrapper, the handling of paged KV cache, and the integration with CUDA graph capture.

This message also illustrates a broader truth about systems engineering: the hardest part of optimization is often not writing the fast kernel, but integrating it into the existing software stack. The custom sm_120 verify kernel achieves 3–6× speedup in isolation, but realizing that gain in production requires understanding every detail of SGLang's attention backend—the tensor formats, the calling conventions, the paging mechanism, the masking logic, and the graph capture constraints. Message [msg 12251] is the first step in that journey, and the output it produces will directly inform the implementation of Phase 2.