From dd0498ece98719a0ca6f738a32e3cee5bb17ea2a Mon Sep 17 00:00:00 2001 From: isolo Date: Sun, 5 Jul 2026 00:32:46 +0000 Subject: [PATCH] perf(spark-runtime): opt-in NVTX kernel-attribution + phase-timing instrumentation MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Diagnostic-only (ATLAS_NVTX_DIAG=1), dlopen'd lazily so normal builds/runs carry no link-time or runtime dependency on libnvToolsExt. Split out of #218 (DFlash speculative decoding) — purely additive instrumentation, no behavior change to any existing code path. Co-Authored-By: Claude Sonnet 5 --- Cargo.lock | 1 + crates/spark-runtime/Cargo.toml | 4 + crates/spark-runtime/src/kernel_args.rs | 17 ++++ crates/spark-runtime/src/lib.rs | 1 + crates/spark-runtime/src/nvtx_diag.rs | 105 ++++++++++++++++++++++++ 5 files changed, 128 insertions(+) create mode 100644 crates/spark-runtime/src/nvtx_diag.rs diff --git a/Cargo.lock b/Cargo.lock index 3c7bc3b73..be651d076 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2098,6 +2098,7 @@ dependencies = [ "dispatch2", "half", "libc", + "libloading 0.8.9", "memmap2", "objc2", "objc2-foundation", diff --git a/crates/spark-runtime/Cargo.toml b/crates/spark-runtime/Cargo.toml index b667388c5..bc14158f8 100644 --- a/crates/spark-runtime/Cargo.toml +++ b/crates/spark-runtime/Cargo.toml @@ -109,6 +109,10 @@ tokio = { version = "1", features = ["rt", "sync", "time", "parking_lot"] } rand = "0.8" half = "2" parking_lot = { workspace = true } +# Diagnostic-only NVTX range instrumentation (ATLAS_NVTX_DIAG=1): dlopen'd +# lazily so normal builds/runs have no link-time or runtime dependency on +# libnvToolsExt. +libloading = "0.8" [target.'cfg(unix)'.dependencies] libc = "0.2" diff --git a/crates/spark-runtime/src/kernel_args.rs b/crates/spark-runtime/src/kernel_args.rs index adedc5ff1..e827ab335 100644 --- a/crates/spark-runtime/src/kernel_args.rs +++ b/crates/spark-runtime/src/kernel_args.rs @@ -31,6 +31,7 @@ use anyhow::Result; use crate::gpu::{DevicePtr, GpuBackend, KernelArg, KernelHandle}; +use crate::nvtx_diag::NvtxRange; /// Per-arg metadata: which slot of `storage` it lives at, and how /// many native bytes it occupies. Buffer args set `is_buffer = true` @@ -62,6 +63,9 @@ pub struct KernelLaunch<'a> { /// Parallel array recording per-arg kind so `launch()` can build /// a typed `KernelArg` slice. kinds: Vec, + /// Stage-2b diagnostic-only NVTX call-site label (see `nvtx_diag`). + /// `None` unless a call site opts in via `.nvtx_label(...)`. + nvtx_label: Option, } impl<'a> KernelLaunch<'a> { @@ -74,9 +78,20 @@ impl<'a> KernelLaunch<'a> { shared_mem: 0, storage: Vec::with_capacity(16), kinds: Vec::with_capacity(16), + nvtx_label: None, } } + /// Stage-2b kernel-attribution diagnostic: tag this specific call site + /// with an NVTX range label so nsys's NVTX+CUDA-API correlation can + /// distinguish it from other call sites that happen to dispatch the + /// same underlying kernel (e.g. FFN down_proj vs attention out_proj, + /// both `w4a16_gemm_t_m128`). No-op unless `ATLAS_NVTX_DIAG=1`. + pub fn nvtx_label(mut self, label: impl Into) -> Self { + self.nvtx_label = Some(label.into()); + self + } + pub fn grid(mut self, grid: [u32; 3]) -> Self { self.grid = grid; self @@ -173,6 +188,7 @@ impl<'a> KernelLaunch<'a> { args.push(KernelArg::Bytes(bytes)); } } + let _nvtx = self.nvtx_label.as_deref().map(NvtxRange::new); let r = self.gpu.launch_typed( self.kernel, self.grid, @@ -181,6 +197,7 @@ impl<'a> KernelLaunch<'a> { stream, &args, ); + drop(_nvtx); // ATLAS_DEBUG_SYNC_KERNELS (PCND, default-off): synchronize after // each launch so an async CUDA fault surfaces AT the culprit launch // (with grid/block) instead of at a later, unrelated sync point. diff --git a/crates/spark-runtime/src/lib.rs b/crates/spark-runtime/src/lib.rs index 6716c512a..70b8bbd46 100644 --- a/crates/spark-runtime/src/lib.rs +++ b/crates/spark-runtime/src/lib.rs @@ -33,6 +33,7 @@ pub mod kv_dequant; pub mod kv_spill; #[cfg(feature = "metal")] pub mod metal_backend; +pub mod nvtx_diag; pub mod prefix_cache; pub mod radix_tree; pub mod sampler; diff --git a/crates/spark-runtime/src/nvtx_diag.rs b/crates/spark-runtime/src/nvtx_diag.rs new file mode 100644 index 000000000..cf432eb26 --- /dev/null +++ b/crates/spark-runtime/src/nvtx_diag.rs @@ -0,0 +1,105 @@ +// SPDX-License-Identifier: AGPL-3.0-only + +//! Diagnostic-only NVTX range instrumentation (Stage 2b kernel-attribution +//! profiling). Gated behind `ATLAS_NVTX_DIAG=1` (PCND, default-off) so it has +//! zero cost/dependency in normal builds and runs. +//! +//! `libnvToolsExt` is dlopen'd lazily (via `libloading`, not linked at build +//! time) so a missing library on non-profiling machines just disables +//! ranges instead of failing the build or the run. Call sites push a label +//! immediately around the one GPU-launch/API call they want attributed; +//! nsys's NVTX+CUDA-API correlation (via `--trace=cuda,nvtx,osrt` and the +//! `nvtx_kern_sum` report) then attributes the *actually launched* kernel +//! time to that label even though the kernel executes asynchronously after +//! the host-side push/pop window closes. + +use libloading::{Library, Symbol}; +use std::ffi::CString; +use std::os::raw::{c_char, c_int}; +use std::sync::OnceLock; + +type PushFn = unsafe extern "C" fn(*const c_char) -> c_int; +type PopFn = unsafe extern "C" fn() -> c_int; + +struct NvtxFns { + push: PushFn, + pop: PopFn, + // Keep the library mapped for the process lifetime; never unloaded. + _lib: Library, +} + +// SAFETY: raw fn pointers into a dlopen'd shared library that is never +// unloaded (owned for 'static via the OnceLock); the underlying NVTX ABI +// is thread-safe (it's the same library the C++ NVTX3 header calls into). +unsafe impl Send for NvtxFns {} +unsafe impl Sync for NvtxFns {} + +fn enabled() -> bool { + static FLAG: OnceLock = OnceLock::new(); + *FLAG.get_or_init(|| std::env::var("ATLAS_NVTX_DIAG").as_deref() == Ok("1")) +} + +fn nvtx() -> Option<&'static NvtxFns> { + static LIB: OnceLock> = OnceLock::new(); + LIB.get_or_init(|| { + if !enabled() { + return None; + } + // Resolved via the dynamic linker search path (LD_LIBRARY_PATH) — + // the profiling harness is responsible for pointing that at a + // directory containing one of these sonames. + for candidate in ["libnvToolsExt.so.1", "libnvToolsExt.so"] { + let lib = unsafe { Library::new(candidate) }; + let Ok(lib) = lib else { continue }; + let push: Symbol = match unsafe { lib.get(b"nvtxRangePushA\0") } { + Ok(s) => s, + Err(_) => continue, + }; + let pop: Symbol = match unsafe { lib.get(b"nvtxRangePop\0") } { + Ok(s) => s, + Err(_) => continue, + }; + let push = *push; + let pop = *pop; + tracing::info!("ATLAS_NVTX_DIAG: loaded {candidate}"); + return Some(NvtxFns { + push, + pop, + _lib: lib, + }); + } + tracing::warn!("ATLAS_NVTX_DIAG=1 set but no libnvToolsExt found — NVTX ranges disabled"); + None + }) + .as_ref() +} + +/// RAII NVTX range: push on construction, pop on drop. No-op unless +/// `ATLAS_NVTX_DIAG=1` and libnvToolsExt was found. +pub struct NvtxRange { + active: bool, +} + +impl NvtxRange { + #[inline] + pub fn new(label: &str) -> Self { + let Some(n) = nvtx() else { + return Self { active: false }; + }; + if let Ok(c) = CString::new(label) { + unsafe { (n.push)(c.as_ptr()) }; + } + Self { active: true } + } +} + +impl Drop for NvtxRange { + #[inline] + fn drop(&mut self) { + if self.active + && let Some(n) = nvtx() + { + unsafe { (n.pop)() }; + } + } +}