[None][feat] add GPUDirect RDMA draft offload for speculative decoding

zhaoyangwang-nvidia · zhaoyangwang-nvidia · commit e051e60d26b1 · 2026-04-28T00:42:26.000-07:00
Implement CPU-initiated libibverbs RDMA transport to offload the draft
  model to a separate GPU, replacing local draft inference with a remote
  RDMA peer. The target model (TRT-LLM) writes accumulated output tokens
  to the draft server via RDMA Write; the draft server returns speculative
  tokens via RDMA Write back.

  Key changes:
  - rdma_draft_offload.py: ibverbs RC QP client with GPUDirect MR registration,
    QP state machine (RESET-&gt;INIT-&gt;RTR-&gt;RTS), and per-round request/response
  - rdma_draft_protocol.py: fixed-size binary protocol (32B header + 256B tokens,
    MAGIC-checked, 4096B total) for target&lt;-&gt;draft RDMA buffers
  - draft_target.py: RDMA offload path in DraftTargetOneModelWorker.forward(),
    output token history accumulation, warmup pre-connection
  - llm_args.py: DraftTargetDecodingConfig RDMA fields; allow speculative_model=None
    when draft_offload_enabled=True
  - model_loader.py: skip draft weight loading when draft_offload_enabled
  - modeling_speculative.py: skip draft model instantiation; thread is_warmup
  - _util.py: skip separate draft KV cache when draft_offload_enabled
  - model_engine.py: pass is_warmup flag through to model forward inputs
  - .gitignore: ignore cmake-created symlinks deep_ep/deep_gemm/flash_mla

Signed-off-by: ZhaoyangWang &lt;zhaoyangw@nvidia.com&gt;
diff --git a/.gitignore b/.gitignore
@@ -110,3 +110,8 @@ enroot/tensorrt_llm.devel.sqsh
 .claude/agent-memory/
 .claude/agent-tests/perf-test-sync/report.html
 .claude/agent-tests/perf-test-sync/results.json
+
+# Runtime third-party dependencies: symlinks created by cmake build, not part of this repo
+tensorrt_llm/deep_ep
+tensorrt_llm/deep_gemm
+tensorrt_llm/flash_mla
diff --git a/examples/llm-api/llm_rdma_draft_offload.py b/examples/llm-api/llm_rdma_draft_offload.py
@@ -0,0 +1,96 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES.
+# SPDX-License-Identifier: Apache-2.0
+"""Minimal target-side RDMA draft offload example.
+
+Start the fake draft server first, then run this script. The target model is a
+real TensorRT-LLM LLM; the draft model is temporarily replaced by the fake RDMA
+peer.
+"""
+
+from __future__ import annotations
+
+import argparse
+import os
+import sys
+from pathlib import Path
+
+_REPO_ROOT = Path(__file__).resolve().parents[2]
+if str(_REPO_ROOT) not in sys.path:
+    sys.path.insert(0, str(_REPO_ROOT))
+
+DEFAULT_MODEL = "/scratch.trt_llm_data/llm-models/Qwen3/Qwen3-8B"
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--model", default=DEFAULT_MODEL)
+    parser.add_argument("--prompt", default="Explain GPUDirect RDMA in one short sentence.")
+    parser.add_argument("--max-tokens", type=int, default=16)
+    parser.add_argument("--max-draft-len", type=int, default=5)
+    parser.add_argument("--draft-host", default="127.0.0.1")
+    parser.add_argument("--draft-port", type=int, default=47320)
+    parser.add_argument("--ib-dev", default="mlx5_0")
+    parser.add_argument(
+        "--gpu-id",
+        type=int,
+        default=0,
+        help="Physical GPU ID for RDMA memory registration (target side)",
+    )
+    parser.add_argument(
+        "--cuda-visible-devices",
+        default="0",
+        help="Set CUDA_VISIBLE_DEVICES before importing TensorRT-LLM.",
+    )
+    parser.add_argument("--max-batch-size", type=int, default=1)
+    parser.add_argument("--max-seq-len", type=int, default=512)
+    parser.add_argument("--max-num-tokens", type=int, default=512)
+    parser.add_argument("--kv-cache-max-tokens", type=int, default=512)
+    parser.add_argument("--kv-cache-free-gpu-memory-fraction", type=float, default=0.05)
+    return parser.parse_args()
+
+
+def main() -> None:
+    args = parse_args()
+    if args.cuda_visible_devices is not None:
+        os.environ["CUDA_VISIBLE_DEVICES"] = args.cuda_visible_devices
+
+    from tensorrt_llm import LLM, SamplingParams
+    from tensorrt_llm.llmapi import DraftTargetDecodingConfig, KvCacheConfig
+
+    spec_config = DraftTargetDecodingConfig(
+        max_draft_len=args.max_draft_len,
+        draft_offload_enabled=True,
+        draft_offload_nic_name=args.ib_dev,
+        draft_offload_server_host=args.draft_host,
+        draft_offload_server_port=args.draft_port,
+        draft_offload_gpu_id=args.gpu_id,
+    )
+    # Disable CUDA graphs: RDMA draft calls are Python-side operations and
+    # would not be re-executed on each CUDA graph replay.
+    # cuda_graph_config=None disables CUDA graphs entirely (empty CudaGraphConfig
+    # does NOT disable graphs because the validator auto-fills batch sizes).
+
+    llm = LLM(
+        model=args.model,
+        speculative_config=spec_config,
+        disable_overlap_scheduler=True,
+        tensor_parallel_size=1,
+        cuda_graph_config=None,
+        max_batch_size=args.max_batch_size,
+        max_seq_len=args.max_seq_len,
+        max_num_tokens=args.max_num_tokens,
+        kv_cache_config=KvCacheConfig(
+            max_tokens=args.kv_cache_max_tokens,
+            free_gpu_memory_fraction=args.kv_cache_free_gpu_memory_fraction,
+        ),
+    )
+    output = llm.generate(
+        args.prompt,
+        SamplingParams(max_tokens=args.max_tokens),
+        use_tqdm=False,
+    )
+    print(output.outputs[0].text)
+
+
+if __name__ == "__main__":
+    main()
diff --git a/tensorrt_llm/_torch/models/modeling_speculative.py b/tensorrt_llm/_torch/models/modeling_speculative.py
@@ -1027,8 +1027,11 @@ def __init__(self, model: TModel, model_config: ModelConfig[TConfig]):
         spec_config = getattr(model_config, 'spec_config', None)
         self.spec_config = spec_config
         if spec_config and spec_config.spec_dec_mode.use_one_engine():
+            draft_offload_enabled = bool(
+                getattr(spec_config, "draft_offload_enabled", False))
             # Only create draft_model for modes MTP, Eagle3 (not SA)
-            if not spec_config.spec_dec_mode.is_sa():
+            if not spec_config.spec_dec_mode.is_sa(
+            ) and not draft_offload_enabled:
                 if spec_config.spec_dec_mode.is_eagle3_one_model():
                     if spec_config.eagle3_model_arch == "mistral_large3":
                         from tensorrt_llm._torch.models.checkpoints.mistral.config_loader import \
@@ -1105,6 +1108,7 @@ def forward(
         return_context_logits: bool = False,
         spec_metadata: Optional[SpecMetadata] = None,
         resource_manager=None,
+        is_warmup: bool = False,
         **kwargs,
     ) -> torch.Tensor:
         hidden_states = self.model(
@@ -1150,7 +1154,8 @@ def forward(
                                     attn_metadata=attn_metadata,
                                     spec_metadata=spec_metadata,
                                     draft_model=self.draft_model,
-                                    resource_manager=resource_manager)
+                                    resource_manager=resource_manager,
+                                    is_warmup=is_warmup)
         else:
             logits = self.logits_processor.forward(
                 hidden_states,
diff --git a/tensorrt_llm/_torch/pyexecutor/_util.py b/tensorrt_llm/_torch/pyexecutor/_util.py
@@ -624,6 +624,8 @@ def _should_create_separate_draft_kv_cache(self) -> bool:
                 "Attention DP is enabled, separate draft KV cache is not supported."
             )
             return False
+        if getattr(self._speculative_config, "draft_offload_enabled", False):
+            return False
         return should_use_separate_draft_kv_cache(self._speculative_config)
 
     def _get_effective_draft_config(self) -> ModelConfig:
diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py
@@ -2200,6 +2200,7 @@ def _apply_incremental_update_target(
             'inputs_embeds': None,
             "multimodal_params": [],
             'resource_manager': resource_manager,
+            'is_warmup': self.is_warmup,
         }
 
         if bool(lora_params):
@@ -3052,6 +3053,7 @@ def previous_seq_slots_device():
             'inputs_embeds': None,
             "multimodal_params": multimodal_params_list,
             'resource_manager': resource_manager,
+            'is_warmup': self.is_warmup,
         }
 
         if bool(lora_params):
@@ -3224,6 +3226,7 @@ def _prepare_tp_inputs_no_cache(
             'inputs_embeds': None,
             "multimodal_params": multimodal_params_list,
             'resource_manager': resource_manager,
+            'is_warmup': self.is_warmup,
         }
 
         if bool(lora_params):
@@ -3492,6 +3495,7 @@ def _prepare_star_attention_inputs(
             'position_ids': self.position_ids_cuda[:num_tokens].unsqueeze(0),
             'inputs_embeds': None,
             'resource_manager': resource_manager,
+            'is_warmup': self.is_warmup,
         }, gather_ids if is_spec_decode else None
 
     def _get_lora_params_from_requests(
diff --git a/tensorrt_llm/_torch/pyexecutor/model_loader.py b/tensorrt_llm/_torch/pyexecutor/model_loader.py
@@ -394,8 +394,9 @@ def init_meta_tensor(t: torch.Tensor):
                 self._call_load_weights(model.load_weights, weights,
                                         self.weight_mapper)
 
-                if self.spec_config is not None and self.spec_config.spec_dec_mode.need_load_draft_weights(
-                ):
+                if (self.spec_config is not None and self.spec_config.
+                        spec_dec_mode.need_load_draft_weights() and not getattr(
+                            self.spec_config, "draft_offload_enabled", False)):
                     weights = checkpoint_loader.load_weights(
                         self.spec_config.speculative_model,
                         mapping=self.mapping)
@@ -414,8 +415,9 @@ def init_meta_tensor(t: torch.Tensor):
                 self.weight_mapper = checkpoint_loader.get_initialized_weight_mapper(
                     model, config)
                 initialize_dummy_weights(model)
-                if self.spec_config is not None and self.spec_config.spec_dec_mode.need_load_draft_weights(
-                ):
+                if (self.spec_config is not None and self.spec_config.
+                        spec_dec_mode.need_load_draft_weights() and not getattr(
+                            self.spec_config, "draft_offload_enabled", False)):
                     model.draft_model.load_weights_from_target_model(model)
 
             elif load_format == LoadFormat.VISION_ONLY:
diff --git a/tensorrt_llm/_torch/speculative/draft_target.py b/tensorrt_llm/_torch/speculative/draft_target.py
@@ -20,13 +20,15 @@
 layers are integrated into the target model's KV cache and run in a single forward pass.
 """
 
+import os
 from dataclasses import dataclass
 from typing import TYPE_CHECKING, Optional
 
 import torch
 from torch import nn
 
 from tensorrt_llm._utils import prefer_pinned
+from tensorrt_llm.logger import logger
 from tensorrt_llm.mapping import Mapping
 
 from ..attention_backend import AttentionMetadata
@@ -38,6 +40,13 @@
     from ...llmapi.llm_args import DraftTargetDecodingConfig
 
 
+def _env_enabled(name: str, default: bool = False) -> bool:
+    value = os.environ.get(name)
+    if value is None:
+        return default
+    return str(value).strip().lower() in {"1", "true", "yes", "on"}
+
+
 @dataclass
 class DraftTargetOneModelSpecMetadata(SpecMetadata):
     """
@@ -96,6 +105,54 @@ def __init__(
         super().__init__(use_separate_draft_kv_cache)
         self.spec_config = spec_config
         self.mapping = mapping
+        self._rdma_offload_enabled = bool(
+            getattr(spec_config, "draft_offload_enabled", False)
+            or _env_enabled("TLLM_DRAFT_RDMA_OFFLOAD")
+        )
+        self._rdma_draft_client = None
+        # Accumulates ALL output tokens across decode rounds so the draft server
+        # can reconstruct the full generation context (prompt tokens are prepended
+        # by the server side using the known prompt text).
+        self._rdma_output_history: list[int] = []
+        if self._rdma_offload_enabled:
+            if getattr(mapping, "tp_size", 1) != 1 or getattr(mapping, "pp_size", 1) != 1:
+                raise RuntimeError(
+                    "RDMA draft offload target path currently supports only "
+                    "single-rank TP/PP. Disable draft_offload_enabled for "
+                    "multi-rank runs."
+                )
+            from .rdma_draft_offload import RdmaDraftOffloadClient, RdmaDraftOffloadConfig
+
+            self._rdma_draft_client = RdmaDraftOffloadClient(
+                RdmaDraftOffloadConfig(
+                    nic_name=getattr(
+                        spec_config,
+                        "draft_offload_nic_name",
+                        os.environ.get("TLLM_DRAFT_RDMA_NIC", "mlx5_0"),
+                    ),
+                    server_host=getattr(
+                        spec_config,
+                        "draft_offload_server_host",
+                        os.environ.get("TLLM_DRAFT_RDMA_HOST", "127.0.0.1"),
+                    ),
+                    server_port=int(
+                        getattr(
+                            spec_config,
+                            "draft_offload_server_port",
+                            os.environ.get("TLLM_DRAFT_RDMA_PORT", "47320"),
+                        )
+                    ),
+                    gpu_id=getattr(spec_config, "draft_offload_gpu_id", None),
+                    max_draft_len=int(spec_config.max_draft_len),
+                    buffer_size=int(getattr(spec_config, "draft_offload_buffer_size", 4096)),
+                )
+            )
+            logger.info(
+                "DraftTarget RDMA draft offload enabled: host=%s port=%s nic=%s",
+                self._rdma_draft_client.config.server_host,
+                self._rdma_draft_client.config.server_port,
+                self._rdma_draft_client.config.nic_name,
+            )
 
     @property
     def max_draft_len(self) -> int:
@@ -162,6 +219,7 @@ def forward(
         spec_metadata: DraftTargetOneModelSpecMetadata,
         draft_model: nn.Module,
         resource_manager=None,
+        is_warmup: bool = False,
     ):
         """
         Technically incorrect at the moment.
@@ -184,6 +242,46 @@ def forward(
             logits, attn_metadata, spec_metadata
         )
 
+        if self._rdma_offload_enabled:
+            if bool(is_warmup):
+                # Warmup: initialize RDMA connection and exercise one real
+                # round-trip (tokens=[] → draft server sees only the prompt).
+                # This pre-warms the QP, GPU buffers, and NIC queues so the
+                # first real decode step does not pay connection-setup latency.
+                # Output history is NOT updated; the result is discarded.
+                self._rdma_draft_client.request(
+                    tokens=[],
+                    position=0,
+                    max_draft_len=self.max_draft_len,
+                    device=logits.device,
+                )
+                next_draft_tokens = torch.zeros(
+                    (batch_size, self.max_draft_len), dtype=torch.int32, device=logits.device
+                )
+            else:
+                next_draft_tokens = self._rdma_offload_draft_tokens(
+                    accepted_tokens=accepted_tokens,
+                    num_accepted_tokens=num_accepted_tokens,
+                    position_ids=position_ids,
+                    logits=logits,
+                    batch_size=batch_size,
+                )
+            next_new_tokens = self._prepare_next_new_tokens(
+                accepted_tokens,
+                next_draft_tokens,
+                spec_metadata.batch_indices_cuda,
+                batch_size,
+                num_accepted_tokens,
+            )
+            attn_metadata.use_spec_decoding = True
+            return {
+                "logits": raw_logits,
+                "new_tokens": accepted_tokens,
+                "new_tokens_lens": num_accepted_tokens,
+                "next_draft_tokens": next_draft_tokens,
+                "next_new_tokens": next_new_tokens,
+            }
+
         # Prepare attention metadata for speculative decoding and save state for restore
         self._prepare_attn_metadata_for_draft_target(attn_metadata, spec_metadata)
 
@@ -297,6 +395,56 @@ def forward(
             "next_new_tokens": next_new_tokens,
         }
 
+    def _rdma_offload_draft_tokens(
+        self,
+        *,
+        accepted_tokens: torch.Tensor,
+        num_accepted_tokens: torch.Tensor,
+        position_ids: Optional[torch.Tensor],
+        logits: torch.Tensor,
+        batch_size: int,
+    ) -> torch.Tensor:
+        if self._rdma_draft_client is None:
+            raise RuntimeError("RDMA draft offload client was not initialized")
+        if int(batch_size) != 1:
+            raise RuntimeError("RDMA draft offload target path currently supports batch_size=1")
+
+        accepted_count = int(num_accepted_tokens[0].detach().cpu().item())
+        accepted_count = max(1, min(accepted_count, accepted_tokens.shape[1]))
+
+        if position_ids is None or int(position_ids.numel()) == 0:
+            position = 0
+        else:
+            position = int(position_ids.reshape(-1)[-1].detach().cpu().item())
+
+        # Accumulate all accepted output tokens for full-context draft inference.
+        for i in range(accepted_count):
+            self._rdma_output_history.append(int(accepted_tokens[0, i].detach().cpu().item()))
+        # Cap at MAX_TOKENS (64) to fit in the RDMA buffer.
+        tokens_to_send = self._rdma_output_history[-64:]
+
+        logger.info(
+            "[RDMA] _rdma_offload_draft_tokens: round=%d pos=%d ctx_len=%d",
+            self._rdma_draft_client.round_seq,
+            position,
+            len(tokens_to_send),
+        )
+        draft_tokens = self._rdma_draft_client.request(
+            tokens=tokens_to_send,
+            position=position,
+            max_draft_len=self.max_draft_len,
+            device=logits.device,
+        )
+        logger.info("[RDMA] got draft tokens: %s", draft_tokens)
+        if not draft_tokens:
+            draft_tokens = [0]
+        if len(draft_tokens) < self.max_draft_len:
+            draft_tokens = draft_tokens + [draft_tokens[-1]] * (
+                self.max_draft_len - len(draft_tokens)
+            )
+        draft_tokens = draft_tokens[: self.max_draft_len]
+        return torch.tensor([draft_tokens], dtype=torch.int32, device=logits.device)
+
     def sample_and_accept_draft_tokens(
         self,
         logits: torch.Tensor,
diff --git a/tensorrt_llm/_torch/speculative/rdma_draft_offload.py b/tensorrt_llm/_torch/speculative/rdma_draft_offload.py
diff --git a/tensorrt_llm/_torch/speculative/rdma_draft_protocol.py b/tensorrt_llm/_torch/speculative/rdma_draft_protocol.py
diff --git a/tensorrt_llm/llmapi/llm_args.py b/tensorrt_llm/llmapi/llm_args.py

Original file line number	Diff line number	Diff line change
`@@ -624,6 +624,8 @@ def _should_create_separate_draft_kv_cache(self) -> bool:`
`624`	`624`	`"Attention DP is enabled, separate draft KV cache is not supported."`
`625`	`625`	`)`
`626`	`626`	`return False`
	`627`	`+ if getattr(self._speculative_config, "draft_offload_enabled", False):`
	`628`	`+ return False`
`627`	`629`	`return should_use_separate_draft_kv_cache(self._speculative_config)`
`628`	`630`
`629`	`631`	`def _get_effective_draft_config(self) -> ModelConfig:`