cherenkov/qwen4_exp/gpu/
memory.rs1use super::*;
4use objc2_metal::MTLDevice;
5use serde::Serialize;
6
7#[derive(Debug, Default, Clone, Copy, Serialize)]
10pub struct MemoryStats {
11 pub metal_allocated_bytes_observed: u64,
12 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 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 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}