chore: import upstream snapshot with attribution

python/ray/llm/tests/serve/cpu/configs/test_models.py

@@ -0,0 +1,560 @@
+import sys
+from pathlib import Path
+from unittest.mock import MagicMock, patch
+
+import pydantic
+import pytest
+
+from ray.llm._internal.common.utils.download_utils import NodeModelDownloadable
+from ray.llm._internal.serve.core.configs.accelerators import (
+    CPUAccelerator,
+    CPUConfig,
+    GPUAccelerator,
+    GPUConfig,
+    TPUAccelerator,
+    TPUConfig,
+)
+from ray.llm._internal.serve.core.configs.llm_config import (
+    LLMConfig,
+    LoraConfig,
+    ModelLoadingConfig,
+)
+from ray.llm._internal.serve.engines.vllm.vllm_models import VLLMEngineConfig
+
+CONFIG_DIRS_PATH = str(Path(__file__).parent / "configs")
+
+
+class TestModelConfig:
+    def test_construction(self):
+        """Test construct an LLMConfig doesn't error out and has correct attributes."""
+        llm_config = LLMConfig(
+            model_loading_config=ModelLoadingConfig(
+                model_id="llm_model_id",
+            ),
+            accelerator_type="A100-40G",  # Dash instead of underscore when specifying accelerator type
+            deployment_config={
+                "autoscaling_config": {
+                    "min_replicas": 3,
+                    "max_replicas": 7,
+                }
+            },
+        )
+        assert llm_config.deployment_config["autoscaling_config"]["min_replicas"] == 3
+        assert llm_config.deployment_config["autoscaling_config"]["max_replicas"] == 7
+        assert llm_config.model_loading_config.model_id == "llm_model_id"
+        assert llm_config.accelerator_type == "A100-40G"
+
+    def test_construction_requires_model_loading_config(self):
+        """Test that constructing an LLMConfig without model_loading_config errors out"""
+        with pytest.raises(
+            pydantic.ValidationError,
+        ):
+            LLMConfig(
+                accelerator_type="L4",
+            )
+
+    def test_accelerator_type_optional(self):
+        """Test that accelerator_type is optional when initializing LLMConfig."""
+        llm_config = LLMConfig(
+            model_loading_config=ModelLoadingConfig(model_id="test_model")
+        )
+        assert llm_config.model_loading_config.model_id == "test_model"
+        assert llm_config.accelerator_type is None
+
+    def test_invalid_accelerator_type(self):
+        """Test that invalid accelerator types raise validation errors."""
+        with pytest.raises(pydantic.ValidationError):
+            LLMConfig(
+                model_loading_config=ModelLoadingConfig(model_id="test_model"),
+                accelerator_type="INVALID_GPU",  # Invalid string value
+            )
+
+        # Test invalid numeric value
+        with pytest.raises(pydantic.ValidationError):
+            LLMConfig(
+                model_loading_config=ModelLoadingConfig(model_id="test_model"),
+                accelerator_type=123,  # Must be a string
+            )
+
+        # Test that underscore is not supported in accelerator type
+        with pytest.raises(pydantic.ValidationError):
+            LLMConfig(
+                model_loading_config=ModelLoadingConfig(model_id="test_model"),
+                accelerator_type="A100_40G",  # Should use A100-40G instead
+            )
+
+    def test_model_loading_config_forbids_extra_fields(self):
+        """Test that ModelLoadingConfig rejects extra fields."""
+
+        with pytest.raises(pydantic.ValidationError, match="engine_kwargs"):
+            ModelLoadingConfig(
+                model_id="test_model",
+                model_source="test_source",
+                engine_kwargs={"max_model_len": 8000},  # This should be rejected
+            )
+
+        valid_config = ModelLoadingConfig(
+            model_id="test_model", model_source="test_source"
+        )
+        assert valid_config.model_id == "test_model"
+        assert valid_config.model_source == "test_source"
+
+    def test_invalid_generation_config(self, disable_placement_bundles):
+        """Test that passing an invalid generation_config raises an error."""
+        with pytest.raises(
+            pydantic.ValidationError,
+        ):
+            LLMConfig(
+                model_loading_config=ModelLoadingConfig(model_id="test_model"),
+                accelerator_type="L4",
+                generation_config="invalid_config",  # Should be a dictionary, not a string
+            )
+
+    def test_deployment_type_checking(self, disable_placement_bundles):
+        """Test that deployment config type checking works."""
+        with pytest.raises(
+            pydantic.ValidationError,
+        ):
+            LLMConfig(
+                model_loading_config=ModelLoadingConfig(model_id="test_model"),
+                deployment_config={
+                    "max_ongoing_requests": -1,
+                },
+                accelerator_type="L4",
+            )
+
+    def test_autoscaling_type_checking(self, disable_placement_bundles):
+        """Test that autoscaling config type checking works."""
+        with pytest.raises(
+            pydantic.ValidationError,
+        ):
+            LLMConfig(
+                model_loading_config=ModelLoadingConfig(model_id="test_model"),
+                deployment_config={
+                    "autoscaling_config": {
+                        "min_replicas": -1,
+                    },
+                },
+                accelerator_type="L4",
+            )
+
+    def test_deployment_unset_fields_are_not_included(self, disable_placement_bundles):
+        """Test that unset fields are not included in the deployment config."""
+        llm_config = LLMConfig(
+            model_loading_config=ModelLoadingConfig(model_id="test_model"),
+            accelerator_type="L4",
+        )
+        assert "max_ongoing_requests" not in llm_config.deployment_config
+        assert "graceful_shutdown_timeout_s" not in llm_config.deployment_config
+
+    def test_autoscaling_unset_fields_are_not_included(self, disable_placement_bundles):
+        """Test that unset fields are not included in the autoscaling config."""
+        llm_config = LLMConfig(
+            model_loading_config=ModelLoadingConfig(model_id="test_model"),
+            deployment_config={
+                "autoscaling_config": {
+                    "min_replicas": 3,
+                    "max_replicas": 7,
+                },
+            },
+            accelerator_type="L4",
+        )
+        assert (
+            "metrics_interval_s"
+            not in llm_config.deployment_config["autoscaling_config"]
+        )
+        assert (
+            "upscaling_factor" not in llm_config.deployment_config["autoscaling_config"]
+        )
+
+    def test_engine_config_cached(self):
+        """Test that the engine config is cached and not recreated when calling
+        get_engine_config so the attributes on the engine will be persisted."""
+
+        llm_config = LLMConfig(
+            model_loading_config=ModelLoadingConfig(
+                model_id="llm_model_id",
+            ),
+        )
+        old_engine_config = llm_config.get_engine_config()
+        old_engine_config.hf_model_id = "fake_hf_model_id"
+        new_engine_config = llm_config.get_engine_config()
+        assert new_engine_config is old_engine_config
+
+    def test_remote_model_source_uses_model_id_as_hf_model_id(self):
+        """A remote model_source must not leak its URI into hf_model_id.
+
+        Using the URI verbatim propagates the scheme and slashes into the HF
+        cache directory name (e.g. ``models--s3:----bucket--...``). The URI
+        should instead be carried by mirror_config while hf_model_id falls back
+        to the user-supplied model_id.
+        """
+        bucket_uri = "s3://my-bucket/my-model"
+        llm_config = LLMConfig(
+            model_loading_config=ModelLoadingConfig(
+                model_id="llm_model_id",
+                model_source=bucket_uri,
+            ),
+        )
+        engine_config = llm_config.get_engine_config()
+        assert engine_config.hf_model_id == "llm_model_id"
+        assert engine_config.mirror_config is not None
+        assert engine_config.mirror_config.bucket_uri == bucket_uri
+
+    def test_hf_model_source_used_as_hf_model_id(self):
+        """A plain HuggingFace model_source is used directly as hf_model_id."""
+        llm_config = LLMConfig(
+            model_loading_config=ModelLoadingConfig(
+                model_id="llm_model_id",
+                model_source="facebook/opt-1.3b",
+            ),
+        )
+        engine_config = llm_config.get_engine_config()
+        assert engine_config.hf_model_id == "facebook/opt-1.3b"
+        assert engine_config.mirror_config is None
+
+    def test_no_model_source_falls_back_to_model_id(self):
+        """With no model_source, hf_model_id falls back to model_id."""
+        llm_config = LLMConfig(
+            model_loading_config=ModelLoadingConfig(
+                model_id="llm_model_id",
+            ),
+        )
+        engine_config = llm_config.get_engine_config()
+        assert engine_config.hf_model_id == "llm_model_id"
+        assert engine_config.mirror_config is None
+
+    def test_experimental_configs(self):
+        """Test that `experimental_configs` can be used."""
+        # Test with a valid dictionary can be used.
+        experimental_configs = {
+            "experimental_feature1": "value1",
+            "experimental_feature2": "value2",
+        }
+        llm_config = LLMConfig(
+            model_loading_config=ModelLoadingConfig(
+                model_id="llm_model_id",
+            ),
+            experimental_configs=experimental_configs,
+        )
+        assert llm_config.experimental_configs == experimental_configs
+
+        # test with invalid dictionary will raise a validation error.
+        with pytest.raises(
+            pydantic.ValidationError,
+        ):
+            LLMConfig(
+                model_loading_config=ModelLoadingConfig(
+                    model_id="llm_model_id",
+                ),
+                experimental_configs={123: "value1"},
+            )
+
+    def test_log_engine_metrics_disable_log_stats_validation(self):
+        """Test that log_engine_metrics=True prevents disable_log_stats=True."""
+        with pytest.raises(
+            pydantic.ValidationError,
+            match="disable_log_stats cannot be set to True when log_engine_metrics is enabled",
+        ):
+            LLMConfig(
+                model_loading_config=ModelLoadingConfig(model_id="test_model"),
+                log_engine_metrics=True,
+                engine_kwargs={"disable_log_stats": True},
+            )
+
+    @pytest.mark.parametrize(
+        "load_format,expected_download_model",
+        [
+            ("runai_streamer", NodeModelDownloadable.NONE),
+            ("runai_streamer_sharded", NodeModelDownloadable.NONE),
+            ("tensorizer", NodeModelDownloadable.NONE),
+            (None, NodeModelDownloadable.MODEL_AND_TOKENIZER),
+        ],
+    )
+    def test_load_format_callback_context(self, load_format, expected_download_model):
+        """Test that different load_format values set correct worker_node_download_model in callback context."""
+        engine_kwargs = {"load_format": load_format} if load_format is not None else {}
+
+        llm_config = LLMConfig(
+            model_loading_config=ModelLoadingConfig(model_id="test_model"),
+            engine_kwargs=engine_kwargs,
+        )
+
+        # Get the callback instance which should trigger the context setup
+        callback = llm_config.get_or_create_callback()
+
+        # Check that the callback context has the correct worker_node_download_model value
+        assert hasattr(callback, "ctx"), "Callback should have ctx attribute"
+        assert callback.ctx.worker_node_download_model == expected_download_model
+
+
+class TestFieldValidators:
+    """Test the field validators for dict validation."""
+
+    def test_model_loading_config_dict_validation(self):
+        """Test that model_loading_config accepts and validates dict input."""
+        config_dict = {"model_id": "microsoft/DialoGPT-medium"}
+
+        llm_config = LLMConfig(model_loading_config=config_dict, llm_engine="vLLM")
+
+        assert isinstance(llm_config.model_loading_config, ModelLoadingConfig)
+        assert llm_config.model_loading_config.model_id == "microsoft/DialoGPT-medium"
+
+    def test_model_loading_config_validation_error(self):
+        """Test that invalid dict raises proper validation error."""
+        with pytest.raises(pydantic.ValidationError) as exc_info:
+            LLMConfig(
+                model_loading_config={"invalid_field": "value"}, llm_engine="vLLM"
+            )
+
+        assert "Invalid model_loading_config" in str(exc_info.value)
+
+    def test_lora_config_dict_validation(self):
+        """Test that lora_config accepts and validates dict input."""
+        llm_config = LLMConfig(
+            model_loading_config={"model_id": "test"},
+            lora_config=None,
+            llm_engine="vLLM",
+        )
+
+        assert llm_config.lora_config is None
+
+        lora_dict = {
+            "dynamic_lora_loading_path": "s3://bucket/lora",
+            "max_num_adapters_per_replica": 8,
+        }
+
+        llm_config2 = LLMConfig(
+            model_loading_config={"model_id": "test"},
+            lora_config=lora_dict,
+            llm_engine="vLLM",
+        )
+
+        assert isinstance(llm_config2.lora_config, LoraConfig)
+        assert llm_config2.lora_config.max_num_adapters_per_replica == 8
+        assert llm_config2.lora_config.dynamic_lora_loading_path == "s3://bucket/lora"
+
+    def test_lora_config_validation_error(self):
+        """Test that invalid lora config dict raises proper validation error."""
+        with pytest.raises(pydantic.ValidationError) as exc_info:
+            LLMConfig(
+                model_loading_config={"model_id": "test"},
+                lora_config={"max_num_adapters_per_replica": "invalid_string"},
+                llm_engine="vLLM",
+            )
+
+        assert "Invalid lora_config" in str(exc_info.value)
+
+
+class TestAcceleratorConfigLogic:
+    """Test the accelerator_config logic and its interaction with accelerator_type."""
+
+    def test_accelerator_config_field_basic(self):
+        """Test that accelerator_config field works with basic values."""
+        # Test CPU config
+        llm_config_cpu = LLMConfig(
+            model_loading_config=ModelLoadingConfig(model_id="test_model"),
+            accelerator_config={"kind": "cpu"},
+        )
+        assert llm_config_cpu.accelerator_config.kind == "cpu"
+        engine_config = llm_config_cpu.get_engine_config()
+        assert engine_config.accelerator_config.kind == "cpu"
+
+        # Test GPU config
+        llm_config_gpu = LLMConfig(
+            model_loading_config=ModelLoadingConfig(model_id="test_model"),
+            accelerator_config={"kind": "gpu"},
+        )
+        assert llm_config_gpu.accelerator_config.kind == "gpu"
+        engine_config_gpu = llm_config_gpu.get_engine_config()
+        assert engine_config_gpu.accelerator_config.kind == "gpu"
+
+    def test_accelerator_type_with_cpu_config_raises_error(self):
+        """Test that accelerator_type with CPU config raises a validation error."""
+        with pytest.raises(
+            pydantic.ValidationError,
+            match="accelerator_type='L4' cannot be used with CPU-only configurations",
+        ):
+            LLMConfig(
+                model_loading_config=ModelLoadingConfig(model_id="test_model"),
+                accelerator_config={"kind": "cpu"},
+                accelerator_type="L4",
+            )
+
+    def test_accelerator_type_with_cpu_only_placement_group_raises_error(self):
+        """Test that accelerator_type with CPU-only placement_group_config raises error."""
+        with pytest.raises(
+            pydantic.ValidationError,
+            match="accelerator_type='L4' cannot be used with CPU-only configurations",
+        ):
+            LLMConfig(
+                model_loading_config=ModelLoadingConfig(model_id="test_model"),
+                accelerator_type="L4",
+                placement_group_config={"bundles": [{"CPU": 4}]},
+            )
+
+    def test_accelerator_type_with_empty_bundles_raises_error(self):
+        """Test that accelerator_type with empty bundles list raises error."""
+        with pytest.raises(
+            pydantic.ValidationError,
+            match="accelerator_type='L4' cannot be used with CPU-only configurations",
+        ):
+            LLMConfig(
+                model_loading_config=ModelLoadingConfig(model_id="test_model"),
+                accelerator_type="L4",
+                placement_group_config={"bundles": []},
+            )
+
+    def test_accelerator_type_with_gpu_placement_group_succeeds(self):
+        """Test that accelerator_type with GPU-containing placement_group_config succeeds."""
+        llm_config = LLMConfig(
+            model_loading_config=ModelLoadingConfig(model_id="test_model"),
+            accelerator_type="L4",
+            placement_group_config={"bundles": [{"GPU": 1, "CPU": 4}]},
+        )
+        assert llm_config.accelerator_type == "L4"
+
+    def test_accelerator_type_with_gpu_config_succeeds(self):
+        """Test that accelerator_type with GPU config succeeds."""
+        llm_config = LLMConfig(
+            model_loading_config=ModelLoadingConfig(model_id="test_model"),
+            accelerator_type="L4",
+            accelerator_config={"kind": "gpu"},
+        )
+        assert llm_config.accelerator_type == "L4"
+        engine_config = llm_config.get_engine_config()
+        assert engine_config.accelerator_type == "L4"
+
+    def test_vllm_engine_config_accelerator_type_with_cpu_config_raises_error(self):
+        """Test that VLLMEngineConfig rejects accelerator_type with CPU config."""
+        with pytest.raises(
+            pydantic.ValidationError,
+            match="accelerator_type='L4' cannot be used with CPU-only configurations",
+        ):
+            VLLMEngineConfig(
+                model_id="test-model",
+                accelerator_type="L4",
+                accelerator_config=CPUConfig(kind="cpu"),
+            )
+
+    def test_vllm_engine_config_accelerator_type_with_gpu_config_succeeds(self):
+        """Test that VLLMEngineConfig accepts accelerator_type with GPU config."""
+        engine_config = VLLMEngineConfig(
+            model_id="test-model",
+            accelerator_type="L4",
+            accelerator_config=GPUConfig(kind="gpu"),
+        )
+
+        assert engine_config.accelerator_type == "L4"
+
+    def test_llm_config_accelerator_type_hardware_mismatch(self):
+        """Test that passing a GPU accelerator_type with a TPU config raises an error."""
+        with pytest.raises(
+            pydantic.ValidationError,
+            match="Hardware mismatch",
+        ):
+            LLMConfig(
+                model_loading_config={"model_id": "test_model"},
+                accelerator_type="L4",
+                accelerator_config={"kind": "tpu", "topology": "4x4"},
+            )
+
+    def test_engine_config_infers_tpu_from_accelerator_type_string(self):
+        """Test that the engine config infers a TPU backend directly from the accelerator_type string."""
+        llm_config = LLMConfig(
+            model_loading_config=ModelLoadingConfig(model_id="test_model"),
+            accelerator_type="TPU-V6E",
+        )
+
+        # Validate engine correctly inferred the TPU backend
+        engine_config = llm_config.get_engine_config()
+
+        assert isinstance(engine_config.accelerator, TPUAccelerator)
+        assert engine_config.accelerator_type == "TPU-V6E"
+
+    def test_requires_deferred_placement_group(self):
+        """Test that requires_deferred_placement_group correctly identifies deferred PG requirements."""
+        cpu_accel = CPUAccelerator()
+        assert cpu_accel.requires_deferred_placement_group is False
+
+        gpu_accel = GPUAccelerator()
+        assert gpu_accel.requires_deferred_placement_group is False
+
+        tpu_accel_no_topo = TPUAccelerator(TPUConfig(kind="tpu"))
+        assert tpu_accel_no_topo.requires_deferred_placement_group is False
+
+        tpu_accel_with_topo = TPUAccelerator(TPUConfig(kind="tpu", topology="4x4"))
+        assert tpu_accel_with_topo.requires_deferred_placement_group is True
+
+    @pytest.mark.parametrize(
+        "topology,num_devices,accelerator_type_str,expected_bundles_count,expected_chips_per_host",
+        [
+            ("1x1", 1, "TPU-V6E", 1, 1),
+            ("1x1", 1, "TPU-V7X", 1, 1),
+            ("4x4", 16, "TPU-V6E", 4, 4),
+            ("2x2x2", 8, "TPU-V5P", 2, 4),
+            ("2x2", 4, "TPU-V5LITEPOD", 1, 4),
+            ("2x2x1", 4, "TPU-V4", 1, 4),
+            ("2x4", 8, "TPU-V6E", 1, 8),
+        ],
+    )
+    def test_default_bundles_topology(
+        self,
+        topology,
+        num_devices,
+        accelerator_type_str,
+        expected_bundles_count,
+        expected_chips_per_host,
+    ):
+        """Test that different topologies return correct per-host bundles."""
+        tpu_accel = TPUAccelerator(TPUConfig(kind="tpu", topology=topology))
+        bundles = tpu_accel.default_bundles(
+            num_devices=num_devices, accelerator_type_str=accelerator_type_str
+        )
+
+        assert len(bundles) == expected_bundles_count
+        for bundle in bundles:
+            assert bundle["TPU"] == expected_chips_per_host
+            assert f"accelerator_type:{accelerator_type_str}" in bundle
+
+    def test_default_bundles_topology_missing_accelerator_type_raises(self):
+        """Test that ValueError is raised when topology is present but accelerator type is missing."""
+        tpu_accel = TPUAccelerator(TPUConfig(kind="tpu", topology="4x4"))
+        with pytest.raises(
+            ValueError,
+            match="`accelerator_type` must be specified when `topology` is present",
+        ):
+            tpu_accel.default_bundles(num_devices=16, accelerator_type_str=None)
+
+    def test_default_bundles_topology_non_multiple_num_devices_raises(self):
+        """Test that ValueError is raised when num_devices is not a multiple of chips_per_host."""
+        tpu_accel = TPUAccelerator(TPUConfig(kind="tpu", topology="4x4"))
+        with pytest.raises(ValueError, match="must be a multiple of chips_per_host"):
+            tpu_accel.default_bundles(num_devices=6, accelerator_type_str="TPU-V6E")
+
+
+class TestCheckpointInfo:
+    def test_apply_checkpoint_info_uses_autoconfig_and_threads_trust_remote_code(self):
+        """apply_checkpoint_info uses AutoConfig (not PretrainedConfig) and forwards
+        trust_remote_code to every HF config load call."""
+        llm_config = LLMConfig(
+            model_loading_config=ModelLoadingConfig(model_id="test_model")
+        )
+        mock_hf_config = MagicMock(spec=["architectures", "vision_config"])
+        mock_hf_config.architectures = ["LlavaForCausalLM"]
+
+        with patch(
+            "transformers.AutoConfig.from_pretrained", return_value=mock_hf_config
+        ) as mock_auto:
+            llm_config.apply_checkpoint_info("vision/model", trust_remote_code=True)
+
+        assert all(
+            call.kwargs["trust_remote_code"] is True
+            for call in mock_auto.call_args_list
+        )
+        assert llm_config._supports_vision is True
+        assert llm_config._model_architecture == "LlavaForCausalLM"
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main(["-v", __file__]))

python/ray/llm/tests/serve/cpu/configs/test_multi_node_placement_groups.py

@@ -0,0 +1,486 @@
+from typing import Any, Dict
+
+import pydantic
+import pytest
+
+from ray.llm._internal.serve.core.server.llm_server import LLMServer
+from ray.llm._internal.serve.engines.vllm.vllm_models import VLLMEngineConfig
+from ray.llm._internal.serve.serving_patterns.data_parallel.dp_server import DPServer
+from ray.serve.llm import LLMConfig, ModelLoadingConfig
+
+
+def get_llm_config_with_placement_group(
+    tensor_parallel_size: int = 1,
+    pipeline_parallel_size: int = 1,
+    placement_group_config: Dict[str, Any] = None,
+) -> LLMConfig:
+    """Create LLMConfig with specified parallelism parameters and placement group config."""
+    return LLMConfig(
+        model_loading_config=ModelLoadingConfig(
+            model_id="test_model",
+            model_source="facebook/opt-1.3b",
+        ),
+        deployment_config=dict(
+            autoscaling_config=dict(
+                min_replicas=1,
+                max_replicas=1,
+            ),
+        ),
+        engine_kwargs=dict(
+            tensor_parallel_size=tensor_parallel_size,
+            pipeline_parallel_size=pipeline_parallel_size,
+            distributed_executor_backend="ray",
+        ),
+        placement_group_config=placement_group_config,
+        runtime_env=None,
+    )
+
+
+@pytest.mark.parametrize(
+    "tp_size,pp_size,placement_strategy",
+    [
+        (2, 4, "PACK"),  # Multi-node PP+TP with PACK
+        (4, 2, "PACK"),  # Multi-node PP+TP with PACK
+        (8, 1, "SPREAD"),  # Multi-node TP with SPREAD
+        (1, 8, "SPREAD"),  # Multi-node PP with SPREAD
+    ],
+)
+def test_llm_serve_custom_placement_group(tp_size, pp_size, placement_strategy):
+    """Test Ray Serve LLM with custom placement group configurations."""
+    total_gpus = tp_size * pp_size
+
+    # Create custom placement group configuration
+    placement_group_config = {
+        "bundles": [{"GPU": 1, "CPU": 1}] * total_gpus,
+        "strategy": placement_strategy,
+    }
+
+    llm_config = get_llm_config_with_placement_group(
+        tensor_parallel_size=tp_size,
+        pipeline_parallel_size=pp_size,
+        placement_group_config=placement_group_config,
+    )
+
+    # Verify the configuration is properly set
+    assert llm_config.placement_group_config == placement_group_config
+    assert llm_config.engine_kwargs["tensor_parallel_size"] == tp_size
+    assert llm_config.engine_kwargs["pipeline_parallel_size"] == pp_size
+
+    # Test that serve options are generated correctly
+    serve_options = LLMServer.get_deployment_options(llm_config)
+    assert "placement_group_bundles" in serve_options
+    assert "placement_group_strategy" in serve_options
+    assert serve_options["placement_group_strategy"] == placement_strategy
+    assert len(serve_options["placement_group_bundles"]) == total_gpus
+
+
+@pytest.mark.parametrize(
+    "tp_size,pp_size",
+    [
+        (2, 1),  # TP-only should use PACK by default
+        (1, 2),  # PP-only should use PACK by default
+        (2, 2),  # TP+PP should use PACK by default
+    ],
+)
+def test_llm_serve_default_placement_strategy(tp_size, pp_size):
+    """Test that Ray Serve LLM uses PACK strategy by default for all configurations."""
+    llm_config = get_llm_config_with_placement_group(
+        tensor_parallel_size=tp_size,
+        pipeline_parallel_size=pp_size,
+        placement_group_config=None,  # Use defaults
+    )
+
+    serve_options = LLMServer.get_deployment_options(llm_config)
+    # All configurations should default to PACK strategy
+    assert serve_options["placement_group_strategy"] == "PACK"
+    assert len(serve_options["placement_group_bundles"]) == tp_size * pp_size
+
+
+def test_llm_serve_placement_group_validation():
+    """Test validation of placement group configurations."""
+
+    # Test missing both bundle_per_worker and bundles
+    with pytest.raises(
+        ValueError, match="must specify either 'bundle_per_worker' or 'bundles'"
+    ):
+        llm_config = get_llm_config_with_placement_group(
+            placement_group_config={"strategy": "PACK"}
+        )
+        LLMServer.get_deployment_options(llm_config)
+
+    # Test missing strategy (should default to PACK, not fail)
+    llm_config = get_llm_config_with_placement_group(
+        placement_group_config={"bundles": [{"GPU": 1}]}
+    )
+    serve_options = LLMServer.get_deployment_options(llm_config)
+    assert serve_options["placement_group_strategy"] == "PACK"
+
+
+def test_llm_serve_multi_gpu_per_bundle_passes_through():
+    """Test multiple GPUs per bundle pass through Serve validation.
+
+    Serve allows GPU>1 per bundle in placement_group_config. vLLM will enforce
+    its own GPU<=1 restriction during engine creation (not tested here).
+    This confirms Serve doesn't block it, allowing vLLM to manage its constraints.
+    """
+    llm_config = get_llm_config_with_placement_group(
+        tensor_parallel_size=1,
+        pipeline_parallel_size=1,
+        placement_group_config={
+            "bundles": [{"GPU": 2, "CPU": 4}],
+            "strategy": "PACK",
+        },
+    )
+
+    # Serve should accept and pass through GPU=2 to placement group
+    # First bundle gets CPU: 4 (from config) + 1 (replica actor) = 5
+    serve_options = LLMServer.get_deployment_options(llm_config)
+    assert serve_options["placement_group_bundles"][0]["GPU"] == 2
+    assert serve_options["placement_group_bundles"][0]["CPU"] == 5
+
+    # vLLM will reject this during actual engine creation with a validation error
+    # (not tested here since this is a config-only CPU test)
+
+
+@pytest.mark.parametrize(
+    "tp_size,pp_size,expected_bundles",
+    [
+        (1, 1, 1),
+        (2, 1, 2),
+        (1, 2, 2),
+        (2, 2, 4),
+        (4, 2, 8),
+        (2, 4, 8),
+    ],
+)
+def test_llm_serve_bundle_count(tp_size, pp_size, expected_bundles):
+    """Test that correct number of bundles are created for different TP/PP configs."""
+    llm_config = get_llm_config_with_placement_group(
+        tensor_parallel_size=tp_size,
+        pipeline_parallel_size=pp_size,
+    )
+
+    serve_options = LLMServer.get_deployment_options(llm_config)
+    assert len(serve_options["placement_group_bundles"]) == expected_bundles
+
+
+def test_llm_serve_accelerator_and_resource_merging():
+    """Test accelerator type injection and replica actor resource merging."""
+    placement_group_config = {
+        "bundles": [{"GPU": 1, "CPU": 1}] * 2,
+        "strategy": "PACK",
+    }
+
+    llm_config = LLMConfig(
+        model_loading_config=ModelLoadingConfig(
+            model_id="test_model",
+            model_source="facebook/opt-1.3b",
+        ),
+        deployment_config=dict(
+            autoscaling_config=dict(min_replicas=1, max_replicas=1),
+            ray_actor_options=dict(
+                num_cpus=2,
+                num_gpus=1,
+                memory=1000000000,  # 1GB
+            ),
+        ),
+        engine_kwargs=dict(
+            tensor_parallel_size=2,
+            pipeline_parallel_size=1,
+            distributed_executor_backend="ray",
+        ),
+        accelerator_type="L4",
+        placement_group_config=placement_group_config,
+    )
+
+    serve_options = LLMServer.get_deployment_options(llm_config)
+
+    # First bundle: merged replica actor resources
+    # CPU: 1 (from bundle) + 2 (from replica actor) = 3
+    # GPU: Already 1 in both
+    first_bundle = serve_options["placement_group_bundles"][0]
+    assert first_bundle["CPU"] == 3
+    assert first_bundle["GPU"] == 2  # 1 from bundle + 1 from replica actor
+    assert "memory" in first_bundle
+    assert "accelerator_type:L4" in first_bundle
+
+    # Tail bundles: original config + accelerator type
+    for bundle in serve_options["placement_group_bundles"][1:]:
+        assert bundle["CPU"] == 1
+        assert bundle["GPU"] == 1
+        assert "accelerator_type:L4" in bundle
+        assert bundle["accelerator_type:L4"] == 0.001
+
+
+def test_llm_serve_data_parallel_placement_override():
+    """Test that data parallel deployments override placement group strategy to STRICT_PACK."""
+    placement_group_config = {
+        "bundles": [{"GPU": 1, "CPU": 1}] * 2,
+        "strategy": "SPREAD",  # This should be overridden
+    }
+
+    llm_config = LLMConfig(
+        model_loading_config=ModelLoadingConfig(
+            model_id="test_model",
+            model_source="facebook/opt-1.3b",
+        ),
+        # For DP correctness, do not set autoscaling_config; DP size fixes replicas
+        deployment_config=dict(),
+        engine_kwargs=dict(
+            tensor_parallel_size=2,
+            pipeline_parallel_size=1,
+            data_parallel_size=2,  # Enable data parallelism
+            distributed_executor_backend="ray",
+        ),
+        placement_group_config=placement_group_config,
+    )
+
+    serve_options = DPServer.get_deployment_options(llm_config)
+
+    # Data parallel should override to STRICT_PACK regardless of user-specified strategy
+    assert serve_options["placement_group_strategy"] == "STRICT_PACK"
+    # Note: num_replicas is set by build_dp_deployment, not by get_deployment_options
+
+
+def test_fractional_gpu_env_inferred_from_pg():
+    """A fractional placement group should inject VLLM_RAY_PER_WORKER_GPUS."""
+    placement_group_config = {"bundles": [{"GPU": 0.49}]}
+    llm_config = LLMConfig(
+        model_loading_config=ModelLoadingConfig(
+            model_id="test_model",
+            model_source="facebook/opt-125m",
+        ),
+        placement_group_config=placement_group_config,
+        runtime_env=dict(env_vars={"EXTRA_VAR": "1"}),
+    )
+
+    engine_config = VLLMEngineConfig.from_llm_config(llm_config)
+    runtime_env = engine_config.get_runtime_env_with_local_env_vars()
+
+    assert runtime_env["env_vars"]["EXTRA_VAR"] == "1"
+    assert runtime_env["env_vars"]["VLLM_RAY_PER_WORKER_GPUS"] == "0.49"
+
+
+def test_fractional_gpu_env_var_override_preserved():
+    """User-provided env var should be preserved when set."""
+    llm_config = LLMConfig(
+        model_loading_config=ModelLoadingConfig(
+            model_id="test_model",
+            model_source="facebook/opt-125m",
+        ),
+        placement_group_config={"bundles": [{"GPU": 0.4}]},
+        runtime_env=dict(
+            env_vars={
+                "VLLM_RAY_PER_WORKER_GPUS": "0.6",
+            }
+        ),
+    )
+
+    engine_config = VLLMEngineConfig.from_llm_config(llm_config)
+    runtime_env = engine_config.get_runtime_env_with_local_env_vars()
+
+    assert runtime_env["env_vars"]["VLLM_RAY_PER_WORKER_GPUS"] == "0.6"
+
+
+@pytest.mark.parametrize(
+    "tp_size,pp_size",
+    [
+        (2, 1),
+        (1, 2),
+        (2, 2),
+        (4, 2),
+    ],
+)
+def test_bundle_per_worker_expands_correctly(tp_size, pp_size):
+    """Test that bundle_per_worker is auto-replicated based on tp*pp."""
+    expected_bundles = tp_size * pp_size
+
+    llm_config = get_llm_config_with_placement_group(
+        tensor_parallel_size=tp_size,
+        pipeline_parallel_size=pp_size,
+        placement_group_config={"bundle_per_worker": {"GPU": 1, "CPU": 2}},
+    )
+
+    # Verify the configuration is properly set
+    assert llm_config.placement_group_config == {
+        "bundle_per_worker": {"GPU": 1, "CPU": 2}
+    }
+
+    # Test that serve options are generated correctly
+    serve_options = LLMServer.get_deployment_options(llm_config)
+    assert "placement_group_bundles" in serve_options
+    assert len(serve_options["placement_group_bundles"]) == expected_bundles
+
+    # Each bundle should have the specified resources (plus accelerator hint injection)
+    for i, bundle in enumerate(serve_options["placement_group_bundles"]):
+        # First bundle gets merged with replica actor resources
+        if i == 0:
+            assert bundle["GPU"] >= 1
+            assert bundle["CPU"] >= 2
+        else:
+            assert bundle["GPU"] == 1
+            assert bundle["CPU"] == 2
+
+
+def test_bundle_per_worker_conflict_with_bundles():
+    """Test that specifying both bundle_per_worker and bundles raises error."""
+    with pytest.raises(ValueError, match="Cannot specify both"):
+        llm_config = LLMConfig(
+            model_loading_config=ModelLoadingConfig(
+                model_id="test_model",
+                model_source="facebook/opt-125m",
+            ),
+            placement_group_config={
+                "bundle_per_worker": {"GPU": 1, "CPU": 1},
+                "bundles": [{"GPU": 1}],
+            },
+        )
+        # Validation happens when VLLMEngineConfig is created
+        VLLMEngineConfig.from_llm_config(llm_config)
+
+
+def test_bundle_per_worker_uses_pack_strategy():
+    """Test that bundle_per_worker defaults to PACK strategy."""
+    llm_config = get_llm_config_with_placement_group(
+        tensor_parallel_size=2,
+        pipeline_parallel_size=1,
+        placement_group_config={"bundle_per_worker": {"GPU": 1}},
+    )
+
+    serve_options = LLMServer.get_deployment_options(llm_config)
+    assert serve_options["placement_group_strategy"] == "PACK"
+
+
+def test_bundle_per_worker_injects_accelerator_type():
+    """Test that bundle_per_worker bundles get accelerator type hint injected."""
+    llm_config = LLMConfig(
+        model_loading_config=ModelLoadingConfig(
+            model_id="test_model",
+            model_source="facebook/opt-125m",
+        ),
+        engine_kwargs=dict(
+            tensor_parallel_size=2,
+            pipeline_parallel_size=1,
+            distributed_executor_backend="ray",
+        ),
+        accelerator_type="L4",
+        placement_group_config={"bundle_per_worker": {"GPU": 1, "CPU": 2}},
+    )
+
+    serve_options = LLMServer.get_deployment_options(llm_config)
+
+    # All bundles should have accelerator type hint injected
+    for bundle in serve_options["placement_group_bundles"]:
+        assert "accelerator_type:L4" in bundle
+        assert bundle["accelerator_type:L4"] == 0.001
+
+
+def test_bundle_per_worker_fractional_gpu_env_var():
+    """Test that fractional GPU in bundle_per_worker injects VLLM_RAY_PER_WORKER_GPUS."""
+    llm_config = LLMConfig(
+        model_loading_config=ModelLoadingConfig(
+            model_id="test_model",
+            model_source="facebook/opt-125m",
+        ),
+        placement_group_config={"bundle_per_worker": {"GPU": 0.5, "CPU": 1}},
+    )
+
+    engine_config = VLLMEngineConfig.from_llm_config(llm_config)
+    runtime_env = engine_config.get_runtime_env_with_local_env_vars()
+
+    assert runtime_env["env_vars"]["VLLM_RAY_PER_WORKER_GPUS"] == "0.5"
+
+
+def test_bundle_per_worker_non_fractional_gpu_no_env_var():
+    """Test that non-fractional GPU in bundle_per_worker does not inject env var."""
+    llm_config = LLMConfig(
+        model_loading_config=ModelLoadingConfig(
+            model_id="test_model",
+            model_source="facebook/opt-125m",
+        ),
+        placement_group_config={"bundle_per_worker": {"GPU": 1, "CPU": 1}},
+    )
+
+    engine_config = VLLMEngineConfig.from_llm_config(llm_config)
+    runtime_env = engine_config.get_runtime_env_with_local_env_vars()
+
+    assert "VLLM_RAY_PER_WORKER_GPUS" not in runtime_env.get("env_vars", {})
+
+
+def test_llm_serve_placement_group_explicit_none():
+    """Test that explicitly setting bundle_per_worker key to None does not crash."""
+    llm_config = LLMConfig(
+        model_loading_config=ModelLoadingConfig(
+            model_id="test_model",
+            model_source="facebook/opt-1.3b",
+        ),
+        placement_group_config={
+            "bundle_per_worker": None,
+            "bundles": [{"GPU": 1}],
+        },
+    )
+
+    # This should succeed fall back to the GPU bundles
+    serve_options = LLMServer.get_deployment_options(llm_config)
+    assert len(serve_options["placement_group_bundles"]) > 0
+
+
+class TestAcceleratorTypeValidation:
+    """Test accelerator_type validation with CPU-only configurations."""
+
+    def test_llm_config_accelerator_type_with_cpu_config_raises_error(self):
+        """LLMConfig raises error with accelerator_type and CPU config."""
+        with pytest.raises(
+            pydantic.ValidationError,
+            match="cannot be used with CPU-only configurations",
+        ):
+            LLMConfig(
+                model_loading_config={"model_id": "test_model"},
+                accelerator_type="L4",
+                accelerator_config={"kind": "cpu"},
+            )
+
+    def test_llm_config_accelerator_type_with_cpu_only_bundles_raises_error(self):
+        """LLMConfig raises error with accelerator_type and CPU-only bundles."""
+        with pytest.raises(
+            pydantic.ValidationError,
+            match="cannot be used with CPU-only configurations",
+        ):
+            LLMConfig(
+                model_loading_config={"model_id": "test_model"},
+                accelerator_type="L4",
+                placement_group_config={"bundles": [{"CPU": 4}]},
+            )
+
+    def test_llm_config_accelerator_type_with_empty_bundles_raises_error(self):
+        """LLMConfig raises error with accelerator_type and empty bundles."""
+        with pytest.raises(
+            pydantic.ValidationError,
+            match="cannot be used with CPU-only configurations",
+        ):
+            LLMConfig(
+                model_loading_config={"model_id": "test_model"},
+                accelerator_type="L4",
+                placement_group_config={"bundles": []},
+            )
+
+    def test_llm_config_accelerator_type_with_gpu_bundles_succeeds(self):
+        """Test that LLMConfig succeeds when accelerator_type is set with GPU bundles."""
+        config = LLMConfig(
+            model_loading_config={"model_id": "test_model"},
+            accelerator_type="L4",
+            placement_group_config={"bundles": [{"GPU": 1, "CPU": 4}]},
+        )
+        assert config.accelerator_type == "L4"
+
+    def test_llm_config_accelerator_type_default_uses_gpu(self):
+        """Test that LLMConfig with accelerator_type defaults to GPU."""
+        config = LLMConfig(
+            model_loading_config={"model_id": "test_model"},
+            accelerator_type="L4",
+        )
+        assert config.accelerator_type == "L4"
+
+
+if __name__ == "__main__":
+    pytest.main(["-v", __file__])

python/ray/llm/tests/serve/cpu/configs/test_openai_api_models_backends.py

@@ -0,0 +1,232 @@
+"""Tests for openai_api_models.py engine-agnostic import behaviour.
+
+SGLang fallback tests live in the llm_serve_sglang_e2e release test.
+"""
+
+import importlib
+import sys
+
+import pytest
+
+_OAI_MODELS_MOD = "ray.llm._internal.serve.core.configs.openai_api_models"
+
+
+class _VLLMImportBlocker:
+    """Meta-path finder that simulates vLLM not being installed.
+
+    Raises ModuleNotFoundError with .name set to mirror what Python raises
+    when a package is genuinely absent, so raise_llm_engine_import_error
+    can distinguish "not installed" from "installed but broken".
+    """
+
+    def find_spec(self, fullname, path=None, target=None):
+        if fullname == "vllm" or fullname.startswith("vllm."):
+            err = ModuleNotFoundError(f"Mocked: {fullname} is not installed")
+            err.name = fullname
+            raise err
+        return None
+
+
+class _VLLMBrokenInstallBlocker:
+    """Meta-path finder that simulates vLLM installed but broken at runtime
+    (e.g. missing libcudart.so or a missing transitive dependency).
+    """
+
+    def __init__(self, error: ImportError):
+        self._error = error
+
+    def find_spec(self, fullname, path=None, target=None):
+        if fullname == "vllm" or fullname.startswith("vllm."):
+            raise self._error
+        return None
+
+
+class TestVLLMBackend:
+    def test_wrapper_classes_inherit_from_vllm(self):
+        from ray.llm._internal.serve.core.configs.openai_api_models import (
+            ChatCompletionRequest,
+            CompletionRequest,
+            ErrorInfo,
+            ErrorResponse,
+        )
+
+        assert "vllm" in ChatCompletionRequest.__mro__[1].__module__
+        assert "vllm" in CompletionRequest.__mro__[1].__module__
+        assert "vllm" in ErrorInfo.__mro__[1].__module__
+        assert "vllm" in ErrorResponse.__mro__[1].__module__
+
+    def test_error_response_round_trip(self):
+        from ray.llm._internal.serve.core.configs.openai_api_models import (
+            ErrorInfo,
+            ErrorResponse,
+        )
+
+        info = ErrorInfo(message="something broke", code=500, type="InternalError")
+        resp = ErrorResponse(error=info)
+        assert resp.error.message == "something broke"
+        assert resp.error.code == 500
+        assert resp.error.type == "InternalError"
+
+        dumped = resp.model_dump()
+        assert dumped["error"]["message"] == "something broke"
+
+    def test_transcription_request_has_request_id(self):
+        from ray.llm._internal.serve.core.configs.openai_api_models import (
+            TranscriptionRequest,
+        )
+
+        assert "request_id" in TranscriptionRequest.model_fields
+
+    def _reload_oai_models_with_blocker(self, blocker):
+        """Helper: evict vllm + the target module, install blocker, reimport."""
+        saved = {
+            k: sys.modules.pop(k)
+            for k in list(sys.modules)
+            if k == "vllm" or k.startswith("vllm.")
+        }
+        sys.modules.pop(_OAI_MODELS_MOD, None)
+        sys.meta_path.insert(0, blocker)
+        try:
+            importlib.import_module(_OAI_MODELS_MOD)
+        finally:
+            sys.meta_path.remove(blocker)
+            sys.modules.pop(_OAI_MODELS_MOD, None)
+            sys.modules.update(saved)
+
+    def test_import_error_when_vllm_blocked(self):
+        """SGLang is not installed here either, so blocking vLLM means neither
+        backend is available."""
+        with pytest.raises(ImportError, match="Neither vLLM nor SGLang"):
+            self._reload_oai_models_with_blocker(_VLLMImportBlocker())
+
+    def test_vllm_installed_but_broken_cuda(self):
+        """Plain ImportError (e.g. missing libcudart.so) → clear message that
+        vLLM is installed but failed to load, not 'not installed'."""
+        cuda_err = ImportError(
+            "libcudart.so.12: cannot open shared object file: No such file or directory"
+        )
+        blocker = _VLLMBrokenInstallBlocker(cuda_err)
+        with pytest.raises(ImportError, match="vLLM is installed but failed to import"):
+            self._reload_oai_models_with_blocker(blocker)
+
+    def test_vllm_installed_but_missing_transitive_dep(self):
+        """ModuleNotFoundError for a *dependency* of vLLM (not vllm itself)
+        must also be reported as 'installed but broken', not 'not installed'."""
+        dep_err = ModuleNotFoundError("No module named 'msgpack'")
+        dep_err.name = "msgpack"
+        blocker = _VLLMBrokenInstallBlocker(dep_err)
+        with pytest.raises(ImportError, match="vLLM is installed but failed to import"):
+            self._reload_oai_models_with_blocker(blocker)
+
+
+class TestSanitizeChatCompletionRequest:
+    def test_serializes_tool_calls_iterator(self):
+        from ray.llm._internal.serve.core.configs.openai_api_models import (
+            ChatCompletionRequest,
+        )
+        from ray.llm._internal.serve.core.ingress.ingress import (
+            _sanitize_chat_completion_request,
+        )
+
+        tool_calls = [
+            {
+                "id": "call_1",
+                "type": "function",
+                "function": {"name": "f", "arguments": "{}"},
+            },
+        ]
+
+        request = ChatCompletionRequest(
+            model="test-model",
+            messages=[
+                {"role": "user", "content": "hello"},
+                {"role": "assistant", "content": None, "tool_calls": iter(tool_calls)},
+            ],
+        )
+
+        result = _sanitize_chat_completion_request(request)
+        assistant_msg = result.messages[1]
+        assert isinstance(assistant_msg["tool_calls"], list)
+        assert len(assistant_msg["tool_calls"]) == 1
+
+    def test_handles_no_tool_calls(self):
+        from ray.llm._internal.serve.core.configs.openai_api_models import (
+            ChatCompletionRequest,
+        )
+        from ray.llm._internal.serve.core.ingress.ingress import (
+            _sanitize_chat_completion_request,
+        )
+
+        request = ChatCompletionRequest(
+            model="test-model",
+            messages=[
+                {"role": "user", "content": "hello"},
+                {"role": "assistant", "content": "world"},
+            ],
+        )
+
+        result = _sanitize_chat_completion_request(request)
+        assert result is request
+        # Sanitizer must not inject tool_calls onto messages that never had them.
+        assistant_msg = result.messages[1]
+        if isinstance(assistant_msg, dict):
+            assert assistant_msg.get("tool_calls") is None
+        else:
+            assert getattr(assistant_msg, "tool_calls", None) is None
+
+    def test_serializes_content_iterator(self):
+        """When `content` is sent as a list of content parts on any message,
+        Pydantic stores it as a ValidatorIterator against the `Iterable[...]`
+        arm and the request becomes unpicklable until sanitized."""
+        import pickle
+
+        from ray.llm._internal.serve.core.configs.openai_api_models import (
+            ChatCompletionRequest,
+        )
+        from ray.llm._internal.serve.core.ingress.ingress import (
+            _sanitize_chat_completion_request,
+        )
+
+        request = ChatCompletionRequest(
+            model="test-model",
+            messages=[
+                {"role": "system", "content": [{"text": "sys", "type": "text"}]},
+                {"role": "user", "content": [{"text": "hi", "type": "text"}]},
+                {
+                    "role": "assistant",
+                    "content": [{"text": "step", "type": "text"}],
+                },
+                {
+                    "role": "tool",
+                    "content": [{"text": "r", "type": "text"}],
+                    "tool_call_id": "c0",
+                },
+            ],
+        )
+
+        result = _sanitize_chat_completion_request(request)
+        for msg in result.messages:
+            assert isinstance(msg, dict)
+            assert isinstance(msg["content"], list)
+            assert type(msg["content"]).__name__ != "ValidatorIterator"
+
+        pickle.dumps(result)
+
+
+class TestLLMServerLazyImport:
+    def test_default_engine_cls_is_none(self):
+        mod_name = "ray.llm._internal.serve.core.server.llm_server"
+        old_mod = sys.modules.pop(mod_name, None)
+        try:
+            mod = importlib.import_module(mod_name)
+            assert mod.LLMServer._default_engine_cls is None
+        finally:
+            if old_mod is not None:
+                sys.modules[mod_name] = old_mod
+
+    def test_ray_serve_llm_importable(self):
+        import ray.serve.llm  # noqa: F401
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main(["-v", __file__]))

python/ray/llm/tests/serve/cpu/deployments/conftest.py

@@ -0,0 +1,55 @@
+import pytest
+
+import ray
+from ray.tests.conftest import _ray_start_cluster
+
+
+@pytest.fixture
+def llm_config_with_mock_engine(llm_config):
+    # Make sure engine is mocked.
+    if llm_config.runtime_env is None:
+        llm_config.runtime_env = {}
+    llm_config.runtime_env.setdefault("env_vars", {})[
+        "RAYLLM_VLLM_ENGINE_CLS"
+    ] = "ray.llm.tests.serve.mocks.mock_vllm_engine.MockVLLMEngine"
+    yield llm_config
+
+
+@pytest.fixture(scope="module")
+def ray_tpu_cluster():
+    """
+    Simulates a Ray cluster with a multi-host TPU v6e-16 slice (4x4 topology).
+    """
+    pod_type = "v6e-16"
+    topology = "4x4"
+
+    with _ray_start_cluster() as cluster:
+        # A 4x4 v6e slice has 16 chips. We simulate 4 hosts with 4 chips each.
+        for i in range(4):
+            env_vars = {
+                "TPU_NAME": "test-slice",
+                "TPU_WORKER_ID": str(i),
+                "TPU_ACCELERATOR_TYPE": pod_type,
+                "TPU_TOPOLOGY": topology,
+            }
+            labels = {
+                "ray.io/tpu-slice-name": "test-slice",
+                "ray.io/tpu-worker-id": str(i),
+                "ray.io/tpu-pod-type": pod_type,
+            }
+            resources = {"TPU": 4, "accelerator_type:TPU-V6E": 4}
+
+            # The first node is the "head" of the slice
+            if i == 0:
+                resources[f"TPU-{pod_type}-head"] = 1
+
+            cluster.add_node(
+                num_cpus=8,
+                resources=resources,
+                labels=labels,
+                env_vars=env_vars,
+            )
+
+        ray.init(address=cluster.address)
+        yield cluster
+        ray.shutdown()

python/ray/llm/tests/serve/cpu/deployments/data_parallel/test_dp_direct_streaming.py

@@ -0,0 +1,66 @@
+import sys
+
+import pytest
+
+from ray.llm._internal.serve.core.configs.llm_config import (
+    LLMConfig,
+    ModelLoadingConfig,
+)
+from ray.llm._internal.serve.serving_patterns.data_parallel.builder import (
+    build_dp_openai_app,
+)
+from ray.llm.tests.serve.cpu.deployments.utils.direct_streaming_utils import (
+    consistent_hash_deployment_config,
+    requires_direct_streaming,
+    run_app_through_haproxy,
+    session_chat_response,
+)
+
+
+@requires_direct_streaming
+class TestDPDirectStreamingConsistentHashRouting:
+    """Session affinity over the DP direct-streaming path.
+
+    The DPServer is the ingress LLMRouter pins via ConsistentHashRouter, so a
+    request flows through HAProxy and the ``/internal/route`` decision to one
+    DPServer replica. The session id reaches the chosen replica, and one session
+    pins to one replica.
+    """
+
+    @pytest.fixture(name="llm_config")
+    def _llm_config(self):
+        return LLMConfig(
+            model_loading_config=ModelLoadingConfig(model_id="test-model"),
+            engine_kwargs={"data_parallel_size": 2},
+        )
+
+    @pytest.fixture(name="base_url")
+    def run_dp_app(
+        self,
+        llm_config_with_mock_engine,
+        shutdown_ray_and_serve,
+        disable_placement_bundles,
+    ):
+        llm_config = llm_config_with_mock_engine
+        llm_config.deployment_config = consistent_hash_deployment_config()
+        yield run_app_through_haproxy(build_dp_openai_app({"llm_config": llm_config}))
+
+    def test_session_affinity(self, base_url):
+        replicas = {
+            session_chat_response(base_url, "test-session-id").headers["x-replica-id"]
+            for _ in range(10)
+        }
+        assert len(replicas) == 1
+
+    def test_different_sessions_spread(self, base_url):
+        replicas = {
+            session_chat_response(base_url, f"test-session-id-{i}").headers[
+                "x-replica-id"
+            ]
+            for i in range(10)
+        }
+        assert len(replicas) > 1
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main(["-v", __file__]))

python/ray/llm/tests/serve/cpu/deployments/data_parallel/test_dp_server.py

@@ -0,0 +1,200 @@
+import asyncio
+import sys
+from copy import deepcopy
+from unittest.mock import patch
+
+import pytest
+
+from ray.llm._internal.serve.core.configs.llm_config import LLMConfig
+from ray.llm._internal.serve.serving_patterns.data_parallel.dp_server import (
+    DPServer,
+    GangMasterInfoRegistry,
+)
+from ray.serve.config import (
+    GangPlacementStrategy,
+    GangRuntimeFailurePolicy,
+    GangSchedulingConfig,
+)
+
+
+class TestGetDeploymentOptions:
+    """Mirrors test_dp_server.py but verifies gang scheduling config."""
+
+    @pytest.mark.parametrize(
+        "data_parallel_size,num_replicas",
+        [
+            (None, 1),
+            (2, None),
+            (1, 1),
+            (2, 4),
+        ],
+    )
+    def test_num_replicas_dp_validation(self, data_parallel_size, num_replicas):
+        engine_kwargs = (
+            {}
+            if data_parallel_size is None
+            else {"data_parallel_size": data_parallel_size}
+        )
+        deployment_config = (
+            {} if num_replicas is None else {"num_replicas": num_replicas}
+        )
+        llm_config = LLMConfig(
+            model_loading_config=dict(model_id="test_model"),
+            engine_kwargs=deepcopy(engine_kwargs),
+            deployment_config=deepcopy(deployment_config),
+        )
+
+        opts = DPServer.get_deployment_options(llm_config)
+        dp_size = data_parallel_size or 1
+        if dp_size > 1:
+            expected_replicas = (
+                num_replicas * dp_size if num_replicas is not None else dp_size
+            )
+            assert opts["num_replicas"] == expected_replicas
+            assert isinstance(opts["gang_scheduling_config"], GangSchedulingConfig)
+            assert opts["gang_scheduling_config"].gang_size == dp_size
+            assert (
+                opts["gang_scheduling_config"].gang_placement_strategy
+                == GangPlacementStrategy.PACK
+            )
+            assert (
+                opts["gang_scheduling_config"].runtime_failure_policy
+                == GangRuntimeFailurePolicy.RESTART_GANG
+            )
+        else:
+            assert "gang_scheduling_config" not in opts
+
+    def test_autoscaling_config(self):
+        llm_config = LLMConfig(
+            model_loading_config=dict(model_id="test_model"),
+            engine_kwargs={"data_parallel_size": 4},
+            deployment_config={
+                "autoscaling_config": {
+                    "target_ongoing_requests": 10,
+                    "min_replicas": 2,
+                    "max_replicas": 8,
+                    "initial_replicas": 3,
+                }
+            },
+        )
+        opts = DPServer.get_deployment_options(llm_config)
+        assert isinstance(opts["gang_scheduling_config"], GangSchedulingConfig)
+        assert opts["gang_scheduling_config"].gang_size == 4
+        # Autoscaling config should have min/max/initial replicas multiplied by dp_size
+        autoscaling_config = opts["autoscaling_config"]
+        assert autoscaling_config["target_ongoing_requests"] == 10
+        assert autoscaling_config["min_replicas"] == 2 * 4
+        assert autoscaling_config["max_replicas"] == 8 * 4
+        assert autoscaling_config["initial_replicas"] == 3 * 4
+
+
+class TestGangMasterInfoRegistry:
+    _KV_MODULE = "ray.llm._internal.serve.serving_patterns.data_parallel.dp_server"
+
+    def _make_kv_store(self):
+        # Mocks GCS KV store
+        store = {}
+        return (
+            store,
+            lambda key, value, overwrite=False: store.__setitem__(key, value),
+            lambda key: store.get(key),
+            lambda key: store.pop(key, None) is not None,
+            lambda key: key in store,
+        )
+
+    @patch(f"{_KV_MODULE}._internal_kv_get")
+    @patch(f"{_KV_MODULE}._internal_kv_put")
+    def test_get_timeout(self, mock_put, mock_get):
+        mock_get.return_value = None
+        with pytest.raises(TimeoutError, match="Timed out"):
+            asyncio.get_event_loop().run_until_complete(
+                GangMasterInfoRegistry.get(
+                    "gang-missing", timeout=0.5, poll_interval=0.1
+                )
+            )
+
+    @patch(f"{_KV_MODULE}._internal_kv_get")
+    @patch(f"{_KV_MODULE}._internal_kv_put")
+    def test_gang_isolation(self, mock_put, mock_get):
+        _, fake_put, fake_get, _, _ = self._make_kv_store()
+        mock_put.side_effect = fake_put
+        mock_get.side_effect = fake_get
+
+        GangMasterInfoRegistry.register("gang-1", "10.0.0.1", 1111, "node-1")
+        GangMasterInfoRegistry.register("gang-2", "10.0.0.2", 2222, "node-2")
+
+        loop = asyncio.get_event_loop()
+        addr1, port1, node1 = loop.run_until_complete(
+            GangMasterInfoRegistry.get("gang-1")
+        )
+        addr2, port2, node2 = loop.run_until_complete(
+            GangMasterInfoRegistry.get("gang-2")
+        )
+
+        assert (addr1, port1, node1) == ("10.0.0.1", 1111, "node-1")
+        assert (addr2, port2, node2) == ("10.0.0.2", 2222, "node-2")
+
+
+class TestBundleIndices:
+    @pytest.mark.parametrize(
+        "engine_kwargs,placement_group_config,dp_rank,sorted_indices,expected",
+        [
+            # TP=1: 1 bundle per replica, identity ordering
+            ({"tensor_parallel_size": 1}, None, 0, list(range(4)), "0"),
+            ({"tensor_parallel_size": 1}, None, 3, list(range(4)), "3"),
+            (
+                {"tensor_parallel_size": 1},
+                {"bundles": [{"GPU": 1, "CPU": 1}]},
+                2,
+                list(range(4)),
+                "2",
+            ),
+            # TP=2: 2 bundles per replica, identity ordering
+            ({"tensor_parallel_size": 2}, None, 0, list(range(8)), "0,1"),
+            ({"tensor_parallel_size": 2}, None, 2, list(range(8)), "4,5"),
+            (
+                {"tensor_parallel_size": 2},
+                {"bundles": [{"GPU": 1, "CPU": 1}, {"GPU": 1}]},
+                1,
+                list(range(4)),
+                "2,3",
+            ),
+            # TP=2, PP=2: 4 bundles per replica, identity ordering
+            (
+                {"tensor_parallel_size": 2, "pipeline_parallel_size": 2},
+                None,
+                0,
+                list(range(8)),
+                "0,1,2,3",
+            ),
+            (
+                {"tensor_parallel_size": 2, "pipeline_parallel_size": 2},
+                None,
+                1,
+                list(range(8)),
+                "4,5,6,7",
+            ),
+            # Out-of-order sorted_indices: bundles reordered by node
+            ({"tensor_parallel_size": 2}, None, 1, [0, 2, 1, 3], "1,3"),
+            ({"tensor_parallel_size": 1}, None, 0, [2, 0, 3, 1], "2"),
+        ],
+    )
+    def test_bundle_indices(
+        self, engine_kwargs, placement_group_config, dp_rank, sorted_indices, expected
+    ):
+        llm_config = LLMConfig(
+            model_loading_config=dict(model_id="test_model"),
+            engine_kwargs=engine_kwargs,
+            placement_group_config=placement_group_config,
+        )
+        engine_config = llm_config.get_engine_config()
+        bundles_per_replica = len(engine_config.placement_bundles)
+
+        result = DPServer._compute_bundle_indices(
+            dp_rank, bundles_per_replica, sorted_indices
+        )
+        assert result == expected
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main(["-v", __file__]))

python/ray/llm/tests/serve/cpu/deployments/llm/multiplex/test_lora_deployment_base_client.py

@@ -0,0 +1,251 @@
+import sys
+from copy import deepcopy
+from typing import List, Set
+
+import pytest
+from fastapi import HTTPException
+
+from ray import serve
+from ray.llm._internal.serve.core.configs.openai_api_models import (
+    ModelCard,
+    to_model_metadata,
+)
+from ray.llm._internal.serve.core.server.llm_server import LLMServer
+from ray.llm.tests.serve.mocks.mock_vllm_engine import MockVLLMEngine
+from ray.serve.handle import DeploymentHandle
+from ray.serve.llm import LLMConfig, LoraConfig
+from ray.serve.llm.ingress import OpenAiIngress, make_fastapi_ingress
+
+VLLM_APP_DEF = """
+model_loading_config:
+  model_id: meta-llama/Llama-2-7b-hf
+
+llm_engine: vLLM
+
+engine_kwargs:
+  trust_remote_code: True
+  max_model_len: 4096
+  tensor_parallel_size: 1
+
+accelerator_type: A10G
+
+deployment_config:
+  autoscaling_config:
+    min_replicas: 1
+    initial_replicas: 1
+    max_replicas: 8
+    target_ongoing_requests: 5
+    metrics_interval_s: 10.0
+    look_back_period_s: 30.0
+    smoothing_factor: 1.0
+    downscale_delay_s: 300.0
+    upscale_delay_s: 60.0
+  max_ongoing_requests: 15
+
+"""
+
+
+VLLM_APP = LLMConfig.parse_yaml(VLLM_APP_DEF)
+
+
+# TODO (shrekris): add test for querying fine-tuned weights stored in the
+# cloud.
+
+
+def get_mocked_llm_deployments(llm_configs) -> List[DeploymentHandle]:
+    llm_deployments = []
+    for llm_config in llm_configs:
+        deployment_args = LLMServer.get_deployment_options(llm_config)
+        deployment = serve.deployment(LLMServer).options(**deployment_args)
+        llm_deployments.append(
+            deployment.bind(
+                llm_config=llm_config,
+                engine_cls=MockVLLMEngine,
+            )
+        )
+    return llm_deployments
+
+
+def make_ingress_app(llm_deployments, llm_configs, **kwargs):
+    ingress_options = OpenAiIngress.get_deployment_options(llm_configs)
+    ingress_cls = make_fastapi_ingress(OpenAiIngress)
+    deployments_by_id = {c.model_id: d for c, d in zip(llm_configs, llm_deployments)}
+    model_cards = {c.model_id: to_model_metadata(c.model_id, c) for c in llm_configs}
+    lora_paths = {
+        c.model_id: c.lora_config.dynamic_lora_loading_path
+        for c in llm_configs
+        if c.lora_config is not None
+    }
+    return (
+        serve.deployment(ingress_cls)
+        .options(**ingress_options)
+        .bind(
+            llm_deployments=deployments_by_id,
+            model_cards=model_cards,
+            lora_paths=lora_paths,
+            **kwargs,
+        )
+    )
+
+
+@pytest.mark.asyncio
+async def test_lora_unavailable_base_model(
+    shutdown_ray_and_serve, disable_placement_bundles
+):
+    """Getting the handle for an unavailable model should return a 404."""
+    llm_config = VLLM_APP.model_copy(deep=True)
+    llm_deployments = get_mocked_llm_deployments([llm_config])
+    app = make_ingress_app(llm_deployments, llm_configs=[llm_config])
+    router_handle = serve.run(app)
+
+    with pytest.raises(HTTPException) as e:
+        await router_handle._get_configured_serve_handle.remote("anyscale-lora")
+
+    assert e.value.status_code == 404
+
+
+@pytest.mark.asyncio
+async def test_lora_get_model(shutdown_ray_and_serve, disable_placement_bundles):
+    """Test behavior when getting a LoRA model."""
+
+    base_model_id = "meta-llama/Llama-2-7b-hf"
+
+    llm_config = VLLM_APP.model_copy(deep=True)
+    llm_config.model_loading_config.model_id = base_model_id
+    llm_deployments = get_mocked_llm_deployments([llm_config])
+    app = make_ingress_app(llm_deployments, llm_configs=[llm_config])
+    router_handle = serve.run(app)
+
+    # Case 1: model does not exist.
+    not_found_config = await router_handle.model.remote("not_found")
+    assert not_found_config is None
+
+    # Case 2: Model has only the base model config.
+    base_model_config = await router_handle.model.remote(base_model_id)
+    assert isinstance(base_model_config, ModelCard)
+    base_model_data = base_model_config.model_dump()
+    assert base_model_data["id"] == base_model_id
+    base_model_config = base_model_data["metadata"]
+
+    # Case 3: model has a multiplex config in the cloud.
+    llm_config = VLLM_APP.model_copy(deep=True)
+    llm_config.lora_config = LoraConfig(dynamic_lora_loading_path="s3://base_path")
+    lora_model = "meta-llama/Llama-2-7b-hf:suffix:1234"
+    llm_deployments = get_mocked_llm_deployments([llm_config])
+
+    async def fake_get_lora_model_metadata(*args, **kwargs):
+        return {
+            "model_id": lora_model,
+            "base_model_id": base_model_id,
+            "max_request_context_length": 4096,
+        }
+
+    app = make_ingress_app(
+        llm_deployments,
+        llm_configs=[llm_config],
+        _get_lora_model_metadata_func=fake_get_lora_model_metadata,
+    )
+    router_handle = serve.run(app)
+
+    lora_model_config = await router_handle.model.remote(lora_model)
+    assert isinstance(lora_model_config, ModelCard)
+    lora_model_data = lora_model_config.model_dump()
+    assert lora_model_data["id"] == lora_model
+    lora_metadata = lora_model_data["metadata"]
+    assert lora_metadata["model_id"] == lora_model
+    assert lora_metadata["base_model_id"] == base_model_id
+    assert lora_metadata["max_request_context_length"] == 4096
+
+
+@pytest.mark.asyncio
+async def test_lora_list_base_model(shutdown_ray_and_serve, disable_placement_bundles):
+    """Test model-listing behavior when only the base model is available."""
+    base_model_id = "base_model"
+    llm_config = VLLM_APP.model_copy(deep=True)
+    llm_config.model_loading_config.model_id = base_model_id
+    llm_deployments = get_mocked_llm_deployments([llm_config])
+    app = make_ingress_app(llm_deployments, llm_configs=[llm_config])
+    router_handle = serve.run(app)
+
+    models = (await router_handle.models.remote()).data
+    assert len(models) == 1
+
+    base_model = models[0]
+    base_model_data = base_model.model_dump()
+    assert base_model_data["id"] == base_model_id
+
+
+@pytest.mark.parametrize(
+    ("dynamic_lora_loading_path", "base_model_id", "expected_model_ids"),
+    [
+        # Case 1: test a path that exists in the cloud. The LoRA adapters
+        # must be included.
+        (
+            "s3://anonymous@air-example-data/rayllm-ossci/lora-checkpoints/meta-llama/Llama-2-7b-chat-hf",
+            "meta-llama/Llama-2-7b-chat-hf",
+            [
+                "meta-llama/Llama-2-7b-chat-hf:gen-config-but-no-context-len:1234",
+                "meta-llama/Llama-2-7b-chat-hf:with-context-len-and-gen-config:1234",
+                "meta-llama/Llama-2-7b-chat-hf:long-context-model:1234",
+                "meta-llama/Llama-2-7b-chat-hf",
+            ],
+        ),
+        # Case 2: test a path with the same model provider (meta-llama in this
+        # case). But test a different model. Ensure that only this model's
+        # LoRA adapters are returned.
+        (
+            "s3://anonymous@air-example-data/rayllm-ossci/lora-checkpoints/meta-llama/Llama-2-13b-chat-hf",
+            "meta-llama/Llama-2-13b-chat-hf",
+            [
+                "meta-llama/Llama-2-13b-chat-hf:pre-long-context-model:1234",
+                "meta-llama/Llama-2-13b-chat-hf",
+            ],
+        ),
+        # Case 3: test a path that doesn't exist in the cloud. Only the
+        # base model_id should be included.
+        (
+            "s3://anonymous@air-example-data/rayllm-ossci/path-does-not-exist/",
+            "meta-llama/Llama-2-7b-chat-hf",
+            ["meta-llama/Llama-2-7b-chat-hf"],
+        ),
+    ],
+)
+@pytest.mark.asyncio
+async def test_lora_include_adapters_in_list_models(
+    shutdown_ray_and_serve,
+    disable_placement_bundles,
+    dynamic_lora_loading_path: str,
+    base_model_id: str,
+    expected_model_ids: List[str],
+):
+    """Check that LoRA adapters are included in the models list.
+
+    This test pulls real configs from an S3 bucket located in
+    `anyscale-legacy-work` account.
+
+    This test is similar to test_lora_list_base_model. It checks that
+    the LoRA adapters are included in the list of models.
+    """
+    config = deepcopy(VLLM_APP)
+    config.model_loading_config.model_id = base_model_id
+    config.lora_config = LoraConfig(dynamic_lora_loading_path=dynamic_lora_loading_path)
+
+    llm_deployments = get_mocked_llm_deployments([config])
+    app = make_ingress_app(llm_deployments, llm_configs=[config])
+    router_handle = serve.run(app)
+
+    models = (await router_handle.models.remote()).data
+    assert {model.id for model in models} == set(expected_model_ids)
+
+    # Confirm that all expected model IDs exist.
+    expected_model_ids_set: Set[str] = set(expected_model_ids)
+    for model in models:
+        model_data = model.model_dump()
+        assert model_data["id"] in expected_model_ids_set
+        expected_model_ids_set.discard(model_data["id"])
+
+    assert len(expected_model_ids_set) == 0
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main(["-v", __file__]))

python/ray/llm/tests/serve/cpu/deployments/llm/multiplex/test_lora_model_loader.py

@@ -0,0 +1,202 @@
+import asyncio
+import sys
+from unittest.mock import Mock, patch
+
+import pytest
+
+from ray.llm._internal.common.utils.cloud_utils import LoraMirrorConfig
+from ray.llm._internal.serve.core.configs.llm_config import (
+    LLMConfig,
+    LLMEngine,
+    LoraConfig,
+    ModelLoadingConfig,
+)
+from ray.llm._internal.serve.utils.lora_serve_utils import LoraModelLoader
+
+
+class TestLoRAModelLoader:
+    """Test suite for the LoraModelLoader class."""
+
+    @pytest.fixture
+    def model_loader(self):
+        """Provides a LoraModelLoader instance for tests."""
+        return LoraModelLoader("/tmp/ray/lora/cache", max_tries=3)
+
+    @pytest.fixture
+    def llm_config(self, disable_placement_bundles):
+        """Common LLM config used across tests."""
+        return LLMConfig(
+            model_loading_config=ModelLoadingConfig(model_id="llm_model_id"),
+            llm_engine=LLMEngine.vLLM,
+            accelerator_type="L4",
+            lora_config=LoraConfig(
+                dynamic_lora_loading_path="s3://fake-bucket-uri-abcd"
+            ),
+        )
+
+    @pytest.fixture
+    def lora_model_id(self):
+        """Common LoRA model ID used across tests."""
+        return "base_model:lora_id"
+
+    @pytest.fixture
+    def lora_mirror_config(self, lora_model_id):
+        """Common LoRA mirror config used across tests."""
+        return LoraMirrorConfig(
+            lora_model_id=lora_model_id,
+            bucket_uri="s3://fake-bucket-uri-abcd",
+            max_total_tokens=4096,
+        )
+
+    @pytest.mark.asyncio
+    async def test_basic_loading(
+        self, model_loader, llm_config, lora_model_id, lora_mirror_config
+    ):
+        """Test basic model loading functionality."""
+        # Create a simple mock for sync_model
+        mock_sync_model = Mock()
+
+        with patch(
+            "ray.llm._internal.serve.utils.lora_serve_utils.sync_files_with_lock",
+            side_effect=mock_sync_model,
+        ):
+            # First load should download the model
+            disk_multiplex_config = await model_loader.load_model(
+                lora_model_id=lora_model_id,
+                lora_mirror_config=lora_mirror_config,
+            )
+
+            # Verify sync_files_with_lock was called with correct parameters
+            mock_sync_model.assert_called_once_with(
+                "s3://fake-bucket-uri-abcd",
+                "/tmp/ray/lora/cache/lora_id",
+                timeout=model_loader.download_timeout_s,
+            )
+            mock_sync_model.reset_mock()
+
+            # Second time we don't load from S3 - should use cache
+            new_disk_config = await model_loader.load_model(
+                lora_model_id=lora_model_id,
+                lora_mirror_config=lora_mirror_config,
+            )
+            assert new_disk_config == disk_multiplex_config
+            mock_sync_model.assert_not_called()
+
+    @pytest.mark.asyncio
+    async def test_retry_logic(
+        self, model_loader, llm_config, lora_model_id, lora_mirror_config
+    ):
+        """Test that the lora model load task is properly retried on failure."""
+        # Counter to track number of sync_model calls
+        attempt_count = 0
+
+        # Create a mock for sync_files_with_lock that tracks calls and fails initially
+        def mock_sync_model(bucket_uri, local_path, timeout=None):
+            nonlocal attempt_count
+            attempt_count += 1
+
+            # Fail on first attempt, succeed on second
+            if attempt_count == 1:
+                raise RuntimeError("Simulated download failure")
+            # Success on subsequent attempts
+            return None
+
+        with patch(
+            "ray.llm._internal.serve.utils.lora_serve_utils.sync_files_with_lock",
+            side_effect=Mock(side_effect=mock_sync_model),
+        ):
+            # First load should trigger a retry
+            disk_multiplex_config = await model_loader.load_model(
+                lora_model_id=lora_model_id,
+                lora_mirror_config=lora_mirror_config,
+            )
+
+            # Verify retry happened exactly once
+            assert attempt_count == 2
+
+            # Reset counter
+            attempt_count = 0
+
+            # Load again (should use cache, no download attempts)
+            new_disk_config = await model_loader.load_model(
+                lora_model_id=lora_model_id,
+                lora_mirror_config=lora_mirror_config,
+            )
+
+            # Verify no new download attempts
+            assert attempt_count == 0
+
+            # Verify cached config is returned
+            assert new_disk_config == disk_multiplex_config
+
+    @pytest.mark.asyncio
+    async def test_concurrent_loading(
+        self, model_loader, llm_config, lora_model_id, lora_mirror_config
+    ):
+        """Test that concurrent loads only trigger one download process."""
+        # Counter to track number of sync_model calls
+        attempt_count = 0
+
+        # Create a mock for sync_files_with_lock that tracks calls and fails initially
+        def mock_sync_model(bucket_uri, local_path, timeout=None):
+            nonlocal attempt_count
+            attempt_count += 1
+
+            # Fail on first attempt, succeed on second
+            if attempt_count == 1:
+                raise RuntimeError("Simulated download failure")
+            # Success on subsequent attempts
+            return None
+
+        with patch(
+            "ray.llm._internal.serve.utils.lora_serve_utils.sync_files_with_lock",
+            side_effect=Mock(side_effect=mock_sync_model),
+        ):
+            # Clear cache to force download
+            model_loader.disk_cache.clear()
+
+            # Create multiple concurrent tasks
+            tasks = [
+                asyncio.create_task(
+                    model_loader.load_model(
+                        lora_model_id=lora_model_id,
+                        lora_mirror_config=lora_mirror_config,
+                    )
+                )
+                for _ in range(3)
+            ]
+
+            # Wait for all tasks to complete
+            results = await asyncio.gather(*tasks)
+
+            # Verify retry happened exactly once across all tasks
+            assert attempt_count == 2
+
+            # All tasks should return the same result
+            assert all(result == results[0] for result in results)
+
+    @pytest.mark.asyncio
+    async def test_max_retries_exhaustion(
+        self, model_loader, llm_config, lora_model_id, lora_mirror_config
+    ):
+        """Test that an error is raised when max retries are exhausted."""
+        # Mock that always fails
+        def mock_sync_model_always_fails(*args, **kwargs):
+            raise RuntimeError("Simulated persistent failure")
+
+        with patch(
+            "ray.llm._internal.serve.utils.lora_serve_utils.sync_files_with_lock",
+            side_effect=Mock(side_effect=mock_sync_model_always_fails),
+        ):
+            # Should fail after max_tries (3) attempts
+            with pytest.raises(RuntimeError) as excinfo:
+                await model_loader.load_model(
+                    lora_model_id=lora_model_id,
+                    lora_mirror_config=lora_mirror_config,
+                )
+
+            assert "Simulated persistent failure" in str(excinfo.value)
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main(["-v", __file__]))

python/ray/llm/tests/serve/cpu/deployments/llm/multiplex/test_multiplex_utils.py

@@ -0,0 +1,165 @@
+import sys
+from unittest.mock import Mock, call, patch
+
+import pytest
+
+from ray.llm._internal.common.utils.lora_utils import (
+    retry_with_exponential_backoff,
+)
+
+
+def test_retry_success_first_try():
+    """Test that the function works normally when no exceptions occur."""
+    # Use a simple counter to track calls instead of a mock
+
+    def success_function(*args, **kwargs):
+        return "success"
+
+    # Apply our retry decorator
+    decorated_fn = retry_with_exponential_backoff(
+        max_tries=3, exception_to_check=ValueError
+    )(success_function)
+
+    # Call the decorated function
+    result = decorated_fn("arg1", "arg2", kwarg1="kwarg1")
+
+    # Verify it returned the expected result
+    assert result == "success"
+
+
+def test_retry_success_after_retries():
+    """Test that the function retries and eventually succeeds."""
+    # Create a mock that raises ValueError twice then succeeds
+    mock_fn = Mock(side_effect=[ValueError("error1"), ValueError("error2"), "success"])
+
+    # Apply our retry decorator with a small delay for faster testing
+    decorated_fn = retry_with_exponential_backoff(
+        max_tries=3, exception_to_check=ValueError, base_delay=0.01
+    )(mock_fn)
+
+    # Call the decorated function
+    result = decorated_fn()
+
+    # Verify it returned the expected result
+    assert result == "success"
+
+    # Verify the mock was called three times
+    assert mock_fn.call_count == 3
+
+
+def test_retry_exhaustion():
+    """Test that the function gives up after max_tries."""
+    # Create a mock that always raises ValueError
+    error = ValueError("persistent error")
+    mock_fn = Mock(side_effect=error)
+
+    # Apply our retry decorator with a small delay for faster testing
+    decorated_fn = retry_with_exponential_backoff(
+        max_tries=3, exception_to_check=ValueError, base_delay=0.01
+    )(mock_fn)
+
+    # Call the decorated function and expect it to raise ValueError
+    with pytest.raises(ValueError) as excinfo:
+        decorated_fn()
+
+    # Verify the error is the same one our mock raised
+    assert excinfo.value is error
+
+    # Verify the mock was called max_tries times
+    assert mock_fn.call_count == 3
+
+
+def test_retry_wrong_exception():
+    """Test that the function doesn't retry for non-matching exceptions."""
+    # Create a mock that raises TypeError
+    error = TypeError("wrong error type")
+    mock_fn = Mock(side_effect=error)
+
+    # Apply our retry decorator
+    decorated_fn = retry_with_exponential_backoff(
+        max_tries=3, exception_to_check=ValueError
+    )(mock_fn)
+
+    # Call the decorated function and expect it to raise TypeError immediately
+    with pytest.raises(TypeError) as excinfo:
+        decorated_fn()
+
+    # Verify the error is the same one our mock raised
+    assert excinfo.value is error
+
+    # Verify the mock was called only once
+    mock_fn.assert_called_once()
+
+
+def test_retry_backoff_timing():
+    """Test that the function backs off with the expected delays."""
+    # Create a mock that always raises ValueError
+    mock_fn = Mock(side_effect=ValueError("error"))
+
+    # Patch time.sleep so we can verify the delay
+    with patch("time.sleep") as mock_sleep:
+        # Apply our retry decorator with specific backoff parameters
+        decorated_fn = retry_with_exponential_backoff(
+            max_tries=4,
+            exception_to_check=ValueError,
+            base_delay=1,
+            max_delay=8,
+            exponential_base=2,
+        )(mock_fn)
+
+        # Call the decorated function and expect it to raise ValueError
+        with pytest.raises(ValueError):
+            decorated_fn()
+
+        # Verify sleep was called with the expected delays
+        # First retry: 1s, Second retry: 2s, Third retry: 4s
+        mock_sleep.assert_has_calls([call(1), call(2), call(4)])
+
+
+def test_retry_max_delay():
+    """Test that the delay is capped at max_delay."""
+    # Create a mock that always raises ValueError
+    mock_fn = Mock(side_effect=ValueError("error"))
+
+    # Patch time.sleep so we can verify the delay
+    with patch("time.sleep") as mock_sleep:
+        # Apply our retry decorator with max_delay=3
+        decorated_fn = retry_with_exponential_backoff(
+            max_tries=4,
+            exception_to_check=ValueError,
+            base_delay=2,
+            max_delay=3,
+            exponential_base=2,
+        )(mock_fn)
+
+        # Call the decorated function and expect it to raise ValueError
+        with pytest.raises(ValueError):
+            decorated_fn()
+
+        # Verify sleep was called with delays capped at max_delay
+        # First retry: 2s, Second retry: 3s (not 4s), Third retry: 3s (not 8s)
+        mock_sleep.assert_has_calls([call(2), call(3), call(3)])
+
+
+def test_retry_preserves_function_metadata():
+    """Test that the decorator preserves the function's metadata."""
+    # Define a function with docstring and name
+    def test_function(x, y):
+        """Test function docstring."""
+        return x + y
+
+    # Apply our retry decorator
+    decorated_fn = retry_with_exponential_backoff(
+        max_tries=3, exception_to_check=ValueError
+    )(test_function)
+
+    # Verify the metadata is preserved
+    assert decorated_fn.__name__ == "test_function"
+    assert decorated_fn.__doc__ == "Test function docstring."
+
+    # Verify the function still works correctly
+    assert decorated_fn(1, 2) == 3
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main(["-v", __file__]))

python/ray/llm/tests/serve/cpu/deployments/llm/test_builder_llm_server.py

@@ -0,0 +1,123 @@
+import sys
+
+import pytest
+
+from ray import serve
+from ray.llm._internal.serve.constants import DEFAULT_MAX_TARGET_ONGOING_REQUESTS
+from ray.llm._internal.serve.core.configs.llm_config import (
+    LLMConfig,
+    ModelLoadingConfig,
+)
+from ray.llm._internal.serve.core.server.builder import (
+    build_llm_deployment,
+)
+
+
+class TestBuildVllmDeployment:
+    def test_build_llm_deployment(
+        self,
+        llm_config_with_mock_engine,
+        shutdown_ray_and_serve,
+        disable_placement_bundles,
+    ):
+        """Test `build_llm_deployment` can build a vLLM deployment."""
+
+        app = build_llm_deployment(llm_config_with_mock_engine)
+        assert isinstance(app, serve.Application)
+        handle = serve.run(app)
+        assert handle.deployment_name.startswith("LLMServer")
+
+    def test_build_llm_deployment_with_name_prefix(
+        self,
+        llm_config_with_mock_engine,
+        shutdown_ray_and_serve,
+        disable_placement_bundles,
+    ):
+        """Test `build_llm_deployment` can build a vLLM deployment with name prefix."""
+
+        _name_prefix_for_test = "test_name_prefix"
+        app = build_llm_deployment(
+            llm_config_with_mock_engine, name_prefix=_name_prefix_for_test
+        )
+        assert isinstance(app, serve.Application)
+        handle = serve.run(app)
+        assert handle.deployment_name.startswith(_name_prefix_for_test)
+
+    def test_build_llm_deployment_name_prefix_along_with_deployment_config(
+        self,
+        llm_config_with_mock_engine,
+        shutdown_ray_and_serve,
+        disable_placement_bundles,
+    ):
+        """Test `build_llm_deployment` can build a vLLM deployment with name prefix and deployment config."""
+
+        config_with_name: LLMConfig = llm_config_with_mock_engine.model_copy(deep=True)
+        _deployment_name = "deployment_name_from_config"
+        _name_prefix_for_test = "test_name_prefix"
+        config_with_name.deployment_config["name"] = _deployment_name
+        app = build_llm_deployment(config_with_name, name_prefix=_name_prefix_for_test)
+        assert isinstance(app, serve.Application)
+        handle = serve.run(app)
+        assert handle.deployment_name == _name_prefix_for_test + _deployment_name
+
+    def test_default_autoscaling_config_included_without_num_replicas(
+        self, disable_placement_bundles
+    ):
+        """Test that default autoscaling_config with target_ongoing_requests is included
+        when num_replicas is not specified.
+        """
+        llm_config = LLMConfig(
+            model_loading_config=ModelLoadingConfig(model_id="test-model"),
+        )
+        app = build_llm_deployment(llm_config)
+
+        deployment = app._bound_deployment
+        autoscaling_config = deployment._deployment_config.autoscaling_config
+        assert autoscaling_config is not None
+        assert (
+            autoscaling_config.target_ongoing_requests
+            == DEFAULT_MAX_TARGET_ONGOING_REQUESTS
+        )
+
+    def test_autoscaling_config_removed_from_defaults_when_num_replicas_specified(
+        self, disable_placement_bundles
+    ):
+        """Test that autoscaling_config from defaults is removed when user specifies
+        num_replicas, since Ray Serve does not allow both.
+        """
+        llm_config = LLMConfig(
+            model_loading_config=ModelLoadingConfig(model_id="test-model"),
+            deployment_config={
+                "num_replicas": 2,
+            },
+        )
+        app = build_llm_deployment(llm_config)
+
+        deployment = app._bound_deployment
+        assert deployment._deployment_config.num_replicas == 2
+        # autoscaling_config should be None since num_replicas is set
+        assert deployment._deployment_config.autoscaling_config is None
+
+    def test_user_target_ongoing_requests_respected(self, disable_placement_bundles):
+        """Test that user-specified target_ongoing_requests is respected and not
+        overridden by defaults.
+        """
+        user_target = 50
+        llm_config = LLMConfig(
+            model_loading_config=ModelLoadingConfig(model_id="test-model"),
+            deployment_config={
+                "autoscaling_config": {
+                    "target_ongoing_requests": user_target,
+                },
+            },
+        )
+        app = build_llm_deployment(llm_config)
+
+        deployment = app._bound_deployment
+        autoscaling_config = deployment._deployment_config.autoscaling_config
+        assert autoscaling_config is not None
+        assert autoscaling_config.target_ongoing_requests == user_target
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main(["-v", __file__]))

python/ray/llm/tests/serve/cpu/deployments/llm/test_llm_engine.py

@@ -0,0 +1,187 @@
+"""This tests the LLM engine by testing the mocked implementations directly.
+
+This implicitly tests the consistency of the engine API through time.
+Also tests that our Mock is behaving as expected to ensure that the downstream tests using Mocks are correct from Mock implementation perspective.
+
+
+We have the following Mock:
+
+- An engine that returns a string of form "test_i" for i in range(max_tokens)
+"""
+
+import sys
+from typing import Optional
+
+import pytest
+
+from ray.llm.tests.serve.mocks.mock_vllm_engine import MockVLLMEngine
+from ray.llm.tests.serve.utils.testing_utils import LLMResponseValidator
+
+
+class TestMockLLMEngine:
+    @pytest.mark.parametrize("api_type", ["chat", "completion"])
+    @pytest.mark.parametrize("stream", [False, True])
+    @pytest.mark.parametrize("max_tokens", [5])
+    @pytest.mark.asyncio
+    async def test_unified_llm_engine(
+        self,
+        mock_llm_config,
+        mock_chat_request,
+        mock_completion_request,
+        api_type: str,
+        stream: bool,
+        max_tokens: int,
+    ):
+        """Unified test for both chat and completion APIs, streaming and non-streaming."""
+        # Create and start the engine
+        engine = MockVLLMEngine(mock_llm_config)
+        await engine.start()
+
+        # Create request based on API type
+        if api_type == "chat":
+            request = mock_chat_request
+            response_generator = engine.chat(request)
+        elif api_type == "completion":
+            request = mock_completion_request
+            response_generator = engine.completions(request)
+
+        print(
+            f"\n\n_____ {api_type.upper()} ({'STREAMING' if stream else 'NON-STREAMING'}) max_tokens={max_tokens} _____\n\n"
+        )
+
+        if stream:
+            # Collect streaming chunks
+            chunks = []
+            async for chunk in response_generator:
+                assert isinstance(chunk, str)
+                chunks.append(chunk)
+
+            # Validate streaming response
+            LLMResponseValidator.validate_streaming_chunks(chunks, api_type, max_tokens)
+        else:
+            # Validate non-streaming response
+            async for response in response_generator:
+                LLMResponseValidator.validate_non_streaming_response(
+                    response, api_type, max_tokens
+                )
+
+    @pytest.mark.parametrize("dimensions", [None, 512])
+    @pytest.mark.asyncio
+    async def test_embedding_mock_engine(
+        self, mock_llm_config, mock_embedding_request, dimensions: Optional[int]
+    ):
+        """Test embedding API with different dimensions."""
+        # Create and start the engine
+        engine = MockVLLMEngine(mock_llm_config)
+        await engine.start()
+
+        # Create embedding request
+        request = mock_embedding_request
+
+        print(f"\n\n_____ EMBEDDING dimensions={dimensions} _____\n\n")
+
+        async for response in engine.embeddings(request):
+            LLMResponseValidator.validate_embedding_response(response, dimensions)
+
+    @pytest.mark.parametrize("stream", [False, True])
+    @pytest.mark.parametrize("temperature", [0.0])
+    @pytest.mark.parametrize("language", ["en", "hi"])
+    @pytest.mark.asyncio
+    async def test_transcription_mock_engine(
+        self,
+        mock_llm_config,
+        mock_transcription_request,
+        stream: bool,
+        temperature: float,
+        language: Optional[str],
+    ):
+        """Test transcription API with different language and temperature, streaming and non-streaming."""
+
+        engine = MockVLLMEngine(mock_llm_config)
+        await engine.start()
+
+        request = mock_transcription_request
+        response_generator = engine.transcriptions(request)
+
+        print(
+            f"\n\n_____ TRANSCRIPTION ({'STREAMING' if stream else 'NON-STREAMING'}) language={language} temperature={temperature} _____\n\n"
+        )
+
+        if stream:
+            # Collect streaming chunks
+            chunks = []
+            async for chunk in response_generator:
+                assert isinstance(chunk, str)
+                chunks.append(chunk)
+
+            # Validate streaming response
+            LLMResponseValidator.validate_transcription_response(
+                chunks, temperature, language
+            )
+        else:
+            # Validate non-streaming response
+            async for response in response_generator:
+                LLMResponseValidator.validate_transcription_response(
+                    response, temperature, language
+                )
+
+    @pytest.mark.asyncio
+    async def test_score_mock_engine(self, mock_llm_config, mock_score_request):
+        """Test score API for text similarity."""
+        # Create and start the engine
+        engine = MockVLLMEngine(mock_llm_config)
+        await engine.start()
+
+        # Create score request
+        request = mock_score_request
+
+        print("\n\n_____ SCORE _____\n\n")
+
+        async for response in engine.score(request):
+            LLMResponseValidator.validate_score_response(response)
+
+    @pytest.mark.parametrize("return_token_strs", [False, True])
+    @pytest.mark.asyncio
+    async def test_tokenize_mock_engine(
+        self, mock_llm_config, mock_tokenize_request, return_token_strs: bool
+    ):
+        """Test tokenize API."""
+        # Create and start the engine
+        engine = MockVLLMEngine(mock_llm_config)
+        await engine.start()
+
+        # Create tokenize request
+        request = mock_tokenize_request
+
+        print(f"\n\n_____ TOKENIZE return_token_strs={return_token_strs} _____\n\n")
+
+        async for response in engine.tokenize(request):
+            LLMResponseValidator.validate_tokenize_response(
+                response,
+                expected_prompt="Hello, world!",
+                return_token_strs=return_token_strs,
+            )
+
+    @pytest.mark.asyncio
+    async def test_detokenize_mock_engine(
+        self, mock_llm_config, mock_detokenize_request
+    ):
+        """Test detokenize API."""
+        # Create and start the engine
+        engine = MockVLLMEngine(mock_llm_config)
+        await engine.start()
+
+        # Create detokenize request
+        request = mock_detokenize_request
+
+        print("\n\n_____ DETOKENIZE _____\n\n")
+
+        async for response in engine.detokenize(request):
+            LLMResponseValidator.validate_detokenize_response(
+                response,
+                expected_text="Hello",  # [72, 101, 108, 108, 111] = "Hello"
+            )
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main(["-v", __file__]))

python/ray/llm/tests/serve/cpu/deployments/llm/test_llm_engine_tpu.py

@@ -0,0 +1,368 @@
+import sys
+
+import pytest
+
+import ray
+from ray import serve
+from ray.llm._internal.serve.core.configs.accelerators import (
+    CPUAccelerator,
+    CPUConfig,
+    GPUAccelerator,
+    GPUConfig,
+    TPUAccelerator,
+    TPUConfig,
+)
+from ray.llm._internal.serve.core.server.llm_server import LLMServer
+from ray.llm.tests.serve.mocks.mock_vllm_engine import PGCreationMockEngine
+from ray.serve.llm import LLMConfig, ModelLoadingConfig
+from ray.util.placement_group import PlacementGroup, placement_group_table
+
+
+def test_tpu_slice_placement_group_creation_default_resources(ray_tpu_cluster):
+    """
+    Verifies that requesting a multi-host TPU topology correctly intercepts
+    standard PG creation and returns a PACK SlicePlacementGroup.
+    """
+    llm_config = LLMConfig(
+        model_loading_config=ModelLoadingConfig(model_id="test-tpu-model"),
+        accelerator_type="TPU-V6E",
+        accelerator_config=TPUConfig(kind="tpu", topology="4x4"),
+        engine_kwargs={"tensor_parallel_size": 16},
+    )
+
+    engine_config = llm_config.get_engine_config()
+
+    pg = None
+    try:
+        pg = engine_config.get_or_create_pg()
+
+        assert isinstance(pg, PlacementGroup)
+
+        pg_table = placement_group_table(pg)
+        assert pg_table["strategy"] == "PACK"
+
+        # 4x4 v6e = 16 chips. We default to 4 TPU chips per bundle (per-host).
+        assert len(pg_table["bundles"]) == 4
+        for bundle in pg_table["bundles"].values():
+            assert "TPU" in bundle
+            assert bundle["TPU"] == 4.0
+    finally:
+        # Let the backend tear down its own resources if it has any
+        engine_config.accelerator.shutdown()
+        if pg is not None:
+            try:
+                ray.util.remove_placement_group(pg)
+            except Exception:
+                pass
+
+
+def test_tpu_slice_placement_group_creation_host_resources(ray_tpu_cluster):
+    """
+    Verifies that explicitly providing host-level bundles via
+    placement_group_config correctly overrides the 1-chip default.
+    """
+    llm_config = LLMConfig(
+        model_loading_config=ModelLoadingConfig(model_id="test-tpu-model"),
+        accelerator_type="TPU-V6E",
+        accelerator_config=TPUConfig(kind="tpu", topology="4x4"),
+        placement_group_config={
+            "strategy": "STRICT_SPREAD",
+            "bundles": [{"TPU": 4}] * 4,
+        },
+    )
+
+    engine_config = llm_config.get_engine_config()
+
+    pg = None
+    try:
+        pg = engine_config.get_or_create_pg()
+
+        assert isinstance(pg, PlacementGroup)
+
+        pg_table = placement_group_table(pg)
+        assert pg_table["strategy"] == "STRICT_SPREAD"
+        # We should provision 4 host-level bundles instead of the default 16 chip-level bundles.
+        assert len(pg_table["bundles"]) == 4
+        for bundle in pg_table["bundles"].values():
+            assert "TPU" in bundle
+            assert bundle["TPU"] == 4
+    finally:
+        # Let the backend tear down its own resources if it has any
+        engine_config.accelerator.shutdown()
+        if pg is not None:
+            try:
+                ray.util.remove_placement_group(pg)
+            except Exception:
+                pass
+
+
+def test_single_tpu_fallback(ray_tpu_cluster):
+    """
+    Verifies that requesting a TPU without a topology gracefully
+    falls back to standard single-host bundle packing.
+    """
+    llm_config = LLMConfig(
+        model_loading_config=ModelLoadingConfig(model_id="test-tpu-model"),
+        accelerator_type="TPU-V6E",
+    )
+
+    engine_config = llm_config.get_engine_config()
+    pg = engine_config.get_or_create_pg()
+
+    pg_table = placement_group_table(pg)
+
+    # Verify it falls back to the default PACK strategy for 1 GPU/TPU
+    assert len(pg_table["bundles"]) == 1
+    assert pg_table["strategy"] == "PACK"
+
+    # Let the backend tear down its own resources if it has any
+    engine_config.accelerator.shutdown()
+    try:
+        ray.util.remove_placement_group(pg)
+    except Exception:
+        pass  # Already cleaned up by the wrapper
+
+
+def test_tpu_slice_placement_group_creation_bundle_per_worker(ray_tpu_cluster):
+    """
+    Verifies that specifying bundle_per_worker correctly expands to bundles,
+    includes the accelerator hint for TPU, and correctly identifies TPU usage.
+    """
+    llm_config = LLMConfig(
+        model_loading_config=ModelLoadingConfig(model_id="test-tpu-model"),
+        accelerator_type="TPU-V6E",
+        accelerator_config={"kind": "tpu", "topology": "4x4"},
+        placement_group_config={
+            "bundle_per_worker": {"TPU": 1},
+        },
+        engine_kwargs={
+            "tensor_parallel_size": 2,
+        },
+    )
+
+    engine_config = llm_config.get_engine_config()
+
+    # Validate the accelerator backend was correctly inferred
+    assert isinstance(engine_config.accelerator, TPUAccelerator)
+
+    bundles = engine_config.placement_bundles
+    assert len(bundles) == 2
+    for bundle in bundles:
+        assert bundle["TPU"] == 1
+        assert "accelerator_type:TPU-V6E" in bundle
+        assert bundle["accelerator_type:TPU-V6E"] == 0.001
+
+
+def test_accelerator_inference_logic():
+    """
+    Verifies that LLMConfig correctly infers the accelerator config
+    when no explicit accelerator_config is provided, and passes it
+    correctly to the engine.
+    """
+    # TPU string correctly infers TPUConfig and TPUAccelerator
+    cfg1 = LLMConfig(
+        model_loading_config={"model_id": "test"},
+        accelerator_type="TPU-V6E",
+    )
+    assert isinstance(cfg1.accelerator_config, TPUConfig)
+    assert isinstance(cfg1.get_engine_config().accelerator, TPUAccelerator)
+
+    # GPU string falls back to GPUConfig and GPUAccelerator
+    cfg2 = LLMConfig(
+        model_loading_config={"model_id": "test"},
+        accelerator_type="A10G",
+    )
+    assert isinstance(cfg2.accelerator_config, GPUConfig)
+    assert isinstance(cfg2.get_engine_config().accelerator, GPUAccelerator)
+
+    # No accelerator hints falls back to GPU by default
+    cfg3 = LLMConfig(model_loading_config={"model_id": "test"})
+    assert isinstance(cfg3.accelerator_config, GPUConfig)
+    assert isinstance(cfg3.get_engine_config().accelerator, GPUAccelerator)
+
+    # Explicit CPU config correctly yields CPUAccelerator
+    cfg4 = LLMConfig(
+        model_loading_config={"model_id": "test"},
+        accelerator_config={"kind": "cpu"},
+    )
+    assert isinstance(cfg4.accelerator_config, CPUConfig)
+    assert isinstance(cfg4.get_engine_config().accelerator, CPUAccelerator)
+
+
+def test_tpu_slice_placement_group_creation_heterogeneous_tpu_bundles_fail():
+    """
+    Verifies that a ValueError is raised when heterogeneous TPU bundles are provided.
+    """
+    accelerator = TPUAccelerator(TPUConfig(kind="tpu", topology="4x4"))
+
+    with pytest.raises(ValueError, match="Heterogeneous TPU bundles are not supported"):
+        accelerator.create_placement_group(
+            bundles=[{"TPU": 4}, {"TPU": 2}],
+            strategy="PACK",
+            name="test-pg",
+            accelerator_type_str="TPU-V6E",
+        )
+
+
+def test_tpu_slice_placement_group_creation_cpu_driver_homogeneous_tpu_bundles_pass(
+    ray_tpu_cluster,
+):
+    """
+    Verifies that CPU-only driver bundles are ignored and do not trigger an error
+    if subsequent TPU bundles are homogeneous.
+    """
+    accelerator = TPUAccelerator(TPUConfig(kind="tpu", topology="4x4"))
+
+    pg = accelerator.create_placement_group(
+        bundles=[{"CPU": 2}, {"TPU": 4}, {"TPU": 4}],
+        strategy="PACK",
+        name="test-pg",
+        accelerator_type_str="TPU-V6E",
+    )
+
+    # Verify valid PG creation
+    assert isinstance(pg, PlacementGroup)
+
+    accelerator.shutdown()
+    try:
+        ray.util.remove_placement_group(pg)
+    except Exception:
+        pass
+
+
+def test_tpu_serve_deployment_default_host_level_bundles(ray_tpu_cluster):
+    """
+    Verifies that a Serve deployment created for a multi-host TPU slice defaults
+    to host-level bundles when no placement_group_config is specified.
+    """
+    llm_config = LLMConfig(
+        model_loading_config=ModelLoadingConfig(model_id="test-tpu-model"),
+        accelerator_type="TPU-V6E",
+        accelerator_config=TPUConfig(kind="tpu", topology="4x4"),
+        engine_kwargs={"tensor_parallel_size": 16},
+    )
+
+    app = serve.deployment(LLMServer).bind(llm_config, engine_cls=PGCreationMockEngine)
+    serve.run(app)
+
+    pg_table = ray.util.placement_group_table()
+    active_pgs = list(
+        {k: v for k, v in pg_table.items() if v["state"] == "CREATED"}.values()
+    )
+
+    assert (
+        len(active_pgs) == 2
+    ), "Expected 2 PGs - one for TPU Head, one for worker bundles"
+
+    tpu_head_resource = "TPU-v6e-16-head"
+    head_pgs = [
+        pg
+        for pg in active_pgs
+        if len(pg["bundles"]) == 1
+        and tpu_head_resource in list(pg["bundles"].values())[0]
+    ]
+    assert len(head_pgs) == 1
+
+    worker_pg = [pg for pg in active_pgs if pg not in head_pgs][0]
+
+    assert worker_pg["strategy"] == "PACK"
+    # 4x4 topology = 16 chips. Default is 4 bundles of 4 TPUs (per-host).
+    assert len(worker_pg["bundles"]) == 4
+    for bundle in worker_pg["bundles"].values():
+        assert bundle.get("TPU", 0) == 4
+
+    serve.shutdown()
+
+
+def test_tpu_serve_deployment_explicit_host_level_bundles(ray_tpu_cluster):
+    """
+    Verifies that a user can explicitly request host-level bundles (4 TPUs per bundle)
+    for a Serve deployment via placement_group_config.
+    """
+    llm_config = LLMConfig(
+        model_loading_config=ModelLoadingConfig(model_id="test-tpu-model"),
+        accelerator_type="TPU-V6E",
+        accelerator_config=TPUConfig(kind="tpu", topology="4x4"),
+        placement_group_config={"bundle_per_worker": {"TPU": 4}},
+        engine_kwargs={"tensor_parallel_size": 16},
+    )
+
+    app = serve.deployment(LLMServer).bind(llm_config, engine_cls=PGCreationMockEngine)
+    serve.run(app)
+
+    pg_table = ray.util.placement_group_table()
+    active_pgs = list(
+        {k: v for k, v in pg_table.items() if v["state"] == "CREATED"}.values()
+    )
+
+    assert (
+        len(active_pgs) == 2
+    ), "Expected 2 PGs - one for TPU Head, one for worker bundles"
+
+    tpu_head_resource = "TPU-v6e-16-head"
+    head_pgs = [
+        pg
+        for pg in active_pgs
+        if len(pg["bundles"]) == 1
+        and tpu_head_resource in list(pg["bundles"].values())[0]
+    ]
+    assert len(head_pgs) == 1
+
+    worker_pg = [pg for pg in active_pgs if pg not in head_pgs][0]
+
+    assert worker_pg["strategy"] == "PACK"
+    # 4x4 topology = 16 chips. With 4 TPUs per bundle, expect exactly 4 bundles.
+    assert len(worker_pg["bundles"]) == 4
+    for bundle in worker_pg["bundles"].values():
+        assert bundle.get("TPU", 0) == 4
+
+    serve.shutdown()
+
+
+def test_tpu_serve_deployment_explicit_per_chip_bundles(ray_tpu_cluster):
+    """
+    Verifies that a user can explicitly request chip-level bundles (1 TPU per bundle)
+    for a full multi-host TPU slice via placement_group_config.
+    """
+    llm_config = LLMConfig(
+        model_loading_config=ModelLoadingConfig(model_id="test-tpu-model"),
+        accelerator_type="TPU-V6E",
+        accelerator_config=TPUConfig(kind="tpu", topology="4x4"),
+        placement_group_config={"bundle_per_worker": {"TPU": 1}},
+        engine_kwargs={"tensor_parallel_size": 16},
+    )
+
+    app = serve.deployment(LLMServer).bind(llm_config, engine_cls=PGCreationMockEngine)
+    serve.run(app)
+
+    pg_table = ray.util.placement_group_table()
+    active_pgs = list(
+        {k: v for k, v in pg_table.items() if v["state"] == "CREATED"}.values()
+    )
+
+    assert (
+        len(active_pgs) == 2
+    ), "Expected 2 PGs - one for TPU Head, one for worker bundles"
+
+    tpu_head_resource = "TPU-v6e-16-head"
+    head_pgs = [
+        pg
+        for pg in active_pgs
+        if len(pg["bundles"]) == 1
+        and tpu_head_resource in list(pg["bundles"].values())[0]
+    ]
+    assert len(head_pgs) == 1
+
+    worker_pg = [pg for pg in active_pgs if pg not in head_pgs][0]
+
+    assert worker_pg["strategy"] == "PACK"
+    # 4x4 topology = 16 chips. Explicitly requested 16 bundles of 1 TPU.
+    assert len(worker_pg["bundles"]) == 16
+    for bundle in worker_pg["bundles"].values():
+        assert bundle.get("TPU", 0) == 1.0
+
+    serve.shutdown()
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main(["-v", __file__]))

python/ray/llm/tests/serve/cpu/deployments/llm/test_llm_server.py

@@ -0,0 +1,789 @@
+import asyncio
+import sys
+import time
+from types import SimpleNamespace
+from typing import AsyncGenerator, Optional
+from unittest.mock import patch
+
+import numpy as np
+import pytest
+
+from ray import serve
+from ray.llm._internal.serve.core.configs.llm_config import (
+    LLMConfig,
+    LoraConfig,
+    ModelLoadingConfig,
+)
+from ray.llm._internal.serve.core.configs.openai_api_models import CompletionRequest
+from ray.llm._internal.serve.core.protocol import RawRequestInfo
+from ray.llm._internal.serve.core.server.llm_server import LLMServer
+from ray.llm._internal.serve.engines.vllm.vllm_engine import (
+    _canonicalize_request_id_header,
+)
+from ray.llm.tests.serve.mocks.mock_vllm_engine import (
+    FakeLoraModelLoader,
+    MockVLLMEngine,
+)
+from ray.llm.tests.serve.utils.testing_utils import LLMResponseValidator
+
+
+@pytest.fixture
+def serve_handle(mock_llm_config, stream_batching_interval_ms=0):
+    mock_llm_config.experimental_configs = {
+        "stream_batching_interval_ms": stream_batching_interval_ms,
+    }
+
+    app = serve.deployment(LLMServer).bind(mock_llm_config, engine_cls=MockVLLMEngine)
+    handle = serve.run(app)
+    # We set stream=True because the interfaces are async generators regardless
+    # of the stream flag on request.
+    handle = handle.options(stream=True)
+    yield handle
+    serve.shutdown()
+
+
+@pytest.fixture
+def multiplexed_serve_handle(mock_llm_config, stream_batching_interval_ms=0):
+    mock_llm_config.experimental_configs = {
+        "stream_batching_interval_ms": stream_batching_interval_ms,
+    }
+    # Set minimal lora_config to enable multiplexing but avoid telemetry S3 calls
+    mock_llm_config.lora_config = LoraConfig(
+        dynamic_lora_loading_path=None,  # No S3 path = no telemetry S3 calls
+        download_timeout_s=60,
+        max_download_tries=3,
+    )
+
+    app = serve.deployment(LLMServer).bind(
+        mock_llm_config,
+        engine_cls=MockVLLMEngine,
+        model_downloader=FakeLoraModelLoader,
+    )
+    handle = serve.run(app)
+    handle = handle.options(stream=True, multiplexed_model_id="test_model_id")
+    yield handle
+    serve.shutdown()
+
+
+async def count_tpot_ms_from_stream(stream: AsyncGenerator) -> list[float]:
+    all_tpots_in_ms = []
+    start = None
+    async for _ in stream:
+        now = time.perf_counter()
+        if start is not None:
+            all_tpots_in_ms.append((now - start) * 1e3)
+        start = now
+    return all_tpots_in_ms
+
+
+class TestLLMServer:
+    @pytest.mark.parametrize("api_type", ["chat", "completion"])
+    @pytest.mark.parametrize("stream", [False, True])
+    @pytest.mark.parametrize("max_tokens", [5])
+    @pytest.mark.parametrize("stream_batching_interval_ms", [0, 10000])
+    @pytest.mark.asyncio
+    async def test_unified_llm_server(
+        self,
+        serve_handle,
+        mock_llm_config,
+        mock_chat_request,
+        mock_completion_request,
+        api_type: str,
+        stream: bool,
+        max_tokens: int,
+        stream_batching_interval_ms: int,
+    ):
+        """Unified test for both chat and completion APIs, streaming and non-streaming."""
+
+        # Create request based on API type
+        if api_type == "chat":
+            request = mock_chat_request
+            batched_chunks = serve_handle.chat.remote(request)
+        elif api_type == "completion":
+            request = mock_completion_request
+            batched_chunks = serve_handle.completions.remote(request)
+
+        print(
+            f"\n\n_____ {api_type.upper()} ({'STREAMING' if stream else 'NON-STREAMING'}) max_tokens={max_tokens} batching_interval_ms={stream_batching_interval_ms} _____\n\n"
+        )
+
+        if stream:
+            # Collect responses from the stream
+            chunks = []
+            async for batch in batched_chunks:
+                chunks.extend(batch)
+
+            # Check that we got responses
+            assert len(chunks) > 0
+
+            # Validate streaming response
+            LLMResponseValidator.validate_streaming_chunks(chunks, api_type, max_tokens)
+        else:
+            # Collect non-streaming response
+            chunks = []
+            async for batch in batched_chunks:
+                chunks.append(batch)
+
+            # Check that we got one response
+            assert len(chunks) == 1
+
+            # Validate non-streaming response
+            LLMResponseValidator.validate_non_streaming_response(
+                chunks[0], api_type, max_tokens
+            )
+
+    @pytest.mark.parametrize("dimensions", [None, 512])
+    @pytest.mark.asyncio
+    async def test_embedding_llm_server(
+        self,
+        serve_handle,
+        mock_llm_config,
+        mock_embedding_request,
+        dimensions: Optional[int],
+    ):
+        """Test embedding API from LLMServer perspective."""
+
+        # Create embedding request
+        request = mock_embedding_request
+
+        print(f"\n\n_____ EMBEDDING SERVER dimensions={dimensions} _____\n\n")
+
+        # Get the response
+        batched_chunks = serve_handle.embeddings.remote(request)
+
+        # Collect responses (should be just one)
+        chunks = []
+        async for batch in batched_chunks:
+            chunks.append(batch)
+
+        # Check that we got one response
+        assert len(chunks) == 1
+
+        # Validate embedding response
+        LLMResponseValidator.validate_embedding_response(chunks[0], dimensions)
+
+    @pytest.mark.parametrize("stream", [False, True])
+    @pytest.mark.parametrize("temperature", [0.0])
+    @pytest.mark.parametrize("language", ["en", "hi"])
+    @pytest.mark.asyncio
+    async def test_transcription_llm_server(
+        self,
+        serve_handle,
+        mock_llm_config,
+        mock_transcription_request,
+        stream: bool,
+        temperature: float,
+        language: Optional[str],
+    ):
+        """Test transcription API from LLMServer perspective."""
+
+        # Create transcription request
+        request = mock_transcription_request
+
+        print(
+            f"\n\n_____ TRANSCRIPTION SERVER ({'STREAMING' if stream else 'NON-STREAMING'}) language={language} temperature={temperature} _____\n\n"
+        )
+
+        # Get the response
+        batched_chunks = serve_handle.transcriptions.remote(request)
+
+        if stream:
+            # Collect streaming responses
+            chunks = []
+            async for batch in batched_chunks:
+                if isinstance(batch, list):
+                    chunks.extend(batch)
+                else:
+                    chunks.append(batch)
+
+            # Check that we got responses
+            assert len(chunks) > 0
+
+            # Validate streaming response
+            LLMResponseValidator.validate_transcription_response(
+                chunks, temperature, language
+            )
+        else:
+            # Collect non-streaming response
+            chunks = []
+            async for batch in batched_chunks:
+                chunks.append(batch)
+
+            # Check that we got one response
+            assert len(chunks) == 1
+
+            # Validate non-streaming response
+            LLMResponseValidator.validate_transcription_response(
+                chunks[0], temperature, language
+            )
+
+    @pytest.mark.asyncio
+    async def test_score_llm_server(
+        self,
+        serve_handle,
+        mock_llm_config,
+        mock_score_request,
+    ):
+        """Test score API from LLMServer perspective."""
+
+        # Create score request
+        request = mock_score_request
+
+        print("\n\n_____ SCORE SERVER _____\n\n")
+
+        # Get the response
+        batched_chunks = serve_handle.score.remote(request)
+
+        # Collect responses (should be just one)
+        chunks = []
+        async for batch in batched_chunks:
+            chunks.append(batch)
+
+        # Check that we got one response
+        assert len(chunks) == 1
+
+        # Validate score response
+        LLMResponseValidator.validate_score_response(chunks[0])
+
+    @pytest.mark.parametrize("return_token_strs", [False, True])
+    @pytest.mark.asyncio
+    async def test_tokenize_llm_server(
+        self,
+        serve_handle,
+        mock_llm_config,
+        mock_tokenize_request,
+        return_token_strs: bool,
+    ):
+        """Test tokenize API from LLMServer perspective."""
+
+        # Create tokenize request
+        request = mock_tokenize_request
+
+        print(
+            f"\n\n_____ TOKENIZE SERVER return_token_strs={return_token_strs} _____\n\n"
+        )
+
+        # Get the response
+        batched_chunks = serve_handle.tokenize.remote(request)
+
+        # Collect responses (should be just one)
+        chunks = []
+        async for batch in batched_chunks:
+            chunks.append(batch)
+
+        # Check that we got one response
+        assert len(chunks) == 1
+
+        # Validate tokenize response
+        LLMResponseValidator.validate_tokenize_response(
+            chunks[0],
+            expected_prompt="Hello, world!",
+            return_token_strs=return_token_strs,
+        )
+
+    @pytest.mark.asyncio
+    async def test_detokenize_llm_server(
+        self,
+        serve_handle,
+        mock_llm_config,
+        mock_detokenize_request,
+    ):
+        """Test detokenize API from LLMServer perspective."""
+
+        # Create detokenize request
+        request = mock_detokenize_request
+
+        print("\n\n_____ DETOKENIZE SERVER _____\n\n")
+
+        # Get the response
+        batched_chunks = serve_handle.detokenize.remote(request)
+
+        # Collect responses (should be just one)
+        chunks = []
+        async for batch in batched_chunks:
+            chunks.append(batch)
+
+        # Check that we got one response
+        assert len(chunks) == 1
+
+        # Validate detokenize response
+        LLMResponseValidator.validate_detokenize_response(
+            chunks[0],
+            expected_text="Hello",  # [72, 101, 108, 108, 111] = "Hello"
+        )
+
+    @pytest.mark.asyncio
+    async def test_check_health(self, mock_llm_config):
+        """Test health check functionality."""
+
+        # Mock the engine's check_health method
+        class LocalMockEngine(MockVLLMEngine):
+            def __init__(self, *args, **kwargs):
+                super().__init__(*args, **kwargs)
+                self.check_health_called = False
+
+            async def check_health(self):
+                self.check_health_called = True
+
+        # Create a server with a mocked engine
+        server = LLMServer.sync_init(mock_llm_config, engine_cls=LocalMockEngine)
+        await server.start()
+
+        # Perform the health check, no exceptions should be raised
+        await server.check_health()
+
+        # Check that the health check method was called
+        assert server.engine.check_health_called
+
+    @pytest.mark.asyncio
+    async def test_reset_prefix_cache(self, mock_llm_config):
+        """Test reset prefix cache functionality."""
+
+        # Mock the engine's reset_prefix_cache method
+        class LocalMockEngine(MockVLLMEngine):
+            def __init__(self, *args, **kwargs):
+                super().__init__(*args, **kwargs)
+                self.reset_prefix_cache_called = False
+
+            async def reset_prefix_cache(self):
+                self.reset_prefix_cache_called = True
+
+        # Create a server with a mocked engine
+        server = LLMServer.sync_init(mock_llm_config, engine_cls=LocalMockEngine)
+        await server.start()
+
+        # Reset prefix cache, no exceptions should be raised
+        await server.reset_prefix_cache()
+
+        # Check that the reset prefix cache method was called
+        assert server.engine.reset_prefix_cache_called
+
+    @pytest.mark.asyncio
+    async def test_start_profile(self, mock_llm_config):
+        """Test start profile functionality."""
+
+        # Mock the engine's start_profile method
+        class LocalMockEngine(MockVLLMEngine):
+            def __init__(self, *args, **kwargs):
+                super().__init__(*args, **kwargs)
+                self.start_profile_called = False
+
+            async def start_profile(self):
+                self.start_profile_called = True
+
+        # Create a server with a mocked engine
+        server = LLMServer.sync_init(mock_llm_config, engine_cls=LocalMockEngine)
+        await server.start()
+
+        # Start profile, no exceptions should be raised
+        await server.start_profile()
+
+        # Check that the start profile method was called
+        assert server.engine.start_profile_called
+
+    @pytest.mark.asyncio
+    async def test_stop_profile(self, mock_llm_config):
+        """Test stop profile functionality."""
+
+        # Mock the engine's stop_profile method
+        class LocalMockEngine(MockVLLMEngine):
+            def __init__(self, *args, **kwargs):
+                super().__init__(*args, **kwargs)
+                self.stop_profile_called = False
+
+            async def stop_profile(self):
+                self.stop_profile_called = True
+
+        # Create a server with a mocked engine
+        server = LLMServer.sync_init(mock_llm_config, engine_cls=LocalMockEngine)
+        await server.start()
+
+        # Stop profile, no exceptions should be raised
+        await server.stop_profile()
+
+        # Check that the stop profile method was called
+        assert server.engine.stop_profile_called
+
+    @pytest.mark.asyncio
+    async def test_llm_config_property(self, mock_llm_config):
+        """Test the llm_config property."""
+        server = LLMServer.sync_init(mock_llm_config, engine_cls=MockVLLMEngine)
+        await server.start()
+        llm_config = await server.llm_config()
+        assert isinstance(llm_config, type(mock_llm_config))
+
+    @pytest.mark.parametrize("stream", [False])
+    @pytest.mark.parametrize("max_tokens", [5])
+    @pytest.mark.asyncio
+    async def test_request_id_handling(
+        self,
+        serve_handle,
+        mock_llm_config,
+        mock_chat_request,
+        stream: bool,
+        max_tokens: int,
+    ):
+        """Test that the request id is handled correctly."""
+
+        # Create a chat completion request
+        # We should patch get_server_request_id to return a test_request_id
+        serve.context._serve_request_context.set(
+            serve.context._RequestContext(**{"request_id": "test_request_id"})
+        )
+        # Get the response
+        chunks = []
+        async for chunk in serve_handle.chat.remote(mock_chat_request):
+            chunks.append(chunk)
+
+        assert len(chunks) == 1
+        assert chunks[0].id == "test_request_id"
+
+    @pytest.mark.parametrize("api_type", ["chat", "completion"])
+    @pytest.mark.parametrize("stream", [False, True])
+    @pytest.mark.parametrize("max_tokens", [5])
+    @pytest.mark.parametrize("stream_batching_interval_ms", [0, 10000])
+    @pytest.mark.asyncio
+    async def test_multiplexed_request_handling(
+        self,
+        multiplexed_serve_handle,
+        mock_chat_request,
+        mock_completion_request,
+        api_type: str,
+        stream: bool,
+        max_tokens: int,
+        stream_batching_interval_ms: int,
+    ):
+        """Unified test for multiplexed (LoRA) requests - both chat and completion APIs, streaming and non-streaming."""
+
+        # Create request based on API type and set model ID for multiplexing
+        if api_type == "chat":
+            request = mock_chat_request
+            batched_chunks = multiplexed_serve_handle.chat.remote(request)
+        elif api_type == "completion":
+            request = mock_completion_request
+            batched_chunks = multiplexed_serve_handle.completions.remote(request)
+
+        request.model = "test_model_id"
+        print(
+            f"\n\n_____ MULTIPLEXED {api_type.upper()} ({'STREAMING' if stream else 'NON-STREAMING'}) max_tokens={max_tokens} batching_interval_ms={stream_batching_interval_ms} _____\n\n"
+        )
+
+        if stream:
+            # Collect responses from the stream
+            chunks = []
+            async for batch in batched_chunks:
+                if isinstance(batch, list):
+                    chunks.extend(batch)
+                else:
+                    chunks.append(batch)
+
+            # Check that we got responses
+            assert len(chunks) > 0
+
+            # Validate streaming response with LoRA model ID
+            LLMResponseValidator.validate_streaming_chunks(
+                chunks, api_type, max_tokens, lora_model_id=request.model
+            )
+        else:
+            # Collect non-streaming response
+            chunks = []
+            async for batch in batched_chunks:
+                if isinstance(batch, list):
+                    chunks.extend(batch)
+                else:
+                    chunks.append(batch)
+
+            # Check that we got one response
+            assert len(chunks) == 1
+
+            # Validate non-streaming response with LoRA model ID
+            LLMResponseValidator.validate_non_streaming_response(
+                chunks[0], api_type, max_tokens, lora_model_id=request.model
+            )
+
+    @pytest.mark.asyncio
+    async def test_push_telemetry(self, mock_llm_config):
+        """Test that the telemetry push is called properly."""
+        with patch(
+            "ray.llm._internal.serve.core.server.llm_server.push_telemetry_report_for_all_models"
+        ) as mock_push_telemetry:
+            server = LLMServer.sync_init(mock_llm_config, engine_cls=MockVLLMEngine)
+            await server.start()
+            mock_push_telemetry.assert_called_once()
+
+    @pytest.mark.parametrize("api_type", ["chat", "completions"])
+    @pytest.mark.parametrize("stream", [True])
+    @pytest.mark.parametrize("max_tokens", [64])
+    @pytest.mark.parametrize("concurrency", [1, 16])
+    @pytest.mark.parametrize("stream_batching_interval_ms", [0])
+    @pytest.mark.asyncio
+    async def test_stable_streaming_tpot(
+        self,
+        serve_handle,
+        mock_llm_config,
+        mock_chat_request,
+        mock_completion_request,
+        api_type: str,
+        stream: bool,
+        max_tokens: int,
+        concurrency: int,
+        stream_batching_interval_ms: int,
+    ):
+        """Test that the streaming TPOT is stable when batching is disabled."""
+
+        # Create request based on API type
+        if api_type == "chat":
+            request = mock_chat_request
+        elif api_type == "completions":
+            request = mock_completion_request
+        batched_chunks: list[AsyncGenerator] = [
+            getattr(serve_handle, api_type).remote(request) for _ in range(concurrency)
+        ]
+
+        print(
+            f"\n\n_____ {api_type.upper()} ({'STREAMING' if stream else 'NON-STREAMING'}) max_tokens={max_tokens} batching_interval_ms={stream_batching_interval_ms} _____\n\n"
+        )
+
+        # Collect responses from llm_server
+        tpots_ms = await asyncio.gather(
+            *[
+                count_tpot_ms_from_stream(server_stream)
+                for server_stream in batched_chunks
+            ]
+        )
+        mean_llm_server = np.mean(tpots_ms)
+        std_var_llm_server = np.std(tpots_ms)
+
+        # Run same request with vllm engine
+        vllm_engine = MockVLLMEngine(llm_config=mock_llm_config)
+        await vllm_engine.start()
+        engine_streams: list[AsyncGenerator] = [
+            getattr(vllm_engine, api_type)(request) for _ in range(concurrency)
+        ]
+        tpots_ms_engine = await asyncio.gather(
+            *[
+                count_tpot_ms_from_stream(engine_stream)
+                for engine_stream in engine_streams
+            ]
+        )
+        mean_engine = np.mean(tpots_ms_engine)
+        std_var_engine = np.std(tpots_ms_engine)
+
+        assert np.isclose(
+            mean_llm_server, mean_engine, rtol=0.1
+        ), f"{mean_llm_server=}, {mean_engine=}"
+        assert np.isclose(
+            std_var_llm_server, std_var_engine, atol=1.0
+        ), f"{std_var_llm_server=}, {std_var_engine=}"
+
+
+class TestGetDeploymentOptions:
+    def test_placement_group_config(self):
+        """Test that placement_group_config is correctly parsed."""
+
+        # Test the default resource bundle
+        llm_config = LLMConfig(
+            model_loading_config=dict(model_id="test_model"),
+            engine_kwargs=dict(tensor_parallel_size=3, pipeline_parallel_size=2),
+        )
+        serve_options = LLMServer.get_deployment_options(llm_config)
+
+        assert serve_options["placement_group_bundles"] == [{"CPU": 1, "GPU": 1}] + [
+            {"GPU": 1} for _ in range(5)
+        ]
+
+        # Test the custom placement group config
+        # Note: The first bundle gets merged with replica actor resources (CPU: 1, GPU: 0)
+        llm_config = LLMConfig(
+            model_loading_config=dict(model_id="test_model"),
+            engine_kwargs=dict(tensor_parallel_size=3, pipeline_parallel_size=2),
+            placement_group_config={
+                "bundles": [{"CPU": 1, "XPU": 1}] + [{"XPU": 1}] * 5,
+                "strategy": "PACK",
+            },
+        )
+        serve_options = LLMServer.get_deployment_options(llm_config)
+        # First bundle has replica actor resources merged in (CPU: 1 from config + 1 from replica = 2)
+        # Bundles are validated via PlacementGroupConfig: unset CPU/GPU are omitted from the output due to exclude_unset=True.
+        assert serve_options["placement_group_bundles"] == [
+            {"CPU": 2.0, "GPU": 0, "XPU": 1}
+        ] + [{"XPU": 1} for _ in range(5)]
+        assert serve_options["placement_group_strategy"] == "PACK"
+
+    def test_get_serve_options_with_accelerator_type(self):
+        """Test that get_serve_options returns the correct options when accelerator_type is set."""
+        llm_config = LLMConfig(
+            model_loading_config=ModelLoadingConfig(model_id="test_model"),
+            accelerator_type="A100-40G",
+            deployment_config={
+                "autoscaling_config": {
+                    "min_replicas": 0,
+                    "initial_replicas": 1,
+                    "max_replicas": 10,
+                },
+            },
+            runtime_env={"env_vars": {"FOO": "bar"}},
+        )
+
+        serve_options = LLMServer.get_deployment_options(llm_config)
+
+        # Test the core functionality without being strict about Ray's automatic runtime env additions
+        assert serve_options["autoscaling_config"] == {
+            "min_replicas": 0,
+            "initial_replicas": 1,
+            "max_replicas": 10,
+        }
+        assert serve_options["placement_group_bundles"] == [
+            {"CPU": 1, "GPU": 1, "accelerator_type:A100-40G": 0.001},
+        ]
+        # Default strategy is PACK (cross-node allowed by default)
+        assert serve_options["placement_group_strategy"] == "PACK"
+
+        # Check that our custom env vars are present
+        assert (
+            serve_options["ray_actor_options"]["runtime_env"]["env_vars"]["FOO"]
+            == "bar"
+        )
+        assert (
+            "worker_process_setup_hook"
+            in serve_options["ray_actor_options"]["runtime_env"]
+        )
+
+    def test_get_serve_options_without_accelerator_type(self):
+        """Test that get_serve_options returns the correct options when accelerator_type is not set."""
+        llm_config = LLMConfig(
+            model_loading_config=ModelLoadingConfig(model_id="test_model"),
+            deployment_config={
+                "autoscaling_config": {
+                    "min_replicas": 0,
+                    "initial_replicas": 1,
+                    "max_replicas": 10,
+                },
+            },
+            runtime_env={"env_vars": {"FOO": "bar"}},
+        )
+        serve_options = LLMServer.get_deployment_options(llm_config)
+
+        # Test the core functionality without being strict about Ray's automatic runtime env additions
+        assert serve_options["autoscaling_config"] == {
+            "min_replicas": 0,
+            "initial_replicas": 1,
+            "max_replicas": 10,
+        }
+        assert serve_options["placement_group_bundles"] == [{"CPU": 1, "GPU": 1}]
+        # Default strategy is PACK (cross-node allowed by default)
+        assert serve_options["placement_group_strategy"] == "PACK"
+
+        # Check that our custom env vars are present
+        assert (
+            serve_options["ray_actor_options"]["runtime_env"]["env_vars"]["FOO"]
+            == "bar"
+        )
+        assert (
+            "worker_process_setup_hook"
+            in serve_options["ray_actor_options"]["runtime_env"]
+        )
+
+    def test_deferred_placement_group_for_tpu_topology(self):
+        """Test that Serve skips PG creation when deferred placement group is required."""
+        llm_config = LLMConfig(
+            model_loading_config=ModelLoadingConfig(model_id="test-tpu-model"),
+            accelerator_type="TPU-V6E",
+            accelerator_config={"kind": "tpu", "topology": "4x4"},
+            llm_engine="vLLM",
+        )
+
+        serve_options = LLMServer.get_deployment_options(llm_config)
+
+        assert "placement_group_bundles" not in serve_options
+        assert "placement_group_strategy" not in serve_options
+
+
+class TestCanonicalizeRequestIdHeader:
+    """Unit tests for the X-Request-Id header canonicalization helper."""
+
+    def test_uncanonical_variants_dropped(self):
+        """Any case/separator variant of the header is dropped and replaced by a
+        single canonical ``x-request-id`` equal to ``request.request_id``."""
+        request = SimpleNamespace(request_id="canonical-id")
+        raw = RawRequestInfo(
+            headers={
+                "X-Request-ID": "stale-upper",
+                "x_request_id": "stale-underscore",
+                "content-type": "application/json",
+            }
+        )
+        out = _canonicalize_request_id_header(request, raw)
+
+        rid_keys = [
+            k for k in out.headers if k.replace("_", "-").lower() == "x-request-id"
+        ]
+        assert rid_keys == ["x-request-id"], rid_keys
+        assert out.headers["x-request-id"] == "canonical-id"
+        # Unrelated headers are preserved.
+        assert out.headers["content-type"] == "application/json"
+
+    def test_noop_when_request_id_unset(self):
+        """With no request_id the helper is a no-op (returns the same object)."""
+        raw = RawRequestInfo(headers={"x-request-id": "keep"})
+        assert (
+            _canonicalize_request_id_header(SimpleNamespace(request_id=None), raw)
+            is raw
+        )
+
+
+class TestMaybeAddRequestId:
+    """``_maybe_add_request_id_to_request`` fills the Serve request id for a
+    defaulted request_id but never clobbers one the caller set explicitly."""
+
+    def _set_ctx(self, request_id):
+        serve.context._serve_request_context.set(
+            serve.context._RequestContext(request_id=request_id)
+        )
+
+    @pytest.mark.asyncio
+    async def test_defaulted_request_id_is_overwritten_with_serve_id(self):
+        server = LLMServer.__new__(LLMServer)
+        req = CompletionRequest(model="m", prompt="hi")  # request_id defaulted
+        assert "request_id" not in req.model_fields_set
+        self._set_ctx("serve-ctx-id")
+        try:
+            await server._maybe_add_request_id_to_request(req)
+        finally:
+            serve.context._serve_request_context.set(serve.context._RequestContext())
+        assert req.request_id == "serve-ctx-id"
+
+    @pytest.mark.asyncio
+    async def test_explicit_request_id_is_preserved(self):
+        server = LLMServer.__new__(LLMServer)
+        req = CompletionRequest(model="m", prompt="hi", request_id="caller-set-id")
+        assert "request_id" in req.model_fields_set
+        self._set_ctx("serve-ctx-id")
+        try:
+            await server._maybe_add_request_id_to_request(req)
+        finally:
+            serve.context._serve_request_context.set(serve.context._RequestContext())
+        # Caller's id wins; the Serve context id does not clobber it.
+        assert req.request_id == "caller-set-id"
+
+    @pytest.mark.asyncio
+    async def test_request_without_request_id_field_is_skipped(self):
+        """Request types without a request_id field (e.g. tokenize/detokenize)
+        must be handled gracefully, not raise."""
+        from pydantic import BaseModel
+
+        class _NoRequestId(BaseModel):
+            pass
+
+        server = LLMServer.__new__(LLMServer)
+        req = _NoRequestId()
+        self._set_ctx("serve-ctx-id")
+        try:
+            await server._maybe_add_request_id_to_request(req)
+        finally:
+            serve.context._serve_request_context.set(serve.context._RequestContext())
+        assert not hasattr(req, "request_id")
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main(["-v", __file__]))

python/ray/llm/tests/serve/cpu/deployments/llm/vllm/kv_transfer_backends/__init__.py

@@ -0,0 +1 @@
+# Test package for KV transfer backends

python/ray/llm/tests/serve/cpu/deployments/llm/vllm/kv_transfer_backends/test_factory.py

@@ -0,0 +1,185 @@
+import sys
+from contextlib import contextmanager
+from typing import Any
+
+import pytest
+
+from ray import serve
+from ray.llm._internal.serve.engines.vllm.kv_transfer.base import (
+    BaseConnectorBackend,
+    DefaultConnectorBackend,
+    DefaultPDProtocolMixin,
+)
+from ray.llm._internal.serve.engines.vllm.kv_transfer.factory import (
+    KVConnectorBackendFactory,
+)
+from ray.serve.llm import LLMConfig
+
+
+@contextmanager
+def registered_backend(name: str, backend_class_or_path: Any):
+    KVConnectorBackendFactory.register_backend(name, backend_class_or_path)
+    try:
+        yield
+    finally:
+        if KVConnectorBackendFactory.is_registered(name):
+            KVConnectorBackendFactory.unregister_backend(name)
+
+
+@pytest.fixture
+def test_deployment_handle():
+    """Fixture that creates a Serve deployment for testing cross-process registry access."""
+
+    # This ensures proper serialization when sent to child processes
+    class TestCrossProcessConnector(DefaultPDProtocolMixin, BaseConnectorBackend):
+        def setup(self):
+            pass
+
+    # Register the backend in the driver process and ensure cleanup
+    with registered_backend("TestCrossProcessConnector", TestCrossProcessConnector):
+        # Create a Serve deployment that will run in a different process than the
+        # driver process
+        @serve.deployment
+        class TestDeployment:
+            def __init__(self):
+                # This runs in a child process - should be able to access the registered backend
+                self.connector_class = KVConnectorBackendFactory.get_backend_class(
+                    "TestCrossProcessConnector"
+                )
+
+            def __call__(self):
+                """Return the connector class to verify it's correct."""
+                return self.connector_class
+
+        # Deploy and yield the handle and connector class
+        app = TestDeployment.bind()
+        handle = serve.run(app)
+        try:
+            yield handle, TestCrossProcessConnector
+        finally:
+            try:
+                serve.shutdown()
+            except RuntimeError:
+                # Handle case where event loop is already closed
+                pass
+
+
+class TestKVConnectorBackendFactory:
+    """Test suite for KVConnectorBackendFactory."""
+
+    def test_get_backend_class_success(self):
+        """Test successful retrieval of a registered backend class."""
+        backend_class = KVConnectorBackendFactory.get_backend_class(
+            "LMCacheConnectorV1"
+        )
+        assert backend_class is not None
+        assert hasattr(backend_class, "setup")
+
+    def test_get_backend_class_not_registered_returns_default(self):
+        """Getting a non-registered backend returns a concrete default backend.
+
+        ``BaseConnectorBackend`` is abstract, so the fallback must be a concrete
+        subclass (``DefaultConnectorBackend``) that is instantiable and provides
+        the default P/D protocol policy.
+        """
+        backend_class = KVConnectorBackendFactory.get_backend_class(
+            "UnregisteredConnector"
+        )
+        assert backend_class is DefaultConnectorBackend
+        assert issubclass(backend_class, BaseConnectorBackend)
+        # Concrete: can be instantiated.
+        DefaultConnectorBackend(llm_config=None)
+
+    def test_create_backend_success(self):
+        """Test successful creation of a backend instance."""
+        llm_config = LLMConfig(
+            model_loading_config=dict(model_id="test-model"),
+            engine_kwargs=dict(
+                kv_transfer_config=dict(
+                    kv_connector="LMCacheConnectorV1",
+                    kv_role="kv_both",
+                )
+            ),
+        )
+        backend = KVConnectorBackendFactory.create_backend(
+            "LMCacheConnectorV1", llm_config
+        )
+        assert isinstance(backend, BaseConnectorBackend)
+        assert backend.llm_config == llm_config
+
+    @pytest.mark.parametrize(
+        "connector_name",
+        ["LMCacheConnectorV1", "NixlConnector", "MultiConnector"],
+    )
+    def test_all_registered_backends_can_be_loaded(self, connector_name):
+        """Test that all pre-registered backends can be loaded."""
+        backend_class = KVConnectorBackendFactory.get_backend_class(connector_name)
+        assert backend_class is not None
+        assert issubclass(backend_class, BaseConnectorBackend)
+
+    def test_get_backend_class_import_error_handling(self):
+        """Test that ImportError during backend loading is handled with clear message."""
+        # Register a backend with a non-existent module path
+        with registered_backend("BadBackend", "non.existent.module:NonExistentClass"):
+            with pytest.raises(
+                ImportError, match="Failed to load connector backend 'BadBackend'"
+            ):
+                KVConnectorBackendFactory.get_backend_class("BadBackend")
+
+    def test_register_backend_with_class_directly(self):
+        """Test registering a backend class directly."""
+
+        class CustomBackend(DefaultPDProtocolMixin, BaseConnectorBackend):
+            def setup(self):
+                pass
+
+        with registered_backend("CustomBackend", CustomBackend):
+            assert KVConnectorBackendFactory.is_registered("CustomBackend")
+            retrieved = KVConnectorBackendFactory.get_backend_class("CustomBackend")
+            assert retrieved == CustomBackend
+
+    def test_register_backend_with_module_path(self):
+        """Test registering a backend via module path string."""
+        # Register using module:class format
+        with registered_backend(
+            "LMCacheViaPath",
+            "ray.llm._internal.serve.engines.vllm.kv_transfer.lmcache:LMCacheConnectorV1Backend",
+        ):
+            assert KVConnectorBackendFactory.is_registered("LMCacheViaPath")
+            backend_class = KVConnectorBackendFactory.get_backend_class(
+                "LMCacheViaPath"
+            )
+            assert backend_class is not None
+            assert issubclass(backend_class, BaseConnectorBackend)
+
+    def test_unregistered_connector_with_llm_config_setup(self):
+        """Test that unregistered connectors work with LLMConfig.setup_engine_backend()."""
+        llm_config = LLMConfig(
+            model_loading_config=dict(model_id="test-model"),
+            engine_kwargs=dict(
+                kv_transfer_config=dict(
+                    kv_connector="SharedStorageConnector",
+                    kv_role="kv_both",
+                )
+            ),
+        )
+        # Should not raise an error
+        llm_config.setup_engine_backend()
+
+    @pytest.mark.asyncio
+    async def test_cross_process_registry_access(self, test_deployment_handle):
+        """Test that registrations made in driver are accessible in Ray Serve child processes."""
+        handle, TestCrossProcessConnector = test_deployment_handle
+
+        # Verify it's registered in driver
+        assert KVConnectorBackendFactory.is_registered("TestCrossProcessConnector")
+
+        result = await handle.remote()
+
+        # Verify it's the correct class
+        assert result == TestCrossProcessConnector
+        assert issubclass(result, BaseConnectorBackend)
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main(["-v", __file__]))

python/ray/llm/tests/serve/cpu/deployments/llm/vllm/kv_transfer_backends/test_lmcache_connector_v1.py

@@ -0,0 +1,123 @@
+import sys
+from unittest.mock import patch
+
+import pytest
+
+from ray.llm._internal.serve.engines.vllm.kv_transfer.lmcache import (
+    LMCacheConnectorV1Backend,
+)
+from ray.serve.llm import LLMConfig
+
+
+class TestLMCacheConnectorV1Backend:
+    @pytest.fixture(autouse=True)
+    def mock_lmcache_check(self):
+        """Mock the lmcache installation check for all tests."""
+        with patch(
+            "ray.llm._internal.serve.engines.vllm.kv_transfer.lmcache._check_lmcache_installed"
+        ):
+            yield
+
+    @pytest.fixture
+    def lmcache_backend_basic(self):
+        """Fixture for basic LMCacheConnectorV1Backend."""
+        return LMCacheConnectorV1Backend(
+            llm_config=LLMConfig(
+                model_loading_config=dict(
+                    model_id="Qwen/Qwen3-0.6B",
+                ),
+                engine_kwargs=dict(
+                    kv_transfer_config=dict(
+                        kv_connector="LMCacheConnectorV1",
+                        kv_role="kv_both",
+                    )
+                ),
+            ),
+        )
+
+    @pytest.fixture
+    def lmcache_backend_with_extra(self):
+        """Fixture for LMCacheConnectorV1Backend with extra config."""
+        return LMCacheConnectorV1Backend(
+            llm_config=LLMConfig(
+                model_loading_config=dict(
+                    model_id="Qwen/Qwen3-0.6B",
+                ),
+                engine_kwargs=dict(
+                    kv_transfer_config=dict(
+                        kv_connector="LMCacheConnectorV1",
+                        kv_role="kv_both",
+                        kv_connector_extra_config={},
+                    )
+                ),
+            ),
+        )
+
+    @pytest.fixture
+    def lmcache_backend_with_port(self):
+        """Fixture for LMCacheConnectorV1Backend with port config."""
+        return LMCacheConnectorV1Backend(
+            llm_config=LLMConfig(
+                model_loading_config=dict(
+                    model_id="Qwen/Qwen3-0.6B",
+                ),
+                engine_kwargs=dict(
+                    kv_transfer_config=dict(
+                        kv_connector="LMCacheConnectorV1",
+                        kv_role="kv_both",
+                        kv_connector_extra_config={
+                            "lmcache_rpc_port": LMCacheConnectorV1Backend.DEFAULT_LMCACHE_RPC_PORT_NAME,
+                        },
+                    )
+                ),
+            ),
+        )
+
+    def test_setup_basic_config(self, lmcache_backend_basic):
+        """Test setup with basic configuration (no kv_connector_extra_config)."""
+        lmcache_backend_basic.setup()
+
+        # Configuration should remain unchanged
+        assert (
+            "kv_connector_extra_config" not in lmcache_backend_basic.kv_transfer_config
+        )
+
+    def test_setup_with_extra_config_no_port(self, lmcache_backend_with_extra):
+        """Test setup with extra config but no lmcache_rpc_port."""
+        lmcache_backend_with_extra.setup()
+
+        # Should add lmcache_rpc_port with default DEFAULT_LMCACHE_RPC_PORT_NAME + random string
+        assert (
+            "lmcache_rpc_port"
+            in lmcache_backend_with_extra.kv_transfer_config[
+                "kv_connector_extra_config"
+            ]
+        )
+        port_value = lmcache_backend_with_extra.kv_transfer_config[
+            "kv_connector_extra_config"
+        ]["lmcache_rpc_port"]
+        assert port_value.startswith(
+            LMCacheConnectorV1Backend.DEFAULT_LMCACHE_RPC_PORT_NAME
+        )
+        assert len(port_value) > len(
+            LMCacheConnectorV1Backend.DEFAULT_LMCACHE_RPC_PORT_NAME
+        )  # Should have random string appended
+
+    def test_setup_with_existing_port(self, lmcache_backend_with_port):
+        """Test setup with existing lmcache_rpc_port configuration."""
+        original_port = lmcache_backend_with_port.kv_transfer_config[
+            "kv_connector_extra_config"
+        ]["lmcache_rpc_port"]
+
+        lmcache_backend_with_port.setup()
+
+        # Should modify the existing port by appending random string
+        new_port = lmcache_backend_with_port.kv_transfer_config[
+            "kv_connector_extra_config"
+        ]["lmcache_rpc_port"]
+        assert new_port.startswith(original_port)
+        assert len(new_port) > len(original_port)  # Should have random string appended
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main(["-v", __file__]))

python/ray/llm/tests/serve/cpu/deployments/llm/vllm/kv_transfer_backends/test_moriio_connector.py

@@ -0,0 +1,369 @@
+import re
+import sys
+from types import SimpleNamespace
+from unittest.mock import patch
+
+import pytest
+
+from ray.llm._internal.serve.engines.vllm.kv_transfer.base import (
+    BaseConnectorBackend,
+)
+from ray.llm._internal.serve.engines.vllm.kv_transfer.factory import (
+    KVConnectorBackendFactory,
+)
+from ray.llm._internal.serve.engines.vllm.kv_transfer.moriio import (
+    _DECODE_ZMQ_RE,
+    _PREFILL_ZMQ_RE,
+    DEFAULT_HANDSHAKE_PORT_BASE,
+    DEFAULT_NOTIFY_PORT_BASE,
+    MoRIIOConnectorBackend,
+    parse_peer_zmq,
+    parse_zmq_address,
+)
+from ray.serve.llm import LLMConfig
+from ray.serve.schema import ReplicaRank
+
+_TEST_HOST = "10.0.0.5"
+
+
+def _replica_context(global_rank: int) -> SimpleNamespace:
+    return SimpleNamespace(
+        rank=ReplicaRank(rank=global_rank, node_rank=0, local_rank=global_rank)
+    )
+
+
+def _make_backend(
+    read_mode: bool = False,
+    extra_exp: dict = None,
+    dp_rank: int = None,
+    dp_size: int = None,
+    tp_size: int = None,
+):
+    extra_config = {}
+    if read_mode:
+        extra_config["read_mode"] = "true"
+    engine_kwargs = dict(
+        kv_transfer_config=dict(
+            kv_connector="MoRIIOConnector",
+            kv_role="kv_both",
+            kv_connector_extra_config=extra_config,
+        )
+    )
+    if dp_rank is not None:
+        engine_kwargs["data_parallel_rank"] = dp_rank
+    if dp_size is not None:
+        engine_kwargs["data_parallel_size"] = dp_size
+    if tp_size is not None:
+        engine_kwargs["tensor_parallel_size"] = tp_size
+    return MoRIIOConnectorBackend(
+        llm_config=LLMConfig(
+            model_loading_config=dict(model_id="Qwen/Qwen3-0.6B"),
+            engine_kwargs=engine_kwargs,
+            experimental_configs=extra_exp or {},
+        ),
+    )
+
+
+def _setup(backend, rank: int = 0):
+    with (
+        patch.dict("os.environ", {}, clear=True),
+        patch("ray.util.get_node_ip_address", return_value=_TEST_HOST),
+        patch("ray.serve.get_replica_context", return_value=_replica_context(rank)),
+    ):
+        backend.setup()
+
+
+class TestMoRIIOConnectorBackendSetup:
+    def test_setup_sets_ports_zmq_and_extra_config(self):
+        backend = _make_backend()
+        _setup(backend, rank=0)
+
+        extra = backend.kv_transfer_config["kv_connector_extra_config"]
+        assert extra["handshake_port"] == str(DEFAULT_HANDSHAKE_PORT_BASE)
+        assert extra["notify_port"] == str(DEFAULT_NOTIFY_PORT_BASE)
+        assert extra["proxy_ip"] == ""
+        assert extra["proxy_ping_port"] == "0"
+        assert "http_port" in extra
+        assert extra["read_mode"] == "false"
+        # Routable node IP passed to vLLM's MoRIIO connector (cross-node fix,
+        # vllm-project/vllm#45488), matching the advertised zmq host.
+        assert extra["host_ip"] == _TEST_HOST
+
+        zmq = backend._zmq_address
+        host, hs, notify = parse_zmq_address(zmq)
+        assert host == _TEST_HOST
+        assert host == extra["host_ip"]
+        assert hs == DEFAULT_HANDSHAKE_PORT_BASE
+        assert notify == DEFAULT_NOTIFY_PORT_BASE
+
+    def test_setup_port_offset_uses_replica_rank(self):
+        backend = _make_backend()
+        num_devices = backend.llm_config.get_engine_config().num_devices
+        _setup(backend, rank=2)
+        extra = backend.kv_transfer_config["kv_connector_extra_config"]
+        assert extra["handshake_port"] == str(
+            DEFAULT_HANDSHAKE_PORT_BASE + 2 * num_devices
+        )
+        assert extra["notify_port"] == str(DEFAULT_NOTIFY_PORT_BASE + 2 * num_devices)
+
+    def test_setup_respects_overridden_bases(self):
+        backend = _make_backend(
+            extra_exp={
+                "MORI_HANDSHAKE_PORT_BASE": 7000,
+                "MORI_NOTIFY_PORT_BASE": 62000,
+            }
+        )
+        _setup(backend, rank=0)
+        extra = backend.kv_transfer_config["kv_connector_extra_config"]
+        assert extra["handshake_port"] == "7000"
+        assert extra["notify_port"] == "62000"
+
+    def test_requires_peer_binding(self):
+        assert MoRIIOConnectorBackend.requires_peer_binding is True
+
+    def test_concurrent_handoff_write_vs_read(self):
+        write_backend = _make_backend(read_mode=False)
+        read_backend = _make_backend(read_mode=True)
+        assert write_backend.concurrent_handoff is True
+        assert read_backend.concurrent_handoff is False
+        assert write_backend._read_mode is False
+        assert read_backend._read_mode is True
+
+    @pytest.mark.parametrize(
+        "value,expected_read",
+        [
+            ("true", True),
+            ("True", True),
+            ("1", True),
+            ("false", False),
+            ("0", False),
+            ("", False),
+        ],
+    )
+    def test_read_mode_parsing(self, value, expected_read):
+        backend = MoRIIOConnectorBackend(
+            llm_config=LLMConfig(
+                model_loading_config=dict(model_id="Qwen/Qwen3-0.6B"),
+                engine_kwargs=dict(
+                    kv_transfer_config=dict(
+                        kv_connector="MoRIIOConnector",
+                        kv_connector_extra_config={"read_mode": value},
+                    )
+                ),
+            ),
+        )
+        assert backend._read_mode is expected_read
+
+    def test_replica_metadata_returns_zmq(self):
+        backend = _make_backend()
+        _setup(backend, rank=0)
+        meta = backend.replica_metadata()
+        assert meta["mori_zmq_address"] == backend._zmq_address
+        # Parallelism is published so the orchestrator can address remote workers.
+        assert meta["dp_rank"] == 0 and meta["dp_size"] == 1 and meta["tp_size"] == 1
+
+    def test_replica_metadata_publishes_dp_tp(self):
+        backend = _make_backend(dp_rank=2, dp_size=4, tp_size=8)
+        _setup(backend, rank=0)
+        meta = backend.replica_metadata()
+        assert meta["dp_rank"] == 2
+        assert meta["dp_size"] == 4
+        assert meta["tp_size"] == 8
+
+    def test_replica_metadata_default_empty(self):
+        # The default backend publishes nothing (concrete default on the base).
+        assert BaseConnectorBackend.replica_metadata(None) == {}
+
+
+class TestMoRIIORequestId:
+    def _prepared_pair(self, backend, request, peer):
+        prefill = backend.prepare_prefill_request(request=request, peer=peer)
+        # prefill_response is unused in WRITE mode; pass a dummy with no params.
+        decode = backend.prepare_decode_request(
+            request=request,
+            peer=peer,
+            prefill_response=SimpleNamespace(kv_transfer_params=None),
+        )
+        return prefill, decode
+
+    def _request_with_copy(self, request_id="user-req-123"):
+        class _Req:
+            def __init__(self, rid):
+                self.request_id = rid
+                self.kv_transfer_params = None
+                self.max_tokens = 128
+                self.max_completion_tokens = 128
+                self.stream = True
+                self.stream_options = {"include_usage": True}
+
+            def model_copy(self, deep=False):
+                new = _Req(self.request_id)
+                return new
+
+        return _Req(request_id)
+
+    def test_prefill_and_decode_share_request_id_and_transfer_id(self):
+        backend = _make_backend(read_mode=False)
+        _setup(backend, rank=0)
+        decode_zmq = backend._zmq_address
+        prefill_zmq = "host:10.0.0.9,handshake:6301,notify:61005"
+        peer = {"mori_zmq_address": prefill_zmq}
+
+        req = self._request_with_copy("user-req-123")
+        prefill, decode = self._prepared_pair(backend, req, peer)
+
+        assert prefill.request_id == decode.request_id
+        assert (
+            prefill.kv_transfer_params["transfer_id"]
+            == decode.kv_transfer_params["transfer_id"]
+        )
+
+        # Round-trips to the right peer zmq via the vLLM regexes.
+        assert _PREFILL_ZMQ_RE.search(prefill.request_id).group(1) == prefill_zmq
+        assert _DECODE_ZMQ_RE.search(prefill.request_id) is not None
+        assert parse_peer_zmq(prefill.request_id, is_producer=False) == prefill_zmq
+        assert parse_peer_zmq(prefill.request_id, is_producer=True) == decode_zmq
+
+        # transfer_id format tx-<32hex>.
+        assert re.fullmatch(
+            r"tx-[0-9a-f]{32}", prefill.kv_transfer_params["transfer_id"]
+        )
+
+    def test_id_is_deterministic_across_calls(self):
+        backend = _make_backend(read_mode=False)
+        _setup(backend, rank=0)
+        peer = {"mori_zmq_address": "host:10.0.0.9,handshake:6301,notify:61005"}
+
+        p1 = backend.prepare_prefill_request(
+            request=self._request_with_copy("R"), peer=peer
+        )
+        p2 = backend.prepare_prefill_request(
+            request=self._request_with_copy("R"), peer=peer
+        )
+        assert p1.request_id == p2.request_id
+        assert (
+            p1.kv_transfer_params["transfer_id"] == p2.kv_transfer_params["transfer_id"]
+        )
+
+    def test_prefill_kv_params_write(self):
+        backend = _make_backend(read_mode=False)
+        _setup(backend, rank=0)
+        peer = {"mori_zmq_address": "host:10.0.0.9,handshake:6301,notify:61005"}
+        prefill = backend.prepare_prefill_request(
+            request=self._request_with_copy(), peer=peer
+        )
+        assert prefill.kv_transfer_params["do_remote_decode"] is True
+        assert prefill.kv_transfer_params["do_remote_prefill"] is False
+        assert prefill.max_tokens == 1
+        assert prefill.stream is False
+        # vLLM reads "tp_size" (not "remote_tp_size"); DP defaults at dp_size=1.
+        assert "remote_tp_size" not in prefill.kv_transfer_params
+        assert prefill.kv_transfer_params["tp_size"] == 1
+        assert prefill.kv_transfer_params["remote_dp_rank"] == 0
+        assert prefill.kv_transfer_params["remote_dp_size"] == 1
+
+    def test_decode_kv_params_write(self):
+        backend = _make_backend(read_mode=False)
+        _setup(backend, rank=0)
+        peer = {"mori_zmq_address": "host:10.0.0.9,handshake:6301,notify:61005"}
+        decode = backend.prepare_decode_request(
+            request=self._request_with_copy(),
+            peer=peer,
+            prefill_response=SimpleNamespace(kv_transfer_params=None),
+        )
+        assert decode.kv_transfer_params["do_remote_prefill"] is True
+        assert decode.kv_transfer_params["do_remote_decode"] is False
+        assert decode.kv_transfer_params["remote_block_ids"] is None
+        assert "remote_tp_size" not in decode.kv_transfer_params
+        assert decode.kv_transfer_params["tp_size"] == 1
+        assert decode.kv_transfer_params["remote_dp_rank"] == 0
+        assert decode.kv_transfer_params["remote_dp_size"] == 1
+
+    def test_dp_routing_targets_correct_ranks(self):
+        """With DP>1, the prefill request addresses the decode (this
+        orchestrator) rank; the decode request addresses the selected prefill
+        peer's rank (read from peer metadata)."""
+        # This orchestrator is decode dp_rank=1 of a 2-way DP decode group.
+        backend = _make_backend(read_mode=False, dp_rank=1, dp_size=2, tp_size=4)
+        _setup(backend, rank=0)
+        # The selected prefill peer is dp_rank=3 of a 4-way DP prefill group.
+        peer = {
+            "mori_zmq_address": "host:10.0.0.9,handshake:6301,notify:61005",
+            "dp_rank": 3,
+            "dp_size": 4,
+            "tp_size": 4,
+        }
+        prefill = backend.prepare_prefill_request(
+            request=self._request_with_copy(), peer=peer
+        )
+        decode = backend.prepare_decode_request(
+            request=self._request_with_copy(),
+            peer=peer,
+            prefill_response=SimpleNamespace(kv_transfer_params=None),
+        )
+        # Prefill engine's remote == this decode orchestrator (rank 1 of 2).
+        assert prefill.kv_transfer_params["remote_dp_rank"] == 1
+        assert prefill.kv_transfer_params["remote_dp_size"] == 2
+        # Decode engine's remote == the selected prefill peer (rank 3 of 4).
+        assert decode.kv_transfer_params["remote_dp_rank"] == 3
+        assert decode.kv_transfer_params["remote_dp_size"] == 4
+        assert decode.kv_transfer_params["tp_size"] == 4
+
+    def test_decode_kv_params_read_forwards_prefill_params(self):
+        backend = _make_backend(read_mode=True)
+        _setup(backend, rank=0)
+        peer = {"mori_zmq_address": "host:10.0.0.9,handshake:6301,notify:61005"}
+        prefill_resp = SimpleNamespace(
+            kv_transfer_params={
+                "remote_block_ids": [1, 2, 3],
+                "remote_engine_id": "eng-7",
+            }
+        )
+        decode = backend.prepare_decode_request(
+            request=self._request_with_copy(), peer=peer, prefill_response=prefill_resp
+        )
+        assert decode.kv_transfer_params["do_remote_prefill"] is True
+        assert decode.kv_transfer_params["remote_block_ids"] == [1, 2, 3]
+        assert decode.kv_transfer_params["remote_engine_id"] == "eng-7"
+        assert "transfer_id" in decode.kv_transfer_params
+        # READ also stamps the remote (prefill peer) routing.
+        assert decode.kv_transfer_params["remote_dp_rank"] == 0
+        assert decode.kv_transfer_params["remote_dp_size"] == 1
+        assert decode.kv_transfer_params["tp_size"] == 1
+
+    def test_decode_read_fallback_when_no_remote_params(self):
+        backend = _make_backend(read_mode=True)
+        _setup(backend, rank=0)
+        peer = {"mori_zmq_address": "host:10.0.0.9,handshake:6301,notify:61005"}
+        decode = backend.prepare_decode_request(
+            request=self._request_with_copy(),
+            peer=peer,
+            prefill_response=SimpleNamespace(kv_transfer_params=None),
+        )
+        assert decode.kv_transfer_params is None
+
+
+class TestMoRIIOZmqValidation:
+    def test_missing_peer_zmq_raises(self):
+        """A missing/empty peer mori_zmq_address must raise a clear error, not
+        silently build a request id containing "None"."""
+        backend = _make_backend(read_mode=False)
+        _setup(backend, rank=0)
+        request = TestMoRIIORequestId()._request_with_copy("user-req-123")
+        for peer in (None, {}, {"mori_zmq_address": ""}):
+            with pytest.raises(ValueError, match="mori_zmq_address"):
+                backend.prepare_prefill_request(request=request, peer=peer)
+
+
+class TestMoRIIOFactory:
+    def test_registered(self):
+        assert KVConnectorBackendFactory.is_registered("MoRIIOConnector")
+
+    def test_create_backend_returns_class(self):
+        backend_class = KVConnectorBackendFactory.get_backend_class("MoRIIOConnector")
+        assert backend_class is MoRIIOConnectorBackend
+        assert issubclass(backend_class, BaseConnectorBackend)
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main(["-v", __file__]))

python/ray/llm/tests/serve/cpu/deployments/llm/vllm/kv_transfer_backends/test_multi_connector.py

@@ -0,0 +1,181 @@
+import sys
+from unittest.mock import MagicMock, patch
+
+import pytest
+
+from ray.llm._internal.serve.engines.vllm.kv_transfer.base import (
+    BaseConnectorBackend,
+)
+from ray.llm._internal.serve.engines.vllm.kv_transfer.factory import (
+    KVConnectorBackendFactory,
+)
+from ray.llm._internal.serve.engines.vllm.kv_transfer.multi_connector import (
+    MultiConnectorBackend,
+)
+from ray.serve.llm import LLMConfig
+
+
+class TestMultiConnectorBackend:
+    """Test suite for MultiConnectorBackend."""
+
+    @pytest.fixture
+    def basic_llm_config(self):
+        """Fixture for basic LLM config with MultiConnector."""
+        return LLMConfig(
+            model_loading_config=dict(model_id="test-model"),
+            engine_kwargs=dict(
+                kv_transfer_config=dict(
+                    kv_connector="MultiConnector",
+                    kv_connector_extra_config=dict(
+                        connectors=[
+                            {"kv_connector": "LMCacheConnectorV1"},
+                            {"kv_connector": "NixlConnector"},
+                        ]
+                    ),
+                )
+            ),
+        )
+
+    @pytest.fixture
+    def multi_backend(self, basic_llm_config):
+        """Fixture for MultiConnectorBackend."""
+        return MultiConnectorBackend(basic_llm_config)
+
+    def test_multi_connector_initialization(self, multi_backend):
+        """Test that MultiConnectorBackend can be initialized."""
+        assert isinstance(multi_backend, MultiConnectorBackend)
+        assert isinstance(multi_backend, BaseConnectorBackend)
+
+    def test_setup_rejects_empty_connectors(self):
+        """An empty `connectors` list is a misconfig: setup() fails fast rather
+        than crash later when the orchestrator delegates to a missing primary."""
+        llm_config = LLMConfig(
+            model_loading_config=dict(model_id="test-model"),
+            engine_kwargs=dict(
+                kv_transfer_config=dict(
+                    kv_connector="MultiConnector",
+                    kv_connector_extra_config=dict(connectors=[]),
+                )
+            ),
+        )
+        with pytest.raises(ValueError, match="at least one sub-connector"):
+            MultiConnectorBackend(llm_config).setup()
+
+    def test_setup_calls_all_connectors(self, multi_backend):
+        """Test that setup calls setup on all configured connectors."""
+        mock_backend1 = MagicMock(spec=BaseConnectorBackend)
+        mock_backend2 = MagicMock(spec=BaseConnectorBackend)
+
+        with patch.object(
+            KVConnectorBackendFactory,
+            "create_backend",
+            side_effect=[mock_backend1, mock_backend2],
+        ) as mock_create:
+            multi_backend.setup()
+
+            assert mock_create.call_count == 2
+            mock_backend1.setup.assert_called_once()
+            mock_backend2.setup.assert_called_once()
+
+    def test_setup_raises_error_when_connector_missing_kv_connector(self):
+        """Test that setup raises ValueError when a connector is missing kv_connector."""
+        llm_config = LLMConfig(
+            model_loading_config=dict(model_id="test-model"),
+            engine_kwargs=dict(
+                kv_transfer_config=dict(
+                    kv_connector="MultiConnector",
+                    kv_connector_extra_config=dict(
+                        connectors=[
+                            {"some_other_key": "value"},
+                        ]
+                    ),
+                )
+            ),
+        )
+        backend = MultiConnectorBackend(llm_config)
+
+        with pytest.raises(ValueError, match="kv_connector is not set"):
+            backend.setup()
+
+    def test_setup_with_nested_multi_connector_raises_error(self):
+        """Test that nesting MultiConnector raises a ValueError."""
+        llm_config = LLMConfig(
+            model_loading_config=dict(model_id="test-model"),
+            engine_kwargs=dict(
+                kv_transfer_config=dict(
+                    kv_connector="MultiConnector",
+                    kv_connector_extra_config=dict(
+                        connectors=[
+                            {"kv_connector": "MultiConnector"},
+                        ]
+                    ),
+                )
+            ),
+        )
+        backend = MultiConnectorBackend(llm_config)
+        with pytest.raises(ValueError, match="Nesting MultiConnector"):
+            backend.setup()
+
+    def test_setup_passes_isolated_config_to_sub_connectors(self):
+        """Test that sub-connectors inherit parent config and receive their specific settings."""
+        llm_config = LLMConfig(
+            model_loading_config=dict(model_id="test-model"),
+            engine_kwargs=dict(
+                kv_transfer_config=dict(
+                    kv_connector="MultiConnector",
+                    engine_id="test-engine-123",
+                    kv_role="kv_both",
+                    kv_connector_extra_config=dict(
+                        connectors=[
+                            {
+                                "kv_connector": "LMCacheConnectorV1",
+                                "custom_param": "value1",
+                            },
+                            {"kv_connector": "NixlConnector", "custom_param": "value2"},
+                        ]
+                    ),
+                )
+            ),
+        )
+
+        captured_configs = []
+
+        def capture_config(name, config):
+            captured_configs.append((name, config.engine_kwargs["kv_transfer_config"]))
+            return MagicMock(spec=BaseConnectorBackend)
+
+        with patch.object(
+            KVConnectorBackendFactory, "create_backend", side_effect=capture_config
+        ):
+            MultiConnectorBackend(llm_config).setup()
+
+        assert len(captured_configs) == 2
+
+        # Verify each connector gets: inherited parent fields + its own specific config
+        expected_configs = [
+            (
+                "LMCacheConnectorV1",
+                {"kv_connector": "LMCacheConnectorV1", "custom_param": "value1"},
+            ),
+            (
+                "NixlConnector",
+                {"kv_connector": "NixlConnector", "custom_param": "value2"},
+            ),
+        ]
+
+        for (actual_name, actual_config), (expected_name, expected_specific) in zip(
+            captured_configs, expected_configs
+        ):
+            assert actual_name == expected_name
+            # Check inherited parent fields
+            assert actual_config["engine_id"] == "test-engine-123"
+            assert actual_config["kv_role"] == "kv_both"
+            # Check connector-specific fields
+            for key, value in expected_specific.items():
+                assert actual_config[key] == value
+            # Verify kv_connector_extra_config is not passed to sub-connectors
+            assert "kv_connector_extra_config" not in actual_config
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main(["-v", __file__]))

python/ray/llm/tests/serve/cpu/deployments/llm/vllm/kv_transfer_backends/test_nixl_connector.py

@@ -0,0 +1,126 @@
+import os
+import sys
+import uuid
+from types import SimpleNamespace
+from unittest.mock import patch
+
+import pytest
+
+from ray.llm._internal.serve.engines.vllm.kv_transfer.nixl import (
+    NixlConnectorBackend,
+)
+from ray.serve.exceptions import RayServeException
+from ray.serve.llm import LLMConfig
+from ray.serve.schema import ReplicaRank
+
+
+@pytest.fixture
+def engine_id():
+    """Fixture for the engine ID."""
+    return str(uuid.uuid4())
+
+
+class TestNixlConnectorBackend:
+    @pytest.fixture
+    def nixl_backend(self, engine_id: str):
+        """Fixture for the NixlConnectorBackend."""
+        return NixlConnectorBackend(
+            llm_config=LLMConfig(
+                model_loading_config=dict(
+                    model_id="Qwen/Qwen3-0.6B",
+                ),
+                engine_kwargs=dict(
+                    kv_transfer_config=dict(
+                        kv_connector="NixlConnector",
+                        kv_role="kv_both",
+                        engine_id=engine_id,
+                    )
+                ),
+            ),
+        )
+
+    @pytest.mark.parametrize(
+        "env_vars",
+        [
+            {},
+            {"VLLM_NIXL_SIDE_CHANNEL_PORT": "8080"},
+            {"VLLM_NIXL_SIDE_CHANNEL_HOST": "127.0.0.1"},
+            {
+                "VLLM_NIXL_SIDE_CHANNEL_PORT": "8080",
+                "VLLM_NIXL_SIDE_CHANNEL_HOST": "127.0.0.1",
+            },
+        ],
+    )
+    def test_setup_environment_variables(self, nixl_backend, env_vars, engine_id: str):
+        """Test that setup configures environment variables and overrides engine_id correctly."""
+        with patch.dict("os.environ", env_vars, clear=True), patch(
+            "ray.serve.get_replica_context",
+            return_value=self._replica_context(0),
+        ):
+            nixl_backend.setup()
+            assert "VLLM_NIXL_SIDE_CHANNEL_PORT" in os.environ
+            assert "VLLM_NIXL_SIDE_CHANNEL_HOST" in os.environ
+            assert engine_id in nixl_backend.kv_transfer_config["engine_id"]
+
+    @staticmethod
+    def _backend_with_port_base(base: int = 30000) -> NixlConnectorBackend:
+        return NixlConnectorBackend(
+            llm_config=LLMConfig(
+                model_loading_config=dict(model_id="Qwen/Qwen3-0.6B"),
+                engine_kwargs=dict(
+                    kv_transfer_config=dict(
+                        kv_connector="NixlConnector",
+                        kv_role="kv_both",
+                    )
+                ),
+                experimental_configs={"NIXL_SIDE_CHANNEL_PORT_BASE": base},
+            ),
+        )
+
+    @staticmethod
+    def _replica_context(global_rank: int) -> SimpleNamespace:
+        """Fake serve.get_replica_context() result with the given global rank."""
+        return SimpleNamespace(
+            rank=ReplicaRank(rank=global_rank, node_rank=0, local_rank=global_rank)
+        )
+
+    def test_default_side_channel_port_uses_configured_base(self):
+        """Rank 0 -> zero offset -> port equals the configured base."""
+        backend = self._backend_with_port_base(30000)
+        with (
+            patch.dict("os.environ", {}, clear=True),
+            patch(
+                "ray.serve.get_replica_context",
+                return_value=self._replica_context(0),
+            ),
+        ):
+            backend.setup()
+            assert os.environ["VLLM_NIXL_SIDE_CHANNEL_PORT"] == "30000"
+
+    def test_side_channel_port_offset_uses_replica_rank(self):
+        """Nonzero rank shifts the port by rank * num_devices (disjoint port blocks)."""
+        backend = self._backend_with_port_base(30000)
+        num_devices = backend.llm_config.get_engine_config().num_devices
+        with (
+            patch.dict("os.environ", {}, clear=True),
+            patch(
+                "ray.serve.get_replica_context",
+                return_value=self._replica_context(2),
+            ),
+        ):
+            backend.setup()
+            assert os.environ["VLLM_NIXL_SIDE_CHANNEL_PORT"] == str(
+                30000 + 2 * num_devices
+            )
+
+    def test_side_channel_port_requires_replica_context(self):
+        """Outside a replica, get_replica_context() raises -> setup fails loudly
+        instead of silently using a colliding 0 offset."""
+        backend = self._backend_with_port_base(30000)
+        with patch.dict("os.environ", {}, clear=True):
+            with pytest.raises(RayServeException):
+                backend.setup()
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main(["-v", __file__]))

python/ray/llm/tests/serve/cpu/deployments/llm/vllm/kv_transfer_backends/test_pd_protocol.py

@@ -0,0 +1,165 @@
+"""Tests for the P/D coordination protocol on the KV connector backends.
+
+Proves that NIXL, LMCache, and the default backend implement the abstract
+``BaseConnectorBackend`` protocol via ``DefaultPDProtocolMixin``, that Multi delegates
+the protocol to its top-most sub-connector, and that the abstract base itself cannot
+be instantiated.
+"""
+
+import sys
+from types import SimpleNamespace
+from unittest.mock import patch
+
+import pytest
+
+from ray.llm._internal.serve.core.configs.openai_api_models import (
+    ChatCompletionRequest,
+    CompletionRequest,
+)
+from ray.llm._internal.serve.engines.vllm.kv_transfer.base import (
+    BaseConnectorBackend,
+    DefaultConnectorBackend,
+    DefaultPDProtocolMixin,
+)
+from ray.llm._internal.serve.engines.vllm.kv_transfer.lmcache import (
+    LMCacheConnectorV1Backend,
+)
+from ray.llm._internal.serve.engines.vllm.kv_transfer.multi_connector import (
+    MultiConnectorBackend,
+)
+from ray.llm._internal.serve.engines.vllm.kv_transfer.nixl import (
+    NixlConnectorBackend,
+)
+from ray.serve.llm import LLMConfig
+
+
+def _llm_config(kv_connector: str) -> LLMConfig:
+    return LLMConfig(
+        model_loading_config=dict(model_id="Qwen/Qwen3-0.6B"),
+        engine_kwargs=dict(
+            kv_transfer_config=dict(
+                kv_connector=kv_connector,
+                kv_role="kv_both",
+            )
+        ),
+    )
+
+
+def test_base_connector_backend_is_abstract():
+    """``BaseConnectorBackend`` is abstract: its ``prepare_*`` methods are
+    abstractmethods, so direct instantiation raises ``TypeError``."""
+    with pytest.raises(TypeError):
+        BaseConnectorBackend(llm_config=None)
+
+
+@pytest.mark.parametrize(
+    "backend_factory",
+    [
+        lambda: NixlConnectorBackend(llm_config=_llm_config("NixlConnector")),
+        lambda: LMCacheConnectorV1Backend(llm_config=_llm_config("LMCacheConnectorV1")),
+        lambda: DefaultConnectorBackend(llm_config=None),
+    ],
+    ids=["nixl", "lmcache", "default"],
+)
+class TestConcreteBackendsProtocol:
+    """The concrete backends expose the default P/D protocol shaping."""
+
+    def test_is_concrete_subclass(self, backend_factory):
+        be = backend_factory()
+        assert isinstance(be, BaseConnectorBackend)
+        assert isinstance(be, DefaultPDProtocolMixin)
+        # Default flags == standard (no-peer, sequential) policy.
+        assert be.requires_peer_binding is False
+        assert be.concurrent_handoff is False
+
+    def test_prepare_prefill_shaping(self, backend_factory):
+        be = backend_factory()
+        request = ChatCompletionRequest(
+            model="test-model",
+            messages=[{"role": "user", "content": "hello"}],
+            max_completion_tokens=32,
+            stream=True,
+            stream_options={"include_usage": True},
+        )
+        prefill = be.prepare_prefill_request(request=request, peer=None)
+        assert prefill.kv_transfer_params["do_remote_decode"] is True
+        assert prefill.kv_transfer_params["do_remote_prefill"] is False
+        assert prefill.max_tokens == 1
+        assert prefill.max_completion_tokens == 1
+        assert prefill.stream is False
+        assert prefill.stream_options is None
+        # Original request untouched.
+        assert request.max_completion_tokens == 32
+        assert request.stream is True
+
+    def test_prepare_decode_forwards_params(self, backend_factory):
+        be = backend_factory()
+        request = CompletionRequest(model="test-model", prompt="hi")
+        chunk = SimpleNamespace(kv_transfer_params={"remote_engine_id": "p1"})
+        decode = be.prepare_decode_request(
+            request=request, peer=None, prefill_response=chunk
+        )
+        assert decode.kv_transfer_params == {"remote_engine_id": "p1"}
+
+    def test_prepare_decode_none_prefill_response_no_crash(self, backend_factory):
+        """Concurrent-handoff mode passes ``prefill_response=None``: must not
+        crash and must leave kv_transfer_params unset (the gemini None-guard)."""
+        be = backend_factory()
+        request = CompletionRequest(model="test-model", prompt="hi")
+        decode = be.prepare_decode_request(
+            request=request, peer=None, prefill_response=None
+        )
+        assert getattr(decode, "kv_transfer_params", None) is None
+
+
+class TestMultiConnectorDelegation:
+    """MultiConnectorBackend delegates the P/D protocol to its top-most
+    sub-connector rather than inheriting the default mixin."""
+
+    def _multi_with_primary(self, primary):
+        multi = MultiConnectorBackend(llm_config=_llm_config("MultiConnector"))
+        multi._connector_backends = [primary]
+        return multi
+
+    def test_no_subconnectors_flags_false(self):
+        multi = MultiConnectorBackend(llm_config=_llm_config("MultiConnector"))
+        assert multi.requires_peer_binding is False
+        assert multi.concurrent_handoff is False
+
+    def test_delegates_prepare_to_primary(self):
+        multi = self._multi_with_primary(DefaultConnectorBackend(llm_config=None))
+        request = ChatCompletionRequest(
+            model="test-model",
+            messages=[{"role": "user", "content": "hello"}],
+            max_completion_tokens=8,
+        )
+        prefill = multi.prepare_prefill_request(request=request, peer=None)
+        assert prefill.kv_transfer_params["do_remote_decode"] is True
+        assert prefill.max_tokens == 1
+        chunk = SimpleNamespace(kv_transfer_params={"remote_engine_id": "p1"})
+        decode = multi.prepare_decode_request(
+            request=CompletionRequest(model="test-model", prompt="hi"),
+            peer=None,
+            prefill_response=chunk,
+        )
+        assert decode.kv_transfer_params == {"remote_engine_id": "p1"}
+
+    def test_flags_follow_primary(self):
+        # A primary that opts into peer binding + concurrent handoff governs the group.
+        primary = SimpleNamespace(requires_peer_binding=True, concurrent_handoff=True)
+        multi = self._multi_with_primary(primary)
+        assert multi.requires_peer_binding is True
+        assert multi.concurrent_handoff is True
+
+
+@patch(
+    "ray.llm._internal.serve.engines.vllm.kv_transfer.lmcache._check_lmcache_installed"
+)
+def test_lmcache_setup_still_works(_mock_check):
+    """The P/D protocol must not break the connector-specific setup() behavior."""
+    be = LMCacheConnectorV1Backend(llm_config=_llm_config("LMCacheConnectorV1"))
+    be.setup()  # no-op path (no kv_connector_extra_config), must not raise
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main(["-v", __file__]))

python/ray/llm/tests/serve/cpu/deployments/llm/vllm/kv_transfer_backends/test_registry.py

@@ -0,0 +1,111 @@
+"""Unit tests for ComponentRegistry."""
+
+import sys
+
+import pytest
+
+from ray.llm._internal.serve.utils.registry import ComponentRegistry, get_registry
+
+
+class TestComponentRegistry:
+    """Test suite for ComponentRegistry."""
+
+    def test_register_and_get_direct_class(self):
+        """Test registering and retrieving a class directly."""
+        registry = ComponentRegistry("test_category")
+        test_class = type("TestClass", (), {})
+
+        registry.register("test_component", test_class)
+        assert registry.contains("test_component")
+        retrieved = registry.get("test_component")
+        assert retrieved == test_class
+
+    def test_register_and_get_module_path(self):
+        """Test registering and retrieving via module path."""
+        registry = ComponentRegistry("test_category")
+
+        registry.register(
+            "test_component",
+            "ray.llm._internal.serve.utils.registry:ComponentRegistry",
+        )
+        assert registry.contains("test_component")
+        retrieved = registry.get("test_component")
+        assert retrieved == ComponentRegistry
+
+    def test_get_nonexistent_component_raises(self):
+        """Test that getting a non-existent component raises ValueError."""
+        registry = ComponentRegistry("test_category")
+
+        with pytest.raises(ValueError, match="not found"):
+            registry.get("nonexistent")
+
+    def test_invalid_string_format_raises(self):
+        """Test that registering with invalid string format raises ValueError."""
+        registry = ComponentRegistry("test_category")
+
+        with pytest.raises(ValueError, match="Invalid format"):
+            registry.register("test_comp", "invalid_format_no_colon")
+
+    def test_double_registration_raises(self):
+        """Test that double registration raises ValueError."""
+        registry = ComponentRegistry("test_category")
+        test_class1 = type("TestClass1", (), {})
+        test_class2 = type("TestClass2", (), {})
+
+        registry.register("test_component", test_class1)
+
+        with pytest.raises(ValueError, match="already registered"):
+            registry.register("test_component", test_class2)
+
+        # Verify original registration is still intact
+        assert registry.get("test_component") == test_class1
+
+    def test_reregister_after_unregister(self):
+        """Test that unregistering allows re-registration."""
+        registry = ComponentRegistry("test_category")
+        test_class1 = type("TestClass1", (), {})
+        test_class2 = type("TestClass2", (), {})
+
+        registry.register("test_component", test_class1)
+        registry.unregister("test_component")
+        registry.register("test_component", test_class2)
+
+        assert registry.get("test_component") == test_class2
+
+    def test_get_registry_singleton(self):
+        """Test that get_registry returns the same instance for the same category."""
+        registry1 = get_registry("test_category")
+        registry2 = get_registry("test_category")
+
+        assert registry1 is registry2
+        assert registry1.category == "test_category"
+
+    def test_get_registry_different_categories(self):
+        """Test that get_registry returns different instances for different categories."""
+        registry1 = get_registry("category1")
+        registry2 = get_registry("category2")
+
+        assert registry1 is not registry2
+        assert registry1.category == "category1"
+        assert registry2.category == "category2"
+
+    def test_unregister(self):
+        """Test unregistering a component."""
+        registry = ComponentRegistry("test_category")
+        test_class = type("TestClass", (), {})
+
+        # Register and verify it exists
+        registry.register("test_component", test_class)
+        assert registry.contains("test_component")
+
+        # Unregister and verify it's removed
+        registry.unregister("test_component")
+        assert not registry.contains("test_component")
+
+        # Verify get raises ValueError
+        with pytest.raises(ValueError, match="not found"):
+            registry.get("test_component")
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main(["-v", __file__]))

python/ray/llm/tests/serve/cpu/deployments/prefill_decode_disagg/test_pd_direct_streaming.py

@@ -0,0 +1,79 @@
+import sys
+
+import pytest
+
+from ray.llm._internal.serve.core.configs.llm_config import (
+    LLMConfig,
+    ModelLoadingConfig,
+)
+from ray.llm._internal.serve.serving_patterns.prefill_decode.builder import (
+    build_pd_openai_app,
+)
+from ray.llm.tests.serve.cpu.deployments.utils.direct_streaming_utils import (
+    consistent_hash_deployment_config,
+    requires_direct_streaming,
+    run_app_through_haproxy,
+    session_chat_response,
+)
+
+
+@requires_direct_streaming
+class TestPDDirectStreamingConsistentHashRouting:
+    """Session affinity over the full PD direct-streaming path.
+
+    The decode server is the ingress; LLMRouter pins it via ConsistentHashRouter
+    and forwards the session id to the prefill handle, so both legs pin per
+    session. The decode replica comes from the ``x-replica-id`` header; the
+    prefill replica from ``kv_transfer_params.remote_engine_id`` (stamped by the
+    prefill, echoed by the decode).
+    """
+
+    @pytest.fixture(name="llm_config")
+    def _llm_config(self):
+        return LLMConfig(model_loading_config=ModelLoadingConfig(model_id="test-model"))
+
+    @pytest.fixture(name="base_url")
+    def run_pd_app(
+        self,
+        llm_config_with_mock_engine,
+        shutdown_ray_and_serve,
+        disable_placement_bundles,
+    ):
+        def _pd_config():
+            config = llm_config_with_mock_engine.model_copy(deep=True)
+            config.engine_kwargs = {
+                "kv_transfer_config": {
+                    "kv_connector": "NixlConnector",
+                    "kv_role": "kv_both",
+                }
+            }
+            config.deployment_config = consistent_hash_deployment_config()
+            return config
+
+        yield run_app_through_haproxy(
+            build_pd_openai_app(
+                {"prefill_config": _pd_config(), "decode_config": _pd_config()}
+            )
+        )
+
+    def _serving_replicas(self, base_url, session_id):
+        """Return the (decode, prefill) replicas that served a ``session_id`` request."""
+        resp = session_chat_response(base_url, session_id)
+        decode_replica = resp.headers["x-replica-id"]
+        prefill_replica = resp.json()["kv_transfer_params"]["remote_engine_id"]
+        return decode_replica, prefill_replica
+
+    def test_session_affinity(self, base_url):
+        pairs = {self._serving_replicas(base_url, "test-session-id") for _ in range(10)}
+        assert len(pairs) == 1
+
+    def test_different_sessions_spread(self, base_url):
+        pairs = [
+            self._serving_replicas(base_url, f"test-session-id-{i}") for i in range(10)
+        ]
+        assert len({decode for decode, _ in pairs}) > 1
+        assert len({prefill for _, prefill in pairs}) > 1
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main(["-v", __file__]))

python/ray/llm/tests/serve/cpu/deployments/prefill_decode_disagg/test_prefill_decode_disagg.py

@@ -0,0 +1,905 @@
+import asyncio
+import sys
+import warnings
+from types import SimpleNamespace
+from unittest.mock import patch
+
+import pytest
+
+from ray.llm._internal.serve.core.configs.llm_config import (
+    LLMConfig,
+    ModelLoadingConfig,
+)
+from ray.llm._internal.serve.core.configs.openai_api_models import (
+    ChatCompletionRequest,
+    CompletionRequest,
+)
+from ray.llm._internal.serve.core.ingress.builder import (
+    IngressClsConfig,
+)
+from ray.llm._internal.serve.core.ingress.ingress import OpenAiIngress
+from ray.llm._internal.serve.core.server.llm_server import LLMServer
+from ray.llm._internal.serve.serving_patterns.prefill_decode.builder import (
+    PDServingArgs,
+    build_pd_openai_app,
+)
+from ray.llm._internal.serve.serving_patterns.prefill_decode.pd_server import (
+    PDDecodeServer,
+    PDPrefillServer,
+)
+from ray.serve._private.http_util import SERVE_SESSION_ID
+
+
+async def _aiter(items):
+    for item in items:
+        yield item
+
+
+class _FakePrefillHandle:
+    """Fake prefill DeploymentHandle. Records each .chat/.completions remote
+    call with the session_id from any preceding ``.options(session_id=...)``,
+    and yields one chunk with kv_transfer_params back to the orchestrator."""
+
+    def __init__(self, calls=None, session_id=None):
+        self.calls = calls if calls is not None else []
+        self.session_id = session_id
+
+    def options(self, **kwargs):
+        return _FakePrefillHandle(
+            calls=self.calls,
+            session_id=kwargs.get("session_id", self.session_id),
+        )
+
+    def _method(self, name):
+        def remote(request, raw_request_info):
+            self.calls.append(
+                {"method": name, "request": request, "session_id": self.session_id}
+            )
+            return _aiter(
+                [SimpleNamespace(kv_transfer_params={"remote_engine_id": "prefill-1"})]
+            )
+
+        return SimpleNamespace(remote=remote)
+
+    @property
+    def chat(self):
+        return self._method("chat")
+
+    @property
+    def completions(self):
+        return self._method("completions")
+
+
+class TestPDServingArgs:
+    """Test suite for PDServingArgs data model."""
+
+    @pytest.fixture
+    def pd_configs(self):
+        """Prefill and decode configs with required kv_transfer_config."""
+        base_config = {
+            "model_loading_config": {
+                "model_id": "test-model",
+                "model_source": "test-source",
+            },
+            "engine_kwargs": {
+                "kv_transfer_config": {
+                    "kv_connector": "NixlConnector",
+                    "kv_role": "kv_both",
+                },
+            },
+        }
+        prefill = LLMConfig.model_validate(base_config)
+        decode = LLMConfig.model_validate(base_config)
+        return prefill, decode
+
+    def test_basic_creation_and_defaults(self, pd_configs):
+        """Test creation with minimal config and verify defaults."""
+        prefill, decode = pd_configs
+        args = PDServingArgs(prefill_config=prefill, decode_config=decode)
+
+        # Verify configs
+        assert isinstance(args.prefill_config, LLMConfig)
+        assert isinstance(args.decode_config, LLMConfig)
+
+        # TODO(Kourosh): Deprecated, remove in Ray 2.58.
+        assert args.proxy_cls_config is None
+        assert args.proxy_deployment_config is None
+        assert isinstance(args.ingress_cls_config, IngressClsConfig)
+        assert args.ingress_cls_config.ingress_cls == OpenAiIngress
+        assert args.ingress_deployment_config == {}
+
+    def test_flexible_input_types(self):
+        """Test accepts dicts for prefill and decode configs."""
+        config_dict = {
+            "model_loading_config": {
+                "model_id": "test-model",
+                "model_source": "test-source",
+            },
+            "engine_kwargs": {
+                "kv_transfer_config": {
+                    "kv_connector": "NixlConnector",
+                    "kv_role": "kv_both",
+                },
+            },
+        }
+        args = PDServingArgs(prefill_config=config_dict, decode_config=config_dict)
+        assert isinstance(args.prefill_config, LLMConfig)
+        assert isinstance(args.decode_config, LLMConfig)
+
+    # TODO(Kourosh): Deprecated, remove in Ray 2.58.
+    def test_proxy_config_deprecated(self, pd_configs):
+        """Test proxy_cls_config and proxy_deployment_config emit deprecation warnings."""
+        prefill, decode = pd_configs
+
+        # proxy_cls_config as dict should warn
+        with warnings.catch_warnings(record=True) as w:
+            warnings.simplefilter("always")
+            PDServingArgs(
+                prefill_config=prefill,
+                decode_config=decode,
+                proxy_cls_config={"proxy_extra_kwargs": {"key": "value"}},
+            )
+            deprecation_msgs = [
+                str(warning.message)
+                for warning in w
+                if issubclass(warning.category, DeprecationWarning)
+            ]
+            assert any(
+                "proxy_cls_config is deprecated" in msg for msg in deprecation_msgs
+            )
+
+        # proxy_deployment_config should warn
+        with warnings.catch_warnings(record=True) as w:
+            warnings.simplefilter("always")
+            PDServingArgs(
+                prefill_config=prefill,
+                decode_config=decode,
+                proxy_deployment_config={"num_replicas": 2},
+            )
+            deprecation_msgs = [
+                str(warning.message)
+                for warning in w
+                if issubclass(warning.category, DeprecationWarning)
+            ]
+            assert any(
+                "proxy_deployment_config is deprecated" in msg
+                for msg in deprecation_msgs
+            )
+
+    def test_ingress_config_flexibility(self, pd_configs):
+        """Test ingress_cls_config: defaults, dict input, object input, and class loading."""
+        prefill, decode = pd_configs
+
+        # Test defaults
+        args_default = PDServingArgs(prefill_config=prefill, decode_config=decode)
+        assert isinstance(args_default.ingress_cls_config, IngressClsConfig)
+        assert args_default.ingress_cls_config.ingress_cls == OpenAiIngress
+        assert args_default.ingress_cls_config.ingress_extra_kwargs == {}
+
+        # Test as dict with custom kwargs
+        args_dict = PDServingArgs(
+            prefill_config=prefill,
+            decode_config=decode,
+            ingress_cls_config={"ingress_extra_kwargs": {"key": "value"}},
+        )
+        assert isinstance(args_dict.ingress_cls_config, IngressClsConfig)
+        assert args_dict.ingress_cls_config.ingress_extra_kwargs == {"key": "value"}
+
+        # Test as object
+        args_obj = PDServingArgs(
+            prefill_config=prefill,
+            decode_config=decode,
+            ingress_cls_config=IngressClsConfig(ingress_extra_kwargs={"key": "value"}),
+        )
+        assert isinstance(args_obj.ingress_cls_config, IngressClsConfig)
+        assert args_obj.ingress_cls_config.ingress_extra_kwargs == {"key": "value"}
+
+        # Test class loading from string
+        args_str = PDServingArgs(
+            prefill_config=prefill,
+            decode_config=decode,
+            ingress_cls_config={
+                "ingress_cls": "ray.llm._internal.serve.core.ingress.ingress:OpenAiIngress"
+            },
+        )
+        assert args_str.ingress_cls_config.ingress_cls == OpenAiIngress
+
+    def test_validation_rules(self):
+        """Test validation: matching model IDs and required kv_transfer_config."""
+        # Mismatched model IDs
+        prefill = LLMConfig(
+            model_loading_config=ModelLoadingConfig(
+                model_id="model-1", model_source="source"
+            ),
+            engine_kwargs={"kv_transfer_config": {"kv_connector": "NixlConnector"}},
+        )
+        decode = LLMConfig(
+            model_loading_config=ModelLoadingConfig(
+                model_id="model-2", model_source="source"
+            ),
+            engine_kwargs={"kv_transfer_config": {"kv_connector": "NixlConnector"}},
+        )
+        with pytest.raises(ValueError, match="P/D model id mismatch"):
+            PDServingArgs(prefill_config=prefill, decode_config=decode)
+
+        # Missing kv_transfer_config
+        prefill_no_kv = LLMConfig(
+            model_loading_config=ModelLoadingConfig(
+                model_id="test-model", model_source="test-source"
+            )
+        )
+        decode_no_kv = LLMConfig(
+            model_loading_config=ModelLoadingConfig(
+                model_id="test-model", model_source="test-source"
+            )
+        )
+        with pytest.raises(ValueError, match="kv_transfer_config is required"):
+            PDServingArgs(prefill_config=prefill_no_kv, decode_config=decode_no_kv)
+
+
+class TestServingArgsParsing:
+    @pytest.mark.parametrize("kv_connector", ["NixlConnector", "LMCacheConnectorV1"])
+    def test_parse_dict(self, kv_connector: str):
+        prefill_config = LLMConfig(
+            model_loading_config=dict(
+                model_id="qwen-0.5b",
+                model_source="Qwen/Qwen2.5-0.5B-Instruct",
+            ),
+            deployment_config=dict(
+                autoscaling_config=dict(
+                    min_replicas=2,
+                    max_replicas=2,
+                )
+            ),
+            engine_kwargs=dict(
+                tensor_parallel_size=1,
+                kv_transfer_config=dict(
+                    kv_connector=kv_connector,
+                    kv_role="kv_both",
+                ),
+            ),
+        )
+
+        decode_config = LLMConfig(
+            model_loading_config=dict(
+                model_id="qwen-0.5b",
+                model_source="Qwen/Qwen2.5-0.5B-Instruct",
+            ),
+            deployment_config=dict(
+                autoscaling_config=dict(
+                    min_replicas=1,
+                    max_replicas=1,
+                )
+            ),
+            engine_kwargs=dict(
+                tensor_parallel_size=1,
+                kv_transfer_config=dict(
+                    kv_connector=kv_connector,
+                    kv_role="kv_both",
+                ),
+            ),
+        )
+
+        pd_config = {"prefill_config": prefill_config, "decode_config": decode_config}
+
+        app = build_pd_openai_app(pd_config)
+        assert app is not None
+
+
+class TestPDOrchestratorMixin:
+    def test_prepare_prefill_request_limits_chat_to_one_token(self):
+        from ray.llm._internal.serve.engines.vllm.kv_transfer.base import (
+            DefaultConnectorBackend,
+        )
+
+        request = ChatCompletionRequest(
+            model="test-model",
+            messages=[{"role": "user", "content": "hello"}],
+            max_completion_tokens=32,
+            stream=True,
+            stream_options={"include_usage": True},
+        )
+
+        be = DefaultConnectorBackend(llm_config=None)
+        prefill_request = be.prepare_prefill_request(request=request, peer=None)
+
+        assert prefill_request.max_tokens == 1
+        assert prefill_request.max_completion_tokens == 1
+        assert prefill_request.stream is False
+        assert prefill_request.stream_options is None
+        assert request.max_completion_tokens == 32
+        assert request.stream is True
+
+    @pytest.mark.asyncio
+    @pytest.mark.parametrize(
+        "method,path,body",
+        [
+            (
+                "chat",
+                "/v1/chat/completions",
+                {"messages": [{"role": "user", "content": "hi"}]},
+            ),
+            ("completions", "/v1/completions", {"prompt": "hi"}),
+        ],
+    )
+    async def test_direct_streaming_http_runs_pd_orchestration(
+        self, method, path, body
+    ):
+        """HTTP traffic to PDDecodeServer's direct-streaming ASGI app must
+        flow through PD orchestration (remote prefill, then local decode),
+        propagate the session-id header to the prefill handle, and pass
+        the prefill's kv_transfer_params to the local decode call.
+        Regression for https://github.com/ray-project/ray/pull/63679.
+        """
+        from fastapi.testclient import TestClient
+
+        from ray.llm.tests.serve.mocks.mock_vllm_engine import MockVLLMEngine
+
+        server = PDDecodeServer.__new__(PDDecodeServer)
+        server._prefill_handle = _FakePrefillHandle()
+        server._llm_config = LLMConfig(
+            model_loading_config=ModelLoadingConfig(model_id="test-model")
+        )
+        # Engine init stores the backend on the config; the orchestrator reads it.
+        from ray.llm._internal.serve.engines.vllm.kv_transfer.base import (
+            DefaultConnectorBackend,
+        )
+
+        server._llm_config._kv_connector_backend = DefaultConnectorBackend(
+            server._llm_config
+        )
+        # The direct-streaming app starts from the engine-native ASGI app, so
+        # the decode server needs a (mock) engine. PD only re-points the
+        # chat/completions routes at the orchestrator, patched below.
+        server.engine = MockVLLMEngine(server._llm_config)
+        await server.engine.start()
+
+        decode_calls = []
+
+        async def _fake_decode(self, req, raw_info):
+            decode_calls.append(req)
+            return _aiter([['data: {"ok":true}\n\n']])
+
+        app = await server.__serve_build_asgi_app__()
+        with patch.object(LLMServer, method, _fake_decode):
+            with TestClient(app) as client:
+                resp = client.post(
+                    path,
+                    json={"model": "test-model", "stream": True, **body},
+                    headers={SERVE_SESSION_ID: "session-a"},
+                )
+
+        assert resp.status_code == 200, resp.text
+        assert server._prefill_handle.calls[0]["method"] == method
+        assert server._prefill_handle.calls[0]["session_id"] == "session-a"
+        assert decode_calls[0].kv_transfer_params == {"remote_engine_id": "prefill-1"}
+
+
+class _ChooseReplicaPrefillHandle:
+    """Fake prefill DeploymentHandle exercising the choose_replica/dispatch path.
+
+    Mirrors the connector-protocol opt-in flow: the orchestrator opens a
+    ``choose_replica`` async context manager, reads ``selection.replica_metadata``,
+    then calls ``dispatch(selection, request, raw_info)``.
+    """
+
+    def __init__(self, calls=None, replica_metadata=None):
+        self.calls = calls if calls is not None else []
+        self._replica_metadata = (
+            replica_metadata if replica_metadata is not None else {"peer": "prefill-7"}
+        )
+
+    def options(self, **kwargs):
+        return self
+
+    def _method(self, name):
+        handle = self
+
+        class _Selection:
+            replica_metadata = handle._replica_metadata
+
+        class _Ctx:
+            async def __aenter__(self_inner):
+                return _Selection()
+
+            async def __aexit__(self_inner, *exc):
+                return False
+
+        def choose_replica(request):
+            handle.calls.append({"phase": "choose_replica", "method": name})
+            return _Ctx()
+
+        def dispatch(selection, request, raw_request_info):
+            handle.calls.append(
+                {"phase": "dispatch", "method": name, "request": request}
+            )
+            return _aiter(
+                [SimpleNamespace(kv_transfer_params={"remote_engine_id": "prefill-7"})]
+            )
+
+        return SimpleNamespace(choose_replica=choose_replica, dispatch=dispatch)
+
+    @property
+    def chat(self):
+        return self._method("chat")
+
+    @property
+    def completions(self):
+        return self._method("completions")
+
+
+class TestConnectorProtocolHook:
+    """The orchestrator delegates request shaping + handoff to the backend."""
+
+    def test_base_connector_backend_is_abstract(self):
+        """``BaseConnectorBackend`` is abstract and cannot be instantiated:
+        ``prepare_prefill_request`` / ``prepare_decode_request`` are abstract."""
+        from ray.llm._internal.serve.engines.vllm.kv_transfer.base import (
+            BaseConnectorBackend,
+        )
+
+        with pytest.raises(TypeError):
+            BaseConnectorBackend(llm_config=None)
+
+    def test_default_protocol_mixin_shaping(self):
+        """The ``DefaultPDProtocolMixin`` policy: prefill stamps the standard
+        kv_transfer_params + clamps to one non-streaming token; decode forwards
+        the prefill response's kv_transfer_params, and tolerates a None prefill
+        response (concurrent-handoff mode) without crashing."""
+        from ray.llm._internal.serve.engines.vllm.kv_transfer.base import (
+            DefaultConnectorBackend,
+        )
+
+        request = ChatCompletionRequest(
+            model="test-model",
+            messages=[{"role": "user", "content": "hello"}],
+            max_completion_tokens=16,
+            stream=True,
+            stream_options={"include_usage": True},
+        )
+
+        be = DefaultConnectorBackend(llm_config=None)
+
+        prefill = be.prepare_prefill_request(
+            request=request.model_copy(deep=True), peer=None
+        )
+        assert prefill.kv_transfer_params["do_remote_decode"] is True
+        assert prefill.kv_transfer_params["do_remote_prefill"] is False
+        assert prefill.max_tokens == 1
+        assert prefill.max_completion_tokens == 1
+        assert prefill.stream is False
+        assert prefill.stream_options is None
+        # Original untouched.
+        assert request.max_completion_tokens == 16
+        assert request.stream is True
+
+        chunk = SimpleNamespace(kv_transfer_params={"remote_engine_id": "p1"})
+        decode = be.prepare_decode_request(
+            request=request.model_copy(deep=True), peer=None, prefill_response=chunk
+        )
+        assert decode.kv_transfer_params == {"remote_engine_id": "p1"}
+
+        # None prefill_response (concurrent mode) must not crash and leaves
+        # kv_transfer_params unset.
+        decode_none = be.prepare_decode_request(
+            request=request.model_copy(deep=True), peer=None, prefill_response=None
+        )
+        assert getattr(decode_none, "kv_transfer_params", None) is None
+
+    def test_get_connector_backend_returns_stored_backend(self):
+        """``_get_connector_backend`` returns the backend that engine init stored
+        on the LLMConfig (and caches it); asserts if none was stored."""
+        from ray.llm._internal.serve.engines.vllm.kv_transfer.base import (
+            BaseConnectorBackend,
+        )
+        from ray.llm._internal.serve.engines.vllm.kv_transfer.nixl import (
+            NixlConnectorBackend,
+        )
+
+        server = PDDecodeServer.__new__(PDDecodeServer)
+        server._llm_config = LLMConfig(
+            model_loading_config=ModelLoadingConfig(model_id="test-model"),
+            engine_kwargs={
+                "kv_transfer_config": {
+                    "kv_connector": "NixlConnector",
+                    "kv_role": "kv_both",
+                }
+            },
+        )
+
+        # No backend stored (engine init / setup_engine_backend didn't run) -> assert.
+        with pytest.raises(AssertionError):
+            server._get_connector_backend()
+
+        # The backend setup_engine_backend would store is returned.
+        stored = NixlConnectorBackend(llm_config=server._llm_config)
+        assert isinstance(stored, BaseConnectorBackend)
+        server._llm_config._kv_connector_backend = stored
+        assert server._get_connector_backend() is stored
+
+        # Cached on first access: a later config change isn't re-read.
+        server._llm_config._kv_connector_backend = None
+        assert server._get_connector_backend() is stored
+
+    @pytest.mark.asyncio
+    async def test_peer_binding_concurrent_handoff_takes_choose_replica_path(self):
+        """A backend opting into requires_peer_binding + concurrent_handoff must
+        drive the orchestrator down the choose_replica/dispatch + concurrent
+        local-decode path, calling the backend's prepare_* with the peer."""
+        from ray.llm._internal.serve.engines.vllm.kv_transfer.base import (
+            BaseConnectorBackend,
+        )
+
+        seen = {}
+
+        class _DummyBackend(BaseConnectorBackend):
+            requires_peer_binding = True
+            concurrent_handoff = True
+
+            def prepare_prefill_request(self, *, request, peer):
+                seen["prefill_peer"] = peer
+                out = request.model_copy(deep=True)
+                out.kv_transfer_params = {"role": "prefill", "peer": peer}
+                return out
+
+            def prepare_decode_request(self, *, request, peer, prefill_response):
+                seen["decode_peer"] = peer
+                seen["prefill_response"] = prefill_response
+                out = request.model_copy(deep=True)
+                out.kv_transfer_params = {"role": "decode", "peer": peer}
+                return out
+
+        server = PDDecodeServer.__new__(PDDecodeServer)
+        server._llm_config = LLMConfig(
+            model_loading_config=ModelLoadingConfig(model_id="test-model")
+        )
+        dummy_backend = _DummyBackend(server._llm_config)
+        server._llm_config._kv_connector_backend = dummy_backend
+        prefill = _ChooseReplicaPrefillHandle(replica_metadata={"peer": "prefill-7"})
+        server._prefill_handle = prefill
+
+        decode_calls = []
+
+        async def _fake_super_completions(self, req, raw_info):
+            decode_calls.append(req)
+            return _aiter(["decode-chunk"])
+
+        request = CompletionRequest(model="test-model", prompt="hi")
+
+        # Patch the super() local-decode target (LLMServer.completions in the MRO).
+        with patch.object(LLMServer, "completions", _fake_super_completions):
+            chunks = [c async for c in server._pd_handle_request(request, None)]
+
+        # choose_replica + dispatch were used (not .remote()).
+        phases = [c["phase"] for c in prefill.calls]
+        assert phases == ["choose_replica", "dispatch"], phases
+        # Backend saw the peer metadata from the selection on both prepares.
+        assert seen["prefill_peer"] == {"peer": "prefill-7"}
+        assert seen["decode_peer"] == {"peer": "prefill-7"}
+        # Concurrent handoff -> no prefill chunk captured before decode.
+        assert seen["prefill_response"] is None
+        # Local decode ran with the backend-shaped decode request.
+        assert decode_calls[0].kv_transfer_params == {
+            "role": "decode",
+            "peer": {"peer": "prefill-7"},
+        }
+        assert chunks == ["decode-chunk"]
+
+    @pytest.mark.asyncio
+    async def test_default_nixl_backend_shapes_prefill_and_forwards_decode(self):
+        """End-to-end (mock handle) default path through a resolved NIXL backend:
+        prefill is shaped (max_tokens=1, do_remote_decode), and decode forwards
+        the prefill chunk's kv_transfer_params."""
+        server = PDDecodeServer.__new__(PDDecodeServer)
+        server._llm_config = LLMConfig(
+            model_loading_config=ModelLoadingConfig(model_id="test-model"),
+            engine_kwargs={
+                "kv_transfer_config": {
+                    "kv_connector": "NixlConnector",
+                    "kv_role": "kv_both",
+                }
+            },
+        )
+        # Engine init (setup_engine_backend) stores the backend on the config;
+        # the orchestrator reads it from there.
+        from ray.llm._internal.serve.engines.vllm.kv_transfer.nixl import (
+            NixlConnectorBackend,
+        )
+
+        server._llm_config._kv_connector_backend = NixlConnectorBackend(
+            server._llm_config
+        )
+        prefill = _FakePrefillHandle()
+        server._prefill_handle = prefill
+
+        decode_calls = []
+
+        async def _fake_super_completions(self, req, raw_info):
+            decode_calls.append(req)
+            return _aiter(["decode-chunk"])
+
+        request = CompletionRequest(model="test-model", prompt="hi")
+
+        with patch.object(LLMServer, "completions", _fake_super_completions):
+            chunks = [c async for c in server._pd_handle_request(request, None)]
+
+        # Standard (non choose_replica) path: .remote() was used.
+        sent_prefill = prefill.calls[0]["request"]
+        assert sent_prefill.max_tokens == 1
+        assert sent_prefill.kv_transfer_params["do_remote_decode"] is True
+        # Decode forwarded prefill's kv_transfer_params.
+        assert decode_calls[0].kv_transfer_params == {"remote_engine_id": "prefill-1"}
+        assert chunks == ["decode-chunk"]
+
+    @pytest.mark.asyncio
+    async def test_concurrent_handoff_cancels_prefill_on_decode_failure(self):
+        """In concurrent-handoff mode, if local decode raises, the background
+        prefill task must be cancelled (no leak)."""
+        from ray.llm._internal.serve.engines.vllm.kv_transfer.base import (
+            BaseConnectorBackend,
+        )
+
+        prefill_started = asyncio.Event()
+        prefill_cancelled = {"value": False}
+
+        class _SlowPrefillHandle:
+            def __init__(self):
+                self.calls = []
+
+            def options(self, **kwargs):
+                return self
+
+            def _method(self, name):
+                async def _gen():
+                    prefill_started.set()
+                    try:
+                        await asyncio.sleep(100)
+                        yield SimpleNamespace(kv_transfer_params={})
+                    except asyncio.CancelledError:
+                        prefill_cancelled["value"] = True
+                        raise
+
+                def remote(request, raw_request_info):
+                    self.calls.append({"method": name})
+                    return _gen()
+
+                return SimpleNamespace(remote=remote)
+
+            @property
+            def completions(self):
+                return self._method("completions")
+
+        class _ConcurrentBackend(BaseConnectorBackend):
+            requires_peer_binding = False
+            concurrent_handoff = True
+
+            def prepare_prefill_request(self, *, request, peer):
+                out = request.model_copy(deep=True)
+                out.kv_transfer_params = {"do_remote_decode": True}
+                return out
+
+            def prepare_decode_request(self, *, request, peer, prefill_response):
+                return request.model_copy(deep=True)
+
+        server = PDDecodeServer.__new__(PDDecodeServer)
+        server._llm_config = LLMConfig(
+            model_loading_config=ModelLoadingConfig(model_id="test-model")
+        )
+        server._llm_config._kv_connector_backend = _ConcurrentBackend(
+            server._llm_config
+        )
+        server._prefill_handle = _SlowPrefillHandle()
+
+        async def _failing_super_completions(self, req, raw_info):
+            await prefill_started.wait()
+
+            async def _gen():
+                raise RuntimeError("decode boom")
+                yield  # pragma: no cover
+
+            return _gen()
+
+        request = CompletionRequest(model="test-model", prompt="hi")
+
+        with patch.object(LLMServer, "completions", _failing_super_completions):
+            with pytest.raises(RuntimeError, match="decode boom"):
+                async for _ in server._pd_handle_request(request, None):
+                    pass
+
+        assert prefill_cancelled["value"] is True
+
+    @pytest.mark.asyncio
+    async def test_concurrent_handoff_surfaces_prefill_error(self):
+        """In concurrent-handoff mode, a prefill ErrorResponse must surface to
+        the client (and abort the hung local decode) instead of being only
+        logged — decode may be waiting on KV that will never arrive."""
+        from ray.llm._internal.serve.core.configs.openai_api_models import (
+            ErrorInfo,
+            ErrorResponse,
+        )
+        from ray.llm._internal.serve.engines.vllm.kv_transfer.base import (
+            BaseConnectorBackend,
+        )
+
+        prefill_error = ErrorResponse(
+            error=ErrorInfo(message="prefill boom", code=500, type="InternalError")
+        )
+        decode_aborted = {"value": False}
+
+        class _ErrorPrefillHandle:
+            def options(self, **kwargs):
+                return self
+
+            @property
+            def completions(self):
+                async def _gen():
+                    yield prefill_error
+
+                def remote(request, raw_request_info):
+                    return _gen()
+
+                return SimpleNamespace(remote=remote)
+
+        class _ConcurrentBackend(BaseConnectorBackend):
+            requires_peer_binding = False
+            concurrent_handoff = True
+
+            def prepare_prefill_request(self, *, request, peer):
+                return request.model_copy(deep=True)
+
+            def prepare_decode_request(self, *, request, peer, prefill_response):
+                return request.model_copy(deep=True)
+
+        server = PDDecodeServer.__new__(PDDecodeServer)
+        server._llm_config = LLMConfig(
+            model_loading_config=ModelLoadingConfig(model_id="test-model")
+        )
+        server._llm_config._kv_connector_backend = _ConcurrentBackend(
+            server._llm_config
+        )
+        server._prefill_handle = _ErrorPrefillHandle()
+
+        async def _hanging_super_completions(self, req, raw_info):
+            async def _gen():
+                try:
+                    # Decode never produces output (waiting on KV that the
+                    # failed prefill will never push).
+                    await asyncio.sleep(100)
+                    yield "never"
+                except (asyncio.CancelledError, GeneratorExit):
+                    decode_aborted["value"] = True
+                    raise
+
+            return _gen()
+
+        request = CompletionRequest(model="test-model", prompt="hi")
+
+        with patch.object(LLMServer, "completions", _hanging_super_completions):
+            chunks = [c async for c in server._pd_handle_request(request, None)]
+
+        assert chunks == [prefill_error]
+        assert decode_aborted["value"] is True
+
+    @pytest.mark.asyncio
+    async def test_prewarm_skipped_for_peer_binding_backend(self):
+        """Pre-warm broadcasts a peerless prefill request, which a peer-binding
+        connector (e.g. MoRIIO) cannot shape -- so it must be skipped rather
+        than crash decode-replica init."""
+        from ray.llm._internal.serve.engines.vllm.kv_transfer.base import (
+            BaseConnectorBackend,
+        )
+
+        class _PeerBindingBackend(BaseConnectorBackend):
+            requires_peer_binding = True
+
+            def prepare_prefill_request(self, *, request, peer):
+                raise AssertionError("prepare_prefill_request must not be called")
+
+            def prepare_decode_request(self, *, request, peer, prefill_response):
+                raise AssertionError("prepare_decode_request must not be called")
+
+        server = PDDecodeServer.__new__(PDDecodeServer)
+        server._llm_config = LLMConfig(
+            model_loading_config=ModelLoadingConfig(model_id="test-model"),
+            experimental_configs={"_prewarm_prefill_decode": True},
+        )
+        server._llm_config._kv_connector_backend = _PeerBindingBackend(
+            server._llm_config
+        )
+        # Must return without raising (and without touching prepare_*).
+        await server._maybe_prewarm()
+
+
+class TestBuildPDOpenaiApp:
+    """Test suite for build_pd_openai_app function."""
+
+    @pytest.fixture
+    def pd_configs(self):
+        """Prefill and decode configs with required kv_transfer_config."""
+        base_config = {
+            "model_loading_config": {
+                "model_id": "test-model",
+                "model_source": "test-source",
+            },
+            "engine_kwargs": {
+                "kv_transfer_config": {
+                    "kv_connector": "NixlConnector",
+                    "kv_role": "kv_both",
+                },
+            },
+        }
+        prefill = LLMConfig.model_validate(base_config)
+        decode = LLMConfig.model_validate(base_config)
+        return prefill, decode
+
+    def test_3_tier_graph_structure(self, pd_configs):
+        """Test that build_pd_openai_app creates a 3-tier graph:
+        ingress -> PDDecodeServer -> PDPrefillServer.
+        """
+        prefill, decode = pd_configs
+        app = build_pd_openai_app({"prefill_config": prefill, "decode_config": decode})
+
+        # The app should have an ingress deployment bound to the decode deployment
+        ingress_deployment = app._bound_deployment
+        llm_deployments = ingress_deployment.init_kwargs["llm_deployments"]
+        # Single model id -> single decode app (P/D shares the same model_id).
+        assert len(llm_deployments) == 1
+        decode_app = next(iter(llm_deployments.values()))
+        decode_deployment = decode_app._bound_deployment
+
+        assert decode_deployment.func_or_class is PDDecodeServer
+
+        # Decode should have a prefill_server in its bind kwargs
+        assert "prefill_server" in decode_deployment.init_kwargs
+
+        # The prefill_server should be a PDPrefillServer Application
+        prefill_app = decode_deployment.init_kwargs["prefill_server"]
+        prefill_deployment = prefill_app._bound_deployment
+        assert prefill_deployment.func_or_class is PDPrefillServer
+
+    def test_ingress_deployment_config(self, pd_configs):
+        """Test that ingress deployment configs are properly applied."""
+        prefill, decode = pd_configs
+        app = build_pd_openai_app(
+            {
+                "prefill_config": prefill,
+                "decode_config": decode,
+                "ingress_deployment_config": {
+                    "num_replicas": 5,
+                    "ray_actor_options": {
+                        "num_cpus": 8,
+                        "memory": 4096,
+                    },
+                    "max_ongoing_requests": 300,
+                },
+            }
+        )
+
+        ingress_deployment = app._bound_deployment
+        assert ingress_deployment._deployment_config.num_replicas == 5
+        assert ingress_deployment.ray_actor_options["num_cpus"] == 8
+        assert ingress_deployment.ray_actor_options["memory"] == 4096
+        assert ingress_deployment._deployment_config.max_ongoing_requests == 300
+
+    # TODO(Kourosh): Deprecated, remove in Ray 2.58.
+    def test_deprecated_proxy_config_ignored(self, pd_configs):
+        """Test that deprecated proxy configs are accepted but ignored."""
+        prefill, decode = pd_configs
+
+        with warnings.catch_warnings(record=True):
+            warnings.simplefilter("always")
+            app = build_pd_openai_app(
+                {
+                    "prefill_config": prefill,
+                    "decode_config": decode,
+                    "proxy_deployment_config": {
+                        "num_replicas": 99,
+                    },
+                }
+            )
+            # App should still be valid — proxy config is just ignored
+            assert app is not None
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main(["-v", __file__]))

python/ray/llm/tests/serve/cpu/deployments/prefill_decode_disagg/test_tokenize_once.py

@@ -0,0 +1,242 @@
+"""Unit tests for the P/D tokenize-once renderer wrap.
+
+vLLM is optional: the real-renderer and orchestrator checks skip when it (or the
+Ray serve import) is unavailable; the rest run with a stubbed renderer.
+"""
+import asyncio
+import contextvars
+import sys
+import types
+
+import pytest
+
+from ray.llm._internal.common.patches.vllm.tokenize_once import (
+    _reused_token_ids,
+    install,
+    reuse_prompt_token_ids,
+)
+
+
+def test_reuse_is_noop_on_falsy():
+    for falsy in (None, [], 0):
+        with reuse_prompt_token_ids(falsy):
+            assert _reused_token_ids.get() is None
+
+
+def test_reuse_sets_private_copy_then_clears():
+    src = [1, 2, 3]
+    with reuse_prompt_token_ids(src):
+        got = _reused_token_ids.get()
+        assert got == [1, 2, 3]
+        assert got is not src  # a private per-task copy, not the caller's list
+    assert _reused_token_ids.get() is None
+
+
+def test_reuse_teardown_is_cross_context_safe():
+    # Enter in one asyncio Context and exit in another, as when a peeked-then-
+    # abandoned decode generator is finalized off-task. reset(token) would raise
+    # "Token was created in a different Context". set(None) must not.
+    cm = reuse_prompt_token_ids([1, 2, 3])
+    contextvars.copy_context().run(cm.__enter__)
+    cm.__exit__(None, None, None)
+    assert _reused_token_ids.get() is None
+
+
+@pytest.fixture
+def fake_renderer(monkeypatch):
+    """Install the wrap onto a stub BaseRenderer whose orig echoes its args."""
+
+    class FakeBaseRenderer:
+        async def tokenize_prompts_async(self, prompts, params):
+            return ("orig", prompts, params)
+
+    base = types.ModuleType("vllm.renderers.base")
+    base.BaseRenderer = FakeBaseRenderer
+    monkeypatch.setitem(sys.modules, "vllm", types.ModuleType("vllm"))
+    monkeypatch.setitem(
+        sys.modules, "vllm.renderers", types.ModuleType("vllm.renderers")
+    )
+    monkeypatch.setitem(sys.modules, "vllm.renderers.base", base)
+    return FakeBaseRenderer
+
+
+def test_install_is_idempotent(fake_renderer):
+    assert install() is True
+    assert install() is True
+    assert getattr(fake_renderer.tokenize_prompts_async, "_pd_tokonce_wrapped", False)
+
+
+def test_install_returns_false_without_vllm(monkeypatch):
+    monkeypatch.setitem(sys.modules, "vllm.renderers.base", None)
+    assert install() is False
+
+
+def test_install_returns_false_when_method_missing(monkeypatch):
+    # A vLLM whose BaseRenderer lacks tokenize_prompts_async must fail safe
+    # (return False), not raise AttributeError and crash replica startup.
+    class RendererWithoutMethod:
+        pass
+
+    base = types.ModuleType("vllm.renderers.base")
+    base.BaseRenderer = RendererWithoutMethod
+    monkeypatch.setitem(sys.modules, "vllm", types.ModuleType("vllm"))
+    monkeypatch.setitem(
+        sys.modules, "vllm.renderers", types.ModuleType("vllm.renderers")
+    )
+    monkeypatch.setitem(sys.modules, "vllm.renderers.base", base)
+    assert install() is False
+
+
+def _tokenize(renderer_cls, prompts, ids):
+    async def run():
+        with reuse_prompt_token_ids(ids):
+            return await renderer_cls.tokenize_prompts_async(
+                renderer_cls(), prompts, "PARAMS"
+            )
+
+    return asyncio.run(run())
+
+
+def test_injects_ids_and_preserves_other_fields(fake_renderer):
+    install()
+    prompt = {"prompt": "hi", "multi_modal_data": {"image": object()}}
+    tag, seen, params = _tokenize(fake_renderer, [dict(prompt)], [7, 8, 9])
+    assert tag == "orig" and params == "PARAMS"
+    # ids injected so vLLM skips the encode; multi_modal_data preserved.
+    assert seen == [{**prompt, "prompt_token_ids": [7, 8, 9]}]
+
+
+@pytest.mark.parametrize(
+    "prompts, ids",
+    [
+        ([{"prompt": "hi"}], None),  # reuse not set
+        ([{"prompt": "a"}, {"prompt": "b"}], [1, 2]),  # batched
+        ([{"prompt_token_ids": [1]}], [9]),  # already tokenized
+        ([{"prompt_embeds": object()}], [9]),  # embeds
+        ([{"encoder_prompt": {}}], [9]),  # encoder-decoder
+    ],
+)
+def test_falls_through_untouched(fake_renderer, prompts, ids):
+    install()
+    expected = [dict(p) for p in prompts]
+    _tag, seen, _params = _tokenize(fake_renderer, prompts, ids)
+    assert seen == expected
+
+
+def test_concurrent_tasks_do_not_cross_contaminate(fake_renderer):
+    # Two requests run on separate asyncio tasks with different reuse ids and
+    # interleave at the await. contextvars are task-local, so each injection must
+    # see only its own ids (this is the property the design relies on for safety).
+    install()
+    Renderer = fake_renderer
+
+    async def one(ids):
+        with reuse_prompt_token_ids(ids):
+            await asyncio.sleep(0)  # yield so the two tasks interleave
+            _tag, seen, _params = await Renderer.tokenize_prompts_async(
+                Renderer(), [{"prompt": "hi"}], "PARAMS"
+            )
+            return seen
+
+    async def run():
+        return await asyncio.gather(one([1, 1, 1]), one([2, 2, 2]))
+
+    a, b = asyncio.run(run())
+    assert a == [{"prompt": "hi", "prompt_token_ids": [1, 1, 1]}]
+    assert b == [{"prompt": "hi", "prompt_token_ids": [2, 2, 2]}]
+
+
+def test_patch_applies_to_real_vllm_renderer():
+    # Confirms install() wraps the *real* vLLM BaseRenderer (correct import path
+    # and method name for this vLLM version), and that with reuse ids set the
+    # rendered text prompt reaches per-prompt tokenization carrying
+    # ``prompt_token_ids`` so vLLM's skip-check avoids the encode. This is the
+    # guard against a vLLM bump silently moving the seam the wrap depends on.
+    vbase = pytest.importorskip("vllm.renderers.base")
+    original = vbase.BaseRenderer.tokenize_prompts_async
+    try:
+        assert install() is True
+        assert getattr(
+            vbase.BaseRenderer.tokenize_prompts_async, "_pd_tokonce_wrapped", False
+        )
+        assert install() is True  # idempotent against the real class
+
+        seen = []
+
+        class StubRenderer:
+            # The real tokenize_prompts_async fans out to tokenize_prompt_async
+            # per prompt; that is the only hook we need to observe injection.
+            async def tokenize_prompt_async(self, prompt, params):
+                seen.append(prompt)
+                return prompt
+
+        async def drive(ids):
+            seen.clear()
+            with reuse_prompt_token_ids(ids):
+                await vbase.BaseRenderer.tokenize_prompts_async(
+                    StubRenderer(), [{"prompt": "hi"}], "PARAMS"
+                )
+
+        asyncio.run(drive([11, 22, 33]))
+        assert seen == [{"prompt": "hi", "prompt_token_ids": [11, 22, 33]}]
+
+        asyncio.run(drive(None))  # no reuse -> untouched, real encode would run
+        assert seen == [{"prompt": "hi"}]
+    finally:
+        vbase.BaseRenderer.tokenize_prompts_async = original
+
+
+def _orchestrator(pd_tokenize_once):
+    """A PDOrchestratorMixin instance with only the reuse flag set.
+
+    Skips when the Ray serve LLM stack (or vLLM) is unavailable.
+    """
+    pd_server = pytest.importorskip(
+        "ray.llm._internal.serve.serving_patterns.prefill_decode.pd_server"
+    )
+    obj = pd_server.PDOrchestratorMixin.__new__(pd_server.PDOrchestratorMixin)
+    obj._pd_tokenize_once = pd_tokenize_once
+    return obj
+
+
+@pytest.mark.parametrize(
+    "enabled, chunk, expected",
+    [
+        # chat: ids echoed top-level on the response
+        (True, types.SimpleNamespace(prompt_token_ids=[1, 2, 3]), [1, 2, 3]),
+        # completions: ids echoed on the first choice
+        (
+            True,
+            types.SimpleNamespace(
+                prompt_token_ids=None,
+                choices=[types.SimpleNamespace(prompt_token_ids=[4, 5])],
+            ),
+            [4, 5],
+        ),
+        # disabled: never reuse, even when prefill echoed ids
+        (False, types.SimpleNamespace(prompt_token_ids=[1, 2, 3]), None),
+    ],
+)
+def test_decode_reuse_ids(enabled, chunk, expected):
+    assert _orchestrator(enabled)._decode_reuse_ids(chunk) == expected
+
+
+@pytest.mark.parametrize(
+    "enabled, expected",
+    [(True, True), (False, False)],  # only request the echo when enabled
+)
+def test_request_prefill_token_ids_gating(enabled, expected):
+    req = types.SimpleNamespace(return_token_ids=False)
+    _orchestrator(enabled)._request_prefill_token_ids(req)
+    assert req.return_token_ids is expected
+
+
+def test_request_prefill_token_ids_noop_when_field_absent():
+    # An older vLLM request without return_token_ids must not crash.
+    req = types.SimpleNamespace()
+    _orchestrator(True)._request_prefill_token_ids(req)
+    assert not hasattr(req, "return_token_ids")
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main(["-v", __file__]))

python/ray/llm/tests/serve/cpu/deployments/routers/kv/test_config_files/llm_kv_aware_deployment.yaml

@@ -0,0 +1,21 @@
+applications:
+  - name: kv-llm
+    route_prefix: /
+    import_path: ray.serve.llm:build_openai_app
+    runtime_env:
+      env_vars:
+        RAY_SERVE_LLM_ENABLE_DIRECT_STREAMING: "1"
+        RAY_SERVE_ENABLE_HA_PROXY: "1"
+    args:
+      llm_configs:
+        - model_loading_config:
+            model_id: qwen3-0.6b
+            model_source: Qwen/Qwen3-0.6B
+          accelerator_type: null
+          deployment_config:
+            autoscaling_config:
+              min_replicas: 0
+              initial_replicas: 0
+              max_replicas: 1
+            request_router_config:
+              request_router_class: ray.serve.llm.request_router.KVAwareRouter

python/ray/llm/tests/serve/cpu/deployments/routers/kv/test_kv_aware_router.py

@@ -0,0 +1,529 @@
+"""KVRouterActor attachment and live replica-membership tracking.
+
+Attachment is covered two ways: ``build_openai_app`` with a Python ``LLMConfig``,
+and a declarative YAML config deployed via ``serve deploy`` (the dotted-string
+router class only YAML can express). Membership tracking is covered by deploying
+a dummy multi-replica deployment and asserting the actor's LongPoll listener
+stays in sync with the live replicas across scale up/down.
+"""
+
+import os
+import subprocess
+import sys
+from typing import List
+from unittest import mock
+
+import pytest
+
+import ray
+from ray import serve
+from ray._common.test_utils import wait_for_condition
+from ray.llm._internal.serve.core.configs.llm_config import LLMConfig
+from ray.llm._internal.serve.core.ingress.builder import (
+    LLMServingArgs,
+    build_openai_app,
+)
+from ray.llm._internal.serve.core.ingress.tokenizer import REQUEST_TOKEN_IDS_KWARG
+from ray.llm._internal.serve.routing_policies.kv_aware.kv_aware_actor import (
+    KV_ROUTER_ACTOR_NAME,
+    KVRouterActor,
+    get_worker_id,
+)
+from ray.serve._private.common import (
+    REPLICA_ID_FULL_ID_STR_PREFIX,
+    DeploymentID,
+    DeploymentTargetInfo,
+    ReplicaID,
+    RequestMetadata,
+    RunningReplicaInfo,
+)
+from ray.serve._private.constants import SERVE_DEPLOYMENT_ACTOR_PREFIX, SERVE_NAMESPACE
+from ray.serve._private.request_router import PendingRequest
+from ray.serve.config import DeploymentActorConfig
+from ray.serve.llm.request_router import KVAwareRouter
+from ray.util.state import list_actors
+
+
+def get_kv_actor_configs(deployment):
+    return [
+        cfg
+        for cfg in (deployment._deployment_config.deployment_actors or [])
+        if (cfg["name"] if isinstance(cfg, dict) else cfg.name) == KV_ROUTER_ACTOR_NAME
+    ]
+
+
+def build_test_llm_config(experimental_configs=None) -> LLMConfig:
+    return LLMConfig(
+        model_loading_config={
+            "model_id": "qwen3-0.6b",
+            "model_source": "Qwen/Qwen3-0.6B",
+        },
+        accelerator_type=None,
+        deployment_config={
+            "autoscaling_config": {"min_replicas": 1, "max_replicas": 1},
+            "request_router_config": {"request_router_class": KVAwareRouter},
+        },
+        experimental_configs=experimental_configs or {},
+    )
+
+
+def build_non_kv_llm_config(**engine_kwargs) -> LLMConfig:
+    """An LLMConfig whose request router is the default (not a KVAwareRouter)."""
+    return LLMConfig(
+        model_loading_config={
+            "model_id": "qwen3-0.6b",
+            "model_source": "Qwen/Qwen3-0.6B",
+        },
+        accelerator_type=None,
+        deployment_config={
+            "autoscaling_config": {"min_replicas": 1, "max_replicas": 1}
+        },
+        engine_kwargs=engine_kwargs,
+    )
+
+
+def get_kv_actor_names(app_name: str) -> list:
+    prefix = f"{SERVE_DEPLOYMENT_ACTOR_PREFIX}{app_name}::"
+    suffix = f"::{KV_ROUTER_ACTOR_NAME}"
+    return [
+        a["name"]
+        for a in list_actors(filters=[("state", "=", "ALIVE")])
+        if a["name"] and a["name"].startswith(prefix) and a["name"].endswith(suffix)
+    ]
+
+
+def discover_deployment_actor(app_name, deployment_name, actor_name):
+    """Handle to a deployment-scoped actor by app/deployment/logical name."""
+    prefix = f"{SERVE_DEPLOYMENT_ACTOR_PREFIX}{app_name}::{deployment_name}::"
+    suffix = f"::{actor_name}"
+    for entry in ray.util.list_named_actors(all_namespaces=True):
+        name = entry.get("name") or ""
+        if (
+            entry.get("namespace") == SERVE_NAMESPACE
+            and name.startswith(prefix)
+            and (name.endswith(suffix))
+        ):
+            return ray.get_actor(name, namespace=SERVE_NAMESPACE)
+    return None
+
+
+def get_candidate_ids(app_name):
+    handle = discover_deployment_actor(
+        app_name, "ReplicaTrackingDeployment", KV_ROUTER_ACTOR_NAME
+    )
+    assert handle is not None
+    return ray.get(handle.get_candidate_worker_ids.remote())
+
+
+def get_live_replica_worker_ids(app_name, deployment_name="ReplicaTrackingDeployment"):
+    """Worker ids derived directly from the deployment's alive replica actors."""
+    prefix = f"{REPLICA_ID_FULL_ID_STR_PREFIX}{app_name}#{deployment_name}#"
+    return {
+        get_worker_id(a["name"][len(prefix) :])
+        for a in list_actors(filters=[("state", "=", "ALIVE")])
+        if a["name"] and a["name"].startswith(prefix)
+    }
+
+
+@pytest.fixture(autouse=True)
+def enable_direct_streaming(monkeypatch):
+    monkeypatch.setattr(
+        "ray.llm._internal.serve.core.ingress.builder."
+        "RAY_SERVE_LLM_ENABLE_DIRECT_STREAMING",
+        True,
+    )
+
+
+@pytest.fixture(scope="module")
+def serve_instance():
+    if not ray.is_initialized():
+        ray.init(address="auto")
+    yield
+    serve.shutdown()
+
+
+def test_build_openai_app_attaches_kv_actor():
+    """A KVAwareRouter on the LLMConfig attaches the KVRouterActor."""
+    app = build_openai_app(LLMServingArgs(llm_configs=[build_test_llm_config()]))
+
+    configs = get_kv_actor_configs(app._bound_deployment)
+    assert len(configs) == 1
+    actor_cfg = configs[0]
+    assert actor_cfg.get_actor_class().__ray_actor_class__ is KVRouterActor
+    assert actor_cfg.actor_options["num_cpus"] == 0
+    assert actor_cfg.init_kwargs == {"indexer_threads": 4}
+
+
+def test_configurable_indexer_threads():
+    llm_config = build_test_llm_config(experimental_configs={"KV_INDEXER_THREADS": 8})
+    app = build_openai_app(LLMServingArgs(llm_configs=[llm_config]))
+
+    actor_cfg = get_kv_actor_configs(app._bound_deployment)[0]
+    assert actor_cfg.init_kwargs["indexer_threads"] == 8
+
+
+def test_non_kv_router_warns_kv_events_config():
+    """Without a KVAwareRouter no KVRouterActor is attached and a user-provided
+    kv_events_config is left untouched (just unused), with a warning pointing at
+    how to consume the engine's KV events."""
+    kv_events_config = {
+        "enable_kv_cache_events": True,
+        "publisher": "zmq",
+        "endpoint": "tcp://*:5557",
+    }
+    llm_config = build_non_kv_llm_config(kv_events_config=kv_events_config)
+
+    with mock.patch(
+        "ray.llm._internal.serve.routing_policies.kv_aware.utils.logger"
+    ) as logger:
+        app = build_openai_app(LLMServingArgs(llm_configs=[llm_config]))
+
+    assert get_kv_actor_configs(app._bound_deployment) == []
+    assert llm_config.engine_kwargs["kv_events_config"] == kv_events_config
+    logger.warning.assert_called_once()
+    assert "KVAwareRouter" in logger.warning.call_args.args[0]
+
+
+def test_yaml_config_attaches_kv_actor(serve_instance):
+    """Deploying a YAML config that selects KVAwareRouter creates the KVRouterActor."""
+    config_file = os.path.join(
+        os.path.dirname(__file__), "test_config_files", "llm_kv_aware_deployment.yaml"
+    )
+    app_name = "kv-llm"
+
+    subprocess.check_output(["serve", "deploy", config_file], stderr=subprocess.STDOUT)
+    try:
+        wait_for_condition(lambda: len(get_kv_actor_names(app_name)) == 1, timeout=60)
+    finally:
+        serve.delete(app_name, _blocking=True)
+
+
+class _TestKVRouterActor(KVRouterActor):
+    """KVRouterActor augmented with test-only introspection."""
+
+    async def get_candidate_worker_ids(self) -> List[int]:
+        """The workers currently tracked from running replicas.
+
+        Async so it runs on the actor's event loop, serialized with
+        ``_on_deployment_targets`` which mutates the same map on that loop.
+        """
+        return sorted(self._replica_id_by_worker)
+
+
+@serve.deployment(
+    num_replicas=4,
+    deployment_actors=[
+        DeploymentActorConfig(
+            name=KV_ROUTER_ACTOR_NAME,
+            actor_class=ray.remote(_TestKVRouterActor),
+            actor_options={"num_cpus": 0},
+            init_kwargs={},
+        ),
+    ],
+)
+class ReplicaTrackingDeployment:
+    """Dummy deployment with a KVRouterActor deployment actor.
+
+    Advertises a per-replica KV-events endpoint via ``record_routing_stats`` as a
+    real engine would, so the selection service tracks each replica as a worker.
+    """
+
+    async def __call__(self) -> str:
+        return "ok"
+
+    async def record_routing_stats(self) -> dict:
+        rank = serve.get_replica_context().rank.local_rank
+        return {
+            "kv_event_metadata": {
+                "endpoint": f"tcp://{ray.util.get_node_ip_address()}:{25000 + rank}",
+                "block_size": 16,
+                "max_num_batched_tokens": 8192,
+                "dp_rank": 0,
+            }
+        }
+
+
+class TestReplicaTrackingIntegration:
+    def test_tracks_running_replicas(self, serve_instance):
+        """KVRouterActor's LongPollClient receives the running replicas."""
+        app_name = "kv-replica-tracking"
+        serve.run(
+            ReplicaTrackingDeployment.bind(), name=app_name, route_prefix="/kv_track"
+        )
+        try:
+            wait_for_condition(
+                lambda: len(get_candidate_ids(app_name)) == 4, timeout=30
+            )
+            # The tracked workers are exactly those of the live replica actors.
+            assert set(get_candidate_ids(app_name)) == get_live_replica_worker_ids(
+                app_name
+            )
+        finally:
+            serve.delete(app_name, _blocking=True)
+
+    def test_membership_broadcast_on_scale(self, serve_instance):
+        """A scale up then down is broadcast over LongPoll; the actor re-syncs to
+        exactly the live replica set each time.
+        """
+        app_name = "kv-replica-scale"
+
+        def tracks_live_replicas(expected):
+            # The tracked workers match the live replica actors by their actual
+            # ids (a stale handle is possible while the deployment is updated).
+            try:
+                tracked = set(get_candidate_ids(app_name))
+            except ray.exceptions.RayActorError:
+                return False
+            return len(tracked) == expected and tracked == get_live_replica_worker_ids(
+                app_name
+            )
+
+        def scale(num_replicas):
+            serve.run(
+                ReplicaTrackingDeployment.options(num_replicas=num_replicas).bind(),
+                name=app_name,
+                route_prefix="/kv_scale",
+            )
+
+        scale(2)
+        try:
+            wait_for_condition(lambda: tracks_live_replicas(2), timeout=30)
+            scale(4)  # upscale: the new replicas are picked up over LongPoll.
+            wait_for_condition(lambda: tracks_live_replicas(4), timeout=30)
+            scale(2)  # downscale: the departed replicas are dropped.
+            wait_for_condition(lambda: tracks_live_replicas(2), timeout=30)
+        finally:
+            serve.delete(app_name, _blocking=True)
+
+
+class _LocalKVRouterActor(_TestKVRouterActor):
+    """In-process KVRouterActor with the selection service and LongPoll disabled,
+    to drive ``_on_deployment_targets`` directly with synthetic snapshots.
+    """
+
+    def _create_selection_service(self) -> None:
+        self._svc = None  # reconcile membership without dynamo
+
+    def _start_replica_tracking(self) -> None:
+        pass
+
+    def _schedule(self, coro) -> None:
+        coro.close()  # _svc is None, so the scheduled upsert is a no-op
+
+
+def make_target_info(unique_ids):
+    """A DeploymentTargetInfo whose replicas advertise a KV-events endpoint via
+    routing_stats, exactly as the controller broadcasts it over LongPoll."""
+    deployment_id = DeploymentID(name="d", app_name="app")
+    running_replicas = [
+        RunningReplicaInfo(
+            replica_id=ReplicaID(unique_id=uid, deployment_id=deployment_id),
+            node_id="node",
+            node_ip="10.0.0.1",
+            availability_zone="az",
+            actor_name=f"actor-{uid}",
+            max_ongoing_requests=1,
+            routing_stats={
+                "kv_event_metadata": {
+                    "endpoint": "tcp://10.0.0.1:25000",
+                    "block_size": 16,
+                    "max_num_batched_tokens": 8192,
+                    "dp_rank": 0,
+                }
+            },
+        )
+        for uid in unique_ids
+    ]
+    return DeploymentTargetInfo(is_available=True, running_replicas=running_replicas)
+
+
+class TestOnDeploymentTargets:
+    async def test_reconciles_added_and_removed_workers(self):
+        actor = _LocalKVRouterActor()
+        actor._on_deployment_targets(make_target_info(["a", "b"]))
+        assert set(await actor.get_candidate_worker_ids()) == {
+            get_worker_id("a"),
+            get_worker_id("b"),
+        }
+        # "a" departs and "c" joins: the tracked set follows the new snapshot.
+        actor._on_deployment_targets(make_target_info(["b", "c"]))
+        assert set(await actor.get_candidate_worker_ids()) == {
+            get_worker_id("b"),
+            get_worker_id("c"),
+        }
+
+
+class _StubReplica:
+    """RunningReplica stand-in exposing only replica_id.unique_id."""
+
+    def __init__(self, unique_id: str):
+        self.replica_id = ReplicaID(
+            unique_id=unique_id, deployment_id=DeploymentID(name="d", app_name="app")
+        )
+
+
+class _SelectWorkerStub:
+    def __init__(self, worker_id: int):
+        self._worker_id = worker_id
+        self.token_ids = None
+        self.allowed = None
+
+    async def remote(self, request_id, token_ids, allowed_worker_ids):
+        self.token_ids = token_ids
+        self.allowed = allowed_worker_ids
+        return {
+            "worker_id": self._worker_id,
+            "dp_rank": 0,
+            "overlap_tokens": 1,
+            "effective_prefill_tokens": len(token_ids),
+        }
+
+
+class _KVRouterActorStub:
+    def __init__(self, worker_id: int):
+        self.select_worker = _SelectWorkerStub(worker_id)
+
+
+class _StubKVAwareRouter(KVAwareRouter):
+    """KVAwareRouter with the scorer actor injected, bypassing actor discovery."""
+
+    def __init__(self, kv_router_actor):
+        self._kv_router_actor = kv_router_actor
+
+
+def _build_kv_aware_router(worker_id: int) -> KVAwareRouter:
+    return _StubKVAwareRouter(_KVRouterActorStub(worker_id))
+
+
+@pytest.mark.asyncio
+async def test_select_worker_requires_tokens():
+    actor = KVRouterActor.__new__(KVRouterActor)
+    actor._svc = object()
+
+    with pytest.raises(ValueError, match="non-empty token_ids"):
+        await actor.select_worker("req-empty", [], [get_worker_id("r1")])
+
+
+@pytest.mark.asyncio
+async def test_select_worker_without_dynamo_raises():
+    """Without ai-dynamo the actor cannot score, so it raises a clear error
+    instead of silently degrading to a non-KV-aware pick."""
+    actor = KVRouterActor.__new__(KVRouterActor)
+    actor._svc = None
+
+    with pytest.raises(RuntimeError, match="ai-dynamo is not installed"):
+        await actor.select_worker("req", [1, 2, 3], [get_worker_id("r1")])
+
+
+@pytest.mark.asyncio
+async def test_choose_replicas_routes_to_selected_worker():
+    """choose_replicas maps candidates to worker ids, asks the actor to select,
+    and returns the chosen worker's replica."""
+    replicas = [_StubReplica("r1"), _StubReplica("r2")]
+    worker_ids = [get_worker_id("r1"), get_worker_id("r2")]
+
+    router = _build_kv_aware_router(worker_ids[1])
+    pending = PendingRequest(
+        args=[],
+        kwargs={REQUEST_TOKEN_IDS_KWARG: [10, 11, 12]},
+        metadata=RequestMetadata(request_id="req-1", internal_request_id="int-1"),
+    )
+
+    groups = await router.choose_replicas(replicas, pending)
+
+    # The actor selected r2's worker, so r2 is returned.
+    assert groups == [[replicas[1]]]
+    # choose_replicas forwarded the prompt token ids and the full candidate set.
+    select = router._kv_router_actor.select_worker
+    assert select.token_ids == [10, 11, 12]
+    assert sorted(select.allowed) == sorted(worker_ids)
+
+
+@pytest.mark.asyncio
+async def test_missing_token_ids_picks_random_replica():
+    """Token-less requests (batch prompts, truncated bodies) route to a single
+    random replica so they spread."""
+    replicas = [_StubReplica("r1"), _StubReplica("r2")]
+    router = _build_kv_aware_router(get_worker_id("r1"))
+
+    picked = set()
+    for _ in range(50):
+        pending = PendingRequest(
+            args=[],
+            kwargs={},
+            metadata=RequestMetadata(request_id="req", internal_request_id="int"),
+        )
+        groups = await router.choose_replicas(replicas, pending)
+        assert len(groups) == 1 and len(groups[0]) == 1
+        assert groups[0][0] in replicas
+        picked.add(groups[0][0].replica_id.unique_id)
+
+    # The picked replica varies across calls, so load spreads (not stuck on one).
+    assert picked == {"r1", "r2"}
+    assert router._kv_router_actor.select_worker.token_ids is None
+
+
+@pytest.mark.asyncio
+async def test_tokenize_call_picks_random_replica():
+    """The pre-routing /tokenize RPC is routed through choose_replicas before any
+    token ids exist; it must resolve so KV routing can bootstrap, and picks a random
+    replica without scoring."""
+    replicas = [_StubReplica("r1"), _StubReplica("r2")]
+
+    router = _build_kv_aware_router(get_worker_id("r2"))
+    pending = PendingRequest(
+        args=[],
+        kwargs={},
+        metadata=RequestMetadata(
+            request_id="req-tokenize",
+            internal_request_id="int-tokenize",
+            call_method="tokenize",
+        ),
+    )
+
+    groups = await router.choose_replicas(replicas, pending)
+
+    assert len(groups) == 1 and len(groups[0]) == 1
+    assert groups[0][0] in replicas
+    assert router._kv_router_actor.select_worker.token_ids is None
+
+
+@pytest.mark.asyncio
+async def test_empty_token_ids_picks_random_replica():
+    """Empty token ids carry no KV signal, so pick a random replica instead of
+    handing an empty prompt to the Dynamo selection service (which rejects it)."""
+    replicas = [_StubReplica("r1"), _StubReplica("r2")]
+
+    router = _build_kv_aware_router(get_worker_id("r2"))
+    pending = PendingRequest(
+        args=[],
+        kwargs={REQUEST_TOKEN_IDS_KWARG: []},
+        metadata=RequestMetadata(
+            request_id="req-empty", internal_request_id="int-empty"
+        ),
+    )
+
+    groups = await router.choose_replicas(replicas, pending)
+
+    assert len(groups) == 1 and len(groups[0]) == 1
+    assert groups[0][0] in replicas
+    assert router._kv_router_actor.select_worker.token_ids is None
+
+
+@pytest.mark.asyncio
+async def test_no_pending_request_picks_random_replica():
+    """Serve may ask again after route metadata has been consumed; pick a random
+    replica (nothing to score on)."""
+    replicas = [_StubReplica("r1"), _StubReplica("r2")]
+
+    router = _build_kv_aware_router(get_worker_id("r1"))
+
+    groups = await router.choose_replicas(replicas, pending_request=None)
+
+    assert len(groups) == 1 and len(groups[0]) == 1
+    assert groups[0][0] in replicas
+    assert router._kv_router_actor.select_worker.token_ids is None
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main(["-v", __file__]))

python/ray/llm/tests/serve/cpu/deployments/routers/kv/test_kv_events.py

@@ -0,0 +1,122 @@
+import sys
+from types import SimpleNamespace
+from unittest import mock
+
+import pytest
+
+import ray
+from ray.llm._internal.serve.core.configs.llm_config import LLMConfig
+from ray.llm._internal.serve.routing_policies.kv_aware.vllm.kv_events import (
+    assign_replica_kv_events_endpoint,
+    configure_kv_events_for_kv_routing,
+    get_kv_event_routing_stats,
+    resolve_kv_event_source_endpoint,
+)
+from ray.serve.llm.request_router import KVAwareRouter
+
+
+def make_kv_aware_llm_config(**kwargs) -> LLMConfig:
+    return LLMConfig(
+        model_loading_config={
+            "model_id": "qwen3-0.6b",
+            "model_source": "Qwen/Qwen3-0.6B",
+        },
+        accelerator_type=None,
+        deployment_config={
+            "autoscaling_config": {"min_replicas": 1, "max_replicas": 1},
+            "request_router_config": {"request_router_class": KVAwareRouter},
+        },
+        **kwargs,
+    )
+
+
+@pytest.fixture(scope="module")
+def ray_instance():
+    started = not ray.is_initialized()
+    if started:
+        ray.init()
+    yield
+    if started:
+        ray.shutdown()
+
+
+class TestConfigureKvEvents:
+    def test_configure_enables_events_and_pins_seed(self):
+        """KV-aware config turns on engine ZMQ KV events and pins the hash seed."""
+        llm_config = make_kv_aware_llm_config()
+        configure_kv_events_for_kv_routing(llm_config)
+
+        assert llm_config.engine_kwargs["kv_events_config"] == {
+            "enable_kv_cache_events": True,
+            "publisher": "zmq",
+            "endpoint": "tcp://*:5557",
+            "replay_endpoint": "tcp://*:6557",
+        }
+        assert llm_config.runtime_env["env_vars"]["PYTHONHASHSEED"] == "0"
+
+    @pytest.mark.parametrize(
+        "engine_kwargs, local_rank, expected_port, expected_replay_port",
+        [
+            # Non-DP: offset the base port by the replica's node-local rank so
+            # colocated replicas don't bind the same ZMQ PUB port.
+            ({}, 2, 5559, 6559),
+            # DP: data_parallel_rank set -> offset 0 (the engine offsets the
+            # bound port by dp_rank itself), so local_rank must be ignored.
+            ({"data_parallel_rank": 2}, 2, 5557, 6557),
+        ],
+    )
+    def test_assign_replica_endpoint_offsets_port(
+        self, engine_kwargs, local_rank, expected_port, expected_replay_port
+    ):
+        """Per-replica endpoint offset: by node-local rank without DP, 0 with DP."""
+        llm_config = make_kv_aware_llm_config(engine_kwargs=dict(engine_kwargs))
+        configure_kv_events_for_kv_routing(llm_config)  # base ports 5557 / 6557
+        replica_context = SimpleNamespace(rank=SimpleNamespace(local_rank=local_rank))
+        with mock.patch("ray.serve.get_replica_context", return_value=replica_context):
+            assign_replica_kv_events_endpoint(llm_config)
+        kv_events_config = llm_config.engine_kwargs["kv_events_config"]
+        assert kv_events_config["endpoint"] == f"tcp://*:{expected_port}"
+        assert kv_events_config["replay_endpoint"] == f"tcp://*:{expected_replay_port}"
+
+    def test_resolve_endpoint_is_node_routable(self, ray_instance):
+        """The advertised endpoint is the replica's node IP."""
+        llm_config = make_kv_aware_llm_config()
+        configure_kv_events_for_kv_routing(llm_config)
+
+        endpoint = resolve_kv_event_source_endpoint(llm_config)
+        node_ip = ray.util.get_node_ip_address()
+        assert endpoint == f"tcp://{node_ip}:5557"
+
+    def test_routing_stats_advertise_endpoint(self, ray_instance):
+        """The replica advertises its node-routable endpoint plus the engine
+        facts the selection service needs to schedule it via record_routing_stats."""
+        llm_config = make_kv_aware_llm_config()
+        configure_kv_events_for_kv_routing(llm_config)
+
+        stats = get_kv_event_routing_stats(
+            llm_config, block_size=16, max_num_batched_tokens=4096
+        )
+        node_ip = ray.util.get_node_ip_address()
+        assert stats == {
+            "kv_event_metadata": {
+                "endpoint": f"tcp://{node_ip}:5557",
+                "block_size": 16,
+                "max_num_batched_tokens": 4096,
+                "dp_rank": 0,
+                "replay_endpoint": f"tcp://{node_ip}:6557",
+            }
+        }
+
+    def test_routing_stats_empty_without_kv_events(self):
+        """Nothing to advertise when KV-cache events are not enabled."""
+        llm_config = make_kv_aware_llm_config()
+        assert (
+            get_kv_event_routing_stats(
+                llm_config, block_size=16, max_num_batched_tokens=4096
+            )
+            == {}
+        )
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main(["-v", __file__]))

python/ray/llm/tests/serve/cpu/deployments/routers/kv/test_request_lifecycle_tracking.py

@@ -0,0 +1,807 @@
+import asyncio
+import random
+import sys
+from collections import OrderedDict
+from dataclasses import asdict
+from types import SimpleNamespace
+
+import pytest
+from vllm.outputs import CompletionOutput, RequestOutput
+from vllm.sampling_params import RequestOutputKind, SamplingParams
+
+import ray
+import ray.cloudpickle
+from ray import serve
+from ray.llm._internal.serve.core.configs.llm_config import LLMConfig
+from ray.llm._internal.serve.routing_policies.kv_aware.constants import (
+    REQUEST_TRACKING_TTL_S,
+)
+from ray.llm._internal.serve.routing_policies.kv_aware.kv_aware_actor import (
+    KV_ROUTER_ACTOR_NAME,
+    KVRouterActor,
+    get_worker_id,
+)
+from ray.llm._internal.serve.routing_policies.kv_aware.kv_aware_router import (
+    is_kv_aware,
+)
+from ray.llm._internal.serve.routing_policies.kv_aware.vllm.token_tracking import (
+    enable_token_tracking,
+)
+from ray.llm.tests.serve.mocks.mock_vllm_engine import MockAsyncLLM
+from ray.serve.llm.request_router import KVAwareRouter
+
+# The @ray.remote actors below are pickled by reference, so a worker must import
+# this module -- but pytest imports it under a bare name the worker cannot import
+# (ModuleNotFoundError). Pickling by value ships the class bodies instead.
+ray.cloudpickle.register_pickle_by_value(sys.modules[__name__])
+
+REPLICA_UNIQUE_ID = "test-replica-uid"
+WORKER_ID = get_worker_id(REPLICA_UNIQUE_ID)
+# A pre-tokenized prompt, as vLLM's serving layer always passes to generate.
+PROMPT = {"prompt_token_ids": [1, 2, 3]}
+# SamplingParams.max_tokens the engine reports as the request's expected output.
+MAX_TOKENS = 20
+
+
+@pytest.fixture(scope="module", autouse=True)
+def ray_cluster():
+    if not ray.is_initialized():
+        ray.init()
+
+
+def request_output(token_counts, prompt_len=5, finished=False):
+    """A real vLLM RequestOutput: one CompletionOutput per entry in token_counts."""
+    return RequestOutput(
+        request_id="r",
+        prompt=None,
+        prompt_token_ids=list(range(prompt_len)),
+        prompt_logprobs=None,
+        outputs=[
+            CompletionOutput(
+                index=i,
+                text="",
+                token_ids=list(range(n)),
+                cumulative_logprob=None,
+                logprobs=None,
+            )
+            for i, n in enumerate(token_counts)
+        ],
+        finished=finished,
+    )
+
+
+def delta_steps(num_tokens, prompt_len=5):
+    """A DELTA-kind stream: one new token per step, last step finished."""
+    return [
+        request_output([1], prompt_len=prompt_len, finished=i == num_tokens - 1)
+        for i in range(num_tokens)
+    ]
+
+
+class MockSelectionService:
+    """Records the selection-service reservation calls the actor's lifecycle
+    hooks make, standing in for the Dynamo selection service. ``add_output_block``
+    is synchronous as in the real binding; the rest are async."""
+
+    def __init__(self):
+        self.calls = []
+        self.reservations = []
+
+    async def create_reservation(self, request):
+        self.reservations.append(dict(request))
+        self.calls.append(
+            (
+                "create_reservation",
+                request["reservation_id"],
+                request["worker_id"],
+                len(request["token_ids"]),
+                request.get("expected_output_tokens"),
+            )
+        )
+
+    async def prefill_complete(self, reservation_id):
+        self.calls.append(("prefill_complete", reservation_id))
+
+    def add_output_block(self, reservation_id, *, decay_fraction=None):
+        self.calls.append(("add_output_block", reservation_id, decay_fraction))
+
+    async def free_reservation(self, reservation_id):
+        self.calls.append(("free_reservation", reservation_id))
+
+    async def delete_worker(self, worker_id):
+        self.calls.append(("delete_worker", worker_id))
+
+
+@ray.remote(num_cpus=0)
+class RecordingKVRouterActor(KVRouterActor):
+    """KVRouterActor that records the lifecycle events it receives for tests."""
+
+    def __init__(self, block_size):
+        self._block_size = block_size
+        self._replica_id_by_worker = {}
+        self._requests = OrderedDict()
+        self._request_ids_by_worker = {}
+        self._effective_prefill_tokens_by_request = {}
+        self._pending_tasks = set()
+        self._svc = MockSelectionService()
+        self._event_log = []
+
+    async def on_lifecycle_events(self, events):
+        self._event_log.extend(events)
+        await super().on_lifecycle_events(events)
+
+    def get_event_log(self):
+        return self._event_log
+
+    def get_selection_service_calls(self):
+        return self._svc.calls
+
+    async def get_request_lifecycle(self, request_id):
+        """Return a snapshot of an in-flight request's state, or ``None``."""
+        state = self._requests.get(request_id)
+        if state is None:
+            return None
+        snapshot = asdict(state)
+        snapshot.pop("created_at", None)  # internal TTL bookkeeping, not asserted on
+        return snapshot
+
+    async def get_active_request_ids(self):
+        """Return ids of the in-flight requests."""
+        return list(self._requests)
+
+    async def get_worker_active_load(self, worker_id):
+        """Return the number of in-flight requests attributed to ``worker_id``."""
+        return sum(1 for s in self._requests.values() if s.worker_id == worker_id)
+
+
+@ray.remote(num_cpus=0)
+class RaisingActor:
+    """A KV-router stand-in whose event ingest always raises, to prove the
+    engine token stream is never disrupted."""
+
+    async def on_lifecycle_events(self, events):
+        raise RuntimeError("actor down")
+
+
+class LocalKVRouterActor(KVRouterActor):
+    """In-process KVRouterActor with the event plane + Serve LongPoll stripped."""
+
+    def __init__(self, block_size):
+        self._block_size = block_size
+        self._replica_id_by_worker = {}
+        self._requests = OrderedDict()
+        self._request_ids_by_worker = {}
+        self._effective_prefill_tokens_by_request = {}
+        self._pending_tasks = set()
+        self._svc = MockSelectionService()
+
+    async def get_request_lifecycle(self, request_id):
+        """Return a snapshot of an in-flight request's state, or ``None``."""
+        state = self._requests.get(request_id)
+        if state is None:
+            return None
+        snapshot = asdict(state)
+        snapshot.pop("created_at", None)  # internal TTL bookkeeping, not asserted on
+        return snapshot
+
+    async def get_active_request_ids(self):
+        """Return ids of the in-flight requests."""
+        return list(self._requests)
+
+    async def get_worker_active_load(self, worker_id):
+        """Return the number of in-flight requests attributed to ``worker_id``."""
+        return sum(1 for s in self._requests.values() if s.worker_id == worker_id)
+
+
+@pytest.fixture
+def build_token_tracking_engine(monkeypatch):
+    def _build(script, actor, **engine_kwargs):
+        def get_deployment_actor(name):
+            assert name == KV_ROUTER_ACTOR_NAME
+            return actor
+
+        monkeypatch.setattr(serve, "get_deployment_actor", get_deployment_actor)
+        monkeypatch.setattr(
+            serve,
+            "get_replica_context",
+            lambda: SimpleNamespace(
+                replica_id=SimpleNamespace(unique_id=REPLICA_UNIQUE_ID)
+            ),
+        )
+        return enable_token_tracking(MockAsyncLLM)(script, **engine_kwargs)
+
+    return _build
+
+
+async def consume(stream, limit=None):
+    """Drain ``stream``, optionally closing it early after ``limit`` outputs."""
+    outputs = []
+    async for output in stream:
+        outputs.append(output)
+        if limit is not None and len(outputs) == limit:
+            await stream.aclose()
+    return outputs
+
+
+async def drain(engine):
+    """Wait for the engine forwarder's queued lifecycle batches to land."""
+    await engine._lifecycle_forwarder.flush()
+
+
+def decode_counts(events):
+    return [args[1] for name, args in events if name == "on_decode_progress"]
+
+
+def op_names(calls):
+    return [c[0] for c in calls]
+
+
+@pytest.mark.asyncio
+async def test_basic_lifecycle(build_token_tracking_engine):
+    """A streamed request reports add -> prefill -> exact decode counts -> done."""
+    actor = RecordingKVRouterActor.remote(block_size=16)
+    engine = build_token_tracking_engine(delta_steps(3, prompt_len=10), actor)
+
+    prompt = {"prompt_token_ids": list(range(10))}
+    outputs = await consume(
+        engine.generate(
+            prompt,
+            SamplingParams(output_kind=RequestOutputKind.DELTA, max_tokens=MAX_TOKENS),
+            "req-1",
+        )
+    )
+    await drain(engine)
+
+    assert ray.get(actor.get_event_log.remote()) == [
+        ("on_request_added", ("req-1", WORKER_ID, list(range(10)), MAX_TOKENS)),
+        ("on_prefill_complete", ("req-1",)),
+        ("on_decode_progress", ("req-1", 1)),
+        ("on_decode_progress", ("req-1", 2)),
+        ("on_decode_progress", ("req-1", 3)),
+        ("on_request_completed", ("req-1",)),
+    ]
+    assert outputs == engine.script
+
+
+@pytest.mark.asyncio
+async def test_lifecycle_uses_serve_request_id(build_token_tracking_engine):
+    """Lifecycle events use the same Serve request id used by routing, even if
+    vLLM's engine-level id is different."""
+    actor = RecordingKVRouterActor.remote(block_size=16)
+    engine = build_token_tracking_engine(delta_steps(1, prompt_len=10), actor)
+
+    prompt = {"prompt_token_ids": list(range(10))}
+    serve.context._serve_request_context.set(
+        serve.context._RequestContext(request_id="serve-route-id")
+    )
+    try:
+        await consume(
+            engine.generate(
+                prompt,
+                SamplingParams(
+                    output_kind=RequestOutputKind.DELTA, max_tokens=MAX_TOKENS
+                ),
+                "chatcmpl-serve-route-id",
+            )
+        )
+    finally:
+        serve.context._serve_request_context.set(serve.context._RequestContext())
+    await drain(engine)
+
+    assert ray.get(actor.get_event_log.remote()) == [
+        (
+            "on_request_added",
+            ("serve-route-id", WORKER_ID, list(range(10)), MAX_TOKENS),
+        ),
+        ("on_prefill_complete", ("serve-route-id",)),
+        ("on_decode_progress", ("serve-route-id", 1)),
+        ("on_request_completed", ("serve-route-id",)),
+    ]
+
+
+@pytest.mark.asyncio
+async def test_in_order_reports(build_token_tracking_engine):
+    """Back-to-back reports reach the actor in submission order."""
+    actor = RecordingKVRouterActor.remote(block_size=16)
+    engine = build_token_tracking_engine(delta_steps(200), actor)
+
+    await consume(
+        engine.generate(
+            PROMPT, SamplingParams(output_kind=RequestOutputKind.DELTA), "r"
+        )
+    )
+    await drain(engine)
+
+    events = ray.get(actor.get_event_log.remote())
+    assert events[0][0] == "on_request_added"
+    assert events[1][0] == "on_prefill_complete"
+    assert events[-1] == ("on_request_completed", ("r",))
+    assert decode_counts(events) == list(range(1, 201))
+
+
+@pytest.mark.asyncio
+async def test_streaming_accumulates_decode_progress(build_token_tracking_engine):
+    """A DELTA (streaming) request sums each step's new tokens into a running
+    total reported as decode progress."""
+    actor = RecordingKVRouterActor.remote(block_size=16)
+    # Steps carry only new tokens: 1, then 2, then 1 -> cumulative 1, 3, 4.
+    script = [request_output([n]) for n in (1, 2, 1)]
+    engine = build_token_tracking_engine(script, actor)
+
+    await consume(
+        engine.generate(
+            PROMPT, SamplingParams(output_kind=RequestOutputKind.DELTA), "r"
+        )
+    )
+    await drain(engine)
+
+    assert decode_counts(ray.get(actor.get_event_log.remote())) == [1, 3, 4]
+
+
+@pytest.mark.asyncio
+async def test_non_streaming_reports_full_output_once(build_token_tracking_engine):
+    """A FINAL_ONLY (non-streaming) request arrives as one finished chunk, with a
+    CompletionOutput per candidate, so progress is reported once at the summed
+    token count across candidates."""
+    actor = RecordingKVRouterActor.remote(block_size=16)
+    # FINAL_ONLY n=3: a single finished chunk carrying every candidate's output.
+    script = [request_output([2, 3, 4], finished=True)]
+    engine = build_token_tracking_engine(script, actor)
+
+    await consume(
+        engine.generate(
+            PROMPT, SamplingParams(output_kind=RequestOutputKind.FINAL_ONLY), "r"
+        )
+    )
+    await drain(engine)
+
+    events = ray.get(actor.get_event_log.remote())
+    assert decode_counts(events) == [9]  # 2 + 3 + 4 summed across candidates
+    assert [name for name, _ in events] == [
+        "on_request_added",
+        "on_prefill_complete",
+        "on_decode_progress",
+        "on_request_completed",
+    ]
+
+
+@pytest.mark.asyncio
+async def test_cumulative_skips_tracking(build_token_tracking_engine):
+    """CUMULATIVE repeats output-so-far each chunk; tracking is skipped (not
+    summed) to avoid over-counting, so no lifecycle events are reported."""
+    actor = RecordingKVRouterActor.remote(block_size=16)
+    engine = build_token_tracking_engine(delta_steps(3), actor)
+
+    outputs = await consume(
+        engine.generate(
+            PROMPT, SamplingParams(output_kind=RequestOutputKind.CUMULATIVE), "r"
+        )
+    )
+
+    assert len(outputs) == 3  # stream still passes through untouched
+    assert ray.get(actor.get_event_log.remote()) == []
+
+
+@pytest.mark.asyncio
+async def test_empty_steps_ignored(build_token_tracking_engine):
+    """Token-less outputs (e.g. a finish-only chunk) emit no progress hooks."""
+    actor = RecordingKVRouterActor.remote(block_size=16)
+    script = [
+        request_output([1]),
+        request_output([0]),  # structural chunk: no new tokens
+        request_output([1], finished=True),
+    ]
+    engine = build_token_tracking_engine(script, actor)
+
+    await consume(
+        engine.generate(
+            PROMPT, SamplingParams(output_kind=RequestOutputKind.DELTA), "r"
+        )
+    )
+    await drain(engine)
+
+    events = ray.get(actor.get_event_log.remote())
+    assert decode_counts(events) == [1, 2]
+    assert [e for e in events if e[0] == "on_prefill_complete"] == [
+        ("on_prefill_complete", ("r",))
+    ]
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("early_drop", [False, True])
+async def test_completed_exactly_once(early_drop, build_token_tracking_engine):
+    """Completion fires exactly once on normal end and on early stream close."""
+    actor = RecordingKVRouterActor.remote(block_size=16)
+    engine = build_token_tracking_engine(delta_steps(3), actor)
+
+    stream = engine.generate(
+        PROMPT, SamplingParams(output_kind=RequestOutputKind.DELTA), "r"
+    )
+    await consume(stream, limit=1 if early_drop else None)
+    await drain(engine)
+
+    events = ray.get(actor.get_event_log.remote())
+    assert [e for e in events if e[0] == "on_request_completed"] == [
+        ("on_request_completed", ("r",))
+    ]
+
+
+@pytest.mark.asyncio
+async def test_engine_error_still_completes(build_token_tracking_engine):
+    """A mid-stream engine error propagates but still frees the request."""
+    actor = RecordingKVRouterActor.remote(block_size=16)
+    engine = build_token_tracking_engine(delta_steps(3), actor, error_after=1)
+
+    with pytest.raises(RuntimeError, match="engine failure"):
+        await consume(
+            engine.generate(
+                PROMPT, SamplingParams(output_kind=RequestOutputKind.DELTA), "r"
+            )
+        )
+    await drain(engine)
+
+    events = ray.get(actor.get_event_log.remote())
+    assert events[-1] == ("on_request_completed", ("r",))
+
+
+@pytest.mark.asyncio
+async def test_zero_token_request(build_token_tracking_engine):
+    """An output-less request (e.g. validation abort) is added and freed only."""
+    actor = RecordingKVRouterActor.remote(block_size=16)
+    engine = build_token_tracking_engine([request_output([0], finished=True)], actor)
+
+    prompt = {"prompt_token_ids": [1, 2, 3]}
+    await consume(
+        engine.generate(
+            prompt,
+            SamplingParams(output_kind=RequestOutputKind.DELTA, max_tokens=MAX_TOKENS),
+            "r",
+        )
+    )
+    await drain(engine)
+
+    assert ray.get(actor.get_event_log.remote()) == [
+        ("on_request_added", ("r", WORKER_ID, [1, 2, 3], MAX_TOKENS)),
+        ("on_request_completed", ("r",)),
+    ]
+
+
+@pytest.mark.asyncio
+async def test_actor_failure_isolation(build_token_tracking_engine):
+    """A failing actor never disrupts the engine's output stream."""
+    engine = build_token_tracking_engine(delta_steps(2), RaisingActor.remote())
+
+    outputs = await consume(
+        engine.generate(
+            PROMPT, SamplingParams(output_kind=RequestOutputKind.DELTA), "r"
+        )
+    )
+    await drain(engine)  # the failed batches are dropped without raising
+
+    assert len(outputs) == 2
+
+
+def test_decorator_returns_subclass():
+    """The decorator returns an isinstance-compatible subclass."""
+    assert issubclass(enable_token_tracking(MockAsyncLLM), MockAsyncLLM)
+
+
+@pytest.mark.asyncio
+async def test_passthrough_without_actor(monkeypatch):
+    """Outside a replica (no actor resolvable) the engine is a pure pass-through."""
+
+    def _raise(name):
+        raise RuntimeError("no actor")
+
+    monkeypatch.setattr(serve, "get_deployment_actor", _raise)
+    engine = enable_token_tracking(MockAsyncLLM)(delta_steps(2))
+
+    outputs = await consume(
+        engine.generate(
+            PROMPT, SamplingParams(output_kind=RequestOutputKind.DELTA), "r"
+        )
+    )
+
+    assert len(outputs) == 2
+    assert engine._lifecycle_forwarder is None  # resolution failed; retried next call
+
+
+@pytest.mark.parametrize(
+    "request_router_config, expected",
+    [
+        ({"request_router_class": KVAwareRouter}, True),
+        ({}, False),  # default (non-KV) router
+        (None, False),  # no router configured
+    ],
+)
+def test_is_kv_aware(request_router_config, expected):
+    """The engine wraps with token tracking only on KVAwareRouter deployments,
+    so non-KV deployments never retry a missing actor lookup per request."""
+    deployment_config = {}
+    if request_router_config is not None:
+        deployment_config["request_router_config"] = request_router_config
+    llm_config = LLMConfig(
+        model_loading_config={
+            "model_id": "qwen3-0.6b",
+            "model_source": "Qwen/Qwen3-0.6B",
+        },
+        accelerator_type=None,
+        deployment_config=deployment_config,
+    )
+    assert is_kv_aware(llm_config) is expected
+
+
+@pytest.mark.asyncio
+async def test_block_boundary_crossings():
+    """Each ceil((prompt+output)/block_size) increase advances total_blocks."""
+    actor = LocalKVRouterActor(block_size=16)
+    await actor.on_request_added("r", 1, list(range(10)))
+    assert (await actor.get_request_lifecycle("r"))["total_blocks"] == 1  # ceil(10/16)
+
+    await actor.on_decode_progress("r", 6)  # 10+6=16 -> still 1 block
+    assert (await actor.get_request_lifecycle("r"))["total_blocks"] == 1
+
+    await actor.on_decode_progress("r", 7)  # 17 -> crosses into block 2
+    assert (await actor.get_request_lifecycle("r"))["total_blocks"] == 2
+
+    await actor.on_decode_progress("r", 39)  # 49 -> ceil=4, crosses two more at once
+    snapshot = await actor.get_request_lifecycle("r")
+    assert snapshot["total_blocks"] == 4
+    assert snapshot["output_tokens"] == 39
+
+
+@pytest.mark.asyncio
+async def test_active_load_tracking():
+    """Active load is per-worker; completion evicts the request entirely."""
+    actor = LocalKVRouterActor(block_size=16)
+    await actor.on_request_added("a", 1, list(range(8)))
+    await actor.on_request_added("b", 1, [])
+    await actor.on_request_added("c", 2, [])
+    assert await actor.get_worker_active_load(1) == 2
+    assert await actor.get_worker_active_load(2) == 1
+
+    await actor.on_prefill_complete("a")
+    await actor.on_decode_progress("a", 5)
+    assert await actor.get_worker_active_load(1) == 2  # still active while decoding
+    assert await actor.get_request_lifecycle("a") == {
+        "worker_id": 1,
+        "prompt_tokens": 8,
+        "expected_output_tokens": None,
+        "prefill_completed": True,
+        "output_tokens": 5,
+        "total_blocks": 1,
+    }
+
+    # Completion evicts (bounding memory to in-flight requests).
+    await actor.on_request_completed("a")
+    assert await actor.get_worker_active_load(1) == 1
+    assert set(await actor.get_active_request_ids()) == {"b", "c"}
+    assert await actor.get_request_lifecycle("a") is None
+
+    # Hooks for an unknown request id are ignored.
+    await actor.on_prefill_complete("missing")
+    await actor.on_decode_progress("missing", 3)
+    await actor.on_request_completed("missing")
+    assert await actor.get_request_lifecycle("missing") is None
+
+
+@pytest.mark.asyncio
+async def test_tracking_drains_under_churn():
+    """Memory chaos: a long run of interleaved submissions and completions across
+    workers grows the in-flight state and drains it back to nothing."""
+    rng = random.Random(20240708)
+    actor = LocalKVRouterActor(block_size=16)
+    workers = [1, 2, 3]
+    total = 400
+    inflight: set = set()
+    launched = 0
+    peak = 0
+
+    # Randomly interleave admitting new requests and completing live ones.
+    for _ in range(total * 3):
+        if launched < total and (not inflight or rng.random() < 0.6):
+            request_id = f"r{launched}"
+            launched += 1
+            # A routed request carries its effective prefill tokens from select();
+            # admission must drain that map too, not just _requests.
+            actor._effective_prefill_tokens_by_request[request_id] = rng.randint(0, 40)
+            await actor.on_request_added(
+                request_id,
+                rng.choice(workers),
+                list(range(rng.randint(1, 40))),
+                expected_output_tokens=rng.choice([None, 32]),
+            )
+            inflight.add(request_id)
+        elif inflight:
+            request_id = rng.choice(list(inflight))
+            if rng.random() < 0.5:
+                await actor.on_prefill_complete(request_id)
+                await actor.on_decode_progress(request_id, rng.randint(1, 80))
+            await actor.on_request_completed(request_id)
+            inflight.discard(request_id)
+        peak = max(peak, len(actor._requests))
+
+    for request_id in list(inflight):
+        await actor.on_request_completed(request_id)
+
+    assert launched == total
+    assert peak > 1  # state actually accumulated under concurrent load
+    # Everything drained: no request state, no index entries, no active load.
+    assert actor._requests == {}
+    assert actor._request_ids_by_worker == {}
+    assert actor._effective_prefill_tokens_by_request == {}
+    assert await actor.get_active_request_ids() == []
+    for worker_id in workers:
+        assert await actor.get_worker_active_load(worker_id) == 0
+
+
+@pytest.mark.asyncio
+async def test_remove_worker_evicts_requests():
+    """A departed replica's in-flight request state is purged: its completion
+    events can never arrive, so the entries would otherwise leak forever."""
+    actor = LocalKVRouterActor(block_size=16)
+    await actor.on_request_added("a", 1, list(range(8)))
+    await actor.on_request_added("b", 1, list(range(8)))
+    await actor.on_request_added("c", 2, list(range(8)))
+
+    actor.remove_worker(1)
+
+    assert set(await actor.get_active_request_ids()) == {"c"}
+    assert await actor.get_worker_active_load(1) == 0
+    assert await actor.get_worker_active_load(2) == 1
+    # The reverse index drops the departed worker and keeps the survivor.
+    assert 1 not in actor._request_ids_by_worker
+    assert actor._request_ids_by_worker.get(2) == {"c"}
+    # remove_worker schedules delete_worker; drain the task and confirm.
+    await asyncio.gather(*list(actor._pending_tasks))
+    assert ("delete_worker", 1) in actor._svc.calls
+
+
+@pytest.mark.asyncio
+async def test_stale_request_evicted_after_ttl():
+    """Backstop for a lost completion on a live replica: a request tracked past
+    the TTL is evicted and its reservation freed, so it cannot accumulate."""
+    actor = LocalKVRouterActor(block_size=16)
+    await actor.on_request_added("stale", 1, list(range(8)))
+    await actor.on_request_added("fresh_untriggered", 1, list(range(8)))
+    # Backdate the stale request's admission beyond the TTL (lost completion).
+    actor._requests["stale"].created_at -= REQUEST_TRACKING_TTL_S + 1
+
+    # A new admission triggers the lazy sweep of the oldest entries.
+    await actor.on_request_added("trigger", 2, list(range(8)))
+
+    assert "stale" not in await actor.get_active_request_ids()
+    assert ("free_reservation", "stale") in actor._svc.calls
+    # A still-fresh request is left untouched (sweep stops at the first fresh one).
+    assert set(await actor.get_active_request_ids()) == {"fresh_untriggered", "trigger"}
+    assert ("free_reservation", "fresh_untriggered") not in actor._svc.calls
+
+
+@pytest.mark.asyncio
+async def test_admission_race_frees_reservation():
+    """If remove_worker evicts a request while create_reservation is in flight,
+    the orphaned reservation is freed rather than leaking in the service."""
+    actor = LocalKVRouterActor(block_size=16)
+
+    # Simulate the LongPoll remove_worker firing during the create_reservation await.
+    book_reservation = actor._svc.create_reservation
+
+    async def racing_create_reservation(request):
+        await book_reservation(request)
+        actor.remove_worker(request["worker_id"])
+
+    actor._svc.create_reservation = racing_create_reservation
+
+    await actor.on_request_added("r", 1, list(range(8)))
+
+    assert await actor.get_active_request_ids() == []  # evicted mid-flight
+    assert ("free_reservation", "r") in actor._svc.calls  # reservation not orphaned
+
+
+@pytest.mark.asyncio
+async def test_tracks_streamed_request_state(build_token_tracking_engine):
+    """End-to-end: exact token counts land as actor block state over ``.remote``."""
+    actor = RecordingKVRouterActor.remote(block_size=8)
+    # prompt 12, block_size 8: baseline ceil(12/8)=2; boundaries at cumulative
+    # output 5 and 13 -> 9 generated tokens cross only the first.
+    engine = build_token_tracking_engine(delta_steps(9, prompt_len=12), actor)
+
+    prompt = {"prompt_token_ids": list(range(12))}
+    stream = engine.generate(
+        prompt,
+        SamplingParams(output_kind=RequestOutputKind.DELTA, max_tokens=MAX_TOKENS),
+        "req-e2e",
+    )
+    for _ in range(9):
+        await stream.__anext__()
+    await drain(engine)
+
+    # All outputs consumed but the stream is still open: the request is
+    # tracked with its exact final counts.
+    assert await actor.get_request_lifecycle.remote("req-e2e") == {
+        "worker_id": WORKER_ID,
+        "prompt_tokens": 12,
+        "expected_output_tokens": MAX_TOKENS,
+        "prefill_completed": True,
+        "output_tokens": 9,
+        "total_blocks": 3,
+    }
+    assert await actor.get_worker_active_load.remote(WORKER_ID) == 1
+
+    # Stream end fires completion, which evicts the request.
+    with pytest.raises(StopAsyncIteration):
+        await stream.__anext__()
+    await drain(engine)
+    assert await actor.get_request_lifecycle.remote("req-e2e") is None
+    assert await actor.get_active_request_ids.remote() == []
+    assert await actor.get_worker_active_load.remote(WORKER_ID) == 0
+
+
+@pytest.mark.asyncio
+async def test_lifecycle_books_selection_service_load(build_token_tracking_engine):
+    """A streamed request books create_reservation -> prefill_complete -> one
+    add_output_block per crossed decode block -> free_reservation, in order."""
+    actor = RecordingKVRouterActor.remote(block_size=8)
+    # prompt 12 (2 blocks); cumulative output crosses into block 3 at 5 tokens
+    # and block 4 at 13 tokens -> exactly two output blocks over 20 tokens.
+    engine = build_token_tracking_engine(delta_steps(20, prompt_len=12), actor)
+
+    prompt = {"prompt_token_ids": list(range(12))}
+    await consume(
+        engine.generate(
+            prompt,
+            SamplingParams(output_kind=RequestOutputKind.DELTA, max_tokens=MAX_TOKENS),
+            "req-1",
+        )
+    )
+    await drain(engine)
+
+    calls = ray.get(actor.get_selection_service_calls.remote())
+    assert op_names(calls) == (
+        ["create_reservation", "prefill_complete"]
+        + ["add_output_block"] * 2
+        + ["free_reservation"]
+    )
+    assert calls[0] == ("create_reservation", "req-1", WORKER_ID, 12, MAX_TOKENS)
+    assert calls[-1] == ("free_reservation", "req-1")
+
+
+@pytest.mark.asyncio
+async def test_decode_blocks_book_add_output_block():
+    """Each crossed decode block books one add_output_block in the service."""
+    actor = LocalKVRouterActor(block_size=16)
+    await actor.on_request_added("r", WORKER_ID, list(range(10)))  # 1 prompt block
+    await actor.on_decode_progress("r", 6)  # 16 -> still 1 block
+    await actor.on_decode_progress("r", 7)  # 17 -> crosses into block 2
+    await actor.on_decode_progress("r", 39)  # 49 -> ceil=4, crosses two more
+
+    assert (
+        op_names(actor._svc.calls) == ["create_reservation"] + ["add_output_block"] * 3
+    )
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+    "expected_output_tokens, expected_decay",
+    [
+        pytest.param(40, pytest.approx(1.0 - 8 / 40), id="with-estimate"),
+        pytest.param(None, None, id="no-estimate"),
+    ],
+)
+async def test_expected_output_tokens_sets_decay_fraction(
+    expected_output_tokens, expected_decay
+):
+    """With an output-length estimate each booked decode block decays by the
+    remaining fraction; without one the block carries no decay."""
+    actor = LocalKVRouterActor(block_size=8)
+    await actor.on_request_added(
+        "r", WORKER_ID, list(range(8)), expected_output_tokens=expected_output_tokens
+    )
+    await actor.on_decode_progress("r", 8)  # total 16 -> crosses into block 2
+
+    block_calls = [c for c in actor._svc.calls if c[0] == "add_output_block"]
+    assert block_calls == [("add_output_block", "r", expected_decay)]
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main(["-v", __file__]))

python/ray/llm/tests/serve/cpu/deployments/routers/test_builder_ingress.py

@@ -0,0 +1,805 @@
+import os
+import re
+import signal
+import subprocess
+import sys
+import tempfile
+
+import pytest
+import yaml
+
+from ray import serve
+from ray._common.test_utils import wait_for_condition
+from ray.llm._internal.serve.constants import DEFAULT_MAX_TARGET_ONGOING_REQUESTS
+from ray.llm._internal.serve.core.configs.llm_config import (
+    LLMConfig,
+    ModelLoadingConfig,
+)
+from ray.llm._internal.serve.core.ingress.builder import (
+    IngressClsConfig,
+    LLMServingArgs,
+    build_openai_app,
+)
+from ray.llm._internal.serve.core.ingress.ingress import OpenAiIngress
+from ray.llm._internal.serve.serving_patterns.data_parallel.builder import (
+    build_dp_openai_app,
+)
+from ray.llm._internal.serve.serving_patterns.data_parallel.dp_server import (
+    DPServer,
+)
+from ray.llm._internal.serve.serving_patterns.prefill_decode.builder import (
+    build_pd_openai_app,
+)
+from ray.llm._internal.serve.serving_patterns.prefill_decode.pd_server import (
+    DPPDDecodeServer,
+    DPPDPrefillServer,
+    PDDecodeServer,
+    PDPrefillServer,
+)
+from ray.serve._private.http_util import ASGIAppReplicaWrapper
+from ray.serve.config import AutoscalingConfig, RequestRouterConfig
+from ray.serve.experimental.consistent_hash_router import ConsistentHashRouter
+from ray.serve.experimental.round_robin_router import RoundRobinRouter
+
+
+@pytest.fixture
+def get_llm_serve_args(llm_config_with_mock_engine):
+    yield LLMServingArgs(llm_configs=[llm_config_with_mock_engine])
+
+
+@pytest.fixture()
+def serve_config_separate_model_config_files():
+    config_dir = tempfile.mkdtemp()
+    serve_config_filename = "llm_app_separate_model_config_files.yaml"
+    config_root = os.path.join(os.path.dirname(__file__), "test_config_files")
+    serve_config_src = os.path.join(config_root, serve_config_filename)
+    serve_config_dst = os.path.join(config_dir, serve_config_filename)
+
+    with open(serve_config_src, "r") as f:
+        serve_config_yaml = yaml.safe_load(f)
+
+    for application in serve_config_yaml["applications"]:
+        llm_configs = application["args"]["llm_configs"]
+        tmp_llm_config_files = []
+        for llm_config in llm_configs:
+            llm_config_src = llm_config.replace(".", config_root, 1)
+            llm_config_dst = llm_config.replace(".", config_dir, 1)
+            tmp_llm_config_files.append(llm_config_dst)
+
+            with open(llm_config_src, "r") as f:
+                llm_config_yaml = yaml.safe_load(f)
+
+            # Make sure engine is mocked.
+            if llm_config_yaml.get("runtime_env", None) is None:
+                llm_config_yaml["runtime_env"] = {}
+            llm_config_yaml["runtime_env"]["env_vars"] = {
+                "RAYLLM_VLLM_ENGINE_CLS": "ray.llm.tests.serve.mocks.mock_vllm_engine.MockVLLMEngine"
+            }
+
+            # Explicitly set accelerator_type to None to avoid GPU placement groups
+            llm_config_yaml["accelerator_type"] = None
+
+            # Use placement_group_config to specify CPU-only bundles
+            llm_config_yaml["placement_group_config"] = {
+                "bundles": [{"CPU": 1, "GPU": 0}]
+            }
+
+            os.makedirs(os.path.dirname(llm_config_dst), exist_ok=True)
+            with open(llm_config_dst, "w") as f:
+                yaml.dump(llm_config_yaml, f)
+
+        application["args"]["llm_configs"] = tmp_llm_config_files
+
+    with open(serve_config_dst, "w") as f:
+        yaml.dump(serve_config_yaml, f)
+
+    yield serve_config_dst
+
+
+class TestLLMServingArgs:
+    """Test suite for LLMServingArgs data model."""
+
+    @pytest.fixture
+    def llm_config(self):
+        """Basic LLMConfig for testing."""
+        return LLMConfig(
+            model_loading_config=ModelLoadingConfig(
+                model_id="test-model", model_source="test-source"
+            )
+        )
+
+    def test_basic_creation_and_defaults(self, llm_config):
+        """Test creation with minimal config and verify defaults."""
+        args = LLMServingArgs(llm_configs=[llm_config])
+
+        # Verify llm_configs
+        assert len(args.llm_configs) == 1
+        assert isinstance(args.llm_configs[0], LLMConfig)
+
+        # Verify defaults
+        assert isinstance(args.ingress_cls_config, IngressClsConfig)
+        assert args.ingress_cls_config.ingress_cls == OpenAiIngress
+        assert args.ingress_deployment_config == {}
+
+    def test_flexible_input_types(self, llm_config):
+        """Test accepts dicts, objects, and mixed types for llm_configs."""
+        config_dict = {
+            "model_loading_config": {
+                "model_id": "test-model-2",
+                "model_source": "test-source-2",
+            }
+        }
+        args = LLMServingArgs(llm_configs=[llm_config, config_dict])
+        assert len(args.llm_configs) == 2
+        assert all(isinstance(c, LLMConfig) for c in args.llm_configs)
+
+    def test_ingress_config_flexibility(self, llm_config):
+        """Test ingress_cls_config: defaults, dict input, object input, and class loading."""
+        # Test defaults
+        args_default = LLMServingArgs(llm_configs=[llm_config])
+        assert isinstance(args_default.ingress_cls_config, IngressClsConfig)
+        assert args_default.ingress_cls_config.ingress_cls == OpenAiIngress
+        assert args_default.ingress_cls_config.ingress_extra_kwargs == {}
+
+        # Test as dict with custom kwargs
+        args_dict = LLMServingArgs(
+            llm_configs=[llm_config],
+            ingress_cls_config={"ingress_extra_kwargs": {"key": "value"}},
+        )
+        assert isinstance(args_dict.ingress_cls_config, IngressClsConfig)
+        assert args_dict.ingress_cls_config.ingress_extra_kwargs == {"key": "value"}
+
+        # Test as object
+        args_obj = LLMServingArgs(
+            llm_configs=[llm_config],
+            ingress_cls_config=IngressClsConfig(ingress_extra_kwargs={"key": "value"}),
+        )
+        assert isinstance(args_obj.ingress_cls_config, IngressClsConfig)
+        assert args_obj.ingress_cls_config.ingress_extra_kwargs == {"key": "value"}
+
+        # Test class loading from string
+        args_str = LLMServingArgs(
+            llm_configs=[llm_config],
+            ingress_cls_config={
+                "ingress_cls": "ray.llm._internal.serve.core.ingress.ingress:OpenAiIngress"
+            },
+        )
+        assert args_str.ingress_cls_config.ingress_cls == OpenAiIngress
+
+    def test_validation_rules(self):
+        """Test validation: unique model IDs and non-empty list."""
+        # Duplicate model IDs
+        config1 = LLMConfig(
+            model_loading_config=ModelLoadingConfig(
+                model_id="same-id", model_source="source1"
+            )
+        )
+        config2 = LLMConfig(
+            model_loading_config=ModelLoadingConfig(
+                model_id="same-id", model_source="source2"
+            )
+        )
+        with pytest.raises(ValueError, match="Duplicate models found"):
+            LLMServingArgs(llm_configs=[config1, config2])
+
+        # Empty list
+        with pytest.raises(ValueError, match="List of models is empty"):
+            LLMServingArgs(llm_configs=[])
+
+
+class TestBuildOpenaiApp:
+    @pytest.fixture
+    def llm_config(self):
+        """Basic LLMConfig for testing."""
+        return LLMConfig(
+            model_loading_config=ModelLoadingConfig(
+                model_id="test-model", model_source="test-source"
+            )
+        )
+
+    def test_build_openai_app(
+        self, get_llm_serve_args, shutdown_ray_and_serve, disable_placement_bundles
+    ):
+        """Test `build_openai_app` can build app and run it with Serve."""
+
+        app = build_openai_app(
+            get_llm_serve_args,
+        )
+        assert isinstance(app, serve.Application)
+        serve.run(app)
+
+    def test_build_openai_app_with_config(
+        self,
+        serve_config_separate_model_config_files,
+        shutdown_ray_and_serve,
+        disable_placement_bundles,
+    ):
+        """Test `build_openai_app` can be used in serve config."""
+
+        def deployments_healthy():
+            status_response = subprocess.check_output(["serve", "status"])
+            print("[TEST] Status response: ", status_response)
+            applications = extract_applications_from_output(status_response)
+
+            if "llm-endpoint" not in applications:
+                print("[TEST] Application 'llm-endpoint' not found.")
+                return False
+
+            llm_endpoint_status = applications["llm-endpoint"]
+            if len(llm_endpoint_status["deployments"]) != 2:
+                print(
+                    f"[TEST] Expected 2 deployments, found {len(llm_endpoint_status['deployments'])}"
+                )
+                return False
+
+            deployment_status = llm_endpoint_status["deployments"].values()
+            if not all([status["status"] == "HEALTHY" for status in deployment_status]):
+                print(f"[TEST] Not all deployments healthy: {deployment_status}")
+                return False
+
+            print("[TEST] All deployments healthy.")
+            return True
+
+        p = subprocess.Popen(["serve", "run", serve_config_separate_model_config_files])
+        wait_for_condition(deployments_healthy, timeout=60, retry_interval_ms=1000)
+
+        p.send_signal(signal.SIGINT)  # Equivalent to ctrl-C
+        p.wait()
+
+    def test_router_built_with_autoscaling_configs(self, disable_placement_bundles):
+        """Test that the router is built with the correct autoscaling configs that
+        will scale.
+        """
+        llm_config_no_autoscaling_configured = LLMConfig(
+            model_loading_config=ModelLoadingConfig(model_id="model_id_1"),
+            accelerator_type="L4",
+        )
+        llm_config_autoscaling_default = LLMConfig(
+            model_loading_config=ModelLoadingConfig(model_id="model_id_2"),
+            accelerator_type="L4",
+            deployment_config={"autoscaling_config": AutoscalingConfig()},
+        )
+        llm_config_autoscaling_non_default = LLMConfig(
+            model_loading_config=ModelLoadingConfig(model_id="model_id_3"),
+            accelerator_type="L4",
+            deployment_config={
+                "autoscaling_config": AutoscalingConfig(
+                    min_replicas=2,
+                    initial_replicas=3,
+                    max_replicas=4,
+                )
+            },
+        )
+
+        app = build_openai_app(
+            LLMServingArgs(
+                llm_configs=[
+                    llm_config_no_autoscaling_configured,
+                    llm_config_autoscaling_default,
+                    llm_config_autoscaling_non_default,
+                ],
+                ingress_deployment_config={
+                    "autoscaling_config": {
+                        "min_replicas": 8,
+                        "initial_replicas": 10,
+                        "max_replicas": 12,
+                        "target_ongoing_requests": 10,
+                    }
+                },
+            )
+        )
+        router_autoscaling_config = (
+            app._bound_deployment._deployment_config.autoscaling_config
+        )
+        assert router_autoscaling_config.min_replicas == 8  # (1 + 1 + 2) * 2
+        assert router_autoscaling_config.initial_replicas == 10  # (1 + 1 + 3) * 2
+        assert router_autoscaling_config.max_replicas == 12  # (1 + 1 + 4) * 2
+        assert router_autoscaling_config.target_ongoing_requests == 10
+
+    def test_ingress_deployment_config_merging(
+        self, llm_config, disable_placement_bundles
+    ):
+        """Test that ingress_deployment_config is properly merged with default options.
+
+        This test ensures that deep_merge_dicts return value is properly assigned
+        and that nested dictionaries are properly deep-merged without losing default values.
+        """
+        # Build app with custom ingress deployment config including nested options
+        app = build_openai_app(
+            dict(
+                llm_configs=[llm_config],
+                ingress_deployment_config={
+                    "num_replicas": 3,
+                    "ray_actor_options": {
+                        "num_cpus": 4,
+                        "memory": 1024,
+                    },
+                    "max_ongoing_requests": 200,  # Override default
+                },
+            )
+        )
+
+        # Verify the custom config was applied
+        deployment = app._bound_deployment
+        assert deployment._deployment_config.num_replicas == 3
+        assert deployment.ray_actor_options["num_cpus"] == 4
+        assert deployment.ray_actor_options["memory"] == 1024
+        assert deployment._deployment_config.max_ongoing_requests == 200
+
+    def test_default_autoscaling_config_included_without_num_replicas(
+        self, llm_config, disable_placement_bundles
+    ):
+        """Test that default autoscaling_config with target_ongoing_requests is included
+        when num_replicas is not specified.
+        """
+        app = build_openai_app(
+            dict(
+                llm_configs=[llm_config],
+            )
+        )
+
+        deployment = app._bound_deployment
+        autoscaling_config = deployment._deployment_config.autoscaling_config
+        assert autoscaling_config is not None
+        assert (
+            autoscaling_config.target_ongoing_requests
+            == DEFAULT_MAX_TARGET_ONGOING_REQUESTS
+        )
+
+    def test_autoscaling_config_removed_from_defaults_when_num_replicas_specified(
+        self, llm_config, disable_placement_bundles
+    ):
+        """Test that autoscaling_config from defaults is removed when user specifies
+        num_replicas, since Ray Serve does not allow both.
+        """
+        app = build_openai_app(
+            dict(
+                llm_configs=[llm_config],
+                ingress_deployment_config={
+                    "num_replicas": 2,
+                },
+            )
+        )
+
+        deployment = app._bound_deployment
+        assert deployment._deployment_config.num_replicas == 2
+        # autoscaling_config should be None since num_replicas is set
+        assert deployment._deployment_config.autoscaling_config is None
+
+    def test_user_target_ongoing_requests_respected(
+        self, llm_config, disable_placement_bundles
+    ):
+        """Test that user-specified target_ongoing_requests is respected and not
+        overridden by defaults.
+        """
+        user_target = 50
+        app = build_openai_app(
+            dict(
+                llm_configs=[llm_config],
+                ingress_deployment_config={
+                    "autoscaling_config": {
+                        "target_ongoing_requests": user_target,
+                    },
+                },
+            )
+        )
+
+        deployment = app._bound_deployment
+        autoscaling_config = deployment._deployment_config.autoscaling_config
+        assert autoscaling_config is not None
+        assert autoscaling_config.target_ongoing_requests == user_target
+
+    def test_direct_streaming_builds_ingress_with_router_attached(
+        self, llm_config, disable_placement_bundles, monkeypatch
+    ):
+        monkeypatch.setattr(
+            "ray.llm._internal.serve.core.ingress.builder."
+            "RAY_SERVE_LLM_ENABLE_DIRECT_STREAMING",
+            True,
+        )
+
+        app = build_openai_app(LLMServingArgs(llm_configs=[llm_config]))
+        ingress_request_router = app._ingress_request_router
+
+        assert app._bound_deployment.name == "LLMServer:test-model"
+        assert issubclass(app._bound_deployment.func_or_class, ASGIAppReplicaWrapper)
+        assert ingress_request_router is not None
+        assert ingress_request_router._bound_deployment.name == "LLMRouter"
+        assert ingress_request_router._bound_deployment.init_kwargs["server"] is app
+
+        # `RequestRouterConfig._serialize_request_router_cls` normalizes the
+        # class to its import path at config-build time.
+        request_router_config = (
+            app._bound_deployment._deployment_config.request_router_config
+        )
+        assert request_router_config.request_router_class == (
+            f"{RoundRobinRouter.__module__}.{RoundRobinRouter.__name__}"
+        )
+
+    def test_direct_streaming_user_request_router_config_wins(
+        self, llm_config, disable_placement_bundles, monkeypatch
+    ):
+        """A user-supplied ``request_router_config`` on ``LLMConfig`` must
+        survive direct-streaming wiring rather than being overwritten with the
+        default.
+        """
+        monkeypatch.setattr(
+            "ray.llm._internal.serve.core.ingress.builder."
+            "RAY_SERVE_LLM_ENABLE_DIRECT_STREAMING",
+            True,
+        )
+        llm_config.deployment_config["request_router_config"] = RequestRouterConfig(
+            request_router_class=ConsistentHashRouter,
+        )
+
+        app = build_openai_app(LLMServingArgs(llm_configs=[llm_config]))
+        request_router_config = (
+            app._bound_deployment._deployment_config.request_router_config
+        )
+        assert request_router_config.request_router_class == (
+            f"{ConsistentHashRouter.__module__}.{ConsistentHashRouter.__name__}"
+        )
+
+    def test_direct_streaming_rejects_multiple_llm_configs(
+        self, llm_config, disable_placement_bundles, monkeypatch
+    ):
+        monkeypatch.setattr(
+            "ray.llm._internal.serve.core.ingress.builder."
+            "RAY_SERVE_LLM_ENABLE_DIRECT_STREAMING",
+            True,
+        )
+        other_llm_config = LLMConfig(
+            model_loading_config=ModelLoadingConfig(model_id="other-model")
+        )
+
+        with pytest.raises(
+            ValueError,
+            match="currently supports exactly one LLM config",
+        ):
+            build_openai_app(LLMServingArgs(llm_configs=[llm_config, other_llm_config]))
+
+    @pytest.mark.parametrize(
+        ("builder_kwargs", "match"),
+        [
+            (
+                {"ingress_deployment_config": {"num_replicas": 2}},
+                "does not support ingress_deployment_config",
+            ),
+            (
+                {"ingress_cls_config": {"ingress_extra_kwargs": {"key": "value"}}},
+                "does not support ingress_cls_config",
+            ),
+        ],
+    )
+    def test_direct_streaming_rejects_ingress_config(
+        self,
+        llm_config,
+        disable_placement_bundles,
+        monkeypatch,
+        builder_kwargs,
+        match,
+    ):
+        monkeypatch.setattr(
+            "ray.llm._internal.serve.core.ingress.builder."
+            "RAY_SERVE_LLM_ENABLE_DIRECT_STREAMING",
+            True,
+        )
+
+        with pytest.raises(ValueError, match=match):
+            build_openai_app(LLMServingArgs(llm_configs=[llm_config], **builder_kwargs))
+
+
+class TestDirectStreamingDP:
+    """Direct-streaming wiring tests for the data-parallel builder.
+
+    Mirrors the ``test_direct_streaming_*`` tests on ``TestBuildOpenaiApp``
+    but exercises ``build_dp_openai_app`` so that regressions in the DP
+    wiring (deployment class, default request router) are caught at CPU
+    unit-test speed instead of in GPU integration / release tests.
+    """
+
+    @pytest.fixture
+    def llm_config(self):
+        return LLMConfig(
+            model_loading_config=ModelLoadingConfig(
+                model_id="test-model", model_source="test-source"
+            )
+        )
+
+    def _enable_direct_streaming(self, monkeypatch):
+        monkeypatch.setattr(
+            "ray.llm._internal.serve.serving_patterns.data_parallel.builder."
+            "RAY_SERVE_LLM_ENABLE_DIRECT_STREAMING",
+            True,
+        )
+
+    def test_dp_builds_dpserver_ingress_with_router_attached(
+        self, llm_config, disable_placement_bundles, monkeypatch
+    ):
+        self._enable_direct_streaming(monkeypatch)
+
+        app = build_dp_openai_app({"llm_config": llm_config})
+        ingress_request_router = app._ingress_request_router
+
+        assert app._bound_deployment.name == "DPServer:test-model"
+        assert issubclass(app._bound_deployment.func_or_class, ASGIAppReplicaWrapper)
+        assert issubclass(app._bound_deployment.func_or_class, DPServer)
+        assert ingress_request_router is not None
+        assert ingress_request_router._bound_deployment.name == "LLMRouter"
+        assert ingress_request_router._bound_deployment.init_kwargs["server"] is app
+
+        request_router_config = (
+            app._bound_deployment._deployment_config.request_router_config
+        )
+        assert request_router_config.request_router_class == (
+            f"{RoundRobinRouter.__module__}.{RoundRobinRouter.__name__}"
+        )
+
+    def test_dp_user_request_router_config_wins(
+        self, llm_config, disable_placement_bundles, monkeypatch
+    ):
+        """A user-supplied ``request_router_config`` on ``LLMConfig`` must
+        survive DP direct-streaming wiring rather than being overwritten with
+        the default ``RoundRobinRouter``.
+        """
+        self._enable_direct_streaming(monkeypatch)
+        llm_config.deployment_config["request_router_config"] = RequestRouterConfig(
+            request_router_class=ConsistentHashRouter,
+        )
+
+        app = build_dp_openai_app({"llm_config": llm_config})
+        request_router_config = (
+            app._bound_deployment._deployment_config.request_router_config
+        )
+        assert request_router_config.request_router_class == (
+            f"{ConsistentHashRouter.__module__}.{ConsistentHashRouter.__name__}"
+        )
+
+
+class TestDirectStreamingPD:
+    """Direct-streaming wiring tests for the prefill/decode builder.
+
+    Covers the decode-class selection (``PDDecodeServer`` vs
+    ``DPPDDecodeServer`` based on ``decode_dp_size``), the prefill binding
+    into decode's init kwargs, and the ``LLMRouter`` ingress-request-router
+    hookup.
+    """
+
+    @pytest.fixture
+    def pd_configs(self):
+        """Prefill and decode configs with required kv_transfer_config."""
+        base_config = {
+            "model_loading_config": {
+                "model_id": "test-model",
+                "model_source": "test-source",
+            },
+            "engine_kwargs": {
+                "kv_transfer_config": {
+                    "kv_connector": "NixlConnector",
+                    "kv_role": "kv_both",
+                },
+            },
+        }
+        prefill = LLMConfig.model_validate(base_config)
+        decode = LLMConfig.model_validate(base_config)
+        return prefill, decode
+
+    def _enable_direct_streaming(self, monkeypatch):
+        monkeypatch.setattr(
+            "ray.llm._internal.serve.serving_patterns.prefill_decode.builder."
+            "RAY_SERVE_LLM_ENABLE_DIRECT_STREAMING",
+            True,
+        )
+
+    @staticmethod
+    def _set_dp_size(llm_config, size):
+        llm_config.engine_kwargs["data_parallel_size"] = size
+
+    @pytest.mark.parametrize(
+        ("prefill_dp", "decode_dp", "expected_prefill_cls", "expected_decode_cls"),
+        [
+            (1, 1, PDPrefillServer, PDDecodeServer),
+            (1, 4, PDPrefillServer, DPPDDecodeServer),
+            (4, 1, DPPDPrefillServer, PDDecodeServer),
+            (4, 4, DPPDPrefillServer, DPPDDecodeServer),
+        ],
+    )
+    def test_pd_decode_class_selection(
+        self,
+        pd_configs,
+        disable_placement_bundles,
+        monkeypatch,
+        prefill_dp,
+        decode_dp,
+        expected_prefill_cls,
+        expected_decode_cls,
+    ):
+        """Verify the DP-vs-non-DP variants are picked based on
+        ``data_parallel_size`` for both prefill and decode legs.
+        """
+        self._enable_direct_streaming(monkeypatch)
+        prefill, decode = pd_configs
+        self._set_dp_size(prefill, prefill_dp)
+        self._set_dp_size(decode, decode_dp)
+
+        app = build_pd_openai_app({"prefill_config": prefill, "decode_config": decode})
+
+        decode_deployment = app._bound_deployment
+        assert issubclass(decode_deployment.func_or_class, ASGIAppReplicaWrapper)
+        assert issubclass(decode_deployment.func_or_class, expected_decode_cls)
+
+        prefill_app = decode_deployment.init_kwargs["prefill_server"]
+        prefill_deployment = prefill_app._bound_deployment
+        assert prefill_deployment.func_or_class is expected_prefill_cls
+
+    def test_pd_ingress_request_router_is_llmrouter(
+        self, pd_configs, disable_placement_bundles, monkeypatch
+    ):
+        self._enable_direct_streaming(monkeypatch)
+        prefill, decode = pd_configs
+
+        app = build_pd_openai_app({"prefill_config": prefill, "decode_config": decode})
+        ingress_request_router = app._ingress_request_router
+
+        assert ingress_request_router is not None
+        assert ingress_request_router._bound_deployment.name == "LLMRouter"
+        assert ingress_request_router._bound_deployment.init_kwargs["server"] is app
+
+        request_router_config = (
+            app._bound_deployment._deployment_config.request_router_config
+        )
+        assert request_router_config.request_router_class == (
+            f"{RoundRobinRouter.__module__}.{RoundRobinRouter.__name__}"
+        )
+
+    def test_pd_user_request_router_config_wins(
+        self, pd_configs, disable_placement_bundles, monkeypatch
+    ):
+        """A user-supplied ``request_router_config`` on the decode
+        ``LLMConfig`` must survive PD direct-streaming wiring rather than
+        being overwritten with the default ``RoundRobinRouter``.
+        """
+        self._enable_direct_streaming(monkeypatch)
+        prefill, decode = pd_configs
+        decode.deployment_config["request_router_config"] = RequestRouterConfig(
+            request_router_class=ConsistentHashRouter,
+        )
+
+        app = build_pd_openai_app({"prefill_config": prefill, "decode_config": decode})
+        request_router_config = (
+            app._bound_deployment._deployment_config.request_router_config
+        )
+        assert request_router_config.request_router_class == (
+            f"{ConsistentHashRouter.__module__}.{ConsistentHashRouter.__name__}"
+        )
+
+
+class TestIngressScaleToZero:
+    """Tests for ingress scale-to-zero behavior when all models have min_replicas=0."""
+
+    def test_all_models_scale_to_zero(self, disable_placement_bundles):
+        """When all models have min_replicas=0, ingress should also have min_replicas=0."""
+        llm_cfg_dict_autoscaling = LLMConfig(
+            model_loading_config=ModelLoadingConfig(model_id="model_a"),
+            accelerator_type="L4",
+            deployment_config={
+                "autoscaling_config": {
+                    "min_replicas": 0,
+                    "max_replicas": 2,
+                }
+            },
+        )
+        llm_cfg_obj_autoscaling = LLMConfig(
+            model_loading_config=ModelLoadingConfig(model_id="model_b"),
+            accelerator_type="L4",
+            deployment_config={
+                "autoscaling_config": AutoscalingConfig(
+                    min_replicas=0,
+                    max_replicas=4,
+                )
+            },
+        )
+
+        app = build_openai_app(
+            LLMServingArgs(
+                llm_configs=[llm_cfg_dict_autoscaling, llm_cfg_obj_autoscaling],
+            )
+        )
+        autoscaling_config = app._bound_deployment._deployment_config.autoscaling_config
+        assert autoscaling_config.min_replicas == 0
+
+    def test_mixed_min_replicas_keeps_default(self, disable_placement_bundles):
+        """When some models have min_replicas>0, ingress should keep default min_replicas."""
+        llm_cfg_zero = LLMConfig(
+            model_loading_config=ModelLoadingConfig(model_id="model_a"),
+            accelerator_type="L4",
+            deployment_config={
+                "autoscaling_config": {
+                    "min_replicas": 0,
+                    "max_replicas": 2,
+                }
+            },
+        )
+        llm_cfg_nonzero = LLMConfig(
+            model_loading_config=ModelLoadingConfig(model_id="model_b"),
+            accelerator_type="L4",
+            deployment_config={
+                "autoscaling_config": AutoscalingConfig(
+                    min_replicas=1,
+                    max_replicas=4,
+                )
+            },
+        )
+
+        app = build_openai_app(
+            LLMServingArgs(
+                llm_configs=[llm_cfg_zero, llm_cfg_nonzero],
+            )
+        )
+        autoscaling_config = app._bound_deployment._deployment_config.autoscaling_config
+        # Default min_replicas from AutoscalingConfig is 1
+        assert autoscaling_config.min_replicas == 1
+
+    def test_no_autoscaling_config_keeps_default(self, disable_placement_bundles):
+        """When models don't have autoscaling_config, ingress should keep default."""
+        llm_cfg = LLMConfig(
+            model_loading_config=ModelLoadingConfig(model_id="model_a"),
+            accelerator_type="L4",
+        )
+
+        app = build_openai_app(
+            LLMServingArgs(llm_configs=[llm_cfg]),
+        )
+        autoscaling_config = app._bound_deployment._deployment_config.autoscaling_config
+        assert autoscaling_config.min_replicas == 1
+
+    def test_user_override_takes_precedence(self, disable_placement_bundles):
+        """User-specified ingress min_replicas should override scale-to-zero logic."""
+        llm_cfg = LLMConfig(
+            model_loading_config=ModelLoadingConfig(model_id="model_a"),
+            accelerator_type="L4",
+            deployment_config={
+                "autoscaling_config": {
+                    "min_replicas": 0,
+                    "max_replicas": 2,
+                }
+            },
+        )
+
+        app = build_openai_app(
+            LLMServingArgs(
+                llm_configs=[llm_cfg],
+                ingress_deployment_config={
+                    "autoscaling_config": {
+                        "min_replicas": 3,
+                        "max_replicas": 5,
+                    }
+                },
+            )
+        )
+        autoscaling_config = app._bound_deployment._deployment_config.autoscaling_config
+        assert autoscaling_config.min_replicas == 3
+
+
+def extract_applications_from_output(output: bytes) -> dict:
+    """
+    Extracts the 'applications' block from mixed output and returns it as a dict.
+    """
+    # 1. Decode bytes to string
+    text = output.decode("utf-8", errors="ignore")
+
+    # 2. Regex to find the 'applications:' block and its indented content
+    #    This matches 'applications:' and all following lines that are indented (YAML block)
+    match = re.search(r"(^applications:\n(?:^(?: {2,}|\t).*\n?)+)", text, re.MULTILINE)
+    if not match:
+        raise ValueError("Could not find 'applications:' block in output.")
+
+    applications_block = match.group(1)
+
+    # 3. Parse the YAML block
+    applications_dict = yaml.safe_load(applications_block)
+    return applications_dict["applications"]
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main(["-v", __file__]))

python/ray/llm/tests/serve/cpu/deployments/routers/test_config_files/llm_app_separate_model_config_files.yaml

@@ -0,0 +1,7 @@
+applications:
+  - args:
+      llm_configs:
+        - ./model_config/llm_config.yaml
+    import_path: ray.serve.llm:build_openai_app
+    name: llm-endpoint
+    route_prefix: /

python/ray/llm/tests/serve/cpu/deployments/routers/test_config_files/model_config/llm_config.yaml

@@ -0,0 +1,2 @@
+model_loading_config:
+  model_id: model1

python/ray/llm/tests/serve/cpu/deployments/routers/test_dev_ingress.py

@@ -0,0 +1,322 @@
+"""Tests for DevIngress control plane endpoints.
+
+This module tests the HTTP endpoints exposed by DevIngress:
+- POST /sleep, POST /wakeup, GET /is_sleeping
+- POST /pause, POST /resume, GET /is_paused
+- POST /reset_prefix_cache
+
+These tests verify:
+1. Endpoints are correctly registered and accessible
+2. Broadcast API correctly broadcasts to replicas
+3. Sleep/wakeup and pause/resume isolation between different models
+"""
+
+import sys
+
+import httpx
+import pytest
+
+import ray
+from ray import serve
+from ray.llm._internal.serve.core.configs.openai_api_models import to_model_metadata
+from ray.llm._internal.serve.core.ingress.dev_ingress import DEV_ENDPOINTS, DevIngress
+from ray.llm._internal.serve.core.ingress.ingress import make_fastapi_ingress
+from ray.llm._internal.serve.core.server.llm_server import LLMServer
+from ray.llm.tests.serve.mocks.mock_vllm_engine import MockVLLMEngine
+from ray.serve.llm import LLMConfig, ModelLoadingConfig
+
+
+@pytest.fixture(scope="module")
+def ray_instance():
+    """Initialize Ray for the module."""
+    if not ray.is_initialized():
+        ray.init()
+    yield
+    serve.shutdown()
+    ray.shutdown()
+
+
+@pytest.fixture
+def single_model_dev_ingress(ray_instance, disable_placement_bundles):
+    """Start a Serve app with one model and DevIngress endpoints."""
+    model_id = "test-model-1"
+    llm_config = LLMConfig(
+        model_loading_config=ModelLoadingConfig(model_id=model_id),
+        runtime_env={},
+        log_engine_metrics=False,
+    )
+
+    # Create LLMServer deployment with mock engine
+    llm_deployment = serve.deployment(LLMServer).bind(
+        llm_config, engine_cls=MockVLLMEngine
+    )
+
+    # Create DevIngress with the dev endpoints
+    ingress_cls = make_fastapi_ingress(DevIngress, endpoint_map=DEV_ENDPOINTS)
+    ingress_options = DevIngress.get_deployment_options([llm_config])
+    ingress_app = serve.deployment(ingress_cls, **ingress_options).bind(
+        llm_deployments={model_id: llm_deployment},
+        model_cards={model_id: to_model_metadata(model_id, llm_config)},
+    )
+
+    serve.run(ingress_app, name="single-model-app")
+    yield model_id
+    serve.delete("single-model-app", _blocking=True)
+
+
+@pytest.fixture
+def two_model_dev_ingress(ray_instance, disable_placement_bundles):
+    """Start a Serve app with TWO model deployments to test isolation."""
+    model_id_1 = "test-model-1"
+    model_id_2 = "test-model-2"
+
+    llm_config_1 = LLMConfig(
+        model_loading_config=ModelLoadingConfig(model_id=model_id_1),
+        runtime_env={},
+        log_engine_metrics=False,
+    )
+    llm_config_2 = LLMConfig(
+        model_loading_config=ModelLoadingConfig(model_id=model_id_2),
+        runtime_env={},
+        log_engine_metrics=False,
+    )
+
+    # Create LLMServer deployments with mock engine
+    llm_deployment_1 = serve.deployment(LLMServer).bind(
+        llm_config_1, engine_cls=MockVLLMEngine
+    )
+    llm_deployment_2 = serve.deployment(LLMServer).bind(
+        llm_config_2, engine_cls=MockVLLMEngine
+    )
+
+    # Create DevIngress with the dev endpoints
+    ingress_cls = make_fastapi_ingress(DevIngress, endpoint_map=DEV_ENDPOINTS)
+    ingress_options = DevIngress.get_deployment_options([llm_config_1, llm_config_2])
+    ingress_app = serve.deployment(ingress_cls, **ingress_options).bind(
+        llm_deployments={
+            model_id_1: llm_deployment_1,
+            model_id_2: llm_deployment_2,
+        },
+        model_cards={
+            model_id_1: to_model_metadata(model_id_1, llm_config_1),
+            model_id_2: to_model_metadata(model_id_2, llm_config_2),
+        },
+    )
+
+    serve.run(ingress_app, name="two-model-app")
+    yield model_id_1, model_id_2
+    serve.delete("two-model-app", _blocking=True)
+
+
+class TestDevIngressEndpoints:
+    """Test DevIngress endpoints."""
+
+    @pytest.mark.asyncio
+    async def test_reset_prefix_cache_endpoint(self, single_model_dev_ingress):
+        """Test POST /reset_prefix_cache endpoint."""
+        model_id = single_model_dev_ingress
+
+        async with httpx.AsyncClient(timeout=60.0) as client:
+            response = await client.post(
+                "http://localhost:8000/reset_prefix_cache",
+                json={"model": model_id},
+            )
+            assert response.status_code == 200
+
+    @pytest.mark.asyncio
+    async def test_sleep_wakeup_cycle(self, single_model_dev_ingress):
+        """Test full sleep -> is_sleeping -> wakeup -> is_sleeping cycle."""
+        model_id = single_model_dev_ingress
+
+        async with httpx.AsyncClient(timeout=60.0) as client:
+            # Initial state - should not be sleeping
+            response = await client.get(
+                f"http://localhost:8000/is_sleeping?model={model_id}",
+            )
+            assert response.status_code == 200
+            assert response.json().get("is_sleeping") is False
+
+            # Sleep the engine
+            response = await client.post(
+                "http://localhost:8000/sleep",
+                json={"model": model_id, "options": {"level": 1}},
+            )
+            assert response.status_code == 200
+
+            # Check is_sleeping - should be True
+            response = await client.get(
+                f"http://localhost:8000/is_sleeping?model={model_id}",
+            )
+            assert response.status_code == 200
+            assert response.json().get("is_sleeping") is True
+
+            # Wake up the engine
+            response = await client.post(
+                "http://localhost:8000/wakeup",
+                json={"model": model_id, "options": {}},
+            )
+            assert response.status_code == 200
+
+            # Check is_sleeping - should be False again
+            response = await client.get(
+                f"http://localhost:8000/is_sleeping?model={model_id}",
+            )
+            assert response.status_code == 200
+            assert response.json().get("is_sleeping") is False
+
+    @pytest.mark.asyncio
+    async def test_pause_resume_cycle(self, single_model_dev_ingress):
+        """Test full pause -> is_paused -> resume -> is_paused cycle."""
+        model_id = single_model_dev_ingress
+
+        async with httpx.AsyncClient(timeout=60.0) as client:
+            # Initial state - should not be paused
+            response = await client.get(
+                f"http://localhost:8000/is_paused?model={model_id}",
+            )
+            assert response.status_code == 200
+            assert response.json().get("is_paused") is False
+
+            # Pause the engine
+            response = await client.post(
+                "http://localhost:8000/pause",
+                json={"model": model_id, "options": {"clear_cache": True}},
+            )
+            assert response.status_code == 200
+
+            # Check is_paused - should be True
+            response = await client.get(
+                f"http://localhost:8000/is_paused?model={model_id}",
+            )
+            assert response.status_code == 200
+            assert response.json().get("is_paused") is True
+
+            # Resume the engine
+            response = await client.post(
+                "http://localhost:8000/resume",
+                json={"model": model_id, "options": {}},
+            )
+            assert response.status_code == 200
+
+            # Check is_paused - should be False again
+            response = await client.get(
+                f"http://localhost:8000/is_paused?model={model_id}",
+            )
+            assert response.status_code == 200
+            assert response.json().get("is_paused") is False
+
+
+class TestDevIngressModelIsolation:
+    """Test that control plane operations are isolated per model."""
+
+    @pytest.mark.asyncio
+    async def test_sleep_wakeup_isolation(self, two_model_dev_ingress):
+        """Test that sleeping model_1 does NOT affect model_2."""
+        model_1, model_2 = two_model_dev_ingress
+
+        async with httpx.AsyncClient(timeout=60.0) as client:
+            # Both models should start awake
+            response = await client.get(
+                f"http://localhost:8000/is_sleeping?model={model_1}",
+            )
+            assert response.json().get("is_sleeping") is False
+
+            response = await client.get(
+                f"http://localhost:8000/is_sleeping?model={model_2}",
+            )
+            assert response.json().get("is_sleeping") is False
+
+            # Sleep model_1 only
+            response = await client.post(
+                "http://localhost:8000/sleep",
+                json={"model": model_1, "options": {"level": 1}},
+            )
+            assert response.status_code == 200
+
+            # model_1 should be sleeping
+            response = await client.get(
+                f"http://localhost:8000/is_sleeping?model={model_1}",
+            )
+            assert response.json().get("is_sleeping") is True
+
+            # model_2 should NOT be sleeping
+            response = await client.get(
+                f"http://localhost:8000/is_sleeping?model={model_2}",
+            )
+            assert response.json().get("is_sleeping") is False
+
+            # Wake up model_1
+            response = await client.post(
+                "http://localhost:8000/wakeup",
+                json={"model": model_1, "options": {}},
+            )
+            assert response.status_code == 200
+
+            # Both should now be awake
+            response = await client.get(
+                f"http://localhost:8000/is_sleeping?model={model_1}",
+            )
+            assert response.json().get("is_sleeping") is False
+
+            response = await client.get(
+                f"http://localhost:8000/is_sleeping?model={model_2}",
+            )
+            assert response.json().get("is_sleeping") is False
+
+    @pytest.mark.asyncio
+    async def test_pause_resume_isolation(self, two_model_dev_ingress):
+        """Test that pausing model_1 does NOT affect model_2."""
+        model_1, model_2 = two_model_dev_ingress
+
+        async with httpx.AsyncClient(timeout=60.0) as client:
+            # Both models should start unpaused
+            response = await client.get(
+                f"http://localhost:8000/is_paused?model={model_1}",
+            )
+            assert response.json().get("is_paused") is False
+
+            response = await client.get(
+                f"http://localhost:8000/is_paused?model={model_2}",
+            )
+            assert response.json().get("is_paused") is False
+
+            # Pause model_1 only
+            response = await client.post(
+                "http://localhost:8000/pause",
+                json={"model": model_1, "options": {"clear_cache": True}},
+            )
+            assert response.status_code == 200
+
+            # model_1 should be paused
+            response = await client.get(
+                f"http://localhost:8000/is_paused?model={model_1}",
+            )
+            assert response.json().get("is_paused") is True
+
+            # model_2 should NOT be paused
+            response = await client.get(
+                f"http://localhost:8000/is_paused?model={model_2}",
+            )
+            assert response.json().get("is_paused") is False
+
+            # Resume model_1
+            response = await client.post(
+                "http://localhost:8000/resume",
+                json={"model": model_1, "options": {}},
+            )
+            assert response.status_code == 200
+
+            # Both should now be unpaused
+            response = await client.get(
+                f"http://localhost:8000/is_paused?model={model_1}",
+            )
+            assert response.json().get("is_paused") is False
+
+            response = await client.get(
+                f"http://localhost:8000/is_paused?model={model_2}",
+            )
+            assert response.json().get("is_paused") is False
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main(["-v", __file__]))

python/ray/llm/tests/serve/cpu/deployments/routers/test_make_fastapi_ingress.py

@@ -0,0 +1,163 @@
+"""Tests for make_fastapi_ingress function."""
+
+import inspect
+import sys
+
+import pytest
+from fastapi import FastAPI, Request
+from fastapi.routing import APIRoute
+
+from ray.llm._internal.serve.core.ingress.ingress import (
+    DEFAULT_ENDPOINTS,
+    OpenAiIngress,
+    make_fastapi_ingress,
+)
+
+
+class TestMakeFastapiIngress:
+    """Test suite for make_fastapi_ingress."""
+
+    def test_subclass_inherits_endpoints(self):
+        """Test that subclassing OpenAiIngress works with make_fastapi_ingress."""
+
+        class MyCustomIngress(OpenAiIngress):
+            """Custom ingress that inherits all OpenAI endpoints."""
+
+            pass
+
+        app = FastAPI()
+        # Create the ingress class - should not raise
+        ingress_cls = make_fastapi_ingress(MyCustomIngress, app=app)
+
+        # Verify the ingress class was created successfully
+        assert ingress_cls is not None
+
+        # Verify routes are registered (inherited from OpenAiIngress)
+        route_paths = [
+            route.path for route in app.routes if isinstance(route, APIRoute)
+        ]
+        assert "/v1/models" in route_paths
+        assert "/v1/completions" in route_paths
+
+    def test_subclass_with_custom_method(self):
+        """Test that custom methods added by subclass are also properly handled."""
+
+        class MyCustomIngress(OpenAiIngress):
+            """Custom ingress with an additional endpoint."""
+
+            async def custom_endpoint(self, request: Request):
+                """A custom endpoint added by the subclass."""
+                return {"status": "ok"}
+
+        custom_endpoints = {
+            "custom_endpoint": lambda app: app.post("/custom"),
+            **DEFAULT_ENDPOINTS,
+        }
+
+        app = FastAPI()
+        ingress_cls = make_fastapi_ingress(
+            MyCustomIngress, endpoint_map=custom_endpoints, app=app
+        )
+
+        # Verify the class was created and the custom route is registered
+        assert ingress_cls is not None
+        route_paths = [
+            route.path for route in app.routes if isinstance(route, APIRoute)
+        ]
+        assert "/custom" in route_paths
+
+    def test_routes_registered_correctly(self):
+        """Test that routes are registered with the FastAPI app."""
+
+        class MyCustomIngress(OpenAiIngress):
+            pass
+
+        app = FastAPI()
+        make_fastapi_ingress(MyCustomIngress, app=app)
+
+        # Get all registered routes
+        route_paths = [
+            route.path for route in app.routes if isinstance(route, APIRoute)
+        ]
+
+        # Check that default endpoints are registered
+        assert "/v1/models" in route_paths
+        assert "/v1/completions" in route_paths
+        assert "/v1/chat/completions" in route_paths
+
+    def test_custom_endpoint_map_overrides_defaults(self):
+        """Test that custom endpoint_map can override default endpoints."""
+
+        class MyCustomIngress(OpenAiIngress):
+            async def models(self):
+                """Override the models endpoint."""
+                return {"custom": True}
+
+        # Only register models endpoint with a custom path
+        custom_endpoints = {
+            "models": lambda app: app.get("/custom/models"),
+        }
+
+        app = FastAPI()
+        make_fastapi_ingress(MyCustomIngress, endpoint_map=custom_endpoints, app=app)
+
+        route_paths = [
+            route.path for route in app.routes if isinstance(route, APIRoute)
+        ]
+
+        # Should have custom path, not default
+        assert "/custom/models" in route_paths
+        assert "/v1/models" not in route_paths
+
+    def test_deeply_nested_inheritance(self):
+        """Test that deeply nested inheritance works correctly."""
+
+        class IntermediateIngress(OpenAiIngress):
+            """Intermediate class in inheritance chain."""
+
+            async def intermediate_method(self, request: Request):
+                return {"level": "intermediate"}
+
+        class FinalIngress(IntermediateIngress):
+            """Final class in inheritance chain."""
+
+            async def final_method(self, request: Request):
+                return {"level": "final"}
+
+        custom_endpoints = {
+            "intermediate_method": lambda app: app.post("/intermediate"),
+            "final_method": lambda app: app.post("/final"),
+            **DEFAULT_ENDPOINTS,
+        }
+
+        app = FastAPI()
+        make_fastapi_ingress(FinalIngress, endpoint_map=custom_endpoints, app=app)
+
+        # Verify all routes are registered
+        route_paths = [
+            route.path for route in app.routes if isinstance(route, APIRoute)
+        ]
+        assert "/intermediate" in route_paths
+        assert "/final" in route_paths
+        assert "/v1/completions" in route_paths
+
+    def test_method_signature_preserved(self):
+        """Test that method signatures are preserved after decoration."""
+
+        class MyCustomIngress(OpenAiIngress):
+            pass
+
+        ingress_cls = make_fastapi_ingress(MyCustomIngress)
+
+        # Get the completions method and check its signature
+        completions_method = ingress_cls.completions
+        sig = inspect.signature(completions_method)
+        param_names = list(sig.parameters.keys())
+
+        # Should have 'self' and 'body' parameters
+        assert "self" in param_names
+        assert "body" in param_names
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main(["-v", __file__]))

python/ray/llm/tests/serve/cpu/deployments/routers/test_router.py

@@ -0,0 +1,563 @@
+import sys
+from contextlib import asynccontextmanager
+from types import SimpleNamespace
+from typing import Optional
+from unittest.mock import AsyncMock, MagicMock, patch
+
+import openai
+import pytest
+from fastapi import HTTPException
+from starlette.datastructures import Headers
+
+from ray import serve
+from ray.llm._internal.serve.core.configs.llm_config import (
+    LLMConfig,
+    ModelLoadingConfig,
+)
+from ray.llm._internal.serve.core.configs.openai_api_models import to_model_metadata
+from ray.llm._internal.serve.core.ingress import router as router_module
+from ray.llm._internal.serve.core.ingress.ingress import (
+    OpenAiIngress,
+    make_fastapi_ingress,
+)
+from ray.llm._internal.serve.core.ingress.router import (
+    LLMRouter,
+    _parse_routing_payload,
+)
+from ray.llm._internal.serve.core.server.llm_server import LLMServer
+from ray.llm.tests.serve.mocks.mock_vllm_engine import MockVLLMEngine
+from ray.serve._private.common import DeploymentID
+from ray.serve.exceptions import DeploymentUnavailableError
+
+
+class _DirectRouterReplicaId:
+    def __init__(self, unique_id: str, full_id: Optional[str] = None):
+        self.unique_id = unique_id
+        self._full_id = full_id or unique_id
+
+    def to_full_id_str(self) -> str:
+        return self._full_id
+
+
+class _FakeRequest:
+    def __init__(self, body: bytes, headers: Optional[dict] = None):
+        self._body = body
+        self.headers = Headers(headers or {})
+
+    async def body(self) -> bytes:
+        return self._body
+
+
+class _DirectRouterReplica:
+    """RunningReplica stand-in for ``LLMRouter._pick_replica`` tests."""
+
+    def __init__(
+        self,
+        unique_id: str,
+        full_id: Optional[str] = None,
+        endpoint: Optional[tuple] = ("127.0.0.1", 8000),
+    ):
+        self.replica_id = _DirectRouterReplicaId(unique_id, full_id)
+        self.backend_http_endpoint = endpoint
+
+
+def _new_direct_router(handle=None):
+    router = LLMRouter.__new__(LLMRouter)
+    router._handle = handle or MagicMock()
+    # Routing tests don't exercise tokenization; that lives in test_tokenizer.py.
+    router._tokenizer = None
+    return router
+
+
+def _selection_for(replica):
+    """Build a ``ReplicaSelection``-shaped mock that ``_pick_replica`` reads."""
+    return MagicMock(replica_id=replica.replica_id.unique_id, _replica=replica)
+
+
+def _choose_replica_returning(*replicas):
+    """Patch ``handle.choose_replica`` to yield the given replicas in order.
+
+    Each call to ``choose_replica`` consumes one replica from the sequence and
+    yields its ``_DirectRouterReplica`` wrapped as a selection.
+    """
+    selections = iter(_selection_for(r) for r in replicas)
+
+    @asynccontextmanager
+    async def fake_choose_replica(*args, **kwargs):
+        yield next(selections)
+
+    return fake_choose_replica
+
+
+@pytest.fixture(name="llm_config")
+def create_llm_config(stream_batching_interval_ms: Optional[int] = None):
+
+    if stream_batching_interval_ms is not None:
+        return LLMConfig(
+            model_loading_config=ModelLoadingConfig(
+                model_id="llm_model_id",
+            ),
+            experimental_configs={
+                "stream_batching_interval_ms": stream_batching_interval_ms,
+            },
+        )
+    else:
+        return LLMConfig(
+            model_loading_config=ModelLoadingConfig(
+                model_id="llm_model_id",
+            ),
+        )
+
+
+@pytest.fixture(name="client")
+def create_oai_client(llm_config: LLMConfig):
+    ServerDeployment = serve.deployment(LLMServer)
+
+    ingress_options = OpenAiIngress.get_deployment_options(llm_configs=[llm_config])
+    ingress_cls = make_fastapi_ingress(OpenAiIngress)
+    RouterDeployment = serve.deployment(ingress_cls, **ingress_options)
+    server = ServerDeployment.bind(llm_config, engine_cls=MockVLLMEngine)
+    router = RouterDeployment.bind(
+        llm_deployments={llm_config.model_id: server},
+        model_cards={
+            llm_config.model_id: to_model_metadata(llm_config.model_id, llm_config)
+        },
+    )
+    serve.run(router)
+
+    client = openai.Client(base_url="http://localhost:8000/v1", api_key="foo")
+    yield client
+
+    serve.shutdown()
+
+
+class TestDirectStreamingLLMRouter:
+    @pytest.mark.asyncio
+    async def test_route_parses_body_into_routing_payload(self):
+        """A parseable body becomes a routing payload passed positionally."""
+        router = _new_direct_router()
+        router._pick_replica = AsyncMock(
+            return_value=("127.0.0.1", 9001, "DeploymentName#replica")
+        )
+
+        body = b'{"model":"x","messages":[{"role":"user","content":"hi"}]}'
+        request = _FakeRequest(body)
+
+        result = await router.route(request)
+
+        assert result == {
+            "host": "127.0.0.1",
+            "port": 9001,
+            "replica_id": "DeploymentName#replica",
+        }
+        _, kwargs = router._pick_replica.call_args
+        assert kwargs["handle"] is router._handle
+        payload = kwargs["routing_payload"]
+        assert isinstance(payload, SimpleNamespace)
+        assert payload.messages == [{"role": "user", "content": "hi"}]
+        # The whole body is exposed, so a router can read any field.
+        assert payload.model == "x"
+        assert not hasattr(payload, "prompt")
+        # A parseable body must not trip the "no routing key" warning.
+        assert router._warned_no_routing_key is False
+
+    @pytest.mark.asyncio
+    async def test_route_truncated_body_yields_no_payload_and_warns_once(self):
+        """A truncated body derives no key. ``route`` forwards ``None`` and
+        warns once per replica."""
+        router = _new_direct_router()
+        router._pick_replica = AsyncMock(
+            return_value=("127.0.0.1", 9001, "DeploymentName#replica")
+        )
+
+        # Truncated prefix is not valid JSON so json.loads fails.
+        body = b'{"model":"x","prompt":"' + (b"x" * 1024)
+        request = _FakeRequest(body, headers={"x-body-truncated": "1058/90000"})
+
+        with patch.object(router_module.logger, "warning") as mock_warning:
+            await router.route(request)
+            await router.route(request)
+
+        # routing_payload is None on both calls. Warning fires once.
+        for call in router._pick_replica.call_args_list:
+            assert call.kwargs["routing_payload"] is None
+        assert mock_warning.call_count == 1
+        assert router._warned_no_routing_key is True
+
+    @pytest.mark.asyncio
+    async def test_route_returns_503_on_pick_failure(self):
+        router = _new_direct_router()
+        router._pick_replica = AsyncMock(side_effect=RuntimeError("no replicas"))
+
+        with pytest.raises(HTTPException) as exc_info:
+            await router.route(_FakeRequest(b"{}"))
+        assert exc_info.value.status_code == 503
+        assert "no replicas" in exc_info.value.detail
+
+    @pytest.mark.asyncio
+    async def test_route_returns_400_on_bad_routing_request(self):
+        router = _new_direct_router()
+        router._pick_replica = AsyncMock(side_effect=ValueError("empty prompt"))
+
+        with pytest.raises(HTTPException) as exc_info:
+            await router.route(_FakeRequest(b"{}"))
+        assert exc_info.value.status_code == 400
+        assert "empty prompt" in exc_info.value.detail
+
+    @pytest.mark.asyncio
+    async def test_route_returns_503_on_deployment_unavailable(self):
+        err = DeploymentUnavailableError(DeploymentID(name="LLMServer:test"))
+        router = _new_direct_router()
+        router._pick_replica = AsyncMock(side_effect=err)
+
+        with pytest.raises(HTTPException) as exc_info:
+            await router.route(_FakeRequest(b"{}"))
+        assert exc_info.value.status_code == 503
+        assert "LLMServer:test" in exc_info.value.detail
+
+    @pytest.mark.asyncio
+    async def test_pick_replica_returns_backend_endpoint_from_handle(self):
+        """``_pick_replica`` reads the endpoint off the selection's replica."""
+        replica = _DirectRouterReplica(
+            "r1",
+            full_id="DeploymentName#r1",
+            endpoint=("10.0.0.1", 8123),
+        )
+        handle = MagicMock()
+        handle.choose_replica = _choose_replica_returning(replica)
+        router = _new_direct_router(handle)
+
+        host, port, replica_id = await router._pick_replica(handle=handle)
+
+        assert (host, port, replica_id) == ("10.0.0.1", 8123, "DeploymentName#r1")
+
+    @pytest.mark.asyncio
+    async def test_pick_replica_forwards_payload_positionally(self):
+        """A routing payload reaches ``choose_replica`` as the first positional
+        arg, alongside the ``_reserve=False`` fast-path flag."""
+        replica = _DirectRouterReplica("r1", full_id="d#r1")
+
+        captured = {}
+
+        @asynccontextmanager
+        async def fake_choose_replica(*args, **kwargs):
+            captured["args"] = args
+            captured["kwargs"] = kwargs
+            yield _selection_for(replica)
+
+        handle = MagicMock()
+        handle.choose_replica = fake_choose_replica
+        router = _new_direct_router(handle)
+
+        payload = SimpleNamespace(messages=[{"role": "user", "content": "hi"}])
+        await router._pick_replica(handle=handle, routing_payload=payload)
+
+        assert captured["args"] == (payload,)
+        assert captured["kwargs"] == {"_reserve": False}
+
+    @pytest.mark.asyncio
+    async def test_pick_replica_omits_positional_arg_when_no_payload(self):
+        """With no routing payload, nothing is forwarded positionally. The
+        configured router then sees empty args and load-balances."""
+        replica = _DirectRouterReplica("r1", full_id="d#r1")
+
+        captured = {}
+
+        @asynccontextmanager
+        async def fake_choose_replica(*args, **kwargs):
+            captured["args"] = args
+            captured["kwargs"] = kwargs
+            yield _selection_for(replica)
+
+        handle = MagicMock()
+        handle.choose_replica = fake_choose_replica
+        router = _new_direct_router(handle)
+
+        await router._pick_replica(handle=handle, routing_payload=None)
+
+        assert captured["args"] == ()
+        assert captured["kwargs"] == {"_reserve": False}
+
+    @pytest.mark.asyncio
+    async def test_pick_replica_raises_when_endpoint_missing(self):
+        """If the picked replica has no backend HTTP endpoint, surface a 503
+        via ``RuntimeError`` (same error contract as before)."""
+        replica = _DirectRouterReplica("r1", endpoint=None)
+        handle = MagicMock()
+        handle.choose_replica = _choose_replica_returning(replica)
+        router = _new_direct_router(handle)
+
+        with pytest.raises(RuntimeError, match="no backend HTTP endpoint"):
+            await router._pick_replica(handle=handle)
+
+
+class TestRoutingPayload:
+    """Unit coverage for wrapping a body as a routing namespace."""
+
+    def test_parses_chat_messages(self):
+        body = b'{"model":"x","messages":[{"role":"user","content":"hi"}]}'
+        payload = _parse_routing_payload(body)
+        assert isinstance(payload, SimpleNamespace)
+        assert payload.messages == [{"role": "user", "content": "hi"}]
+        # A chat body exposes no `prompt`, so `_extract_text_from_request`
+        # resolves it as a chat request. Other fields are still exposed.
+        assert not hasattr(payload, "prompt")
+        assert payload.model == "x"
+
+    def test_parses_completion_prompt(self):
+        payload = _parse_routing_payload(b'{"model":"x","prompt":"hello"}')
+        assert isinstance(payload, SimpleNamespace)
+        assert payload.prompt == "hello"
+        assert not hasattr(payload, "messages")
+
+    @pytest.mark.parametrize(
+        "body",
+        [
+            b"",  # empty
+            b'{"model":"x","prompt":"' + (b"x" * 64),  # truncated, invalid JSON
+            b"not json",  # unparseable
+            b"[1, 2, 3]",  # valid JSON but not an object
+            b'{"model":"x","max_tokens":8}',  # object without messages or prompt
+            b'{"messages":[]}',  # empty messages carry no routing signal
+            b'{"prompt":""}',  # empty prompt carries no routing signal
+            b'{"model":"x","input":"hello"}',  # other request type, no routing key
+        ],
+    )
+    def test_returns_none_when_no_key_derivable(self, body):
+        assert _parse_routing_payload(body) is None
+
+    @pytest.mark.asyncio
+    async def test_payload_satisfies_prefix_router_contract(self):
+        """The normalized payload is read by the real
+        ``PrefixCacheAffinityRouter._extract_text_from_request``, the consumer
+        that regressed in #64326.
+
+        Async so a running event loop exists for the ``PendingRequest`` default
+        ``asyncio.Future``.
+        """
+        from ray.llm._internal.serve.routing_policies.prefix_aware.prefix_aware_router import (  # noqa: E501
+            PrefixCacheAffinityRouter,
+        )
+        from ray.serve._private.request_router.common import PendingRequest
+
+        # __new__ avoids the tree-actor setup in __init__. The method under test
+        # only uses self for the pure `_normalize_prompt_to_string` helper.
+        router = PrefixCacheAffinityRouter.__new__(PrefixCacheAffinityRouter)
+
+        chat = _parse_routing_payload(
+            b'{"messages":[{"role":"user","content":"hello world"}]}'
+        )
+        pr = PendingRequest(args=[chat], kwargs={}, metadata=MagicMock())
+        assert router._extract_text_from_request(pr) == "hello world"
+
+        completion = _parse_routing_payload(b'{"prompt":"hello world"}')
+        pr = PendingRequest(args=[completion], kwargs={}, metadata=MagicMock())
+        assert router._extract_text_from_request(pr) == "hello world"
+
+
+class TestOpenAiIngress:
+    @pytest.mark.parametrize("stream_batching_interval_ms", [None, 0, 10000])
+    @pytest.mark.parametrize("stream", [True, False])
+    @pytest.mark.asyncio
+    async def test_chat(self, stream_batching_interval_ms, client, stream):
+        """Tests chat streaming with different stream_batching_interval_ms values.
+
+        0ms super fast batching (no batching)
+        10000ms basically should be equivalent to non-streaming
+        None is default, which is some fixed non-zero value.
+        """
+
+        # Generate 1000 chunks
+        n_tokens = 1000
+
+        response = client.chat.completions.create(
+            model="llm_model_id",
+            messages=[dict(role="user", content="Hello")],
+            stream=stream,
+            max_tokens=n_tokens,
+        )
+
+        if stream:
+            text = ""
+            role = None
+            for chunk in response:
+                if chunk.choices[0].delta.role is not None and role is None:
+                    role = chunk.choices[0].delta.role
+                if chunk.choices[0].delta.content:
+                    text += chunk.choices[0].delta.content
+        else:
+            text = response.choices[0].message.content
+            role = response.choices[0].message.role
+
+        assert role == "assistant"
+        assert text.strip() == " ".join([f"test_{i}" for i in range(n_tokens)])
+
+    @pytest.mark.parametrize("stream_batching_interval_ms", [None, 0, 10000])
+    @pytest.mark.parametrize("stream", [True, False])
+    @pytest.mark.asyncio
+    async def test_completion(self, stream_batching_interval_ms, client, stream):
+        """Tests text completions streaming with different stream_batching_interval_ms values."""
+
+        # Generate tokens
+        n_tokens = 1000
+
+        response = client.completions.create(
+            model="llm_model_id",
+            prompt="Hello",
+            stream=stream,
+            max_tokens=n_tokens,
+        )
+
+        if stream:
+            text = ""
+            for chunk in response:
+                text += chunk.choices[0].text
+        else:
+            text = response.choices[0].text
+
+        # The mock engine produces "test_0 test_1 test_2 ..." pattern
+        expected_text = " ".join([f"test_{i}" for i in range(n_tokens)])
+        assert text.strip() == expected_text
+
+    @pytest.mark.parametrize("stream", [True, False])
+    @pytest.mark.asyncio
+    async def test_tool_call(self, client, stream):
+        response = client.chat.completions.create(
+            model="llm_model_id",
+            messages=[
+                {
+                    "role": "user",
+                    "content": "Can you tell me what the temperate will be in Dallas, in fahrenheit?",
+                },
+                {
+                    "content": None,
+                    "role": "assistant",
+                    "tool_calls": [
+                        {
+                            "id": "RBS92VTjJ",
+                            "function": {
+                                "arguments": '{"city": "Dallas", "state": "TX", "unit": "fahrenheit"}',
+                                "name": "get_current_weather",
+                            },
+                            "type": "function",
+                        }
+                    ],
+                },
+                {
+                    "role": "tool",
+                    "content": "The weather in Dallas, TX is 85 degrees fahrenheit. It is partly cloudly, with highs in the 90's.",
+                    "tool_call_id": "n3OMUpydP",
+                },
+            ],
+            stream=stream,
+            max_tokens=200,
+        )
+
+        if stream:
+            text = ""
+            role = None
+            for chunk in response:
+                if chunk.choices[0].delta.role is not None and role is None:
+                    role = chunk.choices[0].delta.role
+                if chunk.choices[0].delta.content:
+                    text += chunk.choices[0].delta.content
+        else:
+            text = response.choices[0].message.content
+            role = response.choices[0].message.role
+
+        assert text
+
+    @pytest.mark.asyncio
+    async def test_check_health(self, llm_config: LLMConfig):
+        """Test health check functionality."""
+
+        server = MagicMock()
+        server.check_health = MagicMock()
+        server.check_health.remote = AsyncMock()
+
+        router = OpenAiIngress(
+            llm_deployments={llm_config.model_id: server},
+            model_cards={
+                llm_config.model_id: to_model_metadata(llm_config.model_id, llm_config)
+            },
+        )
+
+        await router.check_health()
+
+    @pytest.mark.asyncio
+    async def test_raw_request_info_passed_to_deployment_handle(
+        self, llm_config: LLMConfig
+    ):
+        """Test that raw_request_info is passed to the deployment handle."""
+        from ray.llm._internal.serve.core.configs.openai_api_models import (
+            ChatCompletionRequest,
+            ChatCompletionResponse,
+        )
+        from ray.llm._internal.serve.core.protocol import RawRequestInfo
+
+        # Track if raw_request_info was received
+        captured_raw_request_infos = []
+
+        # Create a mock deployment handle that captures raw_request_info
+        async def mock_chat_generator(request, raw_request_info):
+            captured_raw_request_infos.append(raw_request_info)
+            # Return a valid response
+            yield ChatCompletionResponse(
+                id="test_id",
+                choices=[
+                    {
+                        "index": 0,
+                        "message": {"role": "assistant", "content": "Hello!"},
+                        "finish_reason": "stop",
+                    }
+                ],
+                model="llm_model_id",
+                object="chat.completion",
+                usage={
+                    "prompt_tokens": 1,
+                    "completion_tokens": 1,
+                    "total_tokens": 2,
+                },
+            )
+
+        mock_handle = MagicMock()
+        mock_handle.chat = MagicMock()
+        mock_handle.chat.remote = mock_chat_generator
+        # Make options() return the same mock so chat.remote is preserved
+        mock_handle.options.return_value = mock_handle
+
+        # Create router with mock handle
+        router = OpenAiIngress(
+            llm_deployments={llm_config.model_id: mock_handle},
+            model_cards={
+                llm_config.model_id: to_model_metadata(llm_config.model_id, llm_config)
+            },
+        )
+
+        # Create a mock FastAPI request
+        from starlette.datastructures import Headers
+
+        mock_request = MagicMock()
+        mock_headers = {
+            "content-type": "application/json",
+            "x-ray-serve-llm-test-header": "router-raw-request-info",
+        }
+        mock_request.headers = Headers(mock_headers)
+
+        # Make a request through the router
+        request_body = ChatCompletionRequest(
+            model="llm_model_id",
+            messages=[{"role": "user", "content": "Hello"}],
+            stream=False,
+        )
+
+        await router.chat(request_body, mock_request)
+
+        # Verify that raw_request_info was passed to the deployment handle
+        assert len(captured_raw_request_infos) == 1
+        assert isinstance(captured_raw_request_infos[0], RawRequestInfo)
+        assert captured_raw_request_infos[0].headers == mock_headers
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main(["-v", __file__]))

python/ray/llm/tests/serve/cpu/deployments/routers/test_router_direct_streaming.py

@@ -0,0 +1,65 @@
+import sys
+
+import pytest
+
+from ray.llm._internal.serve.core.configs.llm_config import (
+    LLMConfig,
+    ModelLoadingConfig,
+)
+from ray.llm._internal.serve.core.ingress.builder import (
+    LLMServingArgs,
+    build_openai_app,
+)
+from ray.llm.tests.serve.cpu.deployments.utils.direct_streaming_utils import (
+    consistent_hash_deployment_config,
+    requires_direct_streaming,
+    run_app_through_haproxy,
+    session_chat_response,
+)
+
+
+@requires_direct_streaming
+class TestDirectStreamingConsistentHashRouting:
+    """Session affinity over the full direct-streaming path.
+
+    A request flows through HAProxy and the LLMRouter ``/internal/route``
+    decision (ConsistentHashRouter) to a backend replica. The session id
+    reaches the chosen replica, and one session pins to one replica.
+    """
+
+    @pytest.fixture(name="llm_config")
+    def _llm_config(self):
+        return LLMConfig(model_loading_config=ModelLoadingConfig(model_id="test-model"))
+
+    @pytest.fixture(name="base_url")
+    def run_direct_streaming_app(
+        self,
+        llm_config_with_mock_engine,
+        shutdown_ray_and_serve,
+        disable_placement_bundles,
+    ):
+        llm_config = llm_config_with_mock_engine
+        llm_config.deployment_config = consistent_hash_deployment_config()
+        yield run_app_through_haproxy(
+            build_openai_app(LLMServingArgs(llm_configs=[llm_config]))
+        )
+
+    def test_session_affinity(self, base_url):
+        replicas = {
+            session_chat_response(base_url, "test-session-id").headers["x-replica-id"]
+            for _ in range(10)
+        }
+        assert len(replicas) == 1
+
+    def test_different_sessions_spread(self, base_url):
+        replicas = {
+            session_chat_response(base_url, f"test-session-id-{i}").headers[
+                "x-replica-id"
+            ]
+            for i in range(10)
+        }
+        assert len(replicas) > 1
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main(["-v", __file__]))

python/ray/llm/tests/serve/cpu/deployments/routers/test_tokenizer.py

@@ -0,0 +1,282 @@
+import sys
+from typing import Optional
+from unittest.mock import AsyncMock, MagicMock
+
+import pytest
+from fastapi import HTTPException
+from starlette.datastructures import Headers
+
+from ray.llm._internal.serve.core.configs.llm_config import LLMConfig
+from ray.llm._internal.serve.core.configs.openai_api_models import (
+    ErrorInfo,
+    ErrorResponse,
+    TokenizeChatRequest,
+    TokenizeCompletionRequest,
+)
+from ray.llm._internal.serve.core.ingress.builder import (
+    LLMServingArgs,
+    build_openai_app,
+)
+from ray.llm._internal.serve.core.ingress.router import LLMRouter
+from ray.llm._internal.serve.core.ingress.tokenizer import TokenizeError, Tokenizer
+from ray.serve.experimental.round_robin_router import RoundRobinRouter
+from ray.serve.llm.request_router import KVAwareRouter
+
+
+class _TokenizeResponse:
+    def __init__(self, tokens):
+        self.tokens = tokens
+
+
+async def _tokenize_stream(response):
+    yield response
+
+
+def _handle_returning(response):
+    """A DeploymentHandle whose /tokenize streams ``response``; captures the
+    Tokenize* request it was called with under ``captured``."""
+    captured = {}
+
+    def tokenize_remote(tok_req, _):
+        captured["request"] = tok_req
+        return _tokenize_stream(response)
+
+    handle = MagicMock()
+    handle.options.return_value.tokenize.remote = tokenize_remote
+    return handle, captured
+
+
+class TestTokenizer:
+    @pytest.mark.parametrize(
+        "payload",
+        [
+            {"model": "m", "prompt": ["a", "b"]},  # batch of prompts
+            {"model": "m", "prompt": [1, 2, 3]},  # pre-tokenized token ids
+            {"model": "m"},  # neither messages nor prompt
+        ],
+    )
+    @pytest.mark.asyncio
+    async def test_untokenizable_payload_returns_none(self, payload):
+        """A parsed payload with no single-string prompt yields None."""
+        assert await Tokenizer(MagicMock()).tokenize(payload) is None
+
+    @pytest.mark.parametrize(
+        "payload, expected_request_type",
+        [
+            (
+                {"model": "m", "messages": [{"role": "user", "content": "hi"}]},
+                TokenizeChatRequest,
+            ),
+            ({"model": "m", "prompt": "hello"}, TokenizeCompletionRequest),
+        ],
+    )
+    @pytest.mark.asyncio
+    async def test_tokenizes_chat_and_completion(self, payload, expected_request_type):
+        """A chat or completion payload is sent to /tokenize as the right
+        Tokenize* request and its returned token ids are surfaced."""
+        handle, captured = _handle_returning(_TokenizeResponse([5, 6, 7]))
+        tokens = await Tokenizer(handle).tokenize(payload)
+        assert tokens == [5, 6, 7]
+        assert isinstance(captured["request"], expected_request_type)
+
+    @pytest.mark.parametrize(
+        "payload, expected",
+        [
+            (  # chat: template-rendering fields + request-provided prompt flags
+                {
+                    "model": "m",
+                    "messages": [{"role": "user", "content": "hi"}],
+                    "tools": [
+                        {
+                            "type": "function",
+                            "function": {"name": "f", "parameters": {}},
+                        }
+                    ],
+                    "chat_template": "TEMPLATE",
+                    "chat_template_kwargs": {"enable_thinking": False},
+                    "mm_processor_kwargs": {"num_crops": 4},
+                    "add_generation_prompt": False,
+                    "continue_final_message": True,
+                    "temperature": 0.7,
+                },
+                {
+                    "chat_template": "TEMPLATE",
+                    "chat_template_kwargs": {"enable_thinking": False},
+                    "mm_processor_kwargs": {"num_crops": 4},
+                    "add_generation_prompt": False,
+                    "continue_final_message": True,
+                },
+            ),
+            (  # completion: add_special_tokens comes from the request
+                {
+                    "model": "m",
+                    "prompt": "hi",
+                    "add_special_tokens": False,
+                    "temperature": 0.7,
+                },
+                {"add_special_tokens": False},
+            ),
+        ],
+    )
+    @pytest.mark.asyncio
+    async def test_forwards_prompt_fields_only(self, payload, expected):
+        """Prompt-rendering fields come from the request (not hardcoded) and
+        sampling params are dropped, so routing ids match prefill."""
+        handle, captured = _handle_returning(_TokenizeResponse([1, 2]))
+        await Tokenizer(handle).tokenize(payload)
+        request = captured["request"]
+        for attr, value in expected.items():
+            assert getattr(request, attr) == value
+        assert "temperature" not in (request.model_extra or {})
+
+    @pytest.mark.asyncio
+    async def test_error_response_raises(self):
+        """A /tokenize ErrorResponse surfaces as a TokenizeError carrying vLLM's
+        status code, message, and type."""
+        err = ErrorResponse(
+            error=ErrorInfo(message="bad model", type="NotFoundError", code=404)
+        )
+        handle, _ = _handle_returning(err)
+        with pytest.raises(TokenizeError) as exc_info:
+            await Tokenizer(handle).tokenize({"model": "m", "prompt": "hi"})
+        assert exc_info.value.status_code == 404
+        assert exc_info.value.message == "bad model"
+        assert exc_info.value.type == "NotFoundError"
+
+    @pytest.mark.asyncio
+    async def test_empty_response_raises(self):
+        """An empty /tokenize stream raises rather than returning no tokens."""
+
+        async def _empty(*_args):
+            for _ in ():
+                yield
+
+        handle = MagicMock()
+        handle.options.return_value.tokenize.remote = _empty
+        with pytest.raises(TokenizeError) as exc_info:
+            await Tokenizer(handle).tokenize({"model": "m", "prompt": "hi"})
+        assert exc_info.value.status_code == 500
+
+
+class TestRoute:
+    @pytest.mark.asyncio
+    async def test_no_tokenizer_forwards_none(self):
+        # A non-KV router has no tokenizer, so route forwards request_token_ids=None.
+        router = LLMRouter.__new__(LLMRouter)
+        router._handle = MagicMock()
+        router._tokenizer = None
+        router._pick_replica = AsyncMock(return_value=("h", 1, "rid"))
+
+        request = MagicMock()
+        request.body = AsyncMock(return_value=b'{"model": "m", "prompt": "hi"}')
+        request.headers = Headers({})
+        await router.route(request)
+        assert router._pick_replica.call_args.kwargs["request_token_ids"] is None
+
+    @pytest.mark.asyncio
+    async def test_forwards_token_ids(self):
+        # A successful tokenization forwards its token ids to _pick_replica.
+        router = LLMRouter.__new__(LLMRouter)
+        router._handle = MagicMock()
+        router._tokenizer = MagicMock()
+        router._tokenizer.tokenize = AsyncMock(return_value=[5, 6, 7])
+        router._pick_replica = AsyncMock(return_value=("h", 1, "rid"))
+
+        request = MagicMock()
+        request.body = AsyncMock(return_value=b'{"model": "m", "prompt": "hi"}')
+        request.headers = Headers({})
+        await router.route(request)
+        assert router._pick_replica.call_args.kwargs["request_token_ids"] == [5, 6, 7]
+
+    @pytest.mark.asyncio
+    async def test_unparseable_body_skips_tokenization(self):
+        # A truncated/unparseable body derives no routing payload, so the
+        # tokenizer is never called and request_token_ids stays None.
+        router = LLMRouter.__new__(LLMRouter)
+        router._handle = MagicMock()
+        router._tokenizer = MagicMock()
+        router._tokenizer.tokenize = AsyncMock(return_value=[5, 6, 7])
+        router._pick_replica = AsyncMock(return_value=("h", 1, "rid"))
+
+        request = MagicMock()
+        # Truncated prefix: not valid JSON, so it can't be parsed or tokenized.
+        request.body = AsyncMock(return_value=b'{"model": "m", "prompt": "' + b"x" * 8)
+        request.headers = Headers({"x-body-truncated": "8/90000"})
+        await router.route(request)
+
+        router._tokenizer.tokenize.assert_not_called()
+        assert router._pick_replica.call_args.kwargs["request_token_ids"] is None
+
+    @pytest.mark.asyncio
+    async def test_tokenize_error_becomes_http_error(self):
+        # A /tokenize rejection becomes an HTTPException with the same status
+        # code, and routing is not attempted.
+        router = LLMRouter.__new__(LLMRouter)
+        router._handle = MagicMock()
+        router._tokenizer = MagicMock()
+        router._tokenizer.tokenize = AsyncMock(
+            side_effect=TokenizeError(
+                "bad model", status_code=404, type="NotFoundError"
+            )
+        )
+        router._pick_replica = AsyncMock()
+
+        request = MagicMock()
+        request.body = AsyncMock(return_value=b'{"model": "m", "prompt": "hi"}')
+        request.headers = Headers({})
+        with pytest.raises(HTTPException) as exc_info:
+            await router.route(request)
+        assert exc_info.value.status_code == 404
+        assert exc_info.value.detail == "bad model"
+        router._pick_replica.assert_not_called()
+
+
+def _build_llm_app(request_router_class):
+    """Build a direct-streaming OpenAI app, optionally pinning a router class."""
+    deployment_config = {"autoscaling_config": {"min_replicas": 1, "max_replicas": 1}}
+    if request_router_class is not None:
+        deployment_config["request_router_config"] = {
+            "request_router_class": request_router_class
+        }
+    llm_config = LLMConfig(
+        model_loading_config={
+            "model_id": "qwen3-0.6b",
+            "model_source": "Qwen/Qwen3-0.6B",
+        },
+        accelerator_type=None,
+        deployment_config=deployment_config,
+    )
+    return build_openai_app(LLMServingArgs(llm_configs=[llm_config]))
+
+
+def _pre_routing_tokenization(app) -> Optional[bool]:
+    init_kwargs = app._ingress_request_router._bound_deployment.init_kwargs
+    return init_kwargs["pre_routing_tokenization"]
+
+
+class TestPreRoutingTokenization:
+    """build_openai_app enables pre-routing tokenization iff the router is KV-aware."""
+
+    @pytest.fixture(autouse=True)
+    def enable_direct_streaming(self, monkeypatch):
+        monkeypatch.setattr(
+            "ray.llm._internal.serve.core.ingress.builder."
+            "RAY_SERVE_LLM_ENABLE_DIRECT_STREAMING",
+            True,
+        )
+
+    @pytest.mark.parametrize(
+        "request_router_class, expected",
+        [
+            (KVAwareRouter, True),
+            (None, False),
+            (RoundRobinRouter, False),
+        ],
+    )
+    def test_enabled_only_for_kv_aware_router(self, request_router_class, expected):
+        app = _build_llm_app(request_router_class)
+        assert _pre_routing_tokenization(app) is expected
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main(["-v", __file__]))

python/ray/llm/tests/serve/cpu/deployments/test_openai_json_wrapper.py

@@ -0,0 +1,44 @@
+import sys
+
+import pytest
+
+from ray.llm._internal.serve.core.ingress.ingress import _openai_json_wrapper
+
+
+async def _async_gen_from_list(items):
+    for item in items:
+        yield item
+
+
+@pytest.mark.asyncio
+class TestOpenAIJsonWrapper:
+    async def test_no_duplicate_done_when_upstream_sends_done(self):
+        """When the upstream generator already yields 'data: [DONE]', the
+        wrapper must not append a second one."""
+        upstream = ['data: {"id": 1}\n\n', "data: [DONE]\n\n"]
+        chunks = [c async for c in _openai_json_wrapper(_async_gen_from_list(upstream))]
+        assert chunks == upstream
+        assert chunks.count("data: [DONE]\n\n") == 1
+
+    async def test_done_appended_when_upstream_does_not_send_done(self):
+        """When the upstream generator does not yield 'data: [DONE]', the
+        wrapper must append it."""
+        upstream = ['data: {"id": 1}\n\n']
+        chunks = [c async for c in _openai_json_wrapper(_async_gen_from_list(upstream))]
+        assert chunks == ['data: {"id": 1}\n\n', "data: [DONE]\n\n"]
+
+    async def test_done_appended_for_empty_stream(self):
+        """An empty upstream stream should still produce a [DONE] sentinel."""
+        chunks = [c async for c in _openai_json_wrapper(_async_gen_from_list([]))]
+        assert chunks == ["data: [DONE]\n\n"]
+
+    async def test_no_duplicate_done_when_upstream_sends_done_batched(self):
+        """When the upstream generator yields a batch containing 'data: [DONE]',
+        the wrapper must not append a second one."""
+        upstream = [['data: {"id": 1}\n\n', "data: [DONE]\n\n"]]
+        chunks = [c async for c in _openai_json_wrapper(_async_gen_from_list(upstream))]
+        assert chunks == ['data: {"id": 1}\n\ndata: [DONE]\n\n']
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main(["-v", __file__]))

python/ray/llm/tests/serve/cpu/deployments/test_prefix_aware_request_router.py

@@ -0,0 +1,508 @@
+import asyncio
+import time
+
+import pytest
+
+import ray
+from ray._common.utils import get_or_create_event_loop
+from ray.llm._internal.serve.routing_policies.prefix_aware.prefix_aware_router import (
+    PrefixCacheAffinityRouter,
+)
+from ray.llm._internal.serve.routing_policies.prefix_aware.prefix_tree import (
+    PrefixTreeActor,
+)
+from ray.serve._private.common import (
+    DeploymentHandleSource,
+    DeploymentID,
+    RequestMetadata,
+)
+from ray.serve._private.request_router.common import PendingRequest
+from ray.serve._private.test_utils import MockTimer
+from ray.serve._private.utils import generate_request_id
+from ray.serve.tests.unit.test_pow_2_request_router import (
+    FakeRunningReplica,
+)  # Reuse the FakeRunningReplica from the Pow2 test
+
+TIMER = MockTimer()
+DEFAULT_MAX_ONGOING_REQUESTS = 10
+
+
+# === Fixtures ===
+
+
+@pytest.fixture
+def tree_actor():
+    """Create a fresh PrefixTreeActor instance."""
+    actor = PrefixTreeActor.options(name="PrefixTreeActor").remote()
+    yield actor
+    ray.kill(actor)
+
+
+@pytest.fixture
+def prefix_request_router(tree_actor, request):
+    """Create a fresh PrefixCacheAffinityRouter with connected tree_actor."""
+    params = getattr(request, "param", {})
+
+    async def construct_request_router(loop: asyncio.AbstractEventLoop):
+        request_router = PrefixCacheAffinityRouter(
+            deployment_id=DeploymentID(name="TEST_DEPLOYMENT"),
+            handle_source=DeploymentHandleSource.REPLICA,
+            use_replica_queue_len_cache=False,
+            get_curr_time_s=TIMER.time,
+        )
+        return request_router
+
+    request_router = asyncio.new_event_loop().run_until_complete(
+        construct_request_router(get_or_create_event_loop())
+    )
+    request_router.initialize_state(
+        imbalanced_threshold=params.get("imbalanced_threshold", float("inf")),
+        match_rate_threshold=params.get("match_rate_threshold", 0.1),
+        do_eviction=params.get("do_eviction", False),
+        eviction_threshold_chars=params.get("eviction_threshold_chars"),
+        eviction_target_chars=params.get("eviction_target_chars"),
+        eviction_interval_secs=params.get("eviction_interval_secs"),
+        tree_actor=tree_actor,
+    )
+
+    yield request_router
+    assert request_router.curr_num_routing_tasks == 0
+    assert request_router.num_pending_requests == 0
+
+
+# === Helpers ===
+
+
+class PromptRequest:
+    def __init__(self, prompt: str):
+        self.prompt = prompt
+
+
+class ChatRequest:
+    def __init__(self, messages):
+        self.messages = messages
+
+
+def fake_pending_request(prompt=None, messages=None) -> PendingRequest:
+    if prompt is not None:
+        args = [PromptRequest(prompt)]
+    elif messages is not None:
+        args = [ChatRequest(messages)]
+    else:
+        args = []
+
+    return PendingRequest(
+        args=args,
+        kwargs={},
+        metadata=RequestMetadata(
+            request_id=generate_request_id(),
+            internal_request_id=generate_request_id(),
+            multiplexed_model_id="",
+        ),
+        created_at=time.time(),
+    )
+
+
+# === Tests ===
+class TestPow2FallbackBehavior:
+    """Tests fallback to Pow2 when prefix-aware logic should be skipped."""
+
+    @pytest.mark.asyncio
+    async def test_fallback_when_no_prompt(self, prefix_request_router):
+        """No args → prefix logic skipped → falls back to least busy replica."""
+        r1 = FakeRunningReplica("r1")
+        r1.set_queue_len_response(0)
+        r2 = FakeRunningReplica("r2")
+        r2.set_queue_len_response(5)
+        prefix_request_router.update_replicas([r1, r2])
+
+        tenant_to_char_count = ray.get(
+            prefix_request_router._tree_actor.getattr.remote("tenant_to_char_count")
+        )
+        assert tenant_to_char_count == {
+            r1.replica_id.to_full_id_str(): 0,
+            r2.replica_id.to_full_id_str(): 0,
+        }
+
+        req = fake_pending_request()
+        for _ in range(10):
+            chosen = await prefix_request_router._choose_replica_for_request(req)
+            assert chosen == r1
+
+    @pytest.mark.parametrize(
+        "prefix_request_router", [{"imbalanced_threshold": 2}], indirect=True
+    )
+    @pytest.mark.asyncio
+    async def test_fallback_when_imbalanced(self, prefix_request_router):
+        """If load is imbalanced beyond threshold, prefix matching is skipped."""
+        r1 = FakeRunningReplica("r1")
+        r1.set_queue_len_response(0)
+        r2 = FakeRunningReplica("r2")
+        r2.set_queue_len_response(10)
+        prefix_request_router.update_replicas([r1, r2])
+
+        ray.get(
+            prefix_request_router._tree_actor.insert.remote(
+                "hello world", r2.replica_id.to_full_id_str(), time.time()
+            )
+        )
+
+        tenant_to_char_count = ray.get(
+            prefix_request_router._tree_actor.getattr.remote("tenant_to_char_count")
+        )
+        assert tenant_to_char_count == {
+            r1.replica_id.to_full_id_str(): 0,
+            r2.replica_id.to_full_id_str(): 11,
+        }
+
+        matched_text, matched_tenants = ray.get(
+            prefix_request_router._tree_actor.prefix_match.remote("hello world")
+        )
+        assert matched_text == "hello world"
+        assert matched_tenants == [r2.replica_id.to_full_id_str()]
+
+        req = fake_pending_request(prompt="hello world")
+        for _ in range(10):
+            chosen = await prefix_request_router._choose_replica_for_request(req)
+            # Even though r2 has a higher match rate, it is not chosen because the load is imbalanced
+            assert chosen == r1
+
+
+class TestPrefixAwareLogic:
+    """Tests that exercise actual prefix-aware request routing logic."""
+
+    @pytest.mark.asyncio
+    async def test_high_match_rate_selects_matching_replica(
+        self, prefix_request_router
+    ):
+        """High match rate → use matched replica instead of Pow2."""
+        r1 = FakeRunningReplica("r1")
+        r1.set_queue_len_response(0)
+        r2 = FakeRunningReplica("r2")
+        r2.set_queue_len_response(0)
+        prefix_request_router.update_replicas([r1, r2])
+        ray.get(
+            prefix_request_router._tree_actor.insert.remote(
+                "Hello", r2.replica_id.to_full_id_str(), time.time()
+            )
+        )
+        # Verify prefix match and smallest tenants
+        matched_text, matched_tenants = ray.get(
+            prefix_request_router._tree_actor.prefix_match.remote("Hello world")
+        )
+        assert matched_text == "Hello"
+        assert matched_tenants == [r2.replica_id.to_full_id_str()]
+
+        tenant_counts = ray.get(
+            prefix_request_router._tree_actor.getattr.remote("tenant_to_char_count")
+        )
+        assert tenant_counts[r1.replica_id.to_full_id_str()] == 0
+        assert tenant_counts[r2.replica_id.to_full_id_str()] == 5
+
+        prompt_req = fake_pending_request(prompt="Hello world")
+        for _ in range(10):
+            chosen = await prefix_request_router._choose_replica_for_request(prompt_req)
+            assert chosen == r2
+        chat_req = fake_pending_request(
+            messages=[{"content": "Hello"}, {"content": " world"}]
+        )
+        for _ in range(10):
+            chosen = await prefix_request_router._choose_replica_for_request(chat_req)
+            assert chosen == r2
+
+    @pytest.mark.asyncio
+    async def test_low_match_rate_uses_smallest_tree(self, prefix_request_router):
+        """Low match rate → use replica with least total inserted characters."""
+        r1 = FakeRunningReplica("r1")
+        r1.set_queue_len_response(0)
+        r2 = FakeRunningReplica("r2")
+        r2.set_queue_len_response(0)
+        prefix_request_router.update_replicas([r1, r2])
+
+        # Make r2 "bigger" tenant
+        ray.get(
+            prefix_request_router._tree_actor.insert.remote(
+                "hi", r1.replica_id.to_full_id_str(), time.time()
+            )
+        )
+        ray.get(
+            prefix_request_router._tree_actor.insert.remote(
+                "longtext", r2.replica_id.to_full_id_str(), time.time()
+            )
+        )
+
+        # Verify tenant character counts
+        tenant_counts = ray.get(
+            prefix_request_router._tree_actor.getattr.remote("tenant_to_char_count")
+        )
+        assert tenant_counts[r1.replica_id.to_full_id_str()] == 2  # "hi"
+        assert tenant_counts[r2.replica_id.to_full_id_str()] == 8  # "longtext"
+
+        prompt_req = fake_pending_request(prompt="z")
+        for _ in range(10):
+            # Both tenants have 0% match rate, so the smaller tenant (r1) is chosen
+            assert (
+                await prefix_request_router._choose_replica_for_request(prompt_req)
+                == r1
+            )
+
+        chat_req = fake_pending_request(messages=[{"content": "z"}])
+        for _ in range(10):
+            # Both tenants have 0% match rate, so the smaller tenant (r1) is chosen
+            assert (
+                await prefix_request_router._choose_replica_for_request(chat_req) == r1
+            )
+
+
+class TestEvictionBehavior:
+    """Tests for prefix tree eviction behavior."""
+
+    @pytest.mark.parametrize(
+        "prefix_request_router",
+        [
+            {
+                "do_eviction": True,
+                "eviction_threshold_chars": 10,
+                "eviction_target_chars": 5,
+                "eviction_interval_secs": 1.0,
+            }
+        ],
+        indirect=True,
+    )
+    @pytest.mark.asyncio
+    async def test_eviction_task_creation(self, prefix_request_router):
+        """Test that eviction task is only created after update_replicas."""
+        # Before update_replicas
+        assert not prefix_request_router._eviction_loop_running
+
+        # After update_replicas
+        r1 = FakeRunningReplica("r1")
+        prefix_request_router.update_replicas([r1])
+        assert prefix_request_router._eviction_loop_running
+
+        # After stop_eviction_loop
+        ray.get(prefix_request_router._tree_actor.stop_eviction_loop.remote())
+        await asyncio.sleep(0.1)
+
+
+class TestPromptNormalization:
+    """Tests for input normalization in the prefix-aware router."""
+
+    def test_normalize_prompt_string(self, prefix_request_router):
+        req = fake_pending_request(prompt="Hello world")
+        normalized = prefix_request_router._extract_text_from_request(req)
+        assert normalized == "Hello world"
+
+    def test_normalize_messages_list_of_strings(self, prefix_request_router):
+        req = fake_pending_request(messages=["Hello", " ", "world"])
+        normalized = prefix_request_router._extract_text_from_request(req)
+        assert normalized == "Hello world"
+
+    def test_normalize_messages_dict_content_string(self, prefix_request_router):
+        req = fake_pending_request(
+            messages=[
+                {"content": "Hello"},
+                {"content": " world"},
+            ]
+        )
+        normalized = prefix_request_router._extract_text_from_request(req)
+        assert normalized == "Hello world"
+
+    def test_normalize_messages_dict_content_list_of_dicts_text(
+        self, prefix_request_router
+    ):
+        req = fake_pending_request(
+            messages=[
+                {
+                    "content": [
+                        {"type": "text", "text": "Hello"},
+                        {"type": "text", "text": " world"},
+                    ]
+                }
+            ]
+        )
+        normalized = prefix_request_router._extract_text_from_request(req)
+        assert normalized == "Hello world"
+
+    def test_normalize_messages_dict_content_list_of_strings(
+        self, prefix_request_router
+    ):
+        req = fake_pending_request(messages=[{"content": ["Hello", " ", "world"]}])
+        normalized = prefix_request_router._extract_text_from_request(req)
+        assert normalized == "Hello world"
+
+    def test_normalize_unsupported_returns_empty(self, prefix_request_router):
+        # For now, unsupported multimodal parts should be ignored, resulting in empty string
+        req = fake_pending_request(
+            messages=[
+                {
+                    "content": [
+                        {
+                            "type": "image_url",
+                            "image_url": {"url": "http://example.com"},
+                        },
+                    ]
+                }
+            ]
+        )
+        normalized = prefix_request_router._extract_text_from_request(req)
+        assert normalized == ""
+
+    def test_extract_raises_when_no_prompt_or_messages(self, prefix_request_router):
+        with pytest.raises(ValueError):
+            _ = prefix_request_router._extract_text_from_request(fake_pending_request())
+
+    @pytest.mark.parametrize(
+        "prefix_request_router",
+        [
+            {
+                "do_eviction": True,
+                "eviction_threshold_chars": 10,
+                "eviction_target_chars": 5,
+                "eviction_interval_secs": 1.0,
+            }
+        ],
+        indirect=True,
+    )
+    @pytest.mark.asyncio
+    async def test_eviction_threshold_behavior(self, prefix_request_router):
+        """Test that eviction reduces tree size below threshold after interval."""
+        r1 = FakeRunningReplica("r1")
+        prefix_request_router.update_replicas([r1])
+
+        # Insert text that exceeds eviction_threshold_chars
+        ray.get(
+            prefix_request_router._tree_actor.insert.remote(
+                "verylongtext", r1.replica_id.to_full_id_str(), time.time()
+            )
+        )
+        ray.get(
+            prefix_request_router._tree_actor.insert.remote(
+                "anotherlongtext", r1.replica_id.to_full_id_str(), time.time()
+            )
+        )
+
+        # Verify initial size exceeds eviction_threshold_chars
+        tenant_counts = ray.get(
+            prefix_request_router._tree_actor.getattr.remote("tenant_to_char_count")
+        )
+        assert tenant_counts[r1.replica_id.to_full_id_str()] > 10
+
+        # Wait for eviction interval
+        await asyncio.sleep(1.1)
+
+        # Verify size is reduced below eviction_target_chars
+        tenant_counts = ray.get(
+            prefix_request_router._tree_actor.getattr.remote("tenant_to_char_count")
+        )
+        assert tenant_counts[r1.replica_id.to_full_id_str()] <= 5
+
+        ray.get(prefix_request_router._tree_actor.stop_eviction_loop.remote())
+        await asyncio.sleep(0.1)
+
+
+class TestMultiDeploymentIsolation:
+    """Tests that multiple deployments get isolated prefix tree actors."""
+
+    @pytest.mark.asyncio
+    async def test_two_deployments_get_separate_tree_actors(self):
+        """Verify that two deployments using PrefixCacheAffinityRouter get
+        deployment-specific prefix tree actors to avoid replica ID conflicts."""
+
+        # Create separate tree actors for each deployment
+        prefill_tree = PrefixTreeActor.options(name="PrefillTree").remote()
+        decode_tree = PrefixTreeActor.options(name="DecodeTree").remote()
+
+        # Create two routers for different deployments (e.g., Prefill and Decode in PD setup)
+        async def construct_router(deployment_name: str, tree_actor):
+            router = PrefixCacheAffinityRouter(
+                deployment_id=DeploymentID(name=deployment_name),
+                handle_source=DeploymentHandleSource.REPLICA,
+                use_replica_queue_len_cache=False,
+                get_curr_time_s=TIMER.time,
+            )
+            router.initialize_state(tree_actor=tree_actor)
+            return router
+
+        prefill_router = await construct_router("Prefill:deepseek", prefill_tree)
+        decode_router = await construct_router("Decode:deepseek", decode_tree)
+
+        # Create replicas for each deployment
+        prefill_r1 = FakeRunningReplica("prefill_r1")
+        prefill_r1.set_queue_len_response(0)
+        prefill_r2 = FakeRunningReplica("prefill_r2")
+        prefill_r2.set_queue_len_response(0)
+
+        decode_r1 = FakeRunningReplica("decode_r1")
+        decode_r1.set_queue_len_response(0)
+        decode_r2 = FakeRunningReplica("decode_r2")
+        decode_r2.set_queue_len_response(0)
+
+        # Update replicas for each router
+        prefill_router.update_replicas([prefill_r1, prefill_r2])
+        decode_router.update_replicas([decode_r1, decode_r2])
+
+        # Verify replicas are tracked independently in each tree
+        prefill_tenants = ray.get(
+            prefill_router._tree_actor.getattr.remote("tenant_to_char_count")
+        )
+        decode_tenants = ray.get(
+            decode_router._tree_actor.getattr.remote("tenant_to_char_count")
+        )
+
+        # Each tree should only know about its own replicas
+        assert set(prefill_tenants.keys()) == {
+            prefill_r1.replica_id.to_full_id_str(),
+            prefill_r2.replica_id.to_full_id_str(),
+        }
+        assert set(decode_tenants.keys()) == {
+            decode_r1.replica_id.to_full_id_str(),
+            decode_r2.replica_id.to_full_id_str(),
+        }
+
+        # Insert text into prefill tree
+        ray.get(
+            prefill_router._tree_actor.insert.remote(
+                "prefill text", prefill_r1.replica_id.to_full_id_str(), time.time()
+            )
+        )
+
+        # Insert text into decode tree
+        ray.get(
+            decode_router._tree_actor.insert.remote(
+                "decode text", decode_r1.replica_id.to_full_id_str(), time.time()
+            )
+        )
+
+        # Verify routing works correctly for both deployments without KeyErrors
+        prefill_req = fake_pending_request(prompt="prefill text continued")
+        chosen_prefill = await prefill_router._choose_replica_for_request(prefill_req)
+        assert chosen_prefill == prefill_r1
+
+        decode_req = fake_pending_request(prompt="decode text continued")
+        chosen_decode = await decode_router._choose_replica_for_request(decode_req)
+        assert chosen_decode == decode_r1
+
+        # Verify trees remain isolated
+        prefill_tenants_after = ray.get(
+            prefill_router._tree_actor.getattr.remote("tenant_to_char_count")
+        )
+        decode_tenants_after = ray.get(
+            decode_router._tree_actor.getattr.remote("tenant_to_char_count")
+        )
+
+        assert prefill_tenants_after[prefill_r1.replica_id.to_full_id_str()] > 0
+        assert prefill_tenants_after[prefill_r2.replica_id.to_full_id_str()] == 0
+        assert decode_tenants_after[decode_r1.replica_id.to_full_id_str()] > 0
+        assert decode_tenants_after[decode_r2.replica_id.to_full_id_str()] == 0
+
+        # Cleanup
+        ray.kill(prefill_router._tree_actor)
+        ray.kill(decode_router._tree_actor)
+
+
+if __name__ == "__main__":
+    import sys
+
+    exit_code = pytest.main(["-vs", __file__])
+    sys.exit(exit_code)

python/ray/llm/tests/serve/cpu/deployments/test_prefix_tree.py

@@ -0,0 +1,999 @@
+import asyncio
+from typing import List, Set
+
+import pytest
+
+import ray
+from ray.llm._internal.serve.routing_policies.prefix_aware.prefix_tree import (
+    Node,
+    PrefixTree,
+    PrefixTreeActor,
+)
+
+
+# Fixtures
+@pytest.fixture
+def tree() -> PrefixTree:
+    """Create a fresh PrefixTree instance for each local test."""
+    return PrefixTree()
+
+
+@pytest.fixture
+def tree_actor():
+    """Create a fresh PrefixTreeActor instance for each ray.remote test."""
+    return PrefixTreeActor.remote()
+
+
+# Helper to get LRU chain texts
+def get_lru_texts_from_tree(tree: PrefixTree, tenant_id: str) -> List[str]:
+    """Gets LRU chain texts directly from a PrefixTree instance."""
+    chain = tree._get_lru_chain(tenant_id)
+    return [node.text for node in chain]
+
+
+async def get_lru_texts_from_tree_actor(
+    tree_actor: PrefixTreeActor, tenant_id: str
+) -> List[str]:
+    """Gets LRU chain texts from a PrefixTreeActor."""
+    chain = ray.get(tree_actor._get_lru_chain.remote(tenant_id))
+    return [node.text for node in chain]
+
+
+class TestPrefixTreeInitialization:
+    """Tests for the PrefixTree class initialization and basic tenant management."""
+
+    def test_initial_state(self, tree: PrefixTree) -> None:
+        """Test the initial state of a new PrefixTree."""
+        assert tree.tenant_to_char_count == {}
+        assert tree.tenant_to_lru_tail == {}
+        assert tree.root is not None
+        assert tree.root.text == ""
+        assert tree.root.parent is None
+        assert tree.root.tenant_to_last_access_time == {}
+        assert tree.root.edge_label_to_child == {}
+
+    def test_add_tenant(self, tree: PrefixTree) -> None:
+        """Test adding a new tenant via add_tenants."""
+        tree.add_tenants(["tenant_1"], 0)
+        assert tree.tenant_to_char_count == {"tenant_1": 0}
+        assert tree.tenant_to_lru_tail.get("tenant_1") == tree.root
+        assert tree.root.tenant_to_last_access_time == {"tenant_1": 0}
+        assert get_lru_texts_from_tree(tree, "tenant_1") == [""]
+
+    def test_add_existing_tenant_noop(self, tree: PrefixTree) -> None:
+        """Test that adding an existing tenant via add_tenants is a no-op."""
+        tree.add_tenants(["tenant_1"], 0)
+        assert tree.tenant_to_char_count == {"tenant_1": 0}
+        assert tree.tenant_to_lru_tail.get("tenant_1") == tree.root
+        assert tree.root.tenant_to_last_access_time == {"tenant_1": 0}
+        assert get_lru_texts_from_tree(tree, "tenant_1") == [""]
+
+        tree.add_tenants(["tenant_1"], 0)  # Add again
+
+        assert tree.tenant_to_char_count == {"tenant_1": 0}
+        assert tree.tenant_to_lru_tail.get("tenant_1") == tree.root
+        assert tree.root.tenant_to_last_access_time == {"tenant_1": 0}
+        assert get_lru_texts_from_tree(tree, "tenant_1") == [""]
+
+    def test_add_multiple_tenants(self, tree: PrefixTree) -> None:
+        """Test adding multiple tenants at once."""
+        tree.add_tenants(["tenant_1", "tenant_2", "tenant_3"], 0)
+
+        assert tree.tenant_to_char_count == {
+            "tenant_1": 0,
+            "tenant_2": 0,
+            "tenant_3": 0,
+        }
+        for tenant in ["tenant_1", "tenant_2", "tenant_3"]:
+            assert tree.tenant_to_lru_tail.get(tenant) == tree.root
+            assert tree.root.tenant_to_newer_node.get(tenant) is None
+            assert tree.root.tenant_to_older_node.get(tenant) is None
+            assert tree.root.tenant_to_last_access_time == {
+                "tenant_1": 0,
+                "tenant_2": 0,
+                "tenant_3": 0,
+            }
+            assert get_lru_texts_from_tree(tree, tenant) == [""]
+
+    def test_add_multiple_tenants_with_existing(self, tree: PrefixTree) -> None:
+        """Test adding multiple tenants when some already exist."""
+        tree.add_tenants(["tenant_1"], 0)
+        assert tree.root.tenant_to_last_access_time == {"tenant_1": 0}
+        assert tree.tenant_to_char_count == {"tenant_1": 0}
+        assert "tenant_1" in tree.tenant_to_lru_tail
+
+        # Add a mix of new and existing tenants
+        tree.add_tenants(["tenant_1", "tenant_2", "tenant_3"], 0)
+        # Existing tenants should remain unchanged
+        assert tree.root.tenant_to_last_access_time == {
+            "tenant_1": 0,
+            "tenant_2": 0,
+            "tenant_3": 0,
+        }
+        assert tree.tenant_to_char_count == {
+            "tenant_1": 0,
+            "tenant_2": 0,
+            "tenant_3": 0,
+        }
+        assert all(
+            tenant in tree.tenant_to_lru_tail
+            for tenant in ["tenant_1", "tenant_2", "tenant_3"]
+        )
+
+
+class TestPrefixTreeInsert:
+    def test_insert_non_existent_tenant(self, tree: PrefixTree) -> None:
+        """Test inserting a string for a non-existent tenant fails."""
+        # Insert without adding tenant first
+        tree.insert("hello", "nonexistent", 1)
+
+        # Verify insert did nothing since tenant doesn't exist
+        assert "nonexistent" not in tree.tenant_to_char_count
+        assert get_lru_texts_from_tree(tree, "nonexistent") == []
+        assert "h" not in tree.root.edge_label_to_child
+
+    def test_insert_single_string(self, tree: PrefixTree) -> None:
+        """Test inserting a single string after adding a tenant."""
+        tree.add_tenants(["tenant_1"], 0)
+        tree.insert("hello", "tenant_1", 1)
+        assert tree.tenant_to_char_count == {"tenant_1": 5}
+        assert get_lru_texts_from_tree(tree, "tenant_1") == ["", "hello"]
+
+        root_node = tree.root
+        assert root_node.tenant_to_last_access_time == {"tenant_1": 1}
+        assert set(root_node.edge_label_to_child.keys()) == {"h"}
+
+        hello_node = root_node.edge_label_to_child["h"]
+        assert hello_node.text == "hello"
+        assert hello_node.parent == root_node
+        assert hello_node.tenant_to_last_access_time == {"tenant_1": 1}
+        assert hello_node.edge_label_to_child == {}
+
+    def test_insert_duplicate_string(self, tree: PrefixTree) -> None:
+        """Test inserting a duplicate string for the same tenant."""
+        tree.add_tenants(["tenant_1"], 0)
+        tree.insert("hello", "tenant_1", 1)  # Initial insert
+        tree.insert("hello", "tenant_1", 1)  # Duplicate insert with the same timestamp
+
+        assert tree.tenant_to_char_count == {"tenant_1": 5}  # Char count unchanged
+        assert get_lru_texts_from_tree(tree, "tenant_1") == [
+            "",
+            "hello",
+        ]  # LRU order same
+
+        hello_node = tree.root.edge_label_to_child["h"]
+        assert tree.root.tenant_to_last_access_time == {"tenant_1": 1}
+        assert hello_node.tenant_to_last_access_time == {"tenant_1": 1}
+
+        tree.insert("hello", "tenant_1", 2)  # Duplicate insert with new timestamp
+
+        assert tree.tenant_to_char_count == {"tenant_1": 5}  # Char count unchanged
+        assert get_lru_texts_from_tree(tree, "tenant_1") == [
+            "",
+            "hello",
+        ]  # LRU order same
+
+        hello_node = tree.root.edge_label_to_child["h"]
+        assert tree.root.tenant_to_last_access_time == {"tenant_1": 2}
+        assert hello_node.tenant_to_last_access_time == {"tenant_1": 2}
+
+    def test_insert_multiple_tenants(self, tree: PrefixTree) -> None:
+        """Test inserting the same string for different tenants."""
+        tree.add_tenants(["tenant_1", "tenant_2"], 0)
+        tree.insert("hello", "tenant_1", 1)
+        tree.insert("hello", "tenant_2", 2)
+
+        assert tree.tenant_to_char_count == {"tenant_1": 5, "tenant_2": 5}
+        assert get_lru_texts_from_tree(tree, "tenant_1") == ["", "hello"]
+        assert get_lru_texts_from_tree(tree, "tenant_2") == ["", "hello"]
+
+        hello_node = tree.root.edge_label_to_child["h"]
+        assert tree.root.tenant_to_last_access_time == {"tenant_1": 1, "tenant_2": 2}
+        assert hello_node.tenant_to_last_access_time == {"tenant_1": 1, "tenant_2": 2}
+
+    def test_insert_node_split(self, tree: PrefixTree) -> None:
+        """Test insertion that causes an existing node to split due to differing suffixes."""
+        tree.add_tenants(["tenant_1", "tenant_2"], 0)
+        tree.insert("helloworld", "tenant_1", 1)
+        tree.insert("hellothere", "tenant_2", 2)  # "hello" is common prefix
+
+        assert tree.tenant_to_char_count == {"tenant_1": 10, "tenant_2": 10}
+        assert get_lru_texts_from_tree(tree, "tenant_1") == ["", "hello", "world"]
+        assert get_lru_texts_from_tree(tree, "tenant_2") == ["", "there", "hello"]
+
+        hello_node = tree.root.edge_label_to_child["h"]
+        assert hello_node.text == "hello"
+        assert hello_node.tenant_to_last_access_time == {"tenant_1": 1, "tenant_2": 2}
+        assert set(hello_node.edge_label_to_child.keys()) == {"w", "t"}
+
+        world_node = hello_node.edge_label_to_child["w"]
+        assert world_node.text == "world"
+        assert world_node.tenant_to_last_access_time == {"tenant_1": 1}
+
+        there_node = hello_node.edge_label_to_child["t"]
+        assert there_node.text == "there"
+        assert there_node.tenant_to_last_access_time == {"tenant_2": 2}
+
+    def test_insert_longer_string_with_shared_prefix(self, tree: PrefixTree) -> None:
+        """Test inserting a longer string that shares a prefix with an existing node string."""
+        tree.add_tenants(["tenant_1", "tenant_2"], 0)
+        tree.insert("hello", "tenant_1", 1)
+        tree.insert("helloworld", "tenant_2", 2)  # "hello" is prefix of "helloworld"
+
+        assert tree.tenant_to_char_count == {"tenant_1": 5, "tenant_2": 10}
+        assert get_lru_texts_from_tree(tree, "tenant_1") == ["", "hello"]
+        assert get_lru_texts_from_tree(tree, "tenant_2") == ["", "world", "hello"]
+
+        hello_node = tree.root.edge_label_to_child["h"]
+        assert hello_node.text == "hello"
+        assert hello_node.tenant_to_last_access_time == {"tenant_1": 1, "tenant_2": 2}
+        assert set(hello_node.edge_label_to_child.keys()) == {"w"}
+
+        world_node = hello_node.edge_label_to_child["w"]
+        assert world_node.text == "world"
+        assert world_node.tenant_to_last_access_time == {"tenant_2": 2}
+
+        # Ensure no empty non-root nodes created
+        empty_text_nodes: List[Node] = []
+        nodes_to_check: List[Node] = [tree.root]
+        visited_nodes: Set[Node] = {tree.root}
+        while nodes_to_check:
+            node: Node = nodes_to_check.pop()
+            if node.text == "" and node != tree.root:  # check for non-root empty nodes
+                empty_text_nodes.append(node)
+            for child in node.edge_label_to_child.values():
+                if child not in visited_nodes:
+                    nodes_to_check.append(child)
+                    visited_nodes.add(child)
+        assert not empty_text_nodes
+
+    def test_insert_shorter_string_with_shared_prefix(self, tree: PrefixTree) -> None:
+        """Test inserting a shorter string that is a prefix of an existing longer string, causing split."""
+        tree.add_tenants(["tenant_1", "tenant_2"], 0)
+        tree.insert("helloworld", "tenant_1", 1)
+        tree.insert(
+            "hello", "tenant_2", 2
+        )  # "hello" is prefix, causes "helloworld" to split
+
+        assert tree.tenant_to_char_count == {"tenant_1": 10, "tenant_2": 5}
+        assert get_lru_texts_from_tree(tree, "tenant_1") == ["", "hello", "world"]
+        assert get_lru_texts_from_tree(tree, "tenant_2") == ["", "hello"]
+
+        hello_node = tree.root.edge_label_to_child["h"]
+        assert hello_node.text == "hello"
+        assert hello_node.tenant_to_last_access_time == {"tenant_1": 1, "tenant_2": 2}
+        assert set(hello_node.edge_label_to_child.keys()) == {"w"}
+
+        world_node = hello_node.edge_label_to_child["w"]
+        assert world_node.text == "world"
+        assert world_node.tenant_to_last_access_time == {"tenant_1": 1}
+
+
+class TestPrefixTreeMatch:
+    def test_prefix_match_empty_tree(self, tree: PrefixTree) -> None:
+        """Test prefix_match on an empty tree returns empty string and None tenants."""
+        matched_text, matched_tenants = tree.prefix_match("hello")
+        assert matched_text == ""
+        assert matched_tenants is None
+
+    def test_prefix_match_no_match(self, tree: PrefixTree) -> None:
+        """Test prefix_match for a non-matching prefix returns empty string and all tenants."""
+        tree.add_tenants(["tenant_1", "tenant_2"], 0)
+        tree.insert("hello", "tenant_1", 1)
+        tree.insert("world", "tenant_2", 2)
+        matched_text, matched_tenants = tree.prefix_match("foobar")
+        assert matched_text == ""
+        assert matched_tenants is not None
+        assert sorted(matched_tenants) == sorted(["tenant_1", "tenant_2"])
+
+    def test_prefix_match_query_longer_than_stored_strings(
+        self, tree: PrefixTree
+    ) -> None:
+        """Test prefix_match where query is longer than any stored string but matches a full path."""
+        tree.add_tenants(["tenant_1", "tenant_2"], 0)
+        tree.insert("helloworld", "tenant_1", 1)
+        tree.insert("hellothere", "tenant_2", 2)
+        matched_text, matched_tenants = tree.prefix_match("hellothereextra")
+        assert matched_text == "hellothere"
+        assert matched_tenants == ["tenant_2"]
+
+    def test_prefix_match_exact_match(self, tree: PrefixTree) -> None:
+        """Test prefix_match with an exact match for a single tenant."""
+        tree.add_tenants(["tenant_1"], 0)
+        tree.insert("hello", "tenant_1", 1)
+        matched_text, matched_tenants = tree.prefix_match("hello")
+        assert matched_text == "hello"
+        assert matched_tenants == ["tenant_1"]
+
+    def test_prefix_match_partial_match(self, tree: PrefixTree) -> None:
+        """Test prefix_match with a partial query matching the longest common part of a branch."""
+        tree.add_tenants(["tenant_1", "tenant_2"], 0)
+        tree.insert("apple", "tenant_1", 1)
+        tree.insert("apricot", "tenant_2", 2)
+        matched_text, matched_tenants = tree.prefix_match("application")
+        assert matched_text == "appl"  # Longest of ("appl", "ap")
+        assert matched_tenants == ["tenant_1"]
+
+    def test_prefix_match_with_tenant_filter(self, tree: PrefixTree) -> None:
+        """Test prefix_match with a tenant filter selecting a specific branch."""
+        tree.add_tenants(["tenant_1", "tenant_2"], 0)
+        tree.insert("apple", "tenant_1", 1)
+        tree.insert("apricot", "tenant_2", 2)
+        matched_text, matched_tenants = tree.prefix_match("application", ["tenant_2"])
+        assert matched_text == "ap"
+        assert matched_tenants == ["tenant_2"]
+
+    def test_prefix_match_with_shared_prefix_tenant_filter(
+        self, tree: PrefixTree
+    ) -> None:
+        """Test prefix_match with a tenant filter when one tenant has a prefix of a longer string."""
+        tree.add_tenants(["tenant_1", "tenant_2"], 0)
+        tree.insert("apple", "tenant_1", 1)
+        tree.insert("applepie", "tenant_2", 2)
+
+        # Match the longer string but only allow tenant_1
+        matched_text, matched_tenants = tree.prefix_match("applepie", ["tenant_1"])
+
+        # Should only match up to "apple" as that's what tenant_1 owns
+        assert matched_text == "apple"
+        assert matched_tenants == ["tenant_1"]
+
+        # Verify that using both tenants would match the full string for tenant_2 only
+        matched_text, matched_tenants = tree.prefix_match(
+            "applepie", ["tenant_1", "tenant_2"]
+        )
+        assert matched_text == "applepie"
+        assert matched_tenants == ["tenant_2"]
+
+        # And both tenants should be returned for "apple"
+        matched_text, matched_tenants = tree.prefix_match(
+            "apple", ["tenant_1", "tenant_2"]
+        )
+        assert matched_text == "apple"
+        assert set(matched_tenants) == {"tenant_1", "tenant_2"}
+
+    def test_prefix_match_with_non_existent_tenant_filter(
+        self, tree: PrefixTree
+    ) -> None:
+        """Test prefix_match with a filter for a non-existent tenant returns no match."""
+        tree.add_tenants(["tenant_1"], 0)
+        tree.insert("apple", "tenant_1", 1)
+        matched_text, matched_tenants = tree.prefix_match(
+            "application", ["non_existent_tenant"]
+        )
+        assert matched_text == ""
+        assert matched_tenants is None
+
+
+class TestPrefixTreeRemove:
+    def test_remove_single_leaf_node_pruned(self, tree: PrefixTree) -> None:
+        """Test _remove_tenant_single_node for a leaf node; node should be pruned."""
+        tree.add_tenants(["tenant_1"], 0)
+        tree.insert("hello", "tenant_1", 1)
+        hello_node = tree.root.edge_label_to_child["h"]
+        assert hello_node.tenant_to_last_access_time == {"tenant_1": 1}
+        assert tree.tenant_to_char_count == {"tenant_1": 5}
+        assert tree.root.edge_label_to_child == {"h": hello_node}
+
+        removed_chars = tree._remove_tenant_single_node("tenant_1", hello_node)
+        assert removed_chars == 5
+        assert hello_node.tenant_to_last_access_time == {}
+        assert tree.tenant_to_char_count == {"tenant_1": 0}
+        assert tree.root.edge_label_to_child == {}  # Node pruned
+
+    def test_remove_single_leaf_node_not_pruned(self, tree: PrefixTree) -> None:
+        """Test _remove_tenant_single_node for a leaf node; node should not be pruned."""
+        tree.add_tenants(["tenant_1", "tenant_2"], 0)
+        tree.insert("hello", "tenant_1", 1)
+        tree.insert("hello", "tenant_2", 2)
+        hello_node = tree.root.edge_label_to_child["h"]
+        assert hello_node.tenant_to_last_access_time == {"tenant_1": 1, "tenant_2": 2}
+        assert tree.tenant_to_char_count == {"tenant_1": 5, "tenant_2": 5}
+        assert tree.root.edge_label_to_child == {"h": hello_node}
+
+        removed_chars = tree._remove_tenant_single_node("tenant_1", hello_node)
+        assert removed_chars == 5
+        assert hello_node.tenant_to_last_access_time == {"tenant_2": 2}
+        assert tree.tenant_to_char_count == {"tenant_1": 0, "tenant_2": 5}
+        assert tree.root.edge_label_to_child == {"h": hello_node}  # Node not pruned
+
+    def test_remove_single_node_with_non_existent_tenant(
+        self, tree: PrefixTree
+    ) -> None:
+        """Test _remove_tenant_single_node for a non-existent tenant is a no-op."""
+        tree.add_tenants(["tenant_1"], 0)
+        tree.insert("hello", "tenant_1", 1)
+        hello_node = tree.root.edge_label_to_child["h"]
+        removed_chars = tree._remove_tenant_single_node(
+            "non_existent_tenant", hello_node
+        )
+        assert removed_chars == 0
+
+    def test_remove_single_node_with_non_matching_tenant(
+        self, tree: PrefixTree
+    ) -> None:
+        """Test _remove_tenant_single_node if node doesn't belong to specified tenant is a no-op."""
+        tree.add_tenants(["tenant_1", "tenant_2"], 0)
+        tree.insert("hello", "tenant_1", 1)
+        tree.insert("world", "tenant_2", 2)  # Node for tenant_2
+        hello_node = tree.root.edge_label_to_child["h"]  # Belongs to tenant_1
+        removed_chars = tree._remove_tenant_single_node(
+            "tenant_2", hello_node
+        )  # Try removing tenant_2 from tenant_1's node
+        assert removed_chars == 0
+
+    def test_remove_tenant(self, tree: PrefixTree) -> None:
+        """Test remove_tenant for a tree with multiple tenants only removes the specified tenant."""
+        tree.add_tenants(["tenant_1", "tenant_2"], 0)
+        tree.insert("hello", "tenant_1", 1)
+        tree.insert("foobar", "tenant_1", 2)
+        tree.insert("helloworld", "tenant_2", 3)
+        removed_chars = tree.remove_tenants(["tenant_1"])
+        assert removed_chars == {"tenant_1": 11}
+        hello_node = tree.root.edge_label_to_child["h"]
+        assert hello_node.tenant_to_last_access_time == {"tenant_2": 3}
+        assert tree.tenant_to_char_count == {"tenant_2": 10}
+        assert set(tree.tenant_to_lru_tail.keys()) == {"tenant_2"}
+        tenant_2_lru_texts = get_lru_texts_from_tree(tree, "tenant_2")
+        assert tenant_2_lru_texts == ["", "world", "hello"]
+
+    def test_remove_non_existent_tenant(self, tree: PrefixTree) -> None:
+        """Test remove_tenant for a non-existent tenant returns 0."""
+        tree.add_tenants(["tenant_1"], 0)
+        tree.insert("hello", "tenant_1", 1)
+        removed_chars = tree.remove_tenants(["non_existent_tenant"])
+        assert removed_chars == {"non_existent_tenant": 0}
+
+    def test_remove_tenant_prunes_nodes(self, tree: PrefixTree) -> None:
+        """Test remove_tenant prunes nodes that become tenant-less and childless."""
+        tree.add_tenants(["tenant_1", "tenant_2"], 0)
+        tree.insert("helloworld", "tenant_1", 1)  # Creates "helloworld"
+        tree.insert(
+            "hellothere", "tenant_2", 2
+        )  # Splits into "hello" -> "world" and "hello" -> "there"
+
+        tree.remove_tenants(["tenant_1"])
+
+        # "world" node should be pruned. "hello" and "there" remain for tenant_2.
+        hello_node = tree.root.edge_label_to_child["h"]
+        assert set(hello_node.edge_label_to_child.keys()) == {"t"}
+        assert hello_node.edge_label_to_child["t"].text == "there"
+        assert hello_node.edge_label_to_child["t"].tenant_to_last_access_time == {
+            "tenant_2": 2
+        }
+
+    def test_remove_tenants(self, tree: PrefixTree) -> None:
+        """Test remove_tenants for multiple tenants with different structures."""
+        tree.add_tenants(["tenant_1", "tenant_2", "tenant_3"], 0)
+        tree.insert("hello", "tenant_1", 1)  # 5 chars
+        tree.insert("foobar", "tenant_1", 2)  # 6 chars
+        tree.insert("helloworld", "tenant_2", 3)  # 10 chars
+        tree.insert("test", "tenant_3", 4)  # 4 chars
+
+        removed_chars = tree.remove_tenants(["tenant_1", "tenant_3"])
+
+        # Check return value contains correct char counts
+        assert removed_chars == {"tenant_1": 11, "tenant_3": 4}
+
+        # Check tree state is correct
+        assert "tenant_1" not in tree.tenant_to_char_count
+        assert "tenant_3" not in tree.tenant_to_char_count
+        assert "tenant_2" in tree.tenant_to_char_count
+        assert tree.tenant_to_char_count == {"tenant_2": 10}
+
+        # Check nodes are correctly maintained
+        assert (
+            "h" in tree.root.edge_label_to_child
+        )  # hello node still exists for tenant_2
+        assert "t" not in tree.root.edge_label_to_child  # test node removed
+        assert "f" not in tree.root.edge_label_to_child  # foobar node removed
+
+        # Check LRU structure
+        assert set(tree.tenant_to_lru_tail.keys()) == {"tenant_2"}
+        tenant_2_lru_texts = get_lru_texts_from_tree(tree, "tenant_2")
+        assert tenant_2_lru_texts == ["", "world", "hello"]
+
+    def test_remove_tenants_with_nonexistent(self, tree: PrefixTree) -> None:
+        """Test remove_tenants with a mix of existing and non-existent tenants."""
+        tree.add_tenants(["tenant_1", "tenant_2"], 0)
+        tree.insert("hello", "tenant_1", 1)
+        tree.insert("world", "tenant_2", 2)
+
+        removed_chars = tree.remove_tenants(["tenant_1", "nonexistent", "alsonotfound"])
+
+        # Check return value
+        assert removed_chars == {"tenant_1": 5, "nonexistent": 0, "alsonotfound": 0}
+
+        # Check tree state
+        assert "tenant_1" not in tree.tenant_to_char_count
+        assert tree.tenant_to_char_count == {"tenant_2": 5}
+        assert "h" not in tree.root.edge_label_to_child  # hello node removed
+        assert "w" in tree.root.edge_label_to_child  # world node still exists
+
+
+class TestPrefixTreeEviction:
+    def test_eviction_non_existent_tenant(self, tree: PrefixTree) -> None:
+        """Test evict_tenant_by_lru for a non-existent tenant returns 0."""
+        assert tree.evict_tenant_by_lru("nonexistent_tenant", 5) == 0
+
+    def test_eviction_exact_min_remove_size_single_node(self, tree: PrefixTree) -> None:
+        """Test evicting exactly min_remove_size characters from a single oldest node."""
+        tree.add_tenants(["tenant_1"], 0)
+        tree.insert("a", "tenant_1", 1)  # Oldest (1 char)
+        tree.insert("bb", "tenant_1", 2)
+        tree.insert("ccc", "tenant_1", 3)
+        assert get_lru_texts_from_tree(tree, "tenant_1") == ["", "ccc", "bb", "a"]
+
+        evicted_count = tree.evict_tenant_by_lru("tenant_1", 1)  # Evict "a"
+        assert evicted_count == 1
+        assert tree.tenant_to_char_count == {"tenant_1": 5}  # 6 - 1
+        assert get_lru_texts_from_tree(tree, "tenant_1") == ["", "ccc", "bb"]
+
+    def test_eviction_exceed_min_remove_size_single_node(
+        self, tree: PrefixTree
+    ) -> None:
+        """Test evicting more than min_remove_size characters from a single oldest node."""
+        tree.add_tenants(["tenant_1"], 0)
+        tree.insert("aaa", "tenant_1", 1)  # Oldest (2 chars)
+        tree.insert("bb", "tenant_1", 2)
+        tree.insert("c", "tenant_1", 3)
+        assert get_lru_texts_from_tree(tree, "tenant_1") == ["", "c", "bb", "aaa"]
+
+        evicted_count = tree.evict_tenant_by_lru("tenant_1", 1)  # Evict "aaa"
+        assert evicted_count == 3
+        assert tree.tenant_to_char_count == {"tenant_1": 3}  # 6 - 3
+        assert get_lru_texts_from_tree(tree, "tenant_1") == ["", "c", "bb"]
+
+    def test_eviction_multiple_nodes(self, tree: PrefixTree) -> None:
+        """Test evicting multiple oldest nodes to meet min_remove_size."""
+        tree.add_tenants(["tenant_1"], 0)
+        tree.insert("a", "tenant_1", 1)  # Oldest (1 char)
+        tree.insert("bb", "tenant_1", 2)  # Next oldest (2 chars)
+        tree.insert("ccc", "tenant_1", 3)
+        assert get_lru_texts_from_tree(tree, "tenant_1") == ["", "ccc", "bb", "a"]
+
+        evicted_count = tree.evict_tenant_by_lru("tenant_1", 2)  # Evict "a" and "b"
+        assert evicted_count == 3  # 1 ("a") + 2 ("b")
+        assert tree.tenant_to_char_count["tenant_1"] == 3  # 6 - 3
+        assert get_lru_texts_from_tree(tree, "tenant_1") == ["", "ccc"]
+
+    def test_eviction_same_timestamps(self, tree: PrefixTree) -> None:
+        """Test evicting more than min_remove_size if multiple nodes share the oldest timestamp."""
+        tree.add_tenants(["tenant_1", "tenant_2"], 0)
+        tree.insert("helloworld", "tenant_1", 1)
+        tree.insert("hellothere", "tenant_2", 2)
+        assert get_lru_texts_from_tree(tree, "tenant_1") == ["", "hello", "world"]
+        assert get_lru_texts_from_tree(tree, "tenant_2") == ["", "there", "hello"]
+
+        # Should remove both "hello" and "world" because they have the same timestamp
+        evicted_count = tree.evict_tenant_by_lru("tenant_1", 1)  # Request 1 char
+        assert evicted_count == 10  # Removes "hello" and "world"
+        assert tree.tenant_to_char_count == {"tenant_1": 0, "tenant_2": 10}
+        assert get_lru_texts_from_tree(tree, "tenant_1") == [""]
+        assert get_lru_texts_from_tree(tree, "tenant_2") == ["", "there", "hello"]
+
+    def test_eviction_insufficient_chars_evicts_all(self, tree: PrefixTree) -> None:
+        """Test evicting when min_remove_size is larger than available; evicts all."""
+        tree.add_tenants(["tenant_1"], 0)
+        tree.insert("xyz", "tenant_1", 1)  # 3 chars available
+        evicted_count = tree.evict_tenant_by_lru("tenant_1", 10)
+        assert evicted_count == 3
+        assert tree.tenant_to_char_count == {"tenant_1": 0}
+        assert get_lru_texts_from_tree(tree, "tenant_1") == [""]
+
+
+class TestPrefixTreeGetSmallestTenants:
+    """Tests for the get_smallest_tenants method."""
+
+    def test_get_smallest_tenants(self, tree: PrefixTree) -> None:
+        """Test get_smallest_tenants identifies the tenant with the fewest characters."""
+        tree.add_tenants(["tenant_1", "tenant_2", "tenant_3"], 0)
+        tree.insert("aaaa", "tenant_1", 1)  # 4 chars
+        tree.insert("bb", "tenant_2", 2)  # 2 chars
+        tree.insert("c", "tenant_3", 3)  # 1 char
+        smallest_tenants = tree.get_smallest_tenants()
+        assert smallest_tenants == ["tenant_3"]
+
+    def test_get_smallest_tenants_empty_tree(self, tree: PrefixTree) -> None:
+        """Test get_smallest_tenants on an empty tree returns None."""
+        assert tree.get_smallest_tenants() is None
+
+    def test_get_smallest_tenants_after_update(self, tree: PrefixTree) -> None:
+        """Test get_smallest_tenants after removing the current smallest tenant."""
+        tree.add_tenants(["tenant_1", "tenant_2", "tenant_3"], 0)
+        tree.insert("aaaa", "tenant_1", 1)
+        tree.insert("bb", "tenant_2", 2)
+        tree.insert("c", "tenant_3", 3)
+        tree.remove_tenants(["tenant_3"])  # Remove "c" (1 char)
+        smallest_tenants = tree.get_smallest_tenants()
+        assert smallest_tenants == ["tenant_2"]  # "bb" (2 chars) is now smallest
+
+    def test_get_smallest_tenants_with_ties(self, tree: PrefixTree) -> None:
+        """Test get_smallest_tenants when multiple tenants have the same minimum count."""
+        tree.add_tenants(["tenant_1", "tenant_2", "tenant_3"], 0)
+        tree.insert("aa", "tenant_1", 1)  # 2 chars
+        tree.insert("bb", "tenant_2", 2)  # 2 chars
+        tree.insert("cccc", "tenant_3", 3)  # 4 chars
+        smallest_tenants = tree.get_smallest_tenants()
+        assert set(smallest_tenants) == {"tenant_1", "tenant_2"}
+
+
+class TestPrefixTreeComprehensive:
+    """Comprehensive tests for the PrefixTree"""
+
+    def test_tree_structure_multiple_insertions(self, tree: PrefixTree) -> None:
+        """Test tree structure after multiple insertions."""
+        tree.add_tenants(["tenant_1", "tenant_2"], 0)
+        tree.insert("helloworld", "tenant_1", 1)
+        tree.insert("hellothere", "tenant_2", 2)
+        tree.insert("hellothomas", "tenant_2", 3)
+
+        # Access tree directly
+        root: Node = tree.root
+
+        # Test tree structure - validate each node
+        # Root node
+        assert root.text == ""
+        assert root.parent is None
+        assert root.tenant_to_last_access_time == {"tenant_1": 1, "tenant_2": 3}
+        assert set(root.edge_label_to_child.keys()) == {"h"}
+
+        # Hello node
+        hello_node: Node = root.edge_label_to_child["h"]
+        assert hello_node.text == "hello"
+        assert hello_node.parent.text == ""
+        assert hello_node.tenant_to_last_access_time == {"tenant_1": 1, "tenant_2": 3}
+        assert set(hello_node.edge_label_to_child.keys()) == {"w", "t"}
+
+        # World node
+        world_node: Node = hello_node.edge_label_to_child["w"]
+        assert world_node.text == "world"
+        assert world_node.parent.text == "hello"
+        assert world_node.tenant_to_last_access_time == {"tenant_1": 1}
+        assert set(world_node.edge_label_to_child.keys()) == set()
+
+        # Th node
+        th_node: Node = hello_node.edge_label_to_child["t"]
+        assert th_node.text == "th"
+        assert th_node.parent.text == "hello"
+        assert th_node.tenant_to_last_access_time == {"tenant_2": 3}
+        assert set(th_node.edge_label_to_child.keys()) == {"e", "o"}
+
+        # Ere node
+        ere_node: Node = th_node.edge_label_to_child["e"]
+        assert ere_node.text == "ere"
+        assert ere_node.parent.text == "th"
+        assert ere_node.tenant_to_last_access_time == {"tenant_2": 2}
+        assert set(ere_node.edge_label_to_child.keys()) == set()
+
+        # Omas node
+        omas_node: Node = th_node.edge_label_to_child["o"]
+        assert omas_node.text == "omas"
+        assert omas_node.parent.text == "th"
+        assert omas_node.tenant_to_last_access_time == {"tenant_2": 3}
+        assert set(omas_node.edge_label_to_child.keys()) == set()
+
+    def test_multiple_evictions_maintains_lru_order(self, tree: PrefixTree) -> None:
+        """Test multiple evictions maintain LRU order."""
+        tree.add_tenants(["tenant_1", "tenant_2"], 0)
+        tree.insert("helloworld", "tenant_1", 1)
+        tree.insert("hellothere", "tenant_2", 2)
+        tree.insert("hellothomas", "tenant_2", 3)
+        assert tree.tenant_to_char_count == {"tenant_1": 10, "tenant_2": 14}
+        assert get_lru_texts_from_tree(tree, "tenant_1") == ["", "hello", "world"]
+        assert get_lru_texts_from_tree(tree, "tenant_2") == [
+            "",
+            "omas",
+            "th",
+            "hello",
+            "ere",
+        ]
+
+        # Eviction 1 (tenant_1): min_remove_size=1. "hello" and "world" removed.
+        evicted_1 = tree.evict_tenant_by_lru("tenant_1", 1)
+        assert evicted_1 == 10
+        assert tree.tenant_to_char_count == {"tenant_1": 0, "tenant_2": 14}
+        assert get_lru_texts_from_tree(tree, "tenant_1") == [""]
+        assert get_lru_texts_from_tree(tree, "tenant_2") == [
+            "",
+            "omas",
+            "th",
+            "hello",
+            "ere",
+        ]  # T2 unchanged
+
+        # Eviction 2 (tenant_2): min_remove_size=1. "ere" is oldest timestamp, removed.
+        evicted_2 = tree.evict_tenant_by_lru("tenant_2", 1)
+        assert evicted_2 == 3  # "ere" is 3 chars
+        assert tree.tenant_to_char_count == {"tenant_1": 0, "tenant_2": 11}  # 14 - 3
+        assert get_lru_texts_from_tree(tree, "tenant_2") == ["", "omas", "th", "hello"]
+
+        # Eviction 3 (tenant_2): min_remove_size=1. "omas"(ts3), "th"(ts3), "hello"(ts3) removed.
+        evicted_3 = tree.evict_tenant_by_lru("tenant_2", 1)
+        assert evicted_3 == 11  # 4+2+5 chars
+        assert tree.tenant_to_char_count == {"tenant_1": 0, "tenant_2": 0}
+        assert get_lru_texts_from_tree(tree, "tenant_2") == [""]
+
+
+@pytest.mark.asyncio
+class TestPrefixTreeActorComprehensive:
+    """Comprehensive tests for the PrefixTreeActor"""
+
+    async def test_tree_structure_multiple_insertions_actor(
+        self, tree_actor: PrefixTreeActor
+    ) -> None:
+        # Add tenants and insert strings in specified order
+        ray.get(tree_actor.add_tenants.remote(["tenant_1", "tenant_2"], 0))
+        ray.get(tree_actor.insert.remote("helloworld", "tenant_1", 1))
+        ray.get(tree_actor.insert.remote("hellothere", "tenant_2", 2))
+        ray.get(tree_actor.insert.remote("hellothomas", "tenant_2", 3))
+        assert await get_lru_texts_from_tree_actor(tree_actor, "tenant_1") == [
+            "",
+            "hello",
+            "world",
+        ]
+
+        # Access tree directly
+        root: Node = ray.get(tree_actor.getattr.remote("root"))
+
+        # Test tree structure - validate each node
+        # Root node
+        assert root.text == ""
+        assert root.parent is None
+        assert root.tenant_to_last_access_time == {"tenant_1": 1, "tenant_2": 3}
+        assert set(root.edge_label_to_child.keys()) == {"h"}
+
+        # Hello node
+        hello_node: Node = root.edge_label_to_child["h"]
+        assert hello_node.text == "hello"
+        assert hello_node.parent.text == ""
+        assert hello_node.tenant_to_last_access_time == {"tenant_1": 1, "tenant_2": 3}
+        assert set(hello_node.edge_label_to_child.keys()) == {"w", "t"}
+
+        # World node
+        world_node: Node = hello_node.edge_label_to_child["w"]
+        assert world_node.text == "world"
+        assert world_node.parent.text == "hello"
+        assert world_node.tenant_to_last_access_time == {"tenant_1": 1}
+        assert set(world_node.edge_label_to_child.keys()) == set()
+
+        # Th node
+        th_node: Node = hello_node.edge_label_to_child["t"]
+        assert th_node.text == "th"
+        assert th_node.parent.text == "hello"
+        assert th_node.tenant_to_last_access_time == {"tenant_2": 3}
+        assert set(th_node.edge_label_to_child.keys()) == {"e", "o"}
+
+        # Ere node
+        ere_node: Node = th_node.edge_label_to_child["e"]
+        assert ere_node.text == "ere"
+        assert ere_node.parent.text == "th"
+        assert ere_node.tenant_to_last_access_time == {"tenant_2": 2}
+        assert set(ere_node.edge_label_to_child.keys()) == set()
+
+        # Omas node
+        omas_node: Node = th_node.edge_label_to_child["o"]
+        assert omas_node.text == "omas"
+        assert omas_node.parent.text == "th"
+        assert omas_node.tenant_to_last_access_time == {"tenant_2": 3}
+        assert set(omas_node.edge_label_to_child.keys()) == set()
+
+    async def test_multiple_evictions_maintains_lru_order_actor(
+        self, tree_actor: PrefixTreeActor
+    ) -> None:
+        """Test multiple evictions maintain LRU order."""
+        # Add tenants and insert test data
+        ray.get(tree_actor.add_tenants.remote(["tenant_1", "tenant_2"], 0))
+        ray.get(tree_actor.insert.remote("helloworld", "tenant_1", 1))
+        ray.get(tree_actor.insert.remote("hellothere", "tenant_2", 2))
+        ray.get(tree_actor.insert.remote("hellothomas", "tenant_2", 3))
+        assert ray.get(tree_actor.getattr.remote("tenant_to_char_count")) == {
+            "tenant_1": 10,
+            "tenant_2": 14,
+        }
+        assert await get_lru_texts_from_tree_actor(tree_actor, "tenant_1") == [
+            "",
+            "hello",
+            "world",
+        ]
+        assert await get_lru_texts_from_tree_actor(tree_actor, "tenant_2") == [
+            "",
+            "omas",
+            "th",
+            "hello",
+            "ere",
+        ]
+
+        # Eviction 1 (tenant_1): min_remove_size=1. "hello" and "world" removed.
+        evicted_1 = ray.get(tree_actor.evict_tenant_by_lru.remote("tenant_1", 1))
+        assert evicted_1 == 10
+        assert ray.get(tree_actor.getattr.remote("tenant_to_char_count")) == {
+            "tenant_1": 0,
+            "tenant_2": 14,
+        }
+        assert await get_lru_texts_from_tree_actor(tree_actor, "tenant_1") == [""]
+        assert await get_lru_texts_from_tree_actor(tree_actor, "tenant_2") == [
+            "",
+            "omas",
+            "th",
+            "hello",
+            "ere",
+        ]  # T2 unchanged
+
+        # Eviction 2 (tenant_2): min_remove_size=1. "ere" is oldest timestamp, removed.
+        evicted_2 = ray.get(tree_actor.evict_tenant_by_lru.remote("tenant_2", 1))
+        assert evicted_2 == 3  # "ere" is 3 chars
+        assert ray.get(tree_actor.getattr.remote("tenant_to_char_count")) == {
+            "tenant_1": 0,
+            "tenant_2": 11,
+        }  # 14 - 3
+        assert await get_lru_texts_from_tree_actor(tree_actor, "tenant_2") == [
+            "",
+            "omas",
+            "th",
+            "hello",
+        ]
+
+        # Eviction 3 (tenant_2): min_remove_size=1. "omas"(ts3), "th"(ts3), "hello"(ts3) removed.
+        evicted_3 = ray.get(tree_actor.evict_tenant_by_lru.remote("tenant_2", 1))
+        assert evicted_3 == 11  # 4+2+5 chars
+        assert ray.get(tree_actor.getattr.remote("tenant_to_char_count")) == {
+            "tenant_1": 0,
+            "tenant_2": 0,
+        }
+        assert await get_lru_texts_from_tree_actor(tree_actor, "tenant_2") == [""]
+
+
+@pytest.mark.asyncio
+class TestPrefixTreeActorEvictionLoop:
+    """Tests for the automatic eviction loop in PrefixTreeActor"""
+
+    async def test_eviction_loop_triggers_automatically(
+        self, tree_actor: PrefixTreeActor
+    ) -> None:
+        """Test that the eviction loop automatically evicts data when threshold is exceeded."""
+        # Set up eviction parameters
+        eviction_threshold = 10  # Low threshold for testing
+        eviction_target = 8  # Target to evict down to
+        interval_secs = 0.1  # Short interval for testing
+
+        # Start the eviction loop
+        ray.get(
+            tree_actor.start_eviction_loop.remote(
+                eviction_threshold, eviction_target, interval_secs
+            )
+        )
+
+        # Add tenant and insert data over the threshold
+        ray.get(tree_actor.add_tenants.remote(["tenant_1"], 0))
+        ray.get(tree_actor.insert.remote("hello", "tenant_1", 1))  # 5 chars
+        ray.get(
+            tree_actor.insert.remote("excess", "tenant_1", 2)
+        )  # 6 more chars, total: 11
+
+        # Verify initial count
+        assert ray.get(tree_actor.getattr.remote("tenant_to_char_count")) == {
+            "tenant_1": 11
+        }
+
+        # Wait for eviction loop to run (interval + small buffer)
+        await asyncio.sleep(interval_secs + 0.2)
+
+        # Verify data was automatically evicted down to target (8 chars)
+        # The eviction should have removed 5 chars, so we should be at 6, which is <= 8
+        char_count = ray.get(tree_actor.getattr.remote("tenant_to_char_count"))
+        assert char_count["tenant_1"] == 6
+
+    async def test_eviction_loop_multiple_tenants(
+        self, tree_actor: PrefixTreeActor
+    ) -> None:
+        """Test that eviction loop evicts from each tenant that exceeds the threshold."""
+        # Set up eviction parameters
+        eviction_threshold = 10
+        eviction_target = 8
+        interval_secs = 0.1
+
+        # Start the eviction loop
+        ray.get(
+            tree_actor.start_eviction_loop.remote(
+                eviction_threshold, eviction_target, interval_secs
+            )
+        )
+
+        # Add two tenants with data over threshold
+        ray.get(tree_actor.add_tenants.remote(["tenant_1", "tenant_2"], 0))
+        ray.get(tree_actor.insert.remote("hello", "tenant_1", 1))  # 5 chars
+        ray.get(
+            tree_actor.insert.remote("excess", "tenant_1", 2)
+        )  # 6 more chars, total: 11
+        ray.get(tree_actor.insert.remote("bigstring", "tenant_2", 3))  # 9 chars
+        ray.get(
+            tree_actor.insert.remote("more", "tenant_2", 4)
+        )  # 4 more chars, total: 13
+
+        # Verify initial counts
+        initial_count = ray.get(tree_actor.getattr.remote("tenant_to_char_count"))
+        assert initial_count["tenant_1"] == 11
+        assert initial_count["tenant_2"] == 13
+
+        # Wait for eviction loop to run
+        await asyncio.sleep(interval_secs + 0.2)
+
+        # Verify both tenants were evicted to target
+        char_count = ray.get(tree_actor.getattr.remote("tenant_to_char_count"))
+
+        # Tenant 1 should have "hello" evicted, so 11 - 5 = 6
+        assert char_count["tenant_1"] == 6
+        # Tenant 2 should have "bigstring" evicted, so 13 - 9 = 4
+        assert char_count["tenant_2"] == 4
+
+    async def test_eviction_loop_respects_threshold(
+        self, tree_actor: PrefixTreeActor
+    ) -> None:
+        """Test that eviction loop only evicts tenants that exceed the threshold."""
+        # Set up eviction parameters
+        eviction_threshold = 10
+        eviction_target = 8
+        interval_secs = 0.1
+
+        # Start the eviction loop
+        ray.get(
+            tree_actor.start_eviction_loop.remote(
+                eviction_threshold, eviction_target, interval_secs
+            )
+        )
+
+        # Add two tenants - one over threshold, one under
+        ray.get(tree_actor.add_tenants.remote(["over_tenant", "under_tenant"], 0))
+        ray.get(tree_actor.insert.remote("hello", "over_tenant", 1))  # 5 chars
+        ray.get(
+            tree_actor.insert.remote("excess", "over_tenant", 2)
+        )  # 6 more chars, total: 11
+        ray.get(tree_actor.insert.remote("small", "under_tenant", 3))  # 5 chars
+
+        # Verify initial counts
+        initial_count = ray.get(tree_actor.getattr.remote("tenant_to_char_count"))
+        assert initial_count["over_tenant"] == 11
+        assert initial_count["under_tenant"] == 5
+
+        # Wait for eviction loop to run
+        await asyncio.sleep(interval_secs + 0.2)
+
+        # Verify only the tenant over threshold was evicted
+        char_count = ray.get(tree_actor.getattr.remote("tenant_to_char_count"))
+        # Tenant 1 should have "hello" evicted, so 11 - 5 = 6
+        assert char_count["over_tenant"] == 6
+        # Tenant 2 should be unchanged
+        assert char_count["under_tenant"] == 5
+
+    async def test_eviction_loop_can_be_started_multiple_times(
+        self, tree_actor: PrefixTreeActor
+    ) -> None:
+        """Test that only the first call to start_eviction_loop starts a new loop."""
+        # Call start_eviction_loop multiple times
+        eviction_task_1 = ray.get(tree_actor.start_eviction_loop.remote(10, 8, 0.1))
+        eviction_task_2 = ray.get(tree_actor.start_eviction_loop.remote(10, 0, 0.1))
+        assert eviction_task_1 and not eviction_task_2
+
+        # Add tenant and insert data over the threshold
+        ray.get(tree_actor.add_tenants.remote(["tenant_1"], 0))
+        ray.get(tree_actor.insert.remote("hello", "tenant_1", 1))  # 5 chars
+        ray.get(
+            tree_actor.insert.remote("excess", "tenant_1", 2)
+        )  # 6 more chars, total: 11
+
+        # Wait for eviction loop to run
+        await asyncio.sleep(0.3)
+
+        # Verify the first eviction_target_chars is respected.
+        # Should evict "hello" to bring the char count down from 11 to 6.
+
+        char_count = ray.get(tree_actor.getattr.remote("tenant_to_char_count"))
+        assert char_count["tenant_1"] == 6
+
+
+if __name__ == "__main__":
+    import sys
+
+    exit_code = pytest.main(["-v", __file__])
+    sys.exit(exit_code)

python/ray/llm/tests/serve/cpu/deployments/utils/direct_streaming_utils.py

@@ -0,0 +1,68 @@
+"""
+Shared helpers for direct-streaming session-affinity tests.
+"""
+
+import httpx
+import pytest
+
+from ray import serve
+from ray._common.test_utils import wait_for_condition
+from ray.llm._internal.serve.constants import RAY_SERVE_LLM_ENABLE_DIRECT_STREAMING
+from ray.serve._private.constants import RAY_SERVE_ENABLE_HA_PROXY, SERVE_SESSION_ID
+from ray.serve._private.test_utils import check_running, get_application_url
+from ray.serve.config import RequestRouterConfig
+
+CONSISTENT_HASH_ROUTER = (
+    "ray.serve.experimental.consistent_hash_router:ConsistentHashRouter"
+)
+
+# Skip unless the direct-streaming + HAProxy env is set
+requires_direct_streaming = pytest.mark.skipif(
+    not (RAY_SERVE_ENABLE_HA_PROXY and RAY_SERVE_LLM_ENABLE_DIRECT_STREAMING),
+    reason="Direct streaming requires RAY_SERVE_ENABLE_HA_PROXY=1 and "
+    "RAY_SERVE_LLM_ENABLE_DIRECT_STREAMING=1.",
+)
+
+
+def consistent_hash_deployment_config() -> dict:
+    return {
+        "num_replicas": 4,
+        "ray_actor_options": {"num_cpus": 0.1},
+        "request_router_config": RequestRouterConfig(
+            request_router_class=CONSISTENT_HASH_ROUTER,
+            request_router_kwargs={
+                "num_virtual_nodes": 100,
+                "num_fallback_replicas": 2,
+            },
+        ),
+    }
+
+
+def run_app_through_haproxy(app, timeout_s: int = 60) -> str:
+    """Run ``app`` and return its (HAProxy) URL once all replicas are RUNNING."""
+    serve.run(app)
+    wait_for_condition(check_running, timeout=timeout_s)
+    return get_application_url(use_localhost=True)
+
+
+def session_chat_response(base_url: str, session_id: str, model: str = "test-model"):
+    """POST a one-token chat request carrying ``session_id`` through HAProxy.
+
+    Asserts the request succeeded and the session id survived the HAProxy hop to
+    the serving replica. Returns the response so callers can read the serving
+    replica from the ``x-replica-id`` header (and, for P/D, the prefill replica
+    from ``kv_transfer_params.remote_engine_id``).
+    """
+    resp = httpx.post(
+        f"{base_url}/v1/chat/completions",
+        json={
+            "model": model,
+            "messages": [{"role": "user", "content": "hi"}],
+            "max_tokens": 1,
+        },
+        headers={SERVE_SESSION_ID: session_id},
+        timeout=30,
+    )
+    assert resp.status_code == 200, resp.text
+    assert resp.headers["x-serve-session-id"] == session_id
+    return resp

python/ray/llm/tests/serve/cpu/deployments/utils/test_batcher.py

@@ -0,0 +1,236 @@
+import asyncio
+import sys
+import time
+from typing import List, Optional
+
+import numpy as np
+import pytest
+
+from ray.llm._internal.serve.constants import MODEL_RESPONSE_BATCH_TIMEOUT_MS
+from ray.llm._internal.serve.utils.batcher import Batcher
+
+TEXT_VALUE = "foo"
+FINAL_TEXT_VALUE = "bar"
+
+
+async def fake_generator():
+    """Returns 100 responses with no delay"""
+    for _i in range(100):
+        yield dict(num_generated_tokens=1, generated_text=TEXT_VALUE)
+
+
+async def fake_generator_slow(num_batches: int):
+    """Returns 100 responses with small delay.
+
+    Delay is set such that the responses are batched into roughly num_batches
+    batches.
+    """
+
+    for _i in range(100):
+        await asyncio.sleep(MODEL_RESPONSE_BATCH_TIMEOUT_MS / 1000 / num_batches)
+        yield dict(num_generated_tokens=1, generated_text=TEXT_VALUE)
+
+
+async def fake_generator_slow_last_return_immediate():
+    """Returns 11 responses with small delay, aside from the last one which is immediate"""
+    for _i in range(10):
+        await asyncio.sleep(MODEL_RESPONSE_BATCH_TIMEOUT_MS / 1000)
+        yield dict(num_generated_tokens=1, generated_text=TEXT_VALUE)
+    yield dict(num_generated_tokens=1, generated_text=FINAL_TEXT_VALUE)
+
+
+async def count_interval_ms_from_stream(stream) -> list[float]:
+    output_intervals: list[float] = []
+    start = None
+    async for _ in stream:
+        if start is None:
+            start = time.perf_counter()
+        else:
+            end = time.perf_counter()
+            output_intervals.append((end - start) * 1e3)
+            start = end
+    return output_intervals
+
+
+class TestBatcher(Batcher):
+    def _merge_results(self, results: List[dict]) -> dict:
+        merged_result = {"num_generated_tokens": 0, "generated_text": ""}
+        for result in results:
+            for key, value in result.items():
+                merged_result[key] += value
+        return merged_result
+
+
+class TestBatching:
+    @pytest.mark.asyncio
+    async def test_batch(self):
+        count = 0
+        batcher = TestBatcher(fake_generator())
+        async for x in batcher.stream():
+            count += 1
+            assert x["num_generated_tokens"] == 100
+            assert x["generated_text"] == TEXT_VALUE * 100
+
+        # Should only have been called once
+        assert count == 1
+        assert batcher.queue.empty()
+
+    @pytest.mark.asyncio
+    async def test_batch_timing(self):
+        count = 0
+        batcher = TestBatcher(fake_generator_slow(num_batches=10))
+        async for _x in batcher.stream():
+            count += 1
+
+        assert 9 <= count <= 12, (
+            "Count should have been called between 9 and 12 times, "
+            "because each iteration takes 1/10th of an interval to yield."
+        )
+        assert batcher.queue.empty()
+
+    @pytest.mark.asyncio
+    async def test_batch_last_return_is_immediate(self):
+        """Test that we don't wait the entire interval for
+        the last response if it returns quickly."""
+        count = 0
+        token_count = 0
+        batcher = TestBatcher(fake_generator_slow_last_return_immediate())
+        last_response = None
+        async for _x in batcher.stream():
+            count += 1
+            token_count += _x["num_generated_tokens"]
+            last_response = _x
+
+        assert (
+            last_response["generated_text"] == TEXT_VALUE + FINAL_TEXT_VALUE
+        ), "the last generated response should be batched with previous one"
+        assert token_count == 11, "token_count should be exactly 11"
+        assert (
+            count == 10
+        ), "Count should have been called exactly 10 times (as many as we generated - 1)"
+        assert batcher.queue.empty()
+
+    @pytest.mark.asyncio
+    async def test_batch_no_interval(self):
+        """Check that the class creates only one batch if there's no interval."""
+
+        batcher = TestBatcher(fake_generator_slow(num_batches=10), interval_ms=None)
+
+        count = 0
+        async for _x in batcher.stream():
+            count += 1
+
+        assert count == 1
+        assert batcher.queue.empty()
+
+    @pytest.mark.asyncio
+    @pytest.mark.parametrize("interval_ms", [100, None])
+    async def test_exception_propagation(self, interval_ms: Optional[float]):
+        """Test that exceptions are propagated correctly to parent."""
+
+        async def generator_should_raise():
+            for _i in range(100):
+                await asyncio.sleep(0.01)
+                yield dict(num_generated_tokens=1, generated_text=TEXT_VALUE)
+                raise ValueError()
+
+        count = 0
+        batched = TestBatcher(generator_should_raise(), interval_ms=interval_ms)
+
+        async def parent():
+            nonlocal count
+            nonlocal batched
+            async for _x in batched.stream():
+                count += 1
+
+        task = asyncio.create_task(parent())
+        await asyncio.sleep(0.2)
+
+        with pytest.raises(ValueError):
+            task.result()
+        assert count == 1
+
+    @pytest.mark.asyncio
+    @pytest.mark.parametrize("interval_ms", [100, None])
+    @pytest.mark.parametrize("to_cancel", ["parent", "inner", "stream"])
+    async def test_cancellation(self, interval_ms: Optional[float], to_cancel: str):
+        """There are 3 ways cancellation can happen:
+        1. The parent is cancelled
+        2. The generator is cancelled
+        3. The stream task is directly cancelled.
+
+        Make sure all associated tasks are cancelled in each instance.
+        """
+
+        async def generator_should_raise():
+            with pytest.raises(asyncio.CancelledError):
+                for _i in range(100):
+                    await asyncio.sleep(0.01)
+                    yield dict(num_generated_tokens=1, generated_text=TEXT_VALUE)
+                    if to_cancel == "inner":
+                        raise asyncio.CancelledError()
+
+        batched = TestBatcher(generator_should_raise(), interval_ms=interval_ms)
+
+        async def parent():
+            nonlocal batched
+            async for _x in batched.stream():
+                pass
+
+        task = asyncio.create_task(parent())
+        await asyncio.sleep(0.2)
+
+        cancel_task = {
+            "parent": task,
+            "stream": batched.read_task,
+        }.get(to_cancel)
+
+        if cancel_task:
+            assert not task.done()
+            assert not batched.read_task.done()
+            cancel_task.cancel()
+
+        await asyncio.sleep(0.3)
+        assert batched.read_task.done(), "Read task should be completed"
+        assert task.done(), "All tasks should be done"
+
+        # Inner task is checked automatically with pytest.raises
+
+    @pytest.mark.asyncio
+    async def test_stable_streaming(self):
+        """Test that the batcher does not add jitter to the stream when interval_ms is 0"""
+
+        async def generator():
+            for i in range(100):
+                await asyncio.sleep(0.01)
+                yield i
+
+        concurrency = 10
+
+        output_intervals = await asyncio.gather(
+            *[
+                count_interval_ms_from_stream(
+                    Batcher(generator(), interval_ms=0).stream()
+                )
+                for _ in range(concurrency)
+            ]
+        )
+        mean_batcher_interval = np.mean(output_intervals)
+        std_batcher_interval = np.std(output_intervals)
+
+        generator_intervals = await asyncio.gather(
+            *[count_interval_ms_from_stream(generator()) for _ in range(concurrency)]
+        )
+        mean_generator_interval = np.mean(generator_intervals)
+        std_generator_interval = np.std(generator_intervals)
+
+        assert np.isclose(
+            mean_batcher_interval, mean_generator_interval, rtol=0.1
+        ), f"{mean_batcher_interval=}, {mean_generator_interval=}"
+        assert np.isclose(
+            std_batcher_interval, std_generator_interval, atol=0.1
+        ), f"{std_batcher_interval=}, {std_generator_interval=}"
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main(["-v", __file__]))

python/ray/llm/tests/serve/cpu/deployments/utils/test_broadcast.py

@@ -0,0 +1,163 @@
+import sys
+import time
+
+import pytest
+
+import ray
+from ray import serve
+from ray.llm._internal.serve.utils.broadcast import broadcast
+
+
+# Define a simple deployment for testing
+@serve.deployment(num_replicas=2)
+class MockLLMDeployment:
+    def __init__(self):
+        self.reset_count = 0
+        self.id = id(self)
+
+    async def reset_prefix_cache(self):
+        self.reset_count += 1
+        return self.id, self.reset_count
+
+    async def get_reset_count(self):
+        return self.id, self.reset_count
+
+    async def echo(self, msg, repeat=1):
+        return f"{self.id}:{msg * repeat}"
+
+    async def self_destruct(self):
+        """Kill this replica's actor. Used for testing dead replica handling."""
+        import os
+
+        os._exit(1)
+
+
+@pytest.fixture(scope="module")
+def serve_instance():
+    # Start ray and serve once for the module
+    if not ray.is_initialized():
+        ray.init()
+    yield
+    serve.shutdown()
+    ray.shutdown()
+
+
+@pytest.fixture
+def mock_handle(serve_instance, request):
+    # Ensure deployment is up and running.
+    # serve.run waits for the deployment to be ready by default unless _blocking=False.
+    app_name = f"mock-llm-{request.node.name}"
+    route_prefix = f"/{app_name}"
+    handle = serve.run(
+        MockLLMDeployment.bind(), name=app_name, route_prefix=route_prefix
+    )
+    yield handle
+    serve.delete(app_name, _blocking=True)
+
+
+@pytest.mark.asyncio
+async def test_dispatch_basic(mock_handle):
+    """Test basic dispatch without combine."""
+    # We can use get_reset_count which doesn't modify state
+    results = broadcast(mock_handle, "get_reset_count")
+
+    assert len(results) == 2
+    # Verify we got unique IDs back
+    ids = {r[0] for r in results}
+    assert len(ids) == 2
+
+
+@pytest.mark.asyncio
+async def test_dispatch_with_combine(mock_handle):
+    """Test dispatch with a combine function."""
+    # First, increment count so we have something to sum
+    broadcast(mock_handle, "reset_prefix_cache")
+
+    def sum_counts(results):
+        # results is list of (id, count)
+        return sum(r[1] for r in results)
+
+    # Get counts using dispatch and combine
+    total_count = broadcast(mock_handle, "get_reset_count", combine=sum_counts)
+
+    # We have 2 replicas, each should have reset_count=1 after one reset call
+    assert total_count == 2
+    assert isinstance(total_count, int)
+
+
+@pytest.mark.asyncio
+async def test_dispatch_args_kwargs(mock_handle):
+    """Test dispatch passing args and kwargs."""
+    results = broadcast(mock_handle, "echo", args=("hello",), kwargs={"repeat": 2})
+
+    assert len(results) == 2
+    for r in results:
+        # Format is "id:msg"
+        msg_part = r.split(":")[1]
+        assert msg_part == "hellohello"
+
+
+@pytest.mark.asyncio
+async def test_dispatch_callable_args(mock_handle):
+    """Test dispatch with callable args generator."""
+
+    def arg_gen(replica):
+        # replica has unique_id or similar
+        return (f"msg-{replica.unique_id}",)
+
+    results = broadcast(mock_handle, "echo", args=arg_gen)
+
+    assert len(results) == 2
+    msgs = set()
+    for r in results:
+        msg_part = r.split(":")[1]
+        msgs.add(msg_part)
+
+    assert len(msgs) == 2
+    for msg in msgs:
+        assert msg.startswith("msg-")
+
+
+@pytest.mark.asyncio
+async def test_dispatch_handles_dead_replica(serve_instance, request):
+    """Test that dispatch gracefully handles a dead replica.
+
+    This test verifies that if one replica dies, dispatch still completes
+    successfully and returns results from the remaining live replicas.
+    """
+    app_name = f"mock-llm-{request.node.name}"
+    route_prefix = f"/{app_name}"
+
+    # Deploy with 2 replicas
+    handle = serve.run(
+        MockLLMDeployment.bind(), name=app_name, route_prefix=route_prefix
+    )
+
+    # First, verify dispatch works with all replicas alive
+    results_before = broadcast(handle, "get_reset_count")
+    assert len(results_before) == 2, "Should have 2 results from 2 replicas"
+
+    # Kill one replica by calling self_destruct through the handle.
+    # This sends an RPC to one replica which will kill itself.
+    # We use options to not wait for response since the actor will die.
+    try:
+        handle.self_destruct.remote()
+    except Exception:
+        # The call may raise if the actor dies mid-request
+        pass
+
+    # Give Serve a moment to detect the dead replica
+    time.sleep(2)
+
+    # Dispatch should still work with the remaining replica(s)
+    # The dead replica will be skipped (ValueError caught in dispatch)
+    results_after = broadcast(handle, "get_reset_count")
+
+    # Should get at least 1 result from the surviving replica
+    # (The killed replica may or may not be in the replica set depending
+    # on timing of Serve's failure detection)
+    assert len(results_after) >= 1, "Should have at least 1 result from live replica"
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main(["-v", __file__]))

python/ray/llm/tests/serve/cpu/observability/test_metrics_middleware.py

@@ -0,0 +1,64 @@
+"""Tests for the LLM Serve metrics middleware route resolution."""
+import sys
+
+import pytest
+from fastapi import APIRouter, FastAPI
+
+from ray.llm._internal.serve.observability.metrics.middleware import (
+    _get_route_details,
+)
+
+
+def _scope(
+    path: str, method: str = "GET", scope_type: str = "http", root_path: str = ""
+):
+    return {
+        "type": scope_type,
+        "method": method,
+        "path": path,
+        "headers": [],
+        "app": None,  # set by caller
+        "path_params": {},
+        "root_path": root_path,
+    }
+
+
+def test_get_route_details_include_router():
+    """Routes added via `include_router` must resolve without crashing (#64245).
+
+    On FastAPI >= 0.137 these routes are nested under an `_IncludedRouter` node
+    that has no `.path` attribute; accessing it previously raised
+    ``AttributeError: '_IncludedRouter' object has no attribute 'path'``.
+    """
+    app = FastAPI()
+
+    @app.get("/direct")
+    def direct():
+        return {}
+
+    router = APIRouter(prefix="/api")
+
+    @router.get("/items/{item_id}")
+    def get_item(item_id: str):
+        return item_id
+
+    app.include_router(router)
+
+    # Directly decorated route.
+    scope = _scope("/direct")
+    scope["app"] = app
+    assert _get_route_details(scope) == "/direct"
+
+    # Route registered via include_router (the #64245 regression).
+    scope = _scope("/api/items/123")
+    scope["app"] = app
+    assert _get_route_details(scope) == "/api/items/{item_id}"
+
+    # Unmatched path resolves to None (unchanged behavior).
+    scope = _scope("/does-not-exist")
+    scope["app"] = app
+    assert _get_route_details(scope) is None
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main(["-v", "-s", __file__]))

python/ray/llm/tests/serve/cpu/observability/usage_telemetry/test_usage.py

@@ -0,0 +1,301 @@
+import sys
+
+import pytest
+
+import ray
+from ray._common.usage.usage_lib import TagKey
+from ray.llm._internal.serve.core.configs.llm_config import (
+    LLMConfig,
+    LLMEngine,
+    LoraConfig,
+    ModelLoadingConfig,
+)
+from ray.llm._internal.serve.observability.usage_telemetry.usage import (
+    HardwareUsage,
+    _get_or_create_telemetry_agent,
+    _retry_get_telemetry_agent,
+    push_telemetry_report_for_all_models,
+)
+
+
+@ray.remote(num_cpus=0)
+class TelemetryRecorder:
+    def __init__(self):
+        self._telemetry = {}
+
+    def record(self, key, value):
+        self._telemetry[key] = value
+
+    def telemetry(self):
+        return self._telemetry
+
+
+def test_push_telemetry_report_for_all_models(disable_placement_bundles):
+    recorder = TelemetryRecorder.remote()
+
+    def record_tag_func(key, value):
+        ray.get(recorder.record.remote(key, value))
+
+    telemetry_agent = _get_or_create_telemetry_agent()
+    telemetry_agent._reset_models.remote()
+    telemetry_agent._update_record_tag_func.remote(record_tag_func)
+
+    dynamic_lora_loading_path = "s3://fake_bucket/fake_path"
+    llm_config_model = LLMConfig(
+        model_loading_config=ModelLoadingConfig(
+            model_id="llm_model_id",
+        ),
+        llm_engine=LLMEngine.vLLM,
+        accelerator_type="L4",
+    )
+    llm_config_model._set_model_architecture(model_architecture="llm_model_arch")
+    llm_config_autoscale_model = LLMConfig(
+        model_loading_config=ModelLoadingConfig(
+            model_id="llm_config_autoscale_model_id",
+        ),
+        llm_engine=LLMEngine.vLLM,
+        accelerator_type="A10G",
+        deployment_config=dict(
+            autoscaling_config=dict(
+                min_replicas=2,
+                max_replicas=3,
+            ),
+        ),
+    )
+    llm_config_autoscale_model._set_model_architecture(
+        model_architecture="llm_config_autoscale_model_arch"
+    )
+    llm_config_json_mode_model = LLMConfig(
+        model_loading_config=ModelLoadingConfig(
+            model_id="llm_config_json_model_id",
+        ),
+        llm_engine=LLMEngine.vLLM,
+        accelerator_type="A10G",
+    )
+    llm_config_json_mode_model._set_model_architecture(
+        model_architecture="llm_config_json_model_arch"
+    )
+    llm_config_lora_model = LLMConfig(
+        model_loading_config=ModelLoadingConfig(
+            model_id="llm_config_lora_model_id",
+        ),
+        llm_engine=LLMEngine.vLLM,
+        accelerator_type="A10G",
+        lora_config=LoraConfig(dynamic_lora_loading_path=dynamic_lora_loading_path),
+    )
+    llm_config_lora_model._set_model_architecture(
+        model_architecture="llm_config_lora_model_arch"
+    )
+    llm_config_no_accelerator_type = LLMConfig(
+        model_loading_config=ModelLoadingConfig(
+            model_id="llm_config_no_accelerator_type_id",
+        ),
+    )
+    llm_config_no_accelerator_type._set_model_architecture(
+        model_architecture="llm_config_no_accelerator_type_arch"
+    )
+    all_models = [
+        llm_config_model,
+        llm_config_autoscale_model,
+        llm_config_json_mode_model,
+        llm_config_lora_model,
+        llm_config_no_accelerator_type,
+    ]
+
+    def fake_get_lora_model_ids(dynamic_lora_loading_path, base_model_id):
+        return ["lora_model_id_1", "lora_model_id_2"]
+
+    def fake_get_gpu_type(*args, **kwargs):
+        return ["Intel Xeon", "L40S"]
+
+    # Ensure that the telemetry is empty before pushing the reports.
+    telemetry = ray.get(recorder.telemetry.remote())
+    assert telemetry == {}
+    push_telemetry_report_for_all_models(
+        all_models=all_models,
+        get_lora_model_func=fake_get_lora_model_ids,
+        get_hardware_fn=fake_get_gpu_type,
+    )
+
+    # Ensure that the telemetry is correct after pushing the reports.
+    telemetry = ray.get(recorder.telemetry.remote())
+    assert telemetry == {
+        TagKey.LLM_SERVE_SERVE_MULTIPLE_MODELS: "1",
+        TagKey.LLM_SERVE_SERVE_MULTIPLE_APPS: "0",
+        TagKey.LLM_SERVE_LORA_BASE_MODELS: "llm_config_lora_model_arch",
+        TagKey.LLM_SERVE_INITIAL_NUM_LORA_ADAPTERS: "2",
+        TagKey.LLM_SERVE_AUTOSCALING_ENABLED_MODELS: "llm_config_autoscale_model_arch",
+        TagKey.LLM_SERVE_AUTOSCALING_MIN_REPLICAS: "2",
+        TagKey.LLM_SERVE_AUTOSCALING_MAX_REPLICAS: "3",
+        TagKey.LLM_SERVE_TENSOR_PARALLEL_DEGREE: "1,1,1,1,1",
+        TagKey.LLM_SERVE_NUM_REPLICAS: "1,2,1,1,1",
+        TagKey.LLM_SERVE_MODELS: "llm_model_arch,llm_config_autoscale_model_arch,llm_config_json_model_arch,llm_config_lora_model_arch,llm_config_no_accelerator_type_arch",
+        TagKey.LLM_SERVE_GPU_TYPE: "L4,A10G,A10G,A10G,L40S",
+        TagKey.LLM_SERVE_NUM_GPUS: "1,1,1,1,1",
+    }
+
+
+@ray.remote(num_cpus=0)
+class Replica:
+    def wait_for_init(self):
+        """
+        When this method returns, the actor initialization is guaranteed
+        to be complete.
+
+        This is used for synchronization between multiple replicas,
+        increasing the chance for get_telemetry_agent() to be called
+        at the same time.
+        """
+        pass
+
+    def get_telemetry_agent(self):
+        return _retry_get_telemetry_agent()
+
+
+def test_telemetry_race_condition():
+    replicas = [Replica.remote() for _ in range(30)]
+    init_refs = [replica.wait_for_init.remote() for replica in replicas]
+    ray.get(init_refs)
+
+    get_refs = [replica.get_telemetry_agent.remote() for replica in replicas]
+    telemetry_agents = ray.get(get_refs)
+    for telemetry_agent in telemetry_agents:
+        assert telemetry_agent is not None
+    assert len(set(telemetry_agents)) == 1
+
+
+def test_infer_gpu_from_hardware():
+    # Test with a valid GPU type
+    def fake_get_gpu_type(*args, **kwargs):
+        return ["Intel Xeon", "A10G"]
+
+    result = HardwareUsage(fake_get_gpu_type).infer_gpu_from_hardware()
+    assert result == "A10G"
+
+    # Test with an unsupported GPU type
+    def fake_get_gpu_type(*args, **kwargs):
+        return ["Intel Xeon", "G"]
+
+    result = HardwareUsage(fake_get_gpu_type).infer_gpu_from_hardware()
+    assert result == "UNSPECIFIED"
+
+
+def test_telemetry_dedups_replicas_and_restarts(disable_placement_bundles):
+    """The same model reported by many replicas/restarts collapses to one entry."""
+    recorder = TelemetryRecorder.remote()
+
+    def record_tag_func(key, value):
+        ray.get(recorder.record.remote(key, value))
+
+    telemetry_agent = _get_or_create_telemetry_agent()
+    telemetry_agent._reset_models.remote()
+    telemetry_agent._update_record_tag_func.remote(record_tag_func)
+
+    config = LLMConfig(
+        model_loading_config=ModelLoadingConfig(model_id="dup_model_id"),
+        llm_engine=LLMEngine.vLLM,
+        accelerator_type="L4",
+    )
+    config._set_model_architecture(model_architecture="dup_arch")
+
+    # Simulate three replicas (or restarts) of the SAME model reporting.
+    for _ in range(3):
+        push_telemetry_report_for_all_models(
+            all_models=[config],
+            get_hardware_fn=lambda *a, **k: ["L4"],
+        )
+
+    telemetry = ray.get(recorder.telemetry.remote())
+    assert telemetry[TagKey.LLM_SERVE_MODELS] == "dup_arch"
+    assert telemetry[TagKey.LLM_SERVE_NUM_REPLICAS] == "1"
+    assert telemetry[TagKey.LLM_SERVE_GPU_TYPE] == "L4"
+
+
+def test_telemetry_reports_fixed_num_replicas(disable_placement_bundles):
+    """A fixed (non-autoscaling) num_replicas is reported, not hardcoded to 1."""
+    recorder = TelemetryRecorder.remote()
+
+    def record_tag_func(key, value):
+        ray.get(recorder.record.remote(key, value))
+
+    telemetry_agent = _get_or_create_telemetry_agent()
+    telemetry_agent._reset_models.remote()
+    telemetry_agent._update_record_tag_func.remote(record_tag_func)
+
+    config = LLMConfig(
+        model_loading_config=ModelLoadingConfig(model_id="fixed_replicas_model"),
+        llm_engine=LLMEngine.vLLM,
+        accelerator_type="L4",
+        deployment_config=dict(num_replicas=4),
+    )
+    config._set_model_architecture(model_architecture="fixed_arch")
+
+    push_telemetry_report_for_all_models(
+        all_models=[config],
+        get_hardware_fn=lambda *a, **k: ["L4"],
+    )
+
+    telemetry = ray.get(recorder.telemetry.remote())
+    assert telemetry[TagKey.LLM_SERVE_NUM_REPLICAS] == "4"
+
+
+def test_telemetry_reports_zero_num_replicas(disable_placement_bundles):
+    """An explicit num_replicas=0 is reported as 0, not coerced to 1."""
+    recorder = TelemetryRecorder.remote()
+
+    def record_tag_func(key, value):
+        ray.get(recorder.record.remote(key, value))
+
+    telemetry_agent = _get_or_create_telemetry_agent()
+    telemetry_agent._reset_models.remote()
+    telemetry_agent._update_record_tag_func.remote(record_tag_func)
+
+    config = LLMConfig(
+        model_loading_config=ModelLoadingConfig(model_id="zero_replicas_model"),
+        llm_engine=LLMEngine.vLLM,
+        accelerator_type="L4",
+        deployment_config=dict(num_replicas=0),
+    )
+    config._set_model_architecture(model_architecture="zero_arch")
+
+    push_telemetry_report_for_all_models(
+        all_models=[config],
+        get_hardware_fn=lambda *a, **k: ["L4"],
+    )
+
+    telemetry = ray.get(recorder.telemetry.remote())
+    assert telemetry[TagKey.LLM_SERVE_NUM_REPLICAS] == "0"
+
+
+def test_telemetry_reports_auto_num_replicas(disable_placement_bundles):
+    """num_replicas="auto" is reported as autoscaling, not dropped."""
+    recorder = TelemetryRecorder.remote()
+
+    def record_tag_func(key, value):
+        ray.get(recorder.record.remote(key, value))
+
+    telemetry_agent = _get_or_create_telemetry_agent()
+    telemetry_agent._reset_models.remote()
+    telemetry_agent._update_record_tag_func.remote(record_tag_func)
+
+    config = LLMConfig(
+        model_loading_config=ModelLoadingConfig(model_id="auto_replicas_model"),
+        llm_engine=LLMEngine.vLLM,
+        accelerator_type="L4",
+        deployment_config=dict(num_replicas="auto"),
+    )
+    config._set_model_architecture(model_architecture="auto_arch")
+
+    push_telemetry_report_for_all_models(
+        all_models=[config],
+        get_hardware_fn=lambda *a, **k: ["L4"],
+    )
+
+    telemetry = ray.get(recorder.telemetry.remote())
+    # Recorded as autoscaling with an integer replica count (not the string "auto").
+    assert telemetry[TagKey.LLM_SERVE_AUTOSCALING_ENABLED_MODELS] == "auto_arch"
+    assert telemetry[TagKey.LLM_SERVE_NUM_REPLICAS].isdigit()
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main(["-v", __file__]))