Skip to content

[TP=2] Quantized MoE expert parallelism on heterogeneous consumer CUDA GPUs — a working NVFP4 port, bit-exact vs TP=1 #520

Open
Open
[TP=2] Quantized MoE expert parallelism on heterogeneous consumer CUDA GPUs — a working NVFP4 port, bit-exact vs TP=1#520

Description

@Yuepixel
Yuepixel
opened on Sep 18, 2026

[TP=2] Quantized MoE expert parallelism on heterogeneous consumer CUDA GPUs — a working NVFP4 port, bit-exact vs TP=1

Model: Qwen3.6-35B-A3B-NVFP4 Engine: FreeToken 0.1.3 Backend: CUDA / WSL2
GPUs: RTX 3090 Ti (sm_86, 24 GB) + RTX 4060 Ti (sm_89, 16 GB) Date: 2026-09-18

TL;DR

We ran a quantized MoE model with --tensor-parallel-size 2 across two different consumer GPUs — different architectures and different VRAM. At the time of writing, FreeToken rejects TP > 1 for every quantized expert format (tp_ok=False), and the Qwen3.6 weight loader is explicitly TP=1-only. We implemented per-rank intermediate sharding for the NVFP4 Triton MoE experts, made the GatedDeltaNet / full-attention geometry rank-local, and added a segment-aware checkpoint loader, then verified bit-identical greedy completions vs TP=1 on 5 prompts.

On this machine the heterogeneous pair is ~2× slower than one RTX 3090 Ti alone (the slower 4060 Ti bounds decode), but the two cards raise the context window the engine can allocate from 8,198 → 272,785 tokens (~33×). The point of this report is feasibility + correctness, not throughput.

This is a local proof of feasibility on one machine, not an upstream-ready PR. Caveats and the full patch are below.

1. Background

What is new in this report, as far as we can tell, is the combination: NVFP4 expert sharding, on heterogeneous sm_86 + sm_89 consumer cards, under CUDA/WSL, with a bit-exact correctness check.

2. Environment and model

  • WSL2 Ubuntu-24.04, CUDA toolkit 13.3, nvcc 13.3, Python 3.12.3.
  • rank 0: RTX 3090 Ti, sm_86, 24 GB. rank 1: RTX 4060 Ti, sm_89, 16 GB.
  • TVM_FFI_CUDA_ARCH_LIST="8.6 8.9" — a short two-value list that covers both cards without tripping the CUTLASS-DSL path from Feature request / guidance: tensor-parallel across heterogeneous and idle consumer GPUs (quantized MoE experts) - a hobbyist request, with reference to #478 #510.
  • Model: NVFP4 (modelopt mixed-precision) checkpoint, ~22.3 GB of weights:
    • 40 layers = 30 Qwen3_5GatedDeltaNet (linear attention) + 10 full attention (interval 4); hidden 2048.
    • MoE: 256 experts, top-8, moe_intermediate_size 512, shared expert 512, 1 MTP layer.
    • GDN: 16 k-heads / 32 v-heads, head dim 128, conv kernel 4.
    • Full attention: 16 q-heads / 2 kv-heads, head dim 256.
    • Quantization: experts + lm_head = NVFP4 W4A16 (group size 16); attention and linear-attention projections = FP8; KV cache FP8.
  • Important caveat: sm_86 and sm_89 have no native FP4 tensor cores, so NVFP4 executes as W4A16 (dequantize → fp16 matmul), not as native FP4.

3. The walls

(a) Every quantized expert kernel declares tp_ok=False.
MoEKernel._common_reject(...) returns "TP > 1 is not supported for this expert format" for all quantized formats: TritonNvfp4MoEKernel, TritonMxfp4MoEKernel, MarlinNvfp4MoEKernel (needs vLLM and max_slots), and B12xNvfp4MoEKernel (needs sm_120+). On sm_86/sm_89 only the Triton path is usable, and it was the one rejecting TP.

The machinery for expert TP already existed: MoEConfig carries tp_rank/tp_size and a local_intermediate = intermediate // tp_size, MoELayer calls _maybe_all_reduce(...) at tp_size > 1, and the BF16 unquantized.py kernel already uses local_intermediate. The NVFP4 kernel simply wasn't using the local geometry, so layout()/pack() sized the banks with the full intermediate and _common_reject refused anyway.

(b) Heterogeneous CUDA.
Two different compute capabilities and very different VRAM. engine._sync_get_memory() raises "Memory across TP ranks are imbalanced" when max_free - min_free > 2 GiB; a 24 GB + 16 GB pair trips this immediately. We relaxed the local gate to 16 GiB (environment-side change, not in the patch below). Also, per #510, a single TVM_FFI_CUDA_ARCH_LIST value produces "no kernel image", while a long multi-value list trips CUTLASS-DSL; "8.6 8.9" threaded the needle for us.

(c) The Qwen3.6 loader was TP=1-only, and the packed q|k|v tensors cannot be cut naively.
iter_weights (and the FP8/bf16 expert resident paths) raise NotImplementedError("qwen3_5_moe weight loading supports TP=1 only"). Beyond removing that, the dense reader emits full checkpoint shapes, while the (now local) modules expect local shapes — load_state_dict asserts exact shape equality.

Worse, two tensors pack q|k|v into one row axis: GDN in_proj_qkv and the depthwise conv1d.weight, both with row layout [key, key, value] where value == 2 * key. A plain contiguous row-chunk splits across the q/k boundary. Our first run produced exactly that symptom — fluent-looking but meaningless token salad — until the slicer became segment-aware.

4. The patch

Five files, 376-line unified diff (full text at the bottom).

File Change
layers/quantization/moe/nvfp4.py tp_ok=True; layout() / pack() use cfg.local_intermediate
models/nvfp4_banks.py _shard_expert_tensor() slices each expert per rank, .clone()d
models/qwen3_5_moe/weight.py drop the TP raise; _shard() for col/row/lm_head; segment-aware _split_qkv_rows(); passthrough slicing
models/qwen3_5_moe/gdn.py local head/conv geometry; out_proj: LinearReplicated → LinearRowParallel
models/qwen3_5_moe/attention.py local q/kv geometry; o_proj: LinearReplicated → LinearRowParallel

Expert sharding (NVFP4). gate/up and their block scales lead with the intermediate dim → split on dim 0; down and its scale trail with it (the packed input dim) → split on the last dim; the per-tensor global scales are scalars and stay whole. The Triton bank layout is derived from cfg.local_intermediate, so each rank builds exactly its half. The partial outputs are summed by the existing MoELayer._maybe_all_reduce. The math is exact: a MoE output is Σ_i silu(gate_i)·up_i @ down_i, which is linear in the intermediate index, so splitting i and all-reducing is safe.

Local geometry (GDN / attention). The linear layers are constructed with the full sizes and shard internally (LinearColParallelMerged splits each output segment with div_even; LinearRowParallel splits the input and all-reduces). But the forward reshape/split sites, the state pools, and the conv1d/A_log/dt_bias tensors must use the local head counts. out_proj/o_proj move from replicated to row-parallel so the head-partitioned input is reduced correctly.

Segment-aware slicer. For the packed [key, key, 2*key] row axis we narrow the three segments separately and concatenate, so rank r gets q-heads r·step, k-heads r·step, and v-heads 2·r·step — matching the contiguous head halves the fla kernels expect. torch.cat allocates fresh storage, so the result does not pin the parent.

.clone(). Following PR #70, every per-rank slice is .clone()d. Besides avoiding a leaked parent allocation (+5.1 GiB/rank there), this had a measurable upside here: after the fix the TP=2 KV budget rose from 208,273 to 272,785 tokens.

Kernel-level pre-check. Before touching the model we verified that the fla GDN kernels are shape-inferred and split-symmetric: full vs. split-and-concat on random bf16 inputs gave max_abs_diff = 4.2e-3 for prefill (bf16 rounding) and 0.0 (bit-identical) for decode. So the kernels themselves did not need modification.

5. Verification

Greedy decoding, temperature=0, top_p=1, seed=0. Five prompts, each run twice to confirm determinism; sha256 of the completion text compared between TP=1 and TP=2.

Prompt sha256[:16] (TP=1 = TP=2)
The capital of France is d5374e6e252ea9d6
1, 2, 3, d575a4226611b8cc
def fibonacci(n): 4ff3698b9a0bc851
The quick brown fox jumps over the lazy dog. 2763c7341483f2fa
Q: What is 17 times 24?\nA: f65f09decf72eacf

All repeats were deterministic and the TP=1 / TP=2 sets were identical. (For completeness: the first run before the segment-aware fix produced incoherent output — a useful negative control that the check is actually sensitive.)

6. Results

Config Decode Context limit KV pool
TP=1, single 3090 Ti 118–170 tok/s 8,198 tokens 0.16 GiB (moe_cache_size=9756)
TP=2, 3090 Ti + 4060 Ti 59–73 tok/s 272,785 tokens 2.60 GiB/rank (moe_cache_size=10240)
  • TP=2 measurements: short prompt 66, medium 64, long prompt 62, sustained 512-token generation 72.7, ~28K-token context 59.6 tok/s. No Prefill hangs on RTX 4060 Ti (SM89) for prompts beyond warmup length #72-style Ada prefill hang was observed.
  • TP=1 exceeds its 8,198-token context and the request is dropped: Input sequence length 24214 exceeds 8198.
  • Interpretation: on this uneven pair decode is bounded by the slower card, so TP=2 is ~2× slower than the lone 3090 Ti; the reason to use both cards here is the ~33× larger context window, made possible because each rank holds only half the weights.

7. Reproduction

  1. Install FreeToken 0.1.3 into a venv; confirm ft serve runs TP=1 on the fast card.

  2. Apply the patch below to the installed freetoken/ package tree (python -m py_compile each file).

  3. Start the server with a short mixed arch list and explicit GPU UUIDs:

    export TVM_FFI_CUDA_ARCH_LIST="8.6 8.9"
ft serve --model-path /path/to/Qwen3.6-35B-A3B-NVFP4 \
         --tensor-parallel-size 2 \
         --gpu GPU-<3090ti-uuid>,GPU-<4060ti-uuid> \
         --port 18020 --host 127.0.0.1 --text-model-only

  4. If rank free-memory differs by more than the built-in 2 GiB gate, relax engine._sync_get_memory locally (we used 16 GiB).

  5. On a WSL host, raise the VM memory so the pinned expert-bank budget fits (the bank for this model is 16.93 GiB; WSL pins ~40% of guest RAM). We used memory=48GB in .wslconfig.

  6. Run greedy comparisons between TP=1 and TP=2.

Full patch — unified diff, 5 files, 376 lines (applies to freetoken 0.1.3 source) 
--- a/freetoken/layers/quantization/moe/nvfp4.py
+++ b/freetoken/layers/quantization/moe/nvfp4.py
@@ -41,14 +41,14 @@
     cpu_format = "nvfp4"
 
     def unusable_reason(self, cfg: MoEConfig) -> str | None:
-        reason = self._common_reject(cfg, resident_ok=False, tp_ok=False, cpu_ok=True, plain_silu_only=False)
+        reason = self._common_reject(cfg, resident_ok=False, tp_ok=True, cpu_ok=True, plain_silu_only=False)
         if reason:
             return reason
         reason = gated_epilogue_reason(cfg)
         return f"triton nvfp4 MoE kernel: {reason}" if reason else None
 
     def layout(self, cfg: MoEConfig) -> dict[str, BankSpec]:
-        i, h = cfg.intermediate, cfg.hidden
+        i, h = cfg.local_intermediate, cfg.hidden
         return {
             "gate_up": BankSpec((2 * i, h // 2), torch.uint8),
             "gate_up_scale": BankSpec((2 * i, h // GROUP), FP8),
@@ -61,7 +61,7 @@
     def pack(self, pieces, cfg: MoEConfig, out):
         out["gate_up"].copy_(fused_piece(pieces, "gate_up"))
         out["gate_up_scale"].copy_(fused_piece(pieces, "gate_up_scale"))
-        out["gate_up_global"].copy_(fused_global(pieces, cfg.intermediate))
+        out["gate_up_global"].copy_(fused_global(pieces, cfg.local_intermediate))
         out["down"].copy_(pieces["down"])
         out["down_scale"].copy_(pieces["down_scale"])
         out["down_global"].copy_(global_rows(pieces["down_global"], cfg.hidden))
--- a/freetoken/models/nvfp4_banks.py
+++ b/freetoken/models/nvfp4_banks.py
@@ -8,6 +8,7 @@
 
 import safetensors
 import torch
+from freetoken.distributed import get_tp_info
 from freetoken.utils import download_hf_weight
 from tqdm import tqdm
 
@@ -63,6 +64,27 @@
     return {"weight": "", "weight_scale": "_scale", "weight_scale_2": "_global"}[kind]
 
 
+def _shard_expert_tensor(role: str, tensor: torch.Tensor, tp_rank: int, tp_size: int) -> torch.Tensor:
+    """Keep this rank's contiguous slice of an expert's intermediate dimension.
+
+    ``gate`` / ``up`` and their per-block ``_scale`` companions lead with the intermediate
+    dim, so they split on dim 0; ``down`` and its scale trail with it (the packed input
+    dim), so they split on the last dim. The per-tensor ``_global`` scales are scalars and
+    stay whole. Disabled at ``tp_size == 1`` so single-rank loading is byte-identical.
+    """
+    if tp_size <= 1:
+        return tensor
+    # ``.clone()``: a bare narrow() view keeps the whole per-expert checkpoint tensor alive
+    # (upstream TP PR #70 measured +5.1 GiB/rank from unpinned views); each rank owns a copy.
+    if role in ("gate", "up", "gate_scale", "up_scale"):
+        step = tensor.shape[0] // tp_size
+        return tensor.narrow(0, tp_rank * step, step).clone()
+    if role in ("down", "down_scale"):
+        step = tensor.shape[-1] // tp_size
+        return tensor.narrow(-1, tp_rank * step, step).clone()
+    return tensor
+
+
 def iter_nvfp4_expert_pieces(
     model_path: str,
     config,
@@ -88,6 +110,7 @@
     drop = drop_page_cache or _drop
     folder = download_hf_weight(model_path)
     weight_map = safetensors_weight_map(folder)
+    tp = get_tp_info()
 
     wanted: dict[str, tuple[int, int, str]] = {}
     for name in weight_map:
@@ -119,18 +142,20 @@
             with safetensors.safe_open(path, framework="pt", device="cpu") as f:
                 for name in by_shard[shard]:
                     tensor = f.get_tensor(name)
-                    if wanted[name][2].endswith("_global"):
+                    role = wanted[name][2]
+                    if role.endswith("_global"):
                         tensor = _ingest_global(spec, tensor)
-                    yield name, tensor
+                    yield name, _shard_expert_tensor(role, tensor, tp.rank, tp.size)
             drop(path)
 
     def _parallel():
         from freetoken.models.weight import iter_expert_tensors_parallel
 
         for name, tensor in iter_expert_tensors_parallel(folder, lambda n: n in wanted, workers=workers, chunk=chunk):
-            if wanted[name][2].endswith("_global"):
+            role = wanted[name][2]
+            if role.endswith("_global"):
                 tensor = _ingest_global(spec, tensor)
-            yield name, tensor
+            yield name, _shard_expert_tensor(role, tensor, tp.rank, tp.size)
 
     return per_expert_pieces(_parallel() if parallel else _serial(), wanted.get, tensors_per_expert=9)
 
--- a/freetoken/models/qwen3_5_moe/weight.py
+++ b/freetoken/models/qwen3_5_moe/weight.py
@@ -80,6 +80,46 @@
     return flat.contiguous()
 
 
+# TP splits of the checkpoint's fused linears: col-parallel modules lead with the
+# (intermediate / head) output dim, row-parallel modules trail with the input dim, and
+# lm_head splits its vocab rows. Replicated modules (router, norms) are left whole.
+_TP_COL_TARGETS = (".in_proj_qkvz", ".in_proj_ba", ".in_proj", ".qkv_proj", ".gate_up_proj")
+_TP_ROW_TARGETS = (".out_proj", ".o_proj", ".down_proj")
+# unquantized tensors the loader passes through: renamed to the module's local shape here
+_TP_PASSTHROUGH = ("embed_tokens.weight", "conv1d.weight", "A_log", "dt_bias")
+
+
+def _tp_narrow(tensor: torch.Tensor, dim: int, tp) -> torch.Tensor:
+    size = tensor.shape[dim]
+    if size % tp.size:
+        raise ValueError(f"cannot split dim {dim} of size {size} across {tp.size} TP ranks")
+    step = size // tp.size
+    # ``.clone()`` is mandatory: a bare narrow() is a view that keeps the whole checkpoint
+    # tensor's storage alive (upstream TP PR #70 measured +5.1 GiB per rank from exactly this).
+    return tensor.narrow(dim, tp.rank * step, step).clone()
+
+
+def _split_qkv_rows(tensor: torch.Tensor, tp) -> torch.Tensor:
+    """q|k|v packed rows ``[key, key, value]`` (value == 2*key for this family) -> one rank's
+    contiguous head halves. A plain row-chunk would cut across the q/k boundary, so the three
+    segments are narrowed separately and re-concatenated (the #478 segment-aware slicer).
+    ``torch.cat`` allocates fresh storage, so the result does not pin the parent tensor."""
+    rows = tensor.shape[0]
+    if rows % 4:
+        raise ValueError(f"packed q|k|v rows {rows} are not 4 * key_dim as GatedDeltaNet expects")
+    key = rows // 4
+    if key % tp.size:
+        raise ValueError(f"q/k rows {key} are not divisible across {tp.size} TP ranks")
+    step = key // tp.size
+    start = tp.rank * step
+    segs = (
+        tensor.narrow(0, start, step),
+        tensor.narrow(0, key + start, step),
+        tensor.narrow(0, 2 * key + 2 * start, 2 * step),
+    )
+    return torch.cat(segs, dim=0)
+
+
 def _dequant_nvfp4(weight: torch.Tensor, weight_scale: torch.Tensor, weight_global: torch.Tensor) -> torch.Tensor:
     """Packed NVFP4 -> bf16 on CUDA (the kernel is GPU-only, the converter reads on CPU), returned on the caller's device."""
     device = weight.device
@@ -189,6 +229,9 @@
             parts = [self._check(target, stored, part) for part in parts]
             if self.scheme(target) is None:
                 parts = [{"weight": _dequant(stored, part)} for part in parts]
+        tp = get_tp_info()
+        if tp.size > 1:
+            parts = [self._shard(target, i, part, tp) for i, part in enumerate(parts)]
         out = []
         for role in parts[0]:
             tensors = [part[role] for part in parts]
@@ -232,6 +275,36 @@
             out["input_scale"] = part["input_scale"].reshape(()).to(torch.float32)
         return out
 
+    def _shard(self, target: str, idx: int, part: dict[str, torch.Tensor], tp) -> dict[str, torch.Tensor]:
+        """Cut one fused part to this rank's TP slice (no-op at tp_size == 1).
+
+        Col-parallel targets: ``weight`` / ``weight_scale`` / ``weight_global`` split on
+        dim 0 (``input_scale`` is a scalar). Row-parallel targets: ``weight`` splits on
+        dim 1, and so do the 2-D block scales; a 1-D per-output-row scale (fp8 scalar
+        expanded by ``_check``) applies to the unsplit output rows and stays whole.
+        ``lm_head`` splits its vocab rows on dim 0. GatedDeltaNet's ``in_proj_qkv`` part
+        packs q|k|v in one tensor, so it needs the segment-aware slicer on dim 0.
+        """
+        if tp.size <= 1:
+            return part
+        if target.endswith(".in_proj_qkvz") and idx == 0:
+            for role in ("weight", "weight_scale", "weight_scale_inv", "weight_global"):
+                t = part.get(role)
+                if t is not None and t.dim() >= 1:
+                    part[role] = _split_qkv_rows(t, tp)
+        elif target.endswith(_TP_ROW_TARGETS):
+            part["weight"] = _tp_narrow(part["weight"], 1, tp)
+            for role in ("weight_scale", "weight_scale_inv"):
+                t = part.get(role)
+                if t is not None and t.dim() >= 2:
+                    part[role] = _tp_narrow(t, 1, tp)
+        elif target.endswith(_TP_COL_TARGETS) or target.endswith("lm_head"):
+            for role in ("weight", "weight_scale", "weight_scale_inv", "weight_global"):
+                t = part.get(role)
+                if t is not None and t.dim() >= 1:
+                    part[role] = _tp_narrow(t, 0, tp)
+        return part
+
 
 def iter_weights(
     model_path: str,
@@ -245,8 +318,6 @@
 
     Per-expert NVFP4 experts always come from the offload cache's expert reader.
     """
-    if get_tp_info().size > 1:
-        raise NotImplementedError("qwen3_5_moe weight loading supports TP=1 only")
     hf_config = cached_load_hf_config(model_path)
     config = parse_config(hf_config)
     stacked = include_moe_experts and config.is_moe and config.expert_quant == "none"
@@ -258,6 +329,7 @@
 
 
 def _iter_shards(model_path: str, device: torch.device, reader: _DenseReader | None, *, stacked: bool, include_vision: bool):
+    tp = get_tp_info()
     for file in tqdm(iter_weight_files(model_path), desc="Loading weights", disable=not get_tp_info().is_primary()):
         with safetensors.safe_open(file, framework="pt", device=str(device)) as f:
             for raw_name in f.keys():
@@ -279,6 +351,11 @@
                 elif _is_gemma_norm(name):
                     yield name, tensor + 1.0  # (1 + weight) baked into the stored norm weight
                 else:
+                    if tp.size > 1 and name.endswith(_TP_PASSTHROUGH):
+                        tensor = (
+                            _split_qkv_rows(tensor, tp) if name.endswith("conv1d.weight")
+                            else _tp_narrow(tensor, 0, tp)
+                        )
                     yield name, tensor
     if reader is not None and reader.pending:
         lines = reader.missing()
--- a/freetoken/models/qwen3_5_moe/gdn.py
+++ b/freetoken/models/qwen3_5_moe/gdn.py
@@ -3,9 +3,11 @@
 import torch
 import torch.nn.functional as F
 from freetoken.core import get_global_ctx
+from freetoken.distributed import get_tp_info
 from freetoken.kernel.causal_conv1d import causal_conv1d_decode, causal_conv1d_varlen
-from freetoken.layers import BaseOP, GatedRMSNorm, LinearColParallelMerged, LinearReplicated
+from freetoken.layers import BaseOP, GatedRMSNorm, LinearColParallelMerged, LinearRowParallel
 from freetoken.layers.quantization import QuantConfig
+from freetoken.utils import div_even
 
 from .gdn_kernels import gdn_decode_fla, gdn_prefill_chunk_fla
 
@@ -40,12 +42,22 @@
         assert head_k_dim == head_v_dim, (
             f"GatedDeltaNet requires head_k_dim == head_v_dim, got {head_k_dim} != {head_v_dim}"
         )
-        self.num_k_heads = num_k_heads
-        self.num_v_heads = num_v_heads
+        tp = get_tp_info()
+        # Per-rank geometry: the head dims the state pool, conv and fla reshape/split sites
+        # use are local, while the Linear layers are built with the FULL sizes and shard
+        # internally (ColParallel splits output rows, RowParallel splits input cols). At
+        # tp_size == 1 local == full, so the single-rank path is unchanged.
         self.head_k_dim = head_k_dim
         self.head_v_dim = head_v_dim
-        self.key_dim = num_k_heads * head_k_dim
-        self.value_dim = num_v_heads * head_v_dim
+        self.full_num_k_heads = num_k_heads
+        self.full_num_v_heads = num_v_heads
+        self.full_key_dim = num_k_heads * head_k_dim
+        self.full_value_dim = num_v_heads * head_v_dim
+        self.full_conv_dim = 2 * self.full_key_dim + self.full_value_dim
+        self.num_k_heads = div_even(num_k_heads, tp.size)
+        self.num_v_heads = div_even(num_v_heads, tp.size)
+        self.key_dim = self.num_k_heads * head_k_dim
+        self.value_dim = self.num_v_heads * head_v_dim
         self.conv_dim = 2 * self.key_dim + self.value_dim
         self.conv_kernel_size = conv_kernel_size
         # quantized checkpoints quantize qkv|z but not b|a, so the fusion splits into a qkvz GEMM and a ba GEMM with their own schemes (matches sglang / vLLM)
@@ -53,10 +65,12 @@
             quant_config is not None and quant_config.scheme_for(f"{prefix}.in_proj_qkvz") is not None
         )
 
-        self._in_proj_split = [self.conv_dim, self.value_dim, num_v_heads, num_v_heads]
+        # local forward split; the constructors get the FULL sizes (ColParallel shards them)
+        self._in_proj_split = [self.conv_dim, self.value_dim, self.num_v_heads, self.num_v_heads]
+        full_in_proj_split = [self.full_conv_dim, self.full_value_dim, num_v_heads, num_v_heads]
         if self._split_in_proj:
             self.in_proj_qkvz = LinearColParallelMerged(
-                hidden_size, [self.conv_dim, self.value_dim], has_bias=False,
+                hidden_size, [self.full_conv_dim, self.full_value_dim], has_bias=False,
                 quant_config=quant_config, prefix=f"{prefix}.in_proj_qkvz",
             )
             self.in_proj_ba = LinearColParallelMerged(
@@ -66,7 +80,7 @@
         else:
             # Fused input projection (one GEMM instead of four): qkv | z | b | a.
             self.in_proj = LinearColParallelMerged(
-                hidden_size, self._in_proj_split, has_bias=False,
+                hidden_size, full_in_proj_split, has_bias=False,
                 quant_config=quant_config, prefix=f"{prefix}.in_proj",
             )
         self.conv1d = _DepthwiseConv1d(self.conv_dim, conv_kernel_size)
@@ -74,11 +88,12 @@
         # and the fla kernel reads them as fp32) -- matches HF/sglang, and avoids a
         # per-call .float() upcast in the decode wrapper. The weight loader exempts
         # *.A_log / *.dt_bias from the model-dtype downcast.
-        self.dt_bias = torch.empty(num_v_heads, dtype=torch.float32)
-        self.A_log = torch.empty(num_v_heads, dtype=torch.float32)
+        self.dt_bias = torch.empty(self.num_v_heads, dtype=torch.float32)
+        self.A_log = torch.empty(self.num_v_heads, dtype=torch.float32)
         self.norm = GatedRMSNorm(head_v_dim, eps=rms_norm_eps)
-        self.out_proj = LinearReplicated(
-            self.value_dim, hidden_size, has_bias=False,
+        # row-parallel: the (value) input is split per rank, outputs summed by all_reduce
+        self.out_proj = LinearRowParallel(
+            self.full_value_dim, hidden_size, has_bias=False,
             quant_config=quant_config, prefix=f"{prefix}.out_proj",
         )
 
--- a/freetoken/models/qwen3_5_moe/attention.py
+++ b/freetoken/models/qwen3_5_moe/attention.py
@@ -4,9 +4,10 @@
 
 import torch
 from freetoken.core import get_global_ctx
-from freetoken.layers import BaseOP, GemmaRMSNorm, LinearColParallelMerged, LinearReplicated
+from freetoken.distributed import get_tp_info
+from freetoken.layers import BaseOP, GemmaRMSNorm, LinearColParallelMerged, LinearRowParallel
 from freetoken.layers.rotary import get_rope
-from freetoken.utils import nvtx_annotate
+from freetoken.utils import div_even, nvtx_annotate
 
 if TYPE_CHECKING:
     from freetoken.models.config import ModelConfig
@@ -21,24 +22,36 @@
         attn = paged_attention(q, k, v)
         out = o_proj(attn * sigmoid(gate))
 
-    TP note: uses replicated linears (tp=1 correctness milestone); swap to
-    column/row-parallel for tensor parallelism later.
+    TP note: q/k/v are column-parallel (the fused qkv_proj shards each of the q/k/v
+    segments across ranks) and o_proj is row-parallel (its input is the local head
+    slice; partial outputs are summed by all_reduce). q/k norms, rope and paged
+    attention all work per local head, so they need no communication.
     """
 
     def __init__(self, config: ModelConfig, layer_id: int, *, prefix: str = ""):
         head_dim = config.head_dim
         self.layer_id = layer_id
-        self.num_q = config.num_qo_heads
-        self.num_kv = config.num_kv_heads
         self.head_dim = head_dim
+        tp = get_tp_info()
+        # Full geometry (drives the packed qkv output sizes) vs local geometry (drives
+        # the forward splits, views and o_proj input). Equal at tp_size == 1.
+        self.full_num_q = config.num_qo_heads
+        self.full_num_kv = config.num_kv_heads
+        self.full_qo_attn_dim = self.full_num_q * head_dim
+        self.full_kv_attn_dim = self.full_num_kv * head_dim
+        self.num_q = div_even(self.full_num_q, tp.size)
+        self.num_kv = div_even(self.full_num_kv, tp.size, allow_replicate=True)
         self.qo_attn_dim = self.num_q * head_dim
         self.kv_attn_dim = self.num_kv * head_dim
 
         # Fused q/k/v projection (one GEMM instead of three); q half is 2x for the
-        # output gate. Split sizes: [num_q*head_dim*2, num_kv*head_dim, num_kv*head_dim].
+        # output gate. Full split sizes [num_q*head_dim*2, num_kv*head_dim, num_kv*head_dim]
+        # are handed to the column-parallel layer (it shards each segment); the forward
+        # split uses the resulting local sizes.
+        full_qkv_split = [self.full_num_q * head_dim * 2, self.full_kv_attn_dim, self.full_kv_attn_dim]
         self._qkv_split = [self.num_q * head_dim * 2, self.kv_attn_dim, self.kv_attn_dim]
         self.qkv_proj = LinearColParallelMerged(
-            config.hidden_size, self._qkv_split, has_bias=False,
+            config.hidden_size, full_qkv_split, has_bias=False,
             quant_config=config.quant, prefix=f"{prefix}.qkv_proj",
         )
         # Qwen3.5 uses Gemma-style (1+weight) RMSNorm; the weight loader bakes the +1
@@ -62,8 +75,8 @@
             ),
             mrope_layout=config.rotary_config.mrope_layout,
         )
-        self.o_proj = LinearReplicated(
-            self.qo_attn_dim, config.hidden_size, has_bias=False,
+        self.o_proj = LinearRowParallel(
+            self.full_qo_attn_dim, config.hidden_size, has_bias=False,
             quant_config=config.quant, prefix=f"{prefix}.o_proj",
         )

8. Limitations and caveats

  • This is a local feasibility result on a single machine, not an upstream PR. The patch is a unified diff against the installed 0.1.3 source tree; it has not been rebased on main and does not include tests.
  • Performance is not a win on this pair. Decode is bounded by the slowest card; we did not attempt an asymmetric shard ratio weighted toward the 3090 Ti, nor any communication tuning.
  • The rank-imbalance gate relaxation (2 GiB → 16 GiB) and the WSL memory setting are environment-side changes, not part of the patch.
  • Only the Triton NVFP4 expert kernel was made TP-capable. Marlin and B12x remain tp_ok=False (and are not usable on sm_86/sm_89 anyway).
  • The NVFP4 experts run as W4A16 here because these cards lack native FP4 tensor cores.
  • The vision tower was not exercised (--text-model-only); MTP was left alone.
  • Correctness evidence is bit-identical greedy output on 5 prompts, which is a strong but finite check. Independent reproduction and broader regression testing would be valuable.

9. Acknowledgements

This work builds directly on the community's reports and analysis, and we want to be precise about who contributed what:

Any mistakes in the above are ours.

10. Availability / next steps

The complete patch is embedded above. We're happy to:

This report is intended as a reproducible data point for the "quantized experts + heterogeneous GPUs" case, not as a claim that the current implementation is production-ready.

Activity

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Projects

    No projects

      Milestone

      No milestone

      Relationships

      None yet

      Development

      No branches or pull requests

      Issue actions