feat(moe): owner-local expert parallelism and tensor parallelism by leiyu1980 · Pull Request #447 · FlashML-org/FreeToken

leiyu1980 · 2026-09-11T10:15:08Z

Adds tensor parallelism for Qwen3.8-Flash-Next (qwen4_exp) on the offload MoE
backend, together with owner-local expert parallelism: the routed-expert group is
partitioned across the TP ranks, so each rank owns a contiguous slice of the
experts and keeps every expert whole.

Addresses #62 (offloaded MoE ignores tensor parallelism).

Related but not fixed here: #29 is the same gap for qwen3_5_moe, whose loader
still rejects TP>1 — only qwen4_exp grows a tp_shard reader path in this PR
(models/weight.py forwards tp_shard only to readers that declare it and fails
fast otherwise). The qwen3_5_moe changes here are limited to the owner-EP
shared+routed fusion and an explicit refusal on its block-FP8 expert banks.

Why owner-local experts

main sets tp_ok=False for the MoE kernels in layers/quantization/moe/
(nvfp4.py, fp8_block.py, mxfp4.py), so running one expert across ranks means
changing that contract and sizing each rank's bank from a local intermediate —
which is what #385 does, and it works (see below). This PR takes the other axis
instead: partition the experts, not each expert. Every expert GEMM then stays
at full intermediate and unsharded, and a MoE layer needs one all-reduce over the
ranks' partial outputs rather than an activation dispatch. No expert is ever
dispatched over the network.

We are not claiming this is the only workable design. The two are close in
capacity — a slot holds a whole expert on one design and half of every expert on
the other, so both roughly double what fits — and the honest trade-off is:

In favour of owner-local: no kernel contract change, and the routed output
is a full sum per expert rather than a split down-projection.
Against it: routing is skewed, so a rank that owns hot experts stalls the
others every layer. We have not measured that imbalance, and it is the main
open question about this design. feat(qwen4_exp): tensor parallelism for Qwen3.8-Flash-Next (offload backend) #385's intermediate-axis split is perfectly
balanced by construction.

What is in it

Owner-local expert parallelism

moe/ownership.py (new) — ExpertOwnership / OwnedRoute /
OwnerCacheGeometry / OwnerCacheAdapter. Three namespaces are kept strictly
apart: global expert id -> local bank row -> local flat id -> cache slot. A
route entry whose expert is remote reuses a row the same route already owns,
with weight 0, so nothing is ever dispatched and no -1 reaches a kernel.
moe/offload_cache.py — OwnerOffloadMoeCache wraps the global-ID cache;
ensure_route_graph admits a route at a fixed shape with zero host
synchronisation, which is what keeps owner-EP decode CUDA-graph capturable. The
eager ensure_route stays for diagnosis.
models/nvfp4_banks.py, moe/expert_banks.py, moe/expert_pieces.py — the
expert stream is filtered to this rank's experts and renumbered into the local
bank rows [0, local_num_experts), so the pieces land where the cache actually
allocates; the bank's expert dimension becomes the local count rather than the
layer's global routing count. A disagreement between ownership.global_num_experts
and the checkpoint's num_experts is rejected, and a reader that cannot serve an
owner-local bank raises NotImplementedError — probed with inspect.signature, so
unrelated readers are untouched.
moe/__init__.py — exports ExpertOwnership / OwnerCacheGeometry /
OwnerCacheUpdate.
layers/moe.py, layers/quantization/moe/base.py — expert_tp_size is plumbed
through MoEConfig.from_layer, so the expert GEMM dimensions come from the owner
geometry rather than from the tensor-parallel group.
layers/linear.py — the row-parallel linears take reduce=False, so the shared
and routed partial sums can be combined into one all-reduce per MoE layer
instead of one each.
models/qwen4_exp/moe.py, models/qwen3_5_moe/moe.py — that shared+routed
fusion, guarded by NotImplementedError when the shared down projection is not
row-parallel.
models/deepseek_v4/moe.py — the owner adapter must see the raw global route,
so the global-cache-only short-prefill shortcut is bypassed when it is active.
models/weight.py — tp_shard / tp_config are forwarded only when the target
reader declares them (inspect.signature), and TP>1 against a reader that does
not fails fast. models/qwen3_5_moe/weight.py refuses owner EP on block-FP8
expert banks explicitly.
moe/route_trace.py (new) — ordered route trace capture for offline LRU/EP
replay.
engine/engine.py — the owner-EP wiring: _owner_ep_enabled /
_validate_owner_ep_config fail before any model or bank allocation unless the
initial topology is an explicit same-group TP2+EP2 on the offload backend with an
explicit or auto cache size; _owner_graph_safe picks the graph-safe route
admission; _resolve_auto_moe_cache_size now solves --moe-cache-auto against
the owner-local expert geometry, so auto sizing works under owner EP instead
of demanding a hand-tuned slot count.
engine/config.py, models/config.py, server/args.py — --moe-ep-size,
--moe-collect-decode-freq, --moe-trace-route.
benchmarks/bench_offload_cache_copy.py — a qwen3.8-flash-next model profile
(48 MoE layers, 512 experts, top-10, H=2560, moe_inter=640 → 2,772,480 B/expert,
matching the served cache's unit_bytes).

Dense tensor parallelism

models/qwen4_exp/weight.py — shard_qwen4_exp_dense_tensor slices the RAW
checkpoint tensors before fusion, so the fused qkv/o_proj buffers keep their
head boundaries. LinearColParallelMerged is handed GLOBAL output sizes because
it shards each output segment itself.

Two companion changes (not part of TP/EP itself)

Both are here because they were needed to validate and run this work. Say the word
and I will split either one out into its own PR.

Context ceiling. kvcache/cache_status.py, server/api_server.py,
server/stats.py, server/openai_api.py, api_models.py report the
effective limit, min(model max_position, KV pool tokens), instead of the
checkpoint limit — so a client sizes its window from what the server will
actually accept, rather than getting a hard 400 on prompts the model card said
were fine.
MoE observability. server/stats.py, scheduler/scheduler.py,
message/frontend.py, message/tokenizer.py, tokenizer/server.py,
control_cli.py — a slot-cache snapshot (residency, miss rate, routing
concentration) surfaced through /v1/stats and ft stats. The device syncs are
throttled to ~1/s and a failed snapshot never breaks the reply stream.

Tested on


GPU	2 × RTX 4090 24 GiB, sm_89, no direct P2P (`cudaDeviceCanAccessPeer=0`); NCCL runs over host shared-memory/PCIe (`NCCL_P2P_DISABLE=1`)
CPU / RAM	2 × Xeon Platinum 8375C @ 2.90 GHz / 503 GiB
Driver	580.142
CUDA / PyTorch / Triton	13.0 / 2.11.0+cu130 / 3.6.0
Python	3.12
Checkpoint	`RadixArk/Qwen3.8-Flash-Next-NVFP4` (base `Qwen/Qwen3.8-Flash-Next`, converted with NVIDIA modelopt 0.46.0)

Exact command

ft serve \
  --model-path /path/to/Qwen3.8-Flash-Next-NVFP4 \
  --gpu 4,5 --host 127.0.0.1 --port 1930 \
  --tensor-parallel-size 2 --moe-ep-size 2 \
  --moe-strategy offload --ple-backend disk \
  --num-tokens 262144 --moe-cache-auto --kv-reserve-tokens 262144 \
  --cuda-graph-max-bs 2 --max-running-requests 2

Results

1. Weight loading — exact

The model state dict and the sharded reader agree key by key on shape and
dtype: 722 exact + 72 widened (the GDN A_log/dt_bias pair, declared fp32 on
purpose), 0 shape mismatches, 0 dtype mismatches, 0 keys never loaded — at both
TP1 and TP2+EP2. This check is what caught the qkv_proj double-sharding bug
above; neither a syntax check nor an import check would have.

2. CPU suite

PYTHONNOUSERSITE=1 PYTHONPATH=<worktree>/python CUDA_VISIBLE_DEVICES="" \
  pytest tests/ -q -m "not slow" -p no:cacheprovider --ignore=tests/e2e
# 1319 passed, 2 failed, 425 skipped

Both failures are tests/models/test_quant_config.py probing local HF
checkpoints that are absent here, and reproduce on a clean main.

3. End-to-end A/B — main vs this branch

Same machine, same card pair, same model, same prompt, same sampling, and
identical flags except --tensor-parallel-size / --moe-ep-size:

python tools/bench/ab_main_vs_branch.py --max-tokens 512
# 511 completion tokens, 262144 context, --moe-cache-auto, graphs [1,2], bs<=2

	code	config	steady-state Decode	max	wall (511 tok)
A	`main`	TP1	21.8 tok/s	23.9	27.4 s
B	this branch	TP1	21.7 tok/s	24.3	26.8 s
C	this branch	TP2 + EP2	42.8 tok/s	50.2	15.0 s

B vs A = −0.5% → the existing single-card path is unchanged.
C vs A = 1.96× at the same 262144-token context, wall clock cut by ~45%.
Decode is the /v1/stats 5-second sliding-window steady-state rate, i.e. the
same number stats.sh prints. It deliberately excludes prefill and TTFT, so it
is not comparable with completion_tokens / wall_clock.
TTFT was not measured, and prefill is affected (the dense projections are
sharded), so the A/B above does not cover it. --moe-cache-auto also resolves
to a different slot count per config, so the cache is not held constant across
the three legs.
Absolute values depend on how busy the shared PCIe/host path is (another user's
training was saturating GPU1/2 during part of this work); only the ratios
measured in one session are meaningful, which is why all three legs were run
back to back on the same pair.

4. Correctness — and an honest caveat

TP1 and TP2+EP2 produce the same answers on the gate prompts, but their token
streams are not byte-identical: 2 of 5 prompts differ by a single word.

We localized this down to individual logits rather than leaving it as a
hand-wave. The experiment dumps, for every decode step, the sampled token and
the top-8 (value, index) pairs, then compares the numbers instead of the text.

Controls first:

control	result
same config, two identical runs	bit-identical (both configs, both prompts)
within TP2, rank0 vs rank1	bit-identical — every rank computes the full logits after the all-reduce and argmaxes independently, so the ranks agreeing rules out a sharding/desync defect

So each configuration is deterministic and TP2 is internally consistent; the
divergence is a systematic difference between the two sharding degrees, not
run-to-run noise.

Then the actual flips. Both land on steps where the model is essentially
indifferent:

	prompt 0	prompt 1
first differing token at step	11	41
TP1 top-1 / top-2 gap there	0.2500	0.0000 (an exact tie)
TP2 top-1 / top-2 gap there	0.6250	0.3750
cross-config delta on the two candidates	+0.375 / −0.500	+0.375 / 0.000

At prompt 1 step 41 the two candidates are exactly equal in TP1
(23.8750 vs 23.8750) and argmax tie-breaks to the lower id; TP2 has the same
two tokens 0.375 apart and picks the other one. The flips occur exactly where
the tie gap (0.25, 0.00) is at or below the ordinary cross-config difference
(0.375) — which is what a tie-break looks like, and the opposite of a computation
that is quietly wrong.

Conclusion: this is tie-breaking under reduction-order differences, not a
defect. The residual ~0.1–0.4 logit difference comes from summing a TP-sharded
reduction instead of one matmul, plus the MoE experts being summed across ranks.
Answers stay correct (TP2_OK, 4, Red). Byte exactness across TP degrees is
not achievable and we are not claiming it.

Relationship to other open PRs

feat(qwen4_exp): tensor parallelism for Qwen3.8-Flash-Next (offload backend) #385 (qwen4_exp tensor parallelism, offload backend) covers the same axis
by sharding each expert along its intermediate dimension, and reports better
numbers on a 48 GiB box where the whole model fits. It is the other half of
this design space, and the trade-off between the two is described above. Happy
to coordinate on which one lands, or on landing them together.
feat(kvcache): store the KV cache as fp8 e4m3 codes (--kv-cache-dtype… #354 (fp8 KV cache) is not included here. It is independent and composes
with this work; keeping it out makes this diff self-contained and keeps every
line of it something we have run ourselves.
The --moe-strategy naming and the frozen-dataclass / quant_config+prefix
style follow upstream conventions.

Known limitations

--moe-ep-size > 1 currently requires TP2+EP2 in the same group, the offload
backend, native NVFP4 experts, and either an explicit cache size or
--moe-cache-auto. _validate_owner_ep_config fails fast otherwise.
--moe-cache-rate is rejected under owner EP: it takes a ratio of the global
expert count, which is meaningless once experts are partitioned.
Headroom is tight on 24 GiB: owner-EP initialisation wants ≥ ~2 GiB free per
card. With 262144 KV tokens, --moe-cache-auto resolves to ~4466 slots/card on
this hardware; push it higher and activation allocations OOM inside the request.
Expert load imbalance under skewed routing is not measured (see above).

Adds tensor parallelism for Qwen3.8-Flash-Next (qwen4_exp) on the offload MoE backend, together with owner-local expert parallelism: the routed-expert group is partitioned across the TP ranks, so each rank owns a contiguous slice of the experts and keeps every expert whole. Why owner-local: every NVFP4 MoE kernel in the tree reports tp_ok=False, so an expert cannot be split along its own dimensions. Partitioning the *experts* instead of each expert's rows leaves each expert GEMM unsharded and needs one all-reduce per MoE layer rather than an activation dispatch. Highlights: * moe/ownership.py -- ExpertOwnership / OwnedRoute / OwnerCacheGeometry / OwnerCacheAdapter. Three namespaces are kept strictly apart: global expert id -> local bank row -> local flat id -> cache slot. A route entry whose expert is remote reuses a row the same route already owns, with weight 0, so nothing is ever dispatched and no -1 reaches a kernel. * moe/offload_cache.py -- OwnerOffloadMoeCache wraps the global-ID cache; ensure_route_graph admits a route at a fixed shape with zero host synchronisation, which is what makes owner-EP decode CUDA-graph capturable. The eager ensure_route stays for diagnosis. * layers/moe.py, layers/quantization/moe/base.py -- expert_tp_size is plumbed through MoEConfig.from_layer so the expert GEMM dimensions come from the owner geometry rather than from the tensor-parallel group. * models/qwen4_exp/weight.py -- shard_qwen4_exp_dense_tensor slices the RAW checkpoint tensors before fusion, so the fused qkv/o_proj buffers keep their head boundaries. LinearColParallelMerged is handed GLOBAL output sizes because it shards each output segment itself. * models/nvfp4_banks.py, moe/expert_banks.py, moe/expert_pieces.py -- owner filtering plus local row renumbering when building the expert banks. * engine/config.py, server/args.py -- --moe-ep-size, --moe-collect-decode-freq, --moe-trace-route. * moe/route_trace.py -- ordered route trace capture for offline LRU/EP replay. * server/stats.py, server/openai_api.py, api_models.py -- report the effective context limit (min(model max_position, KV pool tokens)) instead of the checkpoint limit, so a client sizes its window from what the server will actually accept. Validation on 2x RTX 4090 (sm_89), torch 2.11.0+cu130, driver 580.142: * The model state dict and the sharded reader agree on 794/794 keys (shape and dtype) at both TP1 and TP2+EP2, with zero keys never loaded. * CPU suite: 1319 passed, 2 failed -- both failures reproduce on a clean main (tests/models/test_quant_config.py probes local HF checkpoints that are absent here). * Same-card A/B at 262144 context, identical flags except --tensor-parallel-size and --moe-ep-size: TP1 is unchanged from main (-0.5%, within noise) and TP2+EP2 reaches 1.96x the single-card steady-state decode rate. * TP1 and TP2+EP2 produce the same answers, but their token streams are NOT byte-identical: two of the five gate prompts differ by a single word. Both configurations are individually deterministic (two independent trees produce the same text), so this is floating-point reduction order -- a TP-sharded reduction and a near-tie router/sampling decision -- not a routing defect.

Copilot

🟡 Changes recommended

Critical TP/EP loader, FTW-bank, cache, vocabulary-shard, trace-path, and stats issues remain unresolved.

Once you've addressed the issues Copilot identified, you can request another Copilot review.

Pull request overview

Adds owner-local expert and tensor parallelism for the Qwen4Exp NVFP4 MoE backend, plus route tracing and effective context-limit reporting.

Changes:

Adds owner-local expert ownership, routing, cache geometry, and bank filtering.
Adds Qwen4Exp TP weight sharding and runtime configuration.
Adds route tracing, server limit reporting, and regression coverage.

File summaries

File	Reviewed changes / final review note
`tests/server/test_stats_limits.py`	Tests effective cache/context-limit reporting.
`tests/server/test_openai_api.py`	Tests effective model-card context limits.
`tests/scheduler/test_abort_inflight_prefill.py`	Updates scheduler abort/prefill fixtures.
`tests/moe/test_route_trace.py`	Critical (3 votes): absolute workstation path prevents portable/CI collection.
`tests/moe/test_ownership.py`	Tests ownership and cache namespace contracts.
`tests/moe/test_offload.py`	Tests owner-cache CUDA and prefill paths.
`tests/models/qwen4_exp/test_weight.py`	Critical (1 vote): short final vocabulary shards can fail strict TP shape loading.
`tests/models/qwen4_exp/test_skeleton.py`	Covers padded TP vocabulary shards.
`tests/models/qwen4_exp/test_qsa_backend.py`	Reviewed; no final comment.
`tests/models/qwen4_exp/test_gdn.py`	Reviewed; no final comment.
`tests/models/qwen4_exp/test_config.py`	Reviewed; no final comment.
`tests/kvcache/test_linear_state_pool_alloc.py`	Reviewed; no final comment.
`python/freetoken/tokenizer/server.py`	Reviewed; no final comment.
`python/freetoken/server/stats.py`	Moderate (1 vote): `/v1/stats` can expose the raw model ceiling instead of the effective limit.
`python/freetoken/server/openai_api.py`	Publishes effective model context limits.
`python/freetoken/server/args.py`	Adds MoE EP, decode-frequency, and route-trace options.
`python/freetoken/server/api_server.py`	Publishes effective runtime limits.
`python/freetoken/server/api_models.py`	Adds model context metadata.
`python/freetoken/scheduler/scheduler.py`	Reviewed; no final comment.
`python/freetoken/moe/route_trace.py`	Moderate (1 vote): shared TP trace paths can race and corrupt trace output.
`python/freetoken/moe/ownership.py`	Defines ownership and global/local/cache namespace mappings.
`python/freetoken/moe/offload_cache.py`	Critical (1 vote): owner CPU/hybrid settings are ignored; Critical (1 vote): stale pending state can replay prior routes; Moderate (1 vote): rebuild can desynchronize prefill-overlap geometry.
`python/freetoken/moe/expert_pieces.py`	Filters and renumbers owner-local expert pieces.
`python/freetoken/moe/expert_banks.py`	Critical (1 vote): FTW loading omits ownership and returns global expert rows.
`python/freetoken/moe/__init__.py`	Exports ownership helpers.
`python/freetoken/models/weight.py`	Adds TP-shard loader plumbing.
`python/freetoken/models/qwen4_exp/weight.py`	Critical (1 vote): short final vocabulary shards can fail strict TP shape loading.
`python/freetoken/models/qwen4_exp/moe.py`	Integrates owner-local Qwen4Exp expert execution.
`python/freetoken/models/qwen4_exp/gdn.py`	Reviewed; no final comment.
`python/freetoken/models/qwen4_exp/config.py`	Carries Qwen4Exp configuration.
`python/freetoken/models/qwen4_exp/attention.py`	Reviewed; no final comment.
`python/freetoken/models/qwen3_5_moe/weight.py`	Reviewed; no final comment.
`python/freetoken/models/qwen3_5_moe/moe.py`	Supports owner-local shared/routed expert fusion.
`python/freetoken/models/nvfp4_banks.py`	Builds local NVFP4 expert banks.
`python/freetoken/models/deepseek_v4/moe.py`	Adjusts owner-local MoE behavior.
`python/freetoken/models/config.py`	Carries EP geometry into model configuration.
`python/freetoken/message/tokenizer.py`	Reviewed; no final comment.
`python/freetoken/message/frontend.py`	Reviewed; no final comment.
`python/freetoken/layers/quantization/moe/base.py`	Plumbs expert TP sizing.
`python/freetoken/layers/moe.py`	Routes owner-local experts and reductions.
`python/freetoken/layers/linear.py`	Reviewed; no final comment.
`python/freetoken/kvcache/cache_status.py`	Reviewed; no final comment.
`python/freetoken/kernel/pynccl.py`	Reviewed; no final comment.
`python/freetoken/engine/engine.py`	Critical (3 votes): unconditional `tp_shard` breaks loaders without that parameter. Nit (1 vote): `decode_freq` warning/docs are stale. Moderate (1 vote): owner EP ignores `--moe-cpu-layers`. Critical (1 vote): FTW loading omits ownership.
`python/freetoken/engine/config.py`	Adds MoE runtime configuration.
`python/freetoken/control_cli.py`	Reviewed; no final comment.
`benchmarks/bench_offload_cache_copy.py`	Reviewed; no final comment.

Review details

Suppressed comments (5)

python/freetoken/engine/engine.py:488

decode_freq is updated by device-side scatter_add_ before the cache kernel, and the flag is assigned before GraphRunner capture. That operation is captured and replayed with every decode step, so this warning incorrectly tells users to disable graphs and claims the histogram is stale; remove the warning and update the matching CLI/config documentation to describe the graph-safe accumulation.

        # Set after graph capture: the decode_freq histogram is scattered host-side
        # before the kernel rewrites expert ids to slots, so a captured decode graph
        # replays without it -- warn that the routing stats need graphs off.

python/freetoken/engine/engine.py:745

Owner EP currently reaches this cache construction without carrying the resolved decode_target, while OwnerOffloadMoeCache hard-codes its inner cache to decode_target="gpu". Consequently --moe-cpu-layers is accepted for owner EP but is silently ignored: _decode_owner runs before the normal CPU-layer branch. Either implement the owner-local CPU path or reject this option during owner-EP validation rather than serving a different configuration than requested.

            cache = OwnerOffloadMoeCache(
                owner_geometry,
                self.device,
                cache_policy=config.moe_cache_policy,
                prefill_hit_d2d=config.moe_prefill_hit_d2d,
                quant_format=banks.quant_format,
                graph_safe=_owner_graph_safe(config),
                layout=layout,
                max_slots=max_slots,
            )

python/freetoken/moe/offload_cache.py:1229

__getattr__ forwards runtime rebuild() calls to the inner cache, whose implementation disables prefill_overlap when the new size is below 2 * local_num_experts. geometry.prefill_overlap and geometry.cache_size remain unchanged, so materialize_layer() still takes the overlap path and later calls wait_prefill_layer() against an inner cache that has overlap disabled. Synchronize the owner geometry after rebuild or reject resize requests that invalidate it.

    def __getattr__(self, name):
        # Keep the wrapper small while preserving the existing cache's read-only reports and
        # bank-view helpers. Explicit route methods below prevent unsafe legacy admission.
        cache = object.__getattribute__(self, "_cache")
        return getattr(cache, name)

python/freetoken/moe/route_trace.py:81

Every TP rank constructs this recorder with the same configured path, and wb truncates the body while each rank also rewrites the same metadata file. In TP2+EP2 the writers race, so the resulting binary/metadata pair can be truncated or contain only one rank's trace. Use a rank-specific trace path (or restrict recording to one rank) before opening the file.

        self._f = open(path, "wb")

python/freetoken/server/stats.py:210

Although the new limits.max_seq_len is effective, the model.ctx field returned by /v1/stats still comes from derive_model_card(config) and therefore remains the raw checkpoint ceiling. launch._stats_context_length() uses that field as its fallback for client window sizing, so a smaller KV pool can still make clients send prompts that the scheduler rejects. Set the stats model card's ctx to effective_max_seq_len, while retaining the raw value in limits.model_max_seq_len.

Files reviewed: 47/47 changed files
Comments generated: 7
Review effort level: Lite

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

leiyu1980 · 2026-09-11T11:33:13Z

                config.model_path,
                self.device,
                include_moe_experts=not is_offload_moe_strategy(config.moe_strategy),
+                tp_shard=config.tp_info.size > 1,


Fixed in 81d034d. Confirmed a real regression, and wider than it
looks: only qwen4_exp declares tp_shard, while llama, qwen2, qwen3,
qwen3_moe, mistral, gpt_oss and minimax_m2 shard inside iter_weights via
shard_tensor(rank=tp_info.rank, world_size=tp_info.size). tp_shard appears zero
times in models/weight.py on main, so TP>1 worked there and broke here.

The flag is now forwarded only to readers that declare it. FTW treats it as the no-op
it is — it stores post-shard weights, so tp_shard=True used to raise there.
Readers that genuinely cannot shard still fail in load_state_dict's shape check,
exactly as on main. Covered by tests/models/test_weight_tp_shard.py.

leiyu1980 · 2026-09-11T11:33:19Z

                parallel=expert_parallel,
                decode_target=("cpu" if decode_target in ("cpu", "hybrid") else "gpu"),
                layer_residency=requested_residency,
+                ownership=ownership,


Fixed in 81d034d, by rejecting rather than implementing. You are
right that this cannot work as written: load_ftw_banks has no ownership parameter
and its docstring states it rebuilds [num_experts, ...] global rows, so binding those
to a local geometry is impossible without slicing and renumbering inside the FTW
loader.

Owner EP now fails fast for FTW in _validate_owner_ep_config (before any model or
bank allocation), and load_expert_banks raises too so a converter or tool call cannot
reach the same inconsistent state. Covered by
tests/engine/test_owner_ep_config.py::test_ftw_checkpoints_are_rejected.

leiyu1980 · 2026-09-11T11:33:26Z

+    if key in {"model.embed_tokens.weight", "lm_head.weight"}:
+        rows = div_ceil(tensor.shape[0], world_size)
+        start = rank * rows
+        return tensor[start : min(start + rows, tensor.shape[0])].contiguous()


Fixed in 81d034d — the tail is zero-padded to div_ceil(V, tp).

Worth noting this was reachable, and the model side already expected padding:
test_skeleton.py::test_parallel_lm_head_gathers_vocab_shards_in_rank_order copies
torch.cat((full_weight[4:], torch.zeros(1, 3))) into rank 1 for a 7-row vocabulary
and its docstring says "restore rank-major vocab rows and trim padding". Two
existing TP reassembly tests were encoding the truncating contract (the fixture
vocabulary is not divisible by 2) and now assert the padded one.

The production checkpoint (vocab_size=248320) is divisible by 2/4/8, so this is a
no-op there — the weight gate reports shape=0 at both TP1 and TP2+EP2.

leiyu1980 · 2026-09-11T11:33:30Z

    decode_target: str = "gpu",
    layer_sink=None,
    layer_residency: list[str] | None = None,
+    ownership=None,
 ) -> ExpertBanks:
    """Load (or fabricate, with ``dummy=True``) the expert banks. Two paths, both returning
    the same normalized ``ExpertBanks`` and both pinning after fill:


Same fix as the engine.py thread: rejected in 81d034d.
load_expert_banks now raises NotImplementedError when an owner-local bank is
requested over an FTW checkpoint instead of returning global rows, and
_validate_owner_ep_config rejects the combination before any allocation.

leiyu1980 · 2026-09-11T11:33:34Z

+            prefill_overlap=geometry.prefill_overlap,
+            prefill_hit_d2d=prefill_hit_d2d,
+            quant_format=quant_format,
+            decode_target="gpu",


Fixed in 81d034d, by rejecting the option. Your reading is exactly
right: the owner constructor does not accept decode_target at all (it hard-codes
"gpu"), and _decode_owner is selected before the is_cpu_layer branch, so
--moe-cpu-layers was accepted and then had no effect.

_validate_owner_ep_config now rejects it with a message that says why, and the
constructor documents the GPU-only contract. Covered by
test_owner_ep_config.py::test_cpu_layers_are_rejected_instead_of_silently_ignored
(parametrised over 0, 0,1, 0-3).

leiyu1980 · 2026-09-11T11:33:40Z

+            if self._cache._pending_src_layer is None:
+                return


Fixed in 81d034d. Confirmed: grep finds no assignment of
_pending_src_layer back to None anywhere in the file, so is not None stopped
being a usable proxy. The eager ensure_route only sets _pending_owned = True inside
if owned_positions.numel():, so an all-remote route skips the inner admission while
_decode_owner still calls copy_missing() — and the fallback then replayed the
previous layer's src_indices/evict_slots.

copy_missing now consumes the staged state exactly once (the whole_layer flag is
captured into a local before clearing), so a second call with nothing staged is an
explicit AssertionError rather than a silent replay. Covered by
test_offload.py::test_copy_missing_consumes_the_staged_layer_exactly_once.

leiyu1980 · 2026-09-11T11:33:44Z

+# Load route_trace.py directly by path: it is stdlib-only, so importing it without
+# the freetoken.moe package __init__ (which pulls torch/transformers) keeps these
+# tests runnable on a bare interpreter.
+_MOD = Path("/home/zhanglei/code/freetoken/FreeToken/python/freetoken/moe/route_trace.py")


Fixed in 81d034d — resolved from __file__ now:

_MOD = Path(__file__).resolve().parents[2] / "python" / "freetoken" / "moe" / "route_trace.py"

It had been pointing at an old fork checkout, so the module failed to load at
collection time in any other checkout.

… tests Review follow-ups for the owner-EP/TP change (Copilot review on FlashML-org#447). The functional fixes, in severity order: * `load_weight` forwarded `tp_shard` to EVERY reader at TP>1 and raised for readers without that parameter. Seven readers shard internally instead (llama, qwen2, qwen3, qwen3_moe, mistral, gpt_oss, minimax_m2 call `shard_tensor` with `tp_info.rank`/`tp_info.size`), so TP>1 on any of them — and on FTW checkpoints, which store post-shard weights — failed at startup. The flag is now forwarded only to readers that declare it; readers that cannot shard still fail in `load_state_dict`'s shape check, as before. * `OffloadMoeCache.copy_missing` never cleared `_pending_src_layer`, so a later call with nothing freshly staged replayed the PREVIOUS layer's `src_indices`/`evict_slots` and could overwrite slots reassigned to another layer. The owner adapter's "nothing staged" guard depends on one-shot consumption. Staged state is now consumed exactly once. * The Qwen4Exp vocabulary shard truncated its final rank instead of padding it. `VocabParallelEmbedding`/`ParallelLMHead` always allocate `div_ceil(V, tp)` rows and trim the padding when gathering (see `test_skeleton.py::test_parallel_lm_head_gathers_vocab_shards_in_rank_order`), so a checkpoint whose vocabulary is not divisible by TP failed strict shape loading. The tail is now zero-padded. * `--moe-cpu-layers` and FTW checkpoints are now REJECTED under owner EP instead of being accepted and silently ignored (the owner decode path runs before the CPU-layer branch, and `load_ftw_banks` rebuilds global expert rows with no ownership filter). `load_expert_banks` guards the FTW case too. * `OwnerOffloadMoeCache.rebuild` now re-derives its frozen geometry from the inner cache, which disables `prefill_overlap` when the new size cannot hold two local layers; the geometry kept the old value, so `materialize_layer` still took the overlap path and waited on buffers that no longer existed. * `--moe-trace-route` writes one file per rank. Every rank was handed the same configured path and opened it `wb`, so the TP writers truncated each other. * `/v1/stats` reports the ENFORCED ceiling in `model.ctx`, not the checkpoint ceiling: `launch._stats_context_length` reads that field to size a client window, which re-introduced the over-long-prompt 400 that `limits` exists to prevent. The raw ceiling stays in `limits.model_max_seq_len`. * The `--moe-collect-decode-freq` warning and its config comment claimed the histogram is host-side and needs graphs off. It is a device tensor accumulated by a device-side `scatter_add_`, so a captured graph replays it; the warning is gone and the comment now describes the graph-safe accumulation. Tests: `tests/models/test_weight_tp_shard.py` (forwarding contract, FTW no-op), `tests/engine/test_owner_ep_config.py` (the two new rejections plus the existing guards), route-trace per-rank paths, owner rebuild geometry sync, the one-shot pending state, `/v1/stats` `model.ctx`, and padded vocabulary shards. The two existing TP reassembly tests asserted the old truncating contract and now assert the padded one.

leiyu1980 · 2026-09-11T11:00:53Z

All nine review points are addressed in 81d034d. Every one checked out against the
code, including the two that were suppressed — no false positives. Per-point mapping:

Blocking

1. engine.py — tp_shard unconditional (3 votes). Confirmed a real regression.
Only qwen4_exp declares tp_shard; seven readers shard internally instead
(llama, qwen2, qwen3, qwen3_moe, mistral, gpt_oss, minimax_m2 call
shard_tensor with tp_info.rank/tp_info.size), and tp_shard does not appear in
models/weight.py on main at all — so TP>1 worked there and broke here. FTW was
worse: load_weight raised outright, and FTW stores post-shard weights, so TP>1 was
valid. The flag is now forwarded only to readers that declare it, and FTW treats it as
the no-op it is. Readers that genuinely cannot shard still fail in load_state_dict's
shape check, exactly as on main.

2. offload_cache.py — stale _pending_src_layer. Confirmed: grep finds no
assignment back to None anywhere, and the eager ensure_route only sets
_pending_owned = True inside if owned_positions.numel(): — so an all-remote route
skips the inner admission, _decode_owner still calls copy_missing(), and the
fallback replays the previous layer's src_indices/evict_slots. copy_missing now
consumes the staged state exactly once.

3. tests/moe/test_route_trace.py — absolute author path. Confirmed, and it was
pointing at the old fork checkout. Now resolved from __file__.

Non-blocking

4. FTW + owner EP (2 votes). Confirmed: load_ftw_banks has no ownership
parameter and documents [num_experts, ...] global rows. Rejected now — both in
_validate_owner_ep_config (fail-fast, before any allocation) and in
load_expert_banks, so a converter/tool call cannot reach the same state.

5. --moe-cpu-layers + owner EP. Confirmed: the owner constructor hard-codes
decode_target="gpu" and does not accept the argument, and _decode_owner runs before
the is_cpu_layer branch. Rejected in validation rather than served differently than
requested.

6. Short final vocabulary shard. Confirmed, and this one was reachable: the test
fixture's vocabulary is not divisible by 2, so two existing TP reassembly tests were
encoding the truncating contract. The model side already expects padding —
test_skeleton.py::test_parallel_lm_head_gathers_vocab_shards_in_rank_order copies
cat(full[4:], zeros(1, 3)) into rank 1 and its docstring says "and trim padding". The
tail is now zero-padded and those two tests assert the padded contract.

7. Shared trace path across ranks. Confirmed: every rank got the same configured
path and opened it wb, plus the same .meta.json. One file per rank now (no suffix at
TP=1, so the offline replay tooling is unaffected).

8. /v1/stats model.ctx. Confirmed — and it undercut this PR's own companion fix:
launch._stats_context_length reads payload["model"]["ctx"], which came from
derive_model_card → config.max_seq_len, i.e. the raw checkpoint ceiling. ctx now
carries the enforced value; the raw ceiling stays in limits.model_max_seq_len.

9. decode_freq warning. Confirmed inverted: decode_freq is created with
device=self.device and accumulated by a device-side scatter_add_ at the raw-ids
point, so a captured graph replays it. Warning removed and the engine/config.py
comment now describes the graph-safe accumulation.

Verification

CPU suite (tests/, -m "not slow", no e2e):
  1347 passed, 2 failed, 425 skipped

The two failures are tests/models/test_quant_config.py probing local HF checkpoints
that are absent here, and they reproduce on a clean main. The count is 1319 → 1347
(+28 new tests): tests/models/test_weight_tp_shard.py,
tests/engine/test_owner_ep_config.py, plus additions to the offload, route-trace,
/v1/stats and Qwen4Exp weight suites.

Weight-loading gate on 2× RTX 4090 (idle pair, fp8 KV, --cuda-graph-max-bs 0):

tp1 (TP1):    exact=722/794 widened=72 shape=0 dtype=0 missing=0 unexpected=0 -> PASS
tp2 (TP2+EP2):exact=722/794 widened=72 shape=0 dtype=0 missing=0 unexpected=0 -> PASS

End-to-end on the same branch, 2× RTX 4090, TP2+EP2, 262144 context, graphs [1,2],
bs≤2: TP2_OK, 2 + 2 = 4, and a Chinese MoE question all answer correctly;
moe_cache_size=4466, miss_rate=0.162, missing/layer=0.827, KV 4096 pages × 64 =
262144 tokens, steady-state decode 63.6 tok/s (max 65.1), i.e. unchanged from the
pre-fix 64.3. Auth: LAN no-key 401, bad-key 401, good-key 200, loopback 200.

Copilot AI lite review requested due to automatic review settings September 11, 2026 10:15

Copilot started reviewing on behalf of leiyu1980 September 11, 2026 10:15 View session

Copilot AI reviewed Sep 11, 2026

View reviewed changes

github-actions Bot mentioned this pull request Sep 11, 2026

FreeToken ROCm Watch — notification channel Castoff995/FreeToken-rocm-test#1

Draft

Navigation Menu

feat(moe): owner-local expert parallelism and tensor parallelism - #447

feat(moe): owner-local expert parallelism and tensor parallelism#447
leiyu1980 wants to merge 2 commits into
FlashML-org:mainfrom
leiyu1980:pr/tp-ep-clean

leiyu1980 commented Sep 11, 2026 •

edited

Loading

Uh oh!

Copilot AI left a comment

Uh oh!

leiyu1980 Sep 11, 2026

Uh oh!

leiyu1980 Sep 11, 2026

Uh oh!

leiyu1980 Sep 11, 2026

Uh oh!

leiyu1980 Sep 11, 2026

Uh oh!

leiyu1980 Sep 11, 2026

Uh oh!

leiyu1980 Sep 11, 2026

Uh oh!

leiyu1980 Sep 11, 2026

Uh oh!

leiyu1980 commented Sep 11, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Conversation

leiyu1980 commented Sep 11, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Why owner-local experts

What is in it

Owner-local expert parallelism

Dense tensor parallelism

Two companion changes (not part of TP/EP itself)

Tested on

Exact command

Results

1. Weight loading — exact

2. CPU suite

3. End-to-end A/B — main vs this branch

4. Correctness — and an honest caveat

Relationship to other open PRs

Known limitations

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

🟡 Changes recommended

Suppressed comments (5)

Uh oh!

leiyu1980 Sep 11, 2026

Choose a reason for hiding this comment

Uh oh!

leiyu1980 Sep 11, 2026

Choose a reason for hiding this comment

Uh oh!

leiyu1980 Sep 11, 2026

Choose a reason for hiding this comment

Uh oh!

leiyu1980 Sep 11, 2026

Choose a reason for hiding this comment

Uh oh!

leiyu1980 Sep 11, 2026

Choose a reason for hiding this comment

Uh oh!

leiyu1980 Sep 11, 2026

Choose a reason for hiding this comment

Uh oh!

leiyu1980 Sep 11, 2026

Choose a reason for hiding this comment

Uh oh!

leiyu1980 commented Sep 11, 2026

Blocking

Non-blocking

Verification

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

leiyu1980 commented Sep 11, 2026 •

edited

Loading