kernels/cuda: paged-attention kernel, dispatch, pinned host memory
CUDA layer for the paged-KV + swap work: - csrc: new paged_attention.cu plus updates across attention/gemm/norm/ activation/embedding/reduce kernels and common.cuh. - xserv-kernels: new dispatch module and kernel-binding updates. - xserv-cuda: cudaMallocHost/FreeHost bindings + PinnedBuffer (host swap pool backing) and offset-aware D2H/H2D copies used to move KV blocks between the GPU pool and pinned host memory. Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
This commit is contained in:
@@ -116,6 +116,56 @@ impl GpuBuffer {
|
||||
})
|
||||
}
|
||||
|
||||
/// Async copy `count` bytes from `src` at `src_offset` to `self` at `dst_offset` on `stream`.
|
||||
pub fn copy_from_device_at_async(&mut self, src: &GpuBuffer, src_offset: usize, dst_offset: usize, count: usize, stream: &CudaStream) -> Result<()> {
|
||||
assert!(src_offset + count <= src.len);
|
||||
assert!(dst_offset + count <= self.len);
|
||||
error::check(unsafe {
|
||||
ffi::cudaMemcpyAsync(
|
||||
self.ptr.add(dst_offset),
|
||||
src.ptr.add(src_offset),
|
||||
count,
|
||||
ffi::CUDA_MEMCPY_D2D,
|
||||
stream.as_raw(),
|
||||
)
|
||||
})
|
||||
}
|
||||
|
||||
/// Copy `count` bytes from this GPU buffer at `src_offset` to a host slice (D2H).
|
||||
pub fn copy_to_host_at(&self, dst: &mut [u8], src_offset: usize, count: usize) -> Result<()> {
|
||||
assert!(src_offset + count <= self.len, "src range out of bounds");
|
||||
assert!(count <= dst.len(), "host dst too small");
|
||||
error::check(unsafe {
|
||||
ffi::cudaMemcpy(
|
||||
dst.as_mut_ptr(),
|
||||
self.ptr.add(src_offset),
|
||||
count,
|
||||
ffi::CUDA_MEMCPY_D2H,
|
||||
)
|
||||
})
|
||||
}
|
||||
|
||||
/// Copy `count` bytes from a host slice to this GPU buffer at `dst_offset` (H2D).
|
||||
pub fn copy_from_host_at(&mut self, src: &[u8], dst_offset: usize, count: usize) -> Result<()> {
|
||||
assert!(dst_offset + count <= self.len, "dst range out of bounds");
|
||||
assert!(count <= src.len(), "host src too small");
|
||||
error::check(unsafe {
|
||||
ffi::cudaMemcpy(
|
||||
self.ptr.add(dst_offset),
|
||||
src.as_ptr(),
|
||||
count,
|
||||
ffi::CUDA_MEMCPY_H2D,
|
||||
)
|
||||
})
|
||||
}
|
||||
|
||||
/// Async zero fill on stream.
|
||||
pub fn zero_async(&mut self, stream: &CudaStream) -> Result<()> {
|
||||
error::check(unsafe {
|
||||
ffi::cudaMemsetAsync(self.ptr, 0, self.len, stream.as_raw())
|
||||
})
|
||||
}
|
||||
|
||||
/// Consume the buffer without freeing GPU memory. Returns the raw pointer and length.
|
||||
/// Caller is responsible for eventually calling cudaFree.
|
||||
pub fn into_raw(self) -> (*mut u8, usize) {
|
||||
|
||||
Reference in New Issue
Block a user