fix(qwen4_exp): fix nvfp4 ftw weight loading, qsa indexer, and ple tables

[3Simplex]

I wanted to post this as a reference of the hacks needed to run oakmindai/Qwen3.8-Flash-Next-NVFP4-FTW which were generated by Gemini 3.7 flash.

Summary

Fixes weight loading, tensor layout mismatches, and host table initialization when running Qwen3.8-Flash-Next-NVFP4-FTW (and Qwen4Exp hybrid models) from FTW checkpoints.

Key Changes

1. python/freetoken/models/qwen4_exp/gdn.py:

   • Initialized out_proj as LinearReplicated to match unquantized BF16 weights in NVFP4 checkpoints where linear attention is excluded from quantization (quantization_config['ignore']).

2. python/freetoken/models/qwen4_exp/attention.py:

   • Added flexible load_state_dict() in Qwen4ExpIndexer to resolve checkpoints storing indexer weights flat under self_attn.index_* (e.g. index_qk_proj.weight, index_q_norm.weight) in addition to nested self_attn.indexer.*.

3. python/freetoken/engine/engine.py:

   • Added dynamic fusion in _load_weight_state_dict to concatenate and 16-row pad separate input_mix_weight_down.weight and block_inject_weight.weight tensors into input_mix_weight_down_block_inject.weight when loading FTW checkpoints.

4. python/freetoken/models/qwen4_exp/ple.py & model.py:

   • Added load_state_dict() override to NGramEmbedding so derived hashing metadata is not strictly expected in dense checkpoint state dicts.
   • Replaced .copy_() onto meta tensors with direct CUDA int64 tensor initialization using torch.cuda.current_device().
   • Added safe runtime device migration check in NGramEmbedding.row_ids().

5. python/freetoken/models/qwen4_exp/weight.py:

   • Added direct binary loader support in load_ple_table() for unified qwen4_ngram.bin + qwen4_ngram.json files.
   • Added FREETOKEN_ZERO_PLE=1 environment variable support allowing systems without 100+ GB of host RAM to attach ZeroTable.

---

Test Environment

• GPU: NVIDIA GeForce RTX 5060 Ti (16 GB VRAM)
• CPU: 16 physical cores / 16 threads (AVX-512 VNNI / AVX512BF16)
• System RAM: 128 GiB
• OS: NixOS / Linux

---

Hardware Bandwidth Benchmarks (ft bench bw --dtype nvfp4)

  host nixos-clayton   gpu cuda:0 (NVIDIA GeForce RTX 5060 Ti)   cpu 16c/16t
  ceilings: CPU STREAM read 41.3  |  PCIe linear H2D 28.0  D2H 28.3  GB/s   (threshold 2.0x)

  per-dtype (tuning — what the runtime backend pick matches on)
    format      expert       CPU-MoE   PCIe-gather  CPU/PCIe  backend
    nvfp4      7.61 MB     32.0 GB/s     26.2 GB/s     1.22x  offload
       overlapped: CPU-MoE 28.7 + PCIe 15.5 GB/s -> hybrid fetches 35.0% of misses

---

Generation Benchmark (ft serve / OpenAI API)

• Model: oakmindai/Qwen3.8-Flash-Next-NVFP4-FTW
• Backend: qsa_sparse attention + offload MoE
• Server Flags: FREETOKEN_ZERO_PLE=1 ft serve --max-running-req 1 --cuda-graph-max-bs 1

=== Generation Benchmark Results ===
- Tokens Generated: 127
- Time to First Token (TTFT): 9.76 ms
- Total Time: 9.65 s
- Decode Speed: 13.17 tokens/s

---

Verification Checklist

• Tested regression on KAT-Coder-V2.5-Dev-FTW (clean graph capture and generation).
• Tested full end-to-end inference on oakmindai/Qwen3.8-Flash-Next-NVFP4-FTW.
• Verified on real hardware with both interactive ft shell and API ft serve.

Environment Details (Auto-detected)

---

• OS: NixOS 26.11 (Zokor) (Linux 7.2.0)
• CPU: AMD Ryzen 9 9950X 16-Core Processor
• System RAM: 128 GiB & 78GiB used to serve model
• GPU 0: NVIDIA GeForce RTX 5060 Ti (16311 MiB VRAM)
• NVIDIA Driver: 595.91.07
• CUDA Runtime: 13.0
• PyTorch: 2.11.0+cu130
• Python: 3.12.14

[benwilson]

Data point from a different checkpoint and a card nobody has reported yet, in case it helps scope which of these hacks are FTW-specific: the standard (non-FTW) NVFP4 checkpoint loads and serves at 4b94bdc with none of these changes, on an RTX 3090.

Environment: FreeToken git 4b94bdc (HEAD 2026-08-30, no local patches), installed with pip install --no-build-isolation "freetoken[accel] @ git+…@4b94bdc"; torch 2.11.0+cu130, CUDA toolkit 13.0, driver 595.71.05; RTX 3090 24 GB (sm_86) on a rented vast.ai host (503 GiB RAM); image pytorch/pytorch:2.11.0-cuda13.0-cudnn9-devel; checkpoint dealignai/Qwen3.8-Flash-Next-ABLITERATED-NVFP4 @ be794b99 (standard compressed-tensors NVFP4 layout, byte-identical quant config to the RadixArk stock checkpoint — not the FTW repack).

ft serve --max-seq-len-override 262144 --kv-reserve-tokens 262144 --moe-backend offload --moe-cache-auto --nvfp4-backend triton --max-running-requests 1 --max-prefill-length 8192 --sampling-defaults model --reasoning-parser qwen3 --tool-call-parser qwen3_coder

• weights 206/206 in 7 s; PLE table 47.7 GB → pinned host RAM at ~540 MB/s; experts 63.3 GB at ~1 GB/s; --moe-cache-auto resolved moe_cache_size=1677 num_pages=4097; KV 262,208 tokens = 6.19 GiB; 23.28 GiB free before load, 2.01 GiB after CUDA-graph capture (bs=1, 100.8 s cold, 2 s on a warm JIT cache)
• worker RSS at ready: 113 GiB (so FREETOKEN_ZERO_PLE from this PR is the only path for <128 GiB hosts)
• needle-in-haystack found at 38,587 / 146,337 / 192,587 prompt tokens, zero warnings
• prefill in 8192-token chunks: 1,000–1,290 tok/s steady (380–480 on the cold first chunk); 192.6k tokens in 2 min 52 s
• decode 23–25 tok/s, flat from 38k to 192k context (engine counters); a 55k-token Claude Code turn: cold 56.6 s, warm 7.1 s from the radix cache

So on sm_86 with the Triton NVFP4 path the loader, indexer and PLE init are fine for the compressed-tensors layout; the FTW checkpoint's flat self_attn.index_* keys and split input_mix_weight_down / block_inject tensors look like the actual differentiator. Happy to run the FTW repack on the same box if a direct A/B would help.