Module ngram_embed 

2026-09-25: N-gram embedding (arXiv 2601.21204): the dimensions, the host-side polynomial-rolling-hash ids that select rows in the K * (N-1) lookup tables (ids), the tables (table) and the GPU embedding (embed). ids_match_python_reference checks the ids against bench/ngram_ref/ngram_id_fixtures.json, written by bench/ngram_ref/make_fixtures.py.

Table index = (i-2)*K + j for n-gram size i and split j:

  T      = ratio * vocab_size + 2*index + 1        (table row count)
  mods   = [V^1 mod T, V^2 mod T, ..., V^(i-1) mod T]
  id_t   = ( x_t + Σ_{d=1..i-1} shift_d(x)_t * mods[d-1] ) mod T

where shift_d is a right-shift by d that resets at document boundaries: a position within d tokens of a segment start (a segment ends at an EOS token, inclusive) contributes token id 0. The ids depend only on token ids.

Owner: model-layers (n-gram embedding). Invariants:

• ngram_ids sums in u64 and reduces once per id, so the ids are exact only while neighbor_num * vocab_size * T < 2^64; the code does not check this.

Structs

NgramDims

2026-09-25: The n-gram dimensions, read from ModelConfig.

NgramEmbedding

2026-09-25: GPU n-gram embedding from batched_embed gathers, dense_gemm_bf16_pipelined projections and bf16_scaled_add accumulation: per call 1 + T gathers, T GEMMs and T + 1 scaled adds, T = num_tables().

Enums

NgramTable

2026-09-25: One n-gram lookup table. Fp8 holds E4M3 rows with a per-row f32 scale, which the gather (batched_embed_fp8) applies.

Functions

ngram_ids

2026-09-25: Row ids of every table over ctx. Returns num_tables() vectors of ctx.len() ids each, in table index order ((ngram-2)*K + split).