From 7543f2ff3bb0feaad1bc736cb4d7e20551d2c7b1 Mon Sep 17 00:00:00 2001 From: Didan Deng <33117903+wtomin@users.noreply.github.com> Date: Wed, 11 Mar 2026 10:49:42 +0800 Subject: [PATCH] [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> --- .../diffusion/attention/test_attention_sp.py | 44 +++++++++++++------ vllm_omni/diffusion/forward_context.py | 26 +++++++++-- vllm_omni/diffusion/registry.py | 6 +++ 3 files changed, 60 insertions(+), 16 deletions(-) diff --git a/tests/diffusion/attention/test_attention_sp.py b/tests/diffusion/attention/test_attention_sp.py index 234ac1563..c99ba9772 100644 --- a/tests/diffusion/attention/test_attention_sp.py +++ b/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, diff --git a/vllm_omni/diffusion/forward_context.py b/vllm_omni/diffusion/forward_context.py index d2e68c8be..db574db1d 100644 --- a/vllm_omni/diffusion/forward_context.py +++ b/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 diff --git a/vllm_omni/diffusion/registry.py b/vllm_omni/diffusion/registry.py index fe874c9e2..1b48c9c7d 100644 --- a/vllm_omni/diffusion/registry.py +++ b/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. "