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),
	)