qwen3_5_moe is TP=1-only because the model never matched the engine's per-rank GDN state contract — patch makes TP=2 work on 2x RX 9070 XT (ROCm) #478

Activity

1. No-LetterHead4141 commented on Sep 14, 2026

   

   No-LetterHead4141

   on Sep 14, 2026

   Author

   More actions

   Cross-referencing related work I found after opening this, so it's clear where this sits: #447, #385 (draft), #440.

   What #447 covers, and what it leaves open. #447 removes exactly one TP gate — qwen4_exp weight loading supports TP=1 only. For qwen3_5_moe it keeps the existing gate and adds a second restriction:

   if ownership is not None:
       raise NotImplementedError("block-FP8 expert banks do not support owner-local TP+EP")
   if get_tp_info().size > 1:
       raise NotImplementedError("qwen3_5_moe fp8 expert banks support TP=1 only")

   It also scopes its GDN TP to the unquantized path:

   if self._split_in_proj and tp_size > 1:
       raise NotImplementedError("qwen4_exp dense TP currently supports the BF16 GDN path only")

   So this issue is the complementary gap rather than a duplicate: qwen3_5_moe at TP>1, with block-FP8 experts and the quantized (split_in_proj) GDN path that an FP8 checkpoint like Qwen3.6-35B-A3B-FP8 takes.

   Independent convergence on the GDN design. I implemented this for qwen3_5_moe before reading #447, and the two arrived at the same shape, which is probably worth something as cross-validation — same approach, different model family, different vendor:

   • div_even(num_k_heads/num_v_heads, tp_size, allow_replicate=True) for rank-local head counts
   • LinearReplicated → LinearOProj on out_proj
   • the fused input GEMM built from global sizes, with the forward splitting by rank-local sizes
   • a segment-aware loader shard that splits in_proj_qkv / conv1d into q/k/v, shards each, and re-concatenates (a plain contiguous half of the fused block would hand rank 0 [q|k] and rank 1 [v])
   • A_log / dt_bias sized to local num_v_heads

   The framing in my original report still holds as the underlying reason all of this is needed: _linear_local_dims already allocates GDN conv and recurrent state per rank, so the engine has always expected head-sharded GDN, and replicating it writes full-head state into half-size buffers.

   What my patch adds on top of #447's direction: qwen3_5_moe GDN and attention; block-FP8 routed experts sharded on the intermediate dim; and per-rank expert bank pools, since alloc_persistent_layer_banks has rank 0 fill and every rank map rank 0's files, which cannot represent a rank-sharded expert layout (and PersistentPool.prepare_dir evicts sibling per-rank pools).

   Happy to rebase onto #447 once it lands rather than propose anything competing — it looks like the better base. Validated on 2× RX 9070 XT (gfx1201, ROCm 7.2.1): greedy output identical to TP=1, 120-request soak with 0 errors, 13.5 → 17.7 tok/s, first-token 10.0 → 5.0 s.
2. CraigStone-Dev commented on Sep 24, 2026

   

   CraigStone-Dev

   on Sep 24, 2026

   More actions

   Data point from the NVIDIA side, in case it helps situate the patch: we ran the qwen3_5_moe TP=2 work that grew out of #104 (RuixiangMa's base TP + a 9-commit stack completing the quantized loaders and a resident NVFP4 fused path) on 2× RTX PRO 4000 Blackwell with nvidia/Qwen3.6-35B-A3B-NVFP4:

   • --moe-backend fused (GPU-resident experts), TP=2: 141 tok/s client wall (engine ~146), 22.4 GB VRAM/GPU, boots clean, coherent output. That's 5–9× the offload (~15) and hybrid (~27) TP=2 paths on the same box, since the experts are fully resident and CUDA graphs apply.
   • @alogotron reproduced it almost exactly on 2× RTX 3090 (Ampere, PCIe): 139.8 tok/s client wall, 17.3 GB/GPU, 568k-token KV capacity — see the feat(models): support TP for qwen3_5_moe #104 thread.

   For reference, that stack is parked at CraigStone-Dev/FreeToken:feat/qwen3-5-tp-support (we held off pushing it forward while #447 and the sharded-bank direction in the #385 thread settled), and it is now rebased onto current main as rebase/tp-20260923 (5 commits, CPU suite green) if anyone wants the quantized-loader side (per-head GDN sharding, block-FP8/mixed checkpoint handling, resident NVFP4 expert banks) as a reference.