metrale_model_layers/layers/ngram_embed/
embed.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: The GPU n-gram embedding: the base word row plus the
4//! K*(N-1) hashed table lookups.
5//!
6//! Owner: model-layers (n-gram embedding).
7//! Invariants:
8//! - `tables` and `projs` each hold `dims.num_tables()` entries (`new`
9//!   refuses otherwise).
10
11use anyhow::{Context, Result};
12use metrale_gpu_runtime::gpu::{DevicePtr, GpuBackend, KernelHandle};
13
14use super::ids::ngram_ids;
15use super::{NgramDims, NgramTable};
16use crate::weight_map::DenseWeight;
17
18/// 2026-09-25: GPU n-gram embedding from `batched_embed` gathers,
19/// `dense_gemm_bf16_pipelined` projections and `bf16_scaled_add`
20/// accumulation: per call 1 + T gathers, T GEMMs and T + 1 scaled adds,
21/// T = `num_tables()`.
22pub struct NgramEmbedding {
23    pub dims: NgramDims,
24    /// 2026-09-25: Base word embedding, `[vocab, hidden]` BF16.
25    pub word: DenseWeight,
26    /// 2026-09-25: The K*(N-1) lookup tables in index order
27    /// `(ngram-2)*K + split`, each `[table_rows(index), table_dim]`.
28    pub tables: Vec<NgramTable>,
29    /// 2026-09-25: Per-table projections, `[hidden, table_dim]` BF16.
30    pub projs: Vec<DenseWeight>,
31
32    pub(super) batched_embed_k: KernelHandle,
33    pub(super) batched_embed_fp8_k: KernelHandle,
34    pub(super) gemm_k: KernelHandle,
35    pub(super) scaled_add_k: KernelHandle,
36
37    /// 2026-09-25: Device staging: ids `[max_tokens]` u32 (reused per
38    /// table), gathered rows `[max_tokens, table_dim]` BF16, projected rows
39    /// `[max_tokens, hidden]` BF16.
40    pub(super) ids_dev: DevicePtr,
41    pub(super) gather_buf: DevicePtr,
42    pub(super) proj_buf: DevicePtr,
43    pub(super) max_tokens: usize,
44}
45
46impl NgramEmbedding {
47    pub fn new(
48        dims: NgramDims,
49        word: DenseWeight,
50        tables: Vec<NgramTable>,
51        projs: Vec<DenseWeight>,
52        max_tokens: usize,
53        gpu: &dyn GpuBackend,
54    ) -> Result<Self> {
55        anyhow::ensure!(tables.len() == dims.num_tables(), "table count");
56        anyhow::ensure!(projs.len() == dims.num_tables(), "proj count");
57        let td = dims.table_dim();
58        Ok(Self {
59            dims,
60            word,
61            tables,
62            projs,
63            batched_embed_k: gpu.kernel("embed_from_argmax", "batched_embed")?,
64            batched_embed_fp8_k: gpu.kernel("embed_from_argmax", "batched_embed_fp8")?,
65            gemm_k: gpu.kernel("gemm", "dense_gemm_bf16_pipelined")?,
66            scaled_add_k: gpu.kernel("residual_add", "bf16_scaled_add")?,
67            ids_dev: gpu.alloc(max_tokens * 4)?,
68            gather_buf: gpu.alloc(max_tokens * td * 2)?,
69            proj_buf: gpu.alloc(max_tokens * dims.hidden_size * 2)?,
70            max_tokens,
71        })
72    }
73
74    /// 2026-09-25: Embedding of the last `seq_len` tokens of `ctx_tokens`
75    /// (earlier context tokens, then the new ones), written to `out` as
76    /// `[seq_len, hidden]` BF16. Errors when `seq_len` exceeds the staging
77    /// or `ctx_tokens`, or an id does not fit u32.
78    pub fn embed(
79        &mut self,
80        ctx_tokens: &[u32],
81        seq_len: usize,
82        out: DevicePtr,
83        gpu: &dyn GpuBackend,
84        stream: u64,
85    ) -> Result<()> {
86        use crate::layers::ops;
87        anyhow::ensure!(
88            seq_len <= self.max_tokens,
89            "ngram embed: seq_len over staging"
90        );
91        anyhow::ensure!(seq_len <= ctx_tokens.len(), "ngram embed: seq_len over ctx");
92        let h = self.dims.hidden_size;
93        let td = self.dims.table_dim();
94        let inv_scale = 1.0f32 / (1 + self.dims.num_tables()) as f32;
95
96        // 2026-09-25: Zero `out`, then add the base row and each table's
97        // projection scaled by 1 / (1 + T).
98        gpu.memset(out, 0, seq_len * h * 2)?;
99
100        // 2026-09-25: Base word rows for the new tokens only.
101        let new_tokens = &ctx_tokens[ctx_tokens.len() - seq_len..];
102        let ids_bytes: Vec<u8> = new_tokens.iter().flat_map(|t| t.to_le_bytes()).collect();
103        gpu.copy_h2d_async(&ids_bytes, self.ids_dev, stream)?;
104        ops::batched_embed(
105            gpu,
106            self.batched_embed_k,
107            self.ids_dev,
108            self.word.weight,
109            self.proj_buf,
110            seq_len as u32,
111            h as u32,
112            stream,
113        )?;
114        ops::scaled_add(
115            gpu,
116            self.scaled_add_k,
117            out,
118            self.proj_buf,
119            inv_scale,
120            (seq_len * h) as u32,
121            stream,
122        )?;
123
124        // 2026-09-25: Host-side hash ids over the full context, then one
125        // table at a time.
126        let all_ids = ngram_ids(&self.dims, ctx_tokens);
127        for (index, ids) in all_ids.iter().enumerate() {
128            let tail = &ids[ids.len() - seq_len..];
129            let id_bytes: Vec<u8> = tail
130                .iter()
131                .map(|&v| u32::try_from(v).context("ngram id exceeds u32"))
132                .collect::<Result<Vec<u32>>>()?
133                .iter()
134                .flat_map(|v| v.to_le_bytes())
135                .collect();
136            gpu.copy_h2d_async(&id_bytes, self.ids_dev, stream)?;
137            // 2026-09-25: An NVMe-backed table maps the ids to arena slots
138            // (`resolve`), which replace the ids before the gather.
139            #[cfg(feature = "cuda")]
140            if let NgramTable::Cached(cache) = &mut self.tables[index] {
141                let mut slots: Vec<u32> = Vec::with_capacity(seq_len);
142                cache.resolve(tail, &mut slots)?;
143                let slot_bytes: Vec<u8> =
144                    slots.iter().flat_map(|v: &u32| v.to_le_bytes()).collect();
145                gpu.copy_h2d_async(&slot_bytes, self.ids_dev, stream)?;
146                let table = DevicePtr(cache.table_dev_va()?);
147                match cache.scale_dev_va()? {
148                    Some(sc) => ops::batched_embed_fp8(
149                        gpu,
150                        self.batched_embed_fp8_k,
151                        self.ids_dev,
152                        table,
153                        DevicePtr(sc),
154                        self.gather_buf,
155                        seq_len as u32,
156                        td as u32,
157                        stream,
158                    )?,
159                    None => ops::batched_embed(
160                        gpu,
161                        self.batched_embed_k,
162                        self.ids_dev,
163                        table,
164                        self.gather_buf,
165                        seq_len as u32,
166                        td as u32,
167                        stream,
168                    )?,
169                }
170                cache.end_batch();
171                ops::dense_gemm_bf16_pipelined(
172                    gpu,
173                    self.gemm_k,
174                    self.gather_buf,
175                    &self.projs[index],
176                    self.proj_buf,
177                    seq_len as u32,
178                    h as u32,
179                    td as u32,
180                    stream,
181                )?;
182                ops::scaled_add(
183                    gpu,
184                    self.scaled_add_k,
185                    out,
186                    self.proj_buf,
187                    inv_scale,
188                    (seq_len * h) as u32,
189                    stream,
190                )?;
191                continue;
192            }
193            match &self.tables[index] {
194                NgramTable::Bf16(w) => ops::batched_embed(
195                    gpu,
196                    self.batched_embed_k,
197                    self.ids_dev,
198                    w.weight,
199                    self.gather_buf,
200                    seq_len as u32,
201                    td as u32,
202                    stream,
203                )?,
204                NgramTable::Fp8(w) => ops::batched_embed_fp8(
205                    gpu,
206                    self.batched_embed_fp8_k,
207                    self.ids_dev,
208                    w.weight,
209                    w.row_scale,
210                    self.gather_buf,
211                    seq_len as u32,
212                    td as u32,
213                    stream,
214                )?,
215                #[cfg(feature = "cuda")]
216                NgramTable::Cached(_) => unreachable!("resolved above"),
217            }
218            ops::dense_gemm_bf16_pipelined(
219                gpu,
220                self.gemm_k,
221                self.gather_buf,
222                &self.projs[index],
223                self.proj_buf,
224                seq_len as u32,
225                h as u32,
226                td as u32,
227                stream,
228            )?;
229            ops::scaled_add(
230                gpu,
231                self.scaled_add_k,
232                out,
233                self.proj_buf,
234                inv_scale,
235                (seq_len * h) as u32,
236                stream,
237            )?;
238        }
239        Ok(())
240    }
241}
metrale_model_layers/layers/ngram_embed/embed.rs

metrale_model_layers/layers/ngram_embed/
embed.rs
