[Diffusion] Add native paged KV backend for diffusion workers (#6102)

Signed-off-by: Acerak01-fy <wfy2003324@163.com>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>

tests/diffusion/diffusion_kv/test_config.py

@@ -3,6 +3,7 @@
 
 import pytest
 
+from vllm_omni.config.omni_config import _DiffusionConfigProjection
 from vllm_omni.diffusion.data import OmniDiffusionConfig
 from vllm_omni.diffusion.diffusion_kv.config import DiffusionKVCacheMode
 
@@ -39,6 +40,8 @@ def test_paged_scheduler_rejects_dense_legacy_kv_receive() -> None:
     with pytest.raises(ValueError, match="does not support imported AR KV"):
         OmniDiffusionConfig.from_kwargs(
             diffusion_kv_mode="paged_scheduler",
+            diffusion_kv_max_rows_per_request=1,
+            max_num_batched_tokens=1,
             omni_kv_config={"need_recv_cache": True},
         )
 
@@ -47,6 +50,33 @@ def test_paged_scheduler_does_not_depend_on_model_registry() -> None:
     config = OmniDiffusionConfig.from_kwargs(
         model_class_name="FutureDiffusionModel",
         diffusion_kv_mode="paged_scheduler",
+        diffusion_kv_max_rows_per_request=1,
+        max_num_batched_tokens=1,
     )
 
     assert config.diffusion_kv_mode is DiffusionKVCacheMode.PAGED_SCHEDULER
+
+
+@pytest.mark.parametrize("config_cls", [OmniDiffusionConfig, _DiffusionConfigProjection])
+def test_paged_scheduler_mode_is_platform_agnostic(config_cls) -> None:
+    config = config_cls.from_kwargs(
+        diffusion_kv_mode="paged_scheduler",
+        diffusion_kv_max_rows_per_request=1,
+    )
+
+    assert config.diffusion_kv_mode is DiffusionKVCacheMode.PAGED_SCHEDULER
+
+
+def test_paged_scheduler_requires_a_worker_row_limit() -> None:
+    with pytest.raises(ValueError, match="requires diffusion_kv_max_rows_per_request"):
+        OmniDiffusionConfig.from_kwargs(diffusion_kv_mode="paged_scheduler")
+    with pytest.raises(ValueError, match="requires diffusion_kv_max_rows_per_request"):
+        _DiffusionConfigProjection.from_kwargs(diffusion_kv_mode="paged_scheduler")
+
+
+@pytest.mark.parametrize("invalid_limit", [0, -1, True, 1.5])
+def test_worker_row_limit_must_be_a_positive_integer(invalid_limit: object) -> None:
+    with pytest.raises((TypeError, ValueError)):
+        OmniDiffusionConfig.from_kwargs(diffusion_kv_max_rows_per_request=invalid_limit)
+    with pytest.raises((TypeError, ValueError)):
+        _DiffusionConfigProjection.from_kwargs(diffusion_kv_max_rows_per_request=invalid_limit)