The Orchestrator's Moment: Wiring Split-K Dispatch into a Custom CUDA Attention Kernel

Introduction

In the high-stakes world of production GPU kernel optimization, the difference between a research prototype and a deployable system often comes down to the orchestration layer — the code that decides which kernel to run, how to configure it, and when to switch between paths. Message <msg id=12564> captures precisely this moment in a months-long campaign to optimize DeepSeek-V4-Flash inference on NVIDIA Blackwell (sm_120) GPUs. The assistant, having already written two new CUDA kernels (a split kernel and a combine kernel) in the previous message <msg id=12563>, now turns to the critical task of refactoring the wrapper function that will compute the adaptive split count (NSPLIT) and dispatch between the single-kernel and split-combine execution paths.

This message is brief in its visible content — a few lines of agent reasoning followed by a read tool call that shows lines 466–477 of a Triton kernel file. But its significance extends far beyond those lines. It represents the architectural seam where algorithmic intent meets hardware reality, where the abstract concept of "split-K parallelization" is translated into concrete dispatch logic that must be correct, performant, and CUDA-graph-safe.

The Context: Why Split-K Matters

To understand this message, one must appreciate the problem it solves. The assistant had previously rewritten the sparse-MLA attention kernel for DeepSeek-V4-Flash, replacing a per-head SIMT (Single-Instruction-Multiple-Thread) kernel with an MMA (Matrix-Matrix-Accumulate) kernel using Triton's tl.dot tensor-core operations. This rewrite was spectacularly successful at higher batch sizes — C=64 throughput doubled from 29.7 to 59.6 tokens/second — but it introduced a regression at single-batch workloads (C=1 dropped from 11.5 to 8.1 tok/s, a 29% loss).

The root cause was occupancy. The new kernel used a head-batched grid of shape (B, H/BLOCK_H) where BLOCK_H=32, meaning at batch size 1 with 64 heads, only 2 blocks were launched on a GPU with 188 streaming multiprocessors (SMs). That's 1% occupancy. The old SIMT kernel had launched 64 blocks (one per head) and achieved far better utilization. The fix — split-K parallelization over the topk dimension — would partition each head's work across multiple blocks, each handling a contiguous chunk of the sparse KV cache tokens, then combine results using a log-sum-exp (LSE) merge.

The user had explicitly chosen this path: when asked whether to proceed with torch.compile on the glue code (the 54% bottleneck) or finish the kernel with split-K first, the user answered "Finish the kernel: add split-K first" (see <msg id=12561>). The assistant had already written the split and combine kernels in <msg id=12563>. Now it needed to wire them into the runtime.

The Message: A Read That Reveals the Architecture

The visible content of <msg id=12564> is deceptively simple:

I need to refactor the wrapper function to calculate NSPLIT and route the computation appropriately—either directly to the single kernel when NSPLIT equals one, or through the split-and-combine path otherwise. Now rewrite the wrapper to compute adaptive NSPLIT and dispatch split+combine: [read] /home/theuser/glm-kimi-sm120-rtx6000bw/dsv4-live/attention/flash_mla_sm120_triton.py

The file read shows lines 466–477, which include kernel parameter declarations: L_ptr (the denominator buffer), sm_scale, page_size, page_bytes, scale_section_off, H, topk, topk_rounded, has_topk_len, and stride parameters. These are the interface between the wrapper and the kernel — the contract that must be satisfied regardless of which execution path is taken.

The agent reasoning reveals the core design decision: the wrapper must compute NSPLIT adaptively based on the batch size B, then branch on whether NSPLIT > 1. If NSPLIT == 1, the existing single-kernel path (which processes all topk tokens in one block) is sufficient. If NSPLIT > 1, the split kernel runs first (producing partial accumulators, max values, and denominator terms for each chunk), then the combine kernel merges them using the LSE trick.

The Thinking Process: Constraints and Tradeoffs

The assistant's reasoning, visible in the preceding messages, reveals a sophisticated understanding of the constraints. The CUDA graph capture mechanism is the most demanding constraint: once a graph is captured, all kernel launches, grid sizes, and buffer addresses are frozen. This means NSPLIT must be deterministic for a given batch size bucket. The wrapper cannot dynamically decide NSPLIT at runtime based on runtime conditions — it must be baked in at graph-capture time.

The adaptive NSPLIT formula targets approximately 188–376 blocks (enough to saturate the 188 SMs with 1–2 waves). For B=1 with 2 head-groups, NSPLIT=8 yields 16 blocks — still far from ideal but much better than 2. For B=16, NSPLIT≈6 yields 192 blocks. For B=64, NSPLIT=2 yields 256 blocks. The formula balances two competing goals: maximizing occupancy (more splits = more blocks) while keeping chunk sizes large enough for efficient tile processing (BLOCK_T must fit within each chunk).

The assistant also considered memory overhead. The scratch buffers for partial accumulators are sized as [B, NHG, NSPLIT, BLOCK_H, D] in bf16, which at worst case is about 6.2 MB per layer, or ~530 MB across all 43 layers per decode step. This is acceptable — a small fraction of the GPU's 48 GB HBM2e memory — but must be allocated once per graph bucket and reused across replays.

Assumptions and Their Validity

The assistant makes several assumptions in this message, most of which are well-justified:

  1. NSPLIT is deterministic per batch bucket: This holds because B is fixed within each CUDA graph bucket, and the formula depends only on B, n_hg, and the target block count (an environment variable). The assistant correctly notes that kernel arguments like nsplit and chunk are Python ints that get baked in at capture time.
  2. The autotuner's BLOCK_T choice is acceptable across chunk sizes: The Triton autotuner runs during warmup on the first call, picking a BLOCK_T based on whatever chunk size that initial call has. For different B values with different chunk sizes, the chosen BLOCK_T may not be optimal, but the tiling logic handles variable chunks correctly. This is a pragmatic tradeoff — optimal tuning per bucket would require separate autotuner runs, adding complexity.
  3. The combine kernel's memory layout is correct: The scratch buffer stores nope values in columns 0:448 and rope values in 448:512. The assistant verified this layout matches the output tensor layout, so reading all 512 dimensions and storing directly is correct without dimension masking.
  4. bf16 precision for partial accumulation is sufficient: The assistant notes that the tradeoff of bf16 precision is acceptable given the tolerance requirements (~1e-3 relative error). This was validated in subsequent testing (see <msg id=12567> where relative errors of 3.8e-3 to 6.7e-3 are reported). One potential mistake is the assumption that the single-kernel path (NSPLIT=1) doesn't need the combine step. The assistant's reasoning in <msg id=12562> considered this: "keeping the existing single-split kernel unchanged." This is correct because when NSPLIT=1, the existing kernel already computes the full attention output without needing a merge. The wrapper simply skips the combine step entirely.

Input Knowledge Required

To fully understand this message, one needs knowledge of:

Output Knowledge Created

This message creates the architectural blueprint for the split-K dispatch. The wrapper function that results from this refactoring becomes the central orchestrator for the attention kernel, responsible for:

  1. Computing NSPLIT adaptively from batch size
  2. Allocating scratch buffers of the correct size
  3. Dispatching to the single-kernel path (NSPLIT=1) or the split+combine path (NSPLIT>1)
  4. Ensuring CUDA graph safety by fixing NSPLIT per bucket The downstream impact is visible in <msg id=12567>, where the correctness test passes with relative errors of 3.8e-3 to 6.7e-3 across B=1, 4, 16, 32 configurations — all exercising the split+combine path since NSPLIT≥4 for B≤32. The split-K implementation was ultimately successful, though the even larger breakthrough (the ~17× throughput gain) came from a different optimization discovered later in the same segment: fixing the indexer's O(max_context) bottleneck by capping context length.

The Broader Significance

This message exemplifies a pattern that recurs throughout production kernel development: the hardest work is often not the kernel itself but the orchestration layer that makes it usable. Writing a split kernel and a combine kernel is intellectually challenging, but wiring them into a dispatch system that is correct, performant, and compatible with infrastructure constraints (CUDA graphs, autotuning, memory management) requires a different kind of engineering rigor.

The assistant's approach — first define the environment knobs, then write the kernels, then refactor the wrapper — follows a sound bottom-up methodology. Each layer is tested independently before integration. The wrapper refactoring in this message is the final integration step, the moment when the two new kernels become part of a coherent system rather than standalone components.

For the reader, this message offers a window into the real work of GPU kernel optimization: not just the flashy kernel launches, but the careful, methodical work of building the infrastructure that decides which kernel to run, with what parameters, under what conditions. It is the difference between a demo and a deployment.