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:
2026-05-28 19:58:36 +08:00
parent 3f1c3d429a
commit 4c3f914459
27 changed files with 581 additions and 32 deletions

View File

@@ -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) {