RASMUS/qwen3.5-35b-a3b-fp8-expert-predictor

Expert predictor for Qwen/Qwen3.5-35B-A3B-FP8

Thirty-four small heads, one for each MoE layer from L06 to L39. When token t has finished decoding, each head predicts which of the 256 routed experts its layer will use for token t+1. An offload engine can copy the predicted experts that are not on the GPU before their layer runs. The predictor never changes routing or weights. In every reported paired test, the generated text was identical with and without it.

This is research output from about 230 experiments. The heads are tied to one model and one revision: Qwen/Qwen3.5-35B-A3B-FP8 at 9d1823d2dee688a6b25e77009dc727688c44936e. They are not a general MoE predictor.

Measured effect

One RTX 4080 (16 GB), WSL2, batch 1, FreeToken at commit 2d8b9cc, ten held-out coding prompts (MBPP test, BigCodeBench tasks 800-804), 256 decode tokens, thinking on. Median decode speed per arm:

Arm	Median decode speed
FreeToken 2d8b9cc, no predictor (shared run settings)	27.74 tok/s
Our engine changes, no predictor	51.16 tok/s
Our engine changes plus this predictor	52.23 tok/s

• Most of the speed-up comes from engine changes, not from the predictor: a 16-stripe demand copy, priority CUDA graphs and pre-touch of the host expert banks. The first row ran in a later GPU session, so that comparison is descriptive.
• The predictor's own effect was measured in matched pairs at equal GPU memory. The control spends the predictor's memory on extra cache slots: 2746 slots against 2669 slots plus the int8 heads. The predictor arm was faster in 8 of 9 pairs that passed the memory gate, with a median gain of +1.41 tok/s (+2.7%). Four earlier experiments on three prompt sets gave paired median gains of +1.0 to +1.8 tok/s.
• Settings of the measured predictor arm: int8 weights (group 128), section caps [2, 2, 2, 2, 2, 3, 3], prefetch from decode step 1, 8242 KV pages.

Offline recall

Recall@K is the share of token t+1's routed experts (8 per layer) that are in the head's top K. The lead is one token: every input exists when token t has finished decoding, and the target is token t+1.

These numbers come from each head's original training run, not from the published files:

• All heads use the same validation set: 447 prompts of trace chunk00 (113,953 token rows). The training runs used it to select the checkpoint, so it is not an untouched test set, and the numbers can be slightly optimistic.
• The 4,000-prompt heads were evaluated in FP32. The later heads were evaluated in BF16 autocast.
• The served bank uses int8 weights and adds biases in FP32 before the BF16 rounding. Recall of the int8 bank was not measured separately; the speed results below are the end-to-end check.
• The training pools differ between head generations (see Training data), so compare rows with care.

Mean recall@16 over the 34 heads: 0.878.

Target layer	Head	Pool prompts	@8	@12	@16	@24	@32	@48	@64
L06	EXP-446	4,000	0.737	0.845	0.891	0.933	0.952	0.971	0.980
L07	EXP-447	4,000	0.699	0.808	0.860	0.911	0.936	0.960	0.972
L08	EXP-447	4,000	0.703	0.813	0.865	0.915	0.938	0.961	0.971
L09	EXP-447	4,000	0.726	0.832	0.880	0.925	0.945	0.965	0.975
L10	EXP-465	10,000	0.747	0.855	0.901	0.942	0.960	0.977	0.985
L11	EXP-447	4,000	0.699	0.805	0.855	0.904	0.929	0.955	0.968
L12	EXP-466	10,000	0.704	0.817	0.870	0.920	0.943	0.965	0.976
L13	EXP-447	4,000	0.691	0.801	0.853	0.902	0.926	0.951	0.963
L14	EXP-447	4,000	0.687	0.797	0.850	0.901	0.926	0.950	0.963
L15	EXP-466	10,000	0.706	0.817	0.869	0.917	0.940	0.962	0.973
L16	EXP-478b	20,000	0.712	0.826	0.878	0.925	0.947	0.967	0.976
L17	EXP-478b	20,000	0.724	0.835	0.885	0.930	0.949	0.968	0.977
L18	EXP-478b	20,000	0.726	0.839	0.888	0.931	0.951	0.969	0.978
L19	EXP-478b	20,000	0.715	0.827	0.878	0.924	0.944	0.965	0.975
L20	EXP-478b	20,000	0.709	0.820	0.871	0.917	0.939	0.961	0.971
L21	EXP-478b	20,000	0.713	0.821	0.871	0.917	0.939	0.961	0.973
L22	EXP-478b	20,000	0.709	0.818	0.868	0.914	0.937	0.959	0.971
L23	EXP-478b	20,000	0.701	0.809	0.860	0.907	0.930	0.955	0.968
L24	EXP-478b	20,000	0.676	0.788	0.843	0.897	0.924	0.950	0.965
L25	EXP-478b	20,000	0.698	0.809	0.861	0.909	0.932	0.955	0.967
L26	EXP-478b	20,000	0.698	0.812	0.863	0.911	0.932	0.954	0.965
L27	EXP-478b	20,000	0.703	0.813	0.862	0.909	0.931	0.953	0.965
L28	EXP-478b	20,000	0.721	0.831	0.878	0.919	0.939	0.959	0.969
L29	EXP-478b	20,000	0.746	0.850	0.892	0.929	0.946	0.964	0.973
L30	EXP-478b	20,000	0.748	0.853	0.895	0.931	0.948	0.966	0.975
L31	EXP-478b	20,000	0.742	0.849	0.892	0.930	0.947	0.965	0.974
L32	EXP-478b	20,000	0.770	0.871	0.909	0.941	0.956	0.972	0.980
L33	EXP-478b	20,000	0.778	0.874	0.909	0.940	0.955	0.970	0.978
L34	EXP-478b	20,000	0.789	0.883	0.917	0.945	0.958	0.972	0.979
L35	EXP-478b	20,000	0.741	0.844	0.888	0.928	0.947	0.966	0.976
L36	EXP-478b	20,000	0.741	0.842	0.885	0.924	0.944	0.964	0.974
L37	EXP-478b	20,000	0.732	0.836	0.882	0.925	0.945	0.966	0.976
L38	EXP-478b	20,000	0.741	0.844	0.890	0.934	0.954	0.974	0.983
L39	EXP-478b	20,000	0.754	0.850	0.891	0.931	0.951	0.970	0.980

Inputs

For target layer L, all from the end of token t (see config.json and reference_scorer.py):

• h (2048): token t's MoE input hidden state at layer L (FP16)
• win (768): token t's routed experts at layers L-1, L-2, L-3 (multi-hot)
• prior (256): EMA over tokens <= t of layer L's routed experts (beta 0.9), rounded to FP16
• pos (1): (decode step + 1) / 256
• rl (768): token t's router logits at layers L-1, L-2, L-3 (FP16; L05 stored as zeros)
• prevwin (1280): token t-1's routed experts at layers L-1 .. L-5 (multi-hot)
• hd (2048): h minus token t's MoE input hidden state at layer L-1, subtracted in FP16
• sticky (256): token t's routed experts at layer L (multi-hot)
• Token t's id and the already-sampled id of token t+1, through a frozen random token table (randn, seed 1234, 248,070 x 32). The table is not taken from the model. It covers the text ids 0 to 248,069. The tokenizer's audio and TTS tokens (248,070 to 248,076) have no row, and reference_scorer.py rejects them. Use the predictor for text generation only.

Each head also applies a frozen linear map to a forecast of its own layer's hidden state. That map is a least-squares fit of the layer's router, made from our routing traces (W, router_b). It matches the true router at FP16 noise level, but it is not bit-identical to it. Load it from this file, not from the model checkpoint.

Files

File	Content
model.safetensors	The BF16 bank as served before int8 quantization; calibration a, b in FP32
config.json	Base model, shapes, feature contract, provenance hashes (schema: spec.py)
bank_loader.py	Loads and validates both files; to_int8 rebuilds the served int8 bank
reference_scorer.py	Plain PyTorch feature assembly and BF16 scoring
parity_fixture.safetensors	9 input sets with the expected scores and probabilities (see below)
manifest.json	SHA-256 of every file and of the source artifacts
LICENSE, NOTICE	Apache-2.0 terms; attribution and the data notice

Requirements: Python 3.10 or later, PyTorch, safetensors, Pydantic 2. Tested with torch 2.11.0+cu130, safetensors 0.8.0 and Pydantic 2.13.4, on the CPU.

What the release gate checked. Before publication, a gate on the CPU compared these files with the code of the measured run. It found byte equality for:

• every BF16 tensor of the bank, and the token table against the original table file;
• every buffer of the int8 bank that to_int8 builds (codes, scales, biases, calibration);
• the BF16 scores and probabilities of reference_scorer.py against the original BF16 bank, on the 9 fixture inputs. Eight inputs are mid-answer tokens with random features and a warmed prior. One input is the first decode step after a reset. All inputs have zero layer-5 router logits, as the deployed tracker stored them.

The gate did not check the outputs of the int8 bank, the GPU kernels, or the feature capture of the runtime. Byte equality holds for the tested environment; other PyTorch versions or devices can differ in the last bit.

Check the download from inside the model folder (hf download RASMUS/qwen3.5-35b-a3b-fp8-expert-predictor --local-dir .), so that Python finds the three modules:

from bank_loader import load
from reference_scorer import probabilities
from safetensors.torch import load_file

cfg, state = load('.')
fixture = load_file('parity_fixture.safetensors')
raw = {k[2:]: v for k, v in fixture.items() if k.startswith('0.') and 'expected' not in k}
p = probabilities(state, raw, cfg.base_model.num_experts, cfg.features.ema_beta)
assert (p == fixture['0.expected_probabilities']).all()  # bitwise on the CPU

The runtime integration (feature capture inside the decode graph, the residency-aware copy policy, the staging buffer) is not in this repository yet. It exists as research hooks for FreeToken 2d8b9cc. A clean FreeToken option is the next step.

Limits

• One model, one GPU class, batch 1, WSL2. Not tested on other engines or hardware.
• Trained on 256-token decodes. The position input is (decode step + 1) / 256. Behavior on longer answers is untested.
• The timed tokens were mostly reasoning text. Task correctness was not tested.
• L05 has no head. The router-logit input stores L05 as zeros, as the deployed tracker did.
• At the first decode step after a reset, the deployed runtime passed all-zero expert ids for the previous token. The one-hot encoding turns them into "expert 0 was routed", not "no expert". A port that wants to match the measured run must do the same.
• check_base in spec.py compares shapes only. A fine-tune of the same architecture passes it and gives meaningless probabilities. Pin the base model repository and revision yourself.
• The held-out check of the benchmark prompts finds exact copies of training prompts only.

Training data

The labels are the experts that Qwen/Qwen3.5-35B-A3B-FP8 routed to while it answered prompts from a mostly coding mix (57,813 coding and 800 general prompts in the full corpus). The heads see hidden states, router logits, routing decisions, and the ids of the consumed and the upcoming token. They never see the prompt text as such. The routing traces are public: RASMUS/commitmoe-qwen35-fp8-layers, RASMUS/commitmoe-qwen35-fp8-layers-v2.

Pool prompts	Trace chunks	Heads	Training prompts	Alpaca prompts in pool	CodeAlpaca prompts in pool
4,000	chunk00	L06, L07, L08, L09, L11, L13, L14	3,142	19	389
10,000	chunk00 + chunk01	L10, L12, L15	7,966	68	976
20,000	chunk00 + chunk01 + chunk_10k	L16, L17, L18, L19, L20, L21, L22, L23, L24, L25, L26, L27, L28, L29, L30, L31, L32, L33, L34, L35, L36, L37, L38, L39	15,950	141	1989

Data notice. Part of each pool comes from datasets with non-commercial terms. Stanford Alpaca is licensed for non-commercial use, and its authors ask that models trained on it be used for research only. CodeAlpaca's original repository states CC BY-NC 4.0, and its Hugging Face page states CC BY 4.0. These heads learn the base model's routing decisions, not the prompt text. How the dataset terms apply to such weights is not settled. The weights are released under Apache-2.0. Users who need certainty for commercial use should check the dataset terms themselves.

Source

The export, the gate and the card generator are in predictor_release/ of the project repository at commit 5ef8fdf7b2c4c8351cfd9f23eaf5d83eecfe2f71. That repository is not public yet. The manifest records the SHA-256 of every source artifact that the export read.

Credit

• FreeToken by FlashML-org (Apache-2.0): the offload engine.
• Qwen (Apache-2.0): the base model.