MHA KV-cache footprint across 20 heads and vocabulary factorization at step 20
Hi Lokesh,
Catching a from-scratch 505M custom causal model right at the step 20 verification stage is great timing. Validating the custom tokenizer and 24-layer tensor wiring early prevents costly restarts later.
Looking at your architecture parameters:
Full MHA KV-cache memory scaling:
With 20 attention heads across 24 layers (head dimension 64) and standard Multi-Head Attention, the model stores distinct key and value tensors for every head.
At context length 2,048, a single FP16 generation stream consumes ~240 MB of KV cache. In batched inference or serving environments, full MHA quickly saturates GPU memory bandwidth.
Switching to Grouped-Query Attention (such as 20 query heads to 4 KV heads) or hybrid linear recurrence drops that memory footprint by 4x to 5x.
In an open architecture project called Maba (101M reference model: https://huggingface.co/AndrewThompson1233/maba-v1-architecture), we pair attention with Gated DeltaNet (75% GDN-2 / 25% GQA). The linear recurrent layers maintain fixed O(1) state memory, slashing active KV cache by over 76% across long contexts.Vocabulary parameter reallocation:
With a 24,000 vocabulary at 1,280 hidden dimension, your token embedding matrix consumes 30.72M parameters (over 6% of your entire 505M footprint).
One transformer block in your configuration costs roughly 20M to 26M parameters. Static lookup tables cost more than an entire layer of sequence compute.
Decoupling the input lookup via low-rank factorization (24,000 -> 128 -> 1,280 = 3.24M parameters) frees ~27.5M weights, enough to fund a 25th physical transformer block within the exact same 505M envelope.Representation depth via block recycling:
Since you are at step 20 and testing base convergence, running representations through your 24 blocks twice with Split RMSNorm (distinct norm vectors for pass 0 and pass 1) expands depth to 48 effective layers at zero parameter cost.
Did memory ceilings dictate keeping context at 2,048, and what hardware cluster are you targeting for the full pretraining run?
Best,
Andrew