KV cache memory footprint at 40k context and vocabulary allocation on 1.7B

#1
by AndrewThompson1233 - opened

Hi team,

Targeting extended reasoning (/think) and tool calling in a 1.7B edge footprint is a great direction, especially with native 40k context support.

Looking at the memory and compute dynamics on local laptop deployments:

  1. The KV cache memory wall at 40k context:
    Your q4_k_m GGUF weighs ~1.1 GB. Assuming standard GQA (2 KV heads, head dimension 128), storing the KV cache across 28 layers at 40,960 tokens in fp16 requires:
    28 * 2 * 2 * 128 * 40,960 * 2 bytes ≈ 1.17 GB.
    At full context or during long /think reasoning traces, the dynamic KV cache exceeds the size of the quantized model weights. On laptop CPUs/iGPUs, decoding throughput stalls heavily once memory bandwidth is dominated by reading this cache every step.

  2. Vocabulary parameter tax:
    With a 151,936 tokenizer at 1,536 hidden width, a single embedding table consumes ~233.4M parameters (13.5% of the 1.72B budget if tied, or ~27% if untied).

In an open architecture project called Maba (101M reference model: https://huggingface.co/AndrewThompson1233/maba-v1-architecture), we handle this compute regime using two complementary mechanisms:
Low-rank factorized embeddings: Projecting 151,936 -> 128 -> 1,536 drops the lookup table from 233.4M to 19.6M parameters, freeing over 210M parameters to allocate into deeper layer transformations.
Hybrid recurrence (75% GDN-2 / 25% GQA): Linear recurrent layers maintain fixed O(1) state memory. At 40k context, this slashes total KV cache size by ~75% (from 1.17 GB down to ~300 MB), preventing token generation from slowing down during lengthy reasoning traces.

How does generation speed hold up on edge devices when /think traces push context beyond 8k tokens?

Best,
Andrew

Computational Open Reasoning Technologies org
edited 1 day ago

Hello Andrew,

While we noted reasoning capability in the model card, full end-to-end reasoning is currently separated until our release of Pico 4 later this week, which natively integrates reasoning capabilities.

Regarding performance at extended context: in our CPU-only benchmarks at FP16, we did observe a reduction in generation speed as context length scales up. Unsurprisingly, the impact is most noticeable during the prefill phase, though decoding throughput is also affected as the KV cache grows.

We appreciate the breakdown of Maba's approach to factorized embeddings and hybrid recurrence. Slashing the KV footprint by ~75% while keeping O(1) state memory is a super compelling approach for edge deployment. We're already deep into the training run for Pico 4, so our architecture is locked in for this release, but we're keeping these exact memory-wall dynamics in mind as we start planning Pico 5 and beyond.

Best,
Kraxis

KraxisX changed discussion status to closed
KraxisX changed discussion status to open

Hi Kraxis,

Thanks for confirming the CPU benchmark telemetry. Prefill stalls and memory bandwidth pressure on growing KV caches are practically unavoidable on standard GQA once context scales out.

Excited to see the Pico 4 release later this week.

When you start drafting the architectural specs for Pico 5, feel free to ping me or open a thread on the Maba repo (https://huggingface.co/AndrewThompson1233/maba-v1-architecture). Happy to share profiling data and kernel setups for the GDN-2 recurrent blocks on edge hardware.

Best of luck with the Pico 4 launch!

Best,
Andrew

Sign up or log in to comment