The Split-K Dispatch: A Single Edit That Completed the Attention Kernel

Message 12566: [edit] /home/theuser/glm-kimi-sm120-rtx6000bw/dsv4-live/attention/flash_mla_sm120_triton.py — "Edit applied successfully."

At first glance, message 12566 appears to be the most mundane event in a coding session: a tool result confirming that an edit was applied. The message contains exactly three words of substance — "Edit applied successfully" — and a file path. Yet this brief confirmation marks a pivotal moment in a multi-day kernel optimization campaign for DeepSeek-V4-Flash on NVIDIA Blackwell (sm_120) GPUs. Understanding why this particular edit matters, what it accomplished, and the chain of reasoning that led to it reveals the intricate decision-making process behind high-performance GPU kernel engineering.

The Context: A Kernel Campaign in Full Swing

To grasp the significance of message 12566, one must understand the trajectory of the optimization effort. The assistant had been working on deploying and optimizing DeepSeek-V4-Flash, a large language model, across 8× RTX PRO 6000 Blackwell GPUs. The core challenge was that the model's sparse Multi-head Latent Attention (MLA) decode kernel was running on CUDA-core SIMT units instead of the tensor cores, re-reading the KV cache 64× redundantly per head and consuming 57% of decode GPU time.

The assistant had already designed and implemented a custom MMA (matrix multiply-accumulate) sparse-MLA decode kernel using Triton's tl.dot tensor-core operations, replacing the per-head SIMT kernel. This initial rewrite delivered a dramatic 6.3× speedup for the attention kernel itself and doubled overall throughput at C=64 (29.7 → 59.6 tok/s). However, it introduced a regression at C=1 (11.5 → 8.1 tok/s) due to under-occupancy — the grid was only (B, 2) blocks, leaving the GPU's 188 SMs mostly idle at small batch sizes.

The user was presented with a choice: proceed to torch.compile on the glue code (the now-dominant 54% bottleneck) or finish the kernel work by adding split-K parallelization to fix the occupancy regression. The user chose "Finish the kernel: add split-K first" ([msg 12561]). Message 12566 is the culmination of that decision.

What the Edit Actually Did

The edit confirmed in message 12566 was the final piece of the split-K implementation. The assistant had already written two new Triton kernels in a previous edit ([msg 12563]): a split kernel that divides the topk token dimension across multiple thread blocks (grid (B, n_hg, NSPLIT)), each computing partial unnormalized attention outputs, and a combine kernel that merges these partial results using the log-sum-exp (LSE) trick for numerical stability.

What remained was to refactor the _run_mma_sparse_decode wrapper function — the entry point that orchestrates kernel launches — to compute the adaptive NSPLIT value and dispatch between two paths: the original single-kernel path when NSPLIT equals 1, and the new split+combine path otherwise. Message 12565 shows the assistant reading the wrapper function to understand its current structure. Message 12566 is the result of the edit that rewrote this wrapper.

The edit was the architectural glue that connected the new kernels to the existing inference pipeline. Without it, the split and combine kernels existed as dead code — implemented but unreachable. The wrapper rewrite made them live, completing the split-K feature.

The Reasoning Behind the Design

The assistant's thinking, visible in the extended reasoning of message 12562, reveals a sophisticated understanding of GPU occupancy, numerical stability, and CUDA graph capture constraints. Several key design decisions were made:

Adaptive NSPLIT: Rather than using a fixed number of splits, the assistant designed NSPLIT to adapt to the batch size. The goal was to target approximately 188–376 blocks for good occupancy of the 188 SMs. At B=1 with 2 head groups, NSPLIT=16 yields 32 blocks — still underutilized but far better than the original 2 blocks. At B=16, NSPLIT≈6 yields ~192 blocks. At B=64, NSPLIT=2 yields 256 blocks. This adaptive formula ensures the split-K overhead (additional scratch memory, combine kernel launch) scales naturally with batch size, avoiding unnecessary work when occupancy is already high.

CUDA Graph Safety: A critical constraint was that the kernel must be capturable into CUDA graphs for production deployment. The assistant recognized that since B is fixed per graph bucket and NSPLIT is deterministic from B, the scratch allocation size remains constant across graph replays. The branch on NSPLIT is also deterministic per bucket, making it safe during graph capture. This attention to the production deployment context — not just correctness but also deployability — distinguishes this work from a purely academic kernel exercise.

Numerical Precision Tradeoffs: The assistant considered using bf16 for the partial accumulator scratch buffers to halve memory traffic, calculating that the worst-case per-layer traffic of ~6.2 MB written and read, across 43 layers totaling ~530 MB per step, would add negligible latency. The decision to use f32 initially (with bf16 as a future optimization) reflects a pragmatic "correctness first, optimize later" approach.

The LSE Combine Strategy: The combine kernel uses the standard log-sum-exp trick for numerically stable merging of partial softmax results: loading all NSPLIT partitions of the max values (m), finding the global maximum per head, computing correction factors with exp2(m_local - m_global), and normalizing the accumulated values by the corrected denominator sums. This is a well-known technique in the attention kernel literature (used in FlashAttention and similar algorithms), and its application here demonstrates the assistant's familiarity with established numerical methods.

Assumptions and Potential Pitfalls

Several assumptions underpinned this edit, some of which carried risk:

The autotuner assumption: The assistant noted that the Triton autotuner runs during warmup on the first call, picking a BLOCK_T configuration based on whatever chunk size that initial call uses. This configuration might not be optimal for other chunk sizes across different B buckets. The assistant judged this acceptable because "the tiling logic handles variable chunks correctly so it's still correct" — prioritizing correctness over peak optimality.

The NSPLIT=1 bypass: For B≥128 (where NSPLIT would equal 1), the wrapper bypasses the split+combine path entirely and uses the original single-kernel path. This assumption was safe because B≥128 is outside the cuda-graph range and the single-kernel path was already validated, but it meant the split+combine path was never tested at its boundary condition.

Scratch buffer sizing: The assistant assumed that the product B×NSPLIT stays roughly constant around the target block count, keeping peak memory manageable. This is a reasonable heuristic but could fail for pathological batch sizes where the adaptive formula produces unexpected NSPLIT values.

Input and Output Knowledge

Input knowledge required to understand this message includes: the architecture of the MMA sparse decode kernel (head-group batching, topk dimension, KV cache layout), the concept of split-K parallelization for attention, the log-sum-exp numerical trick, CUDA graph capture constraints, Triton's autotuner behavior, and the overall optimization campaign that preceded this moment.

Output knowledge created by this message is the completion of the split-K feature. The immediate next message ([msg 12567]) shows the assistant running correctness tests across multiple batch sizes (B=1, 4, 16, 32) and topk values (128, 512), validating that the split+combine path produces outputs with relative error ≤ 6.7×10⁻³ against the original SIMT kernel — the same tolerance as the original MMA kernel. The split-K implementation was then ready for deployment.

The Broader Significance

Message 12566 is a reminder that in complex engineering work, the most consequential moments are often the least dramatic. A three-word confirmation of a file edit doesn't look like a breakthrough — but it represents the resolution of a deliberate, multi-step decision process: profiling to identify the C=1 regression, weighing split-K against torch.compile, presenting the tradeoff to the user, designing the adaptive NSPLIT formula, implementing two new kernels, and finally wiring them into the production path.

The edit itself was the last in a chain of three edits to the same file ([msg 12562], [msg 12563], [msg 12566]), each building on the previous one: first adding environment configuration knobs, then inserting the kernel implementations, and finally rewriting the wrapper to orchestrate them. Message 12566 is the capstone — the moment the design became operational.

In the broader arc of the optimization campaign, this split-K work would later be superseded by a far more dramatic breakthrough: the discovery that the "glue bottleneck" was actually the indexer computing scores over the full 1M-token max context every decode step. Capping context length to 8192 would deliver a 17.9× throughput gain at C=64, dwarfing the split-K improvements. But that discovery lay in the future. At the moment of message 12566, the split-K completion was the right next step — a disciplined response to a known regression, executed with careful attention to numerical correctness, production safety, and architectural coherence.