[Feature][KV Transfer] Support multiple memfabric TransDataOpType and A5

docs/source/developer_guide/Design_Documents/layerwise_and_sparse_kv_cache_offloading.md

@@ -71,6 +71,12 @@ flowchart LR
 into Decode-owned host memory and indexer data into rank-local NPU memory
 through MemFabric. Optional LIC8 scale data follows the indexer destination.
 
+The MemFabric data path is selected at launch through
+`kv_connector_extra_config["memfabric_transfer_protocol"]` instead of hardware
+detection: `sdma` (default) and `device_rdma` target A3 series nodes, while
+`device_urma` targets A5 series nodes. Prefill and Decode must use the same
+protocol.
+
 ### Transfer granularity
 
 The data path uses different transfer granularities according to when the

docs/source/user_guide/feature_guide/layerwise_and_sparse_kv_cache_offloading.md

@@ -22,8 +22,9 @@ Other sparse-attention models have not been validated.
 
 ## 1. Install Dependencies
 
-The installation steps are grouped by hardware. Only A3 series is currently
-supported.
+The installation steps are grouped by hardware. A3 and A5 series are supported.
+On A5 nodes, set the MemFabric transfer protocol to `device_urma` as described
+in sections 2 and 3.
 
 ### Prefill Build Dependencies
 
@@ -194,6 +195,13 @@ Do not set `sparse_kv_offload_config` on Prefill. The
 | `layerwise_num_shared_buffers` | Number of reusable NPU buffers. Start with two to four and tune for memory and transfer bandwidth. |
 | `layerwise_independent_layers` | Layers that keep dedicated buffers. The default is `[0]`; `"all"` disables cross-layer reuse. |
 
+The `SfaRemoteD2HConnector` entry accepts the following options:
+
+| Parameter | Description |
+| :--- | :--- |
+| `transfer_backend` | Transfer backend. `memfabric` is the only supported value. |
+| `memfabric_transfer_protocol` | MemFabric data-path protocol: `sdma` (default) and `device_rdma` for A3 series, `device_urma` for A5 series. Must be set to the same value on Prefill and Decode. Invalid values abort startup. |
+
 The following log confirms that buffer reuse is enabled:
 
 ```text
@@ -233,6 +241,9 @@ On Decode, reserve
 `decode_data_parallel_size * decode_tensor_parallel_size` consecutive ports
 starting from `kv_port`.
 
+On A5 nodes, add `"memfabric_transfer_protocol": "device_urma"` to
+`kv_connector_extra_config` on both Prefill and Decode.
+
 | Parameter | Description |
 | :--- | :--- |
 | `topk_buffer_size` | Device hot-buffer size. It must be at least `index_topk` and divisible by `block_size`. Twice `index_topk` is a practical starting point. |
@@ -263,6 +274,9 @@ For multi-node deployment, advertise reachable addresses instead of
 - Context parallelism has not been validated with Layerwise Prefill Offload.
 - Sparse Decode Offload supports DP and TP; CP and PP are not supported.
 - MemFabric is the only supported `SfaRemoteD2HConnector` transfer backend.
+- The MemFabric data-path protocol is selected by launch configuration instead
+  of hardware detection: use `sdma` (default) or `device_rdma` on A3 series and
+  `device_urma` on A5 series, identically on Prefill and Decode.
 - Layerwise buffer reuse cannot currently be combined with
   `MooncakeLayerwiseConnector` because per-buffer transfer completion gating is
   not yet implemented. Support is planned in a follow-up update.

tests/ut/kv_offload/test_memfabric_transfer_engine.py

@@ -1,6 +1,7 @@
 """Unit tests for the isolated MemFabric transfer-engine singleton."""
 
 import sys
+from enum import Enum
 from types import ModuleType
 from unittest.mock import MagicMock, call, patch
 
@@ -15,11 +16,18 @@
 )
 
 
+class _FakeTransDataOpType(Enum):
+    SDMA = 1
+    DEVICE_RDMA = 2
+    DEVICE_URMA = 3
+
+
 def _fake_memfabric(raw_engine: MagicMock) -> ModuleType:
     module = ModuleType("memfabric_hybrid")
     module.TransferEngine = MagicMock(return_value=raw_engine)  # type: ignore[attr-defined]
     module.set_conf_store_tls = MagicMock()  # type: ignore[attr-defined]
     module.set_log_level = MagicMock()  # type: ignore[attr-defined]
+    module.TransDataOpType = _FakeTransDataOpType  # type: ignore[attr-defined]
     return module
 
 
@@ -44,6 +52,7 @@ def test_memfabric_initialization_publishes_session_from_engine_port():
         MEMFABRIC_ROLE_PREFILL,
         3,
         store_server_role=MEMFABRIC_ROLE_PREFILL,
+        data_op_type=_FakeTransDataOpType.SDMA,
     )
     assert raw_engine.method_calls[:2] == [
         call.initialize(
@@ -52,6 +61,7 @@ def test_memfabric_initialization_publishes_session_from_engine_port():
             MEMFABRIC_ROLE_PREFILL,
             3,
             store_server_role=MEMFABRIC_ROLE_PREFILL,
+            data_op_type=_FakeTransDataOpType.SDMA,
         ),
         call.get_rpc_port(),
     ]
@@ -60,13 +70,15 @@ def test_memfabric_initialization_publishes_session_from_engine_port():
 def test_memfabric_configuration_is_idempotent_but_role_bound():
     manager = GlobalMemfabricTE()
 
-    manager.configure(role=MEMFABRIC_ROLE_DECODE, device_id=0)
-    manager.configure(role=MEMFABRIC_ROLE_DECODE, device_id=0)
+    manager.configure(role=MEMFABRIC_ROLE_DECODE, device_id=0, transfer_protocol="device_urma")
+    manager.configure(role=MEMFABRIC_ROLE_DECODE, device_id=0, transfer_protocol="device_urma")
 
     with pytest.raises(RuntimeError, match="already configured"):
-        manager.configure(role=MEMFABRIC_ROLE_PREFILL, device_id=0)
+        manager.configure(role=MEMFABRIC_ROLE_PREFILL, device_id=0, transfer_protocol="device_urma")
     with pytest.raises(RuntimeError, match="already configured"):
-        manager.configure(role=MEMFABRIC_ROLE_DECODE, device_id=1)
+        manager.configure(role=MEMFABRIC_ROLE_DECODE, device_id=1, transfer_protocol="device_urma")
+    with pytest.raises(RuntimeError, match="already configured"):
+        manager.configure(role=MEMFABRIC_ROLE_DECODE, device_id=0, transfer_protocol="sdma")
 
 
 def test_memfabric_engine_is_bound_to_initial_hostname():
@@ -122,3 +134,45 @@ def test_memfabric_initialization_failure_does_not_publish_session():
         manager.get_transfer_engine("127.0.0.1")
     with pytest.raises(RuntimeError, match="has not been initialized"):
         _ = manager.unique_id
+
+
+def test_memfabric_transfer_protocol_defaults_to_sdma():
+    raw_engine = MagicMock()
+    raw_engine.get_rpc_port.return_value = 23456
+    raw_engine.initialize.return_value = 0
+    manager = GlobalMemfabricTE()
+    manager.configure(role=MEMFABRIC_ROLE_DECODE, device_id=0)
+
+    with patch.dict(sys.modules, {"memfabric_hybrid": _fake_memfabric(raw_engine)}):
+        manager.get_transfer_engine("127.0.0.1")
+
+    assert raw_engine.initialize.call_args.kwargs["data_op_type"] == _FakeTransDataOpType.SDMA
+
+
+@pytest.mark.parametrize(
+    "protocol,expected",
+    [
+        ("sdma", _FakeTransDataOpType.SDMA),
+        ("device_rdma", _FakeTransDataOpType.DEVICE_RDMA),
+        ("device_urma", _FakeTransDataOpType.DEVICE_URMA),
+        (" DEVICE_URMA ", _FakeTransDataOpType.DEVICE_URMA),
+    ],
+)
+def test_memfabric_transfer_protocol_from_configure(protocol, expected):
+    raw_engine = MagicMock()
+    raw_engine.get_rpc_port.return_value = 23456
+    raw_engine.initialize.return_value = 0
+    manager = GlobalMemfabricTE()
+    manager.configure(role=MEMFABRIC_ROLE_PREFILL, device_id=1, transfer_protocol=protocol)
+
+    with patch.dict(sys.modules, {"memfabric_hybrid": _fake_memfabric(raw_engine)}):
+        manager.get_transfer_engine("127.0.0.1")
+
+    assert raw_engine.initialize.call_args.kwargs["data_op_type"] == expected
+
+
+def test_memfabric_transfer_protocol_rejects_unknown_value():
+    manager = GlobalMemfabricTE()
+
+    with pytest.raises(ValueError, match="Invalid MemFabric transfer_protocol"):
+        manager.configure(role=MEMFABRIC_ROLE_DECODE, device_id=0, transfer_protocol="rocev2")

vllm_ascend/distributed/kv_transfer/kv_p2p/sfa_pd_rd2h/worker.py

@@ -95,6 +95,16 @@ def _resolve_kv_transfer_backend(vllm_config: VllmConfig) -> str:
     return backend
 
 
+def _resolve_memfabric_transfer_protocol(vllm_config: VllmConfig) -> str | None:
+    """Read the optional MemFabric data-path protocol.
+
+    Read from ``kv_connector_extra_config["memfabric_transfer_protocol"]``:
+    ``sdma``/``device_rdma`` for A3 nodes, ``device_urma`` for A5 nodes.
+    """
+    extra = vllm_config.kv_transfer_config.kv_connector_extra_config or {}
+    return extra.get("memfabric_transfer_protocol")
+
+
 def _validate_tcp_port(port: int, *, description: str) -> None:
     if not MIN_TCP_PORT <= port <= MAX_TCP_PORT:
         raise ValueError(f"{description} must be in [{MIN_TCP_PORT}, {MAX_TCP_PORT}], got {port}")
@@ -158,6 +168,7 @@ def _ensure_engine(self):
             global_memfabric_te.configure(
                 role=MEMFABRIC_ROLE_DECODE,
                 device_id=torch.npu.current_device(),
+                transfer_protocol=_resolve_memfabric_transfer_protocol(self.vllm_config),
             )
             self.engine = global_memfabric_te.get_transfer_engine(self.side_channel_host)
         return self.engine
@@ -444,6 +455,7 @@ def __init__(self, vllm_config: VllmConfig, kv_cache_config: KVCacheConfig, engi
         global_memfabric_te.configure(
             role=MEMFABRIC_ROLE_PREFILL,
             device_id=torch.npu.current_device(),
+            transfer_protocol=_resolve_memfabric_transfer_protocol(vllm_config),
         )
         self.vllm_config = vllm_config
         self.kv_cache_config = kv_cache_config

vllm_ascend/distributed/kv_transfer/utils/memfabric_transfer_engine.py

@@ -17,6 +17,11 @@
 MEMFABRIC_ROLE_PREFILL = "Prefill"
 MEMFABRIC_ROLE_DECODE = "Decode"
 _VALID_MEMFABRIC_ROLES = (MEMFABRIC_ROLE_PREFILL, MEMFABRIC_ROLE_DECODE)
+# MemFabric data-path protocol, selected via kv_connector_extra_config so this
+# module never needs to branch on machine type: sdma/device_rdma for A3 nodes,
+# device_urma for A5 nodes.
+_VALID_MEMFABRIC_TRANSFER_PROTOCOLS = ("sdma", "device_rdma", "device_urma")
+_DEFAULT_MEMFABRIC_TRANSFER_PROTOCOL = "sdma"
 
 
 class MemfabricBackend:
@@ -62,6 +67,7 @@ def __init__(self):
         self._engine: MemfabricBackend | None = None
         self._role: str | None = None
         self._device_id: int | None = None
+        self._transfer_protocol: str | None = None
         self._hostname: str | None = None
         self._unique_id: str | None = None
         self._is_buffer_registered = False
@@ -74,23 +80,50 @@ def unique_id(self) -> str:
             raise RuntimeError("MemFabric transfer engine has not been initialized")
         return self._unique_id
 
-    def configure(self, *, role: str, device_id: int) -> None:
-        """Bind this process singleton to one MemFabric role and device."""
+    def configure(
+        self,
+        *,
+        role: str,
+        device_id: int,
+        transfer_protocol: str | None = None,
+    ) -> None:
+        """Bind this process singleton to one MemFabric role, device and protocol.
+
+        ``transfer_protocol`` comes from
+        ``kv_connector_extra_config["memfabric_transfer_protocol"]`` and selects
+        the data path: ``sdma``/``device_rdma`` for A3 nodes, ``device_urma``
+        for A5 nodes. It defaults to ``sdma``, the memfabric_hybrid library
+        default, so launch scripts pick the value per machine type instead of
+        this module branching on hardware.
+        """
         if role not in _VALID_MEMFABRIC_ROLES:
             raise ValueError(f"Invalid MemFabric role {role!r}; expected one of {_VALID_MEMFABRIC_ROLES}")
         if device_id < 0:
             raise ValueError(f"MemFabric device_id must be non-negative, got {device_id}")
+        protocol = (transfer_protocol or _DEFAULT_MEMFABRIC_TRANSFER_PROTOCOL).strip().lower()
+        if protocol not in _VALID_MEMFABRIC_TRANSFER_PROTOCOLS:
+            raise ValueError(
+                f"Invalid MemFabric transfer_protocol={transfer_protocol!r}; "
+                f"expected one of {_VALID_MEMFABRIC_TRANSFER_PROTOCOLS}"
+            )
 
         with self._engine_lock:
             configured = self._role is not None
-            if configured and (role, device_id) != (self._role, self._device_id):
+            if configured and (role, device_id, protocol) != (
+                self._role,
+                self._device_id,
+                self._transfer_protocol,
+            ):
                 raise RuntimeError(
                     "MemFabric transfer engine is already configured for "
-                    f"role={self._role}, device_id={self._device_id}; cannot "
-                    f"reconfigure it for role={role}, device_id={device_id}"
+                    f"role={self._role}, device_id={self._device_id}, "
+                    f"transfer_protocol={self._transfer_protocol}; cannot "
+                    f"reconfigure it for role={role}, device_id={device_id}, "
+                    f"transfer_protocol={protocol}"
                 )
             self._role = role
             self._device_id = device_id
+            self._transfer_protocol = protocol
 
     def get_transfer_engine(self, hostname: str) -> MemfabricBackend:
         with self._engine_lock:
@@ -105,6 +138,26 @@ def get_transfer_engine(self, hostname: str) -> MemfabricBackend:
                 )
             return self._engine
 
+    def _get_transfer_protocol(self):
+        """Map the configured protocol name to its ``TransDataOpType`` value.
+
+        Unknown names fail fast here as well: a wrong protocol otherwise only
+        surfaces later as an obscure engine initialization failure.
+        """
+        from memfabric_hybrid import TransDataOpType  # type: ignore
+
+        protocol_map = {
+            "sdma": TransDataOpType.SDMA,
+            "device_rdma": TransDataOpType.DEVICE_RDMA,
+            "device_urma": TransDataOpType.DEVICE_URMA,
+        }
+        protocol = self._transfer_protocol or _DEFAULT_MEMFABRIC_TRANSFER_PROTOCOL
+        if protocol not in protocol_map:
+            raise ValueError(
+                f"Invalid MemFabric transfer_protocol={protocol!r}; expected one of {sorted(protocol_map)}"
+            )
+        return protocol_map[protocol]
+
     def _build_engine(self, hostname: str) -> MemfabricBackend:
         try:
             from memfabric_hybrid import (  # type: ignore
@@ -123,19 +176,22 @@ def _build_engine(self, hostname: str) -> MemfabricBackend:
         raw_engine = TransferEngine()
         store_url = f"tcp://{hostname}"
 
+        data_op_type = self._get_transfer_protocol()
         logger.info(
-            "MemFabric TransferEngine initialize: store_url=%s, unique_id=%s, role=%s, device_id=%s",
+            "MemFabric TransferEngine initialize: store_url=%s, unique_id=%s, role=%s, device_id=%s, data_op_type=%s",
             store_url,
             hostname,
             self._role,
             self._device_id,
+            getattr(data_op_type, "name", data_op_type),
         )
         ret = raw_engine.initialize(
             store_url,
             hostname,
             self._role,
             self._device_id,
             store_server_role=MEMFABRIC_ROLE_PREFILL,
+            data_op_type=data_op_type,
         )
         if ret != 0:
             raise RuntimeError(