spark_runtime/cuda_backend/
gpu_impl.rs

1// SPDX-License-Identifier: AGPL-3.0-only
2
3//! `impl GpuBackend for AtlasCudaBackend` — production CUDA backend trait body.
4//!
5//! ## Safety contract for the `unsafe { cu*(...) }` calls below
6//!
7//! Every unsafe block in this file wraps a single CUDA Driver API call.
8//! The invariants the driver requires are uniform:
9//!
10//! - **Context bound**: a CUDA primary context for the device is current
11//!   on the calling thread. `AtlasCudaBackend::new` binds it once via
12//!   `cuCtxSetCurrent`, and we never run on a thread that hasn't been
13//!   bound.
14//! - **Pointer provenance**: every `DevicePtr` came from a prior
15//!   successful `cuMemAlloc_v2` / `cuMemAllocHost_v2` /
16//!   `cuMemAllocManaged` and has not yet been freed. `DevicePtr(0)` is
17//!   treated as "not allocated" by callers.
18//! - **Sizes in bytes**: every `bytes: usize` argument is the exact
19//!   byte count of the allocation (callers compute it from typed
20//!   sizes); the driver does no bounds-checking.
21//! - **Stream / event lifetimes**: handles are owned by `Self` and
22//!   freed in `Drop` after `cuStreamSynchronize`, so they outlive every
23//!   in-flight launch that captured them.
24//! - **`extern "C"` ABI**: matches the cudarc-generated bindings used
25//!   in `super::*` imports; see `cudarc` for the full ABI surface.
26//!
27//! Per-site `// SAFETY:` comments are omitted because the contract is
28//! identical for every call. Anything that *deviates* from this
29//! contract gets a per-site `// SAFETY:` comment explaining the
30//! exception.
31
32use std::ffi::c_void;
33use std::sync::OnceLock;
34
35use anyhow::{Result, bail};
36use atlas_core::registry::{RawCudaFunc, cuda_error_text};
37use cudarc::driver::LaunchConfig;
38
39use super::{
40    AtlasCudaBackend, cuMemAlloc_v2, cuMemAllocManaged, cuMemFree_v2, cuMemGetInfo_v2,
41    cuMemcpyDtoDAsync_v2, cuMemcpyDtoHAsync_v2, cuMemcpyHtoDAsync_v2, cuStreamSynchronize,
42};
43use crate::gpu::{DevicePtr, GpuBackend, GraphHandle, KernelHandle};
44
45/// D2H call counter + one-shot caller identification
46/// (`ATLAS_D2H_TRACE=<N>`: log a backtrace on the Nth call, and the running
47/// count on every 10000th).
48///
49/// Every `copy_d2h*` below pairs its async copy with a `cuStreamSynchronize`,
50/// so each call BLOCKS the host until the GPU drains. An nsys trace of a 1K
51/// Laguna prefill counted 32,343 D2H + 32,533 syncs inside the prefill span,
52/// accounting for 212.8 ms of 306 ms of GPU starvation (58% idle). This exists
53/// to name whoever is issuing them.
54static D2H_COUNT: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
55
56fn d2h_trace_tick() {
57    use std::sync::atomic::Ordering;
58    let n = D2H_COUNT.fetch_add(1, Ordering::Relaxed) + 1;
59    let Ok(target) = std::env::var("ATLAS_D2H_TRACE") else {
60        return;
61    };
62    let target: u64 = target.parse().unwrap_or(0);
63    if target != 0 && n == target {
64        tracing::warn!(
65            "ATLAS_D2H_TRACE: call #{n} backtrace:\n{}",
66            std::backtrace::Backtrace::force_capture()
67        );
68    }
69    if n.is_multiple_of(10_000) {
70        tracing::warn!("ATLAS_D2H_TRACE: {n} D2H copies so far (each forces a stream sync)");
71    }
72}
73
74/// Enqueue an H2D copy on `stream` and return without waiting. Shared by both
75/// async H2D entry points so the two differ ONLY in the ordering they add
76/// afterwards, never in the copy itself.
77fn h2d_enqueue(src: &[u8], dst: DevicePtr, stream: u64) -> Result<()> {
78    let status =
79        unsafe { cuMemcpyHtoDAsync_v2(dst.0, src.as_ptr() as *const c_void, src.len(), stream) };
80    if status != 0 {
81        bail!("cuMemcpyHtoDAsync_v2 failed: status {status}");
82    }
83    Ok(())
84}
85
86/// Say once, loudly, that a page-locked buffer reached the transient H2D path.
87///
88/// This is the tripwire the whole `pinned_hosts` registry exists to arm. It is a
89/// warning and not a `bail!` because the copy is still CORRECT — the sync above
90/// restores the guarantee — but it is a real, silent latency regression, and the
91/// call site almost certainly wants `copy_h2d_async_retained` instead.
92fn warn_pinned_transient_source() {
93    static ONCE: std::sync::Once = std::sync::Once::new();
94    ONCE.call_once(|| {
95        tracing::warn!(
96            "copy_h2d_async was handed a PAGE-LOCKED source. That copy is genuinely \
97             asynchronous, so the promise that the caller may drop the buffer on return \
98             is now being paid for with a cuStreamSynchronize on every such call. If the \
99             source outlives the next sync, switch the call site to \
100             copy_h2d_async_retained; if it does not, this sync is what keeps it from \
101             being a use-after-free."
102        );
103    });
104}
105
106impl GpuBackend for AtlasCudaBackend {
107    #[track_caller]
108    fn alloc(&self, bytes: usize) -> Result<DevicePtr> {
109        let site = std::panic::Location::caller();
110        let mut dptr: u64 = 0;
111        // A55 RED ZONE (`ATLAS_REDZONE=<bytes>`, default 0 = off). Over-allocate by `pad`
112        // and hand the caller the base, so the buffer it sees is unchanged and correctly
113        // aligned (cuMemAlloc is 256-byte aligned; padding the TAIL keeps that). The pad is
114        // poisoned at birth and read back by `scan_redzones`.
115        //
116        // This is the detector compute-sanitizer could not be: Atlas suballocates from pools,
117        // so an overrun that stays inside a pooled block is invisible to memcheck but lands
118        // squarely in a red zone here.
119        let seq = super::ALLOC_SEQ.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
120        let pad = if seq >= super::redzone_min_idx() {
121            super::redzone_bytes()
122        } else {
123            0
124        };
125        let status = unsafe { cuMemAlloc_v2(&mut dptr, bytes + pad) };
126        if status != 0 {
127            let mut free: usize = 0;
128            let mut total: usize = 0;
129            unsafe { cuMemGetInfo_v2(&mut free, &mut total) };
130            bail!(
131                "cuMemAlloc_v2 failed: status {status}, requested {bytes} bytes \
132                 (device reports {:.1} MB free / {:.1} GB total)",
133                free as f64 / (1024.0 * 1024.0),
134                total as f64 / (1024.0 * 1024.0 * 1024.0),
135            );
136        }
137        if pad > 0 {
138            let st = unsafe {
139                super::cuMemsetD8Async(dptr + bytes as u64, super::redzone_fill(), pad, 0)
140            };
141            if st != 0 {
142                bail!("ATLAS_REDZONE: poisoning the guard band failed: status {st}");
143            }
144            self.record_redzone(dptr, bytes, pad, seq);
145            // One line per guarded allocation, in creation order. The bisection reports an
146            // INDEX; this is what turns that index into a buffer you can name by its size,
147            // and `ATLAS_REDZONE_TRACE_IDX` adds the call site for the one that matters.
148            tracing::info!("redzone: alloc#{seq} bytes={bytes} ptr={dptr:#x}");
149            if super::redzone_trace_idx() == Some(seq) {
150                tracing::error!(
151                    "redzone: alloc#{seq} bytes={bytes} backtrace:\n{}",
152                    std::backtrace::Backtrace::force_capture()
153                );
154            }
155        }
156        self.record_alloc(DevicePtr(dptr), bytes, site);
157        // Large-allocation tracing for memory attribution (GB10 unified
158        // memory: every cuMemAlloc consumes host RAM, and a runtime alloc
159        // outside the util pledge is how the box ends up in swap). Debug
160        // level so production INFO stays quiet; RUST_LOG=spark_runtime=debug
161        // turns the trail on.
162        if bytes >= 32 * 1024 * 1024 {
163            tracing::debug!(
164                "alloc {:.1} MB (device ptr {dptr:#x})",
165                bytes as f64 / (1024.0 * 1024.0)
166            );
167        }
168        Ok(DevicePtr(dptr))
169    }
170
171    fn scan_redzones(&self) -> Result<usize> {
172        if super::redzone_bytes() == 0 {
173            return Ok(0);
174        }
175        AtlasCudaBackend::scan_redzones(self)
176    }
177
178    fn poison_redzones(&self, lo: usize, hi: usize) -> Result<()> {
179        if super::redzone_bytes() == 0 {
180            return Ok(());
181        }
182        AtlasCudaBackend::poison_redzones(self, lo, hi)
183    }
184
185    #[track_caller]
186    fn alloc_managed(&self, bytes: usize) -> Result<DevicePtr> {
187        let site = std::panic::Location::caller();
188        let mut dptr: u64 = 0;
189        const CU_MEM_ATTACH_GLOBAL: u32 = 0x1;
190        let status = unsafe { cuMemAllocManaged(&mut dptr, bytes, CU_MEM_ATTACH_GLOBAL) };
191        if status != 0 {
192            bail!(
193                "cuMemAllocManaged failed: status {status}, requested {bytes} bytes. \
194                 Check system swap space: swapon --show"
195            );
196        }
197        self.record_alloc(DevicePtr(dptr), bytes, site);
198        Ok(DevicePtr(dptr))
199    }
200
201    fn free(&self, ptr: DevicePtr) -> Result<()> {
202        if ptr.is_null() {
203            return Ok(());
204        }
205        // Off the ledger BEFORE the free: an entry that survives a successful
206        // free would be double-freed at teardown.
207        self.forget_alloc(ptr);
208        if super::redzone_bytes() > 0 {
209            self.forget_redzone(ptr.0);
210        }
211        let status = unsafe { cuMemFree_v2(ptr.0) };
212        // A context that is already being destroyed reports every free as
213        // failing, and at process exit that is the normal case, not an error:
214        // the driver has reclaimed the allocation by definition. Two other
215        // free paths in this crate already consult `is_teardown_noop`; this one
216        // did not, so wiring `Model::teardown` into shutdown turned a benign
217        // status 4 into `ERROR model teardown reported a failure` on every
218        // clean exit — the exact species of false alarm this work set out to
219        // remove.
220        if status != 0 && !atlas_core::registry::is_teardown_noop(status) {
221            bail!("cuMemFree_v2 failed: status {status}, ptr {ptr}");
222        }
223        Ok(())
224    }
225
226    fn live_bytes(&self) -> Option<usize> {
227        Some(AtlasCudaBackend::live_bytes(self))
228    }
229
230    fn alloc_report(&self, top_n: usize, min_mb: usize) -> Option<String> {
231        Some(AtlasCudaBackend::alloc_report(self, top_n, min_mb))
232    }
233
234    fn sweep_unreleased(&self) -> usize {
235        AtlasCudaBackend::sweep_unreleased(self)
236    }
237
238    fn copy_h2d(&self, src: &[u8], dst: DevicePtr) -> Result<()> {
239        AtlasCudaBackend::copy_h2d_impl(self, src, dst)
240    }
241
242    fn copy_d2h(&self, src: DevicePtr, dst: &mut [u8]) -> Result<()> {
243        d2h_trace_tick();
244        AtlasCudaBackend::copy_d2h_impl(self, src, dst)
245    }
246
247    fn copy_d2h_on_stream(&self, src: DevicePtr, dst: &mut [u8], stream: u64) -> Result<()> {
248        d2h_trace_tick();
249        AtlasCudaBackend::copy_d2h_on_stream_impl(self, src, dst, stream)
250    }
251
252    fn copy_d2h_async(&self, src: DevicePtr, dst: &mut [u8], stream: u64) -> Result<()> {
253        // Deliberately NO cuStreamSynchronize — that is the entire point.
254        // `copy_d2h`/`copy_d2h_on_stream` drain the stream inside every call,
255        // so a multi-chunk gather pays one full drain per chunk (the SSM spill's
256        // 60 chunks × 66 MB measured ~400 ms = ~165 MB/s, vs ~28 ms for the
257        // async H2D scatter of the same bytes). The caller MUST issue exactly
258        // one `synchronize(stream)` before touching `dst`.
259        d2h_trace_tick();
260        let status = unsafe {
261            cuMemcpyDtoHAsync_v2(dst.as_mut_ptr() as *mut c_void, src.0, dst.len(), stream)
262        };
263        if status != 0 {
264            bail!("cuMemcpyDtoHAsync_v2 (async) failed: status {status}");
265        }
266        Ok(())
267    }
268
269    fn copy_d2d(&self, src: DevicePtr, dst: DevicePtr, bytes: usize) -> Result<()> {
270        if crate::launch_trace::on() {
271            crate::launch_trace::record(crate::launch_trace::Entry {
272                kind: "d2d",
273                func: 0,
274                grid: [0, 0, 0],
275                block: [0, 0, 0],
276                smem: 0,
277                args: vec![src.0, dst.0, bytes as u64],
278            });
279        }
280        AtlasCudaBackend::copy_d2d_impl(self, src, dst, bytes)
281    }
282
283    fn launch(
284        &self,
285        func: KernelHandle,
286        grid: [u32; 3],
287        block: [u32; 3],
288        shared_mem: u32,
289        stream: u64,
290        params: &mut [*mut c_void],
291    ) -> Result<()> {
292        let raw_func = RawCudaFunc(func.0 as *mut c_void);
293        let cfg = LaunchConfig {
294            grid_dim: (grid[0], grid[1], grid[2]),
295            block_dim: (block[0], block[1], block[2]),
296            shared_mem_bytes: shared_mem,
297        };
298        let registry = self.registry();
299        unsafe { registry.launch_on_stream(raw_func, cfg, stream, params) }.map_err(|e| {
300            // A launch failure may have destroyed the CUDA context. Probe and
301            // latch before the error is flattened into a string and bubbled —
302            // see `fault_probe`. This does not change control flow: the caller
303            // still receives its error either way.
304            super::fault_probe::note_failure("kernel launch", &e.to_string());
305            anyhow::anyhow!("Kernel launch failed: {e}")
306        })
307    }
308
309    fn stream_is_capturing(&self, stream: u64) -> bool {
310        // SCALE's libcuda does not export cuStreamIsCapturing; report
311        // not-capturing there (gfx1151 telemetry taps then sample eagerly —
312        // acceptable for a default-off measurement knob).
313        #[cfg(atlas_scale)]
314        {
315            let _ = stream;
316            false
317        }
318        #[cfg(not(atlas_scale))]
319        {
320            let mut status: u32 = 0;
321            // CU_STREAM_CAPTURE_STATUS_NONE = 0; treat query failure as
322            // capturing (conservative: the tap skips its sample).
323            let rc = unsafe { super::cuStreamIsCapturing(stream, &mut status) };
324            rc != 0 || status != 0
325        }
326    }
327
328    fn synchronize(&self, stream: u64) -> Result<()> {
329        let status = unsafe { cuStreamSynchronize(stream) };
330        if status != 0 {
331            bail!("cuStreamSynchronize failed: {}", cuda_error_text(status));
332        }
333        Ok(())
334    }
335
336    fn default_stream(&self) -> u64 {
337        self.default_stream
338    }
339
340    fn op_cache(&self) -> &crate::op_cache::OpCache {
341        &self.op_cache
342    }
343
344    fn debug_sync_kernels(&self) -> bool {
345        AtlasCudaBackend::debug_sync_kernels(self)
346    }
347
348    fn kernel_registry(&self) -> Option<std::sync::Arc<atlas_core::registry::AtlasRegistry>> {
349        Some(self.registry().clone())
350    }
351
352    #[track_caller]
353    fn kernel(&self, module: &str, func_name: &str) -> Result<KernelHandle> {
354        // The DISPATCH SITE, not this line: `#[track_caller]` here and on the
355        // trait declaration carries the `.kernel(…)` / `try_kernel(…)` caller's
356        // `file:line` through, which is the only part of an unresolved-lookup
357        // report an operator can act on.
358        let site = std::panic::Location::caller();
359        // Ephemeral OnceLock — no cross-call caching, but kernel() is only
360        // called at model init time. Layers store the returned KernelHandle.
361        let cache: OnceLock<RawCudaFunc> = OnceLock::new();
362        let registry = self.registry();
363        match registry.raw_function_cached(&cache, module, func_name) {
364            Ok(raw) => {
365                crate::kernel_audit::record(module, func_name, true, site);
366                crate::launch_trace::name_kernel(raw.0 as u64, module, func_name);
367                Ok(KernelHandle(raw.0 as u64))
368            }
369            Err(e) => {
370                // Optional kernels (try_kernel) land here and fall back silently;
371                // the audit makes that visible in the startup kernel table.
372                crate::kernel_audit::record(module, func_name, false, site);
373                Err(anyhow::anyhow!("Kernel lookup {module}::{func_name}: {e}"))
374            }
375        }
376    }
377
378    fn copy_h2d_async(&self, src: &[u8], dst: DevicePtr, stream: u64) -> Result<()> {
379        h2d_enqueue(src, dst, stream)?;
380        // The trait promises the caller may drop `src` right now. From PAGEABLE
381        // memory the driver already made that true by staging the bytes before
382        // returning. From PAGE-LOCKED memory it did not — the DMA engine reads
383        // these pages after the enqueue — so buy the same guarantee with an
384        // explicit wait rather than let ~90 call sites that drop their source
385        // immediately turn into use-after-frees the day a buffer gets pinned.
386        //
387        // Costs nothing on the path everything takes today: no Atlas call site
388        // reaches here with a pinned source (the ones that own pinned staging
389        // use `copy_h2d_async_retained`), so `is_pinned` is a lock-free-ish read
390        // of a three-entry table that says "no".
391        if crate::pinned_hosts::is_pinned(src) {
392            warn_pinned_transient_source();
393            let sync = unsafe { cuStreamSynchronize(stream) };
394            if sync != 0 {
395                bail!(
396                    "cuStreamSynchronize after pinned-source H2D failed: {}",
397                    cuda_error_text(sync)
398                );
399            }
400        }
401        Ok(())
402    }
403
404    fn copy_h2d_async_retained(&self, src: &[u8], dst: DevicePtr, stream: u64) -> Result<()> {
405        // The caller has promised `src` outlives the next sync on `stream`, so
406        // no implicit ordering is added — that is the whole reason this variant
407        // exists (a 60-chunk pinned scatter must not pay 60 stream drains).
408        h2d_enqueue(src, dst, stream)
409    }
410
411    fn copy_d2d_async(
412        &self,
413        src: DevicePtr,
414        dst: DevicePtr,
415        bytes: usize,
416        stream: u64,
417    ) -> Result<()> {
418        let status = unsafe { cuMemcpyDtoDAsync_v2(dst.0, src.0, bytes, stream) };
419        if status != 0 {
420            // See copy_d2d_impl: on 901 the backtrace names the reporter,
421            // which bounds where the capture-poisoning op ran.
422            tracing::error!(
423                "copy_d2d_async failed (status {status}) at:\n{}",
424                std::backtrace::Backtrace::force_capture()
425            );
426            bail!("cuMemcpyDtoDAsync_v2 (copy_d2d_async) failed: status {status}");
427        }
428        Ok(())
429    }
430
431    fn copy_d2d_2d_async(
432        &self,
433        src: DevicePtr,
434        src_pitch: usize,
435        dst: DevicePtr,
436        dst_pitch: usize,
437        width_bytes: usize,
438        height: usize,
439        stream: u64,
440    ) -> Result<()> {
441        // One pitched copy (cudaMemcpyDeviceToDevice = 3) on the caller's stream,
442        // replacing a per-row copy_d2d_async loop. cudart is linked (cutlass/
443        // flashinfer use the runtime API); a CUstream handle is a valid
444        // cudaStream_t.
445        unsafe extern "C" {
446            fn cudaMemcpy2DAsync(
447                dst: *mut c_void,
448                dpitch: usize,
449                src: *const c_void,
450                spitch: usize,
451                width: usize,
452                height: usize,
453                kind: i32,
454                stream: u64,
455            ) -> i32;
456        }
457        let status = unsafe {
458            cudaMemcpy2DAsync(
459                dst.0 as *mut c_void,
460                dst_pitch,
461                src.0 as *const c_void,
462                src_pitch,
463                width_bytes,
464                height,
465                3,
466                stream,
467            )
468        };
469        if status != 0 {
470            bail!("cudaMemcpy2DAsync failed: status {status}");
471        }
472        Ok(())
473    }
474
475    fn begin_capture(&self, stream: u64) -> Result<()> {
476        self.begin_capture_cu(stream)
477    }
478    fn end_capture(&self, stream: u64) -> Result<GraphHandle> {
479        self.end_capture_cu(stream)
480    }
481
482    fn abort_capture_if_active(&self, stream: u64) {
483        self.abort_capture_if_active_cu(stream)
484    }
485
486    fn launch_graph(&self, graph: GraphHandle, stream: u64) -> Result<()> {
487        self.launch_graph_cu(graph, stream)
488    }
489    fn destroy_graph(&self, graph: GraphHandle) -> Result<()> {
490        self.destroy_graph_cu(graph)
491    }
492    fn memset(&self, ptr: DevicePtr, value: u8, bytes: usize) -> Result<()> {
493        self.memset_cu(ptr, value, bytes)
494    }
495    fn memset_async(&self, ptr: DevicePtr, value: u8, bytes: usize, stream: u64) -> Result<()> {
496        if crate::launch_trace::on() {
497            crate::launch_trace::record(crate::launch_trace::Entry {
498                kind: "memset",
499                func: 0,
500                grid: [0, 0, 0],
501                block: [0, 0, 0],
502                smem: 0,
503                args: vec![ptr.0, value as u64, bytes as u64],
504            });
505        }
506        self.memset_async_cu(ptr, value, bytes, stream)
507    }
508    fn total_memory(&self) -> Result<usize> {
509        self.total_memory_cu()
510    }
511    fn free_memory(&self) -> Result<usize> {
512        self.free_memory_cu()
513    }
514    fn device_free_memory(&self) -> Result<usize> {
515        self.device_free_memory_cu()
516    }
517    fn live_alloc_count(&self) -> usize {
518        self.live_alloc_len()
519    }
520    fn sm_count(&self) -> Result<u32> {
521        self.sm_count_cu()
522    }
523    fn create_stream(&self) -> Result<u64> {
524        self.create_stream_cu()
525    }
526    fn bind_to_thread(&self) -> Result<()> {
527        self.bind_to_thread_cu()
528    }
529    fn create_event(&self) -> Result<u64> {
530        self.create_event_cu()
531    }
532    fn record_event(&self, event: u64, stream: u64) -> Result<()> {
533        self.record_event_cu(event, stream)
534    }
535    fn stream_wait_event(&self, stream: u64, event: u64) -> Result<()> {
536        self.stream_wait_event_cu(stream, event)
537    }
538    fn event_synchronize(&self, event: u64) -> Result<()> {
539        self.event_synchronize_cu(event)
540    }
541    fn destroy_event(&self, event: u64) -> Result<()> {
542        self.destroy_event_cu(event)
543    }
544    fn host_ptr_to_device(&self, host: *mut u8) -> Result<DevicePtr> {
545        let mut dptr: u64 = 0;
546        let status =
547            unsafe { super::cuMemHostGetDevicePointer_v2(&mut dptr, host as *mut c_void, 0) };
548        if status != 0 {
549            bail!("cuMemHostGetDevicePointer_v2 failed: status {status}");
550        }
551        Ok(DevicePtr(dptr))
552    }
553
554    fn alloc_host_pinned(&self, bytes: usize) -> Result<*mut u8> {
555        if bytes >= 32 * 1024 * 1024 {
556            tracing::debug!(
557                "alloc_host_pinned {:.1} MB",
558                bytes as f64 / (1024.0 * 1024.0)
559            );
560        }
561        self.alloc_host_pinned_cu(bytes)
562    }
563    fn free_host_pinned(&self, ptr: *mut u8, _bytes: usize) -> Result<()> {
564        self.free_host_pinned_cu(ptr, _bytes)
565    }
566}
spark_runtime/cuda_backend/gpu_impl.rs

spark_runtime/cuda_backend/
gpu_impl.rs
