sglang 0.5.0rc1__py3-none-any.whl → 0.5.1__py3-none-any.whl

sglang/srt/mem_cache/storage/nixl/hicache_nixl.py

@@ -3,7 +3,7 @@ import logging
 import os
 import time
 import uuid
-from typing import Dict, List, Optional, Tuple, Union
+from typing import Any, Dict, List, Optional, Tuple, Union
 
 import torch
 
@@ -28,6 +28,8 @@ class HiCacheNixl(HiCacheStorage):
 
     def __init__(self, file_path: str = "/tmp/hicache_storage", plugin: str = "auto"):
         """Initialize NIXL storage connector."""
+        # Might be better to be unified across HiCache backends and moved to HiCacheController
+        file_path = os.getenv("SGLANG_HICACHE_NIXL_BACKEND_STORAGE_DIR", file_path)
         self.file_manager = (
             NixlFileManager(file_path)
             if plugin not in NixlBackendSelection.OBJ_PLUGINS
@@ -44,59 +46,109 @@ class HiCacheNixl(HiCacheStorage):
 
         self.registration = NixlRegistration(self.agent)
 
+    def register_buffers(
+        self, buffers: Union[torch.Tensor, List[torch.Tensor], List[tuple]]
+    ) -> Optional[Any]:
+        """Register tensor(s) or target locations in host memory (list of addr,len tuples) with NIXL."""
+        if isinstance(buffers[0], tuple):
+            tuples = [(x[0], x[1], 0, "") for x in buffers]
+            return self.registration._register_memory(tuples, "DRAM")
+        else:
+            return self.registration._register_memory(buffers)
+
+    def register_files(
+        self, file_paths: List[str], open_file: Optional[bool] = True
+    ) -> Optional[Any]:
+        """Register files with NIXL."""
+        tuples = self.file_manager.files_to_nixl_tuples(file_paths)
+        return self.registration._register_memory(tuples, "FILE")
+
+    def register_objects(
+        self, keys: List[str], sizes: Optional[List[int]] = None
+    ) -> Optional[Any]:
+        """Register objects with NIXL."""
+        if not keys:
+            return None
+        tuples = [(0, 0, key, "") for key in keys]
+        return self.registration._register_memory(tuples, "OBJ")
+
     def _execute_transfer(
-        self, tensors: List[torch.Tensor], keys: List[str], direction: str
+        self,
+        buffers: Optional[List[torch.Tensor | tuple]],
+        keys: List[str],
+        direction: str,
     ) -> bool:
-        if len(tensors) != len(keys):
-            logger.error("Mismatch between number of tensors and files/objects")
+        if len(buffers) != len(keys):
+            logger.error("Mismatch between number of tensors/buffers and files/objects")
             return False
 
-        if not self.registration.register_buffers(tensors):
-            logger.error("Failed to register tensors")
-            return False
-
-        # Get transfer tuples based on backend type
-        tensor_sizes = [tensor.element_size() * tensor.numel() for tensor in tensors]
+        # Registering file and object keys per transfer, to be updated when
+        # pre-registration for file and object is added to HiCache.
         if self.backend_selector.mem_type == "FILE":
-            file_tuples = self.file_manager.files_to_nixl_tuples(keys)
-            if not file_tuples or not self.registration.register_files(file_tuples):
+            tuples = self.file_manager.files_to_nixl_tuples(keys)
+            if not tuples or not self.registration._register_memory(tuples, "FILE"):
                 logger.error("Failed to prepare files for transfer")
                 return False
-            transfer_tuples = [
-                (x[0], s, x[2]) for x, s in zip(file_tuples, tensor_sizes)
-            ]
-        else:
-            if not self.registration.register_objects(keys, tensors):
+        else:  # mem_type == "OBJ"
+            tuples = [(0, 0, key, "") for key in keys]
+            if not tuples or not self.registration._register_memory(tuples, "OBJ"):
                 logger.error("Failed to register objects")
                 return False
-            transfer_tuples = [(0, s, key) for s, key in zip(tensor_sizes, keys)]
 
+        # Prepare transfer descriptors
+        if isinstance(buffers[0], torch.Tensor):
+            tensor_sizes = [
+                tensor.element_size() * tensor.numel() for tensor in buffers
+            ]
+            storage_tuples = [(x[0], s, x[2]) for x, s in zip(tuples, tensor_sizes)]
+            host_descs = self.agent.get_xfer_descs(buffers)
+        elif isinstance(buffers[0], tuple):
+            storage_tuples = [(x[0], y[1], x[2]) for x, y in zip(tuples, buffers)]
+            host_descs = self.agent.get_xfer_descs(
+                [(x[0], x[1], 0) for x in buffers], "DRAM"
+            )
+        else:
+            return False
+
+        storage_descs = self.agent.get_xfer_descs(
+            storage_tuples, self.backend_selector.mem_type
+        )
+
+        if (host_descs is None) or (storage_descs is None):
+            logger.error("Failed to get transfer descriptors")
+            return False
+
+        # Initialize transfer, default assumption that tensor was registered
         try:
-            # Get transfer descriptors
-            if (tensor_descs := self.agent.get_xfer_descs(tensors)) is None or (
-                file_descs := self.agent.get_xfer_descs(
-                    transfer_tuples, self.backend_selector.mem_type
-                )
-            ) is None:
-                logger.error("Failed to get transfer descriptors")
+            xfer_req = self.agent.initialize_xfer(
+                direction, host_descs, storage_descs, self.agent_name
+            )
+        except Exception:
+            # Check if it was due to missing pre-registration
+            if not self.register_buffers(buffers):
+                logger.error("Failed to register tensors/buffers")
                 return False
 
-            # Initialize and execute transfer
-            if (
-                xfer_req := self.agent.initialize_xfer(
-                    direction, tensor_descs, file_descs, self.agent_name
+            try:
+                xfer_req = self.agent.initialize_xfer(
+                    direction, host_descs, storage_descs, self.agent_name
                 )
-            ) is None:
-                logger.error("Failed to create transfer request")
+            except Exception as e:
+                logger.error(f"Failed to create transfer request: {e}")
                 return False
 
+        # Execute transfer and wait for its completion
+        try:
             state = self.agent.transfer(xfer_req)
             while state != "DONE":
                 state = self.agent.check_xfer_state(xfer_req)
                 if state == "ERR":
+                    self.agent.release_xfer_handle(xfer_req)
                     logger.error("Transfer failed")
                     return False
-            time.sleep(0.0001)  # Can be changed to os.sched_yield() or parametrized
+                time.sleep(0.0001)  # Can be changed to os.sched_yield() or parametrized
+
+            self.agent.release_xfer_handle(xfer_req)
             return True
 
         except Exception as e:
@@ -106,45 +158,87 @@ class HiCacheNixl(HiCacheStorage):
             logger.error(f"Traceback: {traceback.format_exc()}")
             return False
 
-    def batch_set(self, keys: List[str], values: List[torch.Tensor]) -> bool:
-        if not keys:
-            return True
-
-        if self.backend_selector.mem_type == "FILE":
-            file_paths = []
-            for key in keys:
-                tensor_path = self.file_manager.get_file_path(key)
-                if not self.file_manager.create_file(tensor_path):
-                    logger.error(f"Failed to create file {tensor_path}")
-                    return False
-                file_paths.append(tensor_path)
-            return self._execute_transfer(values, file_paths, "WRITE")
-        else:
-            return self._execute_transfer(values, keys, "WRITE")
-
-    def set(self, key: str, value: torch.Tensor) -> bool:
-        return self.batch_set([key], [value])
-
     def get(
-        self, key: str, dst_tensor: Optional[torch.Tensor] = None
+        self,
+        key: str,
+        target_location: Optional[torch.Tensor | int] = None,
+        target_sizes: Optional[int] = None,
     ) -> torch.Tensor | None:
-        if dst_tensor is None:  # To be removed, being compatible with the current API
+        # To be removed, being compatible with the current API
+        if target_location is None:
             return None
-        result = self.batch_get([key], [dst_tensor])
+        if target_sizes:
+            result = self.batch_get([key], [target_location], [target_sizes])
+        else:
+            result = self.batch_get([key], [target_location])
         return result[0] if result else None
 
     def batch_get(
-        self, keys: List[str], dst_tensors: List[torch.Tensor]
-    ) -> List[Optional[torch.Tensor]]:
+        self,
+        keys: List[str],
+        target_locations: Optional[List[torch.Tensor | int]] = None,
+        target_sizes: Optional[List[int]] = None,
+    ) -> List[torch.Tensor | None]:
         if not keys:
             return []
 
+        # To be removed, being compatible with the current API
+        if not target_locations:
+            return [None] * len(keys)
+
+        if target_sizes and (len(target_sizes) != len(target_locations)):
+            logger.error("Mismatch between number of target_locations and target_sizes")
+            return [None] * len(keys)
+        if target_sizes:
+            dest = list(zip(target_locations, target_sizes))
+        else:
+            dest = target_locations
+
         if self.backend_selector.mem_type == "FILE":
             file_paths = [self.file_manager.get_file_path(key) for key in keys]
-            success = self._execute_transfer(dst_tensors, file_paths, "READ")
+            success = self._execute_transfer(dest, file_paths, "READ")
         else:
-            success = self._execute_transfer(dst_tensors, keys, "READ")
-        return dst_tensors if success else [None] * len(keys)
+            success = self._execute_transfer(dest, keys, "READ")
+        return target_locations if success and not target_sizes else [None] * len(keys)
+
+    def set(
+        self,
+        key: str,
+        value: Optional[torch.Tensor] = None,
+        target_location: Optional[int] = None,
+        target_sizes: Optional[int] = None,
+    ) -> bool:
+        if target_location and target_sizes:
+            return self.batch_set([key], None, [target_location], [target_sizes])
+        else:
+            return self.batch_set([key], [value])
+
+    def batch_set(
+        self,
+        keys: List[str],
+        values: Optional[List[torch.Tensor]] = None,
+        target_locations: Optional[List[int]] = None,
+        target_sizes: Optional[List[int]] = None,
+    ) -> bool:
+        if not keys or (not values and (not target_locations or not target_sizes)):
+            logger.error("Keys or values were not passed")
+            return False
+
+        if not values:
+            values = list(zip(target_locations, target_sizes))
+
+        if self.backend_selector.mem_type == "FILE":
+            file_paths = []
+            for key in keys:
+                file_path = self.file_manager.get_file_path(key)
+                # New file per set, to be updated when partial writes is added to HiCache
+                if not self.file_manager.create_file(file_path):
+                    logger.error(f"Failed to create file {file_path}")
+                    return False
+                file_paths.append(file_path)
+            return self._execute_transfer(values, file_paths, "WRITE")
+        else:  # mem_type == "OBJ"
+            return self._execute_transfer(values, keys, "WRITE")
 
     def exists(self, key: str) -> bool:
         tuples = self.registration.create_query_tuples(

sglang/srt/mem_cache/storage/nixl/nixl_utils.py

@@ -109,66 +109,35 @@ class NixlRegistration:
             return [(0, 0, key)]
 
     def _register_memory(
-        self, items: Union[List[tuple], List[torch.Tensor]], mem_type: str, desc: str
+        self,
+        items: Union[List[tuple], torch.Tensor, List[torch.Tensor]],
+        mem_type: Optional[str] = None,
     ) -> Optional[Any]:
         """Common registration logic for files, objects, and buffers.
         Args:
             items: List of tuples or tensors to register
-            mem_type: Memory type ("FILE", "OBJ", "DRAM", "VRAM")
-            desc: Description for logging
+            mem_type: Memory type ("FILE", "OBJ") or None for tensor or list of tensors
         """
-        try:
-            if not items:
-                return None
-
-            reg_descs = self.agent.get_reg_descs(items, mem_type)
-            if reg_descs is None:
-                logger.error("Failed to create registration descriptors")
-                return None
-
-            registered_memory = self.agent.register_memory(reg_descs)
-            if registered_memory:
-                return registered_memory
-            else:
-                logger.error("Failed to register with NIXL")
-                return None
-
-        except Exception as e:
-            logger.error(f"Failed to register {desc}: {e}")
+        if isinstance(items, list) and not items:
             return None
 
-    def register_buffers(
-        self, buffers: Union[torch.Tensor, List[torch.Tensor]]
-    ) -> Optional[Any]:
-        """Register tensors/buffers with NIXL."""
-        if isinstance(buffers, torch.Tensor):
-            buffers = [buffers]
-
-        if not buffers:
+        reg_descs = self.agent.get_reg_descs(items, mem_type)
+        if reg_descs is None:
+            logger.error("Failed to create registration descriptors")
             return None
 
-        # Determine memory type based on tensor device
-        mem_type = "VRAM" if buffers[0].device.type == "cuda" else "DRAM"
-        return self._register_memory(buffers, mem_type, "buffers")
-
-    def register_files(self, tuples: List[tuple]) -> Optional[Any]:
-        """Register files with NIXL using (0, 0, fd, file_path) tuples."""
-        return self._register_memory(tuples, "FILE", "files")
-
-    def register_objects(
-        self, keys: List[str], tensors: Optional[List[torch.Tensor]] = None
-    ) -> Optional[Any]:
-        """Register objects with NIXL."""
-        if not keys:
+        try:
+            registered_memory = self.agent.register_memory(reg_descs)
+            return registered_memory  # Could be None in case of error
+        except Exception as e:
+            if not mem_type:
+                logger.error(f"Failed to register Tensors with NIXL: {e}")
+            else:
+                logger.error(
+                    f"Failed to register memory of type {mem_type} with NIXL: {e}"
+                )
             return None
 
-        # Create object tuples with proper sizes
-        tuples = [
-            (0, tensor.element_size() * tensor.numel() if tensor else 0, key)
-            for key, tensor in zip(keys, tensors or [None] * len(keys))
-        ]
-        return self._register_memory(tuples, "OBJ", "objects")
-
 
 class NixlFileManager:
     """Handles file system operations for NIXL."""
@@ -221,12 +190,9 @@ class NixlFileManager:
             return False
 
     def files_to_nixl_tuples(
-        self, file_paths: List[str], open_file: bool = True
+        self, file_paths: List[str]
     ) -> List[Tuple[int, int, int, str]]:
         """Create NIXL tuples (offset, length, fd, file_path) for given files."""
-        if not open_file:
-            return [(0, 0, 0, path) for path in file_paths]
-
         tuples = []
         for path in file_paths:
             if (fd := self.open_file(path)) is None:

sglang/srt/mem_cache/storage/nixl/test_hicache_nixl_storage.py

@@ -7,8 +7,11 @@ from unittest.mock import MagicMock
 
 import torch
 
-from sglang.srt.mem_cache.nixl.hicache_nixl import HiCacheNixl
-from sglang.srt.mem_cache.nixl.nixl_utils import NixlFileManager, NixlRegistration
+from sglang.srt.mem_cache.storage.nixl.hicache_nixl import HiCacheNixl
+from sglang.srt.mem_cache.storage.nixl.nixl_utils import (
+    NixlFileManager,
+    NixlRegistration,
+)
 
 
 class TestNixlUnified(unittest.TestCase):
@@ -88,8 +91,27 @@ class TestNixlUnified(unittest.TestCase):
 
         # Test get
         retrieved = self.hicache.get(key, dst_tensor)
+        self.verify_tensors_equal(value, dst_tensor)
         self.verify_tensors_equal(value, retrieved)
 
+        # Same test in addr,len mode with another key and dst_tensor
+        key2 = "test_key2"
+        dst_tensor2 = torch.zeros_like(value, device="cpu")
+        src_addr, src_len = value.data_ptr(), value.numel() * value.element_size()
+        dst_addr, dst_len = (
+            dst_tensor2.data_ptr(),
+            dst_tensor2.numel() * dst_tensor2.element_size(),
+        )
+
+        # Test set
+        self.assertTrue(self.hicache.set(key, None, src_addr, src_len))
+        self.assertTrue(self.hicache.exists(key))
+
+        # Test get
+        retrieved2 = self.hicache.get(key, dst_addr, dst_len)
+        self.assertTrue(retrieved2 == None)
+        self.verify_tensors_equal(value, dst_tensor2)
+
     def test_batch_set_get(self):
         """Test batch tensor set/get operations."""
         keys = ["key1", "key2", "key3"]
@@ -108,6 +130,23 @@ class TestNixlUnified(unittest.TestCase):
         retrieved = self.hicache.batch_get(keys, dst_tensors)
         self.verify_tensor_lists_equal(values, retrieved)
 
+        # Same test in addr,len mode with another key and dst_tensor
+        keys2 = ["key4", "key5", "key6"]
+        dst_tensors2 = [torch.zeros_like(v, device="cpu") for v in values]
+        src_addrs = [v.data_ptr() for v in values]
+        src_lens = [v.numel() * v.element_size() for v in values]
+        dst_addrs = [dt.data_ptr() for dt in dst_tensors2]
+        dst_lens = [dt.numel() * dt.element_size() for dt in dst_tensors2]
+
+        # Test batch set
+        self.assertTrue(self.hicache.batch_set(keys2, None, src_addrs, src_lens))
+        self.assertTrue(all(self.hicache.exists(key) for key in keys2))
+
+        # Test batch get
+        retrieved2 = self.hicache.batch_get(keys, dst_addrs, dst_lens)
+        self.assertTrue(all(ret == None for ret in retrieved2))
+        self.verify_tensor_lists_equal(values, dst_tensors2)
+
     def test_mixed_operations(self):
         """Test mixing single and batch operations."""
         # Test interleaved set/get operations
@@ -170,7 +209,7 @@ class TestNixlUnified(unittest.TestCase):
         self.file_manager.create_file(test_file)
 
         # Test tuple creation
-        tuples = self.file_manager.files_to_nixl_tuples([test_file], False)
+        tuples = self.file_manager.files_to_nixl_tuples([test_file])
         self.assertIsNotNone(tuples)
         self.assertTrue(len(tuples) > 0)
 
@@ -190,11 +229,11 @@ class TestNixlUnified(unittest.TestCase):
         tensor = torch.randn(10, 10)
 
         # Test buffer registration
-        self.assertIsNotNone(self.registration.register_buffers(tensor))
+        self.assertIsNotNone(self.hicache.register_buffers(tensor))
 
         # Test batch registration
         tensors = [torch.randn(5, 5) for _ in range(3)]
-        self.assertIsNotNone(self.registration.register_buffers(tensors))
+        self.assertIsNotNone(self.hicache.register_buffers(tensors))
 
     def test_register_files_with_tuples(self):
         """Test registration of files using NIXL tuples."""
@@ -203,8 +242,8 @@ class TestNixlUnified(unittest.TestCase):
             self.file_manager.create_file(file)
 
         # Create tuples and register
-        tuples = self.file_manager.files_to_nixl_tuples(files, False)
-        self.registration.register_files(tuples)
+        tuples = self.file_manager.files_to_nixl_tuples(files)
+        self.hicache.register_files(tuples)
 
         # Verify tuples
         self.assertEqual(len(tuples), len(files))

sglang/srt/model_executor/cuda_graph_runner.py

@@ -34,9 +34,10 @@ from sglang.srt.distributed.device_communicators.pynccl_allocator import (
 )
 from sglang.srt.distributed.parallel_state import GroupCoordinator, graph_capture
 from sglang.srt.layers.dp_attention import (
-    DPPaddingMode,
+    DpPaddingMode,
     get_attention_tp_rank,
     get_attention_tp_size,
+    set_dp_buffer_len,
 )
 from sglang.srt.layers.logits_processor import LogitsProcessorOutput
 from sglang.srt.layers.torchao_utils import save_gemlite_cache
@@ -239,6 +240,8 @@ class CudaGraphRunner:
     def __init__(self, model_runner: ModelRunner):
         # Parse args
         self.model_runner = model_runner
+        self.device = model_runner.device
+        self.device_module = torch.get_device_module(self.device)
         self.graphs = {}
         self.output_buffers = {}
         self.enable_torch_compile = model_runner.server_args.enable_torch_compile
@@ -304,13 +307,15 @@ class CudaGraphRunner:
             self.model_runner.lora_manager.init_cuda_graph_batch_info(self.max_bs)
 
         # Graph inputs
-        with torch.device("cuda"):
+        with torch.device(self.device):
             self.input_ids = torch.zeros((self.max_num_token,), dtype=torch.int64)
             self.req_pool_indices = torch.zeros((self.max_bs,), dtype=torch.int32)
             self.seq_lens = torch.full(
                 (self.max_bs,), self.seq_len_fill_value, dtype=torch.int32
             )
-            self.out_cache_loc = torch.zeros((self.max_num_token,), dtype=torch.int64)
+            self.out_cache_loc = torch.zeros(
+                (self.max_num_token,), dtype=self._cache_loc_dtype()
+            )
             self.positions = torch.zeros((self.max_num_token,), dtype=torch.int64)
             self.mrope_positions = torch.zeros((3, self.max_bs), dtype=torch.int64)
             self.num_token_non_padded = torch.zeros((1,), dtype=torch.int32)
@@ -349,30 +354,15 @@ class CudaGraphRunner:
                     self.global_num_tokens_for_logprob_gpu = torch.zeros(
                         (self.dp_size,), dtype=torch.int32
                     )
-                    self.gathered_buffer = torch.zeros(
-                        (
-                            self.max_num_token * self.dp_size,
-                            self.model_runner.model_config.hidden_size,
-                        ),
-                        dtype=self.model_runner.dtype,
-                    )
                 else:
                     assert self.require_attn_tp_gather
                     self.global_num_tokens_gpu = torch.zeros((1,), dtype=torch.int32)
                     self.global_num_tokens_for_logprob_gpu = torch.zeros(
                         (1,), dtype=torch.int32
                     )
-                    self.gathered_buffer = torch.zeros(
-                        (
-                            self.max_num_token,
-                            self.model_runner.model_config.hidden_size,
-                        ),
-                        dtype=self.model_runner.dtype,
-                    )
             else:
                 self.global_num_tokens_gpu = None
                 self.global_num_tokens_for_logprob_gpu = None
-                self.gathered_buffer = None
 
             self.custom_mask = torch.ones(
                 (
@@ -380,12 +370,12 @@ class CudaGraphRunner:
                     * self.num_tokens_per_bs
                 ),
                 dtype=torch.bool,
-                device="cuda",
+                device=self.device,
             )
             self.next_token_logits_buffer = torch.zeros(
                 (self.max_num_token, self.model_runner.model_config.vocab_size),
                 dtype=torch.float,
-                device="cuda",
+                device=self.device,
             )
 
         # Capture
@@ -397,6 +387,9 @@ class CudaGraphRunner:
                 f"Capture cuda graph failed: {e}\n{CUDA_GRAPH_CAPTURE_FAILED_MSG}"
             )
 
+    def _cache_loc_dtype(self):
+        return torch.int64
+
     def can_run(self, forward_batch: ForwardBatch):
         if self.require_mlp_tp_gather:
             cuda_graph_bs = (
@@ -516,8 +509,16 @@ class CudaGraphRunner:
             )
             logger.info(log_message)
 
+    def _capture_graph(self, graph, pool, stream, run_once_fn):
+        with self.device_module.graph(graph, pool=pool, stream=stream):
+            out = run_once_fn()
+        return out
+
+    def _create_device_graph(self):
+        return torch.cuda.CUDAGraph()
+
     def capture_one_batch_size(self, bs: int, forward: Callable):
-        graph = torch.cuda.CUDAGraph()
+        graph = self._create_device_graph()
         stream = self.stream
         num_tokens = bs * self.num_tokens_per_bs
 
@@ -556,7 +557,7 @@ class CudaGraphRunner:
                     device=input_ids.device,
                 )
             )
-            gathered_buffer = self.gathered_buffer[: num_tokens * self.dp_size]
+            global_dp_buffer_len = num_tokens * self.dp_size
         elif self.require_attn_tp_gather:
             self.global_num_tokens_gpu.copy_(
                 torch.tensor(
@@ -572,9 +573,9 @@ class CudaGraphRunner:
                     device=input_ids.device,
                 )
             )
-            gathered_buffer = self.gathered_buffer[:num_tokens]
+            global_dp_buffer_len = num_tokens
         else:
-            gathered_buffer = None
+            global_dp_buffer_len = None
 
         spec_info = self.get_spec_info(num_tokens)
         if self.capture_hidden_mode != CaptureHiddenMode.FULL:
@@ -607,8 +608,8 @@ class CudaGraphRunner:
             positions=positions,
             global_num_tokens_gpu=self.global_num_tokens_gpu,
             global_num_tokens_for_logprob_gpu=self.global_num_tokens_for_logprob_gpu,
-            dp_padding_mode=DPPaddingMode.get_default_mode_in_cuda_graph(),
-            gathered_buffer=gathered_buffer,
+            dp_padding_mode=DpPaddingMode.get_default_mode_in_cuda_graph(),
+            global_dp_buffer_len=global_dp_buffer_len,
             mrope_positions=mrope_positions,
             spec_algorithm=self.model_runner.spec_algorithm,
             spec_info=spec_info,
@@ -637,6 +638,7 @@ class CudaGraphRunner:
         def run_once():
             # Clean intermediate result cache for DP attention
             forward_batch.dp_local_start_pos = forward_batch.dp_local_num_tokens = None
+            set_dp_buffer_len(global_dp_buffer_len, num_tokens)
 
             kwargs = {}
             if (
@@ -656,19 +658,17 @@ class CudaGraphRunner:
             return logits_output_or_pp_proxy_tensors
 
         for _ in range(2):
-            torch.cuda.synchronize()
+            self.device_module.synchronize()
             self.model_runner.tp_group.barrier()
-
             run_once()
 
         if get_global_graph_memory_pool() is None:
-            set_global_graph_memory_pool(torch.cuda.graph_pool_handle())
+            set_global_graph_memory_pool(self.device_module.graph_pool_handle())
         # Set graph pool id globally to be able to use symmetric memory
         set_graph_pool_id(get_global_graph_memory_pool())
-        with torch.cuda.graph(
-            graph, pool=get_global_graph_memory_pool(), stream=stream
-        ):
-            out = run_once()
+        out = self._capture_graph(
+            graph, get_global_graph_memory_pool(), stream, run_once
+        )
 
         return graph, out