The Shared Memory Ceiling: A Moment of Architectural Insight in Custom CUDA Kernel Development

Introduction

In the high-stakes world of custom GPU kernel engineering, the difference between a working benchmark and a silent crash often comes down to a single number: the shared memory limit. This article examines a pivotal message in an opencode coding session where an AI assistant, deep in the process of building a custom flash-attention verify kernel for NVIDIA's Blackwell RTX PRO 6000 GPUs (sm_120 architecture), pauses to add a critical guard before syncing code to the target server. The message at index 12234 captures a moment of architectural realization—the naive oracle kernel, used as the ground-truth reference for correctness testing, cannot physically run at the longer prefix lengths the flash kernel is designed to handle. This insight, and the guard it spawns, transforms a potential benchmark crash into a deliberate demonstration of the new kernel's superior capability.

The Broader Context: Building a Custom sm_120 Verify Kernel

To understand this message, one must appreciate the project's scope. The assistant is developing a custom CUDA implementation of a "verify attention" kernel—a critical component in speculative decoding pipelines (specifically DFlash with DDTree) for large language models like Kimi K2.6. The target hardware is the RTX PRO 6000 Blackwell GPU, which uses the sm_120 architecture. This architecture is a consumer/professional variant of Blackwell, distinct from the datacenter-grade Blackwell (sm_100a/sm_103a) that benefits from advanced features like TMA (Tensor Memory Accelerator) and wgmma instructions.

The project had already discovered that all optimized MLA (Multi-head Latent Attention) kernels—FlashMLA, cutlass-MLA, flashinfer-MLA—are compiled only for sm_90a, sm_100a, and sm_103a. None support sm_120. This forced the team to build an owned kernel from scratch, documented in the plan at plans/0002-sm120-verify-kernel-defrag.md. The kernel's core innovation is a "flash" design that tiles the KV cache into shared memory, amortizing the expensive global memory reads of the MLA latent across multiple heads. Where the naive kernel reads the KV cache once per head (64 reads total), the flash kernel reads it once per head-tile (8 reads with HT=8), yielding an 8× reduction in global memory traffic—and potentially much more when combined with coalesced access patterns and online softmax.

The Message: Reasoning Before Deployment

The subject message begins with the assistant's internal reasoning:

I need to sync the repository to CT200 and run the build and tests there. Let me figure out the best way to transfer the files—I'll use rsync over SSH to push the source directories while excluding the build artifacts and large reference files. I'm keeping GPU memory constraints in mind as I prepare to run this.

This opening reveals the assistant's immediate context: the code has been written locally, and the next step is deployment to the CT200 server where the RTX PRO 6000 GPUs live. The mention of "keeping GPU memory constraints in mind" signals a disciplined approach—the assistant is proactively thinking about resource limits before running anything.

The reasoning then dives into a detailed memory budget analysis:

Looking at the memory situation, the service is consuming about 94GB per GPU with only 4GB free. The kernel tests use small shapes, so they should be fine, but the benchmarks go up to prefix 65536 which could be tight. Let me work through the largest benchmark allocation: with 4 streams, prefix 65536, and the various tensor allocations, I'm looking at around 673MB total, which fits within the 4GB headroom.

This is a careful accounting of GPU memory. The production SGLang service is already consuming 94GB of the available memory, leaving only 4GB of headroom. The assistant mentally computes the worst-case benchmark allocation—673MB—and confirms it fits. But then comes the crucial insight:

However, there's a critical limitation—the naive kernel requires full-score-row in shared memory, and at prefix 65536 that would need 262KB, which exceeds the 100KB shared memory maximum. So the naive kernel will actually fail to launch at prefix 16384 and beyond due to this shared memory constraint.

This is the moment of architectural realization. The naive oracle kernel—the reference implementation used to validate correctness—stores the full attention score row in shared memory. For a prefix of 65,536 tokens, with each score stored as a 4-byte float, that's 256KB of shared memory. But sm_120's shared memory limit is approximately 100KB (the exact figure varies by configuration, but the assistant uses 99,000 bytes as the guard threshold). The naive kernel simply cannot launch for prefixes beyond roughly 25,000 tokens.

The Guard: Turning a Limitation into a Demonstration

The assistant's response is pragmatic and elegant:

I need to guard the naive kernel launch by checking if the shared memory requirement exceeds the limit. I'll compute the naive shared memory as (Dl+Dr+kv_len)*4 bytes, and if that exceeds around 99000 bytes, I'll skip running the naive kernel and show "n/a" instead. This way the benchmark will gracefully handle the limitation and actually demonstrate that flash can handle long prefixes where the naive kernel can't.

This transforms a potential failure into a feature. Instead of crashing or producing misleading results, the benchmark will explicitly show "n/a" for the naive kernel at long prefixes, while the flash kernel continues to produce valid results. This directly demonstrates the flash kernel's key advantage: its online softmax algorithm never stores the full score row, instead maintaining only the running maximum and normalization sum, which are constant-size regardless of prefix length.

The actual edit applied is:

[edit] /home/theuser/glm-kimi-sm120-rtx6000bw/kdtree-engine/tests/bench_kernels.cu
Edit applied successfully.

The edit adds a computed guard: (Dl + Dr + kv_len) * 4 > 99000. If true, the naive kernel is skipped and "n/a" is reported. The flash kernel, which uses online softmax and tiles the KV cache, has no such limitation and runs at all prefix lengths.

Technical Deep Dive: Why Shared Memory Matters

The shared memory limit is a fundamental architectural constraint on GPUs. Each streaming multiprocessor (SM) has a fixed pool of on-chip memory that can be configured as L1 cache, shared memory, or a split between the two. On sm_120 (RTX PRO 6000 Blackwell), the maximum shared memory per block is approximately 100KB. This is generous compared to older architectures (e.g., 48KB on Maxwell, 64KB on Pascal), but it's still a hard ceiling.

The naive verify attention kernel computes attention scores for all KV positions simultaneously. For each query position, it computes a dot product between the query and every key, producing a score vector of length kv_len. These scores must be stored temporarily to compute the softmax normalization and then the weighted sum of values. The straightforward implementation allocates this score row in shared memory because shared memory is fast (low latency, high bandwidth) and accessible by all threads in a block.

The problem is that the score row size scales linearly with kv_len. At 65,536 tokens, the score row alone is 256KB (65,536 × 4 bytes for float32). Add the key and value cache slices (each 512-dimensional, so 2KB per tile), the query vectors, and the output accumulators, and the total far exceeds 100KB.

The flash kernel solves this with online softmax, also known as "streaming softmax" or "tiled softmax." Instead of storing all scores, it processes the KV cache in tiles. For each tile, it computes scores for just that tile's keys, then immediately updates a running maximum and normalization sum. The key insight is that the softmax of the full set can be computed incrementally: when a new tile's maximum exceeds the running max, all previously accumulated values are rescaled by the difference. This requires only two scalars per head (max and sum) plus the output accumulator vector—all constant-size regardless of kv_len.

Assumptions and Correctness

The assistant makes several assumptions in this message, all of which are sound:

  1. The naive kernel's shared memory requirement is (Dl + Dr + kv_len) * 4 bytes. This assumes the naive kernel stores the full score row plus the key and value dimensions in shared memory. This is consistent with the oracle implementation seen earlier in the session.
  2. The shared memory limit is approximately 99,000 bytes. This is a conservative threshold for sm_120's ~100KB limit, leaving a small margin.
  3. The flash kernel has no such limitation. This is correct because online softmax eliminates the kv_len-scaling allocation.
  4. Skipping the naive kernel and showing "n/a" is better than crashing. This is a design judgment that prioritizes graceful degradation and informative output over silent failure. One subtle point: the guard formula (Dl + Dr + kv_len) * 4 may not capture all shared memory allocations in the naive kernel. There could be additional temporary buffers for partial reductions, position embeddings, or mask storage. However, the dominant term is kv_len * 4, and the guard is intentionally conservative (99000 < 100000), so it likely catches all over-limit cases. If the naive kernel has additional shared memory usage beyond the score row, the guard might trigger earlier than strictly necessary, but this is a safe over-approximation.

Input and Output Knowledge

Input knowledge required to understand this message includes:

The Thinking Process: A Window into Engineering Discipline

What makes this message remarkable is the thinking process it reveals. The assistant is not simply executing a pre-planned sequence of steps. It is actively reasoning about constraints, computing budgets, and identifying failure modes before they occur. The chain of thought flows naturally:

  1. "I need to sync to CT200" → practical next step
  2. "Let me consider GPU memory" → proactive constraint checking
  3. "The service uses 94GB, 4GB free" → situational awareness
  4. "Benchmark allocations are 673MB" → quantitative budget
  5. "But wait—the naive kernel needs score row in shared memory" → architectural insight
  6. "At 65536 prefix, that's 262KB" → quantitative realization
  7. "That exceeds the 100KB limit" → identification of failure
  8. "I'll add a guard to skip and show n/a" → elegant solution
  9. "This demonstrates flash can handle lengths the oracle can't" → reframing limitation as advantage This is a textbook example of defensive engineering: anticipating failure modes, understanding root causes, and implementing graceful degradation. The assistant could have simply synced the code, run the benchmarks, and discovered the crash at runtime. Instead, it invested a few minutes of reasoning to prevent the crash entirely and turn the situation into a positive demonstration.

Conclusion

The message at index 12234 captures a moment of genuine engineering insight. In the midst of deploying a custom CUDA kernel for speculative decoding on Blackwell GPUs, the assistant identifies a fundamental architectural constraint—the shared memory ceiling—that would cause the naive reference kernel to fail at long prefix lengths. Rather than letting this become a runtime crash or a confusing benchmark artifact, the assistant adds a guard that gracefully skips the naive kernel and reports "n/a," transforming the limitation into a demonstration of the flash kernel's superior scalability.

This moment illustrates a broader truth about GPU kernel development: the difference between a working system and a failing one often comes down to understanding architectural limits. Shared memory, register pressure, warp occupancy, memory bandwidth—these are the invisible constraints that shape what is possible. The assistant's disciplined reasoning—computing budgets, checking limits, and designing for graceful failure—is exactly the kind of thinking that separates robust engineering from fragile hacking. In the end, the flash kernel would go on to achieve 3–6× decode speedups over the Triton baseline, but this message captures the moment when the assistant ensured those benchmarks would actually run correctly.