test: multimodal tensor-transport + EPD disaggregation coverage

.github/workflows/e2e-gpu-job.yml

@@ -126,6 +126,13 @@ jobs:
       - name: Download models
         run: bash scripts/ci_download_model.sh --gpu-tier ${{ inputs.gpu_tier }}
 
+      # The TokenSpeed EPD model is skip_tier_download (large, engine-specific), so
+      # the tier step above skips it; pull it by id only for the 4-GPU tokenspeed
+      # lane that runs the EPD e2e — other lanes never touch it.
+      - name: Download EPD model (tokenspeed only)
+        if: inputs.engine == 'tokenspeed' && inputs.gpu_tier == '4'
+        run: bash scripts/ci_download_model.sh "Qwen/Qwen3.6-35B-A3B-FP8"
+
       # Run tests
       - name: Run E2E tests
         timeout-minutes: ${{ inputs.test_timeout }}

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

@@ -710,12 +710,22 @@ 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
+            # The inner pytest step has its own timeout (default 25m); the four EPD
+            # topologies each relaunch a worker set, so give pytest more room too.
+            test_timeout: 65
     uses: ./.github/workflows/e2e-gpu-job.yml
     with:
       engine: ${{ matrix.engine }}
       gpu_tier: "4"
       runner: 4-gpu-h100
       timeout: ${{ matrix.timeout }}
+      test_timeout: ${{ matrix.test_timeout || 25 }}
       test_dirs: e2e_test/chat_completions
     secrets: inherit
 

e2e_test/chat_completions/test_epd_multimodal.py

@@ -0,0 +1,132 @@
+"""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.
+
+Like the PD KV-transfer tests (``test_pd_mooncake``/``test_pd_nixl``), these do
+NOT stop at "a plausible answer came back" — a single-worker fallback would pass
+that. They assert a worker-side signal that the disaggregation actually happened:
+the encode worker's own per-request accept log.
+
+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
+import os
+import tempfile
+from pathlib import Path
+
+import pytest
+from infra.pd_logs import assert_worker_logs_captured, wait_for_marker, worker_log_dir
+
+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"
+
+# Router + worker logs land here (per-pid) via the gateway marker below. Worker
+# logs actually go to E2E_LOG_DIR in CI; ``worker_log_dir`` resolves both.
+_LOG_DIR = Path(tempfile.gettempdir()) / f"smg-e2e-epd-{os.getpid()}"
+# Emitted once per Encode RPC by the TokenSpeed encode servicer
+# (grpc_servicer/.../tokenspeed/encoder_servicer.py). Its presence proves the
+# image reached a dedicated encode worker — the defining EPD step.
+ENCODE_ACCEPTED_MARKER = "EPD encode: accepted"
+
+
+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 a URL or data URL string.
+
+    Local file paths must be pre-converted via ``_image_to_base64_url``.
+    """
+    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. The (encode, prefill, decode) counts ride in the param tuple, not
+# a marker: setup_backend is class-scoped, so per-param marks aren't visible there.
+_EPD_TOPOLOGIES = [
+    pytest.param(("epd_grpc", (1, 1, 1)), id="1e1p1d"),
+    pytest.param(("epd_grpc", (1, 2, 1)), id="1e2p1d"),
+    pytest.param(("epd_grpc", (2, 1, 1)), id="2e1p1d"),
+    pytest.param(("epd_grpc", (1, 1, 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.gateway(log_level="debug", log_dir=str(_LOG_DIR))
+@pytest.mark.parametrize("setup_backend", _EPD_TOPOLOGIES, indirect=True)
+class TestEPDMultimodal:
+    """Verify the image really flows encode -> prefill -> decode.
+
+    A naive content check can't distinguish a real 3-worker EPD pipeline from a
+    single-worker fallback; the encode worker's own log can — so that's what this
+    asserts, mirroring how the PD tests assert the KV transfer from logs.
+    """
+
+    def test_single_image_base64(self, model, setup_backend):
+        """One dog image through encode -> prefill -> decode, with the encode
+        worker's participation verified from its logs."""
+        _, _, 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,
+        )
+
+        # (1) The model saw the image and described it correctly.
+        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}"
+        )
+        # (2) The image was tokenized INTO the prompt: the bare text question is
+        # ~10 tokens, so a large prompt confirms the vision tokens were spliced in
+        # (encode -> prefill actually delivered the image), not dropped.
+        assert response.usage.prompt_tokens > 50, (
+            f"prompt_tokens={response.usage.prompt_tokens} is too low; "
+            "the image tokens were likely not delivered to prefill"
+        )
+        assert response.usage.completion_tokens > 0
+
+        # (3) REAL EPD: the encode worker itself logged accepting the dispatch, so
+        # the vision stage ran on a separate encode worker rather than degrading to
+        # a single-worker path. This is the EPD analog of the PD KV-transfer check.
+        worker_dir = worker_log_dir(_LOG_DIR)
+        worker_logs = wait_for_marker(worker_dir, "worker-*.log", ENCODE_ACCEPTED_MARKER)
+        assert_worker_logs_captured(worker_logs, "EPD encode dispatch")
+        assert ENCODE_ACCEPTED_MARKER in worker_logs, (
+            "encode worker never logged accepting the request — the image did not "
+            f"flow through the EPD encode stage; checked {worker_dir}/worker-*.log"
+        )
+        logger.info("EPD single image (encode worker engaged): %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/setup_backend.py

@@ -107,6 +107,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,12 +116,21 @@ 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
+      - EPD topology: pass a ("epd_grpc", (encode, prefill, decode)) param
       - ``@pytest.mark.gateway(policy=..., timeout=..., extra_args=...)``: Gateway config
 
     Returns:
         Tuple of ``(backend_name, model_path, client, gateway)``
     """
-    backend_name: str = request.param
+    # EPD topologies pass a ("epd_grpc", (n_encode, n_prefill, n_decode)) tuple;
+    # every other backend passes a bare protocol string. Read the topology from
+    # request.param, NOT a marker — this fixture is class-scoped, so per-param
+    # marks on the generated test items aren't visible on request.node.
+    param = request.param
+    if isinstance(param, tuple):
+        backend_name, epd_topology = param
+    else:
+        backend_name, epd_topology = param, None
 
     if os.environ.get(ENV_SKIP_BACKEND_SETUP, "").lower() in ("1", "true", "yes"):
         pytest.skip(f"{ENV_SKIP_BACKEND_SETUP} is set")
@@ -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,18 @@ def setup_backend(request: pytest.FixtureRequest):
 
     gateway = Gateway()
     try:
-        if is_pd:
+        if is_epd:
+            yield from _setup_epd(
+                model_id,
+                model_path,
+                engine,
+                connection_mode,
+                epd_topology or (1, 1, 1),
+                gateway_config,
+                gateway,
+                log_dir,
+            )
+        elif is_pd:
             yield from _setup_pd(
                 model_id,
                 model_path,
@@ -299,6 +321,98 @@ def _setup_pd(
         stop_workers(all_workers)
 
 
+# ---------------------------------------------------------------------------
+# EPD disaggregation backend (TokenSpeed)
+# ---------------------------------------------------------------------------
+
+
+def _setup_epd(
+    model_id,
+    model_path,
+    engine,
+    connection_mode,
+    epd_topology,
+    gateway_config,
+    gateway,
+    log_dir,
+):
+    """Launch encode + prefill + decode workers + EPD gateway, yield, tear down.
+
+    Mirrors ``_setup_pd``. ``epd_topology`` is ``(n_encode, n_prefill, n_decode)``.
+    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, num_prefill, num_decode = epd_topology
+    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
 # ---------------------------------------------------------------------------