dasLLAMA: dig the last two board reds green — lcpp-shape mx4 kernels + GPU PLE pre-step, fresh M1 board

doc/source/conf.py

@@ -276,6 +276,8 @@
 \DeclareUnicodeCharacter{21D2}{\ensuremath{\Rightarrow}}
 \DeclareUnicodeCharacter{00D7}{\ensuremath{\times}}
 \DeclareUnicodeCharacter{03C0}{\ensuremath{\pi}}
+\DeclareUnicodeCharacter{03B8}{\ensuremath{\theta}}
+\DeclareUnicodeCharacter{2212}{\ensuremath{-}}
 \DeclareUnicodeCharacter{2248}{\ensuremath{\approx}}
 \DeclareUnicodeCharacter{2208}{\ensuremath{\in}}
 \DeclareUnicodeCharacter{2209}{\ensuremath{\notin}}

modules/dasLLAMA/dasllama/dasllama_common.das

@@ -235,6 +235,18 @@ def active_prefill_override : string {
     return g_prefill_override_name
 }
 
+//! The GPU-PLE pre-step gate: a prefill override that can ALSO build the E-series PLE side input
+//! on its device registers a probe here. On true, forward_prefill stashes the token rows and skips
+//! the CPU gather+finish; the override consumes the stash (ple_pre_prefill = late-decline fallback).
+typedef PleGpuGateFn = function<(t : Model; s : Session; npos : int64; start_pos : int64) : bool>
+var private g_ple_gpu_gate : PleGpuGateFn
+var private g_ple_gpu_gate_set = false
+
+def register_ple_gpu_gate(fn : PleGpuGateFn) {
+    g_ple_gpu_gate = fn
+    g_ple_gpu_gate_set = true
+}
+
 // set by an override that ALSO produced the last position's logits; reset before every invocation
 var private g_prefill_logits_done = false
 
@@ -4520,6 +4532,8 @@ struct Session {
     ple_inp : array<float>      // the built per-(position, layer) side input (npos x n_layers x ple)
     ple_proj : array<float>     // model-proj scratch for the pre-step (npos x n_layers x ple)
     ple_gate : array<float>     // per-(position) gate vector for the layer block (npos x ple)
+    ple_gpu_tokens : array<int64>  // stashed token rows for a driver-owned pre-step (see register_ple_gpu_gate)
+    ple_gpu_pending : bool      // forward_prefill skipped the CPU pre-step — the engaged driver builds it
     // qwen35 deltanet recurrent state (persists across tokens; zeroed at position 0) + scratch.
     // Empty unless config.recr_mask != 0. State is forward-only — dn_pos guards against rewind.
     dn_pos : int64                 // the next position the deltanet stack expects
@@ -8644,8 +8658,9 @@ def private ple_pre_decode(t : Model; var s : Session; token : int64) {
     ple_pre_finish(t, s, s.x, 1l)
 }
 
-// Prefill pre-step: gather every position's per-layer embedding row, then finish (batched).
-def private ple_pre_prefill(t : Model; var s : Session; tokens : array<int64>; npos : int64) {
+//! Prefill pre-step: gather every position's per-layer embedding row, then finish (batched).
+//! Public as the late-decline fallback for a driver that claimed the pre-step via the gate.
+def ple_pre_prefill(t : Model; var s : Session; tokens : array<int64>; npos : int64) {
     let all = t.config.n_layers * t.config.n_embd_per_layer
     for (p in range64(npos)) {
         ple_gather_row(t, tokens[p], s.ple_inp, p * all)
@@ -8680,6 +8695,8 @@ def private ple_block_decode(t : Model; var s : Session; l : int64) {
 // Per-layer block (prefill, batched over npos). ple_inp is position-major so the layer-l side vector
 // of position p sits at p*(layers*ple) + l*ple.
 def private ple_block_prefill(t : Model; var s : Session; l : int64; npos : int64) {
+    // a broken gate promise would read an empty/stale ple_inp here — fail loud, never mismeasure
+    if (s.ple_gpu_pending) panic("ple_block_prefill: GPU pre-step promised but the CPU body engaged")
     let c = t.config
     let dim = c.dim
     let ple = c.n_embd_per_layer
@@ -11952,7 +11969,15 @@ def forward_prefill(t : Model; var s : Session; tokens : array<int64>; npos, sta
     prof_add("embed", ts_embed)
     if (has_ple(c)) {  // gemma-4 E-series: build each position's per-(layer) side input before the stack
         let ts_ple = ref_time_ticks()
-        ple_pre_prefill(t, s, tokens, npos)
+        if (g_ple_gpu_gate_set && invoke(g_ple_gpu_gate, t, s, npos, start_pos)) {
+            s.ple_gpu_tokens |> resize(int(npos))
+            for (p in range64(npos)) {
+                s.ple_gpu_tokens[p] = tokens[p]
+            }
+            s.ple_gpu_pending = true
+        } else {
+            ple_pre_prefill(t, s, tokens, npos)
+        }
         prof_add("ple", ts_ple)
     }
 

modules/dasLLAMA/dasllama/dasllama_metal_common.das

@@ -267,7 +267,6 @@ var g_pso_moe_wscale : MetalComputePipeline?   // gemma4 per-expert down-scale f
 var g_pso_swiglu_oai : MetalComputePipeline?   // gpt-oss clamped swiglu (stage 4)
 var g_pso_moe_gemv_mx4 : MetalComputePipeline?  // gpt-oss MXFP4 expert GEMV
 var g_pso_moe_gemv_q51 : MetalComputePipeline?  // Q5_1 expert GEMV (per-32 planes)
-var g_e8m0_tab : MetalBuffer?                   // [256] exact e8m0_half scales
 var g_mx4_vtab : MetalBuffer?                   // [16] signed doubled-e2m1 values
 var g_pso_dn_conv : MetalComputePipeline?      // deltanet (Wave D): causal conv + SiLU
 var g_pso_dn_hist : MetalComputePipeline?      // conv-history advance
@@ -673,6 +672,15 @@ def blob_of(dev : MetalDevice?; t : Model; off : int64) : tuple<buf : MetalBuffe
     }
 }
 
+//! The kept-bf16 plane (E-series model_proj) + one tensor's byte offset within it — zero-copy
+//! page wrap on image-mapped models, one-time resident copy otherwise (cached per base).
+def bf16_of(dev : MetalDevice?; t : Model; off : int64) : tuple<buf : MetalBuffer?; boff : uint64> {
+    unsafe {
+        return (buf = plane_buffer(dev, addr < void? >(t.bf16blob[0]), uint64(long_length(t.bf16blob) * 2l), t.image_map != null),
+                boff = uint64(off * 2l))
+    }
+}
+
 //! The native-MXFP4 expert planes (gpt-oss) + one stack's byte offsets within them: nibble
 //! bytes at off/2, E8M0 scale bytes at off/32 (disk order — the mx4_repacked gate keeps the
 //! CPU laneq repack out). Region-relative uint32 indices stay safe per layer stack.

modules/dasLLAMA/dasllama/dasllama_metal_kernels.das

@@ -5861,13 +5861,6 @@ def metal_decode_init : bool {
     g_pso_dn_deint = compile_pso(metal_dn_deint_msl, metal_dn_deint_msl_entry, metal_dn_deint_msl_fastmath, ok)
     g_pso_sigmul = compile_pso(metal_sigmul_msl, metal_sigmul_msl_entry, metal_sigmul_msl_fastmath, ok)
     g_pso_axpy_sig = compile_pso(metal_axpy_sig_msl, metal_axpy_sig_msl_entry, metal_axpy_sig_msl_fastmath, ok)
-    g_e8m0_tab = metal_new_buffer(g_dev, 1024ul)
-    unsafe {
-        var ep = reinterpret<uint?>(metal_buffer_contents(g_e8m0_tab))
-        for (ei in range(256)) {
-            ep[ei] = ei < 2 ? (0x00200000u << uint(ei)) : (uint(ei - 1) << 23u)
-        }
-    }
     g_mx4_vtab = metal_new_buffer(g_dev, 64ul)
     unsafe {
         var vp = reinterpret<float?>(metal_buffer_contents(g_mx4_vtab))
@@ -6213,6 +6206,7 @@ def enc_moe_gemv_mx4(enc : MetalComputeEncoder?; t : Model; stack_off, rows, k,
                      bx, by, bn, bd, bsel, beblk, bxs, bss, bbxs, bbys, bwb, bhasb : MetalBuffer?; yoff : uint64 = 0ul) {
     let mp = mx4_of(g_dev, t, stack_off)
     kn_pipeline(enc, g_pso_moe_gemv_mx4)
+    kn_tgmem(enc, metal_moe_gemv_mx4_msl_tgmem, 0)
     kn_buffer(enc, mp.qbuf, mp.qoff, 0)
     kn_buffer(enc, mp.sbuf, mp.soff, 1)
     kn_buffer(enc, bx, 0ul, 2)
@@ -6227,7 +6221,6 @@ def enc_moe_gemv_mx4(enc : MetalComputeEncoder?; t : Model; stack_off, rows, k,
     kn_buffer(enc, bbys, 0ul, 11)
     kn_buffer(enc, bwb, 0ul, 12)
     kn_buffer(enc, bhasb, 0ul, 13)
-    kn_buffer(enc, g_e8m0_tab, 0ul, 14)
     kn_buffer(enc, g_mx4_vtab, 0ul, 15)
     hz_read(bx)
     hz_read(bsel)
@@ -6939,10 +6932,6 @@ def metal_kernels_release {
         metal_release(g_mx4_vtab)
         g_mx4_vtab = null
     }
-    if (g_e8m0_tab != null) {
-        metal_release(g_e8m0_tab)
-        g_e8m0_tab = null
-    }
     if (g_pso_qkvrs16 != null) {
         metal_release(g_pso_qkvrs16)
         g_pso_qkvrs16 = null

modules/dasLLAMA/dasllama/dasllama_metal_llama.das

@@ -1877,6 +1877,9 @@ def private release_step(var r : StepRes) {
         pool_release(g_upool, r.u_moe_k, 4ul)
         pool_release(g_upool, r.u_moe_nfe, 4ul)
         pool_release(g_upool, r.u_moe_knfe, 4ul)
+        pool_release(g_upool, r.u_moe_ss, 4ul)
+        pool_release(g_upool, r.u_moe_kdim, 4ul)
+        pool_release(g_upool, r.u_moe_tot, 4ul)
         pool_release(g_upool, r.u_moe_eblk, 4ul)
         pool_release(g_upool, r.u_moe_esb, 4ul)
         pool_release(g_upool, r.u_moe_xs_dn, 4ul)