feat(qwen3_5_moe): tensor-parallel weight sharding for multi-GPU serving

[zhangnju]

Summary

This change enables tensor-parallel (TP > 1) serving of the Qwen3.5 and Qwen3.6 MoE models. Every class of weight is sharded across ranks at load time, so a checkpoint that does not fit on one GPU can be served across several. It covers bf16, fp8-block, and NVFP4 experts, and it is communication-backend agnostic (it works over either pynccl or torch-nccl).

What is sharded

Each weight is split along the dimension that matches how it is used in the forward pass:

• Dense linears: the attention q/k/v and the MLP gate/up projections are column-parallel (split on their output dimension); o_proj and down_proj are row-parallel (split on their input dimension); the embedding and the lm_head are split along the vocabulary dimension.
• GDN mixer (Qwen3_5GatedDeltaNet): each rank keeps its local key/value heads; out_proj is row-parallel (with an all-reduce); and dt_bias, A_log, and conv1d are allocated at the per-rank sizes. The linear-state pool is already TP-local, so it needs no change.
• Gated attention (Qwen3_5Attention): each rank keeps its local query and key/value heads (with GQA-aware key/value replication when there are fewer KV heads than ranks); the o_proj is row-parallel.
• Vision tower (Qwen3-VL): qkv and fc1 are column-parallel (including their biases); proj and fc2 are row-parallel.
• MoE experts: the unified expert banks are sharded along the intermediate dimension in build_expert_banks._shard_piece. The gate and up tensors (and their scales) are split along their output rows; the down tensor and its scale are split along their input columns. This path handles bf16 stacked experts, fp8-block experts (whose 128-block gate_up_scale and down_scale companions are split the same way), and NVFP4 experts (packed e2m1 codes, a block scale, and the global scales).

Implementation notes

• Quantized formats carry companion tensors that must be sharded consistently with the weight they belong to. An fp8-block weight carries a 128-block weight_scale_inv (named gate_up_scale / down_scale in the expert banks). An NVFP4 weight carries a per-output-row weight_global of shape [N], plus an optional per-tensor input_scale. The sharding follows a single rule: a companion is split along the same axis as its weight when it has that axis, and is replicated otherwise. For example, a per-output-row global is split by head for a column-parallel in_proj; the same global is replicated for a row-parallel out_proj, because out_proj's output dimension is not sharded; and a per-tensor scalar is always replicated.
• The NVFP4 expert bank layout() and pack() size the banks by cfg.local_intermediate rather than the full intermediate, so each rank allocates only its own slice. Without this, every rank would allocate full-size banks and tensor parallelism would not reduce expert memory. (The marlin and b12x NVFP4 kernels are CUDA-only and still size by the full intermediate; they are untested at TP > 1 and left as a follow-up.)
• LinearColParallelMerged divides each output segment's width by the TP size internally, so the mixers pass the full segment widths to the layer and split the output using the local widths in the forward pass.

Validation

• CUDA (2 x RTX 4090): Qwen3.5-9B at TP=2 produces greedy output that is byte-identical to TP=1 on two prompts, over both pynccl and --disable-pynccl (torch-nccl). This confirms the sharding is platform-agnostic.
• ROCm (8 x R9700, gfx1201):
  • bf16 Qwen3.6-35B-A3B serves coherently at TP=4 and TP=8.
  • fp8 Qwen3.6-35B-A3B-FP8 produces byte-identical greedy output across TP=1 (offload), TP=2, and TP=4 on five prompts.
  • NVFP4 Qwen3.6-35B-A3B-NVFP4 serves coherently at TP=2. (This requires a TP-capable NVFP4 kernel at serve time; the stock triton NVFP4 kernel rejects TP > 1.)
• After rebasing onto the current main, NVFP4 at TP=2 and fp8 at TP=4 were both re-checked and still serve coherently.

Notes

• The change is communication-backend agnostic: all collectives go through DistributedCommunicator. Enabling RCCL/pynccl on AMD is a separate concern and is not part of this PR.
• On RDNA, serving the MoE paths depends on the ROCm-7.14 / torch-2.14 compatibility fixes (notably the moe_align torch fallback) from the ROCm-enablement PR. This PR is intended to sit on top of that stack.

Test plan

• CUDA: Qwen3.5-9B TP=2 output byte-identical to TP=1 (pynccl and torch-nccl)
• ROCm: bf16 Qwen3.6-35B-A3B coherent at TP=4 and TP=8
• ROCm: fp8 Qwen3.6-35B-A3B byte-identical across TP=1 / TP=2 / TP=4
• ROCm: NVFP4 Qwen3.6-35B-A3B coherent at TP=2
• marlin / b12x NVFP4 at TP > 1 (CUDA) — not covered by this PR