minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 10 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).
test/registered/kernels/ops/kvcache/test_hicache_page_first_write_back.py
+153
View File
@@ -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"]))