Module ngram_embed 

N-gram embedding id math — LongCat-Flash-Lite / Qwen3.8-Flash-Next family.

Host-side, kernel-independent core of the n-gram embedding path (arxiv 2601.21204): the polynomial-rolling-hash ids that select rows in the K * (N-1) giant lookup tables. Split from the (future) GPU module so the integer math has a pure unit-test surface — the Rust ids are checked BIT-EXACT against the Python reference via bench/ngram_ref/ngram_id_fixtures.json (generated by bench/ngram_ref/make_fixtures.py from the line-faithful numpy port of the HF modeling_longcat_ngram.py).

Mechanism (table index = (i-2)*K + j for n-gram size i, 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 (segments end at an EOS token, inclusive) contributes token id 0 instead of crossing the boundary. Ids depend ONLY on token ids — never on hidden state — which is what makes the lookups deterministic, prefetchable and speculative-decode-friendly.

Overflow contract: x * mod < 2^17 * 2^24 < 2^41 and at most N-1 terms accumulate, so the running sum fits comfortably in i64/u64 WITHOUT intermediate reduction at LongCat scale (V=131072, T≈10.2M). The per-term products must still be computed in 64-bit — 32-bit would overflow — and mods themselves must be built with 64-bit modmul.

Structs

NgramDims

The n-gram trio + derived dims, extracted from ModelConfig.

NgramEmbedding

GPU n-gram embedding: base word embedding fused with the K*(N-1) hashed-table lookups, composed entirely from existing kernels (batched_embed gathers + dense_gemm_bf16_pipelined projections + bf16_scaled_add accumulation). A fused single-kernel version is a later optimization; per-token cost here is 1+T gathers, T tiny GEMMs and T+1 scaled adds (T = 12 for LongCat-Lite) — negligible next to a

Enums

NgramTable

One n-gram lookup table on device: BF16 as shipped, or FP8-quantized at load (per-row E4M3 + f32 scale — halves the ~63 GB table footprint; embeddings tolerate this well and the gather dequantizes on read).

Functions

ngram_ids

Compute the row ids for EVERY table over ctx (the n-1 cached context tokens followed by the new tokens). Returns num_tables vectors of ctx.len() ids each, table-major in reference index order ((ngram-2)*K + split); callers slice the last seq_len entries.