[AMD] Enable JIT staged HiCache write-back and fix CPU-index crash (#28534)
Co-authored-by: Duyi-Wang <duyi.wang@amd.com>
This commit is contained in:
AMD-yanfeiwang
co-authored by
Duyi-Wang
Duyi-Wang
parent
61602b95fb
commit
d74619b373
@@ -37,44 +37,86 @@ inline constexpr auto get_mem_package() {
|
||||
template <int kUnit>
|
||||
using PackageType = decltype(get_mem_package<kUnit>());
|
||||
|
||||
// NVIDIA exposes an explicit "do not allocate in L1" cache hint via PTX. ROCm
|
||||
// has no equivalent PTX, but non-temporal (streaming) loads/stores express the
|
||||
// same intent for one-shot HiCache write-back traffic that should not pollute
|
||||
// the cache. Guard the PTX behind USE_ROCM so the JIT module also compiles with
|
||||
// hipcc; see python/sglang/jit_kernel/utils.py for the ROCm build flags.
|
||||
#ifdef USE_ROCM
|
||||
// Native Clang vector types so a single __builtin_nontemporal_{load,store} maps
|
||||
// to one vectorized global_{load,store}_dwordx{2,4}. Issuing N independent
|
||||
// 32-bit nontemporal ops instead leaves merging to the LoadStoreVectorizer,
|
||||
// which is not guaranteed and may drop the nontemporal hint, throttling HiCache
|
||||
// bandwidth. uint2/uint4 already carry 8B/16B alignment matching the vector
|
||||
// types, so the pointer reinterpret_casts stay correctly aligned.
|
||||
typedef uint32_t native_uint2 __attribute__((ext_vector_type(2)));
|
||||
typedef uint32_t native_uint4 __attribute__((ext_vector_type(4)));
|
||||
#endif
|
||||
|
||||
SGL_DEVICE uint1 load_nc(const uint1* __restrict__ src) {
|
||||
#ifndef USE_ROCM
|
||||
uint32_t tmp;
|
||||
asm volatile("ld.global.L1::no_allocate.b32 %0,[%1];" : "=r"(tmp) : "l"(src));
|
||||
return uint1{tmp};
|
||||
#else
|
||||
return uint1{__builtin_nontemporal_load(&src->x)};
|
||||
#endif
|
||||
}
|
||||
|
||||
SGL_DEVICE uint2 load_nc(const uint2* __restrict__ src) {
|
||||
#ifndef USE_ROCM
|
||||
uint32_t tmp0, tmp1;
|
||||
asm volatile("ld.global.L1::no_allocate.v2.b32 {%0,%1},[%2];" : "=r"(tmp0), "=r"(tmp1) : "l"(src));
|
||||
return uint2{tmp0, tmp1};
|
||||
#else
|
||||
native_uint2 tmp = __builtin_nontemporal_load(reinterpret_cast<const native_uint2*>(src));
|
||||
return __builtin_bit_cast(uint2, tmp);
|
||||
#endif
|
||||
}
|
||||
|
||||
SGL_DEVICE uint4 load_nc(const uint4* __restrict__ src) {
|
||||
#ifndef USE_ROCM
|
||||
uint32_t tmp0, tmp1, tmp2, tmp3;
|
||||
asm volatile("ld.global.L1::no_allocate.v4.b32 {%0,%1,%2,%3},[%4];"
|
||||
: "=r"(tmp0), "=r"(tmp1), "=r"(tmp2), "=r"(tmp3)
|
||||
: "l"(src));
|
||||
return uint4{tmp0, tmp1, tmp2, tmp3};
|
||||
#else
|
||||
native_uint4 tmp = __builtin_nontemporal_load(reinterpret_cast<const native_uint4*>(src));
|
||||
return __builtin_bit_cast(uint4, tmp);
|
||||
#endif
|
||||
}
|
||||
|
||||
SGL_DEVICE void store_nc(uint1* __restrict__ dst, const uint1& value) {
|
||||
#ifndef USE_ROCM
|
||||
uint32_t tmp = value.x;
|
||||
asm volatile("st.global.L1::no_allocate.b32 [%0],%1;" ::"l"(dst), "r"(tmp));
|
||||
#else
|
||||
__builtin_nontemporal_store(value.x, &dst->x);
|
||||
#endif
|
||||
}
|
||||
|
||||
SGL_DEVICE void store_nc(uint2* __restrict__ dst, const uint2& value) {
|
||||
#ifndef USE_ROCM
|
||||
uint32_t tmp0 = value.x;
|
||||
uint32_t tmp1 = value.y;
|
||||
asm volatile("st.global.L1::no_allocate.v2.b32 [%0],{%1,%2};" ::"l"(dst), "r"(tmp0), "r"(tmp1));
|
||||
#else
|
||||
__builtin_nontemporal_store(__builtin_bit_cast(native_uint2, value), reinterpret_cast<native_uint2*>(dst));
|
||||
#endif
|
||||
}
|
||||
|
||||
SGL_DEVICE void store_nc(uint4* __restrict__ dst, const uint4& value) {
|
||||
#ifndef USE_ROCM
|
||||
uint32_t tmp0 = value.x;
|
||||
uint32_t tmp1 = value.y;
|
||||
uint32_t tmp2 = value.z;
|
||||
uint32_t tmp3 = value.w;
|
||||
asm volatile(
|
||||
"st.global.L1::no_allocate.v4.b32 [%0],{%1,%2,%3,%4};" ::"l"(dst), "r"(tmp0), "r"(tmp1), "r"(tmp2), "r"(tmp3));
|
||||
#else
|
||||
__builtin_nontemporal_store(__builtin_bit_cast(native_uint4, value), reinterpret_cast<native_uint4*>(dst));
|
||||
#endif
|
||||
}
|
||||
|
||||
} // namespace details
|
||||
@@ -256,18 +298,18 @@ struct HiCacheKernel {
|
||||
TensorMatcher({-1, D}) //
|
||||
.with_strides({N, 1})
|
||||
.with_dtype(cache_dtype)
|
||||
.with_device<kDLCUDA, kDLCUDAHost, kDLCPU>()
|
||||
.with_device<kDLGPU, kDLGPUHost, kDLCPU>()
|
||||
.verify(k_cache_src)
|
||||
.verify(v_cache_src);
|
||||
TensorMatcher({-1, D}) //
|
||||
.with_strides({M, 1})
|
||||
.with_dtype(cache_dtype)
|
||||
.with_device<kDLCUDA, kDLCUDAHost, kDLCPU>()
|
||||
.with_device<kDLGPU, kDLGPUHost, kDLCPU>()
|
||||
.verify(k_cache_dst)
|
||||
.verify(v_cache_dst);
|
||||
TensorMatcher({L}) //
|
||||
.with_dtype<int32_t, int64_t>(indices_dtype)
|
||||
.with_device<kDLCUDA>(indices_device)
|
||||
.with_device<kDLGPU>(indices_device)
|
||||
.verify(indices_src)
|
||||
.verify(indices_dst);
|
||||
|
||||
@@ -323,14 +365,14 @@ struct HiCacheKernel {
|
||||
|
||||
TensorMatcher({N}) //
|
||||
.with_dtype<uint64_t>()
|
||||
.with_device<kDLCUDA>(device_)
|
||||
.with_device<kDLGPU>(device_)
|
||||
.verify(k_ptr_src)
|
||||
.verify(v_ptr_src)
|
||||
.verify(k_ptr_dst)
|
||||
.verify(v_ptr_dst);
|
||||
TensorMatcher({L}) //
|
||||
.with_dtype<int32_t, int64_t>(dtype_)
|
||||
.with_device<kDLCUDA>(device_)
|
||||
.with_device<kDLGPU>(device_)
|
||||
.verify(indices_src)
|
||||
.verify(indices_dst);
|
||||
|
||||
@@ -381,16 +423,16 @@ struct HiCacheKernel {
|
||||
TensorMatcher({-1, D}) //
|
||||
.with_strides({N, 1})
|
||||
.with_dtype(cache_dtype)
|
||||
.with_device<kDLCUDA, kDLCUDAHost, kDLCPU>()
|
||||
.with_device<kDLGPU, kDLGPUHost, kDLCPU>()
|
||||
.verify(cache_src);
|
||||
TensorMatcher({-1, D}) //
|
||||
.with_strides({M, 1})
|
||||
.with_dtype(cache_dtype)
|
||||
.with_device<kDLCUDA, kDLCUDAHost, kDLCPU>()
|
||||
.with_device<kDLGPU, kDLGPUHost, kDLCPU>()
|
||||
.verify(cache_dst);
|
||||
TensorMatcher({L}) //
|
||||
.with_dtype<int32_t, int64_t>(indices_dtype)
|
||||
.with_device<kDLCUDA>(indices_device)
|
||||
.with_device<kDLGPU>(indices_device)
|
||||
.verify(indices_src)
|
||||
.verify(indices_dst);
|
||||
|
||||
@@ -441,12 +483,12 @@ struct HiCacheKernel {
|
||||
|
||||
TensorMatcher({N}) //
|
||||
.with_dtype<uint64_t>()
|
||||
.with_device<kDLCUDA>(device_)
|
||||
.with_device<kDLGPU>(device_)
|
||||
.verify(ptr_src)
|
||||
.verify(ptr_dst);
|
||||
TensorMatcher({L}) //
|
||||
.with_dtype<int32_t, int64_t>(dtype_)
|
||||
.with_device<kDLCUDA>(device_)
|
||||
.with_device<kDLGPU>(device_)
|
||||
.verify(indices_src)
|
||||
.verify(indices_dst);
|
||||
|
||||
|
||||
@@ -210,41 +210,41 @@ struct HiCacheStagedWriteBackKernel {
|
||||
|
||||
TensorMatcher({T, N, D}) //
|
||||
.with_dtype(cache_dtype)
|
||||
.with_device<kDLCUDA>(device_)
|
||||
.with_device<kDLGPU>(device_)
|
||||
.verify(staging_k);
|
||||
if constexpr (!kIsMLA) {
|
||||
TensorMatcher({T, N, D}) //
|
||||
.with_dtype(cache_dtype)
|
||||
.with_device<kDLCUDA>(device_)
|
||||
.with_device<kDLGPU>(device_)
|
||||
.verify(staging_v);
|
||||
}
|
||||
TensorMatcher({-1, N, D}) //
|
||||
.with_dtype(cache_dtype)
|
||||
.with_device<kDLCPU, kDLCUDAHost>()
|
||||
.with_device<kDLCPU, kDLGPUHost>()
|
||||
.verify(k_cache_dst);
|
||||
if constexpr (!kIsMLA) {
|
||||
TensorMatcher({-1, N, D}) //
|
||||
.with_dtype(cache_dtype)
|
||||
.with_device<kDLCPU, kDLCUDAHost>()
|
||||
.with_device<kDLCPU, kDLGPUHost>()
|
||||
.verify(v_cache_dst);
|
||||
}
|
||||
TensorMatcher({N}) //
|
||||
.with_dtype<uint64_t>()
|
||||
.with_device<kDLCUDA>(device_)
|
||||
.with_device<kDLGPU>(device_)
|
||||
.verify(k_ptr_src);
|
||||
if constexpr (!kIsMLA) {
|
||||
TensorMatcher({N}) //
|
||||
.with_dtype<uint64_t>()
|
||||
.with_device<kDLCUDA>(device_)
|
||||
.with_device<kDLGPU>(device_)
|
||||
.verify(v_ptr_src);
|
||||
}
|
||||
TensorMatcher({P}) //
|
||||
.with_dtype<int32_t, int64_t>(indices_dtype)
|
||||
.with_device<kDLCUDA>(device_)
|
||||
.with_device<kDLGPU>(device_)
|
||||
.verify(page_indices_src);
|
||||
TensorMatcher({T}) //
|
||||
.with_dtype<int64_t>(dst_indices_dtype)
|
||||
.with_device<kDLCPU, kDLCUDAHost>()
|
||||
.with_device<kDLCPU, kDLGPUHost>()
|
||||
.verify(dst_indices_cpu);
|
||||
|
||||
RuntimeCheck(page_size > 0, "HiCache staged relayout: page_size must be positive");
|
||||
|
||||
@@ -89,8 +89,10 @@ using fp32x4_t = float4;
|
||||
// DLPack device type for the current platform
|
||||
#ifndef USE_ROCM
|
||||
inline constexpr auto kDLGPU = kDLCUDA;
|
||||
inline constexpr auto kDLGPUHost = kDLCUDAHost;
|
||||
#else
|
||||
inline constexpr auto kDLGPU = kDLROCM;
|
||||
inline constexpr auto kDLGPUHost = kDLROCMHost;
|
||||
#endif
|
||||
|
||||
namespace device {
|
||||
|
||||
@@ -675,7 +675,14 @@ class HiCacheController:
|
||||
return
|
||||
|
||||
op = CacheOperation.merge_ops(self.write_queue)
|
||||
# Page-first write-back JIT kernels can keep destination host indices on CPU.
|
||||
# Kernel write-back keeps host indices on CPU only for page_first AND only
|
||||
# when the staged JIT write-back kernel is available (it stages through
|
||||
# device memory and accepts CPU destination indices). Otherwise we fall back
|
||||
# to the plain transfer kernel, whose CUDA/HIP implementation requires
|
||||
# device-resident destination indices -- so the indices must be moved to the
|
||||
# device first. Without the can_use_write_back_jit check this crashes on
|
||||
# backends where the JIT kernel is unavailable, with
|
||||
# "Destination indices must be a CUDA tensor".
|
||||
if (
|
||||
self.io_backend == "kernel"
|
||||
and self.mem_pool_host.layout == "page_first"
|
||||
|
||||
@@ -94,7 +94,11 @@ class MLATokenToKVPoolHost(HiSparseHostPoolMixin, HostKVCache):
|
||||
device,
|
||||
allocator_type,
|
||||
)
|
||||
self.can_use_jit = _is_cuda and can_use_hicache_jit_kernel(
|
||||
# The JIT HiCache kernels also build with hipcc (ROCm): the PTX-only
|
||||
# helpers in hicache.cuh are guarded by USE_ROCM and the staged
|
||||
# write-back kernel has a ROCm path, so enable them on HIP too. This
|
||||
# keeps the ROCm write-back path consistent with CUDA.
|
||||
self.can_use_jit = (_is_cuda or _is_hip) and can_use_hicache_jit_kernel(
|
||||
element_size=self.kv_cache_dim * self.dtype.itemsize
|
||||
)
|
||||
|
||||
@@ -214,7 +218,11 @@ class MLATokenToKVPoolHost(HiSparseHostPoolMixin, HostKVCache):
|
||||
if self.layout != "page_first" or (_is_npu or _is_xpu or _is_mps):
|
||||
return
|
||||
|
||||
self.can_use_write_back_jit = _is_cuda and can_use_write_back_jit_kernel(
|
||||
# The staged write-back JIT kernel builds with hipcc and has a ROCm
|
||||
# path, so enable it on HIP too (consistent with the CUDA path).
|
||||
self.can_use_write_back_jit = (
|
||||
_is_cuda or _is_hip
|
||||
) and can_use_write_back_jit_kernel(
|
||||
element_size=self.kv_cache_dim * self.dtype.itemsize,
|
||||
)
|
||||
if not self.can_use_write_back_jit:
|
||||
|
||||
@@ -88,7 +88,11 @@ class MHATokenToKVPoolHost(HostKVCache):
|
||||
allocator_type,
|
||||
)
|
||||
self.element_dim = self.device_pool.head_num * self.device_pool.head_dim
|
||||
self.can_use_jit = _is_cuda and can_use_hicache_jit_kernel(
|
||||
# The JIT HiCache kernels also build with hipcc (ROCm): the PTX-only
|
||||
# helpers in hicache.cuh are guarded by USE_ROCM and the staged
|
||||
# write-back kernel has a ROCm path, so enable them on HIP too. This
|
||||
# keeps the ROCm write-back path consistent with CUDA.
|
||||
self.can_use_jit = (_is_cuda or _is_hip) and can_use_hicache_jit_kernel(
|
||||
element_size=self.element_dim * self.dtype.itemsize
|
||||
)
|
||||
|
||||
@@ -170,7 +174,11 @@ class MHATokenToKVPoolHost(HostKVCache):
|
||||
if self.layout != "page_first" or (_is_npu or _is_xpu or _is_mps):
|
||||
return
|
||||
|
||||
self.can_use_write_back_jit = _is_cuda and can_use_write_back_jit_kernel(
|
||||
# The staged write-back JIT kernel builds with hipcc and has a ROCm
|
||||
# path, so enable it on HIP too (consistent with the CUDA path).
|
||||
self.can_use_write_back_jit = (
|
||||
_is_cuda or _is_hip
|
||||
) and can_use_write_back_jit_kernel(
|
||||
element_size=self.element_dim * self.dtype.itemsize,
|
||||
)
|
||||
if not self.can_use_write_back_jit:
|
||||
|
||||
@@ -0,0 +1,259 @@
|
||||
"""Unit tests for the page_first + ``kernel`` JIT HiCache write-back / load path.
|
||||
|
||||
This file specifically exercises the JIT staged write-back and load kernels that
|
||||
accept a CPU-resident destination index and stage through device memory
|
||||
(``staged_write_back.cuh`` / ``hicache.cuh``). Unlike ``test_hicache.py`` (which
|
||||
is registered CUDA-only), this file is also registered for the AMD PR-CI kernel
|
||||
suite so the ROCm/HIP build and execution of those kernels are validated on AMD
|
||||
hardware, not just CUDA.
|
||||
"""
|
||||
|
||||
import sys
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from sglang.jit_kernel.hicache import can_use_write_back_jit_kernel
|
||||
from sglang.srt.mem_cache.memory_pool import MHATokenToKVPool, MLATokenToKVPool
|
||||
from sglang.srt.mem_cache.memory_pool_host import MLATokenToKVPoolHost
|
||||
from sglang.srt.mem_cache.pool_host.common import (
|
||||
ALLOC_MEMORY_FUNCS,
|
||||
alloc_with_pin_memory,
|
||||
)
|
||||
from sglang.srt.mem_cache.pool_host.mha import MHATokenToKVPoolHost
|
||||
from sglang.srt.utils import is_cuda, is_hip, is_npu, is_xpu
|
||||
from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
|
||||
|
||||
register_cuda_ci(est_time=15, stage="base-b-kernel-unit", runner_config="1-gpu-large")
|
||||
register_amd_ci(est_time=30, stage="jit-kernel-unit", runner_config="amd")
|
||||
|
||||
pytestmark = pytest.mark.skipif(
|
||||
not torch.cuda.is_available()
|
||||
or is_npu()
|
||||
or is_xpu()
|
||||
or not (is_cuda() or is_hip()),
|
||||
reason="HiCache JIT write-back tests require CUDA/ROCm.",
|
||||
)
|
||||
|
||||
DEVICE = "cuda"
|
||||
PAGE_SIZE = 1 if is_hip() else 16
|
||||
NUM_LAYERS = 2
|
||||
MHA_ELEMENT_DIMS = [128, 512]
|
||||
MLA_ELEMENT_DIMS = [576]
|
||||
# Include counts around and above the staging capacity so both the single-pass
|
||||
# and the multi-chunk staged relayout branches are exercised.
|
||||
PAGE_COUNTS = [1, 64, 65, 129]
|
||||
|
||||
|
||||
def _token_indices_for_pages(
|
||||
pages: torch.Tensor,
|
||||
device: str = DEVICE,
|
||||
dtype: torch.dtype = torch.int64,
|
||||
) -> torch.Tensor:
|
||||
parts = [
|
||||
torch.arange(
|
||||
int(page) * PAGE_SIZE,
|
||||
(int(page) + 1) * PAGE_SIZE,
|
||||
device=device,
|
||||
dtype=dtype,
|
||||
)
|
||||
for page in pages.tolist()
|
||||
]
|
||||
return torch.cat(parts, dim=0)
|
||||
|
||||
|
||||
def _pinned_host_pool(host_pool_cls, **kwargs):
|
||||
original_alloc = ALLOC_MEMORY_FUNCS[DEVICE]
|
||||
ALLOC_MEMORY_FUNCS[DEVICE] = alloc_with_pin_memory
|
||||
try:
|
||||
return host_pool_cls(
|
||||
host_to_device_ratio=2.0,
|
||||
host_size=0,
|
||||
page_size=PAGE_SIZE,
|
||||
pin_memory=True,
|
||||
device="cpu",
|
||||
**kwargs,
|
||||
)
|
||||
finally:
|
||||
ALLOC_MEMORY_FUNCS[DEVICE] = original_alloc
|
||||
|
||||
|
||||
def _fill_with_offset(tensor: torch.Tensor, offset: int) -> None:
|
||||
data = torch.arange(
|
||||
tensor.numel(), device=tensor.device, dtype=tensor.dtype
|
||||
).view_as(tensor)
|
||||
tensor.copy_(data + offset)
|
||||
|
||||
|
||||
def _assert_pages_equal(host_ref, device_ref, host_pages, device_pages) -> None:
|
||||
for host_page, device_page in zip(host_pages.tolist(), device_pages.tolist()):
|
||||
host_start = host_page * PAGE_SIZE
|
||||
device_start = device_page * PAGE_SIZE
|
||||
assert torch.equal(
|
||||
host_ref[host_start : host_start + PAGE_SIZE].cpu(),
|
||||
device_ref[device_start : device_start + PAGE_SIZE].cpu(),
|
||||
)
|
||||
|
||||
|
||||
def _run_mha(element_dim: int, page_count: int) -> None:
|
||||
pool_size = PAGE_SIZE * (page_count + 8)
|
||||
device_pool = MHATokenToKVPool(
|
||||
size=pool_size,
|
||||
page_size=PAGE_SIZE,
|
||||
head_num=element_dim // 128,
|
||||
head_dim=128,
|
||||
dtype=torch.bfloat16,
|
||||
layer_num=NUM_LAYERS,
|
||||
device=DEVICE,
|
||||
enable_memory_saver=False,
|
||||
)
|
||||
host_pool = _pinned_host_pool(
|
||||
MHATokenToKVPoolHost, device_pool=device_pool, layout="page_first"
|
||||
)
|
||||
assert can_use_write_back_jit_kernel(
|
||||
element_size=element_dim * host_pool.dtype.itemsize,
|
||||
)
|
||||
# page_first + kernel staged write-back JIT path must be enabled.
|
||||
assert host_pool.can_use_write_back_jit
|
||||
|
||||
for layer_id in range(NUM_LAYERS):
|
||||
_fill_with_offset(device_pool.k_buffer[layer_id], layer_id)
|
||||
_fill_with_offset(device_pool.v_buffer[layer_id], layer_id + 100)
|
||||
|
||||
device_pages = torch.arange(2, 2 + page_count, device=DEVICE, dtype=torch.int64)
|
||||
host_pages = torch.arange(page_count, 0, -1, dtype=torch.int64)
|
||||
device_indices = _token_indices_for_pages(device_pages)
|
||||
# host_indices stay on the CPU: this is the case the staged JIT kernel must
|
||||
# accept (kDLCPU / kDLGPUHost destination indices).
|
||||
host_indices = _token_indices_for_pages(host_pages, device="cpu")
|
||||
assert not host_indices.is_cuda
|
||||
|
||||
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.k_data_refs[layer_id],
|
||||
device_pool.k_buffer[layer_id],
|
||||
host_pages,
|
||||
device_pages,
|
||||
)
|
||||
_assert_pages_equal(
|
||||
host_pool.v_data_refs[layer_id],
|
||||
device_pool.v_buffer[layer_id],
|
||||
host_pages,
|
||||
device_pages,
|
||||
)
|
||||
|
||||
# Load path (prefix-cache hit): exercises the hicache.cuh load matchers.
|
||||
if not host_pool.can_use_jit:
|
||||
return
|
||||
for layer_id in range(NUM_LAYERS):
|
||||
device_pool.k_buffer[layer_id].zero_()
|
||||
device_pool.v_buffer[layer_id].zero_()
|
||||
|
||||
load_pages = torch.arange(1, 1 + page_count, device=DEVICE, dtype=torch.int64)
|
||||
load_indices = _token_indices_for_pages(load_pages)
|
||||
host_indices_device = host_indices.to(DEVICE)
|
||||
for layer_id in range(NUM_LAYERS):
|
||||
host_pool.load_to_device_per_layer(
|
||||
device_pool, host_indices_device, load_indices, layer_id, "kernel"
|
||||
)
|
||||
torch.cuda.synchronize()
|
||||
|
||||
for layer_id in range(NUM_LAYERS):
|
||||
_assert_pages_equal(
|
||||
host_pool.k_data_refs[layer_id],
|
||||
device_pool.k_buffer[layer_id],
|
||||
host_pages,
|
||||
load_pages,
|
||||
)
|
||||
_assert_pages_equal(
|
||||
host_pool.v_data_refs[layer_id],
|
||||
device_pool.v_buffer[layer_id],
|
||||
host_pages,
|
||||
load_pages,
|
||||
)
|
||||
|
||||
|
||||
def _run_mla(element_dim: int, page_count: int) -> None:
|
||||
pool_size = PAGE_SIZE * (page_count + 8)
|
||||
device_pool = MLATokenToKVPool(
|
||||
size=pool_size,
|
||||
page_size=PAGE_SIZE,
|
||||
kv_lora_rank=element_dim - 64,
|
||||
qk_rope_head_dim=64,
|
||||
dtype=torch.bfloat16,
|
||||
layer_num=NUM_LAYERS,
|
||||
device=DEVICE,
|
||||
enable_memory_saver=False,
|
||||
)
|
||||
host_pool = _pinned_host_pool(
|
||||
MLATokenToKVPoolHost, device_pool=device_pool, layout="page_first"
|
||||
)
|
||||
assert can_use_write_back_jit_kernel(
|
||||
element_size=element_dim * host_pool.dtype.itemsize,
|
||||
)
|
||||
assert host_pool.can_use_write_back_jit
|
||||
|
||||
for layer_id in range(NUM_LAYERS):
|
||||
_fill_with_offset(device_pool.kv_buffer[layer_id], layer_id)
|
||||
|
||||
device_pages = torch.arange(2, 2 + page_count, device=DEVICE, dtype=torch.int64)
|
||||
host_pages = torch.arange(page_count, 0, -1, dtype=torch.int64)
|
||||
device_indices = _token_indices_for_pages(device_pages)
|
||||
host_indices = _token_indices_for_pages(host_pages, device="cpu")
|
||||
assert not host_indices.is_cuda
|
||||
|
||||
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],
|
||||
host_pages,
|
||||
device_pages,
|
||||
)
|
||||
|
||||
if not host_pool.can_use_jit:
|
||||
return
|
||||
for layer_id in range(NUM_LAYERS):
|
||||
device_pool.kv_buffer[layer_id].zero_()
|
||||
|
||||
load_pages = torch.arange(1, 1 + page_count, device=DEVICE, dtype=torch.int64)
|
||||
load_indices = _token_indices_for_pages(load_pages)
|
||||
host_indices_device = host_indices.to(DEVICE)
|
||||
for layer_id in range(NUM_LAYERS):
|
||||
host_pool.load_to_device_per_layer(
|
||||
device_pool, host_indices_device, load_indices, layer_id, "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],
|
||||
host_pages,
|
||||
load_pages,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("element_dim", MHA_ELEMENT_DIMS)
|
||||
@pytest.mark.parametrize("page_count", PAGE_COUNTS)
|
||||
def test_page_first_staged_write_back_mha(element_dim: int, page_count: int) -> None:
|
||||
_run_mha(element_dim, page_count)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("element_dim", MLA_ELEMENT_DIMS)
|
||||
@pytest.mark.parametrize("page_count", PAGE_COUNTS)
|
||||
def test_page_first_staged_write_back_mla(element_dim: int, page_count: int) -> None:
|
||||
_run_mla(element_dim, page_count)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(pytest.main([__file__, "-v", "-s"]))
|
||||
Reference in New Issue
Block a user