From 549adf5adfe9a599fcc2a9254d09a980e9d3951b Mon Sep 17 00:00:00 2001 From: Ketor Date: Fri, 28 Aug 2026 09:47:57 +0800 Subject: [PATCH] fix(vllm): support heterogeneous mamba cache groups --- CHANGELOG.md | 11 ++++++ docs/CONNECTORS.md | 7 ++-- integration/vllm/src/dfkv_vllm/connector.py | 9 +---- integration/vllm/src/dfkv_vllm/dfkv_client.py | 3 +- integration/vllm/src/dfkv_vllm/worker.py | 1 + .../vllm/tests/test_connector_validation.py | 31 +++++++++++++++ integration/vllm/tests/test_dfkv_client.py | 38 +++++++++++++++++++ 7 files changed, 88 insertions(+), 12 deletions(-) create mode 100644 integration/vllm/tests/test_connector_validation.py diff --git a/CHANGELOG.md b/CHANGELOG.md index 300d6f23..a21081b5 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -2,6 +2,17 @@ ## Unreleased +### vLLM hybrid state and explicit TCP staging + +- Accepted align-mode hybrid cache groups whose Mamba block size differs from + the attention cache block size. The connector already resolves the scheduler + LCM and preserves each group's physical block geometry; the old equality + gate incorrectly rejected Qwen3.8-Flash-Next (`400` vs `4`) at startup. +- Added opt-in `require_rdma=false` for the vLLM connector. The default remains + fail-closed GPUDirect RDMA; the opt-in path uses libdfkv's bounded host + staging plus final CUDA publication and enables a correctness-first fallback + when a platform's inbound GPUDirect GET path is unusable. + ### LMCache rank-local rail affinity - Added opt-in `rail_affinity` and `rail_affinity_fallbacks` plugin settings to diff --git a/docs/CONNECTORS.md b/docs/CONNECTORS.md index b8e72f5e..1f2918df 100644 --- a/docs/CONNECTORS.md +++ b/docs/CONNECTORS.md @@ -750,7 +750,7 @@ namespace/key 不一致是预期 cold miss。**空环 / MDS 不可达**可直接 | env | 默认 | 推荐 | 说明 | |---|---|---|---| -| `DFKV_RDMA` / `DFKV_RDMA_DEV` | **无；`DFKV_RDMA=1` 必填** | `1` / 本机有序全轨列表 | vLLM GPU 指针连接器仅支持 RDMA；`rail_affinity=true` 时 connector 在每个 worker 内收窄列表 | +| `DFKV_RDMA` / `DFKV_RDMA_DEV` | 无 | `1` / 本机有序全轨列表 | 默认 `require_rdma=true` 时必须启用可用 RDMA；显式 `require_rdma=false` 可用 TCP staging，`rail_affinity=true` 时 connector 在每个 worker 内收窄 rail 列表 | | `DFKV_RDMA_DEPTH` | `4` | 保持生产已验证值 | depth-flat（§1.2） | | `DFKV_RDMA_NUMA` | `0` | affinity 开启时保持 `0` | connector 为严格 primary/fallback 映射显式设置 0；affinity 关闭时才按需启用 NUMA 动态选轨 | | `DFKV_LIB` / `DFKV_BUILD` | — | so 路径 | 被 extra_config `lib` 覆盖 | @@ -767,12 +767,13 @@ namespace/key 不一致是预期 cold miss。**空环 / MDS 不可达**可直接 | `mds_endpoints` | — | `ip:port,...`（dfkv_mds 层） | **生产首选**；设了即走 MDS 动态发现，省略 `members` | | `mds_group` | `default` | 如 `glm` | MDS 成员组名，= `dfkv_server --group` | | `mds_poll_ms` | `3000` | 默认即可 | MDS 轮询间隔（ms） | -| `members` | —（与 mds_endpoints 二选一） | `n=ip:rdma-port,...` | **端口 = server `--rdma-port`** | +| `members` | —（与 mds_endpoints 二选一） | `n=ip:port,...` | `require_rdma=true` 时端口 = server `--rdma-port`；`false` 时端口 = server `--port` | +| `require_rdma` | `True` | 生产保持 True | `False` 显式允许 TCP host staging + CUDA publication；保留正确性但失去 GPUDirect zero-copy，适合 inbound RDMA GET 不可用时降级 | | `lib` | env 兜底 | so 绝对路径 | | | `batch_concurrency` | `8` | **大池可调高到 ≈ 节点数** | 跨节点 fan-out，**真正的吞吐杠杆**（depth 是平的） | | `rail_affinity` | `False` | 多 rank、多 rail 生产设 `true` | 按 vLLM world-group per-host `local_rank` 选择 primary；在 native client 创建前设置每进程独立 rail 环境 | | `rail_affinity_fallbacks` | `1` | `1` | 相邻有序 fallback 数；`0`=严格单 rail，超出可用 rail 数时自动收敛 | -| `load_async` | `True` | 保持 True | 异步 load，走 `WAITING_FOR_REMOTE_KVS`、不占关键路径 | +| `load_async` | `True` | 普通 attention 保持 True；hybrid recurrent 模型设 `False` | `False` 在 forward 前同步完成 load，避免 recurrent-state compute 与远端 GPU 写重叠 | | `transfer_queue_capacity` | `256` | 保持默认，按压测调 | 每个 worker、每个方向的排队上限（`1..65536`）。满队列时非阻塞拒绝新任务：save 立即释放 finish/free fence，load 标记失败并重算；非法值启动即失败。 | | `enable_cross_layers_blocks` | `False` | 默认 False | 仅当引擎分页布局层内交错时开 | | `lookup_rpc_port` | ipc 自动 | 一般不设 | rank0 前缀查询 RPC，仅 socket 名冲突时设 | diff --git a/integration/vllm/src/dfkv_vllm/connector.py b/integration/vllm/src/dfkv_vllm/connector.py index f58416c4..98494f2d 100644 --- a/integration/vllm/src/dfkv_vllm/connector.py +++ b/integration/vllm/src/dfkv_vllm/connector.py @@ -103,7 +103,7 @@ def prefer_cross_layer_blocks(self) -> bool: def _validate_kv_cache_config( vllm_config: VllmConfig, kv_cache_config: KVCacheConfig ) -> None: - from vllm.v1.kv_cache_interface import CrossAttentionSpec, MambaSpec + from vllm.v1.kv_cache_interface import CrossAttentionSpec unsupported: list[str] = [] cache_block_size = vllm_config.cache_config.block_size @@ -111,13 +111,6 @@ def _validate_kv_cache_config( spec = g.kv_cache_spec if isinstance(spec, CrossAttentionSpec): unsupported.append(f"group {g_idx}: CrossAttentionSpec") - # Enforce Mamba align mode - if isinstance(spec, MambaSpec) and spec.block_size != cache_block_size: - unsupported.append( - f"group {g_idx}: MambaSpec with block_size=" - f"{spec.block_size} != cache_config.block_size=" - f"{cache_block_size} (mamba_cache_mode != 'align')" - ) # NOTE: multi-group (hybrid attention, e.g. GLM-5.2 DSA = MLA + sparse # indexer groups) together with PCP/DCP > 1 is now supported: each group # is normalized to scheduler_block_size in the worker and every rank diff --git a/integration/vllm/src/dfkv_vllm/dfkv_client.py b/integration/vllm/src/dfkv_vllm/dfkv_client.py index d6100e5a..0dc013dd 100644 --- a/integration/vllm/src/dfkv_vllm/dfkv_client.py +++ b/integration/vllm/src/dfkv_vllm/dfkv_client.py @@ -252,6 +252,7 @@ def __init__( client_heartbeat_ms: int = 10000, model: str = "", cache_role: str = "", + require_rdma: bool = True, ): if not key_namespace: raise ValueError("DfkvDeviceClient requires a non-empty key namespace") @@ -297,7 +298,7 @@ def __init__( "dfkv_transport_mode failed; set DFKV_RDMA=1 and configure " "a usable RDMA device" ) from exc - if mode != b"rdma": + if require_rdma and mode != b"rdma": rejected_handle = self._h self._h = None self._lib.dfkv_close(rejected_handle) diff --git a/integration/vllm/src/dfkv_vllm/worker.py b/integration/vllm/src/dfkv_vllm/worker.py index cddd9857..03632ed7 100644 --- a/integration/vllm/src/dfkv_vllm/worker.py +++ b/integration/vllm/src/dfkv_vllm/worker.py @@ -1540,6 +1540,7 @@ def __init__( client_info=client_info, model=model_identity, cache_role=str(self.kv_role), + require_rdma=_tcfg.truthy(extra.get("require_rdma", True)), ) # Phase 2a (issue #111): producer non-participants skip the client # (opt-in via DFKV_CONNECTOR_CLIENT_ELIDE=1; layout-clamped). The saved diff --git a/integration/vllm/tests/test_connector_validation.py b/integration/vllm/tests/test_connector_validation.py new file mode 100644 index 00000000..4002bd88 --- /dev/null +++ b/integration/vllm/tests/test_connector_validation.py @@ -0,0 +1,31 @@ +"""Startup validation for vLLM cache-group layouts.""" + +import sys +from pathlib import Path +from types import SimpleNamespace + +import torch + +sys.path.insert(0, str(Path(__file__).resolve().parents[1] / "src")) +from vllm.v1.kv_cache_interface import KVCacheGroupSpec, MambaSpec # noqa: E402 + +from dfkv_vllm.connector import DfkvStoreConnector # noqa: E402 + + +def test_accepts_align_mamba_groups_with_distinct_block_size(): + spec = MambaSpec( + block_size=400, + shapes=((1,),), + dtypes=(torch.float16,), + mamba_cache_mode="align", + ) + vllm_config = SimpleNamespace( + cache_config=SimpleNamespace(block_size=4, mamba_cache_mode="align") + ) + kv_cache_config = SimpleNamespace( + kv_cache_groups=[KVCacheGroupSpec(["model.layers.0.mamba"], spec)] + ) + + DfkvStoreConnector._validate_kv_cache_config( + vllm_config, kv_cache_config + ) diff --git a/integration/vllm/tests/test_dfkv_client.py b/integration/vllm/tests/test_dfkv_client.py index 3fead591..7332f09b 100644 --- a/integration/vllm/tests/test_dfkv_client.py +++ b/integration/vllm/tests/test_dfkv_client.py @@ -82,6 +82,44 @@ def dfkv_close(self, _handle): assert captured["close_calls"] == 1 +def test_explicit_tcp_staging_mode_is_allowed(monkeypatch): + closed = [] + + class FakeLib: + def dfkv_open_v2(self, _ptr): + return 0xBEEF + + def dfkv_transport_mode(self, _handle): + return b"tcp(rdma-not-requested)" + + def dfkv_version(self): + return b"2.0.0" + + def dfkv_close(self, handle): + closed.append(handle) + + monkeypatch.setattr(client_module, "load_lib", lambda _path: FakeLib()) + monkeypatch.setattr( + client_module._push_metrics, "configure", lambda *a, **k: None) + monkeypatch.setattr( + client_module._push_tracing, "configure", lambda *a, **k: None) + monkeypatch.setattr(client_module._alog, "configure", lambda *a, **k: None) + monkeypatch.setattr(client_module._hot_config, "register", lambda *a, **k: None) + monkeypatch.setattr(client_module._hot_config, "start", lambda *a, **k: None) + monkeypatch.setattr(client_module._hot_config, "stop", lambda *a, **k: None) + monkeypatch.setenv("DFKV_CLIENT_STATS_POLL_S", "0") + + client = DfkvDeviceClient( + members="n1=127.0.0.1:28000", + key_namespace=b"dfkv/model/v1/test/tcp-staging", + require_rdma=False, + ) + + assert client.transport_mode == "tcp(rdma-not-requested)" + client.close() + assert closed == [0xBEEF] + + def test_telemetry_setup_failure_releases_handle_and_lifecycle(monkeypatch): calls = []