Skip to main content

cherenkov/qwen4_exp/gpu/
memory.rs

1//! Runtime resource counters shared by CLI telemetry and control status.
2
3use super::*;
4use objc2_metal::MTLDevice;
5use serde::Serialize;
6
7/// Capacities and observations, not a partition of process memory. CPU prefix
8/// checkpoints and session history are reported separately by the server.
9#[derive(Debug, Default, Clone, Copy, Serialize)]
10pub struct MemoryStats {
11    pub metal_allocated_bytes_observed: u64,
12    /// Length of the shared weight buffer, counted once for all tensor views.
13    pub mapped_weight_buffer_bytes: usize,
14    pub kv_index_capacity_bytes: usize,
15    pub context_capacity_tokens: usize,
16    pub mtp_enabled: bool,
17    /// Configured pool capacity; the residency-set backend may use fewer slots.
18    pub expert_pool_bytes: usize,
19    pub device_working_set_bytes: usize,
20    pub allocation_limit_bytes: Option<usize>,
21    pub prefill_reserved_bytes: usize,
22    pub expert_pool_slots: usize,
23    pub resident_experts: usize,
24}
25
26impl Gpu<'_> {
27    pub fn allocated_bytes(&self) -> u64 {
28        self.ctx.device.currentAllocatedSize() as u64
29    }
30
31    pub fn memory_stats(&self) -> MemoryStats {
32        MemoryStats {
33            metal_allocated_bytes_observed: self.allocated_bytes(),
34            mapped_weight_buffer_bytes: self.dense.length(),
35            kv_index_capacity_bytes: self.kv_index_capacity_bytes(),
36            context_capacity_tokens: self.max_t,
37            mtp_enabled: self.has_mtp(),
38            expert_pool_bytes: self.pool_bytes(),
39            device_working_set_bytes: self.ctx.device.recommendedMaxWorkingSetSize() as usize,
40            allocation_limit_bytes: self.ctx.allocation_limit.get(),
41            prefill_reserved_bytes: self.prefill_reserved_bytes,
42            expert_pool_slots: self.pool_slots(),
43            resident_experts: self.pool_resident(),
44        }
45    }
46
47    /// Reserved KV and attention-index bytes, including the enabled MTP layer.
48    fn kv_index_capacity_bytes(&self) -> usize {
49        self.layers
50            .iter()
51            .chain(self.mtp.iter().map(|m| &m.layer))
52            .map(|layer| match &layer.mix {
53                Mix::Attn(a) => a.kc.length() + a.vc.length() + a.ikc.length() + a.blk.length(),
54                Mix::Delta(_) => 0,
55            })
56            .sum()
57    }
58}