The Benchmark Ladder is Behind Us
Two weeks ago, the GPU training path was a GEMM benchmark and a prayer. Today it is an eleven-stage ladder that takes a custom Vulkan/C++ language model from raw matrix multiplication through a full GRU recurrence, attention, classifier, loss, and optimizer step. Every incremental stage now passes its CPU numerical contract, and the end-to-end benchmark is more than twenty times faster per step than the CPU path on the same machine.
The shift is not just speed. It is trust. Each stage compares the GPU against the same model running on the CPU, so the error is quantified before we ask the training loop to make decisions. That discipline is what made a recent run of the full benchmark report all green lights while still carrying a warning that the work is not yet finished.
*Measured on an RTX 3080 Laptop GPU against a CPU AVX2+FMA reference on the same box, using the built-in benchmark fixture. Speedup varies with batch size, sequence length, and whether verification data is downloaded from the device.
What Broke Through
Three changes moved the boundary between "interesting demo" and "training loop candidate."
Bounds-checked GEMM shared loads
Cooperative-matrix tiles do not always line up with tensor dimensions. When the padding path was missing, out-of-bounds shared-memory reads could inject NaN straight into a gradient. A scalar load-and-pad wrapper before the cooperative tile read now keeps the GEMM path finite for any problem shape.
Non-finite gradient guards
The optimizer now refuses to update weights if a pre- or post-clip gradient norm is NaN or zero. This does not fix the NaN source, but it stops a bad step from corrupting the model, which is the difference between a recoverable training run and a silent crash.
Diagnostic parity gating
Hard-failing a benchmark on aggregate gradient-norm difference would have hidden the real signal: the CPU/GPU loss agrees to well under one percent. The benchmark now treats loss as a hard gate and gradient-norm divergence as a diagnostic warning, so progress is visible even when scaling is not yet perfect.
The Two Streams in Dual-Stream
The "dual-stream" in this model is a token stream and a context stream. The token stream is the familiar language-model path: embedding, GRU layers, classification. The context stream is a sliding, attention-weighted context buffer that can read from earlier positions in the sequence without re-feeding them through the recurrent gate. In practice this means the model can carry longer-range dependencies without exploding the per-timestep recurrence, but it also means the backward pass has an extra place for gradients to accumulate and an extra place for numerical edge cases to hide.
What is Still in the Way
Two obstacles keep the headline from being a victory lap.
Both problems are symptoms of the same broader challenge: moving from a numerically controlled benchmark to an open-ended data distribution. Benchmarks use fixed seeds and short sequences. Real data has variable length, irregular tokens, and rare cases where a single long sentence can amplify a small instability by the time it reaches the optimizer.
What the Wider Field Suggests
These issues are not unique to a custom engine. Recent work on small-language-model training and low-precision stability reinforces the same themes we are seeing locally.
- GNMR frames low-precision LLM training as a runtime control problem: most steps are safe, but local units can become risky in short windows, and a good controller should detect and recover rather than assume a static recipe is sufficient.
- GRUfinity demonstrates parallel associative scan and fixed-memory hidden-state passing for GRUs, suggesting that how we carry state through time may matter as much as the kernels that process it.
- The Khronos Vulkan ML Compute tutorial treats memory layout, workgroup sizing, and precision as first-class decisions. GPU ML is not just shader writing; it is memory-bandwidth and synchronization engineering.
- NVIDIA's mixed-precision training guide emphasizes loss scaling and un-scaling gradients before clipping. When different paths scale the loss or the gradient at different points, the same effective update can report very different gradient norms, exactly the kind of comparison trap we are seeing.
There is also a practical lesson from the broader ML tooling ecosystem: gradient-norm comparisons are fragile. Work in PyTorch and torchtune on tensor-parallel and distributed training has repeatedly shown that two correct implementations can produce nearly identical loss while reporting very different gradient-norm values, because norm is an aggregation and aggregation is sensitive to order, precision, and scaling. We are now designing the comparison so the warning carries that context, rather than declaring a false failure.
The Next Sprints
The near-term plan is to separate the two open problems and attack them with the same staged discipline that made the previous parts work.
Length bisect
Cap training sequences at 64, 128, 256, and 512 tokens to find the smallest length that triggers the context-stream NaN. This tells us whether the bug is in long backprop chains, attention over long context, or a dimension-specific shader edge case.
CPU/GPU norm audit
Compare the same long sequence on both paths, token by token and tensor by tensor, to discover whether the divergence is a per-token scale difference, a sum-ordering difference, or a missing unscale step in the comparison code.
Per-token guard
As an interim mitigation, consider clamping or capping the context-stream vocabulary gradient per token so that training can make progress while the exact backward fix is isolated and verified.
The longer goal is to make the training loop robust by design, not by accident: every numerical assertion, every guard, and every diagnostic warning is a place where the pipeline admits what it does not yet know. That is a healthier posture than a green benchmark that lies about real data.