more diags (by DS), but MTP output still garbage

Author: Ljubomir Josifovski

ggml/src/ggml-metal/ggml-metal.metal

@@ -2756,7 +2756,9 @@ template [[host_name("kernel_gated_delta_net_f32_4")]] kernel kernel_gated_delta
 // state_in:   [S*S*H, n_seqs]
 // output (no intermediates):  [S*H, n_tokens*n_seqs + S*n_seqs] (activations + final state)
 // output (keep_intermediates): [S*H, n_tokens*n_seqs + S*S*H*n_seqs + (n_seq_tokens-1)*S*S*H*n_seqs]
-//   layout: [activations] [final_state] [snapshots: t=1 state, t=2 state, ...] (reverse order, skip t=0)
+//   layout: [activations] [final_state] [snapshots: state after confirmed token, state after draft skip, ...]
+//   saves state after each token EXCEPT the last (which is stored as final state)
+//   snap_idx = t_in_seq (0-indexed per sequence), skip t_end-1
 // grid=(S/NSG, S/4, H*n_seqs), threadgroup=(NW, 4, 1)
 template<short NSG, bool KEEP_INTERMEDIATES>
 kernel void kernel_gated_linear_attn_impl(
@@ -2832,20 +2834,18 @@ kernel void kernel_gated_linear_attn_impl(
             dst_out[(t - t_start) * S * H + h_idx * S + j] = out_j;
         }
 
-        // Store intermediate state for rollback (skip t=0, reverse-index per CPU semantics)
-        if (KEEP_INTERMEDIATES) {
+        // Store intermediate state for rollback (save state after each token except the last)
+        // The snapshot at index (t - t_start) is the state AFTER confirmed tokens,
+        // BEFORE the draft token. On draft rejection, this is what we roll back to.
+        if (KEEP_INTERMEDIATES && t < t_end - 1) {
             const short t_in_seq = t - t_start;
-            if (t_in_seq > 0) {
-                const short snap_idx = T_per_seq - 1 - t_in_seq;
-                // Snapshot layout: after final state, indexed by (snap_idx * n_seqs + seq_idx) * H * S * S + h_idx * S * S
-                const int64_t state_elems_all = (int64_t)S * S * H * n_seqs;
-                device float * snap = (device float *) dst + state_base + state_elems_all
-                    + ((int64_t)snap_idx * n_seqs + seq_idx) * H * S * S + h_idx * S * S;
-                FOR_UNROLL (short i = 0; i < NSG; i++) {
-                    const short ii = i_start + tx*NSG + i;
-                    if (ii < S && j < S) {
-                        snap[ii * S + j] = state_vals[i];
-                    }
+            const int64_t state_elems_all = (int64_t)S * S * H * n_seqs;
+            device float * snap = (device float *) dst + state_base + state_elems_all
+                + ((int64_t)t_in_seq * n_seqs + seq_idx) * H * S * S + h_idx * S * S;
+            FOR_UNROLL (short i = 0; i < NSG; i++) {
+                const short ii = i_start + tx*NSG + i;
+                if (ii < S && j < S) {
+                    snap[ii * S + j] = state_vals[i];
                 }
             }
         }