The Execution Phase: From Plan to 3–6× Decode Speedup on Blackwell Consumer GPUs

Introduction

In the lifecycle of any high-performance engineering project, there comes a moment when planning must yield to building. The analysis has been done, the root causes identified, the architectural constraints mapped, and the strategic decisions made. What remains is the hard work of implementation—writing code, running benchmarks, debugging failures, and iterating toward a solution. This article examines the execution phase of a complex optimization effort to deploy the Kimi K2.6 model with DFlash speculative decoding on NVIDIA RTX PRO 6000 Blackwell GPUs (sm_120 architecture). The work documented in this chunk transformed a carefully reasoned plan into a deployed, validated, 3–6× faster inference pipeline—and in doing so, revealed the next bottleneck waiting to be tackled.

The story begins with the user's decisive directive at [msg 12214]: "Write down a plan/spec for each phase, then execute." This seven-word message was the pivot point. Before it, the assistant had been investigating, analyzing, and planning. After it, the assistant shifted into execution mode, systematically reading the existing codebase, writing the formal specification, and then implementing the custom sm_120 verify attention kernel that would become the centerpiece of the optimization.

The Transition: From Planning to Execution

The assistant's response to the user's directive was immediate and methodical. At [msg 12215], the reasoning block explicitly names the cognitive shift: "I'm shifting into execution mode now. I need to set up a structured plan document first, then work through each phase methodically." This metacognitive awareness—consciously switching from "planning mode" to "build mode"—is a hallmark of disciplined engineering practice.

But the assistant did not dive straight into implementation. Instead, it began a systematic survey of the existing codebase, reading each relevant file in dependency order:

  1. The repository structure ([msg 12215]): A ls -R exploration to understand the directory layout, locating kernel source files, test harnesses, build scripts, and Python modules.
  2. The naive kernel implementation (verify_attn.cu, [msg 12216]): The existing correct-but-slow attention kernel that serves as the correctness oracle. Reading this file revealed the exact data structures, tensor layouts, and calling conventions the new kernel must match.
  3. The kernel header (verify_attn.cuh, [msg 12217]): The interface contract documenting the mathematical operation—MLA-absorb attention with a custom tree mask.
  4. The test harness (test_verify_attn.cu, [msg 12218]): The validation framework that compares kernel output against Python-generated golden bundles (.kdtr files).
  5. The build script (build_nvcc.sh, [msg 12219]): The nvcc-based compilation pipeline targeting sm_120 architecture, producing a shared library (libkdtree_kernels_c.so) alongside test and benchmark binaries.
  6. The reference generator (gen_verify_attn_refs.py, [msg 12220]): The script that produces golden test vectors from a NumPy reference implementation.
  7. The NumPy reference (mla_attn_ref.py, [msg 12221]): The mathematical ground truth—the exact computation the CUDA kernel must reproduce.
  8. The benchmark harness (bench_kernels.cu, [msg 12222]): The microbenchmark that measures kernel latency across realistic shape sweeps.
  9. The test runner (run_tests_nvcc.sh, [msg 12223]): The script that orchestrates test execution.
  10. The C API wrapper (capi.cu, [msg 12224]): The bridge between Python (ctypes) and the CUDA kernels. This reading order was not accidental. The assistant was building a mental model of the entire verification pipeline in topological order—from the outermost structure inward to the core mathematical definition. Each file answered specific questions that informed the next. By the time the assistant had read all ten files, it possessed a complete, grounded understanding of the codebase it was about to modify.

The Formal Specification

With the codebase fully understood, the assistant wrote the formal plan document (plans/0002-sm120-verify-kernel-defrag.md). This specification encoded the strategic decisions made in the preceding messages:

The Core Design: A Streaming Flash Attention Kernel

The kernel design itself was the product of deep architectural reasoning. The assistant's analysis at [msg 12213] had identified the central algorithmic insight: the naive verify_attn.cu kernel re-reads the MLA latent representation 64 times—once per attention head—because it processes heads in the outer loop. For a 200k-token prefix, this means 64×200k latent reads from global memory, a catastrophic waste of bandwidth that explains the ~14 GB/s effective throughput.

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. This amortizes the latent read across all heads, reducing the number of global memory loads from 64×L to just L (where L is the prefix length). For a 200k-token prefix, this is the difference between 12.8 million latent reads and 200,000—a 64× reduction in memory traffic.

The design also accounted for sm_120's specific constraints:

Making the Kernel CUDA Graph Capture-Safe

One of the most challenging requirements was CUDA graph capture safety. CUDA graphs allow a sequence of GPU operations to be captured and replayed with minimal CPU overhead, dramatically reducing launch latency for short decode steps. However, graph capture imposes strict constraints: all operations within the captured region must use static buffers, no host-device synchronization is permitted, and no dynamic memory allocation (cudaMalloc) can occur during capture.

The assistant's initial kernel implementation used host-side synchronization and dynamic allocations that violated these constraints. The fix required a fundamental rewrite: instead of synchronizing between kernel launches, the kernel consumed SGLang's native static buffers directly. The NSPLIT parameter (which controls how many KV tiles are processed per thread block) was fixed at compile time rather than being dynamically determined. A torch-allocated workspace provided scratch space without cudaMalloc calls.

The result was a kernel that could be captured into a CUDA graph in approximately 1.5 seconds, and whose replay produced generations matching the Triton baseline token-for-token. This was a critical validation gate: if the graph-captured kernel produced different results from the non-captured version, the optimization would be unusable in production.

The Optimization That Unlocked 3–6× Speedup

With the kernel working correctly, the assistant turned to performance tuning. The profiler had revealed that the verify attention was occupancy-starved in the TP8 regime—with only 8 heads per rank (64 heads ÷ 8 GPUs), there was insufficient parallelism to hide memory latency. Two targeted fixes addressed this:

1. Increasing NSPLIT from 16 to 64. NSPLIT controls how many KV tiles each thread block processes before writing its partial result. A higher NSPLIT means each thread block does more work, improving occupancy by keeping more warps active. However, higher NSPLIT also increases register pressure and shared memory usage. The assistant found that NSPLIT=64 was the sweet spot for sm_120's architecture, providing enough work per block to achieve high occupancy without exceeding the 100 KB shared memory budget.

2. Adding 128-bit vectorized bf16 KV loads. The naive kernel loaded KV data using 32-bit scalar loads, achieving only a fraction of the memory bandwidth. By switching to 128-bit vectorized loads (using uint4 or float4 intrinsics), the kernel could issue four times fewer memory requests for the same data, reducing instruction overhead and improving bandwidth utilization. This was particularly impactful for the MLA latent, which has a contiguous layout in memory that vectorized loads can exploit.

The combined effect was dramatic: the optimized kernel delivered a 3–6× end-to-end decode speedup over the Triton baseline across all context lengths (4k–65k tokens). At 65k context, decode step time dropped from approximately 700ms to under 150ms. The effective bandwidth rose from ~14 GB/s to over 100 GB/s—still below the 1.8 TB/s peak, but a massive improvement that brought the system into a usable regime.

Tier 0 KV Defragmentation

The third major workstream was KV cache defragmentation. The assistant had identified that DDTree's page_size=1 requirement scattered each token's KV cache across physical memory, and that server churn (requests being allocated and freed) exacerbated this fragmentation over time.

Tier 0 defragmentation was implemented by monkeypatching SGLang's TokenToKVPoolAllocator to force need_sort=True. This caused the allocator's merge_and_sort_free() function to return ascending/contiguous slot indices, improving allocation contiguity without any data movement. The change was confirmed active on all 8 TP workers via a runtime probe.

The assistant correctly assessed that Tier 0 was the low-risk, high-impact first step. Tier 1 (live relocation of fragmented KV slots under a copy budget) was deferred with clear reasoning: single-request KV was already contiguous after the Tier 0 fix, and the bottleneck had shifted to MoE expert imbalance—a structural ceiling beyond the verify kernel's scope.

The Remaining Bottleneck: MoE Imbalance

After the verify attention optimization, the assistant ran GPU utilization probes and analyzed the user's screenshots. The results revealed a clear picture:

The Methodology: Evidence-Driven Engineering

Throughout this execution phase, a consistent methodological thread emerges: the assistant's disciplined use of instrumentation to replace assumptions with evidence. Several moments stand out:

The profiler disproved the CPU-orchestration theory. The assistant had initially suspected that CPU-side overhead (tree-building at 1.8ms, mask-building at 0.18ms) was a significant contributor to decode latency. The profiler revealed that these costs were negligible compared to the verify attention kernel itself, which was occupancy-starved in the TP8 regime. This redirected the optimization effort from "reduce CPU overhead" to "improve kernel occupancy."

The benchmark validated the tuning choices. The assistant did not assume that NSPLIT=64 was optimal—it tested multiple values and measured the impact on latency. The 128-bit vectorized loads were not a speculative optimization but a targeted response to the profiler's revelation that memory bandwidth utilization was low.

The GPU utilization screenshots confirmed the bottleneck shift. After the attention fix, the assistant analyzed the user's screenshots to confirm that prefill remained healthy (all GPUs busy) while decode was now MoE-imbalanced. This diagnosis was grounded in empirical observation, not speculation.

This methodology—measure, diagnose, fix, measure again—is the hallmark of effective systems optimization. Each intervention was guided by data, and each result was validated against the same data sources.

Conclusion

The execution phase documented in this chunk transformed a carefully reasoned plan into a deployed, validated optimization. The custom sm_120 verify attention kernel, made CUDA graph capture-safe and tuned for occupancy with NSPLIT=64 and 128-bit vectorized loads, delivered a 3–6× end-to-end decode speedup over the Triton baseline. Tier 0 KV defragmentation was implemented and confirmed active on all 8 TP workers. The bottleneck cleanly shifted from the verify attention to MoE expert imbalance—a structural ceiling beyond the kernel's scope.

The work exemplifies several principles of effective engineering: systematic code comprehension before modification, evidence-driven optimization, disciplined risk management, and clear communication of results. The assistant's methodical approach—reading every relevant file, writing a formal specification, implementing in phases with explicit validation gates, and using instrumentation to guide decisions—provides a template for how to approach complex GPU kernel development in production environments.

The 3–6× speedup is not the end of the story. The MoE imbalance bottleneck remains, and addressing it will require either batching or expert parallelism. But the attention optimization has removed one critical bottleneck, bringing the system closer to its hardware potential. The next phase of work will tackle the MoE problem with the same evidence-driven methodology that made this phase a success.## References

[1] "Diagnosing the Decode Bottleneck: How a Data-Driven Investigation Uncovered a Structural Ceiling in SGLang's DDTree" — Analyzes message 12193, the root cause diagnosis of the 130× bandwidth gap.

[2] "The Two-Question Redirect: How a Brief User Message Reshaped an Engineering Investigation" — Analyzes message 12194, the user's directive to investigate kernel integration and K/V defrag.

[3] "The Architecture Detective: Diagnosing a 130× GPU Bottleneck Through Evidence-Driven Kernel Investigation" — Analyzes message 12195, the launch of parallel investigations into kernel optimization and defrag.

[4] "The Art of the Diagnostic Grep: Investigating SGLang's Verify Attention Backend for DDTree Optimization" — Analyzes message 12196, the first SSH probes into SGLang's attention backend architecture.

[5] "Peeling Back the Layers: Diagnosing SGLang's DDTree Verify Attention Backend" — Analyzes message 12197, reading the resolve_dflash_verify_mask_policy function.

[6] "Mapping the Kernel Frontier: A Systematic Investigation of SGLang's Attention Backends for DDTree Optimization" — Analyzes message 12198, discovering the skip-custom-mask backend set.

[7] "The Art of Systematic Investigation: Mapping the Optimization Landscape for DDTree Verify Attention" — Analyzes message 12199, reading the allocator and FlashMLA signatures.

[8] "The Anatomy of a Kernel Investigation: Tracing the Path from Diagnosis to Solution in ML Inference Optimization" — Analyzes message 12200, synthesizing findings into the split-verify design.

[9] "The Architecture Detective: Uncovering FlashMLA's Potential for Blackwell DDTree Verification" — Analyzes message 12201, discovering FlashMLA's capabilities and the sm_120 risk.

[10] "The Architecture Gate: Tracing Blackwell Kernel Support in SGLang's MLA Backends" — Analyzes message 12202, grepping for architecture gating in backend files.

[11] "The Architecture Gate: Diagnosing Blackwell sm_120 Attention Backend Support in SGLang" — Analyzes message 12203, reading the attention registry for backend selection logic.

[12] "The Architecture Gap: Tracing SGLang's Triton Dependency on Blackwell Consumer GPUs" — Analyzes message 12204, discovering that sm_120 falls outside optimized backend coverage.

[13] "The Decisive Inspection: How cuobjdump Resolved the sm_120 Kernel Compatibility Question" — Analyzes message 12205, the definitive binary inspection that confirmed no sm_120 support.

[14] "The Architecture Wall: When Optimized GPU Kernels Can't Run on Your Hardware" — Analyzes message 12206, the ISA incompatibility analysis.

[15] "The Architecture Decision That Unlocked Blackwell Consumer GPUs" — Analyzes message 12207, the user's greenlight to compile for sm_120.

[16] "Decoding GPU Architecture Constraints: The sm_120 Compatibility Analysis for Custom MLA Verify Kernels" — Analyzes message 12208, the shared memory and ISA constraint analysis.

[17] "The Pivot Point: Examining Source Code to Determine sm_120 Kernel Portability" — Analyzes message 12209, examining the local sgl-kernel source tree.

[18] "The Architecture Detective: How a Deep-Dive Investigation Unlocked sm_120 GPU Kernel Optimization" — Analyzes message 12210, the CMakeLists.txt analysis revealing sm_120a in general gencode but not FlashMLA.

[19] "The Architecture Detective: Tracing CUDA Kernel Compatibility Across GPU Generations" — Analyzes message 12211, the PTX JIT theory and discovery of hidden common_ops libraries.

[20] "The Architecture That Wasn't There: Building a Custom sm_120 Verify Kernel for Blackwell Consumer GPUs" — Analyzes message 12212, the comprehensive investigation report and plan.

[21] "The Architecture of a Pivot: Building an Owned sm_120 Verify Kernel for SGLang" — Analyzes message 12213, the refined execution-ready plan.

[22] "The Green Light: How a Three-Word Directive Transformed an AI's Plan Into Action" — Analyzes message 12214, the user's "write down a plan/spec for each phase, then execute" directive.

[23] "From Planning to Execution: The Pivot Point in Building a Custom sm_120 Verify Kernel" — Analyzes message 12215, the shift into execution mode.

[24] "The Moment of Execution: Reading the Oracle Kernel" — Analyzes message 12216, reading the naive verify_attn.cu kernel.

[25] "Reading the Blueprint: How One File-Read Message Anchored a Complex Kernel Implementation" — Analyzes message 12217, reading the verify_attn.cuh header.

[26] "Reading the Test Harness: A Methodical Foundation for Custom CUDA Kernel Development" — Analyzes message 12218, reading test_verify_attn.cu.

[27] "Reading the Build Script: A Methodical Transition from Planning to Execution" — Analyzes message 12219, reading build_nvcc.sh.

[28] "Reading the Reference: How a Single File-Read Message Anchored a Complex CUDA Kernel Implementation" — Analyzes message 12220, reading gen_verify_attn_refs.py.

[29] "Ground Truth in Silicon: The Role of the Python Reference in Custom CUDA Kernel Development" — Analyzes message 12221, reading mla_attn_ref.py.

[30] "Reading the Benchmark: How a Single File Read Shaped a Custom CUDA Kernel Design" — Analyzes message 12222, reading bench_kernels.cu.