[Core] Add opt-in serialized HWR build admission

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

docs/design/feature/host_weight_runtime.md

@@ -313,6 +313,16 @@ One process per exact identity owns a build; other workers wait and then acquire
 leases for the published artifact. Publication is invisible until all payloads
 and metadata are validated, hashed, fsynced, and atomically renamed.
 
+By default, different identities may build concurrently. A domain configured
+with `CapacityPolicy(build_admission="serialized", max_store_bytes=...)` admits
+only one producer/publication operation at a time. Warm hits remain concurrent,
+and admission uses the existing coordination deadline. All workers must agree
+on the serialized domain's schema-2 capacity policy; use a new root when opting
+in, since in-place migration of an active domain is unsupported. This prevents
+competing producers from racing allocation checks but does not provide a
+filesystem-wide byte quota or automatic stale-data reclamation. See the
+[module capacity contract](../module/host_weight_runtime.md#validation-and-capacity).
+
 `coordination_timeout_seconds` bounds filesystem lock acquisition. It does not
 cancel synchronous validation, a producer that has already started, or atomic
 publication. A hung in-process producer therefore blocks its owning process and

docs/design/module/host_weight_runtime.md

@@ -311,9 +311,41 @@ point-in-time evidence only; kernel locks are not a persistent owner registry.
 Capacity policy covers all store-owned bytes, including ready, temporary, and
 quarantined data. The local writer preflights `max_artifact_bytes`,
 `max_store_bytes`, and `min_free_bytes`, and preallocates payloads where the
-filesystem supports it. `ENOSPC` remains a normal typed store failure because
-concurrent preflight is inherently racy. Automatic eviction and strict
-concurrent reservations are not implemented.
+filesystem supports it. `ENOSPC` remains a normal typed store failure.
+
+`CapacityPolicy.build_admission` defaults to `concurrent`, preserving parallel
+production for different identities and best-effort capacity preflight. An
+explicit `serialized` policy, which requires `max_store_bytes`, admits one
+cooperative build/publication operation per domain through
+`locks/domain-build.lock`. Production lock order becomes domain admission,
+per-key build, then artifact. Validated warm hits bypass admission; queued
+same-key callers still recheck and join a completed artifact. Admission waiting
+uses the caller's existing coordination deadline and returns a retryable build
+timeout without interrupting its owner.
+
+The kernel releases admission on process exit. An explicit same-key retry uses
+the existing stale-temp recovery; files from unrelated dead builds continue to
+count toward capacity. A synchronous producer that hangs still requires
+external process supervision. Admission release errors are logged without
+revising an already determined publication result.
+
+Concurrent domains retain their exact schema-1 capacity document. Serialized
+domains use a schema-2 capacity document with `build_admission=serialized`, so
+mixed policies and old readers fail initialization rather than bypass admission.
+Opt into serialization using a new root with an agreed policy; in-place policy
+migration while old workers may exist is unsupported. For example:
+
+```python
+from vllm_omni.host_weight_runtime import CapacityPolicy
+
+capacity = CapacityPolicy(max_store_bytes=192 * 1024**3, build_admission="serialized")
+```
+
+Serialization prevents cooperative producers from racing each other's
+allocation checks. It is not a filesystem quota: control metadata, finalization,
+and outside writers can still affect total usage. Automatic eviction and
+parallel byte reservations are not implemented. This policy trades parallel
+cold-build throughput for predictable producer admission.
 
 For tmpfs, artifact bytes are host-memory consumption and may consume swap;
 they must not be reported operationally as ordinary disk capacity. The store

tests/diffusion/attention/test_ulysses_uaa.py

@@ -27,6 +27,202 @@ 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(
@@ -634,6 +830,187 @@ 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,6 +1113,120 @@ 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."""

tests/host_weight_runtime/test_build_admission.py

@@ -0,0 +1,305 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""CPU process and policy coverage for domain-wide producer admission."""
+
+from __future__ import annotations
+
+import errno
+import json
+import multiprocessing as mp
+import os
+import time
+from multiprocessing.connection import Connection
+from pathlib import Path
+from typing import Literal
+
+import pytest
+import torch
+
+from tests.host_weight_runtime.test_filesystem_store import FakeProducer, _identity, _make_store, _publish_test_artifact
+from vllm_omni.host_weight_runtime import (
+    BuildRequest,
+    CapacityPolicy,
+    FailureCode,
+    HostWeightError,
+    HostWeightRuntime,
+    HostWeightRuntimeConfig,
+    ProductionMetadata,
+    ResolutionOutcome,
+    RuntimeMode,
+    StoreResult,
+    StoreStatus,
+    TensorWriteSpec,
+    ValidationLevel,
+    WaitPolicy,
+    WeightArtifactIdentity,
+    WeightProductionSpec,
+)
+from vllm_omni.host_weight_runtime.filesystem.locks import FileLock, lock_is_active
+from vllm_omni.host_weight_runtime.protocols import ArtifactWriter
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+class AdmissionProducer:
+    def __init__(self, identity: WeightArtifactIdentity, control: Connection | None = None) -> None:
+        self._spec = FakeProducer(identity).spec
+        self.control = control
+
+    @property
+    def spec(self) -> WeightProductionSpec:
+        return self._spec
+
+    def produce(self, writer: ArtifactWriter) -> ProductionMetadata:
+        action = "produce"
+        if self.control is not None:
+            self.control.send("entered")
+            if not self.control.poll(60):
+                raise TimeoutError("test did not release producer")
+            action = self.control.recv()
+        spec = TensorWriteSpec("payload", (65536,), torch.bfloat16)
+        with writer.open_tensor_file("weights.safetensors", (spec,)) as output:
+            output.write_tensor("payload", torch.arange(65536, dtype=torch.float32).to(torch.bfloat16))
+            if action == "crash":
+                os._exit(27)
+        return ProductionMetadata("test-producer-v1", "test-restorer-v1")
+
+
+def _admission_process(root: Path, policy: CapacityPolicy, rank: int, control: Connection) -> None:
+    try:
+        store = _make_store(root, capacity=policy)
+        identity = _identity(tp_rank=rank)
+        original_lookup = store.lookup
+        first_lookup = True
+
+        def notify_lookup(
+            identity: WeightArtifactIdentity, *, validation: ValidationLevel, deadline: float | None = None
+        ) -> StoreResult:
+            nonlocal first_lookup
+            result = original_lookup(identity, validation=validation, deadline=deadline)
+            if first_lookup:
+                # Prove the caller observed a miss before the parent lets the
+                # active producer publish, independent of process scheduling.
+                control.send("looked_up")
+                first_lookup = False
+            return result
+
+        store.lookup = notify_lookup  # type: ignore[method-assign]
+        control.send("ready")
+        assert control.poll(60) and control.recv() == "start"
+        result = store.get_or_build(
+            BuildRequest(identity),
+            AdmissionProducer(identity, control),
+            validation=ValidationLevel.FULL_CHECKSUM,
+            deadline=time.monotonic() + 30,
+        )
+        if result.lease is not None:
+            result.lease.close()
+        control.send((result.status.value, result.failure.code.value if result.failure is not None else None))
+    finally:
+        control.close()
+
+
+@pytest.mark.parametrize("admission", ["concurrent", "serialized"])
+@pytest.mark.parametrize("same_key", [False, True])
+def test_cross_identity_admission_and_capacity(
+    tmp_path: Path, admission: Literal["concurrent", "serialized"], same_key: bool
+) -> None:
+    policy = CapacityPolicy(max_store_bytes=192 * 1024, build_admission=admission)
+    store = _make_store(tmp_path / "store", capacity=policy)
+    ctx = mp.get_context("spawn")
+    pairs = [ctx.Pipe() for _ in range(2)]
+    processes = [
+        ctx.Process(target=_admission_process, args=(store.root, policy, 0 if same_key else rank, pair[1]))
+        for rank, pair in enumerate(pairs)
+    ]
+    try:
+        for process, pair in zip(processes, pairs, strict=True):
+            process.start()
+            pair[1].close()
+        first, second = [pair[0] for pair in pairs]
+        for connection in (first, second):
+            assert connection.poll(60) and connection.recv() == "ready"
+        first.send("start")
+        assert first.poll(10) and first.recv() == "looked_up"
+        assert first.poll(10) and first.recv() == "entered"
+        second.send("start")
+        assert second.poll(10) and second.recv() == "looked_up"
+        if admission == "concurrent" and not same_key:
+            assert second.poll(10) and second.recv() == "entered"
+        else:
+            assert not second.poll(0.2), "a second producer entered while coordination should exclude it"
+        first.send("produce")
+        assert first.poll(30) and first.recv() == (StoreStatus.BUILT.value, None)
+        if same_key:
+            assert second.poll(30) and second.recv() == (StoreStatus.JOINED.value, None)
+        else:
+            if admission == "serialized":
+                assert second.poll(10) and second.recv() == "entered"
+            second.send("produce")
+            assert second.poll(30)
+            assert second.recv() == (StoreStatus.FAILED.value, FailureCode.STORE_LIMIT_EXCEEDED.value)
+        assert policy.max_store_bytes is not None
+        assert store.inspect_domain().store_bytes < policy.max_store_bytes
+        assert not list(store.tmp_dir.iterdir())
+        assert {path.name for path in store.artifacts_dir.iterdir()} == {_identity(tp_rank=0).key}
+        for process in processes:
+            process.join(10)
+            assert process.exitcode == 0
+    finally:
+        for process in processes:
+            if process.pid is not None:
+                if process.is_alive():
+                    process.kill()
+                process.join(5)
+                process.close()
+        for pair in pairs:
+            for connection in pair:
+                connection.close()
+
+
+def test_admission_timeout_preserves_owner_and_warm_hit_bypasses_lock(tmp_path: Path) -> None:
+    policy = CapacityPolicy(max_store_bytes=512 * 1024, build_admission="serialized")
+    store = _make_store(tmp_path / "store", capacity=policy)
+    identity, _ = _publish_test_artifact(store)
+    other = _identity(tp_rank=1)
+    lock_path = store.locks_dir / "domain-build.lock"
+    with FileLock(lock_path, exclusive=True, deadline=None):
+        hit = store.get_or_build(
+            BuildRequest(identity),
+            FakeProducer(identity),
+            validation=ValidationLevel.FULL_CHECKSUM,
+            deadline=time.monotonic() + 0.02,
+        )
+        assert hit.status is StoreStatus.HIT and hit.lease is not None
+        hit.lease.close()
+        timeout = store.get_or_build(
+            BuildRequest(other),
+            FakeProducer(other),
+            validation=ValidationLevel.FULL_CHECKSUM,
+            deadline=time.monotonic() + 0.02,
+        )
+        assert timeout.status is StoreStatus.TIMEOUT
+        assert timeout.failure is not None and timeout.failure.code is FailureCode.ACTIVE_BUILD_TIMEOUT
+        assert timeout.failure.retryable
+        assert lock_is_active(lock_path)
+    _publish_test_artifact(store, other)
+
+
+def test_producer_failure_releases_admission(tmp_path: Path) -> None:
+    store = _make_store(
+        tmp_path / "store", capacity=CapacityPolicy(max_store_bytes=16384, build_admission="serialized")
+    )
+    identity = _identity()
+    failed = store.get_or_build(
+        BuildRequest(identity),
+        FakeProducer(identity, write_mode="incomplete"),
+        validation=ValidationLevel.FULL_CHECKSUM,
+        deadline=time.monotonic() + 5,
+    )
+    assert failed.status is StoreStatus.FAILED
+    assert not lock_is_active(store.locks_dir / "domain-build.lock")
+    _publish_test_artifact(store, identity)
+
+
+@pytest.mark.parametrize(
+    ("mode", "expected"),
+    [(RuntimeMode.PREFERRED, ResolutionOutcome.CANONICAL_FALLBACK), (RuntimeMode.REQUIRED, ResolutionOutcome.FAILED)],
+)
+def test_admission_timeout_obeys_runtime_mode(tmp_path: Path, mode: RuntimeMode, expected: ResolutionOutcome) -> None:
+    policy = CapacityPolicy(max_store_bytes=16384, build_admission="serialized")
+    store = _make_store(tmp_path / "store", capacity=policy)
+    runtime = HostWeightRuntime.from_config(
+        HostWeightRuntimeConfig(mode=mode, domain=store.domain_policy, capacity=policy, wait=WaitPolicy(0.01))
+    )
+    identity = _identity()
+    with FileLock(store.locks_dir / "domain-build.lock", exclusive=True, deadline=None):
+        resolution = runtime.resolve(identity, producer=FakeProducer(identity))
+    assert resolution.report.outcome is expected
+    failure = resolution.report.attempts[-1].failure
+    assert failure is not None and failure.code is FailureCode.ACTIVE_BUILD_TIMEOUT
+
+
+def test_crashed_producer_releases_admission_and_same_key_retry_recovers(tmp_path: Path) -> None:
+    policy = CapacityPolicy(max_store_bytes=192 * 1024, build_admission="serialized")
+    store = _make_store(tmp_path / "store", capacity=policy)
+    ctx = mp.get_context("spawn")
+    parent, child = ctx.Pipe()
+    process = ctx.Process(target=_admission_process, args=(store.root, policy, 0, child))
+    process.start()
+    child.close()
+    try:
+        assert parent.poll(60) and parent.recv() == "ready"
+        parent.send("start")
+        assert parent.poll(10) and parent.recv() == "looked_up"
+        assert parent.poll(10) and parent.recv() == "entered"
+        parent.send("crash")
+        process.join(10)
+        assert process.exitcode == 27
+    finally:
+        if process.is_alive():
+            process.kill()
+        process.join(5)
+        process.close()
+        parent.close()
+    assert not lock_is_active(store.locks_dir / "domain-build.lock")
+    assert list(store.tmp_dir.iterdir())
+    identity = _identity(tp_rank=0)
+    recovered = store.get_or_build(
+        BuildRequest(identity),
+        AdmissionProducer(identity),
+        validation=ValidationLevel.FULL_CHECKSUM,
+        deadline=time.monotonic() + 10,
+    )
+    assert recovered.status is StoreStatus.BUILT and recovered.lease is not None
+    assert torch.equal(recovered.lease.tensors["payload"], torch.arange(65536, dtype=torch.float32).to(torch.bfloat16))
+    recovered.lease.close()
+    assert not list(store.tmp_dir.iterdir())
+
+
+def test_admission_release_error_preserves_publication(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None:
+    store = _make_store(
+        tmp_path / "store", capacity=CapacityPolicy(max_store_bytes=16384, build_admission="serialized")
+    )
+    original_close = FileLock.close
+
+    def fail_admission_close(lock: FileLock) -> None:
+        original_close(lock)
+        if lock.path.name == "domain-build.lock":
+            raise OSError(errno.EIO, "injected admission release failure")
+
+    monkeypatch.setattr(FileLock, "close", fail_admission_close)
+    _publish_test_artifact(store)
+    assert not lock_is_active(store.locks_dir / "domain-build.lock")
+
+
+def test_admission_policy_is_authoritative_and_legacy_document_is_unchanged(tmp_path: Path) -> None:
+    concurrent = CapacityPolicy(max_store_bytes=16384)
+    serialized = CapacityPolicy(max_store_bytes=16384, build_admission="serialized")
+    legacy = _make_store(tmp_path / "legacy", capacity=concurrent)
+    path = legacy.root / "domain-policy.json"
+    original = path.read_bytes()
+    assert json.loads(original) == {
+        "schema_version": 1,
+        "policy_version": 1,
+        "max_artifact_bytes": None,
+        "max_store_bytes": 16384,
+        "min_free_bytes": 0,
+        "eviction": "none",
+    }
+    with pytest.raises(HostWeightError, match="incompatible schema"):
+        _make_store(legacy.root, capacity=serialized)
+    assert path.read_bytes() == original
+    strict = _make_store(tmp_path / "serialized", capacity=serialized)
+    document = json.loads((strict.root / "domain-policy.json").read_bytes())
+    assert document["schema_version"] == 2 and document["build_admission"] == "serialized"
+    assert _make_store(strict.root, capacity=serialized).domain_uuid == strict.domain_uuid
+    with pytest.raises(HostWeightError, match="incompatible schema"):
+        _make_store(strict.root, capacity=concurrent)
+
+
+def test_admission_policy_validation() -> None:
+    with pytest.raises(ValueError, match="requires max_store_bytes"):
+        CapacityPolicy(build_admission="serialized")
+    with pytest.raises(ValueError, match="must be concurrent or serialized"):
+        CapacityPolicy(build_admission="unknown")  # type: ignore[arg-type]

vllm_omni/diffusion/attention/parallel/ulysses.py

@@ -14,7 +14,11 @@ 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_ulysses_mode
+from vllm_omni.diffusion.forward_context import (
+    get_forward_context,
+    get_ulysses_mode,
+    is_forward_context_available,
+)
 
 logger = init_logger(__name__)
 
@@ -142,7 +146,7 @@ def _ulysses_all_to_all_any_qkv(
 
     input_split_sizes = [s_local] * world_size
     output_split_sizes = seq_lens
-    s_global = int(sum(output_split_sizes))
+    s_global = 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(
@@ -177,7 +181,7 @@ def _ulysses_all_to_all_any_o(
         return x
 
     bsz, s_global, head_cnt_local, head_dim = x.shape
-    s_local = int(local_seq_len)
+    s_local = local_seq_len
 
     # (B, S_global, H_local, D) -> (S_global, B, H_local, D)
     x_t = x.permute(1, 0, 2, 3).contiguous()
@@ -411,9 +415,22 @@ class UlyssesParallelAttention:
                     f"(got scatter_idx={self._scatter_idx}, gather_idx={self._gather_idx})."
                 )
 
-            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))
+            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)
 
             # 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
@@ -571,8 +588,8 @@ class UlyssesParallelAttention:
             joint_len=joint_len,
             joint_strategy=joint_strategy,
             use_uaa=(mode == "advanced_uaa"),
-            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,
+            uaa_seq_lens=tuple(seq_lens) if mode == "advanced_uaa" else (),
+            uaa_local_seq_len=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
+from dataclasses import dataclass, field
 from typing import TYPE_CHECKING, Any
 
 import torch
@@ -66,6 +66,10 @@ 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:
@@ -88,6 +92,15 @@ 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,6 +264,7 @@ 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):
@@ -277,10 +278,13 @@ 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():
-            get_forward_context()._sp_shard_depth += 1
+            ctx = get_forward_context()
+            ctx._sp_shard_depth += 1
+            ctx._sp_equal_pad_stack.append(equal_rank_seq_lens)
 
         return output_list[0] if is_tensor else type(output)(output_list)
 
@@ -553,6 +557,8 @@ 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/config.py

@@ -8,6 +8,7 @@ import math
 from dataclasses import dataclass, field
 from enum import Enum
 from pathlib import Path
+from typing import Literal
 
 
 class RuntimeMode(str, Enum):
@@ -66,6 +67,7 @@ class CapacityPolicy:
     max_store_bytes: int | None = None
     min_free_bytes: int = 0
     eviction: str = "none"
+    build_admission: Literal["concurrent", "serialized"] = "concurrent"
 
     def __post_init__(self) -> None:
         for name, value in (
@@ -78,6 +80,10 @@ class CapacityPolicy:
             raise ValueError("min_free_bytes must not be negative")
         if not isinstance(self.eviction, str) or self.eviction != "none":
             raise ValueError("automatic host weight eviction is not implemented")
+        if self.build_admission not in ("concurrent", "serialized"):
+            raise ValueError("build_admission must be concurrent or serialized")
+        if self.build_admission == "serialized" and self.max_store_bytes is None:
+            raise ValueError("serialized build admission requires max_store_bytes")
 
 
 @dataclass(frozen=True)

vllm_omni/host_weight_runtime/filesystem/store.py

@@ -627,7 +627,7 @@ class FilesystemHostWeightStore:
                 self.policy_version = 1
 
     def _capacity_document(self, *, policy_version: int) -> dict[str, object]:
-        return {
+        document: dict[str, object] = {
             "schema_version": 1,
             "policy_version": policy_version,
             "max_artifact_bytes": self.capacity_policy.max_artifact_bytes,
@@ -635,6 +635,9 @@ class FilesystemHostWeightStore:
             "min_free_bytes": self.capacity_policy.min_free_bytes,
             "eviction": self.capacity_policy.eviction,
         }
+        if self.capacity_policy.build_admission == "serialized":
+            document.update(schema_version=2, build_admission="serialized")
+        return document
 
     def _build_lock_path(self, key: str) -> Path:
         return self.locks_dir / f"{key}.build.lock"
@@ -1013,6 +1016,40 @@ class FilesystemHostWeightStore:
         if initial.status in {StoreStatus.TIMEOUT, StoreStatus.FAILED}:
             return initial
 
+        if self.capacity_policy.build_admission == "concurrent":
+            return self._build_after_miss(identity, producer, validation=validation, deadline=deadline)
+        try:
+            admission = FileLock(self.locks_dir / "domain-build.lock", exclusive=True, deadline=deadline)
+        except FileLockTimeoutError as exc:
+            return StoreResult(
+                StoreStatus.TIMEOUT,
+                failure=_failure(
+                    ResolutionStage.PRODUCTION, FailureCode.ACTIVE_BUILD_TIMEOUT, str(exc), retryable=True
+                ),
+            )
+        except OSError as exc:
+            return StoreResult(
+                StoreStatus.FAILED,
+                failure=_failure(ResolutionStage.DOMAIN, FailureCode.DOMAIN_UNAVAILABLE, str(exc), retryable=True),
+            )
+        try:
+            return self._build_after_miss(identity, producer, validation=validation, deadline=deadline)
+        finally:
+            try:
+                admission.close()
+            except OSError:
+                # Releasing coordination must not revise the completed build's
+                # publication outcome or replace its primary failure.
+                logger.warning("Failed to release domain build admission under %s", self.root, exc_info=True)
+
+    def _build_after_miss(
+        self,
+        identity: WeightArtifactIdentity,
+        producer: WeightProducer,
+        *,
+        validation: ValidationLevel,
+        deadline: float,
+    ) -> StoreResult:
         try:
             build_lock = FileLock(self._build_lock_path(identity.key), exclusive=True, deadline=deadline)
         except FileLockTimeoutError as exc: