feat(models): support TP for qwen3_5_moe

[RuixiangMa]

Problem

support TP for Qwen3.5 MoE model in RoadMap #79

Qwen3.5 MoE model (Qwen3.6-35B-A3B) only supported TP=1. The weight loader raised NotImplementedError("qwen3_5_moe weight loading currently supports TP=1 only"), and all linear layers used LinearReplicated (no sharding, no all-reduce).

Implementation

Weight Sharding

Two sharding functions handle all weight types:

• _shard_tp: Simple chunk along a given dimension (for o_proj, embed_tokens, lm_head, A_log, dt_bias, expert down_proj).
• _shard_tp_parts: Per-part column-parallel sharding with optional GQA KV head replication. When num_kv_heads < tp_size, KV heads are replicated using rank * num_heads // world_size (matching shard_tensor in loader.py).

Qwen3.5's checkpoint stores pre-fused weights (qkv_proj, in_proj, gate_up_proj). Simple chunk on a fused weight gives incorrect results — e.g., chunking [gate, up] gives rank 0 the entire gate and rank 1 the entire up, instead of each rank getting half of both. The fix shards each sub-part independently:

Model Layers

• Attention (attention.py): TP-local head counts, _qkv_split computed from local heads with allow_replicate=True for KV. o_proj changed from make_replicated to make_row_parallel (row-parallel + all-reduce). .view() → .reshape() for non-contiguous tensors after torch.split.

• GDN (gdn.py): TP-local dimensions for heads, key_dim, value_dim, conv_dim. in_proj uses full sizes for LinearColParallelMerged (which handles its own sharding). conv1d, A_log, dt_bias use TP-local sizes. out_proj changed from make_replicated_quant to make_row_parallel_quant. FP8 path also uses TP-local dims.

• MoE (moe.py): OffloadMoELayer._maybe_all_reduce is a no-op (expert banks are full on each rank in offload mode). MoELayer (resident/fused) retains its own all-reduce.

Test Results

UT

tests/models/test_qwen3_5_tp.py::test_shard_tp PASSED
tests/models/test_qwen3_5_tp.py::test_shard_tp_parts PASSED
tests/models/test_qwen3_5_tp.py::test_shard_tp_parts_replicate PASSED

E2E

Offload mode (default):

ft serve --model Qwen/Qwen3.6-35B-A3B --tp-size 1/2/4

TP	Output	Status
1	"Here's a thinking process: 1. Analyze User Input: - Question: 'What is 2+3?'"	✓
2	Same	✓
4	Same	✓

Resident (fused) mode:

ft serve --model Qwen/Qwen3.6-35B-A3B --tp-size 1/2/4 --moe-backend fused

TP	Output	Status
1	Same as above	✓
2	Same	✓
4	Same	✓

GPU Memory per Card

TP	Mode	Expert	KV cache	Non-expert(Attn+GDN+Shared et al.)
1	offload	0 (host RAM)	75.9 GiB	~6 GB
2	offload	0	78.5 GiB	~3 GB
4	offload	0	80.0 GiB	~1.5 GB
1	Resident	~60 GB	19.0 GiB	~6 GB
2	Resident	~30 GB	51.5 GiB	~3 GB
4	Resident	~15 GB	68.0 GiB	~1.5 GB

CC @jason-fxz @andy-yang-1