[perf] Harden Ulysses agreement and tune bounded H3 transfers

[shaoxiongduan]

Problem and behavior

Rank-local construction branches and cached call verdicts could leave peers in different collectives and hang Ulysses. The helper now agrees on construction capability and hostname collectively, then exchanges the complete current call contract on every call using reusable CPU vote buffers.

For contiguous BF16 H3 operands on four GB200 GPUs (56 global heads, D128), the helper additionally uses a 144-CTA launch and transfers one batch plane at a time. Packed QKV/QKVG at 250k tokens can use a 896,000,000-byte registered window instead of exceeding the 1 GiB cap. Outputs retain independent storage, and backward reuses the plan saved by its forward.

Implementation

• Expand the call contract to 12 fields, including chunking and launch size. Rank disagreement declines collectively and can recover on a later compatible call.
• Reserve one NCCL barrier slot per CTA for the maximum 144-CTA launch. Preserve the original 36-CTA native entrypoint and compatibility with older kernel wheels.
• Limit automatic tuning to the measured GB200/SP4/BF16/H3 geometry. Unsupported layouts, topologies, capture and compiled regions retain existing fallback behavior.
• Keep the original policy for grad-tracked planes larger than 512 MiB by default. FASTVIDEO_ULYSSES_A2A_LONG_TRAINING=chunked enables long-training chunks after the activation-memory budget has been validated. No-grad inference uses tuned chunks directly.
• Add policy, backward-plan, native execution, output-ownership and rank-recovery tests, plus usage documentation.

Review and validation

Reviewed head: 1568d3d1f7f2399dafa236df2341beaa985367dc. No blocking production-code findings remain. The review fixed two test gaps: numerical parity could pass without the tuned path running, and default-policy tests depended on the caller's environment. The native gate now asserts actual forward/backward execution plans; a negative control deliberately disables tuning and must fail those assertions.

Fresh validation on four GB200 GPUs, PyTorch 2.12.0+cu130, NCCL 2.29.7, Slurm 6629:

• 37 CPU tests passed with each ambient policy setting (auto and chunked; three GPU pytest cases deliberately deselected for separate execution).
• The SP4 native gate passed at 32k, 128k and 250k, including exact gradients, actual 36/144-CTA helper plans, long-training opt-in, owned outputs and rank-disagreement recovery.
• SP2 native forward/backward parity and 100-round stress passed.
• The untuned negative control was rejected as expected.
• Pre-commit passed; repository-configured exclusions for test/kernel paths were respected.

The production Python and CUDA code is byte-identical to the full-model benchmark source 89836e8a98aad2b2b0c173d085950fbbb8602685; the review commit changes tests only. Native library SHA-256: ccc4b3f3d6610bfb29b14aca41eddafdaaa225dcc679babfd02c49e98a89ad79.

Earlier validation retained from this PR: SP2/SP4 construction, capture, lifecycle and recovery coverage (55 passed, one eight-GPU skip, plus six final capture/recovery cases; Slurm 6464/6468). Six selected lifecycle/fault scenarios were also rerun against the new native build in Slurm 6541. Multi-node topology rejection was fault-injected; multi-node performance was not measured.

H3 full-model performance

Complete DiT measurements on four GB200 GPUs, SP4/FSDP4, FP32 masters, mixed BF16 compute, full activation checkpointing, native VSA at 90% sparsity and resident AdamW. Synthetic inputs use real H3 packing and SFT loss; LR and weight decay are zero to keep weights identical across routes. Reported times are medians of three rotated rank-max wall-time trials. Baseline is the corrected PR helper at a06e63827, using its original 36-CTA path.

Packed tokens	Inference baseline → tuned	Training baseline → tuned	Training throughput gain
32k	0.6388 → 0.6115 s	3.2176 → 3.1034 s	+3.68%
64k	1.2895 → 1.2373 s	6.4591 → 6.2406 s	+3.50%
128k	3.0901 → 2.8997 s	17.3864 → 16.5507 s	+5.05%

At 250k, no-grad inference improved 7.4343 → 7.0460 s (+5.51% throughput). Training exposed activation-allocation pressure, so automatic long-training tuning remains conservative. A separate H3 integration offloads saved activations to pinned CPU memory at 200k packed tokens; with expandable allocator segments, its comparison was:

250k training treatment	Step time	Allocator retries per trial
Original PR helper	71.366 s	14 / 13 / 15
Original helper + activation offload	46.906 s	1 / 0 / 1
Tuned chunks + activation offload	42.693 s	0 / 0 / 0

The combined +67.16% training-throughput result requires the separate H3 activation-offload integration, which is outside this public PR. Tuned transport adds 9.87% throughput relative to the offload-only control. Merely enabling long-training chunks does not establish that combined gain.

All full-model outputs and actual fused transports matched NCCL exactly. Whole-model gradient variation passed a threshold calibrated from two repeated NCCL controls; gradients were not bitwise identical. There are 54 timing records and 26 parity records across the short/long phases (Slurm 6546/6547). These are relative step-throughput measurements, not MFU, full PDD rollout results, compiled-training results or complete video-generation times.

Reproduction and evidence

After rebuilding fastvideo-kernel with NCCL 2.29+ device headers:

FASTVIDEO_ULYSSES_A2A=auto torchrun --standalone --nproc_per_node=4 \
  fastvideo/tests/distributed/check_ulysses_h3_native.py

FASTVIDEO_ULYSSES_A2A=auto pytest -v -rA \
  fastvideo/tests/distributed/test_ulysses_h3_policy.py \
  fastvideo/tests/distributed/test_ulysses_a2a_parity.py \
  fastvideo/tests/distributed/test_ulysses_fault_injection.py

Runtime receipts, pinned source archives, configs, raw metrics, traces and checksums are retained under:

• /mnt/lustre/vlm-wlsaidhi/fastvideo/ulysses_pr1807_review_20260907 — final review validation, run 20260907T014553Z-ulysses-pr1807-review.
• /mnt/lustre/vlm-wlsaidhi/fastvideo/h3_ulysses_production_20260906 — full-model results, methods, plot and source equivalence proof.
• /mnt/lustre/vlm-wlsaidhi/fastvideo/ulysses_pr1807_20260906 — original agreement/fault tests and collective microbenchmarks.

[mergify]

Merge Protections

🔴 1 of 1 protections blocking · waiting on 👀 reviews and 🤖 CI

	Protection	Waiting on
🔴	PR merge requirements	👀 reviews and 🤖 CI

🔴 PR merge requirements

Waiting for

• #approved-reviews-by>=1
• check-success=full-suite-passed

This rule is failing.

• #approved-reviews-by>=1
• check-success=full-suite-passed
• check-success=fastcheck-passed
• check-success~=pre-commit
• title~=(?i)^\[(feat|feature|bugfix|fix|refactor|perf|ci|doc|docs|misc|chore|kernel|new.?model|skill|skills|infra)\]