Host RAM: keep lookup-only tables (token and per-layer embeddings) quantized and memory-mapped instead of resident copies (saves ~1 GB on Qwen3.5-2B, ~2.9 GB on Gemma 4 E2B) #807

On 16 GB Ryzen AI laptops, most of FLM's host-RAM footprint is tables that are only read one row per token. Two models show it most clearly:

With gemma4-it:e2b loaded, the PC's available memory fell from 6,477 MiB to 479 MiB, so a 16 GB machine is close to paging with nothing else open.

What the code does today

• Gemma 4: gemma4e_npu builds embedding_q8_0 for both model.embed_tokens and model.per_layer_token_embd. Each forward(idx) dequantizes a single row. But init_weights copies the whole int8 table and its F32 scales into owned buffer<>s through Q4NX::load_weights (a std::ifstream read). So all 2,952 MiB stays resident, although a token touches about 12 KB of it.
• Qwen3.5: qwen3_5vl_npu.hpp includes modules/embedding.hpp. Its Embedding resizes a full vocab_size × d_model bf16 vdtype and copies one row per forward, and the 970 MiB block is exactly that size. Qwen3.5-2B has tie_word_embeddings: true, and its model.q4nx has no embed_tokens tensor, only lm_head.weight (I8 [7760, 8, 8704]). So the 970 MiB bf16 table looks like a dequantized copy of weights FLM already holds.

Proposal, in order of payoff

1. Map the lookup-only tensors from model.q4nx read-only, instead of copying them. That means CreateFileMappingW + MapViewOfFile(FILE_MAP_READ) on Windows, and mmap(PROT_READ, MAP_PRIVATE) on Linux.
   • embedding_q8_0's tensor_2d views would point into the mapping, not into owned buffers.
   • The pages are clean and file-backed, so only the rows of tokens actually seen stay resident. Under pressure the OS can drop them and read them back.
   • For Gemma 4 E2B, that's about 2.9 GiB less private memory: a few 4 KB pages per distinct token, so tens of MB for a long conversation instead of 2,952 MiB.
2. For tied-embedding models on the bf16 Embedding path (Qwen3.5), don't materialize the vocab × d bf16 table. Dequantize the one row per token from the int8 tied weights, as embedding_q8_0 already does for Gemma 4. That saves 970 MiB on Qwen3.5-2B, and the same holds for any other tied-embedding model on that path.
3. A question, not a request: is the private lm_head.weight copy (515 MiB on Qwen3.5-2B, 240 MiB on Gemma 4 E2B) still used once the weights are on the NPU?
   • If it's only staging for the upload, freeing it (or mapping it per item 1) saves that too.
   • If the row lookup in item 2 reads from it, it should stay, or be mapped.
4. Smaller, independent of mapping: store per_layer_token_embd.weight.scale as bf16 or fp16 instead of F32. That's 280 → 140 MiB for Gemma 4 E2B.

The cost is a page fault the first time a cold row is used. Right after load the file is usually still in the standby cache, so that fault doesn't touch the disk.

Related

#709 asks for weight BOs not to be MAP_LOCKED on Linux, so cold MoE experts can page. This issue is about the CPU-side lookup tables on any OS, which aren't NPU buffers at all.

Environment and method

• FLM: v1.0.7.
• PC: Windows 11 24H2 (26100), Ryzen AI 7 350 (XDNA 2), 16 GB, NPU driver 32.0.20102.3930.
• Method: one flm serve <model> --ctx-len 4096 at a time, nothing else on the NPU, after one short chat request. Three readings:
  • flm.exe's private bytes and its largest private regions, from VirtualQueryEx over the process;
  • NPU memory, from the "GPU Process Memory" shared-usage counter for the NPU adapter;
  • available memory, from \Memory\Available MBytes.
• The tensor names, dtypes and shapes come from the model.q4nx headers.

Drafted with an AI assistant; the measurements and header readings above are from the runs described.