fix(models): load NVFP4 expert banks from checkpoints without a safetensors index

[paulp83]

Problem

Launching Qwen3.8-Flash-Next (the #257 model) with the offload MoE backend crashes at startup:

FileNotFoundError: [Errno 2] No such file or directory: '.../RadixArk--Qwen3.8-Flash-Next-NVFP4/snapshots/.../model.safetensors.index.json'
  File "python/freetoken/models/nvfp4_banks.py", line 222, in load_nvfp4_expert_source_banks_parallel

Reproduced with:

ft serve --model RadixArk/Qwen3.8-Flash-Next-NVFP4 \
    --moe-backend offload --moe-cache-auto \
    --max-running-requests 1 --kv-reserve-tokens 50000

The dense weights load fine ("Loading weights: 206/206"), then the scheduler dies when the offload cache builds the expert banks.

Root cause

The RadixArk Qwen3.8-Flash-Next-NVFP4 checkpoint ships no model.safetensors.index.json. Its weights are split across three shard families:

• model-bf16-*.safetensors — dense bf16 weights
• model-plefp8-*.safetensors — the n-gram PLE table
• layer-{L}-experts-{A}-{B}.safetensors — routed NVFP4 experts (48 layers × 512 experts)

The dense path (iter_weights → iter_weight_files) and the PLE path (_ple_table_files) both already handle a missing index by globbing *.safetensors and reading shard headers — which is why the dense load succeeds. But both NVFP4 expert-bank loaders in models/nvfp4_banks.py (serial and parallel) hard-opened the index file to build their tensor→shard map, so they raise FileNotFoundError as soon as the bank build starts. (iter_expert_tensors_parallel, which the parallel path delegates to, already has the no-index fallback — but it's only reached after the outer function crashes on the index read.)

Fix

One helper in models/nvfp4_banks.py, used by both loaders:

• _weight_map(folder) — returns the name→shard map from model.safetensors.index.json when present, else builds it from each *.safetensors shard's header (the same fallback iter_expert_tensors_parallel already uses).

Checkpoints that do carry an index keep using it unchanged; this only adds the fallback.

Tests

Added two tests in tests/models/qwen4_exp/test_weight.py (serial + parallel) that build a no-index checkpoint whose only tensors are routed NVFP4 experts — split across two shards named like the real layer-*-experts-* files — and assert the full 6-bank placement (gate/up fused on the output-row axis, down separate, scalar weight_scale_2 broadcast into the per-row global banks). Both fail with the launch crash's exact FileNotFoundError before the fix, pass after.

• tests/models/qwen4_exp/test_weight.py: 24/24 pass.
• tests/models/ + tests/moe/: identical results before/after (the pre-existing failures in this environment are all FileNotFoundError: 'ninja' from JIT-compiling CUDA extensions; 56 failed / 192 passed on baseline → 56 failed / 194 passed with the fix, +2 = the new tests).

Validation against the real checkpoint

_weight_map on the actual snapshot dir maps all 296,475 tensors across the 206 shards (matching the dense loader's 206/206), and the NVFP4 expert spec regex then finds exactly 48 × 512 × 3 × 3 = 221,184 expert tensors — every tensor the bank loaders need.

🤖 Generated with Claude Code