Repository navigation
[TP=2] Quantized MoE expert parallelism on heterogeneous consumer CUDA GPUs — a working NVFP4 port, bit-exact vs TP=1 #520
Description
[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
- Issue 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 (2× RX 9070 XT / ROCm / FP8) showed that
Qwen3.6-35B-A3B-FP8can be run TP=2 and be made bit-exact against TP=1. We used it as the methodological blueprint (segment-aware loader slicing, GDN head splitting, per-rank expert banks). - Issue Feature request / guidance: tensor-parallel across heterogeneous and idle consumer GPUs (quantized MoE experts) - a hobbyist request, with reference to #478 #510 (hybrid TP on mixed GPUs /
TVM_FFI_CUDA_ARCH_LIST) framed the architecture-list problem. The deep diagnosis in that thread (single value → "no kernel image"; a long multi-value list → the CUTLASS-DSL crash; root cause in flashinferrmsnorm.pychoosingenable_pdlfromget_device_capability()while compilation uses the sharedARCH_LIST) is what let us pick a working arch list. - PR [1/3] feat(dsv4): add tensor-parallel DeepSeek-V4 runtime #70 (
[1/3] feat(dsv4): add tensor-parallel runtime, open) is the reference implementation for explicit per-tensor ownership. Two of its lessons were directly useful here: (1) weight and quant scale must be sharded along compatible axes, and (2) a barenarrow()view keeps the parent allocation alive — each per-rank slice must be.clone()d. - A community patch for a dual-3080 SGLang setup (
kk-pcl/sglang-qwen38-dual-3080-patch) is prior art for the same class of model and helped confirm the expected failure modes (unsharded FC weights, per-head KV quantization, Mamba/GDN state slots).
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,nvcc13.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_size512, 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.
- 40 layers = 30
- Important caveat:
sm_86andsm_89have 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
-
Install FreeToken 0.1.3 into a venv; confirm
ft serveruns TP=1 on the fast card. -
Apply the patch below to the installed
freetoken/package tree (python -m py_compileeach file). -
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
-
If rank free-memory differs by more than the built-in 2 GiB gate, relax
engine._sync_get_memorylocally (we used 16 GiB). -
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=48GBin.wslconfig. -
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
mainand 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.
MarlinandB12xremaintp_ok=False(and are not usable onsm_86/sm_89anyway). - 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:
- Issue 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 (
No-LetterHead4141) — the actual working TP=2 run on ROCm/FP8 and the methodological blueprint we ported: segment-aware loader slicing, GDN head splitting, per-rank expert banks. - Issue Feature request / guidance: tensor-parallel across heterogeneous and idle consumer GPUs (quantized MoE experts) - a hobbyist request, with reference to #478 #510 and
kenwheeler77— the hybrid-TP /TVM_FFI_CUDA_ARCH_LISTinvestigation. The root-cause analysis (single vs. multi-value arch list, the flashinferrmsnorm.pyenable_pdl/get_device_capability()interaction, and the CUTLASS-DSL crash) directly shaped our arch-list choice and saved a lot of blind debugging. - PR [1/3] feat(dsv4): add tensor-parallel DeepSeek-V4 runtime #70 (
calvarado2004) — the tensor-parallel ownership rules that inspired this patch's structure, in particular thenarrow()view-retention finding (+5.1 GiB/rank) and its.clone()fix. - Issue Support for force selecting a specific GPU on multi-GPU systems #99 / merged PR feat(server)!: add --gpu to choose the GPU on multi-GPU machines #117 — the
--gpudevice-selection flag, without which mixed-card selection is painful. - The dual-3080 SGLang patch (
kk-pcl/sglang-qwen38-dual-3080-patch) — prior art for the same model family; useful as a sanity check on expected failure modes. - The FreeToken maintainers and the community write-ups on the surrounding walls (
#44,#55,#72,#75,#335,#425,#442,#496) — those reports are what let us anticipate the prefill-hang, WDDM memory, and context-cap issues instead of rediscovering them.
Any mistakes in the above are ours.
10. Availability / next steps
The complete patch is embedded above. We're happy to:
- open a PR against a branch if maintainers prefer that form,
- split the change to match PR [1/3] feat(dsv4): add tensor-parallel DeepSeek-V4 runtime #70's structure once it lands,
- and share the benchmark scripts on request.
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