The Read That Preceded a Breakthrough: Context-Gathering in the Split-K Optimization Campaign

Introduction

In the midst of an intense GPU kernel optimization campaign for DeepSeek-V4-Flash on 8× NVIDIA RTX PRO 6000 Blackwell (sm_120) GPUs, a single message stands out for what it represents rather than what it says. Message [msg 12565] is deceptively simple: a read tool call that displays lines 648–659 of a Triton kernel file. On its surface, it is nothing more than an assistant reading its own source code. But in the context of the surrounding engineering effort, this read is a pivotal moment of context-gathering — the assistant pausing to verify the interface boundary between two kernel implementations before performing a surgical refactor that would ultimately contribute to a 17× throughput improvement on the target model.

This article examines that single message in depth: why it was written, what decisions it reflects, the reasoning that led to it, and the knowledge it both consumed and produced.

The Broader Campaign: From SIMT Bottleneck to MMA Victory

To understand message [msg 12565], one must first understand the optimization trajectory that led to it. The assistant had been systematically profiling and optimizing the sparse Multi-head Latent Attention (MLA) decode kernel for DeepSeek-V4-Flash on Blackwell GPUs. The original production kernel used a SIMT (Single-Instruction, Multiple-Thread) approach that re-read the KV cache 64× redundantly per head and ran on CUDA cores rather than tensor cores. Profiling revealed this single kernel consumed 57.1% of total decode GPU time — over 5.3 seconds per profile window.

The assistant designed and implemented a custom MMA (Matrix-Matrix Accumulate) kernel using Triton's tl.dot tensor-core operations, replacing the per-head SIMT kernel with a head-batched grid that gathers KV data once per tile and reuses it across a head group. The result was dramatic: the attention kernel dropped from 5,350 ms to 850 ms — a 6.3× speedup — and total GPU time halved from 9,364 ms to 4,828 ms ([msg 12560]).

But this victory reshuffled the bottleneck hierarchy. With attention no longer dominant, the profile revealed a new culprit: unfused elementwise/copy/reduce glue operations consuming ~54% of GPU time, alongside forced-FP32 SIMT GEMMs in the indexer at 17%. The assistant faced a strategic decision: proceed immediately to torch.compile to fuse the glue operations (the biggest remaining lever for throughput at C=16–64), or first add split-K parallelization to the attention kernel to fix a C=1 regression and improve occupancy at moderate batch sizes.

The User's Decision and the Split-K Mandate

The assistant presented the trade-offs to the user in a structured question ([msg 12561]), outlining two paths:

  1. torch.compile first: Fuse the ~6,000 elementwise/copy/reduce launches consuming 54% of GPU time — the biggest throughput lever at the user's target concurrency range of C=16–64.
  2. Split-K first: Add parallelization over the topk dimension to fix the C=1 regression (11.5 → 8.1 tok/s, a 29% degradation) and boost occupancy at C=16/32 where only 32–64 of 188 SMs were busy. The user chose option 2: "Finish the kernel: add split-K first." This decision set the stage for the code changes that message [msg 12565] would support.

What Split-K Entails

Split-K is a parallelization strategy for attention kernels. In the standard approach, each CUDA block processes all topk tokens for a given batch element and head group — computing QK scores, softmax, and the weighted sum of values in a single program. When the batch size is small (e.g., C=1), the grid has only (B, n_hg) = (1, 2) = 2 blocks, leaving 186 of 188 SMs idle. This under-occupancy is why the C=1 throughput regressed.

Split-K divides the topk dimension into NSPLIT chunks, each handled by a separate CUDA block. Each block computes partial accumulators and local statistics (max score m and sum of exponentials l). A separate combine kernel then reads all NSPLIT partials, finds the global max per head, applies correction factors via exp2(m_local - m_global), and normalizes the accumulated values. The grid becomes (B, n_hg, NSPLIT), which at C=1 with NSPLIT=8 yields 16 blocks — still modest but 8× better than 2.

The design is constrained by CUDA graph capture: because NSPLIT is deterministic from the batch size B (which is fixed per graph bucket), the scratch buffer allocation size is constant across replays, making graph capture safe. The assistant designed an adaptive NSPLIT formula targeting 188–376 blocks (full SM occupancy on the RTX PRO 6000's 188 SMs), with environment-variable knobs for tuning.

The Message Itself: Reading the Interface Boundary

Message [msg 12565] shows the assistant executing a read on the file /home/theuser/glm-kimi-sm120-rtx6000bw/dsv4-live/attention/flash_mla_sm120_triton.py, displaying lines 648 through 659. The content reveals two critical pieces of code:

Lines 648–650 (end of the combine kernel):

        mask=h_mask[:, None],
    )
    tl.store(LSE_ptr + bid * H + h_global, lse, mask=h_mask)

This is the tail of the combine kernel that was inserted in message [msg 12563]. It stores the final log-sum-exp (LSE) values — the normalization denominator after merging across all splits — back to a global buffer. The LSE values are needed by downstream operations (e.g., for speculative decoding verification or for computing token probabilities).

Lines 653–659 (beginning of the wrapper function):

def _run_mma_sparse_decode(
    q: torch.Tensor,  # [B, 1, H, D] bf16
    k_cache: torch.Tensor,
    indices: torch.Tensor,
    topk_length: Optional[torch.Tensor],
    softmax_scale: float,
) -> Tuple[torch.Tensor,...

This is the public interface to the entire sparse decode operation. The assistant is reading this to understand exactly what the wrapper expects and returns, so it can correctly refactor the dispatch logic. The wrapper needs to be modified to:

  1. Compute adaptive NSPLIT from the batch size B and head-group count n_hg.
  2. Route execution: if NSPLIT == 1, call the original single-kernel path; otherwise, allocate scratch buffers, launch the split kernel with grid (B, n_hg, NSPLIT), then launch the combine kernel to merge partial results.
  3. Return the combined output tensor with the same shape and semantics as the original. The read is a verification step — the assistant confirming that the function signature it needs to modify matches its mental model before applying the edit.

The Thinking Process Behind the Read

The assistant's reasoning in message [msg 12562] reveals extensive deliberation about the split-K design. Several key considerations emerge:

Occupancy targeting: The assistant explicitly calculates block counts for different batch sizes. At C=1 with n_hg=2, NSPLIT=8 gives 16 blocks — still underutilizing 188 SMs, but far better than 2. At C=16, NSPLIT≈6 yields 192 blocks (full occupancy). At C=64, NSPLIT=2 yields 256 blocks. The adaptive formula naturally scales to keep the GPU busy.

Scratch buffer sizing: The assistant considers memory traffic carefully. With bf16 for the accumulator scratch, per-layer traffic is ~6.2 MB written and read. Across 43 layers (the DeepSeek-V4 architecture), that's ~530 MB per decode step — negligible compared to the KV cache size. The assistant explicitly considers the precision trade-off (bf16 vs f32) and deems it acceptable given the tolerance requirements.

CUDA graph safety: A critical constraint. The assistant verifies that NSPLIT is deterministic per batch bucket, so the branch on nsplit is safe during graph capture. Kernel arguments like nsplit and chunk are Python ints baked in at capture time, derived from the fixed B value. The autotuner runs during warmup on the first call, picking a BLOCK_T config — the assistant notes this might not be optimal for different chunk sizes across buckets, but the tiling logic handles variable chunks correctly.

Combine kernel algorithm: The assistant designs a two-pass approach within the combine kernel: first pass loads max values across all splits, second pass accumulates normalized values and log-sum-exp. This avoids extra kernel launches while keeping the algorithm numerically stable.

What This Message Produces

The output of message [msg 12565] is knowledge — specifically, confirmed knowledge about the codebase state. The assistant now knows:

Assumptions and Potential Pitfalls

The assistant makes several assumptions in this read:

  1. File consistency: It assumes the file on disk matches its mental model after the previous edits (messages [msg 12562] and [msg 12563]). In a collaborative environment with multiple agents or concurrent edits, this could be wrong — but in this single-threaded session, it's safe.
  2. Interface stability: It assumes the wrapper function signature won't need to change to accommodate split-K. In practice, the split-K path needs scratch buffers, which could be allocated inside the wrapper or passed as parameters. The assistant chose internal allocation, keeping the public interface unchanged.
  3. Autotuner optimality: The assistant acknowledges that the autotuner's BLOCK_T selection, based on the first call's chunk size, might not be optimal for all chunk sizes across batch buckets. This is a known limitation accepted as "still correct."

Conclusion

Message [msg 12565] is a deceptively simple read operation that reveals the meticulous, methodical nature of high-performance GPU kernel engineering. It is not a moment of creation but a moment of verification — the assistant pausing to confirm its understanding of the codebase before making a surgical modification. The read sits at the intersection of several converging threads: the successful MMA kernel rewrite that reshuffled the bottleneck hierarchy, the user's strategic decision to pursue split-K over torch.compile, and the complex design of adaptive parallelization with CUDA graph safety constraints.

In the broader narrative of the optimization campaign, this message is the calm before the storm. The edit that follows ([msg 12566]) would implement the split-K dispatch, and subsequent profiling would reveal that the real bottleneck was not in the attention kernel at all but in the indexer's O(max_context) computation — leading to a 17× breakthrough. But that breakthrough was only possible because the assistant first completed the kernel work the user asked for, and that work began with a careful read of the code.