The Architecture That Wasn't There: Building a Custom sm_120 Verify Kernel for Blackwell Consumer GPUs

Introduction

In the high-stakes world of large language model inference, the difference between a system that works and one that performs well often comes down to a single number: the GPU compute capability version. When that version number doesn't match the compiled kernels, the result is not a graceful degradation but a catastrophic collapse in throughput—from gigabytes-per-second memory bandwidth down to tens of megabytes. This article examines a pivotal message in an opencode coding session where an AI assistant, after days of investigation, presents a comprehensive plan to solve exactly such a collapse on NVIDIA's RTX PRO 6000 "Blackwell" consumer GPUs.

The message at index 12212 is a turning point in the conversation. It synthesizes weeks of debugging, binary inspection, and architectural analysis into a coherent strategy for replacing the bottlenecked Triton-based attention kernel with a custom, owned implementation tuned for the sm_120 architecture. The message is notable not just for its technical depth but for its disciplined reasoning process, its honest assessment of trade-offs, and its effective use of user decision points to align on direction before committing to implementation.

This article explores why this message was written, the investigative trail that led to it, the assumptions and decisions embedded within it, and the thinking process that shaped its structure. It is a case study in how to navigate the treacherous waters of GPU kernel compatibility when deploying cutting-edge models on non-datacenter hardware.

The Context: A Structural Bottleneck Revealed

To understand the significance of message 12212, we must first understand the problem it addresses. The conversation centers on deploying the Kimi K2.6 model with DFlash speculative decoding on a machine equipped with eight RTX PRO 6000 Blackwell GPUs. These GPUs use the sm_120 compute architecture, which is NVIDIA's consumer/professional variant of the Blackwell generation. Unlike the datacenter Blackwell GPUs (sm_100 and sm_103) that ship with 228 KB of shared memory per SM and support advanced tensor core instructions like tcgen05 and wgmma, the sm_120 architecture is closer to Ada Lovelace: it has only ~100 KB of shared memory and relies on mma.sync and cp.async instructions.

This architectural difference has profound consequences for software compatibility. The DDTree speculative decoding algorithm requires a custom attention mask during its verify phase—a visibility mask that determines which tree positions each token can attend to. SGLang, the inference engine being used, supports this custom mask through its Triton-based attention backend. Triton, being a JIT compiler, generates code for whatever GPU it runs on, so it works on sm_120. But it works slowly—so slowly that decode throughput at 128K context length was measured at 0.7 tokens per second, with GPU tensor core utilization hovering around 3%.

The root cause, as the assistant had painstakingly established in preceding messages, was that DDTree forces --page-size 1 and a custom tree mask, which locks the verify attention to the Triton MLA kernel. Every optimized MLA kernel—FlashMLA, cutlass-MLA, flashinfer-MLA—ships compiled only for sm_90a (Hopper), sm_100a (Blackwell datacenter), and sm_103a (Blackwell datacenter variant). None include sm_120a code. Worse, these libraries contain no embedded PTX that the driver could JIT-compile to sm_120 at load time. The optimized kernels use Hopper's wgmma and TMA instructions or Blackwell-DC's tcgen05 tensor core instructions, none of which exist on sm_120's Ada-like ISA. Recompilation is not a flag flip; it is a porting effort.

The assistant's own naive verify_attn.cu kernel, written as a correctness oracle, was not a solution either: it used a three-pass design that stored full score rows in shared memory (impossible for long prefixes given the 100 KB limit), and it re-read the MLA latent representation 64 times—once per head—instead of reading it once and amortizing across heads. The effective memory bandwidth was measured at ~14 GB/s against a hardware peak of 1.8 TB/s, a 130× gap.

This was the situation entering message 12212: a known structural bottleneck, confirmed at the binary level, with no off-the-shelf fix available. The assistant needed a plan, and it needed the user's input on strategic decisions before proceeding.

The Investigation: Binary-Level Proof

The assistant's reasoning section in message 12212 reveals a careful forensic investigation. The key insight came from examining the SGLang kernel package layout:

sgl_kernel/
  flashmla_ops.abi3.so     # sm_90a only
  flash_ops.abi3.so        # sm_80, sm_86, sm_90a
  spatial_ops.abi3.so      # has sm_120a
  sm90/common_ops.abi3.so  # Hopper-specific
  sm100/common_ops.abi3.so # Blackwell-DC-specific

There was no sm120/ directory. The architecture-specific common_ops libraries only shipped for sm_90 and sm_100. The flash_ops and flashmla_ops libraries, which contain the optimized MLA decode kernels, were compiled exclusively for sm_90a with no sm_120a support and no embedded PTX for forward compatibility.

This raised a puzzle: how had the Marlin MoE kernel benchmarked successfully on sm_120 in earlier sessions? The assistant reasoned through this carefully. The Marlin kernel uses mma.sync instructions, which exist on sm_120. If the loader selected the sm_100 variant of common_ops at runtime and the driver found PTX embedded in those binaries, it could JIT-compile the PTX to sm_120. But the assistant had checked flash_ops and flashmla_ops and found no PTX. The real answer, which the assistant arrived at, was that the service runs on Triton for attention because Triton JIT-compiles to sm_120 at runtime—not because the precompiled sgl_kernel ops work on that architecture. The Marlin kernel, being in a different binary with different compilation flags, might indeed have PTX support that the flash libraries lack.

This investigative thread illustrates a crucial methodological principle: when a system partially works (Marlin runs) but partially doesn't (MLA decode is slow), the explanation is rarely a blanket "sm_120 is unsupported." Instead, it is a fine-grained question of which specific instructions each kernel uses, which binaries contain which architecture variants, and whether PTX fallback paths exist. The assistant's disciplined use of cuobjdump, nm, and find to inspect binaries directly—rather than relying on documentation or assumptions—is a model of evidence-based debugging.

The Plan: Structure and Strategy

The centerpiece of message 12212 is a detailed plan document titled "Plan: sm_120 optimized DDTree verify attention + K/V defrag." This plan is notable for its clarity, its honest framing of trade-offs, and its separation of findings from decisions.

Where We Stand

The plan opens with a crisp summary of findings, each backed by specific evidence:

  1. Decode is latency-bound: 3% tensor core active, 133W power draw, memory controller at 1.2%, effective bandwidth ~14 GB/s vs 1.8 TB/s peak. These are not symptoms of a busy GPU; they are symptoms of a stalled one.
  2. No optimized kernel works on sm_120: Every MLA kernel (FlashMLA, cutlass-MLA, flashinfer-MLA) ships only for sm_90a/sm_100a/sm_103a. The ISA gap is fundamental: Hopper's wgmma/TMA and Blackwell-DC's tcgen05 do not exist on sm_120's mma.sync/cp.async instruction set. Reuse requires a port, not a recompile.
  3. The naive kernel is a correctness oracle only: It cannot handle long prefixes (score rows don't fit in 100 KB shared memory), it re-reads the MLA latent 64 times, and it uses no flash/streaming or tensor core instructions.
  4. Reusable primitives exist: merge_state_v2 (a simple LSE merge kernel) and SGLang's AttentionBackend abstract base class plus the DFlash verify plumbing (which passes custom_mask, kv_indices, and qo_indptr) are all usable as-is.
  5. K/V fragmentation is a secondary concern: The allocator returns free_pages[:n], which is contiguous when the pool is cold but scatters after churn. No defrag exists today, but even the near-contiguous benchmark still hit only 14 GB/s, confirming the kernel is the primary bottleneck. This framing is strategically important. By establishing that the kernel is the primary culprit and defrag is secondary, the assistant sets expectations about where the effort should go and what the likely returns will be. It also honestly acknowledges that even a perfectly optimized kernel won't eliminate all overhead—the MoE expert imbalance at batch size 1 is a structural ceiling beyond the verify kernel's scope.

The Core Design: Split Verify

The plan's technical centerpiece is the "split verify" design for the attention kernel. Per transformer layer, the verify phase processes a block consisting of a long prefix (L tokens, potentially 100K+) followed by a small tree of candidate tokens (q ≤ 9). The split verify divides this into three stages:

  1. Dense prefix: A flash-streaming MLA attention over the L prefix keys. The key innovation is reading the MLA latent representation once and amortizing it across all 64 heads, rather than re-reading it per head as the naive kernel does. It uses mma.sync for matrix multiplication and cp.async for double-buffered data movement, with tiles sized to fit within the 100 KB shared memory constraint. No mask is needed here because the prefix is fully visible.
  2. Masked tail: A tiny q×q (≤9×9) attention over the tree keys using the custom visibility bitmask. This is small enough that a naive implementation is fine—the computational cost is negligible compared to the prefix.
  3. Merge: Using merge_state_v2 to combine the output and log-sum-exp (LSE) from the prefix and tail into a single result. This design kills the two main pathologies of the naive kernel: the 64× redundant latent read and the latency-bound scatter access pattern. The assistant estimates this should collapse decode step time from ~700 ms at 128K context to something in the range of 5–45 ms, approaching the memory bandwidth floor.

The Decision Framework: R1 vs R2

One of the most impressive aspects of the plan is its honest framing of the kernel sourcing decision. The assistant presents two options:

The User's Answers: Setting Direction

The plan concludes with four questions to the user, each with recommended options. The user's answers, appended at the end of the message, provide clear direction:

  1. Integration target: "Custom SGLang backend (Recommended)" — the fastest path to a real serving win, reusing the existing K2.6/TP8/MoE infrastructure.
  2. Kernel sourcing: "Go straight to R2 (own it)" — the user chose to skip the R1 spike and go directly to building the owned kernel. This is a vote of confidence in the assistant's assessment that the port would likely fail, and a willingness to invest in the more robust solution.
  3. Defrag scope: "Tier 0 + Tier 1 relocation" — the user wants both the cheap free-list sorting and the more ambitious live relocation pass.
  4. Dev box: "Develop directly on CT200" — the user chose the production machine over the local 5070 Ti, despite the 1T model being harder to iterate on. These answers transform the plan from a proposal into an execution roadmap. The assistant now knows exactly which path to take, and the subsequent messages in the conversation (covered in chunk 1 of the segment summary) show the systematic execution of this plan: making the kernel CUDA-graph capture-safe, optimizing occupancy with NSPLIT tuning and vectorized loads, implementing Tier 0 defrag, and ultimately achieving a 3–6× decode speedup over the Triton baseline.

Assumptions, Mistakes, and Honest Assessments

No plan is perfect, and message 12212 contains several implicit assumptions worth examining.

The assumption that Triton is the only viable fallback: The assistant states that "the service runs on triton purely because triton is the only thing that JIT-compiles to sm_120 and accepts the tree mask." This is true given the current software stack, but it assumes that no other JIT-based attention system (e.g., ThunderKittens, which the assistant had explored in earlier segments) could be adapted. The assumption is reasonable given the integration effort required, but it does narrow the solution space.

The assumption about MoE imbalance being a separate concern: The plan acknowledges that MoE expert imbalance at batch size 1 is a structural ceiling, but treats it as outside the verify kernel's scope. This is a correct scoping decision—you cannot fix every problem at once—but it means the 3–6× speedup from the kernel optimization will still leave throughput below what a batched or expert-parallel deployment could achieve.

The honest assessment of defrag's value: The assistant is refreshingly candid about K/V defragmentation. Tier 0 (free-list sorting) is cheap and low-risk. Tier 1 (live relocation) is more complex, and the assistant notes that "a good paged kernel already tolerates scatter via kv_indices, so this mainly raises sustained bandwidth toward peak and caps tail latency on busy multi-tenant pools—not a fresh-pool win." This is the kind of honest framing that builds trust: the assistant is not overselling defrag as a silver bullet but positioning it as a complementary optimization with specific, limited benefits.

The risk assessment: The plan identifies five key risks: CUDA graph capture safety, the 100 KB shared memory tile pressure on sm_120, exact MLA-absorb layout matching with SGLang's KV pool, per-rank tensor parallelism correctness, and keeping the naive kernel as a test oracle. All of these proved to be real concerns in the subsequent execution, and the assistant addressed them systematically.

The Thinking Process: Evidence-Based Architecture

What makes message 12212 particularly interesting as a subject of study is the thinking process visible in its reasoning section. The assistant works through several puzzles in real time:

The Marlin MoE puzzle: How does Marlin run on sm_120 if common_ops only ships sm90 and sm100 builds? The assistant reasons that the loader must select one variant at runtime and JIT-compile PTX, but it hasn't checked the PTX in common_ops itself. This is a moment of intellectual honesty—acknowledging a gap in the evidence rather than glossing over it.

The shared memory constraint: The user had warned that sm_120 has only 100 KB of shared memory versus 228 KB on datacenter Blackwell. The assistant internalizes this constraint throughout the plan, noting that tiles must be sized for the smaller pool and that the naive kernel's full-score-row approach is impossible at long contexts.

The ISA gap: The assistant clearly understands the architectural difference between sm_120 and sm_90/sm_100 at the instruction level. It knows that wgmma and tcgen05 don't exist on sm_120, and that mma.sync and cp.async are the available primitives. This knowledge shapes the entire kernel design.

The validation strategy: The assistant initially proposes developing on the local 5070 Ti (also sm_120) because "correctness+perf are shape-driven and don't need the 1T model." When the user chooses CT200 instead, the assistant adapts without complaint—the validation principles remain the same even if the hardware is different.

Input and Output Knowledge

To understand message 12212, a reader needs knowledge of:

Conclusion

Message 12212 is a masterclass in technical communication and evidence-based planning. It takes a complex, multi-dimensional problem—slow speculative decoding on unsupported GPU hardware—and decomposes it into clear findings, a coherent strategy, a specific design, and actionable decisions. The assistant's reasoning process is transparent, its assumptions are explicit, and its risk assessments are honest.

The message succeeds because it respects the user's intelligence and agency. It does not present a single solution as the only option; it lays out the evidence, explains the trade-offs, makes a recommendation, and then asks for direction. The user's answers then transform the plan from a proposal into an execution mandate, and the subsequent work (documented in chunk 1 of the segment) shows the plan being executed systematically: CUDA graph capture achieved, kernel optimized for 3–6× speedup, Tier 0 defrag implemented, and the bottleneck cleanly identified as MoE imbalance.

For anyone working on GPU kernel development, inference engine optimization, or deploying LLMs on non-datacenter hardware, this message offers a template for how to approach a seemingly intractable performance problem: gather binary-level evidence, understand the architectural constraints, design a targeted solution, present options honestly, and align with stakeholders before executing. The result is not just a faster kernel but a shared understanding of what the problem is, why it exists, and how to fix it.