metal: pre-M5 decode + prefill speedups — 46 t/s decode, +4% prefill on M3 Ultra (bit-exact)

ds4.h

@@ -488,6 +488,20 @@ int ds4_session_set_logits(ds4_session *s, const float *logits, int n);
 void ds4_session_gpu_warmup(ds4_session *s);
 int ds4_session_eval(ds4_session *s, int token, char *err, size_t errlen);
 
+/* Greedy chain decode for sessions (Metal, greedy only): keeps the next token
+ * id on-device and encodes ahead, removing the per-token host sync.
+ * on_token receives the seed first, then every confirmed id in order; a false
+ * return stops the burst early without committing that token.  Returns the
+ * number of approved tokens (>= 0) or -1 with err set.  When *completed is
+ * false the burst stopped early and s->logits is stale; when true, s->logits
+ * holds the logits after the last approved token, as after ds4_session_eval. */
+bool ds4_session_chain_greedy_supported(const ds4_session *s);
+int ds4_session_eval_chain_greedy(ds4_session *s, int max_tokens,
+                                  bool (*on_token)(void *ctx, int token),
+                                  void *on_token_ctx,
+                                  bool *completed,
+                                  char *err, size_t errlen);
+
 typedef struct {
     ds4_session *session;
     int token;

ds4_cuda.cu

@@ -26000,7 +26000,20 @@ extern "C" int ds4_gpu_routed_moe_owned_packed_combine_tensor(
 extern "C" int ds4_gpu_routed_moe_one_tensor(ds4_gpu_tensor *out, ds4_gpu_tensor *gate, ds4_gpu_tensor *up, ds4_gpu_tensor *mid, ds4_gpu_tensor *down, const void *model_map, uint64_t model_size, uint64_t gate_offset, uint64_t up_offset, uint64_t down_offset, uint32_t gate_type, uint32_t down_type, uint64_t gate_expert_bytes, uint64_t gate_row_bytes, uint64_t down_expert_bytes, uint64_t down_row_bytes, uint32_t expert_in_dim, uint32_t expert_mid_dim, uint32_t out_dim, const ds4_gpu_tensor *selected, const ds4_gpu_tensor *weights, uint32_t n_total_expert, uint32_t n_expert, float clamp, const ds4_gpu_tensor *x,
         const ds4_gpu_tensor *add_in,
         uint32_t layer_index,
-        bool force_resident) {
+        bool force_resident,
+        const ds4_gpu_tensor *hc_shared_out,
+        const ds4_gpu_tensor *hc_residual,
+        const ds4_gpu_tensor *hc_split,
+        ds4_gpu_tensor *hc_out,
+        int *hc_fused_out) {
+    /* The decode MoE+HC tail fusion is Metal-only; CUDA callers always pass
+     * NULL hc tensors. */
+    (void)hc_shared_out; (void)hc_residual; (void)hc_split;
+    if (hc_fused_out) *hc_fused_out = 0;
+    if (hc_out) {
+        fprintf(stderr, "ds4: routed MoE HC tail fusion is Metal-only\n");
+        return 0;
+    }
     if (add_in) {
         if (!ds4_gpu_add_tensor(out, out, add_in,
                                 (uint32_t)(out->bytes / sizeof(float)))) return 0;
@@ -26014,8 +26027,9 @@ extern "C" int ds4_gpu_routed_moe_one_tensor(ds4_gpu_tensor *out, ds4_gpu_tensor
                              selected, weights, n_total_expert, n_expert, clamp, x,
                              layer_index, 1, force_resident ? 0 : 1, 0);
 }
-extern "C" int ds4_gpu_routed_moe_batch_tensor(ds4_gpu_tensor *out, ds4_gpu_tensor *gate, ds4_gpu_tensor *up, ds4_gpu_tensor *mid, ds4_gpu_tensor *down, const void *model_map, uint64_t model_size, uint64_t gate_offset, uint64_t up_offset, uint64_t down_offset, uint32_t gate_type, uint32_t down_type, uint64_t gate_expert_bytes, uint64_t gate_row_bytes, uint64_t down_expert_bytes, uint64_t down_row_bytes, uint32_t expert_in_dim, uint32_t expert_mid_dim, uint32_t out_dim, const ds4_gpu_tensor *selected, const ds4_gpu_tensor *weights, uint32_t n_total_expert, uint32_t n_expert, float clamp, const ds4_gpu_tensor *x, uint32_t layer_index, uint32_t n_tokens, bool *mid_is_f16, bool force_resident) {
+extern "C" int ds4_gpu_routed_moe_batch_tensor(ds4_gpu_tensor *out, ds4_gpu_tensor *gate, ds4_gpu_tensor *up, ds4_gpu_tensor *mid, ds4_gpu_tensor *down, const void *model_map, uint64_t model_size, uint64_t gate_offset, uint64_t up_offset, uint64_t down_offset, uint32_t gate_type, uint32_t down_type, uint64_t gate_expert_bytes, uint64_t gate_row_bytes, uint64_t down_expert_bytes, uint64_t down_row_bytes, uint32_t expert_in_dim, uint32_t expert_mid_dim, uint32_t out_dim, const ds4_gpu_tensor *selected, const ds4_gpu_tensor *weights, uint32_t n_total_expert, uint32_t n_expert, float clamp, const ds4_gpu_tensor *x, uint32_t layer_index, uint32_t n_tokens, bool *mid_is_f16, bool defer_sum6, bool force_resident) {
     (void)force_resident;
+    if (defer_sum6) return 0;
     if (mid_is_f16) *mid_is_f16 = false;
     return routed_moe_launch(out, gate, up, mid, down, model_map, model_size,
                              gate_offset, up_offset, down_offset,

ds4_eval.c

@@ -3634,6 +3634,77 @@ static void eval_prefill_progress(void *ud, const char *event, int current, int
     if (paused_sec > 0.0) ui->phase_start_sec += paused_sec;
 }
 
+/* Greedy-chain burst decode: carries the per-token bookkeeping of the classic
+ * decode loop so a chained burst emits exactly the same stream, stop checks,
+ * and think-close records. */
+typedef struct {
+    ds4_engine *engine;
+    eval_ui *ui;
+    int idx;
+    int generation_limit;
+    byte_buf *raw;
+    bool *generation_in_think;
+    bool *plain_in_think;
+    eval_think_close_info *think_close;
+    double *t0;
+    bool tty;
+    bool use_plain_color;
+    bool stop;
+    bool quit;
+    bool switch_case;
+} eval_chain_ctx;
+
+static bool eval_chain_on_token(void *vctx, int token) {
+    eval_chain_ctx *c = (eval_chain_ctx *)vctx;
+    eval_ui *ui = c->ui;
+    if (c->tty) {
+        tui_consume_input(ui);
+        if (tui_has_quit_request(ui)) { c->quit = true; return false; }
+        if (tui_has_switch_request(ui, c->idx)) { c->switch_case = true; return false; }
+        double paused_sec = tui_wait_if_paused(ui, ui->in_think ? "thinking" : "answer");
+        if (paused_sec > 0.0) {
+            ui->phase_start_sec += paused_sec;
+            *c->t0 += paused_sec;
+        }
+        if (tui_has_quit_request(ui)) { c->quit = true; return false; }
+        if (tui_has_switch_request(ui, c->idx)) { c->switch_case = true; return false; }
+    }
+
+    if (ds4_token_is_stop(c->engine, token)) { c->stop = true; return false; }
+    size_t len = 0;
+    char *text = ds4_token_text(c->engine, token, &len);
+    buf_append(c->raw, text, len);
+    ui->generated++;
+    ui->generated_tokens[c->idx] = ui->generated;
+    tui_run_clock_tick(ui);
+    if (*c->generation_in_think && c->raw->v && strstr(c->raw->v, "</think>")) {
+        *c->generation_in_think = false;
+        if (c->think_close->kind == EVAL_THINK_CLOSE_NONE) {
+            c->think_close->kind = EVAL_THINK_CLOSE_NATURAL;
+            c->think_close->token_index = ui->generated;
+            c->think_close->remaining_budget =
+                c->generation_limit - ui->generated + 1;
+            c->think_close->rank = 0;
+        }
+    }
+    double elapsed = now_sec() - ui->phase_start_sec;
+    ui->speed_tps = elapsed > 0.001 ? (double)ui->generated / elapsed : 0.0;
+
+    if (c->tty) {
+        stream_append_token_text(ui, text, len, false);
+        tui_refresh(ui, ui->in_think ? "thinking" : "answer");
+    } else {
+        if (*c->plain_in_think && strstr(c->raw->v ? c->raw->v : "", "</think>")) {
+            *c->plain_in_think = false;
+            plain_reset_color(c->use_plain_color);
+        }
+        fwrite(text, 1, len, stdout);
+        fflush(stdout);
+    }
+    free(text);
+    return true;
+}
+
 static eval_run_result run_one_case(ds4_engine *engine, ds4_session *session,
                                     const eval_config *cfg, eval_ui *ui,
                                     FILE *trace, int idx, uint64_t *rng) {
@@ -3774,6 +3845,12 @@ static eval_run_result run_one_case(ds4_engine *engine, ds4_session *session,
 
     double t0 = ui->phase_start_sec;
     int forced_close_pos = -1;
+    /* Greedy chain bursts keep the token id on-device between evals (same
+     * machinery as the CLI chain decode).  Checked once per case; the burst
+     * branch below re-verifies per iteration that the think-close controller
+     * cannot intervene, so chained and classic decode pick identical tokens. */
+    const bool chain_burst_ok =
+        cfg->temperature <= 0.0f && ds4_session_chain_greedy_supported(session);
     for (int i = 0; i < generation_limit; i++) {
         if (tty) {
             tui_consume_input(ui);
@@ -3870,6 +3947,58 @@ static eval_run_result run_one_case(ds4_engine *engine, ds4_session *session,
                 }
             }
         }
+        /* Chained greedy burst: while the think-close controller cannot
+         * intervene (outside its reply-budget window), decode a run of argmax
+         * tokens with the id kept on-device.  The window is where the classic
+         * path may force a non-argmax token, so bursts stop ahead of it. */
+        if (token < 0 && chain_burst_ok && remaining_budget >= 2) {
+            int burst = remaining_budget;
+            if (generation_in_think && think_close_tokens.len > 0) {
+                const int window = think_close_tokens.len == 1
+                    ? cfg->soft_limit_reply_budget
+                    : cfg->hard_limit_reply_budget;
+                burst = remaining_budget - window;
+            }
+            if (burst >= 2) {
+                eval_chain_ctx cctx = {
+                    .engine = engine,
+                    .ui = ui,
+                    .idx = idx,
+                    .generation_limit = generation_limit,
+                    .raw = &raw,
+                    .generation_in_think = &generation_in_think,
+                    .plain_in_think = &plain_in_think,
+                    .think_close = &think_close,
+                    .t0 = &t0,
+                    .tty = tty,
+                    .use_plain_color = use_plain_color,
+                };
+                const int n = ds4_session_eval_chain_greedy(session, burst,
+                                                            eval_chain_on_token,
+                                                            &cctx, NULL,
+                                                            err, sizeof(err));
+                if (n < 0) {
+                    plain_reset_color(use_plain_color);
+                    ui->generated_tokens[idx] = ui->generated;
+                    tui_run_clock_stop(ui);
+                    fprintf(stderr, "ds4-eval: decode failed for %s: %s\n",
+                            tc->id, err);
+                    trace_write_case(trace, cfg, tc, idx, ui->ncases, "ERROR",
+                                     err, system, question,
+                                     raw.v ? raw.v : "", think_mode,
+                                     prompt_tokens, ui->generated,
+                                     now_sec() - t0, "?", &think_close);
+                    free(question);
+                    ds4_tokens_free(&think_close_tokens);
+                    buf_free(&raw);
+                    return EVAL_RUN_ERROR;
+                }
+                if (cctx.stop) break;
+                if (cctx.quit || cctx.switch_case) continue; /* handled above */
+                i += n - 1; /* keep i in lockstep with ui->generated */
+                continue;
+            }
+        }
         if (token < 0)
             token = ds4_session_sample(session, cfg->temperature, 0,
                                        cfg->top_p, cfg->min_p, rng);

ds4_gpu.h

@@ -110,6 +110,13 @@ int ds4_gpu_tensor_read_after_selected_event(const ds4_gpu_tensor *tensor,
 int ds4_gpu_end_commands(void);
 int ds4_gpu_synchronize(void);
 
+/* Diagnostic GPU stage-counter profiler: timestamp samples taken at decode
+ * stage boundaries without ending the batch command buffer. */
+int ds4_gpu_stage_counters_enabled(void);
+int ds4_gpu_stage_counter_sample(const char *label);
+void ds4_gpu_stage_counter_reset(void);
+void ds4_gpu_stage_counter_report(uint32_t pos);
+
 int ds4_gpu_set_model_map(const void *model_map, uint64_t model_size);
 int ds4_gpu_set_model_fd(int fd);
 int ds4_gpu_set_model_fd_for_map(int fd, const void *model_map);
@@ -186,7 +193,11 @@ enum {
     DS4_GPU_TEST_MXFP4_DOWN_HALF_LUT = 1u << 4,
     DS4_GPU_TEST_OUTPUT_HC_WEIGHTS4 = 1u << 5,
     DS4_GPU_TEST_HC_RMS_SCALE_PROJ = 1u << 6,
+    DS4_GPU_TEST_ATTN_OUT_LOW_Q8_STATIC = 1u << 7,
+    DS4_GPU_TEST_INDEXED_ATTN_PREFILL_RB4 = 1u << 8,
+    DS4_GPU_TEST_BATCH_ATTN_OUT_HC_FUSION = 1u << 9,
 };
+int ds4_gpu_test_get_quality(void);
 void ds4_gpu_test_set_flags(uint32_t flags);
 void ds4_gpu_release_zero_prefix_prefill_mask_cache(void);
 #else
@@ -1119,6 +1130,28 @@ int ds4_gpu_head_rms_norm_rope_tail_tensor(
         float             beta_slow,
         float             eps);
 
+int ds4_gpu_dsv4_batch_qnorm_rope_kv_finalize_available(void);
+
+int ds4_gpu_dsv4_batch_qnorm_rope_kv_finalize_tensor(
+        ds4_gpu_tensor *q,
+        ds4_gpu_tensor *kv,
+        ds4_gpu_tensor *raw_cache,
+        uint32_t        raw_cap,
+        uint32_t        n_tok,
+        uint32_t        n_head,
+        uint32_t        head_dim,
+        uint32_t        n_rot,
+        uint32_t        pos0,
+        uint32_t        n_ctx_orig,
+        bool            inverse,
+        float           freq_base,
+        float           freq_scale,
+        float           ext_factor,
+        float           attn_factor,
+        float           beta_fast,
+        float           beta_slow,
+        float           eps);
+
 int ds4_gpu_attn_q_b_f16_head_rms_rope_tail_tensor(
         ds4_gpu_tensor       *out,
         ds4_gpu_tensor       *q_half,
@@ -2250,6 +2283,31 @@ int ds4_gpu_attention_output_q8_batch_tensor(
         uint64_t                out_dim,
         const ds4_gpu_tensor *heads,
         uint32_t                n_tokens);
+
+#ifdef __APPLE__
+/* Optional pre-M5 batch Q8 attention-output B + HC4 epilogue. Returns 1
+ * when fused work was encoded, 0 when unsupported without encoding work,
+ * and -1 after an attempted-path failure. */
+int ds4_gpu_attention_output_q8_batch_hc_tensor(
+        ds4_gpu_tensor       *out,
+        ds4_gpu_tensor       *out_hc,
+        const ds4_gpu_tensor *residual_hc,
+        const ds4_gpu_tensor *split,
+        ds4_gpu_tensor       *low,
+        ds4_gpu_tensor       *group_tmp,
+        ds4_gpu_tensor       *low_tmp,
+        const void             *model_map,
+        uint64_t                model_size,
+        uint64_t                out_a_offset,
+        uint64_t                out_b_offset,
+        uint64_t                group_dim,
+        uint64_t                rank,
+        uint32_t                n_groups,
+        uint64_t                out_dim,
+        const ds4_gpu_tensor *heads,
+        uint32_t                n_tokens,
+        uint32_t                n_hc);
+#endif
 int ds4_gpu_attention_output_q4_K_batch_tensor(
         ds4_gpu_tensor       *out,
         ds4_gpu_tensor       *low,
@@ -2385,6 +2443,27 @@ int ds4_gpu_router_select_tensor(
         bool                    hash_mode,
         const ds4_gpu_tensor *logits);
 
+/* One-token router select with the token id read from a device buffer
+ * (greedy chain decode); identical kernels to ds4_gpu_router_select_tensor. */
+int ds4_gpu_router_select_tensor_devtoken(
+        ds4_gpu_tensor       *selected,
+        ds4_gpu_tensor       *weights,
+        ds4_gpu_tensor       *probs,
+        const void             *model_map,
+        uint64_t                model_size,
+        uint64_t                bias_offset,
+        uint64_t                hash_offset,
+        uint32_t                hash_rows,
+        const ds4_gpu_tensor *token_dev,
+        uint32_t                n_expert,
+        uint32_t                n_expert_used,
+        float                   expert_weight_scale,
+        uint32_t                n_expert_groups,
+        uint32_t                n_group_used,
+        bool                    has_bias,
+        bool                    hash_mode,
+        const ds4_gpu_tensor *logits);
+
 int ds4_gpu_router_select_batch_tensor(
         ds4_gpu_tensor       *selected,
         ds4_gpu_tensor       *weights,
@@ -2687,7 +2766,12 @@ int ds4_gpu_routed_moe_one_tensor(
         const ds4_gpu_tensor *x,
         const ds4_gpu_tensor *add_in,
         uint32_t                layer_index,
-        bool                    force_resident);
+        bool                    force_resident,
+        const ds4_gpu_tensor *hc_shared_out,
+        const ds4_gpu_tensor *hc_residual,
+        const ds4_gpu_tensor *hc_split,
+        ds4_gpu_tensor       *hc_out,
+        int                    *hc_fused_out);
 
 int ds4_gpu_routed_moe_batch_tensor(
         ds4_gpu_tensor       *out,
@@ -2718,6 +2802,7 @@ int ds4_gpu_routed_moe_batch_tensor(
         uint32_t                layer_index,
         uint32_t                n_tokens,
         bool                   *mid_is_f16,
+        bool                    defer_sum6,
         bool                    force_resident);
 
 /* =========================================================================
@@ -2899,6 +2984,23 @@ int ds4_gpu_hc_expand_add_split_tensor(
         uint32_t                n_embd,
         uint32_t                n_hc);
 
+#ifdef __APPLE__
+int ds4_gpu_moe_sum6_hc_expand_available(void);
+int ds4_gpu_moe_sum6_hc_expand_split_tensor(
+        ds4_gpu_tensor       *out_hc,
+        const ds4_gpu_tensor *expert_down,
+        const ds4_gpu_tensor *shared_out,
+        const ds4_gpu_tensor *residual_hc,
+        const ds4_gpu_tensor *split,
+        uint32_t                n_embd,
+        uint32_t                n_hc);
+int ds4_gpu_test_moe_sum6_tensor(
+        ds4_gpu_tensor       *out,
+        const ds4_gpu_tensor *expert_down,
+        uint32_t                n_embd,
+        uint32_t                n_tokens);
+#endif
+
 int ds4_gpu_hc_expand_add_split_half_add_tensor(
         ds4_gpu_tensor       *out_hc,
         const ds4_gpu_tensor *block_out,