metrale_model_layers/layers/ple/
layer.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: The PLE layer: ids -> NVMe row gather -> projections -> gate -> dilated
4//! conv -> highway add.
5//!
6//! Owner: model-layers (PLE).
7//! Invariants:
8//! - `forward` advances the token history only on the path where it hashes
9//!   the ids itself; a consumed prestage has already advanced it.
10//! - Under CUDA-graph capture, `forward` returns an error instead of taking
11//!   the D2H id readback or the un-prestaged gather.
12//!
13//! It runs on the layers `ple_layer_ids` names and adds into the
14//! `hc_mult`-wide highway before the layer's first hyper-connection site, as
15//! `Qwen4ExpTextDecoderLayer.forward` adds `self.ple(...)` before
16//! `attn_hyper_connection`.
17
18use anyhow::{Context, Result};
19use metrale_gpu_runtime::gpu::{DevicePtr, GpuBackend, KernelHandle};
20
21use super::ids::{PleIdDims, ple_ngram_ids};
22use crate::layer::ForwardContext;
23use crate::layers::ngram_embed::NgramTable;
24use crate::layers::ops;
25use crate::weight_map::DenseWeight;
26
27/// 2026-09-25: Per-sequence carry: the dilated conv's state and the token history the
28/// id hash needs. Owned by the sequence's [`crate::layer::SsmLayerState`].
29pub struct PleSeqState {
30    /// 2026-09-25: `[(k-1)*dilation, hc_mult*hidden]` FP32, device.
31    conv: DevicePtr,
32    /// 2026-09-25: The last `context_len` token ids, EOS-filled at a sequence start.
33    history: Vec<u32>,
34    /// 2026-09-25: Set by `prestage`: the n-gram table's device VA, recorded when the
35    /// step's host work (hash, fault-in, slot upload) has already run before
36    /// graph replay or capture. `forward` consumes it and enqueues kernels only.
37    prestaged_va: Option<u64>,
38    /// 2026-09-25: The last VA `prestage` staged; only `release_seq_state` clears it.
39    /// `rearm` restores it when a failed capture attempt re-runs the step
40    /// eagerly: the slots are still in `slots_dev` and history has already
41    /// advanced, so hashing again would count the token twice.
42    last_staged_va: u64,
43}
44
45pub struct PleLayer {
46    dims: PleIdDims,
47    head_dim: usize,
48    hidden: usize,
49    hc_mult: usize,
50    state_len: usize,
51    k_size: usize,
52    dilation: usize,
53    eps: f32,
54
55    key_proj: DenseWeight,
56    value_proj: DenseWeight,
57    norm_key: DenseWeight,
58    norm_query: DenseWeight,
59    norm_conv: DenseWeight,
60    conv1d: DenseWeight,
61    /// 2026-09-25: Behind a mutex because the NVMe cache resolves (and faults, and
62    /// evicts) on the forward path, which needs `&mut`, while layers are
63    /// invoked through `&self`.
64    table: std::sync::Mutex<NgramTable>,
65
66    embed_k: KernelHandle,
67    gemm_k: KernelHandle,
68    gate_k: KernelHandle,
69    conv_k: KernelHandle,
70    add_k: KernelHandle,
71
72    /// 2026-09-25: Scratch, sized once for `max_tokens`.
73    emb: DevicePtr,
74    key: DevicePtr,
75    value: DevicePtr,
76    gated: DevicePtr,
77    gated_normed: DevicePtr,
78    out: DevicePtr,
79    slots_dev: DevicePtr,
80    max_tokens: usize,
81}
82
83impl PleLayer {
84    #[allow(clippy::too_many_arguments)]
85    pub fn new(
86        dims: PleIdDims,
87        head_dim: usize,
88        hidden: usize,
89        hc_mult: usize,
90        k_size: usize,
91        dilation: usize,
92        eps: f32,
93        weights: PleWeights,
94        table: NgramTable,
95        max_tokens: usize,
96        gpu: &dyn GpuBackend,
97    ) -> Result<Self> {
98        dims.validate()?;
99        let heads = dims.ngram_heads();
100        anyhow::ensure!(
101            heads * head_dim == hidden,
102            "PLE: {heads} heads x {head_dim} dims = {} != ple_embed_dim {hidden}. \
103             The head slices are CONCATENATED (not summed as LongCat's are), so \
104             this product is the embedding width and a mismatch means the \
105             geometry is not what we think.",
106            heads * head_dim
107        );
108        let c = hc_mult * hidden;
109        let state_len = (k_size - 1) * dilation;
110        Ok(Self {
111            dims,
112            head_dim,
113            hidden,
114            hc_mult,
115            state_len,
116            k_size,
117            dilation,
118            eps,
119            key_proj: weights.key_proj,
120            value_proj: weights.value_proj,
121            norm_key: weights.norm_key,
122            norm_query: weights.norm_query,
123            norm_conv: weights.norm_conv,
124            conv1d: weights.conv1d,
125            table: std::sync::Mutex::new(table),
126            embed_k: gpu.kernel("embed_from_argmax", "batched_embed")?,
127            gemm_k: gpu.kernel("gemm", "dense_gemm_bf16_pipelined")?,
128            gate_k: gpu.kernel("ple", "ple_gate")?,
129            conv_k: gpu.kernel("ple", "ple_conv")?,
130            add_k: gpu.kernel("ple", "ple_add_highway")?,
131            emb: gpu.alloc(max_tokens * hidden * 2)?,
132            key: gpu.alloc(max_tokens * c * 2)?,
133            value: gpu.alloc(max_tokens * hidden * 2)?,
134            gated: gpu.alloc(max_tokens * c * 4)?,
135            gated_normed: gpu.alloc(max_tokens * c * 4)?,
136            out: gpu.alloc(max_tokens * c * 4)?,
137            slots_dev: gpu.alloc(max_tokens * heads * 4)?,
138            max_tokens,
139        })
140    }
141
142    /// 2026-09-25: Allocate one sequence's PLE carry (conv buffer + empty history).
143    /// The conv buffer is not initialised here; the first `forward` or
144    /// `prestage` sees the empty history and runs `reset`.
145    pub fn new_seq_state(&self, gpu: &dyn GpuBackend) -> Result<PleSeqState> {
146        Ok(PleSeqState {
147            conv: gpu.alloc(self.state_len * self.hc_mult * self.hidden * 4)?,
148            history: Vec::new(),
149            prestaged_va: None,
150            last_staged_va: 0,
151        })
152    }
153
154    // 2026-09-25: `reset`, `prestage`, `release_seq_state`, `snapshot_aux` and
155    // `restore_aux` are in `aux_state.rs`.
156
157    /// 2026-09-25: Restore the prestaged state after a failed CUDA-graph capture attempt
158    /// (the eager re-run calls `forward` again, and the first call consumed
159    /// `prestaged_va`).
160    pub fn rearm(&self, st: &mut PleSeqState) {
161        if st.last_staged_va != 0 {
162            st.prestaged_va = Some(st.last_staged_va);
163        }
164    }
165
166    /// 2026-09-25: One highway row with an explicit id: the multi-seq decode entry
167    /// (`ctx.host_token_ids` holds the whole batch; the caller slices). It
168    /// never consumes a prestage.
169    pub fn forward_row(
170        &self,
171        st: &mut PleSeqState,
172        highway_row: DevicePtr,
173        ids: &[u32],
174        ctx: &ForwardContext,
175        stream: u64,
176    ) -> Result<()> {
177        self.forward_with_ids(st, highway_row, 1, false, Some(ids), ctx, stream)
178    }
179
180    /// 2026-09-25: Inject into `highway` `[T, hc_mult*hidden]` FP32, in place.
181    /// `fresh` starts a new sequence (prefill from position 0).
182    pub fn forward(
183        &self,
184        st: &mut PleSeqState,
185        highway: DevicePtr,
186        num_tokens: usize,
187        fresh: bool,
188        ctx: &ForwardContext,
189        stream: u64,
190    ) -> Result<()> {
191        self.forward_with_ids(st, highway, num_tokens, fresh, None, ctx, stream)
192    }
193
194    #[allow(clippy::too_many_arguments)]
195    fn forward_with_ids(
196        &self,
197        st: &mut PleSeqState,
198        highway: DevicePtr,
199        num_tokens: usize,
200        fresh: bool,
201        ids_override: Option<&[u32]>,
202        ctx: &ForwardContext,
203        stream: u64,
204    ) -> Result<()> {
205        anyhow::ensure!(
206            num_tokens <= self.max_tokens,
207            "PLE: {num_tokens} tokens exceeds the {} this layer was sized for. \
208             Raise METRALE_PLE_MAX_TOKENS (costs tokens*10240*14 bytes of \
209             scratch) or lower the prefill chunk size.",
210            self.max_tokens
211        );
212        let c = self.hc_mult * self.hidden;
213        let heads = self.dims.ngram_heads();
214        let gpu = ctx.gpu;
215
216        // 2026-09-25: The ids are hashed on the host from the token ids. The
217        // caller's `ctx.host_token_ids` is preferred over reading the device
218        // copy back: the readback is a blocking copy, and it cannot run inside
219        // a CUDA-graph capture.
220        let tokens: Vec<u32> = if let Some(ov) = ids_override {
221            anyhow::ensure!(ov.len() == num_tokens, "PLE: ids_override length");
222            ov.to_vec()
223        } else if let Some(host) = ctx.host_token_ids {
224            anyhow::ensure!(
225                host.len() >= num_tokens,
226                "PLE: host_token_ids has {} ids for {num_tokens} tokens",
227                host.len()
228            );
229            host[..num_tokens].to_vec()
230        } else {
231            // 2026-09-25: Fallback for passes that did not pass the host slice.
232            // Refused under capture rather than recorded.
233            anyhow::ensure!(
234                !ctx.graph_capture,
235                "PLE: no host_token_ids and a D2H readback is \
236                 capture-unsupported; thread the host ids through this pass"
237            );
238            let tok_dev = ctx.token_ids.ok_or_else(|| {
239                anyhow::anyhow!(
240                    "PLE needs token ids (host or device); this pass staged \
241                     neither"
242                )
243            })?;
244            let mut raw = vec![0u8; num_tokens * 4];
245            gpu.copy_d2h(tok_dev, &mut raw)?;
246            raw.chunks_exact(4)
247                .map(|b| u32::from_le_bytes([b[0], b[1], b[2], b[3]]))
248                .collect()
249        };
250
251        // 2026-09-25: Only a single-token, non-fresh decode without an ids
252        // override consumes a prestage. A replayed decode step never runs this
253        // forward, so a `Some` seen by a prefill or by `forward_row` is a
254        // leftover staged for an earlier token. Consuming it would inject that
255        // token's rows and skip the history advance. `.take()` clears it in
256        // every case.
257        let prestaged = st
258            .prestaged_va
259            .take()
260            .filter(|_| num_tokens == 1 && !fresh && ids_override.is_none());
261        if fresh || st.history.len() != self.dims.context_len() {
262            self.reset(st, gpu, stream)?;
263        }
264
265        if let Some(table_va) = prestaged {
266            // 2026-09-25: The host half already ran from `decode_prestage`, before
267            // graph replay or capture: slots sit in `slots_dev`, history has
268            // advanced. Only the capture-safe kernel half remains.
269            self.gather_embed(table_va, num_tokens, heads, gpu, stream)?;
270        } else {
271            anyhow::ensure!(
272                !ctx.graph_capture,
273                "PLE: un-prestaged forward inside CUDA graph capture — the \
274                 pageable slot upload would invalidate the recording (901); \
275                 the scheduler must call decode_prestage every step"
276            );
277            // 2026-09-25: history ++ tokens, hashed together, then keep the new
278            // tokens' rows: the same slice the reference takes with
279            // `[:, -input_ids.shape[1]:]`.
280            let mut window = st.history.clone();
281            window.extend_from_slice(&tokens);
282            let all = ple_ngram_ids(&self.dims, &window);
283            let rows = &all[all.len() - num_tokens..];
284            let flat: Vec<u64> = rows.iter().flat_map(|r| r.iter().copied()).collect();
285
286            self.gather(&flat, num_tokens, heads, gpu, stream)?;
287
288            let keep = self.dims.context_len();
289            st.history = window[window.len() - keep..].to_vec();
290        }
291
292        // 2026-09-25: `gemm_k` is the pipelined kernel, so it must go through
293        // `ops::dense_gemm_bf16_pipelined` (grid [ceil(n,128), ceil(m,128)],
294        // block 256), not `ops::dense_gemm` (grid [ceil(n,16), ceil(m,16)],
295        // block 16x16): a kernel launched with the other wrapper's grid reads
296        // out of bounds.
297        ops::dense_gemm_bf16_pipelined(
298            gpu,
299            self.gemm_k,
300            self.emb,
301            &self.key_proj,
302            self.key,
303            num_tokens as u32,
304            c as u32,
305            self.hidden as u32,
306            stream,
307        )
308        .context("PLE key_proj")?;
309        ops::dense_gemm_bf16_pipelined(
310            gpu,
311            self.gemm_k,
312            self.emb,
313            &self.value_proj,
314            self.value,
315            num_tokens as u32,
316            self.hidden as u32,
317            self.hidden as u32,
318            stream,
319        )
320        .context("PLE value_proj")?;
321
322        ops::ple_gate(
323            gpu,
324            self.gate_k,
325            highway,
326            self.key,
327            self.value,
328            self.norm_query.weight,
329            self.norm_key.weight,
330            self.norm_conv.weight,
331            self.gated,
332            self.gated_normed,
333            num_tokens as u32,
334            self.hidden as u32,
335            self.hc_mult as u32,
336            self.eps,
337            stream,
338        )?;
339        ops::ple_conv(
340            gpu,
341            self.conv_k,
342            self.gated_normed,
343            self.gated,
344            self.conv1d.weight,
345            st.conv,
346            self.out,
347            num_tokens as u32,
348            c as u32,
349            self.k_size as u32,
350            self.dilation as u32,
351            stream,
352        )?;
353        ops::ple_add_highway(
354            gpu,
355            self.add_k,
356            self.out,
357            highway,
358            (num_tokens * c) as u32,
359            stream,
360        )?;
361
362        Ok(())
363    }
364
365    /// 2026-09-25: Resolve row ids to cache slots and gather them into `self.emb`.
366    ///
367    /// `T * ngram_heads` rows of `head_dim` land contiguously, which is the
368    /// `[T, ngram_heads * head_dim]` concatenation the projections read, so
369    /// the generic `batched_embed` kernel serves.
370    fn gather(
371        &self,
372        ids: &[u64],
373        num_tokens: usize,
374        heads: usize,
375        gpu: &dyn GpuBackend,
376        stream: u64,
377    ) -> Result<()> {
378        let table_va = self.gather_host(ids, gpu, stream)?;
379        self.gather_embed(table_va, num_tokens, heads, gpu, stream)
380    }
381
382    /// 2026-09-25: The host half of `gather`: NVMe fault-in + slot upload into the
383    /// stable `slots_dev` buffer. The upload reads pageable host memory, so
384    /// under CUDA graphs it runs from `prestage`, before replay or capture.
385    /// Returns the table's device VA for the kernel half.
386    fn gather_host(&self, ids: &[u64], gpu: &dyn GpuBackend, stream: u64) -> Result<u64> {
387        let mut table = self
388            .table
389            .lock()
390            .map_err(|_| anyhow::anyhow!("PLE table mutex poisoned"))?;
391        let table_va = match &mut *table {
392            #[cfg(feature = "cuda")]
393            NgramTable::Cached(cache) => {
394                // 2026-09-25: The host resolves row -> slot and faults missing rows
395                // off NVMe into the pinned, GPU-addressable arena. The gather
396                // kernel then reads the arena by slot.
397                let mut slots = Vec::with_capacity(ids.len());
398                let (h0, m0, _) = cache.stats();
399                let t0 = std::time::Instant::now();
400                cache.resolve(ids, &mut slots)?;
401                // 2026-09-25: Gathers of more than 64 ids (prefill scale) log the
402                // hit/miss profile at info; smaller ones (a decode step is
403                // `ngram_heads` ids) at debug.
404                let (h1, m1, _) = cache.stats();
405                let (dh, dm) = (h1 - h0, m1 - m0);
406                let us = t0.elapsed().as_micros();
407                if ids.len() > 64 {
408                    tracing::info!(
409                        "PLE gather: {} ids, {dh} hits / {dm} misses, resolve {us}us",
410                        ids.len()
411                    );
412                } else {
413                    tracing::debug!(
414                        "PLE gather: {} ids, {dh} hits / {dm} misses, resolve {us}us",
415                        ids.len()
416                    );
417                }
418                let bytes: Vec<u8> = slots.iter().flat_map(|s| s.to_le_bytes()).collect();
419                gpu.copy_h2d_async(&bytes, self.slots_dev, stream)?;
420                let va = cache.table_dev_va()?;
421                // 2026-09-25: This releases the pins before `gather_embed` issues
422                // the kernel that reads these slots (under CUDA graphs, before
423                // the replay). `NgramRowCache::resolve` documents `end_batch` as
424                // "once the gather has been issued", the order
425                // `ngram_embed/embed.rs` uses. A later `resolve` may evict one of
426                // these slots and write its new row from the host, which is not
427                // ordered with the stream.
428                cache.end_batch();
429                DevicePtr(va)
430            }
431            NgramTable::Bf16(w) => {
432                // 2026-09-25: Fully resident table: the slot is the row id, so the
433                // ids are uploaded truncated to u32.
434                let bytes: Vec<u8> = ids.iter().flat_map(|r| (*r as u32).to_le_bytes()).collect();
435                gpu.copy_h2d_async(&bytes, self.slots_dev, stream)?;
436                w.weight
437            }
438            NgramTable::Fp8(_) => anyhow::bail!(
439                "PLE: FP8 n-gram tables are not wired. This checkpoint ships BF16 \
440                 rows, which are both simpler and more accurate (on LongCat, BF16 \
441                 measured 0.0050 error vs FP8's 0.0247)."
442            ),
443        };
444        Ok(table_va.0)
445    }
446
447    /// 2026-09-25: The kernel half of `gather`: it reads `slots_dev` and the table
448    /// arena, both stable device addresses, so it is graph-capture-safe.
449    fn gather_embed(
450        &self,
451        table_va: u64,
452        num_tokens: usize,
453        heads: usize,
454        gpu: &dyn GpuBackend,
455        stream: u64,
456    ) -> Result<()> {
457        ops::batched_embed(
458            gpu,
459            self.embed_k,
460            self.slots_dev,
461            DevicePtr(table_va),
462            self.emb,
463            (num_tokens * heads) as u32,
464            self.head_dim as u32,
465            stream,
466        )
467        .context("PLE row gather")
468    }
469}
470
471/// 2026-09-25: The dense weights of one PLE site.
472pub struct PleWeights {
473    pub key_proj: DenseWeight,
474    pub value_proj: DenseWeight,
475    pub norm_key: DenseWeight,
476    pub norm_query: DenseWeight,
477    pub norm_conv: DenseWeight,
478    pub conv1d: DenseWeight,
479}
480
481// 2026-09-25: A child module (not a sibling), because the aux fns read
482// `PleLayer`'s private fields.
483#[path = "aux_state.rs"]
484mod aux_state;
metrale_model_layers/layers/ple/layer.rs

metrale_model_layers/layers/ple/
layer.rs
