The Float4 Vectorization: A Single Edit That Unlocked 3–6× Decode Throughput on Blackwell

Introduction

In the course of building a custom CUDA verify-attention kernel for speculative decoding with the Kimi K2.6 model on NVIDIA RTX PRO 6000 Blackwell GPUs (sm_120 architecture), a single edit to a kernel source file marked the culmination of an extensive reasoning chain that transformed a promising but underperforming implementation into a production-ready component delivering 3–6× end-to-end decode speedup over the existing Triton baseline. The message in question — message 12245 — is deceptively simple:

[edit] /home/theuser/glm-kimi-sm120-rtx6000bw/kdtree-engine/src/kernels/verify_attn_flash.cu Edit applied successfully.

This terse confirmation, barely a line of text, represents the final step in a multi-stage optimization journey spanning several rounds of architectural redesign, occupancy analysis, bandwidth calculation, and production-context reasoning. To understand the significance of this edit, one must trace the reasoning chain that led to it — a chain that reveals deep insights into GPU kernel optimization, the tension between occupancy and memory efficiency, and the importance of considering deployment topology when optimizing.

The Journey to v6: From Occupancy-Starved to Bandwidth-Bound

The story begins several messages earlier, when the assistant's initial flash-attention kernel (v5) achieved only 1.6–1.9× speedup over the naive baseline, despite incorporating KV-splitting and head-tile amortization. The assistant performed a detailed bottleneck analysis in message 12242, working through the theoretical math: ~1×10¹⁰ FMAs total, which at ~90 TFLOP fp32 should take only 0.22ms at peak compute — yet the kernel was taking 11.6ms. The culprit was not memory bandwidth but occupancy: with 70KB of shared memory per block, only one block could fit per streaming multiprocessor (SM), and with only 128 threads per block, the SM achieved just 6% occupancy — a starvation of warps that left every memory stall blocking the entire SM.

The assistant correctly identified that for decode — which is latency-bound rather than throughput-bound — maximizing occupancy matters far more than minimizing redundant latent reads. This insight drove a complete architectural pivot to a per-head, small-shared-memory design (v6) with approximately 23KB of shared memory per block, allowing 4 blocks per SM and 67% occupancy. The design used warp-per-key dot products to keep all 256 threads active during the score computation phase, and KV-splitting to handle arbitrary sequence lengths.

When v6 was tested, the results were encouraging: 1.3–1.5× faster on short sequences and 2.7–3.3× faster at 16384-token prefixes, with the ability to handle 65536 tokens that the naive kernel could not. But the assistant's analysis in message 12244 revealed a new bottleneck: the kernel was now memory-throughput-bound. With 9 splits creating 5184 blocks and high occupancy (32 warps per SM), the achieved bandwidth was only 403 GB/s — roughly 22% of the 1.8 TB/s peak of the RTX PRO 6000 Blackwell GPU. The scalar float loads used to read KV cache tiles were coalesced but not fully utilizing the memory subsystem's capacity.

The Float4 Decision: Why This Edit Matters

The edit in message 12245 applies float4 vectorized loads to the KV tile and query data — the dominant memory traffic in the kernel. By replacing four separate 32-bit scalar loads with a single 128-bit vectorized load, the kernel can better utilize the memory bus width and reduce the number of load instructions, improving achieved bandwidth. This is a well-known GPU optimization pattern: on NVIDIA architectures, 128-bit (float4) loads can be coalesced into a single memory transaction, whereas four scalar loads may require separate transactions even when they target contiguous addresses.

The assistant's reasoning in message 12244 reveals a crucial production-context insight: in the TP8 (tensor parallelism across 8 ranks) deployment scenario, each rank processes only 8 heads instead of the full 64. This means the per-head design reads the KV cache only 8 times per rank rather than 64 — an 8× reduction in memory traffic compared to the single-GPU benchmark. The float4 vectorization therefore provides benefits in both scenarios: it improves bandwidth utilization in the benchmark (where the full 64-head read is the bottleneck) and further reduces latency in production (where the reduced head count already makes memory traffic less dominant).

Assumptions and Correctness

The assistant made several assumptions that proved correct. First, that the v6 per-head design's high occupancy would outweigh the cost of redundant KV cache reads — benchmark results confirmed this. Second, that the kernel was now memory-throughput-bound rather than occupancy-bound — the 22% bandwidth utilization confirmed this diagnosis. Third, that float4 vectorization would be a "cheap, high-impact win" — the chunk summary confirms that this optimization, combined with the NSPLIT increase from 16 to 64, delivered the dramatic 3–6× end-to-end speedup.

One assumption that the assistant explicitly reconsidered was the relevance of the single-GPU benchmark to production performance. Initially, the 5.96ms execution time at 16384 tokens seemed disappointing relative to the 1.8 TB/s peak bandwidth. But the assistant recognized that the benchmark represented a worst case (64 heads, no tensor parallelism), and that in TP8 production the per-head design would naturally read 8× less KV data. This correct interpretation prevented premature abandonment of a well-suited architecture.

Input and Output Knowledge

To understand this message, one needs knowledge of: CUDA memory hierarchy and the distinction between scalar and vectorized loads; the sm_120 architecture of the RTX PRO 6000 Blackwell GPU and its 1.8 TB/s memory bandwidth; the MLA (Multi-head Latent Attention) structure of the Kimi K2.6 model with its 64 heads and 512-dimensional latent; the concept of tensor parallelism (TP8) and how it distributes heads across ranks; and the occupancy-latency tradeoff in GPU kernel design.

The output knowledge created by this edit is a verify-attention kernel that: (a) is token-exact correct against the naive oracle, (b) handles arbitrary sequence lengths up to at least 65536 tokens, (c) achieves 3–6× decode throughput over Triton in production with CUDA graphs enabled, (d) is capture-safe for CUDA graph replay, and (e) operates within the sm_120 architecture's constraints (no wgmma/TMA/tcgen05 instructions available on consumer Blackwell). This kernel became the centerpiece of a deployed SGLang service with 200k context-length support, Tier 0 KV defragmentation active, and the remaining bottleneck cleanly identified as MoE expert imbalance — a structural ceiling beyond the verify kernel's scope.

The Broader Significance

Message 12245 is a testament to the value of systematic, evidence-driven optimization. The assistant did not jump to float4 vectorization as a first step; it arrived there only after: (1) identifying occupancy as the primary bottleneck through theoretical analysis, (2) redesigning the kernel architecture to maximize occupancy, (3) confirming through benchmarks that the new design was memory-throughput-bound, and (4) selecting the most impactful optimization for the current bottleneck. Each step was validated with empirical measurements, and incorrect assumptions (such as the relevance of the single-GPU benchmark) were corrected through careful reasoning about the production deployment topology.

This single edit — the addition of float4 vectorized loads — represents the moment when weeks of architectural iteration, kernel debugging, and performance analysis crystallized into a production-ready component. It is a reminder that in GPU kernel optimization, the most impactful changes are often not the most complex ones, but rather the ones that correctly target the current bottleneck after the preceding bottlenecks have been systematically eliminated.