The Moment Attention Ceased to Matter: A Kernel Optimization Milestone on Blackwell GPUs

Introduction

In the high-stakes world of large language model inference on cutting-edge hardware, optimization is a game of diminishing returns—until it isn't. Every kernel engineer knows the pattern: you attack the dominant bottleneck, squeeze it until it shrinks, and then the next bottleneck reveals itself. But occasionally, there comes a moment where the improvement is so dramatic that it fundamentally reshapes the performance profile of the entire system. Message [msg 12571] captures exactly such a moment.

This message, written by an AI assistant during an intensive optimization campaign for DeepSeek-V4-Flash running on 8× NVIDIA RTX PRO 6000 Blackwell GPUs (sm_120 architecture), documents the aftermath of deploying a custom split-K sparse-MLA attention kernel. The numbers are staggering: the attention kernel, which had consumed 57.1% of all GPU time in the original implementation, dropped to just 9.6%. But more importantly, this message reveals the subtle, counterintuitive dynamics that emerge when a system is pushed to its limits—and the strategic reasoning required to navigate them.

The Road to Split-K

To understand the significance of message [msg 12571], we must first understand what came before. The optimization campaign had been progressing through a planned sequence of improvements. The original per-head SIMT (Single-Instruction, Multiple-Thread) kernel for the sparse multi-head latent attention (MLA) decode was consuming a staggering 57% of GPU time. The assistant had already replaced it with a custom MMA (matrix multiply-accumulate) kernel using Triton's tl.dot tensor-core operations, which eliminated the 64× redundant KV cache re-reading inherent in the SIMT approach. That first rewrite cut attention's share to 17.6% and roughly doubled throughput at higher concurrencies.

But there was a problem: at low concurrency (C=1, a single request), the new MMA kernel actually regressed—throughput dropped from 11.47 to 8.13 tokens/second. The root cause was under-occupancy: the grid was launching only (B, 2) blocks, leaving the GPU's 188 Streaming Multiprocessors mostly idle. The assistant proposed a fix: split-K parallelization over the topk dimension, where each attention computation would be split across multiple thread blocks that each handle a chunk of tokens, then combined via a log-sum-exp (LSE) merge.

The user, when presented with a choice between proceeding to torch.compile on the glue code or finishing the kernel with split-K, chose the latter ([msg 12561]). The assistant then designed and implemented the split-K kernel: a two-phase approach with a split kernel writing unnormalized partials across grid (B, n_hg, NSPLIT) and a combine kernel performing LSE reduction across splits. The NSPLIT parameter was made adaptive—targeting ~188–376 blocks for good occupancy, capped at 16, with a minimum chunk size of 32 tokens—ensuring CUDA graph safety by making it deterministic per batch-size bucket ([msg 12562]).

After validation showing relative error ≤ 6.7e-3 (the LSE merge is numerically exact), the assistant deployed the kernel and ran benchmarks ([msg 12567], [msg 12568]). The results were dramatic: C=1 throughput jumped from 8.13 to 22.43 tokens/second (a 176% improvement over the non-split MMA version, and +96% over the original), C=16 went from 40.97 to 56.46 (+38%), and C=64 from 59.56 to 61.98 (+4%). The C=1 regression was not just fixed—it was obliterated, with per-token latency dropping from 122ms to 30ms, a 4× improvement.

The Message: Reading the Profile After Victory

Message [msg 12571] begins with the assistant examining a fresh C=32 profile taken with the split-K kernel deployed. The profile reveals the new landscape:

Looking at the C=32 profile with split-K, the CUTLASS SIMT kernel is now the top performer at 19.2%, while the elementwise operations and glue code collectively dominate at over 60%. The sparse decode kernel has dropped to under 10%, which is a significant improvement from the previous 57%.

The attention kernel has been tamed. The assistant presents a before-and-after comparison table that tells the story of the entire campaign:

| kernel | original | MMA | MMA+split-K | |---|---|---|---| | sparse-MLA decode | 57.1% | 17.6% | 9.6% | | FP32 SIMT GEMM (indexer+MHC) | 8.8% | 17.0% | 19.2% | | elementwise/copy/reduce glue | ~28% | ~54% | ~60% | | FP4 MoE GEMM | 1.9% | 3.4% | 3.2% |

Attention went from the dominant cost to a minor contributor—a 6× reduction in relative share. But the table also reveals something troubling: the FP32 SIMT GEMM and the elementwise glue code have both grown as a share of total time. The FP32 GEMM, which performs the indexer batch matrix multiply and the MHC (multi-head cache) pre-projection, jumped from 8.8% to 19.2%, becoming the single largest individual kernel. The glue code—thousands of small elementwise operations, copies, and reductions—ballooned from ~28% to ~60%.

The Counterintuitive Finding

The assistant notices something odd:

But there's something odd here — the attention kernel time nearly halved (850ms down to 505ms), yet the total GPU time actually increased slightly (4828ms to 5240ms). The glue kernels like direct_copy and reduce all went up a bit, which is counterintuitive.

This is a subtle but important observation. The split-K kernel, by splitting attention work across more thread blocks, introduced additional synchronization and data movement overhead. The combine kernel itself doesn't show up in the top 12 kernels (it's negligible), but the supporting operations—the direct_copy and elementwise_kernel launches that move data between the split kernel's outputs and the combine kernel's inputs—increased. The step counts are similar between profiles, so this isn't an artifact of capturing different numbers of decode iterations.

The assistant correctly identifies the dynamic: "The split-K approach seems to be adding some overhead in the supporting kernels that's offsetting the attention gains." This is a classic systems optimization trade-off: you optimize one component so aggressively that the overhead of the optimization itself becomes visible in other parts of the system. The split-K kernel uses scratch buffers for partial accumulators, and those buffers must be written by the split kernel and read by the combine kernel. The memory traffic for these scratch buffers—even though the assistant had optimized it to bf16 to halve the footprint—adds pressure on the memory subsystem, and the additional kernel launches add launch latency.

Yet despite this increase in total GPU time, the end-to-end throughput improved. This is possible because the attention kernel was the critical path—it was the operation that everything else waited on. By making attention faster, even at the cost of slightly more total work, the pipeline could process tokens faster overall. The assistant correctly judges: "the split-K trade-off is still worth it."

Strategic Re-Evaluation: The Fork in the Road

The most interesting part of the message is the assistant's strategic reasoning about what to do next. The original plan, agreed with the user, was step "2 then 3": finish the kernel (step 2, now complete with split-K), then move to torch.compile on the glue code (step 3). But the profile reveals a more nuanced picture.

The FP32 SIMT GEMM at 19.2% is now the single largest individual kernel. It's a CUTLASS SIMT sgemm (single-precision GEMM) running on CUDA cores rather than tensor cores—a "fallback" kernel that doesn't use the Blackwell architecture's tensor memory accelerators. The assistant had identified this earlier as a "cheap bf16 flip"—converting the operation from FP32 to bf16 would let it run on tensor cores, likely delivering a significant speedup with minimal risk.

The glue code at ~60%, on the other hand, is fragmented across approximately 6,000 kernel launches per decode step. These are small elementwise operations—multiplications, additions, clamp operations, copies—interspersed with custom kernels for sparse attention, FP4 MoE, and the indexer. The assistant recognizes that torch.compile, which fuses operations by tracing through the PyTorch computation graph, is likely to struggle here:

torch.compile on a custom model like DSv4 with sparse attention, MoE, and custom indexing ops is risky — it'll likely hit graph breaks or fail to trace entirely. The glue code is fragmented across many small elementwise ops interspersed with those custom kernels, so torch.compile's fusion benefits are limited by all those breaks.

This is a crucial insight about the practical limitations of compiler-based optimization. torch.compile works best on "clean" models with standard operations that can be captured in a single graph. DeepSeek-V4-Flash, with its custom sparse attention, FP4 quantized MoE routing, and indexer logic, is anything but clean. The graph breaks at each custom kernel boundary would prevent the compiler from seeing the full picture, limiting fusion opportunities to small islands of elementwise operations between custom kernels.

Meanwhile, the FP32→bf16 GEMM flip is "surgical, targets the single biggest kernel now at 19%, and is low-risk." The assistant weighs these options and decides to present the trade-off to the user rather than unilaterally proceeding:

I think I should present this trade-off and recommend the bf16 GEMM as the more concrete next win, with torch.compile as a more speculative attempt at the fragmented glue.

This is sound engineering judgment. When faced with two possible paths—one high-risk, high-uncertainty (torch.compile on a fragmented graph) and one low-risk, high-confidence (flip a few lines of code to change the dtype)—the prudent choice is to take the sure win first. The bf16 GEMM flip would likely deliver a meaningful speedup with near-zero risk of breaking the model, while torch.compile could consume significant engineering time with uncertain returns.

The Documentation Step

The message concludes with the assistant appending a detailed entry to PROFILE_FINDINGS.md on the remote server. This documentation captures:

The Deeper Significance

Message [msg 12571] is remarkable not for any single insight but for the density of engineering judgment it contains. In a single message, the assistant:

  1. Interprets a complex profile — reading past the headline numbers to notice the counterintuitive increase in total GPU time despite attention halving, correctly attributing it to split-K overhead.
  2. Quantifies the milestone — producing the before/after/now table that shows attention dropping from 57% to 9.6%, giving the team a clear sense of progress.
  3. Re-evaluates the strategy — recognizing that the planned next step (torch.compile) may not be the best use of effort given the profile shift and the risk of graph breaks.
  4. Formulates a recommendation — proposing the FP32→bf16 GEMM flip as a more concrete, lower-risk alternative, while keeping torch.compile as a speculative option.
  5. Documents everything — ensuring the knowledge is preserved for future reference. This is the kind of reasoning that separates mechanical optimization from genuine systems engineering. The assistant isn't just following a script; it's reading the system's behavior, forming hypotheses about what it's seeing, questioning its own assumptions (the "something odd here" observation), and adapting the strategy based on evidence. The message also illustrates a fundamental truth about optimization: the bottleneck always moves. You optimize one component, and another component becomes the new bottleneck. The attention kernel went from 57% to 9.6%, but the FP32 GEMM grew from 8.8% to 19.2%, and the glue code expanded from ~28% to ~60%. The system is like a balloon—squeeze one part and another bulges out. The engineer's job is not to eliminate all bottlenecks (an impossible task) but to identify which bottleneck, if addressed, would yield the greatest improvement in the metric that matters.

The Assumptions and Knowledge at Play

To fully understand this message, several pieces of input knowledge are required:

The Thinking Process

The assistant's reasoning in this message follows a discernible pattern:

  1. Observe: "Looking at the C=32 profile with split-K, the CUTLASS SIMT kernel is now the top performer at 19.2%..."
  2. Question: "But there's something odd here — the attention kernel time nearly halved... yet the total GPU time actually increased slightly."
  3. Hypothesize: "The split-K approach seems to be adding some overhead in the supporting kernels that's offsetting the attention gains."
  4. Verify consistency: "The step counts look roughly similar across both profiles, so it's not just a difference in how many decode iterations were captured."
  5. Accept the trade-off: "Despite the slight GPU time increase, the end-to-end throughput actually improved, so the split-K trade-off is still worth it."
  6. Re-evaluate priorities: "The next optimization targets remain the same: tackle the glue overhead through torch.compile and flip the FP32 GEMM to bf16."
  7. Question the plan: "Actually, reconsidering the risk here: torch.compile on a custom model like DSv4... is risky."
  8. Formulate alternative: "Meanwhile, the FP32→bf16 GEMM flip is surgical, targets the single biggest kernel now at 19%, and is low-risk."
  9. Recommend: "I think I should present this trade-off and recommend the bf16 GEMM as the more concrete next win." This is textbook engineering reasoning: observe, question, hypothesize, verify, accept, re-evaluate, recommend. The assistant doesn't treat its own plans as immutable; it constantly re-examines them in light of new evidence.

Conclusion

Message [msg 12571] captures a pivotal moment in a high-stakes optimization campaign. The split-K kernel has delivered on its promise, fixing the C=1 regression and driving attention's share of GPU time from 57% down to 9.6%. But victory brings new challenges: the profile has shifted, revealing the FP32 GEMM and the fragmented glue code as the new bottlenecks. The assistant's response—careful analysis, strategic re-evaluation, clear documentation, and a well-reasoned recommendation—exemplifies the kind of disciplined engineering thinking that separates successful optimization campaigns from endless tweaking.

The message also serves as a reminder that optimization is never truly "done." Every improvement reveals the next constraint. The attention kernel is now thoroughly optimized, but the FP32 GEMM awaits its bf16 conversion, and the glue code's 6,000 kernel launches still need to be fused. The campaign continues, but this message marks the moment when the team could look at the profile and say: "Attention is no longer the problem."