[Core] Add explicit HWR registration bypass and progress logs

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

docs/design/feature/offloader/distributed_layerwise_offload.md

@@ -289,6 +289,13 @@ store, tensor ownership, or H2D payload. On success, each tensor view copies
 directly into the existing rotating HBM block buffers and the two private host
 staging slots are not allocated.
 
+`--dlo-host-registration-mode` defaults to `auto`, preserving this behavior.
+Select `disabled` to bypass HWR mapping registration and use bounded host staging
+regardless of mapping size or budget. This leaves the existing pin-memory policy
+for staging buffers unchanged. The equivalent offline/stage configuration field
+is `dlo_host_registration_mode`; an explicit disable requires enabled HWR with
+no-AllGather DLO.
+
 `--dlo-host-registration-limit-gib` is an optional per-worker preflight ceiling
 over page-aligned registered bytes. Zero adds no ceiling. A disabled pinned
 memory policy, unsupported platform/capability, over-budget mapping, or fully
@@ -296,6 +303,14 @@ rolled-back registration error selects the existing two-slot staging path. A
 partial registration that cannot be rolled back aborts startup because closing
 the lease would unmap memory still owned by the platform.
 
+Automatic registration logs its entry with device and budget, then each CUDA
+region's ordinal/total and byte size before calling `cudaHostRegister`. A
+completion message follows each successful call. These identify the last
+entered boundary if progress stops; they do not impose a timeout. A synchronous
+CUDA call cannot be safely cancelled by a Python timer, and a stuck process
+still requires external supervision. The explicit disable option avoids this
+registration path; it does not repair the underlying driver hang.
+
 Direct checkpoint mmap remains unchanged and continues to use staging. It may
 require loader-owned per-block transforms, while the HWR artifact already
 contains final runtime bytes. DLO AllGather never receives an HWR final-layout

tests/config/test_omni_config.py

@@ -193,6 +193,39 @@ def test_resolve_execution_mode_rejects_unknown_execution_type():
         omni_config_module._resolve_execution_mode("unknown_execution_type")
 
 
+@pytest.mark.parametrize("registration_mode", ["auto", "disabled"])
+def test_hwr_registration_policy_reaches_offloader(tmp_path: Path, registration_mode: str):
+    from vllm_omni.diffusion.data import OmniDiffusionConfig
+    from vllm_omni.diffusion.offloader.base import OffloadConfig
+
+    options = {
+        "enable_distributed_layerwise_offload": True,
+        "dlo_use_allgather": False,
+        "host_weight_runtime_mode": "preferred",
+        "host_weight_runtime_root": str(tmp_path / "unused-store"),
+        "dlo_host_registration_mode": registration_mode,
+    }
+    config = _from_pipeline_key("dreamzero", deploy_config_path="dreamzero_tp1_cfg2", cli_overrides=options)
+    stage = config.stage_by_id(0)
+    assert isinstance(stage, VllmOmniDiffusionStageConfig)
+    assert stage.diffusion_config.dlo_host_registration_mode == registration_mode
+    offline = OmniDiffusionConfig(**options)
+    transport = OffloadConfig.from_od_config(offline)
+    assert transport.dlo_host_registration_mode == registration_mode
+    assert transport.pin_cpu_memory
+    assert transport.dlo_host_registration_limit_gib == 0.0
+    assert not (tmp_path / "unused-store").exists()
+
+
+@pytest.mark.parametrize("registration_mode", ["unknown", None, "disabled"])
+def test_hwr_registration_policy_rejects_invalid_or_ineligible_config(registration_mode):
+    from vllm_omni.diffusion.data import OmniDiffusionConfig
+
+    for config_class in (OmniDiffusionConfig, omni_config_module._DiffusionConfigProjection):
+        with pytest.raises(ValueError, match="dlo_host_registration_mode"):
+            config_class(dlo_host_registration_mode=registration_mode)
+
+
 def test_from_pipeline_config_preserves_current_pipeline_config_object():
     omni_config = _from_pipeline_key("minicpmo_4_5")
     pipeline = _resolve_pipeline_or_skip("minicpmo_4_5")

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/diffusion/offloader/test_cuda_host_registration.py

@@ -124,6 +124,33 @@ def test_registration_rejects_writable_or_over_budget_before_cuda(monkeypatch: p
     assert runtime.registered == []
 
 
+@pytest.mark.parametrize("second_result", [0, 7])
+def test_registration_logs_entry_before_each_runtime_call(monkeypatch: pytest.MonkeyPatch, second_result: int) -> None:
+    runtime = _FakeRuntime([0, second_result])
+    messages: list[str] = []
+    monkeypatch.setattr(registration_module.logger, "info", lambda message, *args: messages.append(message % args))
+    original_register = runtime.cudaHostRegister
+
+    def register(pointer: int, size: int, flags: int) -> int:
+        ordinal = len(runtime.registered) + 1
+        assert messages[-1] == f"Registering HWR host range {ordinal}/2: {size} bytes"
+        return original_register(pointer, size, flags)
+
+    monkeypatch.setattr(runtime, "cudaHostRegister", register)
+    monkeypatch.setattr(registration_module.torch.cuda, "cudart", lambda: runtime)
+    monkeypatch.setattr(registration_module, "_consume_last_cuda_error", lambda _runtime, _error: None)
+    regions = (_region("first", 0x1000, 4096), _region("second", 0x3000, 4096))
+    if second_result:
+        with pytest.raises(HostRegistrationError, match="error-7"):
+            CudaHostRegistration.create(regions, max_bytes=None)
+        assert runtime.unregistered == [0x1000]
+    else:
+        registration = CudaHostRegistration.create(regions, max_bytes=None)
+        assert registration.close() == ()
+    assert sum(message.startswith("Registering ") for message in messages) == 2
+    assert sum(message.startswith("Registered ") for message in messages) == (1 if second_result else 2)
+
+
 def test_registration_requires_read_only_capability(monkeypatch: pytest.MonkeyPatch) -> None:
     runtime = _FakeRuntime([0], read_only_supported=False)
     monkeypatch.setattr(registration_module.torch.cuda, "cudart", lambda: runtime)

tests/diffusion/offloader/test_distributed_layerwise_backend.py

@@ -96,7 +96,7 @@ class TestDistributedLayerwiseOffloadHook:
         hook.offload_layer()
         assert not hook.is_materialized
 
-    def test_initialize_failure_keeps_next_block_materialized(self, monkeypatch):
+    def test_initialize_failure_keeps_next_block_materialized(self, monkeypatch, patched_offload_runtime):
         current_block = nn.Linear(2, 2)
         next_block = nn.Linear(2, 2)
         expected = {name: tensor.detach().clone() for name, tensor in next_block.state_dict().items()}
@@ -1373,10 +1373,14 @@ class TestMmapWeightLoading:
 
         synchronize.assert_called()
 
-    def test_hwr_registration_failure_falls_back_to_bounded_staging(
+    @pytest.mark.parametrize("registration_mode", ["auto", "disabled"])
+    @pytest.mark.parametrize("limit_gib", [0.0, 1.5])
+    def test_hwr_registration_selection_uses_bounded_staging(
         self,
         monkeypatch,
         patched_offload_runtime,
+        registration_mode,
+        limit_gib,
     ):
         plan, carrier, lease = _fake_hwr_plan("hwr-registration-fallback")
         backend = DistributedLayerwiseOffloadBackend(
@@ -1385,6 +1389,8 @@ class TestMmapWeightLoading:
                 pin_cpu_memory=True,
                 dp_size=1,
                 dlo_use_allgather=False,
+                dlo_host_registration_mode=registration_mode,
+                dlo_host_registration_limit_gib=limit_gib,
             ),
             torch.device("cpu"),
             host_weight_plan=plan,
@@ -1394,12 +1400,14 @@ class TestMmapWeightLoading:
 
         def fail_registration(*args, **kwargs):
             del args, kwargs
+            assert registration_mode == "auto", "disabled mode entered platform registration"
             raise HostRegistrationError("registration unavailable in CPU test")
 
         def allocate_unpinned_staging(hooks, resident_group=None):
             # Registration selection still observes pin_cpu_memory=True. The
             # CPU-only test avoids asking the host for CUDA-pinned allocations.
             for hook in hooks:
+                assert hook.pin_memory, "registration policy disabled staging pinning"
                 hook.pin_memory = False
             return original_staging_allocator(hooks, resident_group)
 
@@ -1412,6 +1420,7 @@ class TestMmapWeightLoading:
         assert backend._using_rank_local_mmap
         assert not backend._using_registered_mmap
         assert backend._host_registration is None
+        assert backend.config.pin_cpu_memory
         assert all(len(hook.cpu_staging_buffers) == 2 for group in backend._all_hook_groups for hook in group)
         assert not lease.closed
 
@@ -2630,7 +2639,9 @@ class TestDistributedComponentSelection:
             assert registry is None or registry.get_hook("distributed_layerwise_offload") is None
             torch.testing.assert_close(block.weight, expected)
 
-    def test_multirank_enable_failure_cleanup_skips_restore_collective(self, monkeypatch, mocker):
+    def test_multirank_enable_failure_cleanup_skips_restore_collective(
+        self, monkeypatch, mocker, patched_offload_runtime
+    ):
         backend = DistributedLayerwiseOffloadBackend(
             OffloadConfig(
                 strategy=OffloadStrategy.DISTRIBUTED_LAYER_WISE,
@@ -2811,7 +2822,7 @@ class TestDistributedComponentSelection:
         with pytest.raises(ValueError, match="not declared replicated"):
             backend.enable(pipeline)
 
-    def test_encoder_allgather_rejects_stub_rank_before_block_discovery(self):
+    def test_encoder_allgather_rejects_stub_rank_before_block_discovery(self, patched_offload_runtime):
         """Every rank must reject an unsafe encoder group, including stub ranks."""
         backend = DistributedLayerwiseOffloadBackend(
             OffloadConfig(

tests/entrypoints/test_async_omni_diffusion_config.py

@@ -507,6 +507,8 @@ def test_serve_cli_forwards_hwr_policy_for_no_allgather_dlo():
             "preferred",
             "--host-weight-runtime-root",
             "/var/cache/vllm-omni/hwr",
+            "--dlo-host-registration-mode",
+            "disabled",
             "--dlo-host-registration-limit-gib",
             "80",
         ]
@@ -522,6 +524,8 @@ def test_serve_cli_forwards_hwr_policy_for_no_allgather_dlo():
     assert engine_args["host_weight_runtime_mode"] == "preferred"
     assert engine_args["host_weight_runtime_root"] == "/var/cache/vllm-omni/hwr"
     assert engine_args["dlo_host_registration_limit_gib"] == 80
+    assert explicit_kwargs["dlo_host_registration_mode"] == "disabled"
+    assert engine_args["dlo_host_registration_mode"] == "disabled"
 
 
 def test_serve_cli_accepts_diffusion_compile_controls():

vllm_omni/config/omni_config.py

@@ -739,6 +739,7 @@ class _DiffusionConfigProjection:
     dlo_resident_layers: int = Field(default=0, ge=0)
     host_weight_runtime_mode: Literal["disabled", "preferred", "required"] = "disabled"
     host_weight_runtime_root: str | None = None
+    dlo_host_registration_mode: str = "auto"
     dlo_host_registration_limit_gib: float = Field(default=0.0, ge=0)
     pin_cpu_memory: bool = True
     diffusion_compile_granularity: Literal["regional", "full"] = "regional"
@@ -895,6 +896,7 @@ class _DiffusionConfigProjection:
         )
         self.dlo_host_registration_limit_gib = validate_dlo_host_registration_options(
             limit_gib=self.dlo_host_registration_limit_gib,
+            mode=self.dlo_host_registration_mode,
             enable_dlo=self.enable_distributed_layerwise_offload,
             use_allgather=self.dlo_use_allgather,
             hwr_mode=self.host_weight_runtime_mode,

vllm_omni/config/stage_config.py

@@ -483,6 +483,7 @@ class StageDeployConfig:
     dlo_resident_layers: int | None = None
     host_weight_runtime_mode: str | None = None
     host_weight_runtime_root: str | None = None
+    dlo_host_registration_mode: str | None = None
     dlo_host_registration_limit_gib: float | None = None
     # Diffusion-specific debug and observability knobs.
     enable_diffusion_pipeline_profiler: bool | None = None

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

@@ -119,8 +119,15 @@ def validate_dlo_host_registration_options(
     enable_dlo: bool,
     use_allgather: bool,
     hwr_mode: object,
+    mode: object = "auto",
 ) -> float:
     """Validate the optional transport budget without probing CUDA or HWR."""
+    if mode not in ("auto", "disabled"):
+        raise ValueError("dlo_host_registration_mode must be auto or disabled")
+    if mode == "disabled" and (not enable_dlo or use_allgather or hwr_mode == "disabled"):
+        raise ValueError(
+            "disabled dlo_host_registration_mode requires enabled no-AllGather DLO and Host Weight Runtime"
+        )
     if not isinstance(limit_gib, (int, float, str)):
         raise TypeError(f"dlo_host_registration_limit_gib must be a number; got {type(limit_gib).__name__}")
     value = float(limit_gib)
@@ -833,6 +840,7 @@ class OmniDiffusionConfig:
     host_weight_runtime_root: str | None = None
     # Optional per-worker ceiling for registering final-layout HWR mappings.
     # Zero adds no ceiling; pin_cpu_memory controls whether registration is tried.
+    dlo_host_registration_mode: str = "auto"
     dlo_host_registration_limit_gib: float = 0.0
 
     pin_cpu_memory: bool = True  # Use pinned memory for faster transfers when offloading
@@ -1244,6 +1252,7 @@ class OmniDiffusionConfig:
         )
         self.dlo_host_registration_limit_gib = validate_dlo_host_registration_options(
             limit_gib=self.dlo_host_registration_limit_gib,
+            mode=self.dlo_host_registration_mode,
             enable_dlo=self.enable_distributed_layerwise_offload,
             use_allgather=self.dlo_use_allgather,
             hwr_mode=self.host_weight_runtime_mode,

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/diffusion/offloader/base.py

@@ -73,12 +73,15 @@ class OffloadConfig:
     # Optional per-worker ceiling for registering an HWR mmap. Zero means no
     # additional ceiling; pin_cpu_memory controls whether registration is tried.
     dlo_host_registration_limit_gib: float = 0.0
+    dlo_host_registration_mode: str = "auto"
     # ``None`` preserves the model's legacy plan-driven component topology;
     # a frozenset is an explicit compact-API selection.
     components: frozenset[str] | None = None
     dlo_transfers: dict[str, DLOTransfer] | None = None
 
     def __post_init__(self) -> None:
+        if self.dlo_host_registration_mode not in ("auto", "disabled"):
+            raise ValueError("dlo_host_registration_mode must be auto or disabled")
         if self.components is not None:
             self.components = parse_offload_components(self.components)
         if self.dlo_transfers is None:
@@ -174,7 +177,9 @@ class OffloadConfig:
         dlo_transfers = dict(resolved.transfers)
         dlo_resident_layers = resolved.resident_layers
         dit_uses_allgather = resolved.uses_allgather(DIT_COMPONENT)
+        registration_mode = getattr(od_config, "dlo_host_registration_mode", "auto")
         dlo_host_registration_limit_gib = validate_dlo_host_registration_options(
+            mode=registration_mode,
             limit_gib=getattr(od_config, "dlo_host_registration_limit_gib", 0.0),
             enable_dlo=enable_distributed_layerwise_offload,
             use_allgather=dit_uses_allgather,
@@ -210,6 +215,7 @@ class OffloadConfig:
             dlo_use_allgather=dit_uses_allgather,
             dlo_resident_layers=dlo_resident_layers,
             dlo_host_registration_limit_gib=dlo_host_registration_limit_gib,
+            dlo_host_registration_mode=registration_mode,
             components=components,
             dlo_transfers=dlo_transfers,
         )

vllm_omni/diffusion/offloader/cuda_host_registration.py

@@ -6,12 +6,14 @@ from __future__ import annotations
 
 import ctypes
 import mmap
+import time
 from collections import defaultdict
 from collections.abc import Sequence
 from dataclasses import dataclass
 from typing import Protocol
 
 import torch
+from vllm.logger import init_logger
 
 from vllm_omni.host_weight_runtime import MappedHostRegion
 
@@ -21,6 +23,8 @@ from .host_registration import (
     HostRegistrationError,
 )
 
+logger = init_logger(__name__)
+
 _CUDA_HOST_REGISTER_READ_ONLY = 0x08
 _CUDA_DEVICE_ATTRIBUTE_HOST_REGISTER_READ_ONLY_SUPPORTED = 113
 
@@ -187,7 +191,9 @@ class CudaHostRegistration:
 
         registered: list[_AddressRange] = []
         try:
-            for region in mapped:
+            for index, region in enumerate(mapped, start=1):
+                logger.info("Registering HWR host range %d/%d: %d bytes", index, len(mapped), region.size)
+                started = time.perf_counter()
                 error = runtime.cudaHostRegister(region.start, region.size, _CUDA_HOST_REGISTER_READ_ONLY)
                 if int(error) != 0:
                     raise HostRegistrationError(
@@ -195,6 +201,9 @@ class CudaHostRegistration:
                         f"[{region.start:#x}, {region.end:#x}): {_handled_error_message(runtime, error)}"
                     )
                 registered.append(region)
+                logger.info(
+                    "Registered HWR host range %d/%d in %.3f s", index, len(mapped), time.perf_counter() - started
+                )
         except Exception as exc:
             rollback_errors: list[str] = []
             rollback_failed: list[_AddressRange] = []

vllm_omni/diffusion/offloader/distributed_layerwise_backend.py

@@ -1080,6 +1080,9 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend):
         lease = self._host_weight_lease
         if lease is None:
             return False
+        if self.config.dlo_host_registration_mode == "disabled":
+            logger.info("HWR mmap registration disabled by transport policy; using bounded host staging")
+            return False
         if not self.config.pin_cpu_memory:
             logger.info("HWR mmap registration disabled by pin_cpu_memory=False; using bounded host staging")
             return False
@@ -1089,6 +1092,12 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend):
 
         limit_gib = self.config.dlo_host_registration_limit_gib
         max_bytes = int(limit_gib * 1024**3) if limit_gib > 0 else None
+        logger.info(
+            "Starting HWR mmap registration on %s: %d source range(s), budget_bytes=%s",
+            self.device,
+            len(lease.mapped_regions),
+            max_bytes,
+        )
         started = time.perf_counter()
         try:
             registration = register_host_mappings(

vllm_omni/engine/arg_utils.py

@@ -582,6 +582,7 @@ class OrchestratorArgs:
     dlo_resident_layers: int = 0
     host_weight_runtime_mode: str = "disabled"
     host_weight_runtime_root: str | None = None
+    dlo_host_registration_mode: str = "auto"
     dlo_host_registration_limit_gib: float = 0.0
     boundary_ratio: float | None = None
     flow_shift: float | None = None

vllm_omni/engine/async_omni_engine.py

@@ -1145,6 +1145,7 @@ class AsyncOmniEngine:
             "dlo_resident_layers": kwargs.get("dlo_resident_layers", 0),
             "host_weight_runtime_mode": kwargs.get("host_weight_runtime_mode", "disabled"),
             "host_weight_runtime_root": kwargs.get("host_weight_runtime_root"),
+            "dlo_host_registration_mode": kwargs.get("dlo_host_registration_mode", "auto"),
             "dlo_host_registration_limit_gib": kwargs.get("dlo_host_registration_limit_gib", 0.0),
             "enforce_eager": False if kwargs.get("enforce_eager") is None else kwargs.get("enforce_eager"),
             "diffusion_compile_granularity": (

vllm_omni/entrypoints/cli/serve.py

@@ -780,6 +780,17 @@ class OmniServeCommand(CLISubcommand):
                 "the same persistent path for population and serving."
             ),
         )
+        omni_config_group.add_argument(
+            "--dlo-host-registration-mode",
+            choices=("auto", "disabled"),
+            default="auto",
+            help=(
+                "HWR transport registration policy for eligible no-AllGather DLO. "
+                "auto attempts registered direct H2D under the existing pin-memory and budget policy; "
+                "disabled selects bounded staging without registering the HWR mappings. "
+                "This preserves the pin-memory policy for staging buffers."
+            ),
+        )
         omni_config_group.add_argument(
             "--dlo-host-registration-limit-gib",
             type=float,