Skip to content

Instantly share code, notes, and snippets.

Show Gist options
  • Select an option

  • Save randomvariable/e9f7a2c32a95d5551a71027c68e1acb2 to your computer and use it in GitHub Desktop.

Select an option

Save randomvariable/e9f7a2c32a95d5551a71027c68e1acb2 to your computer and use it in GitHub Desktop.
Qwen3.8-Flash-Next: serving the 100 GB per-layer embedding table

Qwen3.8-Flash-Next: serving the 100 GB per-layer embedding table

Identity

The model the community calls Qwen3.8-Flash-Next lands in transformers as model_type = "qwen4_exp", under src/transformers/models/qwen4_exp/, with @auto_docstring(checkpoint="Qwen/Qwen4-Exp"). It is a VLM (Qwen4ExpForConditionalGeneration) over a hybrid text stack: linear-attention layers, MoE, hyper-connections, and Qwen sparse attention (QSA) on the layers that indexer_n_heads enables.

One tensor dominates the deployment question. A per-layer embedding (PLE) table of 51.2 billion parameters sits on selected linear-attention layers, and upstream excludes it from the device map by default:

# modeling_qwen4_exp.py
_no_placement_params = ["ple.ple_embedding.ngram_embedding.weight"]

The upstream comment is explicit that this is a fallback, not a rule: "if it fits on accelerator (e.g. huge B200 gpus), then it will not be skipped, and will be put on device". Everything below concerns the case where it does not fit.

Gemma 3n shipped the same idea first, and Google documents the intent directly: PLE data can be generated outside model memory, cached to fast storage, and joined to each layer as it runs. A 2026 paper studies the pattern as STEM. Both are in References.

The gather path

flowchart TD
    T[Token and 2 predecessors] -->|hash into| R[16 row numbers]
    R -->|index| M[Table in host memory]
    M -->|gather| V[16 rows of 160]
    V -->|concatenate into| E[1 vector of 2560]
    E -->|project through| P[key_proj and value_proj]
    P -->|gate into| S[Residual stream]
Loading

The gather path from token identifiers to the residual stream.

The hash is cheap and has no finaliser on the serve path:

mixed2 = (t0 * m0) XOR (t1 * m1)                      bigram order
mixed3 = mixed2 XOR (t2 * m2)                         trigram order
row[h] = mixed[order(h)] % vocab_size[h] + offset[h]  for h in 0..15

Three multiplies and two XOR operations produce both mixed values, because the bigram mix is a prefix of the trigram mix. Sixteen 64-bit modulo operations produce the sixteen rows.

The sixteen heads are two n-gram orders of eight hash heads each. Upstream computes this as ngram_heads = (ngram_size - 1) * heads_per_ngram, so heads_per_ngram = 8 with ngram_size = 3 gives 16, and head_dim_per_ngram = ple_embed_dim // ngram_heads gives 160. Anyone reading heads_per_ngram = 8 against a 16-entry vocabulary array will otherwise assume a config error.

The per-position multipliers m0, m1, m2 are constructed once, not per token. SplitMix64 appears only there:

base_seed   = seed + 10007 * ple_layer_index
m[i]        = 2 * (splitmix64(base_seed + 0x9E3779B97F4A7C15 * (i+1)) % half) + 1
half        = max(1, ((2**63 - 1) // vocab_size) // 2)

The multipliers are odd by construction and bounded so m * token_id cannot overflow a signed 64-bit integer. Inverting this recovers vocab_size = 248320 uniquely from the three published multipliers, because one modulus must satisfy all three residues at once.

Head vocabularies are _find_nth_prime_after(ngram_vocab_size_base - 1, global_head_idx + 1) with global_head_idx = ple_layer_index * 16 + head_idx. For the published checkpoint they are the first sixteen primes after 19,999,999.

Checkpoint geometry

Quantity Value
Logical rows (sum of 16 head vocabularies) 320,001,446
Physical rows after padding 320,001,536, which is 128 x 2,500,012
Row width 160 at bf16, so 320 bytes
Rows gathered per token 16
Table parameters 51,200,245,760
Table size 102.4 GB, which is 95.37 GiB
split_ngram_parts in this checkpoint 128, against a default of 512
Useful bytes per token 5,120
Bytes moved per token at 4 KiB pages 65,536

The upstream tensor-parallel comment rounds this to "~45B == 90 GiB". The exact figure is 51.2 B parameters at 95.37 GiB, so budget against the computed value.

PLE runs only on the layers named in ple_layer_ids, which is one-indexed (config.ple_layer_ids.index(layer_idx + 1)), and only on linear-attention layers. Each PLE layer owns its own table with its own primes and multipliers. The published checkpoint has exactly one: 48 layers x 512 experts x 3 x 2560 x 640 is 241.6 GB at bf16, which with one 102.4 GB table and 16 GB of attention, embedding and norm tensors accounts for the 360 GB total_size. A second table would need 446.4 GB.

Why the access pattern is hostile

flowchart TD
    P[One 4 KiB page] -->|contains| U[320 useful bytes]
    P -->|contains| W[3776 unused bytes]
    U -->|repeats for| N[16 rows per token]
    N -->|moves| M[64 KiB per token]
    M -->|delivers| D[5 KiB of data]
Loading

Sixteen scattered rows move 64 KiB to deliver 5 KiB.

The sixteen rows have no fixed relationship. Each order reduces one mixed value under eight distinct primes, so the positions diverge as the mixed value changes. No permutation co-locates them for every input, which rules out the obvious fix of interleaving the sixteen head slices into one contiguous 5,120-byte row.

Page utilisation is 7.8 percent at bf16, for 12.8x read amplification. Halving the row width to FP8 halves utilisation to 3.9 percent, because the page count per token is fixed at 16 by the addresses, not by the row size.

The fault budget

Choose a budget as a share of one decode step, then divide by 16.

Decode step Budget at 10 percent Cost allowed for one read
10.0 ms 1.00 ms 62.5 µs
17.2 ms 1.72 ms 107.5 µs
30.0 ms 3.00 ms 187.5 µs

Inverted, against a 17.2 ms step:

Cost of one read Added time per token Share of the step
1 µs 0.02 ms 0.1 percent
10 µs 0.16 ms 0.9 percent
50 µs 0.80 ms 4.7 percent
107.5 µs 1.72 ms 10.0 percent
500 µs 8.00 ms 46.5 percent
5 ms 80.00 ms 465.1 percent

Caution

Do not derive a per-read cost from a whole-range refault measurement. A refault of a large mapped range measures sequential page installation and runs orders of magnitude faster than a random single-page fault. Measure the random case directly.

Resident-set policy

flowchart TD
    N[New row] -->|enters| W[Window under LRU]
    W -->|nominates victim to| A[Frequency test]
    C[Main cache] -->|nominates victim to| A
    A -->|admits winner to| C
    A -->|evicts loser| X[Discarded]
Loading

A frequency test arbitrates between the window victim and the main-cache victim.

The kernel admits a page on first touch. That is the wrong rule here, because most rows arrive once and never return. Adaptive W-TinyLFU is the published state of the art for this shape, and its authors report that it equals or exceeds every policy they compared across database, file-system, search-engine and storage traces.

Component Function
Window cache LRU over new arrivals. Published default is 1 percent of the budget
Main cache SLRU over proven rows. Victim comes from a 20 percent probation segment
Count-min sketch Frequency estimate. Published sample size is 10x the cache size
Reset Halves every counter and clears the doorkeeper at the sample size
Doorkeeper Bloom filter that absorbs single-touch rows, keeping them out of the sketch
Indicator Derives the window boundary at run time from measured skew

Two properties of this workload make each part load-bearing. The doorkeeper matters because a row touched once per sample costs one bit rather than a counter, and this table generates a long single-touch tail. The window matters because sixteen rows arrive together and a burst of unrelated prompts yields rows with no history, which a frequency-only admission rejects.

The published sample-size ratio was measured on caches of thousands to millions of objects. A resident set drawn from 320,001,536 rows sits outside that range, so counter width and sample size need validation rather than adoption.

Shrinking the table

flowchart TD
    X[Row of 160 values] -->|rotate by Q| Y[Rotated row]
    Y -->|quantise| Z[Stored row]
    Z -->|multiply by| W[W times Q transpose]
    W -->|equals| R[Unrotated result]
Loading

A block-diagonal rotation cancels against pre-multiplied projection weights.

QuaRot names the property this depends on: a rotation removes outliers from the hidden state without changing the output. Applied per 160-dim head slice, the inverse folds into the projection weights offline, so serve-time cost is zero.

Important

The fold requires that no stage sits between the gather and the projection. In qwen4_exp the concatenated 2560-vector feeds key_proj (2560 to 10,240, which is hidden_size * hc_count with hc_count = 4) and value_proj (2560 to 2560), both nn.Linear with bias=False, and norm_key is applied to key_proj output. Those two weights total 32.8 M parameters, so pre-multiplication is cheap. Re-verify this on any other checkpoint before planning a rotation.

Bits per value Table size Row bytes
16, bf16 102.4 GB 320
8, FP8 51.2 GB 160
4 25.6 GB 80
3.5 22.4 GB 70
2.5 16.0 GB 50
1 6.4 GB 20

The 3.5 and 2.5 entries are TurboQuant's published operating points for a KV cache, and TurboQuant needs no calibration data. Treat them as a starting point for a different tensor rather than a transferable result: TurboQuant targets MSE distortion, whereas RaBitQ bounds distance-estimation error, and a PLE row is a matmul operand rather than a search key.

Layout and codec rules

Rule Reason
Store no index The address is base + row * row_bytes. One multiply replaces any index structure
Store rows at fixed width A codec that codes a group jointly forces partial decode of neighbours

A block codec is the default answer for a large float array and the wrong answer here. The engine reads one row of 160 values, so a codec with a 1024-value coding vector touches about six neighbouring rows per read, and the coding vector itself destroys single-row addressing.

Mapping and registration

Map the file PROT_READ and MAP_SHARED. On one unified-memory accelerator, registering a 1 GiB writable mapping for device access took 1170 ms against 40 ms to 50 ms read-only, and the read-only path returned a device pointer equal to the host pointer. MAP_PRIVATE is not portable for file-backed device registration.

Registration does not pin. A registered range reported no locked or pinned pages, madvise(MADV_DONTNEED) succeeded over it, and the following launch still read correctly after a refault. Two consequences: an engine cannot assume residency from registration, and MADV_DONTNEED is a usable release mechanism when a host-memory arbiter needs bytes back.

Tensor parallelism

Upstream shards the table on dim 1, not dim 0:

"layers.*.ple.ple_embedding.ngram_embedding": "colwise_gather_output",

The upstream comment gives the reason: the checkpoint shards are on dim 0, so sharding TP on dim 1 keeps concatenation and TP off the same axis. A host-resident implementation that consolidates the 128 shards into one flat file inherits this constraint, because dim 0 is the row axis the gather addresses.

Terms

Term Meaning
Doorkeeper Bloom filter placed before a frequency sketch to absorb single-touch keys
Hash head One of several independent hash functions over the same n-gram
Indicator Statistic that sets the W-TinyLFU window boundary at run time
PLE Per-layer embedding. A table one decoder layer gathers from per token
QSA Qwen sparse attention, enabled per layer by indexer_n_heads
Read amplification Bytes moved divided by bytes needed
Resident set The rows an engine chooses to hold in memory, as against page-cache residency
W-TinyLFU LRU window in front of a frequency-tested SLRU main cache

Reproduce the geometry

Two published configuration fields determine the whole layout.

import math

BASE = 20_000_000   # ngram_vocab_size_base
HEADS = 16          # (ngram_size - 1) * heads_per_ngram
DIM = 160           # ple_embed_dim // HEADS
SHARDS = 128        # split_ngram_parts

def is_prime(v):
    if v < 2:
        return False
    if v % 2 == 0:
        return v == 2
    return all(v % d for d in range(3, math.isqrt(v) + 1, 2))

def primes_after(start, count):
    out, p = [], start
    while len(out) < count:
        p += 1
        if is_prime(p):
            out.append(p)
    return out

sizes = primes_after(BASE - 1, HEADS)
total = sum(sizes)
padded = math.ceil(total / SHARDS) * SHARDS

print("head vocabularies:", sizes[0], sizes[1], sizes[2], "...")
print("logical rows:", total)
print("padded rows:", padded, "= %d x %d" % (SHARDS, padded // SHARDS))
print("parameters:", padded * DIM)
print("bytes at bf16:", padded * DIM * 2)
print("useful bytes per token:", HEADS * DIM * 2)
print("bytes moved per token:", HEADS * 4096)

Output:

head vocabularies: 20000003 20000023 20000033 ...
logical rows: 320001446
padded rows: 320001536 = 128 x 2500012
parameters: 51200245760
bytes at bf16: 102400491520
useful bytes per token: 5120
bytes moved per token: 65536

References

The architecture pattern

  • Gemma 3n model overview. Per-layer embedding data can live outside model memory, on fast storage, and join each layer as it runs.
  • Sadhukhan, Cao, Dong, Zhao, Purpura-Pontoniere, Tian, Liu and Chen, "STEM: Scaling Transformers with Embedding Modules", arXiv:2601.10639, 2026. Replaces the FFN up-projection with a layer-local embedding lookup and reports CPU offload with asynchronous prefetch.

Hashing

  • Weinberger, Dasgupta, Attenberg, Langford and Smola, "Feature Hashing for Large Scale Multitask Learning", arXiv:0902.2206, 2009. Tail bounds for hashing a large vocabulary into fixed buckets.
  • Svenstrup, Hansen and Winther, "Hash Embeddings for Efficient Word Representations", arXiv:1709.03933, 2017. The multi-head form, in which each token draws k vectors from a shared pool.
  • Steele, Lea and Flood, "Fast splittable pseudorandom number generators", OOPSLA 2014, pages 453 to 472, doi:10.1145/2660193.2660195. Source of SplitMix64, used here only to derive per-position constants.

Cache policy

  • Einziger, Friedman and Manes, "TinyLFU: A Highly Efficient Cache Admission Policy", ACM Transactions on Storage 13(4), 2017, doi:10.1145/3149371, also arXiv:1512.00727. W-TinyLFU, the doorkeeper, the reset rule, and the published window and segment proportions.
  • Einziger, Eytan, Friedman and Manes, "Adaptive Software Cache Management", Middleware 2018, doi:10.1145/3274808.3274816. Adapting the window beats adapting the sketch increment, and the indicator converges faster than hill climbing.

Quantisation

  • Ashkboos, Mohtashami, Croci, Li, Cameron, Jaggi, Alistarh, Hoefler and Hensman, "QuaRot: Outlier-Free 4-Bit Inference in Rotated LLMs", arXiv:2404.00456, 2024. Rotation removes outliers without changing the output, which is the invariance the offline fold relies on.
  • Chee, Cai, Kuleshov and De Sa, "QuIP: 2-Bit Quantization of Large Language Models With Guarantees", arXiv:2307.13304, 2023. Incoherence processing by random orthogonal multiplication.
  • Zandieh, Daliri, Hadian and Mirrokni, "TurboQuant: Online Vector Quantization with Near-optimal Distortion Rate", arXiv:2504.19874, 2025. A rotation induces a concentrated Beta marginal, after which a fixed scalar quantiser approaches the distortion bound without calibration data.
  • Gao and Long, "RaBitQ: Quantizing High-Dimensional Vectors with a Theoretical Error Bound for Approximate Nearest Neighbor Search", arXiv:2405.12497, 2024. The bound is on distance estimation, so it does not transfer to a projection operand without argument.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment