[Bugfix] Skip HWR store-size scans when no limit is configured

Signed-off-by: Hongsheng Liu <liuhongsheng4@huawei.com>

benchmarks/host_weight_runtime/README.md

@@ -1,57 +0,0 @@
-# Host weight dependency memory diagnostic
-
-`safetensors_retention.py` isolates repeated CPU `get_tensor()` calls from
-access to one cached tensor view. It imports neither vLLM nor HWR and uses a
-synthetic float32 payload controlled by `--tensor-elements` (default: 64).
-Tensor size can affect retention: a negative result for 16 elements does not
-rule out growth for 64 elements. The selected size is recorded in the JSON
-arguments and used for both file creation and the correctness check.
-HWR itself calls `get_tensor()` when
-acquiring a lease, then exposes cached views; this probe does not measure HWR
-requests or repeated lease acquisition.
-
-Run one mode per fresh process in the same Python environment and CPU affinity:
-
-```bash
-git rev-parse HEAD > revision.txt
-git status --short > working-tree.txt
-git diff > working-tree.patch
-# Select an allowed CPU from your process affinity; 0 is only an example.
-CUDA_VISIBLE_DEVICES='' taskset -c 0 timeout 120s python \
-  benchmarks/host_weight_runtime/safetensors_retention.py \
-  --mode get_tensor --iterations 10000 --sample-every 5000 > feasibility.json
-```
-
-For a size comparison, run 16 and 64 elements twice each in fresh processes,
-with cached reuse as a control at both sizes:
-
-```bash
-for repetition in 1 2; do
-  for elements in 16 64; do
-    for mode in get_tensor reuse; do
-      CUDA_VISIBLE_DEVICES='' taskset -c 0 timeout 120s python \
-        benchmarks/host_weight_runtime/safetensors_retention.py \
-        --mode "$mode" --tensor-elements "$elements" \
-        > "$mode-$elements-$repetition.json"
-    done
-  done
-done
-```
-
-The diagnostic sets Torch to one CPU thread, warms up before the baseline,
-and reports preparation separately from timed loop work. Checkpoints run GC
-and report Linux private/anonymous memory, RSS, and open descriptors. It checks
-that transient tensor objects and the final cached view are released. Linux
-`/proc/self/smaps_rollup` must be readable; invalid arguments and missing proc
-support fail explicitly. Keep raw JSON and the repository snapshot with results.
-
-Compare the final loop sample with the warmed baseline, then inspect the closed
-sample. Private-memory growth after Python tensor objects disappear is evidence
-of retention in the dependency/allocator path, not proof of a leak's root cause.
-RSS can include file-backed pages; it is not a private-memory metric. Samples
-also include small diagnostic bookkeeping allocations. Do not assert a portable
-memory threshold or turn this synthetic loop into a per-request service estimate.
-
-Do not drop shared caches, trim allocators, or add serving-time GC to make the
-numbers look smaller. A dependency replacement or version change requires its
-own reproducible comparison and correctness/lifecycle validation.

benchmarks/host_weight_runtime/safetensors_retention.py

@@ -1,120 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
-"""Linux CPU diagnostic: repeated safetensors materialization versus view reuse."""
-
-from __future__ import annotations
-
-import argparse
-import gc
-import json
-import os
-import platform
-import tempfile
-import time
-import weakref
-from pathlib import Path
-
-
-def positive_int(value: str) -> int:
-    parsed = int(value)
-    if parsed <= 0:
-        raise argparse.ArgumentTypeError("must be positive")
-    return parsed
-
-
-def sample(phase: str, iterations: int, loop_seconds: float) -> dict[str, object]:
-    gc.collect()
-    counters = {}
-    for line in Path("/proc/self/smaps_rollup").read_text().splitlines():
-        fields = line.split()
-        if len(fields) == 3 and fields[2] == "kB":
-            counters[fields[0].removesuffix(":")] = int(fields[1]) * 1024
-    return {
-        "phase": phase,
-        "iterations": iterations,
-        "loop_seconds": loop_seconds,
-        "rss_bytes": counters["Rss"],
-        "private_bytes": counters["Private_Clean"] + counters["Private_Dirty"],
-        "anonymous_bytes": counters["Anonymous"],
-        "fd_count": len(tuple(Path("/proc/self/fd").iterdir())),
-    }
-
-
-def main() -> None:
-    parser = argparse.ArgumentParser(description=__doc__)
-    parser.add_argument("--mode", choices=("get_tensor", "reuse"), required=True)
-    parser.add_argument(
-        "--tensor-elements", type=positive_int, default=64, help="Number of float32 elements in the synthetic tensor"
-    )
-    parser.add_argument("--iterations", type=positive_int, default=1_000_000)
-    parser.add_argument("--sample-every", type=positive_int, default=100_000)
-    parser.add_argument("--warmup", type=positive_int, default=1000)
-    args = parser.parse_args()
-    if not Path("/proc/self/smaps_rollup").is_file():
-        parser.error("requires Linux /proc/self/smaps_rollup")
-
-    preparation_start = time.monotonic()
-    import safetensors
-    import torch
-    from safetensors import safe_open
-    from safetensors.torch import save_file
-
-    torch.set_num_threads(1)
-    torch.set_num_interop_threads(1)
-    samples = [sample("before_file", 0, 0.0)]
-    with tempfile.TemporaryDirectory(prefix="safetensors-retention-") as directory:
-        path = Path(directory) / "synthetic.safetensors"
-        save_file({"weight": torch.arange(args.tensor_elements, dtype=torch.float32)}, str(path))
-        with safe_open(path, framework="pt", device="cpu") as reader:
-            cached = reader.get_tensor("weight")
-            assert torch.equal(cached, torch.arange(args.tensor_elements, dtype=torch.float32))
-            for _ in range(args.warmup):
-                tensor = reader.get_tensor("weight") if args.mode == "get_tensor" else cached
-                del tensor
-            samples.append(sample("baseline", 0, 0.0))
-            preparation_seconds = time.monotonic() - preparation_start
-            completed = 0
-            loop_seconds = 0.0
-            while completed < args.iterations:
-                count = min(args.sample_every, args.iterations - completed)
-                start = time.monotonic()
-                for _ in range(count):
-                    tensor = reader.get_tensor("weight") if args.mode == "get_tensor" else cached
-                last_tensor = weakref.ref(tensor)
-                del tensor
-                loop_seconds += time.monotonic() - start
-                completed += count
-                snapshot = sample("loop", completed, loop_seconds)
-                snapshot["last_tensor_alive"] = last_tensor() is not None
-                if args.mode == "get_tensor":
-                    assert last_tensor() is None, "diagnostic retained a transient tensor"
-                samples.append(snapshot)
-            del cached
-        del reader
-        samples.append(sample("closed", completed, loop_seconds))
-        assert last_tensor() is None, "last tensor survived reader/view teardown"
-
-    print(
-        json.dumps(
-            {
-                "schema_version": 1,
-                "arguments": vars(args),
-                "python": platform.python_version(),
-                "torch": torch.__version__,
-                "safetensors": safetensors.__version__,
-                "kernel": platform.release(),
-                "hostname": platform.node(),
-                "pid": os.getpid(),
-                "cpu_affinity": sorted(os.sched_getaffinity(0)),
-                "torch_threads": torch.get_num_threads(),
-                "pythonmalloc": os.environ.get("PYTHONMALLOC", "default"),
-                "preparation_seconds": preparation_seconds,
-                "samples": samples,
-            },
-            indent=2,
-        )
-    )
-
-
-if __name__ == "__main__":
-    main()

docs/design/feature/host_weight_runtime.md

@@ -42,11 +42,6 @@ V1 does not include:
 
 ## Motivation and use cases
 
-For investigating CPU memory retained by dependency tensor materialization,
-see the [standalone safetensors diagnostic](https://github.com/vllm-project/vllm-omni/blob/main/benchmarks/host_weight_runtime/README.md).
-It distinguishes repeated dependency calls from reuse of cached views and does
-not establish per-request HWR leakage.
-
 Model loading can create the same final host representation repeatedly. This
 is especially expensive when loading performs checkpoint decoding, tensor
 renaming, TP slicing, quantization, packing, or scale construction before GPU

tests/diffusion/attention/test_ulysses_uaa.py

@@ -27,202 +27,6 @@ from vllm_omni.platforms import current_omni_platform
 pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
 
 
-@pytest.mark.core_model
-@pytest.mark.cpu
-def test_uaa_equal_rank_contract_skips_length_gather(
-    monkeypatch: pytest.MonkeyPatch,
-) -> None:
-    from vllm_omni.diffusion.attention.parallel import ulysses
-
-    class FakeGroup:
-        ulysses_group = object()
-        ulysses_world_size = 2
-        ulysses_rank = 0
-        ring_world_size = 1
-
-    observed = []
-
-    def fake_all_to_all(pg, tensor, **kwargs):
-        observed.append((kwargs["seq_lens"], kwargs["padded_head_cnt"]))
-        return tensor, tensor.shape[2]
-
-    monkeypatch.setattr(ulysses, "_ulysses_all_to_all_any_qkv", fake_all_to_all)
-    monkeypatch.setattr(ulysses, "get_ulysses_mode", lambda **kwargs: "advanced_uaa")
-    monkeypatch.setattr(
-        ulysses,
-        "_all_gather_int",
-        lambda *args, **kwargs: pytest.fail("length gather must be skipped"),
-    )
-    strategy = ulysses.UlyssesParallelAttention(
-        FakeGroup(),
-        scatter_idx=2,
-        gather_idx=1,
-        use_sync=False,
-    )
-
-    with set_forward_context():
-        get_forward_context()._sp_equal_pad_stack.append(True)
-        strategy.pre_attention(
-            torch.zeros(1, 3, 28, 4),
-            torch.zeros(1, 3, 7, 4),
-            torch.zeros(1, 3, 7, 4),
-            None,
-        )
-
-    assert observed == [([3, 3], 32), ([3, 3], 8), ([3, 3], 8)]
-
-
-@pytest.mark.core_model
-@pytest.mark.cpu
-def test_uaa_hybrid_equal_rank_contract_skips_ulysses_gather_keeps_ring_check(
-    monkeypatch: pytest.MonkeyPatch,
-) -> None:
-    """Contract + ring: the fast path skips only the Ulysses gather.
-
-    The ring equality check is a separate collective on ring_group and must
-    still run, receiving s_global derived from the contract.
-    """
-    from vllm_omni.diffusion.attention.parallel import ulysses
-
-    class FakeGroup:
-        ulysses_group = object()
-        ring_group = object()
-        ulysses_world_size = 2
-        ulysses_rank = 0
-        ring_world_size = 2
-
-    sp_group = FakeGroup()
-    observed = []
-    ring_checks = []
-
-    def fake_all_to_all(pg, tensor, **kwargs):
-        observed.append((kwargs["seq_lens"], kwargs["padded_head_cnt"]))
-        return tensor, tensor.shape[2]
-
-    def fake_gather_int(pg, value, *, device):
-        if pg is sp_group.ulysses_group:
-            pytest.fail("Ulysses length gather must be skipped under the contract")
-        assert pg is sp_group.ring_group
-        ring_checks.append(int(value))
-        return [int(value), int(value)]
-
-    monkeypatch.setattr(ulysses, "_ulysses_all_to_all_any_qkv", fake_all_to_all)
-    monkeypatch.setattr(ulysses, "_all_gather_int", fake_gather_int)
-    monkeypatch.setattr(ulysses, "get_ulysses_mode", lambda **kwargs: "advanced_uaa")
-    strategy = ulysses.UlyssesParallelAttention(
-        sp_group,
-        scatter_idx=2,
-        gather_idx=1,
-        use_sync=False,
-    )
-
-    with set_forward_context():
-        get_forward_context()._sp_equal_pad_stack.append(True)
-        strategy.pre_attention(
-            torch.zeros(1, 3, 28, 4),
-            torch.zeros(1, 3, 7, 4),
-            torch.zeros(1, 3, 7, 4),
-            None,
-        )
-
-    # Ring check ran: s_global = ulysses_world_size * local = 2 * 3.
-    assert ring_checks == [6]
-    assert observed == [([3, 3], 32), ([3, 3], 8), ([3, 3], 8)]
-
-
-@pytest.mark.core_model
-@pytest.mark.cpu
-def test_uaa_without_contract_keeps_dynamic_length_gather(
-    monkeypatch: pytest.MonkeyPatch,
-) -> None:
-    """No auto_pad boundary means shards may differ, so the gather must stay."""
-    from vllm_omni.diffusion.attention.parallel import ulysses
-
-    class FakeGroup:
-        ulysses_group = object()
-        ulysses_world_size = 2
-        ulysses_rank = 0
-        ring_world_size = 1
-
-    observed = []
-    monkeypatch.setattr(
-        ulysses,
-        "_all_gather_int",
-        lambda *args, **kwargs: [3, 4],
-    )
-    monkeypatch.setattr(ulysses, "get_ulysses_mode", lambda **kwargs: "advanced_uaa")
-
-    def fake_all_to_all(pg, tensor, **kwargs):
-        observed.append(kwargs["seq_lens"])
-        return tensor, tensor.shape[2]
-
-    monkeypatch.setattr(ulysses, "_ulysses_all_to_all_any_qkv", fake_all_to_all)
-    strategy = ulysses.UlyssesParallelAttention(
-        FakeGroup(),
-        scatter_idx=2,
-        gather_idx=1,
-        use_sync=False,
-    )
-
-    with set_forward_context():
-        strategy.pre_attention(
-            torch.zeros(1, 3, 28, 4),
-            torch.zeros(1, 3, 7, 4),
-            torch.zeros(1, 3, 7, 4),
-            None,
-        )
-
-    assert observed == [[3, 4], [3, 4], [3, 4]]
-
-
-@pytest.mark.core_model
-@pytest.mark.cpu
-def test_uaa_mixed_boundaries_keep_the_dynamic_length_gather(
-    monkeypatch: pytest.MonkeyPatch,
-) -> None:
-    """One manual boundary inside an auto_pad region voids the contract."""
-    from vllm_omni.diffusion.attention.parallel import ulysses
-
-    class FakeGroup:
-        ulysses_group = object()
-        ulysses_world_size = 2
-        ulysses_rank = 0
-        ring_world_size = 1
-
-    monkeypatch.setattr(ulysses, "get_ulysses_mode", lambda **kwargs: "advanced_uaa")
-    gathered = []
-
-    def fake_gather(*args, **kwargs):
-        gathered.append(args)
-        return [3, 4]
-
-    monkeypatch.setattr(ulysses, "_all_gather_int", fake_gather)
-    monkeypatch.setattr(
-        ulysses,
-        "_ulysses_all_to_all_any_qkv",
-        lambda pg, tensor, **kwargs: (tensor, tensor.shape[2]),
-    )
-    strategy = ulysses.UlyssesParallelAttention(
-        FakeGroup(),
-        scatter_idx=2,
-        gather_idx=1,
-        use_sync=False,
-    )
-
-    with set_forward_context():
-        ctx = get_forward_context()
-        ctx._sp_equal_pad_stack.extend([True, False])
-        assert not ctx.sp_rank_local_seq_lens_equal
-        strategy.pre_attention(
-            torch.zeros(1, 3, 28, 4),
-            torch.zeros(1, 3, 7, 4),
-            torch.zeros(1, 3, 7, 4),
-            None,
-        )
-
-    assert len(gathered) == 1
-
-
 @pytest.mark.core_model
 @pytest.mark.cpu
 def test_uaa_gqa_head_padding_preserves_the_query_to_kv_ratio(
@@ -830,187 +634,6 @@ def test_ulysses_uaa_matches_baseline(
                 pass
 
 
-def _run_contract_skip_case(
-    local_rank: int,
-    world_size: int,
-    init_method: str,
-    input_file: str,
-    num_heads: int,
-    head_size: int,
-    num_kv_heads: int | None,
-) -> None:
-    """Per-rank worker: slow path vs contract fast path on real rank pairs.
-
-    Phase 1 leaves the equal-pad stack empty so the Ulysses length all-gather
-    really runs (counted); phase 2 pushes the split hook's auto_pad marker and
-    replaces _all_gather_int with a raising guard. With ring_degree=1 any call
-    to it is the Ulysses length gather, so the guard proves the skip, and the
-    per-rank output comparison proves the fast path is numerically equivalent
-    to the slow path it replaces.
-    """
-    from vllm_omni.diffusion.attention.parallel import ulysses as ulysses_mod
-
-    device = torch.device(f"{current_omni_platform.device_type}:{local_rank}")
-    current_omni_platform.set_device(device)
-
-    os.environ["DIFFUSION_ATTENTION_BACKEND"] = "TORCH_SDPA"
-
-    init_distributed_environment(world_size=world_size, rank=local_rank, distributed_init_method=init_method)
-    initialize_model_parallel(
-        data_parallel_size=1,
-        cfg_parallel_size=1,
-        sequence_parallel_size=world_size,
-        ulysses_degree=world_size,
-        ring_degree=1,
-        tensor_parallel_size=1,
-        pipeline_parallel_size=1,
-    )
-
-    parallel_config = DiffusionParallelConfig(
-        pipeline_parallel_size=1,
-        data_parallel_size=1,
-        tensor_parallel_size=1,
-        sequence_parallel_size=world_size,
-        ulysses_degree=world_size,
-        ring_degree=1,
-        cfg_parallel_size=1,
-        ulysses_mode="advanced_uaa",
-    )
-    od_config = OmniDiffusionConfig(model="test", dtype=torch.float32, parallel_config=parallel_config)
-
-    with set_forward_context(omni_diffusion_config=od_config), set_current_diffusion_config(od_config):
-        attn = Attention(
-            num_heads=num_heads,
-            head_size=head_size,
-            causal=False,
-            softmax_scale=1.0 / (head_size**0.5),
-            num_kv_heads=num_kv_heads,
-        ).to(device=device, dtype=torch.float32)
-
-        with np.load(input_file, allow_pickle=False) as payload:
-            q_full = torch.from_numpy(payload["q"]).to(device=device)
-            k_full = torch.from_numpy(payload["k"]).to(device=device)
-            v_full = torch.from_numpy(payload["v"]).to(device=device)
-
-        def make_shards() -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
-            # Fresh contiguous shards for each phase; seq_len is divisible by
-            # world_size, so every rank gets the same local length.
-            return (
-                torch.tensor_split(q_full, world_size, dim=1)[local_rank].contiguous(),
-                torch.tensor_split(k_full, world_size, dim=1)[local_rank].contiguous(),
-                torch.tensor_split(v_full, world_size, dim=1)[local_rank].contiguous(),
-            )
-
-        # The Attention layer only enables SP communication when
-        # ForwardContext.sp_active is True; in production the split hook sets
-        # this (and pushes the auto_pad equal-length marker onto
-        # _sp_equal_pad_stack). Here we drive it directly.
-        ctx = get_forward_context()
-        ctx._sp_shard_depth = 1
-
-        def make_sdp_ctx() -> object:
-            # Fresh instance per phase: the @contextmanager instance releases
-            # its state on first exit and cannot be re-entered.
-            return torch.backends.cuda.sdp_kernel(enable_flash=False, enable_mem_efficient=False, enable_math=True)
-
-        # Phase 1: contract absent -> the length all-gather must really run.
-        real_gather = ulysses_mod._all_gather_int
-        gather_calls: list[object] = []
-
-        def counting_gather(pg: object, value: int, *, device: torch.device) -> list[int]:
-            gather_calls.append(pg)
-            return real_gather(pg, value, device=device)
-
-        ulysses_mod._all_gather_int = counting_gather
-        try:
-            q, k, v = make_shards()
-            with make_sdp_ctx():
-                out_slow = attn(q, k, v, attn_metadata=None).contiguous()
-        finally:
-            ulysses_mod._all_gather_int = real_gather
-        assert len(gather_calls) >= 1, "slow path must call the Ulysses length all-gather"
-
-        # Phase 2: replicate the split hook's auto_pad marker so the fast
-        # path fires; any _all_gather_int call in this phase is the Ulysses
-        # length gather and must not happen.
-        ctx._sp_equal_pad_stack.append(True)
-
-        def fail_on_gather(pg: object, value: int, *, device: torch.device) -> list[int]:
-            raise AssertionError("the Ulysses length all-gather must be skipped under the equal-pad contract")
-
-        ulysses_mod._all_gather_int = fail_on_gather
-        try:
-            q, k, v = make_shards()
-            with make_sdp_ctx():
-                out_fast = attn(q, k, v, attn_metadata=None).contiguous()
-        finally:
-            ulysses_mod._all_gather_int = real_gather
-        ctx._sp_equal_pad_stack.pop()
-
-        torch.testing.assert_close(out_fast, out_slow, atol=1e-5, rtol=1e-5)
-
-    destroy_distributed_env()
-
-
-@pytest.mark.parametrize(
-    "num_heads,num_kv_heads",
-    [
-        (3, None),  # MHA: heads padded 3 -> 4 by ulysses_degree=2
-        (28, 7),  # GQA: derived-Q padding 28/7 -> 32/8, the main case for the skip
-    ],
-)
-def test_ulysses_uaa_contract_skips_length_gather_matches_slow_path(num_heads: int, num_kv_heads: int | None) -> None:
-    """Real 2-rank parity: contract fast path == slow path, gather skipped.
-
-    The CPU-only contract tests replace the all-to-all with a fake; this test
-    runs the real strategy on two ranks. Each rank first computes the slow
-    path (empty equal-pad stack, real length all-gather), then the fast path
-    with the split hook's auto_pad marker pushed and _all_gather_int replaced
-    by a raising guard -- which also proves the collective was truly skipped,
-    not just silently replaced.
-    """
-    sp_world_size = 2
-    if current_omni_platform.get_device_count() < sp_world_size:
-        pytest.skip(f"Test requires {sp_world_size} GPUs")
-
-    batch_size = 2
-    head_size = 8
-    seq_len = 6  # divisible by 2 so both ranks hold equal-length shards
-
-    init_method = get_distributed_init_method()
-
-    with tempfile.NamedTemporaryFile(delete=False, suffix=".npz") as f_in:
-        input_file = f_in.name
-
-    try:
-        torch.manual_seed(0)
-        kv_heads = num_kv_heads if num_kv_heads is not None else num_heads
-        np.savez(
-            input_file,
-            q=torch.randn(batch_size, seq_len, num_heads, head_size, dtype=torch.float32).numpy(),
-            k=torch.randn(batch_size, seq_len, kv_heads, head_size, dtype=torch.float32).numpy(),
-            v=torch.randn(batch_size, seq_len, kv_heads, head_size, dtype=torch.float32).numpy(),
-        )
-
-        torch.multiprocessing.spawn(
-            _run_contract_skip_case,
-            args=(
-                sp_world_size,
-                init_method,
-                input_file,
-                num_heads,
-                head_size,
-                num_kv_heads,
-            ),
-            nprocs=sp_world_size,
-        )
-    finally:
-        try:
-            os.remove(input_file)
-        except OSError:
-            pass
-
-
 @pytest.mark.parametrize(
     "num_heads,num_kv_heads,joint_len",
     [

tests/diffusion/distributed/test_sp_plan_hooks.py

@@ -1113,120 +1113,6 @@ class TestStrictModeSplitValidation:
                 hook._prepare_sp_input(x, metadata["x"], (), {})
 
 
-@pytest.mark.cpu
-class TestEqualPadContract:
-    """The stack of per-boundary auto_pad flags that lets attention skip a gather.
-
-    A boundary is recorded as equal-length only when the framework padded it,
-    because only then is every rank's shard guaranteed the same size.
-    """
-
-    def _split_hook(self, *, auto_pad: bool, monkeypatch):
-        from vllm_omni.diffusion.attention import selector
-        from vllm_omni.diffusion.distributed import parallel_state, sp_sharding
-        from vllm_omni.diffusion.distributed.sp_plan import SequenceParallelConfig
-        from vllm_omni.diffusion.hooks.sequence_parallel import (
-            SequenceParallelSplitHook,
-        )
-
-        class MaskBackend:
-            supports_attention_mask = True
-
-            @staticmethod
-            def get_name():
-                return "test"
-
-        monkeypatch.setattr(parallel_state, "get_sequence_parallel_world_size", lambda: 2)
-        monkeypatch.setattr(parallel_state, "get_sequence_parallel_rank", lambda: 1)
-        monkeypatch.setattr(parallel_state, "get_ring_parallel_world_size", lambda: 1)
-        monkeypatch.setattr(selector, "get_attn_backend_for_role", lambda **_: (MaskBackend(), None))
-        # sp_shard, used by the manual path, binds these at import time.
-        monkeypatch.setattr(sp_sharding, "get_sequence_parallel_world_size", lambda: 2)
-        monkeypatch.setattr(sp_sharding, "get_sequence_parallel_rank", lambda: 1)
-
-        hook = SequenceParallelSplitHook(
-            {
-                0: SequenceParallelInput(
-                    split_dim=1,
-                    expected_dims=3,
-                    split_output=True,
-                    auto_pad=auto_pad,
-                )
-            },
-            SequenceParallelConfig(ulysses_degree=2),
-        )
-        hook.initialize_hook(nn.Identity())
-        return hook
-
-    def test_auto_pad_split_claims_equal_rank_lengths(self, monkeypatch):
-        from vllm_omni.diffusion.forward_context import (
-            get_forward_context,
-            set_forward_context,
-        )
-
-        hook = self._split_hook(auto_pad=True, monkeypatch=monkeypatch)
-        with set_forward_context():
-            hook.pre_forward(nn.Identity())
-            shard = hook.post_forward(nn.Identity(), torch.arange(5).reshape(1, 5, 1))
-            ctx = get_forward_context()
-
-            # 5 tokens padded to 6, so both ranks hold 3.
-            assert shard.shape[1] == 3
-            assert ctx._sp_equal_pad_stack == [True]
-            assert ctx.sp_rank_local_seq_lens_equal
-
-    def test_manual_split_does_not_claim_equal_rank_lengths(self, monkeypatch):
-        from vllm_omni.diffusion.forward_context import (
-            get_forward_context,
-            set_forward_context,
-        )
-
-        hook = self._split_hook(auto_pad=False, monkeypatch=monkeypatch)
-        with set_forward_context():
-            hook.pre_forward(nn.Identity())
-            hook.post_forward(nn.Identity(), torch.arange(6).reshape(1, 6, 1))
-            ctx = get_forward_context()
-
-            assert ctx._sp_equal_pad_stack == [False]
-            assert not ctx.sp_rank_local_seq_lens_equal
-
-    def test_gather_pops_only_its_own_boundary(self, monkeypatch):
-        """A stack, not a counter: the gather must not resurrect the contract.
-
-        With one auto_pad and one manual boundary open, a counter could not
-        tell which one a gather closes and would wrongly report equal lengths.
-        """
-        from vllm_omni.diffusion.distributed.sp_plan import SequenceParallelConfig
-        from vllm_omni.diffusion.forward_context import (
-            get_forward_context,
-            set_forward_context,
-        )
-        from vllm_omni.diffusion.hooks import sequence_parallel
-        from vllm_omni.diffusion.hooks.sequence_parallel import (
-            SequenceParallelGatherHook,
-        )
-
-        monkeypatch.setattr(
-            sequence_parallel,
-            "sp_gather",
-            lambda tensor, dim, validate: torch.cat([tensor, tensor], dim=dim),
-        )
-        hook = SequenceParallelGatherHook(
-            SequenceParallelOutput(gather_dim=1, expected_dims=3),
-            SequenceParallelConfig(ulysses_degree=2),
-        )
-
-        with set_forward_context():
-            ctx = get_forward_context()
-            ctx._sp_shard_depth = 2
-            ctx._sp_equal_pad_stack.extend([False, True])
-            hook.post_forward(nn.Identity(), torch.zeros(1, 3, 4))
-
-            assert ctx._sp_shard_depth == 1
-            assert ctx._sp_equal_pad_stack == [False]
-            assert not ctx.sp_rank_local_seq_lens_equal
-
-
 @pytest.mark.cpu
 class TestSequenceParallelConfig:
     """Test SequenceParallelConfig dataclass."""

vllm_omni/diffusion/attention/parallel/ulysses.py

@@ -14,11 +14,7 @@ from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata
 from vllm_omni.diffusion.attention.parallel.base import ParallelAttentionContext
 from vllm_omni.diffusion.distributed.comm import SeqAllToAll4D
 from vllm_omni.diffusion.distributed.group_coordinator import SequenceParallelGroupCoordinator
-from vllm_omni.diffusion.forward_context import (
-    get_forward_context,
-    get_ulysses_mode,
-    is_forward_context_available,
-)
+from vllm_omni.diffusion.forward_context import get_ulysses_mode
 
 logger = init_logger(__name__)
 
@@ -146,7 +142,7 @@ def _ulysses_all_to_all_any_qkv(
 
     input_split_sizes = [s_local] * world_size
     output_split_sizes = seq_lens
-    s_global = sum(output_split_sizes)
+    s_global = int(sum(output_split_sizes))
 
     out = torch.empty((s_global, bsz, head_cnt_local, head_dim), device=x.device, dtype=x.dtype)
     dist.all_to_all_single(
@@ -181,7 +177,7 @@ def _ulysses_all_to_all_any_o(
         return x
 
     bsz, s_global, head_cnt_local, head_dim = x.shape
-    s_local = local_seq_len
+    s_local = int(local_seq_len)
 
     # (B, S_global, H_local, D) -> (S_global, B, H_local, D)
     x_t = x.permute(1, 0, 2, 3).contiguous()
@@ -415,22 +411,9 @@ class UlyssesParallelAttention:
                     f"(got scatter_idx={self._scatter_idx}, gather_idx={self._gather_idx})."
                 )
 
-            if is_forward_context_available() and get_forward_context().sp_rank_local_seq_lens_equal:
-                # auto_pad already made every rank's shard the same length, so
-                # the lengths are known without a collective. Keep local_seq_len
-                # as a SymInt (no int()) so torch.compile(dynamic=True) does not
-                # specialize on it, and skip the host sync + graph break that
-                # _all_gather_int would introduce in every attention call.
-                local_seq_len = query.shape[1]
-                seq_lens = [local_seq_len] * ulysses_world_size
-            else:
-                local_seq_len = int(query.shape[1])
-                seq_lens = _all_gather_int(
-                    self._ulysses_pg,
-                    local_seq_len,
-                    device=query.device,
-                )
-            s_global = sum(seq_lens)
+            local_seq_len = int(query.shape[1])
+            seq_lens = _all_gather_int(self._ulysses_pg, local_seq_len, device=query.device)
+            s_global = int(sum(seq_lens))
 
             # In hybrid Ulysses+Ring, Ring attention uses P2P send/recv with fixed-shape
             # buffers. This requires all ring ranks to have the same seq_len after the
@@ -588,8 +571,8 @@ class UlyssesParallelAttention:
             joint_len=joint_len,
             joint_strategy=joint_strategy,
             use_uaa=(mode == "advanced_uaa"),
-            uaa_seq_lens=tuple(seq_lens) if mode == "advanced_uaa" else (),
-            uaa_local_seq_len=local_seq_len if mode == "advanced_uaa" else 0,
+            uaa_seq_lens=tuple(int(x) for x in seq_lens) if mode == "advanced_uaa" else (),
+            uaa_local_seq_len=int(local_seq_len) if mode == "advanced_uaa" else 0,
             orig_head_cnt=int(orig_head_cnt) if mode == "advanced_uaa" else 0,
             joint_orig_head_cnt=int(joint_orig_head_cnt) if mode == "advanced_uaa" else 0,
         )

vllm_omni/diffusion/forward_context.py

@@ -4,7 +4,7 @@
 from __future__ import annotations
 
 from contextlib import contextmanager
-from dataclasses import dataclass, field
+from dataclasses import dataclass
 from typing import TYPE_CHECKING, Any
 
 import torch
@@ -66,10 +66,6 @@ class ForwardContext:
     # 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
-    # One entry per active SP split boundary: whether it auto-pads, which makes
-    # every rank's shard the same length and lets attention collectives skip a
-    # runtime length all-gather. Pushed on split, popped on gather.
-    _sp_equal_pad_stack: list[bool] = field(default_factory=list)
 
     @property
     def sp_active(self) -> bool:
@@ -92,15 +88,6 @@ class ForwardContext:
         sp_size = self.omni_diffusion_config.parallel_config.sequence_parallel_size
         return sp_size is not None and sp_size > 1
 
-    @property
-    def sp_rank_local_seq_lens_equal(self) -> bool:
-        """Whether every active SP boundary guarantees equal local shard sizes.
-
-        Only framework-managed auto_pad boundaries provide this contract; a
-        region that also shards manually keeps the dynamic length exchange.
-        """
-        return bool(self._sp_equal_pad_stack) and all(self._sp_equal_pad_stack)
-
     def __post_init__(self):
         pass
 

vllm_omni/diffusion/hooks/sequence_parallel.py

@@ -264,7 +264,6 @@ class SequenceParallelSplitHook(ModelHook):
 
         output_list = [output] if is_tensor else list(output)
         actually_sharded = False
-        equal_rank_seq_lens = True
 
         for index, spm in self.metadata.items():
             if not isinstance(index, int):
@@ -278,13 +277,10 @@ class SequenceParallelSplitHook(ModelHook):
             output_list[index] = self._prepare_sp_input(original, spm, self._last_args, self._last_kwargs)
             if output_list[index] is not original:
                 actually_sharded = True
-                equal_rank_seq_lens &= isinstance(spm, SequenceParallelInput) and spm.auto_pad
 
         # Mark SP as active only if at least one tensor was actually sharded
         if actually_sharded and is_forward_context_available():
-            ctx = get_forward_context()
-            ctx._sp_shard_depth += 1
-            ctx._sp_equal_pad_stack.append(equal_rank_seq_lens)
+            get_forward_context()._sp_shard_depth += 1
 
         return output_list[0] if is_tensor else type(output)(output_list)
 
@@ -557,8 +553,6 @@ class SequenceParallelGatherHook(ModelHook):
         if actually_gathered and is_forward_context_available():
             ctx = get_forward_context()
             ctx._sp_shard_depth = max(0, ctx._sp_shard_depth - 1)
-            if ctx._sp_equal_pad_stack:
-                ctx._sp_equal_pad_stack.pop()
 
         return output[0] if is_tensor else type(output)(output)
 

vllm_omni/host_weight_runtime/filesystem/store.py

@@ -1247,8 +1247,10 @@ class FilesystemHostWeightStore:
                     retryable=True,
                 )
             )
-        store_bytes = self._tree_bytes(self.root)
-        if policy.max_store_bytes is not None and store_bytes + additional_bytes > policy.max_store_bytes:
+        if (
+            policy.max_store_bytes is not None
+            and self._tree_bytes(self.root) + additional_bytes > policy.max_store_bytes
+        ):
             raise HostWeightError(
                 _failure(
                     ResolutionStage.CAPACITY,