feat(qwen4_exp): tensor parallelism for Qwen3.8-Flash-Next (offload backend)

[gdevenyi]

What this adds

ft serve --tp-size 2 for qwen4_exp (Qwen3.8-Flash-Next) on the offload MoE backend. Upstream refuses TP>1 for this architecture; with two 48 GiB cards the model runs as two independent TP=1 instances, each keeping 37% of the NVFP4 experts resident and streaming the rest over PCIe every step.

Per rank, the patch:

• Shards the dense weights at load (weight.py::_shard): attention qkv_proj by head ([q|gate] per head; kv heads split, or replicated when there are fewer than ranks), GDN in_proj as its six parts ([q | k | v | z | b | a]) plus the matching conv1d channels and A_log / dt_bias, shared-expert gate_up_proj per part; o_proj, out_proj and the shared down_proj row-parallel with the all-reduce inside; embed_tokens / lm_head by vocab rows. Router, QSA indexer, norms, hyper-connections and PLE stay replicated, so every rank selects the same sparse blocks and n-gram rows.
• Shards the NVFP4 expert banks along the intermediate axis (nvfp4_banks.py, I=640 -> 320 per rank): packed codes, the 16-wide scale blocks and the per-row globals for the gate/up rows and the down columns. The offload cache then holds half the experts per rank.
• One all-reduce per MoE layer (moe.py): the routed and shared partial sums are combined as routed + sigmoid(gate) * shared before a single reduce, instead of one reduce each.
• LinearColParallelMerged(local_output_sizes=) for the kv-replicated case (the same shape as feat(models): support TP for qwen3_5_moe #104's hunk), and distributed_timeout 60 s -> 1800 s: behind a 100+ GiB load the ranks reach their first collective minutes apart, and 60 s kills the launch.

Everything else in the engine (scheduler, KV pool, PLE table, CUDA graphs) is untouched.

Measurements

2 x RTX 6000 Ada (48 GiB, sm_89, PCIe Gen4 x16, no NVLink), 2 x Xeon Gold 6526Y, 503 GiB RAM. RadixArk/Qwen3.8-Flash-Next-NVFP4, --moe-backend offload --ple-backend pinned --num-tokens 262144 --memory-ratio 0.94 --moe-prefill-hit-d2d.

run (both GPUs, KV 262,144 tokens)	single-stream tok/s	8 concurrent tok/s	expert residency	TTFT, 1k prompt
two TP=1 instances (previous layout)	55.9 / 56.7	266.1 aggregate	36.6% each	2.05 s
TP=2, 8 running requests	88.6	298.8	100%	1.55 s
TP=2, 16 running / graph bs 16	89.8	324.5	94.5%	0.84 s

A 262,144-token prompt reaches its first token in 74 s at TP=2 (116 s at TP=1). The decode step is dominated by the bf16 dense read; TP=2 halves it per GPU and removes the PCIe expert gather because everything fits.

Correctness. An 8-question probe gives identical answers at TP=1 and TP=2. Three raw prompts decoded greedily for 256 tokens: the ~1k-token prompt (QSA over many blocks, GDN state, PLE context) is word-for-word identical across every run; the two short prompts diverge after 13 and 29 words between TP=1 and TP=2, but the TP=2 server diverges from itself at the same points on a second pass (bf16 atomics in the expert kernels), so that is run-to-run noise rather than a sharding error. Greedy output at TP=2 is not bit-exact between passes.

Limits

Offload backend with bf16 dense projections only: fp8_block / nvfp4 dense checkpoints raise under TP (row-parallel FP8 / NVFP4 linears do not exist yet, the same gap #104 has). The hybrid / CPU MoE backends are not sharded.

Related: #62, #29 (TP for offloaded MoE), #104 (TP for qwen3_5_moe, which this reuses the merged-linear hunk from).

Testing

• tests/models/qwen4_exp/test_tp_shard.py: the per-head / per-part row sharding of every fused projection reassembles to the original (CPU).
• tests/models/test_nvfp4_banks_tp.py: the bank placer's per-rank slices of codes, scales and globals cover the intermediate axis exactly once (CPU).
• Served for a day on the machine above at TP=2 with the numbers in the table. tests/models/qwen4_exp on one of its GPUs: 97 passed, 3 failed; the same 3 (test_chunked_prefill_matches_one_shot[*], a bit-exact assertion off by bf16 noise on this torch 2.11 / flashinfer 0.6.18 / triton 3.6 stack) fail on plain main there too.

🤖 Generated with Claude Code

https://claude.ai/code/session_0173pf9k9fSVtwbm3f898HDt

[gdevenyi]

tests/models/qwen4_exp on a RTX 6000 Ada with this branch: 97 passed, 3 failed; plain main (af71ba4) on the same GPU: 94 passed, the same 3 failed (test_qsa_backend.py::test_chunked_prefill_matches_one_shot[*], a bit-exact chunked-vs-one-shot assertion that is off by bf16 noise with torch 2.11.0+cu130 / flashinfer 0.6.18 / triton 3.6.0). So nothing here regresses the package; the second commit fixes the one regression the first had (test_weight.py feeds iter_weights a synthetic checkpoint without model_type, and the loader now reads the config only when TP > 1).