Gahow Wang
4088f49b7d
- Thread-local launch stream (xserv_cuda::stream): every kernel wrapper, cublasSetStream, and NCCL collective now launches on current_stream_raw() — the legacy null stream by default (behavior unchanged), or the capture stream installed via push_stream during graph capture. Capture is impossible on the legacy stream. - Allocator retain mode: blocks freed inside a retain window are quarantined (RetainedBlocks) instead of pooled, so an instantiated graph keeps exclusive ownership of every intermediate buffer it references across replays. - Capture mode GLOBAL -> THREAD_LOCAL: concurrent TP rank threads must not poison each other's captures with their own cudaMallocs. - embedding_device_ids / rope_inplace_device_pos: variants reading token ids / positions from persistent device buffers, replacing the per-call host upload that a captured region cannot contain. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
70 lines
1.9 KiB
Rust
70 lines
1.9 KiB
Rust
use crate::error::{self, Result};
|
|
use crate::ffi;
|
|
|
|
pub struct CudaStream {
|
|
raw: ffi::CudaStream,
|
|
}
|
|
|
|
impl CudaStream {
|
|
pub fn new() -> Result<Self> {
|
|
let mut raw = std::ptr::null_mut();
|
|
error::check(unsafe { ffi::cudaStreamCreate(&mut raw) })?;
|
|
Ok(Self { raw })
|
|
}
|
|
|
|
pub fn synchronize(&self) -> Result<()> {
|
|
error::check(unsafe { ffi::cudaStreamSynchronize(self.raw) })
|
|
}
|
|
|
|
pub fn as_raw(&self) -> ffi::CudaStream {
|
|
self.raw
|
|
}
|
|
}
|
|
|
|
impl Drop for CudaStream {
|
|
fn drop(&mut self) {
|
|
if !self.raw.is_null() {
|
|
unsafe { ffi::cudaStreamDestroy(self.raw) };
|
|
}
|
|
}
|
|
}
|
|
|
|
// Can move across threads, but not shared without synchronization
|
|
unsafe impl Send for CudaStream {}
|
|
|
|
// --- Thread-local launch stream -------------------------------------------
|
|
//
|
|
// Every kernel wrapper in xserv-kernels launches on `current_stream_raw()`,
|
|
// which defaults to the legacy null stream (the historical behavior). CUDA
|
|
// graph capture requires work to be issued on an explicit stream, so capture
|
|
// code installs its stream here for the duration of the captured region via
|
|
// `push_stream` / `StreamGuard`.
|
|
|
|
use std::cell::Cell;
|
|
|
|
thread_local! {
|
|
static CURRENT_STREAM: Cell<ffi::CudaStream> = const { Cell::new(std::ptr::null_mut()) };
|
|
}
|
|
|
|
/// The stream kernel launches on this thread should use (null = legacy default).
|
|
pub fn current_stream_raw() -> ffi::CudaStream {
|
|
CURRENT_STREAM.with(|c| c.get())
|
|
}
|
|
|
|
/// RAII guard that installs a launch stream for the current thread and
|
|
/// restores the previous one on drop.
|
|
pub struct StreamGuard {
|
|
prev: ffi::CudaStream,
|
|
}
|
|
|
|
pub fn push_stream(stream: &CudaStream) -> StreamGuard {
|
|
let prev = CURRENT_STREAM.with(|c| c.replace(stream.as_raw()));
|
|
StreamGuard { prev }
|
|
}
|
|
|
|
impl Drop for StreamGuard {
|
|
fn drop(&mut self) {
|
|
CURRENT_STREAM.with(|c| c.set(self.prev));
|
|
}
|
|
}
|