glm: overlap the indexed full-layer prefill sweep with a pread prepare of the next layer

ds4.c

@@ -30722,7 +30722,15 @@ static bool glm_graph_forward_indexed_tokens(
     const uint32_t drain_interval =
         progress_flush_interval != 0 ? glm_graph_indexed_prefill_drain_interval() : 0u;
     const bool progress_requested = display_progress && work_total > 0;
-    ds4_gpu_set_glm_streaming_prefill_full_layer(false);
+    /* Mirror the plain prefill flow: above the full-layer threshold the
+     * selected-expert addr loader is the wrong tool — a chunk of N tokens
+     * selects far more unique experts per layer than the cache can hold, and
+     * its overflow branch needs the full expert tensors mapped.  Hardcoding
+     * false here broke every generation prompt above ~64 tokens ("failed to
+     * map overflow expert views"), while short prompts keep the addr path. */
+    const bool full_layer_prefill =
+        glm_graph_stream_prefill_full_layer_enabled(g, n_tokens);
+    ds4_gpu_set_glm_streaming_prefill_full_layer(full_layer_prefill);
 
     if (trace) {
         glm_graph_indexed_prefill_tracef(
@@ -30810,15 +30818,58 @@ static bool glm_graph_forward_indexed_tokens(
     } while (0)
     ds4_gpu_tensor *last_indexer_selected = NULL;
     uint32_t last_indexer_selected_count = 0;
+    /* The indexed prefill maps each full layer and then pays the sweep as GPU
+     * demand faults during the kernels (~2.6 GB/s measured on an M5 Pro
+     * against ~7 GB/s the SSD can sustain).  Reuse the full-attention layer
+     * prepare (pread pool warming the next layer while this one computes):
+     * measured 3.5 -> 8.5 t/s prefill on a 271-token prompt, decode
+     * unaffected, identical greedy output.  Opt-in while it soaks. */
+    metal_graph_stream_prepare_slot idx_prepare_slots[DS4_STREAM_PREFILL_MAX_PREPARE_AHEAD];
+    memset(idx_prepare_slots, 0, sizeof(idx_prepare_slots));
+    const bool idx_prepare =
+        g->ssd_streaming &&
+        full_layer_prefill &&
+        getenv("DS4_METAL_ENABLE_GLM_INDEXED_PREFILL_PREPARE") != NULL;
     for (uint32_t il = g->layer_start; ok && il <= g->layer_end; il++) {
         const uint32_t slice_layer_done = il - g->layer_start + 1u;
         if (g->ssd_streaming) {
+            if (idx_prepare &&
+                !metal_graph_stream_prepare_join_layer(NULL,
+                                                       model,
+                                                       weights,
+                                                       il,
+                                                       n_tokens,
+                                                       false,
+                                                       true,
+                                                       false,
+                                                       false,
+                                                       idx_prepare_slots,
+                                                       DS4_STREAM_PREFILL_MAX_PREPARE_AHEAD)) {
+                ok = false;
+                break;
+            }
             ok = glm_graph_stream_map_prefill_layer(g,
                                                     model,
                                                     weights,
                                                     il,
                                                     n_tokens,
-                                                    false);
+                                                    full_layer_prefill);
+            if (ok && idx_prepare && il < g->layer_end) {
+                if (!metal_graph_stream_prepare_start_if_needed(
+                            NULL,
+                            model,
+                            weights,
+                            il + 1,
+                            n_tokens,
+                            false,
+                            true,
+                            false,
+                            false,
+                            idx_prepare_slots,
+                            DS4_STREAM_PREFILL_MAX_PREPARE_AHEAD)) {
+                    ok = false;
+                }
+            }
             if (ok) ok = ds4_gpu_begin_commands() != 0;
         }
         const ds4_layer_weights *l = &weights->layer[il];
@@ -31738,6 +31789,11 @@ static bool glm_graph_forward_indexed_tokens(
         }
         (void)ds4_gpu_synchronize();
     }
+    if (idx_prepare &&
+        !metal_graph_stream_prepare_join_all(idx_prepare_slots,
+                                             DS4_STREAM_PREFILL_MAX_PREPARE_AHEAD)) {
+        ok = false;
+    }
     if (ok) {
         glm_graph_report_prefill_display_progress(display_progress,
                                                   display_progress_ud,

ds4_metal.m

@@ -12859,12 +12859,16 @@ static int ds4_gpu_stream_expert_cache_prepare_selected_batch(
                 continue;
             }
 
-            if (cache_budget != 0 && reserved_entries >= cache_budget) {
-                unique_entries[u] = NULL;
-                continue;
-            }
-
-            const int force_reuse = 0;
+            /* At budget, evict-and-reuse like the decode loader does instead
+             * of spilling to whole-tensor overflow views: the overflow wrap
+             * needs the expert tensors in the mapped view set, which the
+             * decode-expert-cache prefill mapping does not provide, so a
+             * full cache made any checkpoint-extend or small-suffix prefill
+             * fail ("failed to map overflow expert views").  Commands were
+             * just synced by the caller, so reuse cannot race the GPU, and
+             * prepare_load_buffers protects this layer's unique set. */
+            const int force_reuse =
+                cache_budget != 0 && reserved_entries >= cache_budget;
             ds4_gpu_stream_expert_readahead_range(unique_gate_offsets[u],
                                                   gate_expert_bytes);
             ds4_gpu_stream_expert_readahead_range(unique_up_offsets[u],