The Work Moved Down a Layer
Cat Chat began as a language-model experiment inside a custom C++ neural-network library. The first useful question was not “can a GPU run a matrix multiply?” It was “can the same explicit Vulkan machinery that renders a game also carry a training step without hiding the costs?”
That question led to a small family of compute shaders for recurrent language-model training: GEMM, GRU gates, RMSNorm, attention, softmax cross-entropy, gradient reduction, and optimizer updates. The implementation is deliberately close to the metal, but the lesson is broader: GPU ML is less about one impressive kernel than about making data movement, synchronization, precision, and verification agree.
*A local tiled-GEMM benchmark on an RTX 3080 Laptop GPU, not an end-to-end training throughput claim. Results depend on dimensions, dispatch shape, driver, clocks, and verification mode.
One Pipeline, Several Kinds of Work
The current design keeps model state in persistent GPU allocations rather than uploading weights for every operation. A forward pass can move through attention, recurrent layers, normalization, and classification; the backward pass reverses that path while accumulating gradients. The CPU remains involved at synchronization boundaries—feeding batches, collecting loss or gradients, and coordinating optimizer updates—but the expensive tensor operations have a common execution home.
Cooperation Beats Isolation
The central GEMM path uses VK_KHR_cooperative_matrix where the device exposes a supported configuration. Cooperative matrices distribute a matrix fragment across a subgroup, allowing the invocations to participate in a matrix multiply together. The dimensions are not universal: Vulkan exposes a query mechanism because implementations may support different component types, scopes, and tile shapes.
Our shader pairs FP16 inputs with an FP32 accumulator, then performs the final conversion explicitly. That choice follows the same mixed-precision pattern described in the Khronos cooperative-matrix tutorial: save bandwidth and use specialized matrix hardware for the inputs, while retaining more range and precision during accumulation.
For devices without a compatible cooperative-matrix configuration, the library keeps an FP32 route. Portability is not achieved by pretending every GPU is the same; it is achieved by querying capabilities and selecting a legal path.
Tile the data
The benchmark kernel reuses B tiles through workgroup memory and computes multiple output tiles per workgroup. The point is not a magic tile size; it is reducing repeated global-memory traffic while respecting device limits.
Keep state resident
Weights, recurrent state, activations, and gradients are allocated once for the training session. Re-uploading a large weight matrix for every timestep can erase the benefit of faster arithmetic.
Fuse the small work
GRU gate math, normalization, residual handling, reductions, and loss each have their own shader responsibilities, but dispatches are organized around useful training work rather than isolated scalar operations.
The Debugging Lesson Was Numerical
The hardest failures did not look like Vulkan failures. They looked like machine-learning failures: alternating NaNs, exploding recurrent states, or a loss that could not be trusted. Several fixes came from tracing the value through the full pipeline instead of changing the optimizer first.
- A cooperative-matrix accumulator was written directly into an FP16 output buffer. The reliable path was to store into an FP32 scratch buffer, apply bias and activation in FP32, and convert to FP16 at the final elementwise store.
- Reductions for softmax and loss were moved to FP32 shared memory. A low-precision reduction can turn a small numerical difference into a control-flow or gradient problem.
- RMSNorm and recurrent state needed explicit precision boundaries. FP16 is useful for bandwidth and matrix inputs, but recurrent feedback and normalization statistics are sensitive places to retain FP32.
- Every GPU result was compared with a CPU reference using error statistics rather than a single “looks right” value. This makes the benchmark a correctness instrument, not only a speedometer.
The Vulkan ML correctness guidance makes the same general point: CPU and GPU reduction order can differ, so tolerances must account for floating-point behavior without masking real corruption.
From Benchmark to Training Stages
We used an incremental ladder instead of attempting the complete trainer in one pass:
- Establish a raw GEMM and verify it against a CPU calculation.
- Add a GRU gate, then the full cell and recurrent timesteps.
- Stack layers with RMSNorm and residual paths.
- Add attention, classifier, softmax cross-entropy, and gradient flow.
- Exercise the complete path with smoke coverage before tuning performance.
This approach made a shader bug local. When a complete model produces NaNs, it is difficult to know whether the cause is a descriptor, a stride, a barrier, a precision conversion, or a wrong reference. When each stage has a small numerical contract, the search space shrinks.
Why This Belongs in a Game Engine
CatGame is not trying to turn a renderer into a general-purpose framework by accident. The same explicit resource model is valuable when a game needs local inference, learned NPC behavior, procedural content, or an offline training tool. Vulkan gives the project one vocabulary for queues, buffers, descriptors, synchronization, and shader capability discovery.
The cost is real. A GPU delegate that submits tiny, serialized jobs can be slower than a well-tuned CPU path; our earlier experiments measured that directly. The answer was not “GPU is always faster.” It was to batch work, keep weights resident, reduce round trips, and preserve a CPU baseline that can expose regressions.
What Comes Next
The next frontier is integration quality: end-to-end optimizer steps, more complete parity tests, device-specific tuning, and a cleaner boundary between experimental shaders and the production trainer. Cooperative matrices may offer a large arithmetic advantage, but only when the surrounding pipeline feeds them efficiently and verifies their output.
That is the useful discovery so far: ML compute in Vulkan is a systems problem. The multiply is the visible center. The engineering is everything around it.
External references: Vulkan compute guide, shader FP16 reference, and compute validation tutorial.