Merge remote-tracking branch 'origin/main' into qwen-image-2.1

# Conflicts:
#	tests/entrypoints/openai_api/test_image_server.py
#	vllm_omni/entrypoints/openai/api_server.py

tests/diffusion/models/bagel/test_step_execution.py

@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 """Focused CPU tests for BAGEL step execution and packed batching."""
 
 from __future__ import annotations
@@ -18,12 +18,14 @@ from vllm_omni.diffusion.models.bagel.bagel_transformer import Bagel, NaiveCache
 from vllm_omni.diffusion.models.bagel.pipeline_bagel import (
     BagelGenParams,
     BagelPipeline,
+    _bagel_canvas_requested,
     get_bagel_pre_process_func,
 )
 from vllm_omni.diffusion.request import OmniDiffusionRequest
 from vllm_omni.diffusion.sched import StepScheduler
 from vllm_omni.diffusion.worker.input_batch import InputBatch
 from vllm_omni.diffusion.worker.utils import StepRequestState
+from vllm_omni.errors import OmniClientError
 from vllm_omni.inputs.data import OmniDiffusionSamplingParams
 
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
@@ -118,7 +120,7 @@ def test_bagel_step_preprocessor_separates_incompatible_cfg_settings(first_extra
     assert first.batch_compatibility_key != second.batch_compatibility_key
 
 
-def test_bagel_step_preprocessor_buckets_effective_img2img_sizes(tmp_path):
+def _pre_process_with_stub_checkpoint(tmp_path):
     (tmp_path / "config.json").write_text(
         '{"vae_config":{"downsample":8},"latent_patch_size":2,"max_latent_size":32}',
         encoding="utf-8",
@@ -131,9 +133,22 @@ def test_bagel_step_preprocessor_buckets_effective_img2img_sizes(tmp_path):
         '{"weight_map":{"latent_pos_embed.pos_embed":"ema.safetensors"}}',
         encoding="utf-8",
     )
-    pre_process = get_bagel_pre_process_func(
-        types.SimpleNamespace(model=str(tmp_path), revision=None, step_execution=True)
+    return get_bagel_pre_process_func(types.SimpleNamespace(model=str(tmp_path), revision=None, step_execution=True))
+
+
+def _edit_request(source_size, sampling_params, request_id="req"):
+    return OmniDiffusionRequest(
+        prompt={
+            "prompt": "edit this image",
+            "multi_modal_data": {"image": [Image.new("RGB", source_size)]},
+        },
+        sampling_params=sampling_params,
+        request_id=request_id,
     )
+
+
+def test_bagel_step_preprocessor_buckets_effective_img2img_sizes(tmp_path):
+    pre_process = _pre_process_with_stub_checkpoint(tmp_path)
     requests = [
         OmniDiffusionRequest(
             prompt={
@@ -163,6 +178,172 @@ def test_bagel_step_preprocessor_buckets_effective_img2img_sizes(tmp_path):
     assert schedule.num_waiting_reqs == 1
 
 
+def test_bagel_pre_process_keeps_explicitly_requested_canvas(tmp_path):
+    pre_process = _pre_process_with_stub_checkpoint(tmp_path)
+    request = _edit_request((512, 1024), OmniDiffusionSamplingParams(num_inference_steps=2, height=1024, width=1024))
+
+    pre_process(request)
+
+    assert (request.sampling_params.height, request.sampling_params.width) == (1024, 1024)
+    assert _bagel_canvas_requested(request.sampling_params)
+
+
+def test_bagel_pre_process_aligns_source_derived_canvas_for_size_auto(tmp_path):
+    pre_process = _pre_process_with_stub_checkpoint(tmp_path)
+    request = _edit_request(
+        (500, 700),
+        OmniDiffusionSamplingParams(
+            num_inference_steps=2,
+            height=700,
+            width=500,
+            height_not_provided=True,
+            width_not_provided=True,
+        ),
+    )
+
+    pre_process(request)
+
+    assert (request.sampling_params.height, request.sampling_params.width) == (704, 496)
+    assert not _bagel_canvas_requested(request.sampling_params)
+
+
+def test_bagel_pre_process_derives_canvas_when_caller_asked_for_none(tmp_path):
+    pre_process = _pre_process_with_stub_checkpoint(tmp_path)
+    request = _edit_request((800, 400), OmniDiffusionSamplingParams(num_inference_steps=2))
+
+    pre_process(request)
+
+    assert (request.sampling_params.height, request.sampling_params.width) == (400, 800)
+    assert (request.sampling_params.height_not_provided, request.sampling_params.width_not_provided) == (True, True)
+    assert not _bagel_canvas_requested(request.sampling_params)
+
+
+def test_bagel_pre_process_marks_kv_metadata_canvas_as_requested(tmp_path):
+    pre_process = _pre_process_with_stub_checkpoint(tmp_path)
+    request = _edit_request(
+        (500, 700),
+        OmniDiffusionSamplingParams(
+            num_inference_steps=2,
+            height=700,
+            width=500,
+            height_not_provided=True,
+            width_not_provided=True,
+            kv_metadata={"image_shape": (1024, 1024)},
+        ),
+    )
+
+    pre_process(request)
+
+    assert (request.sampling_params.height, request.sampling_params.width) == (1024, 1024)
+    assert (request.sampling_params.height_not_provided, request.sampling_params.width_not_provided) == (False, False)
+    assert _bagel_canvas_requested(request.sampling_params)
+
+
+def test_bagel_pre_process_prefers_explicit_canvas_over_kv_metadata(tmp_path):
+    pre_process = _pre_process_with_stub_checkpoint(tmp_path)
+    request = _edit_request(
+        (1280, 720),
+        OmniDiffusionSamplingParams(
+            num_inference_steps=2,
+            height=1024,
+            width=1024,
+            kv_metadata={"image_shape": (576, 1024)},
+        ),
+    )
+
+    pre_process(request)
+
+    assert (request.sampling_params.height, request.sampling_params.width) == (1024, 1024)
+    assert _bagel_canvas_requested(request.sampling_params)
+
+
+def test_bagel_pre_process_floors_unaligned_explicit_canvas_with_warning(tmp_path):
+    pre_process = _pre_process_with_stub_checkpoint(tmp_path)
+    request = _edit_request(
+        (1280, 720),
+        OmniDiffusionSamplingParams(num_inference_steps=2, height=700, width=1000),
+    )
+
+    with patch("vllm_omni.diffusion.models.bagel.pipeline_bagel.logger.warning") as warning:
+        pre_process(request)
+
+    assert (request.sampling_params.height, request.sampling_params.width) == (688, 992)
+    assert _bagel_canvas_requested(request.sampling_params)
+    warning.assert_called_once()
+    assert "width 1000->992, height 700->688" in warning.call_args.args[-1]
+
+
+def test_bagel_pre_process_keeps_aligned_explicit_canvas_silent(tmp_path):
+    pre_process = _pre_process_with_stub_checkpoint(tmp_path)
+    request = _edit_request(
+        (1280, 720),
+        OmniDiffusionSamplingParams(num_inference_steps=2, height=512, width=1024),
+    )
+
+    with patch("vllm_omni.diffusion.models.bagel.pipeline_bagel.logger.warning") as warning:
+        pre_process(request)
+
+    assert (request.sampling_params.height, request.sampling_params.width) == (512, 1024)
+    warning.assert_not_called()
+
+
+def test_bagel_pre_process_rejects_explicit_canvas_over_checkpoint_limit(tmp_path):
+    pre_process = _pre_process_with_stub_checkpoint(tmp_path)
+    request = _edit_request(
+        (1280, 720),
+        OmniDiffusionSamplingParams(num_inference_steps=2, height=512, width=2048),
+    )
+
+    with pytest.raises(OmniClientError, match="width=2048 exceeds the BAGEL checkpoint limit of 1024") as excinfo:
+        pre_process(request)
+
+    assert excinfo.value.status_code == 400
+
+
+def test_bagel_pre_process_aligns_text2img_canvas_too(tmp_path):
+    pre_process = _pre_process_with_stub_checkpoint(tmp_path)
+    request = OmniDiffusionRequest(
+        prompt="a cat",
+        sampling_params=OmniDiffusionSamplingParams(num_inference_steps=2, height=700, width=1000),
+        request_id="req",
+    )
+
+    pre_process(request)
+
+    assert (request.sampling_params.height, request.sampling_params.width) == (688, 992)
+
+
+def test_bagel_pre_process_skips_canvas_alignment_for_text_output(tmp_path):
+    pre_process = _pre_process_with_stub_checkpoint(tmp_path)
+    request = OmniDiffusionRequest(
+        prompt={
+            "prompt": "describe this image",
+            "modalities": ["text"],
+            "multi_modal_data": {"image": [Image.new("RGB", (1280, 720))]},
+        },
+        sampling_params=OmniDiffusionSamplingParams(num_inference_steps=2, height=700, width=1000),
+        request_id="req",
+    )
+
+    pre_process(request)
+
+    assert (request.sampling_params.height, request.sampling_params.width) == (700, 1000)
+    assert request.use_step_execution is False
+
+
+@pytest.mark.parametrize(
+    ("sampling_kwargs", "requested"),
+    [
+        ({"height": 1024, "width": 1024}, True),
+        ({"height": 1024, "width": 512, "height_not_provided": True, "width_not_provided": True}, False),
+        ({}, False),
+    ],
+    ids=["explicit_size", "size_auto", "unset"],
+)
+def test_bagel_canvas_requested_separates_asked_for_from_derived(sampling_kwargs, requested):
+    assert _bagel_canvas_requested(OmniDiffusionSamplingParams(**sampling_kwargs)) is requested
+
+
 def _cache(length: int, value: float, num_layers: int = 1) -> NaiveCache:
     cache = NaiveCache(num_layers)
     if length:

tests/diffusion/models/hunyuan_image3/test_hunyuan_image3_it2i_multi_image.py

@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 """Multi-image input regression for HunyuanImage3 IT2I prompt construction.
 
 The official HunyuanImage-3.0-Instruct supports up to 3 reference images
@@ -41,7 +41,7 @@ pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 
 
 class FakeTokenizer:
-    """Recording fake tokenizer mirroring the one in test_prompt_utils.
+    """Recording fake tokenizer mirroring the one in test_hunyuan_image3_prompt_utils.
 
     Special token ids: `<|startoftext|>`=1, `<img>`=2, `<think>`=3,
     `<recaption>`=4. encode() returns one id per character starting at

tests/diffusion/models/hunyuan_image3/test_hunyuan_image3_prompt_utils.py

@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 from __future__ import annotations
 
 import ast

tests/diffusion/models/lingbot_world/test_lingbot_world_attention.py

@@ -134,6 +134,8 @@ def _install_vllm_stubs() -> None:
             **kwargs,
         ) -> None:
             super().__init__()
+            self.quant_config = kwargs.pop("quant_config", None)
+            self.prefix = kwargs.pop("prefix", "")
             del kwargs
             self.return_bias = return_bias
             self.weight = nn.Parameter(torch.empty(output_size, input_size))
@@ -164,6 +166,8 @@ def _install_vllm_stubs() -> None:
             **kwargs,
         ) -> None:
             super().__init__()
+            self.quant_config = kwargs.pop("quant_config", None)
+            self.prefix = kwargs.pop("prefix", "")
             del kwargs
             self.num_heads = total_num_heads
             self.num_kv_heads = total_num_kv_heads or total_num_heads

tests/diffusion/models/lingbot_world/test_lingbot_world_transformer.py

@@ -35,6 +35,8 @@ def _tiny_model(
     num_layers: int = 2,
     num_frames_per_block: int = 1,
     sliding_window_num_frames: int = 3,
+    quant_config=None,
+    prefix: str = "",
     attention_head_dim: int = 2,
     rope_max_seq_len: int = 16,
 ):
@@ -55,6 +57,8 @@ def _tiny_model(
         num_frames_per_block=num_frames_per_block,
         sliding_window_num_frames=sliding_window_num_frames,
         local_attn_size=-1,
+        quant_config=quant_config,
+        prefix=prefix,
     )
 
 
@@ -485,11 +489,52 @@ def test_constructor_rejects_non_null_image_embedding_fields(field: str) -> None
 
 
 @pytest.mark.cpu
-def test_constructor_rejects_unsupported_quantization_with_runtime_error() -> None:
+@pytest.mark.parametrize("prefix", ["", "transformer"])
+@pytest.mark.parametrize("quantized", [False, True])
+def test_constructor_propagates_quantization_to_parallel_linears(prefix: str, quantized: bool) -> None:
     module = attention_tests._load_module()
+    quant_config = object() if quantized else None
+    model = _tiny_model(module, num_layers=1, quant_config=quant_config, prefix=prefix)
+    block = model.blocks[0]
+    quantized_linears = [
+        model.c2ws_hidden_states_layer1,
+        model.c2ws_hidden_states_layer2,
+        block.self_attn.qkv,
+        block.self_attn.o,
+        block.cross_attn.q,
+        block.cross_attn.k,
+        block.cross_attn.v,
+        block.cross_attn.o,
+        block.ffn[0],
+        block.ffn[2],
+        block.cam_injector_layer1,
+        block.cam_injector_layer2,
+    ]
 
-    with pytest.raises(RuntimeError, match="quant_config.*not supported"):
-        module.CausalLingBotWorldTransformer3DModel(quant_config=object())
+    assert len(quantized_linears) == 12
+    assert all(layer.quant_config is quant_config for layer in quantized_linears)
+    expected_names = {
+        "c2ws_hidden_states_layer1",
+        "c2ws_hidden_states_layer2",
+        *(
+            f"blocks.0.{name}"
+            for name in (
+                "self_attn.qkv",
+                "self_attn.o",
+                "cross_attn.q",
+                "cross_attn.k",
+                "cross_attn.v",
+                "cross_attn.o",
+                "ffn.0",
+                "ffn.2",
+                "cam_injector_layer1",
+                "cam_injector_layer2",
+            )
+        ),
+    }
+    assert {layer.prefix for layer in quantized_linears} == {
+        f"{prefix}.{name}" if prefix else name for name in expected_names
+    }
 
 
 @pytest.mark.cpu

tests/diffusion/models/lingbot_world/test_pipeline_lingbot_world.py

@@ -737,6 +737,99 @@ def test_component_discovery_uses_official_checkpoint_contract() -> None:
     assert module._loader_state.prefetch_calls == [("checkpoint", ("tokenizer", "text_encoder", "vae"), False)]
 
 
+def _ulysses_config(size: int):
+    parallel_config = _od_config().parallel_config
+    parallel_config.sequence_parallel_size = size
+    parallel_config.ulysses_degree = size
+    return _od_config(parallel_config=parallel_config)
+
+
+def _record_shard_install(monkeypatch, module, *, world_size: int):
+    """Stand in for the Ulysses group and the decoder patch; return what the install was asked for."""
+    installs: list[dict] = []
+    group = object()
+
+    def install(vae, group_arg, split_dim, *, dst):
+        installs.append({"vae": vae, "group": group_arg, "split_dim": split_dim, "dst": dst})
+
+    monkeypatch.setattr(module, "install_wan_spatial_shard_decode", install)
+    monkeypatch.setattr(module.LingBotWorldCausalDMDPipeline, "_vae_shard_group", lambda self: (group, world_size))
+    return installs, group
+
+
+def test_a_multi_rank_deployment_shards_the_decoder_across_the_ulysses_ranks(monkeypatch) -> None:
+    """Multi-rank deployments shard decode by default."""
+    module = _load_pipeline_module()
+    installs, group = _record_shard_install(monkeypatch, module, world_size=2)
+
+    pipeline = module.LingBotWorldCausalDMDPipeline(od_config=_ulysses_config(2))
+
+    # Along the width, and assembled on every rank: each rank's post_decode consumes the frame.
+    assert installs == [{"vae": pipeline.vae, "group": group, "split_dim": "width", "dst": None}]
+    assert pipeline._vae_shard_split_dim == "width"
+
+
+def test_a_single_rank_deployment_leaves_the_decoder_alone(monkeypatch) -> None:
+    module = _load_pipeline_module()
+    installs, _ = _record_shard_install(monkeypatch, module, world_size=1)
+
+    pipeline = module.LingBotWorldCausalDMDPipeline(od_config=_ulysses_config(1))
+
+    assert installs == [] and pipeline._vae_shard_split_dim is None
+
+
+@pytest.mark.parametrize("enabled", [True, False])
+@pytest.mark.parametrize("width, local_width", [(832, 208), (840, 216)])
+def test_streaming_decode_reservation_uses_padded_shards(monkeypatch, enabled, width, local_width):
+    module = _load_pipeline_module()
+    installs, _ = _record_shard_install(monkeypatch, module, world_size=4)
+    config = _ulysses_config(4)
+    config.model_config.update(lingbot_vae_spatial_sharding=enabled, ar_diffusion_height=480, ar_diffusion_width=width)
+    pipeline = module.LingBotWorldCausalDMDPipeline(od_config=config)
+    from vllm_omni.experimental.ar_diffusion.streaming_decode import WanStreamingDecoder
+
+    # Exercise the real byte estimator without constructing a checkpoint VAE.
+    decoder = object.__new__(WanStreamingDecoder)
+    decoder._bytes_per_pixel_fp32 = 16.0
+    monkeypatch.setattr(pipeline, "_streaming_decoder", lambda: decoder)
+    expected_width = local_width if enabled else width
+    assert pipeline._streaming_decode_bytes_per_session() == 16 * 480 * expected_width
+    assert len(installs) == int(enabled)
+    assert pipeline._vae_shard_world_size == (4 if enabled else 1)
+    assert config.parallel_config.sequence_parallel_size == config.parallel_config.ulysses_degree == 4
+
+
+def test_disabling_vae_sharding_does_not_require_a_shard_group(monkeypatch):
+    module = _load_pipeline_module()
+    config = _ulysses_config(4)
+    config.model_config["lingbot_vae_spatial_sharding"] = False
+
+    def unexpected_group(self):
+        pytest.fail("disabled VAE sharding must not access its process group")
+
+    monkeypatch.setattr(module.LingBotWorldCausalDMDPipeline, "_vae_shard_group", unexpected_group)
+    pipeline = module.LingBotWorldCausalDMDPipeline(od_config=config)
+    assert pipeline._vae_shard_split_dim is None
+
+
+@pytest.mark.parametrize("value", ["false", 0, None])
+def test_vae_sharding_switch_requires_a_boolean(value):
+    module = _load_pipeline_module()
+    config = _ulysses_config(4)
+    config.model_config["lingbot_vae_spatial_sharding"] = value
+    with pytest.raises(ValueError, match="lingbot_vae_spatial_sharding must be a boolean"):
+        module.LingBotWorldCausalDMDPipeline(od_config=config)
+
+
+def test_vae_shard_refuses_a_group_of_the_wrong_size(monkeypatch) -> None:
+    module = _load_pipeline_module()
+    installs, _ = _record_shard_install(monkeypatch, module, world_size=4)
+
+    with pytest.raises(RuntimeError, match="sequence_parallel_size=2 but the Ulysses group has 4 ranks"):
+        module.LingBotWorldCausalDMDPipeline(od_config=_ulysses_config(2))
+    assert installs == []
+
+
 @pytest.mark.parametrize(
     ("field", "value", "feature"),
     [
@@ -758,8 +851,9 @@ def test_unsupported_parallel_modes_fail_before_component_loading(field: str, va
     assert module._loader_state.prefetch_calls == []
 
 
-def test_pure_ulysses_parallel_config_is_supported() -> None:
+def test_pure_ulysses_parallel_config_is_supported(monkeypatch) -> None:
     module = _load_pipeline_module()
+    _record_shard_install(monkeypatch, module, world_size=2)
     parallel_config = _od_config().parallel_config
     parallel_config.sequence_parallel_size = 2
     parallel_config.ulysses_degree = 2
@@ -795,13 +889,14 @@ def test_unsupported_sp_config_fails_before_component_loading(overrides):
     assert module._loader_state.prefetch_calls == []
 
 
-def test_unsupported_quantization_fails_before_component_loading() -> None:
+def test_quantization_reaches_transformer_factory() -> None:
     module = _load_pipeline_module()
+    quant_config = object()
 
-    with pytest.raises(NotImplementedError, match="quantization"):
-        module.LingBotWorldCausalDMDPipeline(od_config=_od_config(quantization_config=object()))
+    pipeline = module.LingBotWorldCausalDMDPipeline(od_config=_od_config(quantization_config=quant_config))
 
-    assert module._loader_state.prefetch_calls == []
+    assert pipeline.transformer is not None
+    assert _FakeTransformerFactory.last_call[1:] == (quant_config, "transformer")
 
 
 def test_official_scheduler_config_matches_fixed_dmd_contract() -> None:
@@ -1358,10 +1453,34 @@ def test_first_frame_condition_and_camera_fold_match_transformer_contract() -> N
     assert expected_camera.shape == (1, 384, 3, 2, 2)
 
 
-def test_fixed_dmd_transition_and_cache_commit_trace() -> None:
+@pytest.mark.parametrize("value", ["true", "false", 1, 0, None])
+def test_reuse_last_step_kv_rejects_non_boolean_config(value: object) -> None:
+    config = _od_config()
+    config.model_config["lingbot_reuse_last_step_kv"] = value
+    with pytest.raises(ValueError, match="lingbot_reuse_last_step_kv must be a bool"):
+        _pipeline(_load_pipeline_module(), od_config=config)
+
+
+def test_reuse_last_step_kv_is_fixed_per_instance(monkeypatch: pytest.MonkeyPatch) -> None:
+    monkeypatch.setenv("VLLM_OMNI_LINGBOT_REUSE_LAST_STEP_KV", "1")
+    module = _load_pipeline_module()
+    config = _od_config()
+    config.model_config["lingbot_reuse_last_step_kv"] = True
+    reuse = _pipeline(module, od_config=config)
+    default = _pipeline(module)
+    config.model_config["lingbot_reuse_last_step_kv"] = False
+    assert reuse._dmd_blocks.reuse_last_step_kv is True
+    assert default._dmd_blocks.reuse_last_step_kv is False
+    assert reuse._dmd_blocks is reuse._dmd_blocks
+
+
+@pytest.mark.parametrize("reuse_last_step_kv", [False, True])
+def test_fixed_dmd_transition_and_cache_commit_trace(reuse_last_step_kv: bool) -> None:
     module = _load_pipeline_module()
     transformer = _RecordingTransformer()
-    pipeline = _pipeline(module, transformer=transformer)
+    config = _od_config()
+    config.model_config["lingbot_reuse_last_step_kv"] = reuse_last_step_kv
+    pipeline = _pipeline(module, transformer=transformer, od_config=config)
 
     result = pipeline(_request())
 
@@ -1383,14 +1502,18 @@ def test_fixed_dmd_transition_and_cache_commit_trace() -> None:
             current = x0
     torch.testing.assert_close(result.output, current)
 
+    expected_timesteps = [timestep for timestep, _ in warped_schedule]
+    if not reuse_last_step_kv:
+        expected_timesteps.append(0.0)
     torch.testing.assert_close(
-        torch.cat([call["timestep"] for call in transformer.calls]),
-        torch.tensor([*(timestep for timestep, _ in warped_schedule), 0.0]),
+        torch.cat([call["timestep"] for call in transformer.calls]), torch.tensor(expected_timesteps)
     )
-    assert [call["update_cache"] for call in transformer.calls] == [False, False, False, False, True]
-    assert [call["start_frame"] for call in transformer.calls] == [0, 0, 0, 0, 0]
+    commit_flags = [False] * (len(expected_timesteps) - 1) + [True]
+    assert [call["update_cache"] for call in transformer.calls] == commit_flags
+    assert [call["start_frame"] for call in transformer.calls] == [0] * len(expected_timesteps)
     assert len({call["cache_id"] for call in transformer.calls}) == 1
-    torch.testing.assert_close(transformer.calls[-1]["hidden_states"][:, :16], result.output)
+    committed_latent = result.output + warped_schedule[-1][1] if reuse_last_step_kv else result.output
+    torch.testing.assert_close(transformer.calls[-1]["hidden_states"][:, :16], committed_latent)
 
 
 @pytest.mark.parametrize(
@@ -2389,10 +2512,13 @@ def test_forward_rejects_a_stepwise_camera_action_script(extra_args) -> None:
         pipeline(_request(sampling=sampling))
 
 
-def test_stepwise_progress_metadata_and_commit_trace() -> None:
+@pytest.mark.parametrize("reuse_last_step_kv", [False, True])
+def test_stepwise_progress_metadata_and_commit_trace(reuse_last_step_kv: bool) -> None:
     module = _load_pipeline_module()
     transformer = _RecordingTransformer()
-    pipeline = _pipeline(module, transformer=transformer)
+    config = _od_config()
+    config.model_config["lingbot_reuse_last_step_kv"] = reuse_last_step_kv
+    pipeline = _pipeline(module, transformer=transformer, od_config=config)
     pipeline._ar_height = 16
     pipeline._ar_width = 16
     state = _stepwise_state(num_frames=21)
@@ -2406,8 +2532,9 @@ def test_stepwise_progress_metadata_and_commit_trace() -> None:
     assert [output.chunk_index for output in outputs] == [0, 1]
     assert all(output.total_chunks == 2 for output in outputs)
     assert outputs[-1].finished is True
-    assert [call["update_cache"] for call in transformer.calls] == [False, False, False, False, True] * 2
-    assert [call["start_frame"] for call in transformer.calls] == [0] * 5 + [3] * 5
+    calls_per_chunk = 4 if reuse_last_step_kv else 5
+    assert [call["update_cache"] for call in transformer.calls] == ([False] * (calls_per_chunk - 1) + [True]) * 2
+    assert [call["start_frame"] for call in transformer.calls] == [0] * calls_per_chunk + [3] * calls_per_chunk
     assert fake.commits == ["main", "main"]
     for chunk_index, output in enumerate(outputs):
         metadata = output.output["metadata"]["ar_diffusion"]

tests/diffusion/models/magi2/test_pipeline_magi2.py

@@ -171,7 +171,7 @@ def test_huggingface_url_resolves_pinned_snapshot(tmp_path, monkeypatch):
         else pytest.fail("unexpected snapshot request")
     )
     monkeypatch.setattr(
-        "vllm.transformers_utils.repo_utils.hf_api",
+        "vllm_omni.transformers_utils.repo_utils.hf_api",
         lambda: api,
     )
     assert _resolve_checkpoint_root(

tests/diffusion/models/minimax_h3/test_minimax_h3_contract.py

@@ -3,7 +3,7 @@
 
 import json
 import sys
-from contextlib import contextmanager
+from contextlib import contextmanager, nullcontext
 from multiprocessing.reduction import ForkingPickler
 from types import SimpleNamespace
 from typing import Any, TypeVar
@@ -2837,41 +2837,63 @@ def test_dit_encoder_selection_keeps_vae_resident():
     component.offload_to_cpu.assert_not_called()
 
 
-def test_keyframe_encode_pins_and_restores_cudnn_settings():
+@pytest.mark.parametrize(
+    ("capability", "allow_tf32"),
+    [((8, 6), True), ((9, 0), False), ((10, 0), True), ((10, 3), True), (None, True)],
+)
+@pytest.mark.parametrize("initial_tf32", [False, True])
+@pytest.mark.parametrize("fail_encode", [False, True])
+def test_keyframe_encode_pins_and_restores_cudnn_settings(
+    monkeypatch, capability, allow_tf32, initial_tf32, fail_encode
+):
+    from vllm.platforms.interface import DeviceCapability
+
     from vllm_omni.diffusion.models.minimax_h3.vae import (
         _minimax_h3_keyframe_encode_context,
     )
+    from vllm_omni.platforms import current_omni_platform
+
+    def get_capability(device_id):
+        assert device_id == 1
+        return DeviceCapability(*capability) if capability is not None else None
+
+    monkeypatch.setattr(current_omni_platform, "is_cuda", lambda: True)
+    monkeypatch.setattr(current_omni_platform, "get_device_capability", get_capability)
 
     cudnn = torch.backends.cudnn
-    original = (
-        cudnn.enabled,
-        cudnn.benchmark,
-        cudnn.deterministic,
-        cudnn.allow_tf32,
-    )
-    try:
-        cudnn.enabled = False
-        cudnn.benchmark = True
-        cudnn.deterministic = False
-        cudnn.allow_tf32 = False
-
-        with _minimax_h3_keyframe_encode_context(torch.device("cuda")):
+    with cudnn.flags(enabled=False, benchmark=True, deterministic=False, allow_tf32=initial_tf32):
+        error_context = pytest.raises(RuntimeError, match="encode failed") if fail_encode else nullcontext()
+        with error_context, _minimax_h3_keyframe_encode_context(torch.device("cuda:1")):
             assert cudnn.enabled
             assert not cudnn.benchmark
             assert cudnn.deterministic
-            assert cudnn.allow_tf32
+            assert cudnn.allow_tf32 is allow_tf32
+            if fail_encode:
+                raise RuntimeError("encode failed")
 
         assert not cudnn.enabled
         assert cudnn.benchmark
         assert not cudnn.deterministic
-        assert not cudnn.allow_tf32
-    finally:
-        (
-            cudnn.enabled,
-            cudnn.benchmark,
-            cudnn.deterministic,
-            cudnn.allow_tf32,
-        ) = original
+        assert cudnn.allow_tf32 is initial_tf32
+
+
+def test_keyframe_encode_cpu_keeps_cudnn_settings(monkeypatch):
+    from vllm_omni.diffusion.models.minimax_h3.vae import (
+        _minimax_h3_keyframe_encode_context,
+    )
+    from vllm_omni.platforms import current_omni_platform
+
+    def unexpected_capability_query(device_id):
+        pytest.fail("CPU keyframe encode must not query CUDA capability")
+
+    monkeypatch.setattr(current_omni_platform, "get_device_capability", unexpected_capability_query)
+    cudnn = torch.backends.cudnn
+    with cudnn.flags(enabled=False, benchmark=True, deterministic=False, allow_tf32=False):
+        with _minimax_h3_keyframe_encode_context(torch.device("cpu")):
+            assert not cudnn.enabled
+            assert cudnn.benchmark
+            assert not cudnn.deterministic
+            assert not cudnn.allow_tf32
 
 
 @pytest.mark.parametrize("fail_encode", [False, True])

tests/diffusion/models/minimax_h3/test_minimax_h3_modulation.py

@@ -10,7 +10,6 @@ from vllm_omni.diffusion.layers.indexed_modulation import (
     _MAX_1D_GRID_SIZE,
     _iter_row_chunks,
     _launch_row_chunks,
-    _use_hopper_bf16_affine_semantics,
     indexed_gate,
     indexed_gate_rms_norm_scale_shift,
     indexed_scale_shift_,
@@ -194,52 +193,28 @@ def test_fused_modulation_preserves_bf16_residual_boundary() -> None:
 
 
 @pytest.mark.cuda
-@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
-@pytest.mark.skipif(
-    torch.cuda.is_available() and torch.cuda.get_device_capability()[0] != 9,
-    reason="SM90-specific precision path",
-)
+@pytest.mark.skipif(not current_omni_platform.is_cuda(), reason="CUDA required")
 @pytest.mark.skipif(not HAS_TRITON, reason="Triton required")
-def test_hopper_fused_modulation_matches_pytorch_reference() -> None:
-    torch.manual_seed(42)
-    rows, hidden_size, conditions = 128, 3072, 4
-    residual = torch.randn(rows, hidden_size, device="cuda", dtype=torch.bfloat16)
-    gate = torch.randn(conditions, hidden_size, device="cuda", dtype=torch.bfloat16)
-    branch = torch.randn(rows, hidden_size, device="cuda", dtype=torch.bfloat16)
-    weight = torch.randn(hidden_size, device="cuda", dtype=torch.bfloat16)
-    shift = torch.randn(conditions, hidden_size, device="cuda", dtype=torch.bfloat16)
-    scale = torch.randn(conditions, hidden_size, device="cuda", dtype=torch.bfloat16)
-    indices = torch.arange(rows, device="cuda") % conditions
+@pytest.mark.parametrize("fused_gate", [False, True])
+def test_fused_modulation_keeps_fp32_affine_intermediates(fused_gate: bool) -> None:
+    device = current_omni_platform.get_torch_device()
+    rows, hidden_size = 2, 5376
+    x = torch.ones(rows, hidden_size, device=device, dtype=torch.bfloat16)
+    weight = torch.full((hidden_size,), 1.125, device=device, dtype=x.dtype)
+    shift = torch.zeros(1, hidden_size, device=device, dtype=x.dtype)
+    scale = torch.full_like(shift, 0.09375)
+    indices = torch.zeros(rows, device=device, dtype=torch.int64)
     eps = 1e-6
+    if fused_gate:
+        residual, actual = indexed_gate_rms_norm_scale_shift(
+            x, torch.zeros_like(shift), x, weight, shift, scale, indices, eps
+        )
+        assert torch.equal(residual, x)
+    else:
+        actual = rms_norm_indexed_scale_shift(x, weight, shift, scale, indices, eps)
 
-    reference_residual = (residual + gate.index_select(0, indices) * branch).to(torch.bfloat16)
-    normalized = reference_residual.float()
-    variance = normalized.pow(2).mean(-1, keepdim=True)
-    normalized = (weight.float() * (normalized * torch.rsqrt(variance + eps))).to(torch.bfloat16)
-    reference_modulated = (normalized * (1.0 + scale.index_select(0, indices)) + shift.index_select(0, indices)).to(
-        torch.bfloat16
-    )
-
-    actual_residual, actual_modulated = indexed_gate_rms_norm_scale_shift(
-        residual,
-        gate,
-        branch,
-        weight,
-        shift,
-        scale,
-        indices,
-        eps,
-    )
-
-    torch.testing.assert_close(actual_residual, reference_residual, atol=5e-2, rtol=5e-2)
-    torch.testing.assert_close(actual_modulated, reference_modulated, atol=5e-2, rtol=5e-2)
-
-
-@pytest.mark.cuda
-@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
-def test_hopper_precision_dispatch_is_architecture_scoped(monkeypatch) -> None:
-    tensor = torch.empty(0, device="cuda")
-    monkeypatch.setattr(torch.cuda, "get_device_capability", lambda _device: (9, 0))
-    assert _use_hopper_bf16_affine_semantics(tensor)
-    monkeypatch.setattr(torch.cuda, "get_device_capability", lambda _device: (10, 3))
-    assert not _use_hopper_bf16_affine_semantics(tensor)
+    # Without the RMS epsilon, 1.125 * 1.09375 = 1.23046875 is exactly
+    # halfway between two BF16 values. FP32 normalization keeps the epsilon
+    # and rounds down to 1.2265625. Premature BF16 rounding drops the epsilon
+    # and instead rounds the tie up to 1.234375, changing every output row.
+    assert torch.equal(actual, torch.full_like(x, 1.2265625))

tests/diffusion/models/minimax_h3/test_minimax_h3_quantization_quality.py

@@ -16,6 +16,7 @@ from tests.diffusion.quantization.test_quantization_quality import (
     _maybe_save_output,
 )
 from tests.helpers.mark import hardware_marks
+from tests.helpers.mock import patch_hf_snapshot_download
 
 _MINIMAX_H3_REPO = "MiniMaxAI/MiniMax-H3"
 _MINIMAX_H3_REVISION = "48d93ede732756e404a3b1b2f3b3a9b5a22f6cfc"
@@ -258,13 +259,13 @@ def test_resolve_fl2va_model_ref(tmp_path, monkeypatch):
     fl2va_root.mkdir()
     (fl2va_root / "model_index.json").write_text("{}", encoding="utf-8")
 
-    def fake_snapshot_download(self, *, repo_id, revision, allow_patterns):
+    def fake_snapshot_download(*, repo_id, revision, allow_patterns):
         assert repo_id == _MINIMAX_H3_REPO
         assert revision == _MINIMAX_H3_REVISION
         assert allow_patterns == ["FL2VA/**"]
         return str(tmp_path)
 
-    monkeypatch.setattr("huggingface_hub.HfApi.snapshot_download", fake_snapshot_download)
+    patch_hf_snapshot_download(monkeypatch, fake_snapshot_download, hf_home=tmp_path)
     assert _resolve_fl2va_model_ref() == str(fl2va_root)
 
 

tests/diffusion/models/nextstep_1_1/test_nextstep_rectangular.py

@@ -0,0 +1,52 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+from types import MethodType, SimpleNamespace
+from unittest.mock import Mock
+
+import pytest
+import torch
+
+from vllm_omni.diffusion.models.nextstep_1_1 import pipeline_nextstep_1_1 as pipeline_module
+from vllm_omni.diffusion.models.nextstep_1_1.modeling_nextstep import NextStepModel
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
+
+
+@pytest.mark.parametrize("height,width", [(32, 32), (32, 48), (48, 32)])
+@pytest.mark.parametrize("batch", [1, 2])
+def test_forward_preserves_latent_grid(monkeypatch, height, width, batch):
+    pipeline = object.__new__(pipeline_module.NextStep11Pipeline)
+    torch.nn.Module.__init__(pipeline)
+    pipeline._NextStep11Pipeline__device = torch.device("cpu")
+    pipeline.config = SimpleNamespace(use_gen_pos_embed=False)
+    pipeline.down_factor = 16
+    pipeline.scaling_factor = 2.0
+    pipeline.shift_factor = 0.25
+    grid = SimpleNamespace(config=SimpleNamespace(latent_channels=2, latent_patch_size=2))
+    expected = torch.arange(batch * 2 * (height // 8) * (width // 8), dtype=torch.float32)
+    expected = expected.reshape(batch, 2, height // 8, width // 8)
+    tokens = NextStepModel.patchify(grid, expected)
+    pipeline.model = Mock(
+        return_value=SimpleNamespace(last_hidden_state=torch.zeros(batch, 1, 4), past_key_values=None)
+    )
+    pipeline.model.forward_model.return_value = pipeline.model.return_value
+    pipeline.model.unpatchify = MethodType(NextStepModel.unpatchify, grid)
+    pipeline.tokenizer = Mock(
+        bos_token=None,
+        return_value=SimpleNamespace(
+            input_ids=torch.ones(batch, 2, dtype=torch.long), attention_mask=torch.ones(batch, 2, dtype=torch.long)
+        ),
+    )
+    pipeline._check_input = Mock(return_value=(["teapot"], None))
+    pipeline._build_captions = Mock(return_value=(["teapot"] * batch, None, 1, 1.0))
+    pipeline._add_prefix_ids = lambda _hw, ids, mask: (ids, mask)
+    pipeline.decoding = Mock(return_value=tokens)
+    pipeline.vae = SimpleNamespace(dtype=torch.float32, decode=lambda latent: SimpleNamespace(sample=latent))
+    monkeypatch.setattr(pipeline_module, "StaticCache", lambda **_kwargs: None)
+    params = OmniDiffusionSamplingParams(
+        height=height, width=width, guidance_scale=1.0, num_outputs_per_prompt=batch, seed=None
+    )
+    output = pipeline.forward(SimpleNamespace(prompts=["teapot"], sampling_params=params))
+    torch.testing.assert_close(output.output, expected / 2 + 0.25, rtol=0, atol=0)