minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

[AMD] Fix registered HiCache host pointer aliases (#35233)

Browse Source
This commit is contained in:
AMD-yanfeiwang
2026-09-14 15:49:05 -07:00
committed by GitHub
parent 2fca6d69aa
commit 0163f8ff74
15 changed files with 307 additions and 51 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/kernels/aot/csrc/common_extension.cc
+1
View File
@@ -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
+1
View File
@@ -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
+40 -4
View File
@@ -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,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<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
+2
View File
@@ -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
+5
View File
@@ -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
+9 -8
View File
@@ -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,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 = 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<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 +441,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 = 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<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
+15 -6
View File
@@ -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,11 @@ 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 auto device = LaunchKernel::resolve_device(kernel_device);
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::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.
@@ -748,8 +753,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 +913,19 @@ 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 auto device = LaunchKernel::resolve_device(kernel_device);
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::get_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
+5 -3
View File
@@ -1,5 +1,7 @@
#pragma once
#include <sgl_kernel/runtime.cuh>
#include "hicache.cuh"
#include <algorithm>
#include <cstdint>
@@ -115,8 +117,8 @@ struct TransferMambaKernel {
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::get_device_accessible_ptr(src)),
.dst_base = static_cast<char*>(runtime::get_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()),
@@ -169,7 +171,7 @@ struct TransferMambaKernel {
const auto params = MambaTransferParams{
.src_base = nullptr,
.dst_base = static_cast<char*>(dst.data_ptr()),
.dst_base = static_cast<char*>(runtime::get_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
+16
View File
@@ -50,6 +50,22 @@ namespace sglang {
namespace host::runtime {
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) {
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/srt/mem_cache/memory_pool_host.py
+10 -9
View File
@@ -46,6 +46,7 @@ from sglang.srt.mem_cache.pool_host.base import (
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
@@ -283,10 +284,10 @@ class DeepSeekV4PagedHostPool(HiSparseHostPoolMixin, HostKVCache):
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
@@ -336,7 +337,7 @@ class DeepSeekV4PagedHostPool(HiSparseHostPoolMixin, HostKVCache):
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
@@ -787,10 +788,10 @@ class DeepSeekV4StateHostPool(HostKVCache):
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
python/sglang/srt/mem_cache/pool_host/common.py
+26
View File
@@ -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,
{
python/sglang/srt/mem_cache/pool_host/dsa.py
+5 -4
View File
@@ -24,6 +24,7 @@ from sglang.srt.mem_cache.pool_host.base import (
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 @@ class DSAIndexerPoolHost(HostKVCache):
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(
python/sglang/srt/mem_cache/pool_host/mha.py
+13 -12
View File
@@ -37,6 +37,7 @@ from sglang.srt.mem_cache.pool_host.base import (
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 @@ class MHATokenToKVPoolHost(HostKVCache):
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 @@ class MHATokenToKOnlyPoolHost(HostKVCache):
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):
python/sglang/srt/mem_cache/pool_host/mla.py
+6 -5
View File
@@ -28,6 +28,7 @@ from sglang.srt.mem_cache.pool_host.base import (
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 @@ class MLATokenToKVPoolHost(HiSparseHostPoolMixin, HostKVCache):
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 @@ class MLATokenToKVPoolHost(HiSparseHostPoolMixin, HostKVCache):
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).