diff --git a/docs/configuration/stage_configs.md b/docs/configuration/stage_configs.md index 290c15b94..38eceb634 100644 --- a/docs/configuration/stage_configs.md +++ b/docs/configuration/stage_configs.md @@ -38,11 +38,12 @@ Each entry under `stages:` accepts any `StageDeployConfig` field directly (no ne |-------|------|----------|---------|-------------| | `stage_id` | int | required | — | Stage identity; matched against `PipelineConfig.stages[*].stage_id`. | | `max_num_seqs` | int | optional | `64` | Max concurrent sequences per stage. | -| `gpu_memory_utilization` | float | optional | `0.9` | Per-stage memory budget. | +| `gpu_memory_utilization` | float | optional | `0.9` | Per-stage total memory target; used for automatic KV-cache sizing. | +| `kv_cache_memory_bytes` | int \| null | optional | `null` | Explicit per-rank KV-cache budget in bytes; overrides automatic sizing. | | `tensor_parallel_size` | int | optional | `1` | TP degree for this stage. | | `enforce_eager` | bool | optional | `false` | Disable CUDA graphs. | -| `max_num_batched_tokens` | int | optional | `32768` | Prefill budget. | -| `max_model_len` | int \| null | optional | `null` | Per-stage context length (auto-sets `VLLM_ALLOW_LONG_MAX_MODEL_LEN=1` when larger than HF default). | +| `max_num_batched_tokens` | int | optional | `32768` | Per-stage prefill/token budget; also contributes to the native maximum in-flight token limit. | +| `max_model_len` | int \| null | optional | `null` | Per-sequence context or KV length; `-1` enables native cache-capacity auto-fitting, while values above the HF default auto-set `VLLM_ALLOW_LONG_MAX_MODEL_LEN=1`. | | `async_scheduling` | bool \| null | optional | `null` | Per-stage async scheduling toggle. | | `devices` | str | optional | `"0"` | `CUDA_VISIBLE_DEVICES`-style device list. | | `output_connectors` | dict \| null | optional | `null` | Keyed by `to_stage_`; values are names registered under top-level `connectors:`. | diff --git a/tests/diffusion/diffusion_kv/test_initialization.py b/tests/diffusion/diffusion_kv/test_initialization.py new file mode 100644 index 000000000..126b25d18 --- /dev/null +++ b/tests/diffusion/diffusion_kv/test_initialization.py @@ -0,0 +1,297 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from types import SimpleNamespace +from unittest.mock import Mock + +import pytest +import torch +from vllm.v1.kv_cache_interface import FullAttentionSpec + +import vllm_omni.diffusion.diffusion_kv.initialization as initialization +import vllm_omni.diffusion.vllm_config as diffusion_vllm_config +from vllm_omni.diffusion.diffusion_kv.config import DiffusionKVCacheMode +from vllm_omni.diffusion.sched.request_scheduler import RequestScheduler +from vllm_omni.quantization import build_quant_config + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion] + + +class _Executor: + def __init__(self) -> None: + self.specs = [{"layer0": object()}, {"layer0": object()}] + self.memory = [1024, 1024] + self.configs = None + self.profile_requests = None + + def get_kv_cache_specs(self): + return self.specs + + def determine_available_kv_memory(self, profile_requests): + self.profile_requests = profile_requests + return self.memory + + def set_kv_cache_configs(self, configs) -> None: + self.configs = configs + + +def _od_config(**overrides): + values = dict( + model="test-diffusion", + diffusion_kv_mode=DiffusionKVCacheMode.PAGED_SCHEDULER, + dtype=torch.bfloat16, + max_model_len=64, + quantization_config=None, + tf_model_config=None, + enforce_eager=True, + is_moe=False, + additional_config={}, + profiler_config=None, + parallel_config=SimpleNamespace( + tensor_parallel_size=1, + pipeline_parallel_size=1, + data_parallel_size=1, + enable_expert_parallel=False, + ), + gpu_memory_utilization=0.9, + kv_cache_memory_bytes=None, + max_num_batched_tokens=64, + num_gpus=1, + ) + values.update(overrides) + return SimpleNamespace(**values) + + +def test_control_plane_builds_worker_and_scheduler_configs(monkeypatch) -> None: + executor = _Executor() + od_config = _od_config(num_gpus=2) + vllm_config = SimpleNamespace( + model_config=SimpleNamespace(max_model_len=4096), + max_in_flight_tokens=128, + ) + worker_configs = [object(), object()] + scheduler_kv_cache_config = object() + profile_requests = [object()] + monkeypatch.setattr(initialization, "create_diffusion_vllm_config", lambda *args, **kwargs: vllm_config) + build = Mock(return_value=(worker_configs, scheduler_kv_cache_config)) + monkeypatch.setattr( + initialization, + "build_native_kv_cache_configs", + build, + ) + monkeypatch.setattr(initialization, "resolve_kv_cache_block_sizes", lambda *args: (16, 16)) + + result = initialization.initialize_diffusion_kv_control_plane( + executor, + od_config, + profile_requests=profile_requests, + device=torch.device("cpu"), + ) + + build.assert_called_once_with(vllm_config, executor.specs, executor.memory) + assert result == (scheduler_kv_cache_config, 16, 16, vllm_config) + assert executor.configs == worker_configs + assert executor.profile_requests is profile_requests + + +def test_dense_mode_skips_worker_kv_initialization(monkeypatch) -> None: + executor = _Executor() + od_config = _od_config(diffusion_kv_mode=DiffusionKVCacheMode.DENSE_LEGACY) + get_device = Mock(side_effect=AssertionError("dense mode must not resolve a KV device")) + monkeypatch.setattr(initialization.current_omni_platform, "get_torch_device", get_device) + + result = initialization.initialize_diffusion_kv_control_plane(executor, od_config) + + assert result is None + assert executor.configs is None + get_device.assert_not_called() + + +def test_control_plane_rejects_rank_count_mismatch(monkeypatch) -> None: + executor = _Executor() + executor.memory = [1024] + od_config = _od_config(num_gpus=2) + monkeypatch.setattr(initialization, "create_diffusion_vllm_config", lambda *args, **kwargs: object()) + + with pytest.raises(ValueError, match="rank count mismatch"): + initialization.initialize_diffusion_kv_control_plane( + executor, + od_config, + profile_requests=[object()], + device=torch.device("cpu"), + ) + + +def test_paged_control_plane_requires_prepared_profile_request() -> None: + with pytest.raises(ValueError, match="prepared memory profile requests"): + initialization.initialize_diffusion_kv_control_plane( + _Executor(), + _od_config(), + device=torch.device("cpu"), + ) + + +def test_native_config_builder_produces_matching_worker_and_scheduler_pools() -> None: + od_config = _od_config() + vllm_config = diffusion_vllm_config.create_diffusion_vllm_config(torch.device("cpu"), od_config) + spec = FullAttentionSpec( + block_size=vllm_config.cache_config.block_size, + num_kv_heads=2, + head_size=8, + dtype=torch.bfloat16, + non_causal=True, + ) + num_blocks = 32 + + worker_configs, scheduler_kv_cache_config = initialization.build_native_kv_cache_configs( + vllm_config, + [{"layer0": spec}], + [spec.page_size_bytes * num_blocks], + ) + + assert len(worker_configs) == 1 + assert worker_configs[0].num_blocks == num_blocks + assert scheduler_kv_cache_config.num_blocks == num_blocks + assert worker_configs[0].kv_cache_groups == scheduler_kv_cache_config.kv_cache_groups + + +def test_native_control_plane_config_initializes_scheduler() -> None: + od_config = _od_config() + vllm_config = diffusion_vllm_config.create_diffusion_vllm_config(torch.device("cpu"), od_config) + spec = FullAttentionSpec( + block_size=vllm_config.cache_config.block_size, + num_kv_heads=2, + head_size=8, + dtype=torch.bfloat16, + non_causal=True, + ) + _, scheduler_kv_cache_config = initialization.build_native_kv_cache_configs( + vllm_config, + [{"layer0": spec}], + [spec.page_size_bytes * 32], + ) + + scheduler = RequestScheduler() + scheduler.initialize( + od_config, + kv_cache_config=scheduler_kv_cache_config, + scheduler_block_size=spec.block_size, + hash_block_size=spec.block_size, + kv_vllm_config=vllm_config, + ) + + assert scheduler._diffusion_kv_manager is not None + scheduler.close() + + +def test_dense_config_keeps_new_paged_sizing_inputs_inactive() -> None: + od_config = _od_config( + diffusion_kv_mode=DiffusionKVCacheMode.DENSE_LEGACY, + parallel_config=SimpleNamespace( + tensor_parallel_size=1, + pipeline_parallel_size=4, + data_parallel_size=1, + enable_expert_parallel=False, + ), + gpu_memory_utilization=0.5, + kv_cache_memory_bytes=4096, + max_num_batched_tokens=64, + ) + vllm_config = diffusion_vllm_config.create_base_diffusion_vllm_config(torch.device("cpu"), od_config) + original_pp_size = vllm_config.parallel_config.pipeline_parallel_size + original_memory_utilization = vllm_config.cache_config.gpu_memory_utilization + original_memory_bytes = vllm_config.cache_config.kv_cache_memory_bytes + original_token_budget = vllm_config.scheduler_config.max_num_batched_tokens + + diffusion_vllm_config.configure_diffusion_vllm_config(vllm_config, od_config) + + assert vllm_config.parallel_config.pipeline_parallel_size == original_pp_size + assert vllm_config.cache_config.gpu_memory_utilization == original_memory_utilization + assert vllm_config.cache_config.kv_cache_memory_bytes == original_memory_bytes + assert vllm_config.scheduler_config.max_num_batched_tokens == original_token_budget + + +def test_paged_config_forwards_gpu_memory_utilization_to_native_cache_config() -> None: + vllm_config = diffusion_vllm_config.create_diffusion_vllm_config( + torch.device("cpu"), + _od_config(gpu_memory_utilization=0.42), + ) + + assert vllm_config.cache_config.gpu_memory_utilization == 0.42 + + +def test_diffusion_vllm_model_config_supplies_dtype_for_quant_methods() -> None: + quantization_config = build_quant_config( + { + "quant_method": "modelopt", + "quant_algo": "FP8", + "ignore": [], + } + ) + model_config = diffusion_vllm_config._make_diffusion_vllm_model_config( + _od_config( + diffusion_kv_mode=DiffusionKVCacheMode.DENSE_LEGACY, + quantization_config=quantization_config, + ) + ) + + assert model_config.dtype is torch.bfloat16 + assert model_config.quantization == "modelopt" + assert model_config.quantization_config is quantization_config + assert model_config.is_quantized + assert model_config.is_moe is False + assert model_config.is_nvfp4_quantized() is False + assert model_config.original_max_model_len == model_config.max_model_len + + unquantized_model_config = diffusion_vllm_config._make_diffusion_vllm_model_config( + _od_config( + diffusion_kv_mode=DiffusionKVCacheMode.DENSE_LEGACY, + quantization_config=None, + ) + ) + assert not unquantized_model_config.is_quantized + + +def test_minus_one_max_model_len_uses_model_limit() -> None: + od_config = SimpleNamespace( + max_model_len=-1, + tf_model_config=SimpleNamespace(max_position_embeddings=4096), + diffusion_kv_mode=DiffusionKVCacheMode.PAGED_SCHEDULER, + ) + + assert diffusion_vllm_config.resolve_diffusion_max_model_len(od_config) == 4096 + + model_config = diffusion_vllm_config._make_diffusion_vllm_model_config( + _od_config( + max_model_len=-1, + tf_model_config=SimpleNamespace(max_position_embeddings=4096), + ) + ) + assert model_config.max_model_len == 4096 + assert model_config.original_max_model_len == -1 + + +def test_minus_one_max_model_len_is_auto_fitted_by_native_cache_sizing() -> None: + vllm_config = diffusion_vllm_config.create_diffusion_vllm_config( + torch.device("cpu"), + _od_config( + max_model_len=-1, + tf_model_config=SimpleNamespace(max_position_embeddings=4096), + ), + ) + spec = FullAttentionSpec( + block_size=vllm_config.cache_config.block_size, + num_kv_heads=2, + head_size=8, + dtype=torch.bfloat16, + non_causal=True, + ) + + initialization.build_native_kv_cache_configs( + vllm_config, + [{"layer0": spec}], + [spec.page_size_bytes * 16], + ) + + assert vllm_config.model_config.original_max_model_len == -1 + assert vllm_config.model_config.max_model_len == 256 diff --git a/tests/diffusion/diffusion_kv/test_manager.py b/tests/diffusion/diffusion_kv/test_manager.py new file mode 100644 index 000000000..5152c2a67 --- /dev/null +++ b/tests/diffusion/diffusion_kv/test_manager.py @@ -0,0 +1,211 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import pytest +import torch +import vllm.v1.core.single_type_kv_cache_manager as native_kv_managers +from vllm.v1.kv_cache_interface import FullAttentionSpec, KVCacheConfig, KVCacheGroupSpec, KVCacheTensor + +from vllm_omni.diffusion.diffusion_kv.manager import DiffusionKVAdmissionError, DiffusionKVCacheManager +from vllm_omni.diffusion.diffusion_kv.request import DiffusionKVContext, DiffusionKVRequest + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion] + +BLOCK_SIZE = 4 + + +def _config(num_blocks: int) -> KVCacheConfig: + native_kv_managers.register_all_kvcache_specs(None) + spec = FullAttentionSpec( + block_size=BLOCK_SIZE, + num_kv_heads=2, + head_size=8, + dtype=torch.bfloat16, + ) + return KVCacheConfig( + num_blocks=num_blocks, + kv_cache_tensors=[KVCacheTensor(size=spec.page_size_bytes * num_blocks, shared_by=["layer0"])], + kv_cache_groups=[KVCacheGroupSpec(layer_names=["layer0"], kv_cache_spec=spec)], + ) + + +def _request(public_id: str, sequence_id: int, *, seq_len: int = 8, kv_contexts=()) -> DiffusionKVRequest: + return DiffusionKVRequest( + f"{public_id}/diffusion-kv/{sequence_id}", + sequence_id=sequence_id, + prefix_len=4, + target_len=4, + seq_len=seq_len, + kv_contexts=kv_contexts, + ) + + +def _manager(num_blocks: int, *, max_model_len: int = 64) -> DiffusionKVCacheManager: + return DiffusionKVCacheManager( + _config(num_blocks), + max_model_len=max_model_len, + scheduler_block_size=BLOCK_SIZE, + hash_block_size=BLOCK_SIZE, + ) + + +def test_reserve_and_free_multi_cfg_request() -> None: + manager = _manager(8) + free_before = manager.native_manager.block_pool.get_num_free_blocks() + requests = (_request("public", 0), _request("public", 1)) + + metadata = manager.reserve_request("public", requests) + + assert metadata is not None + assert metadata.request_id == "public" + assert metadata.allocation_generation == 1 + assert [sequence.sequence_id for sequence in metadata.sequences] == [0, 1] + assert [sequence.prefix_len for sequence in metadata.sequences] == [4, 4] + assert [sequence.target_len for sequence in metadata.sequences] == [4, 4] + assert [sequence.seq_len for sequence in metadata.sequences] == [8, 8] + assert [len(sequence.block_ids[0]) for sequence in metadata.sequences] == [2, 2] + assert all(sequence.context_ids == () for sequence in metadata.sequences) + assert metadata.contexts == () + assert metadata == manager.get_metadata("public") + assert manager.native_manager.block_pool.get_num_free_blocks() == free_before - 4 + + manager.free_request("public") + assert manager.native_manager.block_pool.get_num_free_blocks() == free_before + + +def test_impossible_cfg_allocation_rolls_back_and_fails_fast() -> None: + manager = _manager(3) + free_before = manager.native_manager.block_pool.get_num_free_blocks() + + with pytest.raises(DiffusionKVAdmissionError, match="cannot fit even when the block pool is empty"): + manager.reserve_request( + "public", + (_request("public", 0), _request("public", 1)), + ) + + assert manager.has_request("public") is False + assert manager.native_manager.block_pool.get_num_free_blocks() == free_before + + +def test_impossible_cfg_allocation_fails_fast_while_pool_is_busy() -> None: + manager = _manager(5) + assert manager.reserve_request("running", (_request("running", 0),)) is not None + free_before = manager.native_manager.block_pool.get_num_free_blocks() + + with pytest.raises(DiffusionKVAdmissionError, match="required_blocks=6, available_blocks=4"): + manager.reserve_request( + "impossible", + tuple(_request("impossible", sequence_id) for sequence_id in range(3)), + ) + + assert manager.has_request("impossible") is False + assert manager.native_manager.block_pool.get_num_free_blocks() == free_before + + +def test_temporary_capacity_pressure_returns_none() -> None: + manager = _manager(3) + assert manager.reserve_request("running", (_request("running", 0),)) is not None + + allocation = manager.reserve_request("waiting", (_request("waiting", 0),)) + + assert allocation is None + assert manager.has_request("waiting") is False + + +def test_partial_cfg_rollback_can_retry_after_capacity_is_released() -> None: + manager = _manager(5) + assert manager.reserve_request("running", (_request("running", 0),)) is not None + waiting = (_request("waiting", 0), _request("waiting", 1)) + + assert manager.reserve_request("waiting", waiting) is None + assert manager.has_request("waiting") is False + + manager.free_request("running") + metadata = manager.reserve_request("waiting", waiting) + + assert metadata is not None + assert [len(sequence.block_ids[0]) for sequence in metadata.sequences] == [2, 2] + + +def test_reallocated_request_gets_a_new_allocation_generation() -> None: + manager = _manager(3) + request = (_request("public", 0),) + + first = manager.reserve_request("public", request) + assert first is not None + manager.free_request("public") + second = manager.reserve_request("public", request) + + assert second is not None + assert second.allocation_generation == first.allocation_generation + 1 + + +def test_cfg_allocation_exception_rolls_back(monkeypatch) -> None: + manager = _manager(8) + free_before = manager.native_manager.block_pool.get_num_free_blocks() + native_allocate = manager.native_manager.allocate_slots + num_calls = 0 + + def fail_second_allocation(*args, **kwargs): + nonlocal num_calls + num_calls += 1 + if num_calls == 2: + raise RuntimeError("injected allocation failure") + return native_allocate(*args, **kwargs) + + monkeypatch.setattr(manager.native_manager, "allocate_slots", fail_second_allocation) + + with pytest.raises(RuntimeError, match="injected allocation failure"): + manager.reserve_request( + "public", + (_request("public", 0), _request("public", 1)), + ) + + assert manager.has_request("public") is False + assert manager.native_manager.block_pool.get_num_free_blocks() == free_before + + +def test_rejects_independent_context_until_role_routing_is_implemented() -> None: + manager = _manager(8) + context = DiffusionKVContext(context_id="text", cache_role="cross.text", num_tokens=8) + + with pytest.raises(DiffusionKVAdmissionError, match="DiffusionKVContext"): + manager.reserve_request( + "public", + (_request("public", 0, kv_contexts=(context,)),), + ) + + +def test_rejects_sequence_longer_than_admission_bound() -> None: + manager = _manager(8, max_model_len=8) + + with pytest.raises(DiffusionKVAdmissionError, match="exceeds max_model_len"): + manager.reserve_request("public", (_request("public", 0, seq_len=9),)) + + +@pytest.mark.parametrize( + "requests", + [ + (_request("public", 1),), + (_request("public", 1), _request("public", 0)), + (_request("public", 0), _request("public", 2)), + ], +) +def test_rejects_ambiguous_sequence_order(requests) -> None: + manager = _manager(8) + + with pytest.raises(ValueError, match="ordered by contiguous sequence_id"): + manager.reserve_request("public", requests) + + +def test_close_releases_every_public_request() -> None: + manager = _manager(8) + free_before = manager.native_manager.block_pool.get_num_free_blocks() + assert manager.reserve_request("first", (_request("first", 0),)) is not None + assert manager.reserve_request("second", (_request("second", 0),)) is not None + + manager.close() + + assert manager.native_manager.block_pool.get_num_free_blocks() == free_before + assert manager.has_request("first") is False + assert manager.has_request("second") is False diff --git a/tests/diffusion/diffusion_kv/test_request.py b/tests/diffusion/diffusion_kv/test_request.py index 43734ca31..ae5609b2a 100644 --- a/tests/diffusion/diffusion_kv/test_request.py +++ b/tests/diffusion/diffusion_kv/test_request.py @@ -107,11 +107,8 @@ def test_request_rejects_non_context_values() -> None: def _manager(*, enable_caching: bool) -> KVCacheManager: # vLLM 0.26 lazily registers its built-in KV specs. Another test may have # already registered an out-of-tree spec, making the registry non-empty - # before FullAttentionSpec is registered. Initialize the native mappings - # explicitly; vLLM 0.25 has no registration hook and needs no setup. - register_all = getattr(native_kv_managers, "register_all_kvcache_specs", None) - if register_all is not None: - register_all(None) + # before FullAttentionSpec is registered. Initialize the native mappings. + native_kv_managers.register_all_kvcache_specs(None) spec = FullAttentionSpec( block_size=BLOCK_SIZE, diff --git a/tests/diffusion/diffusion_kv/test_spec_discovery.py b/tests/diffusion/diffusion_kv/test_spec_discovery.py new file mode 100644 index 000000000..1d1d6c8d9 --- /dev/null +++ b/tests/diffusion/diffusion_kv/test_spec_discovery.py @@ -0,0 +1,259 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from contextlib import nullcontext +from types import SimpleNamespace +from unittest.mock import Mock + +import pytest +import torch +from torch import nn +from vllm.v1.kv_cache_interface import FullAttentionSpec, KVCacheConfig, KVCacheGroupSpec, KVCacheTensor + +import vllm_omni.diffusion.worker.diffusion_worker as diffusion_worker_module +from vllm_omni.diffusion.attention.layer import Attention +from vllm_omni.diffusion.diffusion_kv.config import DiffusionKVCacheMode +from vllm_omni.diffusion.worker.diffusion_model_runner import DiffusionModelRunner +from vllm_omni.diffusion.worker.diffusion_worker import DiffusionWorker + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion] + + +class _Backend: + @classmethod + def indexes_kv_by_block_stride(cls) -> bool: + return True + + +def _attention(*, enabled: bool) -> Attention: + attention = Attention.__new__(Attention) + nn.Module.__init__(attention) + attention.paged_kv_cache_role = "primary" if enabled else None + attention.paged_kv_cache_dtype = torch.bfloat16 + attention.num_kv_heads = 2 + attention.head_size = 8 + attention.causal = False + attention.attn_backend = _Backend + return attention + + +def _runner(attention: Attention) -> DiffusionModelRunner: + runner = object.__new__(DiffusionModelRunner) + runner.od_config = SimpleNamespace(diffusion_kv_mode=DiffusionKVCacheMode.PAGED_SCHEDULER) + runner.vllm_config = SimpleNamespace( + cache_config=SimpleNamespace(block_size=4), + model_config=SimpleNamespace(dtype=torch.float16), + ) + pipeline = nn.Module() + pipeline.image_attention = attention + runner.pipeline = pipeline + return runner + + +def test_runner_discovers_native_spec_from_loaded_attention() -> None: + specs = _runner(_attention(enabled=True)).get_kv_cache_spec() + + assert set(specs) == {"image_attention"} + spec = specs["image_attention"] + assert isinstance(spec, FullAttentionSpec) + assert spec.block_size == 4 + assert spec.num_kv_heads == 2 + assert spec.head_size == 8 + assert spec.dtype is torch.bfloat16 + assert spec.indexes_kv_by_block_stride is True + assert spec.non_causal is True + + +def test_runner_rejects_paged_mode_without_cache_enabled_attention() -> None: + runner = _runner(_attention(enabled=False)) + + with pytest.raises(RuntimeError, match="no cache-enabled Attention"): + runner.get_kv_cache_spec() + + +def test_runner_retains_matching_rank_local_config() -> None: + runner = _runner(_attention(enabled=True)) + spec = runner.get_kv_cache_spec()["image_attention"] + config = KVCacheConfig( + num_blocks=8, + kv_cache_tensors=[KVCacheTensor(size=spec.page_size_bytes * 8, shared_by=["image_attention"])], + kv_cache_groups=[KVCacheGroupSpec(layer_names=["image_attention"], kv_cache_spec=spec)], + ) + + runner.set_kv_cache_config(config) + + assert runner.kv_cache_config is config + + +def test_runner_rejects_rank_local_config_for_different_layers() -> None: + runner = _runner(_attention(enabled=True)) + spec = runner.get_kv_cache_spec()["image_attention"] + config = KVCacheConfig( + num_blocks=8, + kv_cache_tensors=[KVCacheTensor(size=spec.page_size_bytes * 8, shared_by=["other_attention"])], + kv_cache_groups=[KVCacheGroupSpec(layer_names=["other_attention"], kv_cache_spec=spec)], + ) + + with pytest.raises(ValueError, match="layer mismatch"): + runner.set_kv_cache_config(config) + + +def test_worker_selects_its_rank_local_config() -> None: + worker = object.__new__(DiffusionWorker) + worker.rank = 1 + worker.od_config = SimpleNamespace(num_gpus=2) + worker.model_runner = SimpleNamespace(set_kv_cache_config=lambda config: setattr(worker, "installed", config)) + configs = [object(), object()] + + worker.set_kv_cache_configs(configs) + + assert worker.installed is configs[1] + with pytest.raises(ValueError, match="rank count mismatch"): + worker.set_kv_cache_configs(configs[:1]) + + +def test_worker_honors_explicit_kv_memory_budget(monkeypatch) -> None: + worker = object.__new__(DiffusionWorker) + worker.rank = 0 + worker.vllm_config = SimpleNamespace(cache_config=SimpleNamespace(kv_cache_memory_bytes=4096)) + worker.init_snapshot = SimpleNamespace(free_memory=16384) + worker.requested_memory = 8192 + worker.model_runner = SimpleNamespace(profile_run=Mock(), model_memory_usage=1024) + monkeypatch.setattr(diffusion_worker_module, "_all_gather_rank_values", lambda value: [value]) + profile_request = object() + + assert worker.determine_available_kv_memory(profile_request) == [4096] + worker.model_runner.profile_run.assert_called_once_with(profile_request) + + +def test_worker_treats_zero_kv_memory_budget_as_profiled_auto_sizing(monkeypatch) -> None: + worker = object.__new__(DiffusionWorker) + worker.rank = 0 + worker.init_snapshot = SimpleNamespace(free_memory=900) + worker.requested_memory = 750 + worker.vllm_config = SimpleNamespace( + cache_config=SimpleNamespace( + kv_cache_memory_bytes=0, + gpu_memory_utilization=0.75, + ) + ) + worker.model_runner = SimpleNamespace(profile_run=Mock(), model_memory_usage=100) + monkeypatch.setattr(diffusion_worker_module, "_all_gather_rank_values", lambda value: [value]) + profile_result = SimpleNamespace( + non_kv_cache_memory=400, + after_profile=SimpleNamespace(free_memory=500), + ) + monkeypatch.setattr( + diffusion_worker_module, + "memory_profiling", + lambda snapshot, weights_memory: nullcontext(profile_result), + ) + profile_request = object() + + assert worker.determine_available_kv_memory(profile_request) == [350] + worker.model_runner.profile_run.assert_called_once_with(profile_request) + + +@pytest.mark.parametrize("non_kv_cache_memory", [750, 800]) +def test_worker_rejects_non_positive_profiled_kv_memory(monkeypatch, non_kv_cache_memory: int) -> None: + worker = object.__new__(DiffusionWorker) + worker.rank = 0 + worker.init_snapshot = SimpleNamespace(free_memory=900) + worker.requested_memory = 750 + worker.vllm_config = SimpleNamespace( + cache_config=SimpleNamespace( + kv_cache_memory_bytes=None, + gpu_memory_utilization=0.75, + ) + ) + worker.model_runner = SimpleNamespace(profile_run=Mock(), model_memory_usage=100) + monkeypatch.setattr(diffusion_worker_module, "_all_gather_rank_values", lambda value: [value]) + profile_result = SimpleNamespace( + non_kv_cache_memory=non_kv_cache_memory, + after_profile=SimpleNamespace(free_memory=100), + ) + monkeypatch.setattr( + diffusion_worker_module, + "memory_profiling", + lambda snapshot, weights_memory: nullcontext(profile_result), + ) + profile_request = object() + + with pytest.raises( + RuntimeError, + match=( + "rank 0: RuntimeError: No memory remains for Diffusion KV cache after profiling: " + f"requested_memory=750 bytes, non_kv_cache_memory={non_kv_cache_memory} bytes" + ), + ): + worker.determine_available_kv_memory(profile_request) + + worker.model_runner.profile_run.assert_called_once_with(profile_request) + + +def test_worker_profiles_activation_headroom_instead_of_current_residency(monkeypatch) -> None: + worker = object.__new__(DiffusionWorker) + worker.rank = 0 + worker.init_snapshot = SimpleNamespace(free_memory=1200) + worker.requested_memory = 1000 + worker.vllm_config = SimpleNamespace( + cache_config=SimpleNamespace( + kv_cache_memory_bytes=None, + gpu_memory_utilization=0.8, + ) + ) + worker.model_runner = SimpleNamespace(profile_run=Mock(), model_memory_usage=250) + monkeypatch.setattr(diffusion_worker_module, "_all_gather_rank_values", lambda value: [value]) + profile_result = SimpleNamespace( + non_kv_cache_memory=650, + after_profile=SimpleNamespace(free_memory=550), + ) + profiling_calls = [] + + def profile(snapshot, *, weights_memory): + profiling_calls.append((snapshot, weights_memory)) + return nullcontext(profile_result) + + monkeypatch.setattr( + diffusion_worker_module, + "memory_profiling", + profile, + ) + profile_request = object() + + assert worker.determine_available_kv_memory(profile_request) == [350] + assert profiling_calls == [(worker.init_snapshot, 250)] + worker.model_runner.profile_run.assert_called_once_with(profile_request) + + +def test_worker_rejects_profile_without_preload_snapshot(monkeypatch) -> None: + worker = object.__new__(DiffusionWorker) + worker.init_snapshot = None + worker.requested_memory = None + worker.vllm_config = SimpleNamespace(cache_config=SimpleNamespace(kv_cache_memory_bytes=None)) + worker.model_runner = SimpleNamespace() + monkeypatch.setattr(diffusion_worker_module, "_all_gather_rank_values", lambda value: [value]) + + with pytest.raises(RuntimeError, match="snapshot was not captured before model loading"): + worker.determine_available_kv_memory(object()) + + +def test_rank_probe_gathers_local_failure_before_raising(monkeypatch) -> None: + gathered = [] + + def gather(local_result): + gathered.append(local_result) + return [local_result, (True, "peer-result")] + + def fail_local_probe(): + raise ValueError("local failure") + + monkeypatch.setattr(diffusion_worker_module, "_all_gather_rank_values", gather) + + with pytest.raises(RuntimeError, match="rank 0: ValueError: local failure"): + diffusion_worker_module._run_and_gather_rank_values( + "test probe", + fail_local_probe, + ) + + assert gathered == [(False, "ValueError: local failure")] diff --git a/tests/diffusion/diffusion_kv/test_worker_contract.py b/tests/diffusion/diffusion_kv/test_worker_contract.py index 0c78abbaf..e0f428fb6 100644 --- a/tests/diffusion/diffusion_kv/test_worker_contract.py +++ b/tests/diffusion/diffusion_kv/test_worker_contract.py @@ -85,19 +85,15 @@ def test_new_request_data_carries_scheduler_allocation_atomically() -> None: assert new_req.diffusion_kv_metadata is metadata -def test_paged_request_without_metadata_reaches_request_forward() -> None: +def test_paged_request_without_metadata_fails_before_request_forward() -> None: runner = make_runner(DiffusionKVCacheMode.PAGED_SCHEDULER) - expected_output = DiffusionOutput(output=None) - runner._execute_request_list = Mock( - return_value=SimpleNamespace( - runner_outputs=[SimpleNamespace(result=expected_output)], - ) - ) + runner._execute_request_list = Mock() req = SimpleNamespace(request_id="req-0") - assert runner.execute_model(req) is expected_output + with pytest.raises(ValueError, match="requires Diffusion KV metadata"): + runner.execute_model(req) - runner._execute_request_list.assert_called_once() + runner._execute_request_list.assert_not_called() def test_dense_request_rejects_metadata_before_request_forward() -> None: @@ -141,18 +137,18 @@ def test_request_rpc_rejects_envelope_request_identity_mismatch() -> None: assert calls == [] -def test_batch_path_allows_missing_metadata_before_forward() -> None: +def test_batch_path_rejects_missing_metadata_before_forward() -> None: runner = make_runner(DiffusionKVCacheMode.PAGED_SCHEDULER) - expected_output = object() - runner._execute_request_list = Mock(return_value=expected_output) + runner._execute_request_list = Mock() new_req = NewRequestData( request_id="req-0", req=SimpleNamespace(request_id="req-0"), ) - assert runner.execute_model_batch(make_scheduler_output(new_req), runner.od_config) is expected_output + with pytest.raises(ValueError, match="requires Diffusion KV metadata"): + runner.execute_model_batch(make_scheduler_output(new_req), runner.od_config) - runner._execute_request_list.assert_called_once() + runner._execute_request_list.assert_not_called() def test_batch_path_rejects_envelope_request_identity_mismatch() -> None: @@ -170,7 +166,7 @@ def test_batch_path_rejects_envelope_request_identity_mismatch() -> None: runner._execute_request_list.assert_not_called() -def test_step_path_allows_missing_metadata_before_step_support_check() -> None: +def test_step_path_rejects_missing_metadata_before_step_support_check() -> None: runner = make_runner(DiffusionKVCacheMode.PAGED_SCHEDULER) runner.pipeline = object() runner._supports_step_mode = Mock(return_value=False) @@ -179,10 +175,10 @@ def test_step_path_allows_missing_metadata_before_step_support_check() -> None: req=SimpleNamespace(request_id="req-0"), ) - with pytest.raises(ValueError, match="does not support step execution"): + with pytest.raises(ValueError, match="requires Diffusion KV metadata"): runner.execute_stepwise(make_scheduler_output(new_req)) - runner._supports_step_mode.assert_called_once() + runner._supports_step_mode.assert_not_called() def test_step_path_rejects_envelope_request_identity_mismatch() -> None: @@ -288,3 +284,42 @@ def test_dense_worker_call_does_not_extend_model_runner_signature() -> None: worker.execute_model(req, SimpleNamespace()) assert calls == [(req, None)] + + +def test_dlo_worker_selects_request_and_metadata_from_same_envelope(monkeypatch) -> None: + calls: list[tuple] = [] + + class ModelRunner: + def execute_model(self, req, **kwargs): + calls.append((req, kwargs)) + return DiffusionOutput(output=None) + + worker = object.__new__(DiffusionWorker) + worker.model_runner = ModelRunner() + worker.lora_manager = None + worker._get_profiler = lambda: None + monkeypatch.setattr( + "vllm_omni.diffusion.distributed.parallel_state.get_data_parallel_rank", + lambda: 1, + ) + req_0 = SimpleNamespace(request_id="req-0") + req_1 = SimpleNamespace(request_id="req-1") + metadata_0 = make_metadata("req-0") + metadata_1 = make_metadata("req-1") + envelopes = [ + NewRequestData(request_id="req-0", req=req_0, diffusion_kv_metadata=metadata_0), + NewRequestData(request_id="req-1", req=req_1, diffusion_kv_metadata=metadata_1), + ] + + result = worker.execute_model(envelopes, SimpleNamespace()) + + assert calls == [ + ( + req_1, + { + "kv_prefetch_job": None, + "diffusion_kv_metadata": metadata_1, + }, + ) + ] + assert result["dp_rank"] == 1 diff --git a/tests/diffusion/models/hunyuan_image3/test_image_kv_cache_manager.py b/tests/diffusion/models/hunyuan_image3/test_image_kv_cache_manager.py index 543b93a25..fed56af2d 100644 --- a/tests/diffusion/models/hunyuan_image3/test_image_kv_cache_manager.py +++ b/tests/diffusion/models/hunyuan_image3/test_image_kv_cache_manager.py @@ -118,6 +118,14 @@ def _call_mgr( ) +def test_cache_manager_registers_attention_without_adding_dense_state() -> None: + mgr = _make_cache_mgr() + + assert isinstance(mgr, nn.Module) + assert dict(mgr.named_modules())["attn"] is mgr.attn + assert mgr.state_dict() == {} + + # ============================================================ # Test 1: No AR KV — basic cache → reuse # ============================================================ diff --git a/tests/diffusion/test_diffusion_config_propagation.py b/tests/diffusion/test_diffusion_config_propagation.py index e25d6497c..e456e302b 100644 --- a/tests/diffusion/test_diffusion_config_propagation.py +++ b/tests/diffusion/test_diffusion_config_propagation.py @@ -130,6 +130,21 @@ class TestCreateDefaultDiffusion: assert od.diffusion_kv_mode is DiffusionKVCacheMode.PAGED_SCHEDULER + def test_diffusion_kv_sizing_fields_roundtrip(self, monkeypatch): + monkeypatch.setattr(OmniDiffusionConfig, "_resolve_master_port", lambda self: 29500) + od = _roundtrip_diffusion_config( + model="x", + kv_cache_memory_bytes=4096, + gpu_memory_utilization=0.75, + max_num_batched_tokens=2048, + max_model_len=4096, + ) + + assert od.kv_cache_memory_bytes == 4096 + assert od.gpu_memory_utilization == 0.75 + assert od.max_num_batched_tokens == 2048 + assert od.max_model_len == 4096 + def test_qwen_image_edit_plus_sets_generic_multimodal_limit(): od_config = OmniDiffusionConfig(model="Qwen/Qwen-Image-Edit-2511", model_class_name="QwenImageEditPlusPipeline") diff --git a/tests/diffusion/test_diffusion_engine_cleanup.py b/tests/diffusion/test_diffusion_engine_cleanup.py index fb8858a48..b626f8ac6 100644 --- a/tests/diffusion/test_diffusion_engine_cleanup.py +++ b/tests/diffusion/test_diffusion_engine_cleanup.py @@ -161,9 +161,83 @@ def test_scheduler_initialization_failure_closes_scheduler_and_executor(monkeypa assert exc_info.value is initialization_error custom_scheduler.initialize.assert_called_once_with(od_config) custom_scheduler.close.assert_called_once_with() + + +def test_init_shuts_down_executor_when_kv_control_plane_initialization_fails( + monkeypatch: pytest.MonkeyPatch, +) -> None: + fake_executor = SimpleNamespace(shutdown=Mock()) + monkeypatch.setattr(DiffusionEngine, "_init_process_hooks", lambda self, config: None) + monkeypatch.setattr( + DiffusionEngine, + "_resolve_execution_mode", + lambda self, config: DiffusionExecutionMode.REQUEST_BATCH, + ) + monkeypatch.setattr( + DiffusionEngine, "_init_executor", lambda self, config: setattr(self, "executor", fake_executor) + ) + monkeypatch.setattr( + "vllm_omni.diffusion.diffusion_engine.initialize_diffusion_kv_control_plane", + Mock(side_effect=RuntimeError("KV initialization failed")), + ) + + with pytest.raises(RuntimeError, match="KV initialization failed"): + DiffusionEngine(SimpleNamespace()) + fake_executor.shutdown.assert_called_once_with() +def test_paged_init_profiles_before_scheduler_initialization(monkeypatch: pytest.MonkeyPatch) -> None: + events = [] + expected_profile_requests = [object()] + fake_executor = SimpleNamespace(shutdown=Mock()) + kv_cache_config = object() + kv_vllm_config = object() + od_config = SimpleNamespace(diffusion_kv_mode="paged_scheduler") + + monkeypatch.setattr(DiffusionEngine, "_init_process_hooks", lambda self, config: None) + monkeypatch.setattr( + DiffusionEngine, + "_resolve_execution_mode", + lambda self, config: DiffusionExecutionMode.REQUEST_BATCH, + ) + monkeypatch.setattr( + DiffusionEngine, + "_init_executor", + lambda self, config: setattr(self, "executor", fake_executor), + ) + monkeypatch.setattr( + DiffusionEngine, + "_prepare_diffusion_kv_profile_requests", + lambda self: events.append("prepare-profile") or expected_profile_requests, + ) + + def initialize(executor, config, *, profile_requests): + events.append("profile-workers") + assert executor is fake_executor + assert config is od_config + assert profile_requests is expected_profile_requests + return kv_cache_config, 16, 16, kv_vllm_config + + monkeypatch.setattr( + "vllm_omni.diffusion.diffusion_engine.initialize_diffusion_kv_control_plane", + initialize, + ) + + def init_scheduler(self, config, scheduler, *args, **kwargs): + events.append("initialize-scheduler") + self.scheduler = SimpleNamespace(close=Mock()) + + monkeypatch.setattr(DiffusionEngine, "_init_scheduler", init_scheduler) + monkeypatch.setattr(DiffusionEngine, "_init_runtime_state", lambda self: None) + monkeypatch.setattr(DiffusionEngine, "_init_execute_fn", lambda self: None) + monkeypatch.setattr(DiffusionEngine, "_log_execution_mode", lambda self, config: None) + + DiffusionEngine(od_config) + + assert events == ["prepare-profile", "profile-workers", "initialize-scheduler"] + + @pytest.mark.asyncio async def test_step_compatibility_wrapper_returns_final_batch() -> None: engine = _make_engine() diff --git a/tests/diffusion/test_diffusion_engine_dummy_run.py b/tests/diffusion/test_diffusion_engine_dummy_run.py index 50dde3083..247fee029 100644 --- a/tests/diffusion/test_diffusion_engine_dummy_run.py +++ b/tests/diffusion/test_diffusion_engine_dummy_run.py @@ -1,9 +1,17 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from types import SimpleNamespace +from unittest.mock import Mock + import pytest from vllm_omni.diffusion import io_support +from vllm_omni.diffusion.diffusion_engine import DiffusionEngine, DiffusionExecutionMode +from vllm_omni.diffusion.diffusion_kv.config import DiffusionKVCacheMode +from vllm_omni.diffusion.diffusion_kv.request import DiffusionKVRequest +from vllm_omni.diffusion.request import OmniDiffusionRequest +from vllm_omni.inputs.data import OmniDiffusionSamplingParams pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion] @@ -55,3 +63,151 @@ def test_dummy_run_num_frames_uses_audio_input_fallback(monkeypatch: pytest.Monk ) assert io_support.get_dummy_run_num_frames("unknown", supports_audio_input=True) == 2 + + +def test_dummy_run_image_count_resolves_hunyuan_architecture_alias() -> None: + assert io_support.get_dummy_run_num_image_inputs("HunyuanImage3ForCausalMM") == 3 + assert io_support.get_dummy_run_num_image_inputs("unknown") == 1 + + +def test_dense_mode_does_not_build_kv_profile_request() -> None: + engine = object.__new__(DiffusionEngine) + engine.od_config = SimpleNamespace(diffusion_kv_mode=DiffusionKVCacheMode.DENSE_LEGACY) + engine._make_dummy_request = Mock(side_effect=AssertionError("dense mode must not profile")) + + assert engine._prepare_diffusion_kv_profile_requests() is None + engine._make_dummy_request.assert_not_called() + + +@pytest.mark.parametrize( + ("execution_mode", "uses_dlo_dp", "max_num_seqs", "expected_profile_requests"), + [ + (DiffusionExecutionMode.STEP_BATCH, False, 2, 2), + (DiffusionExecutionMode.REQUEST_BATCH, True, 4, 1), + (DiffusionExecutionMode.STEP_BATCH, True, 3, 3), + ], +) +def test_paged_kv_profile_requests_match_per_rank_batch( + execution_mode: DiffusionExecutionMode, + uses_dlo_dp: bool, + max_num_seqs: int, + expected_profile_requests: int, +) -> None: + engine = object.__new__(DiffusionEngine) + engine.execution_mode = execution_mode + engine.od_config = SimpleNamespace( + diffusion_kv_mode=DiffusionKVCacheMode.PAGED_SCHEDULER, + model_class_name="HunyuanImage3ForCausalMM", + max_num_seqs=max_num_seqs, + parallel_config=SimpleNamespace(data_parallel_size=max_num_seqs if uses_dlo_dp else 1), + enable_distributed_layerwise_offload=uses_dlo_dp, + dlo_use_allgather=True, + ) + request = OmniDiffusionRequest( + prompt="profile", + request_id="profile-request", + sampling_params=OmniDiffusionSamplingParams(num_inference_steps=1), + ) + engine._make_dummy_request = Mock(return_value=request) + prepared_layout = object() + scheduler_kv_state = ( + DiffusionKVRequest( + "profile-request/diffusion-kv/0", + sequence_id=0, + prefix_len=3, + target_len=8, + seq_len=12, + ), + DiffusionKVRequest( + "profile-request/diffusion-kv/1", + sequence_id=1, + prefix_len=4, + target_len=8, + seq_len=13, + ), + ) + + def preprocess(req: OmniDiffusionRequest) -> OmniDiffusionRequest: + req.prepared_layout = prepared_layout + req.diffusion_kv_requests = scheduler_kv_state + return req + + engine._prepare_request_for_admission = Mock(side_effect=preprocess) + + result = engine._prepare_diffusion_kv_profile_requests() + + assert result is not None + assert len(result) == expected_profile_requests + assert len({profile_request.request_id for profile_request in result}) == expected_profile_requests + assert all(profile_request.is_dummy_run() for profile_request in result) + assert all(profile_request.prepared_layout is prepared_layout for profile_request in result) + assert all(profile_request.diffusion_kv_requests is None for profile_request in result) + assert len({id(profile_request.sampling_params) for profile_request in result}) == expected_profile_requests + assert engine._diffusion_kv_profile_limits == (2, 13, 8) + assert len(result) * engine._diffusion_kv_profile_limits[0] == expected_profile_requests * 2 + engine._make_dummy_request.assert_called_once_with( + height=1024, + width=1024, + guidance_scale=5.0, + num_image_inputs=3, + ) + engine._prepare_request_for_admission.assert_called_once_with(request) + + +@pytest.mark.parametrize( + ("num_sequences", "seq_len", "target_len"), + [ + (3, 16, 8), + (1, 17, 8), + (1, 16, 9), + ], +) +def test_paged_kv_admission_rejects_shape_beyond_profile_envelope( + num_sequences: int, + seq_len: int, + target_len: int, +) -> None: + engine = object.__new__(DiffusionEngine) + engine._diffusion_kv_profile_limits = (2, 16, 8) + engine.pre_process_func = lambda request: request + request = OmniDiffusionRequest( + prompt="too large", + request_id="too-large", + sampling_params=OmniDiffusionSamplingParams(num_inference_steps=1), + diffusion_kv_requests=tuple( + DiffusionKVRequest( + f"too-large/diffusion-kv/{sequence_id}", + sequence_id=sequence_id, + prefix_len=0, + target_len=target_len, + seq_len=seq_len, + ) + for sequence_id in range(num_sequences) + ), + ) + + with pytest.raises(ValueError, match="exceeds the startup memory-profile envelope"): + engine._prepare_request_for_admission(request) + + +def test_paged_kv_admission_accepts_shape_at_profile_envelope() -> None: + engine = object.__new__(DiffusionEngine) + engine._diffusion_kv_profile_limits = (2, 16, 8) + engine.pre_process_func = lambda request: request + request = OmniDiffusionRequest( + prompt="fits", + request_id="fits", + sampling_params=OmniDiffusionSamplingParams(num_inference_steps=1), + diffusion_kv_requests=tuple( + DiffusionKVRequest( + f"fits/diffusion-kv/{sequence_id}", + sequence_id=sequence_id, + prefix_len=8, + target_len=8, + seq_len=16, + ) + for sequence_id in range(2) + ), + ) + + assert engine._prepare_request_for_admission(request) is request diff --git a/tests/diffusion/test_diffusion_engine_rpc_routing.py b/tests/diffusion/test_diffusion_engine_rpc_routing.py index 5aef831aa..119939ae2 100644 --- a/tests/diffusion/test_diffusion_engine_rpc_routing.py +++ b/tests/diffusion/test_diffusion_engine_rpc_routing.py @@ -473,6 +473,38 @@ async def test_cancelled_future_does_not_kill_busy_loop(): _stop_engine(engine) +@pytest.mark.asyncio +async def test_native_kv_reservation_error_wakes_stream_without_killing_busy_loop(): + loop = asyncio.get_running_loop() + engine = _make_engine_with_loop(loop) + + def fail_reservation(*args, **kwargs): + raise RuntimeError("native allocation bug") + + engine.scheduler._diffusion_kv_manager = SimpleNamespace( + has_request=lambda request_id: False, + reserve_request=fail_reservation, + free_request=lambda request_id: None, + ) + request = _make_request("kv-error") + request.diffusion_kv_requests = (object(),) + try: + response_stream = engine.async_add_req_and_stream_response(request) + output_queue = engine._out_streams[request.request_id] + for _ in range(300): + if not output_queue.empty(): + break + await asyncio.sleep(0.01) + assert not output_queue.empty(), "terminal KV allocation error was not delivered to the request stream" + + output = await anext(response_stream) + assert output.error == "native allocation bug" + assert output.finished + assert engine.worker_thread.is_alive() + finally: + _stop_engine(engine) + + # ─────────────────────────── _RpcTask basics ─────────────────────────────── diff --git a/tests/diffusion/test_diffusion_model_runner.py b/tests/diffusion/test_diffusion_model_runner.py index b60cec89c..66da14fd5 100644 --- a/tests/diffusion/test_diffusion_model_runner.py +++ b/tests/diffusion/test_diffusion_model_runner.py @@ -3,6 +3,7 @@ from contextlib import contextmanager from types import SimpleNamespace +from unittest.mock import Mock import pytest import torch @@ -512,6 +513,84 @@ def test_execute_model_accepts_bare_diffusion_output_from_single_request_pipelin assert runner.pipeline.last_req.num_reqs == 1 +@pytest.mark.core_model +@pytest.mark.cpu +def test_profile_run_executes_forward_without_scheduler_kv_validation(monkeypatch): + runner = _make_runner(cache_backend=None, cache_backend_name="none") + request = _make_request() + runner._validate_diffusion_kv_metadata = Mock(side_effect=AssertionError("profile must bypass admission")) + record_names = [] + original_execute = runner._execute_request_list + + def execute(*args, **kwargs): + record_names.append(kwargs["record_name"]) + return original_execute(*args, **kwargs) + + runner._execute_request_list = execute + monkeypatch.setattr(model_runner_module, "set_forward_context", _noop_forward_context) + reset_peak_memory_stats = Mock() + monkeypatch.setattr( + model_runner_module.current_omni_platform, + "reset_peak_memory_stats", + reset_peak_memory_stats, + ) + monkeypatch.setattr(model_runner_module.current_omni_platform, "max_memory_reserved", lambda: 0) + monkeypatch.setattr(model_runner_module.current_omni_platform, "max_memory_allocated", lambda: 0) + synchronize = Mock() + monkeypatch.setattr(model_runner_module.current_omni_platform, "synchronize", synchronize) + + runner.profile_run([request]) + + assert runner.pipeline.forward_calls == 1 + assert record_names == ["pipeline_memory_profile"] + runner._validate_diffusion_kv_metadata.assert_not_called() + reset_peak_memory_stats.assert_not_called() + synchronize.assert_called_once_with() + + +@pytest.mark.core_model +@pytest.mark.cpu +def test_profile_run_executes_maximum_step_batch_without_resetting_peak(monkeypatch): + runner = _make_runner(cache_backend=None, cache_backend_name=None) + runner.pipeline = _FinalOnlyStepPipeline() + runner.od_config.step_execution = True + requests = [_make_request(), _make_request()] + requests[0].request_id = "profile-0" + requests[1].request_id = "profile-1" + for request in requests: + request.sampling_params.num_inference_steps = 1 + + observed_batch_rows = [] + original_denoise_step = runner.pipeline.denoise_step + + def denoise_step(input_batch, states): + observed_batch_rows.append((len(states), int(input_batch.latents.shape[0]))) + return original_denoise_step(input_batch, states) + + runner.pipeline.denoise_step = denoise_step + runner._validate_diffusion_kv_metadata = Mock(side_effect=AssertionError("profile must bypass admission")) + monkeypatch.setattr(model_runner_module, "set_forward_context", _noop_forward_context) + reset_peak_memory_stats = Mock() + monkeypatch.setattr( + model_runner_module.current_omni_platform, + "reset_peak_memory_stats", + reset_peak_memory_stats, + ) + monkeypatch.setattr(model_runner_module.current_omni_platform, "max_memory_reserved", lambda: 0) + monkeypatch.setattr(model_runner_module.current_omni_platform, "max_memory_allocated", lambda: 0) + synchronize = Mock() + monkeypatch.setattr(model_runner_module.current_omni_platform, "synchronize", synchronize) + + runner.profile_run(requests) + + assert observed_batch_rows == [(2, 2)] + assert runner.state_cache == {} + assert runner.input_batch is None + runner._validate_diffusion_kv_metadata.assert_not_called() + reset_peak_memory_stats.assert_not_called() + synchronize.assert_called_once_with() + + class _BatchPipeline: """Pipeline returning a configurable list of outputs from forward().""" diff --git a/tests/diffusion/test_diffusion_request.py b/tests/diffusion/test_diffusion_request.py index 259492c1e..7476ec9c1 100644 --- a/tests/diffusion/test_diffusion_request.py +++ b/tests/diffusion/test_diffusion_request.py @@ -52,6 +52,7 @@ def test_dummy_run_request_is_identified_by_reserved_request_id(): assert req.is_dummy_run() assert OmniDiffusionRequest.is_dummy_run_request_id(DUMMY_DIFFUSION_REQUEST_ID) + assert OmniDiffusionRequest.is_dummy_run_request_id(f"{DUMMY_DIFFUSION_REQUEST_ID}/kv-profile-1") def test_non_dummy_request_is_not_identified_as_dummy_run(): diff --git a/tests/diffusion/test_diffusion_scheduler.py b/tests/diffusion/test_diffusion_scheduler.py index 46ec0cbe6..56fc95912 100644 --- a/tests/diffusion/test_diffusion_scheduler.py +++ b/tests/diffusion/test_diffusion_scheduler.py @@ -8,12 +8,14 @@ from types import SimpleNamespace import pytest import torch +import vllm.v1.core.single_type_kv_cache_manager as native_kv_managers from pytest_mock import MockerFixture +from vllm.v1.kv_cache_interface import FullAttentionSpec, KVCacheConfig, KVCacheGroupSpec, KVCacheTensor from vllm_omni.diffusion.data import DiffusionOutput, DiffusionRequestAbortedError from vllm_omni.diffusion.diffusion_engine import DiffusionEngine, DiffusionExecutionMode from vllm_omni.diffusion.diffusion_kv.config import DiffusionKVCacheMode -from vllm_omni.diffusion.diffusion_kv.request import DiffusionKVContext, DiffusionKVRequest +from vllm_omni.diffusion.diffusion_kv.request import DiffusionKVRequest from vllm_omni.diffusion.request import OmniDiffusionRequest from vllm_omni.diffusion.sched import ( BaseScheduler, @@ -87,6 +89,52 @@ def _cached_ids(sched_output) -> list[str]: return list(sched_output.scheduled_cached_reqs.request_ids) +def _initialize_paged_scheduler( + scheduler: BaseScheduler, + *, + num_blocks: int = 64, + max_num_seqs: int = 1, +) -> None: + native_kv_managers.register_all_kvcache_specs(None) + spec = FullAttentionSpec( + block_size=4, + num_kv_heads=2, + head_size=8, + dtype=torch.bfloat16, + ) + config = KVCacheConfig( + num_blocks=num_blocks, + kv_cache_tensors=[KVCacheTensor(size=spec.page_size_bytes * num_blocks, shared_by=["layer0"])], + kv_cache_groups=[KVCacheGroupSpec(layer_names=["layer0"], kv_cache_spec=spec)], + ) + scheduler.initialize( + SimpleNamespace( + diffusion_kv_mode=DiffusionKVCacheMode.PAGED_SCHEDULER, + max_model_len=64, + max_num_seqs=max_num_seqs, + ), + kv_cache_config=config, + scheduler_block_size=4, + hash_block_size=4, + kv_vllm_config=SimpleNamespace( + model_config=SimpleNamespace(max_model_len=64), + max_in_flight_tokens=64, + ), + ) + + +def _attach_diffusion_kv(request: OmniDiffusionRequest, *, seq_len: int = 8) -> None: + request.diffusion_kv_requests = ( + DiffusionKVRequest( + f"{request.request_id}/diffusion-kv/0", + sequence_id=0, + prefix_len=4, + target_len=4, + seq_len=seq_len, + ), + ) + + class _StubScheduler: def __init__(self, request: OmniDiffusionRequest, output) -> None: self._request = request @@ -400,18 +448,16 @@ class TestRequestScheduler: assert self.scheduler.has_requests() is False def test_diffusion_kv_request_moves_to_scheduler_state(self) -> None: - self.scheduler.initialize(SimpleNamespace(diffusion_kv_mode=DiffusionKVCacheMode.PAGED_SCHEDULER)) + _initialize_paged_scheduler(self.scheduler) request = _make_request("diffusion-kv") prepared_layout = object() request.prepared_layout = prepared_layout - text_context = DiffusionKVContext(context_id="text", cache_role="cross.text", num_tokens=8) kv_request = DiffusionKVRequest( "diffusion-kv/diffusion-kv/0", sequence_id=0, prefix_len=4, target_len=8, seq_len=16, - kv_contexts=(text_context,), ) request.diffusion_kv_requests = (kv_request,) @@ -421,12 +467,179 @@ class TestRequestScheduler: assert state is not None assert state.diffusion_kv_requests == (kv_request,) - assert state.diffusion_kv_requests[0].kv_contexts == (text_context,) assert request.diffusion_kv_requests is None assert scheduler_output.scheduled_new_reqs[0].req.prepared_layout is prepared_layout + metadata = scheduler_output.scheduled_new_reqs[0].diffusion_kv_metadata + assert metadata is not None + assert metadata.request_id == request.request_id + assert len(metadata.sequences[0].block_ids[0]) == 4 + + def test_diffusion_kv_capacity_backpressures_fifo_until_blocks_are_freed(self) -> None: + _initialize_paged_scheduler(self.scheduler, num_blocks=3, max_num_seqs=2) + first_request = _make_request("first") + second_request = _make_request("second") + _attach_diffusion_kv(first_request) + _attach_diffusion_kv(second_request) + self.scheduler.add_request(first_request) + self.scheduler.add_request(second_request) + + first_output = self.scheduler.schedule() + + assert _new_ids(first_output) == ["first"] + assert first_output.num_waiting_reqs == 1 + + self.scheduler.update_from_output(first_output, _make_request_output("first")) + second_output = self.scheduler.schedule() + + assert _new_ids(second_output) == ["second"] + assert second_output.num_waiting_reqs == 0 + # The same output carries both the release and the newly admitted block + # table; the future Worker data plane must process finished ids first. + assert second_output.finished_req_ids == {"first"} + assert second_output.scheduled_new_reqs[0].diffusion_kv_metadata is not None + + def test_impossible_diffusion_kv_capacity_finishes_only_that_request(self) -> None: + _initialize_paged_scheduler(self.scheduler, num_blocks=2, max_num_seqs=2) + impossible = _make_request("impossible") + schedulable = _make_request("schedulable") + _attach_diffusion_kv(impossible, seq_len=8) + schedulable.diffusion_kv_requests = ( + DiffusionKVRequest( + "schedulable/diffusion-kv/0", + sequence_id=0, + prefix_len=0, + target_len=4, + seq_len=4, + ), + ) + self.scheduler.add_request(impossible) + self.scheduler.add_request(schedulable) + + sched_output = self.scheduler.schedule() + + failed_state = self.scheduler.get_request_state("impossible") + assert failed_state is not None + assert failed_state.status == DiffusionRequestStatus.FINISHED_ERROR + assert failed_state.error is not None + assert "cannot fit even when the block pool is empty" in failed_state.error + assert sched_output.finished_req_ids == {"impossible"} + assert _new_ids(sched_output) == ["schedulable"] + + finished = self.scheduler.update_from_output( + sched_output, + _make_request_output("schedulable"), + ) + + assert finished == {"impossible", "schedulable"} + + def test_impossible_diffusion_kv_capacity_does_not_block_waiters_under_load(self) -> None: + _initialize_paged_scheduler(self.scheduler, num_blocks=4, max_num_seqs=2) + running = _make_request("running") + _attach_diffusion_kv(running) + self.scheduler.add_request(running) + assert _new_ids(self.scheduler.schedule()) == ["running"] + + impossible = _make_request("impossible") + impossible.diffusion_kv_requests = tuple( + DiffusionKVRequest( + f"impossible/diffusion-kv/{sequence_id}", + sequence_id=sequence_id, + prefix_len=0, + target_len=4, + seq_len=4, + ) + for sequence_id in range(4) + ) + schedulable = _make_request("schedulable") + schedulable.diffusion_kv_requests = ( + DiffusionKVRequest( + "schedulable/diffusion-kv/0", + sequence_id=0, + prefix_len=0, + target_len=4, + seq_len=4, + ), + ) + self.scheduler.add_request(impossible) + self.scheduler.add_request(schedulable) + + sched_output = self.scheduler.schedule() + + failed_state = self.scheduler.get_request_state("impossible") + assert failed_state is not None + assert failed_state.status == DiffusionRequestStatus.FINISHED_ERROR + assert failed_state.error is not None + assert "required_blocks=4, available_blocks=3" in failed_state.error + assert sched_output.finished_req_ids == {"impossible"} + assert _new_ids(sched_output) == ["schedulable"] + + def test_diffusion_kv_internal_allocation_error_finishes_request(self, monkeypatch) -> None: + _initialize_paged_scheduler(self.scheduler) + request = _make_request("native-error") + _attach_diffusion_kv(request) + self.scheduler.add_request(request) + manager = self.scheduler._diffusion_kv_manager + assert manager is not None + + def raise_native_error(*args, **kwargs): + raise ValueError("native allocation bug") + + monkeypatch.setattr( + manager, + "reserve_request", + raise_native_error, + ) + + sched_output = self.scheduler.schedule() + + state = self.scheduler.get_request_state("native-error") + assert state is not None + assert state.status == DiffusionRequestStatus.FINISHED_ERROR + assert state.error == "native allocation bug" + assert sched_output.finished_req_ids == {"native-error"} + assert sched_output.scheduled_request_ids == [] + + def test_diffusion_kv_preemption_retains_allocation(self) -> None: + _initialize_paged_scheduler(self.scheduler, num_blocks=3) + request = _make_request("preempted") + _attach_diffusion_kv(request) + self.scheduler.add_request(request) + self.scheduler.schedule() + manager = self.scheduler._diffusion_kv_manager + assert manager is not None + metadata = manager.get_metadata("preempted") + + assert self.scheduler.preempt_request("preempted") is True + resumed_output = self.scheduler.schedule() + + assert resumed_output.scheduled_new_reqs == [] + assert _cached_ids(resumed_output) == ["preempted"] + assert manager.get_metadata("preempted") == metadata + + @pytest.mark.parametrize( + "status", + [ + DiffusionRequestStatus.FINISHED_COMPLETED, + DiffusionRequestStatus.FINISHED_ABORTED, + DiffusionRequestStatus.FINISHED_ERROR, + ], + ) + def test_diffusion_kv_terminal_status_releases_allocation(self, status: DiffusionRequestStatus) -> None: + _initialize_paged_scheduler(self.scheduler, num_blocks=3) + request = _make_request("terminal") + _attach_diffusion_kv(request) + self.scheduler.add_request(request) + self.scheduler.schedule() + manager = self.scheduler._diffusion_kv_manager + assert manager is not None + assert manager.has_request("terminal") is True + + self.scheduler.finish_requests("terminal", status) + + assert manager.has_request("terminal") is False def test_paged_scheduler_rejects_missing_kv_request(self) -> None: - self.scheduler.initialize(SimpleNamespace(diffusion_kv_mode=DiffusionKVCacheMode.PAGED_SCHEDULER)) + _initialize_paged_scheduler(self.scheduler) request = _make_request("diffusion-kv") with pytest.raises(ValueError, match="did not produce DiffusionKVRequest"): @@ -443,7 +656,7 @@ class TestRequestScheduler: ], ) def test_diffusion_kv_request_rejects_legacy_dense_kv(self, owner_name: str, field_name: str) -> None: - self.scheduler.initialize(SimpleNamespace(diffusion_kv_mode=DiffusionKVCacheMode.PAGED_SCHEDULER)) + _initialize_paged_scheduler(self.scheduler) request = _make_request("diffusion-kv") owner = request if owner_name == "request" else request.sampling_params setattr(owner, field_name, object()) @@ -921,6 +1134,45 @@ class TestDiffusionEngine: assert engine.scheduler.max_num_running_reqs == 1 fake_executor_cls.assert_called_once_with(od_config) + def test_initializes_paged_scheduler_from_native_control_plane( + self, + monkeypatch: pytest.MonkeyPatch, + mocker: MockerFixture, + ) -> None: + od_config = SimpleNamespace(model_class_name="mock_model", streaming_output=False) + fake_executor_cls = mocker.Mock(return_value=mocker.Mock()) + scheduler = mocker.Mock() + kv_cache_config = object() + kv_vllm_config = object() + + monkeypatch.setattr( + "vllm_omni.diffusion.diffusion_engine.get_diffusion_post_process_func", + lambda *args, **kwargs: None, + ) + monkeypatch.setattr( + "vllm_omni.diffusion.diffusion_engine.get_diffusion_pre_process_func", + lambda *args, **kwargs: None, + ) + monkeypatch.setattr( + "vllm_omni.diffusion.diffusion_engine.DiffusionExecutor.get_class", + lambda *args, **kwargs: fake_executor_cls, + ) + monkeypatch.setattr( + "vllm_omni.diffusion.diffusion_engine.initialize_diffusion_kv_control_plane", + lambda *args, **kwargs: (kv_cache_config, 16, 16, kv_vllm_config), + ) + + engine = DiffusionEngine(od_config, scheduler=scheduler) + + assert engine.scheduler is scheduler + scheduler.initialize.assert_called_once_with( + od_config, + kv_cache_config=kv_cache_config, + scheduler_block_size=16, + hash_block_size=16, + kv_vllm_config=kv_vllm_config, + ) + def test_scheduler_alias_keeps_default_request_scheduler(self) -> None: scheduler = Scheduler() scheduler.initialize(SimpleNamespace()) @@ -987,6 +1239,21 @@ class TestDiffusionEngine: assert output.aborted is True assert output.abort_message == "Request req-finalize aborted." + def test_finalize_finished_request_returns_scheduler_admission_error(self) -> None: + engine = DiffusionEngine.__new__(DiffusionEngine) + engine.scheduler = StepScheduler() + engine.scheduler.initialize(SimpleNamespace()) + + req_id = engine.scheduler.add_request(_make_request("req-error")) + engine.scheduler._finish_requests( + {req_id: DiffusionRequestStatus.FINISHED_ERROR}, + {req_id: "KV request cannot fit"}, + ) + + output = engine._finalize_finished_request(req_id) + + assert output.error == "KV request cannot fit" + @pytest.mark.asyncio async def test_streaming_runner_output_notifies_each_chunk(self) -> None: engine = DiffusionEngine.__new__(DiffusionEngine) diff --git a/tests/diffusion/test_diffusion_worker.py b/tests/diffusion/test_diffusion_worker.py index 6a49d0909..0d8d5e92a 100644 --- a/tests/diffusion/test_diffusion_worker.py +++ b/tests/diffusion/test_diffusion_worker.py @@ -13,11 +13,7 @@ import pytest import torch from pytest_mock import MockerFixture -from vllm_omni.diffusion.worker.diffusion_worker import ( - DiffusionWorker, - _create_diffusion_worker_vllm_config, - _make_diffusion_vllm_model_config, -) +from vllm_omni.diffusion.worker.diffusion_worker import DiffusionWorker pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.gpu] @@ -58,34 +54,6 @@ def mock_gpu_worker(mocker: MockerFixture, mock_od_config): return worker -def test_diffusion_vllm_model_config_supplies_dtype_for_quant_methods(): - from types import SimpleNamespace - - from vllm_omni.quantization import build_quant_config - - od_config = SimpleNamespace( - model="dummy", - dtype=torch.bfloat16, - quantization_config=build_quant_config( - { - "quant_method": "modelopt", - "quant_algo": "FP8", - "ignore": [], - } - ), - tf_model_config=SimpleNamespace(), - enforce_eager=True, - is_moe=False, - ) - - model_config = _make_diffusion_vllm_model_config(od_config) - - assert model_config.dtype is torch.bfloat16 - assert model_config.quantization == "modelopt" - assert model_config.quantization_config is od_config.quantization_config - assert model_config.is_quantized() - - class TestDiffusionWorkerSleep: """Test DiffusionWorker.sleep method.""" @@ -329,49 +297,3 @@ class TestDiffusionWorkerWakeUp: mock_buffer2.data.copy_.assert_not_called() assert result is True - - -class TestWorkerVllmConfigAdditionalConfig: - """Test worker-side VllmConfig construction with additional_config.""" - - def test_create_diffusion_worker_vllm_config_passes_additional_config(self, mocker: MockerFixture): - mock_vllm_config = mocker.Mock() - mock_vllm_cls = mocker.patch( - "vllm_omni.diffusion.worker.diffusion_worker.VllmConfig", - return_value=mock_vllm_config, - ) - od_config = mocker.Mock(additional_config={"torchair_graph_config": {"enabled": True}}) - - result = _create_diffusion_worker_vllm_config(torch.device("cpu"), od_config) - - assert result is mock_vllm_config - assert mock_vllm_cls.call_args.kwargs["additional_config"] == od_config.additional_config - - def test_create_diffusion_worker_vllm_config_falls_back_when_constructor_rejects_additional_config( - self, mocker: MockerFixture - ): - mock_vllm_config = mocker.Mock() - mock_vllm_cls = mocker.patch( - "vllm_omni.diffusion.worker.diffusion_worker.VllmConfig", - side_effect=[ - TypeError("VllmConfig.__init__() got an unexpected keyword argument 'additional_config'"), - mock_vllm_config, - ], - ) - od_config = mocker.Mock(additional_config={"ascend_scheduler_config": {"foo": "bar"}}) - - result = _create_diffusion_worker_vllm_config(torch.device("cpu"), od_config) - - assert result is mock_vllm_config - assert mock_vllm_cls.call_count == 2 - assert getattr(mock_vllm_config, "additional_config") == od_config.additional_config - - def test_create_diffusion_worker_vllm_config_reraises_other_type_errors(self, mocker: MockerFixture): - mocker.patch( - "vllm_omni.diffusion.worker.diffusion_worker.VllmConfig", - side_effect=TypeError("additional_config values must be hashable"), - ) - od_config = mocker.Mock(additional_config={"ascend_scheduler_config": {"foo": "bar"}}) - - with pytest.raises(TypeError, match="hashable"): - _create_diffusion_worker_vllm_config(torch.device("cpu"), od_config) diff --git a/tests/diffusion/test_multiproc_engine_concurrency.py b/tests/diffusion/test_multiproc_engine_concurrency.py index aa40d0770..d5c36a093 100644 --- a/tests/diffusion/test_multiproc_engine_concurrency.py +++ b/tests/diffusion/test_multiproc_engine_concurrency.py @@ -413,6 +413,26 @@ class TestRequestModeDispatch: assert [output.result.error for output in result.runner_outputs] == ["A", "B"] executor.collective_rpc.assert_called_once() + def test_dlo_dp_forwards_request_metadata_envelopes_as_one_wave(self): + executor, _, _ = _make_executor(num_gpus=2) + executor.od_config = SimpleNamespace( + step_execution=False, + parallel_config=SimpleNamespace(data_parallel_size=2), + enable_distributed_layerwise_offload=True, + dlo_use_allgather=True, + ) + executor.collective_rpc = Mock(return_value=[_tagged_output("A"), _tagged_output("B")]) + scheduler_output = _make_sched_output("A", "B") + for new_req in scheduler_output.scheduled_new_reqs: + new_req.diffusion_kv_metadata = SimpleNamespace(request_id=new_req.request_id) + + result = executor.execute_request(scheduler_output) + + assert [output.result.error for output in result.runner_outputs] == ["A", "B"] + forwarded_envelopes = executor.collective_rpc.call_args.kwargs["args"][0] + assert forwarded_envelopes is scheduler_output.scheduled_new_reqs + assert [envelope.diffusion_kv_metadata.request_id for envelope in forwarded_envelopes] == ["A", "B"] + @pytest.mark.parametrize( ("field", "value"), [ diff --git a/vllm_omni/config/omni_config.py b/vllm_omni/config/omni_config.py index 8c4a933e2..4a8a54eca 100644 --- a/vllm_omni/config/omni_config.py +++ b/vllm_omni/config/omni_config.py @@ -814,6 +814,10 @@ _DIFFUSION_MOVED_SHARED_FIELDS = frozenset( "omni_kv_config", "cfg_kv_collect_func", "max_num_seqs", + "kv_cache_memory_bytes", + "gpu_memory_utilization", + "max_num_batched_tokens", + "max_model_len", "enable_sleep_mode", "enforce_eager", "enable_multithread_weight_load", diff --git a/vllm_omni/diffusion/attention/backends/abstract.py b/vllm_omni/diffusion/attention/backends/abstract.py index f3a5f99ef..4a34aa765 100644 --- a/vllm_omni/diffusion/attention/backends/abstract.py +++ b/vllm_omni/diffusion/attention/backends/abstract.py @@ -81,6 +81,19 @@ class AttentionBackend(ABC): supported_head_sizes = cls.get_supported_head_sizes() return (not supported_head_sizes) or head_size in supported_head_sizes + @classmethod + def indexes_kv_by_block_stride(cls) -> bool: + """Whether this backend reads K/V pages by the runtime block stride. + + Returning ``True`` means the physical cache layout has ``num_blocks`` + as its outer stride, so native vLLM may safely use page-size padding + when it unifies cache layouts across layers. Dense diffusion backends + conservatively keep the default ``False``; a paged backend should + override this only when its kernel actually follows that layout. + """ + + return False + @dataclass(frozen=True, slots=True) class QueryRange: diff --git a/vllm_omni/diffusion/attention/layer.py b/vllm_omni/diffusion/attention/layer.py index 7eb74a941..ad9fe2b36 100644 --- a/vllm_omni/diffusion/attention/layer.py +++ b/vllm_omni/diffusion/attention/layer.py @@ -11,8 +11,10 @@ from dataclasses import replace import torch import torch.nn as nn +from vllm.config import VllmConfig, set_current_vllm_config from vllm.logger import init_logger from vllm.model_executor.models.utils import extract_layer_index +from vllm.v1.kv_cache_interface import FullAttentionSpec, KVCacheSpec from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata from vllm_omni.diffusion.attention.backends.sdpa import SDPABackend @@ -61,12 +63,23 @@ class Attention(nn.Module): # perf for this layer (e.g. Wan2.2 cross-attn has short sequences and # block-FP8 quant offers no win). Default False = follow global config. disable_kv_quant: bool = False, + # Opt-in marker for Scheduler-managed paged KV. Unmarked diffusion + # attention remains dense and contributes no native KVCacheSpec. + paged_kv_cache_role: str | None = None, + paged_kv_cache_dtype: torch.dtype | None = None, ): super().__init__() self.role = role self.role_category = role_category self.qkv_layout = qkv_layout + if paged_kv_cache_role == "": + raise ValueError("paged_kv_cache_role must be non-empty when provided") + self.paged_kv_cache_role = paged_kv_cache_role + self.paged_kv_cache_dtype = paged_kv_cache_dtype + self.num_heads = num_heads + self.num_kv_heads = num_kv_heads if num_kv_heads is not None else num_heads + self.head_size = head_size # Resolve backend via role-aware config. # The global diffusion config is set during model init via @@ -171,6 +184,25 @@ class Attention(nn.Module): self._disable_kv_quant: bool = disable_kv_quant self._init_kv_cache_quantization(config) + def get_kv_cache_spec(self, vllm_config: VllmConfig) -> KVCacheSpec | None: + """Return native rank-local geometry for an opted-in paged cache.""" + + if self.paged_kv_cache_role is None: + return None + dtype = self.paged_kv_cache_dtype or vllm_config.model_config.dtype + # Keep backend layout discovery under the same config context used by + # upstream vLLM's attention-spec collector. + with set_current_vllm_config(vllm_config): + indexes_kv_by_block_stride = self.attn_backend.indexes_kv_by_block_stride() + return FullAttentionSpec( + block_size=vllm_config.cache_config.block_size, + num_kv_heads=self.num_kv_heads, + head_size=self.head_size, + dtype=dtype, + indexes_kv_by_block_stride=indexes_kv_by_block_stride, + non_causal=not self.causal, + ) + def _get_active_parallel_strategy(self): """Get the parallel strategy based on current SP active state. diff --git a/vllm_omni/diffusion/data.py b/vllm_omni/diffusion/data.py index 0ea85405b..34c8ccb96 100644 --- a/vllm_omni/diffusion/data.py +++ b/vllm_omni/diffusion/data.py @@ -851,6 +851,14 @@ class OmniDiffusionConfig: # Maximum number of sequences to generate in a batch max_num_seqs: int = 1 + # Native vLLM KV-cache sizing inputs used only by paged_scheduler mode. + # These mirror the structured stage cache/scheduler config until all + # diffusion runtime construction consumes VllmOmniConfig directly. + kv_cache_memory_bytes: int | None = None + gpu_memory_utilization: float = 0.9 + max_num_batched_tokens: int | None = None + max_model_len: int | None = None + # Request-mode batch admission: wait briefly for compatible requests to # accumulate in the scheduler waiting queue before the first schedule() of # a wave. Improves fused forward batch sizes under bursty HTTP ingress. @@ -936,6 +944,14 @@ class OmniDiffusionConfig: if not isinstance(self.diffusion_compile_dynamic, bool): raise TypeError(f"diffusion_compile_dynamic must be a bool, got {type(self.diffusion_compile_dynamic)!r}") self.diffusion_kv_mode = parse_diffusion_kv_cache_mode(self.diffusion_kv_mode) + if self.kv_cache_memory_bytes is not None and self.kv_cache_memory_bytes < 0: + raise ValueError("kv_cache_memory_bytes must be non-negative") + if not 0.0 < self.gpu_memory_utilization <= 1.0: + raise ValueError("gpu_memory_utilization must be in (0, 1]") + if self.max_num_batched_tokens is not None and self.max_num_batched_tokens <= 0: + raise ValueError("max_num_batched_tokens must be positive") + if self.max_model_len is not None and self.max_model_len != -1 and self.max_model_len <= 0: + raise ValueError("max_model_len must be positive or -1") if self.omni_kv_config is None: self.omni_kv_config = {} diff --git a/vllm_omni/diffusion/diffusion_engine.py b/vllm_omni/diffusion/diffusion_engine.py index fe9ed5c8c..4812c8925 100644 --- a/vllm_omni/diffusion/diffusion_engine.py +++ b/vllm_omni/diffusion/diffusion_engine.py @@ -5,6 +5,7 @@ from __future__ import annotations import asyncio import concurrent.futures +import copy import inspect import queue import threading @@ -17,18 +18,23 @@ from typing import TYPE_CHECKING, Any import numpy as np import PIL.Image import torch +from vllm.config import VllmConfig from vllm.logger import init_logger from vllm.utils.import_utils import resolve_obj_by_qualname from vllm.v1.engine.exceptions import EngineDeadError +from vllm.v1.kv_cache_interface import KVCacheConfig from vllm_omni.diffusion.data import ( DiffusionOutput, DiffusionRequestAbortedError, OmniDiffusionConfig, ) +from vllm_omni.diffusion.diffusion_kv.config import DiffusionKVCacheMode +from vllm_omni.diffusion.diffusion_kv.initialization import initialize_diffusion_kv_control_plane from vllm_omni.diffusion.executor.abstract import DiffusionExecutor from vllm_omni.diffusion.io_support import ( get_dummy_run_num_frames, + get_dummy_run_num_image_inputs, image_color_format, supports_audio_output, supports_multimodal_input, @@ -184,12 +190,39 @@ class DiffusionEngine: from the resolved execution mode. """ self.od_config = od_config + # Set after the paged-KV profile request has gone through model-owned + # preprocessing. Real requests are admitted only within this measured + # activation envelope: (max execution sequences, max seq_len, + # max target_len). + self._diffusion_kv_profile_limits: tuple[int, int, int] | None = None self._init_process_hooks(od_config) self.execution_mode = self._resolve_execution_mode(od_config) self._init_executor(od_config) try: - self._init_scheduler(od_config, scheduler) + profile_requests = self._prepare_diffusion_kv_profile_requests() + kv_control_plane = initialize_diffusion_kv_control_plane( + self.executor, + od_config, + profile_requests=profile_requests, + ) + if kv_control_plane is None: + self._init_scheduler(od_config, scheduler) + else: + ( + kv_cache_config, + scheduler_block_size, + hash_block_size, + kv_vllm_config, + ) = kv_control_plane + self._init_scheduler( + od_config, + scheduler, + kv_cache_config, + scheduler_block_size=scheduler_block_size, + hash_block_size=hash_block_size, + kv_vllm_config=kv_vllm_config, + ) self._init_runtime_state() self._init_execute_fn() self._log_execution_mode(od_config) @@ -242,6 +275,11 @@ class DiffusionEngine: self, od_config: OmniDiffusionConfig, scheduler: BaseScheduler | None = None, + kv_cache_config: KVCacheConfig | None = None, + *, + scheduler_block_size: int | None = None, + hash_block_size: int | None = None, + kv_vllm_config: VllmConfig | None = None, ) -> None: if scheduler is not None: self.scheduler = scheduler @@ -249,7 +287,16 @@ class DiffusionEngine: self.scheduler = StepScheduler() else: self.scheduler = RequestScheduler() - self.scheduler.initialize(od_config) + if kv_cache_config is None: + self.scheduler.initialize(od_config) + else: + self.scheduler.initialize( + od_config, + kv_cache_config=kv_cache_config, + scheduler_block_size=scheduler_block_size, + hash_block_size=hash_block_size, + kv_vllm_config=kv_vllm_config, + ) def _init_runtime_state(self) -> None: # DP multi-concurrency: allow batching dp_size requests so each @@ -703,9 +750,36 @@ class DiffusionEngine: """Run model-owned preprocessing once, before entering Engine locks.""" pre_process_func = getattr(self, "pre_process_func", None) - if pre_process_func is None: - return request - return pre_process_func(request) + if pre_process_func is not None: + request = pre_process_func(request) + self._validate_diffusion_kv_profile_limits(request) + return request + + def _validate_diffusion_kv_profile_limits(self, request: OmniDiffusionRequest) -> None: + """Keep admitted paged-KV requests within the profiled activation shape.""" + + profile_limits = getattr(self, "_diffusion_kv_profile_limits", None) + if profile_limits is None: + return + kv_requests = request.diffusion_kv_requests + if not kv_requests: + return + + request_limits = ( + len(kv_requests), + max(kv_request.seq_len for kv_request in kv_requests), + max(kv_request.target_len for kv_request in kv_requests), + ) + if any(actual > profiled for actual, profiled in zip(request_limits, profile_limits, strict=True)): + request_sequences, request_seq_len, request_target_len = request_limits + profile_sequences, profile_seq_len, profile_target_len = profile_limits + raise ValueError( + f"Diffusion KV request {request.request_id!r} exceeds the startup memory-profile envelope: " + f"sequences={request_sequences} (profiled={profile_sequences}), " + f"max_seq_len={request_seq_len} (profiled={profile_seq_len}), " + f"max_target_len={request_target_len} (profiled={profile_target_len}). " + "Reduce the request shape or extend the model's paged-KV profile recipe." + ) def _add_prepared_request(self, request: OmniDiffusionRequest) -> str: """Admit a request whose model-owned preprocessing is complete.""" @@ -869,47 +943,111 @@ class DiffusionEngine: self.close() raise e - def _dummy_run(self): - """A dummy run to warm up the model.""" - num_inference_steps = 1 - height = 512 - width = 512 - prompt: OmniTextPrompt = {"prompt": "dummy run"} + def _make_dummy_request( + self, + *, + height: int, + width: int, + guidance_scale: float, + num_image_inputs: int = 1, + ) -> OmniDiffusionRequest | None: + """Build a one-step model request for startup profiling or warmup.""" + prompt: OmniTextPrompt = {"prompt": "dummy run"} supports_image_input, supports_audio_input = supports_multimodal_input(self.od_config) if supports_image_input: - # Provide a dummy image input if the model supports it color_format = image_color_format(self.od_config.model_class_name) - dummy_image = PIL.Image.new(color_format, (width, height)) - prompt.setdefault("multi_modal_data", {})["image"] = dummy_image + images = [PIL.Image.new(color_format, (width, height)) for _ in range(num_image_inputs)] + prompt.setdefault("multi_modal_data", {})["image"] = images[0] if len(images) == 1 else images if supports_audio_input: audio_sr = 16000 - dummy_audio = np.random.randn(audio_sr * 2).astype(np.float32) - prompt.setdefault("multi_modal_data", {})["audio"] = dummy_audio + prompt.setdefault("multi_modal_data", {})["audio"] = np.random.randn(audio_sr * 2).astype(np.float32) num_frames = get_dummy_run_num_frames(self.od_config.model_class_name, supports_audio_input) if num_frames <= 0: - logger.info("Skipping dummy warmup run (num_frames=0)") - return - req = OmniDiffusionRequest( + return None + return OmniDiffusionRequest( prompt=prompt, request_id=DUMMY_DIFFUSION_REQUEST_ID, sampling_params=OmniDiffusionSamplingParams( height=height, width=width, - num_inference_steps=num_inference_steps, + num_inference_steps=1, num_frames=num_frames, - # Keep warmup path minimal and robust across text encoders. - # Some models may fail when warmup implicitly triggers - # classifier-free guidance with an empty negative prompt. - guidance_scale=0.0, + guidance_scale=guidance_scale, num_outputs_per_prompt=1, - # Disable CFG for warmup to avoid triggering CFG parallel - # validation when cfg_parallel_size > 1. extra_args={"cfg_text_scale": 1.0, "cfg_img_scale": 1.0}, ), ) + + def _prepare_diffusion_kv_profile_requests(self) -> list[OmniDiffusionRequest] | None: + """Prepare the per-rank request batch used to profile paged-KV headroom. + + The profile executes directly on each Worker before the Scheduler and + its KV manager exist. It uses the maximum number of requests that one + rank can execute together. DLO+DP is the exception because each rank + executes one request from the collective wave. + + Hunyuan is currently the only model integrated with + ``paged_scheduler``; 1024x1024, enabled CFG, and the maximum advertised + reference-image count exercise its first-step activation peak. + Admission compares each preprocessed request's CFG count and tokenized + sequence/target shape with the resulting per-request profile envelope. + Future paged model integrations must extend this recipe for their + serving limits (for example, video frame count) rather than reusing it + silently. + """ + + if ( + getattr(self.od_config, "diffusion_kv_mode", DiffusionKVCacheMode.DENSE_LEGACY) + is not DiffusionKVCacheMode.PAGED_SCHEDULER + ): + return None + request = self._make_dummy_request( + height=1024, + width=1024, + guidance_scale=5.0, + num_image_inputs=get_dummy_run_num_image_inputs(self.od_config.model_class_name), + ) + if request is None: + raise RuntimeError("paged_scheduler requires a runnable Diffusion KV memory profile request") + request = self._prepare_request_for_admission(request) + kv_requests = request.diffusion_kv_requests + if not kv_requests: + raise RuntimeError("paged_scheduler profile preprocessing must produce Diffusion KV requests") + self._diffusion_kv_profile_limits = ( + len(kv_requests), + max(kv_request.seq_len for kv_request in kv_requests), + max(kv_request.target_len for kv_request in kv_requests), + ) + dlo_dp_request_mode = self.execution_mode is DiffusionExecutionMode.REQUEST_BATCH and _uses_dlo_dp_concurrency( + self.od_config + ) + profile_batch_size = 1 if dlo_dp_request_mode else _max_num_seqs(self.od_config) + profile_requests: list[OmniDiffusionRequest] = [] + for index in range(profile_batch_size): + profile_request = copy.copy(request) + profile_request.request_id = f"{DUMMY_DIFFUSION_REQUEST_ID}/kv-profile-{index}" + profile_request.sampling_params = copy.deepcopy(request.sampling_params) + # Preprocessing owns request geometry and prepared model inputs, + # but native KV requests remain Scheduler-only mutable state. The + # profile bypasses Scheduler admission and must not send them to + # the Worker. + profile_request.diffusion_kv_requests = None + profile_requests.append(profile_request) + return profile_requests + + def _dummy_run(self): + """A dummy run to warm up the model.""" + req = self._make_dummy_request( + height=512, + width=512, + guidance_scale=0.0, + ) + if req is None: + logger.info("Skipping dummy warmup run (num_frames=0)") + return logger.info("dummy run to warm up the model") output = self.add_req_and_wait_for_response(req) if output.error: @@ -1127,4 +1265,7 @@ class DiffusionEngine: if runner_output is not None and runner_output.async_output_id is not None: return DiffusionOutput(async_output_id=runner_output.async_output_id) + if state.status == DiffusionRequestStatus.FINISHED_ERROR and state.error: + return DiffusionOutput(error=state.error) + return DiffusionOutput(error=missing_result_error) diff --git a/vllm_omni/diffusion/diffusion_kv/initialization.py b/vllm_omni/diffusion/diffusion_kv/initialization.py new file mode 100644 index 000000000..854f3aaa5 --- /dev/null +++ b/vllm_omni/diffusion/diffusion_kv/initialization.py @@ -0,0 +1,82 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from __future__ import annotations + +from typing import TYPE_CHECKING + +import torch +from vllm.config import VllmConfig +from vllm.v1.core.kv_cache_utils import ( + generate_scheduler_kv_cache_config, + get_kv_cache_configs, + resolve_kv_cache_block_sizes, +) +from vllm.v1.kv_cache_interface import KVCacheConfig, KVCacheSpec + +from vllm_omni.diffusion.data import OmniDiffusionConfig +from vllm_omni.diffusion.diffusion_kv.config import DiffusionKVCacheMode +from vllm_omni.diffusion.vllm_config import create_diffusion_vllm_config +from vllm_omni.platforms import current_omni_platform + +if TYPE_CHECKING: + from vllm_omni.diffusion.executor.abstract import DiffusionExecutor + from vllm_omni.diffusion.request import OmniDiffusionRequest + + +def build_native_kv_cache_configs( + vllm_config: VllmConfig, + worker_specs: list[dict[str, KVCacheSpec]], + available_memory: list[int], +) -> tuple[list[KVCacheConfig], KVCacheConfig]: + """Build rank-local and Scheduler-native configs using vLLM utilities.""" + + from vllm.v1.core import single_type_kv_cache_manager + + single_type_kv_cache_manager.register_all_kvcache_specs(vllm_config) + worker_configs = get_kv_cache_configs(vllm_config, worker_specs, available_memory) + return worker_configs, generate_scheduler_kv_cache_config(worker_configs) + + +def initialize_diffusion_kv_control_plane( + executor: DiffusionExecutor, + od_config: OmniDiffusionConfig, + *, + profile_requests: list[OmniDiffusionRequest] | None = None, + device: torch.device | None = None, +) -> tuple[KVCacheConfig, int, int, VllmConfig] | None: + """Run the native Worker-spec to Scheduler-config initialization chain.""" + + if ( + getattr(od_config, "diffusion_kv_mode", DiffusionKVCacheMode.DENSE_LEGACY) + is not DiffusionKVCacheMode.PAGED_SCHEDULER + ): + return None + if not profile_requests: + raise ValueError("paged_scheduler Diffusion KV initialization requires prepared memory profile requests") + if device is None: + device = current_omni_platform.get_torch_device(0) + vllm_config = create_diffusion_vllm_config(device, od_config) + worker_specs = executor.get_kv_cache_specs() + available_memory = executor.determine_available_kv_memory(profile_requests) + if len(worker_specs) != od_config.num_gpus or len(available_memory) != od_config.num_gpus: + raise ValueError( + "Diffusion KV initialization rank count mismatch: " + f"expected={od_config.num_gpus}, specs={len(worker_specs)}, memory={len(available_memory)}" + ) + worker_configs, scheduler_kv_cache_config = build_native_kv_cache_configs( + vllm_config, + worker_specs, + available_memory, + ) + scheduler_block_size, hash_block_size = resolve_kv_cache_block_sizes( + scheduler_kv_cache_config, + vllm_config, + ) + executor.set_kv_cache_configs(worker_configs) + return ( + scheduler_kv_cache_config, + scheduler_block_size, + hash_block_size, + vllm_config, + ) diff --git a/vllm_omni/diffusion/diffusion_kv/manager.py b/vllm_omni/diffusion/diffusion_kv/manager.py new file mode 100644 index 000000000..eee51b864 --- /dev/null +++ b/vllm_omni/diffusion/diffusion_kv/manager.py @@ -0,0 +1,191 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from __future__ import annotations + +from collections.abc import Sequence + +from vllm.v1.core.kv_cache_manager import KVCacheManager +from vllm.v1.kv_cache_interface import KVCacheConfig + +from vllm_omni.diffusion.diffusion_kv.metadata import ( + DiffusionKVMetadata, + DiffusionKVSequenceMetadata, +) +from vllm_omni.diffusion.diffusion_kv.request import DiffusionKVRequest + + +class DiffusionKVAdmissionError(ValueError): + """A request-specific KV admission failure safe to return to the caller.""" + + +class DiffusionKVCacheManager: + """Atomic public-request facade over native vLLM KVCacheManager.""" + + def __init__( + self, + kv_cache_config: KVCacheConfig, + *, + max_model_len: int, + scheduler_block_size: int, + hash_block_size: int, + max_in_flight_tokens: int | None = None, + ) -> None: + if max_model_len <= 0: + raise ValueError(f"max_model_len must be positive, got {max_model_len}") + if kv_cache_config.num_blocks <= 0 or not kv_cache_config.kv_cache_groups: + raise ValueError("Diffusion KVCacheConfig must contain a positive block pool and at least one group") + if scheduler_block_size <= 0 or hash_block_size <= 0: + raise ValueError("scheduler_block_size and hash_block_size must be positive") + if max_in_flight_tokens is not None and max_in_flight_tokens <= 0: + raise ValueError(f"max_in_flight_tokens must be positive when provided, got {max_in_flight_tokens}") + self.native_manager = KVCacheManager( + kv_cache_config=kv_cache_config, + max_model_len=max_model_len, + scheduler_block_size=scheduler_block_size, + hash_block_size=hash_block_size, + max_in_flight_tokens=max_in_flight_tokens, + enable_caching=False, + ) + self.max_model_len = max_model_len + # Native vLLM may reserve a null block, so an idle BlockPool does not + # necessarily report ``kv_cache_config.num_blocks`` free blocks. + self._empty_pool_num_free_blocks = self.native_manager.block_pool.get_num_free_blocks() + self._requests: dict[str, tuple[DiffusionKVRequest, ...]] = {} + self._metadata: dict[str, DiffusionKVMetadata] = {} + self._internal_request_ids: set[str] = set() + self._next_allocation_generation = 1 + + def _get_empty_pool_required_blocks(self, requests: Sequence[DiffusionKVRequest]) -> int: + """Return the native full-sequence reservation for one public request. + + vLLM performs this calculation inside ``allocate_slots`` when + ``full_sequence_must_fit`` is enabled. Diffusion CFG adds one semantic + requirement that native requests do not have: every execution sequence + must fit atomically. Sum the native per-sequence result so a request + that can never fit is rejected independently of current pool usage. + """ + + coordinator = self.native_manager.coordinator + empty_blocks = self.native_manager.empty_kv_cache_blocks.blocks + return sum( + coordinator.get_num_blocks_to_allocate( + request_id=request.request_id, + num_tokens=request.num_tokens, + new_computed_blocks=empty_blocks, + num_encoder_tokens=0, + total_computed_tokens=0, + num_local_computed_tokens=0, + num_tokens_main_model=request.num_tokens, + apply_admission_cap=True, + ) + for request in requests + ) + + def has_request(self, public_request_id: str) -> bool: + return public_request_id in self._requests + + def reserve_request( + self, + public_request_id: str, + kv_requests: Sequence[DiffusionKVRequest], + ) -> DiffusionKVMetadata | None: + """Reserve all CFG sequences or roll back the complete public request.""" + + if not public_request_id: + raise ValueError("public_request_id must be non-empty") + if public_request_id in self._requests: + raise ValueError(f"Diffusion KV request {public_request_id!r} is already allocated") + requests = tuple(kv_requests) + if not requests: + raise ValueError("Diffusion KV allocation requires at least one sequence") + + sequence_ids = [request.sequence_id for request in requests] + internal_ids = [request.request_id for request in requests] + expected_sequence_ids = list(range(len(requests))) + if sequence_ids != expected_sequence_ids: + raise ValueError( + "Diffusion KV requests must be ordered by contiguous sequence_id values: " + f"expected={expected_sequence_ids}, got={sequence_ids}" + ) + if len(internal_ids) != len(set(internal_ids)): + raise ValueError(f"Diffusion KV internal request IDs must be unique, got {internal_ids}") + conflicts = self._internal_request_ids.intersection(internal_ids) + if conflicts: + raise ValueError(f"Diffusion KV internal request IDs are already allocated: {sorted(conflicts)}") + for request in requests: + if request.kv_contexts: + raise DiffusionKVAdmissionError("independent DiffusionKVContext allocation is not implemented yet") + if request.seq_len > self.max_model_len: + raise DiffusionKVAdmissionError( + f"Diffusion KV sequence {request.request_id!r} exceeds max_model_len: " + f"seq_len={request.seq_len}, max_model_len={self.max_model_len}" + ) + + required_blocks = self._get_empty_pool_required_blocks(requests) + if required_blocks > self._empty_pool_num_free_blocks: + raise DiffusionKVAdmissionError( + f"Diffusion KV request {public_request_id!r} cannot fit even when the block pool is empty: " + f"required_blocks={required_blocks}, available_blocks={self._empty_pool_num_free_blocks}; " + "increase KV cache capacity or reduce the request sequence count/length" + ) + + allocated: list[DiffusionKVRequest] = [] + sequence_metadata: list[DiffusionKVSequenceMetadata] = [] + try: + for request in requests: + blocks = self.native_manager.allocate_slots( + request, + num_new_tokens=request.seq_len, + delay_cache_blocks=True, + full_sequence_must_fit=True, + ) + if blocks is None: + self._rollback(allocated) + allocated.clear() + return None + allocated.append(request) + sequence_metadata.append( + DiffusionKVSequenceMetadata( + sequence_id=request.sequence_id, + prefix_len=request.prefix_len, + target_len=request.target_len, + seq_len=request.seq_len, + block_ids=blocks.get_block_ids(), + ) + ) + except Exception: + self._rollback(allocated) + raise + + metadata = DiffusionKVMetadata( + request_id=public_request_id, + allocation_generation=self._next_allocation_generation, + sequences=tuple(sequence_metadata), + ) + self._next_allocation_generation += 1 + self._requests[public_request_id] = requests + self._metadata[public_request_id] = metadata + self._internal_request_ids.update(internal_ids) + return metadata + + def get_metadata(self, public_request_id: str) -> DiffusionKVMetadata: + try: + return self._metadata[public_request_id] + except KeyError as exc: + raise KeyError(f"Diffusion KV request {public_request_id!r} is not allocated") from exc + + def free_request(self, public_request_id: str) -> None: + requests = self._requests.pop(public_request_id, ()) + self._metadata.pop(public_request_id, None) + for request in reversed(requests): + self.native_manager.free(request) + self._internal_request_ids.discard(request.request_id) + + def close(self) -> None: + for public_request_id in tuple(self._requests): + self.free_request(public_request_id) + + def _rollback(self, requests: Sequence[DiffusionKVRequest]) -> None: + for request in reversed(requests): + self.native_manager.free(request) diff --git a/vllm_omni/diffusion/diffusion_kv/request.py b/vllm_omni/diffusion/diffusion_kv/request.py index 4db9701a1..fd4740099 100644 --- a/vllm_omni/diffusion/diffusion_kv/request.py +++ b/vllm_omni/diffusion/diffusion_kv/request.py @@ -105,7 +105,7 @@ class DiffusionKVRequest: self.status = RequestStatus.WAITING self.num_preemptions = 0 self.num_in_flight_tokens = 0 - # vLLM 0.26 uses this optional boundary when publishing cached blocks; + # vLLM 0.26+ uses this optional boundary when publishing cached blocks; # zero means that no sparse-retention boundary is active. self.shared_prefix_boundary = 0 diff --git a/vllm_omni/diffusion/executor/abstract.py b/vllm_omni/diffusion/executor/abstract.py index ec6b3486a..e653dc7ce 100644 --- a/vllm_omni/diffusion/executor/abstract.py +++ b/vllm_omni/diffusion/executor/abstract.py @@ -5,10 +5,12 @@ from collections.abc import Callable from typing import TYPE_CHECKING, Any from vllm.utils.import_utils import resolve_obj_by_qualname +from vllm.v1.kv_cache_interface import KVCacheConfig, KVCacheSpec from vllm_omni.diffusion.data import OmniDiffusionConfig if TYPE_CHECKING: + from vllm_omni.diffusion.request import OmniDiffusionRequest from vllm_omni.diffusion.sched.interface import DiffusionSchedulerOutput from vllm_omni.diffusion.worker.utils import BaseRunnerOutput @@ -115,6 +117,39 @@ class DiffusionExecutor(ABC): """ return None + def get_kv_cache_specs(self) -> list[dict[str, KVCacheSpec]]: + """Collect rank-local native specs after every Worker loads its model.""" + + result = self.collective_rpc( + "get_kv_cache_specs", + unique_reply_rank=0, + exec_all_ranks=True, + ) + if not isinstance(result, list): + raise TypeError(f"get_kv_cache_specs returned {type(result).__name__}, expected list") + return result + + def determine_available_kv_memory(self, profile_requests: list[OmniDiffusionRequest]) -> list[int]: + """Profile and collect the KV memory budget on every Worker rank.""" + + result = self.collective_rpc( + "determine_available_kv_memory", + args=(profile_requests,), + unique_reply_rank=0, + exec_all_ranks=True, + ) + if not isinstance(result, list) or not all(isinstance(value, int) for value in result): + raise TypeError("determine_available_kv_memory must return list[int]") + return result + + def set_kv_cache_configs(self, kv_cache_configs: list[KVCacheConfig]) -> None: + """Send each native rank-local config back to all Workers.""" + + # The default control-plane RPC mode executes on every rank and has + # rank 0 return the gathered rank statuses, so failures on nonzero + # ranks are not silently dropped. + self.collective_rpc("set_kv_cache_configs", args=(kv_cache_configs,)) + @abstractmethod def shutdown(self) -> None: """Shutdown the executor and release resources.""" diff --git a/vllm_omni/diffusion/executor/multiproc_executor.py b/vllm_omni/diffusion/executor/multiproc_executor.py index 6f0917e32..ff9253962 100644 --- a/vllm_omni/diffusion/executor/multiproc_executor.py +++ b/vllm_omni/diffusion/executor/multiproc_executor.py @@ -465,16 +465,14 @@ class MultiprocDiffusionExecutor(DiffusionExecutor): for new_req in new_reqs: validate_new_request_data_identity(new_req) - has_diffusion_kv_metadata = any(new_req.diffusion_kv_metadata is not None for new_req in new_reqs) - # DP multi-concurrency: when DLO+AllGather is active and multiple - # requests are scheduled, send ALL requests in one broadcast RPC. - # Each rank picks req[rank % len(reqs)] and computes independently. + # requests are scheduled, send every complete NewRequestData envelope + # in one broadcast RPC. Each rank picks one envelope, keeping its + # request and Diffusion KV metadata inseparable. # All ranks reply (unique_reply_rank=None) so we collect dp_size # responses and match by dp_rank. if ( len(new_reqs) > 1 - and not has_diffusion_kv_metadata and getattr(self.od_config, "enable_distributed_layerwise_offload", False) and getattr(self.od_config, "dlo_use_allgather", True) ): @@ -507,12 +505,11 @@ class MultiprocDiffusionExecutor(DiffusionExecutor): f"empty prompt request IDs: {empty_prompt_ids}." ) - reqs_list = [nr.req for nr in new_reqs] try: results = self.collective_rpc( "execute_model", timeout=_DLO_DP_WAVE_TIMEOUT_S, - args=(reqs_list, self.od_config, scheduler_output.kv_prefetch_job), + args=(new_reqs, self.od_config, scheduler_output.kv_prefetch_job), unique_reply_rank=None, exec_all_ranks=True, ) diff --git a/vllm_omni/diffusion/io_support.py b/vllm_omni/diffusion/io_support.py index 8d1405a6c..368577a0a 100644 --- a/vllm_omni/diffusion/io_support.py +++ b/vllm_omni/diffusion/io_support.py @@ -6,6 +6,7 @@ from __future__ import annotations import inspect from vllm_omni.diffusion.data import OmniDiffusionConfig +from vllm_omni.diffusion.model_metadata import get_diffusion_model_metadata from vllm_omni.diffusion.registry import DiffusionModelRegistry @@ -47,3 +48,9 @@ def get_dummy_run_num_frames(model_class_name: str, supports_audio_input: bool) if model_cls is not None and hasattr(model_cls, "dummy_run_num_frames"): return int(getattr(model_cls, "dummy_run_num_frames")) return 2 if supports_audio_input or supports_audio_output(model_class_name) else 1 + + +def get_dummy_run_num_image_inputs(model_class_name: str) -> int: + """Return the maximum advertised image-input count for profiling.""" + + return get_diffusion_model_metadata(model_class_name).max_multimodal_image_inputs or 1 diff --git a/vllm_omni/diffusion/models/hunyuan_image3/hunyuan_image3_transformer.py b/vllm_omni/diffusion/models/hunyuan_image3/hunyuan_image3_transformer.py index 877fd1c3f..45a94d76d 100644 --- a/vllm_omni/diffusion/models/hunyuan_image3/hunyuan_image3_transformer.py +++ b/vllm_omni/diffusion/models/hunyuan_image3/hunyuan_image3_transformer.py @@ -913,7 +913,7 @@ class HunYuanRotary2DEmbedder: return q, k -class ImageKVCacheManager: +class ImageKVCacheManager(nn.Module): """ Manages specialized caching and updating of KV-Cache for image tokens in multimodal models. """ @@ -932,6 +932,7 @@ class ImageKVCacheManager: image_token_len: Number of tokens per image (including special placeholders), default 4097 (timestamp + 4096 image tokens). """ + super().__init__() # attention related self.num_heads = num_heads self.num_kv_heads = num_kv_heads @@ -953,6 +954,10 @@ class ImageKVCacheManager: softmax_scale=self.scaling, num_kv_heads=self.num_kv_heads, prefix=f"{prefix}.attn" if prefix else "", + paged_kv_cache_role="primary", + # Hunyuan image RoPE produces bfloat16 K and denoising runs under + # bfloat16 autocast, independent of the weight-loading dtype. + paged_kv_cache_dtype=torch.bfloat16, ) @staticmethod @@ -1095,7 +1100,7 @@ class ImageKVCacheManager: ) return key, value - def __call__( + def forward( self, query: torch.Tensor, key: torch.Tensor, diff --git a/vllm_omni/diffusion/request.py b/vllm_omni/diffusion/request.py index ce097a2d9..82ff7ebe2 100644 --- a/vllm_omni/diffusion/request.py +++ b/vllm_omni/diffusion/request.py @@ -14,6 +14,7 @@ if TYPE_CHECKING: from vllm_omni.diffusion.diffusion_kv.request import DiffusionKVRequest DUMMY_DIFFUSION_REQUEST_ID = "dummy_req_id" +_DUMMY_DIFFUSION_REQUEST_ID_PREFIX = f"{DUMMY_DIFFUSION_REQUEST_ID}/" @dataclass @@ -90,4 +91,6 @@ class OmniDiffusionRequest: @classmethod def is_dummy_run_request_id(cls, request_id: str | None) -> bool: - return request_id == DUMMY_DIFFUSION_REQUEST_ID + return request_id == DUMMY_DIFFUSION_REQUEST_ID or ( + isinstance(request_id, str) and request_id.startswith(_DUMMY_DIFFUSION_REQUEST_ID_PREFIX) + ) diff --git a/vllm_omni/diffusion/sched/base_scheduler.py b/vllm_omni/diffusion/sched/base_scheduler.py index b558c1ee6..17291db37 100644 --- a/vllm_omni/diffusion/sched/base_scheduler.py +++ b/vllm_omni/diffusion/sched/base_scheduler.py @@ -7,10 +7,14 @@ from abc import ABC, abstractmethod from collections import deque from dataclasses import fields +from vllm.config import VllmConfig from vllm.logger import init_logger +from vllm.v1.kv_cache_interface import KVCacheConfig from vllm_omni.diffusion.data import OmniDiffusionConfig from vllm_omni.diffusion.diffusion_kv.config import DiffusionKVCacheMode +from vllm_omni.diffusion.diffusion_kv.manager import DiffusionKVAdmissionError, DiffusionKVCacheManager +from vllm_omni.diffusion.diffusion_kv.metadata import DiffusionKVMetadata from vllm_omni.diffusion.request import OmniDiffusionRequest from vllm_omni.diffusion.sched.interface import ( CachedRequestData, @@ -49,9 +53,17 @@ class BaseScheduler(ABC): self._finished_req_ids: set[str] = set() self.max_num_running_reqs: int = 1 self._prefetch_enabled: bool = False - self._diffusion_kv_enabled = False + self._diffusion_kv_manager: DiffusionKVCacheManager | None = None - def initialize(self, od_config: OmniDiffusionConfig) -> None: + def initialize( + self, + od_config: OmniDiffusionConfig, + *, + kv_cache_config: KVCacheConfig | None = None, + scheduler_block_size: int | None = None, + hash_block_size: int | None = None, + kv_vllm_config: VllmConfig | None = None, + ) -> None: self.od_config = od_config self._request_states.clear() self._step_id = 0 @@ -66,10 +78,36 @@ class BaseScheduler(ABC): self.max_num_running_reqs = 1 omni_kv = getattr(od_config, "omni_kv_config", None) or {} self._prefetch_enabled = bool(omni_kv.get("enable_kv_async_prefetch", False)) - self._diffusion_kv_enabled = ( + diffusion_kv_enabled = ( getattr(od_config, "diffusion_kv_mode", DiffusionKVCacheMode.DENSE_LEGACY) is DiffusionKVCacheMode.PAGED_SCHEDULER ) + if diffusion_kv_enabled: + if kv_cache_config is None: + raise ValueError("paged_scheduler Diffusion KV requires a native Scheduler KVCacheConfig") + if scheduler_block_size is None or hash_block_size is None: + raise ValueError("paged_scheduler Diffusion KV requires native scheduler/hash block sizes") + if kv_vllm_config is None: + raise ValueError("paged_scheduler Diffusion KV requires the native VllmConfig used for cache sizing") + self._diffusion_kv_manager = DiffusionKVCacheManager( + kv_cache_config, + max_model_len=kv_vllm_config.model_config.max_model_len, + scheduler_block_size=scheduler_block_size, + hash_block_size=hash_block_size, + max_in_flight_tokens=kv_vllm_config.max_in_flight_tokens, + ) + else: + if any( + value is not None + for value in ( + kv_cache_config, + scheduler_block_size, + hash_block_size, + kv_vllm_config, + ) + ): + raise ValueError("dense_legacy Scheduler received unexpected Diffusion KV cache initialization state") + self._diffusion_kv_manager = None self._reset_scheduler_state() def add_request(self, request: OmniDiffusionRequest) -> str: @@ -104,6 +142,35 @@ class BaseScheduler(ABC): if not self._can_schedule_waiting(state): break + diffusion_kv_metadata: DiffusionKVMetadata | None = None + if self._diffusion_kv_manager is not None: + if self._diffusion_kv_manager.has_request(request_id): + diffusion_kv_metadata = self._diffusion_kv_manager.get_metadata(request_id) + else: + try: + allocation = self._diffusion_kv_manager.reserve_request( + request_id, + state.diffusion_kv_requests, + ) + except Exception as exc: + # schedule() runs outside the Engine execution-error + # wrapper. Convert both expected admission failures and + # unexpected native allocator failures into a terminal + # request error so the busy loop can wake its stream. + if not isinstance(exc, DiffusionKVAdmissionError): + logger.exception( + "Unexpected Diffusion KV allocation failure for request %s", + request_id, + ) + self._finish_requests( + {request_id: DiffusionRequestStatus.FINISHED_ERROR}, + {request_id: str(exc)}, + ) + continue + if allocation is None: + break + diffusion_kv_metadata = allocation + self._waiting.popleft() was_new_request = state.status == DiffusionRequestStatus.WAITING if not self._running: @@ -111,7 +178,12 @@ class BaseScheduler(ABC): state.status = DiffusionRequestStatus.RUNNING self._running.append(request_id) if was_new_request: - scheduled_new_reqs.append(NewRequestData.from_state(state)) + scheduled_new_reqs.append( + NewRequestData.from_state( + state, + diffusion_kv_metadata=diffusion_kv_metadata, + ) + ) else: scheduled_cached_request_ids.append(request_id) @@ -184,6 +256,8 @@ class BaseScheduler(ABC): return self._request_states.get(request_id) def pop_request_state(self, request_id: str) -> SchedulerRequestState | None: + if self._diffusion_kv_manager is not None: + self._diffusion_kv_manager.free_request(request_id) self._pop_extra_request_state(request_id) return self._request_states.pop(request_id, None) @@ -206,6 +280,9 @@ class BaseScheduler(ABC): self._finish_requests({request_id: status for request_id in request_ids}) def close(self) -> None: + if self._diffusion_kv_manager is not None: + self._diffusion_kv_manager.close() + self._diffusion_kv_manager = None self._request_states.clear() self._waiting.clear() self._running.clear() @@ -245,6 +322,8 @@ class BaseScheduler(ABC): self._waiting = deque(request_id for request_id in self._waiting if request_id not in waiting_to_remove) for request_id in finished_req_ids: + if self._diffusion_kv_manager is not None: + self._diffusion_kv_manager.free_request(request_id) state = self._request_states[request_id] status = statuses[request_id] state.status = status @@ -266,7 +345,13 @@ class BaseScheduler(ABC): # update_from_output() processes the runner output. It is already # marked finished at that point, but we still need to surface its id # in this update so the engine can observe the terminal state. + # Also surface admission failures recorded while schedule() built this + # output. Older finished ids retained only for Worker cleanup have + # already been popped by the Engine and are deliberately ignored. finished_req_ids = { + request_id for request_id in sched_output.finished_req_ids if request_id in self._request_states + } + finished_req_ids |= { request_id for request_id in sched_output.scheduled_request_ids if request_id in self._finished_req_ids } finished_req_ids |= self._finish_requests(statuses, errors) @@ -280,7 +365,7 @@ class BaseScheduler(ABC): def _make_request_state(self, request_id: str, request: OmniDiffusionRequest) -> SchedulerRequestState: kv_requests = request.diffusion_kv_requests or () - if self._diffusion_kv_enabled: + if self._diffusion_kv_manager is not None: self._reject_legacy_dense_kv(request) if not kv_requests: raise ValueError("paged_scheduler request preprocessing did not produce DiffusionKVRequest state") diff --git a/vllm_omni/diffusion/vllm_config.py b/vllm_omni/diffusion/vllm_config.py new file mode 100644 index 000000000..816739de4 --- /dev/null +++ b/vllm_omni/diffusion/vllm_config.py @@ -0,0 +1,180 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +"""Build native vLLM configuration views for the Omni diffusion runtime.""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Any + +import torch +from vllm.config import CompilationConfig, DeviceConfig, VllmConfig +from vllm.transformers_utils.config import get_hf_text_config + +from vllm_omni.diffusion.data import OmniDiffusionConfig +from vllm_omni.diffusion.diffusion_kv.config import DiffusionKVCacheMode + + +def resolve_diffusion_max_model_len(od_config: OmniDiffusionConfig) -> int: + """Resolve the native KV admission bound for a diffusion stage.""" + + configured = getattr(od_config, "max_model_len", None) + if configured is not None: + configured = int(configured) + if configured == -1: + configured = None + elif configured <= 0: + raise ValueError(f"max_model_len must be positive for Scheduler-managed Diffusion KV, got {configured}") + else: + return configured + + # ``max_model_len`` keeps the native vLLM meaning: the per-sequence KV + # admission ceiling consumed by ``KVCacheManager``. HunyuanImage3 exposes + # that ceiling directly as ``max_position_embeddings`` because its text, + # reference-image, and generated-image tokens share one Transformer + # sequence. Other diffusion models commonly express their serving limits + # through resolution, duration, or frame-count constraints instead. When + # paged KV support expands to those models, the model integration must + # derive ``max_model_len`` from its serving limits, while model-owned + # preprocessing validates each request and derives its matching + # ``seq_len``. Text-encoder length fields must not be mistaken for the DiT + # KV limit. + hf_config = getattr(od_config, "tf_model_config", None) + for field_name in ("max_position_embeddings", "max_sequence_length", "seq_length"): + value = getattr(hf_config, field_name, None) + if value is None and isinstance(hf_config, dict): + value = hf_config.get(field_name) + if value is not None and int(value) > 0: + return int(value) + + mode = getattr(getattr(od_config, "diffusion_kv_mode", None), "value", None) + if mode == "paged_scheduler": + raise ValueError( + "paged_scheduler Diffusion KV requires max_model_len or a model config " + "with max_position_embeddings/max_sequence_length/seq_length" + ) + # Dense execution does not consume the native KV admission bound, but the + # shared Worker VllmConfig still requires a positive placeholder. + return 1 + + +@dataclass +class _DiffusionVllmModelConfig: + """Private adapter for the native ``VllmConfig.model_config`` surface.""" + + model: str | None + dtype: torch.dtype + max_model_len: int + original_max_model_len: int | None + runner_type: str = "generate" + quantization: str | None = None + quantization_config: Any | None = None + hf_config: Any | None = None + hf_text_config: Any | None = None + multimodal_config: Any | None = None + enforce_eager: bool = False + disable_cascade_attn: bool = False + enable_return_routed_experts: bool = False + enable_prompt_embeds: bool = False + logits_processors: tuple[str, ...] = () + use_mla: bool = False + is_moe: bool = False + + @property + def is_quantized(self) -> bool: + return self.quantization is not None + + def is_nvfp4_quantized(self) -> bool: + return self.quantization == "modelopt_fp4" + + @property + def is_diffusion(self) -> bool: + # This config is also installed as the current VllmConfig for the + # existing Omni diffusion runtime, which does not use vLLM's native + # diffusion runner. Preserve that established backend-selection flag. + return False + + +def _make_diffusion_vllm_model_config(od_config: OmniDiffusionConfig) -> _DiffusionVllmModelConfig: + quant_config = getattr(od_config, "quantization_config", None) + quantization = quant_config.get_name() if quant_config is not None and hasattr(quant_config, "get_name") else None + hf_config = getattr(od_config, "tf_model_config", None) + hf_text_config = get_hf_text_config(hf_config) if hasattr(hf_config, "get_text_config") else hf_config + configured_max_model_len = getattr(od_config, "max_model_len", None) + max_model_len = resolve_diffusion_max_model_len(od_config) + paged_scheduler = ( + getattr(od_config, "diffusion_kv_mode", DiffusionKVCacheMode.DENSE_LEGACY) + is DiffusionKVCacheMode.PAGED_SCHEDULER + ) + return _DiffusionVllmModelConfig( + model=od_config.model, + dtype=od_config.dtype, + max_model_len=max_model_len, + # Preserve the user's -1 sentinel only for the native paged-cache + # sizing flow. Dense keeps its established resolved model-config view. + original_max_model_len=configured_max_model_len if paged_scheduler else max_model_len, + quantization=quantization, + quantization_config=quant_config, + hf_config=hf_config, + hf_text_config=hf_text_config, + enforce_eager=getattr(od_config, "enforce_eager", False), + is_moe=bool(getattr(od_config, "is_moe", False)), + ) + + +def create_base_diffusion_vllm_config( + device: torch.device, + od_config: OmniDiffusionConfig, +) -> VllmConfig: + """Create the native vLLM 0.27 config used by diffusion.""" + + return VllmConfig( + compilation_config=CompilationConfig(), + device_config=DeviceConfig(device=device), + additional_config=od_config.additional_config, + ) + + +def configure_diffusion_vllm_config(vllm_config: VllmConfig, od_config: OmniDiffusionConfig) -> VllmConfig: + """Populate native model, parallel, cache, and Scheduler configuration.""" + + parallel_config = od_config.parallel_config + vllm_config.parallel_config.tensor_parallel_size = parallel_config.tensor_parallel_size + vllm_config.parallel_config.data_parallel_size = parallel_config.data_parallel_size + if parallel_config.enable_expert_parallel and od_config.is_moe: + vllm_config.parallel_config.data_parallel_size = ( + parallel_config.data_parallel_size * parallel_config.cfg_parallel_size + ) + vllm_config.parallel_config.prefill_context_parallel_size = parallel_config.sequence_parallel_size + vllm_config.parallel_config.enable_expert_parallel = parallel_config.enable_expert_parallel + + vllm_config.model_config = _make_diffusion_vllm_model_config(od_config) # type: ignore[assignment] + vllm_config.quant_config = od_config.quantization_config + vllm_config.profiler_config = od_config.profiler_config + if ( + getattr(od_config, "diffusion_kv_mode", DiffusionKVCacheMode.DENSE_LEGACY) + is DiffusionKVCacheMode.PAGED_SCHEDULER + ): + vllm_config.parallel_config.pipeline_parallel_size = parallel_config.pipeline_parallel_size + vllm_config.cache_config.gpu_memory_utilization = float(getattr(od_config, "gpu_memory_utilization", 0.9)) + vllm_config.cache_config.kv_cache_memory_bytes = getattr(od_config, "kv_cache_memory_bytes", None) + # Prefix hashes/publication are intentionally deferred to the next PR. + vllm_config.cache_config.enable_prefix_caching = False + max_num_batched_tokens = getattr(od_config, "max_num_batched_tokens", None) + if max_num_batched_tokens is not None: + vllm_config.scheduler_config.max_num_batched_tokens = int(max_num_batched_tokens) + return vllm_config + + +def create_diffusion_vllm_config( + device: torch.device, + od_config: OmniDiffusionConfig, +) -> VllmConfig: + """Build the same fully populated native config in Engine and Workers.""" + + vllm_config = create_base_diffusion_vllm_config( + device, + od_config, + ) + return configure_diffusion_vllm_config(vllm_config, od_config) diff --git a/vllm_omni/diffusion/worker/diffusion_model_runner.py b/vllm_omni/diffusion/worker/diffusion_model_runner.py index 4a763d651..44993fd2c 100644 --- a/vllm_omni/diffusion/worker/diffusion_model_runner.py +++ b/vllm_omni/diffusion/worker/diffusion_model_runner.py @@ -11,6 +11,7 @@ model-related operations. from __future__ import annotations import copy +import gc import time from collections.abc import Callable from contextlib import AbstractContextManager, nullcontext @@ -21,6 +22,7 @@ from torch.profiler import record_function from vllm.config import LoadConfig, VllmConfig from vllm.logger import init_logger from vllm.utils.mem_utils import DeviceMemoryProfiler, GiB_bytes +from vllm.v1.kv_cache_interface import KVCacheConfig, KVCacheSpec from vllm_omni.diffusion.cache.cachedit import cache_summary from vllm_omni.diffusion.cache.prompt_embed_cache import ( @@ -45,8 +47,10 @@ from vllm_omni.diffusion.offloader import get_offload_backend from vllm_omni.diffusion.registry import _NO_CACHE_ACCELERATION from vllm_omni.diffusion.request import OmniDiffusionRequest from vllm_omni.diffusion.sched.interface import ( + CachedRequestData, DiffusionSchedulerOutput, KVPrefetchJob, + NewRequestData, validate_new_request_data_identity, ) from vllm_omni.diffusion.worker.input_batch import InputBatch, scatter_latents @@ -139,6 +143,11 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin): self.offload_backend: Any | None = None self.prompt_embed_cache: Any | None = None self.input_batch: InputBatch | None = None + self.model_memory_usage = 0 + # Rank-local native config generated by the Engine. Physical tensors + # and BlockTable binding remain the responsibility of the paged + # data-plane initialization path. + self.kv_cache_config: KVCacheConfig | None = None # Cache for per-request stepwise state. self.state_cache: dict[str, StepRequestState] = {} @@ -167,8 +176,8 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin): metadata: DiffusionKVMetadata | None, ) -> None: cache_mode = getattr(self.od_config, "diffusion_kv_mode", DiffusionKVCacheMode.DENSE_LEGACY) - # W0 does not wire the Scheduler allocator or Worker installer yet, so - # validate an allocation snapshot when present without requiring one. + if cache_mode is DiffusionKVCacheMode.PAGED_SCHEDULER and metadata is None: + raise ValueError(f"paged_scheduler request {request_id!r} requires Diffusion KV metadata") if cache_mode is not DiffusionKVCacheMode.PAGED_SCHEDULER and metadata is not None: raise ValueError(f"{cache_mode.value} request {request_id!r} must not carry Diffusion KV metadata") @@ -273,6 +282,7 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin): m.consumed_memory / GiB_bytes, time_after_load - time_before_load, ) + self.model_memory_usage = int(m.consumed_memory) logger.info("Model runner: Model loaded successfully.") if self.od_config.streaming_output and not getattr(self.od_config, "step_execution", False): @@ -368,6 +378,45 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin): logger.info("Model runner: Initialization complete.") + def get_kv_cache_spec(self) -> dict[str, KVCacheSpec]: + """Collect native specs from cache-enabled loaded attention modules.""" + + from vllm_omni.diffusion.attention.layer import Attention + from vllm_omni.diffusion.diffusion_kv.config import DiffusionKVCacheMode + + if self.od_config.diffusion_kv_mode is not DiffusionKVCacheMode.PAGED_SCHEDULER: + return {} + if self.pipeline is None: + raise RuntimeError("Model must be loaded before collecting Diffusion KV cache specs") + + specs: dict[str, KVCacheSpec] = {} + for layer_name, module in self.pipeline.named_modules(): + if not isinstance(module, Attention): + continue + spec = module.get_kv_cache_spec(self.vllm_config) + if spec is not None: + specs[layer_name] = spec + if not specs: + raise RuntimeError( + "paged_scheduler Diffusion KV found no cache-enabled Attention modules " + f"in {type(self.pipeline).__name__}" + ) + return specs + + def set_kv_cache_config(self, kv_cache_config: KVCacheConfig) -> None: + """Retain the Engine-generated rank-local config for the data plane.""" + + configured_layers = { + layer_name for group in kv_cache_config.kv_cache_groups for layer_name in group.layer_names + } + expected_layers = set(self.get_kv_cache_spec()) + if configured_layers != expected_layers: + raise ValueError( + "Rank-local Diffusion KVCacheConfig layer mismatch: " + f"expected={sorted(expected_layers)}, configured={sorted(configured_layers)}" + ) + self.kv_cache_config = kv_cache_config + def clear_prompt_embed_cache(self) -> None: """Evict all cached text-encoder outputs (e.g. between training epochs). @@ -516,6 +565,7 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin): require_request_batch_support: bool, kv_prefetch_job: KVPrefetchJob | None = None, record_name: str, + record_output_peak_memory: bool = True, ) -> BatchRunnerOutput: assert self.pipeline is not None, "Model not loaded. Call load_model() first." if not reqs: @@ -545,7 +595,7 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin): batch = DiffusionRequestBatch(requests=reqs) is_primary = not torch.distributed.is_initialized() or torch.distributed.get_rank() == 0 - if is_primary: + if is_primary and record_output_peak_memory: current_omni_platform.reset_peak_memory_stats() with set_forward_context(vllm_config=self.vllm_config, omni_diffusion_config=od_config): @@ -558,7 +608,7 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin): pipeline_name=type(self.pipeline).__name__, ) - if is_primary and outputs: + if is_primary and outputs and record_output_peak_memory: batch_peak_memory_mb = self._sample_peak_memory_mb() for output in outputs: output.peak_memory_mb = max(output.peak_memory_mb, batch_peak_memory_mb) @@ -626,6 +676,59 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin): assert output is not None return output + def profile_run(self, requests: list[OmniDiffusionRequest]) -> None: + """Run the maximum per-rank request batch for memory profiling. + + This deliberately bypasses Scheduler admission and Diffusion KV + metadata validation because cache capacity has not been sized yet. + It otherwise uses the normal execution-mode path so model inputs, + collective communication, backend workspaces, denoising, and decode + allocations contribute to the observed peak. Step execution profiles + one fused ``InputBatch`` instead of sequential single-request forwards. + """ + + if not requests: + raise ValueError("Diffusion memory profiling requires at least one request.") + + runner_output: BatchRunnerOutput | None = None + request_ids = [request.request_id for request in requests] + try: + if getattr(self.od_config, "step_execution", False): + scheduler_output = DiffusionSchedulerOutput( + step_id=0, + scheduled_new_reqs=[ + NewRequestData(request_id=request.request_id, req=request) for request in requests + ], + scheduled_cached_reqs=CachedRequestData.make_empty(), + finished_req_ids=set(), + num_running_reqs=len(requests), + num_waiting_reqs=0, + ) + runner_output = self._execute_stepwise( + scheduler_output, + validate_kv_metadata=False, + record_output_peak_memory=False, + ) + else: + runner_output = self._execute_request_list( + requests, + od_config=self.od_config, + allow_single_output=len(requests) == 1, + require_request_batch_support=len(requests) > 1, + record_name="pipeline_memory_profile", + # The enclosing native memory_profiling context owns the + # peak counters. Resetting them here would discard request + # preparation allocations and understate the budget. + record_output_peak_memory=False, + ) + current_omni_platform.synchronize() + finally: + for request_id in request_ids: + self.state_cache.pop(request_id, None) + self.input_batch = None + del runner_output + gc.collect() + def execute_model_batch( self, scheduler_output: DiffusionSchedulerOutput, @@ -750,13 +853,29 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin): def execute_stepwise(self, scheduler_output: DiffusionSchedulerOutput) -> BatchRunnerOutput: """Execute one step for one scheduled request and return runner output.""" + return self._execute_stepwise( + scheduler_output, + validate_kv_metadata=True, + record_output_peak_memory=True, + ) + + def _execute_stepwise( + self, + scheduler_output: DiffusionSchedulerOutput, + *, + validate_kv_metadata: bool, + record_output_peak_memory: bool, + ) -> BatchRunnerOutput: + """Execute one step with explicit validation and profiling policy.""" + assert self.pipeline is not None, "Model not loaded. Call load_model() first." for new_req in scheduler_output.scheduled_new_reqs: validate_new_request_data_identity(new_req) - self._validate_diffusion_kv_metadata( - request_id=new_req.req.request_id, - metadata=new_req.diffusion_kv_metadata, - ) + if validate_kv_metadata: + self._validate_diffusion_kv_metadata( + request_id=new_req.req.request_id, + metadata=new_req.diffusion_kv_metadata, + ) if not self._supports_step_mode(): raise ValueError("Current pipeline does not support step execution.") # Stepwise mode only supports the basic state-driven denoise path for now. @@ -771,7 +890,13 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin): had_active_states = bool(self.state_cache) states, new_request_ids = self._update_states(scheduler_output) is_primary = not torch.distributed.is_initialized() or torch.distributed.get_rank() == 0 - if new_request_ids and not had_active_states and is_primary and current_omni_platform.is_available(): + if ( + record_output_peak_memory + and new_request_ids + and not had_active_states + and is_primary + and current_omni_platform.is_available() + ): current_omni_platform.reset_peak_memory_stats() input_batch = self._prepare_batch_inputs(states, new_request_ids) attn_metadata = {} @@ -864,7 +989,7 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin): f"but batched noise_pred has {noise_pred.shape[0]} rows." ) - if is_primary: + if is_primary and record_output_peak_memory: batch_peak_memory_mb = self._sample_peak_memory_mb() states_by_id = {state.request_id: state for state in states} for state in states: diff --git a/vllm_omni/diffusion/worker/diffusion_worker.py b/vllm_omni/diffusion/worker/diffusion_worker.py index 436fcca9a..d9030e891 100644 --- a/vllm_omni/diffusion/worker/diffusion_worker.py +++ b/vllm_omni/diffusion/worker/diffusion_worker.py @@ -16,20 +16,21 @@ import signal import threading import traceback import uuid -from collections.abc import Iterator +from collections.abc import Callable, Iterator from contextlib import AbstractContextManager, contextmanager, nullcontext -from dataclasses import dataclass from typing import Any import torch +import torch.distributed as dist import zmq -from vllm.config import CompilationConfig, DeviceConfig, VllmConfig, set_current_vllm_config +from vllm.config import VllmConfig, set_current_vllm_config from vllm.distributed.device_communicators.shm_broadcast import MessageQueue from vllm.logger import init_logger from vllm.profiler.wrapper import CudaProfilerWrapper, WorkerProfiler -from vllm.transformers_utils.config import get_hf_text_config from vllm.utils.import_utils import resolve_obj_by_qualname -from vllm.utils.mem_utils import GiB_bytes +from vllm.utils.mem_utils import GiB_bytes, MemorySnapshot, format_gib, memory_profiling +from vllm.v1.kv_cache_interface import KVCacheConfig, KVCacheSpec +from vllm.v1.worker.utils import request_memory from vllm.v1.worker.workspace import init_workspace_manager from vllm_omni.diffusion.data import ( @@ -41,6 +42,7 @@ from vllm_omni.diffusion.data import ( OmniSleepTask, OmniWakeTask, ) +from vllm_omni.diffusion.diffusion_kv.config import DiffusionKVCacheMode from vllm_omni.diffusion.diffusion_kv.metadata import DiffusionKVMetadata from vllm_omni.diffusion.distributed.parallel_state import ( destroy_distributed_env, @@ -52,7 +54,13 @@ from vllm_omni.diffusion.ipc import DIFFUSION_RPC_RESULT_ENVELOPE, pack_diffusio from vllm_omni.diffusion.lora.manager import DiffusionLoRAManager, LoRABackend from vllm_omni.diffusion.registry import get_diffusion_ir_op_priority_func from vllm_omni.diffusion.request import OmniDiffusionRequest -from vllm_omni.diffusion.sched.interface import DiffusionSchedulerOutput, KVPrefetchJob +from vllm_omni.diffusion.sched.interface import ( + DiffusionSchedulerOutput, + KVPrefetchJob, + NewRequestData, + validate_new_request_data_identity, +) +from vllm_omni.diffusion.vllm_config import create_diffusion_vllm_config from vllm_omni.diffusion.worker.diffusion_model_runner import DiffusionModelRunner from vllm_omni.diffusion.worker.utils import BaseRunnerOutput, BatchRunnerOutput from vllm_omni.engine.stage_init_utils import set_death_signal @@ -76,49 +84,28 @@ _ASYNC_OUTPUT_DRAIN_TIMEOUT_S = 10.0 _MEMORY_RELEASING_METHODS = frozenset({"sleep", "handle_sleep_task"}) -@dataclass -class _DiffusionVllmModelConfig: - model: str - dtype: torch.dtype - quantization: str | None = None - quantization_config: Any | None = None - hf_config: Any | None = None - hf_text_config: Any | None = None - multimodal_config: Any | None = None - enforce_eager: bool = False - disable_cascade_attn: bool = False - enable_return_routed_experts: bool = False - is_moe: bool = False - - def is_quantized(self) -> bool: - return self.quantization is not None - - def is_model_moe(self) -> bool: - return self.is_moe - - def is_nvfp4_quantized(self) -> bool: - return self.quantization == "modelopt_fp4" - - @property - def is_diffusion(self) -> bool: - return False +def _all_gather_rank_values(value: Any) -> list[Any]: + if not dist.is_available() or not dist.is_initialized(): + return [value] + values: list[Any] = [None] * dist.get_world_size() + dist.all_gather_object(values, value) + return values -def _make_diffusion_vllm_model_config(od_config: OmniDiffusionConfig) -> _DiffusionVllmModelConfig: - quant_config = getattr(od_config, "quantization_config", None) - quantization = quant_config.get_name() if quant_config is not None and hasattr(quant_config, "get_name") else None - hf_config = getattr(od_config, "tf_model_config", None) - hf_text_config = get_hf_text_config(hf_config) if hasattr(hf_config, "get_text_config") else hf_config - return _DiffusionVllmModelConfig( - model=od_config.model, - dtype=od_config.dtype, - quantization=quantization, - quantization_config=quant_config, - hf_config=hf_config, - hf_text_config=hf_text_config, - enforce_eager=getattr(od_config, "enforce_eager", False), - is_moe=bool(getattr(od_config, "is_moe", False)), - ) +def _run_and_gather_rank_values(operation: str, func: Callable[[], Any]) -> list[Any]: + """Run one rank-local probe without stranding peers on local failure.""" + + try: + local_result = (True, func()) + except Exception as exc: + logger.exception("%s failed on this Worker rank", operation) + local_result = (False, f"{type(exc).__name__}: {exc}") + + rank_results = _all_gather_rank_values(local_result) + failures = [f"rank {rank}: {result}" for rank, (ok, result) in enumerate(rank_results) if not ok] + if failures: + raise RuntimeError(f"{operation} failed on " + "; ".join(failures)) + return [result for _, result in rank_results] @contextmanager @@ -152,37 +139,6 @@ def _force_cutlass_fp8_linear_kernel(quant_config: object | None) -> Iterator[No yield -def _is_unexpected_additional_config_type_error(exc: TypeError) -> bool: - """Return True only for constructor rejections of the additional_config kwarg.""" - message = str(exc) - return "unexpected keyword argument" in message and "additional_config" in message - - -def _create_diffusion_worker_vllm_config(device: torch.device, od_config: OmniDiffusionConfig) -> VllmConfig: - """Create a worker-local VllmConfig while preserving additional_config when supported.""" - config_kwargs: dict[str, Any] = { - "compilation_config": CompilationConfig(), - "device_config": DeviceConfig(device=device), - } - if od_config.additional_config: - config_kwargs["additional_config"] = od_config.additional_config - - try: - return VllmConfig(**config_kwargs) - except TypeError as exc: - if not _is_unexpected_additional_config_type_error(exc): - raise - - logger.debug("Worker-local VllmConfig does not accept additional_config in constructor: %s", exc) - config_kwargs.pop("additional_config", None) - vllm_config = VllmConfig(**config_kwargs) - try: - setattr(vllm_config, "additional_config", dict(od_config.additional_config)) - except Exception as set_exc: # pragma: no cover - defensive for older vLLM builds - logger.warning("Failed to attach additional_config to worker VllmConfig: %s", set_exc) - return vllm_config - - def _get_cumem_allocator_class() -> type: from vllm.device_allocator.cumem import CuMemAllocator @@ -223,6 +179,8 @@ class DiffusionWorker: self.device: torch.device | None = None self.vllm_config: VllmConfig | None = None self.model_runner: DiffusionModelRunner | None = None + self.init_snapshot: MemorySnapshot | None = None + self.requested_memory: int | None = None self._sleep_saved_buffers: dict[str, torch.Tensor] = {} self.lora_manager: DiffusionLoRAManager | None = None # Worker-side cache of (lora_request, lora_scale) per scheduled @@ -287,22 +245,7 @@ class DiffusionWorker: # Create vllm_config for parallel configuration. Pass explicit device_config # so DeviceConfig does not rely on current_platform in worker subprocesses. - vllm_config = _create_diffusion_worker_vllm_config(self.device, self.od_config) - parallel_config = self.od_config.parallel_config - vllm_config.parallel_config.tensor_parallel_size = parallel_config.tensor_parallel_size - vllm_config.parallel_config.data_parallel_size = parallel_config.data_parallel_size - if parallel_config.enable_expert_parallel and self.od_config.is_moe: - # Diffusion uses its own DP/CFG/SP groups normally. vLLM groups are - # only remapped for expert-parallel runtimes that consume vLLM's - # FusedMoE/EP semantics. - vllm_config.parallel_config.data_parallel_size = ( - parallel_config.data_parallel_size * parallel_config.cfg_parallel_size - ) - vllm_config.parallel_config.prefill_context_parallel_size = parallel_config.sequence_parallel_size - vllm_config.parallel_config.enable_expert_parallel = parallel_config.enable_expert_parallel - vllm_config.profiler_config = self.od_config.profiler_config - vllm_config.model_config = _make_diffusion_vllm_model_config(self.od_config) # type: ignore[assignment] - vllm_config.quant_config = self.od_config.quantization_config + vllm_config = create_diffusion_vllm_config(self.device, self.od_config) # Since vLLM v0.20.0, IR wraps GPU ops. Set IR op priority preference to enforce GPU op fusion during wrapping. # Also need to log, because vLLM internally logs another line in VllmConfig.__post_init__. Avoid confusion. vllm_config.kernel_config.ir_op_priority = _resolve_ir_op_priority(self.od_config, vllm_config) @@ -340,6 +283,23 @@ class DiffusionWorker: hsdp_replicate_size=parallel_config.hsdp_replicate_size if parallel_config.use_hsdp else 1, enable_expert_parallel=parallel_config.enable_expert_parallel, ) + if ( + getattr(self.od_config, "diffusion_kv_mode", DiffusionKVCacheMode.DENSE_LEGACY) + is DiffusionKVCacheMode.PAGED_SCHEDULER + ): + gc.collect() + current_omni_platform.empty_cache() + self.init_snapshot = MemorySnapshot(device=self.device) + self.requested_memory = request_memory( + self.init_snapshot, + vllm_config.cache_config, + ) + logger.debug( + "Worker %d: Diffusion KV initial memory snapshot: %r; requested=%s GiB", + self.rank, + self.init_snapshot, + format_gib(self.requested_memory), + ) init_workspace_manager(self.device) def _create_profiler(self) -> WorkerProfiler | None: @@ -393,6 +353,96 @@ class DiffusionWorker: if load_format != "dummy": assert self.model_runner.pipeline is not None + def get_kv_cache_specs(self) -> list[dict[str, KVCacheSpec]]: + """Return native rank-local specs for every diffusion Worker.""" + + assert self.model_runner is not None + return _run_and_gather_rank_values( + "Diffusion KV cache spec discovery", + self.model_runner.get_kv_cache_spec, + ) + + def determine_available_kv_memory(self, profile_requests: list[OmniDiffusionRequest]) -> list[int]: + """Profile and return each rank's safe Diffusion KV memory budget.""" + + def determine_local_memory() -> int: + assert self.vllm_config is not None + assert self.model_runner is not None + if self.init_snapshot is None or self.requested_memory is None: + raise RuntimeError("Diffusion KV memory snapshot was not captured before model loading") + override = self.vllm_config.cache_config.kv_cache_memory_bytes + if override: + # Match native vLLM: an explicit cache budget skips automatic + # capacity derivation, but still runs the maximum-shape model + # request so lazy kernels and communication buffers initialize. + self.model_runner.profile_run(profile_requests) + logger.info( + "Worker %d: Initial free memory %s GiB, reserved %s GiB memory for " + "Diffusion KV Cache as specified by kv_cache_memory_bytes config and " + "skipped automatic memory profiling. This does not respect the " + "gpu_memory_utilization config. A profile warmup was still executed.", + self.rank, + format_gib(self.init_snapshot.free_memory), + format_gib(int(override)), + ) + return int(override) + + with memory_profiling( + self.init_snapshot, + weights_memory=self.model_runner.model_memory_usage, + ) as profile_result: + self.model_runner.profile_run(profile_requests) + + available_memory = self.requested_memory - profile_result.non_kv_cache_memory + if available_memory <= 0: + raise RuntimeError( + "No memory remains for Diffusion KV cache after profiling: " + f"requested_memory={self.requested_memory} bytes, " + f"non_kv_cache_memory={profile_result.non_kv_cache_memory} bytes. " + "Increase gpu_memory_utilization or reduce the maximum profile request shape." + ) + free_gpu_memory = profile_result.after_profile.free_memory + unrequested_memory = self.init_snapshot.free_memory - self.requested_memory + logger.debug( + "Worker %d: Initial free memory: %s GiB; Requested memory: %f (util), %s GiB", + self.rank, + format_gib(self.init_snapshot.free_memory), + self.vllm_config.cache_config.gpu_memory_utilization, + format_gib(self.requested_memory), + ) + logger.debug( + "Worker %d: Free memory after profiling: %s GiB (total), %s GiB (within requested)", + self.rank, + format_gib(free_gpu_memory), + format_gib(free_gpu_memory - unrequested_memory), + ) + logger.debug("Worker %d: %r", self.rank, profile_result) + logger.info( + "Worker %d: Available Diffusion KV cache memory: %s GiB", + self.rank, + format_gib(available_memory), + ) + return int(available_memory) + + return [ + int(value) + for value in _run_and_gather_rank_values( + "Diffusion KV memory discovery", + determine_local_memory, + ) + ] + + def set_kv_cache_configs(self, kv_cache_configs: list[KVCacheConfig]) -> None: + """Install this rank's native config without allocating tensors yet.""" + + assert self.model_runner is not None + if len(kv_cache_configs) != self.od_config.num_gpus: + raise ValueError( + "Diffusion KVCacheConfig rank count mismatch: " + f"expected={self.od_config.num_gpus}, got={len(kv_cache_configs)}" + ) + self.model_runner.set_kv_cache_config(kv_cache_configs[self.rank]) + def init_lora_manager(self) -> None: """Initialize the LoRA manager for this worker.""" if self.model_runner.pipeline is None: @@ -465,18 +515,19 @@ class DiffusionWorker: def execute_model( self, - req: OmniDiffusionRequest | list[OmniDiffusionRequest], + req: OmniDiffusionRequest | list[NewRequestData], od_config: OmniDiffusionConfig, kv_prefetch_job: KVPrefetchJob | None = None, diffusion_kv_metadata: DiffusionKVMetadata | None = None, ) -> DiffusionOutput: """Execute a forward pass by delegating to the model runner. - If *req* is a list (DP multi-concurrency), each rank picks one - request based on its distributed rank. AllGather in the layerwise - offload only gathers weight shards (request-independent), so all - ranks stay synchronised at each AllGather call while computing - different activations. + If *req* is a list (DP multi-concurrency), each rank picks one complete + NewRequestData envelope based on its distributed rank. AllGather in + the layerwise offload only gathers weight shards (request-independent), + so all ranks stay synchronised at each AllGather call while computing + different activations. Selecting the envelope keeps Scheduler-issued + KV metadata bound to the request that owns its block tables. Each rank returns its OWN DiffusionOutput (no gather). The executor collects N responses via the per-worker result queues. @@ -492,7 +543,10 @@ class DiffusionWorker: dp_rank = get_data_parallel_rank() idx = dp_rank % len(req) - req = req[idx] + new_req = req[idx] + validate_new_request_data_identity(new_req) + req = new_req.req + diffusion_kv_metadata = new_req.diffusion_kv_metadata if self.lora_manager is not None: try: @@ -1443,7 +1497,7 @@ class WorkerWrapperBase: def execute_model( self, - req: OmniDiffusionRequest, + req: OmniDiffusionRequest | list[NewRequestData], od_config: OmniDiffusionConfig, kv_prefetch_job: KVPrefetchJob | None = None, diffusion_kv_metadata: DiffusionKVMetadata | None = None,