fix: support DFlash tensor-split Meta placement

Keep auto-detected DFlash drafters on tensor-split placement when target shared tensors are Meta-backed, add Meta-safe capture/ring/recurrent fallbacks, and cover the plumbing with focused regression checks.

Author: Anbeeld

CHANGELOG.md

@@ -1,5 +1,11 @@
 # Changelog
 
+## v0.3.2
+
+- Fixed DFlash on target tensor-split / Meta placement. Auto-detected DFlash draft loading now keeps tensor-split placement when shared target output tensors live in a Meta buffer, instead of reloading the drafter onto one GPU and crashing during draft warmup. Explicit incompatible `--spec-draft-device` placement now fails closed with a direct diagnostic.
+- Added the Meta backend support needed by DFlash tensor-split runs: TurboQuant WHT split-state propagation, public Meta device introspection for capture placement, Meta-backed hidden/prefill capture allocation, Meta-safe GPU ring readback fallback, Meta-safe recurrent backup row copies, and a full-logits drafter fallback when compact argmax/top-k cannot run on Meta output.
+- Verified the single-RTX-3090 tensor-split crash proxy with DFlash enabled. Forced `-sm tensor` now completes and generates drafts; normal single-GPU DFlash remains on the fast CUDA path. The forced single-GPU tensor-split path is still a correctness fallback and is slower because recurrent backup and hidden-ring transfer avoid unsafe raw Meta pointers.
+
 ## v0.3.1
 
 - Merged latest upstream llama.cpp master. This pulls in Gemma 4 12B and Gemma 4 unified multimodal support fixes, including non-causal vision, unified audio/vision projector handling, and FPE fixes; Qwen3.5 post-norm hidden-state behavior for MTP; CUDA KV-cache quantization preallocation and PDL race fixes; WebGPU FlashAttention refactoring with standardized quantization support; CPU backend improvements for RVV/SVE; lower-latency Metal command-buffer status polling; Mermaid diagram rendering and preview support in `tools/ui`; updated BoringSSL, SYCL documentation, save/load-state tests, Docker docs, and small CI/release maintenance.

common/speculative.cpp

@@ -4296,6 +4296,14 @@ llama_context * common_speculative_create_ctx_dft(const common_params_speculativ
     if (params.sample_temp > 0.0f) {
         llama_set_dflash_sample_temp(ctx_dft, params.sample_temp);
     }
+    bool draft_reduced_logits = params.draft.backend_sampling;
+    ggml_backend_dev_t draft_output_dev = llama_model_dev_output(params.model_dft);
+    if (draft_output_dev && ggml_backend_dev_type(draft_output_dev) == GGML_BACKEND_DEVICE_TYPE_META) {
+        draft_reduced_logits = false;
+        LOG_WRN("%s",
+                "dflash: draft backend sampling disabled on Meta output placement; using full-logits fallback\n");
+    }
+    llama_set_dflash_consume_reduced(ctx_dft, draft_reduced_logits);
 
     // warmup the draft context
     {

ggml/include/ggml-backend.h

@@ -401,6 +401,9 @@ extern "C" {
     //       express this as a backend registry functionality instead
     GGML_API ggml_backend_dev_t ggml_backend_meta_device(
         ggml_backend_dev_t * devs, size_t n_devs, ggml_backend_meta_get_split_state_t get_split_state, void * get_split_state_ud);
+    GGML_API bool ggml_backend_dev_is_meta(ggml_backend_dev_t dev);
+    GGML_API size_t ggml_backend_meta_dev_n_devs(ggml_backend_dev_t meta_dev);
+    GGML_API ggml_backend_dev_t ggml_backend_meta_dev_simple_dev(ggml_backend_dev_t meta_dev, size_t index);
 
     //
     // Utils

ggml/src/ggml-backend-meta.cpp

@@ -82,8 +82,6 @@ struct ggml_backend_meta_device_context {
     }
 };
 
-static bool ggml_backend_dev_is_meta(ggml_backend_dev_t dev);
-
 static const char * ggml_backend_meta_device_get_name(ggml_backend_dev_t dev) {
     GGML_ASSERT(ggml_backend_dev_is_meta(dev));
     const ggml_backend_meta_device_context * meta_dev_ctx = (const ggml_backend_meta_device_context *) dev->context;
@@ -193,17 +191,17 @@ static const ggml_backend_device_i ggml_backend_meta_device_iface = {
     /* .event_synchronize    = */ nullptr,
 };
 
-static bool ggml_backend_dev_is_meta(ggml_backend_dev_t dev) {
+bool ggml_backend_dev_is_meta(ggml_backend_dev_t dev) {
     return dev != nullptr && dev->iface.get_name == ggml_backend_meta_device_iface.get_name;
 }
 
-static size_t ggml_backend_meta_dev_n_devs(ggml_backend_dev_t meta_dev) {
+size_t ggml_backend_meta_dev_n_devs(ggml_backend_dev_t meta_dev) {
     GGML_ASSERT(ggml_backend_dev_is_meta(meta_dev));
     const ggml_backend_meta_device_context * meta_dev_ctx = (const ggml_backend_meta_device_context *) meta_dev->context;
     return meta_dev_ctx->simple_devs.size();
 }
 
-static ggml_backend_dev_t ggml_backend_meta_dev_simple_dev(ggml_backend_dev_t meta_dev, size_t index) {
+ggml_backend_dev_t ggml_backend_meta_dev_simple_dev(ggml_backend_dev_t meta_dev, size_t index) {
     GGML_ASSERT(ggml_backend_dev_is_meta(meta_dev));
     const ggml_backend_meta_device_context * meta_dev_ctx = (const ggml_backend_meta_device_context *) meta_dev->context;
     GGML_ASSERT(index < meta_dev_ctx->simple_devs.size());
@@ -538,7 +536,10 @@ static struct ggml_backend_meta_split_state ggml_backend_meta_get_split_state(
 
     // Some ops process data on a per-row bases:
     auto handle_per_row = [&](const std::vector<ggml_backend_meta_split_state> & src_ss) -> ggml_backend_meta_split_state {
-        GGML_ASSERT(src_ss[0].axis != GGML_BACKEND_SPLIT_AXIS_0);
+        if (src_ss[0].axis == GGML_BACKEND_SPLIT_AXIS_0) {
+            GGML_ABORT("meta backend per-row op %s (%s) cannot use axis-0 split src0 %s (%s)",
+                    tensor->name, ggml_op_name(tensor->op), tensor->src[0]->name, ggml_op_name(tensor->src[0]->op));
+        }
         return src_ss[0];
     };
 
@@ -782,6 +783,20 @@ static struct ggml_backend_meta_split_state ggml_backend_meta_get_split_state(
         return {GGML_BACKEND_SPLIT_AXIS_0, {0}, {1}, 1};
     };
 
+    auto handle_turbo_wht = [&](const std::vector<ggml_backend_meta_split_state> & src_ss) -> ggml_backend_meta_split_state {
+        ggml_backend_meta_split_state ret = handle_generic(src_ss, /*scalar_only =*/ false);
+        if (ret.axis == GGML_BACKEND_SPLIT_AXIS_0) {
+            const size_t n_bufs = ggml_backend_meta_buffer_n_bufs(tensor->buffer);
+            for (size_t s = 0; s < ret.n_segments; ++s) {
+                for (size_t j = 0; j < n_bufs; ++j) {
+                    const int64_t turbo_wht_split_axis_ne = ret.ne[s*n_bufs + j];
+                    GGML_ASSERT(turbo_wht_split_axis_ne % 128 == 0);
+                }
+            }
+        }
+        return ret;
+    };
+
     auto calculate_split_state = [&]() -> ggml_backend_meta_split_state {
         if (ggml_nelements(tensor) == 0) {
             return {GGML_BACKEND_SPLIT_AXIS_UNKNOWN, {0}, {1}, 1};
@@ -1002,6 +1017,9 @@ static struct ggml_backend_meta_split_state ggml_backend_meta_get_split_state(
             case GGML_OP_GLU: {
                 split_state = handle_generic(src_ss, /*scalar_only =*/ false);
             } break;
+            case GGML_OP_TURBO_WHT: {
+                split_state = handle_turbo_wht(src_ss);
+            } break;
             default: {
                 GGML_ABORT("ggml op not implemented: %s", ggml_op_name(tensor->op));
                 split_state = {GGML_BACKEND_SPLIT_AXIS_UNKNOWN, {0}, {1}, 1};

src/llama-context.cpp

@@ -170,16 +170,64 @@ static void dflash_log_backend_layout(
 
 static ggml_backend_t dflash_backend_for_dev(
         const std::vector<ggml_backend_ptr> & backends,
-        ggml_backend_dev_t want_dev) {
+        ggml_backend_dev_t want_dev,
+        bool allow_meta = false) {
     for (const auto & backend : backends) {
         auto * dev = ggml_backend_get_device(backend.get());
-        if (dev == want_dev && dflash_backend_dev_is_gpu(dev)) {
+        if (dev == want_dev && (dflash_backend_dev_is_gpu(dev) ||
+                (allow_meta && ggml_backend_dev_is_meta(dev)))) {
             return backend.get();
         }
     }
     return nullptr;
 }
 
+struct dflash_capture_backend {
+    ggml_backend_t     backend = nullptr;
+    ggml_backend_dev_t dev = nullptr;
+    bool               from_meta = false;
+};
+
+static dflash_capture_backend dflash_capture_backend_for_layer(
+        const std::vector<ggml_backend_ptr> & backends,
+        ggml_backend_dev_t layer_dev) {
+    dflash_capture_backend result;
+
+    result.backend = dflash_backend_for_dev(backends, layer_dev, layer_dev && ggml_backend_dev_is_meta(layer_dev));
+    if (result.backend) {
+        result.dev = layer_dev;
+        result.from_meta = layer_dev && ggml_backend_dev_is_meta(layer_dev);
+        return result;
+    }
+
+    if (!layer_dev || !ggml_backend_dev_is_meta(layer_dev)) {
+        return result;
+    }
+
+    const size_t n_devs = ggml_backend_meta_dev_n_devs(layer_dev);
+    for (size_t i = 0; i < n_devs; ++i) {
+        ggml_backend_dev_t simple_dev = ggml_backend_meta_dev_simple_dev(layer_dev, i);
+        result.backend = dflash_backend_for_dev(backends, simple_dev);
+        if (result.backend) {
+            result.dev = simple_dev;
+            result.from_meta = true;
+            return result;
+        }
+    }
+
+    return result;
+}
+
+static bool dflash_tensor_buffer_is_meta(const ggml_tensor * tensor) {
+    if (!tensor || !tensor->buffer) {
+        return false;
+    }
+
+    ggml_backend_buffer_type_t buft = ggml_backend_buffer_get_type(tensor->buffer);
+    ggml_backend_dev_t dev = buft ? ggml_backend_buft_get_device(buft) : nullptr;
+    return ggml_backend_dev_is_meta(dev);
+}
+
 static void dflash_capture_add_wait_backend(
         dflash_capture_data & cap,
         ggml_backend_t backend,
@@ -2161,8 +2209,8 @@ void llama_context::allocate_tape_gpu(int n_slots, int max_tokens) {
         for (int li = 0; li < n_rec; ++li) {
             const int il = rec_ids[li];
             ggml_backend_dev_t layer_dev = model.dev_layer(il);
-            ggml_backend_t layer_backend = dflash_backend_for_dev(backends, layer_dev);
-            if (!layer_backend) {
+            const dflash_capture_backend capture_backend = dflash_capture_backend_for_layer(backends, layer_dev);
+            if (!capture_backend.backend) {
                 LLAMA_LOG_WARN("%s: no GPU backend for recurrent layer %d device %s; falling back to CPU tape\n",
                     __func__, il, layer_dev ? ggml_backend_dev_name(layer_dev) : "<null>");
                 dflash_capture->tapes.clear();
@@ -2193,29 +2241,29 @@ void llama_context::allocate_tape_gpu(int n_slots, int max_tokens) {
             tl.qkv  = ggml_new_tensor_2d(tape_ctx, GGML_TYPE_F32, conv_ch, (int64_t)max_tokens);
 
             tl.ctx = tape_ctx;
-            tl.dev = layer_dev;
-            tl.buf = ggml_backend_alloc_ctx_tensors(tape_ctx, layer_backend);
+            tl.dev = capture_backend.dev;
+            tl.buf = ggml_backend_alloc_ctx_tensors(tape_ctx, capture_backend.backend);
 
             if (!tl.buf) {
                 LLAMA_LOG_WARN("%s: failed to allocate GPU tape buffer for slot %d layer %d device %s, falling back to CPU tape\n",
-                    __func__, slot, il, layer_dev ? ggml_backend_dev_name(layer_dev) : "<null>");
+                    __func__, slot, il, capture_backend.dev ? ggml_backend_dev_name(capture_backend.dev) : "<null>");
                 dflash_capture->tapes.clear();
                 return;
             }
 
-            dflash_capture_add_wait_backend(*dflash_capture, layer_backend, ggml_backend_dev_backend_reg(layer_dev));
+            dflash_capture_add_wait_backend(*dflash_capture, capture_backend.backend, ggml_backend_dev_backend_reg(capture_backend.dev));
             total_size += ggml_backend_buffer_get_size(tl.buf);
 
             bool found = false;
             for (auto & dc : dev_counts) {
-                if (dc.dev == layer_dev) {
+                if (dc.dev == capture_backend.dev) {
                     ++dc.count;
                     found = true;
                     break;
                 }
             }
             if (!found) {
-                dev_counts.push_back({ layer_dev, 1 });
+                dev_counts.push_back({ capture_backend.dev, 1 });
             }
         }
 
@@ -2275,8 +2323,8 @@ void llama_context::allocate_hidden_gpu(int n_slots, int max_tokens) {
         for (int i = 0; i < n_layers; ++i) {
             const int il = dflash_capture->layer_ids[i];
             ggml_backend_dev_t layer_dev = model.dev_layer(il);
-            ggml_backend_t layer_backend = dflash_backend_for_dev(backends, layer_dev);
-            if (!layer_backend) {
+            const dflash_capture_backend capture_backend = dflash_capture_backend_for_layer(backends, layer_dev);
+            if (!capture_backend.backend) {
                 LLAMA_LOG_WARN("%s: no GPU backend for hidden layer %d device %s; using callback hidden fallback\n",
                     __func__, il, layer_dev ? ggml_backend_dev_name(layer_dev) : "<null>");
                 dflash_capture->hidden_gpu.clear();
@@ -2296,15 +2344,15 @@ void llama_context::allocate_hidden_gpu(int n_slots, int max_tokens) {
 
             hidden->layers[i] = ggml_new_tensor_2d(hidden_ctx, GGML_TYPE_F32, n_embd, (int64_t) max_tokens);
 
-            ggml_backend_buffer_t buf = ggml_backend_alloc_ctx_tensors(hidden_ctx, layer_backend);
+            ggml_backend_buffer_t buf = ggml_backend_alloc_ctx_tensors(hidden_ctx, capture_backend.backend);
             if (!buf) {
                 LLAMA_LOG_WARN("%s: failed to allocate GPU hidden buffer for slot %d layer %d; using callback hidden fallback\n",
                     __func__, slot, il);
                 dflash_capture->hidden_gpu.clear();
                 return;
             }
             hidden->bufs.push_back(buf);
-            dflash_capture_add_wait_backend(*dflash_capture, layer_backend, ggml_backend_dev_backend_reg(layer_dev));
+            dflash_capture_add_wait_backend(*dflash_capture, capture_backend.backend, ggml_backend_dev_backend_reg(capture_backend.dev));
             total_size += ggml_backend_buffer_get_size(buf);
         }
 
@@ -2361,8 +2409,8 @@ bool llama_context::allocate_prefill_gpu(int n_slots, int max_tokens) {
         for (int i = 0; i < n_layers; ++i) {
             const int il = dflash_capture->layer_ids[i];
             ggml_backend_dev_t layer_dev = model.dev_layer(il);
-            ggml_backend_t layer_backend = dflash_backend_for_dev(backends, layer_dev);
-            if (!layer_backend) {
+            const dflash_capture_backend capture_backend = dflash_capture_backend_for_layer(backends, layer_dev);
+            if (!capture_backend.backend) {
                 LLAMA_LOG_WARN("%s: no GPU backend for prefill layer %d device %s; using callback fallback\n",
                     __func__, il, layer_dev ? ggml_backend_dev_name(layer_dev) : "<null>");
                 dflash_capture->prefill_gpu.clear();
@@ -2384,7 +2432,7 @@ bool llama_context::allocate_prefill_gpu(int n_slots, int max_tokens) {
 
             hidden->layers[i] = ggml_new_tensor_2d(hidden_ctx, GGML_TYPE_F32, n_embd, (int64_t) max_tokens);
 
-            ggml_backend_buffer_t buf = ggml_backend_alloc_ctx_tensors(hidden_ctx, layer_backend);
+            ggml_backend_buffer_t buf = ggml_backend_alloc_ctx_tensors(hidden_ctx, capture_backend.backend);
             if (!buf) {
                 LLAMA_LOG_WARN("%s: failed to allocate prefill GPU buffer for slot %d layer %d; using callback fallback\n",
                     __func__, slot, il);
@@ -2393,7 +2441,7 @@ bool llama_context::allocate_prefill_gpu(int n_slots, int max_tokens) {
                 return false;
             }
             hidden->bufs.push_back(buf);
-            dflash_capture_add_wait_backend(*dflash_capture, layer_backend, ggml_backend_dev_backend_reg(layer_dev));
+            dflash_capture_add_wait_backend(*dflash_capture, capture_backend.backend, ggml_backend_dev_backend_reg(capture_backend.dev));
             total_size += ggml_backend_buffer_get_size(buf);
         }
 
@@ -8792,7 +8840,8 @@ bool llama_context::cross_ring_gpu_write_hidden(void * handle, int layer, int ri
     }
 
     auto * tensor = hgpu->layers[layer];
-    if (!tensor || !tensor->data) {
+    const bool tensor_is_meta = dflash_tensor_buffer_is_meta(tensor);
+    if (!tensor || !tensor->buffer || (!tensor_is_meta && !tensor->data)) {
         return false;
     }
     if (!dflash_gpu_hidden_span_in_bounds(tensor, src_offset, n_tokens, n_embd, __func__)) {
@@ -8801,15 +8850,17 @@ bool llama_context::cross_ring_gpu_write_hidden(void * handle, int layer, int ri
 
     auto * h = (dflash_cross_ring_handle *)handle;
     const size_t src_offset_bytes = (size_t) src_offset * (size_t) n_embd * sizeof(float);
-    const void * src = (const char *) tensor->data + src_offset_bytes;
-    if (h->fn_write_d2d(h->gpu_ring, layer, ring_pos, src, n_tokens, n_embd)) {
-        return true;
+    if (!tensor_is_meta) {
+        const void * src = (const char *) tensor->data + src_offset_bytes;
+        if (h->fn_write_d2d(h->gpu_ring, layer, ring_pos, src, n_tokens, n_embd)) {
+            return true;
+        }
     }
 
     static bool warned_d2h_fallback = false;
     if (!warned_d2h_fallback) {
-        LLAMA_LOG_WARN("%s: GPU hidden D2D ring write unavailable; falling back to GPU readback + H2D ring upload\n",
-            __func__);
+        LLAMA_LOG_WARN("%s: GPU hidden D2D ring write unavailable%s; falling back to backend readback + H2D ring upload\n",
+            __func__, tensor_is_meta ? " for Meta capture tensor" : "");
         warned_d2h_fallback = true;
     }
 
@@ -8843,7 +8894,8 @@ bool llama_context::prefill_gpu_write_hidden(void * handle, int slot, int layer,
     }
 
     auto * tensor = pgpu->layers[layer];
-    if (!tensor || !tensor->data) {
+    const bool tensor_is_meta = dflash_tensor_buffer_is_meta(tensor);
+    if (!tensor || !tensor->buffer || (!tensor_is_meta && !tensor->data)) {
         return false;
     }
     if (!dflash_gpu_hidden_span_in_bounds(tensor, src_offset, n_tokens, n_embd, __func__)) {
@@ -8852,15 +8904,17 @@ bool llama_context::prefill_gpu_write_hidden(void * handle, int slot, int layer,
 
     auto * h = (dflash_cross_ring_handle *)handle;
     const size_t src_offset_bytes = (size_t) src_offset * (size_t) n_embd * sizeof(float);
-    const void * src = (const char *) tensor->data + src_offset_bytes;
-    if (h->fn_write_d2d(h->gpu_ring, layer, ring_pos, src, n_tokens, n_embd)) {
-        return true;
+    if (!tensor_is_meta) {
+        const void * src = (const char *) tensor->data + src_offset_bytes;
+        if (h->fn_write_d2d(h->gpu_ring, layer, ring_pos, src, n_tokens, n_embd)) {
+            return true;
+        }
     }
 
     static bool warned_prefill_d2h_fallback = false;
     if (!warned_prefill_d2h_fallback) {
-        LLAMA_LOG_WARN("%s: prefill GPU D2D ring write unavailable; falling back to GPU readback + H2D ring upload\n",
-            __func__);
+        LLAMA_LOG_WARN("%s: prefill GPU D2D ring write unavailable%s; falling back to backend readback + H2D ring upload\n",
+            __func__, tensor_is_meta ? " for Meta capture tensor" : "");
         warned_prefill_d2h_fallback = true;
     }