[AMD] Fix registered HiCache host pointer aliases

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) -> ()");

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) -> ()");

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<c10::DeviceIndex>(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<int64_t>(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,17 +367,21 @@ 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);
   const int32_t num_blocks = div_up(num_items, items_per_warp * num_warps_per_block);
   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<uintptr_t>() : nullptr;
   const uintptr_t* dst_k_tbl_ptr = dst_k_layers.defined() ? dst_k_layers.data_ptr<uintptr_t>() : nullptr;
   const uintptr_t* src_v_tbl_ptr = IsMLA || !src_v_layers.defined() ? nullptr : src_v_layers.data_ptr<uintptr_t>();

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<at::Tensor>& src_layers,
     std::vector<at::Tensor> dst_layers,

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.
 

python/sglang/kernels/jit/csrc/kvcacheio/hicache.cuh

@@ -3,6 +3,7 @@
 #include <sgl_kernel/tensor.h>
 #include <sgl_kernel/utils.h>
 
+#include <sgl_kernel/runtime.cuh>
 #include <sgl_kernel/utils.cuh>
 #include <sgl_kernel/vec.cuh>
 
@@ -318,17 +319,18 @@ 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();
+    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<uint32_t>(L.unwrap());
     const auto kv_cache_src_stride = static_cast<int64_t>(N.unwrap() * dtype_size);
     const auto kv_cache_dst_stride = static_cast<int64_t>(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 +442,16 @@ 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();
+    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<uint32_t>(L.unwrap());
     const auto cache_src_stride = static_cast<int64_t>(N.unwrap() * dtype_size);
     const auto cache_dst_stride = static_cast<int64_t>(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);

python/sglang/kernels/jit/csrc/kvcacheio/hisparse.cuh

@@ -1,6 +1,7 @@
 #include <sgl_kernel/tensor.h>
 #include <sgl_kernel/utils.h>
 
+#include <sgl_kernel/runtime.cuh>
 #include <sgl_kernel/utils.cuh>
 
 #include <sgl_kernel/deepseek_v4/kvcacheio.cuh>
@@ -731,7 +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();
+  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 : 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.
@@ -748,8 +754,8 @@ void load_cache_to_device_buffer(
         static_cast<int32_t*>(device_buffer_tokens.data_ptr()),
         static_cast<const int64_t*>(host_cache_locs.data_ptr()),
         static_cast<const int32_t*>(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<int32_t*>(top_k_device_locs.data_ptr()),
@@ -908,15 +914,20 @@ 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();
+  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 : runtime::device_accessible_ptr(host_cache_v);
   LaunchKernel(num_blocks, BLOCK_SIZE, device)(
       copy_cache_planned_kernel<BLOCK_SIZE, IsMLA, IsDsv4Layout, SkipIO>,
       static_cast<const int64_t*>(miss_src_locs.data_ptr()),
       static_cast<const int32_t*>(miss_dst_locs.data_ptr()),
       static_cast<const int32_t*>(miss_counts.data_ptr()),
       static_cast<const int32_t*>(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,

python/sglang/kernels/jit/csrc/kvcacheio/transfer_mamba.cuh

@@ -1,5 +1,7 @@
 #pragma once
 
+#include <sgl_kernel/runtime.cuh>
+
 #include "hicache.cuh"
 #include <algorithm>
 #include <cstdint>
@@ -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<int>(std::min(static_cast<int64_t>(kBlockQuotaLoad), num_items));
     dim3 grid(grid_x);
 
     const auto params = MambaTransferParams{
-        .src_base = static_cast<const char*>(src.data_ptr()),
-        .dst_base = static_cast<char*>(dst.data_ptr()),
+        .src_base = static_cast<const char*>(runtime::device_accessible_ptr(src)),
+        .dst_base = static_cast<char*>(runtime::device_accessible_ptr(dst)),
         .layer_ptrs = nullptr,
         .src_indices = static_cast<const int64_t*>(src_indices.data_ptr()),
         .dst_indices = static_cast<const int64_t*>(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<int>(std::min(static_cast<int64_t>(kBlockQuotaBackup), total_work));
     dim3 grid(grid_x);
 
     const auto params = MambaTransferParams{
         .src_base = nullptr,
-        .dst_base = static_cast<char*>(dst.data_ptr()),
+        .dst_base = static_cast<char*>(runtime::device_accessible_ptr(dst)),
         .layer_ptrs = static_cast<const uintptr_t*>(src_ptrs.data_ptr()),
         .src_indices = static_cast<const int64_t*>(src_indices.data_ptr()),
         .dst_indices = static_cast<const int64_t*>(dst_indices.data_ptr()),

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 <typename T>
 inline auto get_blocks_per_sm(T&& kernel, int32_t block_dim, std::size_t dynamic_smem = 0) -> uint32_t {

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

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