feat: add RDMA transport for ViT output

Build a reusable Tensor RDMA transport contract around an owner-side
RdmaExport API and a reader-side RdmaRead API. Define stable descriptors,
per-tensor manifests, lease-based slot ownership, batch reads, provider
selection, and a no-op implementation so applications do not depend on a
specific RDMA library.

Decouple separated ViT communication from its data plane. Keep gRPC as the
control plane for capability negotiation, requests, receipts, deadlines, and
lease release, while routing output delivery through transport backends and
receipt readers. The ViT and multimodal processor business paths now consume a
single transport interface instead of branching on grpc-inline or RDMA.

Adapt the complete ViT output to the common RDMA contract. Export embedding,
position IDs, and extra inputs from registered GPU slots, return descriptors in
the receipt, and reconstruct the original MultimodalOutput on the LLM side.
Oversized outputs are split into ordered slots without changing split_size
semantics.

Keep the ViT exporter and LLM receipt reader in separate build targets. Expose
the exporter through libmm_rdma_exporter.so so a ViT process does not load the
LLM and Embedding engine bindings in libth_transformer.so.

Validate descriptors and manifests, bound in-flight slot memory, release slots
through the control plane, and reclaim abandoned leases with GC. RDMA setup,
export, or READ failures use one bounded grpc-inline fallback. The open-source
build provides the transport contract and fallback; the internal build supplies
the Barex provider. MM_TRANSPORT_MODE=auto tries RDMA first, while
MM_TRANSPORT_MODE=grpc forces inline delivery; invalid modes fail fast. Configure
GPU-NIC affinity during service startup using the physical GPU mapping.

Author: ydshi0

BUILD

@@ -184,6 +184,19 @@ cc_binary(
     }),
 )
 
+cc_binary(
+    name = "mm_rdma_exporter",
+    copts = copts(),
+    linkopts = [
+        "-Wl,-rpath='$$ORIGIN'",
+    ],
+    linkshared = 1,
+    visibility = ["//visibility:public"],
+    deps = [
+        "//rtp_llm/cpp/pybind:mm_rdma_exporter_pybind",
+    ],
+)
+
 cc_binary(
     name = "th_transformer",
     srcs = [

arch_config/arch_select.bzl

@@ -8,6 +8,7 @@ load("@pip_gpu_rocm_torch//:requirements.bzl", requirement_gpu_rocm="requirement
 load("@rtp_llm//bazel:defs.bzl", "copy_so")
 
 def copy_all_so():
+    copy_so("@rtp_llm//:mm_rdma_exporter")
     copy_so("@rtp_llm//:th_transformer")
     copy_so("@rtp_llm//:th_transformer_config")
     copy_so("@rtp_llm//:th_grammar_tokenizer_info")
@@ -55,6 +56,14 @@ def cache_store_deps():
         actual = "@rtp_llm//rtp_llm/cpp/disaggregate/cache_store:cache_store_base_impl"
     )
 
+def rdma_transport_deps():
+    # Open-source builds expose the same factory API but have no RDMA provider.
+    native.alias(
+        name = "rdma_transport_arch_select_impl",
+        actual = "@rtp_llm//rtp_llm/cpp/rdma_transport:rdma_transport_no_impl",
+        visibility = ["//visibility:public"],
+    )
+
 def transfer_rdma_deps():
     native.alias(
         name = "transfer_rdma_impl",

bazel/py_proto.bzl

@@ -17,7 +17,7 @@ def _generate_grpc_proto_impl(ctx):
     # use create_grpc_proto generate proto py files
     ctx.actions.run(
         outputs = [pb2_file, pb2_grpc_file],
-        inputs = [proto_file],
+        inputs = depset([proto_file] + ctx.files.proto_srcs),
         executable = ctx.executable.create_grpc_proto,
         arguments = [proto_file.path, output_dir],
         tools = [ctx.executable.create_grpc_proto]
@@ -43,5 +43,8 @@ generate_grpc_proto = rule(
             cfg = "exec",
             mandatory = True,
         ),
+        "proto_srcs": attr.label(
+            allow_files = True,
+        ),
     },
-)
+)

rtp_llm/config/py_config_modules.py

@@ -263,8 +263,38 @@ def to_string(self):
         )
 
 
+# Keep these transport defaults aligned with cpp/config/ConfigModules.h::MMTransportConfig.
+MM_TRANSPORT_MODE_AUTO = "auto"
+MM_TRANSPORT_MODE_GRPC = "grpc"
+MM_TRANSPORT_MODES = (MM_TRANSPORT_MODE_AUTO, MM_TRANSPORT_MODE_GRPC)
+DEFAULT_MM_TIMEOUT_MS = 120000
+
+
+class MMRdmaConfig:
+    def __init__(self):
+        self.bind_ip: str = ""
+        self.port: int = 0
+        self.connect_timeout_ms: int = 250
+        self.read_timeout_ms: int = 3000
+        self.qp_count: int = 8
+        self.slot_gc_timeout_ms: int = 60 * 1000
+        self.max_slot_bytes: int = 1024 * 1024 * 1024
+
+
+class MMControlConfig:
+    def __init__(self):
+        self.release_timeout_ms: int = 1000
+
+
+class MMTransportConfig:
+    def __init__(self):
+        self.mode: str = MM_TRANSPORT_MODE_AUTO
+        self.control = MMControlConfig()
+        self.rdma = MMRdmaConfig()
+
+
 class VitConfig:
-    DEFAULT_MM_TIMEOUT_MS: int = 120000
+    DEFAULT_MM_TIMEOUT_MS: int = DEFAULT_MM_TIMEOUT_MS
     DEFAULT_MM_IMAGE_MAX_FILE_SIZE_KB: int = 100 * 1024
     DEFAULT_MM_VIDEO_MAX_FILE_SIZE_KB: int = 2 * 1024 * 1024
 
@@ -296,6 +326,7 @@ def __init__(self):
         self.disable_access_log: bool = False
         self.use_local_preprocess: bool = False
         self.vit_proxy_load_balance_strategy: str = "round_robin"
+        self.output_transport = MMTransportConfig()
         # Cross-request GPU batching is inferred from gpu_max_batch_size alone:
         # == 1 -> serial (one request per forward, no wait window); > 1 -> merge
         # compatible requests within gpu_batch_wait_ms. Default 1 keeps the old
@@ -344,6 +375,9 @@ def embedding_scheduler_args(self) -> Dict[str, int]:
         }
 
     def to_string(self):
+        transport = self.output_transport
+        control = transport.control
+        rdma = transport.rdma
         return (
             f"vit_separation: {self.vit_separation}\n"
             f"vit_trt: {self.vit_trt}\n"
@@ -368,6 +402,15 @@ def to_string(self):
             f"disable_access_log: {self.disable_access_log}\n"
             f"use_local_preprocess: {self.use_local_preprocess}\n"
             f"vit_proxy_load_balance_strategy: {self.vit_proxy_load_balance_strategy}\n"
+            f"mm_transport_mode: {transport.mode}\n"
+            f"mm_rdma_bind_ip: {rdma.bind_ip}\n"
+            f"mm_rdma_port: {rdma.port}\n"
+            f"mm_rdma_connect_timeout_ms: {rdma.connect_timeout_ms}\n"
+            f"mm_rdma_read_timeout_ms: {rdma.read_timeout_ms}\n"
+            f"mm_rdma_qp_count: {rdma.qp_count}\n"
+            f"mm_rdma_release_timeout_ms: {control.release_timeout_ms}\n"
+            f"mm_rdma_slot_gc_timeout_ms: {rdma.slot_gc_timeout_ms}\n"
+            f"mm_rdma_max_slot_bytes: {rdma.max_slot_bytes}\n"
             f"gpu_batch_wait_ms: {self.gpu_batch_wait_ms}\n"
             f"gpu_max_batch_size: {self.gpu_max_batch_size}\n"
             f"gpu_max_batch_images: {self.gpu_max_batch_images}\n"

rtp_llm/config/server_config_setup.py

@@ -2,6 +2,7 @@
 import logging
 import os
 import socket
+import subprocess
 from typing import Optional
 
 import torch
@@ -521,6 +522,144 @@ def fetch_model_files_to_local(py_env_configs: PyEnvConfigs):
     )
 
 
+def _physical_gpu_index(local_rank: int) -> Optional[str]:
+    visible_devices = os.environ.get("CUDA_VISIBLE_DEVICES")
+    if not visible_devices:
+        return str(local_rank)
+
+    devices = [device.strip() for device in visible_devices.split(",")]
+    if local_rank < 0 or local_rank >= len(devices) or not devices[local_rank]:
+        logging.warning(
+            "local rank %d is outside CUDA_VISIBLE_DEVICES=%s",
+            local_rank,
+            visible_devices,
+        )
+        return None
+
+    device = devices[local_rank]
+    if device.isdigit():
+        return device
+    if not device.startswith("GPU-"):
+        logging.warning("unsupported CUDA_VISIBLE_DEVICES entry: %s", device)
+        return None
+
+    try:
+        result = subprocess.run(
+            [
+                "nvidia-smi",
+                "--query-gpu=index,uuid",
+                "--format=csv,noheader,nounits",
+            ],
+            check=True,
+            stdout=subprocess.PIPE,
+            stderr=subprocess.PIPE,
+            text=True,
+            timeout=10,
+        )
+    except Exception as e:
+        logging.warning("failed to resolve CUDA device UUID %s: %s", device, e)
+        return None
+
+    for line in result.stdout.splitlines():
+        index, separator, uuid = line.partition(",")
+        if separator and uuid.strip().startswith(device):
+            return index.strip()
+    logging.warning("CUDA device UUID %s was not reported by nvidia-smi", device)
+    return None
+
+
+def load_gpu_nic_affinity() -> bool:
+    """Load the host GPU-NIC map before a service configures its CUDA device."""
+    if os.environ.get("ACCL_NIC_GPU_AFFINITY") is not None:
+        return True
+
+    run_affinity_path = "/usr/local/bin/run_affinity"
+    if not os.path.exists(run_affinity_path):
+        logging.info("get gpu nic affinity failed, %s not exist", run_affinity_path)
+        return False
+
+    try:
+        subprocess.run(
+            [run_affinity_path],
+            check=True,
+            stdout=subprocess.PIPE,
+            stderr=subprocess.PIPE,
+            timeout=30,
+        )
+        with open("npu_nic_affinity.json") as affinity_file:
+            content = affinity_file.read().strip()
+        affinity = json.loads(content)
+        if not isinstance(affinity, dict) or not affinity:
+            raise ValueError("affinity content is not a non-empty object")
+        if not all(
+            isinstance(rank, str) and isinstance(nic, str) and nic
+            for rank, nic in affinity.items()
+        ):
+            raise ValueError("affinity keys and NIC values must be non-empty strings")
+    except Exception as e:
+        logging.info(
+            "get gpu nic affinity failed, run %s failed, exception is %s",
+            run_affinity_path,
+            e,
+        )
+        return False
+
+    os.environ["ACCL_NIC_GPU_AFFINITY"] = content
+    logging.info(
+        "get gpu nic affinity success, set env ACCL_NIC_GPU_AFFINITY to %s",
+        content,
+    )
+    return True
+
+
+def _configure_gpu_nic_affinity(local_rank: int) -> bool:
+    configured_nics = os.environ.get("ACCL_USE_NICS")
+    if configured_nics:
+        logging.info("keep explicit ACCL_USE_NICS=%s", configured_nics)
+        return True
+
+    content = os.environ.get("ACCL_NIC_GPU_AFFINITY")
+    if not content:
+        return False
+
+    try:
+        affinity = json.loads(content)
+    except (json.JSONDecodeError, TypeError) as e:
+        logging.warning(
+            "invalid ACCL_NIC_GPU_AFFINITY; ACCL_USE_NICS will not be derived: %s",
+            e,
+        )
+        return False
+    if not isinstance(affinity, dict):
+        logging.warning(
+            "invalid ACCL_NIC_GPU_AFFINITY; ACCL_USE_NICS will not be derived: expected object"
+        )
+        return False
+
+    physical_gpu = _physical_gpu_index(local_rank)
+    affinity_nic = affinity.get(physical_gpu) if physical_gpu is not None else None
+    # Keep compatibility with existing mappings keyed by process-local rank.
+    if not affinity_nic:
+        affinity_nic = affinity.get(str(local_rank))
+    if not isinstance(affinity_nic, str) or not affinity_nic:
+        logging.warning(
+            "local rank %d (physical GPU %s) get affinity nic failed, content is %s",
+            local_rank,
+            physical_gpu,
+            content,
+        )
+        return False
+
+    os.environ["ACCL_USE_NICS"] = affinity_nic
+    logging.info(
+        "local rank %d maps to physical GPU %s, set ACCL_USE_NICS to %s",
+        local_rank,
+        physical_gpu,
+        affinity_nic,
+    )
+    return True
+
+
 def setup_cuda_device_and_accl_env(local_rank: int) -> None:
     """Apply CUDA device and ACCL env side effects (same as ParallelInfo.from_params)."""
     if torch.cuda.is_available():
@@ -531,27 +670,7 @@ def setup_cuda_device_and_accl_env(local_rank: int) -> None:
         os.environ["ACCL_SELECT_PORT"] = select_port
         logging.info(f"local rank {local_rank} set accl select port to {select_port} ")
 
-    if (
-        os.environ.get("ACCL_USE_NICS") is None
-        and os.environ.get("ACCL_NIC_GPU_AFFINITY") is not None
-    ):
-        content = os.environ.get("ACCL_NIC_GPU_AFFINITY")
-        try:
-            gpu_nic_affinity = json.loads(content)
-            if str(local_rank) in gpu_nic_affinity:
-                affinity_nic = gpu_nic_affinity[str(local_rank)]
-                os.environ["ACCL_USE_NICS"] = affinity_nic
-                logging.info(
-                    f"local rank {local_rank} use cuda device {local_rank} set ACCL_USE_NICS to {affinity_nic}"
-                )
-            else:
-                logging.info(
-                    f"local rank {local_rank} use cuda device {local_rank} get affinity nic failed, content is {content}"
-                )
-        except json.JSONDecodeError:
-            logging.info(
-                f"try decode ACCL_NIC_GPU_AFFINITY failed, content is {content}"
-            )
+    _configure_gpu_nic_affinity(local_rank)
 
 
 def setup_and_configure_server(py_env_configs: PyEnvConfigs):

rtp_llm/cpp/config/BUILD

@@ -17,6 +17,12 @@ cc_library(
     copts = copts(),
 )
 
+cc_library(
+    name = "rdma_config",
+    hdrs = ["RdmaConfig.h"],
+    visibility = ["//visibility:public"],
+)
+
 cc_library(
     name = "config_modules",
     hdrs = ["ConfigModules.h"],
@@ -25,6 +31,8 @@ cc_library(
         "@havenask//aios/autil:env_util",
         ":role_types",
         ":eplb_config",
+        ":mm_transport_mode",
+        ":rdma_config",
         "//rtp_llm/models_py/bindings/core:types",
     ],
     copts = copts(),
@@ -38,6 +46,13 @@ cc_library(
     copts = copts(),
 )
 
+cc_library(
+    name = "mm_transport_mode",
+    hdrs = ["MMTransportMode.h"],
+    visibility = ["//visibility:public"],
+    copts = copts(),
+)
+
 cc_library(
     name = "special_tokens",
     hdrs = ["SpecialTokens.h"],

rtp_llm/cpp/config/ConfigModules.cc

@@ -328,7 +328,15 @@ std::string VitConfig::to_string() const {
             vit_sep_str = "UNKNOWN(" + std::to_string(static_cast<int>(vit_separation)) + ")";
             break;
     }
-    oss << "vit_separation: " << vit_sep_str;
+    const auto& rdma = output_transport.rdma;
+    oss << "vit_separation: " << vit_sep_str << ", mm_transport_mode: " << output_transport.mode
+        << ", mm_rdma_bind_ip: " << rdma.bind_ip << ", mm_rdma_port: " << rdma.port
+        << ", mm_rdma_connect_timeout_ms: " << rdma.connect_timeout_ms
+        << ", mm_rdma_read_timeout_ms: " << rdma.read_timeout_ms
+        << ", mm_rdma_qp_count: " << rdma.qp_count
+        << ", mm_rdma_release_timeout_ms: " << output_transport.control.release_timeout_ms
+        << ", mm_rdma_slot_gc_timeout_ms: " << rdma.slot_gc_timeout_ms
+        << ", mm_rdma_max_slot_bytes: " << rdma.max_slot_bytes;
     return oss.str();
 }