[dghosef]

Up to date with main at 9c6f1ad75 by a merge (#860 has landed).

CFG parallelism (classifier-free guidance, with the conditional and unconditional branches run on different GPUs) had no mesh axis of its own. Wan2.2 borrowed the SP axis for it: an sp_size: 2 instance whose ranks split guidance rows rather than the sequence. The engine sizes attention heads by EngineResourceInfo.head_shards, the node's whole instance size on the base, so Wan's dit declared RaggedAttentionConfig(shard_over_sp=False) to keep every head on those ranks. So sp_size meant two different things, and a CFG-parallel node couldn't also use Ulysses sequence parallelism. Separately, everything the ranks of a lockstep instance must agree on (capture barriers, and the votes through JointGroups.agree: capture decisions, the KV manager's symmetric free-pages check, the read-check vote, the TP-async flag check) ran as one collective per mesh axis, with the axes listed by hand. This PR gives CFG its own outermost axis, cfg_size, moves Wan2.2 onto it, and makes every lockstep agreement one gloo collective on one group per instance.

The mesh is row-major over PARALLEL_AXES, outermost first, and one helper (mesh_groups) builds it; comm groups, rank lists and shapes are keyed by axis name. The axes' own declarations still have one field or keyword each (ShardingConfig.<axis>_enabled_nodes, JointGroups.<axis>_group, get_submodule's sp_group / cfg_group), and which axes split heads is written out in head_shards. That step changes no behaviour. JointGroups.instance is a CommGroup over the instance's ranks, and every agreement is one gloo collective on it: a barrier, or a vote (JointGroups.agree is one CommGroup.agree on it rather than one per axis). It has no device group and no NCCL communicator, so the capture barriers are now gloo barriers. TP-only and SP-only deployments reuse their axis's gloo group and create no new process group. A multi-axis instance's gloo group is made first, on every rank: a rank outside the instance makes a group of just itself, so all members agree on the group's name. On XPU it is made like the XPU axis groups, and a device collective on a group with no device group raises ValueError. Placement refuses a lockstep instance whose ranks are listed out of order ([1, 0] with tp_size: 2): torch orders a process group by global rank, so such a group would gather in a different order from the one its ranks split work by and scramble TP shards or swap CFG branches without any error. Instances themselves may come in any order ([2, 3, 0, 1] is two valid replicas), and no shipped config is affected. CFG is set by a node group's cfg_size, shares the TP and SP code paths for instances, groups and lockstep agreement, and routes as one sharding group. On the wire it is still sent as tp_size (the instance size), so the conductor's messages and the Rust shard templates are unchanged. A model declares ShardingConfig.cfg_enabled_nodes with the most branches each node runs, and placement refuses, before any worker starts, cfg_size on an undeclared node or one that doesn't divide that count; a paged KV cache on a cfg_size > 1 node, because the pool's rank math counts every instance rank as a head shard (branches that hold different KV are separate nodes in a graph Parallel on a cfg_size: 1 node group, as BAGEL's LLM_cfg_text and LLM_cfg_img are); and a shard_dim signal into or out of a node whose node group has sp_size or cfg_size above 1, since the fanout splits such a signal over the whole instance while only the TP ranks split the weights (under TP2 × CFG2 or TP2 × SP2 each rank would get a quarter of the tensor where its weights expect half; that held for SP on the base too, and no shipped model or config declares a shard_dim). The engine manager hands the node cfg_group in get_submodule when it is non-trivial. head_shards goes from the instance size to TP ranks × SP ranks, so CFG ranks keep every head; _build_ragged still splits a shard_over_sp=False attention over TP only (LTX-2.5 SP2's audio and cross attentions). Wan2.2's wan22_cfg_parallel_2gpu.yaml sets cfg_size: 2; the dit splits its [cond, uncond] stack with my_rows(cfg_group, CFG_BRANCHES), and the model declares cfg_enabled_nodes={"dit": 2} and no longer sp_enabled_nodes. It drops shard_over_sp=False: its dit isn't SP-enabled at this PR, so placement refuses sp_size > 1 on it and the flag was a no-op, and #949's Ulysses SP needs the default, which splits heads over SP. Wan's own validate_config_yaml is deleted. A Wan2.2 config that splits guidance with sp_size: 2 must become cfg_size: 2: placement refuses the old shape and says so. Two smaller fixes: CommGroup.barrier raises on a group whose cpu_group isn't set up (torch reads group=None as the whole world), as agree already did; and a node handed two different groups on one axis of one rank is a ValueError that placement reports (it was a RuntimeError). That check also covers the instance group for multi-rank instances: two multi-rank instances of a node that share a rank differ there on some axis, size-1 groups included. It misses a node placed one-rank for one walk and multi-rank for another, since a worker graph with one-rank instances registers no groups; placement's resource-spec shape check still refuses that placement whenever the node belongs to a resource spec, so only a node with no resource spec gets through. That gap predates this PR.

GPU rows: Eras H200s (4× at b58ca20cc, 2× and 4× before the merge of #860), torch 2.12.1+cu132, Rust graph runtime (and control mesh at b58ca20cc, where both arms were built from one dependency set). Throughput rounds each on a fresh server.

Check	Setup	Result
Wan2.2 cfg_size: 2 vs base's sp_size: 2 CFG split, eager dit	b58ca20cc vs #860's 5dcbd78e5, one 81-frame 704×1280 video; also before that merge	Base, base and PR give byte-identical videos, on both trees.
Same, compiled dit	b58ca20cc vs 5dcbd78e5	The PR's video differs from base's by 26.8 dB, but the difference follows the Inductor cache, not the code. On a shared cache, base gave one video (twice) and the PR the other; with a fresh cache per arm the two swapped, byte for byte. Inductor's generated code (TORCH_LOGS=output_code) is identical between the arms on both ranks, and the traced graphs differ only in their generated function names.
Qwen3.5-9B TP2, lockstep agreement on the gloo instance group	b58ca20cc vs 5dcbd78e5, three rounds each	Startup and capture pass; a probe confirms JointGroups.agree and the barriers run on one gloo group over ranks (0, 1). c=32 4797 / 4792 / 4830 vs 4727 / 4788 / 4763 tok/s; c=1 239 / 240 / 251 vs 243 / 246 / 242. Greedy: the PR matches the two base runs on 11/16 and 13/16 prompts, where base matches itself on 13/16. No phase moved by more than run-to-run variance in the profile pass. On TP2 the instance group is the TP gloo group, so the read-check vote is the same single collective as on base; its 2.10 → 2.49 ms at c=32 (0.70 → 0.37 ms at c=1) is a wait phase moving between runs.
Same	Before the merge of #860	Startup and capture pass. c=32 4841 / 4814 vs 4756 / 4836 tok/s; c=1 238 vs 241. Greedy 14/16 identical to base, where base against itself with a cold cache gives 15/16. No moved phase.
Multi-axis instance groups, Cosmos3-Nano tp2_sp2 and sp2	Before the merge of #860, with #705 (now on main)	tp2_sp2 starts. Heads shard over TP × SP: output within base's own cold-cache spread (t2v 23.4 / 23.8 dB against 23.6 dB, t2i 42.1 / 42.4 dB against 45.1 dB).
ruff; CPU Core vs main 9f84f18cd	Eras CPU box, OMP_NUM_THREADS=2, head dc0c20852	ruff clean. CPU Core against main 9f84f18cd: 9091 passed, 1066 skipped, 2 failed, the same 2 as main (the test_eval_compare_multigpu dump_stacks pair); none fail only on this branch. GitHub CI green at dc0c20852.
New CPU tests	Linux CPU	Placement refuses a paged KV cache on a cfg_size: 2 node, and a shard_dim signal on TP × SP and TP × CFG instances (test_validate_deployment); a barrier on a group with no cpu_group raises (test_parallel_groups); head_shards with a CFG axis, including Wan's shipped cfg-only shape (test_ragged_attention_cpu::test_cfg_ranks_keep_every_head, which fails 4 of 5 cases with head_shards as the instance size). The three refusal tests fail without their fixes.
test_sequence_parallel, three-rank gloo case	Eras CPU box	A known intermittent failure: it aborts in some runs, more often on a loaded box. With a private --basetemp, the module passed 4 of 4 runs on both this PR and #860. The cause is under investigation.

Not yet run on GPU at this head: every GPU result above predates the merges of main (8a2095f81, then bb0e0a8d7, c96e02b7d and 9f84f18cd through #860, and 9c6f1ad75) and of #860's heads since 5dcbd78e5, including the change to Wan's attention declaration. The one exception: Qwen3.5-9B, which crashed at load on the previous head because _unchunkable_nodes still called the removed get_sp_config_for_node, boots, captures and serves on H200 with the fix (b3eaf9ddf). To confirm against main: Wan2.2 cfg_size: 2, eager dit, one 81-frame 704×1280 video byte-identical to main's sp_size: 2 CFG split (this confirms that dropping shard_over_sp=False keeps every head); Qwen3.5-9B TP2 startup, capture, c=32 and c=1 throughput over three fresh-server rounds with the phase-profile pass (expected neutral); Cosmos3-Nano tp2_sp2 and sp2 t2v / t2i within base's cold-cache spread. The placement checks, the error type and the barrier refusal change no GPU path, since production instance groups always have a cpu_group (init_dist sets it or raises); that is traced in code, not run. XPU: no hardware; the XPU instance-group path is covered only by a mocked test.

Follow-ups: get_submodule takes one keyword argument per axis (sp_group, cfg_group), and changing that public signature across all models is a separate decision. Per-axis knowledge is still written out in a few places (the ShardingConfig.<axis>_enabled_nodes fields, the JointGroups.<axis>_group fields with cfg_group defaulted last, the engine manager's kwarg tuple, and head_shards = tp × sp); one JointGroups.axes mapping plus each axis's role (splits heads, splits tokens, splits batch rows) beside PARALLEL_AXES would remove them, touching every JointGroups caller and #949's. tp_allreduce sets up its backend on every axis group, CFG groups included, though a CFG group only all-gathers: free with the default nccl, one 512 KiB buffer plus a rendezvous per group otherwise; restricting it to TP needs an audit of which comm_group.all_reduce callers run over SP. my_rows lives in components/distributed/sequence_parallel.py but is the CFG axis's row split too.

AI assistance (Claude) was used for this PR.

🤖 Generated with Claude Code

[vasilevklart]

Chunkable models crash at load; merging main breaks tests (inline)

[vasilevklart]

_unchunkable_nodes (:45) calls the removed get_sp_config_for_node, so Qwen3.5, Qwen3-Omni, Orpheus and BAGEL fail to load. You can use get_joint_group_for_node(node).sp_group (also #949; test_chunked_prefill.py:757's stub hid it). On H200, Qwen3.5-9B raises AttributeError at load; with the fix it boots, captures and serves

[dghosef]

Fixed in b3eaf9ddf - it reads get_joint_group_for_node(node).sp_group now, and the test builds real WorkerParallelGroups so a stub can't hide this again. Ty!

[vasilevklart]

Merging main breaks 5 CPU Core tests: main's test_lora_pool.py:498, test_lora_tp_lockstep.py:419 (and GPU-only test_qwen3_tts_packed_replay.py:85) build JointGroups without instance

[dghosef]

Fixed in the main merge (8020ecc00): those tests pass instance now. Main's test_validate_deployment.py LoRA case also expected the old "TP x SP shape" message - fixed in 5ca1a571f.

[vasilevklart]

Heads-up: #1124 turns this into group.world_size // sp, which counts CFG ranks once this lands, and still calls the removed getters. Whichever merges second, keep it TP x SP proabaly, or TP alone when positions go over SP

[dghosef]

Makes sense - whichever lands second, I'll keep it TP x SP (TP alone when positions go over SP) in #1124.