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