[Bugfix]: SP attention not enabling when _sp_plan hooks are not applied (#1704)

Signed-off-by: Didan Deng <33117903+wtomin@users.noreply.github.com>
Co-authored-by: Canlin Guo <961750412@qq.com>

tests/diffusion/attention/test_attention_sp.py

@@ -1,5 +1,29 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""Tests for Ulysses + Ring sequence-parallel attention correctness.
+
+What is tested
+--------------
+* ``test_sequence_parallel`` verifies that the ``Attention`` layer produces
+  numerically equivalent results whether the sequence is processed on a single
+  rank (baseline, SP=1) or sharded across multiple ranks via Ulysses/Ring SP.
+  The test spawns two separate multi-process runs with ``torch.multiprocessing.spawn``:
+  1. Baseline   – world_size=1, ulysses_degree=1, ring_degree=1.
+  2. SP run     – world_size=ulysses_degree*ring_degree, each rank holds a
+                  contiguous slice of the full sequence.
+  After both runs, rank-0 output tensors are compared element-wise with a
+  tolerance appropriate for the dtype (bfloat16).
+
+SP-plan hooks are NOT applied in this test
+------------------------------------------
+``ForwardContext.sp_plan_hooks_applied`` remains ``False``.  As a result,
+``ForwardContext.sp_active`` falls back to the "no hooks" branch: SP is
+considered active whenever ``parallel_config.sequence_parallel_size > 1``.
+This makes the test self-contained and suitable for standalone CI runs that
+do not exercise the full model-registry pipeline.
+"""
+
 import os
 import pickle
 import tempfile
@@ -27,6 +51,11 @@ def update_environment_variables(envs_dict: dict[str, str]):
         os.environ[k] = v
 
 
+def seed_everything(seed: int):
+    torch.manual_seed(seed)
+    current_omni_platform.manual_seed(seed)
+
+
 class TestAttentionModel(torch.nn.Module):
     """Test model using Attention layer."""
 
@@ -138,11 +167,10 @@ class TestMultiLayerAttentionModel(torch.nn.Module):
 @pytest.mark.parametrize("num_heads", [8])
 @pytest.mark.parametrize("head_size", [8])
 @pytest.mark.parametrize("causal", [False])
-@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])  # [torch.float16, torch.bfloat16]
+@pytest.mark.parametrize("dtype", [torch.bfloat16])
 @pytest.mark.parametrize("use_sync", [False])
 @pytest.mark.parametrize("dynamic", [False])
 @pytest.mark.parametrize("use_compile", [False])
-@pytest.mark.parametrize("attn_backend", ["sdpa", "flash_attn"])
 def test_sequence_parallel(
     ulysses_degree: int,
     ring_degree: int,
@@ -156,7 +184,6 @@ def test_sequence_parallel(
     seq_len: int,
     num_heads: int,
     head_size: int,
-    attn_backend: str,
 ):
     """Test Ulysses attention by comparing with and without SP enabled."""
     sequence_parallel_size = ulysses_degree * ring_degree
@@ -200,7 +227,6 @@ def test_sequence_parallel(
                 model_state_file,
                 input_data_file,
                 True,  # is_baseline
-                attn_backend,
             ),
             nprocs=1,
         )
@@ -228,7 +254,6 @@ def test_sequence_parallel(
                 model_state_file,
                 input_data_file,
                 False,  # is_baseline
-                attn_backend,
             ),
             nprocs=sequence_parallel_size,
         )
@@ -324,12 +349,11 @@ def ulysses_attention_on_test_model(
     model_state_file: str,
     input_data_file: str,
     is_baseline: bool,
-    attn_backend: str,
 ):
     """Run Ulysses attention test on a test model and save results for comparison."""
     # Use fixed seed for reproducibility across baseline and SP runs
     RANDOM_SEED = 42
-    current_omni_platform.seed_everything(RANDOM_SEED)
+    seed_everything(RANDOM_SEED)
 
     mode_str = "Baseline (no SP)" if is_baseline else f"SP (ulysses={ulysses_degree}, ring={ring_degree})"
     print(f"\n[{mode_str}] Rank {local_rank}/{world_size} - Random seed set to {RANDOM_SEED}")
@@ -366,7 +390,6 @@ def ulysses_attention_on_test_model(
         model="test_model",
         dtype=dtype,
         parallel_config=parallel_config,
-        attention_backend=attn_backend,  # Set the attention backend here
     )
 
     # Initialize model parallel
@@ -411,11 +434,6 @@ def ulysses_attention_on_test_model(
             with open(model_state_file, "wb") as f:
                 pickle.dump(model_state, f)
 
-            # Generate and save full input data with fixed seed
-            # Reinitialize RNG to ensure reproducibility
-            torch.manual_seed(42)
-            if torch.cuda.is_available():
-                torch.cuda.manual_seed_all(42)
             full_hidden_states = torch.randn(
                 (batch_size, seq_len, hidden_size),
                 dtype=dtype,

vllm_omni/diffusion/forward_context.py

@@ -31,15 +31,35 @@ class ForwardContext:
     # Original sequence length before padding (for removing padding in gather)
     sp_original_seq_len: int | None = None
 
-    # SP active scope tracking
+    # Set by registry when _sp_plan hooks are applied.
+    # When True, sp_active is determined by _sp_shard_depth (for _sp_plan hooks)
+    # When False, sp_active defaults to True when sequence_parallel_size > 1 (for manual SP, standalone tests, etc.)
+    sp_plan_hooks_applied: bool = False
+    # SP active scope tracking within the _sp_plan hook mechanism.
     # Tracks the depth of SP sharding - incremented on shard, decremented on gather
     # Used by attention layers to determine if SP communication should be enabled
     _sp_shard_depth: int = 0
 
     @property
     def sp_active(self) -> bool:
-        """Returns True when inside an SP sharded region (between shard and gather)."""
-        return self._sp_shard_depth > 0
+        """Returns True when SP attention parallelism should be enabled.
+
+        - If _sp_plan hooks are applied: use _sp_shard_depth (0 = outside sharded region).
+        - If _sp_plan hooks are NOT applied: default to True when sequence_parallel_size > 1,
+          since _sp_shard_depth is only meaningful within the _sp_plan hook mechanism.
+        """
+        if self.sp_plan_hooks_applied:
+            return self._sp_shard_depth > 0
+        # No _sp_plan: assume SP active when configured (manual SP, standalone tests)
+        if self.omni_diffusion_config is None:
+            raise ValueError(
+                "omni_diffusion_config is not set when checking sp_active! "
+                "This usually means set_forward_context() was not called. "
+                "Please call with set_forward_context(omni_diffusion_config=...)."
+            )
+
+        sp_size = self.omni_diffusion_config.parallel_config.sequence_parallel_size
+        return sp_size is not None and sp_size > 1
 
     def __post_init__(self):
         pass

vllm_omni/diffusion/registry.py

@@ -10,6 +10,7 @@ from vllm.model_executor.models.registry import _LazyRegisteredModel, _ModelRegi
 from vllm_omni.diffusion.data import OmniDiffusionConfig
 from vllm_omni.diffusion.distributed.autoencoders.distributed_vae_executor import DistributedVaeMixin
 from vllm_omni.diffusion.distributed.sp_plan import SequenceParallelConfig, get_sp_plan_from_model
+from vllm_omni.diffusion.forward_context import get_forward_context
 from vllm_omni.diffusion.hooks.sequence_parallel import apply_sequence_parallel
 
 logger = init_logger(__name__)
@@ -264,6 +265,11 @@ def _apply_sequence_parallel_if_enabled(model, od_config: OmniDiffusionConfig) -
             apply_sequence_parallel(transformer, sp_config, plan)
             applied_count += 1
 
+        # update forward context sp_plan_hooks_applied
+        ctx = get_forward_context()
+        ctx.sp_plan_hooks_applied = applied_count > 0
+        logger.debug(f"Setting sp_plan_hooks_applied={ctx.sp_plan_hooks_applied} in ``ForwardContext``!")
+
         if applied_count == 0:
             logger.warning(
                 f"Sequence parallelism is enabled (sp_size={sp_size}) but no transformer with _sp_plan found. "