From 62c7995aad402463f3e3a434a002a40520148c23 Mon Sep 17 00:00:00 2001
From: Robert Washbourne <rawashbourne@gmail.com>
Date: Tue, 11 Aug 2026 08:53:55 -0700
Subject: [PATCH] Support current vLLM stateless process groups

Preserve compatibility with older vLLM releases that require the socket argument while constructing newer vLLM process groups without it.

Co-authored-by: OpenAI Codex <codex@openai.com>
---
 checkpoint_engine/distributed/vllm_compat.py | 16 +++++++++
 checkpoint_engine/distributed/vllm_hccl.py   |  8 ++++-
 checkpoint_engine/distributed/vllm_nccl.py   |  8 ++++-
 tests/test_vllm_compat.py                    | 37 ++++++++++++++++++++
 4 files changed, 67 insertions(+), 2 deletions(-)
 create mode 100644 checkpoint_engine/distributed/vllm_compat.py
 create mode 100644 tests/test_vllm_compat.py

diff --git a/checkpoint_engine/distributed/vllm_compat.py b/checkpoint_engine/distributed/vllm_compat.py
new file mode 100644
index 0000000..1649796
--- /dev/null
+++ b/checkpoint_engine/distributed/vllm_compat.py
@@ -0,0 +1,16 @@
+import inspect
+from typing import Any
+
+
+def create_stateless_process_group(
+    group_cls: type,
+    *,
+    rank: int,
+    world_size: int,
+    store: Any,
+) -> Any:
+    """Construct a vLLM stateless process group across supported APIs."""
+    kwargs = {"rank": rank, "world_size": world_size, "store": store}
+    if "socket" in inspect.signature(group_cls).parameters:
+        kwargs["socket"] = None
+    return group_cls(**kwargs)
diff --git a/checkpoint_engine/distributed/vllm_hccl.py b/checkpoint_engine/distributed/vllm_hccl.py
index fbdab0c..b1e2db1 100644
--- a/checkpoint_engine/distributed/vllm_hccl.py
+++ b/checkpoint_engine/distributed/vllm_hccl.py
@@ -19,6 +19,7 @@
 from vllm_ascend.utils import current_stream
 
 from checkpoint_engine.distributed.base import CommGroup, Distributed, _common_all_gather_object
+from checkpoint_engine.distributed.vllm_compat import create_stateless_process_group
 
 
 class HcclCommConfig(ctypes.Structure):
@@ -240,7 +241,12 @@ def init_process_group(
         self.world_size = world_size
         self.device = torch.device("npu", torch.npu.current_device())
 
-        self.pg = StatelessProcessGroup(rank=rank, world_size=world_size, store=store, socket=None)
+        self.pg = create_stateless_process_group(
+            StatelessProcessGroup,
+            rank=rank,
+            world_size=world_size,
+            store=store,
+        )
         self.pyhccl = PyHcclCommunicatorEx(group=self.pg, device=self.device)
         self.comm = self.pyhccl.comm
         self.initialized = True
diff --git a/checkpoint_engine/distributed/vllm_nccl.py b/checkpoint_engine/distributed/vllm_nccl.py
index b2eb1aa..713e6c9 100644
--- a/checkpoint_engine/distributed/vllm_nccl.py
+++ b/checkpoint_engine/distributed/vllm_nccl.py
@@ -14,6 +14,7 @@
 from vllm.distributed.utils import StatelessProcessGroup
 
 from checkpoint_engine.distributed.base import CommGroup, Distributed, _common_all_gather_object
+from checkpoint_engine.distributed.vllm_compat import create_stateless_process_group
 
 
 try:
@@ -153,7 +154,12 @@ def init_process_group(
         self.world_size = world_size
         self.device = torch.device("cuda", torch.cuda.current_device())
 
-        self.pg = StatelessProcessGroup(rank=rank, world_size=world_size, store=store, socket=None)
+        self.pg = create_stateless_process_group(
+            StatelessProcessGroup,
+            rank=rank,
+            world_size=world_size,
+            store=store,
+        )
         self.pynccl = PyNcclCommunicatorEx(group=self.pg, device=self.device)
         self.comm = self.pynccl.comm
         self.initialized = True
diff --git a/tests/test_vllm_compat.py b/tests/test_vllm_compat.py
new file mode 100644
index 0000000..1ce6569
--- /dev/null
+++ b/tests/test_vllm_compat.py
@@ -0,0 +1,37 @@
+from checkpoint_engine.distributed.vllm_compat import create_stateless_process_group
+
+
+class CurrentProcessGroup:
+    def __init__(self, rank: int, world_size: int, store: object):
+        self.arguments = rank, world_size, store
+
+
+class LegacyProcessGroup:
+    def __init__(self, rank: int, world_size: int, store: object, socket: object):
+        self.arguments = rank, world_size, store, socket
+
+
+def test_current_vllm_process_group_signature():
+    store = object()
+
+    group = create_stateless_process_group(
+        CurrentProcessGroup,
+        rank=2,
+        world_size=8,
+        store=store,
+    )
+
+    assert group.arguments == (2, 8, store)
+
+
+def test_legacy_vllm_process_group_signature():
+    store = object()
+
+    group = create_stateless_process_group(
+        LegacyProcessGroup,
+        rank=2,
+        world_size=8,
+        store=store,
+    )
+
+    assert group.arguments == (2, 8, store, None)