kime_cuda/
lib.rs

1//! The NVIDIA backend, from spec/08-cuda.md.
2//!
3//! This first version runs the compat graph with FP16 weights and activations and FP32
4//! accumulation. The residual stream, the head's inputs and everything after the scorer stay in
5//! FP32. The GEMMs go through cuBLASLt, and the other ops are our own kernels in
6//! `kernels/compat.cu`, compiled with NVRTC for the device when the backend starts, so the build
7//! needs no CUDA toolkit and the binary runs on any machine with a driver.
8//!
9//! Every launch is sized for the plan's bucket and every kernel reads the batch's real counts from
10//! a device buffer, so a run is a fixed sequence of launches that a CUDA graph can capture.
11//!
12//! One of the six crates where `unsafe` is allowed. Every block carries a `// SAFETY:` comment
13//! that names the invariant which makes it sound.
14
15mod compat;
16pub mod energy;
17#[cfg(test)]
18mod kernel_tests;
19mod lt;
20mod plan;
21mod tune;
22
23pub use compat::executor;
24pub use plan::{CudaPlan, Weights};
25
26use std::sync::Arc;
27
28use cudarc::driver::{CudaContext, CudaFunction, CudaModule, CudaSlice, CudaStream, DevicePtr};
29use cudarc::nvrtc::{CompileOptions, compile_ptx_with_opts};
30use kime_tensor::{Error, Result};
31
32const SOURCE: &str = include_str!("../kernels/compat.cu");
33
34/// Bytes of scratch cuBLASLt may use.
35const WORKSPACE: usize = 32 << 20;
36
37fn dev(e: impl std::fmt::Debug) -> Error {
38    Error::Device(format!("{e:?}"))
39}
40
41/// The kernels of `kernels/compat.cu`.
42pub(crate) struct Kernels {
43    pub(crate) embed: CudaFunction,
44    /// Indexed by `2 * (input is f16) + (output is f16)`.
45    pub(crate) ln: [CudaFunction; 4],
46    /// f32, f16.
47    pub(crate) bias_act: [CudaFunction; 2],
48    /// f32, f16.
49    pub(crate) rope: [CudaFunction; 2],
50    /// Indexed like `ln`.
51    pub(crate) attention: [CudaFunction; 4],
52    /// Indexed like `ln`.
53    pub(crate) geglu: [CudaFunction; 4],
54    pub(crate) add_type: CudaFunction,
55    /// f32, f16.
56    pub(crate) gather: [CudaFunction; 2],
57    pub(crate) act_features: CudaFunction,
58    pub(crate) to_f16: CudaFunction,
59}
60
61impl Kernels {
62    fn load(m: &Arc<CudaModule>) -> Result<Self> {
63        let f = |name: &str| m.load_function(name).map_err(dev);
64        Ok(Self {
65            embed: f("embed")?,
66            ln: [f("ln_f32_f32")?, f("ln_f32_f16")?, f("ln_f16_f32")?, f("ln_f16_f16")?],
67            bias_act: [f("bias_act_f32")?, f("bias_act_f16")?],
68            rope: [f("rope_f32")?, f("rope_f16")?],
69            attention: [
70                f("attention_f32_f32")?,
71                f("attention_f32_f16")?,
72                f("attention_f16_f32")?,
73                f("attention_f16_f16")?,
74            ],
75            geglu: [
76                f("geglu_f32_f32")?,
77                f("geglu_f32_f16")?,
78                f("geglu_f16_f32")?,
79                f("geglu_f16_f16")?,
80            ],
81            add_type: f("add_type")?,
82            gather: [f("gather_f32")?, f("gather_f16")?],
83            act_features: f("act_features")?,
84            to_f16: f("to_f16")?,
85        })
86    }
87}
88
89/// How much of the graph runs in FP16.
90#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
91pub enum Precision {
92    /// FP32 values, weights and GEMMs. Matches Laya to about 1e-4 in the logits.
93    #[default]
94    F32,
95    /// FP16 weights and GEMM inputs with FP32 accumulation, the residual stream, attention's qkv
96    /// and the heads in FP32.
97    F16,
98}
99
100/// One GPU, with its stream, kernels and cuBLASLt handle.
101pub struct CudaBackend {
102    stream: Arc<CudaStream>,
103    k: Kernels,
104    lt: lt::Handle,
105    workspace: CudaSlice<u8>,
106    name: String,
107    arch: (i32, i32),
108    precision: Precision,
109    picks: tune::Picks,
110}
111
112impl std::fmt::Debug for CudaBackend {
113    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
114        f.debug_struct("CudaBackend")
115            .field("name", &self.name)
116            .field("arch", &self.arch)
117            .field("precision", &self.precision)
118            .finish_non_exhaustive()
119    }
120}
121
122impl CudaBackend {
123    /// Opens GPU `ordinal` and compiles the kernels for it.
124    ///
125    /// # Errors
126    ///
127    /// [`Error::Device`] when there is no such GPU, no driver, or the kernels do not compile.
128    pub fn new(ordinal: usize, precision: Precision) -> Result<Self> {
129        // cudarc panics on first use of a library it cannot load, so check for all three first.
130        // SAFETY: these only try to open the libraries, which runs no code of theirs we rely on.
131        let present = unsafe {
132            [
133                ("the CUDA driver", cudarc::driver::sys::is_culib_present()),
134                ("NVRTC", cudarc::nvrtc::sys::is_culib_present()),
135                ("cuBLASLt", cudarc::cublaslt::sys::is_culib_present()),
136            ]
137        };
138        if let Some((name, _)) = present.iter().find(|p| !p.1) {
139            return Err(Error::Device(format!("{name} is not installed")));
140        }
141        let ctx = CudaContext::new(ordinal).map_err(dev)?;
142        // SAFETY: all work goes to one stream, in order, and the plan synchronizes that stream
143        // before it reads results or frees buffers, so cudarc's per buffer events are not needed.
144        unsafe { ctx.disable_event_tracking() };
145        let stream = ctx.new_stream().map_err(dev)?;
146        let arch = ctx.compute_capability().map_err(dev)?;
147        let opts = CompileOptions { arch: Some(arch_flag(arch)), ..Default::default() };
148        let ptx = compile_ptx_with_opts(SOURCE, opts).map_err(dev)?;
149        let module = ctx.load_module(ptx).map_err(dev)?;
150        let k = Kernels::load(&module)?;
151        let lt = lt::Handle::new()?;
152        let workspace = stream.alloc_zeros::<u8>(WORKSPACE).map_err(dev)?;
153        let name = ctx.name().map_err(dev)?;
154        let picks = tune::Picks::for_gpu(&name);
155        Ok(Self { stream, k, lt, workspace, name, arch, precision, picks })
156    }
157
158    /// The GPU's name, as the driver reports it.
159    #[must_use]
160    pub fn name(&self) -> &str {
161        &self.name
162    }
163
164    /// What runs in FP16.
165    #[must_use]
166    pub fn precision(&self) -> Precision {
167        self.precision
168    }
169
170    /// Compute capability, major and minor.
171    #[must_use]
172    pub fn arch(&self) -> (i32, i32) {
173        self.arch
174    }
175
176    fn workspace_ptr(&self) -> u64 {
177        self.workspace.device_ptr(&self.stream).0
178    }
179}
180
181/// The virtual architecture NVRTC compiles for. The driver JITs the PTX for the actual device,
182/// so a newer GPU than any listed gets the newest one here.
183fn arch_flag((major, minor): (i32, i32)) -> &'static str {
184    match (major, minor) {
185        (..=7, _) => "compute_75",
186        (8, 0) => "compute_80",
187        (8, 6..=8) => "compute_86",
188        (8, _) => "compute_89",
189        _ => "compute_90",
190    }
191}
192
193#[cfg(test)]
194mod tests {
195    use super::*;
196
197    #[test]
198    fn arch_flags() {
199        assert_eq!(arch_flag((7, 5)), "compute_75");
200        assert_eq!(arch_flag((8, 6)), "compute_86");
201        assert_eq!(arch_flag((8, 9)), "compute_89");
202        assert_eq!(arch_flag((12, 0)), "compute_90");
203    }
204}
kime_cuda/lib.rs

kime_cuda/
lib.rs
