hf2q 0.1.3

Pure Rust CLI for converting HuggingFace models to hardware-optimized formats and serving them over an OpenAI-compatible API on Apple Silicon
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232

//! Native DeepSeek-V4 Hyper-Connection output head and vocabulary logits.

use anyhow::{bail, Context, Result};
use mlx_native::ops::argmax::dispatch_argmax_f32;
use mlx_native::ops::deepseek_hyper_connection::{dispatch_hc_head_weights, dispatch_hc_pre};
use mlx_native::{DType, MlxBuffer};

use super::forward_support::{alloc, raw_matmul, rms_params};
use super::Deepseek4Model;

impl Deepseek4Model {
    /// Return a zero-copy one-row view for prompt-to-decode handoff.
    pub fn last_token_state(&self, state: &MlxBuffer) -> Result<MlxBuffer> {
        let hidden = self.cfg.hidden_size as usize;
        let hc = self.cfg.hyper_connection_count as usize;
        if state.dtype() != DType::F32
            || state.shape().len() != 3
            || state.shape()[0] == 0
            || state.shape()[1..] != [hc, hidden]
        {
            bail!(
                "DeepSeek-V4 output state must be F32 [tokens, {hc}, {hidden}], got {} {:?}",
                state.dtype(),
                state.shape()
            );
        }
        let row_elements = hc
            .checked_mul(hidden)
            .context("DeepSeek-V4 output row size overflow")?;
        let row_offset = (state.shape()[0] - 1)
            .checked_mul(row_elements)
            .and_then(|elements| elements.checked_mul(DType::F32.size_of()))
            .context("DeepSeek-V4 output row offset overflow")?;
        state
            .slice_view(
                u64::try_from(row_offset).context("DeepSeek-V4 output row offset exceeds u64")?,
                row_elements,
            )
            .with_shape(vec![1, hc, hidden])
            .context("view DeepSeek-V4 last prompt state")
    }

    /// Collapse completed verifier states and compute vocabulary logits.
    ///
    /// The input is the exact `[tokens, 4, hidden]` F32 state emitted by the
    /// final verifier FFN. This graph does not mutate request cache state, so
    /// callers may safely retry it after a command-buffer failure.
    pub fn forward_logits(&mut self, state: &MlxBuffer) -> Result<MlxBuffer> {
        let hidden = self.cfg.hidden_size as usize;
        let hc = self.cfg.hyper_connection_count as usize;
        let hc_hidden = hc
            .checked_mul(hidden)
            .context("DeepSeek-V4 output HC width overflow")?;
        let vocab = self.cfg.vocab_size as usize;
        if state.dtype() != DType::F32
            || state.shape().len() != 3
            || state.shape()[1..] != [hc, hidden]
        {
            bail!(
                "DeepSeek-V4 output state must be F32 [tokens, {hc}, {hidden}], got {} {:?}",
                state.dtype(),
                state.shape()
            );
        }
        let tokens = state.shape()[0];
        if tokens == 0 {
            bail!("DeepSeek-V4 output state must contain at least one token");
        }
        let tokens_u32 =
            u32::try_from(tokens).context("DeepSeek-V4 output token count exceeds u32")?;
        let hidden_u32 =
            u32::try_from(hidden).context("DeepSeek-V4 output hidden size exceeds u32")?;
        let hc_hidden_u32 =
            u32::try_from(hc_hidden).context("DeepSeek-V4 output HC width exceeds u32")?;

        let hc_fn = self.weights.raw_matrix_ref("output_hc_fn.weight")?;
        let hc_scale = self.weights.f32_state("output_hc_scale.weight")?;
        let hc_base = self.weights.f32_state("output_hc_base.weight")?;
        let output_norm_weight = self.weights.f32_state("output_norm.weight")?;
        let output_weight = self.weights.raw_matrix_ref("output.weight")?;

        let device = self.ctx.device().clone();
        let state_flat = state.with_shape(vec![tokens, hc_hidden])?;
        let state_norm = alloc(
            &device,
            DType::F32,
            vec![tokens, hc_hidden],
            "output HC normalized state",
        )?;
        let mixes = alloc(&device, DType::F32, vec![tokens, hc], "output HC mixes")?;
        let pre = alloc(&device, DType::F32, vec![tokens, hc], "output HC weights")?;
        let collapsed = alloc(
            &device,
            DType::F32,
            vec![tokens, hidden],
            "output HC collapsed state",
        )?;
        let normalized = alloc(
            &device,
            DType::F32,
            vec![tokens, hidden],
            "output normalized state",
        )?;
        let logits = alloc(
            &device,
            DType::F32,
            vec![tokens, vocab],
            "vocabulary logits",
        )?;
        let hc_params = rms_params(
            &device,
            self.cfg.rms_norm_eps,
            hc_hidden,
            "output HC RMS params",
        )?;
        let hidden_params =
            rms_params(&device, self.cfg.rms_norm_eps, hidden, "output RMS params")?;

        let (executor, registry) = self.ctx.split();
        let mut session = executor.begin().context("begin DeepSeek-V4 output head")?;
        session.barrier_between(&[&state_flat], &[&state_norm]);
        session.rms_norm_no_scale_f32(
            registry,
            device.metal_device(),
            &state_flat,
            &state_norm,
            &hc_params,
            tokens_u32,
            hc_hidden_u32,
        )?;
        raw_matmul(
            &mut session,
            registry,
            &device,
            &state_norm,
            &hc_fn,
            &mixes,
            tokens,
            hc,
            hc_hidden,
            "output HC function",
        )?;
        session.barrier_between(&[&mixes, hc_scale, hc_base], &[&pre]);
        dispatch_hc_head_weights(
            session.encoder_mut(),
            registry,
            &device,
            &mixes,
            hc_scale,
            hc_base,
            &pre,
            tokens_u32,
        )?;
        session.barrier_between(&[state, &pre], &[&collapsed]);
        dispatch_hc_pre(
            session.encoder_mut(),
            registry,
            &device,
            state,
            &pre,
            &collapsed,
            tokens_u32,
            hidden_u32,
        )?;
        session.barrier_between(&[&collapsed, output_norm_weight], &[&normalized]);
        session.rms_norm(
            registry,
            device.metal_device(),
            &collapsed,
            output_norm_weight,
            &normalized,
            &hidden_params,
            tokens_u32,
            hidden_u32,
        )?;
        raw_matmul(
            &mut session,
            registry,
            &device,
            &normalized,
            &output_weight,
            &logits,
            tokens,
            vocab,
            hidden,
            "vocabulary projection",
        )?;
        session
            .finish()
            .context("execute DeepSeek-V4 output head")?;
        Ok(logits)
    }

    /// Select the first maximum from one vocabulary-logit row on Metal.
    pub fn greedy_token(&mut self, logits: &MlxBuffer) -> Result<u32> {
        let vocab = self.cfg.vocab_size as usize;
        if logits.dtype() != DType::F32 || logits.shape() != [1, vocab] {
            bail!(
                "DeepSeek-V4 greedy logits must be F32 [1, {vocab}], got {} {:?}",
                logits.dtype(),
                logits.shape()
            );
        }
        let vocab_u32 = u32::try_from(vocab).context("DeepSeek-V4 vocabulary size exceeds u32")?;
        let device = self.ctx.device().clone();
        let out_index = alloc(&device, DType::U32, vec![1], "greedy token index")?;
        let out_value = alloc(&device, DType::F32, vec![1], "greedy token value")?;
        let mut params = alloc(&device, DType::U32, vec![1], "greedy token params")?;
        params.as_mut_slice::<u32>()?[0] = vocab_u32;

        let (executor, registry) = self.ctx.split();
        let mut session = executor
            .begin()
            .context("begin DeepSeek-V4 greedy argmax")?;
        session.barrier_between(&[logits, &params], &[&out_index, &out_value]);
        dispatch_argmax_f32(
            session.encoder_mut(),
            registry,
            device.metal_device(),
            logits,
            &out_index,
            &out_value,
            &params,
            vocab_u32,
        )?;
        session
            .finish()
            .context("execute DeepSeek-V4 greedy argmax")?;
        let token = out_index.as_slice::<u32>()?[0];
        Ok(token)
    }
}