TEST: Cover multi-prefill and multi-decode GPU routing

.github/workflows/pd-gpu.yaml

@@ -28,6 +28,8 @@ jobs:
     env:
       NO_WEB_UI: '1'
       XINFERENCE_TEST_PD_GPU: '1'
+      XINFERENCE_TEST_PD_MULTI_GPU: '1'
+      XINFERENCE_ALLOW_MULTI_REPLICA_PER_GPU: '1'
       XINFERENCE_AUTH_ADVANCED: 'false'
       XINFERENCE_ENABLE_VIRTUAL_ENV: '0'
       VLLM_VERSION: ${{ inputs.vllm_version }}
@@ -45,10 +47,12 @@ jobs:
       - name: Verify PD with actual KV transfer
         run: |
           source .venv-pd/bin/activate
-          python -m pytest --timeout=1800 -vv xinference/model/llm/vllm/xavier/test/test_pd_gpu.py xinference/model/llm/vllm/xavier/test/test_v1_cache_layout.py xinference/model/llm/vllm/xavier/test/test_direct_history.py
+          python -m pytest --basetemp=pd-test-logs --timeout=1800 -vv -rP xinference/model/llm/vllm/xavier/test/test_pd_gpu.py xinference/model/llm/vllm/xavier/test/test_pd_multi_gpu.py xinference/model/llm/vllm/xavier/test/test_v1_cache_layout.py xinference/model/llm/vllm/xavier/test/test_direct_history.py
       - name: Upload runtime logs
         if: always()
         uses: actions/upload-artifact@v4
         with:
           name: pd-runtime-logs
-          path: ~/.xinference/logs/
+          path: |
+            ~/.xinference/logs/
+            pd-test-logs/

xinference/core/pd_model.py

@@ -294,6 +294,13 @@ async def _infer(self, method, inputs, *args, **kwargs):
             raise ValueError("PD KV handoff currently requires n=1")
         prefill = self._prefill_policy.schedule()
         decode = self._decode_policy.schedule()
+        logger.debug(
+            "PD route: request=%s prefill=%s decode=%s backend=%s",
+            request_id,
+            prefill.uid,
+            decode.uid,
+            self._transport_backend,
+        )
         prefill_args = list(copy.deepcopy(args))
         if not prefill_args or prefill_args[0] is None:
             prefill_args = [{}] + prefill_args[1:]

xinference/model/llm/vllm/xavier/README.md

@@ -68,6 +68,81 @@ Repeated workload entries exercise warm prefix reuse; distinct prefixes exercise
 
 ## Two-GPU integration test
 
+### Multiple P/D replicas sharing GPUs
+
+Run the opt-in multi-replica test with a small full-attention model:
+
+```bash
+XINFERENCE_ALLOW_MULTI_REPLICA_PER_GPU=1 XINFERENCE_TEST_PD_MULTI_GPU=1 \
+  python -m pytest -v xinference/model/llm/vllm/xavier/test/test_pd_multi_gpu.py
+```
+
+The test covers 2P1D, 1P2D and 2P2D for Xavier and native NIXL. All producers
+share GPU 0 and all decoders share GPU 1, with a 0.35 engine memory budget per
+replica. The existing `XINFERENCE_TEST_PD_MODEL_PATH`, model name and size
+overrides apply. It checks request-specific answers, streaming, eight concurrent
+requests, actual KV loads and Xavier history restoration with engine prefix
+caching disabled. For 2P2D it also removes and re-registers one decoder route:
+equal-length round robins otherwise cover only two of the four P/D pairs.
+This changes route registration, not the decoder process, and is not a crash
+recovery test. `MULTI_PD_RESULT` records observed pairs and request counts.
+
+The manually triggered **PD GPU integration** workflow includes these cases
+and uploads each deployment's isolated evidence log.
+
+This setup validates multi-replica behavior on two GPUs. The colocated engines
+share compute and memory bandwidth, so its performance does not establish
+four-GPU scaling.
+
+### Multi-replica routing measurements (2026-10-05)
+
+A separate prefix-affinity prototype was compared with Xavier round robin on
+two RTX 3090 Ti GPUs (NVLink NV3), with 2P on GPU 0 and 2D on GPU 1. The model
+was Qwen2.5-0.5B-Instruct FP16, using vLLM 0.21.0, xoscar 0.11.1 and NIXL 1.1.
+Each replica used a 0.35 GPU memory fraction, eager execution, disabled engine
+prefix caching and the default 256 MiB Xavier GPU history. All four transfer
+pairs were warmed before measurement; output length was 16 tokens.
+
+| Workload, concurrency 8 | Round-robin requests/s | Affinity requests/s | Round-robin TTFT p95 (ms) | Affinity TTFT p95 (ms) |
+| --- | ---: | ---: | ---: | ---: |
+| 12 prefixes, about 1,840 input tokens | 33.61–34.68 | 33.85–34.08 | 154.3–159.1 | 156.2–157.3 |
+| 12 prefixes, mixed input lengths | 33.71–34.79 | 33.75–34.01 | 173.7–179.1 | 177.8–182.0 |
+| 20 prefixes, GPU history capacity pressure | 32.19–33.00 | 32.32–32.45 | 173.8–176.0 | 175.7–177.0 |
+
+Ranges are the two repetitions of each policy in ABBA order. The first two
+profiles ran 2,400 concurrent requests each per deployment, after 12 cold and
+36 sequential repeated requests. Mixed inputs were approximately 500, 900,
+1,800 and 3,200 tokens and followed the fixed-length profile without clearing
+history. Capacity pressure used 20 cold, 60 repeated and 1,200 concurrent
+requests per deployment. The 20-prefix KV working set is about 430 MiB: larger
+than one producer's GPU history budget, but smaller than the combined budget.
+
+All 24,704 accepted measurement requests succeeded. Across 805 one-second GPU
+process samples, no foreign processes were observed after excluding explicitly
+recorded idle daemons. Interfered pressure cases were discarded and rerun,
+which introduced extra gaps between the valid ABBA cases. Process sampling is
+not a hardware reservation. DEBUG evidence logging affects absolute rates;
+shared-GPU results do not establish four-GPU scaling. These comparisons concern
+Xavier scheduling policies, not Xavier versus native vLLM NIXL performance.
+
+Sequential fixed-length reuse improved TTFT p95 from 107.4–109.6 ms to
+85.3–85.9 ms. Under capacity pressure, affinity reduced CPU history restores
+from 54,340–54,620 blocks to 36,392–42,298 blocks per 1,200 requests, but total
+GPU-plus-CPU restores also fell slightly. CPU read reduction alone is therefore
+not an end-to-end gain. Concurrent throughput and tail latency did not improve
+consistently, so production routing remains round robin.
+
+Under the sustained workloads, median P RPC time was about 33–39 ms, while the
+interval from P completion to the first D response was about 64–67 ms. That
+interval includes RPC, decode queueing, KV load and first execution; it must not
+be labeled pure transfer time. Further optimization requires separating these
+costs and accounting for actual cache residency and work size in scheduling.
+The [experiment source](https://github.com/qinxuye/inference/blob/d7db34a77849198323e5cf2e08dc7d40e22baa7c/benchmark/tests/test_pd_affinity_gpu.py)
+contains the opt-in runner and GPU contention checks. The Xavier README on
+that branch documents reproduction settings.
+
+### One producer and one decoder
+
 Use two free NVIDIA GPUs and the pinned vLLM environment, with NIXL installed. Run from the repository root:
 
 ```bash

xinference/model/llm/vllm/xavier/test/test_pd_gpu.py

@@ -21,11 +21,7 @@ def pd_cluster(monkeypatch, tmp_path, backend):
     import xoscar as xo
 
     from xinference.api.restful_api import run_in_subprocess as start_api
-    from xinference.conftest import (
-        TEST_FILE_LOGGING_CONF,
-        TEST_LOG_FILE_PATH,
-        api_health_check,
-    )
+    from xinference.conftest import TEST_FILE_LOGGING_CONF, api_health_check
     from xinference.deploy.local import health_check, run_in_subprocess
 
     # Real subprocesses are required for per-replica CUDA_VISIBLE_DEVICES.
@@ -34,9 +30,17 @@ def pd_cluster(monkeypatch, tmp_path, backend):
     monkeypatch.setenv("XINFERENCE_AUTH_ADVANCED", "false")
     # vLLM spawns its own engine process, outside xoscar's logging setup.
     # Capture connector logs there as well, including the GPU cache write.
-    logging_config = deepcopy(TEST_FILE_LOGGING_CONF)
+    # Each deployment needs an isolated, non-rotating evidence log. Shared
+    # rotating logs can lose phase offsets or mix concurrent test sessions.
+    cluster_logging = deepcopy(TEST_FILE_LOGGING_CONF)
+    log_path = str(tmp_path / "pd-cluster.log")
+    cluster_logging["handlers"]["file_handler"].update(
+        filename=log_path, maxBytes=0, backupCount=0
+    )
+    logging_config = deepcopy(cluster_logging)
     logging_config["loggers"]["vllm"] = {
         "handlers": ["stream_handler", "file_handler"],
+        # KV transfer evidence assertions require vLLM DEBUG records.
         "level": "DEBUG",
         "propagate": False,
     }
@@ -52,7 +56,7 @@ def pd_cluster(monkeypatch, tmp_path, backend):
     # Exercise NIXL host discovery without changing Xavier's Gloo interface.
     host = "0.0.0.0" if backend == "nixl" else "127.0.0.1"
     address = f"{host}:{xo.utils.get_next_port()}"
-    cluster = run_in_subprocess(address, None, None, deepcopy(TEST_FILE_LOGGING_CONF))
+    cluster = run_in_subprocess(address, None, None, deepcopy(cluster_logging))
     api = None
     try:
         assert health_check(address=address, max_attempts=20, sleep_interval=1)
@@ -62,10 +66,10 @@ def pd_cluster(monkeypatch, tmp_path, backend):
             address,
             host="127.0.0.1",
             port=port,
-            logging_conf=deepcopy(TEST_FILE_LOGGING_CONF),
+            logging_conf=deepcopy(cluster_logging),
         )
         assert api_health_check(endpoint, max_attempts=20, sleep_interval=1)
-        yield endpoint, TEST_LOG_FILE_PATH
+        yield endpoint, log_path
     finally:
         if api is not None:
             api.kill()

xinference/model/llm/vllm/xavier/test/test_pd_multi_gpu.py

@@ -0,0 +1,190 @@
+# Copyright 2022-2026 Xinference Holdings Pte. Ltd
+# Licensed under the Apache License, Version 2.0.
+"""Exercise multiple P/D engines sharing two physical GPUs."""
+import asyncio
+import json
+import os
+import re
+from concurrent.futures import ThreadPoolExecutor
+from pathlib import Path
+
+import pytest
+
+from .test_pd_gpu import pd_cluster  # noqa: F401
+
+pytestmark = pytest.mark.skipif(
+    os.environ.get("XINFERENCE_TEST_PD_MULTI_GPU") != "1",
+    reason="requires two CUDA GPUs with room for two small engines each",
+)
+
+
+@pytest.mark.parametrize("backend", ["xavier", "nixl"])
+@pytest.mark.parametrize("num_p,num_d", [(2, 1), (1, 2), (2, 2)])
+def test_multi_pd_gpu(pd_cluster, backend, num_p, num_d):  # noqa: F811
+    import openai
+    import torch
+    import xoscar as xo
+
+    from xinference.client import Client
+    from xinference.core.pd_model import PDModelActor
+
+    assert torch.cuda.device_count() >= 2
+    endpoint, log_path = pd_cluster
+    client = Client(endpoint)
+    worker = client.get_workers_info()[0]["work-ip"]
+    uid = f"pd-multi-{backend}-{num_p}p{num_d}d"
+    loop = asyncio.new_event_loop()
+    api = openai.OpenAI(base_url=endpoint + "/v1", api_key="unused", timeout=180)
+
+    def logs(offset):
+        with open(log_path, "rb") as log:
+            log.seek(offset)
+            return log.read().decode(errors="replace")
+
+    def chat(index, stream=False):
+        # Different answers detect cross-request cache contamination. Long shared
+        # prefixes exercise history independently of the engine's prefix cache.
+        word = ["orange", "purple", "silver", "yellow"][index % 4]
+        prompt = (
+            "Read the following background, then answer the final question. "
+            + "Water evaporates and condenses in clouds. " * 40
+            + f"\nThe secret word is {word}. What is the secret word? "
+            "Reply with only the secret word."
+        )
+        result = api.chat.completions.create(
+            model=uid,
+            messages=[{"role": "user", "content": prompt}],
+            max_tokens=16,
+            temperature=0,
+            stream=stream,
+            extra_body={"chat_template_kwargs": {"enable_thinking": False}},
+        )
+        if stream:
+            chunks = list(result)
+            text = "".join(
+                chunk.choices[0].delta.content or ""
+                for chunk in chunks
+                if chunk.choices
+            )
+        else:
+            text = result.choices[0].message.content
+        assert word in text.lower(), (word, text)
+        assert not any(
+            other in text.lower()
+            for other in ("orange", "purple", "silver", "yellow")
+            if other != word
+        ), text
+        return text
+
+    def check_transfers(evidence, count):
+        if backend == "xavier":
+            producers = set(
+                re.findall(r"Register Xavier direct handoff: request=(\S+)", evidence)
+            )
+            loaded = set(
+                re.findall(r"Finished Xavier async KV load: request=(\S+)", evidence)
+            )
+            assert len(producers) == count, len(producers)
+            assert len(loaded - producers) == count, len(loaded - producers)
+        else:
+            loaded = set(
+                re.findall(r"with remote block size \d+ for req (\S+)", evidence)
+            )
+            assert len(loaded) == count, len(loaded)
+            assert re.search(r"and [1-9]\d* requests done recving", evidence)
+
+    try:
+        client.launch_model(
+            model_uid=uid,
+            vllm_transfer_backend_type=backend,
+            model_name=os.environ.get(
+                "XINFERENCE_TEST_PD_MODEL_NAME", "qwen2.5-instruct"
+            ),
+            model_engine="vLLM",
+            model_size_in_billions=os.environ.get(
+                "XINFERENCE_TEST_PD_MODEL_SIZE", "0_5"
+            ),
+            model_path=os.environ.get("XINFERENCE_TEST_PD_MODEL_PATH"),
+            model_format="pytorch",
+            quantization="none",
+            replica=num_p + num_d,
+            enable_virtual_env=False,
+            max_model_len=2048,
+            max_num_seqs=16,
+            gpu_memory_utilization=0.35,
+            enforce_eager=True,
+            dtype="float16",
+            enable_prefix_caching=False,
+            replica_config=[
+                {
+                    "role": role,
+                    "devices": [{"worker_ip": worker, "n_gpu": 1, "gpu_idx": [gpu]}],
+                }
+                for role, gpu in [("prefill", 0)] * num_p + [("decode", 1)] * num_d
+            ],
+        )
+        pd = loop.run_until_complete(
+            xo.actor_ref(
+                address=client._get_supervisor_internal_address(),
+                uid=f"{uid}-{PDModelActor.default_uid()}",
+            )
+        )
+        info = loop.run_until_complete(pd.get_pd_info())
+        assert (info["prefill_count"], info["decode_count"]) == (num_p, num_d)
+        start = Path(log_path).stat().st_size
+        # Both producers see the same prefix; repeated requests must restore it.
+        for i in range(4):
+            chat(i, bool(i % 2))
+        offset = Path(log_path).stat().st_size
+        for i in range(4):
+            chat(i, bool(i % 2))
+        warm_evidence = logs(offset)
+        check_transfers(warm_evidence, 4)
+        if backend == "xavier":
+            assert re.search(r"Restored Xavier history: blocks=[1-9]\d*", warm_evidence)
+
+        def concurrent():
+            offset = Path(log_path).stat().st_size
+            with ThreadPoolExecutor(max_workers=8) as executor:
+                futures = [executor.submit(chat, i, bool(i % 2)) for i in range(8)]
+                for future in futures:
+                    future.result(timeout=180)
+            check_transfers(logs(offset), 8)
+
+        concurrent()
+        requests = 16
+        if num_p == num_d == 2:
+            # Equal-length round robins only visit two diagonal pairs. Removing
+            # and re-registering D0 rotates D, exercising the other two pairs.
+            decoder = info["decode_replica_uids"][0]
+            ref = loop.run_until_complete(pd.get_decode_actor(decoder))
+            loop.run_until_complete(pd.remove_decode_actor(decoder))
+            loop.run_until_complete(pd.add_decode_actor(decoder, ref))
+            concurrent()
+            requests += 8
+        evidence = logs(start)
+        routes = re.findall(
+            r"PD route: request=(\S+) prefill=(\S+) decode=(\S+) backend=(\S+)",
+            evidence,
+        )
+        assert len(routes) == requests, routes
+        pairs = {(p, d) for _, p, d, _ in routes}
+        assert len(pairs) == num_p * num_d, pairs
+        assert {b for _, _, _, b in routes} == {backend}
+        summary = {
+            "backend": backend,
+            "topology": f"{num_p}P{num_d}D",
+            "requests": requests,
+            "pairs": sorted(pairs),
+            "history_restores": len(re.findall(r"Restored Xavier history:", evidence)),
+        }
+        print("MULTI_PD_RESULT " + json.dumps(summary), flush=True)
+        client.terminate_model(uid)
+        assert uid not in client.list_models()
+    finally:
+        api.close()
+        try:
+            if uid in client.list_models():
+                client.terminate_model(uid)
+        finally:
+            loop.close()