Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Inline comments:
In `@tensorrt_llm/_torch/distributed/ops.py`:
- Around line 780-782: The init assigns self.tp_group and self.pg only inside a
narrow branch, causing forward() to sometimes hit AttributeError;
unconditionally cache both group attributes during construction by always
setting self.tp_group = self.mapping.tp_group and self.pg =
getattr(self.mapping, "tp_group_pg", None) (or assigning mapping.tp_group_pg
without condition) so forward() can safely reference them; retain existing logic
that checks self._disable_mpi and self.mnnvl_allreduce when choosing all-reduce
behavior, and apply the same unconditional assignment fix for the other
constructor section covering the code around the forward-related logic (the
region referenced by 880-918) so both branches always define tp_group and pg.

In `@tensorrt_llm/_torch/visual_gen/config.py`:
- Around line 145-146: The n_workers method currently returns only dit_cfg_size
* dit_ulysses_size * dit_tp_size and omits attention2D/ring (CP) factors; update
n_workers to align with total_parallel_size (which already includes
attn2d_row_size/attn2d_col_size or other CP dimensions) by returning or
delegating to total_parallel_size so supported configs like attn2d_row_size=2,
attn2d_col_size=2 are counted correctly.

In `@tensorrt_llm/_torch/visual_gen/mapping.py`:
- Around line 141-147: The dims tuple currently appends a trailing "cp" even
though self._dim_names already contains "cp", producing duplicate names and
breaking device mesh lookups; change the construction of dims and the
corresponding shape before calling init_device_mesh (used when setting
cls.device_mesh) to avoid adding a duplicate "cp" — e.g., build dims as ("pp",)
+ tuple(self._dim_names) and ensure shape mirrors dims (use (1,) +
tuple(self._dim_sizes[d] for d in self._dim_names) + (1,) only if the trailing
wrapper dim is actually not already in self._dim_names), or conditionally append
the wrapper "cp" and its size only when it is missing from self._dim_names so
init_device_mesh receives matching names and sizes.

---

Outside diff comments:
In `@tensorrt_llm/_torch/visual_gen/models/wan/transformer_wan.py`:
- Around line 289-299: The current change constructs Attention with
qkv_mode=QKVMode.SEPARATE_QKV which prevents Attention from wrapping
sequence-parallel backends (Ulysses/Attention2D), causing each rank to only
attend its local chunk when WanTransformer3DModel shards sequence; modify the
Attention construction so it still enables/wraps sequence-parallel backends when
the model is configured for sequence sharding: either remove the conditional
that skips those backends for QKVMode.SEPARATE_QKV in Attention, or add an
explicit flag (e.g., enable_sequence_parallel=True) when creating Attention in
WanTransformer3DModel so that Attention will use Ulysses/Attention2D whenever
ulysses_size>1 (or model_config indicates sequence-parallel) regardless of
qkv_mode.

In `@tensorrt_llm/_torch/visual_gen/pipeline_loader.py`:
- Around line 120-130: _pipeline_loader's _setup_visual_gen_mapping uses local
variables ws/rk only set inside the args branch, causing UnboundLocalError when
args is None on a distributed run; initialize ws and rk (or derive from dist if
initialized) before the args conditional so the fallback path can read them.
Update PipelineLoader._setup_visual_gen_mapping to set ws and rk to sensible
defaults at top (e.g., ws=1, rk=0) and, if dist.is_initialized(), override from
torch.distributed (or compute from args when present) before creating
VisualGenMapping and calling init_pg; ensure you reference VisualGenMapping,
config.visual_gen_mapping, config.mapping, and init_pg in the fix so those uses
no longer read uninitialized ws/rk.