mohith-das/jetson-hybrid-flat-cache-0.5b

jetson-hybrid-flat-cache-0.5b

Phase 2 of the llm_distillery project: an architecture-surgery experiment exploring O(1)-memory KV-cache alternatives to standard softmax attention, trained entirely on a Jetson Orin Nano Super (8GB unified memory).

Starting from `jetson-assistant-0.5b` (Phase 1's finished assistant model), every other attention layer was surgically replaced with one of two mechanisms, alternating layer-by-layer:

• —SWA (Sliding-Window Attention) — exact softmax attention bounded to a 4096-token window, O(window) KV-cache.
• —GLA (Gated Linear Attention) — a RetNet-style linear-attention recurrence with a per-head learnable decay gate, O(1) KV-cache (a fixed-size state matrix instead of a growing cache). Implemented with a chunk-parallel algorithm (chunk_size=256) for roughly 100x the throughput of a naive per-token recurrence on this hardware.

After the attention surgery, the model went through a short Continual Pre-Training (CPT) stabilization pass (200 steps, LoRA r=8/alpha=16 on q/k/v/o projections plus the new GLA decay parameters unfrozen) on a small synthetic corpus, letting the model adapt to its new attention mechanism before the LoRA was merged back in. See `jetson-flat-cache-cpt-adapter` for that adapter standalone.

This repo contains:

• —model.safetensors + config — the full FP16 hybrid model, loads directly with AutoModelForCausalLM.from_pretrained.
• —hybrid_flat_cache.gguf (Q80) — for `llama.cpp` inference. **Note:** llama.cpp has no native GLA kernel yet, so the GGUF's GLA layers fall back to standard sliding-window softmax attention at inference time — the GGUF is structurally valid and runs, but doesn't yet realize the O(1)-cache benefit of the GLA layers. The `logdecay` GLA gate parameters (no GGUF equivalent) were stripped during conversion.

This is a research/learning artifact exploring the architecture, not a production-optimized model — the CPT pass used a tiny synthetic corpus (200 rows) purely to stabilize the new attention layers, not to teach new capabilities.