From d9c505d0fbc66428e359c44da62f081446bd9df2 Mon Sep 17 00:00:00 2001 From: AMD-yanfeiwang Date: Tue, 18 Aug 2026 02:29:05 +0000 Subject: [PATCH 1/3] [AMD] Fix registered HiCache host pointer aliases --- .../kernels/aot/csrc/common_extension.cc | 1 + .../kernels/aot/csrc/common_extension_rocm.cc | 1 + .../kernels/aot/csrc/kvcacheio/transfer.cu | 44 ++++- .../kernels/aot/include/sgl_kernel_ops.h | 2 + .../aot/python/sgl_kernel/kvcacheio.py | 5 + .../kernels/jit/csrc/kvcacheio/hicache.cuh | 16 +- .../kernels/jit/csrc/kvcacheio/hisparse.cuh | 16 +- .../jit/csrc/kvcacheio/transfer_mamba.cuh | 6 +- .../kernels/jit/include/sgl_kernel/utils.cuh | 60 +++++++ .../sglang/srt/mem_cache/memory_pool_host.py | 19 +-- .../sglang/srt/mem_cache/pool_host/common.py | 26 +++ python/sglang/srt/mem_cache/pool_host/dsa.py | 9 +- python/sglang/srt/mem_cache/pool_host/mha.py | 25 +-- python/sglang/srt/mem_cache/pool_host/mla.py | 11 +- .../test_hicache_page_first_write_back.py | 153 ++++++++++++++++++ 15 files changed, 345 insertions(+), 49 deletions(-) diff --git a/python/sglang/kernels/aot/csrc/common_extension.cc b/python/sglang/kernels/aot/csrc/common_extension.cc index 4d032020255a..a55fde927a86 100644 --- a/python/sglang/kernels/aot/csrc/common_extension.cc +++ b/python/sglang/kernels/aot/csrc/common_extension.cc @@ -302,6 +302,7 @@ TORCH_LIBRARY_FRAGMENT(sgl_kernel, m) { "transfer_kv_all_layer_mla_lf_pf(Tensor src_layers, Tensor dst, Tensor src_indices, Tensor dst_indices, " "int item_size, int dst_layout_dim, int num_layers, int block_quota, int num_warps_per_block) -> ()"); m.impl("transfer_kv_all_layer_mla_lf_pf", torch::kCUDA, &transfer_kv_all_layer_mla_lf_pf); + m.def("get_device_accessible_ptr(Tensor tensor, int device_index) -> int", &get_device_accessible_ptr); m.def( "transfer_kv_direct(Tensor[] src_layers, Tensor[] dst_layers, Tensor src_indices, Tensor dst_indices, int " "page_size) -> ()"); diff --git a/python/sglang/kernels/aot/csrc/common_extension_rocm.cc b/python/sglang/kernels/aot/csrc/common_extension_rocm.cc index 2aadfc314440..22e15f8f045c 100644 --- a/python/sglang/kernels/aot/csrc/common_extension_rocm.cc +++ b/python/sglang/kernels/aot/csrc/common_extension_rocm.cc @@ -201,6 +201,7 @@ TORCH_LIBRARY_EXPAND(sgl_kernel, m) { "transfer_kv_all_layer_mla_lf_pf(Tensor src_layers, Tensor dst, Tensor src_indices, Tensor dst_indices, " "int item_size, int dst_layout_dim, int num_layers, int block_quota, int num_warps_per_block) -> ()"); m.impl("transfer_kv_all_layer_mla_lf_pf", torch::kCUDA, &transfer_kv_all_layer_mla_lf_pf); + m.def("get_device_accessible_ptr(Tensor tensor, int device_index) -> int", &get_device_accessible_ptr); m.def( "transfer_kv_direct(Tensor[] src_layers, Tensor[] dst_layers, Tensor src_indices, Tensor dst_indices, int " "page_size) -> ()"); diff --git a/python/sglang/kernels/aot/csrc/kvcacheio/transfer.cu b/python/sglang/kernels/aot/csrc/kvcacheio/transfer.cu index 6b776a202df3..434a409cce9e 100644 --- a/python/sglang/kernels/aot/csrc/kvcacheio/transfer.cu +++ b/python/sglang/kernels/aot/csrc/kvcacheio/transfer.cu @@ -17,6 +17,38 @@ #include "utils.h" // WARP_SIZE #endif +inline void* resolve_device_accessible_ptr(const at::Tensor& tensor) { + void* ptr = tensor.data_ptr(); +#if defined(USE_ROCM) + if (tensor.device().is_cpu()) { + void* device_ptr = nullptr; + C10_CUDA_CHECK(hipHostGetDevicePointer(&device_ptr, ptr, 0)); + return device_ptr; + } +#elif !defined(USE_MUSA) + if (tensor.device().is_cpu()) { + void* device_ptr = nullptr; + C10_CUDA_CHECK(cudaHostGetDevicePointer(&device_ptr, ptr, 0)); + return device_ptr; + } +#endif + return ptr; +} + +#if !defined(USE_MUSA) +int64_t get_device_accessible_ptr(const at::Tensor& tensor, int64_t device_index) { + TORCH_CHECK(device_index >= 0, "Target device index must be non-negative"); + const c10::Device target_device(c10::DeviceType::CUDA, static_cast(device_index)); + const at::cuda::OptionalCUDAGuard device_guard(target_device); + if (tensor.is_cuda()) { + TORCH_CHECK(tensor.device() == target_device, "GPU tensor must be on the target device"); + } else { + TORCH_CHECK(tensor.device().is_cpu(), "Only CPU and target-device tensors are supported"); + } + return reinterpret_cast(resolve_device_accessible_ptr(tensor)); +} +#endif + #if !defined(USE_ROCM) && !defined(USE_MUSA) __device__ __forceinline__ void transfer_item_warp(int32_t lane_id, const void* src_addr, void* dst_addr, int64_t item_size_bytes) { @@ -335,6 +367,10 @@ void transfer_kv_launcher( TORCH_CHECK(src_indices.numel() == dst_indices.numel(), "Source and destination indices must have the same length"); TORCH_CHECK(item_size % 8 == 0, "Item byte size must be divisible by 8"); +#if !defined(USE_MUSA) + const at::cuda::OptionalCUDAGuard device_guard(src_indices.device()); +#endif + auto div_up = [](int64_t x, int64_t y) { return (x + y - 1) / y; }; const int64_t num_items = src_indices.numel(); const int64_t items_per_warp = div_up(num_items, block_quota * num_warps_per_block); @@ -342,10 +378,10 @@ void transfer_kv_launcher( dim3 grid_dim(num_blocks, 1, 1); const int32_t threads_per_block = num_warps_per_block * WARP_SIZE; - const void* src_k_ptr = src_k.defined() ? src_k.data_ptr() : nullptr; - void* dst_k_ptr = dst_k.defined() ? dst_k.data_ptr() : nullptr; - const void* src_v_ptr = IsMLA || !src_v.defined() ? nullptr : src_v.data_ptr(); - void* dst_v_ptr = IsMLA || !dst_v.defined() ? nullptr : dst_v.data_ptr(); + const void* src_k_ptr = src_k.defined() ? resolve_device_accessible_ptr(src_k) : nullptr; + void* dst_k_ptr = dst_k.defined() ? resolve_device_accessible_ptr(dst_k) : nullptr; + const void* src_v_ptr = IsMLA || !src_v.defined() ? nullptr : resolve_device_accessible_ptr(src_v); + void* dst_v_ptr = IsMLA || !dst_v.defined() ? nullptr : resolve_device_accessible_ptr(dst_v); const uintptr_t* src_k_tbl_ptr = src_k_layers.defined() ? src_k_layers.data_ptr() : nullptr; const uintptr_t* dst_k_tbl_ptr = dst_k_layers.defined() ? dst_k_layers.data_ptr() : nullptr; const uintptr_t* src_v_tbl_ptr = IsMLA || !src_v_layers.defined() ? nullptr : src_v_layers.data_ptr(); diff --git a/python/sglang/kernels/aot/include/sgl_kernel_ops.h b/python/sglang/kernels/aot/include/sgl_kernel_ops.h index 846e67efdfae..9e3bb1dedee9 100644 --- a/python/sglang/kernels/aot/include/sgl_kernel_ops.h +++ b/python/sglang/kernels/aot/include/sgl_kernel_ops.h @@ -573,6 +573,8 @@ void transfer_kv_all_layer_mla_lf_pf( int64_t block_quota, int64_t num_warps_per_block); +int64_t get_device_accessible_ptr(const at::Tensor& tensor, int64_t device_index); + void transfer_kv_direct( const std::vector& src_layers, std::vector dst_layers, diff --git a/python/sglang/kernels/aot/python/sgl_kernel/kvcacheio.py b/python/sglang/kernels/aot/python/sgl_kernel/kvcacheio.py index 43e3cdec80cd..abf4a8b514fe 100644 --- a/python/sglang/kernels/aot/python/sgl_kernel/kvcacheio.py +++ b/python/sglang/kernels/aot/python/sgl_kernel/kvcacheio.py @@ -11,6 +11,11 @@ def is_hip() -> bool: _is_hip = is_hip() +def get_device_accessible_ptr(tensor: torch.Tensor, device_index: int) -> int: + """Return the address a kernel on ``device_index`` must use for ``tensor``.""" + return torch.ops.sgl_kernel.get_device_accessible_ptr.default(tensor, device_index) + + def _default_mla_block_quota() -> int: """CU (block) quota for the MLA page_first KV gather kernel. diff --git a/python/sglang/kernels/jit/csrc/kvcacheio/hicache.cuh b/python/sglang/kernels/jit/csrc/kvcacheio/hicache.cuh index bf107730e31a..79b950a897bc 100644 --- a/python/sglang/kernels/jit/csrc/kvcacheio/hicache.cuh +++ b/python/sglang/kernels/jit/csrc/kvcacheio/hicache.cuh @@ -318,17 +318,17 @@ struct HiCacheKernel { const auto element_bytes = D.unwrap() * dtype_size; RuntimeCheck(kElementSize == element_bytes, "HicacheKernel: cache dimension mismatch."); - const auto k_cache_dst_ptr = k_cache_dst.data_ptr(); - const auto v_cache_dst_ptr = v_cache_dst.data_ptr(); - const auto k_cache_src_ptr = k_cache_src.data_ptr(); - const auto v_cache_src_ptr = v_cache_src.data_ptr(); + const auto device = indices_device.unwrap(); + const auto k_cache_dst_ptr = device_accessible_ptr(k_cache_dst, device); + const auto v_cache_dst_ptr = device_accessible_ptr(v_cache_dst, device); + const auto k_cache_src_ptr = device_accessible_ptr(k_cache_src, device); + const auto v_cache_src_ptr = device_accessible_ptr(v_cache_src, device); const auto indices_dst_ptr = indices_dst.data_ptr(); const auto indices_src_ptr = indices_src.data_ptr(); const auto length = static_cast(L.unwrap()); const auto kv_cache_src_stride = static_cast(N.unwrap() * dtype_size); const auto kv_cache_dst_stride = static_cast(M.unwrap() * dtype_size); const auto use_int32 = indices_dtype.unwrap().bits == 32; - const auto device = indices_device.unwrap(); constexpr auto kWorkersPerBlock = kBlockSize / (device::kWarpThreads / kUnroll); const auto num_blocks = std::min(div_ceil(length, kWorkersPerBlock), kBlockQuota); @@ -440,15 +440,15 @@ struct HiCacheKernel { const auto element_bytes = D.unwrap() * dtype_size; RuntimeCheck(kElementSize == element_bytes, "HicacheKernel MLA: cache dimension mismatch."); - const auto cache_dst_ptr = cache_dst.data_ptr(); - const auto cache_src_ptr = cache_src.data_ptr(); + const auto device = indices_device.unwrap(); + const auto cache_dst_ptr = device_accessible_ptr(cache_dst, device); + const auto cache_src_ptr = device_accessible_ptr(cache_src, device); const auto indices_dst_ptr = indices_dst.data_ptr(); const auto indices_src_ptr = indices_src.data_ptr(); const auto length = static_cast(L.unwrap()); const auto cache_src_stride = static_cast(N.unwrap() * dtype_size); const auto cache_dst_stride = static_cast(M.unwrap() * dtype_size); const auto use_int32 = indices_dtype.unwrap().bits == 32; - const auto device = indices_device.unwrap(); constexpr auto kWorkersPerBlock = kBlockSize / (device::kWarpThreads / kUnroll); const auto num_blocks = std::min(div_ceil(length, kWorkersPerBlock), kBlockQuota); diff --git a/python/sglang/kernels/jit/csrc/kvcacheio/hisparse.cuh b/python/sglang/kernels/jit/csrc/kvcacheio/hisparse.cuh index 33bad3ceb24a..70b27ef11c22 100644 --- a/python/sglang/kernels/jit/csrc/kvcacheio/hisparse.cuh +++ b/python/sglang/kernels/jit/csrc/kvcacheio/hisparse.cuh @@ -732,6 +732,10 @@ void load_cache_to_device_buffer( const int64_t top_k_tokens_stride = top_k_tokens.strides()[0]; const int64_t top_k_device_locs_stride = top_k_device_locs.strides()[0]; const auto device = LaunchKernel::resolve_device(top_k_tokens.device()); + const auto kernel_device = top_k_tokens.device(); + const void* const host_cache_k_ptr = device_accessible_ptr(host_cache_k, kernel_device); + const void* const host_cache_v_ptr = + (IsMLA || host_cache_v.ndim() == 0) ? nullptr : device_accessible_ptr(host_cache_v, kernel_device); // Generic lambda: int32/int64 kernel variants are compiled for both // seq_lens and req_pool_indices; the correct combo is selected at runtime. @@ -748,8 +752,8 @@ void load_cache_to_device_buffer( static_cast(device_buffer_tokens.data_ptr()), static_cast(host_cache_locs.data_ptr()), static_cast(device_buffer_locs.data_ptr()), - host_cache_k.data_ptr(), - (IsMLA || host_cache_v.ndim() == 0) ? (const void*)nullptr : host_cache_v.data_ptr(), + host_cache_k_ptr, + host_cache_v_ptr, device_buffer_k.data_ptr(), (IsMLA || device_buffer_v.ndim() == 0) ? (void*)nullptr : device_buffer_v.data_ptr(), static_cast(top_k_device_locs.data_ptr()), @@ -909,14 +913,18 @@ void copy_cache_planned( throw std::runtime_error("copy_cache_planned: miss_src/miss_dst row strides differ"); } const auto device = LaunchKernel::resolve_device(miss_src_locs.device()); + const auto kernel_device = miss_src_locs.device(); + const void* const host_cache_k_ptr = device_accessible_ptr(host_cache_k, kernel_device); + const void* const host_cache_v_ptr = + (IsMLA || host_cache_v.ndim() == 0) ? nullptr : device_accessible_ptr(host_cache_v, kernel_device); LaunchKernel(num_blocks, BLOCK_SIZE, device)( copy_cache_planned_kernel, static_cast(miss_src_locs.data_ptr()), static_cast(miss_dst_locs.data_ptr()), static_cast(miss_counts.data_ptr()), static_cast(num_real_reqs.data_ptr()), - host_cache_k.data_ptr(), - (IsMLA || host_cache_v.ndim() == 0) ? (const void*)nullptr : host_cache_v.data_ptr(), + host_cache_k_ptr, + host_cache_v_ptr, device_buffer_k.data_ptr(), (IsMLA || device_buffer_v.ndim() == 0) ? (void*)nullptr : device_buffer_v.data_ptr(), plan_stride, diff --git a/python/sglang/kernels/jit/csrc/kvcacheio/transfer_mamba.cuh b/python/sglang/kernels/jit/csrc/kvcacheio/transfer_mamba.cuh index 139d00d48be9..10631f30587c 100644 --- a/python/sglang/kernels/jit/csrc/kvcacheio/transfer_mamba.cuh +++ b/python/sglang/kernels/jit/csrc/kvcacheio/transfer_mamba.cuh @@ -115,8 +115,8 @@ struct TransferMambaKernel { dim3 grid(grid_x); const auto params = MambaTransferParams{ - .src_base = static_cast(src.data_ptr()), - .dst_base = static_cast(dst.data_ptr()), + .src_base = static_cast(device_accessible_ptr(src, device)), + .dst_base = static_cast(device_accessible_ptr(dst, device)), .layer_ptrs = nullptr, .src_indices = static_cast(src_indices.data_ptr()), .dst_indices = static_cast(dst_indices.data_ptr()), @@ -169,7 +169,7 @@ struct TransferMambaKernel { const auto params = MambaTransferParams{ .src_base = nullptr, - .dst_base = static_cast(dst.data_ptr()), + .dst_base = static_cast(device_accessible_ptr(dst, device)), .layer_ptrs = static_cast(src_ptrs.data_ptr()), .src_indices = static_cast(src_indices.data_ptr()), .dst_indices = static_cast(dst_indices.data_ptr()), diff --git a/python/sglang/kernels/jit/include/sgl_kernel/utils.cuh b/python/sglang/kernels/jit/include/sgl_kernel/utils.cuh index 53356ff89da3..40cf9f3d82a6 100644 --- a/python/sglang/kernels/jit/include/sgl_kernel/utils.cuh +++ b/python/sglang/kernels/jit/include/sgl_kernel/utils.cuh @@ -248,6 +248,66 @@ inline void RuntimeDeviceCheck(DebugInfo location = {}) { return RuntimeDeviceCheck(::cudaGetLastError(), location); } +class DeviceGuard { + public: + explicit DeviceGuard(int target_device) { +#ifdef USE_ROCM + RuntimeDeviceCheck(::hipGetDevice(&original_device_)); + if (original_device_ != target_device) { + RuntimeDeviceCheck(::hipSetDevice(target_device)); + changed_ = true; + } +#else + RuntimeDeviceCheck(::cudaGetDevice(&original_device_)); + if (original_device_ != target_device) { + RuntimeDeviceCheck(::cudaSetDevice(target_device)); + changed_ = true; + } +#endif + } + + ~DeviceGuard() { + if (changed_) { +#ifdef USE_ROCM + (void)::hipSetDevice(original_device_); +#else + (void)::cudaSetDevice(original_device_); +#endif + } + } + + DeviceGuard(const DeviceGuard&) = delete; + DeviceGuard& operator=(const DeviceGuard&) = delete; + + private: + int original_device_ = 0; + bool changed_ = false; +}; + +inline void* device_accessible_ptr(const tvm::ffi::TensorView& tensor, DLDevice target_device) { + void* ptr = tensor.data_ptr(); + const auto tensor_device_type = tensor.device().device_type; + if (tensor_device_type != kDLCPU && tensor_device_type != kDLGPUHost) { + return ptr; + } + + DeviceGuard guard(target_device.device_id); + void* device_ptr = nullptr; +#ifdef USE_ROCM + RuntimeDeviceCheck(::hipHostGetDevicePointer(&device_ptr, ptr, 0)); +#else + RuntimeDeviceCheck(::cudaHostGetDevicePointer(&device_ptr, ptr, 0)); +#endif + return device_ptr; +} + +inline auto alloc_workspace_tensor(size_t required_bytes, DLDevice device) -> tvm::ffi::Tensor { + if (required_bytes == 0) return {}; + DLDataType u8 = {kDLUInt, 8, 1}; + int64_t shape[] = {static_cast(required_bytes)}; + return ffi::empty(tvm::ffi::ShapeView(shape, 1), u8, device); +} + /** * \brief Kernel launcher with automatic stream resolution and PDL support. * diff --git a/python/sglang/srt/mem_cache/memory_pool_host.py b/python/sglang/srt/mem_cache/memory_pool_host.py index 006d0dbde0ea..eb38f3807d82 100644 --- a/python/sglang/srt/mem_cache/memory_pool_host.py +++ b/python/sglang/srt/mem_cache/memory_pool_host.py @@ -43,6 +43,7 @@ from sglang.srt.mem_cache.pool_host.common import ( ALLOC_MEMORY_FUNCS, get_allocator_from_storage, + make_kernel_ptr_table, ) from sglang.srt.mem_cache.pool_host.hisparse import HiSparseHostPoolMixin @@ -272,10 +273,10 @@ def __init__( device=self.gpu_device, ) self.data_ptrs = ( - torch.tensor( - [x.data_ptr() for x in self.data_refs], - dtype=torch.uint64, - device=self.gpu_device, + make_kernel_ptr_table( + self.data_refs, + self.gpu_device, + host_memory_registered=self.pin_memory, ) if self.data_refs else None @@ -302,7 +303,7 @@ def _init_write_back_staging_buffers(self): def get_contiguous_buf_infos(self): """Return per-layer page-row buffers for PD direct-to-host transfer.""" - data_ptrs = [int(self.data_ptrs[i].item()) for i in range(self.layer_num)] + data_ptrs = [tensor.data_ptr() for tensor in self.data_refs] data_lens = [self.kv_buffer[i].nbytes for i in range(self.layer_num)] item_lens = [self.item_bytes * self.dtype.itemsize] * self.layer_num return data_ptrs, data_lens, item_lens @@ -688,10 +689,10 @@ def __init__( device=self.gpu_device, ) self.data_ptrs = ( - torch.tensor( - [x.data_ptr() for x in self.data_refs], - dtype=torch.uint64, - device=self.gpu_device, + make_kernel_ptr_table( + self.data_refs, + self.gpu_device, + host_memory_registered=self.pin_memory, ) if self.data_refs else None diff --git a/python/sglang/srt/mem_cache/pool_host/common.py b/python/sglang/srt/mem_cache/pool_host/common.py index df724029850a..43616e769c56 100644 --- a/python/sglang/srt/mem_cache/pool_host/common.py +++ b/python/sglang/srt/mem_cache/pool_host/common.py @@ -250,6 +250,32 @@ def alloc_with_pin_memory( return buffer +def make_kernel_ptr_table( + tensors: list[torch.Tensor], + target_device: torch.device | str, + *, + host_memory_registered: bool, +) -> torch.Tensor: + device = torch.device(target_device) + if host_memory_registered and device.type == "cuda": + from sgl_kernel.kvcacheio import get_device_accessible_ptr + + if device.index is None: + device_index = torch.cuda.current_device() + else: + device_index = device.index + pointers = [ + get_device_accessible_ptr(tensor, device_index) for tensor in tensors + ] + else: + pointers = [tensor.data_ptr() for tensor in tensors] + return torch.tensor( + pointers, + dtype=torch.uint64, + device=device, + ) + + ALLOC_MEMORY_FUNCS = defaultdict( lambda: alloc_with_host_register, { diff --git a/python/sglang/srt/mem_cache/pool_host/dsa.py b/python/sglang/srt/mem_cache/pool_host/dsa.py index 7ef693fb08a1..4e6ab2f1659c 100644 --- a/python/sglang/srt/mem_cache/pool_host/dsa.py +++ b/python/sglang/srt/mem_cache/pool_host/dsa.py @@ -24,6 +24,7 @@ from sglang.srt.mem_cache.pool_host.common import ( ALLOC_MEMORY_FUNCS, get_allocator_from_storage, + make_kernel_ptr_table, ) from sglang.srt.utils import is_cuda, is_hip, is_mps, is_npu, is_xpu @@ -170,10 +171,10 @@ def init_kv_buffer(self): self.index_k_data_refs = [ self.index_k_with_scale_buffer[i] for i in range(self.layer_num) ] - self.index_k_data_ptrs = torch.tensor( - [x.data_ptr() for x in self.index_k_data_refs], - dtype=torch.uint64, - device=self.device_pool.device, + self.index_k_data_ptrs = make_kernel_ptr_table( + self.index_k_data_refs, + self.device_pool.device, + host_memory_registered=self.pin_memory, ) elif self.layout in ["page_first", "page_first_direct"]: self.index_k_with_scale_buffer = alloc_func( diff --git a/python/sglang/srt/mem_cache/pool_host/mha.py b/python/sglang/srt/mem_cache/pool_host/mha.py index 2bb1d3776d3b..8a3d488cc757 100644 --- a/python/sglang/srt/mem_cache/pool_host/mha.py +++ b/python/sglang/srt/mem_cache/pool_host/mha.py @@ -37,6 +37,7 @@ from sglang.srt.mem_cache.pool_host.common import ( ALLOC_MEMORY_FUNCS, get_allocator_from_storage, + make_kernel_ptr_table, ) from sglang.srt.utils import is_cuda, is_hip, is_mps, is_npu, is_xpu @@ -116,15 +117,15 @@ def __init__( else: self.k_data_refs = [self.k_buffer[i] for i in range(self.layer_num)] self.v_data_refs = [self.v_buffer[i] for i in range(self.layer_num)] - self.k_data_ptrs = torch.tensor( - [x.data_ptr() for x in self.k_data_refs], - dtype=torch.uint64, - device=self.device_pool.device, + self.k_data_ptrs = make_kernel_ptr_table( + self.k_data_refs, + self.device_pool.device, + host_memory_registered=self.pin_memory, ) - self.v_data_ptrs = torch.tensor( - [x.data_ptr() for x in self.v_data_refs], - dtype=torch.uint64, - device=self.device_pool.device, + self.v_data_ptrs = make_kernel_ptr_table( + self.v_data_refs, + self.device_pool.device, + host_memory_registered=self.pin_memory, ) if self.mtp_draft_device_pools: device_pools = (self.device_pool, *self.mtp_draft_device_pools) @@ -775,10 +776,10 @@ def __init__( self.k_data_refs = [self.k_buffer[i] for i in range(self.layer_num)] else: self.k_data_refs = [] - self.k_data_ptrs = torch.tensor( - [x.data_ptr() for x in self.k_data_refs], - dtype=torch.uint64, - device=self.device_pool.device, + self.k_data_ptrs = make_kernel_ptr_table( + self.k_data_refs, + self.device_pool.device, + host_memory_registered=self.pin_memory, ) def get_size_per_token(self): diff --git a/python/sglang/srt/mem_cache/pool_host/mla.py b/python/sglang/srt/mem_cache/pool_host/mla.py index 1d394aa95017..6c69c6a9946a 100644 --- a/python/sglang/srt/mem_cache/pool_host/mla.py +++ b/python/sglang/srt/mem_cache/pool_host/mla.py @@ -28,6 +28,7 @@ from sglang.srt.mem_cache.pool_host.common import ( ALLOC_MEMORY_FUNCS, get_allocator_from_storage, + make_kernel_ptr_table, ) from sglang.srt.mem_cache.pool_host.hisparse import HiSparseHostPoolMixin from sglang.srt.mem_cache.pool_host.npu_memfabric import ( @@ -132,10 +133,10 @@ def __init__( self.data_refs = [transposed[i] for i in range(self.layer_num)] else: self.data_refs = [self.kv_buffer[i] for i in range(self.layer_num)] - self.data_ptrs = torch.tensor( - [x.data_ptr() for x in self.data_refs], - dtype=torch.uint64, - device=self.device_pool.device, + self.data_ptrs = make_kernel_ptr_table( + self.data_refs, + self.device_pool.device, + host_memory_registered=self.pin_memory, ) if self.mtp_draft_device_pools: device_pools = (self.device_pool, *self.mtp_draft_device_pools) @@ -214,7 +215,7 @@ def get_contiguous_buf_infos(self): for registering host memory with the disaggregation transfer engine.""" if self._is_dummy: return [], [], [] - data_ptrs = [int(self.data_ptrs[i].item()) for i in range(self.layer_num)] + data_ptrs = [tensor.data_ptr() for tensor in self.data_refs] if self.layout == "page_first_kv_split": # data_refs are per-layer views of the k_buffer (page-major), so # take the per-layer slab size instead of kv_buffer[i] (a page slab). diff --git a/test/registered/kernels/ops/kvcache/test_hicache_page_first_write_back.py b/test/registered/kernels/ops/kvcache/test_hicache_page_first_write_back.py index 6bf9d13bba0b..547ebe401a2e 100644 --- a/test/registered/kernels/ops/kvcache/test_hicache_page_first_write_back.py +++ b/test/registered/kernels/ops/kvcache/test_hicache_page_first_write_back.py @@ -78,6 +78,18 @@ def _pinned_host_pool(host_pool_cls, **kwargs): ALLOC_MEMORY_FUNCS[DEVICE] = original_alloc +def _registered_host_pool(host_pool_cls, **kwargs): + return host_pool_cls( + host_to_device_ratio=2.0, + host_size=0, + page_size=PAGE_SIZE, + pin_memory=True, + device="cpu", + allocator_type="default", + **kwargs, + ) + + def _fill_with_offset(tensor: torch.Tensor, offset: int) -> None: data = torch.arange( tensor.numel(), device=tensor.device, dtype=tensor.dtype @@ -255,5 +267,146 @@ def test_page_first_staged_write_back_mla(element_dim: int, page_count: int) -> _run_mla(element_dim, page_count) +def test_registered_mmap_pointer_domains_and_all_layer_transfer() -> None: + from sgl_kernel.kvcacheio import get_device_accessible_ptr + + device_pool = MLATokenToKVPool( + size=PAGE_SIZE * 4, + page_size=PAGE_SIZE, + kv_lora_rank=512, + qk_rope_head_dim=64, + dtype=torch.bfloat16, + layer_num=NUM_LAYERS, + device=DEVICE, + enable_memory_saver=False, + ) + host_pool = _registered_host_pool( + MLATokenToKVPoolHost, + device_pool=device_pool, + layout="layer_first", + ) + try: + device_index = torch.cuda.current_device() + kernel_ptrs = [ + get_device_accessible_ptr(tensor, device_index) + for tensor in host_pool.data_refs + ] + assert host_pool.data_ptrs.cpu().tolist() == kernel_ptrs + + raw_ptrs, _, _ = host_pool.get_contiguous_buf_infos() + assert raw_ptrs == [tensor.data_ptr() for tensor in host_pool.data_refs] + page_ptrs, _ = host_pool.get_page_buffer_meta(torch.arange(PAGE_SIZE)) + assert page_ptrs[0] == host_pool.kv_buffer.data_ptr() + + for layer_id in range(NUM_LAYERS): + _fill_with_offset(device_pool.kv_buffer[layer_id], layer_id + 1) + host_indices = _token_indices_for_pages(torch.tensor([0])) + device_indices = _token_indices_for_pages(torch.tensor([1])) + host_pool.can_use_jit = False + host_pool.backup_from_device_all_layer( + device_pool, host_indices, device_indices, "kernel" + ) + torch.cuda.synchronize() + for layer_id in range(NUM_LAYERS): + _assert_pages_equal( + host_pool.data_refs[layer_id], + device_pool.kv_buffer[layer_id], + torch.tensor([0]), + torch.tensor([1]), + ) + finally: + host_pool.destroy() + + +@pytest.mark.parametrize("pool_kind", ["mha", "mla"]) +def test_registered_mmap_page_first_kernel_operands_and_graph( + pool_kind: str, +) -> None: + if pool_kind == "mha": + device_pool = MHATokenToKVPool( + size=PAGE_SIZE * 4, + page_size=PAGE_SIZE, + head_num=1, + head_dim=128, + dtype=torch.bfloat16, + layer_num=1, + device=DEVICE, + enable_memory_saver=False, + ) + host_pool = _registered_host_pool( + MHATokenToKVPoolHost, + device_pool=device_pool, + layout="page_first", + ) + host_refs = [host_pool.k_data_refs[0], host_pool.v_data_refs[0]] + device_refs = [device_pool.k_buffer[0], device_pool.v_buffer[0]] + else: + device_pool = MLATokenToKVPool( + size=PAGE_SIZE * 4, + page_size=PAGE_SIZE, + kv_lora_rank=512, + qk_rope_head_dim=64, + dtype=torch.bfloat16, + layer_num=1, + device=DEVICE, + enable_memory_saver=False, + ) + host_pool = _registered_host_pool( + MLATokenToKVPoolHost, + device_pool=device_pool, + layout="page_first", + ) + host_refs = [host_pool.data_refs[0]] + device_refs = [device_pool.kv_buffer[0]] + + graph = None + try: + assert host_pool.can_use_jit + for index, host_ref in enumerate(host_refs): + _fill_with_offset(host_ref, index + 3) + + host_indices = _token_indices_for_pages(torch.tensor([0])) + aot_device_indices = _token_indices_for_pages(torch.tensor([1])) + host_pool.can_use_jit = False + host_pool.load_to_device_per_layer( + device_pool, host_indices, aot_device_indices, 0, "kernel" + ) + torch.cuda.synchronize() + for host_ref, device_ref in zip(host_refs, device_refs): + _assert_pages_equal( + host_ref, + device_ref, + torch.tensor([0]), + torch.tensor([1]), + ) + + graph_device_indices = _token_indices_for_pages(torch.tensor([2])) + host_pool.can_use_jit = True + capture_stream = torch.cuda.Stream() + capture_stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(capture_stream): + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + host_pool.load_to_device_per_layer( + device_pool, host_indices, graph_device_indices, 0, "kernel" + ) + torch.cuda.current_stream().wait_stream(capture_stream) + for device_ref in device_refs: + device_ref.index_fill_(0, graph_device_indices, 0) + graph.replay() + torch.cuda.synchronize() + for host_ref, device_ref in zip(host_refs, device_refs): + _assert_pages_equal( + host_ref, + device_ref, + torch.tensor([0]), + torch.tensor([2]), + ) + finally: + torch.cuda.synchronize() + del graph + host_pool.destroy() + + if __name__ == "__main__": sys.exit(pytest.main([__file__, "-v", "-s"])) From 55e12628880a0906a470e692d9cefacd03b4bc87 Mon Sep 17 00:00:00 2001 From: AMD-yanfeiwang Date: Wed, 19 Aug 2026 05:01:08 +0000 Subject: [PATCH 2/3] [AMD] Refine JIT host pointer device handling --- .../kernels/jit/csrc/kvcacheio/hicache.cuh | 15 +++-- .../kernels/jit/csrc/kvcacheio/hisparse.cuh | 15 +++-- .../jit/csrc/kvcacheio/transfer_mamba.cuh | 10 +++- .../jit/include/sgl_kernel/runtime.cuh | 52 ++++++++++++++++ .../kernels/jit/include/sgl_kernel/utils.cuh | 60 ------------------- python/sglang/kernels/ops/kvcache/hicache.py | 12 +++- 6 files changed, 87 insertions(+), 77 deletions(-) diff --git a/python/sglang/kernels/jit/csrc/kvcacheio/hicache.cuh b/python/sglang/kernels/jit/csrc/kvcacheio/hicache.cuh index 79b950a897bc..52bf43948747 100644 --- a/python/sglang/kernels/jit/csrc/kvcacheio/hicache.cuh +++ b/python/sglang/kernels/jit/csrc/kvcacheio/hicache.cuh @@ -3,6 +3,7 @@ #include #include +#include #include #include @@ -319,10 +320,11 @@ struct HiCacheKernel { RuntimeCheck(kElementSize == element_bytes, "HicacheKernel: cache dimension mismatch."); const auto device = indices_device.unwrap(); - const auto k_cache_dst_ptr = device_accessible_ptr(k_cache_dst, device); - const auto v_cache_dst_ptr = device_accessible_ptr(v_cache_dst, device); - const auto k_cache_src_ptr = device_accessible_ptr(k_cache_src, device); - const auto v_cache_src_ptr = device_accessible_ptr(v_cache_src, device); + runtime::DeviceGuard device_guard(device.device_id); + const auto k_cache_dst_ptr = runtime::device_accessible_ptr(k_cache_dst); + const auto v_cache_dst_ptr = runtime::device_accessible_ptr(v_cache_dst); + const auto k_cache_src_ptr = runtime::device_accessible_ptr(k_cache_src); + const auto v_cache_src_ptr = runtime::device_accessible_ptr(v_cache_src); const auto indices_dst_ptr = indices_dst.data_ptr(); const auto indices_src_ptr = indices_src.data_ptr(); const auto length = static_cast(L.unwrap()); @@ -441,8 +443,9 @@ struct HiCacheKernel { RuntimeCheck(kElementSize == element_bytes, "HicacheKernel MLA: cache dimension mismatch."); const auto device = indices_device.unwrap(); - const auto cache_dst_ptr = device_accessible_ptr(cache_dst, device); - const auto cache_src_ptr = device_accessible_ptr(cache_src, device); + runtime::DeviceGuard device_guard(device.device_id); + const auto cache_dst_ptr = runtime::device_accessible_ptr(cache_dst); + const auto cache_src_ptr = runtime::device_accessible_ptr(cache_src); const auto indices_dst_ptr = indices_dst.data_ptr(); const auto indices_src_ptr = indices_src.data_ptr(); const auto length = static_cast(L.unwrap()); diff --git a/python/sglang/kernels/jit/csrc/kvcacheio/hisparse.cuh b/python/sglang/kernels/jit/csrc/kvcacheio/hisparse.cuh index 70b27ef11c22..b84176309266 100644 --- a/python/sglang/kernels/jit/csrc/kvcacheio/hisparse.cuh +++ b/python/sglang/kernels/jit/csrc/kvcacheio/hisparse.cuh @@ -1,6 +1,7 @@ #include #include +#include #include #include @@ -731,11 +732,12 @@ void load_cache_to_device_buffer( const int64_t lru_slot_stride_0 = lru_slots.strides()[0]; const int64_t top_k_tokens_stride = top_k_tokens.strides()[0]; const int64_t top_k_device_locs_stride = top_k_device_locs.strides()[0]; - const auto device = LaunchKernel::resolve_device(top_k_tokens.device()); const auto kernel_device = top_k_tokens.device(); - const void* const host_cache_k_ptr = device_accessible_ptr(host_cache_k, kernel_device); + runtime::DeviceGuard device_guard(kernel_device.device_id); + const auto device = LaunchKernel::resolve_device(kernel_device); + const void* const host_cache_k_ptr = runtime::device_accessible_ptr(host_cache_k); const void* const host_cache_v_ptr = - (IsMLA || host_cache_v.ndim() == 0) ? nullptr : device_accessible_ptr(host_cache_v, kernel_device); + (IsMLA || host_cache_v.ndim() == 0) ? nullptr : runtime::device_accessible_ptr(host_cache_v); // Generic lambda: int32/int64 kernel variants are compiled for both // seq_lens and req_pool_indices; the correct combo is selected at runtime. @@ -912,11 +914,12 @@ void copy_cache_planned( if (miss_dst_locs.strides()[0] != plan_stride) { throw std::runtime_error("copy_cache_planned: miss_src/miss_dst row strides differ"); } - const auto device = LaunchKernel::resolve_device(miss_src_locs.device()); const auto kernel_device = miss_src_locs.device(); - const void* const host_cache_k_ptr = device_accessible_ptr(host_cache_k, kernel_device); + runtime::DeviceGuard device_guard(kernel_device.device_id); + const auto device = LaunchKernel::resolve_device(kernel_device); + const void* const host_cache_k_ptr = runtime::device_accessible_ptr(host_cache_k); const void* const host_cache_v_ptr = - (IsMLA || host_cache_v.ndim() == 0) ? nullptr : device_accessible_ptr(host_cache_v, kernel_device); + (IsMLA || host_cache_v.ndim() == 0) ? nullptr : runtime::device_accessible_ptr(host_cache_v); LaunchKernel(num_blocks, BLOCK_SIZE, device)( copy_cache_planned_kernel, static_cast(miss_src_locs.data_ptr()), diff --git a/python/sglang/kernels/jit/csrc/kvcacheio/transfer_mamba.cuh b/python/sglang/kernels/jit/csrc/kvcacheio/transfer_mamba.cuh index 10631f30587c..e030d39c559b 100644 --- a/python/sglang/kernels/jit/csrc/kvcacheio/transfer_mamba.cuh +++ b/python/sglang/kernels/jit/csrc/kvcacheio/transfer_mamba.cuh @@ -1,5 +1,7 @@ #pragma once +#include + #include "hicache.cuh" #include #include @@ -111,12 +113,13 @@ struct TransferMambaKernel { if (num_items == 0) return; const auto device = device_.unwrap(); + runtime::DeviceGuard device_guard(device.device_id); const int grid_x = static_cast(std::min(static_cast(kBlockQuotaLoad), num_items)); dim3 grid(grid_x); const auto params = MambaTransferParams{ - .src_base = static_cast(device_accessible_ptr(src, device)), - .dst_base = static_cast(device_accessible_ptr(dst, device)), + .src_base = static_cast(runtime::device_accessible_ptr(src)), + .dst_base = static_cast(runtime::device_accessible_ptr(dst)), .layer_ptrs = nullptr, .src_indices = static_cast(src_indices.data_ptr()), .dst_indices = static_cast(dst_indices.data_ptr()), @@ -163,13 +166,14 @@ struct TransferMambaKernel { if (num_items == 0) return; const auto device = device_.unwrap(); + runtime::DeviceGuard device_guard(device.device_id); const int64_t total_work = num_items * num_layers; const int grid_x = static_cast(std::min(static_cast(kBlockQuotaBackup), total_work)); dim3 grid(grid_x); const auto params = MambaTransferParams{ .src_base = nullptr, - .dst_base = static_cast(device_accessible_ptr(dst, device)), + .dst_base = static_cast(runtime::device_accessible_ptr(dst)), .layer_ptrs = static_cast(src_ptrs.data_ptr()), .src_indices = static_cast(src_indices.data_ptr()), .dst_indices = static_cast(dst_indices.data_ptr()), diff --git a/python/sglang/kernels/jit/include/sgl_kernel/runtime.cuh b/python/sglang/kernels/jit/include/sgl_kernel/runtime.cuh index 3e4009df3a67..dbcb0f141c83 100644 --- a/python/sglang/kernels/jit/include/sgl_kernel/runtime.cuh +++ b/python/sglang/kernels/jit/include/sgl_kernel/runtime.cuh @@ -50,6 +50,58 @@ namespace sglang { namespace host::runtime { +class DeviceGuard { + public: + explicit DeviceGuard(int target_device) { +#ifdef USE_ROCM + RuntimeDeviceCheck(::hipGetDevice(&original_device_)); + if (original_device_ != target_device) { + RuntimeDeviceCheck(::hipSetDevice(target_device)); + changed_ = true; + } +#else + RuntimeDeviceCheck(::cudaGetDevice(&original_device_)); + if (original_device_ != target_device) { + RuntimeDeviceCheck(::cudaSetDevice(target_device)); + changed_ = true; + } +#endif + } + + ~DeviceGuard() { + if (changed_) { +#ifdef USE_ROCM + (void)::hipSetDevice(original_device_); +#else + (void)::cudaSetDevice(original_device_); +#endif + } + } + + DeviceGuard(const DeviceGuard&) = delete; + DeviceGuard& operator=(const DeviceGuard&) = delete; + + private: + int original_device_ = 0; + bool changed_ = false; +}; + +inline void* device_accessible_ptr(const tvm::ffi::TensorView& tensor) { + void* ptr = tensor.data_ptr(); + const auto tensor_device_type = tensor.device().device_type; + if (tensor_device_type != kDLCPU && tensor_device_type != kDLGPUHost) { + return ptr; + } + + void* device_ptr = nullptr; +#ifdef USE_ROCM + RuntimeDeviceCheck(::hipHostGetDevicePointer(&device_ptr, ptr, 0)); +#else + RuntimeDeviceCheck(::cudaHostGetDevicePointer(&device_ptr, ptr, 0)); +#endif + return device_ptr; +} + // Return the maximum number of active blocks per SM for the given kernel template inline auto get_blocks_per_sm(T&& kernel, int32_t block_dim, std::size_t dynamic_smem = 0) -> uint32_t { diff --git a/python/sglang/kernels/jit/include/sgl_kernel/utils.cuh b/python/sglang/kernels/jit/include/sgl_kernel/utils.cuh index 40cf9f3d82a6..53356ff89da3 100644 --- a/python/sglang/kernels/jit/include/sgl_kernel/utils.cuh +++ b/python/sglang/kernels/jit/include/sgl_kernel/utils.cuh @@ -248,66 +248,6 @@ inline void RuntimeDeviceCheck(DebugInfo location = {}) { return RuntimeDeviceCheck(::cudaGetLastError(), location); } -class DeviceGuard { - public: - explicit DeviceGuard(int target_device) { -#ifdef USE_ROCM - RuntimeDeviceCheck(::hipGetDevice(&original_device_)); - if (original_device_ != target_device) { - RuntimeDeviceCheck(::hipSetDevice(target_device)); - changed_ = true; - } -#else - RuntimeDeviceCheck(::cudaGetDevice(&original_device_)); - if (original_device_ != target_device) { - RuntimeDeviceCheck(::cudaSetDevice(target_device)); - changed_ = true; - } -#endif - } - - ~DeviceGuard() { - if (changed_) { -#ifdef USE_ROCM - (void)::hipSetDevice(original_device_); -#else - (void)::cudaSetDevice(original_device_); -#endif - } - } - - DeviceGuard(const DeviceGuard&) = delete; - DeviceGuard& operator=(const DeviceGuard&) = delete; - - private: - int original_device_ = 0; - bool changed_ = false; -}; - -inline void* device_accessible_ptr(const tvm::ffi::TensorView& tensor, DLDevice target_device) { - void* ptr = tensor.data_ptr(); - const auto tensor_device_type = tensor.device().device_type; - if (tensor_device_type != kDLCPU && tensor_device_type != kDLGPUHost) { - return ptr; - } - - DeviceGuard guard(target_device.device_id); - void* device_ptr = nullptr; -#ifdef USE_ROCM - RuntimeDeviceCheck(::hipHostGetDevicePointer(&device_ptr, ptr, 0)); -#else - RuntimeDeviceCheck(::cudaHostGetDevicePointer(&device_ptr, ptr, 0)); -#endif - return device_ptr; -} - -inline auto alloc_workspace_tensor(size_t required_bytes, DLDevice device) -> tvm::ffi::Tensor { - if (required_bytes == 0) return {}; - DLDataType u8 = {kDLUInt, 8, 1}; - int64_t shape[] = {static_cast(required_bytes)}; - return ffi::empty(tvm::ffi::ShapeView(shape, 1), u8, device); -} - /** * \brief Kernel launcher with automatic stream resolution and PDL support. * diff --git a/python/sglang/kernels/ops/kvcache/hicache.py b/python/sglang/kernels/ops/kvcache/hicache.py index 66c5c0febcba..6844b34e04a4 100644 --- a/python/sglang/kernels/ops/kvcache/hicache.py +++ b/python/sglang/kernels/ops/kvcache/hicache.py @@ -3,14 +3,22 @@ import logging from typing import TYPE_CHECKING -from sglang.kernels.jit.utils import cache_once, load_jit, make_cpp_args +from sglang.kernels.jit.utils import ( + cache_once, + is_hip_runtime, + load_jit, + make_cpp_args, +) from sglang.kernels.kernel_api_logging import debug_kernel_api if TYPE_CHECKING: import torch from tvm_ffi.module import Module -DEFAULT_BLOCK_QUOTA = 2 +_is_hip = is_hip_runtime() + +# ROCm needs a wider grid to saturate mapped-host transfers; CUDA keeps the legacy quota. +DEFAULT_BLOCK_QUOTA = 32 if _is_hip else 2 @cache_once From 7c10333c940a3c624e35b0dd6e9af0d1b67f3c99 Mon Sep 17 00:00:00 2001 From: AMD-yanfeiwang Date: Fri, 11 Sep 2026 09:57:18 +0000 Subject: [PATCH 3/3] [HiCache] Simplify JIT host pointer resolution --- .../kernels/jit/csrc/kvcacheio/hicache.cuh | 14 +++---- .../kernels/jit/csrc/kvcacheio/hisparse.cuh | 10 ++--- .../jit/csrc/kvcacheio/transfer_mamba.cuh | 8 ++-- .../jit/include/sgl_kernel/runtime.cuh | 38 +------------------ 4 files changed, 14 insertions(+), 56 deletions(-) diff --git a/python/sglang/kernels/jit/csrc/kvcacheio/hicache.cuh b/python/sglang/kernels/jit/csrc/kvcacheio/hicache.cuh index 52bf43948747..3a9a5afa5e4e 100644 --- a/python/sglang/kernels/jit/csrc/kvcacheio/hicache.cuh +++ b/python/sglang/kernels/jit/csrc/kvcacheio/hicache.cuh @@ -320,11 +320,10 @@ struct HiCacheKernel { RuntimeCheck(kElementSize == element_bytes, "HicacheKernel: cache dimension mismatch."); const auto device = indices_device.unwrap(); - runtime::DeviceGuard device_guard(device.device_id); - const auto k_cache_dst_ptr = runtime::device_accessible_ptr(k_cache_dst); - const auto v_cache_dst_ptr = runtime::device_accessible_ptr(v_cache_dst); - const auto k_cache_src_ptr = runtime::device_accessible_ptr(k_cache_src); - const auto v_cache_src_ptr = runtime::device_accessible_ptr(v_cache_src); + const auto k_cache_dst_ptr = runtime::get_device_accessible_ptr(k_cache_dst); + const auto v_cache_dst_ptr = runtime::get_device_accessible_ptr(v_cache_dst); + const auto k_cache_src_ptr = runtime::get_device_accessible_ptr(k_cache_src); + const auto v_cache_src_ptr = runtime::get_device_accessible_ptr(v_cache_src); const auto indices_dst_ptr = indices_dst.data_ptr(); const auto indices_src_ptr = indices_src.data_ptr(); const auto length = static_cast(L.unwrap()); @@ -443,9 +442,8 @@ struct HiCacheKernel { RuntimeCheck(kElementSize == element_bytes, "HicacheKernel MLA: cache dimension mismatch."); const auto device = indices_device.unwrap(); - runtime::DeviceGuard device_guard(device.device_id); - const auto cache_dst_ptr = runtime::device_accessible_ptr(cache_dst); - const auto cache_src_ptr = runtime::device_accessible_ptr(cache_src); + const auto cache_dst_ptr = runtime::get_device_accessible_ptr(cache_dst); + const auto cache_src_ptr = runtime::get_device_accessible_ptr(cache_src); const auto indices_dst_ptr = indices_dst.data_ptr(); const auto indices_src_ptr = indices_src.data_ptr(); const auto length = static_cast(L.unwrap()); diff --git a/python/sglang/kernels/jit/csrc/kvcacheio/hisparse.cuh b/python/sglang/kernels/jit/csrc/kvcacheio/hisparse.cuh index b84176309266..4d1da2c32c1a 100644 --- a/python/sglang/kernels/jit/csrc/kvcacheio/hisparse.cuh +++ b/python/sglang/kernels/jit/csrc/kvcacheio/hisparse.cuh @@ -733,11 +733,10 @@ void load_cache_to_device_buffer( const int64_t top_k_tokens_stride = top_k_tokens.strides()[0]; const int64_t top_k_device_locs_stride = top_k_device_locs.strides()[0]; const auto kernel_device = top_k_tokens.device(); - runtime::DeviceGuard device_guard(kernel_device.device_id); const auto device = LaunchKernel::resolve_device(kernel_device); - const void* const host_cache_k_ptr = runtime::device_accessible_ptr(host_cache_k); + const void* const host_cache_k_ptr = runtime::get_device_accessible_ptr(host_cache_k); const void* const host_cache_v_ptr = - (IsMLA || host_cache_v.ndim() == 0) ? nullptr : runtime::device_accessible_ptr(host_cache_v); + (IsMLA || host_cache_v.ndim() == 0) ? nullptr : runtime::get_device_accessible_ptr(host_cache_v); // Generic lambda: int32/int64 kernel variants are compiled for both // seq_lens and req_pool_indices; the correct combo is selected at runtime. @@ -915,11 +914,10 @@ void copy_cache_planned( throw std::runtime_error("copy_cache_planned: miss_src/miss_dst row strides differ"); } const auto kernel_device = miss_src_locs.device(); - runtime::DeviceGuard device_guard(kernel_device.device_id); const auto device = LaunchKernel::resolve_device(kernel_device); - const void* const host_cache_k_ptr = runtime::device_accessible_ptr(host_cache_k); + const void* const host_cache_k_ptr = runtime::get_device_accessible_ptr(host_cache_k); const void* const host_cache_v_ptr = - (IsMLA || host_cache_v.ndim() == 0) ? nullptr : runtime::device_accessible_ptr(host_cache_v); + (IsMLA || host_cache_v.ndim() == 0) ? nullptr : runtime::get_device_accessible_ptr(host_cache_v); LaunchKernel(num_blocks, BLOCK_SIZE, device)( copy_cache_planned_kernel, static_cast(miss_src_locs.data_ptr()), diff --git a/python/sglang/kernels/jit/csrc/kvcacheio/transfer_mamba.cuh b/python/sglang/kernels/jit/csrc/kvcacheio/transfer_mamba.cuh index e030d39c559b..cafc003c6d64 100644 --- a/python/sglang/kernels/jit/csrc/kvcacheio/transfer_mamba.cuh +++ b/python/sglang/kernels/jit/csrc/kvcacheio/transfer_mamba.cuh @@ -113,13 +113,12 @@ struct TransferMambaKernel { if (num_items == 0) return; const auto device = device_.unwrap(); - runtime::DeviceGuard device_guard(device.device_id); const int grid_x = static_cast(std::min(static_cast(kBlockQuotaLoad), num_items)); dim3 grid(grid_x); const auto params = MambaTransferParams{ - .src_base = static_cast(runtime::device_accessible_ptr(src)), - .dst_base = static_cast(runtime::device_accessible_ptr(dst)), + .src_base = static_cast(runtime::get_device_accessible_ptr(src)), + .dst_base = static_cast(runtime::get_device_accessible_ptr(dst)), .layer_ptrs = nullptr, .src_indices = static_cast(src_indices.data_ptr()), .dst_indices = static_cast(dst_indices.data_ptr()), @@ -166,14 +165,13 @@ struct TransferMambaKernel { if (num_items == 0) return; const auto device = device_.unwrap(); - runtime::DeviceGuard device_guard(device.device_id); const int64_t total_work = num_items * num_layers; const int grid_x = static_cast(std::min(static_cast(kBlockQuotaBackup), total_work)); dim3 grid(grid_x); const auto params = MambaTransferParams{ .src_base = nullptr, - .dst_base = static_cast(runtime::device_accessible_ptr(dst)), + .dst_base = static_cast(runtime::get_device_accessible_ptr(dst)), .layer_ptrs = static_cast(src_ptrs.data_ptr()), .src_indices = static_cast(src_indices.data_ptr()), .dst_indices = static_cast(dst_indices.data_ptr()), diff --git a/python/sglang/kernels/jit/include/sgl_kernel/runtime.cuh b/python/sglang/kernels/jit/include/sgl_kernel/runtime.cuh index dbcb0f141c83..9f1b4f1d869b 100644 --- a/python/sglang/kernels/jit/include/sgl_kernel/runtime.cuh +++ b/python/sglang/kernels/jit/include/sgl_kernel/runtime.cuh @@ -50,43 +50,7 @@ namespace sglang { namespace host::runtime { -class DeviceGuard { - public: - explicit DeviceGuard(int target_device) { -#ifdef USE_ROCM - RuntimeDeviceCheck(::hipGetDevice(&original_device_)); - if (original_device_ != target_device) { - RuntimeDeviceCheck(::hipSetDevice(target_device)); - changed_ = true; - } -#else - RuntimeDeviceCheck(::cudaGetDevice(&original_device_)); - if (original_device_ != target_device) { - RuntimeDeviceCheck(::cudaSetDevice(target_device)); - changed_ = true; - } -#endif - } - - ~DeviceGuard() { - if (changed_) { -#ifdef USE_ROCM - (void)::hipSetDevice(original_device_); -#else - (void)::cudaSetDevice(original_device_); -#endif - } - } - - DeviceGuard(const DeviceGuard&) = delete; - DeviceGuard& operator=(const DeviceGuard&) = delete; - - private: - int original_device_ = 0; - bool changed_ = false; -}; - -inline void* device_accessible_ptr(const tvm::ffi::TensorView& tensor) { +inline void* get_device_accessible_ptr(const tvm::ffi::TensorView& tensor) { void* ptr = tensor.data_ptr(); const auto tensor_device_type = tensor.device().device_type; if (tensor_device_type != kDLCPU && tensor_device_type != kDLGPUHost) {