VULKAN CRASH FIX Gate DFlash graph argmax/top-k to CUDA backend

src/llama-context.cpp

@@ -22,6 +22,7 @@
 #include "ggml-alloc.h"
 
 #include <algorithm>
+#include <atomic>
 #include <cinttypes>
 #include <cmath>
 #include <cstdlib>
@@ -274,6 +275,16 @@ llama_context::llama_context(
                                         1, (int) LLAMA_DFLASH_MAX_SLOTS);
     cparams.dflash_cross_ctx = params.dflash_cross_ctx > 0 ? params.dflash_cross_ctx : 512;
 
+    // DFlash graph argmax/topk is only implemented by the CUDA backend; the stock
+    // GGML_OP_ARGMAX kernels on other backends ignore the extended op_params and
+    // leave most of the compact ids/probs tensor uninitialized.
+    {
+        ggml_backend_dev_t output_dev = model.dev_output();
+        ggml_backend_reg_t output_reg = output_dev ? ggml_backend_dev_backend_reg(output_dev) : nullptr;
+        const char * output_reg_name = output_reg ? ggml_backend_reg_name(output_reg) : nullptr;
+        cparams.dflash_graph_argmax = output_reg_name != nullptr && strcmp(output_reg_name, "CUDA") == 0;
+    }
+
     // Initialize backend samplers here so they are part of the sampling graph
     // before the reserve passes run later in this function. This avoids a later
     // re-reserve when graph nodes change.
@@ -1741,6 +1752,14 @@ void llama_context::set_dflash_topk(int k) {
 
 void llama_context::set_dflash_verify_logits(bool enabled, int top_k) {
     const int clamped_top_k = std::max(1, std::min(top_k, 64));
+    if (enabled && !cparams.dflash_graph_argmax) {
+        static std::atomic<bool> warned{false};
+        if (!warned.exchange(true)) {
+            LLAMA_LOG_WARN("%s: DFlash reduced verify logits requested, but graph argmax/topk is only implemented "
+                    "on the CUDA backend; falling back to full-logits verification\n", __func__);
+        }
+        enabled = false;
+    }
     if (cparams.dflash_verify_logits == enabled && cparams.dflash_verify_topk == clamped_top_k) {
         return;
     }

src/llama-cparams.h

@@ -68,6 +68,12 @@ struct llama_cparams {
     int  dflash_verify_topk = 1;
     bool dflash_reduced_consumer_active = false;
 
+    // DFlash: whether graph-embedded argmax/topk (ggml_argmax_ext/ggml_topk_ext)
+    // produces valid results on the model's output device. Only the CUDA backend
+    // implements the extended semantics; other backends run the stock argmax
+    // kernel and leave the rest of the compact ids/probs tensor uninitialized.
+    bool dflash_graph_argmax = false;
+
     // DFlash: cross-attention window in tokens (how many target hidden states the drafter sees)
     int dflash_cross_ctx = 512;
 

src/models/dflash_draft.cpp

@@ -1085,18 +1085,25 @@ llm_build_dflash_draft::llm_build_dflash_draft(
     cb(cur, "result_output", -1);
     res->t_logits = cur;
 
-    // GPU top-K or argmax — avoids 15.9MB logits transfer + CPU scan for DFlash draft
-    const float sample_temp = cparams.dflash_sample_temp;
-    static std::atomic<uint64_t> gumbel_counter{1};
-    const uint64_t seed = (sample_temp > 0.0f) ? gumbel_counter.fetch_add(1) : 0;
-    const int topk = cparams.dflash_topk;
-    if (topk > 1) {
-        res->t_logits_argmax = ggml_topk_ext(ctx0, cur, topk, sample_temp, seed);
+    // GPU top-K or argmax — avoids 15.9MB logits transfer + CPU scan for DFlash draft.
+    // Only built when the output device backend implements the extended argmax
+    // semantics (CUDA); otherwise leave t_logits_argmax unset so the drafter
+    // falls back to scanning the full logits on the CPU.
+    if (cparams.dflash_graph_argmax) {
+        const float sample_temp = cparams.dflash_sample_temp;
+        static std::atomic<uint64_t> gumbel_counter{1};
+        const uint64_t seed = (sample_temp > 0.0f) ? gumbel_counter.fetch_add(1) : 0;
+        const int topk = cparams.dflash_topk;
+        if (topk > 1) {
+            res->t_logits_argmax = ggml_topk_ext(ctx0, cur, topk, sample_temp, seed);
+        } else {
+            res->t_logits_argmax = ggml_argmax_ext(ctx0, cur, sample_temp, seed);
+        }
+
+        ggml_build_forward_expand(gf, res->t_logits_argmax);
     } else {
-        res->t_logits_argmax = ggml_argmax_ext(ctx0, cur, sample_temp, seed);
+        ggml_build_forward_expand(gf, cur);
     }
-
-    ggml_build_forward_expand(gf, res->t_logits_argmax);
 }
 
 llm_build_dflash_kv_update::llm_build_dflash_kv_update(