cherenkov/qwen4_exp/gpu/
load.rs

1//! Model loading, pipeline creation, and fixed/pool allocation.
2
3use super::*;
4use crate::units::BYTES_PER_GB;
5use objc2_metal::MTLDevice;
6
7impl<'a> Gpu<'a> {
8    pub fn load(p: &'a Packed, max_t: usize, options: &Options) -> Result<Self> {
9        Self::load_bounded(p, max_t, options, 0, None, prefill::MAX_PREFILL_ROWS)
10    }
11
12    pub(crate) fn load_bounded(
13        p: &'a Packed,
14        max_t: usize,
15        options: &Options,
16        reserve: usize,
17        memory_bytes: Option<usize>,
18        prefill_rows: usize,
19    ) -> Result<Self> {
20        options.validate()?;
21        p.manifest.experts.validate_dimensions()?;
22
23        let c = &p.cfg;
24
25        anyhow::ensure!(
26            c.hidden_size.is_multiple_of(256),
27            "fn_norm_prep_b needs hidden % 256 == 0"
28        );
29        anyhow::ensure!(
30            c.indexer_head_dim <= 256 && c.indexer_head_dim.is_multiple_of(32),
31            "indexer head dim must be a multiple of 32 up to 256"
32        );
33        anyhow::ensure!(
34            c.indexer_n_heads * c.indexer_head_dim <= 1024,
35            "indexer query heads too wide for fn_index_score"
36        );
37        anyhow::ensure!(
38            c.indexer_budget.is_multiple_of(c.indexer_compress_ratio),
39            "indexer budget must be a multiple of the block size"
40        );
41        anyhow::ensure!(
42            c.moe_intermediate_size.is_multiple_of(64),
43            "expert width must be a multiple of the 64-code quantization group"
44        );
45        anyhow::ensure!(
46            c.num_experts_per_tok * MAX_NB < SLOT_STRIDE,
47            "slot table too narrow"
48        );
49        anyhow::ensure!(
50            c.output_gate_type == "sigmoid",
51            "only the sigmoid DeltaNet output gate is implemented"
52        );
53
54        let ctx = MetalContext::new()?;
55        let allocation_limit = memory_bytes
56            .map(|bytes| {
57                bytes
58                    .checked_sub(reserve)
59                    .context("cache exceeds memory budget")
60            })
61            .transpose()?;
62
63        ctx.allocation_limit.set(allocation_limit);
64
65        let lib = ctx.compile_library(FORWARD_MSL)?;
66        let blib = ctx.compile_library(BATCH_MSL)?;
67        let per_nb = |base: &str| -> Result<[Pso; MAX_NB]> {
68            Ok([
69                ctx.pipeline(&blib, &format!("{base}1"))?,
70                ctx.pipeline(&blib, &format!("{base}2"))?,
71                ctx.pipeline(&blib, &format!("{base}3"))?,
72                ctx.pipeline(&blib, &format!("{base}4"))?,
73            ])
74        };
75        let expert_qmm = |bits| -> Result<[Pso; 3]> {
76            Ok([
77                ctx.pipeline(&blib, &format!("fn_expert_qmm_q{bits}_n8"))?,
78                ctx.pipeline(&blib, &format!("fn_expert_qmm_q{bits}_n16"))?,
79                ctx.pipeline(&blib, &format!("fn_expert_qmm_q{bits}_n32"))?,
80            ])
81        };
82        let pipes = Pipes {
83            expert_qmm: [expert_qmm(2)?, expert_qmm(3)?],
84            qmv_h: ctx.pipeline(&lib, "qmv_multi_h")?,
85            qmv_hn: ctx.pipeline(&lib, "qmv_multi_hn")?,
86            prep_h: ctx.pipeline(&lib, "deinterleave_bh")?,
87            embed_rows: ctx.pipeline(&lib, "embed_rows")?,
88            qk_norm_rope_b: ctx.pipeline(&lib, "qk_norm_rope_b")?,
89            conv_b: ctx.pipeline(&lib, "conv_b")?,
90            delta_norms: ctx.pipeline(&lib, "delta_norms")?,
91            delta_gates: ctx.pipeline(&lib, "delta_gates")?,
92            delta_scan2: ctx.pipeline(&lib, "delta_scan2")?,
93            kv_append_q8: ctx.pipeline(&lib, "kv_append_q8")?,
94            attn_part2_q8: ctx.pipeline(&lib, "attn_part2_q8")?,
95            attn_combine: ctx.pipeline(&lib, "attn_combine")?,
96            argmax_partial: ctx.pipeline(&lib, "argmax_partial")?,
97            argmax_final: ctx.pipeline(&lib, "argmax_final")?,
98            add: ctx.pipeline(&lib, "add_inplace")?,
99            copy_f32: ctx.pipeline(&lib, "copy_f32")?,
100            zero: ctx.pipeline(&blib, "fn_zero")?,
101            replicate_b: ctx.pipeline(&blib, "fn_replicate_b")?,
102            group_norm_b: ctx.pipeline(&blib, "fn_group_norm_b")?,
103            norm_prep_b: ctx.pipeline(&blib, "fn_norm_prep_b")?,
104            qmv_silu_b: per_nb("fn_qmv_silu_b")?,
105            hc_mix_b: ctx.pipeline(&blib, "fn_hc_mix_b")?,
106            inject_b: ctx.pipeline(&blib, "fn_inject_b")?,
107            bf16_matvec_b: ctx.pipeline(&blib, "fn_bf16_matvec_b")?,
108            topk_softmax_b: ctx.pipeline(&blib, "fn_topk_softmax_b")?,
109            moe_gate_up_b: per_nb("fn_moe_gate_up_b")?,
110            moe_act_b: ctx.pipeline(&blib, "fn_moe_act_b")?,
111            moe_down_b: per_nb("fn_moe_down_b")?,
112            moe_combine_b: ctx.pipeline(&blib, "fn_moe_combine_b")?,
113            ple_gate_b: ctx.pipeline(&blib, "fn_ple_gate_b")?,
114            ple_conv_b: ctx.pipeline(&blib, "fn_ple_conv_b")?,
115            mtp_fold: ctx.pipeline(&blib, "fn_mtp_fold")?,
116            gate_norm_sigmoid_b: ctx.pipeline(&blib, "fn_delta_gate_norm_sigmoid_b")?,
117            index_append: ctx.pipeline(&blib, "fn_index_append")?,
118            index_blocks: ctx.pipeline(&blib, "fn_index_blocks")?,
119            index_q: ctx.pipeline(&blib, "fn_index_q")?,
120            index_score: ctx.pipeline(&blib, "fn_index_score")?,
121            index_select: ctx.pipeline(&blib, "fn_index_select")?,
122            attn_sel: ctx.pipeline(&blib, "fn_attn_part2_q8_sel")?,
123            qmm_n8: ctx.pipeline(&lib, "qmm_af4_n8")?,
124            qmm_n16: ctx.pipeline(&lib, "qmm_af4_n16")?,
125            qmm_w: ctx.pipeline(&lib, "qmm_af4_w")?,
126            silu_mul: ctx.pipeline(&lib, "silu_mul")?,
127            attn_q_stage: ctx.pipeline(&lib, "attn_q_stage")?,
128            attn_kv_stage: ctx.pipeline(&lib, "attn_kv_stage")?,
129            gemm_hh: ctx.pipeline(&lib, "gemm_hh")?,
130            attn_o_scatter: ctx.pipeline(&lib, "attn_o_scatter")?,
131            softmax_sel: ctx.pipeline(&blib, "fn_attn_softmax_sel")?,
132            silu_rows: ctx.pipeline(&blib, "fn_silu_rows")?,
133            gather_rows: ctx.pipeline(&blib, "fn_gather_rows")?,
134            scatter_add_rows: ctx.pipeline(&blib, "fn_scatter_add_rows")?,
135            shared_add_rows: ctx.pipeline(&blib, "fn_shared_add_rows")?,
136            qmv_small_b: ctx.pipeline(&blib, "fn_qmv_small_b")?,
137        };
138        // Gpu borrows Packed for its entire lifetime, keeping this
139        // page-aligned mapping alive while dense kernels can read it.
140        let dense = unsafe { ctx.wrap_mmap(&p.dense)? };
141
142        let q = |prefix: &str| -> Result<Q> {
143            let w = p.manifest.dense(&format!("{prefix}.weight"))?;
144            let s = p.manifest.dense(&format!("{prefix}.scales"))?;
145            let b = p.manifest.dense(&format!("{prefix}.biases"))?;
146
147            anyhow::ensure!(w.dtype == "U32", "{prefix}: not a Q4 weight");
148            anyhow::ensure!(w.offset % 16 == 0, "{prefix}: weight offset not 16-aligned");
149
150            Ok(Q {
151                w: w.offset as usize,
152                s: s.offset as usize,
153                b: b.offset as usize,
154                out: w.shape[0] as u32,
155                inp: (w.shape[1] * 8) as u32,
156            })
157        };
158        let t = |name: &str| -> Result<T> {
159            let e = p.manifest.dense(name)?;
160
161            anyhow::ensure!(e.dtype == "BF16", "{name}: expected BF16");
162
163            Ok(T(e.offset as usize))
164        };
165        let hc = |prefix: &str, inject: bool| -> Result<Hc> {
166            Ok(Hc {
167                norm: t(&format!("{prefix}.hc_norm.weight"))?,
168                down: q(&format!("{prefix}.input_mix_weight_down"))?,
169                up: q(&format!("{prefix}.input_mix_weight_up"))?,
170                inject: if inject {
171                    Some(q(&format!("{prefix}.block_inject_weight"))?)
172                } else {
173                    None
174                },
175            })
176        };
177
178        let h = c.hidden_size;
179        let hh = c.hc_hidden();
180        let kv_row = c.num_key_value_heads * c.head_dim;
181        let kv_side = kv_q8_side(max_t, kv_row).1;
182        let conv_dim = 2 * c.linear_num_key_heads * c.linear_key_head_dim
183            + c.linear_num_value_heads * c.linear_value_head_dim;
184        let v_dim = c.linear_num_value_heads * c.linear_value_head_dim;
185        let state_len = c.linear_num_value_heads * c.linear_key_head_dim * c.linear_value_head_dim;
186        let hist_len = conv_dim * (c.linear_conv_kernel_dim - 1);
187
188        let load_layer =
189            |lp: &str, linear: bool, record_layer: usize, with_ple: bool| -> Result<GLayer> {
190                let mix = if linear {
191                    let d = format!("{lp}.linear_attn");
192
193                    Mix::Delta(Delta {
194                        qkv: q(&format!("{d}.in_proj_qkv"))?,
195                        z: q(&format!("{d}.in_proj_z"))?,
196                        a: q(&format!("{d}.in_proj_a"))?,
197                        b: q(&format!("{d}.in_proj_b"))?,
198                        conv: t(&format!("{d}.conv1d.weight"))?,
199                        a_log: t(&format!("{d}.A_log"))?,
200                        dt_bias: t(&format!("{d}.dt_bias"))?,
201                        norm: t(&format!("{d}.norm.weight"))?,
202                        o: q(&format!("{d}.out_proj"))?,
203                        state: ctx.new_buffer(state_len * 4)?,
204                        hist: ctx.new_buffer(hist_len * 4)?,
205                        mid: ctx.new_buffer(MAX_SNAP * state_len * 4)?,
206                        mid_hist: ctx.new_buffer(MAX_SNAP * hist_len * 4)?,
207                    })
208                } else {
209                    let a = format!("{lp}.self_attn");
210
211                    Mix::Attn(Attn {
212                        q: q(&format!("{a}.q_proj"))?,
213                        k: q(&format!("{a}.k_proj"))?,
214                        v: q(&format!("{a}.v_proj"))?,
215                        o: q(&format!("{a}.o_proj"))?,
216                        qn: t(&format!("{a}.q_norm.weight"))?,
217                        kn: t(&format!("{a}.k_norm.weight"))?,
218                        kc: ctx.new_buffer(kv_side)?,
219                        vc: ctx.new_buffer(kv_side)?,
220                        iqk: q(&format!("{a}.indexer.index_qk_proj"))?,
221                        iqn: t(&format!("{a}.indexer.q_layernorm.weight"))?,
222                        ikn: t(&format!("{a}.indexer.k_layernorm.weight"))?,
223                        ikc: ctx.new_buffer(max_t * c.indexer_head_dim * 4)?,
224                        blk: ctx.new_buffer(
225                            (max_t / c.indexer_compress_ratio + 1) * c.indexer_head_dim * 4,
226                        )?,
227                    })
228                };
229                let moe = Moe {
230                    router: t(&format!("{lp}.mlp.gate.weight"))?,
231                    shared_gate: t(&format!("{lp}.mlp.shared_expert_gate.weight"))?,
232                    sg: q(&format!("{lp}.mlp.shared_expert.gate_proj"))?,
233                    su: q(&format!("{lp}.mlp.shared_expert.up_proj"))?,
234                    sd: q(&format!("{lp}.mlp.shared_expert.down_proj"))?,
235                    record_layer,
236                };
237                let ple = if with_ple {
238                    let pp = format!("{lp}.ple");
239                    let ngram = p.ngram_metadata(&format!("{pp}.ple_embedding"))?;
240                    let kernel = p.shape(&format!("{pp}.conv1d.weight"))?[2] as u32;
241                    let span = (kernel - 1) * c.ngram_size as u32;
242
243                    Some(Ple {
244                        key: q(&format!("{pp}.key_proj"))?,
245                        value: q(&format!("{pp}.value_proj"))?,
246                        norm_key: t(&format!("{pp}.norm_key.weight"))?,
247                        norm_query: t(&format!("{pp}.norm_query.weight"))?,
248                        norm_conv: t(&format!("{pp}.norm_conv.weight"))?,
249                        conv: t(&format!("{pp}.conv1d.weight"))?,
250                        kernel,
251                        dilation: c.ngram_size as u32,
252                        span,
253                        hist: ctx.new_buffer(span as usize * hh * 4)?,
254                        e: ctx.new_buffer(MAX_NB * c.ple_embed_dim * 4)?,
255                        multipliers: ngram.multipliers,
256                        head_offsets: ngram.head_offsets,
257                        head_sizes: ngram.head_sizes,
258                    })
259                } else {
260                    None
261                };
262
263                Ok(GLayer {
264                    attn_hc: hc(&format!("{lp}.attn_hyper_connection"), true)?,
265                    mlp_hc: hc(&format!("{lp}.mlp_hyper_connection"), true)?,
266                    mix,
267                    moe,
268                    ple,
269                })
270            };
271
272        let m = "language_model.model";
273        let mut layers = Vec::with_capacity(c.num_hidden_layers);
274
275        for l in 0..c.num_hidden_layers {
276            layers.push(load_layer(
277                &format!("{m}.layers.{l}"),
278                c.is_linear(l),
279                p.expert_layer(l),
280                c.ple_layer() == Some(l),
281            )?);
282        }
283
284        let mtp = if c.mtp_num_hidden_layers > 0 && options.drafts > 0 {
285            let layer = load_layer("mtp.layers.0", false, p.mtp_expert_layer(0)?, false)?;
286
287            anyhow::ensure!(
288                matches!(layer.mix, Mix::Attn(_)),
289                "MTP layer must be full attention"
290            );
291
292            Some(Mtp {
293                enorm: t("mtp.pre_fc_norm_embedding.weight")?,
294                hnorm: t("mtp.pre_fc_norm_hidden.weight")?,
295                fc_e: q("mtp.fc_embedding")?,
296                fc_h: q("mtp.fc_hidden")?,
297                layer,
298                mixer: hc("mtp.hyper_connection_mixer", false)?,
299            })
300        } else {
301            None
302        };
303
304        let max_blk = max_t.div_ceil(ATTN_TB).max(ATTN_MAX_WG);
305        let vocab = c.vocab_size;
306        let inter = c.moe_intermediate_size;
307
308        anyhow::ensure!(
309            c.shared_expert_intermediate_size == inter,
310            "shared expert width differs from routed experts"
311        );
312
313        // Reduced checkpoints can retain wide attention/DeltaNet heads even
314        // when their residual stream is small. Size scratch for every projection.
315        let max_in = [
316            hh,
317            c.ple_embed_dim,
318            2 * h,
319            c.hc_lowrank,
320            v_dim,
321            c.num_attention_heads * c.head_dim,
322            inter,
323        ]
324        .into_iter()
325        .max()
326        .unwrap();
327        let half_set = |rows: usize, in_dim: usize| -> Result<HalfSet> {
328            Ok(HalfSet {
329                xe: ctx.new_buffer(rows * in_dim / 2 * 2)?,
330                xo: ctx.new_buffer(rows * in_dim / 2 * 2)?,
331                xsum: ctx.new_buffer(rows * in_dim / 32 * 4)?,
332            })
333        };
334        let hc_bufs = || -> Result<HcBufs> {
335            Ok(HcBufs {
336                normed: ctx.new_buffer(MAX_NB * hh * 4)?,
337                d: ctx.new_buffer(MAX_NB * c.hc_lowrank * 4)?,
338                u: ctx.new_buffer(MAX_NB * hh * 4)?,
339                mixed: ctx.new_buffer(MAX_NB * h * 4)?,
340                inj: ctx.new_buffer(MAX_NB * c.hc_count * 4)?,
341                h1: half_set(MAX_NB, max_in)?,
342                h2: half_set(MAX_NB, max_in)?,
343            })
344        };
345        let n_u_max = c.num_experts_per_tok * MAX_NB + 1;
346        let scratch = Scratch {
347            ids: ctx.new_buffer(64 * 4)?,
348            e: ctx.new_buffer(MAX_NB * h * 4)?,
349            hyper: ctx.new_buffer(MAX_NB * hh * 4)?,
350            hc: hc_bufs()?,
351            la: hc_bufs()?,
352            mix_out: ctx.new_buffer(MAX_NB * h * 4)?,
353            qg: ctx.new_buffer(MAX_NB * c.num_attention_heads * c.head_dim * 2 * 4)?,
354            k: ctx.new_buffer(MAX_NB * kv_row * 4)?,
355            v: ctx.new_buffer(MAX_NB * kv_row * 4)?,
356            attn_out: ctx.new_buffer(MAX_NB * c.num_attention_heads * c.head_dim * 4)?,
357            attn_parts: ctx.new_buffer(c.num_attention_heads * max_blk * (2 + c.head_dim) * 4)?,
358            qkv: ctx.new_buffer(MAX_NB * conv_dim * 4)?,
359            z: ctx.new_buffer(MAX_NB * v_dim * 4)?,
360            a: ctx.new_buffer(MAX_NB * c.linear_num_value_heads * 4)?,
361            b: ctx.new_buffer(MAX_NB * c.linear_num_value_heads * 4)?,
362            kqn: ctx.new_buffer(MAX_NB * 2 * c.linear_num_key_heads * c.linear_key_head_dim * 4)?,
363            gbuf: ctx.new_buffer(MAX_NB * c.linear_num_value_heads * 2 * 4)?,
364            delta_y: ctx.new_buffer(MAX_NB * v_dim * 4)?,
365            router: ctx.new_buffer(MAX_NB * c.num_experts * 4)?,
366            topk_idx: ctx.new_buffer(MAX_NB * c.num_experts_per_tok * 4)?,
367            topk_w: ctx.new_buffer(MAX_NB * c.num_experts_per_tok * 4)?,
368            la_router: ctx.new_buffer(MAX_NB * c.num_experts * 4)?,
369            la_idx: ctx.new_buffer(MAX_NB * 32 * 4)?,
370            la_w: ctx.new_buffer(MAX_NB * 32 * 4)?,
371            gate_e: ctx.new_buffer(n_u_max * MAX_NB * 2 * inter * 4)?,
372            hx: half_set(n_u_max * MAX_NB, inter)?,
373            y_e: ctx.new_buffer(n_u_max * MAX_NB * h * 4)?,
374            moe_out: ctx.new_buffer(MAX_NB * h * 4)?,
375            ple_key: ctx.new_buffer(MAX_NB * hh * 4)?,
376            ple_keyn: ctx.new_buffer(MAX_NB * hh * 4)?,
377            ple_value: ctx.new_buffer(MAX_NB * h * 4)?,
378            ple_query: ctx.new_buffer(MAX_NB * hh * 4)?,
379            ple_gated: ctx.new_buffer(MAX_NB * hh * 4)?,
380            ple_gvn: ctx.new_buffer(MAX_NB * hh * 4)?,
381            ple_out: ctx.new_buffer(MAX_NB * hh * 4)?,
382            logits: ctx.new_buffer(MAX_NB * vocab * 4)?,
383            mtp_logits: ctx.new_buffer(MAX_NB * vocab * 4)?,
384            mtp_logits2: ctx.new_buffer(vocab * 4)?,
385            partials: ctx.new_buffer(ARGMAX_TGS * 8)?,
386            mtp_hyper: ctx.new_buffer(MAX_NB * hh * 4)?,
387            fe: ctx.new_buffer(MAX_NB * h * 4)?,
388            fh: ctx.new_buffer(MAX_NB * hh * 4)?,
389            iqk: ctx.new_buffer(MAX_NB * (c.indexer_n_heads + 1) * c.indexer_head_dim * 4)?,
390            iq: ctx.new_buffer(MAX_NB * c.indexer_n_heads * c.indexer_head_dim * 4)?,
391            bscore: ctx.new_buffer(MAX_NB * (max_t / c.indexer_compress_ratio + 1) * 4)?,
392            vis: ctx.new_buffer(MAX_NB * (c.indexer_budget + c.indexer_compress_ratio) * 4)?,
393            nvis: ctx.new_buffer(MAX_NB * 4)?,
394            vmask: ctx
395                .new_buffer(MAX_NB * (max_t / c.indexer_compress_ratio + 1).div_ceil(32) * 4)?,
396        };
397        let n_records = p.manifest.experts.layers * p.manifest.experts.experts;
398        let stride = p.manifest.experts.record_stride as usize;
399        let n_rows = c.num_hidden_layers + 1;
400        let slot_tab = ctx.new_buffer(n_rows * SLOT_STRIDE * 8)?;
401        let wmap = ctx.new_buffer(n_rows * MAX_NB * SLOT_STRIDE * 4)?;
402        let ring = ctx.new_buffer(prefill::RING * stride)?;
403
404        // A capped diagnostic trunk does not retain the normal adjacent-pass layout.
405        let (phase_timer, gpu_timing) = super::phases::PhaseTimer::initialize(
406            &ctx,
407            p.manifest.experts.layers,
408            layer_cap() >= layers.len(),
409        );
410        let mut activity = ExpertActivity::new(&p.manifest.experts);
411        activity.gpu_timestamps_available = phase_timer.is_some();
412        activity.gpu_timing = gpu_timing;
413
414        let prefill_bytes =
415            prefill::allocation::scratch_bytes(c, max_t, prefill_rows.min(max_t).max(1), false)?;
416        let memory = budget::PoolMemory {
417            device: ctx.device.recommendedMaxWorkingSetSize() as usize,
418            fixed: ctx.device.currentAllocatedSize(),
419            host: reserve,
420            allocation_limit,
421            prefill: prefill_bytes,
422        };
423        let mut pool_bytes = memory.bytes(options.pool_gb)?;
424        let shrink = matches!(options.pool_gb, PoolBudget::Max);
425
426        // CHERENKOV_POOL=set puts the pool in a residency set over the
427        // file mapping (page cache as a second tier); the default copies
428        // records into a wired buffer.
429        let copy = std::env::var("CHERENKOV_POOL").as_deref() != Ok("set");
430        // Only one low-bit layout is attached; --miss-experts can differ
431        // from the resident precision only when the latter is four bits.
432        let all_bits = options.experts;
433        let miss_bits = options.miss_bits();
434
435        anyhow::ensure!(
436            copy || (all_bits == 4 && miss_bits == 4),
437            "the residency-set pool supports only 4-bit experts"
438        );
439
440        let low_bits = if all_bits == 3 || all_bits == 2 {
441            all_bits
442        } else {
443            miss_bits
444        };
445        let all_low_bits = all_bits == 3 || all_bits == 2;
446        let mut low_bit_store = None;
447
448        if (low_bits == 3 || low_bits == 2) && copy {
449            // Build the store from experts.bin if it is missing, the wrong
450            // size, or in an older layout. This pass over the 4-bit source
451            // builds only the selected precision; other cached stores stay.
452            let l = crate::qwen4_exp::lowbit::ensure_with_policy(
453                &p.dir,
454                &p.manifest.experts,
455                low_bits,
456                options.repack,
457                options.build_missing_store,
458            )?;
459            let low_path = p.dir.join(format!("experts{low_bits}.bin"));
460            let low_file = std::fs::File::open(&low_path)
461                .with_context(|| format!("opening {}", low_path.display()))?;
462            let st = l;
463
464            anyhow::ensure!(
465                st.stride > 0 && st.stride <= stride,
466                "{low_bits}-bit record stride {} does not fit a 4-bit slot",
467                st.stride
468            );
469
470            {
471                use std::os::unix::io::AsRawFd as _;
472
473                unsafe { libc::fcntl(low_file.as_raw_fd(), libc::F_NOCACHE, 1) };
474            }
475
476            low_bit_store = Some((st, low_file, low_path));
477        }
478
479        // Every record low-bit means a smaller slot and so more of them.
480        let slot_stride = match (&low_bit_store, all_low_bits) {
481            (Some((st, _, _)), true) => st.stride,
482            _ => stride,
483        };
484        let mut res = loop {
485            let minimum = 64.min(n_records);
486            let slots = (pool_bytes / slot_stride).min(n_records);
487
488            ensure!(
489                slots >= minimum,
490                "expert pool budget cannot hold {minimum} records"
491            );
492
493            match residency::Pool::new(
494                &ctx,
495                p.experts.as_ptr(),
496                stride,
497                slot_stride,
498                n_records,
499                slots,
500                copy,
501            ) {
502                Ok(res) => break res,
503                Err(_) if shrink && pool_bytes > BYTES_PER_GB / 2 => {
504                    pool_bytes -= BYTES_PER_GB / 2;
505                }
506                Err(e) => {
507                    return Err(e.context(format!("allocating a {pool_bytes}-byte expert pool")));
508                }
509            }
510        };
511        // Misses read through the page cache (their pages are what the
512        // residency set pins; what it drops stays cached while memory
513        // allows); the prefill ring streams past the cache.
514        let pool_file = std::fs::File::open(p.dir.join("experts.bin"))?;
515        let pool_file_nocache = std::fs::File::open(p.dir.join("experts.bin"))?;
516        let low_bit_store = low_bit_store.map(|(st, low_file, low_path)| {
517            res.set_low_bit_store(
518                low_file,
519                st.stride,
520                if all_low_bits {
521                    if st.bits == 2 { 2 } else { 1 }
522                } else {
523                    0
524                },
525            );
526            eprintln!(
527                "{}-bit store ({}): {} bytes per record, {:.0}% of 4-bit, for {}",
528                st.bits,
529                low_path.display(),
530                st.stride,
531                100.0 * st.stride as f64 / stride as f64,
532                if all_low_bits {
533                    "every record"
534                } else {
535                    "synchronous misses"
536                }
537            );
538
539            st
540        });
541
542        {
543            use std::os::unix::io::AsRawFd as _;
544
545            unsafe { libc::fcntl(pool_file_nocache.as_raw_fd(), libc::F_NOCACHE, 1) };
546        }
547
548        let event_res = {
549            use objc2_metal::MTLDevice as _;
550
551            ctx.device.newSharedEvent().context("shared event")?
552        };
553        let (event, event_cpu) = {
554            use objc2_metal::MTLDevice as _;
555
556            (
557                ctx.device.newSharedEvent().context("shared event")?,
558                ctx.device.newSharedEvent().context("shared event")?,
559            )
560        };
561
562        Ok(Gpu {
563            embed: q(&format!("{m}.embed_tokens"))?,
564            final_mixer: hc(&format!("{m}.hyper_connection_mixer"), false)?,
565            lm_head: q("language_model.lm_head")?,
566            p,
567            ctx,
568            pipes,
569            dense,
570            res,
571            activity,
572            read_tracker: super::activity::reads::ReadTracker::new(p.manifest.experts.layers),
573            activity_started: std::time::Instant::now(),
574            phase_timer,
575            pool_file,
576            pool_file_nocache,
577            low_bit_store,
578            all_low_bits,
579            ngram_file: std::fs::File::open(p.dir.join("ngram.bin"))?,
580            ngram_prefetch: std::cell::RefCell::new(None),
581            step_no: 0,
582            pending: None,
583            event,
584            event_base: 0,
585            event_cpu,
586            event_res,
587            event_cpu_base: 0,
588            slot_tab,
589            wmap,
590            layers,
591            mtp,
592            scratch,
593            max_t,
594            pos: 0,
595            tokens: Vec::new(),
596            batch_pos: 0,
597            batch_nb: 0,
598            batch_snap: false,
599            mtp_len: 0,
600            last_experts: Vec::new(),
601            expert_history: Vec::new(),
602            route_history: Vec::new(),
603            dump_states: std::env::var_os("CHERENKOV_DUMP_STATES").is_some(),
604            state_history: Vec::new(),
605            last_states: Vec::new(),
606            last_routes: Vec::new(),
607            last_route_w: Vec::new(),
608            last_miss: Vec::new(),
609            ngram_gather_s: std::cell::Cell::new(0.0),
610            ngram_ms: Vec::new(),
611            step_ms: Vec::new(),
612            gpu_ms: Vec::new(),
613            io_ms: Vec::new(),
614            step_set_s: 0.0,
615            step_read_s: 0.0,
616            step_warm: 0,
617            warm: Vec::new(),
618            folded_mtp: None,
619            set_ms: Vec::new(),
620            read_ms: Vec::new(),
621            misses: Vec::new(),
622            miss_bytes: Vec::new(),
623            step_misses: 0,
624            step_miss_bytes: 0,
625            lookahead_hit: Vec::new(),
626            lookahead_issued: Vec::new(),
627            la_log: Vec::new(),
628            la_pending: Vec::new(),
629            log_la: std::env::var_os("CHERENKOV_DUMP_LA").is_some(),
630            cut_w: options.cut_weak,
631            step_cut: 0,
632            cut: Vec::new(),
633            inflight: Vec::new(),
634            lookahead: std::env::var("CHERENKOV_LOOKAHEAD").as_deref() != Ok("0"),
635            // Blocking on the event measured ~30 ms/step slower (48 wake-ups)
636            // with no thermal benefit; spinning is the default.
637            spin_wait: std::env::var("CHERENKOV_SPIN").as_deref() != Ok("0"),
638            fake_experts: std::env::var("CHERENKOV_FAKE").as_deref() == Ok("experts"),
639            skip: std::env::var("CHERENKOV_SKIP")
640                .map(|v| v.split(',').map(|s| s.trim().to_string()).collect())
641                .unwrap_or_default(),
642            dispatches: Vec::new(),
643            gpu_idle_ms: Vec::new(),
644            rows: Vec::new(),
645            mtp_ms: Vec::new(),
646            dispatch_count: std::cell::Cell::new(0),
647            pf: None,
648            ring,
649            prefill_reserved_bytes: prefill_bytes,
650            prefill_stats: Vec::new(),
651        })
652    }
653}
cherenkov/qwen4_exp/gpu/load.rs

cherenkov/qwen4_exp/gpu/
load.rs
