[Diffusion] Non-Intrusive Sequence Parallelism (SP) Model Support Abstraction for vLLM-Omni Framework (#779)

Signed-off-by: mxuax <mxuax@connect.ust.hk> Signed-off-by: XU Mingshi <91017482+mxuax@users.noreply.github.com>
2026-01-21 14:43:09 +08:00
1 parent 334d306d9a
commit 25b08b3371
@@ -0,0 +1,703 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project and the HuggingFace Team.
+# All rights reserved.
+#
+# This module is adapted from HuggingFace diffusers library:
+#   diffusers/src/diffusers/hooks/context_parallel.py
+#
+# NOTE: Our "Sequence Parallelism" (SP) corresponds to "Context Parallelism" (CP) in diffusers.
+# We use the term "Sequence Parallelism" to align with vLLM-Omni's existing terminology.
+#
+# Key adaptations from diffusers:
+#   - ModuleForwardMetadata: parameter lookup logic (adapted)
+#   - SequenceParallelSplitHook/GatherHook: hook structure (adapted from ContextParallel*)
+#   - apply_sequence_parallel: registration logic (adapted from apply_context_parallel)
+#
+# Key differences from diffusers:
+#   - Uses vLLM-Omni's SequenceParallelGroupCoordinator instead of DeviceMesh
+#   - Uses sp_shard/sp_gather from sp_sharding.py instead of funcol operations
+#   - Supports Ulysses + Ring hybrid parallelism
+#
+"""Sequence Parallelism hooks for non-intrusive SP support.
+
+This module implements the hook-based mechanism for applying sequence parallelism
+to models without modifying their forward() methods.
+
+Usage:
+    1. Define _sp_plan on your model class (corresponds to diffusers' _cp_plan)
+    2. Call apply_sequence_parallel(model, config, plan) to enable SP
+    3. Call remove_sequence_parallel(model, plan) to disable SP
+
+The hooks automatically shard inputs before forward and gather outputs after,
+based on the plan specification.
+"""
+
+from __future__ import annotations
+
+import inspect
+from dataclasses import dataclass
+from typing import Any
+
+import torch
+import torch.nn as nn
+from vllm.logger import init_logger
+
+from vllm_omni.diffusion.distributed.sp_plan import (
+    AnySequenceParallelInput,
+    SequenceParallelConfig,
+    SequenceParallelInput,
+    SequenceParallelModelPlan,
+    SequenceParallelOutput,
+    SequenceParallelPartialInput,
+)
+from vllm_omni.diffusion.distributed.sp_sharding import sp_gather, sp_shard
+from vllm_omni.diffusion.hooks.base import HookRegistry, ModelHook
+
+logger = init_logger(__name__)
+
+# Hook name templates for identifying SP hooks
+_SP_INPUT_HOOK_TEMPLATE = "sp_input---{}"
+_SP_OUTPUT_HOOK_TEMPLATE = "sp_output---{}"
+
+
+@dataclass
+class ModuleForwardMetadata:
+    """Metadata for mapping forward() parameter names to args/kwargs positions.
+
+    This caches the inspection of a module's forward signature to efficiently
+    locate parameters by name in subsequent calls.
+    """
+
+    cached_parameter_indices: dict[str, int] | None = None
+    _cls: type | None = None
+
+    def _get_parameter_from_args_kwargs(
+        self,
+        identifier: str,
+        args: tuple = (),
+        kwargs: dict | None = None,
+    ) -> tuple[Any, bool, int | None]:
+        """Get a parameter value from args or kwargs by name.
+
+        Args:
+            identifier: The parameter name to look up.
+            args: Positional arguments passed to forward.
+            kwargs: Keyword arguments passed to forward.
+
+        Returns:
+            Tuple of (value, is_kwarg, index).
+            - value: The parameter value (or None if not found)
+            - is_kwarg: True if found in kwargs
+            - index: Position in args if found there
+
+        Raises:
+            ValueError: If parameter not found in signature.
+        """
+        kwargs = kwargs or {}
+
+        # First check kwargs
+        if identifier in kwargs:
+            return kwargs[identifier], True, None
+
+        # Check cached indices
+        if self.cached_parameter_indices is not None:
+            index = self.cached_parameter_indices.get(identifier, None)
+            if index is None:
+                raise ValueError(f"Parameter '{identifier}' not found in cached indices.")
+            if index < len(args):
+                return args[index], False, index
+            return None, False, index
+
+        # Build cache from forward signature
+        if self._cls is None:
+            raise ValueError("Model class is not set for metadata.")
+
+        parameters = list(inspect.signature(self._cls.forward).parameters.keys())
+        parameters = parameters[1:]  # Skip `self`
+        self.cached_parameter_indices = {param: i for i, param in enumerate(parameters)}
+
+        if identifier not in self.cached_parameter_indices:
+            raise ValueError(f"Parameter '{identifier}' not found in function signature.")
+
+        index = self.cached_parameter_indices[identifier]
+
+        if index >= len(args):
+            return None, False, index
+
+        return args[index], False, index
+
+
+def _unwrap_module(module: nn.Module) -> nn.Module:
+    """Unwrap a module from any wrappers to get the original class.
+
+    Args:
+        module: Potentially wrapped module.
+
+    Returns:
+        The unwrapped module.
+    """
+    # Handle common wrappers
+    while hasattr(module, "_modules") and len(module._modules) == 1:
+        inner = next(iter(module._modules.values()))
+        if inner is not None:
+            module = inner
+        else:
+            break
+    return module
+
+
+class SequenceParallelSplitHook(ModelHook):
+    """Hook for splitting inputs before a module's forward pass.
+
+    This hook is registered to modules that need their inputs sharded
+    across sequence parallel ranks. It intercepts the forward call,
+    shards specified inputs according to the plan, and passes the
+    sharded inputs to the original forward.
+
+    For split_output=True inputs, it shards the output instead.
+
+    Supports both SequenceParallelInput (full split) and SequenceParallelPartialInput
+    (partial split for text/image separation).
+
+    Note: This corresponds to `ContextParallelSplitHook` in diffusers.
+    """
+
+    def __init__(
+        self,
+        metadata: dict[str | int, AnySequenceParallelInput | list[AnySequenceParallelInput]],
+        config: SequenceParallelConfig,
+    ) -> None:
+        super().__init__()
+        self.metadata = metadata
+        self.config = config
+        self.module_forward_metadata: ModuleForwardMetadata | None = None
+        # Cache for text lengths resolved from kwargs
+        self._text_len_cache: dict[str, int] = {}
+
+    def initialize_hook(self, module: nn.Module) -> nn.Module:
+        cls = _unwrap_module(module).__class__
+        self.module_forward_metadata = ModuleForwardMetadata(_cls=cls)
+        return module
+
+    def pre_forward(self, module: nn.Module, *args: Any, **kwargs: Any) -> tuple[tuple, dict]:
+        """Shard inputs before forward."""
+        args_list = list(args)
+        # Clear text length cache for this forward pass
+        self._text_len_cache.clear()
+
+        for name, spm in self.metadata.items():
+            # Skip if this is a split_output entry (handled in post_forward)
+            if isinstance(spm, (SequenceParallelInput, SequenceParallelPartialInput)) and spm.split_output:
+                continue
+
+            # Get the parameter value
+            input_val, is_kwarg, index = self.module_forward_metadata._get_parameter_from_args_kwargs(
+                name, args_list, kwargs
+            )
+
+            if input_val is None:
+                continue
+
+            # Shard the input
+            if isinstance(input_val, torch.Tensor):
+                input_val = self._prepare_sp_input(input_val, spm, args_list, kwargs)
+            elif isinstance(input_val, (list, tuple)):
+                # Handle list/tuple of tensors with per-element config
+                if not isinstance(spm, (list, tuple)):
+                    raise ValueError(
+                        f"Expected list/tuple of SequenceParallelInput for parameter '{name}' "
+                        f"which is a list/tuple, but got {type(spm).__name__}"
+                    )
+                if len(input_val) != len(spm):
+                    raise ValueError(f"Expected {len(spm)} elements for parameter '{name}', got {len(input_val)}")
+                sharded_input_val = []
+                for i, x in enumerate(input_val):
+                    if torch.is_tensor(x) and not spm[i].split_output:
+                        x = self._prepare_sp_input(x, spm[i], args_list, kwargs)
+                    sharded_input_val.append(x)
+                input_val = type(input_val)(sharded_input_val)
+            else:
+                raise ValueError(f"Unsupported input type for sharding: {type(input_val).__name__}")
+
+            # Update args or kwargs
+            if is_kwarg:
+                kwargs[name] = input_val
+            elif index is not None and index < len(args_list):
+                args_list[index] = input_val
+            else:
+                raise ValueError(f"Failed to update parameter '{name}' after sharding.")
+
+        # Store kwargs for post_forward to resolve text lengths
+        self._last_kwargs = kwargs
+        self._last_args = tuple(args_list)
+
+        return tuple(args_list), kwargs
+
+    def post_forward(self, module: nn.Module, output: Any) -> Any:
+        """Shard outputs for split_output=True entries."""
+        is_tensor = isinstance(output, torch.Tensor)
+        is_tensor_list = isinstance(output, (list, tuple)) and all(isinstance(x, torch.Tensor) for x in output)
+
+        if not is_tensor and not is_tensor_list:
+            # No tensor outputs to shard
+            return output
+
+        output_list = [output] if is_tensor else list(output)
+
+        for index, spm in self.metadata.items():
+            if not isinstance(index, int):
+                continue
+            if not isinstance(spm, (SequenceParallelInput, SequenceParallelPartialInput)) or not spm.split_output:
+                continue
+            if index >= len(output_list):
+                raise ValueError(f"Index {index} out of bounds for output of length {len(output_list)}.")
+
+            output_list[index] = self._prepare_sp_input(output_list[index], spm, self._last_args, self._last_kwargs)
+
+        return output_list[0] if is_tensor else type(output)(output_list)
+
+    def _resolve_text_len(
+        self,
+        sp_input: SequenceParallelPartialInput,
+        args: tuple,
+        kwargs: dict,
+    ) -> int:
+        """Resolve text length from the source specification."""
+        source = sp_input.text_len_source
+
+        if isinstance(source, int):
+            return source
+
+        # String source - look up from kwargs or args
+        if source in self._text_len_cache:
+            return self._text_len_cache[source]
+
+        # Try to get from kwargs/args
+        try:
+            val, _, _ = self.module_forward_metadata._get_parameter_from_args_kwargs(source, args, kwargs)
+            if val is None:
+                raise ValueError(f"Parameter '{source}' is None, cannot determine text length.")
+            if isinstance(val, torch.Tensor):
+                # TODO: Currently assumes batch_size=1, where shape[0] is sequence length.
+                # For batch inference support, this should be updated to handle
+                # shape (batch_size, seq_len, ...) where text_len varies per sample.
+                text_len = val.shape[0]
+            elif isinstance(val, int):
+                text_len = val
+            else:
+                raise ValueError(f"Cannot determine text length from '{source}' of type {type(val).__name__}")
+            self._text_len_cache[source] = text_len
+            return text_len
+        except ValueError as e:
+            raise ValueError(f"Failed to resolve text_len_source '{source}': {e}") from e
+
+    def _prepare_sp_input(
+        self,
+        x: torch.Tensor,
+        sp_input: AnySequenceParallelInput,
+        args: tuple = (),
+        kwargs: dict | None = None,
+    ) -> torch.Tensor:
+        """Shard a tensor according to the input specification."""
+        kwargs = kwargs or {}
+
+        if sp_input.expected_dims is not None and x.dim() != sp_input.expected_dims:
+            logger.warning_once(f"Expected tensor with {sp_input.expected_dims} dims, got {x.dim()}. Skipping split.")
+            return x
+
+        if isinstance(sp_input, SequenceParallelInput):
+            # Full split with optional auto-padding
+            if sp_input.auto_pad:
+                return self._shard_with_auto_pad(x, sp_input.split_dim)
+            return sp_shard(x, sp_input.split_dim, validate=False)
+        elif isinstance(sp_input, SequenceParallelPartialInput):
+            # Partial split: keep text portion, split image portion
+            text_len = self._resolve_text_len(sp_input, args, kwargs)
+            dim = sp_input.split_dim
+
+            # Split tensor into text and image portions
+            text_part = x.narrow(dim, 0, text_len)
+            image_part = x.narrow(dim, text_len, x.size(dim) - text_len)
+
+            # Only shard the image portion
+            image_part_sharded = sp_shard(image_part, dim, validate=False)
+
+            # Concatenate back: [text_full, image_sharded]
+            return torch.cat([text_part, image_part_sharded], dim=dim)
+        else:
+            raise ValueError(f"Unsupported input config type: {type(sp_input).__name__}")
+
+    def _shard_with_auto_pad(self, x: torch.Tensor, dim: int) -> torch.Tensor:
+        """Shard tensor with automatic padding and attention mask creation.
+
+        When sequence length is not divisible by SP world size, this method:
+        1. Pads the tensor to make it divisible
+        2. Creates an attention mask indicating valid vs padding positions
+        3. Stores the mask and padding info in ForwardContext
+        """
+        from vllm_omni.diffusion.attention.selector import (
+            _BACKENDS_SUPPORT_ATTENTION_MASK,
+            get_attn_backend,
+        )
+        from vllm_omni.diffusion.distributed.parallel_state import (
+            get_ring_parallel_world_size,
+            get_sequence_parallel_rank,
+            get_sequence_parallel_world_size,
+        )
+        from vllm_omni.diffusion.forward_context import get_forward_context, is_forward_context_available
+
+        world_size = get_sequence_parallel_world_size()
+        if world_size == 1:
+            return x
+
+        seq_len = x.size(dim)
+        remainder = seq_len % world_size
+
+        if remainder == 0:
+            # No padding needed
+            return sp_shard(x, dim, validate=False)
+
+        # Check backend compatibility
+        attn_backend = get_attn_backend(-1)
+        if attn_backend.get_name() not in _BACKENDS_SUPPORT_ATTENTION_MASK:
+            raise ValueError(
+                f"Sequence length ({seq_len}) is not divisible by SP world size ({world_size}). "
+                f"Cannot use {attn_backend.get_name()} which does not support attention_mask. "
+                f"Please switch to SDPA or Ascend attention backend."
+            )
+
+        # Ring attention does not support attention_mask
+        if get_ring_parallel_world_size() > 1:
+            raise ValueError(
+                f"Sequence length ({seq_len}) is not divisible by SP world size ({world_size}). "
+                f"Cannot use Ring attention which does not support attention_mask. "
+                f"Please switch to Ulysses SP only."
+            )
+
+        # Calculate padding
+        pad_size = world_size - remainder
+        padded_seq_len = seq_len + pad_size
+
+        # Pad the tensor
+        pad_shape = list(x.shape)
+        pad_shape[dim] = pad_size
+        padding = torch.zeros(pad_shape, dtype=x.dtype, device=x.device)
+        x_padded = torch.cat([x, padding], dim=dim)
+
+        # Store padding info in forward context (only once, for primary tensor)
+        # Attention layers will create masks dynamically using this info
+        if is_forward_context_available():
+            ctx = get_forward_context()
+            # Only set if not already set (first auto_pad tensor wins)
+            if ctx.sp_original_seq_len is None:
+                ctx.sp_padding_size = pad_size
+                ctx.sp_original_seq_len = seq_len
+                logger.debug(
+                    f"Auto-padded sequence from {seq_len} to {padded_seq_len} "
+                    f"(pad_size={pad_size}, world_size={world_size}, dim={dim})"
+                )
+
+        # Shard the padded tensor
+        rank = get_sequence_parallel_rank()
+        return x_padded.chunk(world_size, dim=dim)[rank]
+
+
+class SequenceParallelGatherHook(ModelHook):
+    """Hook for gathering outputs after a module's forward pass.
+
+    This hook is registered to modules that need their outputs gathered
+    from all sequence parallel ranks. It intercepts the output and gathers
+    it according to the plan specification.
+
+    Note: This corresponds to `ContextParallelGatherHook` in diffusers.
+    """
+
+    def __init__(
+        self,
+        metadata: SequenceParallelOutput | list[SequenceParallelOutput],
+        config: SequenceParallelConfig,
+    ) -> None:
+        super().__init__()
+        if isinstance(metadata, SequenceParallelOutput):
+            metadata = [metadata]
+        self.metadata = metadata
+        self.config = config
+
+    def initialize_hook(self, module: nn.Module) -> nn.Module:
+        return module
+
+    def post_forward(self, module: nn.Module, output: Any) -> Any:
+        """Gather outputs after forward and remove padding if applied."""
+        from vllm_omni.diffusion.forward_context import get_forward_context, is_forward_context_available
+
+        is_tensor = isinstance(output, torch.Tensor)
+
+        if is_tensor:
+            output = [output]
+        elif not (isinstance(output, (list, tuple)) and all(isinstance(x, torch.Tensor) for x in output)):
+            raise ValueError(f"Expected tensor or list/tuple of tensors, got {type(output).__name__}")
+
+        output = list(output)
+
+        if len(output) != len(self.metadata):
+            raise ValueError(f"Expected {len(self.metadata)} outputs, got {len(output)}.")
+
+        # Check if padding was applied during split
+        original_seq_len = None
+        if is_forward_context_available():
+            ctx = get_forward_context()
+            original_seq_len = ctx.sp_original_seq_len
+
+        for i, spm in enumerate(self.metadata):
+            if spm is None:
+                continue
+
+            x = output[i]
+            if spm.expected_dims is not None and x.dim() != spm.expected_dims:
+                logger.warning_once(
+                    f"Expected output tensor with {spm.expected_dims} dims, got {x.dim()}. Skipping gather."
+                )
+                continue
+
+            # Gather from all ranks
+            gathered = sp_gather(x, spm.gather_dim, validate=False)
+
+            # Remove padding if it was applied
+            if original_seq_len is not None and gathered.size(spm.gather_dim) > original_seq_len:
+                gathered = gathered.narrow(spm.gather_dim, 0, original_seq_len)
+                logger.debug(f"Removed padding: gathered shape {gathered.shape} (original_seq_len={original_seq_len})")
+
+            output[i] = gathered
+
+        return output[0] if is_tensor else type(output)(output)
+
+
+def _get_submodule_by_name(model: nn.Module, name: str) -> nn.Module | list[nn.Module]:
+    """Get a submodule by dotted name, supporting wildcards.
+
+    Args:
+        model: The root module.
+        name: Dotted path to submodule. Use "*" to match all children
+            of a ModuleList.
+
+    Returns:
+        The submodule or list of submodules if wildcard used.
+
+    Raises:
+        ValueError: If the path is invalid or module not found.
+    """
+    if name.count("*") > 1:
+        raise ValueError("Wildcard '*' can only be used once in the name")
+    return _find_submodule_by_name(model, name)
+
+
+def _find_submodule_by_name(model: nn.Module, name: str) -> nn.Module | list[nn.Module]:
+    """Recursive helper for _get_submodule_by_name."""
+    if name == "":
+        return model
+
+    first_atom, remaining_name = name.split(".", 1) if "." in name else (name, "")
+
+    if first_atom == "*":
+        if not isinstance(model, nn.ModuleList):
+            raise ValueError("Wildcard '*' can only be used with ModuleList")
+        submodules = []
+        for submodule in model:
+            subsubmodules = _find_submodule_by_name(submodule, remaining_name)
+            if not isinstance(subsubmodules, list):
+                subsubmodules = [subsubmodules]
+            submodules.extend(subsubmodules)
+        return submodules
+    else:
+        if hasattr(model, first_atom):
+            submodule = getattr(model, first_atom)
+            return _find_submodule_by_name(submodule, remaining_name)
+        else:
+            raise ValueError(f"'{first_atom}' is not a submodule of '{model.__class__.__name__}'")
+
+
+def apply_sequence_parallel(
+    module: nn.Module,
+    config: SequenceParallelConfig,
+    plan: SequenceParallelModelPlan,
+) -> None:
+    """Apply sequence parallel hooks to a model according to the plan.
+
+    This function registers hooks on the specified submodules to automatically
+    shard inputs and gather outputs for sequence parallelism.
+
+    Note: This corresponds to `apply_context_parallel` in diffusers.
+
+    The complete SP flow is:
+    1. Input sharding (SequenceParallelSplitHook): Split sequence across SP ranks
+    2. Attention parallelism (handled by vLLM-Omni's Attention layer):
+       - Ulysses: All-to-All over Q/K/V heads
+       - Ring: K/V circulation in ring topology
+       - Hybrid: Both (Ulysses handles head redistribution, Ring handles K/V)
+    3. Output gathering (SequenceParallelGatherHook): Gather sequence from SP ranks
+
+    Args:
+        module: The model to apply SP to.
+        config: The sequence parallel configuration.
+        plan: Dictionary mapping module names to input/output specifications.
+
+    Example:
+        config = SequenceParallelConfig(ulysses_degree=2)
+        plan = {
+            "": {"hidden_states": SequenceParallelInput(split_dim=1, expected_dims=3)},
+            "proj_out": SequenceParallelOutput(gather_dim=1, expected_dims=3),
+        }
+        apply_sequence_parallel(model, config, plan)
+
+    Note:
+        vLLM-Omni's Attention layer automatically handles the internal
+        parallelism (Ulysses All-to-All or Ring attention) based on the
+        forward_context configuration. This function only handles input/output
+        sharding for the model as a whole.
+    """
+    logger.debug(
+        f"Applying sequence parallel with config: ulysses={config.ulysses_degree}, "
+        f"ring={config.ring_degree}, plan keys: {list(plan.keys())}"
+    )
+
+    for module_id, sp_model_plan in plan.items():
+        submodule = _get_submodule_by_name(module, module_id)
+        if not isinstance(submodule, list):
+            submodule = [submodule]
+
+        logger.debug(f"Applying SP hooks to '{module_id}' ({len(submodule)} module(s))")
+
+        for m in submodule:
+            if isinstance(sp_model_plan, dict):
+                # Input specification
+                hook = SequenceParallelSplitHook(sp_model_plan, config)
+                hook_name = _SP_INPUT_HOOK_TEMPLATE.format(module_id)
+            elif isinstance(sp_model_plan, (SequenceParallelOutput, list, tuple)):
+                # Output specification
+                if isinstance(sp_model_plan, SequenceParallelOutput):
+                    sp_model_plan = [sp_model_plan]
+                if not all(isinstance(x, SequenceParallelOutput) or x is None for x in sp_model_plan):
+                    raise ValueError(f"Expected SequenceParallelOutput elements, got {sp_model_plan}")
+                hook = SequenceParallelGatherHook(sp_model_plan, config)
+                hook_name = _SP_OUTPUT_HOOK_TEMPLATE.format(module_id)
+            else:
+                raise ValueError(f"Unsupported plan type: {type(sp_model_plan).__name__}")
+
+            registry = HookRegistry.get_or_create(m)
+            registry.register_hook(hook_name, hook)
+
+
+def remove_sequence_parallel(
+    module: nn.Module,
+    plan: SequenceParallelModelPlan,
+) -> None:
+    """Remove sequence parallel hooks from a model.
+
+    Note: This corresponds to `remove_context_parallel` in diffusers.
+
+    Args:
+        module: The model to remove SP from.
+        plan: The same plan used when applying SP.
+    """
+    for module_id, sp_model_plan in plan.items():
+        submodule = _get_submodule_by_name(module, module_id)
+        if not isinstance(submodule, list):
+            submodule = [submodule]
+
+        for m in submodule:
+            registry = getattr(m, "_hook_registry", None)
+            if registry is None:
+                continue
+
+            if isinstance(sp_model_plan, dict):
+                hook_name = _SP_INPUT_HOOK_TEMPLATE.format(module_id)
+            elif isinstance(sp_model_plan, (SequenceParallelOutput, list, tuple)):
+                hook_name = _SP_OUTPUT_HOOK_TEMPLATE.format(module_id)
+            else:
+                continue
+
+            registry.remove_hook(hook_name)
+
+
+def enable_sequence_parallel_for_model(
+    model: nn.Module,
+    config: SequenceParallelConfig | None = None,
+) -> None:
+    """Enable sequence parallelism for a model using its _sp_plan.
+
+    This is a convenience function that reads the model's _sp_plan attribute
+    and applies sequence parallelism automatically.
+
+    Note: This corresponds to `enable_context_parallel_for_model` in diffusers,
+    but uses vLLM-Omni's _sp_plan instead of diffusers' _cp_plan.
+
+    The function performs two main tasks:
+    1. Applies _sp_plan hooks to shard inputs and gather outputs
+    2. Ensures Attention layers are configured for the correct parallel mode
+       (handled automatically by vLLM-Omni's forward_context mechanism)
+
+    Args:
+        model: The model to enable SP for. Must have a _sp_plan attribute.
+        config: Optional config. If None, uses default based on current
+            parallel state.
+
+    Raises:
+        ValueError: If model has no _sp_plan defined.
+
+    Note:
+        vLLM-Omni supports Ulysses + Ring hybrid parallelism:
+        - ulysses_degree > 1: Uses All-to-All communication over Q/K/V heads
+        - ring_degree > 1: Uses Ring attention with K/V passing
+        - Both > 1: Hybrid mode (Ulysses handles head redistribution,
+          Ring handles K/V circulation)
+    """
+    from vllm_omni.diffusion.distributed.parallel_state import (
+        get_ring_parallel_world_size,
+        get_ulysses_parallel_world_size,
+    )
+    from vllm_omni.diffusion.distributed.sp_plan import get_sp_plan_from_model
+
+    plan = get_sp_plan_from_model(model)
+    if plan is None:
+        raise ValueError(
+            f"Model {model.__class__.__name__} has no _sp_plan defined. "
+            f"Define _sp_plan as a class attribute or pass a plan explicitly."
+        )
+
+    if config is None:
+        # Create config from current parallel state
+        ulysses_degree = get_ulysses_parallel_world_size()
+        ring_degree = get_ring_parallel_world_size()
+        config = SequenceParallelConfig(
+            ulysses_degree=ulysses_degree,
+            ring_degree=ring_degree,
+        )
+        if ulysses_degree > 1 and ring_degree > 1:
+            mode = "hybrid"
+        elif ulysses_degree > 1:
+            mode = "ulysses"
+        else:
+            mode = "ring"
+        logger.info(
+            f"Created SP config from parallel state: "
+            f"ulysses_degree={ulysses_degree}, ring_degree={ring_degree}, "
+            f"mode={mode}"
+        )
+
+    apply_sequence_parallel(model, config, plan)
+    logger.info(f"Enabled sequence parallelism for {model.__class__.__name__}")
+
+
+def disable_sequence_parallel_for_model(model: nn.Module) -> None:
+    """Disable sequence parallelism for a model.
+
+    Note: This corresponds to `disable_context_parallel_for_model` in diffusers.
+
+    Args:
+        model: The model to disable SP for.
+    """
+    from vllm_omni.diffusion.distributed.sp_plan import get_sp_plan_from_model
+
+    plan = get_sp_plan_from_model(model)
+    if plan is not None:
+        remove_sequence_parallel(model, plan)