Extending the 256 context window and KV cache economics for Python code
Hi Kiran,
Pumping ~42B cumulative tokens across base, DAPT, and SFT stages into a 125M parameter model shows great training diligence, and testing AST-level metrics for Python is a great touch.
Looking at your architecture and KV profiling table:
Context length is capped at 256 tokens using learned positional embeddings.
In Python, 256 tokens rarely covers multi-function context, class definitions, or imports with docstrings.
Profiling shows KV cache overhead reaching +72 MB at batch 4 / context 256 due to standard 12-head MHA storing full key/value states on every layer.
In an open architecture project called Maba (101M reference model: https://huggingface.co/AndrewThompson1233/maba-v1-architecture), we handle this memory wall using two architectural changes:
- RoPE over learned positions: Switching to Rotary Position Embeddings removes the hard 256-token parameter cap entirely, allowing zero-shot sequence length extension at inference.
- Hybrid recurrence (75% GDN-2 / 25% GQA): Linear recurrent layers maintain fixed O(1) state memory. At 1,024 context, a hybrid model consumes less KV memory than a pure MHA model at 256 context, while preserving global attention where it matters.
Also, given that you trained across 42B tokens, parameter capacity is your main constraint. A deterministic 2-pass physical block recycling setup would expand your 16 physical layers into 32 effective non-linear transformations without adding weights.
Did learned positional embeddings and MHA memory pressure dictate the 256 context cap during initial design?
Best,
Andrew
Hi Andrew,
Thanks for taking the time to go through the architecture and KV-cache analysis in detail.
Yes — the 256-token context was a deliberate constraint in the initial design. The model was designed as a relatively small first-principles SLM, and the initial objective was to keep the architecture simple and make the effects of tokenization, training scale, domain adaptation, and SFT easier to isolate and measure.
Learned positional embeddings were part of that initial design, which makes the 256-token context a hard architectural limit rather than simply an inference-memory choice. The KV-cache profiling was then used to characterize the resulting inference-memory behavior.
I agree that extending context is an important next direction. RoPE and alternative attention/recurrent mechanisms would require a new architectural experiment rather than just a configuration change, so I would prefer to evaluate those separately rather than modify the current checkpoint.
The physical block-recycling idea is also interesting in the context of parameter-constrained SLMs. I'll add it to the architectural experiments to investigate.
Thanks again for the detailed suggestions.
Hi Kiran,
That makes complete sense. Keeping the initial architecture strictly vanilla to cleanly isolate the effects of tokenization, 42B token exposure, and domain adaptation is a very disciplined baseline methodology.
Because learned positional embeddings bake the [256, hidden_dim] weight matrix directly into the state dict, decoupling that via RoPE in a fresh experiment is definitely the cleanest way forward without corrupting your baseline checkpoint.
Looking forward to seeing the next architectural phase when you test longer context windows and block recycling. Best of luck with the continued research!
Best,
Andrew