TitanAI is a full training and inference pipeline for a language model that replaces the transformer attention mechanism with a stack of quantized feed-forward layers, built entirely in HIP for the AMD MI300X. Every core operation — forward projection, backward gradients, normalization, embeddings, and cross-entropy loss — is a hand-written kernel with no dependency on existing deep learning frameworks like PyTorch or JAX. The model uses BitLinear layers with ternary (1.58-bit) weights packed 16 per uint32, combined with int8 activations throughout. Weight updates follow a quantization-aware training approach: full-precision latent weights are maintained alongside the quantized forward-pass weights, updated via straight-through gradient estimation, then re-quantized each step. Logits are computed via weight-tied embeddings and a fused softmax/cross-entropy kernel that produces the loss and gradient in a single pass. The implementation is deliberately hardware-specific rather than a ported CUDA codebase. All warp-level reductions are written for AMD's 64-thread wavefront rather than a 32-thread warp. dp4a is implemented manually through bit-unpacking, since it isn't available as a native intrinsic on this path. Shared memory tiles are padded to avoid LDS bank conflicts, and memory access is structured for coalesced reads across the wavefront, using 128-bit vectorized loads and stores where the data layout allows it. This was built solo in roughly 48 hours with a limited GPU compute budget for the hackathon. The focus was on demonstrating a working, hardware-aware training and inference stack at the kernel level rather than on final model output quality — the current configuration uses a modest hidden size and vocabulary, and generation quality reflects that constraint rather than a limitation of the underlying architecture or kernel design.
Category tags: