FP16 Accumulators for GRU-L3: Block-Diagonal Recurrence, Stochastic Rounding, and the Precision Ladder

The Accumulator Problem

Our GRU-L3 V3a model runs at 75.9M parameters on a single RTX 3080 8GB. The recurrent hidden state is the heart of the model — it carries information across timesteps, and every step adds to it. In FP32, this works fine. In FP16, it diverges.

The problem is not the GEMM. Our Vulkan gemm_mixed.comp shader already uses FP16 tensor core inputs with an FP32 accumulator via GL_KHR_cooperative_matrix — the cooperative matrix multiply-add loads FP16 operands into tensor core tiles and accumulates into coopmat<float, ...>. That part is fast and accurate. The problem is the elementwise GRU gate update that runs after the GEMM:

This is a running sum. Each timestep's h_new becomes the next step's h_prev. In FP16, each step introduces ~2^-10 relative rounding error. Over the V3a model's T=1024 timesteps, those errors compound multiplicatively through the recurrent loop. At the proposed T=4096 context length, the accumulated error can reach the same order as the hidden state itself — the accumfrontier result quantifies this: an FP16 sum of 4096 ones stalls at 2048, because past 2048 the gap between adjacent FP16 numbers exceeds 1 and every +1 rounds away. Unlike a transformer's parallel attention, where each position is computed independently during training forward pass, the GRU's recurrent state is a serial dependency chain — each step's error feeds into the next.

Standard mixed precision — FP16 tensor cores with FP32 storage and accumulation — does not solve this. The GEMM accumulator is FP32, but the recurrent state h_prev is stored as FP32 and the elementwise gates run in FP32. The activation memory dominates VRAM, and activations must be FP32. We tried this. It did not meet the goal.

The Precision Ladder on Ampere

Before designing the fix, we need to understand the compute and memory landscape. The RTX 3080 Ti uses the GA102 die — 80 SMs, 320 third-generation tensor cores, 12 GB GDDR6X at 912 GB/s. Here is the full precision ladder, rated in dense tensor core TFLOPS (theoretical peak from NVIDIA's whitepaper) and bytes per element:

All tensor core figures are dense (no 2:4 sparsity). FP32 runs on CUDA cores, not tensor cores. TF32 is a tensor core compute mode, not a storage type — it occupies 32 bits in memory.

FP16 vs BF16: Same Speed, Different Risk

On Ampere, FP16 and BF16 run at identical tensor core throughput — 68.2 TFLOPS dense, 136.4 TFLOPS with 2:4 structured sparsity. NVIDIA's own whitepaper confirms: "BF16 Tensor Core instructions at the same throughput as FP16." The difference is purely numerical:

For the fastest possible GEMM on the RTX 3080 Ti, the answer is clear: FP16 or BF16 with FP32 accumulation, at 68.2 TFLOPS dense. That is 2x the throughput of TF32 or FP32, at half the memory cost. The question is not which format is fastest — they are tied — but which format can be made numerically safe for the recurrent accumulator.

Why TF32 Does Not Help Here

TF32 looks attractive on paper: FP32 range with FP16 precision, no overflow risk. But it has three problems for our use case:

1. Half the throughput. TF32 runs at 34.1 TFLOPS dense on the RTX 3080 Ti's tensor cores — exactly half of FP16/BF16 (68.2 TFLOPS). TF32 and FP16 both have 10 mantissa bits, but TF32 is stored in 32-bit words (FP32 format with truncated mantissa). The tensor core's data path processes half as many TF32 elements per cycle as FP16 elements because each TF32 element occupies twice the register space. For a memory-bound RNN where the GEMM is already 99% of compute, halving the GEMM throughput halves the training speed.

2. No memory savings. TF32 is a compute mode, not a storage type. Inputs and outputs are stored as FP32 (4 bytes). The activation memory — which dominates VRAM for long sequences — stays at FP32 size. We get no VRAM reduction from TF32.

3. No Vulkan cooperative matrix support. TF32 is a CUDA-only compute mode exposed through cuBLAS and cuDNN. The Vulkan cooperative matrix extension (GL_KHR_cooperative_matrix) does not define a TF32 component type. Our renderer uses Vulkan compute shaders, not CUDA. TF32 is simply not available to us.

The Vulkan Constraint

Our GRU-L3 trainer runs on Vulkan compute shaders, not CUDA. This constrains the available precision paths:

Our existing gemm_mixed.comp shader uses coopmat<float16_t, gl_ScopeSubgroup, TILE_M, TILE_K, gl_MatrixUseA> for the input operands and coopmat<float, gl_ScopeSubgroup, TILE_M, TILE_N, gl_MatrixUseAccumulator> for the accumulator. This is the FP16-input, FP32-accumulate path — 68.2 TFLOPS on the RTX 3080 Ti. BF16 would use the same pattern with bfloat16_t inputs, but requires the newer GL_EXT_bfloat16 extension and driver support validation.

The Four-Part Architecture

The goal is pure FP16 accumulators — no FP32 in the recurrent state. This requires four architectural changes that reduce the error budget so FP16's 10-bit mantissa becomes sufficient:

The Error Budget

With all four changes, the FP16 accumulator error stays bounded. Here is the math:

D=4096, d=d_head=64, T=4096. ULP for FP16 is 2^-10; for paired-FP16 (22 effective bits) it is 2^-22. Coherent error (round-to-nearest) grows as O(T); stochastic rounding reduces this to O(√T) with zero mean.

The block-diagonal split is the highest-impact change — it reduces the GEMM accumulation length by 64x, which is the single biggest contributor to FP16 error. Stochastic rounding converts the remaining coherent bias into zero-mean noise. Paired-FP16 gives enough effective precision to match FP32. Pre-recurrence RMSNorm keeps the state in the sweet spot.

Compute and Memory Estimates

For the target model — D=4096, L=4, T=4096, B=32, d_head=64, N_heads=64 — the per-step FLOP budget is:

1 MAD = 2 FLOPs. Backward is ~2x forward for GEMM. Per step = per layer per timestep × B × T × L.

At 68.2 TFLOPS theoretical peak (FP16 tensor core, FP32 accumulator), the estimated step time is 161.1 / 68.2 = 2.36s at peak. Real throughput is typically 50% of peak for memory-bound kernels — so ~4.7s per step. The FP32 baseline (dense W_hh, 34.1 TFLOPS CUDA cores) would take ~317.0 / 34.1 = 9.3s at peak, ~18.6s at 50% efficiency. Estimated speedup: ~3.9x. These are theoretical estimates — measured benchmarks are required before claiming any real-world speedup.

The Storage Layout

The block-diagonal W_hh changes the weight layout. The current dense layout uses combined_dim = in_dim + D for the fused gate GEMM, where the weight matrix is D × combined_dim with three gates stored as [W_z | W_r | W_h]. With block-diagonal W_hh, the recurrent portion is N_heads × d_head × d_head instead of D × D. The input projection W_x stays full-rank. The fused gate GEMM must be split: a full-rank input GEMM for W_x, and a separate block-diagonal GEMM for W_hh.

Block-diagonal W_hh: 64 independent 64×64 blocks along the diagonal. Each head's recurrent state evolves independently. The input projection W_x (left block) stays full-rank and provides cross-head mixing.

What About BF16?

BF16 has the same tensor core throughput as FP16 on Ampere (68.2 TFLOPS dense) and avoids the overflow problem entirely — its 8-bit exponent matches FP32's range. For most mixed-precision training, BF16 is the safer choice. So why pursue FP16 accumulators at all?

Three reasons:

1. Mantissa precision. FP16 has 10 mantissa bits; BF16 has 7. The paired-FP16 technique gives ~22 effective bits from two FP16 values. The same technique with BF16 gives only ~14 effective bits — not enough to match FP32. For the recurrent state, where every bit of precision matters, FP16's wider mantissa is the better base.

2. Vulkan driver support. FP16 cooperative matrix support is mature and widely deployed. BF16 cooperative matrix support (GL_EXT_bfloat16, SPV_KHR_bfloat16, VK_KHR_shader_bfloat16) was finalized in February 2025. The RTX 3080 Ti's Ampere tensor cores support BF16 natively, but Vulkan driver exposure of the BFloat16CooperativeMatrixKHR capability may lag. FP16 is the proven path.

3. The research question. The project goal is specifically to investigate whether pure FP16 accumulators can be made to work — not to pick the easiest path. If the four-part architecture succeeds with FP16, it will also work with BF16 (with reduced effective precision). The FP16 path is the harder, more interesting problem.

What About INT8 Integer-Only?

INT8 tensor cores on Ampere run at 273 TOPS dense — 4x FP16's 68.2 TFLOPS. The hardware path is sint8 × sint8 → sint32: the tensor core packs 4 INT8 values into one INT32 register lane, performs 4 parallel INT8 × INT8 multiplies, and accumulates into INT32. The accumulation is bit-exact — zero rounding error, ever. This is a fundamental advantage over FP16, where every accumulation step introduces ~2^-10 relative error.

The question: can we redesign the GRU math to work entirely in integer space — no floats anywhere in the forward pass — and train to convergence?

The Integer-Only GRU Cell

The iRNN paper (arXiv 2212.11791) demonstrated integer-only RNN inference using piecewise linear (PWL) activation approximations and quantization-aware training (QAT). No published work has implemented integer-only GRU training with INT8 tensor cores on GPU. The design:

The forward pass is entirely integer. The GEMM runs on INT8 tensor cores at 273 TOPS. Sigmoid and tanh are approximated as 8-piece PWL functions with 256-entry lookup tables — integer-only, ~0.1% approximation error. The recurrent update h_new = (127 - z) * h_prev + z * h_tilde is computed in INT32 (exact) and requantized to INT16 for storage.

The Gradient Problem — Solved by Design

The naive concern is that INT8 quantization destroys gradients: the derivative of the quantization function is zero almost everywhere. But the integer-only design avoids this by using the straight-through estimator (STE). The STE says: d(quantize(x))/dx ≈ 1 for the backward pass. The gradient flows through as if the quantization weren't there. The forward pass has quantization noise, but the backward pass sees clean gradients.

This is the standard QAT approach (Jacob et al. 2017). The STE paper (arXiv 2405.05171) proves it is equivalent to more complex gradient estimators after adjusting the learning rate. The key difference from naive INT8 quantization (which fails): in the integer-only design, there is no float↔int conversion in the forward hot path — everything is already integer. The STE applies once, consistently, across the whole network.

The Error Budget

The GEMM being exact is the key advantage. In FP16, the GEMM accumulation over d_head=64 introduces ~64 ULP of error per timestep. In INT8, it's zero. The error only comes from input quantization and output requantization, not from the accumulation itself.

The Throughput Advantage

Theoretical peak. Real throughput is typically 40-60% of peak for memory-bound kernels. Expected end-to-end: ~6-8x vs FP32 baseline (vs ~3-4x for FP16 plan). INT8 weights use 0.4x the memory of the FP16 plan and 50x less than FP32 baseline.

Measured Vulkan INT8 Benchmark Results

We implemented an optimized INT8 cooperative matrix GEMM shader using the same tiling pattern as our production gemm_mixed.comp (8 subgroups per workgroup, 4 N-tiles per subgroup, double-buffered shared memory, prefetching). Key optimizations: packed uint32_t global loads (4 int8_t per transaction), shared memory stored as packed uint arrays (32-bit aligned, zero bank conflicts), and padded shared memory stride (+1 uint per row) to eliminate bank conflicts on NVIDIA's 32-bank shared memory.

Correctness verified against CPU reference: 0 mismatches, max error 0 — the INT32 accumulation is bit-exact.

Measured on RTX 3080 Laptop GPU (Ampere GA106). INT8 achieves ~19-21 TOPS (7.0-7.7% of 273 TOPS theoretical peak). FP16 achieves ~8 TFLOPS (11.7% of 68.2 TFLOPS peak). FP32 scalar achieves ~0.1 TFLOPS (no tensor cores). The 2.3-2.7x INT8 vs FP16 speedup is consistent across all sizes — compute-bound regime. The 178-210x INT8 vs FP32 speedup demonstrates the tensor core advantage over scalar CUDA cores.

Memory comparison at 4096×4096×4096: INT8 uses 100.7 MB (int8 weights + int32 accumulator), FP16 uses 134.2 MB (fp16 weights + fp32 accumulator), FP32 uses 201.3 MB (fp32 weights + fp32 accumulator). INT8 saves 25% vs FP16 and 50% vs FP32.

The remaining gap to 4x theoretical INT8 vs FP16 is due to the INT8 cooperative matrix tile size (16×16×32) requiring more K iterations than FP16 (16×16×16), and the overhead of unpacking packed uint data for cooperative matrix loads. We tested uvec4 (128-bit) packed loads but found them slower — the 16 scalar shared memory stores per thread outweigh the global load benefit. The packed uint shared memory layout (storing int8_t data as uint arrays with 32-bit alignment) eliminated the bank conflict bottleneck and improved throughput from 16 to 19-21 TOPS.

INT8 GRU Cell Prototype

Beyond raw GEMM throughput, we prototyped a complete integer-only GRU cell using Vulkan compute shaders. The cell implements all four GRU gates (update gate z, reset gate r, new gate n, hidden state h_new) entirely in INT8/INT32 arithmetic — no floating-point operations anywhere in the forward pass. Key design choices:

• Packed INT8 weights: 4 int8 values packed per uint32_t, loaded with 32-bit transactions
• Subgroup reductions: 32-lane subgroup add for dot-product accumulation (INT32, bit-exact)
• PWL sigmoid/tanh: piecewise-linear approximations mapping INT32 accumulators to INT8 range
• Fixed-point requantization: INT32 accumulator → INT8 via right-shift by 16 + signed rounding division by 127
• INT32 hidden state: h stored as int32 to avoid requantization noise across timesteps

Correctness verified against CPU reference at D=64, V=64: 0 mismatches, max error 0 — PASS. The signed rounding helper (round_div127) eliminates the 1-LSB truncation error from naive integer division.

GEMV Benchmark (GRU Inference Shape)

GRU inference uses matrix-vector products (W·x + U·h), not large GEMMs. We benchmarked INT8 GEMV at GRU-relevant dimensions:

GEMV correctness: 0 mismatches, max error 0 at 64×64. The 2048×101 shape (W_xz · x for D=2048, V=101) runs in 0.036 ms — the actual Cat Chat input projection. The 2048×2048 shape (W_hz · h) runs in 0.051 ms — the recurrent projection.

GRU Cell Throughput

Single-timestep GRU cell at D=2048, V=101 (actual Cat Chat dimensions):

*LUT accuracy issue: the 256-entry LUT uses true sigmoid/tanh (scaled), while the PWL uses a coarse 3-segment linear approximation. The LUT needs wider input range mapping to match the accumulator distribution. Throughput is comparable across all variants (~0.095-0.098 ms per cell). Weight memory: 13.22 MB per layer (50% savings vs FP32).

Multi-Timestep Inference

We implemented sequential multi-timestep inference with ping-pong hidden-state buffers. Each timestep's h_out feeds into the next timestep's h_in via descriptor set swapping, with pipeline barriers between dispatches:

Per-step time decreases slightly at T=4096 (0.0645 ms vs 0.0656 ms at T=128) due to better GPU occupancy with longer command buffers. The multi-timestep path is faster than the single-timestep estimate (264 ms vs 400 ms) because the single-timestep benchmark includes vkQueueWaitIdle overhead per cell.

Full 3-Layer GRU-L3 Benchmark

At the actual Cat Chat model dimensions (L=3, D=2048, V=101, T=4096):

Per-token latency: 0.188 ms. Weight memory (3 layers): 39.67 MB (vs ~80 MB for FP32). The 769 ms total for a 4096-token sequence is well within real-time inference bounds for interactive dialogue.

Limitations

• Inference-only: No backward pass. Training in INT8 requires gradients through PWL sigmoid/tanh — a separate, harder problem.
• Canonical reset gate: The 2-pass canonical implementation has 10/64 mismatches (max error 1) from the additional requantization step rh = round_div127(r * h). This is expected for integer-only arithmetic.
• LUT activations: The 256-entry LUT needs wider input range mapping to match the PWL accuracy. Current LUT uses true sigmoid/tanh scaled to [-1, 1] × 8, which doesn't cover the full accumulator range.
• PWL coarseness: The 3-segment PWL sigmoid is coarse. For production, a finer PWL or corrected LUT may be needed for accuracy-sensitive layers.
• Standalone experiment: Not integrated into the trainer. Production FP32 path is untouched.

Why INT8 Is the Harder, More Interesting Problem

The FP16 accumulator plan uses proven numerical techniques (stochastic rounding, error-free transforms, RMSNorm) to make FP16 safe. The INT8 integer-only plan is more novel: no published work has trained an integer-only GRU on GPU with INT8 tensor cores. The iRNN paper did integer-only inference; the FloatSD8 paper (arXiv 2001.08450) trained LSTMs with 8-bit weights but kept accumulations at 16 bits. The research questions are:

1. Does the exact INT32 GEMM accumulation compensate for the INT8 input quantization error over T=4096 timesteps?
2. Does the STE produce stable training gradients through the recurrent loop?
3. Is INT16 fixed-point state better or worse than FP16 state for recurrent dynamics?
4. How many PWL pieces are needed for sigmoid/tanh to maintain training stability?
5. Does the block-diagonal split + spectral norm + RMSNorm combination prevent gradient collapse under INT8 quantization noise?

The full integer-only GRU experiment plan is documented in .devin/plans/INTEGER_ONLY_GRU_PLAN.md — 22 files across 12 phases, starting with a Vulkan INT8 cooperative matrix probe and ending with a drift test at T=4096.

Status and Next Steps

Two implementation plans are now documented:

• .devin/plans/FP16_ACCUMULATOR_PLAN.md — 26 files, FP16 accumulators with block-diagonal + stochastic rounding + paired-FP16 + pre-recurrence RMSNorm. Roasted through 10 rounds, converged. Lower risk, proven techniques, ~3.9x speedup.
• .devin/plans/INTEGER_ONLY_GRU_PLAN.md — 22 files, integer-only GRU with INT8 tensor cores + PWL activations + STE gradients. Higher risk, novel contribution, ~15.9x theoretical speedup.

The verification criteria for both plans are strict:

• Numerical drift < 1% relative error at T=4096 (vs FP32 reference)
• Stochastic rounding determinism (same seed → same output)
• VRAM < 8 GB on RTX 3080 Ti (D=4096, L=4, T=4096, B=32, 8 segments)
• Throughput > 2x vs FP32 baseline (measured, not estimated)
• CPU/GPU parity within tolerance

If the drift test fails, the fallback is to tighten the spectral norm constraint on W_hh (from σ₁ < 1.5 to σ₁ < 1.0) or increase d_head from 64 to 128. If the throughput test fails, the bottleneck is likely memory bandwidth, not compute — and the block-diagonal split's 64x reduction in recurrent parameters should help by reducing weight memory traffic.

Epoch 7 Pretraining Results

After 7 epochs of curriculum pretraining on the Cat Chat corpus (V=101, D=2048, L=3, ASCII character-level):

The teacher-forced accuracy of 61.47% shows the model has learned character-level patterns, but the 8.45% free-running accuracy indicates it cannot yet generate coherent text autoregressively. This is expected for a character-level model with only 7 epochs of curriculum pretraining.

INT8 Inference Mode — Working End-to-End

The INT8 inference path is now working end-to-end on the RTX 3080 Laptop GPU. The key breakthrough was making the requantization shift a per-layer push constant instead of a hardcoded constant:

• Per-layer REQUANT_SHIFT: The shader accepts requant_shift_x and requant_shift_h as push constants, computed from actual weight and input quantization scales
• Separate W_x*x and W_h*h shifts: The embedding input and hidden state have different quantization scales, so they need separate shifts to align to the PWL activation range
• Global embedding scale: Uses a single scale_x = 127 / max_abs_all_embeddings for all tokens, matching the fixed shift
• Hybrid INT8/FP32: GRU cell on GPU (INT8), embeddings/RMSNorm/classifier on CPU (FP32)

All correctness benchmarks pass:

The cat_chat int8-chat CLI command loads the E7 checkpoint, quantizes weights to INT8, and runs interactive generation. Output is currently gibberish (matching the model 8.45% free-running accuracy), but the INT8 inference path itself is numerically correct.

Related

• GRU-L3 V3a: Gradient Checkpointing, 75M Parameters — the previous post in the GRU-L3 series
• GEMM Tiling: CPU and GPU Kernels — our Vulkan cooperative matrix GEMM implementation
• Vulkan ML Compute for Cat Chat — why we chose Vulkan over CUDA for the trainer
• FP16 Batch Interleaving — our CPU FP16 GEMM work for inference
• Backward Pass Vulkan Tiled GEMM — the backward GEMM shader stack