`diffsynth.utils.xfuser.__init__` missing re-export of `get_sequence_parallel_rank` and `get_sp_group`

[calibrator2026]

Summary

diffsynth/utils/xfuser/__init__.py re-exports get_sequence_parallel_world_size but not the companion APIs get_sequence_parallel_rank and get_sp_group, even though all three are imported together from xfuser.core.distributed in xdit_context_parallel.py:5-7.

Problem

In my use case — fusing proxypose 6-DoF tracking with USP (Ulysses Sequence Parallel) on Wan2.1-T2V-14B — I need both missing APIs inside a patched WanModel.forward that combines proxypose's per-token timestep modulation with USP's sequence chunking.

Where: patched_usp_forward.py:132,141,150,179 (a monkey-patch installed onto pipe.dit.forward)

What each function does in this code:

1. get_sequence_parallel_rank() — after USP splits the concatenated [B, 2N, D] sequence (proxypose's y+x tokens) into sp_size chunks, determines which chunk this rank owns, so the correct slice of both x and its aligned t_mod (per-token timestep modulation, shape [B, 2N, 6, D]) is selected:

   sp_size = get_sequence_parallel_world_size()
   sp_rank = get_sequence_parallel_rank()
   
   x_chunks = torch.chunk(x, sp_size, dim=1)
   x = x_chunks[sp_rank]                            # this rank's x slice
   
   t_mod_chunks = torch.chunk(t_mod, sp_size, dim=1)
   t_mod = t_mod_chunks[sp_rank]                    # aligned t_mod slice — y/x boundary info preserved

   The sp_rank index is what keeps x and t_mod aligned per-rank — without it, each rank can't pick its own slice.

2. get_sp_group() — after the transformer blocks + head, all-gathers the per-rank chunks back into the full sequence:

   x = self.head(x, t)
   x = get_sp_group().all_gather(x, dim=1)          # [B, 2N/sp, D_out] → [B, 2N, D_out]

   The all-gather must be on the SP group specifically (not the whole world group) so only SP-participating ranks combine.

Expected import to work:

from diffsynth.utils.xfuser import get_sequence_parallel_rank, get_sp_group
# ImportError: cannot import name 'get_sequence_parallel_rank'

Workaround I'm currently using (couples my code to xfuser.core.distributed internals):

from xfuser.core.distributed import get_sequence_parallel_rank, get_sp_group

This works functionally — in fact, diffsynth/utils/xfuser/xdit_context_parallel.py:5-7 does exactly this import internally — but it bypasses the diffsynth.utils.xfuser abstraction layer. If xfuser ever refactors its internal module paths, code using the workaround and code going through the diffsynth.utils.xfuser abstraction will break together.

Evidence that this is unintentional

xdit_context_parallel.py:5-7 imports all three together:

from xfuser.core.distributed import (
    get_sequence_parallel_rank,
    get_sequence_parallel_world_size,   # ← this one IS re-exported by __init__.py
    get_sp_group,
)

__init__.py re-exports world_size but hides rank and group from the same import statement. Re-exporting world_size while hiding rank and group has no design rationale — they're a coherent set for querying USP distributed state (you almost always need rank and size together, and group is the handle for any collective op like all_gather).

Suggested fix

One-line addition to diffsynth/utils/xfuser/__init__.py:

from .xdit_context_parallel import (
    usp_attn_forward,
    usp_dit_forward,
    usp_vace_forward,
    get_sequence_parallel_world_size,
    get_sequence_parallel_rank,    # ← add
    get_sp_group,                  # ← add
    initialize_usp,
    get_current_chunk,
    gather_all_chunks,
)

No logic change — xdit_context_parallel.py already imports these names; __init__.py just needs to re-export them.

Environment

• diffsynth version: 2.0.16
• xfuser version: 0.4.5
• Model: Wan-AI/Wan2.1-T2V-14B
• Hardware: 2× RTX 4090
• Use case: proxypose 6-DoF tracking + USP fusion inference

Context

• Searched existing issues/PRs for xfuser, get_sequence_parallel_rank, get_sp_group — 14 results, none report this re-export gap.
• The fusion this enables (proxypose per-token timestep + USP sequence parallel) runs end-to-end at 14.04 s/step on dual 4090 with the workaround in place.