Conversation
|
|
Both from @gdevenyi's review on 2 x RTX 6000 Ada, where this is carried on a deploy branch alongside FlashML-org#385 (TP). 1. Attention routed its bf16 fallback through the quantized factories too, which swaps in their generic fallback and drops the tensor-parallel classes FlashML-org#385 needs (per-rank local_output_sizes, row-parallel o_proj). That is exactly the path a rank takes under TP>1, since the block-FP8 linears have no parallel variant. The factories are now used only on the fp8_block branch; every other case keeps LinearColParallelMerged / LinearReplicated as before. 2. config.parse_config and weight._dense_is_block_fp8 read the same declaration through two independent code paths, each with its own copy of _FP8_BLOCK_ALGOS. That is safe only while they cannot disagree, and they can: a rank downgrading under TP>1 must have the modules it BUILDS and the buffers it LOADS downgrade together, or the buffers will not match. Both now resolve through one helper, config.dense_quant_mode, which owns the declaration test and the TP downgrade. The duplicate constant is gone. It reads TP through try_get_tp_info, not get_tp_info: Engine.__init__ sets TP info as its first statement so a rank always knows its size by the time this matters, but config parsing also happens with no engine at all (checkpoint conversion, tooling, tests) where get_tp_info raises. Verified on the modelopt checkpoint: parse_config still yields nvfp4/fp8_block; the two sides agree at TP=1 (both fp8_block) and at TP=2 (both downgraded); attention builds LinearColParallelMerged/LinearReplicated under bf16 and Fp8BlockColMerged/Fp8BlockLinear under fp8_block. tests/models/qwen4_exp/test_config.py + test_weight.py: 30 passed. The whole qwen4_exp suite reports 47 failed / 46 passed / 50 skipped both at the merge-base and with these fixes - identical sets, no regressions. Those failures are pre-existing and are an artefact of this box rather than the code: its single 24 GB card is 23.6 GB occupied serving a model, so the GPU-dependent tests cannot allocate. I have not been able to run them on a free card.
Follow-up to @gdevenyi's note on the previous fix. The bf16 branch kept o_proj as LinearReplicated, which is what main does today and is correct at TP=1, but it is the path a rank falls back to under TP>1 (FlashML-org#385), and there a replicated o_proj is wrong three ways at once: qkv_proj is column-parallel so each rank's attention output is its local head slice, o_proj therefore needs the sharded input dim, and the partial sums need an all-reduce. It also fails quietly, since a missing reduction still decodes to fluent-looking text. LinearOProj does all of that and degenerates to LinearReplicated at TP=1: div_even(x, 1) == x, and the all-reduce is skipped when tp_size == 1. So this is a no-op for main and only changes what FlashML-org#385 finds when the two meet, whichever lands second. It does mean get_tp_info() runs in __init__, but the same branch already does that two lines up through LinearColParallelMerged, so there is no new constraint: this path was engine-only before and still is. Config parsing, the one no-engine path that mattered, stays on try_get_tp_info. The comment above the branch also claimed a row-parallel o_proj that the code did not build; it now describes what is built. Verified on ailab1 (single RTX PRO 4000 Blackwell, TP=1): - tests/models/qwen4_exp/{test_config,test_weight,test_skeleton, test_qsa_backend}.py: 49 passed at 72773b0 and 49 passed with this change. That set includes test_qsa_layer_matches_hf_dense, which builds Qwen4ExpAttention on the card and checks the whole layer against the HF dense reference, so the new o_proj is exercised through a real forward. - A direct check on CPU and CUDA: the built o_proj is a LinearOProj with weight [hidden, qo_attn_dim] and local_input_size == qo_attn_dim, and its forward is bit-identical (max |diff| = 0.0) to a LinearReplicated carrying the same weight. Correction to the previous message: the 47 failures I attributed to VRAM contention were, for these four files at least, a missing ninja on PATH in my throwaway test venv. With the serving venv's bin on PATH the files pass in full.
qwen4_exp is the only family that builds its attention o_proj as LinearReplicated; llama, gpt_oss and minimax_m2 all use LinearOProj. That is correct at TP=1 and wrong under TP>1 three ways at once: qkv_proj is column-parallel, so a rank's attention output is its local head slice rather than the full qo_attn_dim, o_proj therefore has to take the sharded input dim, and the partial sums need an all-reduce. LinearReplicated keeps the full [hidden, qo_attn_dim] weight, expects the unsharded input and reduces nothing. It also fails quietly: a missing all-reduce leaves each rank holding a partial sum that still decodes to fluent-looking text. LinearOProj degenerates to exactly LinearReplicated at TP=1 -- div_even(x, 1) == x, and the all-reduce is skipped when tp_size == 1 -- so this is a no-op for main as it stands and only changes what FlashML-org#385 finds when the two meet. It adds no constraint from calling get_tp_info() in __init__ either, since the same constructor already reaches it two lines up through LinearColParallelMerged. The comment above the branch now describes what is built. Raised by @gdevenyi against the earlier form of this work in FlashML-org#392.
…ackend) Shard the dense weights per rank at load (attention qkv by head, GDN in_proj as its six parts with the matching conv1d channels and A_log/dt_bias, shared-expert gate_up per part; o_proj/out_proj/down_proj row-parallel; embed/lm_head by vocab rows) and the NVFP4 expert banks along the intermediate axis, so every rank holds half the experts and each MoE layer needs one all-reduce (routed + gate * shared are combined before the reduce). Router, QSA indexer, norms, hyper-connections and PLE stay replicated so all ranks select the same blocks and n-gram rows. Also: LinearColParallelMerged(local_output_sizes=) for the kv-replicated case and distributed_timeout 60 -> 1800 s (ranks reach their first collective minutes apart behind a 100+ GiB load). Limits: offload backend with bf16 dense projections; fp8_block / nvfp4 dense checkpoints raise under TP. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_0173pf9k9fSVtwbm3f898HDt
tests/models/qwen4_exp/test_weight.py feeds iter_weights a synthetic checkpoint whose config.json has no model_type; at TP=1 nothing is sharded, so do not touch the config. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_0173pf9k9fSVtwbm3f898HDt
d16d634 to
a99433c
Compare
|
Rebased onto What changed and why. The original sharded inside So the sharding moved to the two places the new architecture actually owns:
Two other resolutions worth flagging:
Testing. Full 🤖 Generated with Claude Code |
|
Re-tested against current main ( Method. This PR's head merged onto main, then the full Result: 1214 passed, 350 skipped, no new failures. The +9 over main are this PR's own tests, and they ran (not skipped). 🤖 Generated with Claude Code |
…build
Qwen4ExpDecoderLayer builds its MoE as `Qwen4ExpMoE(config, layer_id, prefix=...)`,
but this PR's override of __init__ (added to hold the TP communicator) took only
(config, layer_id), so a server boot died with
TypeError: Qwen4ExpMoE.__init__() got an unexpected keyword argument 'prefix'
The whole CPU test suite was green with that bug in place, because every test that
builds a decoder layer is behind requires_cuda -- nothing without a GPU ever
constructed the model. tests/models/qwen4_exp/test_build_cpu.py closes that: it
builds the full model on the meta device (no GPU, no memory) and asserts the state
dict has both layer families, an lm_head, and MoE weights on more than one layer,
so a dropped or shared prefix fails too.
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_0173pf9k9fSVtwbm3f898HDt
…el shard
With the expert piece stream sliced per rank and the banks sized from
MoEConfig.local_intermediate, a rank holds exactly its half of every expert, so this
kernel can serve TP>1 -- the routed output is a partial sum and the MoE layer already
reduces it (_maybe_all_reduce, or the single combined all-reduce in qwen4_exp's block).
Without this the whole selection table is empty under TP=2 on sm_89 and the server
refuses to start:
KernelSelectionError: no usable kernel in table;
triton: TP > 1 is not supported for this expert format;
marlin: vLLM is not installed;
b12x: b12x requires sm_120+, got sm_89
marlin and b12x keep tp_ok=False deliberately: their pack() repacks the native rows and
neither has been verified against a per-rank bank.
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_0173pf9k9fSVtwbm3f898HDt
|
Pushed Integrating this branch with the NVFP4 expert path after the #418/#427 quantization refactor, the engine refused to start: Every NVFP4 MoE kernel carried Worth a second pair of eyes on one thing: this made the bank-sizing bug below silent rather than fatal, so I want the sizing assertion reviewed as much as the flag. Related, and the more interesting failure. A merge of upstream main into the deployment line silently dropped this same Also on this branch since the last comment: 🤖 Generated with Claude Code |
|
End-to-end validation of this PR at TP=2, as a controlled A/B against the branch it replaces. Our deployment branch (upstream
Throughput and residency are identical. TTFT is ~7% better on the mean, and more interestingly it is flat at 0.68 from the first measurement where the old branch drifts 0.79 → 0.70 as its expert cache settles. Accuracy is one question higher with two fewer genuine misses — both inside noise at n=300, but nothing regressed. Two methodology notes, because they cost me real time and may save yours: Do not benchmark this box inside ~10 minutes of start. The aggregate reads 2–3% low while the expert cache fills. My first run on the new branch showed 343.7 against a 352.3 baseline and I nearly reported a 2.4% regression that did not exist; settled, it is 354–357. GSM8K at the default 768-token cap measures truncation as much as reasoning. At 768 the two branches score 95.67% and 96.33%, which looks like a regression; the difference is entirely that one truncated 10 answers and the other 6. At 1536 the ordering reverses. Per-question flip analysis: 9 of 10 truncated answers become correct when the cap is lifted. Since the last comment this branch also carries 🤖 Generated with Claude Code |
|
Closing this. It was written with heavy AI assistance, and the maintainers have indicated that do not want such contributions. The description and the diff stay here for anyone who wants to pick the idea up. |
19 commits and upstream's README is not something a stranger can act on. This adds README.gfx1201.md -- the branch's own documentation -- and prepends a banner to README.md so that whoever lands on the branch page sees it. ⛔ Upstream's README is PREPENDED TO, never replaced. Its text is byte-identical below the banner. What the documentation had to carry, or it would be worse than nothing: - The three settings without which it does not serve, with the reason each one is not discoverable: `--expert-load parallel` (auto picks serial here, ~100x slower load, and the flag exists partly to override a host-RAM guard that a 128 GiB box can never satisfy); PYTORCH_ALLOC_CONF=expandable_segments:False (the engine forces :True, which GPU-FAULTS on gfx1201 -- and torch's own OOM message advises :True); HIP_VISIBLE_DEVICES set explicitly, always. - The defects, named. Image token positions on this branch are wrong, and so is everything after them: the checkpoint declares M-RoPE and this branch feeds 1-D positions. Nothing errors. Text-only is provably exact. ⭐ Upstream merged a correct M-RoPE implementation in FlashML-org#454 on 2026-09-13, so for image quality upstream's vision path is better than this one, and the README says so. Nothing bounds the number of images in one request, so one legal request can exhaust host RAM mid-request; the mitigations that work are listed. TP=2 with images is the least-tested combination here. - What is actually still ours, checked against upstream main at 68a81ff rather than assumed: the RDNA LDS clamp (without it this model serves prompts of at most 15 tokens on RDNA before the worker dies), tensor parallelism for qwen4_exp (upstream's weight loader still raises TP=1-only), the relay handshake, and all-rank agreement on an encode failure. Vision and the multimodal scheduler work have upstream counterparts now and are credited as such. @gdevenyi's PRs FlashML-org#385 and FlashML-org#386 are credited by number: the same shard was arrived at independently on CUDA at the same time. - Host RAM is the gate, not VRAM: the expert banks stay host-resident at ~31.65 GiB per rank, at TP=1 and at TP=2 alike. ⛔ No quality or fidelity claim is made anywhere in it. This work has no fidelity instrument, and the README says that too. llm-server #1018.
19 commits and upstream's README is not something a stranger can act on. This adds README.gfx1201.md -- the branch's own documentation -- and prepends a banner to README.md so that whoever lands on the branch page sees it. ⛔ Upstream's README is PREPENDED TO, never replaced. Its text is byte-identical below the banner. What the documentation had to carry, or it would be worse than nothing: - The three settings without which it does not serve, with the reason each one is not discoverable: `--expert-load parallel` (auto picks serial here, ~100x slower load, and the flag exists partly to override a host-RAM guard that a 128 GiB box can never satisfy); PYTORCH_ALLOC_CONF=expandable_segments:False (the engine forces :True, which GPU-FAULTS on gfx1201 -- and torch's own OOM message advises :True); HIP_VISIBLE_DEVICES set explicitly, always. - The defects, named. Image token positions on this branch are wrong, and so is everything after them: the checkpoint declares M-RoPE and this branch feeds 1-D positions. Nothing errors. Text-only is provably exact. ⭐ Upstream merged a correct M-RoPE implementation in FlashML-org#454 on 2026-09-13, so for image quality upstream's vision path is better than this one, and the README says so. Nothing bounds the number of images in one request, so one legal request can exhaust host RAM mid-request; the mitigations that work are listed. TP=2 with images is the least-tested combination here. - What is actually still ours, checked against upstream main at 68a81ff rather than assumed: the RDNA LDS clamp (without it this model serves prompts of at most 15 tokens on RDNA before the worker dies), tensor parallelism for qwen4_exp (upstream's weight loader still raises TP=1-only), the relay handshake, and all-rank agreement on an encode failure. Vision and the multimodal scheduler work have upstream counterparts now and are credited as such. @gdevenyi's PRs FlashML-org#385 and FlashML-org#386 are credited by number: the same shard was arrived at independently on CUDA at the same time. - Host RAM is the gate, not VRAM: the expert banks stay host-resident at ~31.65 GiB per rank, at TP=1 and at TP=2 alike. ⛔ No quality or fidelity claim is made anywhere in it. This work has no fidelity instrument, and the README says that too. llm-server #1018.
1 participant
What this adds
ft serve --tp-size 2forqwen4_exp(Qwen3.8-Flash-Next) on the offload MoE backend. Upstream refuses TP>1 for this architecture; with two 48 GiB cards the model runs as two independent TP=1 instances, each keeping 37% of the NVFP4 experts resident and streaming the rest over PCIe every step.Per rank, the patch:
weight.py::_shard): attentionqkv_projby head ([q|gate]per head; kv heads split, or replicated when there are fewer than ranks), GDNin_projas its six parts ([q | k | v | z | b | a]) plus the matchingconv1dchannels andA_log/dt_bias, shared-expertgate_up_projper part;o_proj,out_projand the shareddown_projrow-parallel with the all-reduce inside;embed_tokens/lm_headby vocab rows. Router, QSA indexer, norms, hyper-connections and PLE stay replicated, so every rank selects the same sparse blocks and n-gram rows.nvfp4_banks.py, I=640 -> 320 per rank): packed codes, the 16-wide scale blocks and the per-row globals for the gate/up rows and the down columns. The offload cache then holds half the experts per rank.moe.py): the routed and shared partial sums are combined asrouted + sigmoid(gate) * sharedbefore a single reduce, instead of one reduce each.LinearColParallelMerged(local_output_sizes=)for the kv-replicated case (the same shape as feat(models): support TP for qwen3_5_moe #104's hunk), anddistributed_timeout60 s -> 1800 s: behind a 100+ GiB load the ranks reach their first collective minutes apart, and 60 s kills the launch.Everything else in the engine (scheduler, KV pool, PLE table, CUDA graphs) is untouched.
Measurements
2 x RTX 6000 Ada (48 GiB, sm_89, PCIe Gen4 x16, no NVLink), 2 x Xeon Gold 6526Y, 503 GiB RAM.
RadixArk/Qwen3.8-Flash-Next-NVFP4,--moe-backend offload --ple-backend pinned --num-tokens 262144 --memory-ratio 0.94 --moe-prefill-hit-d2d.A 262,144-token prompt reaches its first token in 74 s at TP=2 (116 s at TP=1). The decode step is dominated by the bf16 dense read; TP=2 halves it per GPU and removes the PCIe expert gather because everything fits.
Correctness. An 8-question probe gives identical answers at TP=1 and TP=2. Three raw prompts decoded greedily for 256 tokens: the ~1k-token prompt (QSA over many blocks, GDN state, PLE context) is word-for-word identical across every run; the two short prompts diverge after 13 and 29 words between TP=1 and TP=2, but the TP=2 server diverges from itself at the same points on a second pass (bf16 atomics in the expert kernels), so that is run-to-run noise rather than a sharding error. Greedy output at TP=2 is not bit-exact between passes.
Limits
Offload backend with bf16 dense projections only:
fp8_block/nvfp4dense checkpoints raise under TP (row-parallel FP8 / NVFP4 linears do not exist yet, the same gap #104 has). The hybrid / CPU MoE backends are not sharded.Related: #62, #29 (TP for offloaded MoE), #104 (TP for qwen3_5_moe, which this reuses the merged-linear hunk from).
Testing
tests/models/qwen4_exp/test_tp_shard.py: the per-head / per-part row sharding of every fused projection reassembles to the original (CPU).tests/models/test_nvfp4_banks_tp.py: the bank placer's per-rank slices of codes, scales and globals cover the intermediate axis exactly once (CPU).tests/models/qwen4_expon one of its GPUs: 97 passed, 3 failed; the same 3 (test_chunked_prefill_matches_one_shot[*], a bit-exact assertion off by bf16 noise on this torch 2.11 / flashinfer 0.6.18 / triton 3.6 stack) fail on plainmainthere too.🤖 Generated with Claude Code
https://claude.ai/code/session_0173pf9k9fSVtwbm3f898HDt