Offloaded MoE ignores tensor parallelism: TP=2 is a regression, and the fix is not expert parallelism

[gdevenyi]

Summary

--tensor-parallel-size 2 is currently a regression for offloaded MoE: measured 15.63 tok/s against 18.00 at TP=1 on DeepSeek-V4-Flash (2x RTX 6000 Ada, no NVLink). Two independent single-GPU servers give 2.05x instead, so the hardware scales fine — the engine does not.

I went looking for expert parallelism and concluded that EP is the wrong fix, and the right one is already half-built.

Why TP=2 loses today

The resident MoE path shards by TP:

intermediate_size_per_partition = div_even(intermediate_size, tp_size)
if allocate_experts:
    self._alloc_resident_experts(intermediate_size_per_partition)

OffloadMoELayer passes allocate_experts=False and takes its weights from OffloadMoeCache, which has no TP awareness at all — no tp_size/tp_rank in offload_cache.py, expert_banks.py or models/deepseek_v4/moe.py. Every rank therefore loads and caches whole, unsharded experts.

Confirmed at runtime: moe_cache_size=2477 per rank at TP=2 versus 2559 at TP=1. The cache did not grow. Both ranks stream the same experts over their own PCIe link, so aggregate PCIe traffic doubles for identical work, and then the layer pays an all-reduce over PCIe. Only attention and dense shard usefully, and this model is MoE-dominated.

Why not expert parallelism proper

Sharding experts across ranks (rank r owns e % world == r) needs all-to-all dispatch and combine per MoE layer. Three problems, in order of severity:

1. CUDA-graph capture. Decode graphs are captured at static shapes. All-to-all counts depend on routing, so they vary per step. Workable by padding to the worst case (bs * top_k tokens per rank), but that is real added machinery in the hot path.
2. Load imbalance. Real MoE routing is skewed; a rank owning hot experts stalls the others every layer.
3. New collectives in a codebase that currently only needs all_reduce.

The fix that fits: TP-shard the offload banks

Split the intermediate dimension, exactly as the resident path already does. Rank r holds I/tp of every expert.

• Cache doubles. A slot is half the bytes, so the same VRAM holds ~2x as many distinct experts. Residency here goes from 2559/11008 = 23% to roughly 46%.
• PCIe halves per rank. Each rank fetches only its half of a missed expert, over its own link — two links, half the bytes each.
• Host RAM is unchanged. 71.5 GiB per rank instead of 143 GiB replicated, so TP=2 costs no more host memory than TP=1 (today it costs double).
• No new collectives. Intermediate-dim sharding makes the down projection produce partial sums, and _maybe_all_reduce already combines them. The communication is bs * H * 2 bytes per layer — 8 KiB at bs=1, negligible.
• CUDA graphs unaffected. Shapes stay static.

Extrapolating the measured residency curve (25% -> 50% residency was worth +59% on a comparable model), plus the halved per-rank gather, this looks like 1.4-1.8x on this box — versus 0.87x today.

What it touches

The work is almost entirely in loading, not in the kernels — they already accept whatever I they are handed.

1. Per-format bank slicing at load (moe/expert_banks.py). The awkward part: gate_up is [2I, H] with gate at rows [0, I) and up at [I, 2I), so a rank's slice is two disjoint row ranges, not one. down is [H, I], a column slice.
2. Divisibility. I/tp must stay a multiple of each format's block size — 128 for fp8_block, 32 for the fp4 scale blocks, 32 for q4_0. DSV4 at I=2048, tp=2 gives 1024, which is fine, but this needs a guard rather than an assumption.
3. mxfp4 is transposed (blocks_t [K//2, N]), so its slice is along N and the scale tensor slices differently again. This format has already broken two general assumptions in recent work; it should be last, or explicitly unsupported at first.
4. Thread I_local through OffloadMoeCache construction in engine.py and the _BANK_SCHEMAS shape asserts.
5. The CPU MoE / hybrid path takes I from the cache, so it follows for free — but --moe-backend hybrid under TP needs its own check.

Suggested staging

1. bf16 only, TP=2, behind a guard that refuses unsupported formats. Prove the cache doubles and measure end to end.
2. Add ds_fp4 and nvfp4 (row-major, straightforward slices).
3. fp8_block and q4_0.
4. mxfp4 last, or leave it TP=1-only.

Measurement notes

• ft bench decode cannot test this: LLM.__init__ hardcodes tp_info=DistributedInfo(0, 1). Either that gets a parameter or the comparison has to go through ft serve.
• A server binds PORT and PORT+1 (a multiprocessing.spawn child), so multi-instance comparisons need spaced ports.
• FreeToken JIT-links its own pynccl against a system NCCL; without libnccl-dev installed, TP fails at ld: cannot find -lnccl even though torch bundles libnccl.so.2. Falling back to torch's copy would remove a sharp edge.

[lukascechovic]

We are building the intermediate-dimension shard you describe here, on a different model
(qwen4_exp / Qwen3.8-Flash-Next) and on two AMD R9700s. Adding what we have so far, because your
central claim — that the cache grows and no new collectives are needed — is holding up in
implementation, and because two things bit us that are not obvious from reading the tree.

Your "the cache doubles" argument checks out, and by more than double

Sharding the intermediate dimension per rank, both ranks build 206/206 weights with no shape
mismatch, and the expert cache resolves to:

	TP=1	TP=2 (per rank)
MoE cache slots	5,960	16,306
KV pages / tokens	1720 / 110,080	1720 / 110,080
K + V	2.60 GiB	1.34 GiB per rank
free VRAM after init	3.18 GiB	3.17 GiB

⭐ 2.74x the slots at the same VRAM. Your extrapolation from the residency curve assumed roughly
2x; on this model the slot is small enough relative to the page geometry that it does better.

✅ And the parts you predicted would be free were free. layers/moe.py:94-97 already all-reduces at
tp_size > 1, the fused kernels are dense-in/dense-out, so no all-to-all exists or is needed.
norm_topk_prob defaults True and the router is replicated, so a plain SUM all-reduce is
algebraically exact — not an approximation we are tolerating.

What the reduction costs, measured

We priced the collective before building anything, because it is the dominant risk in the whole case.
On two R9700s at the real decode shape ([1, 2560] bf16 = 5,120 bytes, 48 layers):

• 28.6–41.9 µs per call ⇒ 1.34–2.01 ms per token
• RCCL 2.30.4 chose via P2P/IPC across the root complex, no host staging, despite no XGMI on this
  board

⚠️ Two caveats on how to read that. It is a standalone measurement of the transport — our model
load has not yet reached its own collective, so this is what the reduction costs, not yet what the
design costs. And it is one collective per layer; a design that also shards the attention mixers
pays two, which is where our own numbers now sit.

⚠️ We also found prefill is compute-bound at full chunks on this hardware, not PCIe-bound. So the
second card's prefill win is much smaller than the residency argument suggests — around 8 % at 100k
context on our box. The case for TP here is a decode case, and it should be argued as one.

Two things that will bite anyone trying this on AMD

1. ⛔ Two ranks cannot share one copy of anything pinned, and it is a driver fact, not an allocator
one. svm_range_new() reserves size << PAGE_SHIFT per process against a single driver-global
counter, with no physical-page and no BO identity anywhere in the accounting
(kfd_svm.c:335-352, amdgpu_amdkfd_gpuvm.c:169-268). Two processes registering the same physical
pages are charged twice. The dma-buf escape is closed — "Only VRAM and GTT BOs are supported".

⇒ on ROCm the design must shard, never share. Which is an argument for your axis: a sharded
bank is genuinely half the bytes per process, where a shared one is billed as if it were not. We
established this from the driver source and have not measured it; the ~60-line probe that would
is written but unspent.

2. ⚠️ Ranks are spawned, not forked. server/launch.py:158 sets
mp.set_start_method("spawn", force=True), and the banks are mmap.mmap(-1, n) —
MAP_SHARED|MAP_ANONYMOUS, so shareable by fork but unreachable by name. Any design that hoped to
attach one copy from a sibling process needs those mappings named first. Sharding sidesteps this
entirely, which is another point in its favour.

Where our implementation actually is

Both ranks load and initialise. We are working through host RAM: the expert banks total 111.00 GiB
whether we run TP=1 or TP=2 (55.50 GiB per rank), against 122 GiB on this machine, so we are
tight in a way that is a property of our box rather than of your design. On a host with more RAM the
same branch has much more room.

⛔ Not yet cleared, and we are not claiming otherwise: the model's own collective at tp_size = 2,
CUDA-graph capture under the sharded shape, whether the QSA sparse-attention backend accepts a
per-rank head subset (its indexer is LinearReplicated and emits all heads), and the KFD pin ceiling.

⚠️ One thing worth flagging for qwen4_exp specifically, since it is not DSV4: the TP=1 only gate
in its weight loader is load-bearing. qkv_proj and the GDN in_proj are
LinearColParallelMerged and shard whatever you choose, so the offload MoE cannot be sharded in
isolation on this model — mixers come with it. That may be particular to this architecture, but it is
worth knowing before someone scopes "shard the MoE only" as a small change. #29 is the tracking issue
for that gate.

The question

Is the intermediate-dimension shard the axis you would accept for the offload MoE path?

The reduction is already in the tree and the fused kernels are dense-in/dense-out, so the change looks
confined to the loader and to rank-awareness in moe/, which currently has none. We would rather find
that out before polishing a PR against an axis you do not want.

⚠️ Two references for anyone weighing this: a third party measured PR #104's replicated-bank offload
TP=2 path at 3.1 tok/s against 94 at TP=1, which is the strongest argument that replicated banks
are the wrong axis; and #324 on this tracker runs the same model on two RTX 5090s as two separate
engines, because one model across both is not available. The demand is not AMD-specific.

Platform

AMD Radeon AI PRO R9700 x2 (gfx1201), Linux, 122 GiB host RAM · PyTorch 2.11.0+rocm7.14.0, HIP
7.14.60850, RCCL 2.30.4 · FreeToken main@4b94bdc3 + #132 + #133 + #134 + our own sharding work ·
BLOCK_N clamped to 32 in the QSA attend kernel (filed separately), which does not touch any number
above.

[lukascechovic]

Following up on my comment above with the verdict, because the four things I listed as "not yet
cleared" are now cleared: the intermediate-dimension shard serves, and on this hardware it is worth
about +34 % decode and −40 % TTFT.

It serves

Two AMD R9700s (gfx1201), qwen4_exp / Qwen3.8-Flash-Next NVFP4, offload MoE, --tp-size 2.
Against the same box's two single-card controls, full-box, steady state:

	TP=2 (this work)	vs the same-backend TP=1 disk row	vs the TP=1 row we actually deploy
decode 10k / 100k (tok/s)	35.01 ± 0.08 / 34.63 ± 0.02	18.98 / 19.43	26.10 / 26.35
TTFT 10k / 100k (s)	5.68 / 53.1	9.59 / 88.8	9.39 / 89.0

Both ranks build 206/206 weights with no shape mismatch; CUDA-graph capture completes in 18–20 s;
the QSA sparse-attention backend does accept a per-rank head subset. At 110k the cache resolves
to 16,306 slots / 1,720 pages per rank — the 2.74x-at-equal-VRAM your "the cache doubles"
argument predicted, and then some.

⭐ Context is where it pays most. At --kv-reserve-tokens 262144 we get 14,876 slots / 4,096
pages, KV 3.19 GiB per rank, and decode costs only −0.9 % / −1.7 % against the 110k steady
state. The same jump on our TP=1 row costs −7.4 %. ⇒ the second card does not just raise the
throughput, it makes long context roughly free.

PR #311 is what unblocked this, and it is worth knowing why

We were stuck. At TP=2 both ranks loaded and then died in CUDA-graph capture at 0.38 GiB
MemAvailable — the pinned expert banks total the same 111.42 GiB at TP=1 and TP=2, and a
second rank's overhead was all it took on a 122 GiB host.

#311 (stream the PLE n-gram table from disk, merged 2026-09-01) removed the one bank nobody was
touching. --ple-backend disk pins 63.74 GiB instead of 111.42 and leaves 51.4 GiB of
MemAvailable at ready instead of 5.3 — and returned byte-identical greedy text on our TP=1 control.
That turned a 2.1 GiB deficit into ~45 GiB of margin, and the two-card path went from "cannot capture"
to "serves" with no change of our own to the memory design.

⇒ if anyone reads #62 and concludes the offload-MoE TP path needs a host with more RAM: on this model
it needed #311, which is already merged.

Vision replicates cleanly at TP=2, and it is charged somewhere non-obvious

We since put the vision tower and TP=2 in one process. The ViT is LinearReplicated throughout — no
name in it matches the _TP_SHARDED suffix table — so the 333-tensor, 897,862,112 B BF16 tower
replicates whole on both ranks, and it answers correctly from the image on a two-rank load.

⚠ The tower is paid for in MoE cache slots, not out of free VRAM. It enters weights_bytes, which
is subtracted from the same budget moe_cache_size is computed from: the identical 262k config goes
14,876 → 14,221 slots per rank, and 655 slots x ~1.33 MiB = 0.848 GiB against the 0.836 GiB
tower. Free-VRAM-after-init is bit-identical either way. Worth knowing before someone sizes a
vision + TP=2 config by looking at free VRAM and concludes the tower is free.

The three we did not close

1. ⚠ Greedy text at TP=2 differs from TP=1, and I cannot tell you whether that is numerics or a
bug. All three of our prompts diverge (first divergence at char 9 / 123 / 207), byte-identical
across three TP=2 loads, so it is systematic rather than a race. A SUM all-reduce over a replicated
router is algebraically exact, so reduction order is the benign explanation — but I have not
demonstrated that, and I am not claiming the two configurations serve the same model bit-for-bit.

⛔ I cannot discriminate it, because the API refuses logprobs ("logprobs is not supported") —
which is exactly what PR #224 implements. ⭐ If #224 lands, this becomes a one-afternoon
question: dump raw log_softmax at the first divergent position on both configurations and see
whether the top-2 gap is inside bf16 reduction noise (numerics) or wide (a defect on the sharded
path). That is the one thing I would ask for from this thread.

⚠ Caveat I should state rather than have found for me: our image predates #329, so these runs used
the approximate top-k/top-p kernel. We pinned top_k:1, top_p:1, so this is argmax and #329's
boundary-tie clipping is an unlikely cause across three prompts — but it is not excluded, and a
re-run on current main is the clean test.

2. ⚠ Roughly one TP=2 load in three comes up ~35 % slower, and nothing in the log distinguishes
them. 2 of 7 cold loads decoded at 22.9 tok/s instead of ~35, on a byte-identical command,
with identical moe_cache_size, identical relay handshake, no swap activity, and a fresh Triton
cache in one of the fast cases (so it is not JIT warmth). A row in that state serves **12 % under
the single-card row it replaces while looking perfectly healthy. Mechanism unidentified — flagging it
in case the shape is familiar to you, and because it means no one should bank a single TP=2 load as
a benchmark, including us.

3. ⚠ Cold load is ~18–20 minutes at TP=2 (1077–1190 s to ready), against ~3 minutes at TP=1. Both
ranks read all 206 shards, and the read is CPU-bound and strided rather than I/O-bound. Fine for a
long-lived server, painful for anything that swaps models.

⚠ Separately, at tp_size > 1 the first request after start is lost in the rank-0→rank-N ZeroMQ
PUB/SUB relay — a classic slow-joiner race — which wedges the server rather than erroring. Filed as
#364 with a handshake fix; not repeated here.

Platform

AMD Radeon AI PRO R9700 x2 (gfx1201), Linux, 122 GiB host RAM · PyTorch 2.11.0+rocm7.14.0, HIP
7.14.60850, RCCL 2.30.4 · FreeToken main@4b94bdc3 + #132/#133/#134 + #311 + our sharding work ·
BLOCK_N clamped to 32 in the QSA attend kernel (filed as #349), which does not touch any number
above · half-box controls where stated, otherwise full-box. No quality claim — we have no
fidelity instrument on this machine, which is exactly why item 1 above is open.

[lukascechovic]

Follow-up on the "18-20 min cold load" I listed above as unresolved: that is closed. It was not the expert read being slow, it was OpenMP barrier spin-wait — both ranks default to one intra-op thread per logical CPU, so TP=2 on a 6-core box runs 24 threads on 6 cores. Clamping per rank takes the load from 1175 s to 61 s.

Filed separately as #371, since it is a load-path bug rather than anything to do with this thread's design question.