Module ngram_cache

Module ngram_cache 

Source
Expand description

NVMe-backed row cache for the n-gram embedding tables.

The n-gram tables of the LongCat / Qwen3.8-Flash-Next family are the model’s largest tensors by far (31.4 B params on LongCat-Flash-Lite, ~51 B announced for Flash-Next) and simultaneously its least bandwidth-hungry: a token touches exactly one row per table — 12 rows, ~3 KB — regardless of sequence length. Pure capacity, near-zero bandwidth, which makes them the best demotion candidate in the model.

Design, and why it needs no CUDA kernel change:

  • The cache is a flat PINNED arena of slots × row_stride bytes. On GB10 pinned host memory is GPU-addressable at the SAME virtual address (ExpertArena asserts this), so the arena is a [slots, dim] device-side table.
  • The n-gram row ids are computed HOST-side (they are a pure function of token ids), so a lookup resolves row_id -> slot on the host and hands the gather kernel the SLOT INDEX in place of the row id. batched_embed / batched_embed_fp8 then run verbatim against the arena base.
  • A miss reads the row straight off NVMe into its pinned slot — no cuMemcpyHtoD anywhere on the path.

Eviction is CLOCK (second-chance): O(1), no per-hit bookkeeping, and it approximates LRU well for the power-law access pattern these tables have. Rows touched by the CURRENT batch are pinned so a large prefill can never evict a row it is still about to read.

O_DIRECT requires 4 KiB-aligned reads, while a row is typically 256 B (FP8, dim 256). Reads are therefore issued as the containing 4 KiB block into a bounce buffer and the row copied out — the block is the disk’s minimum transfer anyway, so this costs no extra I/O, only a 256 B host memcpy. Cache capacity stays row-granular, which matters because the hash scatters ids: neighbouring rows in a table are unrelated.

Structs§

NgramRowCache
One table’s on-NVMe backing file plus its resident row cache.