combs_core/lib.rs
1//! # combs-core
2//!
3//! Backend type aliases, device helpers, and the memory-pool facade for the
4//! Combs Engine L0 Rust core.
5//!
6//! Phase 1 runs entirely on the wgpu backend (Metal on Apple Silicon) with
7//! f32 compute. f16 compute and custom CubeCL kernels are later-phase
8//! optimizations; the aliases here are the single place to change.
9
10pub use burn::backend::wgpu;
11
12pub mod quant;
13
14use burn::backend::wgpu::{RuntimeOptions, WgpuDevice, WgpuSetup, graphics::AutoGraphicsApi};
15
16/// The default inference backend: autotuned, fusing wgpu/CubeCL backend
17/// (`Fusion<CubeBackend<WgpuRuntime, f32, i32, u32>>`).
18///
19/// `--features f16` switches to an **unfused f16** backend, which ~halves
20/// weight + KV + activation memory (e.g. a 3B model drops from ~12 GB to
21/// ~6 GB) and is typically faster. The numerically sensitive reductions
22/// (RMS/LayerNorm, attention scores + softmax, gelu) run in f32 regardless
23/// of backend (see `combs-models::precision`), so f16 output stays coherent.
24///
25/// Note: f16 uses the **unfused** `CubeBackend` type directly — burn-fusion
26/// 0.21 panics on reduced-precision tensors, so we bypass the fusion layer
27/// for f16 while keeping f32 fused. bf16 is unavailable (cubecl's matmul has
28/// no bf16 path on Metal/wgpu).
29#[cfg(not(feature = "f16"))]
30pub type CombsBackend = burn::backend::Wgpu<f32, i32, u32>;
31
32/// Always-f32 backend on the same wgpu runtime. The diffusion pipeline is
33/// pinned to it in every build: SD-1.5's UNet/VAE collapse to black output
34/// under f16 (range, not rounding), so image generation computes in f32
35/// regardless of the text stack's dtype.
36pub type CombsBackendF32 = burn::backend::Wgpu<f32, i32, u32>;
37#[cfg(feature = "f16")]
38pub type CombsBackend = burn::backend::wgpu::CubeBackend<
39 burn::backend::wgpu::WgpuRuntime,
40 burn::tensor::f16,
41 i32,
42 u32,
43>;
44
45/// The default device handle type.
46pub type CombsDevice = WgpuDevice;
47
48/// Returns the default wgpu device (best available GPU; on macOS this is the
49/// Metal device). Honors cubecl's `CUBECL_WGPU_DEFAULT_DEVICE` override.
50pub fn init_device() -> CombsDevice {
51 WgpuDevice::default()
52}
53
54/// True when wgpu can see at least one adapter. Cached after the first
55/// probe. Initializing a cubecl device on an adapterless machine (e.g. a
56/// CI runner) panics in a worker thread, so GPU-dependent tests check
57/// this first and skip rather than fail.
58pub fn gpu_available() -> bool {
59 static AVAILABLE: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
60 // `::wgpu` is the raw crate — the bare name resolves to this
61 // module's `burn::backend::wgpu` re-export.
62 *AVAILABLE.get_or_init(|| {
63 let instance = ::wgpu::Instance::default();
64 let adapters =
65 cubecl::future::block_on(instance.enumerate_adapters(::wgpu::Backends::all()));
66 !adapters.is_empty()
67 })
68}
69
70/// Basic information about a wgpu adapter.
71#[derive(Debug, Clone)]
72pub struct DeviceInfo {
73 /// Human-readable adapter name (e.g. "Apple M3 Pro").
74 pub name: String,
75 /// Graphics backend in use (e.g. "Metal").
76 pub backend: String,
77 /// Device type (e.g. "IntegratedGpu").
78 pub device_type: String,
79 /// Driver name + info string.
80 pub driver: String,
81}
82
83/// Hardware capabilities consumed by the application-layer device planner
84/// (sharding, KV budget, prefill chunk sizing). Serialized to JSON across
85/// the FFI boundary.
86#[derive(Debug, Clone, serde::Serialize)]
87pub struct DeviceCaps {
88 /// Human-readable adapter name (e.g. "Apple M3 Pro").
89 pub name: String,
90 /// Graphics backend in use ("Metal", "Vulkan", "Dx12", "Gl", ...).
91 pub backend: String,
92 /// Device type ("IntegratedGpu", "DiscreteGpu", ...).
93 pub device_type: String,
94 /// `max_storage_buffer_binding_size`: the hard cap on a single GPU
95 /// buffer — the sharding limit on mobile devices.
96 pub max_storage_buffer_binding_size: u64,
97 /// `max_buffer_size`: the largest single allocation the driver allows.
98 pub max_buffer_size: u64,
99 /// Largest compute workgroup dimension.
100 pub max_compute_workgroup_size_x: u32,
101 /// `max_compute_invocations_per_workgroup`.
102 pub max_compute_invocations_per_workgroup: u32,
103 /// Debug dump of the adapter's enabled feature set (wgpu 29 no longer
104 /// exposes WebGPU extension features like `SHADER_F16` through the
105 /// public adapter API, so we surface the raw list for the planner).
106 pub features: String,
107}
108
109/// Queries the adapter for its limits/features and returns [`DeviceCaps`].
110///
111/// Like [`device_info`], this performs the cubecl runtime setup for the
112/// device, so it is safe (and cheap) to use the device afterwards.
113pub fn device_caps(device: &CombsDevice) -> DeviceCaps {
114 let setup: WgpuSetup =
115 burn::backend::wgpu::init_setup::<AutoGraphicsApi>(device, RuntimeOptions::default());
116 let info = setup.adapter.get_info();
117 let limits = setup.adapter.limits();
118 let features = setup.adapter.features();
119 DeviceCaps {
120 name: info.name,
121 backend: format!("{:?}", info.backend),
122 device_type: format!("{:?}", info.device_type),
123 max_storage_buffer_binding_size: limits.max_storage_buffer_binding_size as u64,
124 max_buffer_size: limits.max_buffer_size,
125 max_compute_workgroup_size_x: limits.max_compute_workgroup_size_x,
126 max_compute_invocations_per_workgroup: limits.max_compute_invocations_per_workgroup,
127 features: format!("{:?}", features),
128 }
129}
130
131/// GPU allocator state from cubecl's memory manager (authoritative — process
132/// RSS is meaningless for unified-memory GPU accounting).
133#[derive(Debug, Clone, Copy, Default, serde::Serialize)]
134pub struct GpuMemory {
135 /// Bytes referenced by live handles.
136 pub bytes_in_use: u64,
137 /// Bytes reserved by the pool (in-use + cached slabs).
138 pub bytes_reserved: u64,
139 /// Bytes lost to alignment padding.
140 pub bytes_padding: u64,
141 /// Live allocation count.
142 pub number_allocs: u64,
143}
144
145/// Samples the GPU allocator. `memory_usage()` is `submit_blocking` on the
146/// compute stream — call from the engine worker between generations (or
147/// rate-limited), not from request threads during a long prefill.
148pub fn gpu_memory(device: &CombsDevice) -> Option<GpuMemory> {
149 let client =
150 <burn::backend::wgpu::WgpuRuntime as cubecl::prelude::Runtime>::client(device);
151 client.memory_usage().ok().map(|m| GpuMemory {
152 bytes_in_use: m.bytes_in_use,
153 bytes_reserved: m.bytes_reserved,
154 bytes_padding: m.bytes_padding,
155 number_allocs: m.number_allocs,
156 })
157}
158
159/// Initializes the wgpu runtime for `device` and returns adapter information.
160///
161/// Note: this performs the cubecl runtime setup for the device (the same setup
162/// burn performs lazily on first tensor use), so it is safe to use the device
163/// for compute afterwards.
164pub fn device_info(device: &CombsDevice) -> DeviceInfo {
165 let setup: WgpuSetup =
166 burn::backend::wgpu::init_setup::<AutoGraphicsApi>(device, RuntimeOptions::default());
167 let info = setup.adapter.get_info();
168 DeviceInfo {
169 name: info.name,
170 backend: format!("{:?}", info.backend),
171 device_type: format!("{:?}", info.device_type),
172 driver: format!("{} ({})", info.driver, info.driver_info),
173 }
174}
175
176/// Facade over the GPU buffer pool.
177///
178/// # Phase 1 status: documented no-op
179///
180/// The plan's hand-rolled slab/coalescing pool was deliberately replaced by
181/// cubecl's built-in allocator, which already does pooled slab allocation and
182/// reuse (configured through [`burn::backend::wgpu::MemoryConfiguration`]).
183/// burn 0.21 does not publicly expose cubecl 0.10's
184/// `ComputeClient::memory_cleanup`, and `memory_persistent_allocations` does
185/// not exist in cubecl 0.10's public API at all, so this facade currently does
186/// nothing. It exists so that the runtime can call `pool.cleanup()` /
187/// `pool.pin_persistent()` today and Phase 2 can back those calls with real
188/// handles (persistent KV/weight arenas) without changing call sites.
189#[derive(Debug, Default, Clone, Copy)]
190pub struct BufferPool;
191
192impl BufferPool {
193 /// Creates a new pool facade.
194 pub fn new() -> Self {
195 BufferPool
196 }
197
198 /// Pin long-lived allocations (weights, KV arena) so the pool never
199 /// releases them. No-op in Phase 1 — cubecl's pooled allocator keeps
200 /// freed blocks for reuse anyway.
201 pub fn pin_persistent(&self) {
202 // no-op: see type-level docs.
203 }
204
205 /// Release cached free blocks back to the driver. No-op in Phase 1 —
206 /// cubecl 0.10's `memory_cleanup` is not reachable through burn's public
207 /// API.
208 pub fn cleanup(&self) {
209 // no-op: see type-level docs.
210 }
211}