Harden DFlash CUDA multi-GPU fallbacks

Anbeeld · Anbeeld · commit fca6b9350d08 · 2026-05-12T14:00:49.000+02:00
diff --git a/common/speculative.cpp b/common/speculative.cpp
@@ -77,16 +77,22 @@ static bool common_dflash_kv_cache_disabled() {
     return disabled;
 }
 
-static bool common_dflash_gpu_ring_allowed(llama_context * ctx_tgt) {
-    const llama_model * model_tgt = ctx_tgt ? llama_get_model(ctx_tgt) : nullptr;
-    if (!model_tgt) {
-        return true;
+static bool common_dflash_gpu_ring_env_enabled() {
+    const char * env = std::getenv("GGML_DFLASH_GPU_RING");
+    return env == nullptr || std::atoi(env) != 0;
+}
+
+static bool common_dflash_gpu_ring_allowed(llama_context * ctx_tgt, llama_context * ctx_dft) {
+    if (!common_dflash_gpu_ring_env_enabled()) {
+        LOG_INF("dflash: GPU cross ring disabled by GGML_DFLASH_GPU_RING=0; using CPU hidden capture\n");
+        return false;
     }
 
-    const int32_t n_devices = llama_model_n_devices(model_tgt);
-    if (n_devices > 1) {
-        LOG_INF("dflash: multi-GPU target detected (%d devices); disabling GPU cross ring because target hidden capture uses CPU fallback\n",
-                n_devices);
+    const int32_t n_tgt_devices = ctx_tgt ? llama_model_n_devices(llama_get_model(ctx_tgt)) : 1;
+    const int32_t n_dft_devices = ctx_dft ? llama_model_n_devices(llama_get_model(ctx_dft)) : 1;
+    if (n_tgt_devices > 1 || n_dft_devices > 1) {
+        LOG_INF("dflash: multi-GPU placement detected (target=%d devices, drafter=%d devices); disabling GPU cross ring and graph hidden capture\n",
+                n_tgt_devices, n_dft_devices);
         return false;
     }
 
@@ -1621,13 +1627,20 @@ struct common_speculative_state_dflash : public common_speculative_state {
 
         batch_dft = llama_batch_init(block_size, 0, 1);
 
-        // try to allocate GPU ring buffer on drafter's GPU
-        if (common_dflash_gpu_ring_allowed(ctx_tgt)) {
+        // Try to allocate the GPU ring only when both contexts can safely share
+        // single-backend cross tensors. Otherwise keep the target eval callback
+        // active so the CPU ring fallback has hidden states to consume.
+        const bool gpu_ring_requested = common_dflash_gpu_ring_allowed(ctx_tgt, ctx_dft);
+        llama_set_dflash_gpu_capture(ctx_tgt, gpu_ring_requested);
+        if (gpu_ring_requested) {
             gpu_ring_handle = llama_dflash_cross_ring_gpu_init(ctx_dft, n_target_layers, n_embd, cross_ctx);
-            if (gpu_ring_handle) {
-                LOG_INF("dflash: GPU cross ring enabled (%d layers x %d slots x %d embd)\n",
-                        n_target_layers, cross_ctx, n_embd);
-            }
+        }
+        if (gpu_ring_handle) {
+            LOG_INF("dflash: GPU cross ring enabled (%d layers x %d slots x %d embd)\n",
+                    n_target_layers, cross_ctx, n_embd);
+        } else if (gpu_ring_requested) {
+            llama_set_dflash_gpu_capture(ctx_tgt, false);
+            LOG_WRN("dflash: GPU cross ring unavailable; using CPU hidden capture\n");
         }
 
         LOG_INF("dflash: block_size=%d, mask_token=%d, n_target_layers=%d, n_embd=%d\n",
diff --git a/ggml/src/ggml-cuda/ggml-cuda.cu b/ggml/src/ggml-cuda/ggml-cuda.cu
@@ -1389,6 +1389,12 @@ ggml_backend_buffer_type_t ggml_backend_cuda_host_buffer_type() {
     return &ggml_backend_cuda_buffer_type_host;
 }
 
+static bool ggml_cuda_buffer_visible_to_backend(ggml_backend_cuda_context * cuda_ctx, ggml_backend_buffer_type_t buft, bool integrated) {
+    return buft == ggml_backend_cuda_buffer_type(cuda_ctx->device) ||
+           ggml_backend_buft_is_cuda_split(buft) ||
+           (integrated && ggml_backend_buft_is_cuda_host(buft));
+}
+
 //static bool ggml_backend_buffer_is_cuda_host(ggml_backend_buffer_t buffer) {
 //    return buffer->buft->iface.get_name == ggml_backend_cuda_host_buffer_type_name;
 //}
@@ -3101,54 +3107,87 @@ static void ggml_backend_cuda_synchronize(ggml_backend_t backend) {
     GGML_UNUSED(backend);
 }
 
-static bool ggml_cuda_buffer_visible_to_backend(
-        ggml_backend_cuda_context * cuda_ctx,
-        ggml_backend_buffer_type_t buft,
-        bool integrated) {
-    return buft == ggml_backend_cuda_buffer_type(cuda_ctx->device) ||
-           ggml_backend_buft_is_cuda_split(buft) ||
-           (integrated && ggml_backend_buft_is_cuda_host(buft));
-}
-
-#ifndef NDEBUG
 static void ggml_cuda_log_nonlocal_src_buffer(
         ggml_backend_cuda_context * cuda_ctx,
         const ggml_tensor * node,
-        int src_idx) {
-    const ggml_tensor * src = node->src[src_idx];
+        int src_index,
+        bool integrated,
+        const char * action) {
+    const ggml_tensor * src = src_index >= 0 ? node->src[src_index] : node;
+    const void * data = src ? src->data : nullptr;
     int ptr_device = -1;
-    int ptr_type = -1;
+    const char * ptr_type = "unknown";
 
-    if (src && src->data) {
+    if (data) {
         cudaPointerAttributes attr;
-        cudaError_t err = cudaPointerGetAttributes(&attr, src->data);
+        cudaError_t err = cudaPointerGetAttributes(&attr, data);
         if (err == cudaSuccess) {
             ptr_device = attr.device;
 #if CUDART_VERSION >= 10000
-            ptr_type = attr.type;
+            switch (attr.type) {
 #else
-            ptr_type = attr.memoryType;
+            switch (attr.memoryType) {
 #endif
+                case cudaMemoryTypeHost:    ptr_type = "host";    break;
+                case cudaMemoryTypeDevice:  ptr_type = "device";  break;
+                case cudaMemoryTypeManaged: ptr_type = "managed"; break;
+                default:                    ptr_type = "other";   break;
+            }
         } else {
-            (void) cudaGetLastError();
+            cudaGetLastError();
         }
     }
 
     GGML_LOG_ERROR(
-        "%s: source buffer is not visible to CUDA backend device %d: node=%s op=%s src[%d]=%s src_buffer=%s src_ptr=%p ptr_device=%d ptr_type=%d dst_buffer=%s\n",
-        __func__,
-        cuda_ctx->device,
-        node->name,
-        ggml_op_name(node->op),
-        src_idx,
-        src ? src->name : "(null)",
-        src && src->buffer ? ggml_backend_buft_name(src->buffer->buft) : "(no buffer)",
-        src ? src->data : nullptr,
-        ptr_device,
-        ptr_type,
-        ggml_backend_buft_name(ggml_backend_cuda_buffer_type(cuda_ctx->device)));
+            "%s: source buffer is not visible to CUDA backend device; action=%s backend_device=%d integrated=%d node=%s op=%s src[%d]=%s src_buft=%s dst_buft=%s src_data=%p ptr_type=%s ptr_device=%d\n",
+            __func__,
+            action ? action : "fail",
+            cuda_ctx->device,
+            integrated ? 1 : 0,
+            node ? node->name : "(null)",
+            node ? ggml_op_name(node->op) : "(null)",
+            src_index,
+            src ? src->name : "(null)",
+            src && src->buffer ? ggml_backend_buft_name(src->buffer->buft) : "(none)",
+            node && node->buffer ? ggml_backend_buft_name(node->buffer->buft) : "(none)",
+            data,
+            ptr_type,
+            ptr_device);
+}
+
+static bool ggml_cuda_graph_node_buffers_visible(
+        ggml_backend_cuda_context * cuda_ctx,
+        const ggml_tensor * node,
+        bool integrated,
+        bool log_errors) {
+    if (!node || ggml_is_empty(node) || node->op == GGML_OP_RESHAPE || node->op == GGML_OP_TRANSPOSE ||
+            node->op == GGML_OP_VIEW || node->op == GGML_OP_PERMUTE || node->op == GGML_OP_NONE ||
+            (node->flags & GGML_TENSOR_FLAG_COMPUTE) == 0) {
+        return true;
+    }
+
+    if (!node->buffer || node->buffer->buft != ggml_backend_cuda_buffer_type(cuda_ctx->device)) {
+        if (log_errors) {
+            ggml_cuda_log_nonlocal_src_buffer(cuda_ctx, node, -1, integrated, "fail graph compute");
+        }
+        return false;
+    }
+
+    for (int j = 0; j < GGML_MAX_SRC; ++j) {
+        const ggml_tensor * src = node->src[j];
+        if (!src) {
+            continue;
+        }
+        if (!src->buffer || !ggml_cuda_buffer_visible_to_backend(cuda_ctx, src->buffer->buft, integrated)) {
+            if (log_errors) {
+                ggml_cuda_log_nonlocal_src_buffer(cuda_ctx, node, j, integrated, "fail graph compute");
+            }
+            return false;
+        }
+    }
+
+    return true;
 }
-#endif
 
 #ifdef USE_CUDA_GRAPH
 static bool ggml_cuda_graph_check_compability(ggml_backend_cuda_context * cuda_ctx, ggml_cgraph * cgraph) {
@@ -3169,7 +3208,6 @@ static bool ggml_cuda_graph_check_compability(ggml_backend_cuda_context * cuda_c
             if (!src || !src->buffer) {
                 continue;
             }
-
             const ggml_backend_buffer_type_t src_buft = src->buffer->buft;
             if (ggml_backend_buft_is_cuda_split(src_buft)) {
                 use_cuda_graph = false; // Split buffers are not supported by CUDA graph capture
@@ -3189,6 +3227,9 @@ static bool ggml_cuda_graph_check_compability(ggml_backend_cuda_context * cuda_c
                 break;
             }
         }
+        if (!use_cuda_graph) {
+            break;
+        }
 
         // [TAG_MUL_MAT_ID_CUDA_GRAPHS]
         if (node->op == GGML_OP_MUL_MAT_ID) {
@@ -4128,7 +4169,7 @@ static int ggml_cuda_try_fuse(ggml_backend_cuda_context * cuda_ctx, ggml_cgraph
 }
 
 
-static void ggml_cuda_graph_evaluate_and_capture(ggml_backend_cuda_context * cuda_ctx, ggml_cgraph * cgraph, const bool use_cuda_graph, const bool cuda_graph_update_required, const void * graph_key) {
+static enum ggml_status ggml_cuda_graph_evaluate_and_capture(ggml_backend_cuda_context * cuda_ctx, ggml_cgraph * cgraph, const bool use_cuda_graph, const bool cuda_graph_update_required, const void * graph_key) {
     bool graph_evaluated_or_captured = false;
 
     // flag used to determine whether it is an integrated_gpu
@@ -4280,34 +4321,15 @@ static void ggml_cuda_graph_evaluate_and_capture(ggml_backend_cuda_context * cud
                     i += nodes_to_skip;
                     continue;
                 }
-#ifndef NDEBUG
-                const ggml_backend_buffer_type_t local_buft = ggml_backend_cuda_buffer_type(cuda_ctx->device);
-                if (node->buffer->buft != local_buft) {
-                    GGML_LOG_ERROR("%s: node buffer is not local to CUDA backend device %d: node=%s op=%s node_buffer=%s dst_buffer=%s\n",
-                            __func__, cuda_ctx->device, node->name, ggml_op_name(node->op),
-                            ggml_backend_buft_name(node->buffer->buft), ggml_backend_buft_name(local_buft));
-                }
-                assert(node->buffer->buft == local_buft);
-                for (int j = 0; j < GGML_MAX_SRC; j++) {
-                    if (node->src[j] != nullptr) {
-                        assert(node->src[j]->buffer);
-                        const ggml_backend_buffer_type_t src_buft = node->src[j]->buffer->buft;
-                        const bool src_visible = ggml_cuda_buffer_visible_to_backend(cuda_ctx, src_buft, integrated);
-                        if (!src_visible) {
-                            ggml_cuda_log_nonlocal_src_buffer(cuda_ctx, node, j);
-                        }
-                        assert(src_visible);
-                    }
+                if (!ggml_cuda_graph_node_buffers_visible(cuda_ctx, node, integrated, true)) {
+                    return GGML_STATUS_FAILED;
                 }
-#else
-                GGML_UNUSED(integrated);
-#endif  // NDEBUG
 
                 bool ok = ggml_cuda_compute_forward(*cuda_ctx, node);
                 if (!ok) {
                     GGML_LOG_ERROR("%s: op not supported %s (%s)\n", __func__, node->name, ggml_op_name(node->op));
+                    return GGML_STATUS_FAILED;
                 }
-                GGML_ASSERT(ok);
 
                 if (!is_concurrent_event_active) {
                     try_launch_concurrent_event(node);
@@ -4350,6 +4372,8 @@ static void ggml_cuda_graph_evaluate_and_capture(ggml_backend_cuda_context * cud
         graph_evaluated_or_captured = true;
 #endif  // USE_CUDA_GRAPH
     }
+
+    return GGML_STATUS_SUCCESS;
 }
 
 #ifdef USE_CUDA_GRAPH
@@ -4413,6 +4437,13 @@ static enum ggml_status ggml_backend_cuda_graph_compute(ggml_backend_t backend,
     }
 #endif // USE_CUDA_GRAPH
 
+    const bool integrated = ggml_cuda_info().devices[cuda_ctx->device].integrated;
+    for (int i = 0; i < cgraph->n_nodes; ++i) {
+        if (!ggml_cuda_graph_node_buffers_visible(cuda_ctx, cgraph->nodes[i], integrated, true)) {
+            return GGML_STATUS_FAILED;
+        }
+    }
+
     if (use_cuda_graph && cuda_graph_update_required) {
         // Start CUDA graph capture
         {
@@ -4423,7 +4454,10 @@ static enum ggml_status ggml_backend_cuda_graph_compute(ggml_backend_t backend,
         CUDA_CHECK(cudaStreamBeginCapture(cuda_ctx->stream(), cudaStreamCaptureModeRelaxed));
     }
 
-    ggml_cuda_graph_evaluate_and_capture(cuda_ctx, cgraph, use_cuda_graph, cuda_graph_update_required, graph_key);
+    const ggml_status status = ggml_cuda_graph_evaluate_and_capture(cuda_ctx, cgraph, use_cuda_graph, cuda_graph_update_required, graph_key);
+    if (status != GGML_STATUS_SUCCESS) {
+        return status;
+    }
 
     return GGML_STATUS_SUCCESS;
 }
diff --git a/include/llama.h b/include/llama.h
@@ -1088,6 +1088,11 @@ extern "C" {
     // Pass n_layers=0 to disable capture
     LLAMA_API void llama_set_dflash_capture(struct llama_context * ctx, const int32_t * layer_ids, int32_t n_layers);
 
+    // DFlash: enable graph-embedded GPU hidden/tape capture for target decode.
+    // Disable this before decode when the drafter cannot consume GPU cross-ring
+    // tensors directly, so the eval callback keeps CPU hidden buffers populated.
+    LLAMA_API void llama_set_dflash_gpu_capture(struct llama_context * ctx, bool enabled);
+
     // DFlash: set drafter sampling temperature (Gumbel-max trick)
     // temp=0: greedy argmax (default), temp>0: sample from softmax(logits/temp)
     LLAMA_API void llama_set_dflash_sample_temp(struct llama_context * ctx, float temp);
diff --git a/src/llama-context.cpp b/src/llama-context.cpp
@@ -1432,6 +1432,30 @@ void llama_context::set_dflash_capture(const int32_t * layer_ids, int32_t n_laye
     }
 }
 
+void llama_context::set_dflash_gpu_capture(bool enabled) {
+    if (!dflash_capture) {
+        return;
+    }
+
+    dflash_capture->gpu_capture_enabled = enabled;
+
+    if (enabled) {
+        return;
+    }
+
+    dflash_capture->hidden_gpu.clear();
+    dflash_capture->tapes.clear();
+    cparams.tape_gpu = nullptr;
+    cparams.tape_gpu_n_seqs = 0;
+    cparams.hidden_gpu_n_seqs = 0;
+    for (int s = 0; s < (int) LLAMA_DFLASH_MAX_SLOTS; ++s) {
+        cparams.tape_gpu_seqs[s] = nullptr;
+        cparams.hidden_gpu_seqs[s] = nullptr;
+    }
+    cparams.cb_eval = dflash_eval_callback;
+    cparams.cb_eval_user_data = dflash_capture.get();
+}
+
 void llama_context::dflash_reset_hidden_capture() {
     if (!dflash_capture) {
         return;
@@ -1558,6 +1582,12 @@ void llama_context::allocate_tape_gpu(int n_slots, int max_tokens) {
     // populate recurrent-layer metadata if the caller beat set_tape_recording() to it
     dflash_ensure_recurrent_setup();
 
+    if (!dflash_capture->gpu_capture_enabled) {
+        dflash_capture->hidden_gpu.clear();
+        dflash_capture->tapes.clear();
+        return;
+    }
+
     if (model.n_devices() > 1) {
         dflash_capture->hidden_gpu.clear();
         dflash_capture->tapes.clear();
@@ -1673,6 +1703,10 @@ void llama_context::allocate_hidden_gpu(int n_slots, int max_tokens) {
     if (n_slots < 1) {
         n_slots = 1;
     }
+    if (!dflash_capture->gpu_capture_enabled) {
+        dflash_capture->hidden_gpu.clear();
+        return;
+    }
     if (model.n_devices() > 1) {
         dflash_capture->hidden_gpu.clear();
         return;
@@ -4234,7 +4268,7 @@ int llama_context::decode(const llama_batch & batch_inp) {
         // correct layer_hiddens slot. Populate per-seq tape pointers for the
         // graph builder so GPU tape copies target the correct per-slot buffers.
         if (dflash_capture) {
-            const bool dflash_gpu_capture_ready = model.n_devices() <= 1;
+            const bool dflash_gpu_capture_ready = model.n_devices() <= 1 && dflash_capture->gpu_capture_enabled;
             dflash_capture->ubatch = &ubatch;
             cparams.hidden_gpu_n_seqs = 0;
             for (int s = 0; s < (int) LLAMA_DFLASH_MAX_SLOTS; ++s) {
@@ -5890,6 +5924,10 @@ void llama_set_dflash_capture(llama_context * ctx, const int32_t * layer_ids, in
     ctx->set_dflash_capture(layer_ids, n_layers);
 }
 
+void llama_set_dflash_gpu_capture(llama_context * ctx, bool enabled) {
+    ctx->set_dflash_gpu_capture(enabled);
+}
+
 void llama_set_dflash_sample_temp(llama_context * ctx, float temp) {
     ctx->set_dflash_sample_temp(temp);
 }
diff --git a/src/llama-context.h b/src/llama-context.h
@@ -117,6 +117,7 @@ struct dflash_capture_data {
 
     // tape recording (for DeltaNet state rollback)
     bool tape_enabled = false;
+    bool gpu_capture_enabled = true;
     std::vector<int32_t> recurrent_layer_ids;       // model layer indices that are DeltaNet
     std::unordered_map<std::string, std::pair<int, int>> tape_name_map;  // name → (layer_idx, type)
     std::vector<dflash_tape_layer> tape_layers;     // one per recurrent layer (CPU fallback)
@@ -474,6 +475,7 @@ struct llama_context {
 
     // DFlash: configure hidden state capture layers
     void set_dflash_capture(const int32_t * layer_ids, int32_t n_layers);
+    void set_dflash_gpu_capture(bool enabled);
     void set_dflash_sample_temp(float temp);
     void set_dflash_topk(int k);
     void set_dflash_verify_logits(bool enabled, int top_k);
diff --git a/tests/test-dflash-plumbing.cpp b/tests/test-dflash-plumbing.cpp
