[Model] Add LingBot World Ulysses sequence parallelism

docs/design/feature/realtime_ar_diffusion.md

@@ -178,6 +178,24 @@ For LingBot World v2:
 - `lingbot_world/pipeline.py` constructs conditioning, owns small non-KV session state, and
   produces one block plus the standard metadata envelope.
 
+## Ulysses sequence parallelism
+
+LingBot supports pure Ulysses sequence parallelism for both direct and
+AR-Diffusion execution. With Ulysses degree greater than one, hidden tokens,
+camera features, token-expanded timestep modulation, and RoPE tables are sharded
+together. Without SP, timestep modulation retains the frame-broadcast path.
+Self-attention performs the sequence-to-head all-to-all before reading or writing
+paged KV, while static text K/V uses the same local head shard. Text K/V shards
+own compact storage; cross-attention exchanges query/output layouts to use these
+shards. This retains the shared cache geometry and reduces text K/V storage at
+the cost of two all-to-all calls per layer.
+
+The output head projects local tokens before gathering the flow values. For
+the 14B model this reduces the gathered width from 5120 to 64; frame modulation
+uses each shard's global token offset, including shards that split a frame. Only
+`ulysses_mode="strict"` is supported; `advanced_uaa`, Ring, and AllGather-KV modes
+remain unsupported for this model.
+
 ## Non-goals
 
 This contract does not currently provide:

examples/offline_inference/diffusion/lingbot_world_v2.py

@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 """Generate a LingBot-World v2 video from an image and camera trajectory.
 
 The official checkpoint is licensed separately under CC BY-NC-SA and is
@@ -99,6 +99,7 @@ def parse_args(argv: Sequence[str] | None = None) -> argparse.Namespace:
         default=1,
         help="Number of GPUs used for tensor parallelism inside the DiT.",
     )
+    parser.add_argument("--ulysses-degree", type=int, default=1, help="Pure Ulysses sequence parallel degree.")
     parser.add_argument("--flow-shift", type=float, default=5.0, help="Positive FlowUniPC scheduler shift.")
     parser.add_argument("--fps", type=int, default=16, help="Frames per second in the exported MP4.")
     parser.add_argument("--output", default="lingbot_world_v2.mp4", help="Output MP4 path.")
@@ -158,13 +159,16 @@ def build_omni_kwargs(
 
     if args.tensor_parallel_size <= 0:
         raise ValueError("--tensor-parallel-size must be a positive integer.")
+    if args.ulysses_degree <= 0:
+        raise ValueError("--ulysses-degree must be a positive integer.")
     flow_shift = _positive_finite(args.flow_shift, "--flow-shift")
     model_path = Path(args.model).expanduser()
     model = str(model_path.resolve()) if model_path.exists() else args.model
     return {
         "model": model,
         "flow_shift": flow_shift,
         "tensor_parallel_size": args.tensor_parallel_size,
+        "ulysses_degree": args.ulysses_degree,
         "enforce_eager": args.enforce_eager,
         "model_config": {"lingbot_action_root": str(paths.action_root)},
     }

tests/diffusion/models/lingbot_world/test_lingbot_world_attention.py

@@ -1,13 +1,13 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 
 from __future__ import annotations
 
 import importlib.util
 import math
 import sys
 from pathlib import Path
-from types import ModuleType
+from types import ModuleType, SimpleNamespace
 
 import pytest
 import torch
@@ -33,6 +33,10 @@
     "vllm_omni.diffusion",
     "vllm_omni.diffusion.attention",
     "vllm_omni.diffusion.attention.layer",
+    "vllm_omni.diffusion.distributed",
+    "vllm_omni.diffusion.distributed.comm",
+    "vllm_omni.diffusion.distributed.parallel_state",
+    "vllm_omni.diffusion.distributed.sp_plan",
     "vllm_omni.diffusion.layers",
     "vllm_omni.diffusion.layers.norm",
     "vllm_omni.diffusion.layers.rope",
@@ -58,6 +62,47 @@ def _install_vllm_stubs() -> None:
     distributed.get_tensor_model_parallel_world_size = lambda: 1
     distributed.tensor_model_parallel_all_reduce = lambda value: value
 
+    class _SeqAllToAll4D:
+        @staticmethod
+        def apply(group, value, scatter_idx, gather_idx, use_sync=False):
+            del group, scatter_idx, gather_idx, use_sync
+            return value
+
+    setattr(
+        sys.modules["vllm_omni.diffusion.distributed.comm"],
+        "SeqAllToAll4D",
+        _SeqAllToAll4D,
+    )
+
+    def get_sp_group():
+        return SimpleNamespace(
+            ulysses_world_size=1,
+            ulysses_rank=0,
+            ulysses_group=None,
+        )
+
+    setattr(
+        sys.modules["vllm_omni.diffusion.distributed.parallel_state"],
+        "get_sp_group",
+        get_sp_group,
+    )
+
+    class _SequenceParallelInput:
+        def __init__(self, split_dim, expected_dims=None, split_output=False, auto_pad=False):
+            self.split_dim = split_dim
+            self.expected_dims = expected_dims
+            self.split_output = split_output
+            self.auto_pad = auto_pad
+
+    class _SequenceParallelOutput:
+        def __init__(self, gather_dim, expected_dims=None):
+            self.gather_dim = gather_dim
+            self.expected_dims = expected_dims
+
+    sp_plan = sys.modules["vllm_omni.diffusion.distributed.sp_plan"]
+    setattr(sp_plan, "SequenceParallelInput", _SequenceParallelInput)
+    setattr(sp_plan, "SequenceParallelOutput", _SequenceParallelOutput)
+
     def set_weight_attrs(weight: torch.Tensor, attrs: dict) -> None:
         for name, value in attrs.items():
             setattr(weight, name, value)
@@ -541,3 +586,38 @@ def test_tp_rmsnorm_weight_loader_selects_rank_shard(monkeypatch: pytest.MonkeyP
     norm.weight.weight_loader(norm.weight, torch.tensor([10.0, 20.0, 30.0, 40.0]))
 
     torch.testing.assert_close(norm.weight, torch.tensor([30.0, 40.0]))
+
+
+def test_single_token_text_kv_shard_owns_compact_storage(monkeypatch):
+    module = _load_module()
+    monkeypatch.setattr(
+        module,
+        "get_sp_group",
+        lambda: SimpleNamespace(
+            ulysses_world_size=2,
+            ulysses_rank=1,
+            ulysses_group=None,
+        ),
+    )
+    attention = module.LingBotCrossAttention(dim=8, num_heads=4)
+    full = torch.arange(8, dtype=torch.float32).reshape(1, 1, 4, 2)
+    shard = attention.shard_kv_heads(full)
+    torch.testing.assert_close(shard, full[:, :, 2:], rtol=0, atol=0)
+    assert shard.is_contiguous()
+    assert shard.untyped_storage().nbytes() == shard.numel() * shard.element_size()
+
+
+@pytest.mark.parametrize("attention_class", ["LingBotSelfAttention", "LingBotCrossAttention"])
+def test_ulysses_rejects_non_divisible_head_count(monkeypatch, attention_class):
+    module = _load_module()
+    monkeypatch.setattr(
+        module,
+        "get_sp_group",
+        lambda: SimpleNamespace(
+            ulysses_world_size=3,
+            ulysses_rank=0,
+            ulysses_group=None,
+        ),
+    )
+    with pytest.raises(ValueError, match="heads must be divisible"):
+        getattr(module, attention_class)(dim=8, num_heads=4)