Add fast-ulysses attention backend

[Fatemanx]

Summary

• Add fast_ulysses as a sequence-parallel attention backend that wraps fast-ulysses A2A around the existing LightX2V attention module.
• Vendor the A2A-only fast-ulysses native package with attribution to https://github.com/triple-mu/fast-ulysses.
• Fall back to the existing ulysses backend for unsupported paths, native init failures, and cross-node sequence-parallel groups.

Scope

This PR only integrates fast-ulysses for A2A optimization. It intentionally does not include upstream fused QK/RoPE/RMSNorm ops.

Performance

Measured on one H800 node with 4 GPUs , seq_p_size=4, cfg_p_size=1, bf16, and flash_attn3. The fast run used the native A2A path without fallback.

Scope	Shape / metric	ulysses	fast_ulysses	Speedup
Ulysses attention layer	N=27280,H=40,D=128, warmup=10, iters=50	7518.400 us	6762.496 us	1.112x
Wan I2V DiT steady state	average step time excluding first step	1.601850 s	1.540803 s	1.040x

Full single-run e2e wall time is not used as the primary claim because model loading, cache state, and first-step native initialization add unrelated noise.

Validation

• python test_cases/test_fast_ulysses_backend.py
• git diff --cached --check
• Manual distributed A2A correctness against PyTorch all-to-all on 4 GPUs
• Manual distributed attention parity against existing ulysses
• Manual Wan I2V steady-state performance validation with Wan2.1-I2V-14B-480P

[gemini-code-assist]

Code Review

This pull request introduces an optional NVSHMEM-backed sequence-parallel attention backend called fast_ulysses as a native extension (lightx2v_fast_ulysses) with fallback mechanisms to the standard ulysses backend. The review feedback highlights two critical issues: a potential hang risk in subgroup or multi-group settings due to global NVSHMEM initialization on dist.group.WORLD instead of the local process group pg, and a deadlock risk during garbage collection if collective cleanup operations are implicitly triggered inside the C++ destructor.

Important

The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.

[gemini-code-assist]

Critical Hang Risk in Subgroup / Multi-Group Settings

Currently, NVSHMEM is initialized globally using dist.group.WORLD and global ranks. If fast_ulysses is used in a subgroup setting (e.g., sequence parallel size < world size) or if some ranks do not participate in sequence parallelism (e.g., pipeline or tensor parallel ranks), they will never instantiate UlyssesGroup and thus never call this constructor. This will cause the global broadcast on dist.group.WORLD to hang indefinitely.

To resolve this, NVSHMEM should be initialized and scoped exactly to the sequence parallel subgroup pg. Since all ranks in pg participate in the sequence parallel attention, they will collectively call this constructor at the same time, avoiding any hangs.

This requires:

1. Broadcasting the unique ID within pg using the global rank of local rank 0 in pg as the source.
2. Initializing NVSHMEM with the local rank and size of pg.
3. Passing list(range(self.world_size)) as the PE indices to UlyssesGroup since the NVSHMEM world is now exactly pg.

Suggested change

	if dist.get_rank() == 0:
	uid = cls.get_uniqueid()
	else:
	uid = [0] * cls.uniqueid_nints()
	uid_t = torch.tensor(uid, dtype=torch.int64, device=device)
	dist.broadcast(uid_t, src=0, group=dist.group.WORLD)
	cls.init_world(uid_t.tolist(), dist.get_rank(), dist.get_world_size())
	dist.barrier(group=pg)
	self._group = cls(
	[int(r) for r in self.peer_global_ranks],
	int(self.rank),
	int(device.index),
	int(initial_pool_bytes),
	)
	pg_rank = dist.get_rank(pg)
	if pg_rank == 0:
	uid = cls.get_uniqueid()
	else:
	uid = [0] * cls.uniqueid_nints()
	uid_t = torch.tensor(uid, dtype=torch.int64, device=device)
	src_global_rank = dist.get_global_rank(pg, 0)
	dist.broadcast(uid_t, src=src_global_rank, group=pg)
	cls.init_world(uid_t.tolist(), pg_rank, dist.get_world_size(pg))
	dist.barrier(group=pg)
	self._group = cls(
	list(range(self.world_size)),
	int(pg_rank),
	int(device.index),
	int(initial_pool_bytes),
	)

[STwangyingrui]

Thank you for providing the initial benchmark against the legacy ulysses backend. Could you also compare fast_ulysses with the latest ulysses-opt under the same H800 setup, tensor shapes, and attention configuration?

{
    "seq_p_attn_type": "ulysses-opt",
    "seq_p_tensor_fusion": true
}

ulysses-opt independently optimizes the communication pre/post layouts while retaining the standard PyTorch/NCCL transport. This comparison would help clarify how much of the improvement comes from layout optimization and whether the native A2A implementation provides additional performance benefits.

[Fatemanx]

Thanks @STwangyingrui. I re-ran the comparison after merging latest upstream/main, including the new ulysses-opt backend.

Test environment:

• Branch: feat/fast-ulysses-backend
• Base after merge: includes upstream commit 438dd064 Add optimized Ulysses attention path (#1235)
• GPU: 4x H800
• Distributed setup: torch.distributed.run --nproc_per_node=4
• Model: Wan2.1-I2V-14B-480P
• Task: Wan2.1 I2V
• Output setting: target_video_length=81, target_height=480, target_width=832
• Precision: BF16
• Parallel config: seq_p_size=4, cfg_p_size=1
• Input image: assets/inputs/imgs/img_0.jpg
• Prompt/negative prompt: same as scripts/dist_infer/run_wan_i2v_dist_cfg_ulysses.sh

Backends compared:

• ulysses
• ulysses-opt with seq_p_tensor_fusion=true
• fast_ulysses

Attention microbenchmark:

• Local sequence length per rank: 6820
• Global sequence length: 27280
• Heads: 40
• Head dim: 128
• Attention backend: flash_attn3
• Warmup / iterations: 10 / 50
• Correctness check: all backends matched, max_abs_diff=0.000000

Backend	Median us	P90 us	Min us	Correct
ulysses	7592.608	7693.623	7558.784	True
ulysses-opt	7248.064	7306.198	7207.200	True
fast_ulysses	6744.656	6779.843	6706.304	True

Microbenchmark speedup:

• ulysses-opt vs ulysses: 1.0475x
• fast_ulysses vs ulysses: 1.1257x
• fast_ulysses vs ulysses-opt: 1.0746x

E2E Wan2.1 I2V result:

• Metric: per-rank denoising infer_main cost
• Effective samples: 156 values, from 4 ranks x 39 steps
• Each backend successfully generated the output video
• No Traceback, OOM, NCCL error, or fast_ulysses fallback was observed
• fast_ulysses log confirmed native A2A path: fast_ulysses native A2A path active for ranks (0, 1, 2, 3)

Backend	Count	Mean s	Median s	Min s	Max s
ulysses	156	1.598913	1.598196	1.591014	1.619212
ulysses-opt	156	1.573792	1.574747	1.557260	1.580761
fast_ulysses	156	1.538968	1.539087	1.526863	1.554712

E2E speedup by mean:

• ulysses-opt vs ulysses: 1.0160x
• fast_ulysses vs ulysses: 1.0390x
• fast_ulysses vs ulysses-opt: 1.0226x

E2E speedup by median:

• ulysses-opt vs ulysses: 1.0149x
• fast_ulysses vs ulysses: 1.0384x
• fast_ulysses vs ulysses-opt: 1.0232x

Conclusion:
After comparing against the latest ulysses-opt, fast_ulysses still shows a measurable improvement. The gain is larger in the isolated attention benchmark, around 7.46% over ulysses-opt, and smaller but still positive in the full Wan2.1 I2V E2E run, around 2.26% by mean infer_main time. The E2E run also confirms that the native fast path is active and the output video is generated successfully.