feat: Integrate vLLM with Weight Propagation Interface (WPI) for Zero-Copy Weight Transfer

docs/training/rlhf.md

@@ -16,7 +16,7 @@ The following open-source RL libraries use vLLM for fast rollouts (sorted alphab
 - [Unsloth](https://github.com/unslothai/unsloth)
 - [verl](https://github.com/volcengine/verl)
 
-For weight synchronization between training and inference, see the [Weight Transfer](weight_transfer/README.md) documentation, which covers the pluggable backend system with [NCCL](weight_transfer/nccl.md) (multi-GPU) and [IPC](weight_transfer/ipc.md) (same-GPU) engines.
+For weight synchronization between training and inference, see the [Weight Transfer](weight_transfer/README.md) documentation, which covers the pluggable backend system with [NCCL](weight_transfer/nccl.md) (multi-GPU), [IPC](weight_transfer/ipc.md) (same-GPU), and [WPI](weight_transfer/wpi.md) (zero-copy cross-node) engines.
 
 For pipelining generation and training to improve GPU utilization and throughput, see the [Async Reinforcement Learning](async_rl.md) guide, which covers the pause/resume API for safely updating weights mid-flight.
 

docs/training/weight_transfer/README.md

@@ -18,6 +18,7 @@ The weight transfer system follows a **four-phase protocol** with a pluggable ba
 | [NCCL](nccl.md) | NCCL broadcast | Separate GPUs for training and inference |
 | [IPC](ipc.md) | CUDA IPC handles | Colocated training and inference on same GPU |
 | [sparse_nccl](nccl.md#sparse-nccl) | NCCL broadcast | Sparse flat-index weight patches (TP=1/PP=1) |
+| [WPI](wpi.md) | WPI / NCCL | Zero-copy multi-node weight propagation |
 
 ## Configuration
 
@@ -31,7 +32,7 @@ from vllm.config import WeightTransferConfig
 
 llm = LLM(
     model="my-model",
-    weight_transfer_config=WeightTransferConfig(backend="nccl"),  # or "ipc"
+    weight_transfer_config=WeightTransferConfig(backend="nccl"),  # or "ipc" or "wpi"
 )
 ```
 
@@ -42,7 +43,7 @@ vllm serve my-model \
     --weight-transfer-config '{"backend": "nccl"}'
 ```
 
-The `backend` field accepts `"nccl"` (default), `"ipc"`, or `"sparse_nccl"`.
+The `backend` field accepts `"nccl"` (default), `"ipc"`, `"sparse_nccl"`, or `"wpi"`.
 
 ## API Endpoints
 
@@ -63,7 +64,7 @@ When running vLLM as an HTTP server, the following endpoints are available for w
 
 ## Trainer-Side API
 
-Both backends provide static methods that the trainer calls to send weights. The general pattern is:
+All backends provide static methods that the trainer calls to send weights. The general pattern is:
 
 ```python
 # 1. Initialize the transfer engine (backend-specific)
@@ -82,7 +83,7 @@ EngineClass.trainer_send_weights(
 llm.finish_weight_update()
 ```
 
-See the [NCCL](nccl.md) and [IPC](ipc.md) pages for backend-specific trainer APIs and full examples.
+See the [NCCL](nccl.md), [IPC](ipc.md), and [WPI](wpi.md) pages for backend-specific trainer APIs and full examples.
 
 ## Extending the System
 

docs/training/weight_transfer/wpi.md

@@ -0,0 +1,127 @@
+# WPI Engine
+
+The WPI (Weight Propagation Interface) weight transfer engine uses the [Weight Propagation Interface](https://github.com/llm-d-incubation/weight-propagation-interface) to enable high-throughput, cross-node zero-copy weight updates.
+
+Unlike the [NCCL](nccl.md) or [IPC](ipc.md) engines where vLLM itself manages communication channels and CPU/GPU memory allocation, the WPI backend relies on a persistent VRAM buffer and an internal NCCL communicator managed externally by a node-level WPI driver.
+
+## When to Use WPI
+
+- Training and inference run on **separate GPUs** across one or more nodes.
+- You have large models (e.g. 70B+ parameters) and need to avoid the GPU memory allocation overhead of creating duplicate parameter tensors during updates.
+- You want to utilize high-bandwidth hardware fabrics (like InfiniBand or RoCE) at near-line rate.
+- You are using Kubernetes and have the WPI operator and driver daemonset deployed in your cluster.
+
+## How It Works
+
+1. **Persistent VRAM Buffer**: The WPI driver manages a pre-allocated VRAM buffer on each GPU node. This buffer is persistent and reused across training/sync iterations (avoiding allocation/deallocation overhead).
+2. **Zero-Copy Memory Sharing**: The local WPI driver passes the CUDA memory file descriptor (FD) to the containerized trainer/worker processes using UNIX domain sockets and SCM_RIGHTS.
+3. **CUDA Import**: Both trainer and inference processes import the shareable handle and map the memory space into their own GPU memory, wrapping the pointer as a flat PyTorch tensor.
+4. **Hardware Propagation**: The trainer copies weights into the flat buffer and triggers propagation. The WPI driver executes the NCCL broadcast or sharded scatter over InfiniBand/RDMA fabric.
+5. **Worker Sync**: The inference worker waits for a READY signal from the local notify socket, unpacks the parameters from the mapped VRAM buffer using offset metadata, and swaps them into the active model.
+
+## Installation
+
+The WPI engine requires the WPI client Python package to be installed in both the trainer and vLLM environments.
+
+```bash
+# From the WPI repository subdirectory
+pip install weight-propagation-interface/consumer/wpi_client/
+```
+
+## Initialization
+
+The WPI backend requires explicit setup of the persistent buffer and drivers on both the trainer and inference worker processes.
+
+### Inference Side
+
+Call `init_weight_transfer_engine` with `WPIWeightTransferInitInfo`:
+
+```python
+from vllm.config.weight_transfer import WeightTransferConfig
+from vllm.distributed.weight_transfer.base import WeightTransferInitRequest
+
+# Initialize WPI weight transfer engine
+llm.init_weight_transfer_engine(
+    WeightTransferInitRequest(
+        init_info=dict(
+            buffer_id="vllm-weights",
+            buffer_size_bytes=20 * 1024**3,  # 20 GiB (must hold all weights)
+            socket_dir="/run/wpi/sockets",
+            driver_port=50051,
+            # Note: sharded WPI is temporarily disabled; keep total_shards=0
+            # total_shards=0
+        )
+    )
+)
+```
+
+### Trainer Side
+
+Call `WPIWeightTransferEngine.trainer_init` to stage the source buffer and map the trainer GPU's memory:
+
+```python
+from vllm.distributed.weight_transfer.wpi_engine import (
+    WPIWeightTransferEngine,
+)
+
+ctx = WPIWeightTransferEngine.trainer_init(
+    init_info=dict(
+        buffer_id="vllm-weights",
+        buffer_size_bytes=20 * 1024**3,
+        socket_dir="/run/wpi/sockets",
+        driver_port=50051,
+    ),
+    target_node_ids=["10.0.0.2", "10.0.0.3"],  # Inference node IPs
+)
+```
+
+The returned context contains references to the WPI client and the flat mapped VRAM buffer.
+
+## Sending Weights
+
+To perform a weight update, the trainer packs model parameters into the mapped WPI buffer, triggers the WPI driver's NCCL propagation, and sends the layout metadata to vLLM workers (via HTTP or Ray control plane):
+
+```python
+from vllm.distributed.weight_transfer.wpi_engine import (
+    WPIWeightTransferEngine,
+    WPITrainerSendWeightsArgs,
+)
+
+# Start weight update on inference side
+llm.start_weight_update()
+
+# Prepare send arguments with the trainer context
+trainer_args = WPITrainerSendWeightsArgs(
+    send_mode="http",  # 'http' or 'ray'
+    url="http://localhost:8000",
+    trainer_ctx=ctx,
+)
+
+# Pack weights, propagate via WPI NCCL, and deliver layout metadata
+WPIWeightTransferEngine.trainer_send_weights(
+    iterator=model.named_parameters(),
+    trainer_args=trainer_args,
+)
+
+# Finalize weight update on inference side
+llm.finish_weight_update()
+```
+
+See [`WPITrainerSendWeightsArgs`](https://github.com/vllm-project/vllm/blob/main/vllm/distributed/weight_transfer/wpi_engine.py) for the full list of configurable fields.
+
+## Comparison with Other Backends
+
+| Feature | IPC Backend | NCCL Backend | WPI Backend |
+| :--- | :--- | :--- | :--- |
+| **Colocation** | Must be on same GPU/node | Can be separate nodes | Can be separate nodes |
+| **VRAM Allocation** | Zero-copy (shares trainer's memory) | Allocates temporary duplicate buffers | Zero-copy (persistent driver-managed buffer) |
+| **Network Fabric** | CUDA IPC (NVLink/PCIe) | NCCL broadcast (NVLink/InfiniBand) | NCCL broadcast/scatter (InfiniBand/RoCE) |
+| **Connection Mgmt** | Ray / HTTP control plane | vLLM `StatelessProcessGroup` | External WPI DaemonSet |
+| **Orchestration** | Python process orchestration | PyTorch distributed group | Kubernetes Custom Resources (DRA) |
+
+## Performance Benchmarks
+
+The WPI engine achieves high fabric utilization by avoiding CPU staging and runtime memory allocation overheads. Below are baseline performance metrics:
+
+- **Cross-Node (8-GPU Scatter, A4)**: **251 GB/s aggregate throughput** for a 600 GB tensor-parallel sharded scatter over InfiniBand (GPUDirect RDMA enabled).
+- **Cross-Node (A3 Ultra InfiniBand)**: **36.57 GB/s** for a 75 GB model broadcast, achieving over 73% of maximum physical fabric bandwidth (a 565% improvement over standard socket routing).

examples/rl/rlhf_http_wpi.py

@@ -0,0 +1,216 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""
+Demonstrates reinforcement learning from human feedback (RLHF) using vLLM
+via HTTP API, with WPI (Weight Propagation Interface) weight syncing APIs.
+
+Unlike rlhf.py which creates a vLLM instance programmatically, this script
+assumes you have already started a vLLM server using `vllm serve`. It uses:
+- OpenAI-compatible API for inference requests
+- HTTP endpoints for weight transfer control plane
+- WPI (Weight Propagation Interface) for actual weight data transfer
+
+Prerequisites:
+    1. Ensure the WPI driver is running on the node(s).
+    2. Start a vLLM server with WPI weight transfer enabled:
+
+    $ VLLM_SERVER_DEV_MODE=1 vllm serve facebook/opt-125m \
+        --enforce-eager \
+        --weight-transfer-config '{"backend": "wpi"}' \
+        --load-format dummy
+
+    Then run this script:
+
+    $ python rlhf_http_wpi.py
+"""
+
+import requests
+import torch
+from openai import OpenAI
+from transformers import AutoModelForCausalLM
+
+from vllm.distributed.weight_transfer.wpi_engine import (
+    WPITrainerSendWeightsArgs,
+    WPIWeightTransferEngine,
+)
+
+BASE_URL = "http://localhost:8000"
+MODEL_NAME = "facebook/opt-125m"
+
+
+def generate_completions(client: OpenAI, model: str, prompts: list[str]) -> list[str]:
+    """Generate completions using the OpenAI-compatible API."""
+    results = []
+    for prompt in prompts:
+        response = client.completions.create(
+            model=model,
+            prompt=prompt,
+            max_tokens=32,
+            temperature=0,
+        )
+        results.append(response.choices[0].text)
+    return results
+
+
+def init_weight_transfer_engine(
+    base_url: str,
+    buffer_id: str,
+    buffer_size_bytes: int,
+    socket_dir: str,
+    driver_port: int,
+) -> None:
+    """Initialize weight transfer via HTTP endpoint."""
+    url = f"{base_url}/init_weight_transfer_engine"
+    payload = {
+        "init_info": dict(
+            buffer_id=buffer_id,
+            buffer_size_bytes=buffer_size_bytes,
+            socket_dir=socket_dir,
+            driver_port=driver_port,
+            shard_index=-1,
+            total_shards=0,
+        )
+    }
+    response = requests.post(url, json=payload, timeout=60)
+    response.raise_for_status()
+
+
+def start_weight_update(base_url: str) -> None:
+    """Start a weight update via HTTP endpoint."""
+    url = f"{base_url}/start_weight_update"
+    response = requests.post(url, json={}, timeout=60)
+    response.raise_for_status()
+
+
+def finish_weight_update(base_url: str) -> None:
+    """Finish a weight update via HTTP endpoint."""
+    url = f"{base_url}/finish_weight_update"
+    response = requests.post(url, json={}, timeout=60)
+    response.raise_for_status()
+
+
+def pause_generation(base_url: str) -> None:
+    """Pause generation via HTTP endpoint."""
+    url = f"{base_url}/pause"
+    response = requests.post(url, timeout=60)
+    response.raise_for_status()
+
+
+def resume_generation(base_url: str) -> None:
+    """Resume generation via HTTP endpoint."""
+    url = f"{base_url}/resume"
+    response = requests.post(url, timeout=60)
+    response.raise_for_status()
+
+
+def get_world_size(base_url: str) -> int:
+    """Get world size from the vLLM server."""
+    url = f"{base_url}/get_world_size"
+    response = requests.get(url, timeout=10)
+    response.raise_for_status()
+    return response.json()["world_size"]
+
+
+def main():
+    # Get the inference world size from the vLLM server
+    inference_world_size = get_world_size(BASE_URL)
+    device = f"cuda:{inference_world_size}"
+    torch.accelerator.set_device_index(device)
+
+    # Load the training model
+    print(f"Loading training model: {MODEL_NAME}")
+    train_model = AutoModelForCausalLM.from_pretrained(MODEL_NAME, dtype=torch.bfloat16)
+    train_model.to(device)
+
+    # Create OpenAI client pointing to the vLLM server
+    client = OpenAI(
+        base_url=f"{BASE_URL}/v1",
+        api_key="EMPTY",  # vLLM doesn't require an API key by default
+    )
+
+    # Test prompts
+    prompts = [
+        "Hello, my name is",
+        "The president of the United States is",
+        "The capital of France is",
+        "The future of AI is",
+    ]
+
+    # Generate text before weight update. The output is expected to be nonsense
+    # because the server is initialized with dummy weights.
+    print("-" * 50)
+    print("Generating text BEFORE weight update (expect nonsense):")
+    print("-" * 50)
+    outputs = generate_completions(client, MODEL_NAME, prompts)
+    for prompt, generated_text in zip(prompts, outputs):
+        print(f"Prompt: {prompt!r}\nGenerated text: {generated_text!r}")
+        print("-" * 50)
+
+    # WPI specific setup
+    buffer_id = "vllm-weights"
+    socket_dir = "/run/wpi/sockets"
+    driver_port = 50051
+
+    # Calculate model size in bytes with alignment padding
+    total_model_bytes = 0
+    for p in train_model.parameters():
+        itemsize = p.element_size()
+        total_model_bytes = (total_model_bytes + itemsize - 1) // itemsize * itemsize
+        total_model_bytes += p.numel() * itemsize
+    print(f"Calculated model size with alignment: {total_model_bytes} bytes")
+
+    print("Initializing weight transfer on server...")
+
+    # Initialize weight transfer on vLLM server
+    init_weight_transfer_engine(
+        BASE_URL, buffer_id, total_model_bytes, socket_dir, driver_port
+    )
+
+    # Initialize WPI client on trainer side
+    print("Initializing trainer WPI context...")
+    ctx = WPIWeightTransferEngine.trainer_init(
+        dict(
+            buffer_id=buffer_id,
+            buffer_size_bytes=total_model_bytes,
+            socket_dir=socket_dir,
+            driver_port=driver_port,
+        ),
+        target_node_ids=["127.0.0.1"],  # Assuming local testing
+    )
+
+    # Pause generation before weight sync
+    pause_generation(BASE_URL)
+
+    # Start weight update, broadcast via WPI, then finish
+    start_weight_update(BASE_URL)
+
+    # Broadcast all weights from trainer to vLLM workers
+    print("Propagating weights via WPI...")
+    param_iter = ((n, p) for n, p in train_model.named_parameters())
+    args = WPITrainerSendWeightsArgs(
+        send_mode="http",
+        url=BASE_URL,
+        trainer_ctx=ctx,
+    )
+
+    # This will pack weights, trigger WPI propagate, and call /update_weights on server
+    WPIWeightTransferEngine.trainer_send_weights(param_iter, args)
+
+    finish_weight_update(BASE_URL)
+
+    # Resume generation after weight sync
+    resume_generation(BASE_URL)
+
+    # Generate text after weight update. The output is expected to be normal
+    # because the real weights are now loaded.
+    print("-" * 50)
+    print("Generating text AFTER weight update:")
+    print("-" * 50)
+    outputs_updated = generate_completions(client, MODEL_NAME, prompts)
+    for prompt, generated_text in zip(prompts, outputs_updated):
+        print(f"Prompt: {prompt!r}\nGenerated text: {generated_text!r}")
+        print("-" * 50)
+
+
+if __name__ == "__main__":
+    main()