The Architecture Detective: Tracing CUDA Kernel Compatibility Across GPU Generations

Introduction

In the high-stakes world of large language model inference, every microsecond counts. When deploying a model like Kimi K2.6 with speculative decoding on novel hardware, the difference between a working system and a broken one often comes down to a single detail: which GPU architecture a kernel was compiled for. This article examines a pivotal moment in an opencode coding session—message 12211—where an AI assistant grapples with the discovery that the entire optimized attention kernel stack is incompatible with the target hardware, and must chart a path forward through a landscape of architectural constraints, JIT compilation quirks, and shared memory limits.

The message captures a critical inflection point in a months-long effort to deploy the Kimi K2.6 model with DFlash speculative decoding (DDTree) on NVIDIA's RTX PRO 6000 Blackwell GPU. This consumer-grade Blackwell variant uses the sm_120 architecture, which sits in an awkward middle ground: it's newer than Hopper (sm_90) but lacks the specialized tensor core instructions found in datacenter Blackwell (sm_100/sm_103). The assistant has just discovered that FlashMLA, cutlass-MLA, and essentially all optimized MLA (Multi-head Latent Attention) kernels in the SGLang stack are compiled exclusively for architectures the RTX PRO 6000 doesn't support. This message represents the moment of reckoning—and the beginning of a plan to build a custom solution.

The Message in Context

[msg 12211] arrives after a long chain of investigation. In preceding messages, the assistant used cuobjdump to inspect the compiled GPU architectures embedded in SGLang's kernel shared objects. The results were damning: flashmla_ops.abi3.so contained code only for sm_90a, sm_100a, and sm_103a—none for sm_120a. The larger flash_ops.abi3.so was similarly limited to sm_80, sm_86, and sm_90a. The RTX PRO 6000, with its sm_120 architecture, had no compiled kernel path available.

The user had given a greenlight to compile for sm_120 if needed, but with an important caveat: sm_120 has only 100KB of shared memory per SM compared to 228KB on datacenter Blackwell. This constraint would affect any porting effort.

By message 12211, the assistant has already examined the SGLang kernel source tree, finding that FlashMLA and the cutlass_sm100_mla kernels use Hopper-specific instructions (wgmma, TMA) or Blackwell-DC-specific instructions (tcgen05) that simply don't exist on sm_120's Ada-like ISA. A simple recompile won't work—these kernels would need architectural rewrites.

What the Message Actually Says

The message opens with the assistant reasoning about the architecture of a custom attention backend:

So I need to implement a custom attention backend that mirrors what the Triton backend does—building those intermediate tensors like custom_mask, mask_indptr, and the various index pointers, then dispatching through the attention kernel. The key constraint is that it has to support CUDA graph capture, which means the backend needs to handle both the metadata initialization and the graph recording/replay lifecycle.

This immediately establishes the scope: the assistant is not just fixing a kernel flag; it's contemplating building an entire attention backend from scratch, integrated into SGLang's existing architecture, with full CUDA graph support. The reasoning reveals a sophisticated understanding of what SGLang's verify phase requires: a dense MLA flash attention over prefix KV slots, a smaller masked kernel for the tree tail, and a merge step to combine results.

The assistant then lays out the strategic decision:

I need to decide between porting an existing optimized kernel or building our own native sm_120 version, weighing the risk that Hopper-specific instructions won't compile against the reward of reusing a proven implementation.

This is the core tension of the message. The assistant has two paths: attempt to port FlashMLA to sm_120 (risky, given the ISA differences), or build a custom sm_120-tuned kernel from scratch (more work, but full ownership). It leans toward a "quick 1-2 day spike on porting first, then defaulting to our own kernel."

The message also outlines plans for K/V defragmentation, microbenchmarking on a 5070 Ti (which shares the sm_120 architecture), and validation against the Triton baseline. It identifies key constraints: CUDA graph capture safety, the 100KB shared memory limit, MLA layout compatibility, and per-rank TP8 handling.

The Critical Discovery

The most important moment in the message comes when the assistant realizes it has been missing a piece of the puzzle:

I need to track down where the optimized kernels actually live in the .so files — merge_state_v2 and cutlass_mla are torch.ops.sgl, but the .so layout is unclear since flash_ops only has sm_80/86/90a, yet the user ran Marlin MoE successfully on sm_120, so there must be a common_ops or similar with sm_120a support that I haven't fully located yet.

This is a crucial insight. The user had successfully run Marlin MoE (a weight quantization kernel) on sm_120, which seemed to contradict the finding that no optimized kernels were available. The assistant realizes the resolution: flash_ops is a 1GB binary that may embed PTX (parallel thread execution) code alongside SASS (Shader Assembly). PTX is a portable intermediate representation that the NVIDIA driver JIT-compiles to the target architecture at load time. Kernels using portable instructions like mma.sync (used by Marlin) can forward-JIT to sm_120, while kernels using architecture-specific features like wgmma or tcgen05 (used by FlashMLA) cannot—those instructions don't exist on sm_120 at all.

This explains the asymmetry: Marlin works because it uses portable sm_80 PTX that JITs forward, while FlashMLA is locked to sm_90a/sm_100a because it uses non-portable instructions. The assistant then runs a command to verify this theory by checking for PTX embeddings and, crucially, discovers the existence of common_ops.abi3.so files in sm90/ and sm100/ subdirectories—shared objects the assistant hadn't previously found because the initial ls only listed top-level .so files.

The Thinking Process

The assistant's reasoning in this message is a masterclass in systematic debugging. It moves through several layers of analysis:

  1. Architecture design: First, it designs the high-level structure of the custom backend—three-part computation (prefix flash, masked tail, merge), CUDA graph lifecycle, tensor parallelism.
  2. Strategic planning: It weighs two approaches (port vs. build) and establishes a decision framework with a spike-and-pivot timeline.
  3. Defragmentation planning: It considers K/V defragmentation as a secondary optimization, distinguishing between "free-list sorting" (low-risk) and "live relocation" (higher risk).
  4. Validation strategy: It plans a multi-stage validation pipeline—microbenchmarks on 5070 Ti, parity testing on CT200, A/B defrag testing.
  5. Anomaly resolution: It spots the contradiction (Marlin works on sm_120 but flash_ops only shows sm_80/86/90a) and resolves it through the PTX JIT theory.
  6. Evidence gathering: It executes a bash command to verify the theory by checking for PTX embeddings and discovering the hidden common_ops libraries. This layered thinking—moving from architecture to strategy to anomaly detection to evidence collection—shows a disciplined approach to problem-solving. The assistant doesn't just accept the surface-level finding; it actively looks for contradictions and uses them to deepen its understanding.

Assumptions and Their Validity

Several assumptions underpin the assistant's reasoning in this message:

Assumption 1: The Triton backend is the correct reference implementation. The assistant assumes that mirroring the Triton backend's interface is the right approach for a custom backend. This is reasonable—Triton is the fallback path that works on sm_120, and matching its interface ensures compatibility with SGLang's verify phase.

Assumption 2: CUDA graph capture is feasible with a custom backend. The assistant assumes it can make the custom backend CUDA graph-capture-safe by using static buffers and avoiding host synchronization. This is a significant engineering challenge, and the assumption is tested later in the session (and ultimately validated).

Assumption 3: The 100KB shared memory limit is the primary architectural constraint. The assistant correctly identifies this as a key difference between sm_120 and datacenter Blackwell, and factors it into the kernel design.

Assumption 4: The 5070 Ti is a valid development proxy for the RTX PRO 6000. Both use sm_120, so kernel behavior should be identical. The 16GB memory limit is a practical constraint but doesn't affect kernel correctness.

Assumption 5: PTX JIT explains Marlin's compatibility. This is the most important assumption in the message. The assistant theorizes that Marlin works on sm_120 because its sm_80 PTX forward-JITs, while FlashMLA's sm_90a SASS cannot. This is architecturally sound—NVIDIA's driver can JIT PTX to any forward-compatible architecture, but architecture-specific SASS (the "a" suffix) is locked to its target. The assistant runs a command to verify this by checking for PTX embeddings.

Potential Mistakes and Incorrect Assumptions

While the assistant's reasoning is generally sound, several points deserve scrutiny:

The "quick 1-2 day spike" on porting may be optimistic. Porting a kernel as complex as FlashMLA to a different ISA with half the shared memory is not a trivial task. The assistant acknowledges this implicitly by noting that "a simple recompile won't work" and that "these kernels would need actual rewrites." The spike is framed as a low-risk exploration, but the complexity could easily exceed the estimate.

The split approach (prefix flash + masked tail + merge) adds integration complexity. While architecturally clean, this three-part decomposition requires careful handling of intermediate states (log-sum-exp values for online softmax merging) and introduces multiple kernel launch boundaries. Each boundary is a potential source of overhead and correctness issues.

The assistant may be underestimating the CUDA graph capture challenge. CUDA graphs require all operations to be known at capture time, with no host-device synchronization, no dynamic allocations, and no conditional execution paths. Building a custom backend that satisfies these constraints while handling variable-length sequences and custom masks is a significant engineering effort.

The PTX JIT theory, while correct in principle, may not fully explain Marlin's compatibility. Marlin might also have explicit sm_120 SASS compiled in that wasn't visible in the cuobjdump output due to how the tool reports architectures. The assistant's command to check PTX embeddings is designed to resolve this ambiguity.

The assistant assumes that merge_state_v2 already runs on sm_120 via PTX JIT. This is plausible but unverified. If merge_state uses sm_90a-specific features, it too would fail on sm_120, requiring a custom elementwise merge kernel.

Input Knowledge Required

To fully understand this message, a reader needs:

  1. GPU architecture knowledge: Understanding of NVIDIA's compute capability versions (sm_80 = Ampere, sm_86 = Ampere consumer, sm_90 = Hopper, sm_100/sm_103 = Blackwell datacenter, sm_120 = Blackwell consumer), the difference between SASS and PTX, and the concept of architecture-specific ("a" suffix) vs. portable code.
  2. CUDA graph concepts: Understanding of CUDA graph capture and replay, and why static buffers and no host synchronization are required.
  3. MLA (Multi-head Latent Attention): Knowledge of how DeepSeek's MLA attention mechanism works, including the latent representation, KV cache layout, and the absorb-form computation.
  4. SGLang architecture: Familiarity with SGLang's attention backend interface, the verify phase in speculative decoding, and how custom masks flow through the Triton backend.
  5. DDTree speculative decoding: Understanding of how DFlash with DDTree uses a drafter model and a verify pass to accelerate autoregressive decoding, and the role of the custom mask in the verify attention.
  6. The session's history: Knowledge that the RTX PRO 6000 is a sm_120 device, that FlashMLA and cutlass-MLA were found to lack sm_120 support, and that the user approved a compile-for-sm_120 approach with the shared memory caveat.

Output Knowledge Created

This message creates several forms of knowledge:

  1. A concrete architecture plan: The split approach (prefix flash + masked tail + merge) becomes the blueprint for the custom sm_120 verify kernel implemented later in the session.
  2. A decision framework: The port-vs-build tradeoff analysis, with the spike-and-pivot timeline, structures the subsequent engineering effort.
  3. A validation methodology: The multi-stage testing plan (microbenchmark on 5070 Ti, parity on CT200, A/B defrag) becomes the testing strategy.
  4. A key architectural insight: The PTX JIT theory explains why Marlin works on sm_120 while FlashMLA doesn't, and this understanding shapes the entire approach to kernel compatibility.
  5. Discovery of hidden libraries: The common_ops.abi3.so files in sm90/ and sm100/ subdirectories are found, revealing that SGLang has architecture-specific subdirectories the assistant hadn't previously known about.
  6. A roadmap for defragmentation: The tiered approach (free-list sorting first, live relocation later) defines the K/V defrag strategy.

The Bash Command and Its Significance

The message ends with a bash command that the assistant executes to resolve the PTX JIT question:

SK=/root/venv_sglang211/lib/python3.12/site-packages/sgl_kernel
CUO=/usr/local/cuda-13.0/bin/cuobjdump
timeout 60 ssh -o StrictHostKeyChecking=no root@10.1.230.171 "
echo '### all sgl_kernel .so (recursive) ###'
find $SK -name '*.so' 2>/dev/null
echo '### PTX arches embedded (JIT-forward-compat) ###'
for f in flash_ops flashmla_ops; do echo -n \"\$f PTX: \"; $CUO --list-ptx $SK/\$f.abi3.so 2>/dev/null | grep -oE 'sm_[0-9]+a?' | sort -u | tr '\n' ' '; echo; done
echo '### does flash_ops contain merge_state / cutlass_mla / marlin symbols? ###'
/usr/bin/nm -D $SK/flash_ops.abi3.so 2>/dev/null | grep -iE 'merge_state|cutlass_mla|marlin|fwd_kvcache' | head
" 2>&1

This command is significant for several reasons. First, it uses find recursively rather than ls on the top-level directory, which is what reveals the sm90/common_ops.abi3.so and sm100/common_ops.abi3.so files—libraries the assistant had missed in earlier investigations. Second, it uses cuobjdump --list-ptx to check for PTX embeddings, directly testing the JIT-forward-compatibility theory. Third, it uses nm -D to check which symbols are actually in the massive flash_ops.abi3.so binary, helping determine whether merge_state and cutlass_mla live there or elsewhere.

The output shown in the message confirms the recursive find discovered the hidden common_ops libraries, but the PTX and symbol results are cut off. This incomplete output creates a natural tension—the reader (and the assistant) must wait for the next message to see whether the PTX JIT theory is confirmed.

The Broader Significance

This message represents a pivotal moment in the coding session for several reasons:

It marks the transition from diagnosis to action. The preceding messages were about understanding the problem—inspecting binaries, reading source code, checking architecture support. This message is the first where the assistant begins designing the solution.

It demonstrates the importance of architectural knowledge in ML engineering. The difference between sm_90a and sm_120 is not just a version number; it represents a fundamentally different instruction set architecture. Understanding this distinction—and its implications for kernel portability—is what separates a working solution from a wasted effort.

It shows the value of systematic anomaly resolution. The assistant could have accepted the surface-level finding (no sm_120 support) and moved on. Instead, it noticed the contradiction with Marlin's compatibility and dug deeper, leading to the PTX JIT insight and the discovery of the common_ops libraries.

It reveals the complexity of modern ML inference stacks. A single inference request passes through multiple kernel boundaries, each compiled for specific GPU architectures, each with its own compatibility constraints. Understanding this stack requires knowledge of CUDA compilation, GPU architecture, JIT compilation, and the specific kernel implementations.

Conclusion

Message 12211 captures a moment of synthesis in a complex engineering effort. The assistant has gathered evidence, identified contradictions, formed theories, and is now designing a multi-phase plan to build a custom sm_120 attention kernel. The message is notable for its systematic reasoning, its awareness of architectural constraints, and its willingness to challenge surface-level findings.

The PTX JIT theory—that portable sm_80 code can forward-JIT to sm_120 while architecture-specific sm_90a/sm_100a code cannot—is the key insight that resolves the apparent contradiction between Marlin's compatibility and FlashMLA's incompatibility. This understanding shapes everything that follows: the decision to build a custom kernel rather than port an existing one, the focus on sm_120's 100KB shared memory constraint, and the validation strategy using the 5070 Ti as a development proxy.

In the broader narrative of the coding session, this message is the turning point. The assistant has moved from "what's wrong" to "what we're going to build." The plan outlined here—custom attention backend, split prefix/tail/merge computation, CUDA graph support, K/V defragmentation—becomes the roadmap for the subsequent chunks, where the assistant successfully implements a custom sm_120 verify kernel that achieves 3–6× speedup over the Triton baseline.

The message also serves as a case study in disciplined engineering reasoning: start with the architecture, identify constraints, look for contradictions, form hypotheses, gather evidence, and only then commit to a plan. It's a pattern that applies far beyond GPU kernel development, and one that this session demonstrates with remarkable clarity.