The Split-K Decision: When Engineering Discipline Trumps the Obvious Shortcut

Introduction

In the high-stakes world of ML inference optimization on cutting-edge hardware, every decision carries weight. A single wrong turn can waste days of engineering effort; a single correct one can unlock an order-of-magnitude performance gain. The conversation captured in message 12562 of this opencode session represents one of those pivotal moments—a fork in the road where the assistant, having just delivered a stunning 2× throughput improvement on a complex attention kernel, must decide where to invest its next cycle of effort.

This article examines that single message in depth: its reasoning, its assumptions, its technical depth, and the engineering philosophy it embodies. The message is the assistant's response to a user who has just answered a critical question: "Finish the kernel: add split-K first." What follows is not merely an implementation note but a window into how a sophisticated AI agent reasons about GPU kernel design, occupancy modeling, CUDA graph safety, and the tradeoffs between competing optimization strategies.

The Context: A Campaign Nearing Its Climax

To understand message 12562, we must first understand the journey that led to it. The session involved deploying and optimizing the DeepSeek-V4-Flash model (an NVFP4-quantized variant) on a machine with 8× NVIDIA RTX PRO 6000 Blackwell GPUs (sm_120 architecture). The assistant had been engaged in a multi-phase optimization campaign:

Phase 1 — The MMA Kernel Rewrite: The original sparse-MLA decode kernel used a per-head SIMT (Single Instruction, Multiple Thread) design that re-read the KV cache 64× redundantly across heads and relied on CUDA-core fallback operations rather than tensor-core MMA (matrix multiply-accumulate) instructions. The assistant replaced this with a custom Triton kernel using tl.dot tensor-core operations, head-batching the KV cache read so it was shared across a group of heads. This delivered a 6.3× speedup on the attention kernel itself, cutting total GPU decode time in half.

The Profile Shift: The MMA rewrite was a textbook success—it solved the bottleneck it targeted. But as any experienced optimizer knows, solving one bottleneck merely reveals the next. The profile shifted dramatically:

| Kernel Component | Before (SIMT) | After (MMA) | |---|---|---| | Sparse-MLA decode | 57.1% (5350ms) | 17.6% (850ms) | | FP32 SIMT GEMM (indexer+MHC) | 8.8% | 17.0% (822ms) | | Elementwise/copy/reduce glue | ~28% | ~54% (now dominant) | | FP4 MoE GEMM | 1.9% | 3.4% |

The attention bottleneck was solved. But the C=1 concurrency level had regressed from 11.5 tok/s to 8.1 tok/s because the new kernel's grid structure—(B, H/BLOCK_H=2)—produced only 2 blocks at batch size 1, severely under-occupying the GPU's 188 SMs.

The Fork in the Road: The assistant presented the user with a choice: proceed to torch.compile on the now-dominant 54% glue code (the biggest throughput lever at the target C=16-64 regime), or finish the kernel work by adding split-K parallelization to fix the C=1 regression and boost occupancy across the board. The user chose the latter: "Finish the kernel: add split-K first."

Message 12562: The Split-K Design Manifesto

The subject message opens with the assistant's reasoning block—a remarkable piece of technical writing that reads like a design document for a GPU kernel optimization technique. Let us examine it section by section.

The Core Problem: Under-Occupancy

The assistant begins by framing the problem with surgical precision:

"I'm going to implement split-K parallelization across the topk dimension to address the C=1 regression and improve occupancy for larger C values. The current grid structure processes all topk tokens per program, so I'll redesign it to split that work across multiple blocks and add a combine kernel to aggregate the results."

This opening sentence does three things simultaneously: it states the goal (fix C=1 regression and improve occupancy), identifies the root cause (the grid processes all topk tokens per program), and sketches the solution (split the work across blocks, combine results). The assistant is not blindly following a user directive—it is reasoning from first principles about why the current design is suboptimal and how to fix it.

The key insight here is that the current MMA kernel, for all its tensor-core efficiency, suffers from a structural occupancy problem. The grid is (B, n_hg) where n_hg is the number of head groups (2 in this case). At B=1, that's just 2 blocks on 188 SMs—1% occupancy. The GPU is nearly idle. Split-K addresses this by also parallelizing over the token dimension (the "K" in attention terminology), so the grid becomes (B, n_hg, NSPLIT), multiplying the block count by NSPLIT.

The Split-Combine Architecture

The assistant then lays out the mathematical foundation of the split-combine approach:

"Each split-K block will handle a contiguous chunk of topk tokens, computing partial accumulators and local statistics (max and sum of exponentials). Then the combine kernel will read all partials, compute the global maximum, and properly normalize the results using the exponential trick to avoid numerical issues."

This is the classic online-softmax or LSE (Log-Sum-Exp) combine pattern used in flash attention and similar algorithms. Each split computes a partial attention output over its token chunk, along with local statistics (the maximum score m and the sum of exponentials l). The combine kernel then merges these partial results: it finds the global maximum across all splits, adjusts each partial output by the correction factor exp(m_local - m_global), and sums the adjusted outputs weighted by their local l values.

The assistant demonstrates a deep understanding of numerical stability here. The "exponential trick" refers to the fact that directly summing exponentials of large negative numbers can cause underflow; instead, you track the maximum score and compute exp(score - max_score) to keep values in a representable range.

Memory Footprint and Precision Tradeoffs

A particularly sophisticated section deals with the memory implications:

"I'll allocate separate f32 scratch arrays for the accumulator, m, and l values — about 12 MB total when bounded by occupancy constraints, or 8 MB when NSPLIT caps at 1. Using f32 for now to preserve precision in the partial sums, though bf16 for the accumulator could halve memory traffic later if needed."

The assistant is sizing the scratch buffers based on the product B × NSPLIT × D, where D=512 is the MLA dimension. Since NSPLIT is adaptive (larger when B is small, smaller when B is large), the product stays roughly constant, bounding peak memory. The mention of bf16 as a future optimization shows forward thinking: the initial implementation uses fp32 for safety, but once correctness is validated, switching to bf16 could halve memory traffic at the cost of slightly reduced precision in the partial accumulators.

The assistant also notes that this memory footprint is negligible in context: "at worst case around 6.2MB written and read per layer, which across 43 layers is only about 530MB per step, adding negligible latency." This kind of system-level thinking—understanding how a per-layer optimization affects the entire model's memory budget—is characteristic of expert-level inference optimization.

CUDA Graph Safety: The Hidden Constraint

One of the most technically nuanced sections deals with CUDA graph safety:

"For CUDA graph safety, since B is fixed per graph bucket and NSPLIT is deterministic from B, the scratch allocation size stays constant across replays, so recording and replaying the graph is straightforward."

CUDA graphs are a mechanism for capturing a sequence of GPU operations and replaying them without host-side launch overhead. SGLang uses them extensively for the decode path. However, CUDA graphs have strict requirements: all memory allocations, kernel launches, and control flow must be deterministic and repeatable. A kernel whose grid size depends on runtime data (like the actual number of tokens) cannot be captured in a CUDA graph.

The assistant's solution is elegant: NSPLIT is computed from B (the batch size), which is fixed per graph bucket. SGLang captures separate CUDA graphs for each batch size bucket (e.g., B=1, B=2, B=4, ..., B=64). Since B is known at capture time, NSPLIT is deterministic, and the scratch buffer size is fixed. This means the split-K kernel can be safely captured and replayed.

The assistant also considers the autotuner interaction:

"The branch on nsplit is deterministic per bucket since B is fixed, so it's safe during graph capture. The kernel arguments nsplit and chunk are Python ints that get baked in at capture time, which is fine since they're derived from the fixed B value. The autotuner runs during warmup on the first call, picking a BLOCK_T config based on whatever chunk size that initial call has—it might not be optimal for other chunk sizes across different B buckets, but the tiling logic handles variable chunks correctly so it's still correct."

This reveals a deep understanding of Triton's autotuning mechanism and its interaction with CUDA graphs. The autotuner benchmarks multiple kernel variants on the first invocation and caches the best configuration. If the first call happens during graph capture (with a particular chunk size), the cached config might be suboptimal for other chunk sizes. But the assistant judges this acceptable because the tiling logic handles variable chunks correctly—correctness is preserved even if performance isn't perfectly tuned for every bucket.

The Adaptive NSPLIT Strategy

The assistant devotes significant reasoning to the adaptive NSPLIT formula:

"I'm targeting around 188–376 blocks for good occupancy, with a minimum chunk size to keep BLOCK_T tiles valid. For small B like 1 with n_hg=2, a fixed NSPLIT of 8 gives 16 blocks, which is better than 2 but still underutilizes the SMs. I can increase NSPLIT if I allow smaller chunks—say chunk≥32 instead of 64—which would let NSPLIT go up to 16 and get 32 blocks at B=1."

The target of 188-376 blocks is directly tied to the hardware: 188 SMs on the RTX PRO 6000 Blackwell. The assistant wants enough blocks to keep all SMs busy (or oversubscribe them for latency hiding) without creating so many blocks that each one's work is too small to be efficient.

The minimum chunk size constraint comes from Triton's tiling: each block processes BLOCK_T tokens at a time, and BLOCK_T is typically 16, 32, or 64. If the chunk is smaller than BLOCK_T, the kernel wastes threads on masked-out elements. The assistant is weighing the tradeoff between occupancy (more splits = more blocks) and efficiency (larger chunks = better utilization per block).

For larger batch sizes, the formula naturally produces good occupancy:

"B=16 gives NSPLIT≈6 with chunk≈85 (192 blocks total), and B=64 gives NSPLIT=2 (256 blocks). This keeps occupancy high across different batch sizes."

At B=64 with n_hg=2, the base grid is already 128 blocks. Adding NSPLIT=2 brings it to 256 blocks—good occupancy without needing aggressive splitting. The adaptive formula ensures that NSPLIT is large when B is small (where occupancy is worst) and small when B is large (where occupancy is already acceptable).

The Torch Combine Rejection

One of the most revealing passages is the assistant's consideration and rejection of a torch-based combine:

"The torch combine would work but adds extra elementwise kernels and launches, which defeats the purpose of reducing glue overhead. Even with torch.compile fusion, a dedicated Triton combine kernel is cleaner and guarantees a single launch."

This shows the assistant thinking at multiple levels of abstraction. A naive implementation might use PyTorch operations to combine the splits: torch.max to find the global max, torch.exp for the correction factors, torch.sum for the normalization. But this would launch a dozen small elementwise kernels, adding launch latency and memory traffic. A dedicated Triton kernel fuses all of these operations into a single kernel launch, reading the scratch buffers once and writing the final output directly.

The assistant's rejection of torch.compile as a solution is also telling: "Even with torch.compile fusion" — the assistant is aware that torch.compile could theoretically fuse these operations, but it's already been established (in earlier messages) that torch.compile is incompatible with SGLang's CUDA graph capture mechanism. So the assistant correctly identifies that a hand-written Triton combine kernel is the only viable path.

The Valid Topk Masking Problem

A subtle but important detail appears mid-reasoning:

"The split kernel needs to track valid_topk per batch to mask out invalid tokens beyond that range. Tokens in [valid_topk, topk) get zeroed out during the split, and the combine kernel handles the accumulation across splits."

The sparse-MLA attention mechanism doesn't use all topk positions equally—some are padding. The original kernel handled this with a single mask. In the split-K design, each split needs to know which of its tokens are valid, because a split might contain only padding tokens (if valid_topk falls within its chunk). The assistant correctly identifies that zeroing out invalid tokens in the split kernel and letting the combine kernel handle the accumulation is the correct approach.

Environment Knobs for Tunability

The assistant concludes the reasoning by planning for future tunability:

"I should add environment variable knobs for the target block count, max split count, and minimum chunk size to let users tune this behavior."

This is a mark of production-quality engineering. Rather than hardcoding the occupancy target or chunk size, the assistant plans to expose them as environment variables, allowing future users to tune the kernel for different GPU configurations (e.g., a GPU with fewer SMs might need different NSPLIT targets).

Assumptions Made by the Assistant

Throughout this message, the assistant makes several assumptions that are worth examining critically:

1. The combine kernel's overhead is negligible compared to the split kernels. The assistant assumes that a single Triton combine kernel launch (reading all scratch buffers, computing the LSE merge, writing the final output) adds negligible latency compared to the split kernels' work. This is likely true for large B and NSPLIT, but at B=1 with NSPLIT=16, the combine kernel might be a significant fraction of total time. The assistant doesn't quantify this tradeoff.

2. CUDA graph capture will succeed with the split-combine pattern. The assistant assumes that the split and combine kernels, with their deterministic grid sizes per bucket, will be safely captured by CUDA graphs. This is a reasonable assumption but not guaranteed—CUDA graph capture can fail for subtle reasons (e.g., a kernel using device-side synchronization that isn't supported in graph mode). The assistant acknowledges this implicitly by mentioning that the autotuner runs during warmup, but doesn't consider the possibility that the combine kernel itself might fail graph capture.

3. The autotuner's cached config for one chunk size is "good enough" for others. The assistant acknowledges that the autotuner might pick a suboptimal BLOCK_T config for some chunk sizes but considers this acceptable. This is a pragmatic assumption but could leave performance on the table. A more thorough approach would be to run the autotuner separately for each chunk size, but that would complicate the warmup process.

4. bf16 precision for partial accumulators is acceptable. The assistant plans to use fp32 initially but mentions bf16 as a future optimization. The assumption is that bf16's reduced precision (7 bits of mantissa vs 23 bits for fp32) won't cause numerical issues in the partial accumulation. This is likely true for the attention mechanism's tolerance (the assistant previously validated that bf16-MMA vs fp32-SIMT had relative error ≤ 6.7e-3), but it's an assumption that would need empirical validation.

5. The user's choice to prioritize split-K over torch.compile is correct. The assistant had presented the tradeoff clearly: torch.compile on the 54% glue code is the bigger throughput lever at C=16-64, while split-K primarily fixes the C=1 regression. The user chose split-K. The assistant accepts this without re-arguing the case, but the reasoning section reveals that the assistant internally believes torch.compile would deliver more throughput at the target regime. This is a respectful deference to user direction, but it's worth noting that the assistant's own analysis suggested a different priority.

Potential Mistakes and Incorrect Assumptions

While the assistant's reasoning is generally sound, several areas warrant scrutiny:

The C=1 regression may not be worth fixing. The assistant's own analysis shows that C=1 is a latency edge case, not the throughput target. The user's goal is high-throughput serving at C=16-64, where the MMA kernel already delivers 54-100% improvement. Fixing C=1 with split-K adds complexity (a second kernel, scratch buffers, a combine step) that could introduce bugs and maintenance burden. The assistant's original recommendation to proceed with torch.compile was arguably the better engineering decision for the user's stated goals. The user overruled this, and the assistant dutifully complied, but the cost of this decision is real engineering complexity.

The adaptive NSPLIT formula may interact poorly with CUDA graph bucket boundaries. SGLang captures CUDA graphs for specific batch sizes (e.g., B=1, 2, 4, 8, 16, 32, 64). The adaptive NSPLIT formula produces different NSPLIT values for each B. This means each bucket has a different grid size and different scratch buffer size. While the assistant correctly notes that this is deterministic per bucket, it means the warmup phase must run the autotuner for each bucket separately, potentially adding minutes to server startup time. The assistant doesn't discuss this startup cost.

The combine kernel's memory access pattern may be inefficient. The combine kernel reads all NSPLIT scratch buffers for each (batch, head-group) pair. If NSPLIT is large (e.g., 16 at B=1), the combine kernel reads 16× the output size in scratch data. This could become bandwidth-bound if the scratch buffers are in HBM rather than SRAM. The assistant mentions bf16 as a way to halve this traffic but doesn't analyze whether the combine kernel itself could become a bottleneck.

The assumption that "valid_topk masking in the split kernel" works correctly may be wrong for edge cases. If a split chunk contains only invalid tokens (because valid_topk is small), the split kernel produces zero partials. The combine kernel must handle this correctly—it can't divide by zero or produce NaN outputs. The assistant's reasoning doesn't explicitly address this edge case, though the standard LSE combine formula handles it naturally (l=0, m=-inf, output=0).

Input Knowledge Required to Understand This Message

To fully appreciate message 12562, a reader needs substantial background knowledge:

GPU Architecture: Understanding of SM (Streaming Multiprocessor) occupancy, grid/block launch dimensions, and how block count affects GPU utilization. The assistant's target of 188-376 blocks is meaningless without knowing that the RTX PRO 6000 Blackwell has 188 SMs.

Attention Mechanisms: Knowledge of sparse-MLA (Multi-head Latent Attention), the QK·PV decomposition, topk selection, and the role of the indexer in selecting which KV cache positions to attend to.

Numerical Algorithms: Understanding of the online-softmax / LSE combine trick, where partial attention outputs are merged using per-split maximum scores and normalization factors.

CUDA Graphs: Knowledge of how CUDA graphs capture and replay GPU operations, and the constraints they impose (deterministic grid sizes, fixed memory allocations, no host-device synchronization during replay).

Triton Kernel Programming: Understanding of Triton's tiling model, autotuning mechanism, and how tl.dot maps to tensor-core MMA instructions.

SGLang Architecture: Knowledge of how SGLang uses CUDA graphs for the decode path, how batch size buckets work, and how the attention kernel is dispatched.

DeepSeek-V4 Model Architecture: Understanding of the NVFP4 quantization format, the MLA (Multi-head Latent Attention) mechanism with its nope/rope split, and the 43-layer transformer structure.

A reader lacking any of these knowledge areas would find significant portions of the assistant's reasoning opaque.

Output Knowledge Created by This Message

Message 12562 creates substantial output knowledge, both in the conversation record and in the codebase:

A Design Specification for Split-K: The assistant's reasoning serves as a complete design document for the split-K implementation. It specifies:

The Thinking Process: A Window into Expert Reasoning

The most remarkable aspect of message 12562 is the transparency of the assistant's reasoning. Unlike a human engineer who might think through these issues silently and then write code, the assistant externalizes its entire thought process. This provides an extraordinary window into how an AI agent reasons about complex engineering problems.

Iterative Refinement: The assistant doesn't arrive at the final design in one pass. It starts with a high-level approach (split-K with combine), then iteratively refines details:

  1. First: "split across blocks, combine results"
  2. Then: "need scratch buffers for accumulator, m, l"
  3. Then: "CUDA graph safety means NSPLIT must be deterministic per bucket"
  4. Then: "torch combine would add launches, use Triton combine kernel instead"
  5. Then: "adaptive NSPLIT formula targeting 188-376 blocks"
  6. Then: "minimum chunk size constraint from BLOCK_T"
  7. Then: "valid_topk masking in split kernel"
  8. Then: "bf16 precision tradeoff for scratch buffers"
  9. Then: "environment variable knobs for tunability" This reads like a human engineer thinking through a design, starting broad and drilling into details as each consideration raises new questions. Constraint Propagation: The assistant is particularly skilled at propagating constraints. When it identifies that CUDA graphs require deterministic grid sizes, it immediately traces the implications: NSPLIT must be computed from B (which is fixed per bucket), scratch buffer sizes must be constant per bucket, and the autotuner must run during warmup (not during graph capture). This kind of constraint propagation is characteristic of expert system design. Tradeoff Articulation: The assistant consistently articulates tradeoffs rather than presenting single solutions. fp32 vs bf16 for scratch buffers, torch combine vs Triton combine, large NSPLIT vs large chunk size—each choice is presented with its pros and cons. This allows the user (or a future reader) to understand why a particular choice was made. Self-Correction: At one point, the assistant considers using torch for the combine step, then immediately corrects itself: "The torch combine would work but adds extra elementwise kernels and launches, which defeats the purpose of reducing glue overhead." This self-correction shows the assistant evaluating its own ideas and rejecting suboptimal ones—a hallmark of mature reasoning.

The Broader Engineering Philosophy

Beyond the technical details, message 12562 embodies several engineering principles worth highlighting:

Respect for User Direction: The assistant clearly believes torch.compile is the higher-impact next step, but when the user chooses split-K, the assistant commits fully to that path. The reasoning section doesn't contain any passive resistance or re-argument—it's a genuine design effort for the chosen approach. This is a model of how to handle technical disagreement in engineering: state your case, accept the decision, and execute wholeheartedly.

Production-First Thinking: The assistant consistently considers production constraints: CUDA graph safety, autotuner interaction, memory budgets, startup time, tunability via environment variables. This isn't research code—it's code designed to run in a live serving environment.

System-Level Optimization: The assistant never optimizes in isolation. When considering the scratch buffer memory footprint, it calculates the total across all 43 layers. When designing the combine kernel, it considers the interaction with the existing CUDA graph infrastructure. This system-level thinking prevents local optimizations that create global problems.

Documentation as Thinking: The assistant's reasoning section serves as living documentation. A future engineer reading this code will understand not just what the split-K implementation does, but why it was designed that way, what alternatives were considered, and what tradeoffs were made. This is far more valuable than code comments alone.

Conclusion

Message 12562 is a masterclass in GPU kernel design reasoning. It demonstrates how an expert engineer (or AI agent) approaches a complex optimization problem: starting from first principles, iteratively refining the design, propagating constraints, articulating tradeoffs, and always keeping the production deployment in mind.

The split-K implementation itself is a classic technique—parallelizing over the token dimension to improve occupancy—but the assistant's reasoning reveals the depth of consideration required to apply it correctly in a production setting. Every decision, from the choice of combine kernel to the adaptive NSPLIT formula to the environment variable knobs, is grounded in a clear understanding of the hardware, the software stack, and the deployment constraints.

For the reader, this message offers a rare opportunity to observe expert-level reasoning in real time. It's not just about split-K; it's about how to think about GPU kernel optimization systematically. The assistant's transparent reasoning process—its self-corrections, its tradeoff analyses, its constraint propagation—provides a template that human engineers can learn from and emulate.

In the end, whether split-K was the "right" choice compared to torch.compile is almost beside the point. What matters is that the assistant committed to the chosen path with full intellectual honesty, designed a robust solution, and documented its reasoning for posterity. That is engineering at its best.