phase 15: Tensor::empty + CUDA Graph infra — 50.3 tok/s (140% of HF, 45% roofline)

Two optimizations:

1. Tensor::empty() — skip cudaMemset for output tensors
   All kernel wrappers that fully overwrite their output now use
   Tensor::empty() instead of Tensor::zeros(). Eliminates ~756
   cudaMemset calls per decode step (21 per layer × 36 layers).
   Improvement: 46.6 → 50.3 tok/s (+8%).

2. CUDA Graph infrastructure (for future use)
   Added FFI bindings (cudaStreamBeginCapture, cudaGraphInstantiate,
   cudaGraphLaunch) and RAII CudaGraph wrapper. Not yet used in the
   forward pass due to variable kv_len, but provides foundation for
   future graph-based decode optimization.

Ablation (dash5, RTX 5090, Qwen3-8B BF16, serial decode):

| Optimization | tok/s | vs HF | Roofline |
|-------------|-------|-------|----------|
| Phase 14 baseline | 12.9 | 36% | 12% |
| + Fused kernels | 13.2 | 37% | 12% |
| + Batched decode | 13.2 (serial) | 37% | 12% |
| + Custom GEMV | 46.6 | 130% | 42% |
| + Tensor::empty | 50.3 | 140% | 45% |

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
2026-05-22 23:57:34 +08:00
parent e207523e21
commit d5532ef209
13 changed files with 170 additions and 20 deletions

View File

@@ -0,0 +1,98 @@
//! CUDA Graphs: capture a sequence of kernel launches and replay them with
//! near-zero host-side overhead (~3-5 us per launch eliminated).
//!
//! Usage:
//! ```ignore
//! let stream = CudaStream::new()?;
//! let mut graph = CudaGraph::new();
//!
//! // First call: capture
//! graph.begin_capture(&stream)?;
//! // ... launch kernels on `stream` ...
//! graph.end_capture(&stream)?;
//!
//! // Subsequent calls: replay
//! graph.launch(&stream)?;
//! ```
//!
//! Requirements for captured kernels:
//! - All tensor shapes must be identical between capture and replay.
//! - No host-side branching during the captured section.
//! - Memory addresses used during capture must remain valid during replay.
use crate::error::{self, Result};
use crate::ffi;
use crate::stream::CudaStream;
/// RAII wrapper around a captured CUDA graph and its executable instance.
pub struct CudaGraph {
graph: ffi::CudaGraph,
exec: ffi::CudaGraphExec,
}
impl CudaGraph {
/// Create an empty graph handle (not yet captured).
pub fn new() -> Self {
Self {
graph: std::ptr::null_mut(),
exec: std::ptr::null_mut(),
}
}
/// Returns true if a graph has been captured and instantiated.
pub fn is_ready(&self) -> bool {
!self.exec.is_null()
}
/// Begin capturing kernel launches on `stream`.
/// All subsequent kernel launches on this stream are recorded into the
/// graph instead of being executed.
pub fn begin_capture(&mut self, stream: &CudaStream) -> Result<()> {
// If we have an old graph, destroy it first
self.destroy_inner();
error::check(unsafe {
ffi::cudaStreamBeginCapture(
stream.as_raw(),
ffi::CUDA_STREAM_CAPTURE_MODE_GLOBAL,
)
})
}
/// End capture and instantiate the executable graph.
pub fn end_capture(&mut self, stream: &CudaStream) -> Result<()> {
error::check(unsafe {
ffi::cudaStreamEndCapture(stream.as_raw(), &mut self.graph)
})?;
error::check(unsafe {
ffi::cudaGraphInstantiate(&mut self.exec, self.graph, 0)
})
}
/// Replay the captured graph on `stream`.
/// Panics if no graph has been captured yet.
pub fn launch(&self, stream: &CudaStream) -> Result<()> {
assert!(self.is_ready(), "CudaGraph::launch called before capture");
error::check(unsafe {
ffi::cudaGraphLaunch(self.exec, stream.as_raw())
})
}
fn destroy_inner(&mut self) {
if !self.exec.is_null() {
unsafe { ffi::cudaGraphExecDestroy(self.exec) };
self.exec = std::ptr::null_mut();
}
if !self.graph.is_null() {
unsafe { ffi::cudaGraphDestroy(self.graph) };
self.graph = std::ptr::null_mut();
}
}
}
impl Drop for CudaGraph {
fn drop(&mut self) {
self.destroy_inner();
}
}
unsafe impl Send for CudaGraph {}