[jason-fxz]

Summary

--embed-weights host keeps the input token embedding table in pinned host RAM, so its VRAM goes to the MoE expert cache and the KV pool. The default stays --embed-weights gpu, which leaves serving unchanged.

The lookup reads only a few rows per step (one per decode token, one per prompt token in prefill). The GPU reads them over PCIe, and the table never has to sit in VRAM. A table tied to the LM head stays on the GPU because the head reads the whole table every step, so host is refused at start-up when the model config ties them (tie_word_embeddings). It is also refused for a model whose embedding is not a plain table, such as the DF11 table of non-GGUF GLM-4.x checkpoints.

Closes #754

Design

load:   checkpoint tensor --> cudaHostAlloc(Portable | Mapped) --> layer.weight  (CPU tensor, host VA)
lookup: ids (GPU) --> indexing() / ggml.get_rows() --> kernel reads rows at device_alias(table) --> output on the ids device

• The table is created pinned while the state dict is materialized, in _materialize_loaded_weight_state_dict and in the dummy-weight path.
  • FTW and GGUF readers hand the table over in host memory, so it never enters VRAM.
  • HF safetensors readers open the shards on the GPU, so there the table passes through VRAM: it is copied to pinned memory as it arrives, and its VRAM is released before the post-load memory snapshot.
  • Either way the snapshot excludes the table, and --moe-cache-auto gets the freed VRAM without extra code.
• The tables that move are the VocabParallelEmbedding ops of the model; the tie comes from model_config.tie_word_embeddings, which every adapter that can tie its LM head uses to wire it. The op walk is iter_ops in layers/base.py, the walk finalize_quant had, now shared by both.
• The table stays a CPU tensor holding the host address, and the GPU kernels translate the pointer where they use it. This is the same rule fast_index_copy follows for pinned expert banks. host_ptr_identity / device_alias move from fast_index_copy.cuh into a shared header freetoken/pinned.h. On Windows/WDDM, where the device address can differ from the host address, they translate through cudaHostGetDevicePointer.
• The bf16 index kernel (including the TP masked variant) and ggml's op_get_rows accept a host weight. They take the launch device and stream from the output, and allocate the output on the ids device. With --embed-weights gpu they get the same tensors as before.
• The table's bytes count toward the host pin budget (_host_tables_bytes), next to qwen4_exp's PLE table, before the MoE offload cache plans its banks.

Results

1x H100 80GB HBM3 (driver 580.95.05), Xeon Platinum 8480+ (26 vCPUs), CUDA 13.0, torch 2.14.1+cu130, Python 3.14.3.

End-to-end: --embed-weights gpu vs host, 24 GiB free VRAM, AIME25 problem 1, 3 runs per side

A separate process held GPU memory so the server saw 24.00 GiB free. Both sides run this branch; gpu is the default.

checkpoint	table in host RAM	expert cache slots (gpu -> host)	decode tok/s, gpu	decode tok/s, host	answer 70
Qwen/Qwen3.6-35B-A3B (bf16 table)	0.95 GiB	2618 -> 2778	47.12 / 47.08 / 47.02, mean 47.07	51.12 / 50.61 / 51.78, mean 51.17 (+8.7%)	3/3, 3/3
ISTA-DASLab/Qwen3.8-Flash-Next-GSQ-RCO-GGUF, IQ3_XXS (GGUF table)	0.25 GiB	9593 -> 9761	88.20 / 90.74 / 90.51, mean 89.82	89.94 / 91.88 / 91.51, mean 91.11 (+1.4%)	3/3, 3/3

With Flash-Next, both sides produced the same completion token counts in each run (821, 778, 681).

python -m freetoken --model-path <checkpoint> --port 2350 --served-model-name default \
  --max-running-requests 1 --cuda-graph-max-bs 1 --max-seq-len-override 20480 \
  --moe-strategy offload --kv-reserve-tokens 20480 --embed-weights {gpu,host} \
  [--text-model-only for Qwen3.6-35B-A3B] [--ple-backend disk for Flash-Next]

The client sends one streamed /v1/chat/completions request per run:

• The prompt is the problem plus "Please reason step by step, and put your final answer within \boxed{}.".
• Sampling is temperature 1.0, top_p 0.95, top_k 20 (the checkpoints' recommended values), with max_tokens 16384.
• decode tok/s = (completion tokens - 1) / (time from the first to the last streamed token).

TTFT on this short prompt, gpu vs host: Qwen3.6 1.74 / 1.22 / 1.22 s vs 1.71 / 1.22 / 1.22 s; Flash-Next 3.02 / 1.11 / 0.81 s vs 3.02 / 0.83 / 0.81 s.

Lookup time, host vs GPU table

Measured inside a CUDA graph, in microseconds per lookup (host / GPU):

table	1 row	16 rows	512 rows	8192 rows
bf16 [151936, 2048]	2.6 / 1.2	4.1 / 1.3	67.6 / 1.8	1058 / 24.6
bf16 [248320, 2560]	2.6 / 1.3	4.8 / 1.4	85.7 / 2.1	1344 / 30.6
Q4_0 [32768, 2048]	3.8 / 1.6	5.0 / 1.8	66.8 / 4.4	1037 / 50.8
Q8_0 [32768, 2048]	3.9 / 1.6	5.3 / 1.8	104.4 / 4.3	1645 / 51.5

A decode-sized lookup (1 to 16 rows) costs 1.4 to 3.5 us more. An 8192-row prefill lookup costs 1.0 to 1.6 ms more. Long-prompt TTFT was not measured end to end.

Correctness

Host and GPU tables give bit-identical rows, eager and after CUDA graph replay with new ids. This covers bf16 tables with int32 and int64 ids, the TP masked variant on a half-vocab shard, and a Q4_0 GGUF table decoded to bf16 and fp32.

With the first commit rebased onto 20d5d93 (main with the NVFP4 get_rows case), a pinned NVFP4 GGUF table gives rows bit-identical to the same table on the GPU and to gguf-py's dequantize. This covers [64, 320], [32768, 2048] and [16384, 2560] tables, bf16 and fp32 output, eager and after CUDA graph replay with new ids.

Table selection was checked on 21 local checkpoints built on the meta device: qwen3_moe, qwen3_5, qwen4_exp (HF and GGUF), gpt_oss, DSV4, minimax_m2, muse_glimmer, GLM-4.7, and gemma4 (HF and GGUF). The config-based tie check gives the same keys, or the same refusal, as a walk of the actual LM-head wiring. The one difference is GLM-4.7 (DF11), which now gets the DF11 message instead of the tied one.

--embed-weights host on Qwen/Qwen3.5-2B (tied embeddings) stops at start-up with a ValueError that says the table is tied to the LM head.

Tests

PYTHONPATH=python pytest tests/kernels/test_ggml.py tests/kernels/test_pinned_tensor.py tests/moe/test_fused_copy.py \
  tests/moe/test_offload.py tests/engine tests/server tests/layers -m "not slow"   # 837 passed

The end-to-end and lookup numbers were measured with the first commit on c63fc13, before rebasing onto 7974c9c. The three commits in between touch GGUF mmproj vision input, two tests, and block-fp8 MoE decode. After the second commit, a Qwen3.6-35B-A3B start with host at 24 GiB free resolved the same 2778 expert cache slots. The tests above ran on the head of this branch.

ROCm

@LZ-QWQ validated this PR on R9700 (gfx1201) and W7900D (gfx1100) with ROCm 10.1.0 (comment):

• Host lookups of bf16 and GGUF Q4_0 / Q8_0 / NVFP4 tables matched the GPU and CPU references exactly, including the masked lookup and HIP Graph replay. The loading and refusal checks and this PR's tests passed.
• AIME25 problem 1: 3/3 correct per mode on each device, with identical output tokens between gpu and host.
• host reduced resident-weight VRAM by about 1.19 GiB and added 260 expert cache slots on both devices.

Decode tok/s with Qwen3.8-Flash-Next-FP8, 256 input / 128 output tokens, median of six measurements:

GPU	--embed-weights gpu	--embed-weights host
R9700	17.88	20.65
W7900D	21.24	21.65

Known issues

• With HF safetensors checkpoints the table still passes through VRAM while the checkpoint loads (see Design). host frees that VRAM after loading, but does not keep the table out of VRAM during loading.
• The kernel names are unchanged, so the freetoken-kernel-cache has to ship from the same build. A cache with a different git stamp is already rejected, but a dev build without a stamp only compares the release number and can load an older cache; --embed-weights host then fails at the first lookup with a device-check error. gpu is unaffected.

Not covered

• TP > 1 end to end (only the masked kernel was checked), and --dummy-weight with host end to end.
• A main-side run of the end-to-end table.
• An NVFP4 GGUF token_embd end to end: no local checkpoint has one.

[jason-fxz]

@zihaomu @LZ-QWQ could you help review this PR, especially the ROCm side? I have only checked ROCm by reading the code.

[zihaomu]

@LZ-QWQ please take a look

[LZ-QWQ]

@zihaomu @LZ-QWQ could you help review this PR, especially the ROCm side? I have only checked ROCm by reading the code.

@jason-fxz Validated on R9700 (gfx1201) and W7900D (gfx1100) with ROCm 10.1.0.

• Correctness: Dense and GGUF Q4_0/Q8_0/NVFP4 host lookups matched GPU/CPU references exactly, including masked lookup and HIP Graph replay. Loading/rejection checks and the PR's regression tests passed.
• AIME25 problem 1: 3/3 correct per mode on each device, with identical GPU/host output tokens.
• Memory: Host placement reduced resident-weight VRAM by approximately 1.19 GiB and added 260 expert-cache slots on both devices.

Decode performance with Qwen3.8-Flash-Next-FP8 (256 input / 128 output tokens; median of six measurements per configuration):

GPU	GPU embedding	Host embedding
R9700	17.88 tok/s	20.65 tok/s
W7900D	21.24 tok/s	21.65 tok/s

No ROCm-specific issues found in the tested PR paths.