The Architecture of a Pivot: Building an Owned sm_120 Verify Kernel for SGLang

Introduction

In the lifecycle of any high-performance computing project, there comes a moment when investigation must yield to execution. All the diagnostics have been run, the root causes identified, the architectural constraints mapped, and the strategic decisions made. What remains is the plan—the bridge between what has been learned and what will be built. Message 12213 in this opencode session is precisely such a moment. It is the assistant's refined, execution-ready plan for replacing the bottlenecked Triton-based verify attention in a live SGLang inference service with an owned, custom-built CUDA kernel targeting the NVIDIA sm_120 architecture (RTX PRO 6000 Blackwell consumer GPUs), alongside a two-tier KV cache defragmentation strategy.

This message is remarkable not because it introduces novel ideas—the kernel design, the backend integration approach, and the defragmentation tiers had all been discussed in preceding messages—but because it crystallizes a sprawling technical investigation into a concrete, phased, gated execution plan. It reflects the assimilation of dozens of prior tool calls, binary analyses, benchmark runs, and user decisions into a coherent engineering roadmap. The message is a study in how an AI assistant reasons about resource constraints, architectural trade-offs, risk management, and validation strategy when building production GPU kernel software.

To understand this message fully, one must appreciate the journey that led to it. The preceding messages (12211 and 12212) documented a deep forensic investigation into why the SGLang speculative decoding service was achieving only ~0.7 tokens per second at 185k context length on 8× RTX PRO 6000 Blackwell GPUs. The assistant had traced the bottleneck to the DDTree verify attention being locked to the Triton MLA kernel with page_size=1, which caused scattered KV cache accesses achieving only ~14 GB/s effective bandwidth—130× below the 1.8 TB/s peak. The assistant then discovered, through cuobjdump and nm analysis of the sgl_kernel package, that every optimized MLA kernel (FlashMLA, cutlass-MLA, flashinfer-MLA) ships only for sm_90a/sm_100a/sm_103a architectures, with no sm_120a support and no embedded PTX for JIT forward-compatibility. The RTX PRO 6000's sm_120 architecture is Ada-like, lacking the Hopper wgmma/TMA and Blackwell-DC tcgen05 instructions that those optimized kernels depend on. Reuse was not a flag-flip; it was a port.

The user had then answered four critical strategic questions: (1) the kernel should land in a custom SGLang attention backend (not a standalone engine), (2) the dense prefix MLA kernel should be built from scratch (R2) rather than attempting to port FlashMLA (R1), (3) defragmentation should include both Tier 0 (free-list sorting) and Tier 1 (live relocation), and (4) development should happen directly on the CT200 production server rather than on a local 5070 Ti dev box.

Message 12213 is the assistant's response to those decisions. It is the plan.

The Message in Full Context

Before diving into analysis, let us situate the message within the broader session. The session (Segment 66) began with the deployment of 200k context-length on the CT200 server, followed by a comprehensive long-context benchmark that revealed the catastrophic decode slowdown. The user correctly suspected drafter issues or low GPU utilization, but the assistant's investigation confirmed the drafter's sliding-window attention and KV caching were working correctly—the root cause was the verify attention kernel itself. The assistant then discovered the absence of sm_120 support in all optimized MLA kernels, received user approval to build an owned kernel, and implemented a first pass (verify_attn_flash.cu) that passed correctness tests but showed only modest speedups (1.7–2.0× at long prefixes, slower than naive at short prefixes).

The chunk summary for Chunk 1 (which contains message 12213) tells us what happens after this plan is executed: the assistant goes on to make the kernel CUDA-graph capture-safe, optimize it with increased NSPLIT and 128-bit vectorized loads to achieve 3–6× end-to-end decode speedup over Triton, implement Tier 0 defrag, and identify MoE imbalance as the remaining bottleneck. Message 12213 is the fulcrum—the planning message that precedes all that execution.

The Reasoning Process: From Decisions to Architecture

The assistant's reasoning in message 12213 reveals a sophisticated multi-constraint optimization problem being solved in real time. The reasoning section opens by acknowledging the settled strategic direction: "the team has settled on a clear technical direction: we're injecting our own kernels directly into a live SGLang backend rather than trying to reuse FlashMLA, we're building the sm_120 flash-MLA kernel from scratch and pushing it to R2, and we're tackling both lightweight and full KV relocation defragmentation in this phase."

But immediately, the assistant flags a resource constraint that will shape the entire plan: "CT200 is the production box running the live SGLang service across all 8 GPUs with nearly full memory utilization (~94GB per GPU), leaving only ~4GB headroom." This is a critical piece of context. The development environment is also the production environment. Kernel unit tests must compete with a live inference service for GPU memory. The assistant correctly identifies that small synthetic tests can fit in the ~4GB margin, but full-scale parity testing will require a maintenance window. This constraint is not merely noted—it propagates through the entire plan, influencing Phase 0's logistics and the validation strategy.

The reasoning then dives into the core kernel design insight: "by tiling the KV latent across a block or cluster, I can load a chunk of the shared latent representation once into shared memory and compute partial scores for multiple heads simultaneously, avoiding the 64× redundant reads that make the naive approach slow." This is the central algorithmic innovation. The naive verify_attn.cu kernel reads the MLA latent representation 64 times—once per head—because it processes heads in the outer loop. The flash kernel restructures this to KV-tile outer, heads inner: load a tile of the KV latent into shared memory once, then compute partial scores for all 64 heads before moving to the next tile. For a 200k-token prefix, this reduces the number of KV latent loads from 64×200k to just 200k (amortized across heads), which is the difference between 14 GB/s and approaching the memory bandwidth ceiling.

The assistant makes a deliberate engineering choice to prioritize memory efficiency over tensor core utilization: "since the decode phase is memory-bound; the strategy is coalesced vectorized loads with double-buffered cp.async streaming of the KV latent, online softmax rescaling to avoid materializing full score rows, and careful tiling (32–64 token chunks) to fit within the 100KB shared memory budget." This is a principled architectural decision grounded in the roofline model. When decode is memory-bound, the optimization target is not FLOP/s but effective bandwidth. The kernel is designed to maximize bytes-per-load by amortizing latent reads across heads, not to maximize tensor core utilization.

The Architectural Evolution: From Split-Merge to Unified Streaming

One of the most interesting aspects of the reasoning is the assistant's real-time evolution of the kernel design. The message begins with the split-verify architecture from the earlier plan (dense prefix + masked tail + merge_state), but then the assistant has a realization: "rather than splitting into two kernels and merging states, I can extend the same streaming loop to handle both regions in one pass—processing the prefix unmasked then the tail masked within the same online softmax accumulation, which eliminates the need for a separate merge step entirely."

This is a significant architectural improvement. The original three-phase design (flash prefix, naive masked tail, merge_state) required two kernel launches and a merge kernel, adding launch latency, extra shared memory pressure, and complexity in the CUDA graph capture path. The unified streaming kernel processes all keys in a single pass: it streams through the prefix tokens with no masking (full visibility), then when it reaches the tail tokens, it applies the DDTree visibility bitmask. The online softmax accumulation handles both regions seamlessly because the mask simply zeros out the non-visible positions in the final tile's score computation.

The assistant then has a second realization about the backend architecture. Initially, it considers implementing a full AttentionBackend that handles prefill (with causal masking) and verify (with the tree mask). But then it recognizes a cleaner path: "use a hybrid backend that delegates prefill and decode to Triton (which already works well and shows healthy utilization) and only overrides the verify forward with the custom kernel. This minimizes risk and surface area since prefill is compute-bound and already performing fine."

This is a textbook example of the principle of minimal change. The prefill phase already achieves 50–90% tensor core utilization—it is compute-bound and healthy. The decode phase (specifically the verify extend) is the only bottleneck. By subclassing TritonAttnBackend and overriding only the verify method, the assistant avoids reimplementing prefill, avoids duplicating the CUDA graph plumbing for the non-verify path, and limits the risk surface to exactly the code path that needs optimization. The fallback to the parent implementation for non-DDTree inputs means the backend is backward-compatible by construction.

The Four Guiding Engineering Choices

The plan section of the message enumerates four guiding engineering choices, each of which deserves close examination:

"One streaming kernel, no merge_state." This is the architectural decision discussed above. The assistant explicitly notes that merge_state_v2 stays as a fallback, acknowledging that the unified approach might encounter edge cases where separate phases are necessary. This is prudent engineering—the fallback exists but the primary path is simpler and faster.

"Subclass TritonAttnBackend, override only the DDTree verify." This is the integration strategy. The key insight is that SGLang's attention backend architecture supports subclassing, and the Triton backend already handles CUDA graph capture correctly. By inheriting that infrastructure, the custom backend gets capture-safety for free on the delegated paths and only needs to ensure the overridden verify method is also capture-safe.

"Kill the 64× latent re-read." This is the performance thesis. The assistant explicitly states: "This (not tensor cores) is the main win — decode is memory-bound." This is a crucial framing. In GPU kernel optimization, there is always a temptation to optimize for tensor core utilization because that is what benchmark suites measure. But the assistant correctly identifies that decode is memory-bound—the bottleneck is getting data from HBM to the compute units, not the compute itself. The optimization target is therefore bandwidth efficiency, specifically the number of bytes loaded per useful computation.

"sm_120 reality: mma.sync + cp.async double-buffer, fp32 accumulate from bf16." This anchors the implementation in the actual instruction set architecture. The assistant is explicitly acknowledging that sm_120 does not have Hopper's wgmma (warp-group matrix multiply-accumulate) or TMA (tensor memory accelerator) instructions. It has mma.sync (matrix multiply-accumulate synchronous) and cp.async (copy asynchronous)—the standard Ada-class instruction set. The tile sizes are constrained to fit within 100KB of shared memory, which limits each tile to approximately 32–64 tokens × 576 dimensions × 2 bytes (bf16).

The Phased Execution Plan

The plan is structured into five phases, each with explicit gates that block the next phase:

Phase 0 — Scaffolding & Parity Harness. This phase establishes the build infrastructure and correctness oracle. The assistant must confirm that scripts/build_nvcc.sh emits compute_120a,code=sm_120a as the gencode target, add bf16 KV and masked-tail test vectors from the model reference, and reproduce the naive verify_attn kernel bit-for-bit as the oracle. The logistics note about the ~4GB free GPU memory is repeated here, with the explicit acknowledgment that full-scale parity testing requires a maintenance window.

Phase 1 — sm_120 Flash MLA Verify Kernel. This is the core implementation phase. The kernel implements online-softmax streaming over KV tiles, with the latent read once across heads, cp.async double-buffering, and mask application only on the final q-tile. The validation matrix is specified precisely: q ∈ {1, 9, 33, 65}, prefix ∈ {256…200k}, H=64, Dl=512/Dr=64. The performance target is stated as "≥10–30× step-time at long prefix (toward the ~5–45 ms/step floor)." This target is grounded in the earlier analysis that the Triton kernel was achieving ~700ms per step at 128k context, and the memory-bandwidth-limited floor is approximately 5–45ms per step.

Phase 2 — SGLang Custom Backend. This phase implements the KDTreeMLAVerifyBackend(TritonAttnBackend) class, overriding the DDTree verify extend to call the kernel via a C-ABI shared library (libkdtree_kernels_c.so). The backend must be capture-safe (static buffers, no host synchronization during captured region), handle per-TP-rank execution, consume the custom_mask, kv_indices, and qo_indptr tensors, and be registered so it is not in the skip-mask denylist. The gate is token-exact greedy decoding parity with Triton on fixed prompts.

Phase 3 — Live Validation on CT200. This phase re-runs the context-scaling benchmark and GPU utilization probe to confirm that decode step-time collapses and tensor/memory utilization rises. It also verifies that budget > 8 now works (the old Triton kernel crashed on budgets above 8 due to mask handling limitations).

Phase 4 — K/V Defrag. This phase is split into two tiers. Tier 0 enables the allocator's need_sort flag so that merge_and_sort_free() returns ascending/contiguous slot indices—a simple configuration change that improves allocation contiguity without any data movement. Tier 1 implements a background relocation pass that runs under the scheduler lock with a per-step copy budget, moving a fragmented live request's MLA latent (61 layers) into a contiguous block and reindexing req_to_token. The assistant explicitly flags the key risk: radix-cache-shared slots, which are referenced by the radix tree and must be skipped or have their node-to-slot mappings updated during defrag.

Assumptions and Their Validity

The plan rests on several assumptions that deserve scrutiny:

Assumption: Prefill is healthy and does not need optimization. The assistant states that "screenshots show prefill is already healthy at 50–90% tensor." This is based on empirical observation from the live service. If this assumption were wrong—if prefill were also bottlenecked—the custom backend would need to be extended to cover prefill as well. However, the assistant has structured the code so that this extension is straightforward: the same kernel logic with different mask variants can handle prefill (causal mask) and verify (tree mask).

Assumption: The unified streaming kernel can handle both prefix and tail in one pass without numerical instability. The online softmax algorithm is numerically stable for arbitrarily long sequences, but the mask application on the tail tile requires care. The mask must be applied to the scores before the softmax reduction, and the softmax must be rescaled correctly. The assistant's design of applying the mask "only on the final q-key tile" assumes that the prefix tiles are unmasked and the tail fits in a single tile. For very large tree structures (q > tile size), the tail might span multiple tiles, requiring the mask to be applied across multiple tiles.

Assumption: CUDA graph capture will work with the subclassed backend. CUDA graph capture is notoriously finicky. It requires that all operations within the captured region use static buffers, no host-device synchronization, and no dynamic memory allocation. The Triton backend already supports graph capture, but the overridden verify method must also be capture-safe. The assistant's plan to use a C-ABI kernel with pre-allocated torch tensors and no cudaMalloc calls during capture is sound, but the interaction between the parent class's graph capture logic and the subclass's override could introduce subtle issues.

Assumption: The ~4GB free GPU memory on CT200 is sufficient for kernel development. This is a tight constraint. Small synthetic tests (tile-level microbenchmarks) should fit, but full-scale parity tests with 200k-token prefixes will not. The assistant correctly identifies that a maintenance window will be needed for those tests. The risk is that the development cycle becomes dependent on scheduling maintenance windows, which could slow iteration.

Assumption: Tier 1 defrag can handle radix-cache-shared slots. This is flagged as a key risk, and rightly so. The radix cache shares prefix KV slots across requests with common prefixes. If defrag relocates a slot that is shared by multiple requests, it must update all references to that slot, including the radix tree's node-to-slot mappings. This is a complex distributed consistency problem. The assistant's plan to "skip or update the radix node→slot mapping" acknowledges the complexity but does not specify the mechanism.

Input Knowledge Required

To fully understand message 12213, the reader needs knowledge spanning several domains:

GPU Architecture: The distinction between sm_90a (Hopper), sm_100a/sm_103a (Blackwell Data Center), and sm_120a (Blackwell consumer/RTX PRO 6000) is central. The reader must understand that sm_120 lacks wgmma and TMA instructions, has only ~100KB of shared memory, and uses mma.sync + cp.async for matrix operations. The concept of CUDA graph capture—recording a sequence of GPU operations into a reusable graph that can be launched with minimal CPU overhead—is essential for understanding the capture-safety requirements.

MLA (Multi-head Latent Attention): The KV cache stores latent representations rather than full key-value pairs. Each token's KV latent is a vector of dimension Dl=512 (or similar), which is then expanded to full head dimensions via learned projections. This means the "KV read" is actually a read of the latent, followed by an expansion step. The 64× re-read problem arises because the naive kernel reads the latent once per head, when it could read it once and reuse across heads.

Speculative Decoding with DDTree: The verify phase in speculative decoding takes a set of candidate tokens (the "tree") produced by a draft model and computes their likelihoods under the target model in parallel. The DDTree algorithm uses a visibility mask to ensure that each candidate token only attends to tokens that precede it in the tree structure. This mask is what makes the verify attention different from standard causal attention.

SGLang's Attention Backend Architecture: SGLang uses an AttentionBackend abstract base class with methods for init_forward_metadata, forward_prefill, forward_decode, and forward_extend. The TritonAttnBackend is one implementation. The assistant's plan to subclass it requires understanding this class hierarchy and the CUDA graph lifecycle (metadata initialization, graph recording, graph replay).

Memory Allocation in Paged KV Cache: SGLang uses a paged KV cache where the KV pool is divided into pages. The alloc() method returns a list of free page indices. After churn (allocations and deallocations), these indices become non-contiguous, leading to scattered KV access. The need_sort flag causes the free list to be sorted, improving contiguity.

Output Knowledge Created

Message 12213 creates several forms of output knowledge:

An Execution Plan: The primary output is the phased, gated execution plan that will guide the next several rounds of implementation. This plan encodes strategic decisions (R2 over R1, subclassing over full backend, unified streaming over split-merge) into concrete phases with deliverables, validation gates, and risk mitigations.

Architectural Decisions: The message records the rationale for each architectural choice. Why one streaming kernel instead of split-merge? Why subclass TritonAttnBackend? Why KV-tile outer / heads inner? These decisions are documented with their engineering justifications, creating a design record that future developers can reference.

Risk Register: The message explicitly enumerates six top risks: CUDA graph capture-safety, 100KB shared memory tile pressure, MLA-absorb/bf16 numeric parity with Triton, TP8 per-rank correctness, Tier 1 defrag versus radix sharing, and GPU contention with the live service. This risk register informs the validation strategy and helps prioritize testing effort.

Validation Gates: The message establishes a chain of validation gates: oracle parity → kernel token-exact + faster → backend token-exact vs Triton → live step-time win → defrag no-regression + bandwidth gain. Each gate blocks the next phase, ensuring that problems are caught early before they compound.

Resource Constraint Documentation: The message documents the CT200 resource constraint (~4GB free per GPU) and its implications for the development and testing strategy. This is important institutional knowledge that explains why certain tests require maintenance windows.

The Thinking Process: A Window into Engineering Reasoning

The reasoning section of message 12213 is unusually rich, showing the assistant working through multiple design iterations in real time. Let us trace the evolution:

1. Initial framing: The assistant acknowledges the settled strategic direction and immediately flags the CT200 resource constraint. This sets the context for all subsequent decisions.
2. Kernel design insight: The assistant articulates the core algorithmic insight—tiling the KV latent and amortizing across heads—and connects it to the specific sm_120 constraints (100KB shared memory, mma.sync + cp.async).
3. Numerical precision requirement: The assistant notes the need for bf16 KV cache and fp32 accumulation to match Triton's numerical behavior exactly. This is a correctness constraint that will shape the kernel implementation.
4. Architectural breakthrough: The assistant realizes that the split-merge design can be replaced with a unified streaming kernel. This is presented as a realization ("rather than splitting into two kernels and merging states, I can extend the same streaming loop"), suggesting the assistant is working through the design space in real time.
5. Backend architecture evolution: The assistant initially considers a full AttentionBackend implementation, then realizes that subclassing TritonAttnBackend is cleaner. This is another real-time design evolution.
6. Defragmentation analysis: The assistant works through the two tiers of defrag, identifying the key risk (radix-cache-shared slots) and the cost model (moving a 100k-token request = ~6.8GB at ~1.8TB/s bandwidth).
7. Validation strategy: The assistant enumerates the validation gates, connecting each to the phase that produces it.
8. Plan synthesis: All of these considerations are synthesized into the final phased plan. What is striking about this reasoning is its discipline. The assistant does not simply present a plan; it shows the work of arriving at the plan. Each design choice is motivated by a constraint or insight. When the assistant changes its mind (e.g., from split-merge to unified streaming), it acknowledges the change and explains why the new approach is better. This is the hallmark of mature engineering reasoning.

Mistakes and Corrective Evolution

While message 12213 is a planning message and does not contain execution errors, we can identify areas where the assistant's earlier reasoning was incomplete and was corrected in this message:

The split-merge assumption: In message 12211, the assistant's plan was built around a three-phase split-merge design (dense prefix + masked tail + merge_state). By message 12213, the assistant has realized that a unified streaming kernel eliminates the merge step. This is a significant simplification that reduces kernel launch count, shared memory pressure, and numerical complexity.

The full backend assumption: In message 12211, the assistant considered implementing a full AttentionBackend. By message 12213, the assistant has realized that subclassing TritonAttnBackend is simpler and lower-risk. This change reduces the implementation scope from ~5 methods to ~1 method.

The R1 spike assumption: The user explicitly chose to skip the R1 reuse spike and go straight to R2 (own the kernel). The assistant's original recommendation was a 1–2 day R1 spike followed by R2. The user's decision to skip the spike reflects confidence in the binary analysis that showed FlashMLA cannot be ported to sm_120 without a complete rewrite.

The Broader Significance

Message 12213 is significant beyond its immediate technical content because it exemplifies how an AI assistant can function as a senior engineering partner. The assistant does not merely execute instructions; it reasons about trade-offs, identifies risks, proposes alternatives, and synthesizes complex information into actionable plans. The message demonstrates:

Constraint-aware planning: The assistant identifies the CT200 resource constraint and propagates it through the entire plan, adjusting testing strategy and logistics accordingly.

Evidence-based decision-making: Every architectural choice is grounded in evidence from prior investigation—the binary analysis (cuobjdump), the benchmark data (14 GB/s vs 1.8 TB/s), the utilization screenshots (3% tensor active during decode, 50–90% during prefill).

Risk-conscious engineering: The assistant explicitly enumerates risks and connects them to mitigation strategies (validation gates, fallback paths, maintenance window requests).

Design iteration: The assistant does not present a static plan but shows the evolution of its thinking, demonstrating that the plan is the product of genuine engineering reasoning rather than template-following.

Conclusion

Message 12213 is the hinge point of a substantial engineering effort. It captures the moment when investigation crystallizes into execution, when analysis yields to building. The plan it presents is not merely a list of tasks but a coherent engineering argument: here is what we know, here is what we will build, here is how we will know it works, and here are the risks we must manage.

The message's richness lies in its reasoning. We see the assistant working through kernel design alternatives in real time, evolving from a split-merge architecture to a unified streaming kernel, from a full backend implementation to a targeted subclass override. We see the assistant grounding every decision in empirical evidence—binary analysis, benchmark data, utilization metrics. We see the assistant identifying and documenting risks, establishing validation gates, and planning for resource constraints.

For anyone interested in how AI systems can engage in sophisticated engineering reasoning, message 12213 is a compelling case study. It shows an AI assistant functioning not as a code generator but as an engineering partner—one that understands GPU architecture, inference serving systems, memory allocation strategies, and the discipline of phased, gated execution planning. The plan it produces is not perfect (no plan is), but it is thoughtful, evidence-based, and structured for iterative refinement as new information emerges from each phase's execution.

The subsequent chunk summary tells us that the plan succeeded: the kernel was made capture-safe, optimized to 3–6× Triton's throughput, deployed with CUDA graphs and Tier 0 defrag, and the bottleneck was cleanly identified as MoE imbalance—a structural ceiling beyond the verify kernel's scope. Message 12213 was the blueprint that made that success possible.