test: multimodal tensor-transport + EPD disaggregation coverage

.github/workflows/pr-test-rust.yml

@@ -710,6 +710,12 @@ jobs:
             timeout: 30
           - engine: trtllm
             timeout: 45
+          # EPD multimodal (test_epd_multimodal.py): the encode/prefill/decode
+          # TokenSpeed workers span the 4 GPUs at tp=1. Higher timeout — tokenspeed
+          # builds from source (~30m cold) and the four topologies each relaunch
+          # their worker set.
+          - engine: tokenspeed
+            timeout: 75
     uses: ./.github/workflows/e2e-gpu-job.yml
     with:
       engine: ${{ matrix.engine }}

e2e_test/chat_completions/test_epd_multimodal.py

@@ -0,0 +1,88 @@
+"""EPD (Encode-Prefill-Decode) multimodal Chat Completions E2E Tests.
+
+Exercises TokenSpeed's EPD disaggregation on a vision-language model: the
+encode worker runs the vision tower, prefill/decode run the LM, and the
+gateway stitches encode -> prefill -> decode. The point of these tests is
+that a disaggregated encode->prefill->decode path still produces a correct
+multimodal answer.
+
+EPD is TokenSpeed-only. On a small MoE VLM (Qwen3.6-35B-A3B, 3B active) at
+tp=1 per worker, every topology fits the 4-GPU runner: 1e1p1d=3 GPUs and
+1e2p1d/2e1p1d/1e1p2d=4 GPUs (EPD needs >=3 cards since encode/prefill/decode
+are separate workers).
+
+Usage:
+    pytest e2e_test/chat_completions/test_epd_multimodal.py -v
+"""
+
+from __future__ import annotations
+
+import base64
+import logging
+from pathlib import Path
+
+import pytest
+
+logger = logging.getLogger(__name__)
+
+# Local test image (checked into repo) — a black labrador puppy.
+FIXTURES_DIR = Path(__file__).parent.parent / "fixtures" / "images"
+DOG_IMAGE_PATH = FIXTURES_DIR / "dog.jpg"
+
+
+def _image_to_base64_url(path: Path) -> str:
+    """Convert a local image file to a base64 data URL."""
+    data = base64.b64encode(path.read_bytes()).decode("utf-8")
+    return f"data:image/jpeg;base64,{data}"
+
+
+def _make_image_content(image_source: str) -> dict:
+    """Create an image_url content part from either a URL or local path."""
+    return {"type": "image_url", "image_url": {"url": image_source}}
+
+
+# The four EPD topologies to cover. Every worker runs at tp=1 (a 3B-active MoE
+# that fits one card), so 1e1p1d uses 3 GPUs and the rest use 4 — all fit the
+# 4-GPU runner. Each carries an ``epd`` marker consumed by setup_backend's EPD path.
+_EPD_TOPOLOGIES = [
+    pytest.param("epd_grpc", marks=pytest.mark.epd(encode=1, prefill=1, decode=1), id="1e1p1d"),
+    pytest.param("epd_grpc", marks=pytest.mark.epd(encode=1, prefill=2, decode=1), id="1e2p1d"),
+    pytest.param("epd_grpc", marks=pytest.mark.epd(encode=2, prefill=1, decode=1), id="2e1p1d"),
+    pytest.param("epd_grpc", marks=pytest.mark.epd(encode=1, prefill=1, decode=2), id="1e1p2d"),
+]
+
+
+@pytest.mark.engine("tokenspeed")
+@pytest.mark.gpu(4)
+@pytest.mark.e2e
+@pytest.mark.model("Qwen/Qwen3.6-35B-A3B-FP8")
+@pytest.mark.parametrize("setup_backend", _EPD_TOPOLOGIES, indirect=True)
+class TestEPDMultimodal:
+    """Multimodal tests over TokenSpeed EPD disaggregation (4 GPU)."""
+
+    def test_single_image_base64(self, model, setup_backend):
+        """A single dog image travels encode -> prefill -> decode and is described."""
+        _, _, client, *_ = setup_backend
+
+        response = client.chat.completions.create(
+            model=model,
+            messages=[
+                {
+                    "role": "user",
+                    "content": [
+                        {"type": "text", "text": "What animal is in this image?"},
+                        _make_image_content(_image_to_base64_url(DOG_IMAGE_PATH)),
+                    ],
+                }
+            ],
+            temperature=0,
+            max_tokens=100,
+        )
+
+        text = response.choices[0].message.content
+        assert text is not None and len(text) > 0
+        assert any(k in text.lower() for k in ["dog", "puppy", "labrador"]), (
+            f"Expected dog-related content, got: {text}"
+        )
+        assert response.usage.prompt_tokens > 0
+        logger.info("EPD single image base64: %s", text)

e2e_test/chat_completions/test_multimodal.py

@@ -173,6 +173,64 @@ def test_multi_images_mixed(self, model, setup_backend):
         logger.info("Multi image mixed response: %s", text)
 
 
+# =============================================================================
+# /dev/shm tensor-transport verification (1 GPU, vLLM)
+# =============================================================================
+
+
+@pytest.mark.engine("vllm")
+@pytest.mark.gpu(1)
+@pytest.mark.e2e
+@pytest.mark.model("Qwen/Qwen3-VL-8B-Instruct")
+# Force the /dev/shm tensor transport. A low min-bytes so any real image tensor
+# crosses the threshold, making the shm path deterministic regardless of size.
+@pytest.mark.gateway(
+    extra_args=["--multimodal-tensor-transport", "shm", "--multimodal-shm-min-bytes", "1024"]
+)
+@pytest.mark.parametrize("setup_backend", ["grpc"], indirect=True)
+class TestMultimodalShmTransport:
+    """Ground truth that the /dev/shm tensor transport actually engages.
+
+    Gateway and worker are co-located in CI, so they share /dev/shm and the shm
+    path can be used. This asserts not just that multimodal works, but that the
+    pixel tensor traveled over shm (``smg_mm_tensors_total{path="shm"}``) rather
+    than silently falling back to the inline gRPC payload.
+    """
+
+    def test_single_image_uses_shm_transport(self, model, setup_backend):
+        _, _, client, gateway = setup_backend
+
+        before = gateway.metric_sum("smg_mm_tensors_total", path="shm")
+
+        response = client.chat.completions.create(
+            model=model,
+            messages=[
+                {
+                    "role": "user",
+                    "content": [
+                        {"type": "text", "text": "What animal is in this image?"},
+                        _make_image_content(_image_to_base64_url(DOG_IMAGE_PATH)),
+                    ],
+                }
+            ],
+            temperature=0,
+            max_tokens=100,
+        )
+
+        text = _extract_text(response, False)
+        assert any(k in text.lower() for k in ["dog", "puppy", "labrador"]), (
+            f"Expected dog-related content, got: {text}"
+        )
+
+        after = gateway.metric_sum("smg_mm_tensors_total", path="shm")
+        assert after > before, (
+            "expected the /dev/shm tensor transport to be used "
+            f"(smg_mm_tensors_total path=shm before={before} after={after}); "
+            "the pixel tensor silently fell back to the inline gRPC payload"
+        )
+        logger.info("shm transport: path=shm count rose %s -> %s", before, after)
+
+
 # =============================================================================
 # Llama-4-Scout multimodal tests (4 GPU)
 # =============================================================================

e2e_test/fixtures/hooks.py

@@ -55,6 +55,10 @@ def pytest_configure(config: pytest.Config) -> None:
         "workers(count=1, prefill=None, decode=None, gpus=None, extra_engine_args=None): "
         "worker topology configuration",
     )
+    config.addinivalue_line(
+        "markers",
+        "epd(encode=1, prefill=1, decode=1): EPD disaggregation worker topology",
+    )
     config.addinivalue_line(
         "markers",
         "storage(backend): storage backend for cloud tests (memory, oracle-custom)",

e2e_test/fixtures/setup_backend.py

@@ -57,6 +57,14 @@
     "extra_engine_args": None,
 }
 
+# EPD topology defaults (1 encode + 1 prefill + 1 decode), read from
+# ``@pytest.mark.epd(encode=..., prefill=..., decode=...)``.
+_EPD_DEFAULTS = {
+    "encode": 1,
+    "prefill": 1,
+    "decode": 1,
+}
+
 # Track worker startup failures — fail fast after repeated failures
 _worker_start_failures: dict[str, int] = {}  # engine -> count
 _MAX_WORKER_START_FAILURES = 3  # fail fast after this many failures (matches --reruns 2)
@@ -107,6 +115,7 @@ def setup_backend(request: pytest.FixtureRequest):
     Backend type is determined by parametrize value via ``request.param``:
       - ``"http"``, ``"grpc"``: Local workers (SGLang, vLLM, or TRT-LLM)
       - ``"pd_http"``, ``"pd_grpc"``: PD disaggregation workers
+      - ``"epd_grpc"``: EPD (encode-prefill-decode) disaggregation (TokenSpeed)
       - ``"openai"``, ``"xai"``, ``"anthropic"``: Cloud backends (no workers)
 
     Configuration via markers:
@@ -115,6 +124,7 @@ def setup_backend(request: pytest.FixtureRequest):
       - ``@pytest.mark.workers(gpus=2, extra_engine_args=[...])``: Per-worker
         GPU count and extra engine CLI args (local workers only)
       - ``@pytest.mark.workers(prefill=1, decode=1)``: PD worker counts
+      - ``@pytest.mark.epd(encode=1, prefill=1, decode=1)``: EPD worker counts
       - ``@pytest.mark.gateway(policy=..., timeout=..., extra_args=...)``: Gateway config
 
     Returns:
@@ -137,8 +147,9 @@ def setup_backend(request: pytest.FixtureRequest):
         return
 
     # Local backends
+    is_epd = backend_name.startswith("epd_")
     is_pd = backend_name.startswith("pd_")
-    protocol = backend_name.replace("pd_", "")
+    protocol = backend_name.replace("epd_", "").replace("pd_", "")
     connection_mode = ConnectionMode(protocol)
     engine = get_runtime()
     model_path = get_model_spec(model_id)["model"]
@@ -154,7 +165,19 @@ def setup_backend(request: pytest.FixtureRequest):
 
     gateway = Gateway()
     try:
-        if is_pd:
+        if is_epd:
+            epd_config = get_marker_kwargs(request, "epd", defaults=_EPD_DEFAULTS)
+            yield from _setup_epd(
+                model_id,
+                model_path,
+                engine,
+                connection_mode,
+                epd_config,
+                gateway_config,
+                gateway,
+                log_dir,
+            )
+        elif is_pd:
             yield from _setup_pd(
                 model_id,
                 model_path,
@@ -299,6 +322,99 @@ def _setup_pd(
         stop_workers(all_workers)
 
 
+# ---------------------------------------------------------------------------
+# EPD disaggregation backend (TokenSpeed)
+# ---------------------------------------------------------------------------
+
+
+def _setup_epd(
+    model_id,
+    model_path,
+    engine,
+    connection_mode,
+    epd_config,
+    gateway_config,
+    gateway,
+    log_dir,
+):
+    """Launch encode + prefill + decode workers + EPD gateway, yield, tear down.
+
+    Mirrors ``_setup_pd``. The encode worker runs the vision tower at tp=1
+    (``gpus=1``); prefill/decode run the LM at the model spec's tp. GPU offsets
+    are laid out encode-first so co-located workers don't share GPUs.
+    """
+    spec = get_model_spec(model_id)
+    tp = spec.get("tp", 1)
+    num_encode = epd_config.get("encode") or 1
+    num_prefill = epd_config.get("prefill") or 1
+    num_decode = epd_config.get("decode") or 1
+    backend_name = f"epd_{connection_mode.value}"
+    runtime_label = RUNTIME_LABELS.get(engine, engine)
+
+    logger.info(
+        "Starting %s EPD backend: model=%s, %d encode + %d prefill + %d decode",
+        runtime_label,
+        model_id,
+        num_encode,
+        num_prefill,
+        num_decode,
+    )
+
+    all_workers: list = []
+    try:
+        # Encode workers: vision tower at tp=1, one GPU each, starting at GPU 0.
+        encode_workers = _start_workers_tracked(
+            model_id=model_id,
+            engine=engine,
+            mode=connection_mode,
+            count=num_encode,
+            worker_type=WorkerType.ENCODE,
+            log_dir=log_dir,
+            gpus=1,
+        )
+        all_workers.extend(encode_workers)
+
+        # Prefill workers start on GPUs after the (tp=1) encode workers.
+        prefill_gpu_offset = num_encode
+        prefill_workers = _start_workers_tracked(
+            model_id=model_id,
+            engine=engine,
+            mode=connection_mode,
+            count=num_prefill,
+            worker_type=WorkerType.PREFILL,
+            log_dir=log_dir,
+            gpu_offset=prefill_gpu_offset,
+        )
+        all_workers.extend(prefill_workers)
+
+        # Decode workers start on GPUs after encode + prefill.
+        decode_gpu_offset = prefill_gpu_offset + num_prefill * tp
+        decode_workers = _start_workers_tracked(
+            model_id=model_id,
+            engine=engine,
+            mode=connection_mode,
+            count=num_decode,
+            worker_type=WorkerType.DECODE,
+            log_dir=log_dir,
+            gpu_offset=decode_gpu_offset,
+        )
+        all_workers.extend(decode_workers)
+
+        _start_gateway(
+            gateway,
+            gateway_config,
+            encode_workers=encode_workers,
+            prefill_workers=prefill_workers,
+            decode_workers=decode_workers,
+        )
+        logger.info("%s EPD backend ready at %s", runtime_label, gateway.base_url)
+        yield backend_name, model_path, _make_openai_client(gateway), gateway
+    finally:
+        logger.info("Tearing down %s EPD backend", runtime_label)
+        gateway.shutdown()
+        stop_workers(all_workers)
+
+
 # ---------------------------------------------------------------------------
 # Cloud backend
 # ---------------------------------------------------------------------------

e2e_test/infra/constants.py

@@ -17,6 +17,7 @@ class WorkerType(StrEnum):
     REGULAR = "regular"
     PREFILL = "prefill"
     DECODE = "decode"
+    ENCODE = "encode"  # EPD encode worker (vision tower); TokenSpeed-only
 
 
 class Runtime(StrEnum):