The Moment of Truth: Analyzing a Custom MMA Kernel's Performance on Blackwell GPUs

Introduction

In the high-stakes world of large language model inference on cutting-edge hardware, few moments are as revealing as the first benchmark run of a custom GPU kernel. After days of designing, implementing, and debugging a custom Matrix Multiply-Accumulate (MMA) sparse-MLA decode kernel for the DeepSeek-V4-Flash model on 8× NVIDIA RTX PRO 6000 Blackwell (sm_120) GPUs, the assistant in this opencode session finally had performance numbers in hand. Message [msg 12558] captures the critical analysis phase: the assistant interprets benchmark results, diagnoses a subtle occupancy problem, and plans the next optimization step — all while running a fresh profile to confirm the bottleneck has shifted.

This message is a masterclass in performance engineering reasoning. It demonstrates how a kernel developer reads between the lines of raw throughput numbers, connects observed regressions to their root causes in GPU occupancy and kernel launch geometry, and makes data-driven decisions about where to invest optimization effort next. The message is also notable for what it reveals about the assistant's mental model of the system: a deep understanding of Triton's compilation model, CUDA graph capture constraints, and the interaction between batch size, grid dimensions, and SM utilization on Blackwell's 188 streaming multiprocessors.

The Context: A Long Optimization Campaign

To understand the significance of this message, we must situate it within the broader engineering effort. The assistant had been working for days — across multiple segments and dozens of messages — to optimize DeepSeek-V4-Flash inference on Blackwell GPUs. The journey began with a sobering discovery: the production SGLang server was hitting CUDA-core fallback kernels for the attention mechanism, achieving only ~28 tokens/second on a model that should have been capable of hundreds. The root cause was structural — the sparse-MLA (Multi-head Latent Attention) decode kernel was using per-head SIMT (Single Instruction, Multiple Thread) operations that re-read the KV cache 64× redundantly, and was forced into FP32 precision for key operations like the indexer bmm and MHC-pre linear projection.

The assistant's response was bold: design and implement a custom MMA kernel using Triton's tl.dot tensor-core operations, replacing the redundant per-head SIMT approach with a head-batched strategy that shares a single KV cache read across all 64 query heads. This was the subject of messages [msg 12543] through [msg 12557], culminating in a successful deployment and benchmark run. By message [msg 12557], the benchmark harness (sweep_scale.sh) had been executed with the MMA configuration, producing raw throughput numbers that the assistant now needed to interpret.

The Core Analysis: Reading Between the Throughput Numbers

The subject message opens with the assistant's reasoning block, which immediately engages with the benchmark data:

Looking at the performance metrics across different concurrency levels, the MMA approach shows significant improvements over the old SIMT implementation. At C=16, we're seeing a 54% boost in throughput and a 66% improvement in steady-state performance, with gains continuing to scale up at higher concurrency levels.

This is not a casual glance at numbers. The assistant is performing a structured comparison, noting both the headline improvements and the troubling regression at low concurrency. The assistant then constructs a concise performance table that crystallizes the tradeoff:

| C | old SIMT (out tok/s) | MMA | Δ | |---|---|---|---| | 1 | 11.5 | 8.1 | −29% ⚠️ | | 16 | 26.6 | 41.0 | +54% | | 64 | 29.7 | 59.6 | +100% |

This table is the analytical heart of the message. It reveals a clear pattern: the MMA kernel excels at higher concurrency (C=16 and C=64) but regresses at single-batch workloads. The assistant immediately identifies the root cause with surgical precision:

The C=1 regression is exactly the predicted under-occupancy: head-batching makes the grid (B, 2) = only 2 blocks at B=1 (vs 64 before) on 188 SMs.

This diagnosis is the key insight. The old SIMT kernel launched 64 blocks (one per head) regardless of batch size, providing ample work for the GPU's 188 SMs even at single-batch. The new MMA kernel, by sharing KV cache reads across all heads, reduces the grid to just (B, 2) — at B=1, that's only 2 blocks on 188 SMs, a catastrophic 1% occupancy. The GPU is essentially idle 99% of the time, waiting for those two blocks to complete before it can launch more work.

The Occupancy Problem: A Deeper Dive

The assistant's reasoning reveals a nuanced understanding of GPU occupancy and its implications. It notes that even at C=16, the grid is only (16, 2) = 32 blocks, which represents just ~17% occupancy on 188 SMs. This means the throughput regime (C=16 and C=64) is also under-occupied — it just happens that the MMA kernel's superior arithmetic intensity per block compensates enough to show a net gain.

The marginal cost analysis is particularly insightful:

Marginal per-request decode cost dropped from ~30 ms → ~13 ms (asymptote ~33 → ~77 t/s).

This tells us that the MMA kernel's per-request latency has more than halved, from 30 milliseconds to 13 milliseconds. The asymptotic throughput (the theoretical maximum as concurrency approaches infinity) has more than doubled, from ~33 to ~77 tokens/second. The C=1 regression is purely an occupancy artifact — if the assistant can fix the occupancy problem, the MMA kernel should dominate across all concurrency levels.

The Split-K Solution: Reasoning About the Fix

Having diagnosed the problem, the assistant immediately begins reasoning about the solution: split-K parallelization over the topk dimension. The key insight is that the attention computation involves a "topk" dimension (the number of KV cache entries attended to per token), which can be split across multiple blocks. Instead of having 2 blocks compute the full attention over all topk entries, the assistant proposes splitting the topk dimension into NSPLIT chunks, with each block computing a partial result over its assigned chunk. A subsequent combine step reduces the partial results into the final output.

The assistant's reasoning about CUDA graph constraints is particularly sophisticated:

The challenge is that CUDA graphs need fixed grids per batch bucket, so NSPLIT has to be determined at capture time based on the batch size.

This reveals a deep understanding of SGLang's CUDA graph capture mechanism. CUDA graphs record GPU operations (kernel launches, memcpys, etc.) into a replayable graph, which eliminates CPU launch overhead for repeated operations. However, the graph captures fixed grid dimensions — you cannot dynamically change the number of blocks between captures. The assistant recognizes that NSPLIT must be computed at capture time (which occurs during server warmup for each batch size bucket) and baked into the graph.

The assistant also considers the scratch buffer allocation challenge:

For CUDA graph safety, the scratch buffer allocation is tricky — torch.zeros() inside the captured region gets recorded as a fixed pointer from the caching allocator, which works as long as the size is deterministic per batch bucket.

This is a subtle point: within a CUDA graph capture region, memory allocations must be deterministic. If torch.zeros() allocates a new buffer each time the graph is replayed, the graph would break because the recorded pointer would be stale. However, PyTorch's caching allocator tends to return the same buffer for the same size, so as long as NSPLIT is deterministic per batch bucket, the allocation should be safe.

The Decision: Profile Before Optimizing

Despite having a clear plan for split-K, the assistant exercises engineering discipline:

Before implementing split-K, I should re-profile the MMA kernel at C=32 to see if attention is still the bottleneck or if the glue code has become dominant.

This is a critical decision. The assistant recognizes that the optimization landscape may have shifted. The MMA kernel dramatically reduced attention time (from ~57% of decode GPU time to ~10%), which means other operations — the "glue code" of elementwise operations, copies, and FP32 SIMT GEMMs — may now be the dominant bottleneck. If the glue code has become the primary consumer of GPU time, optimizing attention further with split-K would yield diminishing returns. The assistant needs fresh profiling data to confirm the bottleneck before investing in implementation.

The choice of C=32 is strategic. At C=32 with a (32, 2) = 64 block grid on 188 SMs, occupancy is only ~34% — still under-occupied, but high enough that the attention kernel's share of GPU time should be measurable. If attention is still significant at this occupancy level, split-K will help. If the glue code dominates even at 34% occupancy, the assistant may need to pivot to a different optimization strategy (such as torch.compile for the glue operations, as the user had suggested earlier in [msg 12557]).

The Profiling Execution

The message concludes with the assistant executing the profiling command:

timeout 120 ssh -o StrictHostKeyChecking=no root@10.1.230.171 'bash /root/prof_steady.sh 2>&1 | tail -3'

and checking for the output trace file:

timeout 20 ssh -o StrictHostKeyChecking=no root@10.1.230.171 'ls -t /root/prof/*TP-0.trace.json.gz | head -1'

The prof_steady.sh script runs a steady-state profiling workload at C=32, capturing a GPU trace via NVIDIA's profiling tools. The output is a .trace.json.gz file that the assistant will analyze in subsequent messages to determine the new bottleneck distribution.

Assumptions and Potential Pitfalls

The assistant's reasoning in this message rests on several assumptions that deserve scrutiny:

Split-K will help across the throughput regime, not just C=1: The assistant assumes that the under-occupancy at C=16 (17%) and C=64 (34%) is significant enough that split-K will yield measurable gains. This is plausible but not guaranteed — at C=64, the grid is already 128 blocks, which may be sufficient to saturate the GPU's memory bandwidth. Split-K introduces additional overhead (scratch buffer writes and a combine kernel) that could offset the occupancy gains.
CUDA graph capture is compatible with split-K: The assistant assumes that the split-K combine step can be captured in a CUDA graph. This requires the scratch buffer allocation to be deterministic and the combine kernel to have a fixed grid. Both conditions are likely satisfiable, but the assistant has not yet validated this.
The glue code is not the bottleneck: The assistant's decision to profile before optimizing split-K implicitly assumes that attention is still a significant bottleneck. If the profile reveals that glue code dominates (as the user suspected in [msg 12557]), the assistant will need to pivot, potentially wasting the profiling effort.
NSPLIT targeting ~188 blocks is optimal: The assistant proposes scaling NSPLIT to achieve ~188 blocks (matching the SM count). This is a reasonable heuristic but may not be optimal — factors like shared memory usage, register pressure, and the combine kernel overhead could shift the optimal NSPLIT value.

Input Knowledge Required

To fully understand this message, the reader needs knowledge of:

GPU occupancy: The ratio of active warps to maximum possible warps on a streaming multiprocessor. Low occupancy means the GPU cannot hide latency through thread-level parallelism.
CUDA grid/block model: How GPU kernels are launched as grids of blocks, and how grid dimensions affect occupancy.
Triton compiler: The Python-based DSL for writing GPU kernels. The assistant uses Triton's tl.dot for tensor-core operations and tl.arange for index computation.
CUDA graph capture: A mechanism for recording GPU operations into a replayable graph, eliminating CPU launch overhead. Requires deterministic memory allocations and fixed grid dimensions.
Sparse MLA (Multi-head Latent Attention): The attention mechanism used by DeepSeek-V4, which attends to a sparse set of KV cache entries per token.
Split-K parallelization: A technique for parallelizing attention over the key dimension by splitting the computation into independent chunks and combining results.
Blackwell SM count: The RTX PRO 6000 Blackwell has 188 SMs, which determines the target grid size for full occupancy.

Output Knowledge Created

This message produces several valuable outputs:

A clear performance comparison between the old SIMT kernel and the new MMA kernel across three concurrency levels, establishing the tradeoff space.
A root-cause diagnosis of the C=1 regression: head-batching reduces grid size from 64 blocks to 2 blocks, causing catastrophic under-occupancy.
A marginal cost analysis: per-request decode cost dropped from ~30ms to ~13ms, with asymptotic throughput rising from ~33 to ~77 tok/s.
A design plan for split-K: adaptive NSPLIT targeting ~188 blocks, with CUDA graph constraints considered.
A profiling execution at C=32 to determine whether attention or glue code is the current bottleneck.

The Thinking Process: A Window into Performance Engineering

The assistant's reasoning in this message reveals a structured approach to performance engineering that is worth examining in detail.

Step 1: Data Collection and Organization. The assistant begins by extracting the raw benchmark numbers from the previous message ([msg 12557]) and organizing them into a comparison table. This is not merely transcription — the assistant is selecting the relevant metrics (output token throughput) and computing the delta between old and new kernels.

Step 2: Pattern Recognition. The assistant identifies the key pattern: gains at high concurrency, regression at low concurrency. This pattern immediately suggests an occupancy problem, because occupancy effects are most pronounced at low batch sizes where the grid is small.

Step 3: Root-Cause Analysis. The assistant connects the pattern to the kernel's launch geometry. The old kernel launched 64 blocks (one per head), providing full occupancy even at B=1. The new kernel launches (B, 2) blocks, which at B=1 is only 2 blocks on 188 SMs. The diagnosis is precise and quantitative.

Step 4: Solution Design. Having identified the root cause, the assistant designs a solution: split-K over the topk dimension. The reasoning shows an understanding of how split-K interacts with CUDA graph constraints and memory allocation.

Step 5: Verification Before Action. Crucially, the assistant does not immediately implement split-K. Instead, it decides to profile first, to confirm that attention is still the bottleneck. This is the mark of a disciplined engineer who understands that optimization landscapes shift, and that yesterday's bottleneck may not be today's.

Step 6: Execution. The assistant executes the profiling command and checks for the output file, setting up the data needed for the next decision.

Conclusion

Message [msg 12558] captures a pivotal moment in a complex optimization campaign. The assistant has deployed a custom MMA kernel that delivers 2× throughput at high concurrency but regresses at single-batch workloads. Rather than celebrating the win or rushing to fix the regression, the assistant pauses to analyze, diagnose, and plan. The decision to profile before optimizing is a textbook example of data-driven performance engineering.

The message also reveals the assistant's deep technical knowledge: understanding of GPU occupancy, CUDA graph constraints, Triton compilation, and split-K parallelization. But more importantly, it reveals a structured thought process that separates great engineers from good ones: collect data, identify patterns, diagnose root causes, design solutions, verify assumptions, then act.

The profiling data that the assistant is collecting will inform the next phase of optimization. Will split-K fix the C=1 regression and boost throughput across the board? Or will the profile reveal that the glue code has become the new bottleneck, forcing a pivot to torch.compile or manual fusion? The answer awaits in the subsequent messages, but the analytical framework established in this message ensures that whatever path the assistant takes, it will be grounded in data and sound reasoning.