[fix] Sync PAN2 padding parity with latest upstream (#1961)

docs-site/navigation.json

@@ -123,6 +123,10 @@
         ],
         "display": "group"
       },
+      {
+        "label": "PAN 2 text/image to video",
+        "href": "/inference/pan2"
+      },
       {
         "label": "Examples",
         "href": "/inference/examples/examples_inference_index"

docs/inference/pan2.md

@@ -1,6 +1,6 @@
-# PAN2 Text/Image to Video
+# PAN 2 Text/Image to Video
 
-PAN2 is a 29B mixture-of-transformers video DiT with separate video and text streams, a Qwen3.5 text encoder and the
+PAN 2 is a 29B mixture-of-transformers video DiT with separate video and text streams, a Qwen3.5 text encoder and the
 Wan2.2 VAE. One pipeline serves text-to-video and image-to-video: pass `image_path` to condition on a first frame and
 leave it out for text only.
 
@@ -21,7 +21,7 @@ from the `pan2_t2v` and `pan2_i2v` presets:
 | Steps | 50, uniform sigmas from 1 to 0 with flow shift 7 |
 | Guidance | 3.0, classifier-free guidance with an empty negative prompt |
 
-PAN2 is trained on the structured captions its prompt rewriter produces, serialized as one-line JSON. Short raw
+PAN 2 is trained on the structured captions its prompt rewriter produces, serialized as one-line JSON. Short raw
 prompts work, with weaker results.
 
 ## Usage
@@ -37,6 +37,7 @@ generator = VideoGenerator.from_pretrained(
     "IFM/PAN2",
     num_gpus=1,
     dit_cpu_offload=False,
+    dit_layerwise_offload=False,
     text_encoder_cpu_offload=True,
 )
 generator.generate_video(
@@ -55,6 +56,27 @@ With `num_gpus > 1`, FastVideo runs Ulysses sequence parallelism: the video toke
 text tokens are replicated. The number of GPUs must divide the 32 attention heads. Sequence-parallel padding of the
 video tokens is trimmed before attention, so no attention mask is needed.
 
+Padding is always excluded, without an opt-in setting. At 704x1248 with 313 frames, the 67,782 video tokens need
+two padding tokens at SP=4. They are removed before joining the text tokens in every attention block. This provides
+the same exclusion as vLLM-Omni's unconditional padding mask in
+[revision `e0c4c731`](https://github.com/vllm-project/vllm-omni/commit/e0c4c73158a5e4f44eda8eaa8879f7fd9cedbd74).
+The trimmed sequence can use either FlashAttention or Torch SDPA.
+
+```bash
+PAN2_NUM_GPUS=4 python examples/inference/basic/basic_pan2.py
+```
+
+## Memory and offloading
+
+The example keeps the DiT on the GPUs and offloads the text encoder between requests. Sequence parallelism shards
+video tokens; each rank still holds the DiT weights. Increasing `num_gpus` alone does not shard those weights.
+For devices that cannot hold the DiT, see [CPU offloading](offloading.md) for `dit_cpu_offload` and
+`dit_layerwise_offload`. Hardware recipes and timings reported by vLLM-Omni use its own offload implementation;
+they have not been validated for FastVideo.
+
+The full `IFM/PAN2` checkpoint is not public yet. Local parity tests use the public random-weight
+`wuqing157/tiny-pan2-modular-pipe` checkpoint; it verifies the integration and produces noise rather than useful video.
+
 ## Differences from the vLLM-Omni integration
 
 - The optional prompt and video guardrails (`pan2-guardrail`) are not included.

examples/inference/basic/basic_pan2.py

@@ -22,12 +22,12 @@ def main():
         num_gpus=int(os.environ.get("PAN2_NUM_GPUS", "1")),
         use_fsdp_inference=False,
         dit_cpu_offload=False,
+        dit_layerwise_offload=False,
         vae_cpu_offload=False,
         text_encoder_cpu_offload=True,
         pin_cpu_memory=True,
     )
 
-    # Defaults come from the pan2 preset: 704x1248, 313 frames (13 s at 24 fps), 50 steps, guidance 3.0.
     generator.generate_video(
         PROMPT,
         image_path=IMAGE_PATH,

tests/local_tests/pan2/PORT_STATUS.md

@@ -3,7 +3,7 @@
 ## Summary
 - model_family: pan2
 - workload_types: T2V, I2V (one pipeline; I2V when the request carries a first frame)
-- official_ref: vllm-project/vllm-omni#8602 at `59dd8c9a` (the diffusers integration is private)
+- official_ref: vllm-project/vllm-omni#8602 at `e0c4c73158a5e4f44eda8eaa8879f7fd9cedbd74` (the diffusers integration is private)
 - official_ref_dir: none; rebuilt in `tests/local_tests/pan2/reference.py`
 - hf_weights_path: `IFM/PAN2` (not public); tests use `wuqing157/tiny-pan2-modular-pipe@3ff3130b`
 - local_weights_dir: HF cache, or `PAN2_TINY_MODEL_DIR`
@@ -14,7 +14,7 @@
 - phase: 11 (post-parity review and handoff)
 - status: in_progress (GPU and real-weight validation pending, see I001 and I002)
 - owner: orchestrator
-- last_updated: 2026-10-15
+- last_updated: 2026-10-20
 
 ## Component Matrix
 | Component | Type | Reuse/Port | Official Definition | Official Instantiation | FastVideo Target | Prototype | Conversion | Parity | Open Issues |
@@ -36,24 +36,24 @@
 ## Parity Commands
 | Scope | Command | Last Result | Notes |
 |---|---|---|---|
-| DiT | `pytest tests/local_tests/transformers/test_pan2_transformer_parity.py` | 12 passed (CPU) | includes SP=2 vs one rank over gloo |
+| DiT | `pytest tests/local_tests/transformers/test_pan2_transformer_parity.py` | 13 passed (CPU) | SP=2/4 vs one rank over gloo, T2V/I2V, one/two pad tokens; max abs 7.7e-6 |
 | Text encoder | `pytest tests/local_tests/encoders/test_pan2_qwen3_5_parity.py` | 5 passed (CPU) | exact |
-| VAE | `pytest tests/local_tests/vaes/test_pan2_vae_parity.py` | passed (CPU) | exact |
-| Scheduler | `pytest tests/local_tests/schedulers/test_pan2_scheduler_parity.py` | passed (CPU) | exact |
+| VAE | `pytest tests/local_tests/vaes/test_pan2_vae_parity.py` | 9 passed (CPU) | exact |
+| Scheduler | `pytest tests/local_tests/schedulers/test_pan2_scheduler_parity.py` | 9 passed (CPU) | exact |
 | Pipeline | `pytest tests/local_tests/pipelines/test_pan2_pipeline_parity.py` | 3 passed (CPU) | T2V and I2V latents exact |
 
 ## Open Questions
 | ID | Question | Owner | Needed By Phase | Status | Resolution |
 |---|---|---|---|---|---|
 | Q001 | Which reference, with the diffusers integration private? | orchestrator | 2 | resolved | vLLM-Omni #8602, rebuilt from diffusers/transformers blocks in `reference.py` |
-| Q002 | Port the `pan2-guardrail` prompt/video guardrails? | orchestrator | 7 | resolved | Not ported: the package is not public and FastVideo has no guardrail stage. |
+| Q002 | Port the `pan2-guardrail` prompt/video guardrails? | orchestrator | 7 | resolved | Not ported: FastVideo has no guardrail stage. |
 | Q003 | Port Cache-DiT `quality` and CFG-parallel? | orchestrator | 7 | resolved | Not ported: vLLM-Omni serving features. Ulysses SP is supported through `DistributedAttention`. |
 
 ## Issues And Blockers
 | ID | Phase | Component | Severity | Issue | Evidence | Owner | Status | Resolution |
 |---|---|---|---|---|---|---|---|---|
 | I001 | 9-10 | all | medium | No real-weight parity or SSIM regression | `IFM/PAN2` is not public | orchestrator | open | Add `fastvideo/tests/ssim/test_pan2_similarity.py` once weights are public. |
-| I002 | 9 | transformer, text_encoder | medium | No GPU run yet (bf16 + FlashAttention, NCCL SP, 9B encoder speed) | All tests ran on a CPU login node | orchestrator | open | Run the same tests on a GPU node. |
+| I002 | 9 | transformer, text_encoder | medium | No GPU parity run yet (bf16 + FlashAttention, NCCL SP, 9B encoder speed) | Local parity tests ran on a CPU login node | orchestrator | open | Run the same tests on a GPU node. |
 | I003 | 7 | pipeline | low | The shared denoiser casts latents to the model-output dtype in the scheduler step; PAN2 samples on fp32 latents | pipeline parity fails (8e-2) without the fix | orchestrator | resolved | `DenoisingStage.prepare_scheduler_model_output` hook; `PAN2DenoisingStage` returns `noise_pred.float()`. |
 
 ## Escape Hatches

tests/local_tests/pan2/README.md

@@ -7,7 +7,7 @@ not part of the default CI suite.
 
 PAN2's diffusers integration (`PAN2Transformer3DModel`, `PAN2ModularPipeline`) is not public yet, and the full
 `IFM/PAN2` weights are not released. The reference is the vLLM-Omni port,
-[vllm-project/vllm-omni#8602](https://github.com/vllm-project/vllm-omni/pull/8602) (head `59dd8c9a`), which reports
+[vllm-project/vllm-omni#8602](https://github.com/vllm-project/vllm-omni/pull/8602) (head `e0c4c731`), which reports
 parity with the diffusers implementation. `reference.py` rebuilds it from diffusers and transformers building blocks:
 the DiT from `AdaLayerNormZero`, `AdaLayerNormContinuous`, `FeedForward` (gelu-approximate and swiglu), `Timesteps`,
 `TimestepEmbedding` and the diffusers RoPE helpers, using the checkpoint's module names so the checkpoint loads
@@ -36,7 +36,7 @@ when no GPU is present. On a GPU-less host, `import fastvideo` also needs the `f
 |---|---|---|
 | DiT: tiny checkpoint, strict load + forward | `pytest tests/local_tests/transformers/test_pan2_transformer_parity.py -k tiny_checkpoint` | max abs 6.0e-7 |
 | DiT: random weights, batch 1 and 2 | `pytest tests/local_tests/transformers/test_pan2_transformer_parity.py -k random_weights` | max abs 0 and 1.9e-5 (outputs ~16) |
-| DiT: Ulysses SP=2 vs one rank, padded tokens | `pytest tests/local_tests/transformers/test_pan2_transformer_parity.py -k sp2` | max abs 5.7e-6 (gloo on CPU; NCCL on 2 GPUs) |
+| DiT: Ulysses SP=2/4 vs one rank, T2V/I2V with 15/18 tokens | `pytest tests/local_tests/transformers/test_pan2_transformer_parity.py -k sequence_parallel` | max abs 4.8e-6 / 7.7e-6 (gloo on CPU; NCCL on 2/4 GPUs) |
 | Qwen3.5 text encoder | `pytest tests/local_tests/encoders/test_pan2_qwen3_5_parity.py` | all hidden states exact (fp32, bf16, padded batch) |
 | Wan2.2 VAE encode/decode + normalization | `pytest tests/local_tests/vaes/test_pan2_vae_parity.py` | exact; normalization 2.4e-7 |
 | Flow-match scheduler, PAN2 sigmas + shift | `pytest tests/local_tests/schedulers/test_pan2_scheduler_parity.py` | exact |
@@ -48,6 +48,19 @@ confirm the threshold catches real deviations, three injected changes were each
 `hidden_states[-2]`, letting the latents fall to bf16 between steps, and shifting the first-frame crop by one pixel.
 Each moved the latents by about 8e-2 relative L2.
 
+Upstream revision `e0c4c731` always masks video padding, fixing Ulysses=4 drift when padding was left unmasked.
+FastVideo already excludes this padding by passing the original video length to `DistributedAttention`, which trims
+it before RoPE and joint attention, then pads the output for the reverse all-to-all. The distributed test covers one
+pad token (15 tokens at SP=2/4), two pad tokens (18 at SP=4, matching 704x1248x313 modulo 4), and an unpadded
+control (18 at SP=2), with both zero T2V and first-frame I2V conditioning. No masking toggle is required.
+The pipeline math and checkpoint layout are unchanged between upstream revisions `59dd8c9a` and `e0c4c731`.
+
+On 2026-10-20, all 39 PAN2 local tests passed on CPU (Torch 2.12.0, Transformers 5.15.0, Diffusers 0.39.0).
+The broader API/contract suite had 499 passes and two failures also reproduced on clean `main` at `cdda1b753`:
+the schema inventory lacks MiniMax-H3's `audio_scheduler_shift`/`video_scheduler_shift`, and
+`test_load_kandinsky5.py` reads `FASTVIDEO_ATTENTION_BACKEND` outside the environment registry.
+Importing the PAN2 tests and running their fixtures preserves both an unset and a pre-existing attention backend.
+
 ## Not covered
 
 - Real-weight parity and an SSIM regression: `IFM/PAN2` is not public.

tests/local_tests/pan2/reference.py

@@ -6,7 +6,8 @@ assembled from diffusers' and transformers' own building blocks (``AdaLayerNormZ
 ``FeedForward`` with ``gelu-approximate``/``swiglu``, ``Timesteps``/``TimestepEmbedding``,
 ``get_1d_rotary_pos_embed``/``apply_rotary_emb``, ``Qwen3_5TextModel``, ``AutoencoderKLWan``,
 ``FlowMatchEulerDiscreteScheduler``) and follows the vLLM-Omni PAN2 port (vllm-project/vllm-omni#8602), which was
-validated against the diffusers implementation. Module names are the checkpoint's, so checkpoints load strictly.
+validated against the diffusers implementation (upstream revision ``e0c4c731``). Module names are the checkpoint's,
+so checkpoints load strictly. Its unconditional SP padding mask is equivalent to unpadded reference attention.
 """
 
 from __future__ import annotations

tests/local_tests/pipelines/test_pan2_pipeline_parity.py

@@ -12,16 +12,15 @@ Runs on CPU when no GPU is present. Set ``PAN2_TINY_MODEL_DIR`` to a local copy
 
 from __future__ import annotations
 
-import os
 import socket
+from contextlib import ExitStack
 
 import PIL.Image
 import pytest
 import torch
 
-os.environ.setdefault("FASTVIDEO_ATTENTION_BACKEND", "TORCH_SDPA")
-
-from tests.local_tests.pan2.reference import load_reference_components, reference_generate, tiny_model_dir  # noqa: E402
+from fastvideo import envs
+from tests.local_tests.pan2.reference import load_reference_components, reference_generate, tiny_model_dir
 
 HEIGHT, WIDTH, NUM_FRAMES = 64, 96, 9
 NUM_INFERENCE_STEPS = 3
@@ -51,37 +50,35 @@ def pan2_runtime():
     from fastvideo.fastvideo_args import FastVideoArgs
     from fastvideo.pipelines import build_pipeline
 
-    monkeypatch = pytest.MonkeyPatch()
-    if not torch.cuda.is_available():
-        # On a GPU-less Linux host fastvideo resolves the local device to mps; run on CPU instead. Swapping the code
-        # object reaches every module that imported the function by name.
-        def _cpu_device() -> torch.device:
-            return torch.device("cpu")
+    with ExitStack() as stack:
+        stack.enter_context(envs.FASTVIDEO_ATTENTION_BACKEND.override("TORCH_SDPA"))
+        monkeypatch = stack.enter_context(pytest.MonkeyPatch.context())
+        if not torch.cuda.is_available():
+            # Reach modules that imported the device helper by name on a GPU-less Linux host.
+            def _cpu_device() -> torch.device:
+                return torch.device("cpu")
 
-        monkeypatch.setattr(parallel_state.get_local_torch_device, "__code__", _cpu_device.__code__)
-    for name, value in (("MASTER_ADDR", "127.0.0.1"), ("MASTER_PORT", str(_free_port())), ("RANK", "0"),
-                        ("WORLD_SIZE", "1"), ("LOCAL_RANK", "0")):
-        monkeypatch.setenv(name, os.environ.get(name, value))
-    maybe_init_distributed_environment_and_model_parallel(1, 1)
+            monkeypatch.setattr(parallel_state.get_local_torch_device, "__code__", _cpu_device.__code__)
+        for name, value in (("MASTER_ADDR", "127.0.0.1"), ("MASTER_PORT", str(_free_port())), ("RANK", "0"),
+                            ("WORLD_SIZE", "1"), ("LOCAL_RANK", "0")):
+            stack.enter_context(envs.override_external(name, value))
+        stack.callback(cleanup_dist_env_and_memory)
+        maybe_init_distributed_environment_and_model_parallel(1, 1)
 
-    pipeline_config = PAN2PipelineConfig(dit_precision="bf16",
-                                         vae_precision="fp32",
-                                         text_encoder_precisions=("fp32", ))
-    args = FastVideoArgs(model_path=str(model_dir),
-                         num_gpus=1,
-                         pipeline_config=pipeline_config,
-                         dit_cpu_offload=False,
-                         text_encoder_cpu_offload=False,
-                         vae_cpu_offload=False,
-                         use_fsdp_inference=False)
-    pipeline = build_pipeline(args)
-    device = parallel_state.get_local_torch_device()
-    try:
+        pipeline_config = PAN2PipelineConfig(dit_precision="bf16",
+                                             vae_precision="fp32",
+                                             text_encoder_precisions=("fp32", ))
+        args = FastVideoArgs(model_path=str(model_dir),
+                             num_gpus=1,
+                             pipeline_config=pipeline_config,
+                             dit_cpu_offload=False,
+                             dit_layerwise_offload=False,
+                             text_encoder_cpu_offload=False,
+                             vae_cpu_offload=False,
+                             use_fsdp_inference=False)
+        pipeline = build_pipeline(args)
+        device = parallel_state.get_local_torch_device()
         yield pipeline, args, model_dir, device
-    finally:
-        del pipeline
-        cleanup_dist_env_and_memory()
-        monkeypatch.undo()
 
 
 def _fastvideo_generate(pipeline, args, image, guidance_scale: float = GUIDANCE_SCALE):

tests/local_tests/transformers/test_pan2_transformer_parity.py

@@ -17,6 +17,7 @@ import os
 import socket
 import subprocess
 import sys
+from contextlib import ExitStack
 from dataclasses import fields
 from pathlib import Path
 
@@ -25,16 +26,14 @@ import torch
 import torch.nn as nn
 from torch.testing import assert_close
 
-os.environ.setdefault("DISABLE_SP", "1")
-os.environ.setdefault("FASTVIDEO_ATTENTION_BACKEND", "TORCH_SDPA")
-
-from fastvideo.forward_context import set_forward_context  # noqa: E402
-from fastvideo.models.loader.fsdp_load import load_model_from_full_model_state_dict  # noqa: E402
-from fastvideo.models.loader.utils import get_param_names_mapping  # noqa: E402
-from tests.local_tests.pan2.reference import TINY_REPO_ID, ReferencePAN2Transformer, tiny_model_dir  # noqa: E402
+from fastvideo import envs
+from fastvideo.forward_context import set_forward_context
+from fastvideo.models.loader.fsdp_load import load_model_from_full_model_state_dict
+from fastvideo.models.loader.utils import get_param_names_mapping
+from tests.local_tests.pan2.reference import TINY_REPO_ID, ReferencePAN2Transformer, tiny_model_dir
 
 REPO_ROOT = Path(__file__).resolve().parents[3]
-SP_WORLD_SIZE = 2
+SP_WORLD_SIZES = (2, 4)
 
 
 # --------------------------------------------------------------------------------------------------------------------
@@ -69,17 +68,25 @@ def _small_kwargs() -> dict:
 
 
 @pytest.fixture
-def single_rank_dist():
+def single_rank_dist(monkeypatch):
     from fastvideo.distributed import cleanup_dist_env_and_memory, maybe_init_distributed_environment_and_model_parallel
+    import fastvideo.distributed.parallel_state as parallel_state
 
-    os.environ.setdefault("MASTER_ADDR", "127.0.0.1")
-    os.environ.setdefault("MASTER_PORT", str(_free_port()))
-    os.environ.setdefault("RANK", "0")
-    os.environ.setdefault("WORLD_SIZE", "1")
-    os.environ.setdefault("LOCAL_RANK", "0")
-    maybe_init_distributed_environment_and_model_parallel(1, 1)
-    yield
-    cleanup_dist_env_and_memory()
+    if not torch.cuda.is_available():
+        def _cpu_device() -> torch.device:
+            return torch.device("cpu")
+
+        monkeypatch.setattr(parallel_state.get_local_torch_device, "__code__", _cpu_device.__code__)
+    with ExitStack() as stack:
+        stack.enter_context(envs.FASTVIDEO_ATTENTION_BACKEND.override("TORCH_SDPA"))
+        for name, value in (("MASTER_ADDR", "127.0.0.1"), ("MASTER_PORT", str(_free_port())), ("RANK", "0"),
+                            ("WORLD_SIZE", "1"), ("LOCAL_RANK", "0")):
+            stack.enter_context(envs.override_external(name, value))
+        try:
+            maybe_init_distributed_environment_and_model_parallel(1, 1)
+            yield
+        finally:
+            cleanup_dist_env_and_memory()
 
 
 def _build_native(arch_kwargs: dict):
@@ -218,17 +225,16 @@ def test_pan2_random_weights_forward_parity(single_rank_dist, batch_size: int):
 
 
 # --------------------------------------------------------------------------------------------------------------------
-# Sequence parallel: SP=2 must match one rank (video tokens are padded to a multiple of the SP size)
+# Sequence parallel: padded video tokens must never enter joint video/text attention.
 # --------------------------------------------------------------------------------------------------------------------
 
 
-def _run_sp_worker(mode: str, output_path: Path) -> None:
+def _run_sp_worker(sp_size: int, output_path: Path) -> None:
     from fastvideo.configs.models.dits.pan2 import PAN2VideoConfig
     from fastvideo.distributed import cleanup_dist_env_and_memory, maybe_init_distributed_environment_and_model_parallel
     from fastvideo.models.dits.pan2 import PAN2Transformer3DModel
     import torch.distributed as dist
 
-    sp_size = 1 if mode == "single" else SP_WORLD_SIZE
     if torch.cuda.is_available():
         device = torch.device(f"cuda:{int(os.environ.get('LOCAL_RANK', '0'))}")
         torch.cuda.set_device(device)
@@ -254,58 +260,76 @@ def _run_sp_worker(mode: str, output_path: Path) -> None:
         native = PAN2Transformer3DModel(config=config, hf_config=kwargs).float().eval()
         _load_into_native(native, reference.state_dict().items())
         native = native.to(device)
-        # 5x2x6 latents patchify to 15 video tokens, which SP=2 pads to 16.
-        inputs = _make_inputs(kwargs["in_channels"], kwargs["text_embed_dim"])
-        inputs["hidden_states"] = torch.randn(2, kwargs["in_channels"], 5, 2, 6, generator=torch.Generator().manual_seed(5))
-        inputs = {name: value.to(device) for name, value in inputs.items()}
-        with torch.inference_mode(), set_forward_context(current_timestep=0, attn_metadata=None):
-            output = native(**inputs)
+        outputs = {}
+        # 15 tokens need one pad at SP=2/4; 18 need two at SP=4, as at 704x1248x313.
+        for latent_frames in (5, 6):
+            for image_conditioned in (False, True):
+                inputs = _make_inputs(kwargs["in_channels"], kwargs["text_embed_dim"])
+                video = torch.randn(2, kwargs["in_channels"], latent_frames, 2, 6,
+                                    generator=torch.Generator().manual_seed(5))
+                video[:, kwargs["out_channels"]:] = 0
+                if image_conditioned:
+                    video[:, kwargs["out_channels"]:-1, :1] = torch.randn(
+                        2, kwargs["out_channels"], 1, 2, 6, generator=torch.Generator().manual_seed(6))
+                    video[:, -1:, :1] = 1
+                inputs["hidden_states"] = video
+                inputs = {name: value.to(device) for name, value in inputs.items()}
+                with torch.inference_mode(), set_forward_context(current_timestep=0, attn_metadata=None):
+                    output = native(**inputs)
+                outputs[f"{latent_frames}_{image_conditioned}"] = output.detach().cpu()
         if int(os.environ.get("RANK", "0")) == 0:
-            torch.save({"output": output.detach().cpu()}, output_path)
+            torch.save(outputs, output_path)
         dist.barrier()
     finally:
         cleanup_dist_env_and_memory()
 
 
-def _run_torchrun(mode: str, nproc_per_node: int, output_path: Path) -> None:
+def _run_torchrun(nproc_per_node: int, output_path: Path) -> None:
     command = [
         sys.executable, "-m", "torch.distributed.run", "--nnodes", "1", "--nproc_per_node",
         str(nproc_per_node), "--master_port",
         str(_free_port()),
-        str(Path(__file__).resolve()), "--pan2-sp-worker", "--mode", mode, "--output",
+        str(Path(__file__).resolve()), "--pan2-sp-worker", "--sp-size", str(nproc_per_node), "--output",
         str(output_path)
     ]
     environment = os.environ.copy()
     environment["PYTHONPATH"] = os.pathsep.join(filter(None, [str(REPO_ROOT), environment.get("PYTHONPATH")]))
-    environment["DISABLE_SP"] = "0"
     environment["FASTVIDEO_ATTENTION_BACKEND"] = "TORCH_SDPA"
-    process = subprocess.run(command, capture_output=True, text=True, env=environment)
+    process = subprocess.run(command, capture_output=True, text=True, env=environment, timeout=180)
     if process.returncode != 0:
-        raise RuntimeError(f"{mode} worker failed with code {process.returncode}\n"
+        raise RuntimeError(f"SP={nproc_per_node} worker failed with code {process.returncode}\n"
                            f"STDOUT:\n{process.stdout}\nSTDERR:\n{process.stderr}")
 
 
-def test_pan2_sp2_matches_single_rank(tmp_path: Path):
-    """Ulysses SP=2 on an odd video token count (padded to the SP size) must match one rank. Runs on CPU with gloo
-    when no GPU is present."""
-    if torch.cuda.is_available() and torch.cuda.device_count() < SP_WORLD_SIZE:
-        pytest.skip(f"PAN2 SP parity requires {SP_WORLD_SIZE} CUDA devices.")
-    single_path, sp_path = tmp_path / "single.pt", tmp_path / "sp2.pt"
-    _run_torchrun("single", 1, single_path)
-    _run_torchrun("sp", SP_WORLD_SIZE, sp_path)
-    single = torch.load(single_path, map_location="cpu", weights_only=True)["output"]
-    sp = torch.load(sp_path, map_location="cpu", weights_only=True)["output"]
-    diff = (sp - single).abs()
-    print(f"sp2-vs-single diff_max={diff.max().item():.3e}")
-    assert_close(sp, single, atol=1e-5, rtol=1e-5)
+@pytest.fixture(scope="module")
+def single_rank_sp_outputs(tmp_path_factory):
+    path = tmp_path_factory.mktemp("pan2-single") / "single.pt"
+    _run_torchrun(1, path)
+    return torch.load(path, map_location="cpu", weights_only=True)
+
+
+@pytest.mark.parametrize("sp_size", SP_WORLD_SIZES)
+def test_pan2_sequence_parallel_matches_single_rank(tmp_path: Path, single_rank_sp_outputs, sp_size: int):
+    """T2V and I2V joint attention excludes padding at Ulysses SP=2/4. Uses gloo on CPU or NCCL on 2/4 GPUs."""
+    if torch.cuda.is_available() and torch.cuda.device_count() < sp_size:
+        pytest.skip(f"PAN2 SP parity requires {sp_size} CUDA devices.")
+    sp_path = tmp_path / f"sp{sp_size}.pt"
+    _run_torchrun(sp_size, sp_path)
+    sp_outputs = torch.load(sp_path, map_location="cpu", weights_only=True)
+    assert sp_outputs.keys() == single_rank_sp_outputs.keys()
+    for case, single in single_rank_sp_outputs.items():
+        sp = sp_outputs[case]
+        diff = (sp - single).abs()
+        print(f"sp{sp_size}-vs-single {case} diff_max={diff.max().item():.3e}")
+        assert_close(sp, single, atol=1e-5, rtol=1e-5)
 
 
 if __name__ == "__main__":
     parser = argparse.ArgumentParser()
     parser.add_argument("--pan2-sp-worker", action="store_true")
-    parser.add_argument("--mode", choices=["single", "sp"], required=True)
+    parser.add_argument("--sp-size", type=int, choices=[1, *SP_WORLD_SIZES], required=True)
     parser.add_argument("--output", type=str, required=True)
     args = parser.parse_args()
     if not args.pan2_sp_worker:
         raise SystemExit("This module is intended to be run by pytest.")
-    _run_sp_worker(args.mode, Path(args.output))
+    _run_sp_worker(args.sp_size, Path(args.output))