feat(multimodal): wire vLLM RDMA multimodal tensor pull

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

@@ -54,6 +54,11 @@ on:
         type: string
         default: "nixl"
         description: "KV transfer backend for vLLM PD workers: nixl or mooncake"
+      wheel_artifact:
+        required: false
+        type: string
+        default: "smg-wheel"
+        description: "Name of the wheel artifact to install (e.g. smg-wheel-mm-rdma)"
 
 jobs:
   run:
@@ -92,7 +97,7 @@ jobs:
       - name: Download wheel artifact
         uses: actions/download-artifact@v8
         with:
-          name: smg-wheel
+          name: ${{ inputs.wheel_artifact }}
           path: wheel/
 
       - name: Download WASM test fixtures

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

@@ -212,6 +212,37 @@ jobs:
           python3 -c "from smg.smg_rs import Router; print('Rust extension: OK')"
           python3 -m smg.launch_router --help > /dev/null && echo "Entry point: OK"
 
+  build-wheel-mm-rdma:
+    # Dedicated mm-rdma wheel for the RDMA e2e, built natively on the bare Ubuntu
+    # runner (GNU libstdc++). The shared wheel is cross-compiled with
+    # `maturin --zig`, which ships LLVM libc++ and can't link nixl-sys's
+    # GNU-libstdc++ C++ stub; building it separately here keeps that off the
+    # shared wheel so it can't break the other e2e lanes. Runs on the same runner
+    # image as the GPU e2e lanes, so the native wheel loads there unchanged.
+    needs: detect-changes
+    if: >-
+      always()
+      && !cancelled()
+      && (github.event_name != 'pull_request'
+          || (needs.detect-changes.result == 'success'
+              && (needs.detect-changes.outputs.common == 'true'
+                  || needs.detect-changes.outputs.chat-completions == 'true')))
+    runs-on: k8s-runner-cpu
+    permissions:
+      contents: read
+    steps:
+      - uses: actions/checkout@v7
+      - name: Setup Rust
+        uses: ./.github/actions/setup-rust
+      - name: Build mm-rdma wheel
+        run: bash scripts/ci_build_wheel_mm_rdma.sh
+      - name: Upload mm-rdma wheel artifact
+        uses: actions/upload-artifact@v7
+        with:
+          name: smg-wheel-mm-rdma
+          path: bindings/python/dist/*.whl
+          retention-days: 1
+
   python-unit-tests:
     needs: build-wheel
     runs-on: k8s-runner-cpu
@@ -508,6 +539,9 @@ jobs:
           - engine: vllm
             timeout: 24
             test_timeout: 18
+            # The RDMA multimodal test needs the dedicated mm-rdma wheel + a NIXL
+            # runner; it runs in e2e-1gpu-mm-rdma, not this shared-wheel lane.
+            test_filter: "--ignore=e2e_test/chat_completions/test_multimodal_rdma.py"
           - engine: trtllm
             timeout: 32
             test_timeout: 18
@@ -529,6 +563,31 @@ jobs:
       timeout: ${{ matrix.timeout }}
       test_timeout: ${{ matrix.test_timeout }}
       test_dirs: ${{ matrix.test_dirs || 'e2e_test/chat_completions' }}
+      test_filter: ${{ matrix.test_filter || '' }}
+    secrets: inherit
+
+  e2e-1gpu-mm-rdma:
+    # Isolated RDMA multimodal e2e: installs the dedicated mm-rdma wheel and runs
+    # only test_multimodal_rdma.py against a vLLM worker on a NIXL-capable runner.
+    needs: [build-wheel, build-wheel-mm-rdma, detect-changes]
+    if: >-
+      always()
+      && !cancelled()
+      && needs.build-wheel.result == 'success'
+      && needs.build-wheel-mm-rdma.result == 'success'
+      && (github.event_name != 'pull_request'
+          || (needs.detect-changes.result == 'success'
+              && (needs.detect-changes.outputs.common == 'true'
+                  || needs.detect-changes.outputs.chat-completions == 'true')))
+    uses: ./.github/workflows/e2e-gpu-job.yml
+    with:
+      engine: vllm
+      gpu_tier: "1"
+      runner: 1-gpu-h100
+      timeout: 24
+      test_timeout: 18
+      test_dirs: e2e_test/chat_completions/test_multimodal_rdma.py
+      wheel_artifact: smg-wheel-mm-rdma
     secrets: inherit
 
   e2e-1gpu-completions:
@@ -1037,7 +1096,7 @@ jobs:
           path: benchmark_go_bindings/
 
   finish:
-    needs: [pre-commit, python-lint, grpc-proto-build-check, build-wheel, python-unit-tests, unit-tests, benchmarks, e2e-1gpu-chat, e2e-1gpu-completions, e2e-1gpu-embeddings, e2e-1gpu-gateway, e2e-1gpu-responses, e2e-2gpu-pd, e2e-4gpu-chat, e2e-4gpu-gateway, e2e-4gpu-epd, e2e-vendor, go-unit-tests, go-bindings-e2e]
+    needs: [pre-commit, python-lint, grpc-proto-build-check, build-wheel, build-wheel-mm-rdma, python-unit-tests, unit-tests, benchmarks, e2e-1gpu-chat, e2e-1gpu-mm-rdma, e2e-1gpu-completions, e2e-1gpu-embeddings, e2e-1gpu-gateway, e2e-1gpu-responses, e2e-2gpu-pd, e2e-4gpu-chat, e2e-4gpu-gateway, e2e-4gpu-epd, e2e-vendor, go-unit-tests, go-bindings-e2e]
     if: always()
     runs-on: k8s-runner-cpu
     permissions: {}
@@ -1048,10 +1107,12 @@ jobs:
                 "${{ needs.python-lint.result }}" == "failure" || \
                 "${{ needs.grpc-proto-build-check.result }}" == "failure" || \
                 "${{ needs.build-wheel.result }}" == "failure" || \
+                "${{ needs.build-wheel-mm-rdma.result }}" == "failure" || \
                 "${{ needs.python-unit-tests.result }}" == "failure" || \
                 "${{ needs.unit-tests.result }}" == "failure" || \
                 "${{ needs.benchmarks.result }}" == "failure" || \
                 "${{ needs.e2e-1gpu-chat.result }}" == "failure" || \
+                "${{ needs.e2e-1gpu-mm-rdma.result }}" == "failure" || \
                 "${{ needs.e2e-1gpu-completions.result }}" == "failure" || \
                 "${{ needs.e2e-1gpu-embeddings.result }}" == "failure" || \
                 "${{ needs.e2e-1gpu-gateway.result }}" == "failure" || \

bindings/python/Cargo.toml

@@ -36,6 +36,7 @@ workspace = true
 default = ["pyo3/extension-module"]
 opencv-video = ["smg/opencv-video"]
 vendored-openssl = ["smg/vendored-openssl"]
+mm-rdma = ["smg/mm-rdma"]
 
 [profile.ci]
 inherits = "release"

crates/grpc_client/proto/vllm_engine.proto

@@ -360,6 +360,11 @@ message GetServerInfoResponse {
   // the router can verify a shared /dev/shm before using the SHM tensor
   // transport under `auto`. Empty when it can't be determined.
   string shm_namespace_id = 10;
+
+  // Whether this worker can pull RDMA (NIXL) multimodal tensor payloads. The
+  // router only emits a `remote` payload when this is true, so an older worker
+  // that cannot consume it stays on the inline/SHM path.
+  bool supports_rdma_pull = 11;
 }
 
 // =====================

crates/grpc_client/python/pyproject.toml

@@ -8,7 +8,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "smg-grpc-proto"
-version = "0.4.14"
+version = "0.4.15"
 description = "SMG gRPC proto definitions for vLLM, TRT-LLM, MLX, TokenSpeed, and SGLang"
 requires-python = ">=3.10"
 dependencies = [

e2e_test/chat_completions/test_multimodal_rdma.py

@@ -0,0 +1,108 @@
+"""vLLM RDMA multimodal transport E2E test.
+
+Exercises the full RDMA (NIXL) pixel lane end to end: the gateway stages
+``pixel_values`` into a pre-registered arena and the vLLM worker pulls them with a
+one-sided READ instead of receiving them inline. Verifies both that the model
+still understands the image over RDMA AND that the pixels actually travelled the
+remote path (``smg_mm_tensor_bytes_total{path="remote"}`` grew), so a silent
+fall-back to inline fails the test rather than passing on correctness alone.
+
+Requires a gateway built with ``--features mm-rdma`` and a NIXL/UCX-capable host.
+
+Usage:
+    pytest e2e_test/chat_completions/test_multimodal_rdma.py -v
+"""
+
+from __future__ import annotations
+
+import base64
+import logging
+from pathlib import Path
+
+import httpx
+import pytest
+
+logger = logging.getLogger(__name__)
+
+FIXTURES_DIR = Path(__file__).parent.parent / "fixtures" / "images"
+DOG_IMAGE_PATH = FIXTURES_DIR / "dog.jpg"  # Black labrador puppy
+
+
+def _image_to_base64_url(path: Path) -> str:
+    data = base64.b64encode(path.read_bytes()).decode("utf-8")
+    return f"data:image/jpeg;base64,{data}"
+
+
+def _remote_pixel_bytes(metrics_url: str) -> float:
+    """Sum ``smg_mm_tensor_bytes_total`` samples on the RDMA (remote) path."""
+    resp = httpx.get(f"{metrics_url}/metrics", timeout=10)
+    resp.raise_for_status()
+    total = 0.0
+    for line in resp.text.splitlines():
+        if line.startswith("smg_mm_tensor_bytes_total") and 'path="remote"' in line:
+            total += float(line.rsplit(" ", 1)[1])
+    return total
+
+
+@pytest.fixture(scope="class", autouse=True)
+def _rdma_env():
+    """Turn the RDMA lane on for both the gateway and the vLLM worker, which
+    inherit this process's environment when launched locally. Restored on teardown.
+
+    - ``SMG_MM_TENSOR_TRANSPORT=rdma``: the first-class transport switch (also set
+      as a gateway CLI flag) that enables the worker-side puller.
+    - ``SMG_MM_PIXEL_RDMA=1``: the legacy puller switch, belt-and-suspenders.
+    - ``SMG_RDMA_LISTEN_IP=127.0.0.1``: the gateway's NIXL listener (loopback).
+    """
+    with pytest.MonkeyPatch.context() as mp:
+        mp.setenv("SMG_MM_TENSOR_TRANSPORT", "rdma")
+        mp.setenv("SMG_MM_PIXEL_RDMA", "1")
+        mp.setenv("SMG_RDMA_LISTEN_IP", "127.0.0.1")
+        yield
+
+
+@pytest.mark.engine("vllm")
+@pytest.mark.gpu(1)
+@pytest.mark.e2e
+@pytest.mark.model("Qwen/Qwen3-VL-8B-Instruct")
+@pytest.mark.gateway(extra_args=["--multimodal-tensor-transport", "rdma"])
+@pytest.mark.parametrize("setup_backend", ["grpc"], indirect=True)
+class TestMultimodalRdmaQwen3VL:
+    """vLLM multimodal over the RDMA pixel lane via gRPC."""
+
+    def test_single_image_uses_rdma(self, model, setup_backend):
+        _, _, client, gateway = setup_backend
+
+        response = client.chat.completions.create(
+            model=model,
+            messages=[
+                {
+                    "role": "user",
+                    "content": [
+                        {"type": "text", "text": "What animal is in this image?"},
+                        {
+                            "type": "image_url",
+                            "image_url": {"url": _image_to_base64_url(DOG_IMAGE_PATH)},
+                        },
+                    ],
+                }
+            ],
+            temperature=0,
+            max_tokens=100,
+        )
+
+        text = response.choices[0].message.content
+        assert any(k in text.lower() for k in ["dog", "puppy", "labrador"]), (
+            f"Expected dog-related content over the RDMA transport, got: {text}"
+        )
+        logger.info("RDMA multimodal response: %s", text)
+
+        # Critical: the request must have actually used the remote (RDMA) path.
+        # A silent inline fall-back would still answer correctly, so correctness
+        # alone does not prove the lane worked — the transport metric does.
+        remote_bytes = _remote_pixel_bytes(gateway.metrics_url)
+        assert remote_bytes > 0, (
+            'expected smg_mm_tensor_bytes_total{path="remote"} > 0; the pixels '
+            "silently fell back to inline (RDMA lane not exercised)"
+        )
+        logger.info("RDMA remote pixel bytes: %s", remote_bytes)

e2e_test/infra/worker_pool.py

@@ -25,6 +25,7 @@
 
 import atexit
 import logging
+import os
 import threading
 
 from .constants import DEFAULT_STARTUP_TIMEOUT, ConnectionMode, WorkerType
@@ -33,12 +34,24 @@
 logger = logging.getLogger(__name__)
 
 
-# Key is (engine, model_id, mode, worker_type, count, gpus, extra_engine_args).
-# ``count`` is part of the key because a class asking for count=2 after a
-# count=1 class on the same backend would otherwise reuse a 1-worker entry and
-# run with the wrong topology; ``gpus``/``extra_engine_args`` likewise change
-# the launched topology (e.g. --data-parallel-size).
-_PoolKey = tuple[str, str, ConnectionMode, WorkerType, int, int | None, tuple[str, ...] | None]
+# Key is (engine, model_id, mode, worker_type, count, gpus, extra_engine_args,
+# rdma_env). ``count``/``gpus``/``extra_engine_args`` change the launched
+# topology, so a class asking for a different one must not reuse a cached entry.
+# ``rdma_env`` is included because workers inherit the launcher's environment and
+# the RDMA lane is env-gated, so an RDMA-enabled worker must not be reused for an
+# inline test (or vice versa).
+_PoolKey = tuple[
+    str, str, ConnectionMode, WorkerType, int, int | None, tuple[str, ...] | None, tuple[str, ...]
+]
+
+
+def _rdma_env_signature() -> tuple[str, ...]:
+    """RDMA-lane env inherited by workers at launch; part of the pool key."""
+    return (
+        os.environ.get("SMG_MM_TENSOR_TRANSPORT", ""),
+        os.environ.get("SMG_MM_PIXEL_RDMA", ""),
+        os.environ.get("SMG_RDMA_LISTEN_IP", ""),
+    )
 
 
 class WorkerPool:
@@ -123,6 +136,7 @@ def acquire(
                 count,
                 gpus,
                 tuple(extra_engine_args) if extra_engine_args else None,
+                _rdma_env_signature(),
             )
 
             if self._key == key and all(w.is_alive() for w in self._workers):

grpc_servicer/pyproject.toml

@@ -8,7 +8,9 @@ version = "0.6.0"
 description = "SMG gRPC servicer implementations for LLM inference engines (vLLM, MLX, TokenSpeed, SGLang)"
 requires-python = ">=3.10"
 dependencies = [
-    "smg-grpc-proto>=0.4.13",
+    # >=0.4.15 ships GetServerInfoResponse.supports_rdma_pull (field 11); the vLLM
+    # servicer sets it, so older stubs would raise on the unknown keyword.
+    "smg-grpc-proto>=0.4.15",
     "grpcio>=1.81.1",
     "grpcio-reflection>=1.81.1",
     "grpcio-health-checking>=1.81.1",

grpc_servicer/smg_grpc_servicer/mm_rdma.py

@@ -20,6 +20,14 @@
 _LOCAL_IP_CACHE: dict[str, bool] = {}
 
 
+def _rdma_enabled_from_env() -> bool:
+    """Whether the RDMA lane is on: the first-class `SMG_MM_TENSOR_TRANSPORT=rdma`
+    (matching the gateway) or the legacy `SMG_MM_PIXEL_RDMA` flag."""
+    if os.environ.get("SMG_MM_PIXEL_RDMA") in ("1", "true"):
+        return True
+    return os.environ.get("SMG_MM_TENSOR_TRANSPORT", "").strip().lower() == "rdma"
+
+
 @dataclass(frozen=True)
 class _RemotePixelDescriptor:
     remote_addr: int
@@ -121,7 +129,7 @@ def __init__(
         self._landing_slot_bytes = 0
         self._landing_free = None
 
-        if os.environ.get("SMG_MM_PIXEL_RDMA") not in ("1", "true"):
+        if not _rdma_enabled_from_env():
             return
 
         try:
@@ -169,6 +177,11 @@ def __init__(
             self._nixl_agent = None
             self._landing_free = None
 
+    @property
+    def ready(self) -> bool:
+        """True once the NIXL agent + landing pool initialized (RDMA pulls possible)."""
+        return self._nixl_agent is not None and self._landing_free is not None
+
     def _ensure_remote_ready(self, ip: str, port: int, remote, room: int) -> None:
         """One-time metadata handshake per gateway listener."""
         key = (ip, port)