cherenkov/qwen4_exp/gpu/prefill/
allocation.rs

1//! Prefill scratch allocation and chunk capacity.
2
3use super::*;
4use crate::metal::PAGE_SIZE;
5use std::cell::Cell;
6
7// Both sizing and allocation walk this layout, including padding and the
8// attention buffers that grow with context capacity.
9fn scratch<T>(
10    c: &crate::qwen4_exp::Qwen4ExpConfig,
11    max_t: usize,
12    rows: usize,
13    all_logits: bool,
14    allocate: impl Fn(usize) -> Result<T>,
15) -> Result<PrefillScratch<T>> {
16    let h = c.hidden_size;
17    let hh = c.hc_hidden();
18    let kv_row = c.num_key_value_heads * c.head_dim;
19    let conv_dim = 2 * c.linear_num_key_heads * c.linear_key_head_dim
20        + c.linear_num_value_heads * c.linear_value_head_dim;
21    let v_dim = c.linear_num_value_heads * c.linear_value_head_dim;
22    let inter = c.moe_intermediate_size;
23    let k = c.num_experts_per_tok;
24    let rp = rows.div_ceil(32) * 32 + 32;
25    let max_blocks = max_t / c.indexer_compress_ratio + 1;
26    let kl_pad = max_t.div_ceil(32) * 32;
27    let n_rep = c.num_attention_heads / c.num_key_value_heads;
28    let f = |n: usize| allocate(n * 4);
29    let hf = |n: usize| allocate(n * 2);
30
31    Ok(PrefillScratch {
32        allocated_bytes: 0,
33        rows: rp,
34        ids: allocate(rp * 4)?,
35        e: f(rp * h)?,
36        hyper: f(rp * hh)?,
37        normed: f(rp * hh)?,
38        d: f(rp * c.hc_lowrank)?,
39        u: f(rp * hh)?,
40        mixed: f(rp * h)?,
41        inj: f(rp * c.hc_count)?,
42        mix_out: f(rp * h)?,
43        moe_out: f(rp * h)?,
44        qg: f(rp * c.num_attention_heads * 2 * c.head_dim)?,
45        k: f(rp * kv_row)?,
46        v: f(rp * kv_row)?,
47        attn_out: f(rp * c.num_attention_heads * c.head_dim)?,
48        iqk: f(rp * (c.indexer_n_heads + 1) * c.indexer_head_dim)?,
49        iq: f(rp * c.indexer_n_heads * c.indexer_head_dim)?,
50        bscore: f(rp * max_blocks)?,
51        vis: allocate(rp * (c.indexer_budget + c.indexer_compress_ratio) * 4)?,
52        nvis: allocate(rp * 4)?,
53        vmask: allocate(rp * max_blocks.div_ceil(32) * 4)?,
54        ag_qh: hf(c.num_attention_heads * QS * c.head_dim)?,
55        ag_kh: hf(kl_pad * c.head_dim)?,
56        ag_vt: hf(c.head_dim * kl_pad)?,
57        ag_s: f(n_rep * QS * kl_pad)?,
58        ag_p: hf(n_rep * QS * kl_pad)?,
59        ag_o: f(n_rep * QS * c.head_dim)?,
60        qkv: f(rp * conv_dim)?,
61        z: f(rp * v_dim)?,
62        a: f(rp * c.linear_num_value_heads)?,
63        b: f(rp * c.linear_num_value_heads)?,
64        kqn: f(rp * 2 * c.linear_num_key_heads * c.linear_key_head_dim)?,
65        gbuf: f(rp * c.linear_num_value_heads * 2)?,
66        delta_y: f(rp * v_dim)?,
67        router: f(rp * c.num_experts)?,
68        topk_idx: allocate(rp * k * 4)?,
69        topk_w: f(rp * k)?,
70        csr_rows: allocate(rp * k * 4)?,
71        csr_w: f(rp * k)?,
72        xg: f(rp * h)?,
73        ge: f(rp * inter)?,
74        ue: f(rp * inter)?,
75        hg: f(rp * inter)?,
76        ye: f(rp * h)?,
77        mtp_hyper: f(rp * hh)?,
78        fe: f(rp * h)?,
79        fh: f(rp * hh)?,
80        logits_all: if all_logits {
81            Some(f(rp * c.vocab_size)?)
82        } else {
83            None
84        },
85    })
86}
87
88pub(in crate::qwen4_exp::gpu) fn scratch_bytes(
89    c: &crate::qwen4_exp::Qwen4ExpConfig,
90    max_t: usize,
91    rows: usize,
92    all_logits: bool,
93) -> Result<usize> {
94    let total = Cell::new(0usize);
95
96    scratch(c, max_t, rows, all_logits, |bytes| {
97        // Metal may round each resource to a page. Budget that rounding too.
98        let padded = bytes.div_ceil(PAGE_SIZE) * PAGE_SIZE;
99        let sum = total
100            .get()
101            .checked_add(padded)
102            .context("prefill scratch size overflow")?;
103
104        total.set(sum);
105
106        Ok(())
107    })?;
108
109    Ok(total.get())
110}
111
112fn rows_fit(
113    c: &crate::qwen4_exp::Qwen4ExpConfig,
114    max_t: usize,
115    available: usize,
116    all_logits: bool,
117) -> Result<usize> {
118    let (mut low, mut high) = (0, MAX_PREFILL_ROWS.min(max_t));
119
120    while low < high {
121        let mid = low + (high - low).div_ceil(2);
122
123        if scratch_bytes(c, max_t, mid, all_logits)? <= available {
124            low = mid;
125        } else {
126            high = mid - 1;
127        }
128    }
129
130    ensure!(
131        low > 0,
132        "insufficient memory for prefill scratch; reduce expert pool or context capacity"
133    );
134
135    Ok(low)
136}
137
138impl Gpu<'_> {
139    pub(super) fn pf_alloc(&self, rows: usize, all_logits: bool) -> Result<PrefillScratch> {
140        let before = self.allocated_bytes();
141        let mut buffers = scratch(&self.p.cfg, self.max_t, rows, all_logits, |bytes| {
142            self.ctx.new_buffer(bytes)
143        })?;
144        buffers.allocated_bytes = self.allocated_bytes().saturating_sub(before) as usize;
145
146        Ok(buffers)
147    }
148
149    /// Maximum chunk that fits the configured budget, or the device recommendation.
150    pub fn prefill_rows_fit(&self, all_logits: bool) -> Result<usize> {
151        use objc2_metal::MTLDevice as _;
152
153        let device_limit = self.ctx.device.recommendedMaxWorkingSetSize() as usize;
154        let limit = self.ctx.allocation_limit.get().unwrap_or(device_limit);
155        let mut used = self.ctx.device.currentAllocatedSize();
156
157        // Existing scratch is reused or released before its replacement.
158        if let Some(pf) = &self.pf {
159            used = used.saturating_sub(pf.allocated_bytes);
160        }
161
162        rows_fit(
163            &self.p.cfg,
164            self.max_t,
165            limit.saturating_sub(used),
166            all_logits,
167        )
168    }
169}
170
171#[cfg(test)]
172#[path = "../../../../tests/unit/qwen4_exp/gpu/prefill_allocation.rs"]
173mod tests;
cherenkov/qwen4_exp/gpu/prefill/allocation.rs

cherenkov/qwen4_exp/gpu/prefill/
allocation.rs
