diff --git a/src/llama-context.cpp b/src/llama-context.cpp index 2dfa35538b2f..1b9898e94ea0 100644 --- a/src/llama-context.cpp +++ b/src/llama-context.cpp @@ -22,6 +22,7 @@ #include "ggml-alloc.h" #include +#include #include #include #include @@ -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 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; } diff --git a/src/llama-cparams.h b/src/llama-cparams.h index 1643b2b7376c..3918de3ba445 100644 --- a/src/llama-cparams.h +++ b/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; diff --git a/src/models/dflash_draft.cpp b/src/models/dflash_draft.cpp index fbf34cb485f4..719f1fa1a558 100644 --- a/src/models/dflash_draft.cpp +++ b/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 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 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(