diff --git a/.gitignore b/.gitignore index fd4098c..0de31f0 100644 --- a/.gitignore +++ b/.gitignore @@ -29,6 +29,8 @@ weights/ models/ Model/ !Tests/**/Model/ +# Committed parity goldens are small (KiB-scale) .safetensors test fixtures, not weights. +!Tests/Mference/Fixtures/**/*.safetensors # Repacked .gturbo directories (manifest.json/layout.json inside aren't caught by *.bin) *.gturbo/ diff --git a/Scripts/parity/README.md b/Scripts/parity/README.md new file mode 100644 index 0000000..0c49e83 --- /dev/null +++ b/Scripts/parity/README.md @@ -0,0 +1,312 @@ +# qwen38flashnext reference-parity golden harness (bring-up kit W3.2) + +Goldens for the `qwen38flashnext` (upstream `qwen4_exp`) port, captured from the **installed +`transformers` package** running a toy `Qwen4ExpForCausalLM` on CPU in float32. + +Contract being pinned: [`docs/superpowers/specs/2026-09-01-qwen38flashnext-runtime-design.md`](../../docs/superpowers/specs/2026-09-01-qwen38flashnext-runtime-design.md). +Family dossier: [`docs/families/QWEN38_FLASH_NEXT.md`](../../docs/families/QWEN38_FLASH_NEXT.md). + +- Generator: [`qwen4exp_make_goldens.py`](qwen4exp_make_goldens.py) (committed) +- Goldens: `Tests/Mference/Fixtures/qwen4exp/` (committed, 2.5 MiB) — captured from the + float32 init weights +- Goldens: `Tests/Mference/Fixtures/qwen4exp-bf16/` (committed, 2.5 MiB) — captured from the + same weights rounded to bfloat16, i.e. the ones the emitted checkpoint carries. **This is + the set a port gates on**; see "Two golden sets" below +- Toy checkpoints: `scratch/qwen4exp-toy-ckpt{,-prodlayout}/` (**not** committed, regenerable) +- venv: `scratch/qwen4exp-parity-venv/` (**not** committed) + +## Recorded reference version + +| | | +|---|---| +| `transformers` | `5.16.0.dev0` | +| `transformers` git commit | **`4da05482135896a529d5536c3c003102d36528a2`** (`refs/heads/main` at capture time) | +| `torch` | `2.13.0` (CPU) | +| Python | 3.12 | +| attn implementation | `eager` | +| dtype | float32 forward, bfloat16 checkpoint | + +The commit is also carried in `goldens-manifest.json` under `reference.git_commit`, read from +the dist-info `direct_url.json` that pip/uv writes for a VCS install. + +> The runtime design doc quotes `v5.8.0.dev0` (the version the pinned production checkpoint +> declares). The installed `main` is newer, but `models/qwen4_exp/modeling_qwen4_exp.py` and +> `configuration_qwen4_exp.py` are **byte-identical** to the reference copies the doc was read +> from, so there is no semantic drift. Re-check with `diff` after any reinstall. + +## Environment setup + +```bash +cd /path/to/Mference # repo root; scratch/ is gitignored + +uv venv --python 3.12 scratch/qwen4exp-parity-venv +VIRTUAL_ENV=scratch/qwen4exp-parity-venv \ + uv pip install torch --index-url https://download.pytorch.org/whl/cpu +VIRTUAL_ENV=scratch/qwen4exp-parity-venv \ + uv pip install numpy safetensors "git+https://github.com/huggingface/transformers" +``` + +`python -m venv` + `pip install` works identically; `uv` is just faster. To pin the exact +commit these goldens were captured against rather than tracking `main`: + +```bash +VIRTUAL_ENV=scratch/qwen4exp-parity-venv uv pip install \ + "git+https://github.com/huggingface/transformers@4da05482135896a529d5536c3c003102d36528a2" +``` + +Confirm the model classes are present before anything else — if this fails, stop: + +```bash +./scratch/qwen4exp-parity-venv/bin/python -c \ + "from transformers import Qwen4ExpForCausalLM, Qwen4ExpTextConfig; print('ok')" +``` + +## Regenerating the goldens + +```bash +./scratch/qwen4exp-parity-venv/bin/python Scripts/parity/qwen4exp_make_goldens.py --emit-checkpoint +``` + +Output is **byte-reproducible**: same seed, single-threaded, +`torch.use_deterministic_algorithms(True)`. Verify: + +```bash +find Tests/Mference/Fixtures/qwen4exp scratch/qwen4exp-toy-ckpt scratch/qwen4exp-toy-ckpt-prodlayout \ + -type f | sort | xargs shasum -a 256 > /tmp/r1.sha +rm -rf scratch/qwen4exp-toy-ckpt scratch/qwen4exp-toy-ckpt-prodlayout Tests/Mference/Fixtures/qwen4exp +./scratch/qwen4exp-parity-venv/bin/python Scripts/parity/qwen4exp_make_goldens.py --emit-checkpoint +find Tests/Mference/Fixtures/qwen4exp scratch/qwen4exp-toy-ckpt scratch/qwen4exp-toy-ckpt-prodlayout \ + -type f | sort | xargs shasum -a 256 | diff - /tmp/r1.sha && echo OK +``` + +`--print-hashes` prints the golden hashes without writing them. Checkpoint emission is +unconditional (the manifest embeds the checkpoints' tensor names, shapes and sha256), so +`--emit-checkpoint` is retained for the documented interface rather than gating anything. + +## Two golden sets: `--weight-dtype fp32` and `--weight-dtype bf16` + +```bash +# the original set, Tests/Mference/Fixtures/qwen4exp/ (default) +./scratch/qwen4exp-parity-venv/bin/python Scripts/parity/qwen4exp_make_goldens.py +# the checkpoint-faithful set, Tests/Mference/Fixtures/qwen4exp-bf16/ +./scratch/qwen4exp-parity-venv/bin/python Scripts/parity/qwen4exp_make_goldens.py \ + --weight-dtype bf16 +``` + +The default set is captured from the **float32 weights `Qwen4ExpForCausalLM(cfg)` +initialized**. The checkpoint this harness emits is a **lossy bfloat16 copy** of those +weights, so no consumer of the checkpoint can reproduce them. Measured by running the +reference twice, once with each weight dtype: + +| | SHORT | LONG | +|---|---|---| +| logits max-abs, bf16 weights vs the fp32 goldens | `1.16e-3` | `3.09e-2` | + +against the `atol = rtol = 1e-4` gate this manifest recommends — 10x to 300x over. And it +is not only a tolerance problem, because the discrete gates move too: + +* `layer02.router_indices` flips at LONG query 18 (`[2,1]` → `[1,2]`), +* `layer03.indexer_selected` changes at LONG query 28 (`[8,9,10,11,…]` → `[0,1,2,3,…]`), +* the LONG cached-decode greedy rollout diverges at token 7: `…28, 48, 36, 14, 41…` + becomes `…28, 2, 56, 43, 86…`. + +`--weight-dtype bf16` rounds every parameter through bfloat16 and back to float32 **before** +any golden is captured, so the goldens describe the weights the checkpoint carries. The +forward is float32 in both cases; only the stored weight values differ. Rounding is +idempotent, so `emit_checkpoint` writes byte-identical checkpoints either way, and the +eight fp32 golden files still reproduce byte-for-byte. + +Use the `fp32` set as the record of the reference's own arithmetic — it is what the design +contract was read against. Use the `bf16` set to gate a port: it is the one +`FlashNextReferenceParityTests` runs, because that suite loads an install built from the +emitted checkpoint. + +The script **fails loudly** rather than writing bad goldens. It asserts, in-process: + +1. `layer_multipliers`, `ngram_heads_vocab_sizes` and `ngram_heads_offsets` equal an + independent pure-python splitmix64 / prime-search derivation from the spec. +2. Every PLE n-gram embedding row id equals an independent pure-python reimplementation of + `_shift_right_ignore_eos` + the XOR-mix + per-head modulo. +3. The SHORT prompt's indexer selection equals the visible set at **every** query position of + **every** attention layer (dense-equivalent regime). +4. The LONG prompt's selection excludes visible positions at **every** attention layer + (genuinely sparse regime), and never selects a non-visible position. +5. The cached decode rollout equals a no-cache re-prefill rollout, token for token. +6. The bf16 checkpoint round-trips through `from_pretrained` with every tensor unchanged + (this is what proves the `split_ngram_parts` shard concatenation is correct). +7. The fused `gate_up_proj` row order is gate-then-up. +8. Committed goldens stay under 8 MiB. + +## Toy configuration + +Full config is in `goldens-manifest.json` under `config`. Shape summary: `hidden_size` 64, +6 layers `[lin, lin(PLE), lin, qsa, lin, qsa]`, 4 attention heads / 2 KV heads / `head_dim` 16, +`partial_rotary_factor` 0.25 (rotary_dim 4, θ 1e4), GDN `Hk=2 dk=8 / Hv=4 dv=8` conv 4, +`hc_count` 4 / `hc_lowrank` 8 (stream width 256), PLE on one-indexed layer 2 with +`ngram_size` 3 / `heads_per_ngram` 8 / `ngram_vocab_size_base` 97 (padded table 2176×4), +QSA `2` query heads / `1` key head / `head_dim` 8 / `budget` 8 / `compress_ratio` 4 +(`block_topk` 2), MoE 8 experts top-2, `vocab_size` 128, `eos_token_id` 0. + +Three deviations from the bring-up-kit suggestion (also listed in the manifest under +`deviations_from_bringup_kit`); **no config value was rejected by the `@strict` validators**: + +| Suggested | Used | Why | +|---|---|---| +| 4 layers, one attention layer | **6 layers**, two attention layers | The runtime design doc requires ≥2 attention layers so the per-layer indexer key cache and `layer_idx` plumbing are exercised. PLE still lands on one-indexed layer 2 (`layers[1]`), a `linear_attention` layer, as the validator demands. | +| SHORT = 12 tokens | **11 tokens** | With `budget` 8 / `ratio` 4 the `block_topk` is 2, so a 12-token prompt already has 3 complete blocks at its last query and is **not** dense-equivalent. 11 is the largest dense-equivalent length. | +| (unspecified) `mrope_section` | **`[1, 1, 0]`** | `rotary_dim` is 4 → only 2 frequency pairs, so the transformers default `[11, 11, 10]` cannot be expressed. Text-only positions collapse the sections, so this is semantically inert. | + +Prompts (exact ids in the manifest under `prompts`): SHORT is 11 non-EOS tokens; LONG is 48 +tokens with `eos_token_id` at index 20 so the PLE EOS-segmentation path is exercised. Both get +a 16-step greedy decode continuation. + +Selection-regime coverage across the four golden files: + +| | prefill | cached decode | +|---|---|---| +| SHORT | **dense-equivalent** (`selected == visible` everywhere) — this is the byte-gate A/B | sparse from step 0 (context passes 11 tokens immediately) | +| LONG | sparse from query 11 onward on both attention layers | sparse throughout | + +So the goldens cover dense prefill, sparse prefill, and sparse cached decode. There is no +dense-equivalent *decode* case by construction: with `block_topk` 2 any decode step past a +12-token context is already sparse. + +## Golden file contents + +All float tensors are float32, batch dimension dropped. +`T` = 11 (short) / 48 (long); `S` = 256 (`hc_count × hidden`); `H` = 64. + +### `prefill_{short,long}.safetensors` — 65 tensors + +Captured from a single `use_cache=False` forward over the whole prompt. + +| Key | Shape | Contents | +|---|---|---| +| `embed_out` | `[T, H]` | `embed_tokens` output, **before** the `repeat(1,1,hc_count)` | +| `layer{NN}.attn_hc_stream_in` | `[T, S]` | hyper stream entering the attention HC (post-PLE on layer 01) | +| `layer{NN}.attn_hc_mixed` | `[T, H]` | HC mix → attention/GDN block input | +| `layer{NN}.attn_hc_inject` | `[T, 4]` | `2·sigmoid(W_inject·h_n / hc_count)` injection weights | +| `layer{NN}.block_out` | `[T, H]` | GDN output (linear layers) or attention output after `o_proj` and the sigmoid gate (attention layers) | +| `layer{NN}.mlp_hc_stream_in` | `[T, S]` | hyper stream entering the MLP HC | +| `layer{NN}.mlp_hc_mixed` | `[T, H]` | HC mix → MoE block input | +| `layer{NN}.mlp_hc_inject` | `[T, 4]` | MLP injection weights | +| `layer{NN}.router_weights` | `[T, 2]` | top-k probs **after** renormalization | +| `layer{NN}.moe_out` | `[T, H]` | routed experts + gated shared expert | +| `layer{NN}.stream_out` | `[T, S]` | hyper stream leaving the layer | +| `layer01.ple_ngram_embeds` | `[T, 64]` | 16 concatenated n-gram head rows | +| `layer01.ple_out` | `[T, S]` | PLE contribution added to the stream, pre-attention | +| `last_hidden_state` | `[T, H]` | global `hyper_connection_mixer` output (there is no final norm) | +| `logits` | `[T, 128]` | `lm_head(last_hidden_state)` | + +### `decode_{short,long}.safetensors` — 66 tensors + +Greedy decode with a `DynamicCache`. **Row `i` is decode step `i+1`**: the first generated +token comes from the cached prefill leg (`cached_prefill_logits_last`, `[1, 128]`), which is +captured without hooks, so the per-layer decode tensors have `16 - 1 = 15` rows. Keys are the +same as above with `T = 15`, plus `step_logits` `[15, 128]`. + +### `integers_prefill_{short,long}.json`, `integers_decode_{short,long}.json` + +| Key | Contents | +|---|---| +| `layer{03,05}.indexer_selected` | Per query position, the **sorted list of selected KV token indices** — the QSA selection set. This is the gate the Swift indexer must match *exactly as an integer set*, before any attention math. | +| `layer{03,05}.indexer_visible` | Per query position, the causal-mask visible set. `selected ⊆ visible` always; `selected == visible` iff dense-equivalent. | +| `layer{NN}.router_indices` | Per token, the top-2 expert indices (paired with `router_weights`). | +| `layer01.ple_ngram_row_ids` | Per token, the 16 **absolute row indices** into the padded n-gram table (already offset by `ngram_heads_offsets`). Subtract the offset to recover the per-head modulo. | +| `argmax_all_positions`, `next_token` | Prefill only. | +| `generated_token_ids` | Decode only — the 16 greedy tokens from the cached loop. | +| `uncached_rollout_token_ids` | Decode only — the same 16 tokens from a full re-prefill each step. Equality is asserted in-script; it is the reference's own cache-equivalence proof. | + +In the **prefill** files a per-layer value is indexed `[query_position][...]`. In the **decode** +files it gains one outer level, `[decode_step][query_position][...]`, with 15 steps (same +off-by-one as the decode tensors) and exactly one query position per step. So +`integers_decode_long.json["layer03.indexer_selected"][0][0]` is the selected KV set for the +first hooked decode step. `generated_token_ids` and `uncached_rollout_token_ids` are flat lists +of 16 — they include the token produced by the cached prefill leg. + +### `goldens-manifest.json` + +`reference` (versions/commit/flags), `dtype_policy`, `config` (the full toy config), +`prompts`, `deviations_from_bringup_kit`, `checkpoint_naming`, `port_hazards`, +`golden_files` (size + sha256 + description of each file), and `findings` (the results of +every in-script assertion, including the PLE hash derivation and the indexer regime stats). + +## Toy checkpoint (uncommitted, in `scratch/`) + +Two copies, both bf16, both regenerated deterministically: + +- **`scratch/qwen4exp-toy-ckpt/`** — what `Qwen4ExpForCausalLM.save_pretrained` actually emits. + Loadable with `from_pretrained`; the harness round-trips it. +- **`scratch/qwen4exp-toy-ckpt-prodlayout/`** — the same weights renamed and re-fused into the + **production** checkpoint layout, for the repacker. + +⚠️ **The two layouts differ in two ways the repacker cares about**, both recorded in the +manifest under `checkpoint_naming`: + +1. **Prefix.** The text-only `ForCausalLM` emits `model.layers.{L}.*` with **no + `language_model` segment**, because production Qwen4-Exp ships the multimodal + `Qwen4ExpForConditionalGeneration` wrapper and therefore uses + `model.language_model.layers.{L}.*` (see `docs/families/qwen38flashnext.tensors.json`). +2. **Experts.** `save_pretrained` **de-fuses** the MoE experts back to the source layout, + `mlp.experts.{e}.{gate,up,down}_proj.weight` — one tensor per expert — via the `qwen2_moe` + `WeightConverter`. Production ships them **fused**: `mlp.experts.gate_up_proj` + `[E, 2·I, H]` and `mlp.experts.down_proj` `[E, H, I]`. + +The n-gram table is written as `...ngram_embedding.shard_{S}.weight` (`split_ngram_parts` +shards, concatenated on dim 0 at load) in both copies. `layer_multipliers`, +`ngram_heads_vocab_sizes` and `ngram_heads_offsets` stay `I64` even in a bf16 checkpoint. + +## Reference dtype behaviour the port must match + +Determined by reading the installed package, and reflected in `dtype_policy`: + +- **`Qwen3_5RMSNorm` (hence every `Qwen4ExpTextRMSNorm`) upcasts to fp32 internally**: + `_norm(x.float()) * (1.0 + weight.float())`, then `.type_as(x)`. This resolves the open + question in the runtime design doc. The zero-centered `(1 + w)` convention is confirmed, so + the planned `+1` bake at repack is correct — but the fp32 upcast must be preserved. +- **`Qwen3_5RMSNormGated` (`linear_attn.norm`) is different**: it is **ones**-initialized, not + zero-centered — *do not* bake `+1` into it. It upcasts for the variance, applies the weight + in the *input* dtype, then multiplies by `act(gate.to(float32))`. +- **Indexer**: block pooling is `key_groups.float().mean(1)` and then cast **back to the key + dtype** before `k_layernorm`; scoring is fp32 (`q.float() @ k.float()`). RoPE on the pooled + key is applied at the **block's first position**. +- **Router**: `softmax(logits, dtype=float32)` → `topk` **of the probs** → renormalize → cast + back to the logits dtype. +- **eager vs sdpa** on identical fp32 inputs: max abs logit delta `7.45e-8`, argmax identical. + That is the reference's own kernel spread and the floor for any port tolerance; the manifest + recommends `atol = rtol = 1e-4` for the fp32 Swift gates. + +## Port hazards (manifest: `port_hazards`) + +- The n-gram hash mix uses the **full positive int64 range by construction**: + `(vocab_size − 1) × max(layer_multipliers)` has bit length 63 — **zero headroom**. Exact + `Int64`/`UInt64` arithmetic only; `Double` cannot represent these products and `Int32` + overflows immediately. This holds for the production vocab too, since `multiplier_max` is + derived as `(2^63 − 1) // vocab_size`. +- **A query is not guaranteed to attend to itself.** When the query completes a block and that + block loses the top-k, the query's own key is absent from the selected set — observed at LONG + query 47 on both attention layers (8 selected of 48 visible, self excluded). Do not add an + "always keep self" shortcut. +- Only **complete** `compress_ratio` blocks participate; the 1..`ratio`−1 tail is always kept, + and the tail is taken from the *visible index list*, so padding changes which tokens are tail. +- PLE segmentation treats `eos_token_id` as **inclusive**: a shift may not cross an EOS and the + EOS position itself starts the next segment; positions that would read across the boundary + read `eos_token_id`. Production `eos_token_id` is `248044` (`text_config`). +- Attention `q_proj` packs query and gate **per head**: the output is viewed as + `(.., num_heads, 2 · head_dim)` and chunked in 2 on the last dim, so within each head the + first `head_dim` is query and the second is gate. It is **not** a global half/half split. +- **`torch.topk`'s tie-break is `std::nth_element`'s, and it is not lowest-index-first.** + PyTorch's CPU topk takes the `nth_element` branch whenever `k * 64 > n` (always, at indexer + widths), and libc++'s `__nth_element` short-circuits by length: `2` swaps if out of order, + `3` runs `__sort3`, `<= 7` runs `__selection_sort` (first maximum, swapped into place — the + swap displaces whatever was there), larger ranges run a median-of-3 quickselect. So for + block scores `[0, 0, 0.385, 0]` with `k = 2` it returns blocks `{2, 1}`, not `{2, 0}`. + + This bites only on *bit-equal* scores. Across all four golden runs there are exactly **four** + boundary ties, all in the LONG prefill, all with the tied score exactly `0.0` after the ReLU: + `layer03 q11`, `layer05 q11`, `layer05 q17`, `layer05 q18` (3 and 4 candidate blocks). A + lowest-index-first port fails `layer05.indexer_selected` at query 17 and, through the + changed indexer key at token 11, cascades into `layer03`/`layer05` selections at queries 33, + 40 and 47 and into `logits` at `3.6e-2`. `FlashNextReferenceRunner.descendingTopK` + reproduces the length-2/3/<=7 branches and records every boundary tie it sees, so a future + failure can be attributed rather than guessed at. diff --git a/Scripts/parity/qwen4exp_make_goldens.py b/Scripts/parity/qwen4exp_make_goldens.py new file mode 100755 index 0000000..9e32ed8 --- /dev/null +++ b/Scripts/parity/qwen4exp_make_goldens.py @@ -0,0 +1,1083 @@ +#!/usr/bin/env python3 +"""Reference-parity golden harness for the ``qwen38flashnext`` family (bring-up kit W3.2). + +Builds a *toy* ``Qwen4ExpTextConfig`` that exercises every module the Swift port has to +reproduce (hyper-connections, PLE n-gram hashing, the QSA indexer, MoE routing, GDN and +gated full attention), instantiates ``Qwen4ExpForCausalLM`` from it with a fixed seed, and +captures per-module goldens from the *installed* ``transformers`` package. + +Everything is float32 on CPU, single-threaded, ``model.eval()`` under ``torch.no_grad()``. +Goldens are byte-reproducible: running the script twice produces identical output files. + +Usage +----- + Scripts/parity/qwen4exp_make_goldens.py # write goldens + Scripts/parity/qwen4exp_make_goldens.py --emit-checkpoint # ... and the toy bf16 ckpt + Scripts/parity/qwen4exp_make_goldens.py --print-hashes # hash outputs, write nothing + +See Scripts/parity/README.md for the venv setup and the meaning of each golden file. +""" + +from __future__ import annotations + +import argparse +import hashlib +import json +import math +import os +import shutil +import sys +from collections import OrderedDict +from pathlib import Path + +# --------------------------------------------------------------------------------------- +# Determinism: must be configured before torch does any work. +# --------------------------------------------------------------------------------------- +os.environ.setdefault("PYTHONHASHSEED", "0") +os.environ.setdefault("CUBLAS_WORKSPACE_CONFIG", ":4096:8") +os.environ.setdefault("TOKENIZERS_PARALLELISM", "false") + +import torch # noqa: E402 +from safetensors.torch import save_file # noqa: E402 + +SEED = 1234 +REPO_ROOT = Path(__file__).resolve().parents[2] +FIXTURES = REPO_ROOT / "Tests" / "Mference" / "Fixtures" / "qwen4exp" +FIXTURES_BF16 = REPO_ROOT / "Tests" / "Mference" / "Fixtures" / "qwen4exp-bf16" +CKPT_DIR = REPO_ROOT / "scratch" / "qwen4exp-toy-ckpt" +CKPT_PROD_DIR = REPO_ROOT / "scratch" / "qwen4exp-toy-ckpt-prodlayout" + +# --------------------------------------------------------------------------------------- +# Toy configuration. +# +# Deviations from the bring-up-kit suggestion, and why: +# * num_hidden_layers 4 -> 6 with layer_types [lin, lin, lin, qsa, lin, qsa]. The runtime +# design doc requires ">=2 attention layers" so the per-layer indexer key cache and the +# layer_idx plumbing are actually exercised. PLE still lands on one-indexed layer 2 +# (= layers[1]), a linear_attention layer, as the validator demands. +# * SHORT prompt 12 -> 11 tokens. With indexer_budget 8 / compress_ratio 4 the block_topk +# is 2, so a 12-token prompt already has 3 complete blocks at the last query position and +# is NOT dense-equivalent. 11 tokens gives at most 2 complete blocks (+3 tail) at every +# query, which is the dense-equivalent regime the A/B gate needs. +# * rope_parameters.mrope_section [1, 1, 0]. rotary_dim is 4, so there are only 2 frequency +# pairs and the transformers default [11, 11, 10] cannot be expressed; the sections +# collapse for text-only positions anyway. +# Every other value is exactly as specified. No validator rejected anything. +# --------------------------------------------------------------------------------------- +TOY_CONFIG = dict( + vocab_size=128, + hidden_size=64, + num_hidden_layers=6, + layer_types=[ + "linear_attention", + "linear_attention", # <- PLE layer (one-indexed id 2) + "linear_attention", + "qwen_sparse_attention", + "linear_attention", + "qwen_sparse_attention", + ], + num_attention_heads=4, + num_key_value_heads=2, + head_dim=16, + hidden_act="silu", + output_gate_type="sigmoid", + rms_norm_eps=1e-6, + max_position_embeddings=512, + rope_parameters={ + "rope_type": "default", + "rope_theta": 10000.0, + "partial_rotary_factor": 0.25, + "mrope_section": [1, 1, 0], + }, + linear_key_head_dim=8, + linear_num_key_heads=2, + linear_value_head_dim=8, + linear_num_value_heads=4, + linear_conv_kernel_dim=4, + hc_count=4, + hc_lowrank=8, + ple_layer_ids=[2], + ple_embed_dim=64, + ple_conv_kernel_size=4, + ngram_size=3, + heads_per_ngram=8, + ngram_vocab_size_base=97, + make_ngram_vocab_size_divisible_by=128, + split_ngram_parts=2, + seed=SEED, + indexer_n_heads=2, + indexer_kv_heads=1, + indexer_head_dim=8, + indexer_budget=8, + indexer_compress_ratio=4, + num_experts=8, + num_experts_per_tok=2, + moe_intermediate_size=32, + shared_expert_intermediate_size=32, + norm_topk_prob=True, + eos_token_id=0, + bos_token_id=0, + pad_token_id=None, + tie_word_embeddings=False, + initializer_range=0.02, + attention_bias=False, + attention_dropout=0.0, + use_cache=True, +) + +SHORT_LEN = 11 +LONG_LEN = 48 +DECODE_STEPS = 16 +# LONG_EOS_AT places an eos_token_id inside the long prompt so the PLE n-gram shift's +# EOS-segmentation path (`_shift_right_ignore_eos`) is exercised, not just the prefix pad. +LONG_EOS_AT = 20 + + +# --------------------------------------------------------------------------------------- +# Independent (pure-python) reimplementation of the PLE hash path, per the runtime design +# doc. Asserted against the model's own buffers -- if these ever diverge the port is +# reading the spec wrong, or the reference changed. +# --------------------------------------------------------------------------------------- +_MASK64 = (1 << 64) - 1 +_GAMMA = 0x9E3779B97F4A7C15 +_M1 = 0xBF58476D1CE4E5B9 +_M2 = 0x94D049BB133111EB +_PRIME_1 = 10007 + + +def splitmix64(value: int) -> int: + value = (value + _GAMMA) & _MASK64 + value = ((value ^ (value >> 30)) * _M1) & _MASK64 + value = ((value ^ (value >> 27)) * _M2) & _MASK64 + return (value ^ (value >> 31)) & _MASK64 + + +def derive_layer_multipliers(unigram_vocab_size: int, ngram_size: int, ple_layer_index: int, seed: int) -> list[int]: + max_long = (1 << 63) - 1 + multiplier_max = max_long // max(unigram_vocab_size, 1) + half_bound = max(1, multiplier_max // 2) + base_seed = seed + _PRIME_1 * ple_layer_index + out = [] + for i in range(ngram_size): + out.append(2 * (splitmix64((base_seed + _GAMMA * (i + 1)) & _MASK64) % half_bound) + 1) + return out + + +def is_prime(v: int) -> bool: + if v < 2: + return False + if v % 2 == 0: + return v == 2 + for d in range(3, math.isqrt(v) + 1, 2): + if v % d == 0: + return False + return True + + +def nth_prime_after(start: int, count: int) -> int: + p = start + for _ in range(count): + p += 1 + while not is_prime(p): + p += 1 + return p + + +def derive_head_tables(ngram_vocab_size_base: int, ngram_heads: int, ple_layer_index: int) -> tuple[list[int], list[int]]: + sizes, offsets, total = [], [], 0 + for head_idx in range(ngram_heads): + global_head_idx = ple_layer_index * ngram_heads + head_idx + size = nth_prime_after(ngram_vocab_size_base - 1, global_head_idx + 1) + sizes.append(size) + offsets.append(total) + total += size + return sizes, offsets + + +def shift_right_ignore_eos(ids: list[int], shift: int, eos: int) -> list[int]: + """Pure-python mirror of Qwen4ExpTextNGramEmbedding._shift_right_ignore_eos.""" + if shift == 0: + return list(ids) + n = len(ids) + prev_eos_inclusive, cur = [], -1 + for i, t in enumerate(ids): + if t == eos: + cur = i + prev_eos_inclusive.append(cur) + previous_eos = [-1] + prev_eos_inclusive[:-1] + out = [] + for i in range(n): + position_in_segment = i - (previous_eos[i] + 1) + source = i - shift + out.append(ids[source] if (position_in_segment >= shift and source >= 0) else eos) + return out + + +def derive_ngram_row_ids( + token_ids: list[int], + *, + ngram_size: int, + heads_per_ngram: int, + multipliers: list[int], + head_vocab_sizes: list[int], + head_offsets: list[int], + eos: int, +) -> list[list[int]]: + """Pure-python mirror of Qwen4ExpTextNGramEmbedding.forward -> final embedding row ids. + + Returns [seq_len][ngram_heads] absolute row indices into the padded n-gram table. + Prefixes the sequence with (ngram_size - 1) eos tokens, exactly like the fresh-cache path. + """ + context_len = ngram_size - 1 + history = [eos] * context_len + list(token_ids) + shifted = [shift_right_ignore_eos(history, s, eos) for s in range(ngram_size)] + + blocks: list[list[int]] = [] + for ngram in range(2, ngram_size + 1): + start = (ngram - 2) * heads_per_ngram + mixed = [t * multipliers[0] for t in shifted[0]] + for position in range(1, ngram): + mixed = [m ^ (t * multipliers[position]) for m, t in zip(mixed, shifted[position])] + block = [ + [(m % head_vocab_sizes[start + h]) + head_offsets[start + h] for h in range(heads_per_ngram)] + for m in mixed + ] + blocks.append(block) + + rows = [sum((blk[i] for blk in blocks), []) for i in range(len(history))] + return rows[-len(token_ids):] + + +# --------------------------------------------------------------------------------------- +# Capture plumbing. +# --------------------------------------------------------------------------------------- +class Capture: + """Collects float tensors (-> safetensors) and integer structures (-> json).""" + + def __init__(self) -> None: + self.floats: "OrderedDict[str, torch.Tensor]" = OrderedDict() + self.ints: "OrderedDict[str, object]" = OrderedDict() + self.prefix = "" + + def put_float(self, name: str, tensor: torch.Tensor) -> None: + self.floats[self.prefix + name] = tensor.detach().to(torch.float32).contiguous().clone() + + def put_int(self, name: str, value: object) -> None: + self.ints[self.prefix + name] = value + + def append_float(self, name: str, tensor: torch.Tensor) -> None: + """Accumulate one decode step; stacked on flush.""" + key = self.prefix + name + self.floats.setdefault(key, []) + assert isinstance(self.floats[key], list), f"{key} already flushed as a tensor" + self.floats[key].append(tensor.detach().to(torch.float32).contiguous().clone()) + + def append_int(self, name: str, value: object) -> None: + key = self.prefix + name + self.ints.setdefault(key, []) + self.ints[key].append(value) + + def stack_pending(self) -> None: + for key, value in list(self.floats.items()): + if isinstance(value, list): + self.floats[key] = torch.cat(value, dim=0).contiguous() + + +def selected_sets_from_mask(mask: torch.Tensor) -> list[list[int]]: + """Indexer mask -> per-query sorted list of selected kv indices. + + `mask` is (batch, 1, q_len, kv_len); bool for sdpa, float (0 / dtype-min) for eager. + """ + keep = mask if mask.dtype == torch.bool else (mask == 0) + keep = keep[0, 0] # batch 0, broadcast head dim + return [torch.nonzero(row, as_tuple=False).flatten().tolist() for row in keep] + + +def visible_sets_from_causal(mask: torch.Tensor) -> list[list[int]]: + keep = mask if mask.dtype == torch.bool else (mask == 0) + keep = keep[0, 0] + return [torch.nonzero(row, as_tuple=False).flatten().tolist() for row in keep] + + +class Recorder: + """Registers hooks on every module whose output the Swift port must reproduce.""" + + def __init__(self, model, cap: Capture, *, streaming: bool) -> None: + self.model = model + self.cap = cap + self.streaming = streaming # True for the decode loop: append instead of assign + self.handles: list = [] + self._wire() + + def _emit_float(self, name: str, tensor: torch.Tensor) -> None: + if self.streaming: + self.cap.append_float(name, tensor[0]) # drop batch dim; steps concat on dim 0 + else: + self.cap.put_float(name, tensor[0]) + + def _emit_int(self, name: str, value: object) -> None: + if self.streaming: + self.cap.append_int(name, value) + else: + self.cap.put_int(name, value) + + def _hook(self, module, fn): + # A forward hook that returns non-None REPLACES the module output, so every hook + # here is wrapped to swallow its return value. + def wrapped(mod, inputs, output, _fn=fn): + _fn(mod, inputs, output) + return None + + self.handles.append(module.register_forward_hook(wrapped)) + + def _pre_hook(self, module, fn): + def wrapped(mod, inputs, _fn=fn): + _fn(mod, inputs) + return None + + self.handles.append(module.register_forward_pre_hook(wrapped)) + + def _wire(self) -> None: + inner = self.model.model + + self._hook(inner.embed_tokens, lambda m, i, o: self._emit_float("embed_out", o)) + self._hook( + inner.hyper_connection_mixer, + lambda m, i, o: self._emit_float("last_hidden_state", o), + ) + + for layer_idx, layer in enumerate(inner.layers): + p = f"layer{layer_idx:02d}." + + def on_layer(m, i, o, pfx=p): + self._emit_float(pfx + "stream_out", o) + + def on_attn_hc(m, i, o, pfx=p): + self._emit_float(pfx + "attn_hc_mixed", o[0]) + self._emit_float(pfx + "attn_hc_stream_in", o[1]) + self._emit_float(pfx + "attn_hc_inject", o[2]) + + def on_mlp_hc(m, i, o, pfx=p): + self._emit_float(pfx + "mlp_hc_mixed", o[0]) + self._emit_float(pfx + "mlp_hc_stream_in", o[1]) + self._emit_float(pfx + "mlp_hc_inject", o[2]) + + def on_moe(m, i, o, pfx=p): + self._emit_float(pfx + "moe_out", o) + + def on_router(m, i, o, pfx=p): + # router returns (logits, renormalized top-k weights, top-k indices) + self._emit_float(pfx + "router_weights", o[1].unsqueeze(0)) + self._emit_int(pfx + "router_indices", o[2].tolist()) + + self._hook(layer, on_layer) + self._hook(layer.attn_hyper_connection, on_attn_hc) + self._hook(layer.mlp_hyper_connection, on_mlp_hc) + self._hook(layer.mlp, on_moe) + self._hook(layer.mlp.gate, on_router) + + if layer.layer_type == "linear_attention": + + def on_gdn(m, i, o, pfx=p): + self._emit_float(pfx + "block_out", o) + + self._hook(layer.linear_attn, on_gdn) + else: + + def on_attn(m, i, o, pfx=p): + self._emit_float(pfx + "block_out", o[0]) + + def on_indexer(m, i, o, pfx=p): + self._emit_int(pfx + "indexer_selected", selected_sets_from_mask(o)) + self._emit_int(pfx + "indexer_visible", visible_sets_from_causal(i[2])) + + self._hook(layer.self_attn, on_attn) + self._hook(layer.self_attn.indexer, on_indexer) + + if layer.ple is not None: + + def on_ple(m, i, o, pfx=p): + self._emit_float(pfx + "ple_out", o) + + def on_ngram(m, i, o, pfx=p): + self._emit_float(pfx + "ple_ngram_embeds", o) + + def on_ngram_ids(m, i, pfx=p): + self._emit_int(pfx + "ple_ngram_row_ids", i[0][0].tolist()) + + self._hook(layer.ple, on_ple) + self._hook(layer.ple.ple_embedding, on_ngram) + self._pre_hook(layer.ple.ple_embedding.ngram_embedding, on_ngram_ids) + + def close(self) -> None: + for h in self.handles: + h.remove() + self.handles.clear() + + +# --------------------------------------------------------------------------------------- +# Model construction. +# --------------------------------------------------------------------------------------- +def build_config(attn_implementation: str = "eager"): + from transformers import Qwen4ExpTextConfig + + cfg = Qwen4ExpTextConfig(**TOY_CONFIG) + cfg._attn_implementation = attn_implementation + return cfg + + +def build_model(attn_implementation: str = "eager", weight_dtype: str = "fp32"): + """The toy model, in float32. + + ``weight_dtype="bf16"`` rounds every parameter through bfloat16 and back to + float32 **before** any golden is captured. That is not a cosmetic option: + the checkpoint this harness emits is bfloat16, so with the default + ``fp32`` weights the goldens describe a model whose weights **no consumer + can obtain**. A port that loads the emitted checkpoint sees the rounded + weights, and the resulting logit spread is 1.2e-3 (SHORT) / 3.1e-2 (LONG) + max-abs -- 10x to 300x the 1e-4 fp32 gate this manifest recommends. + + So the two sets have distinct jobs: + + * ``fp32`` (``Tests/Mference/Fixtures/qwen4exp``) pins the *reference's own* + arithmetic and is the artifact the design contract was read against. + * ``bf16`` (``Tests/Mference/Fixtures/qwen4exp-bf16``) pins what a port that + loads the shipped checkpoint must reproduce, and is the one a Swift + parity suite can actually gate on at 1e-4. + + The forward is float32 in both cases; only the stored weight values differ. + Rounding is idempotent (bf16 -> fp32 -> bf16 is exact), so + ``emit_checkpoint`` writes byte-identical checkpoints either way. + """ + from transformers import Qwen4ExpForCausalLM + + cfg = build_config(attn_implementation) + torch.manual_seed(SEED) + model = Qwen4ExpForCausalLM(cfg) + model = model.to(torch.float32).eval() + if weight_dtype == "bf16": + model = model.to(torch.bfloat16).to(torch.float32).eval() + elif weight_dtype != "fp32": + raise ValueError(f"unknown weight dtype {weight_dtype!r}") + for p in model.parameters(): + p.requires_grad_(False) + return model + + +def make_prompt(length: int, *, rng_seed: int, eos_at: int | None) -> list[int]: + g = torch.Generator().manual_seed(rng_seed) + ids = torch.randint(1, TOY_CONFIG["vocab_size"], (length,), generator=g).tolist() + if eos_at is not None: + ids[eos_at] = TOY_CONFIG["eos_token_id"] + return ids + + +# --------------------------------------------------------------------------------------- +# Runs. +# --------------------------------------------------------------------------------------- +@torch.no_grad() +def run_prefill(model, token_ids: list[int], cap: Capture, prefix: str): + ids = torch.tensor([token_ids], dtype=torch.long) + attn = torch.ones_like(ids) + cap.prefix = prefix + rec = Recorder(model, cap, streaming=False) + try: + out = model(input_ids=ids, attention_mask=attn, use_cache=False) + finally: + rec.close() + cap.prefix = "" + cap.prefix = prefix + cap.put_float("logits", out.logits[0]) + cap.put_int("argmax_all_positions", out.logits[0].argmax(-1).tolist()) + cap.put_int("next_token", int(out.logits[0, -1].argmax())) + cap.prefix = "" + return out + + +@torch.no_grad() +def run_decode(model, token_ids: list[int], steps: int, cap: Capture, prefix: str): + """Greedy decode with a DynamicCache; per-step goldens captured through the same hooks.""" + from transformers import DynamicCache + + cache = DynamicCache(config=model.config) + ids = torch.tensor([token_ids], dtype=torch.long) + attn = torch.ones_like(ids) + + # Prefill leg (hooks off -- the prefill goldens come from run_prefill's no-cache run). + out = model(input_ids=ids, attention_mask=attn, past_key_values=cache, use_cache=True) + cache = out.past_key_values + next_id = int(out.logits[0, -1].argmax()) + generated = [next_id] + + cap.prefix = prefix + cap.put_float("cached_prefill_logits_last", out.logits[0, -1:]) + rec = Recorder(model, cap, streaming=True) + try: + for _ in range(steps - 1): + step_ids = torch.tensor([[next_id]], dtype=torch.long) + attn = torch.cat([attn, torch.ones((1, 1), dtype=torch.long)], dim=1) + out = model( + input_ids=step_ids, + attention_mask=attn, + past_key_values=cache, + use_cache=True, + ) + cache = out.past_key_values + cap.append_float("step_logits", out.logits[0]) + next_id = int(out.logits[0, -1].argmax()) + generated.append(next_id) + finally: + rec.close() + + cap.stack_pending() + cap.put_int("generated_token_ids", generated) + cap.prefix = "" + return generated + + +@torch.no_grad() +def run_uncached_rollout(model, token_ids: list[int], steps: int) -> list[int]: + """Same greedy rollout with no cache at all -- re-prefills the whole prefix each step.""" + context = list(token_ids) + generated: list[int] = [] + for _ in range(steps): + ids = torch.tensor([context], dtype=torch.long) + out = model(input_ids=ids, attention_mask=torch.ones_like(ids), use_cache=False) + nxt = int(out.logits[0, -1].argmax()) + generated.append(nxt) + context.append(nxt) + return generated + + +# --------------------------------------------------------------------------------------- +# Assertions the harness must not ship without. +# --------------------------------------------------------------------------------------- +def assert_ple_hash_derivation(model, report: dict) -> None: + cfg = model.config + ple = model.model.layers[cfg.ple_layer_ids[0] - 1].ple + emb = ple.ple_embedding + ple_layer_index = emb.ple_layer_index + + want_mult = derive_layer_multipliers(cfg.vocab_size, cfg.ngram_size, ple_layer_index, cfg.seed) + got_mult = emb.layer_multipliers.tolist() + assert got_mult == want_mult, f"layer_multipliers mismatch: model={got_mult} derived={want_mult}" + + ngram_heads = (cfg.ngram_size - 1) * cfg.heads_per_ngram + want_sizes, want_offsets = derive_head_tables(cfg.ngram_vocab_size_base, ngram_heads, ple_layer_index) + assert emb.ngram_heads_vocab_sizes.tolist() == want_sizes, "ngram_heads_vocab_sizes mismatch" + assert emb.ngram_heads_offsets.tolist() == want_offsets, "ngram_heads_offsets mismatch" + + # No int64 overflow is possible with these multipliers: max token id * max multiplier. + headroom = (cfg.vocab_size - 1) * max(want_mult) + assert headroom < (1 << 63), "int64 overflow in the n-gram mix -- multipliers/vocab inconsistent" + + report["ple_hash_derivation"] = { + "ple_layer_index": ple_layer_index, + "layer_multipliers": [str(m) for m in want_mult], + "ngram_heads_vocab_sizes": want_sizes, + "ngram_heads_offsets": want_offsets, + "total_vocab_size": sum(want_sizes), + "padded_vocab_size": int(emb.ngram_embedding.weight.shape[0]), + "head_dim_per_ngram": int(emb.ngram_embedding.weight.shape[1]), + "max_mixed_id_headroom_bits": 63 - headroom.bit_length(), + "verified_against": "pure-python splitmix64 + prime search in this script", + } + + +def assert_indexer_regimes(cap: Capture, short_prefix: str, long_prefix: str, report: dict) -> None: + def gather(prefix: str): + return {k: v for k, v in cap.ints.items() if k.startswith(prefix) and k.endswith("indexer_selected")} + + short_sel = gather(short_prefix) + long_sel = gather(long_prefix) + assert short_sel, "no indexer selections captured for the SHORT prompt" + assert long_sel, "no indexer selections captured for the LONG prompt" + + # SHORT: dense-equivalent -- every visible position selected, at every attention layer. + for key, sel in short_sel.items(): + vis_key = key.replace("indexer_selected", "indexer_visible") + for q, (s, v) in enumerate(zip(sel, cap.ints[vis_key])): + assert set(s) == set(v), ( + f"SHORT prompt is not dense-equivalent at {key} query {q}: " + f"selected {sorted(s)} != visible {sorted(v)}" + ) + + # LONG: genuinely sparse -- at least one query drops a visible position, at every layer. + long_stats = {} + for key, sel in long_sel.items(): + vis_key = key.replace("indexer_selected", "indexer_visible") + sparse_queries = [ + q for q, (s, v) in enumerate(zip(sel, cap.ints[vis_key])) if set(s) != set(v) + ] + assert sparse_queries, f"LONG prompt selection is dense at {key} -- lower the budget or lengthen the prompt" + assert set(sel[-1]).issubset(set(cap.ints[vis_key][-1])), f"{key} selected a non-visible position" + last = len(sel) - 1 + long_stats[key] = { + "first_sparse_query": sparse_queries[0], + "num_sparse_queries": len(sparse_queries), + "last_query_selected": len(sel[last]), + "last_query_visible": len(cap.ints[vis_key][last]), + "last_query_sees_itself": last in sel[last], + } + report["indexer_regimes"] = { + "short_dense_equivalent": True, + "long": long_stats, + } + + +# --------------------------------------------------------------------------------------- +# Checkpoint emission. +# --------------------------------------------------------------------------------------- +def emit_checkpoint(model, report: dict) -> None: + """Write the toy checkpoint twice: HF-native layout and production-name layout. + + Also round-trips the HF-native copy through ``from_pretrained`` to prove the + ``split_ngram_parts`` shard concatenation reassembles the table correctly, and checks + empirically which half of the fused ``gate_up_proj`` rows is gate and which is up. + """ + import copy + + from safetensors import safe_open + from transformers import Qwen4ExpForCausalLM + + bf16 = copy.deepcopy(model).to(torch.bfloat16) + if CKPT_DIR.exists(): + shutil.rmtree(CKPT_DIR) + bf16.save_pretrained(CKPT_DIR, safe_serialization=True) + + native: "OrderedDict[str, tuple[list[int], str]]" = OrderedDict() + for path in sorted(CKPT_DIR.glob("*.safetensors")): + with safe_open(path, framework="pt") as f: + for k in sorted(f.keys()): + sl = f.get_slice(k) + native[k] = (list(sl.get_shape()), sl.get_dtype()) + + # Production layout: `model.language_model.*` prefix and FUSED expert tensors, matching + # docs/families/qwen38flashnext.tensors.json. The text-only ForCausalLM cannot emit this + # itself, so the repacker's real patterns get a fixture too. + prod: "OrderedDict[str, torch.Tensor]" = OrderedDict() + runtime_sd = bf16.state_dict() + for k, v in runtime_sd.items(): + if k.startswith("model."): + prod["model.language_model." + k[len("model."):]] = v.contiguous().clone() + else: + prod[k] = v.contiguous().clone() + # Split the n-gram table back into shards, as the real checkpoint stores it. + shard_keys = [k for k in prod if k.endswith("ple_embedding.ngram_embedding.weight")] + n_shards = model.config.split_ngram_parts + for k in shard_keys: + w = prod.pop(k) + assert w.shape[0] % n_shards == 0, "padded n-gram vocab not divisible by split_ngram_parts" + for s, chunk in enumerate(w.chunk(n_shards, dim=0)): + prod[k[: -len("weight")] + f"shard_{s}.weight"] = chunk.contiguous().clone() + + if CKPT_PROD_DIR.exists(): + shutil.rmtree(CKPT_PROD_DIR) + CKPT_PROD_DIR.mkdir(parents=True) + save_file(dict(sorted(prod.items())), str(CKPT_PROD_DIR / "model.safetensors")) + shutil.copy(CKPT_DIR / "config.json", CKPT_PROD_DIR / "config.json") + + # --- round trip: shards must reassemble, and the reloaded bf16 model must run --- + reloaded = Qwen4ExpForCausalLM.from_pretrained(CKPT_DIR, dtype=torch.bfloat16) + reloaded.eval() + emb_ref = bf16.model.layers[model.config.ple_layer_ids[0] - 1].ple.ple_embedding.ngram_embedding.weight + emb_rt = reloaded.model.layers[model.config.ple_layer_ids[0] - 1].ple.ple_embedding.ngram_embedding.weight + assert emb_rt.shape == emb_ref.shape, f"ngram table shape changed on reload: {emb_rt.shape} vs {emb_ref.shape}" + assert torch.equal(emb_rt, emb_ref), "ngram shard concatenation did not reassemble the table" + rt_missing = [k for k, v in bf16.state_dict().items() if not torch.equal(v, reloaded.state_dict()[k])] + assert not rt_missing, f"checkpoint round trip changed tensors: {rt_missing[:5]}" + + # --- fused gate|up row split, verified rather than assumed --- + experts = bf16.model.layers[0].mlp.experts + inter = model.config.moe_intermediate_size + fused = experts.gate_up_proj[0] + with safe_open(next(CKPT_DIR.glob("*.safetensors")), framework="pt") as f: + saved_gate = f.get_tensor("model.layers.0.mlp.experts.0.gate_proj.weight") + saved_up = f.get_tensor("model.layers.0.mlp.experts.0.up_proj.weight") + gate_is_first = torch.equal(fused[:inter], saved_gate) and torch.equal(fused[inter:], saved_up) + assert gate_is_first, "fused gate_up_proj row order is NOT gate-then-up -- update the repacker" + + report["checkpoint"] = { + "round_trip_verified": True, + "fused_gate_up_row_split": f"gate = rows [0, {inter}), up = rows [{inter}, {2 * inter})", + "hf_native_dir": str(CKPT_DIR.relative_to(REPO_ROOT)), + "hf_native_prefix": "model.", + "hf_native_files": sorted(p.name for p in CKPT_DIR.iterdir()), + "hf_native_tensors": {k: {"shape": s, "dtype": d} for k, (s, d) in native.items()}, + "prod_layout_dir": str(CKPT_PROD_DIR.relative_to(REPO_ROOT)), + "prod_layout_prefix": "model.language_model.", + "prod_layout_tensors": { + k: {"shape": list(v.shape), "dtype": "BF16"} for k, v in sorted(prod.items()) + }, + "sha256": { + **{f"hf_native/{p.name}": sha256_file(p) for p in sorted(CKPT_DIR.iterdir()) if p.is_file()}, + **{f"prod_layout/{p.name}": sha256_file(p) for p in sorted(CKPT_PROD_DIR.iterdir()) if p.is_file()}, + }, + } + + +def sha256_file(path: Path) -> str: + h = hashlib.sha256() + h.update(path.read_bytes()) + return h.hexdigest() + + +def sha256_bytes(b: bytes) -> str: + return hashlib.sha256(b).hexdigest() + + +# --------------------------------------------------------------------------------------- +# Main. +# --------------------------------------------------------------------------------------- +def main() -> int: + ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + ap.add_argument( + "--emit-checkpoint", + action="store_true", + help=( + "write the toy bf16 checkpoints to scratch/ (DEFAULT: always on -- the manifest " + "embeds their tensor names, shapes and sha256, so it cannot be built without them; " + "the flag is kept for the documented bring-up-kit interface)" + ), + ) + ap.add_argument("--print-hashes", action="store_true", help="print sha256 of every output and exit non-writing") + ap.add_argument( + "--weight-dtype", + choices=["fp32", "bf16"], + default="fp32", + help=( + "storage dtype the weights are rounded to before capture (the forward is float32 " + "either way). fp32 pins the reference's own arithmetic; bf16 pins what a port " + "loading the emitted bfloat16 checkpoint must reproduce. See build_model." + ), + ) + ap.add_argument("--out", default=None, help="fixture output directory") + args = ap.parse_args() + if args.out is None: + args.out = str(FIXTURES if args.weight_dtype == "fp32" else FIXTURES_BF16) + + torch.set_num_threads(1) + torch.manual_seed(SEED) + deterministic_algorithms = True + try: + torch.use_deterministic_algorithms(True) + except Exception as exc: # pragma: no cover - recorded, not fatal + deterministic_algorithms = False + print(f"[warn] torch.use_deterministic_algorithms(True) refused: {exc}", file=sys.stderr) + + import transformers + + report: dict = {} + model = build_model("eager", args.weight_dtype) + + short_ids = make_prompt(SHORT_LEN, rng_seed=11, eos_at=None) + long_ids = make_prompt(LONG_LEN, rng_seed=48, eos_at=LONG_EOS_AT) + + assert_ple_hash_derivation(model, report) + + cap_pre = Capture() + run_prefill(model, short_ids, cap_pre, "short.") + run_prefill(model, long_ids, cap_pre, "long.") + assert_indexer_regimes(cap_pre, "short.", "long.", report) + + # PLE row ids: assert the reference against the independent pure-python derivation. + cfg = model.config + ple_prefix = f"layer{cfg.ple_layer_ids[0] - 1:02d}.ple_ngram_row_ids" + mult = derive_layer_multipliers(cfg.vocab_size, cfg.ngram_size, 0, cfg.seed) + ngram_heads = (cfg.ngram_size - 1) * cfg.heads_per_ngram + sizes, offsets = derive_head_tables(cfg.ngram_vocab_size_base, ngram_heads, 0) + for name, ids in (("short", short_ids), ("long", long_ids)): + got = cap_pre.ints[f"{name}.{ple_prefix}"] + want = derive_ngram_row_ids( + ids, + ngram_size=cfg.ngram_size, + heads_per_ngram=cfg.heads_per_ngram, + multipliers=mult, + head_vocab_sizes=sizes, + head_offsets=offsets, + eos=cfg.eos_token_id, + ) + assert got == want, f"PLE n-gram row ids diverge from the independent derivation ({name})" + report["ple_row_ids_match_independent_derivation"] = True + + cap_dec = Capture() + gen_short = run_decode(model, short_ids, DECODE_STEPS, cap_dec, "short.") + gen_long = run_decode(model, long_ids, DECODE_STEPS, cap_dec, "long.") + + # Cache equivalence inside the reference itself. + unc_short = run_uncached_rollout(model, short_ids, DECODE_STEPS) + unc_long = run_uncached_rollout(model, long_ids, DECODE_STEPS) + assert gen_short == unc_short, f"cached decode != uncached rollout (short): {gen_short} vs {unc_short}" + assert gen_long == unc_long, f"cached decode != uncached rollout (long): {gen_long} vs {unc_long}" + cap_dec.ints["short.uncached_rollout_token_ids"] = unc_short + cap_dec.ints["long.uncached_rollout_token_ids"] = unc_long + report["cache_equivalence_in_reference"] = {"short": True, "long": True} + + # Tolerance observation: sdpa vs eager on identical fp32 inputs. + model_sdpa = build_model("sdpa", args.weight_dtype) + with torch.no_grad(): + ids = torch.tensor([long_ids], dtype=torch.long) + lo_eager = model(input_ids=ids, attention_mask=torch.ones_like(ids), use_cache=False).logits + lo_sdpa = model_sdpa(input_ids=ids, attention_mask=torch.ones_like(ids), use_cache=False).logits + report["tolerances"] = { + "dtype": "float32", + "eager_vs_sdpa_logits_max_abs": float((lo_eager - lo_sdpa).abs().max()), + "eager_vs_sdpa_logits_max_rel": float( + ((lo_eager - lo_sdpa).abs() / lo_eager.abs().clamp_min(1e-6)).max() + ), + "eager_vs_sdpa_argmax_identical": bool( + (lo_eager.argmax(-1) == lo_sdpa.argmax(-1)).all() + ), + "recommended_swift_gate_fp32": {"atol": 1e-4, "rtol": 1e-4}, + "note": ( + "Goldens are produced with attn_implementation='eager' (float additive masks). " + "sdpa uses boolean masks; the deltas above are the reference's own fp32 spread " + "between the two kernels and are the floor for any port tolerance." + ), + } + + # Checkpoint emission is unconditional: the manifest embeds its names/shapes/hashes. + emit_checkpoint(model, report) + + # ---- assemble outputs ------------------------------------------------------------- + out_dir = Path(args.out) + files: "OrderedDict[str, bytes]" = OrderedDict() + + def split(cap: Capture, want: str) -> dict: + return {k[len(want):]: v for k, v in cap.floats.items() if k.startswith(want)} + + def split_i(cap: Capture, want: str) -> dict: + return {k[len(want):]: v for k, v in cap.ints.items() if k.startswith(want)} + + payloads = { + "prefill_short.safetensors": split(cap_pre, "short."), + "prefill_long.safetensors": split(cap_pre, "long."), + "decode_short.safetensors": split(cap_dec, "short."), + "decode_long.safetensors": split(cap_dec, "long."), + } + for name, tensors in payloads.items(): + tmp = out_dir / (name + ".tmp") + out_dir.mkdir(parents=True, exist_ok=True) + save_file({k: v for k, v in sorted(tensors.items())}, str(tmp)) + files[name] = tmp.read_bytes() + tmp.unlink() + + int_payloads = { + "integers_prefill_short.json": split_i(cap_pre, "short."), + "integers_prefill_long.json": split_i(cap_pre, "long."), + "integers_decode_short.json": split_i(cap_dec, "short."), + "integers_decode_long.json": split_i(cap_dec, "long."), + } + for name, obj in int_payloads.items(): + files[name] = (json.dumps(obj, indent=1, sort_keys=True) + "\n").encode() + + tf_commit = getattr(transformers, "__commit__", None) or os.environ.get("QWEN4EXP_TRANSFORMERS_COMMIT", "") + manifest = OrderedDict( + schema="mference.qwen4exp.goldens/1", + generated_by="Scripts/parity/qwen4exp_make_goldens.py", + family="qwen38flashnext", + reference=OrderedDict( + package="transformers", + version=transformers.__version__, + git_commit=tf_commit or resolve_transformers_commit(), + model_class="Qwen4ExpForCausalLM", + model_type="qwen4_exp_text", + torch_version=torch.__version__, + device="cpu", + attn_implementation="eager", + deterministic_algorithms=deterministic_algorithms, + num_threads=1, + seed=SEED, + ), + dtype_policy=OrderedDict( + model_weights=( + "float32" + if args.weight_dtype == "fp32" + else "bfloat16 values held in float32 (every parameter rounded through bf16 " + "before capture, so these goldens describe the weights the emitted checkpoint " + "actually carries)" + ), + weight_dtype=args.weight_dtype, + forward="float32", + checkpoint="bfloat16", + rmsnorm=( + "Qwen3_5RMSNorm (and therefore every Qwen4ExpTextRMSNorm) UPCASTS to float32 " + "internally: output = _norm(x.float()) * (1.0 + weight.float()), then .type_as(x). " + "The (1 + w) zero-centered convention is confirmed; the port's `+1` weight bake is " + "correct, but the fp32 upcast must be preserved for bf16 runtimes." + ), + rmsnorm_gated=( + "Qwen3_5RMSNormGated upcasts hidden_states to fp32 for the variance, applies the " + "weight in the INPUT dtype (weight * hs.to(input_dtype)), then multiplies by " + "act(gate.to(float32)) and casts back. Its weight is ones-initialized and is NOT " + "zero-centered -- do not bake +1 into linear_attn.norm.weight." + ), + indexer=( + "Pooling is float32 (key_groups.float().mean(1)) then cast BACK to the key dtype " + "before k_layernorm; scoring is float32 (q.float() @ k.float()). RoPE on pooled " + "keys is applied at the block's FIRST position." + ), + router="softmax in float32 -> topk of the probs -> renormalize -> cast to logits dtype", + ), + config=json_safe_config(model.config), + prompts=OrderedDict( + short=OrderedDict( + token_ids=short_ids, + length=SHORT_LEN, + regime="dense-equivalent (every visible position selected at every attention layer)", + ), + long=OrderedDict( + token_ids=long_ids, + length=LONG_LEN, + eos_token_at=LONG_EOS_AT, + regime="genuinely sparse (selection excludes visible positions)", + ), + decode_steps=DECODE_STEPS, + greedy_rollout_short=gen_short, + greedy_rollout_long=gen_long, + ), + deviations_from_bringup_kit=[ + "num_hidden_layers 4 -> 6, layer_types [lin,lin,lin,qsa,lin,qsa]: the runtime design " + "doc requires >=2 attention layers so the per-layer indexer key cache is exercised. " + "PLE still sits on one-indexed layer 2 (layers[1]), a linear_attention layer.", + "SHORT prompt 12 -> 11 tokens: with indexer_budget 8 / compress_ratio 4 (block_topk 2) " + "a 12-token prompt already has 3 complete blocks at the last query and is NOT " + "dense-equivalent. 11 tokens is the largest dense-equivalent length.", + "rope_parameters.mrope_section set to [1,1,0]: rotary_dim is 4 so there are only 2 " + "frequency pairs; the transformers default [11,11,10] cannot be expressed. Text-only " + "positions collapse the sections, so this is semantically inert.", + "No config value was rejected by the @strict validators.", + ], + checkpoint_naming=OrderedDict( + toy_forcausallm_prefix="model.", + production_prefix="model.language_model.", + prefix_warning=( + "IMPORTANT: text-only Qwen4ExpForCausalLM.save_pretrained emits `model.layers.*` " + "with NO `language_model` segment, because production Qwen4-Exp ships the " + "multimodal Qwen4ExpForConditionalGeneration wrapper. The repacker's production " + "pattern is `model.language_model.layers.{L}.*`. The harness therefore also emits " + "a production-name copy at scratch/qwen4exp-toy-ckpt-prodlayout/." + ), + expert_layout_warning=( + "IMPORTANT: save_pretrained DE-FUSES the MoE experts back to the source layout, " + "`mlp.experts.{e}.gate_proj/up_proj/down_proj.weight` (one tensor per expert), " + "via the qwen2_moe WeightConverter. Production ships FUSED " + "`mlp.experts.gate_up_proj` [E, 2*I, H] and `mlp.experts.down_proj` [E, H, I]. " + "The prod-layout copy carries the fused runtime tensors. Row split confirmed: " + "gate = rows [0, I), up = rows [I, 2I) of gate_up_proj." + ), + ngram_shards=( + "ngram_embedding is written as `...ngram_embedding.shard_{S}.weight`, " + "split_ngram_parts shards concatenated on dim 0 at load. layer_multipliers, " + "ngram_heads_vocab_sizes and ngram_heads_offsets stay I64 in a bf16 checkpoint." + ), + ), + port_hazards=[ + "The n-gram hash mix consumes the FULL positive int64 range by construction: " + "(vocab_size - 1) * max(layer_multipliers) has bit length 63, i.e. zero headroom. " + "The Swift port must do this in Int64/UInt64 exact integer arithmetic. Float64 " + "cannot represent these products and Int32 overflows immediately. This holds for " + "the production vocab too -- multiplier_max is derived as (2^63-1)//vocab_size.", + "A query position is NOT guaranteed to attend to itself. When the query completes a " + "block and that block loses the top-k, the query's own key is absent from the " + "selected set (observed at LONG query 47 on both attention layers: 8 selected out of " + "48 visible, self excluded). Do not add a 'always keep self' shortcut.", + "Only COMPLETE compress_ratio blocks participate; the 1..ratio-1 tail is always kept. " + "The tail is taken from the visible index list, not from raw positions, so a padded " + "batch changes which tokens are 'tail'.", + "PLE segmentation uses eos_token_id INCLUSIVE: a shift may not cross an EOS, and the " + "EOS position itself starts the next segment. Positions that would read across the " + "boundary read eos_token_id instead. The toy LONG prompt has an EOS at index 20 to " + "pin this. Production eos_token_id is 248044 (text_config).", + "Qwen4ExpTextRMSNorm is zero-centered (1 + w) AND upcasts to fp32; " + "Qwen3_5RMSNormGated (linear_attn.norm) is ones-centered and must NOT get the +1 " + "bake. Confirmed by reading both classes in the installed package.", + "Attention q_proj packs query and gate PER HEAD: the output is viewed as " + "(.., num_heads, 2 * head_dim) and chunked in 2 on the last dim, so head h occupies " + "rows [h*2*head_dim, (h+1)*2*head_dim) with q first and gate second WITHIN the head. " + "It is not a global first-half/second-half split.", + ], + golden_files=OrderedDict(), + findings=report, + ) + + manifest["golden_files"] = OrderedDict( + (name, OrderedDict(bytes=len(blob), sha256=sha256_bytes(blob), contents=describe(name))) + for name, blob in files.items() + ) + + manifest_blob = (json.dumps(manifest, indent=1) + "\n").encode() + + if args.print_hashes: + for name, blob in files.items(): + print(f"{sha256_bytes(blob)} {name} ({len(blob)} bytes)") + print(f"{sha256_bytes(manifest_blob)} goldens-manifest.json ({len(manifest_blob)} bytes)") + return 0 + + out_dir.mkdir(parents=True, exist_ok=True) + for name, blob in files.items(): + (out_dir / name).write_bytes(blob) + (out_dir / "goldens-manifest.json").write_bytes(manifest_blob) + + total = sum(len(b) for b in files.values()) + len(manifest_blob) + print(f"wrote {len(files) + 1} files, {total / 1024:.1f} KiB total -> {out_dir}") + for name, blob in files.items(): + print(f" {name:34s} {len(blob):>9d} {sha256_bytes(blob)}") + print(f" {'goldens-manifest.json':34s} {len(manifest_blob):>9d} {sha256_bytes(manifest_blob)}") + + print(f"wrote toy checkpoints -> {CKPT_DIR} and {CKPT_PROD_DIR}") + assert total < 8 * 1024 * 1024, f"committed goldens exceed the 8 MB budget: {total} bytes" + return 0 + + +def describe(name: str) -> str: + kind = "prefill" if "prefill" in name else "decode" + prompt = "SHORT (dense-equivalent)" if "short" in name else "LONG (sparse)" + if name.endswith(".json"): + return ( + f"{kind} integer goldens for the {prompt} prompt: per-attention-layer indexer " + f"selected/visible token index sets (sorted, per query position), per-MoE router " + f"top-k expert indices, PLE n-gram embedding row ids (16 per token), greedy token ids." + ) + return ( + f"{kind} float32 goldens for the {prompt} prompt: embed_out; per layer " + f"stream_out / attn_hc_{{mixed,stream_in,inject}} / mlp_hc_{{mixed,stream_in,inject}} / " + f"block_out / moe_out / router_weights; PLE layer ple_out + ple_ngram_embeds; " + f"last_hidden_state; logits." + ) + + +def json_safe_config(cfg) -> dict: + d = cfg.to_dict() + return json.loads(json.dumps(d, default=str, sort_keys=True)) + + +def resolve_transformers_commit() -> str: + """Read the git commit pip/uv recorded in the dist-info ``direct_url.json``. + + Populated because the venv installs transformers from + ``git+https://github.com/huggingface/transformers``; a wheel install leaves it blank. + """ + import importlib.metadata as md + + try: + dist = md.distribution("transformers") + root = getattr(dist, "_path", None) + candidates = [] + if root is not None: + candidates.append(Path(root) / "direct_url.json") + site = Path(md.distribution("transformers").locate_file("")) + candidates.extend(sorted(site.glob("transformers-*.dist-info/direct_url.json"))) + for path in candidates: + if path.is_file(): + info = json.loads(path.read_text() or "{}") + commit = info.get("vcs_info", {}).get("commit_id", "") + if commit: + return commit + except Exception: # pragma: no cover - the manifest records "" and the README carries it + pass + return "" + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/Sources/Mference/Infrastructure/Metal/MetalContext.swift b/Sources/Mference/Infrastructure/Metal/MetalContext.swift index b1dbe3d..1251b4f 100644 --- a/Sources/Mference/Infrastructure/Metal/MetalContext.swift +++ b/Sources/Mference/Infrastructure/Metal/MetalContext.swift @@ -80,6 +80,10 @@ public final class MetalContext: @unchecked Sendable { "dsv4", "inkling", "dflash2", + "flashnext", + "flashnext_moe", + "flashnext_indexer", + "flashnext_gdn", ] /// Bundle locations for runtime shader modules. @@ -89,6 +93,10 @@ public final class MetalContext: @unchecked Sendable { "dequant_int8": "Metal/Quant", "dflash2": "Metal/DFlash2", "dsv4": "Metal/DSV4", + "flashnext": "Metal/FlashNext", + "flashnext_gdn": "Metal/FlashNext", + "flashnext_indexer": "Metal/FlashNext", + "flashnext_moe": "Metal/FlashNext", "fused": "Metal/Fusions", "gdn": "Metal/GDN", "inkling": "Metal/Inkling", diff --git a/Sources/Mference/Infrastructure/ModelIO/FlashNextResident.swift b/Sources/Mference/Infrastructure/ModelIO/FlashNextResident.swift new file mode 100644 index 0000000..25bb72e --- /dev/null +++ b/Sources/Mference/Infrastructure/ModelIO/FlashNextResident.swift @@ -0,0 +1,358 @@ +import Foundation +import Metal + +/// How the runtime must treat RMSNorm weights for the loaded family. +/// +/// Qwen4-Exp (Flash-Next) initializes every RMSNorm weight at zero and applies +/// `(1 + weight)`. The port's decision is to bake the `+1` into the stored +/// weights so the runtime's standard RMSNorm applies unchanged — the same +/// conversion the Qwen 3.8 MTP attach already performs at repack +/// (`MTPAttachTool`, `addOne`). +/// +/// The W2 install path copies norms verbatim today, so the bake happens here, +/// at load, gated on the family. Once the repacker folds it in, the install +/// publishes `manifest.zeroCenteredNormsBakedAtInstall = true` and this policy +/// resolves to `.storedInFullForm`, leaving the loader a pass-through. The +/// two paths must never both apply: that is what the manifest flag decides. +public enum ZeroCenteredNormPolicy: Sendable, Equatable { + /// Stored weights are already `1 + w`; hand them to the kernels as they are. + case storedInFullForm + /// Stored weights are the zero-centered `w`; widen to `1 + w` at load. + case bakeAtLoad +} + +extension Model { + + // MARK: - Zero-centered norm policy + + /// Families whose checkpoint RMSNorm weights are zero-centered, i.e. whose + /// reference implementation applies `(1 + w)`. + /// + /// Qwen 3.8's own trunk is *not* here: its shipped MLX conversion stores + /// norms in full form already, and only the separately attached MTP shard + /// carries the zero-centered convention, which `MTPAttachTool` converts at + /// attach time. Adding a family to this set changes the bytes every kernel + /// sees, so it is an explicit list rather than a heuristic. + static let zeroCenteredNormFamilies: Set = [.qwen38flashnext] + + /// Whether this load must apply the `(1 + w)` bake itself. + public var zeroCenteredNormPolicy: ZeroCenteredNormPolicy { + guard Self.zeroCenteredNormFamilies.contains(config.family) else { + return .storedInFullForm + } + return manifest.zeroCenteredNormsBakedAtInstall == true + ? .storedInFullForm : .bakeAtLoad + } + + /// Tensor-name suffixes of the norms the reference zero-centers. + /// + /// Every `Qwen4ExpTextRMSNorm` instance in the text stack: the attention + /// per-head q/k norms, the QSA indexer's q/k layernorms, the + /// hyper-connection group norms (per-site and the global mixer's), and the + /// PLE block's three norms. All of them are weight-zero-initialized and + /// applied as `(1 + w)`. + /// + /// Deliberately **excluded**: + /// * `linear_attn.norm.weight` — the gated DeltaNet norm is + /// `Qwen3_5RMSNormGated`, which is **ones**-initialized and applies + /// `w` directly. Confirmed against the reference implementation + /// (2026-09-01 parity harness), not inferred: it is the one norm in + /// this stack that is not zero-centered, and baking it would add one + /// to an already-full-form weight. + /// * the `mtp.*` sidecar's norms: MTP draft decode is out of scope for + /// v1, and its norms need the same treatment when it lands. + /// + /// The reference RMSNorm upcasts internally — + /// `_norm(x.float()) * (1 + w.float())`, cast back afterwards — so a + /// kernel consuming these baked weights must accumulate in fp32 to match. + static let zeroCenteredNormSuffixes: [String] = [ + ".self_attn.q_norm.weight", + ".self_attn.k_norm.weight", + ".self_attn.indexer.q_layernorm.weight", + ".self_attn.indexer.k_layernorm.weight", + ".attn_hyper_connection.hc_norm.weight", + ".mlp_hyper_connection.hc_norm.weight", + ".hyper_connection_mixer.hc_norm.weight", + ".ple.norm_conv.weight", + ".ple.norm_key.weight", + ".ple.norm_query.weight", + ] + + /// Whether `name` is one of the norms the `(1 + w)` bake applies to. + static func isZeroCenteredNorm(_ name: String) -> Bool { + zeroCenteredNormSuffixes.contains { name.hasSuffix($0) } + } + + /// Resolve a norm weight, applying the family's `(1 + w)` bake when the + /// install has not already done it. + /// + /// The resident buffer is a read-only mapping of the install, so the bake + /// materializes a BF16 copy in a private shared-storage buffer, cached by + /// name for the model's lifetime. Norms are `[10240]` at worst — 20 KB — + /// and there are a few hundred of them, so the copies cost single-digit + /// megabytes against a 175 GB install. + /// + /// Names outside `zeroCenteredNormSuffixes`, and every family outside + /// `zeroCenteredNormFamilies`, return the mapped tensor untouched: this + /// accessor is safe to route all norm reads through. + public func normWeight(name: String) throws -> TensorView { + let source = try resident(name: name) + guard zeroCenteredNormPolicy == .bakeAtLoad, + Self.isZeroCenteredNorm(name) else { return source } + guard source.dtype == 1 else { + throw ModelError.indexCorrupt( + detail: "zero-centered norm \(name) is dtype \(source.dtype), not BF16") + } + let key = "1+w:" + name + if let cached = streamersQueue.sync(execute: { convertedBox.views[key] }) { + return cached + } + let count = Int(source.length) / MemoryLayout.stride + guard let buffer = device.makeBuffer( + length: max(1, count * MemoryLayout.stride), + options: .storageModeShared) else { + throw ModelError.residentBufferWrapFailed + } + let src = source.buffer.contents().advanced(by: Int(source.offset)) + .assumingMemoryBound(to: UInt16.self) + let dst = buffer.contents().assumingMemoryBound(to: UInt16.self) + for i in 0.. String { + "\(trunkPrefix)layers.\(L).\(site.rawValue)." + } + + /// The two per-layer hyper-connection sites. Each owns a complete + /// `GatedResidual`: mix down/up, block inject, and a group norm. + public enum HyperConnectionSite: String, Sendable { + case attention = "attn_hyper_connection" + case mlp = "mlp_hyper_connection" + } + + /// `[hcLowRank, hcCount * hidden]` — collapses the normed residual bundle + /// to the low-rank mix vector. + public func hcMixDown(site: HyperConnectionSite, layer L: Int) throws -> TensorView { + try resident(name: hcSite(site, layer: L) + "input_mix_weight_down.weight") + } + /// `[hcCount * hidden, hcLowRank]` — expands the mix vector back to a + /// per-channel sigmoid gate over the bundle. + public func hcMixUp(site: HyperConnectionSite, layer L: Int) throws -> TensorView { + try resident(name: hcSite(site, layer: L) + "input_mix_weight_up.weight") + } + /// `[hcCount, hcCount * hidden]` — the per-stream placement weights the + /// block output is injected with. + public func hcInject(site: HyperConnectionSite, layer L: Int) throws -> TensorView { + try resident(name: hcSite(site, layer: L) + "block_inject_weight.weight") + } + /// `[hcCount * hidden]` — group RMSNorm weight, group size `hidden`: each + /// stream is normalized independently over its own channels against one + /// shared bundle-wide vector. + public func hcNorm(site: HyperConnectionSite, layer L: Int) throws -> TensorView { + try normWeight(name: hcSite(site, layer: L) + "hc_norm.weight") + } + + /// The global mixer that collapses the residual bundle to `hidden` after + /// the last layer. It is a `GatedResidual` without the inject path, and it + /// stands in for the final norm — this family has none. + public var hcGlobalMixDown: TensorView { + get throws { try resident(name: "\(trunkPrefix)hyper_connection_mixer.input_mix_weight_down.weight") } + } + public var hcGlobalMixUp: TensorView { + get throws { try resident(name: "\(trunkPrefix)hyper_connection_mixer.input_mix_weight_up.weight") } + } + public var hcGlobalNorm: TensorView { + get throws { try normWeight(name: "\(trunkPrefix)hyper_connection_mixer.hc_norm.weight") } + } + + // MARK: - QSA indexer + + private func indexerPrefix(layer L: Int) -> String { + "\(trunkPrefix)layers.\(L).self_attn.indexer." + } + + /// `[(indexerNumHeads + indexerNumKVHeads) * indexerHeadDim, hidden]`, + /// fused: the query heads first, then the single key head. + public func indexerQKProj(layer L: Int) throws -> TensorView { + try resident(name: indexerPrefix(layer: L) + "index_qk_proj.weight") + } + /// `[indexerHeadDim]`, zero-centered. + public func indexerQNorm(layer L: Int) throws -> TensorView { + try normWeight(name: indexerPrefix(layer: L) + "q_layernorm.weight") + } + /// `[indexerHeadDim]`, zero-centered. Applied to the pooled block key + /// *after* the float32 mean, not to the raw per-token keys. + /// + /// Note for whatever consumes these: the indexer does **not** guarantee + /// that a query's own block survives selection. There is no "always keep + /// self" rule anywhere in the reference — only the top-k over block scores + /// plus the always-selected incomplete tail — so no layer of this plumbing + /// may add one. + public func indexerKNorm(layer L: Int) throws -> TensorView { + try normWeight(name: indexerPrefix(layer: L) + "k_layernorm.weight") + } + + // MARK: - PLE n-gram embedding + + private func plePrefix(layer L: Int) -> String { + "\(trunkPrefix)layers.\(L).ple." + } + + /// `[hcCount * hidden, hidden]` — projects the gathered n-gram embedding + /// to the per-stream key. + public func pleKeyProj(layer L: Int) throws -> TensorView { + try resident(name: plePrefix(layer: L) + "key_proj.weight") + } + /// `[hidden, hidden]` — the value the per-stream gate scales. + public func pleValueProj(layer L: Int) throws -> TensorView { + try resident(name: plePrefix(layer: L) + "value_proj.weight") + } + /// `[hcCount * hidden, 1, pleConvKernelSize]` — depthwise causal conv over + /// the gated value, dilated by the n-gram size. + public func pleConv1d(layer L: Int) throws -> TensorView { + try resident(name: plePrefix(layer: L) + "conv1d.weight") + } + public func pleNormConv(layer L: Int) throws -> TensorView { + try normWeight(name: plePrefix(layer: L) + "norm_conv.weight") + } + public func pleNormKey(layer L: Int) throws -> TensorView { + try normWeight(name: plePrefix(layer: L) + "norm_key.weight") + } + public func pleNormQuery(layer L: Int) throws -> TensorView { + try normWeight(name: plePrefix(layer: L) + "norm_query.weight") + } + + /// The three splitmix64-derived n-gram hash multipliers, I64 `[3]`. + /// + /// **Loaded, never re-derived.** The reference derives them from a seed and + /// the layer index, but the installed values are the contract: a + /// re-derivation that disagreed by one bit would index a different row and + /// be undetectable at load. Their count is also the n-gram size. + public func pleLayerMultipliers(layer L: Int) throws -> [Int64] { + try residentInt64(name: plePrefix(layer: L) + "ple_embedding.layer_multipliers") + } + /// Per-head base row offsets into the pool, I64 `[ngramHeads]`. Loaded, + /// never re-derived. + public func pleNgramHeadOffsets(layer L: Int) throws -> [Int64] { + try residentInt64(name: plePrefix(layer: L) + "ple_embedding.ngram_heads_offsets") + } + /// Per-head vocabulary sizes (consecutive primes), I64 `[ngramHeads]`. + /// A head's row is `hash % vocabSizes[h] + offsets[h]`. + public func pleNgramHeadVocabSizes(layer L: Int) throws -> [Int64] { + try residentInt64(name: plePrefix(layer: L) + "ple_embedding.ngram_heads_vocab_sizes") + } + + /// N-gram size for a PLE layer, i.e. `layer_multipliers.count`. The token + /// history the runtime must cache is one less than this. + public func pleNgramSize(layer L: Int) throws -> Int { + try pleLayerMultipliers(layer: L).count + } + + /// Read an I64 lookup table out of the resident mapping. + /// + /// These are CPU-side hash tables, not kernel operands: the loader hands + /// back a typed `[Int64]` rather than a `TensorView` so a caller cannot + /// accidentally bind a 64-bit integer buffer as float data. + func residentInt64(name: String) throws -> [Int64] { + let view = try resident(name: name) + guard view.dtype == 4 else { + throw ModelError.indexCorrupt( + detail: "\(name) is dtype \(view.dtype); expected I64 (4)") + } + let count = Int(view.length) / MemoryLayout.stride + guard count > 0, UInt64(count * MemoryLayout.stride) == view.length else { + throw ModelError.tensorSizeMismatch( + name: name, + expected: UInt64(count * MemoryLayout.stride), + actual: view.length) + } + let base = view.buffer.contents().advanced(by: Int(view.offset)) + return (0...stride), + count: MemoryLayout.stride)) + } + return Int64(littleEndian: value) + } + } + + // MARK: - PLE row pool + + /// Validated geometry of the layer's streamed n-gram row pool. + public func plePoolGeometry(layer L: Int) throws -> PleRowPoolGeometry { + guard let pool = manifest.plePool else { + throw ModelError.plePoolMissing(layer: L) + } + guard pool.kind == PleRowPool.supportedKind else { + throw ModelError.plePoolInvalid(detail: "unknown kind \"\(pool.kind)\"") + } + guard let entry = pool.layers.first(where: { $0.layer == L }) else { + throw ModelError.plePoolMissing(layer: L) + } + return try PleRowPoolGeometry(layer: entry, hiddenSize: config.hiddenSize) + } + + /// Open the layer's row pool. The caller owns the returned reader (and its + /// file descriptor and cache slab); one per PLE layer is enough. + /// + /// The installed `ngram_heads_*` tables are cross-checked against the + /// pool's derived head count here, because that is the one place both + /// facts are in hand: a table of the wrong length would otherwise surface + /// as a silently wrong row index at decode. + public func openPleRowPool(layer L: Int, + cacheRows: Int = PleRowPool.defaultCacheRows) throws -> PleRowPool { + let geometry = try plePoolGeometry(layer: L) + let offsets = try pleNgramHeadOffsets(layer: L) + let vocabSizes = try pleNgramHeadVocabSizes(layer: L) + guard offsets.count == geometry.ngramHeads, + vocabSizes.count == geometry.ngramHeads else { + throw ModelError.plePoolInvalid( + detail: "layer \(L): pool geometry implies \(geometry.ngramHeads) " + + "n-gram heads but the checkpoint carries " + + "\(offsets.count) offsets / \(vocabSizes.count) vocab sizes") + } + for head in 0..= 0, vocabSizes[head] > 0, + last <= Int64(geometry.rows) else { + throw ModelError.plePoolInvalid( + detail: "layer \(L): n-gram head \(head) spans rows " + + "[\(offsets[head]), \(last)) outside the pool's " + + "\(geometry.rows) rows") + } + } + return try PleRowPool(directoryURL: directoryURL, + geometry: geometry, + cacheRows: cacheRows) + } + + // MARK: - Sidecars + + /// Whether the install carried an optional source tensor group (`mtp`, + /// `vision`). Absent means the install predates the sidecar policy. + public func sidecarCarried(_ group: String) -> Bool? { + manifest.sidecars?[group]?.carried + } + + /// A routed-expert pool installed outside `packed_experts/` — today only + /// the MTP draft layer's own 512 experts. Nothing reads these yet; the + /// accessor exists so the streaming layer can be pointed at them when MTP + /// draft decode lands, without another manifest change. + public func auxiliaryExpertPool(named name: String) -> ManifestAuxiliaryExpertPool? { + manifest.auxiliaryExpertPools?.first { $0.name == name } + } +} diff --git a/Sources/Mference/Infrastructure/ModelIO/ManifestReader.swift b/Sources/Mference/Infrastructure/ModelIO/ManifestReader.swift index 8551238..ba0a5a1 100644 --- a/Sources/Mference/Infrastructure/ModelIO/ManifestReader.swift +++ b/Sources/Mference/Infrastructure/ModelIO/ManifestReader.swift @@ -87,6 +87,95 @@ public struct ManifestArch: Decodable, Equatable, Sendable { public let routerNormAfterTopK: Bool? public let routerGlobalScale: Bool? public let unpaddedVocabSize: Int? + + // Qwen3.8-Flash-Next extensions. Optional for the same reason: absent + // values validate against the zeroed `FlashNextConfig.none`. Field names + // match what `MferenceRepack`'s `FlashNextAxes` publishes verbatim. + public let hcCount: Int? + public let hcLowRank: Int? + public let indexerNumHeads: Int? + public let indexerHeadDim: Int? + public let indexerNumKVHeads: Int? + public let indexerBudget: Int? + public let indexerCompressRatio: Int? + public let pleLayerIDs: [Int]? + public let pleNgramShardCount: Int? + public let pleNgramVocabSizeBase: Int? + public let pleConvKernelSize: Int? + /// Not emitted by installs predating this axis, so `validateArch` checks + /// it only when present. See `FlashNextConfig.pleEosTokenID`. + public let pleEosTokenID: Int? + /// Axis names the install claims a runner must implement. Advisory only — + /// `ManifestReader.familiesWithoutRunner` is the authority. + public let requiredAxes: [String]? +} + +/// One page-aligned region of a row-lookup pool, mirroring one source shard. +public struct ManifestPleShard: Decodable, Equatable, Sendable { + public let shard: Int + public let rows: Int + /// Byte offset of the region's first block within the pool file. + public let offset: UInt64 + public let size: UInt64 +} + +/// One PLE n-gram row pool: the repacked, page-aligned form of a layer's +/// hashed n-gram embedding table. +/// +/// Row `i` of shard `s` lives at +/// `shards[s].offset + (i / rowsPerBlock) * blockStride +/// + (i % rowsPerBlock) * rowStride`, and a record never straddles a block, +/// so one row costs one page fault and a cached page serves `rowsPerBlock` +/// neighbours. `rows` is the pool-wide total; shard rows are consecutive in +/// shard order. +public struct ManifestPlePoolLayer: Decodable, Equatable, Sendable { + public let layer: Int + /// Path relative to the install directory, e.g. `ple/layer_01_ngram_rows.bin`. + public let file: String + public let rows: Int + public let rowDim: Int + /// `"bf16"` or `"int4AffineG64"`. Chosen from the row width at install: + /// group-64 needs `rowDim % 64 == 0`, which 160 does not satisfy. + public let storage: String + public let weightBits: Int + public let groupSize: Int + public let rowWeightBytes: Int + public let rowScaleBytes: Int + public let rowBiasBytes: Int + public let rowStride: Int + public let rowsPerBlock: Int + public let blockStride: Int + public let fileSize: UInt64 + public let shards: [ManifestPleShard] +} + +public struct ManifestPlePool: Decodable, Equatable, Sendable { + /// Pool format tag. Readers must refuse an unknown kind. + public let kind: String + public let layers: [ManifestPlePoolLayer] +} + +/// An additive routed-expert pool outside `packed_experts/`, for a sidecar +/// (today: the MTP draft layer's own 512 experts at `packed_experts_mtp/`). +/// Kept out of `packed_experts/layout.json` so the shipped layout validator +/// and the routed-expert reader are untouched. +public struct ManifestAuxiliaryExpertPool: Decodable, Equatable, Sendable { + public struct Layer: Decodable, Equatable, Sendable { + public let layer: Int + public let file: String + } + public let name: String + public let directory: String + public let expertsPerLayer: Int + public let expertStride: UInt64 + public let layers: [Layer] +} + +/// Whether an optional tensor group from the source checkpoint was carried +/// into the install (`mtp`) or skipped (`vision`). +public struct ManifestSidecar: Decodable, Equatable, Sendable { + public let carried: Bool + public let tensorCount: Int } public struct ManifestQuantSlot: Decodable, Equatable, Sendable { @@ -133,6 +222,24 @@ public struct Manifest: Decodable, Equatable, Sendable { public let expertsPerLayer: Int public let numLayers: Int public let expertStride: UInt64 + + // Additive install blocks. Every family that predates them emits none of + // these keys, so existing manifests decode byte-identically. + + /// Streamed n-gram row pools, one per PLE layer. + public let plePool: ManifestPlePool? + /// Routed-expert pools outside `packed_experts/` (sidecar draft layers). + public let auxiliaryExpertPools: [ManifestAuxiliaryExpertPool]? + /// Optional source tensor groups and whether the install carried them. + public let sidecars: [String: ManifestSidecar]? + /// True when the installer has already folded the `+1` of the + /// zero-centered `(1 + w)` RMSNorm convention into the stored norm + /// weights. Absent means it has not: the loader applies the bake itself + /// for families whose norms are zero-centered + /// (`ZeroCenteredNormPolicy`). Deliberately outside `arch` — it describes + /// how the bytes were written, not what architecture they describe, so it + /// must not participate in the `archMismatch` field-by-field comparison. + public let zeroCenteredNormsBakedAtInstall: Bool? } public enum ManifestReader { @@ -212,6 +319,12 @@ public enum ManifestReader { if let flashHead = m.flashHead { try validateMapleFlashHead(flashHead, expected: expected) } + if let plePool = m.plePool { + try validatePlePool(plePool, expected: expected) + } else if !expected.flashNext.pleLayerIDs.isEmpty { + throw ModelError.plePoolMissing( + layer: expected.flashNext.pleLayerIndices[0]) + } let pageSize = UInt64(getpagesize()) guard m.expertStride % pageSize == 0 else { throw ModelError.expertStrideNotPageAligned(stride: m.expertStride, @@ -315,6 +428,31 @@ public enum ManifestReader { } return } + if expected.family == .qwen38flashnext { + // Workstream-2 quantize-in-flight: every eligible resident and + // routed tensor is INT4 affine group-64, the router included + // (the shipped families keep an INT8 router because their source + // conversions did; this one is quantized by the repacker itself + // under one uniform policy). + let slots: [(String, ManifestQuantSlot)] = [ + ("embedding", quant.embedding), + ("attention", quant.attention), + ("router", quant.router), + ("sharedExpert", quant.sharedExpert), + ("routedExpert", quant.routedExpert), + ] + for (name, slot) in slots { + guard slot.weightBits == 4, + slot.scheme.lowercased() == "affine", + slot.scaleType.lowercased() == "bf16", + slot.biasType.lowercased() == "bf16", + slot.groupSize == Quantization.groupSize else { + throw ModelError.indexCorrupt( + detail: "unsupported Qwen3.8-Flash-Next quantization for \(name)") + } + } + return + } // Routed experts additionally allow 2-bit: the DeepSeek-V4-Flash // dynamic-quant checkpoint ships Q2 experts under a Q4 core, and the // MoE runtime dispatches on `quant.routedExpert.weightBits`. @@ -336,6 +474,27 @@ public enum ManifestReader { } } + /// Structural gate on `manifest.plePool`. The runtime refuses a pool whose + /// geometry cannot address its own file: every arithmetic assumption the + /// row reader makes (`PleRowPool`) is checked once here rather than being + /// rediscovered per row read. + static func validatePlePool(_ pool: ManifestPlePool, + expected: ArchConfig) throws { + guard pool.kind == PleRowPool.supportedKind else { + throw ModelError.plePoolInvalid( + detail: "unknown kind \"\(pool.kind)\"; this runtime reads " + + "\"\(PleRowPool.supportedKind)\"") + } + for layerIndex in expected.flashNext.pleLayerIndices { + guard pool.layers.contains(where: { $0.layer == layerIndex }) else { + throw ModelError.plePoolMissing(layer: layerIndex) + } + } + for layer in pool.layers { + _ = try PleRowPoolGeometry(layer: layer, hiddenSize: expected.hiddenSize) + } + } + private static func validateMapleFlashHead(_ flashHead: ManifestMapleFlashHead, expected: ArchConfig) throws { guard expected.family == .maple, @@ -526,6 +685,27 @@ public enum ManifestReader { a.routerGlobalScale ?? false, e.routerGlobalScale) try check("unpaddedVocabSize", a.unpaddedVocabSize ?? 0, e.unpaddedVocabSize) + + let fn = e.flashNext + try check("hcCount", a.hcCount ?? 0, fn.hcCount) + try check("hcLowRank", a.hcLowRank ?? 0, fn.hcLowRank) + try check("indexerNumHeads", a.indexerNumHeads ?? 0, fn.indexerNumHeads) + try check("indexerHeadDim", a.indexerHeadDim ?? 0, fn.indexerHeadDim) + try check("indexerNumKVHeads", a.indexerNumKVHeads ?? 0, fn.indexerNumKVHeads) + try check("indexerBudget", a.indexerBudget ?? 0, fn.indexerBudget) + try check("indexerCompressRatio", a.indexerCompressRatio ?? 0, fn.indexerCompressRatio) + try check("pleLayerIDs", + (a.pleLayerIDs ?? []).description, fn.pleLayerIDs.description) + try check("pleNgramShardCount", a.pleNgramShardCount ?? 0, fn.pleNgramShardCount) + try check("pleNgramVocabSizeBase", + a.pleNgramVocabSizeBase ?? 0, fn.pleNgramVocabSizeBase) + try check("pleConvKernelSize", a.pleConvKernelSize ?? 0, fn.pleConvKernelSize) + // Checked only when the install publishes it: the shipped repack path + // predates this axis, and defaulting to the baseline would turn a real + // mismatch into silent agreement. See `FlashNextConfig.pleEosTokenID`. + if let published = a.pleEosTokenID { + try check("pleEosTokenID", published, fn.pleEosTokenID) + } } /// Decode just enough of `manifest.json` to identify the model family, @@ -542,6 +722,17 @@ public enum ManifestReader { detail: "manifest.json size \(size) exceeds metadata cap \(maxBytes)") } let data = try Data(contentsOf: manifestURL) + // The capability gate runs off a minimal decode, before the full + // `Manifest` shape is required. A family the runner cannot execute must + // say so by name even if its manifest carries fields (or omits ones) + // the strict decoder does not expect — otherwise the first thing the + // user sees is a decode error that reads like a corrupt install. + if let declared = try? JSONDecoder().decode(FamilyPeek.self, from: data), + let raw = declared.arch.family, + let missingAxes = familiesWithoutRunner[raw] { + throw ModelError.familyRunnerNotImplemented(family: raw, + missingAxes: missingAxes) + } let manifest: Manifest do { manifest = try JSONDecoder().decode(Manifest.self, from: data) @@ -554,4 +745,38 @@ public enum ManifestReader { } return family } + + /// Just enough of `manifest.json` to read `arch.family`. + private struct FamilyPeek: Decodable { + struct Arch: Decodable { let family: String? } + let arch: Arch + } + + /// Families `MferenceRepack` can install but `RealForwardRunner` cannot + /// execute, mapped to the axes whose kernels are missing. + /// + /// The gate lives here because `peekFamily` is the single funnel every + /// entry point uses — `Model.load`, the CLI, the loopback server, the Mac + /// app's installation probe and the tokenizer loader all call it — so one + /// check makes the failure named and actionable everywhere instead of + /// surfacing as "unknown arch.family" or, worse, matching some other + /// family's runner. + /// + /// The axis names are the runtime's own; the install also publishes them as + /// `manifest.arch.requiredAxes`, but this table is deliberately the + /// authority — a manifest does not get to tell the runtime what it can run. + /// Delete an entry only when the family's runner actually lands. + /// + /// A gated family may still carry an `ArchConfig` baseline, a `ModelFamily` + /// case, tensor accessors and manifest validation — that is what the + /// runner is built *against*. Presence in this table is the single fact + /// that decides whether it can be loaded, and it is checked in + /// `peekFamily` before any of that machinery is reached. + static let familiesWithoutRunner: [String: [String]] = [ + "qwen38flashnext": [ + "hyperConnectionsLowRank", + "attentionIndexer", + "pleNgramEmbedding", + ], + ] } diff --git a/Sources/Mference/Infrastructure/ModelIO/ModelTypes.swift b/Sources/Mference/Infrastructure/ModelIO/ModelTypes.swift index 9d7ba03..63fd8dc 100644 --- a/Sources/Mference/Infrastructure/ModelIO/ModelTypes.swift +++ b/Sources/Mference/Infrastructure/ModelIO/ModelTypes.swift @@ -12,6 +12,11 @@ public enum ModelFamily: String, Sendable, Hashable { case deepseekV4Flash = "deepseekV4Flash" case inklingSmall = "inklingSmall" case maple = "maple" + /// Qwen3.8-Flash-Next. The repacker installs it and the runtime carries a + /// compiled baseline so the install validates, but no runner executes it: + /// `ManifestReader.familiesWithoutRunner` refuses it by name at every + /// entry point until the kernels for its three new axes land. + case qwen38flashnext = "qwen38flashnext" } /// Gated-DeltaNet (linear attention) dimensions. Zeroed for architectures @@ -165,6 +170,133 @@ public struct RelativePositionConfig: Sendable, Equatable { dRel: 0, extent: 0, projDim: 0, logScalingFloor: 0, logScalingAlpha: 0) } +/// Qwen3.8-Flash-Next's three new axes: low-rank hyper-connections, the QSA +/// attention indexer, and the per-layer n-gram embedding (PLE). Zeroed +/// (`.none`) for every family that carries none of them. +/// +/// **Low-rank hyper-connections.** The residual stream is `hcCount` parallel +/// copies of `hiddenSize`, i.e. `hcCount * hiddenSize` wide from the embedding +/// through the last layer. Each sub-block site owns a `GatedResidual`: +/// `input_mix_weight_down` `[hcLowRank, hcCount * hidden]` and +/// `input_mix_weight_up` `[hcCount * hidden, hcLowRank]` factorize the mix, +/// `block_inject_weight` `[hcCount, hcCount * hidden]` places the block output +/// back into all streams, and `hc_norm` is a group RMSNorm with group size +/// `hiddenSize`. This is a *different* mechanism from +/// `HyperConnectionConfig` (DeepSeek V4's manifold-constrained mHC, which +/// carries a Sinkhorn-projected combine matrix and no low-rank factorization), +/// so the two configs are separate axes rather than one shared struct. +/// +/// **QSA indexer.** Per full-attention layer, `index_qk_proj` +/// `[indexerNumHeads * indexerHeadDim + indexerNumKVHeads * indexerHeadDim, +/// hidden]` scores the visible prefix in blocks of `indexerCompressRatio` +/// tokens; the top `indexerBudget / indexerCompressRatio` blocks (plus the +/// incomplete tail, always selected) form the attention mask for that query. +/// Selection sparsity is applied through the mask only — KV entries are never +/// dropped. +/// +/// **PLE.** At each layer id in `pleLayerIDs` (**one-indexed**, so id 2 is +/// `layers[1]`) a hashed n-gram embedding is added to the residual stream +/// before attention. Rows come from the streamed row pool described by +/// `manifest.plePool`, not from the resident buffer. `pleEosTokenID` delimits +/// the segments the n-gram shift must not cross. +/// +/// Facts established against the reference implementation (2026-09-01 parity +/// harness) that constrain any kernel built on these axes: +/// * **RMSNorm upcasts.** Every `Qwen4ExpTextRMSNorm` computes +/// `_norm(x.float()) * (1 + w.float())` and casts back, so the reduction +/// and the scale are fp32 even when the operands are BF16. +/// * **The n-gram mix has zero integer headroom.** The hash values are +/// exactly 63 bits by construction, so the mix must be done in true +/// 64-bit integer arithmetic — never `Double` (53-bit significand) and +/// never a 32-bit intermediate. +/// * **The indexer does not guarantee a query selects its own block.** +/// Nothing in this plumbing may assume a "keep self" entry; the visible +/// set is whatever the top-k over block scores plus the incomplete tail +/// produces. +/// * **Gated `q_proj` packs per head.** The `2 * numHeads * fullHeadDim` +/// rows are `[heads, 2 * headDim]` split on the last dimension, not a +/// global query-half / gate-half row split — which is already the +/// convention the shipped Qwen path implements (`split_q_gate_fp16`), so +/// `attnOutputGate` is reusable here unchanged. +public struct FlashNextConfig: Sendable, Equatable { + /// Parallel residual streams. The residual is `hcCount * hiddenSize` wide. + public let hcCount: Int + /// Rank of the hyper-connection mix factorization (`input_mix_weight_*`). + public let hcLowRank: Int + /// Indexer query heads. + public let indexerNumHeads: Int + /// Indexer head width, shared by the query heads and the single key head. + public let indexerHeadDim: Int + /// Indexer key heads; 1 in production (one pooled key per block). + public let indexerNumKVHeads: Int + /// Tokens the indexer keeps visible per query, before the always-selected + /// tail. `indexerBudget / indexerCompressRatio` blocks are selected. + public let indexerBudget: Int + /// Consecutive tokens pooled into one indexer block key. + public let indexerCompressRatio: Int + /// **One-indexed** layer ids carrying a PLE block: id `n` is `layers[n-1]`. + public let pleLayerIDs: [Int] + /// `split_ngram_parts`: how many shards the source n-gram table arrives + /// in, and how many page-aligned regions the installed row pool holds. + public let pleNgramShardCount: Int + /// `ngram_vocab_size_base` verbatim — the PER-HEAD base vocab, *not* the + /// table's row count. The true row count comes from the shard headers and + /// is published in `manifest.plePool.layers[].rows`; nothing validates one + /// against the other. + public let pleNgramVocabSizeBase: Int + /// Depthwise causal conv width in the PLE mixer (dilation is the n-gram + /// size, which is read from the installed `layer_multipliers` length). + public let pleConvKernelSize: Int + /// Token id that delimits PLE n-gram segments: a shifted token stream must + /// not read across it (`_shift_right_ignore_eos`, EOS inclusive). + /// + /// The installed manifests predating this axis do not publish it, so + /// `validateArch` checks it only when present and otherwise trusts this + /// compiled constant. Emitting `arch.pleEosTokenID` from the repacker is + /// an open follow-up; until it lands, a checkpoint with a different EOS + /// would violate the constant silently. + public let pleEosTokenID: Int + + public init(hcCount: Int, hcLowRank: Int, + indexerNumHeads: Int, indexerHeadDim: Int, + indexerNumKVHeads: Int, indexerBudget: Int, + indexerCompressRatio: Int, + pleLayerIDs: [Int], + pleNgramShardCount: Int, + pleNgramVocabSizeBase: Int, + pleConvKernelSize: Int, + pleEosTokenID: Int) { + self.hcCount = hcCount + self.hcLowRank = hcLowRank + self.indexerNumHeads = indexerNumHeads + self.indexerHeadDim = indexerHeadDim + self.indexerNumKVHeads = indexerNumKVHeads + self.indexerBudget = indexerBudget + self.indexerCompressRatio = indexerCompressRatio + self.pleLayerIDs = pleLayerIDs + self.pleNgramShardCount = pleNgramShardCount + self.pleNgramVocabSizeBase = pleNgramVocabSizeBase + self.pleConvKernelSize = pleConvKernelSize + self.pleEosTokenID = pleEosTokenID + } + + public static let none = FlashNextConfig( + hcCount: 0, hcLowRank: 0, + indexerNumHeads: 0, indexerHeadDim: 0, indexerNumKVHeads: 0, + indexerBudget: 0, indexerCompressRatio: 0, + pleLayerIDs: [], pleNgramShardCount: 0, pleNgramVocabSizeBase: 0, + pleConvKernelSize: 0, pleEosTokenID: 0) + + /// Zero-indexed layer indices carrying a PLE block, derived from the + /// one-indexed `pleLayerIDs` the checkpoint publishes. + public var pleLayerIndices: [Int] { pleLayerIDs.map { $0 - 1 } } + + /// Indexer blocks selected per query: `indexerBudget / compressRatio`. + public var indexerBlockBudget: Int { + indexerCompressRatio > 0 ? indexerBudget / indexerCompressRatio : 0 + } +} + /// Compile-time architecture baseline. `manifest.json -> arch` must match this /// field-by-field at load time; mismatches throw `ModelError.archMismatch`. /// @@ -269,6 +401,9 @@ public struct ArchConfig: Sendable, Equatable { public let routerNormAfterTopK: Bool /// Per-layer learned scalar multiplying the router weights. public let routerGlobalScale: Bool + /// Qwen3.8-Flash-Next's low-rank hyper-connections, QSA indexer and + /// per-layer n-gram embedding. `.none` for every other family. + public let flashNext: FlashNextConfig public init( hiddenSize: Int, @@ -318,7 +453,8 @@ public struct ArchConfig: Sendable, Equatable { routerGateBias: Bool = false, routerNormAfterTopK: Bool = false, routerGlobalScale: Bool = false, - unpaddedVocabSize: Int = 0 + unpaddedVocabSize: Int = 0, + flashNext: FlashNextConfig = .none ) { self.hiddenSize = hiddenSize self.intermediateSize = intermediateSize @@ -368,6 +504,7 @@ public struct ArchConfig: Sendable, Equatable { self.routerNormAfterTopK = routerNormAfterTopK self.routerGlobalScale = routerGlobalScale self.unpaddedVocabSize = unpaddedVocabSize + self.flashNext = flashNext } /// Canonical Gemma 4 26B-A4B baseline, checked against the installed @@ -701,7 +838,94 @@ public struct ArchConfig: Sendable, Equatable { (0..<42).map { $0 % 6 == 5 ? 1 : 0 } } + /// Canonical Qwen3.8-Flash-Next 180B-A3.5B baseline (text stack of the + /// multimodal `qwen4_exp` checkpoint; the vision tower is excluded at + /// repack). 48 layers in the same 3 : 1 hybrid as Qwen 3.6/3.8 — 36 + /// gated-DeltaNet layers and 12 full-attention layers — with 512 routed + /// experts (top-10) of width 640 plus one sigmoid-gated shared expert, + /// gated attention output, per-head q/k norms, partial NeoX RoPE over 64 + /// of 256 dims at θ 1e7, and an untied 248 320-row head. + /// + /// Three axes are new and carried in `flashNext`: the residual is 4 + /// low-rank hyper-connected streams (so the stream is 10 240 wide and + /// there is **no** final norm — the global mixer collapses it before + /// `lm_head`), each full-attention layer carries a QSA indexer, and + /// layer id 2 (`layers[1]`) carries a PLE n-gram embedding whose 320M-row + /// table streams from `manifest.plePool` rather than sitting resident. + /// + /// Values are read from the installed manifest at + /// `scratch/qwen38flashnext.gturbo` (revision `de4b8e4d`); see + /// `docs/families/QWEN38_FLASH_NEXT.md`. The runner does not exist yet: + /// `ManifestReader.familiesWithoutRunner` refuses this family by name at + /// every load path. The baseline exists so the manifest can be validated + /// and toy fixtures built while the kernels are written. + public static let qwen38FlashNext_180B_A3_5B = ArchConfig( + hiddenSize: 2560, + intermediateSize: 640, + moeIntermediateSize: 640, + numHeads: 24, + numKVHeads: 2, + numFullKVHeads: 2, + headDim: 256, + fullHeadDim: 256, + vocabSize: 248_320, + slidingWindow: 0, + finalLogitSoftcap: 0.0, + ropeTheta: 10_000_000.0, + fullRopeTheta: 10_000_000.0, + partialRotaryFactor: 0.25, + numLayers: 48, + numExperts: 512, + topKExperts: 10, + tieWordEmbeddings: false, + attentionKEqV: false, + fullAttentionLayerMask: Self.qwen38FlashNextLayerMask(), + hiddenActivation: "silu", + family: .qwen38flashnext, + attnOutputGate: true, // output_gate_type: sigmoid + attentionScale: 0.0625, // 256^-0.5 + embeddingScaledBySqrtHidden: false, + routerScaled: false, + ffnSandwichNorms: false, + sharedExpertGated: true, + ropeNeoxSubdim: true, + linearAttention: LinearAttentionConfig( + numKHeads: 16, numVHeads: 48, + keyHeadDim: 128, valueHeadDim: 128, + convKernelSize: 4), + // The router path is Qwen 3.6's: softmax over the logits, top-k of the + // probabilities, then renormalize the selected k (`norm_topk_prob`). + // `routerNormAfterTopK` stays false because that flag selects + // Inkling's and Maple's *different* ordering, not this one. + numSharedExperts: 1, + flashNext: FlashNextConfig( + hcCount: 4, + hcLowRank: 320, + indexerNumHeads: 4, + indexerHeadDim: 128, + indexerNumKVHeads: 1, + indexerBudget: 2048, + indexerCompressRatio: 4, + pleLayerIDs: [2], + pleNgramShardCount: 128, + pleNgramVocabSizeBase: 20_000_000, + pleConvKernelSize: 4, + // config.json text_config.eos_token_id @ de4b8e4d. + pleEosTokenID: 248_044) + ) + + private static func qwen38FlashNextLayerMask() -> [UInt8] { + // Same 3:1 hybrid shape as Qwen 3.6/3.8 (full_attention_interval = 4). + var mask = [UInt8](repeating: 2, count: 48) + for i in stride(from: 3, to: 48, by: 4) { mask[i] = 1 } + return mask + } + /// Registry keyed by `manifest.arch.family` for auto-detection at load. + /// + /// A family here has a validated baseline, not necessarily a runner: + /// `ManifestReader.familiesWithoutRunner` is the separate, authoritative + /// gate over whether the runtime can execute one. public static let knownArchitectures: [ModelFamily: ArchConfig] = [ .gemma4: .gemma4_26B_A4B, .qwen36: .qwen36_35B_A3B, @@ -709,6 +933,7 @@ public struct ArchConfig: Sendable, Equatable { .deepseekV4Flash: .deepseekV4Flash_284B_A13B, .inklingSmall: .inklingSmall_276B_A12B, .maple: .maplePreview, + .qwen38flashnext: .qwen38FlashNext_180B_A3_5B, ] /// Resident INT4 GEMV shapes this architecture issues during decode, for @@ -782,6 +1007,25 @@ public struct ArchConfig: Sendable, Equatable { } /// Hash-routed MoE layer: expert selection is `tid2eid[token]`. public func layerIsHashRouted(_ layer: Int) -> Bool { layer < numHashRoutedLayers } + + /// True when the residual is Flash-Next's low-rank hyper-connected bundle + /// of `flashNext.hcCount` streams rather than a single stream. + public var hasLowRankHyperConnections: Bool { flashNext.hcCount > 0 } + /// Width of the residual the layer graph carries between blocks. Blocks + /// themselves always run at `hiddenSize`; only the residual widens. + public var residualStreamWidth: Int { + hasLowRankHyperConnections ? flashNext.hcCount * hiddenSize : hiddenSize + } + /// True when this layer carries a PLE n-gram embedding block. `pleLayerIDs` + /// is one-indexed in the checkpoint, so id 2 answers true for layer 1. + public func layerIsPLE(_ layer: Int) -> Bool { + flashNext.pleLayerIDs.contains(layer + 1) + } + /// True when this layer carries a QSA indexer: full-attention layers of a + /// family that configures one. + public func layerHasAttentionIndexer(_ layer: Int) -> Bool { + flashNext.indexerNumHeads > 0 && layerIsFull(layer) + } } /// Failure modes for the validation gates in `Model.load`. @@ -802,6 +1046,19 @@ enum ModelError: Error, CustomStringConvertible, Equatable { case trustedReceiptInvalid(detail: String) case routedExpertPlanUnavailable(layer: Int) case eagerExpertFillFailed(layer: Int) + /// The repacker can install this family but no runner executes it yet. + /// Named rather than "unknown family" so the failure points at the missing + /// kernels instead of reading as a corrupt install. + case familyRunnerNotImplemented(family: String, missingAxes: [String]) + /// The install declares no `manifest.plePool` block, or none for the layer + /// the arch says carries a PLE n-gram embedding. + case plePoolMissing(layer: Int) + /// The `manifest.plePool` block is present but self-inconsistent (unknown + /// kind, geometry that does not tile the file, shard rows that do not sum + /// to the declared total). + case plePoolInvalid(detail: String) + /// A row index outside the pool's declared row count. + case plePoolRowOutOfRange(row: Int, rows: Int) public var description: String { switch self { @@ -837,6 +1094,15 @@ enum ModelError: Error, CustomStringConvertible, Equatable { return "routed expert fetch plan unavailable for layer \(layer)" case .eagerExpertFillFailed(let layer): return "eager routed expert read failed for layer \(layer); decode step aborted" + case .familyRunnerNotImplemented(let family, let missingAxes): + return "family \(family) is installed but its runner is not implemented; " + + "missing axes: \(missingAxes.joined(separator: ", "))" + case .plePoolMissing(let layer): + return "manifest.plePool has no row-lookup pool for PLE layer \(layer)" + case .plePoolInvalid(let detail): + return "manifest.plePool is invalid: \(detail)" + case .plePoolRowOutOfRange(let row, let rows): + return "PLE row \(row) is outside the pool's \(rows) rows" } } } diff --git a/Sources/Mference/Infrastructure/ModelIO/PleRowPool.swift b/Sources/Mference/Infrastructure/ModelIO/PleRowPool.swift new file mode 100644 index 0000000..04954b7 --- /dev/null +++ b/Sources/Mference/Infrastructure/ModelIO/PleRowPool.swift @@ -0,0 +1,492 @@ +import Darwin +import Foundation + +/// Validated geometry of one PLE n-gram row pool. +/// +/// The pool is the third instantiation of the fixed-aperture idea: weights +/// stream into expert slots, KV streams into pages, and here a 320-million-row +/// embedding table streams into a bounded row cache. Rows are addressed by +/// n-gram hash, a few per token, so the table never needs to be resident. +/// +/// Addressing, from `manifest.plePool` and reproduced here so the row reader +/// is one division and two multiplies: +/// +/// ``` +/// row i of shard s -> shards[s].offset +/// + (i / rowsPerBlock) * blockStride +/// + (i % rowsPerBlock) * rowStride +/// ``` +/// +/// Rows are **not** individually page-aligned — at a 320-byte record that +/// would inflate the ~102 GB table to ~5 TB. Blocks are, and a record never +/// straddles a block, so one row costs one page fault and a cached page serves +/// `rowsPerBlock` neighbours. +public struct PleRowPoolGeometry: Sendable, Equatable { + /// Zero-indexed layer this pool belongs to. + public let layer: Int + /// Pool file path relative to the install directory. + public let file: String + /// Total rows across every shard. + public let rows: Int + /// Row width in elements (160 in production: 16 n-gram heads x 160 = the + /// 2560-wide PLE embedding is assembled from 16 separate rows). + public let rowDim: Int + /// `bf16` in production — group-64 quantization needs `rowDim % 64 == 0`, + /// which 160 does not satisfy, so the pool stays BF16. + public let storage: Storage + /// Bytes of one complete row record, `[weights | scales | biases]`. + public let rowStride: Int + public let rowWeightBytes: Int + public let rowScaleBytes: Int + public let rowBiasBytes: Int + /// Rows packed into one page-aligned block. + public let rowsPerBlock: Int + /// Bytes between block starts (one page). + public let blockStride: Int + public let fileSize: UInt64 + /// Per-shard regions, in shard order. Global rows are consecutive across + /// them, so `firstRow` is the running prefix sum. + public let shards: [Shard] + + public enum Storage: String, Sendable, Equatable { + case bf16 + case int4AffineG64 + } + + public struct Shard: Sendable, Equatable { + public let shard: Int + public let rows: Int + public let offset: UInt64 + public let size: UInt64 + /// Global index of this shard's first row. + public let firstRow: Int + } + + /// How many n-gram head rows are concatenated to form one PLE embedding. + /// Derived rather than stored: the checkpoint publishes the embedding + /// width (the model's `hiddenSize`) and the row width, and the installed + /// `ngram_heads_offsets` / `ngram_heads_vocab_sizes` tables are exactly + /// this long — `FlashNextResident` cross-checks them at load. + public let ngramHeads: Int + + /// Validate `manifest.plePool.layers[]` against the arch and precompute the + /// shard prefix sums. Every assumption the row reader makes is checked + /// here, once, so a per-row read is pure arithmetic. + public init(layer entry: ManifestPlePoolLayer, hiddenSize: Int) throws { + func fail(_ detail: String) -> ModelError { + .plePoolInvalid(detail: "layer \(entry.layer): \(detail)") + } + guard let storage = Storage(rawValue: entry.storage) else { + throw fail("unknown storage \"\(entry.storage)\"") + } + guard entry.rowDim > 0, entry.rows > 0 else { + throw fail("rowDim \(entry.rowDim) and rows \(entry.rows) must be positive") + } + switch storage { + case .bf16: + guard entry.weightBits == 16, entry.groupSize == 0, + entry.rowScaleBytes == 0, entry.rowBiasBytes == 0, + entry.rowWeightBytes == entry.rowDim * MemoryLayout.stride else { + throw fail("bf16 pool must carry \(entry.rowDim * 2)-byte rows " + + "and no scale/bias companions") + } + case .int4AffineG64: + guard entry.weightBits == 4, + entry.groupSize == Quantization.groupSize, + entry.rowDim % Quantization.groupSize == 0, + entry.rowWeightBytes == entry.rowDim / 2 else { + throw fail("int4 pool geometry does not match group-" + + "\(Quantization.groupSize) affine rows") + } + } + let recordBytes = entry.rowWeightBytes + entry.rowScaleBytes + entry.rowBiasBytes + guard entry.rowStride == recordBytes else { + throw fail("rowStride \(entry.rowStride) != record \(recordBytes)") + } + guard entry.rowsPerBlock > 0, entry.blockStride > 0, + entry.rowsPerBlock * entry.rowStride <= entry.blockStride else { + throw fail("\(entry.rowsPerBlock) rows of \(entry.rowStride) B do not " + + "fit in a \(entry.blockStride) B block") + } + guard !entry.shards.isEmpty else { throw fail("no shard regions") } + + var shards: [Shard] = [] + shards.reserveCapacity(entry.shards.count) + var firstRow = 0 + for (position, shard) in entry.shards.sorted(by: { $0.shard < $1.shard }).enumerated() { + guard shard.shard == position else { + throw fail("shard ids are not 0..<\(entry.shards.count)") + } + guard shard.rows > 0 else { throw fail("shard \(shard.shard) has no rows") } + let blocks = (shard.rows + entry.rowsPerBlock - 1) / entry.rowsPerBlock + let needed = UInt64(blocks) * UInt64(entry.blockStride) + guard shard.size >= needed else { + throw fail("shard \(shard.shard) region \(shard.size) B cannot hold " + + "\(shard.rows) rows (needs \(needed) B)") + } + guard shard.offset % UInt64(entry.blockStride) == 0 else { + throw fail("shard \(shard.shard) offset \(shard.offset) is not " + + "block-aligned") + } + guard shard.offset <= entry.fileSize, + shard.size <= entry.fileSize - shard.offset else { + throw fail("shard \(shard.shard) region runs past the \(entry.fileSize) B file") + } + shards.append(Shard(shard: shard.shard, rows: shard.rows, + offset: shard.offset, size: shard.size, + firstRow: firstRow)) + firstRow += shard.rows + } + guard firstRow == entry.rows else { + throw fail("shard rows sum to \(firstRow), not \(entry.rows)") + } + guard hiddenSize > 0, hiddenSize % entry.rowDim == 0 else { + throw fail("rowDim \(entry.rowDim) does not divide hidden \(hiddenSize)") + } + + self.layer = entry.layer + self.file = entry.file + self.rows = entry.rows + self.rowDim = entry.rowDim + self.storage = storage + self.rowStride = entry.rowStride + self.rowWeightBytes = entry.rowWeightBytes + self.rowScaleBytes = entry.rowScaleBytes + self.rowBiasBytes = entry.rowBiasBytes + self.rowsPerBlock = entry.rowsPerBlock + self.blockStride = entry.blockStride + self.fileSize = entry.fileSize + self.shards = shards + self.ngramHeads = hiddenSize / entry.rowDim + } + + /// Byte offset of a global row within the pool file. + public func fileOffset(row: Int) throws -> UInt64 { + guard row >= 0, row < rows else { + throw ModelError.plePoolRowOutOfRange(row: row, rows: rows) + } + let shard = shards[shardIndex(containing: row)] + let local = row - shard.firstRow + return shard.offset + + UInt64(local / rowsPerBlock) * UInt64(blockStride) + + UInt64(local % rowsPerBlock) * UInt64(rowStride) + } + + /// Index of the shard owning a global row. Binary search over the prefix + /// sums rather than a division: the last shard is short in general. + func shardIndex(containing row: Int) -> Int { + var low = 0 + var high = shards.count - 1 + while low < high { + let mid = (low + high + 1) / 2 + if shards[mid].firstRow <= row { low = mid } else { high = mid - 1 } + } + return low + } +} + +/// Bounded LFU cache over PLE rows. +/// +/// Follows the routed-expert slot cache's conventions: a fixed number of +/// equal-sized slots carved out of one contiguous allocation, eviction by use +/// count with a use-clock tiebreak, and no growth under load. The difference +/// is scale — thousands of slots rather than tens — so the eviction victim is +/// found through frequency buckets in O(1) instead of the expert cache's sort +/// over its handful of slots. +/// +/// Ties *within* a frequency class are broken by whichever slot the bucket +/// happens to hold last; the order is deterministic for a given call sequence +/// but is not a documented LRU. Callers that need a specific victim should +/// raise the survivor's use count instead of relying on the tiebreak. +final class PleRowCache { + let capacityRows: Int + let rowStride: Int + + private let slab: UnsafeMutableRawPointer + /// slot -> global row index, or -1 when free. + private var slotRow: [Int] + /// slot -> use count. + private var slotUses: [Int] + /// slot -> position within its frequency bucket. + private var slotBucketPosition: [Int] + /// use count -> slots holding rows used that many times. + private var buckets: [Int: [Int]] = [:] + private var minimumUses = 0 + private var rowSlot: [Int: Int] = [:] + private var freeSlots: [Int] + + private(set) var hits = 0 + private(set) var misses = 0 + private(set) var evictions = 0 + + init(capacityRows: Int, rowStride: Int) { + precondition(capacityRows > 0, "PLE row cache needs at least one row") + precondition(rowStride > 0, "PLE row stride must be positive") + self.capacityRows = capacityRows + self.rowStride = rowStride + self.slab = UnsafeMutableRawPointer.allocate( + byteCount: capacityRows * rowStride, + alignment: MemoryLayout.alignment) + self.slotRow = [Int](repeating: -1, count: capacityRows) + self.slotUses = [Int](repeating: 0, count: capacityRows) + self.slotBucketPosition = [Int](repeating: -1, count: capacityRows) + self.freeSlots = Array((0..(_ row: Int, _ body: (UnsafeRawBufferPointer) -> T) -> T? { + guard let slot = rowSlot[row] else { + misses += 1 + return nil + } + hits += 1 + bumpUses(slot) + return body(UnsafeRawBufferPointer( + start: slab.advanced(by: slot * rowStride), count: rowStride)) + } + + /// Insert a row, evicting the least-frequently-used one when full. + /// Returns the slot's storage so the caller can fill it in place. + func insert(_ row: Int, _ fill: (UnsafeMutableRawBufferPointer) -> Void) { + if let existing = rowSlot[row] { + fill(UnsafeMutableRawBufferPointer( + start: slab.advanced(by: existing * rowStride), count: rowStride)) + bumpUses(existing) + return + } + let slot: Int + if let free = freeSlots.popLast() { + slot = free + } else { + slot = evictVictim() + evictions += 1 + } + fill(UnsafeMutableRawBufferPointer( + start: slab.advanced(by: slot * rowStride), count: rowStride)) + slotRow[slot] = row + rowSlot[row] = slot + slotUses[slot] = 1 + addToBucket(slot, uses: 1) + minimumUses = 1 + } + + /// Diagnostic view: which rows are resident, ascending. + func residentRows() -> [Int] { rowSlot.keys.sorted() } + + /// Use count of a resident row, for eviction-policy tests. + func useCount(of row: Int) -> Int? { rowSlot[row].map { slotUses[$0] } } + + private func bumpUses(_ slot: Int) { + let uses = slotUses[slot] + removeFromBucket(slot, uses: uses) + slotUses[slot] = uses + 1 + addToBucket(slot, uses: uses + 1) + if minimumUses == uses, buckets[uses] == nil { minimumUses = uses + 1 } + } + + private func addToBucket(_ slot: Int, uses: Int) { + slotBucketPosition[slot] = buckets[uses]?.count ?? 0 + buckets[uses, default: []].append(slot) + } + + /// Swap-remove: O(1), at the cost of any order within the frequency class. + private func removeFromBucket(_ slot: Int, uses: Int) { + guard var bucket = buckets[uses] else { return } + let position = slotBucketPosition[slot] + let last = bucket.count - 1 + if position != last { + bucket[position] = bucket[last] + slotBucketPosition[bucket[position]] = position + } + bucket.removeLast() + slotBucketPosition[slot] = -1 + if bucket.isEmpty { buckets[uses] = nil } else { buckets[uses] = bucket } + } + + private func evictVictim() -> Int { + while buckets[minimumUses] == nil { minimumUses += 1 } + let slot = buckets[minimumUses]!.last! + removeFromBucket(slot, uses: minimumUses) + rowSlot.removeValue(forKey: slotRow[slot]) + slotRow[slot] = -1 + slotUses[slot] = 0 + return slot + } +} + +/// `pread`-based reader over one PLE n-gram row pool, with a bounded LFU row +/// cache in front of it. +/// +/// CPU-only by design at this stage: the rows a token needs are known only +/// after the n-gram hash, and the kernels that would consume them on the GPU +/// do not exist yet. When they land, the cache's slab becomes the upload +/// staging buffer exactly as `PreadExpertStreamer`'s slot slab does. +public final class PleRowPool: @unchecked Sendable { + /// The only pool format this runtime reads. An install declaring anything + /// else is refused rather than guessed at. + public static let supportedKind = "rowLookupPoolV1" + + /// 65 536 rows is ~20 MB at the production 320-byte row: a modest default + /// beside the expert slot cache's multiple GB, and roughly 4 000 tokens of + /// distinct n-gram lookups at 16 rows per token. + public static let defaultCacheRows = 65_536 + + public let geometry: PleRowPoolGeometry + public let path: String + + private let fd: Int32 + private let cache: PleRowCache + private let lock = NSLock() + + /// Open the pool file described by `geometry`, relative to the install. + public init(directoryURL: URL, + geometry: PleRowPoolGeometry, + cacheRows: Int = PleRowPool.defaultCacheRows) throws { + let url = directoryURL.appendingPathComponent(geometry.file) + let opened = open(url.path, O_RDONLY) + guard opened >= 0 else { + throw ModelError.posixFailed(call: "open(\(url.path))", errno: errno) + } + var fileStats = stat() + if fstat(opened, &fileStats) == 0, + UInt64(fileStats.st_size) < geometry.fileSize { + close(opened) + throw ModelError.plePoolInvalid( + detail: "\(geometry.file) is \(fileStats.st_size) B; the manifest " + + "declares \(geometry.fileSize) B") + } + self.geometry = geometry + self.path = url.path + self.fd = opened + self.cache = PleRowCache(capacityRows: cacheRows, + rowStride: geometry.rowStride) + } + + deinit { close(fd) } + + /// Raw bytes of one row record, `[weights | scales | biases]`. Served from + /// the LFU cache when resident, otherwise read with one `pread` and + /// inserted. + public func readRowRecord(_ row: Int) throws -> [UInt8] { + let offset = try geometry.fileOffset(row: row) + lock.lock() + defer { lock.unlock() } + if let cached = cache.withCachedRow(row, { Array($0) }) { return cached } + var bytes = [UInt8](repeating: 0, count: geometry.rowStride) + try bytes.withUnsafeMutableBytes { raw in + try Self.readFull(fd: fd, into: raw.baseAddress!, + fileOffset: offset, count: geometry.rowStride) + } + cache.insert(row) { destination in + bytes.withUnsafeBytes { source in + destination.baseAddress!.copyMemory(from: source.baseAddress!, + byteCount: geometry.rowStride) + } + } + return bytes + } + + /// One row widened to float. BF16 pools decode losslessly; an INT4 pool + /// dequantizes its group-64 affine record. + public func readRow(_ row: Int) throws -> [Float] { + let record = try readRowRecord(row) + switch geometry.storage { + case .bf16: + return (0.. [Float] { + guard rows.count == geometry.ngramHeads else { + throw ModelError.plePoolInvalid( + detail: "expected \(geometry.ngramHeads) n-gram head rows, got \(rows.count)") + } + var out: [Float] = [] + out.reserveCapacity(geometry.ngramHeads * geometry.rowDim) + for row in rows { out.append(contentsOf: try readRow(row)) } + return out + } + + /// Cache counters, for tests and diagnostics. + public var cacheStatistics: (hits: Int, misses: Int, evictions: Int, resident: Int) { + lock.lock() + defer { lock.unlock() } + return (cache.hits, cache.misses, cache.evictions, cache.count) + } + + /// Test hook: use count of a cached row, or nil when it is not resident. + func cachedUseCount(of row: Int) -> Int? { + lock.lock() + defer { lock.unlock() } + return cache.useCount(of: row) + } + + /// Test hook: which rows the cache currently holds. + func residentRows() -> [Int] { + lock.lock() + defer { lock.unlock() } + return cache.residentRows() + } + + private static func dequantizeInt4AffineRow( + _ record: [UInt8], geometry: PleRowPoolGeometry + ) -> [Float] { + let groups = geometry.rowDim / Quantization.groupSize + let scaleBase = geometry.rowWeightBytes + let biasBase = scaleBase + geometry.rowScaleBytes + func u16(_ base: Int, _ index: Int) -> UInt16 { + UInt16(record[base + index * 2]) | (UInt16(record[base + index * 2 + 1]) << 8) + } + var out = [Float](repeating: 0, count: geometry.rowDim) + for group in 0..> 4) + out[index] = Float(nibble) * scale + bias + } + } + return out + } + + private static func readFull(fd: Int32, + into destination: UnsafeMutableRawPointer, + fileOffset: UInt64, + count: Int) throws { + var filled = 0 + while filled < count { + let got = pread(fd, destination.advanced(by: filled), count - filled, + off_t(fileOffset) + off_t(filled)) + if got < 0 { throw ModelError.posixFailed(call: "pread", errno: errno) } + if got == 0 { + throw ModelError.plePoolInvalid( + detail: "short read at row offset \(fileOffset): " + + "\(filled)/\(count) bytes") + } + filled += got + } + } +} diff --git a/Sources/Mference/Kernels/FlashNext/FlashNextAttention.swift b/Sources/Mference/Kernels/FlashNext/FlashNextAttention.swift new file mode 100644 index 0000000..b49a59c --- /dev/null +++ b/Sources/Mference/Kernels/FlashNext/FlashNextAttention.swift @@ -0,0 +1,261 @@ +import Foundation +import Metal + +/// Flash-Next gated full attention (`Qwen4ExpTextAttention`) over the KV subset +/// the QSA indexer selected. +/// +/// The math is `Qwen3_5Attention` — gated `q_proj`, per-head q/k RMSNorm, partial +/// NeoX RoPE, GQA, sigmoid output gate — so the shipped Qwen 3.8 kernels do the +/// work: `split_q_gate_fp16`, `PrefillQKVEpilogue`'s fused per-head norm + RoPE +/// (which, unlike the plain `rope_neox_subdim`, advances the position per row and +/// so serves prefill and decode alike), the dense `attention_full` decode kernel, +/// and `sigmoid_gate_mul_fp16`. +/// +/// # How the sparsity is applied +/// +/// Not by masking and not by dropping KV. The full KV cache is written for every +/// token; the selected positions are **gathered into contiguous scratch** and the +/// dense kernel runs over that short run. The design doc is explicit that +/// selection sparsity never removes KV, and gathering keeps the attention kernel +/// itself unmodified: every gathered position is at or before the query, so the +/// gathered run needs no causal mask. +/// +/// The cost is bounded by construction — a selection is at most +/// `indexer_budget + compress_ratio` positions (2051 in production) whatever the +/// context length, which is the whole point of the indexer. Prefill gathers in +/// waves of `gatherSlots` rows so one command buffer can cover several queries +/// without each overwriting the previous one's scratch. +final class FlashNextAttention { + + struct Geometry { + let hidden: Int + let numHeads: Int + let numKVHeads: Int + let headDim: Int + let rotaryDim: Int + let theta: Float + let eps: Float + let scale: Float + var qDim: Int { numHeads * headDim } + var kvDim: Int { numKVHeads * headDim } + } + + struct Weights { + /// `[2 * qDim, hidden]` — packed **per head** as `[query | gate]`. + let q: FlashNextWeightMatrix + let k: FlashNextWeightMatrix // [kvDim, hidden] + let v: FlashNextWeightMatrix // [kvDim, hidden] + let o: FlashNextWeightMatrix // [hidden, qDim] + let qNorm: MTLBuffer + let qNormOffset: Int + let kNorm: MTLBuffer + let kNormOffset: Int + } + + struct Scratch { + let packed: MTLBuffer // [rows, 2 * qDim] FP16 + let queries: MTLBuffer // [rows, qDim] FP16 + let gates: MTLBuffer // [rows, qDim] FP16 + let attnOut: MTLBuffer // [rows, qDim] FP16 + let gatheredK: MTLBuffer // [gatherSlots, maxSelected, kvDim] FP16 + let gatheredV: MTLBuffer + let maxRows: Int + let gatherSlots: Int + let maxSelected: Int + } + + /// The dense KV this family always keeps: 12 layers x 2 heads x 256, the + /// 24 KiB/token the dossier prices. Selection never touches it. + struct KVCache { + let keys: MTLBuffer // [maxTokens, kvDim] FP16 + let values: MTLBuffer // [maxTokens, kvDim] FP16 + let maxTokens: Int + } + + let geometry: Geometry + private let matVec: FlashNextMatVec + private let elementwise: Elementwise + private let epilogue: PrefillQKVEpilogue + private let attention: Attention + + init(context: MetalContext, matVec: FlashNextMatVec, + elementwise: Elementwise, epilogue: PrefillQKVEpilogue, + attention: Attention, geometry: Geometry) { + precondition(geometry.numHeads % geometry.numKVHeads == 0) + precondition(geometry.rotaryDim.isMultiple(of: 2)) + precondition(geometry.rotaryDim <= geometry.headDim) + self.geometry = geometry + self.matVec = matVec + self.elementwise = elementwise + self.epilogue = epilogue + self.attention = attention + } + + // MARK: - Allocation + + func makeScratch(device: MTLDevice, rows: Int, maxSelected: Int, + gatherSlots: Int) throws -> Scratch { + precondition(rows > 0 && gatherSlots > 0 && maxSelected > 0) + let half = MemoryLayout.stride + func buffer(_ elements: Int, _ mode: MTLResourceOptions) throws -> MTLBuffer { + guard let b = device.makeBuffer(length: max(1, elements) * half, + options: mode) else { + throw MetalError.noDevice + } + return b + } + return Scratch( + packed: try buffer(rows * 2 * geometry.qDim, .storageModePrivate), + queries: try buffer(rows * geometry.qDim, .storageModePrivate), + gates: try buffer(rows * geometry.qDim, .storageModePrivate), + attnOut: try buffer(rows * geometry.qDim, .storageModePrivate), + gatheredK: try buffer(gatherSlots * maxSelected * geometry.kvDim, + .storageModePrivate), + gatheredV: try buffer(gatherSlots * maxSelected * geometry.kvDim, + .storageModePrivate), + maxRows: rows, + gatherSlots: gatherSlots, + maxSelected: maxSelected) + } + + func makeKVCache(device: MTLDevice, maxTokens: Int) throws -> KVCache { + let half = MemoryLayout.stride + guard let k = device.makeBuffer( + length: max(1, maxTokens * geometry.kvDim) * half, + options: .storageModePrivate), + let v = device.makeBuffer( + length: max(1, maxTokens * geometry.kvDim) * half, + options: .storageModePrivate) else { + throw MetalError.noDevice + } + k.label = "flashnext.attn.k" + v.label = "flashnext.attn.v" + return KVCache(keys: k, values: v, maxTokens: maxTokens) + } + + // MARK: - Encode + + /// Projections, KV append, per-head q/k norm and partial RoPE. + /// + /// K and V are written **straight into the cache** at the rows this call + /// owns, then normed and roped in place there — the same in-cache norm the + /// Qwen 3.8 decode path uses, which is why the cache holds post-norm, + /// post-RoPE keys and raw values. + /// + /// PERF, not correctness: the three projections are per-row mat-vecs, one + /// compute encoder each. Decode is one row; a wide prefill chunk wants the + /// batched form, which belongs with the perf pass. + func encodeProjectAndCache(commandBuffer: MTLCommandBuffer, + weights w: Weights, + scratch: Scratch, + cache: KVCache, + x: MTLBuffer, xOffset: Int, + rows: Int, + startPosition: Int) { + precondition(rows > 0 && rows <= scratch.maxRows) + precondition(startPosition + rows <= cache.maxTokens, + "KV cache holds \(cache.maxTokens) tokens") + let half = MemoryLayout.stride + let qDim = geometry.qDim + let kvDim = geometry.kvDim + let kvBase = startPosition * kvDim * half + + for row in 0.. 0 && selectedCount <= scratch.maxSelected) + let half = MemoryLayout.stride + let runBytes = scratch.maxSelected * geometry.kvDim * half + attention.encodeFull( + commandBuffer: commandBuffer, + q: scratch.queries, qOffset: row * geometry.qDim * half, + k: scratch.gatheredK, kOffset: slot * runBytes, + v: scratch.gatheredV, vOffset: slot * runBytes, + out: scratch.attnOut, outOffset: row * geometry.qDim * half, + headDim: UInt32(geometry.headDim), + numQHeads: UInt32(geometry.numHeads), + numKVHeads: UInt32(geometry.numKVHeads), + seqLen: UInt32(selectedCount), + scale: geometry.scale) + } + + /// Byte offsets of one gather slot, for `FlashNextIndexer.encodeGatherKV`. + func gatherSlotOffset(_ slot: Int, scratch: Scratch) -> Int { + slot * scratch.maxSelected * geometry.kvDim * MemoryLayout.stride + } + + /// The sigmoid output gate and `o_proj`. The gate multiplies the flattened + /// per-head attention output **before** the projection, and the gate half is + /// never normed or rotated — it comes straight out of `q_proj`. + func encodeGateAndProject(commandBuffer: MTLCommandBuffer, + weights w: Weights, + scratch: Scratch, + out: MTLBuffer, outOffset: Int, + rows: Int) { + let half = MemoryLayout.stride + elementwise.encodeSigmoidGateMul(commandBuffer: commandBuffer, + out: scratch.attnOut, + gate: scratch.gates, + count: rows * geometry.qDim) + for row in 0.. n` +/// — always, at these widths — calling `nth_element(begin, begin + k - 1, end, +/// greater)` and reading the first `k` entries. libc++'s `__nth_element` +/// short-circuits by length: `2` swaps if out of order, `3` runs `__sort3`, +/// `<= 7` runs `__selection_sort` (first maximum, swapped into place — the swap +/// displaces whatever was there, which is why ties do NOT resolve to +/// lowest-index-first), and larger ranges run a median-of-3 quickselect. +/// +/// Only the first three branches are reproduced; above 7 candidates the +/// documented lowest-index-first fallback applies, and it is unobservable unless +/// a tie occurs there. This is a transcription of +/// `FlashNextReferenceRunner.descendingTopK`, which is oracle code and is not +/// modified to share it; `FlashNextIndexerReferenceTieBackTests` discharges the +/// transcription against the runner's own captured selections. +enum FlashNextDescendingTopK { + + /// The indices `torch.topk` puts in the first `k` slots, in its order. + /// + /// Note the `k >= n` short-circuit returns **identity order**, not a + /// descending sort — a real divergence from `torch.topk` that is + /// unobservable here because the indexer sorts the chosen positions + /// afterwards. `RouterWideTopK10Tests` pins the same property. + static func indices(_ values: [Float], k: Int) -> [Int] { + var a = values.enumerated().map { (index: $0.offset, value: $0.element) } + let n = a.count + if k >= n { return a.map(\.index) } + func greater(_ x: Int, _ y: Int) -> Bool { a[x].value > a[y].value } + switch n { + case 0, 1: + break + case 2: + if greater(1, 0) { a.swapAt(0, 1) } + case 3: + if !greater(1, 0) { + if greater(2, 1) { + a.swapAt(1, 2) + if greater(1, 0) { a.swapAt(0, 1) } + } + } else if greater(2, 1) { + a.swapAt(0, 2) + } else { + a.swapAt(0, 1) + if greater(2, 1) { a.swapAt(1, 2) } + } + case 4...7: + for i in 0..<(n - 1) { + var best = i + for m in (i + 1).. $1.value : $0.index < $1.index } + } + return (0.. Bool { + guard values.count > k, k >= 1 else { return false } + let ranked = values.sorted(by: >) + return ranked[k - 1] == ranked[k] + } +} diff --git a/Sources/Mference/Kernels/FlashNext/FlashNextGDN.swift b/Sources/Mference/Kernels/FlashNext/FlashNextGDN.swift new file mode 100644 index 0000000..fb87de9 --- /dev/null +++ b/Sources/Mference/Kernels/FlashNext/FlashNextGDN.swift @@ -0,0 +1,201 @@ +import Foundation +import Metal + +/// The dimension-generic gated DeltaNet decode path. +/// +/// Flash-Next's GDN block is Qwen 3.8's geometry (Hk 16, Hv 48, Dk 128, Dv 128, +/// conv 4), so the production runner takes the shipped `GDN` kernels — the fused +/// Hv=48 decode included — with the gated norm switched to sigmoid. This type is +/// the fallback for a geometry those kernels refuse: they are 32-lane tiled and +/// require `key_head_dim % 32 == 0`, which the parity toy's Dk of 8 is not. +/// +/// See `flashnext_gdn.metal` for the recurrence and the two easy-to-lose details +/// (the l2norm's sum-with-eps-inside, and the ones-centered gated-norm weight). +/// +/// The state layout is `[head][dk][dv]` — the reference's own indexing, and a +/// different order from the shipped kernels' `[Hv, Dv, Dk]`. The two paths are +/// never mixed on one install, but a runner must not switch between them +/// mid-sequence. +final class FlashNextGDN { + + struct Geometry { + let numKHeads: Int + let numVHeads: Int + let keyHeadDim: Int + let valueHeadDim: Int + let convKernel: Int + let eps: Float + + var keyDim: Int { numKHeads * keyHeadDim } + var valueDim: Int { numVHeads * valueHeadDim } + var qkvDim: Int { 2 * keyDim + valueDim } + /// Rows of conv input carried between steps. + var convStateRows: Int { convKernel - 1 } + } + + /// Per-layer decode state. The recurrent state is FP32 and the conv tail + /// FP16, matching what the shipped `GDNStateManager` holds — but allocated + /// here because the element ORDER differs. + struct LayerState { + let recurrent: MTLBuffer // [Hv, Dk, Dv] FP32 + let convTail: MTLBuffer // [K-1, qkvDim] FP16 + } + + struct Scratch { + /// `[(K - 1) + rows, qkvDim]` FP16 — `[carried tail | new rows]`. + let convPadded: MTLBuffer + /// `[rows, qkvDim]` FP16 — post-conv, post-SiLU. + let convOut: MTLBuffer + let rows: Int + } + + /// Largest key head dim the kernel's threadgroup scratch supports. + static let maxKeyHeadDim = 128 + + let geometry: Geometry + private let convPSO: MTLComputePipelineState + private let recurrencePSO: MTLComputePipelineState + + init(context: MetalContext, geometry: Geometry) throws { + precondition(geometry.keyHeadDim <= Self.maxKeyHeadDim, + "generic GDN supports key head dims up to \(Self.maxKeyHeadDim)") + precondition(geometry.numVHeads % geometry.numKHeads == 0, + "numVHeads must be a multiple of numKHeads") + precondition(geometry.convKernel >= 1) + self.geometry = geometry + self.convPSO = try context.pipeline("flashnext_gdn_conv") + self.recurrencePSO = try context.pipeline("flashnext_gdn_recurrence") + } + + // MARK: - Allocation + + func makeState(device: MTLDevice) throws -> LayerState { + let recurrentCount = geometry.numVHeads * geometry.keyHeadDim + * geometry.valueHeadDim + guard let recurrent = device.makeBuffer( + length: max(1, recurrentCount) * MemoryLayout.stride, + options: .storageModePrivate), + let tail = device.makeBuffer( + length: max(1, geometry.convStateRows * geometry.qkvDim) + * MemoryLayout.stride, + options: .storageModePrivate) else { + throw MetalError.noDevice + } + recurrent.label = "flashnext.gdn.state" + tail.label = "flashnext.gdn.convTail" + return LayerState(recurrent: recurrent, convTail: tail) + } + + func makeScratch(device: MTLDevice, rows: Int) throws -> Scratch { + let padded = (geometry.convStateRows + rows) * geometry.qkvDim + guard let p = device.makeBuffer( + length: max(1, padded) * MemoryLayout.stride, + options: .storageModePrivate), + let o = device.makeBuffer( + length: max(1, rows * geometry.qkvDim) + * MemoryLayout.stride, + options: .storageModePrivate) else { + throw MetalError.noDevice + } + return Scratch(convPadded: p, convOut: o, rows: rows) + } + + /// Zero the recurrent state and the conv tail — a fresh sequence. + func encodeReset(commandBuffer: MTLCommandBuffer, state: LayerState) { + guard let blit = commandBuffer.makeBlitCommandEncoder() else { return } + blit.fill(buffer: state.recurrent, range: 0.. 0 && rows <= scratch.rows) + let half = MemoryLayout.stride + let rowBytes = geometry.qkvDim * half + let stateRows = geometry.convStateRows + if let blit = cb.makeBlitCommandEncoder() { + if stateRows > 0 { + blit.copy(from: state.convTail, sourceOffset: 0, + to: scratch.convPadded, destinationOffset: 0, + size: stateRows * rowBytes) + } + blit.copy(from: qkv, sourceOffset: qkvOffset, + to: scratch.convPadded, + destinationOffset: stateRows * rowBytes, + size: rows * rowBytes) + blit.endEncoding() + } + if let enc = cb.makeComputeCommandEncoder() { + enc.setComputePipelineState(convPSO) + enc.setBuffer(scratch.convPadded, offset: 0, index: 0) + enc.setBuffer(scratch.convOut, offset: 0, index: 1) + enc.setBuffer(convWeight, offset: convWeightOffset, index: 2) + var qkvDim = UInt32(geometry.qkvDim) + var kernelWidth = UInt32(geometry.convKernel) + enc.setBytes(&qkvDim, length: 4, index: 3) + enc.setBytes(&kernelWidth, length: 4, index: 4) + let w = min(Int(convPSO.maxTotalThreadsPerThreadgroup), 256) + enc.dispatchThreads( + MTLSize(width: geometry.qkvDim, height: rows, depth: 1), + threadsPerThreadgroup: MTLSize(width: min(w, geometry.qkvDim), + height: 1, depth: 1)) + enc.endEncoding() + } + // Keep the last (K - 1) rows of [tail | rows] as the next tail. A blit + // between two distinct buffers, never in place: blits inside one encoder + // are unordered, so a self-overlapping shift would race whenever + // `rows < stateRows` — which is every decode step. + guard stateRows > 0, let blit = cb.makeBlitCommandEncoder() else { return } + blit.copy(from: scratch.convPadded, + sourceOffset: rows * rowBytes, + to: state.convTail, destinationOffset: 0, + size: stateRows * rowBytes) + blit.endEncoding() + } + + /// One decode step of the recurrence plus the sigmoid gated norm. + func encodeRecurrence(commandBuffer cb: MTLCommandBuffer, + scratch: Scratch, state: LayerState, + z: MTLBuffer, a: MTLBuffer, b: MTLBuffer, + aLog: MTLBuffer, aLogOffset: Int, + dtBias: MTLBuffer, dtBiasOffset: Int, + normWeight: MTLBuffer, normWeightOffset: Int, + out: MTLBuffer) { + guard let enc = cb.makeComputeCommandEncoder() else { return } + enc.setComputePipelineState(recurrencePSO) + enc.setBuffer(scratch.convOut, offset: 0, index: 0) + enc.setBuffer(z, offset: 0, index: 1) + enc.setBuffer(a, offset: 0, index: 2) + enc.setBuffer(b, offset: 0, index: 3) + enc.setBuffer(aLog, offset: aLogOffset, index: 4) + enc.setBuffer(dtBias, offset: dtBiasOffset, index: 5) + enc.setBuffer(normWeight, offset: normWeightOffset, index: 6) + enc.setBuffer(state.recurrent, offset: 0, index: 7) + enc.setBuffer(out, offset: 0, index: 8) + var hk = UInt32(geometry.numKHeads) + var hv = UInt32(geometry.numVHeads) + var dk = UInt32(geometry.keyHeadDim) + var dv = UInt32(geometry.valueHeadDim) + var eps = geometry.eps + enc.setBytes(&hk, length: 4, index: 9) + enc.setBytes(&hv, length: 4, index: 10) + enc.setBytes(&dk, length: 4, index: 11) + enc.setBytes(&dv, length: 4, index: 12) + enc.setBytes(&eps, length: 4, index: 13) + // One threadgroup per value head, one thread per value channel, rounded + // up to a full SIMD group so the reductions are well formed. + let threads = max(32, ((geometry.valueHeadDim + 31) / 32) * 32) + enc.dispatchThreadgroups( + MTLSize(width: geometry.numVHeads, height: 1, depth: 1), + threadsPerThreadgroup: MTLSize( + width: min(threads, Int(recurrencePSO.maxTotalThreadsPerThreadgroup)), + height: 1, depth: 1)) + enc.endEncoding() + } +} diff --git a/Sources/Mference/Kernels/FlashNext/FlashNextHyperConnections.swift b/Sources/Mference/Kernels/FlashNext/FlashNextHyperConnections.swift new file mode 100644 index 0000000..d15f544 --- /dev/null +++ b/Sources/Mference/Kernels/FlashNext/FlashNextHyperConnections.swift @@ -0,0 +1,265 @@ +import Foundation +import Metal + +/// `Qwen4ExpTextGatedResidual` — the low-rank hyper-connection that mixes the +/// `hc_count`-stream residual bundle down to one block input and injects the +/// block's output back into all the streams. +/// +/// Semantics (design doc, "Hyper-connections"), for `hyper` in R^(hc*H): +/// +/// ``` +/// h_n = group_rmsnorm_H(hyper) * hc_norm_w // hc*H +/// m = silu( W_down . h_n / hc ) // lowRank +/// g = sigmoid( W_up . m ) // hc*H +/// mixed = mean over streams of (g * h_n) // H +/// inj = 2 * sigmoid( W_inject . h_n / hc ) // hc +/// hyper = hyper + flatten( block_out * inj ) // raw stream, not h_n +/// ``` +/// +/// The global `hyper_connection_mixer` is the same minus the inject path, which +/// is why `encodeMix` and `encodeInject` are separate entry points: the mixer +/// calls only the first. +/// +/// # Precision +/// +/// Activations are FP16 in memory, per the runtime's convention. Two values are +/// deliberately kept FP32 through the chain because they feed a sigmoid, where +/// FP16 rounding is amplified rather than absorbed: the 10240-wide pre-sigmoid +/// mix gate `W_up . m`, and the four injection scalars. The low-rank vector +/// rounds to FP16 once, on the way into the second GEMV, because the shipped +/// INT4 mat-vec takes a half activation. +final class FlashNextHyperConnections { + + /// Buffers this encoder needs per (rows) batch. Owned by the caller so a + /// runner can size them once for its prefill chunk. + struct Scratch { + /// `[rows * bundle]` FP16 — the group-normed stream. + let normed: MTLBuffer + /// `[rows * lowRank]` FP32 — `W_down . h_n`, before the SiLU. + let lowRankRaw: MTLBuffer + /// `[rows * lowRank]` FP16 — `silu(raw / hc)`. + let lowRank: MTLBuffer + /// `[rows * bundle]` FP32 — `W_up . m`, before the sigmoid. + let mixGate: MTLBuffer + /// `[rows * hcCount]` FP32 — `W_inject . h_n`, before the gate. + let injectRaw: MTLBuffer + /// `[rows * hcCount]` FP32 — `2 * sigmoid(raw / hc)`. + let injectGate: MTLBuffer + } + + struct Weights { + let norm: MTLBuffer // [bundle] BF16, (1 + w) already baked + let normOffset: Int + let mixDown: FlashNextWeightMatrix // [lowRank, bundle] + let mixUp: FlashNextWeightMatrix // [bundle, lowRank] + /// Absent on the global mixer. + let inject: FlashNextWeightMatrix? // [hcCount, bundle] + } + + private let hidden: Int + private let hcCount: Int + private let lowRank: Int + private let eps: Float + + private let rms: RMSNorm + private let matVec: FlashNextMatVec + private let lowRankActivationPSO: MTLComputePipelineState + private let mixPSO: MTLComputePipelineState + private let injectGatePSO: MTLComputePipelineState + private let injectAccumulatePSO: MTLComputePipelineState + private let tileEmbeddingPSO: MTLComputePipelineState + + var bundle: Int { hidden * hcCount } + + init(context: MetalContext, + rms: RMSNorm, + matVec: FlashNextMatVec, + hidden: Int, hcCount: Int, lowRank: Int, eps: Float) throws { + precondition(hidden > 0 && hcCount > 0 && lowRank > 0) + self.hidden = hidden + self.hcCount = hcCount + self.lowRank = lowRank + self.eps = eps + self.rms = rms + self.matVec = matVec + self.lowRankActivationPSO = + try context.pipeline("flashnext_hc_lowrank_activation") + self.mixPSO = try context.pipeline("flashnext_hc_mix") + self.injectGatePSO = try context.pipeline("flashnext_hc_inject_gate") + self.injectAccumulatePSO = + try context.pipeline("flashnext_hc_inject_accumulate") + self.tileEmbeddingPSO = try context.pipeline("flashnext_hc_tile_embedding") + } + + func makeScratch(device: MTLDevice, rows: Int) throws -> Scratch { + func buffer(_ elements: Int, _ stride: Int) throws -> MTLBuffer { + guard let b = device.makeBuffer(length: max(1, elements) * stride, + options: .storageModePrivate) else { + throw MetalError.noDevice + } + return b + } + let half = MemoryLayout.stride + let float = MemoryLayout.stride + return Scratch( + normed: try buffer(rows * bundle, half), + lowRankRaw: try buffer(rows * lowRank, float), + lowRank: try buffer(rows * lowRank, half), + mixGate: try buffer(rows * bundle, float), + injectRaw: try buffer(rows * hcCount, float), + injectGate: try buffer(rows * hcCount, float)) + } + + /// `hidden_states = embed(ids).repeat(1, 1, hc_count)` — a tile, so stream + /// `j` is an exact copy of the embedding row. + func encodeTileEmbedding(commandBuffer: MTLCommandBuffer, + embedding: MTLBuffer, embeddingOffset: Int = 0, + hyper: MTLBuffer, hyperOffset: Int = 0, + rows: Int) { + guard let enc = commandBuffer.makeComputeCommandEncoder() else { return } + enc.setComputePipelineState(tileEmbeddingPSO) + enc.setBuffer(embedding, offset: embeddingOffset, index: 0) + enc.setBuffer(hyper, offset: hyperOffset, index: 1) + var hiddenVar = UInt32(hidden) + var hcVar = UInt32(hcCount) + enc.setBytes(&hiddenVar, length: MemoryLayout.size, index: 2) + enc.setBytes(&hcVar, length: MemoryLayout.size, index: 3) + dispatch2D(enc, pso: tileEmbeddingPSO, width: bundle, height: rows) + enc.endEncoding() + } + + /// The mix half: group norm, the low-rank gate, and the mean over streams. + /// Leaves the normed stream in `scratch.normed` because `encodeInject` reads + /// it — the inject path scores the NORMED stream even though the residual + /// add lands on the raw one. + func encodeMix(commandBuffer: MTLCommandBuffer, + weights: Weights, + scratch: Scratch, + hyper: MTLBuffer, hyperOffset: Int = 0, + mixed: MTLBuffer, mixedOffset: Int = 0, + rows: Int) { + rms.encodeBF16WGrouped(commandBuffer: commandBuffer, + x: hyper, xOffset: hyperOffset, + weight: weights.norm, weightOffset: weights.normOffset, + out: scratch.normed, + groupSize: UInt32(hidden), groups: UInt32(hcCount), + rows: rows, eps: eps) + // The two GEMVs are per-token, so a chunk walks them one row at a time. + // + // PERF, not correctness: that is one compute encoder per row per GEMV. + // Decode (rows == 1) is fine; a wide prefill chunk is not — at 48 layers + // and two hyper-connections each, a 256-token chunk would encode ~49k + // dispatches for these two matrices alone. Batching this into a 2-D + // mat-vec (rows on the grid's second axis, which the BF16 kernel here + // could take directly and `dequant_int4_gemv_simd` would need a variant + // for) is the fix, and it belongs with the production runner rather than + // ahead of it — the shapes are small enough that the dispatch count, + // not the bandwidth, is what it buys back. + for row in 0...stride, + y: scratch.lowRankRaw, + yOffset: row * lowRank * MemoryLayout.stride, + rows: lowRank, cols: bundle, outputFloat32: true) + } + encodeElementwise(commandBuffer: commandBuffer, + pso: lowRankActivationPSO, + input: scratch.lowRankRaw, output: scratch.lowRank, + count: rows * lowRank, divisor: Float(hcCount)) + for row in 0...stride, + y: scratch.mixGate, + yOffset: row * bundle * MemoryLayout.stride, + rows: bundle, cols: lowRank, outputFloat32: true) + } + guard let enc = commandBuffer.makeComputeCommandEncoder() else { return } + enc.setComputePipelineState(mixPSO) + enc.setBuffer(scratch.mixGate, offset: 0, index: 0) + enc.setBuffer(scratch.normed, offset: 0, index: 1) + enc.setBuffer(mixed, offset: mixedOffset, index: 2) + var hiddenVar = UInt32(hidden) + var hcVar = UInt32(hcCount) + enc.setBytes(&hiddenVar, length: MemoryLayout.size, index: 3) + enc.setBytes(&hcVar, length: MemoryLayout.size, index: 4) + dispatch2D(enc, pso: mixPSO, width: hidden, height: rows) + enc.endEncoding() + } + + /// The inject gate, from the normed stream `encodeMix` left behind. + func encodeInjectGate(commandBuffer: MTLCommandBuffer, + weights: Weights, + scratch: Scratch, + rows: Int) { + guard let inject = weights.inject else { + preconditionFailure("the global mixer has no inject path") + } + for row in 0...stride, + y: scratch.injectRaw, + yOffset: row * hcCount * MemoryLayout.stride, + rows: hcCount, cols: bundle, outputFloat32: true) + } + encodeElementwise(commandBuffer: commandBuffer, + pso: injectGatePSO, + input: scratch.injectRaw, output: scratch.injectGate, + count: rows * hcCount, divisor: Float(hcCount)) + } + + /// `hyper += flatten(block_out * inject)`, in place on the raw stream. + func encodeInjectAccumulate(commandBuffer: MTLCommandBuffer, + scratch: Scratch, + hyper: MTLBuffer, hyperOffset: Int = 0, + block: MTLBuffer, blockOffset: Int = 0, + rows: Int) { + guard let enc = commandBuffer.makeComputeCommandEncoder() else { return } + enc.setComputePipelineState(injectAccumulatePSO) + enc.setBuffer(hyper, offset: hyperOffset, index: 0) + enc.setBuffer(block, offset: blockOffset, index: 1) + enc.setBuffer(scratch.injectGate, offset: 0, index: 2) + var hiddenVar = UInt32(hidden) + var hcVar = UInt32(hcCount) + enc.setBytes(&hiddenVar, length: MemoryLayout.size, index: 3) + enc.setBytes(&hcVar, length: MemoryLayout.size, index: 4) + dispatch2D(enc, pso: injectAccumulatePSO, width: bundle, height: rows) + enc.endEncoding() + } + + // MARK: - Dispatch helpers + + private func encodeElementwise(commandBuffer: MTLCommandBuffer, + pso: MTLComputePipelineState, + input: MTLBuffer, output: MTLBuffer, + count: Int, divisor: Float) { + guard count > 0, + let enc = commandBuffer.makeComputeCommandEncoder() else { return } + enc.setComputePipelineState(pso) + enc.setBuffer(input, offset: 0, index: 0) + enc.setBuffer(output, offset: 0, index: 1) + var countVar = UInt32(count) + var divisorVar = divisor + enc.setBytes(&countVar, length: MemoryLayout.size, index: 2) + enc.setBytes(&divisorVar, length: MemoryLayout.size, index: 3) + let width = min(Int(pso.maxTotalThreadsPerThreadgroup), 256) + enc.dispatchThreads(MTLSize(width: count, height: 1, depth: 1), + threadsPerThreadgroup: MTLSize(width: min(width, count), + height: 1, depth: 1)) + enc.endEncoding() + } + + private func dispatch2D(_ enc: MTLComputeCommandEncoder, + pso: MTLComputePipelineState, + width: Int, height: Int) { + let w = min(Int(pso.maxTotalThreadsPerThreadgroup), 256) + enc.dispatchThreads(MTLSize(width: width, height: height, depth: 1), + threadsPerThreadgroup: MTLSize(width: min(w, width), + height: 1, depth: 1)) + } +} diff --git a/Sources/Mference/Kernels/FlashNext/FlashNextIndexer.swift b/Sources/Mference/Kernels/FlashNext/FlashNextIndexer.swift new file mode 100644 index 0000000..3d40572 --- /dev/null +++ b/Sources/Mference/Kernels/FlashNext/FlashNextIndexer.swift @@ -0,0 +1,417 @@ +import Foundation +import Metal + +/// The Flash-Next QSA indexer (`Qwen4ExpTextQSAIndexer`): per-query block +/// selection over the visible prefix, and the gather that hands gated full +/// attention the KV subset it selected. +/// +/// # What is on the GPU and what is not +/// +/// The projection, the query heads, the pooled block keys and the scores are GPU +/// work. The **ranking is not**: scores come back as at most a few thousand FP32 +/// values per query and the top-k runs through `FlashNextDescendingTopK`, the +/// exact `torch.topk` CPU ordering the reference uses. Selection is a support +/// set, not a tensor — a flipped boundary changes which KV a layer may read — +/// and relu-zero ties at the boundary are common enough that no GPU sort with a +/// different tie rule would be safe. The readback is small next to the attention +/// it gates. +/// +/// # Caches +/// +/// Two per attention layer, both FP32 (see `flashnext_indexer.metal` for why the +/// rest of the runtime's FP16 convention is deliberately broken here): +/// +/// * `rawKeys` — append-only, one `headDim` row per token. Un-normed, un-roped. +/// * `blockKeys` — one row per **complete** block of `compressRatio` tokens, +/// written once when the block's last token lands and immutable thereafter. +/// That immutability is the whole point: a block completed at token 40 scores +/// identically for every later query, so decode pools one new row every four +/// tokens instead of re-pooling the prefix. +/// +/// # Prefill +/// +/// Selection is computed **per position**, never lagged. `encodeScores` writes a +/// `[rows, scoreStride]` grid in one dispatch, with each row masked to its own +/// visible block count, so a chunked prefill reproduces the reference's +/// per-position selection exactly. The paged-KV Quest path's lag-one policy is +/// explicitly not reused here — the design doc rules it out. +final class FlashNextIndexer { + + struct Geometry { + let numHeads: Int + let numKVHeads: Int + let headDim: Int + let compressRatio: Int + /// `indexer_budget / compress_ratio` — how many complete blocks survive. + let blockBudget: Int + /// Shared with full attention: `fullHeadDim * partialRotaryFactor`. + let rotaryDim: Int + let theta: Float + let eps: Float + + /// Rows of `index_qk_proj`: the query heads then the single key head. + var projRows: Int { (numHeads + numKVHeads) * headDim } + /// Channel offset of the raw key head inside a projection row. + var keyOffset: Int { numHeads * headDim } + } + + /// Per-forward scratch, sized for the widest chunk and context the runner + /// will drive. Owned by the caller so it is allocated once. + struct Scratch { + /// `[rows, projRows]` FP32 — the whole `index_qk_proj` output. + let projection: MTLBuffer + /// `[rows, numHeads, headDim]` FP32 — normed and roped query heads. + let queries: MTLBuffer + /// `[rows, scoreStride]` FP32, **shared** so the CPU can rank it. + let scores: MTLBuffer + /// `[rows, selectionStride]` UInt32, shared. Row-indexed rather than a + /// single list because the gather reads it at GPU execution time: a + /// prefill wave that overwrote one slot per row before committing would + /// gather the last row's selection for every row. + let selection: MTLBuffer + let maxRows: Int + let scoreStride: Int + /// `blockBudget * compressRatio + compressRatio` — the widest selection + /// any query can produce. Above the budget a query takes exactly + /// `blockBudget` blocks plus at most `compressRatio - 1` tail tokens; + /// below it, everything visible, which is smaller still. + let selectionStride: Int + } + + /// One attention layer's indexer state. + struct LayerCache { + let rawKeys: MTLBuffer // [maxTokens, headDim] FP32 + let blockKeys: MTLBuffer // [maxTokens / ratio, headDim] FP32 + let maxTokens: Int + } + + /// Largest indexer head dim the kernels' thread-local scratch supports. + /// Mirrors `kFlashNextIndexerMaxHeadDim` in `flashnext_indexer.metal`. + static let maxHeadDim = 128 + + let geometry: Geometry + private let matVec: FlashNextMatVec + private let preparePSO: MTLComputePipelineState + private let appendPSO: MTLComputePipelineState + private let poolPSO: MTLComputePipelineState + private let scoresPSO: MTLComputePipelineState + private let gatherPSO: MTLComputePipelineState + + init(context: MetalContext, matVec: FlashNextMatVec, + geometry: Geometry) throws { + precondition(geometry.headDim <= Self.maxHeadDim, + "indexer head dim \(geometry.headDim) exceeds the kernels' " + + "\(Self.maxHeadDim)-wide thread scratch") + precondition(geometry.numKVHeads == 1, + "the indexer has exactly one key head") + precondition(geometry.compressRatio > 0 && geometry.blockBudget > 0) + precondition(geometry.rotaryDim.isMultiple(of: 2)) + precondition(geometry.rotaryDim <= geometry.headDim) + self.geometry = geometry + self.matVec = matVec + self.preparePSO = try context.pipeline("flashnext_indexer_prepare_queries") + self.appendPSO = try context.pipeline("flashnext_indexer_append_raw_keys") + self.poolPSO = try context.pipeline("flashnext_indexer_pool_block_keys") + self.scoresPSO = try context.pipeline("flashnext_indexer_scores") + self.gatherPSO = try context.pipeline("flashnext_indexer_gather_kv") + } + + // MARK: - Allocation + + func makeScratch(device: MTLDevice, rows: Int, maxTokens: Int) throws -> Scratch { + let stride = maxTokens / geometry.compressRatio + 1 + func buffer(_ elements: Int, _ byteStride: Int, + _ mode: MTLResourceOptions) throws -> MTLBuffer { + guard let b = device.makeBuffer(length: max(1, elements) * byteStride, + options: mode) else { + throw MetalError.noDevice + } + return b + } + let float = MemoryLayout.stride + let selectionStride = maxSelected + return Scratch( + projection: try buffer(rows * geometry.projRows, float, .storageModePrivate), + queries: try buffer(rows * geometry.numHeads * geometry.headDim, + float, .storageModePrivate), + scores: try buffer(rows * stride, float, .storageModeShared), + selection: try buffer(rows * selectionStride, + MemoryLayout.stride, .storageModeShared), + maxRows: rows, + scoreStride: stride, + selectionStride: selectionStride) + } + + /// The widest selection any query can produce. + var maxSelected: Int { + geometry.blockBudget * geometry.compressRatio + geometry.compressRatio + } + + func makeLayerCache(device: MTLDevice, maxTokens: Int) throws -> LayerCache { + let float = MemoryLayout.stride + let blocks = maxTokens / geometry.compressRatio + 1 + guard let raw = device.makeBuffer( + length: max(1, maxTokens * geometry.headDim) * float, + options: .storageModePrivate), + let pooled = device.makeBuffer( + length: max(1, blocks * geometry.headDim) * float, + options: .storageModePrivate) else { + throw MetalError.noDevice + } + raw.label = "flashnext.indexer.rawKeys" + pooled.label = "flashnext.indexer.blockKeys" + return LayerCache(rawKeys: raw, blockKeys: pooled, maxTokens: maxTokens) + } + + // MARK: - Encode + + /// `proj = index_qk_proj . x`, one row at a time, straight into FP32. + /// + /// PERF, not correctness: one compute encoder per row. Decode is one row; + /// a wide prefill chunk pays a dispatch per token per attention layer. The + /// fix is a batched mat-vec (rows on the grid's second axis) and it belongs + /// with the perf pass — the same note `FlashNextHyperConnections.encodeMix` + /// carries for the identical reason. + func encodeProjection(commandBuffer: MTLCommandBuffer, + weight: FlashNextWeightMatrix, + x: MTLBuffer, xOffset: Int, + hidden: Int, + scratch: Scratch, + rows: Int) { + precondition(rows <= scratch.maxRows) + for row in 0..