FREETOKEN_FP8_DENSE=1 with a bf16 FTW checkpoint dies with KeyError in_proj_qkvz.weight #8

Environment: RTX 5080 16 GB (sm_120), AMD Ryzen 9 9950X, 123 GB RAM, Ubuntu 24.04, driver 610.57.04, CUDA 13.3, torch 2.11.0+cu130, triton 3.6.0. Checkpoint: nvidia/Qwen3.8-Flash-Next-NVFP4.

Branch: 5080deploy @ 2d56708 (upstream FlashML-org#389 load-time FP8 dense, merged onto main + FlashML-org#408 and others).

Steps:

1. Convert with FP8 off: ft checkpoint --model ~/models/Qwen3.8-Flash-Next-NVFP4 --out ~/models/Qwen3.8-Flash-Next-NVFP4-FTW (bf16 dense). It serves fine without the env vars.
2. Serve that FTW with FP8 on:

FREETOKEN_FP8_DENSE=1 FREETOKEN_FP8_LMHEAD=1 ft serve --model ~/models/Qwen3.8-Flash-Next-NVFP4-FTW --kv-cache-dtype nvfp4 --num-tokens 262144 --max-seq-len-override 262144 --max-running-requests 1 --cuda-graph-max-bs 1 --moe-cache-size 512 --disable-moe-prefill-overlap --memory-ratio 0.95 --max-prefill-length 4096 --skip-prefill-warmup

What happens: startup dies 3 s after loading begins:

File ".../layers/base.py", line 43, in load_state_dict
  item = state_dict.pop(_concat_prefix(prefix, name))
KeyError: 'model.layers.0.linear_attn.in_proj_qkvz.weight'

parse_config reads the env var and builds the FP8 layout (split in_proj_qkvz / in_proj_ba, Fp8Dynamic* layers). The quantization itself lives in qwen4_exp.weight.iter_weights, which the FTW fast path doesn't run. So the FTW's bf16 fused in_proj never becomes the FP8 keys the model expects. The same env vars against the HF checkpoint dir load and serve correctly.

Expected: either quantize at FTW load as well, or detect the mismatch (bf16 FTW + FREETOKEN_FP8_DENSE=1) and say so instead of a bare KeyError. Whether ft checkpoint run with the env vars set produces a loadable FP8 FTW is being tested separately.

Related

• Upstream: feat(qwen4_exp): load-time per-tensor FP8 dense projections (W8A8 via _scaled_mm) FlashML-org/FreeToken#389 (load-time FP8 dense, closed by its author) and feat(qwen4_exp): tensor parallelism for Qwen3.8-Flash-Next (offload backend) FlashML-org/FreeToken#385 (Qwen3.8-Flash-Next TP, which feat(qwen4_exp): load-time per-tensor FP8 dense projections (W8A8 via _scaled_mm) FlashML-org/FreeToken#389 is stacked on); FTW handling for this model in fix(models): serve qwen4_exp FTW checkpoints with PLE streamed from source FlashML-org/FreeToken#405 (open) and fix(checkpoint): write the Qwen3.8-Flash-Next PLE table next to the FTW FlashML-org/FreeToken#420 (in main as 3d919e9).
• Fork PR: 5080deploy 9/9: merge upstream #389 (+#385), load-time FP8 dense projections #17 (merge of feat(qwen4_exp): load-time per-tensor FP8 dense projections (W8A8 via _scaled_mm) FlashML-org/FreeToken#389; documents the FP8 FTW conversion that does load).