feat(models): add Qwen3-Next-80B-A3B support

[akushonkamen]

feat(models): add Qwen3-Next support

Summary

This PR adds serving support for Qwen3-Next-80B-A3B (Qwen3NextForCausalLM) on top of the existing qwen3_5_moe runtime — same GDN + MoE family, flat model.layers.* weights, per-expert gate/up/down. The work is in telling the two checkpoint dialects apart at load time and in the router, not in the runtime itself.

What's included

• Registration: Qwen3NextForCausalLM is registered in models/register.py and reuses the qwen3_5_moe package (Qwen3_5MoEForCausalLM).
• Checkpoint dialect probing (qwen3_5_moe/config.py):
  • _is_ct_storage() probes the first shard to tell modelopt NVFP4 (two-level scales weight_scale/weight_scale_2, bf16 shared_expert, bf16 attention) from compressed-tensors packed FP4 (weight_packed/weight_global_scale).
  • _gdn_split_layout() detects the pre-fused GDN in_proj layout (in_proj_qkvz/in_proj_ba vs the four unfused parts).
  • _shared_expert_quant() probes whether the MoE shared_expert is packed NVFP4 or left bf16 (modelopt ignore list excludes it; the previous unconditional dense_quant = "nvfp4" broke load_state_dict with a missing weight_scale KeyError on those checkpoints).
  • parse_config now takes an optional model_path so every model config gets the extra parameter; the engine passes self.model_path through.
• Weight loading (qwen3_5_moe/weight.py):
  • _gdn_split_reorder() de-interleaves the per-k-head-group qkvz/ba rows into the contiguous [q|k|v|z]/[b|a] split the GDN expects, including the per-128-row fp8 scale blocks that alias head_dim. Loading the interleaved layout as-is scrambles q/k/v/z silently — output is fluent garbage and decode speed looks normal.
  • The NVFP4 expert-key regex is relaxed to accept model.layers.* (no language_model. prefix, which Qwen3-Next checkpoints don't carry).
  • iter_weights is split into a thin reorder/dispatch shell over a new _iter_weights_flat core, so the de-interleave runs once per loaded tensor without duplicating the shard-walk.
• Router (moe/fused.py, layers/moe.py): fused_topk/_torch_fused_topk gain a scoring parameter ("softmax" default, "sigmoid" for Qwen3-Next). Sigmoid scoring has no triton_kernels counterpart and Qwen3-Next's top-10 is not a power of 2, so both cases take the pure-torch path; the layer passes config.moe_scoring_func through.
• GDN module (qwen3_5_moe/gdn.py, model.py): the pre-fused split layout (in_proj_split) takes the same two-GEMM path as fp8, but with a bf16 in_proj_qkvz GEMM.
• bench_decode_moe: --extra-args passthrough to the spawned server.
• bf16 per-expert offload banks (qwen3_5_moe/weight.py): original (un-quantized) Qwen3-Next checkpoints store each routed expert as separate gate_proj/up_proj/down_proj tensors per expert. The dense loader previously binned every .mlp.experts.<int>. key as an NVFP4 offload tensor, so a bf16 checkpoint fed zero expert layers to the bank builder and died with "Missing MoE expert source layers". A _Bf16ExpertPacker now fuses gate|up, stacks all num_experts per layer into [E, 2*I, H] / [E, H, I], and yields the whole-layer experts.gate_up_proj / experts.down_proj sources the bank builder expects — in both the serial and parallel reader paths.
• docs/models.md: Qwen3-Next row added with checkpoint links.
• Tests: tests/models/test_qwen3_next_weights.py covers the in_proj_qkvz/in_proj_ba de-interleave round-trip (bf16 weights + the per-128-row fp8 scale blocks that alias head_dim), pass-through of unrelated names, rejection of ambiguous row counts, and sigmoid-router correctness against a hand-rolled reference.

What's NOT included

• No kernel/pynccl rename changes (the prometheus.NCCLWrapper → freetoken.NCCLWrapper FFI rename is a local artifact; upstream FreeToken already has the correct name).
• No benchbw fallback path.
• No DSV4 ds_fp4 machinery (separate baseline, not FreeToken).
• No MTP / speculative-decoding changes.

Test hardware

• GPU: single NVIDIA RTX 4090 (24 GB)
• Driver: NVIDIA 535.x / CUDA 12.4 (driver version as reported by nvidia-smi)
• Backend: offload (experts in host RAM, LRU on GPU) and hybrid (CPU + PCIe overlap)

Checkpoints

Checkpoint	HF ID
Qwen3-Next-80B-A3B-Instruct (bf16)	Qwen/Qwen3-Next-80B-A3B-Instruct
Qwen3-Next-80B-A3B-Instruct-FP8	Qwen/Qwen3-Next-80B-A3B-Instruct-FP8
Qwen3-Next-80B-A3B-Instruct-NVFP4	nvidia/Qwen3-Next-80B-A3B-Instruct-NVFP4

Run command

# offload (NVFP4)
ft serve --model nvidia/Qwen3-Next-80B-A3B-Instruct-NVFP4 --moe-backend offload

# hybrid (NVFP4)
ft serve --model nvidia/Qwen3-Next-80B-A3B-Instruct-NVFP4 --moe-backend hybrid

# offload (FP8)
ft serve --model Qwen/Qwen3-Next-80B-A3B-Instruct-FP8 --moe-backend offload

Performance

Single RTX 4090, greedy decode, warmup + 3 runs (median):

Checkpoint	Backend	tok/s
Qwen3-Next-80B-A3B NVFP4	offload	96.9
Qwen3-Next-80B-A3B NVFP4	hybrid	33.2
Qwen3-Next-80B-A3B FP8	offload	40.6

Commit breakdown

1. feat(models): add Qwen3-Next support — registration, dialect probing, de-interleave, sigmoid router, bench passthrough, docs, tests.
2. feat(models): bf16 per-expert offload banks for qwen3_5_moe — _Bf16ExpertPacker for plain-bf16 per-expert checkpoints in serial + parallel reader paths.