qwen4_exp: build o_proj row-parallel, as every other family does

[gberasmus87]

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 in 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. The loader already assumes the row-parallel layout — _shard puts o_proj on dim 1.

It also fails quietly. A missing all-reduce leaves each rank holding a partial sum that still decodes to fluent-looking text, so there is no crash to catch it.

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 #385 finds when the two meet. It adds no constraint from calling get_tp_info() in __init__ either — the same constructor already reaches it two lines up through LinearColParallelMerged, so that path was engine-only before and still is.

The comment above the branch also promised a row-parallel o_proj that the code did not build; it now describes what is there.

Verified at TP=1 on an RTX PRO 4000 Blackwell: 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, on CPU and on CUDA. tests/models/qwen4_exp/ shows the same pass/fail set with and without the change.

Raised by @gdevenyi on #392 as a merge hazard against #385. Split out here so it can land on its own — #392's reader half is superseded by #428.

[gdevenyi]

Confirming this from the TP>1 side, which is the half the PR cannot test on one card.

deploy/chatdnp, the branch our production server runs, has built o_proj as LinearOProj at TP=2 since the qwen4_exp TP work went in. Same constructor, same argument order:

self.o_proj = LinearOProj(
    self.qo_attn_dim, config.hidden_size, has_bias=False,
    quant_config=config.quant, prefix=f"{prefix}.o_proj",
)

On 2 x RTX 6000 Ada (sm_89, PCIe, no NVLink), nvidia/Qwen3.8-Flash-Next-NVFP4, TP=2, offload MoE: GSM8K 97.00%, unchanged across every deploy in that period, and 99-105 tok/s single-stream. That number is the useful part of this report. Your own PR text names the failure mode precisely — a missing all-reduce leaves each rank with a partial sum that "still decodes to fluent-looking text" — so a token-level smoke test proves nothing here. An accuracy suite does, and 97.00% is not what a half-summed o_proj produces.

So: TP=1 bit-identity from you, TP=2 accuracy from us. Between them the change is covered in both directions.

Two notes for whoever merges.

It is not a no-op the moment #385 lands — it is a prerequisite. #385 shards qkv_proj by head. Without this change the input dim of o_proj and the width of the tensor reaching it disagree, and depending on the path you get a shape error or a silent partial sum. Landing #429 first is the right order and makes #385 smaller.

make_replicated_quant has the same bug one branch over. The block-fp8 path from #392 still builds o_proj replicated:

self.o_proj = make_replicated_quant("none", mode, self.qo_attn_dim, config.hidden_size)

That is consistent today only because #392 declares itself TP=1-only. If #428 supersedes #392 and drops that restriction, this branch needs the same treatment, and it will fail the same quiet way. Worth a comment there rather than a silent gap.

🤖 Generated with Claude Code

https://claude.ai/code/session_0173pf9k9fSVtwbm3f898HDt

[gberasmus87]

Thanks — GSM8K 97.00% at TP=2 is exactly the half I could not produce, and you are right that it is the load-bearing number. A token-level smoke test cannot distinguish a correct o_proj from a half-summed one, which is the whole reason this failure is worth fixing rather than waiting for someone to notice bad output.

Your make_replicated_quant flag is now moot, and I checked rather than assuming. #428 merged, #392 is closed, and #428 touched only weight.py and the tests — so make_replicated_quant no longer exists anywhere on main (git grep across python/ returns nothing; it only ever lived on #392's branch). On current main there is exactly one o_proj construction site, models/qwen4_exp/attention.py:131, and it is the LinearReplicated this PR replaces. No second branch, no silent gap left behind.

That also tightens your merge-order point. With the block-fp8 branch gone, this is a one-line change to the single site #385 needs corrected, so landing it first genuinely does make #385 smaller rather than just earlier.

To be precise about what is verified when: the make_replicated_quant check above is against current main (46d2743), today. The TP=1 evidence in the PR body — same pass/fail set across the qwen4_exp suite with and without the change, and the built o_proj bit-identical to LinearReplicated (max |diff| 0.0, CPU and CUDA) — was taken against this branch's base. The branch is a few commits behind main now; GitHub still reports it mergeable and clean, and I would rather rebase when there is something to rebase for than churn it under your review.

[gberasmus87]

Follow-up: I swept the rest of the repo for this pairing and it is not unique to qwen4_exp.

LinearQKVMerged shards heads (div_even(num_qo_heads, tp_info.size)) and LinearColParallelMerged shards its output sizes, so any family pairing either with a replicated o_proj has this bug. Seven families pair them with LinearOProj; gemma4, qwen3_5_moe and muse_glimmer do not. gemma4 is the starkest — every other LinearQKVMerged user in the tree pairs it with LinearOProj.

All three are latent exactly as this one is: their loaders raise NotImplementedError(... supports TP=1 only), so nothing is broken today. Opened as #440 rather than folded in here, so this PR stays the single-file change you have already validated.

Left alone deliberately: glm5_next, glm_moe_dsa and minimax_m3 replicate their q/kv projections too, so replicated o_proj is consistent there, and glm4_moe uses LinearDF11.