cuda: infrastructure for whole-step CUDA graph capture
- 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>
This commit is contained in:
@@ -31,3 +31,39 @@ impl Drop for CudaStream {
|
||||
|
||||
// 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));
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user