[Core] Add explicit cleanup for quarantined HWR artifacts

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

docs/design/feature/host_weight_runtime.md

@@ -313,16 +313,6 @@ 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

@@ -299,6 +299,43 @@ Explicit cleanup uses the locked move and retries exact-key
 `.cleanup.*` tombstones left by an interrupted removal; a later build performs
 the same tombstone reconciliation before producing replacement content.
 
+Failure-quarantined entries remain until explicitly selected for removal;
+ordinary `cleanup(identity)` does not remove that diagnostic history.
+`FilesystemHostWeightStore.cleanup_quarantined(storage_name)` accepts one
+quarantined inventory basename and preserves the current artifact, deny marker,
+and other failure-quarantined entries. It takes the same nonblocking build and
+exclusive artifact locks, so an active builder or lease returns a typed refusal.
+Malformed names are rejected before mutation, and symlinks are not followed.
+
+Selected entries move through the existing `.cleanup.*` tombstone protocol.
+Retrying the original name reconciles pending cleanup tombstones for that key,
+even if the original entry disappeared during an earlier attempt. A missing
+selection is otherwise an idempotent success. Existing cleanup intents for the
+key may be completed along with the selection; other quarantine history is
+retained. Parent synchronization is retried even when an earlier deletion
+removed the last tombstone before its sync failed.
+
+For a configured store, operators can inspect and deliberately select an entry:
+
+```python
+from vllm_omni.host_weight_runtime import ArtifactInventoryState, HostWeightError
+
+for entry in store.inspect_domain().inventory:
+    if entry.state is ArtifactInventoryState.QUARANTINED:
+        print(entry.storage_name, entry.size_bytes)
+
+selected_storage_name = input("Quarantined entry to remove: ")
+failure = store.cleanup_quarantined(selected_storage_name)
+if failure is not None:
+    raise HostWeightError(failure)
+```
+
+Use the same authoritative domain and capacity policy as the service. Review
+inventory sizes and inspect again after cleanup before explicitly retrying a
+build or startup. These are logical-byte, point-in-time observations; concurrent
+activity can affect the difference and physical disk/RAM reclamation is not
+implied. This operation adds no automatic retention policy or service retry.
+
 An operational failure while inspecting a noncooperative competing publication
 is a retryable storage failure. The competing entry remains authoritative and
 is not mislabeled as corrupt or quarantined merely because its manifest could
@@ -311,41 +348,9 @@ 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.
-
-`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.
+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.
 
 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,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."""

tests/host_weight_runtime/test_build_admission.py

@@ -1,305 +0,0 @@
-# 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]

tests/host_weight_runtime/test_filesystem_store.py

@@ -440,6 +440,22 @@ def _corrupt_test_payload(artifact: Path) -> None:
     os.chmod(artifact, 0o555)
 
 
+def _quarantine_and_rebuild(store: FilesystemHostWeightStore, identity: WeightArtifactIdentity) -> Path:
+    before = set(store.quarantine_dir.iterdir())
+    _corrupt_test_payload(store.artifacts_dir / identity.key)
+    result = store.get_or_build(
+        BuildRequest(identity),
+        FakeProducer(identity),
+        validation=ValidationLevel.FULL_CHECKSUM,
+        deadline=time.monotonic() + 10,
+    )
+    assert result.status is StoreStatus.BUILT
+    assert result.lease is not None
+    result.lease.close()
+    (quarantined,) = set(store.quarantine_dir.iterdir()) - before
+    return quarantined
+
+
 def _build_process(root: str, counter: str, initial_lookup_barrier: object, result_queue: object) -> None:
     try:
         identity = _identity()
@@ -2496,6 +2512,199 @@ def test_store_and_free_space_capacity_limits_fail_before_publication(tmp_path:
     assert not list(free_limited.artifacts_dir.iterdir())
 
 
+def test_cleanup_quarantined_preserves_current_artifact_other_entries_and_deny(tmp_path: Path) -> None:
+    store = _make_store(tmp_path / "store")
+    identity, artifact = _publish_test_artifact(store)
+    first = _quarantine_and_rebuild(store, identity)
+    second = _quarantine_and_rebuild(store, identity)
+    ready_inode = artifact.stat().st_ino
+    lock_inodes = {path.name: path.stat().st_ino for path in store.locks_dir.iterdir()}
+
+    assert store.cleanup_quarantined(first.name) is None
+    assert set(store.quarantine_dir.iterdir()) == {second}
+    assert artifact.stat().st_ino == ready_inode
+    hit = store.lookup(identity, validation=ValidationLevel.FULL_CHECKSUM)
+    assert hit.lease is not None
+    assert torch.equal(hit.lease.tensors["layer.weight"], _weight())
+    hit.lease.close()
+
+    _corrupt_test_payload(artifact)
+    assert store.lookup(identity, validation=ValidationLevel.FULL_CHECKSUM).status is StoreStatus.INVALID
+    deny = store.deny_dir / f"{identity.key}.json"
+    deny_bytes = deny.read_bytes()
+    assert store.cleanup_quarantined(second.name) is None
+    assert deny.read_bytes() == deny_bytes
+    assert artifact.stat().st_ino == ready_inode
+    assert {path.name: path.stat().st_ino for path in store.locks_dir.iterdir()} == lock_inodes
+    assert store.cleanup_quarantined(first.name) is None
+
+
+@pytest.mark.parametrize("active", ["build", "lease"])
+def test_cleanup_quarantined_refuses_active_key(tmp_path: Path, active: str) -> None:
+    store = _make_store(tmp_path / "store")
+    identity, _ = _publish_test_artifact(store)
+    quarantined = _quarantine_and_rebuild(store, identity)
+    owner: FileLock | HostWeightLease
+    if active == "build":
+        owner = FileLock(store._build_lock_path(identity.key), exclusive=True, deadline=None)
+        expected = FailureCode.ACTIVE_BUILD_TIMEOUT
+    else:
+        lease = store.lookup(identity, validation=ValidationLevel.FULL_CHECKSUM).lease
+        assert lease is not None
+        owner = lease
+        expected = FailureCode.ACTIVE_LEASE
+    try:
+        failure = store.cleanup_quarantined(quarantined.name)
+        assert failure is not None and failure.code is expected and failure.retryable
+        assert quarantined.is_dir()
+    finally:
+        owner.close()
+    assert store.cleanup_quarantined(quarantined.name) is None
+
+
+@pytest.mark.parametrize("recovery", ["cleanup", "build"])
+def test_cleanup_quarantined_retries_post_move_removal_failure(
+    tmp_path: Path, monkeypatch: pytest.MonkeyPatch, recovery: str
+) -> None:
+    store = _make_store(tmp_path / "store")
+    identity, artifact = _publish_test_artifact(store)
+    quarantined = _quarantine_and_rebuild(store, identity)
+    original_rmtree = shutil.rmtree
+
+    def fail_removal(path: str | os.PathLike[str]) -> None:
+        if Path(path).parent == store.quarantine_dir:
+            raise OSError(errno.EIO, "injected removal failure")
+        original_rmtree(path)
+
+    with monkeypatch.context() as fault:
+        fault.setattr(shutil, "rmtree", fail_removal)
+        failure = store.cleanup_quarantined(quarantined.name)
+    assert failure is not None and failure.code is FailureCode.QUARANTINE_FAILED and failure.retryable
+    assert not quarantined.exists()
+    assert len(list(store.quarantine_dir.glob(f"{identity.key}.cleanup.*"))) == 1
+
+    if recovery == "cleanup":
+        assert store.cleanup_quarantined(quarantined.name) is None
+    else:
+        # A real build (not a warm hit) reconciles the same tombstone protocol.
+        _corrupt_test_payload(artifact)
+        result = store.get_or_build(
+            BuildRequest(identity),
+            FakeProducer(identity),
+            validation=ValidationLevel.FULL_CHECKSUM,
+            deadline=time.monotonic() + 10,
+        )
+        assert result.status is StoreStatus.BUILT
+        assert result.lease is not None
+        assert torch.equal(result.lease.tensors["layer.weight"], _weight())
+        result.lease.close()
+    assert not list(store.quarantine_dir.glob(f"{identity.key}.cleanup.*"))
+
+
+@pytest.mark.parametrize("after_removal", [False, True])
+def test_cleanup_quarantined_retries_parent_sync_failure(
+    tmp_path: Path, monkeypatch: pytest.MonkeyPatch, after_removal: bool
+) -> None:
+    store = _make_store(tmp_path / "store")
+    identity, _ = _publish_test_artifact(store)
+    quarantined = _quarantine_and_rebuild(store, identity)
+    original_sync = filesystem_store_module._fsync_directory
+
+    def fail_sync(path: Path) -> None:
+        if path == store.quarantine_dir and (
+            not after_removal or not list(store.quarantine_dir.glob(f"{identity.key}.cleanup.*"))
+        ):
+            raise OSError(errno.EIO, "injected quarantine sync failure")
+        original_sync(path)
+
+    with monkeypatch.context() as fault:
+        fault.setattr(filesystem_store_module, "_fsync_directory", fail_sync)
+        failure = store.cleanup_quarantined(quarantined.name)
+    assert failure is not None and failure.code is FailureCode.QUARANTINE_FAILED and failure.retryable
+    assert not quarantined.exists()
+    assert bool(list(store.quarantine_dir.glob(f"{identity.key}.cleanup.*"))) is not after_removal
+
+    # An already absent selection must still retry the failed durability step.
+    with monkeypatch.context() as fault:
+        fault.setattr(filesystem_store_module, "_fsync_directory", fail_sync)
+        assert store.cleanup_quarantined(quarantined.name) is not None
+    assert store.cleanup_quarantined(quarantined.name) is None
+    assert not list(store.quarantine_dir.glob(f"{identity.key}.cleanup.*"))
+
+
+@pytest.mark.parametrize("entry_kind", ["symlink", "file"])
+def test_cleanup_quarantined_rejects_non_directory_entry(tmp_path: Path, entry_kind: str) -> None:
+    store = _make_store(tmp_path / "store")
+    identity, artifact = _publish_test_artifact(store)
+    entry = store.quarantine_dir / f"{identity.key}.validation_failure.test"
+    if entry_kind == "symlink":
+        entry.symlink_to(artifact, target_is_directory=True)
+    else:
+        entry.write_text("not an artifact directory")
+    failure = store.cleanup_quarantined(entry.name)
+    assert failure is not None and failure.code is FailureCode.QUARANTINE_FAILED
+    assert not failure.retryable
+    assert entry.exists()
+    hit = store.lookup(identity, validation=ValidationLevel.FULL_CHECKSUM)
+    assert hit.lease is not None
+    assert torch.equal(hit.lease.tensors["layer.weight"], _weight())
+    hit.lease.close()
+
+
+@pytest.mark.parametrize(
+    "name",
+    ["../escape", "/tmp/escape", "f" * 64, "g" * 64 + ".x", "f" * 64 + ".", "f" * 64 + ".x/../y", "f" * 64 + ".x\0"],
+)
+def test_cleanup_quarantined_rejects_invalid_selection_before_mutation(tmp_path: Path, name: str) -> None:
+    store = _make_store(tmp_path / "store")
+    locks = set(store.locks_dir.iterdir())
+    with pytest.raises(ValueError, match="inventory basename"):
+        store.cleanup_quarantined(name)
+    assert set(store.locks_dir.iterdir()) == locks
+
+
+def test_cleanup_quarantined_restores_bounded_store_recovery_headroom(tmp_path: Path) -> None:
+    store = _make_store(tmp_path / "store", capacity=CapacityPolicy(max_store_bytes=16384))
+    identity = _identity()
+    for _ in range(32):
+        result = store.get_or_build(
+            BuildRequest(identity),
+            FakeProducer(identity),
+            validation=ValidationLevel.FULL_CHECKSUM,
+            deadline=time.monotonic() + 10,
+        )
+        if result.status is StoreStatus.FAILED:
+            break
+        assert result.lease is not None
+        result.lease.close()
+        _corrupt_test_payload(store.artifacts_dir / identity.key)
+    assert result.failure is not None and result.failure.code is FailureCode.STORE_LIMIT_EXCEEDED
+    quarantined = [
+        entry for entry in store.inspect_domain().inventory if entry.state is ArtifactInventoryState.QUARANTINED
+    ]
+    assert len(quarantined) > 1
+    assert store.cleanup(identity) is None
+    still_full = store.get_or_build(
+        BuildRequest(identity),
+        FakeProducer(identity),
+        validation=ValidationLevel.FULL_CHECKSUM,
+        deadline=time.monotonic() + 10,
+    )
+    assert still_full.failure is not None and still_full.failure.code is FailureCode.STORE_LIMIT_EXCEEDED
+
+    assert store.cleanup_quarantined(quarantined[0].storage_name) is None
+    recovered = store.get_or_build(
+        BuildRequest(identity),
+        FakeProducer(identity),
+        validation=ValidationLevel.FULL_CHECKSUM,
+        deadline=time.monotonic() + 10,
+    )
+    assert recovered.status is StoreStatus.BUILT
+    assert recovered.lease is not None
+    assert torch.equal(recovered.lease.tensors["layer.weight"], _weight())
+    recovered.lease.close()
+
+
 def test_domain_capacity_policy_is_authoritative(tmp_path: Path) -> None:
     root = tmp_path / "store"
     _make_store(root, capacity=CapacityPolicy(max_store_bytes=4096))

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/config.py

@@ -8,7 +8,6 @@ import math
 from dataclasses import dataclass, field
 from enum import Enum
 from pathlib import Path
-from typing import Literal
 
 
 class RuntimeMode(str, Enum):
@@ -67,7 +66,6 @@ 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 (
@@ -80,10 +78,6 @@ 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]:
-        document: dict[str, object] = {
+        return {
             "schema_version": 1,
             "policy_version": policy_version,
             "max_artifact_bytes": self.capacity_policy.max_artifact_bytes,
@@ -635,9 +635,6 @@ 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"
@@ -1016,40 +1013,6 @@ 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:
@@ -1503,43 +1466,45 @@ class FilesystemHostWeightStore:
                         total += file_stat.st_size
         return total
 
+    @contextlib.contextmanager
+    def _cleanup_locks(self, key: str) -> Iterator[None]:
+        """Exclude cooperative builders and leases without waiting for them."""
+        try:
+            build_lock = FileLock(self._build_lock_path(key), exclusive=True, deadline=None, nonblocking=True)
+        except FileLockTimeoutError as exc:
+            raise HostWeightError(
+                _failure(
+                    ResolutionStage.LIFECYCLE,
+                    FailureCode.ACTIVE_BUILD_TIMEOUT,
+                    f"artifact {key} has an active builder",
+                    retryable=True,
+                )
+            ) from exc
+        with build_lock:
+            try:
+                artifact_lock = FileLock(self._artifact_lock_path(key), exclusive=True, deadline=None, nonblocking=True)
+            except FileLockTimeoutError as exc:
+                raise HostWeightError(
+                    _failure(
+                        ResolutionStage.LIFECYCLE,
+                        FailureCode.ACTIVE_LEASE,
+                        f"artifact {key} has an active lease",
+                        retryable=True,
+                    )
+                ) from exc
+            with artifact_lock:
+                yield
+
     def cleanup(self, identity: WeightArtifactIdentity) -> HostWeightFailure | None:
         """Explicitly remove an inactive artifact without waiting for leases."""
         try:
-            try:
-                build_lock = FileLock(
-                    self._build_lock_path(identity.key),
-                    exclusive=True,
-                    deadline=None,
-                    nonblocking=True,
-                )
-            except FileLockTimeoutError:
-                return _failure(
-                    ResolutionStage.LIFECYCLE,
-                    FailureCode.ACTIVE_BUILD_TIMEOUT,
-                    f"artifact {identity.key} has an active builder",
-                    retryable=True,
-                )
-            with build_lock:
-                try:
-                    artifact_lock = FileLock(
-                        self._artifact_lock_path(identity.key),
-                        exclusive=True,
-                        deadline=None,
-                        nonblocking=True,
-                    )
-                except FileLockTimeoutError:
-                    return _failure(
-                        ResolutionStage.LIFECYCLE,
-                        FailureCode.ACTIVE_LEASE,
-                        f"artifact {identity.key} has an active lease",
-                        retryable=True,
-                    )
-                with artifact_lock:
-                    self._repair_quarantine_transitions_locked(identity.key)
-                    self._quarantine_locked(identity.key, reason="cleanup")
-                    self._remove_cleanup_tombstones_locked(identity.key)
-                    self._remove_deny_locked(identity.key)
+            with self._cleanup_locks(identity.key):
+                self._repair_quarantine_transitions_locked(identity.key)
+                self._quarantine_locked(identity.key, reason="cleanup")
+                self._remove_cleanup_tombstones_locked(identity.key)
+                self._remove_deny_locked(identity.key)
+        except HostWeightError as exc:
+            return exc.failure
         except OSError as exc:
             return _failure(
                 ResolutionStage.LIFECYCLE,
@@ -1549,6 +1514,50 @@ class FilesystemHostWeightStore:
             )
         return None
 
+    def cleanup_quarantined(self, storage_name: str) -> HostWeightFailure | None:
+        """Remove one quarantine inventory entry, preserving the current artifact.
+
+        A retry also finishes earlier cleanup tombstones for the same key.
+        Active builders or leases are reported rather than interrupted.
+        """
+        key, separator, suffix = storage_name.partition(".")
+        if (
+            not separator
+            or not suffix
+            or "\0" in storage_name
+            or _ARTIFACT_KEY_RE.fullmatch(key) is None
+            or Path(storage_name).name != storage_name
+        ):
+            raise ValueError("quarantined storage_name must be a single inventory basename with an artifact key")
+        try:
+            with self._cleanup_locks(key):
+                entry = self.quarantine_dir / storage_name
+                exists = _path_exists_without_following(entry)
+                if exists and not stat.S_ISDIR(entry.lstat().st_mode):
+                    return _failure(
+                        ResolutionStage.LIFECYCLE,
+                        FailureCode.QUARANTINE_FAILED,
+                        f"quarantine entry is not a directory: {storage_name}",
+                    )
+                self._repair_quarantine_transitions_locked(key)
+                if exists and not storage_name.startswith(f"{key}.cleanup."):
+                    tombstone = self.quarantine_dir / f"{key}.cleanup.{uuid.uuid4().hex}"
+                    _move_artifact_locked(entry, tombstone)
+                self._remove_cleanup_tombstones_locked(key)
+                # A previous attempt may have removed the tombstone but failed
+                # its parent sync. Retry that sync even when no entry remains.
+                _fsync_directory(self.quarantine_dir)
+        except HostWeightError as exc:
+            return exc.failure
+        except OSError as exc:
+            return _failure(
+                ResolutionStage.LIFECYCLE,
+                FailureCode.QUARANTINE_FAILED,
+                f"failed to remove quarantined entry {storage_name}: {exc}",
+                retryable=True,
+            )
+        return None
+
     def inspect_domain(self) -> DomainInspection:
         usage = shutil.disk_usage(self.root)
         inventory: list[ArtifactInventoryEntry] = []