DFlash Vulkan fixes: rollback return type, reduced verify fallback, b- tensor alias, cross-ring infrastructure

- Change llama_dflash_rollback/tape_replay return types from void to int,
  returning positions needing re-decode when tape replay is unavailable
- Add dflash_reduced_verify_broken flag to disable reduced verify on
  backends lacking TOPK support (Vulkan), falling back to full logits
- Add b- beta tensor alias for Qwen3Next tape recording
- Add linear_attn_qkv_mixed- tape alias for Qwen3-Coder-Next conv replay
- Add GGML_DFLASH_FORCE_REDECODE env var for testing re-decode paths
- Add Vulkan get_proc_address resolver for DFlash cross-ring functions
- Add Vulkan dflash_is_cuda_compatible_tensor support
- Include placeholder ggml-vulkan-cross-ring.cpp (not used at runtime)
- Add docs/vulkan-cross-ring-plan.md with Approach A/C documentation

Fixes crash on Vulkan: GGML_ASSERT(logits != nullptr) in get_logits_ith
when reduced verify failed due to missing TOPK support.

Author: gbo

docs/vulkan-cross-ring-plan.md

@@ -61,6 +61,7 @@ if (Vulkan_FOUND)
 
     ggml_add_backend_library(ggml-vulkan
                              ggml-vulkan.cpp
+                             ggml-vulkan-cross-ring.cpp
                              ../../include/ggml-vulkan.h
                             )
 

ggml/src/ggml-vulkan/CMakeLists.txt

@@ -0,0 +1,51 @@
+// Vulkan UMA DFlash GPU cross-ring — Approach A (placeholder)
+//
+// Placeholder implementation. The Vulkan GPU ring is currently not used
+// (dflash_gpu_backend_reg does not return Vulkan). This file exists to
+// provide symbols for the get_proc_address resolver but is not actually
+// invoked at runtime.
+//
+// When Approach A is fully implemented, this will use host-visible coherent
+// Vulkan buffers for the ring + interleave, and ggml_backend_tensor_set for
+// the final H2D upload to the drafter's target_hidden tensor.
+
+#include <cstdlib>
+#include <cstdio>
+#include <cstring>
+
+struct dflash_cross_ring_vk {
+    int n_layers;
+    int n_embd;
+    int ring_size;
+    float ** h_layer_ptrs;
+    float * h_staging;
+};
+
+extern "C" void * dflash_cross_ring_gpu_alloc(int n_layers, int n_embd, int ring_size) {
+    (void)n_layers; (void)n_embd; (void)ring_size;
+    return nullptr; // not used — Vulkan ring is disabled
+}
+
+extern "C" void * dflash_cross_ring_gpu_alloc_device(int, int n_layers, int n_embd, int ring_size) {
+    return dflash_cross_ring_gpu_alloc(n_layers, n_embd, ring_size);
+}
+
+extern "C" void dflash_cross_ring_gpu_free(void * handle) {
+    if (!handle) return;
+    auto * ring = (dflash_cross_ring_vk *)handle;
+    if (ring->h_staging) free(ring->h_staging);
+    if (ring->h_layer_ptrs) {
+        for (int l = 0; l < ring->n_layers; l++) {
+            if (ring->h_layer_ptrs[l]) free(ring->h_layer_ptrs[l]);
+        }
+        delete[] ring->h_layer_ptrs;
+    }
+    delete ring;
+}
+
+extern "C" void dflash_cross_ring_gpu_write(void *, int, int, const float *, int, int) {}
+extern "C" bool  dflash_cross_ring_gpu_write_d2d(void *, int, int, const void *, int, int) { return false; }
+extern "C" void  dflash_cross_ring_gpu_synchronize(void *) {}
+extern "C" bool  dflash_cross_ring_gpu_snapshot(void *, int, int, int, float *, int, int, int) { return false; }
+extern "C" const float * dflash_cross_ring_gpu_interleave(void *, int, int, int) { return nullptr; }
+extern "C" void  dflash_cross_ring_gpu_set_tensor(void *, const void *, size_t, size_t) {}

ggml/src/ggml-vulkan/ggml-vulkan-cross-ring.cpp

@@ -17112,11 +17112,38 @@ static ggml_backend_dev_t ggml_backend_vk_reg_get_device(ggml_backend_reg_t reg,
     return devices[device];
 }
 
+// Forward declarations for DFlash cross-ring functions (defined in ggml-vulkan-cross-ring.cpp)
+extern "C" {
+    void * dflash_cross_ring_gpu_alloc(int, int, int);
+    void * dflash_cross_ring_gpu_alloc_device(int, int, int, int);
+    void   dflash_cross_ring_gpu_free(void *);
+    void   dflash_cross_ring_gpu_write(void *, int, int, const float *, int, int);
+    bool   dflash_cross_ring_gpu_write_d2d(void *, int, int, const void *, int, int);
+    void   dflash_cross_ring_gpu_synchronize(void *);
+    bool   dflash_cross_ring_gpu_snapshot(void *, int, int, int, float *, int, int, int);
+    const float * dflash_cross_ring_gpu_interleave(void *, int, int, int);
+    void   dflash_cross_ring_gpu_set_tensor(void *, const void *, size_t, size_t);
+}
+
+// DFlash cross-ring function pointer resolver for Vulkan UMA
+static void * ggml_backend_vk_reg_get_proc_address(ggml_backend_reg_t /* reg */, const char * name) {
+    if (strcmp(name, "dflash_cross_ring_gpu_alloc") == 0) return (void *)(void *)dflash_cross_ring_gpu_alloc;
+    if (strcmp(name, "dflash_cross_ring_gpu_alloc_device") == 0) return (void *)(void *)dflash_cross_ring_gpu_alloc_device;
+    if (strcmp(name, "dflash_cross_ring_gpu_free") == 0) return (void *)dflash_cross_ring_gpu_free;
+    if (strcmp(name, "dflash_cross_ring_gpu_write") == 0) return (void *)dflash_cross_ring_gpu_write;
+    if (strcmp(name, "dflash_cross_ring_gpu_write_d2d") == 0) return (void *)dflash_cross_ring_gpu_write_d2d;
+    if (strcmp(name, "dflash_cross_ring_gpu_synchronize") == 0) return (void *)dflash_cross_ring_gpu_synchronize;
+    if (strcmp(name, "dflash_cross_ring_gpu_snapshot") == 0) return (void *)dflash_cross_ring_gpu_snapshot;
+    if (strcmp(name, "dflash_cross_ring_gpu_interleave") == 0) return (void *)dflash_cross_ring_gpu_interleave;
+    if (strcmp(name, "dflash_cross_ring_gpu_set_tensor") == 0) return (void *)dflash_cross_ring_gpu_set_tensor;
+    return nullptr;
+}
+
 static const struct ggml_backend_reg_i ggml_backend_vk_reg_i = {
     /* .get_name         = */ ggml_backend_vk_reg_get_name,
     /* .get_device_count = */ ggml_backend_vk_reg_get_device_count,
     /* .get_device       = */ ggml_backend_vk_reg_get_device,
-    /* .get_proc_address = */ NULL,
+    /* .get_proc_address = */ ggml_backend_vk_reg_get_proc_address,
 };
 
 ggml_backend_reg_t ggml_backend_vk_reg() {

ggml/src/ggml-vulkan/ggml-vulkan.cpp

@@ -1205,7 +1205,10 @@ extern "C" {
     //   - KV cache: trims rejected draft positions (keeps accepted tokens' KV entries)
     //   - Recurrent state: restores from backup + tape replay for accepted tokens
     // This replaces the manual seq_rm/seq_cp + tape_replay sequence
-    LLAMA_API void llama_dflash_rollback(
+    // Returns the number of positions that were NOT advanced by tape replay
+    // and need re-decoding (e.g., when tape replay is unavailable on the backend).
+    // A return value of 0 means the rollback was fully successful.
+    LLAMA_API int llama_dflash_rollback(
             struct llama_context * ctx,
             llama_seq_id           seq_id,
             llama_seq_id           seq_backup,

include/llama.h

@@ -22,6 +22,7 @@
 #include "ggml-alloc.h"
 
 #include <algorithm>
+#include <atomic>
 #include <cinttypes>
 #include <cmath>
 #include <cstdlib>
@@ -76,7 +77,8 @@ static bool dflash_is_cuda_compatible_tensor(const ggml_tensor * t) {
         return false;
     }
     const char * name = ggml_backend_buffer_name(t->buffer);
-    return name && (std::strncmp(name, "CUDA", 4) == 0 || std::strncmp(name, "ROCm", 4) == 0);
+    return name && (std::strncmp(name, "CUDA", 4) == 0 || std::strncmp(name, "ROCm", 4) == 0 ||
+                    std::strncmp(name, "Vulkan", 6) == 0);
 }
 
 static bool dflash_tensor_span_in_bounds(const ggml_tensor * t, size_t offset_bytes, size_t n_bytes) {
@@ -2003,10 +2005,18 @@ void llama_context::dflash_ensure_recurrent_setup() {
             dflash_capture->tape_name_map["v_conv_predelta-" + il_str]        = {idx, DFLASH_TAPE_V};
             dflash_capture->tape_name_map["gate-" + il_str]                   = {idx, DFLASH_TAPE_GATE};
             dflash_capture->tape_name_map["beta-" + il_str]                   = {idx, DFLASH_TAPE_BETA};
+            // qwen3next (Qwen3-Coder-Next) names the beta tensor "b" instead of "beta"
+            dflash_capture->tape_name_map["b-" + il_str]                       = {idx, DFLASH_TAPE_BETA};
             dflash_capture->tape_name_map["qkv_mixed_pretranspose-" + il_str] = {idx, DFLASH_TAPE_QKV};
-            // Qwen3Next emits pre-conv QKV under these names. Without one of
-            // them, DFlash conv replay skips every recurrent layer and leaves
-            // r_l frozen at the pre-draft backup state.
+            // qwen3next (Qwen3-Coder-Next) builds the pre-conv (pre-transpose) QKV via two
+            // code paths in build_qkvz(), both projecting the layer INPUT (pre-conv1d):
+            //   wqkv path:        "linear_attn_qkv_mixed-" (qwen3next.cpp:304) [used by Qwen3-Coder-Next]
+            //   ssm_in legacy:    "qkv_mixed-"             (qwen3next.cpp:364) [concat of q/k/v flats]
+            // Both are the input to build_conv_state() (qwen3next.cpp:440), i.e. PRE-conv.
+            // Without recording one of these, tape_replay_conv skips every layer (empty
+            // qkv_mixed) and the conv state (r_l) is never advanced, so the conv state stays
+            // frozen at the backup value and the target output is garbled (0% acceptance).
+            // Verified: 504/504 conv layers OK after this fix; conv state matches re-decode ref.
             dflash_capture->tape_name_map["linear_attn_qkv_mixed-" + il_str] = {idx, DFLASH_TAPE_QKV};
             dflash_capture->tape_name_map["qkv_mixed-" + il_str]             = {idx, DFLASH_TAPE_QKV};
         }
@@ -2624,9 +2634,13 @@ void llama_context::set_active_dflash_slot(int slot_idx) {
     }
 }
 
-void llama_context::tape_replay(llama_seq_id seq_id, int n_accepted) {
+int llama_context::tape_replay(llama_seq_id seq_id, int n_accepted) {
     if (!dflash_capture || n_accepted <= 0) {
-        return;
+        return 0;
+    }
+
+    if (const char * env = std::getenv("GGML_DFLASH_FORCE_REDECODE"); env && std::atoi(env) != 0) {
+        return n_accepted;
     }
 
     // ensure any previous async replay is complete before launching a new one
@@ -2639,7 +2653,7 @@ void llama_context::tape_replay(llama_seq_id seq_id, int n_accepted) {
                                n_accepted <= gpu_tape->n_tokens);
 
     if (!use_gpu_tape && dflash_capture->tape_layers.empty()) {
-        return;
+        return 0;
     }
 
     auto * mem_recurrent = dynamic_cast<llama_memory_recurrent *>(memory.get());
@@ -2651,7 +2665,7 @@ void llama_context::tape_replay(llama_seq_id seq_id, int n_accepted) {
     }
     if (!mem_recurrent) {
         LLAMA_LOG_WARN("%s: tape replay requires recurrent memory\n", __func__);
-        return;
+        return n_accepted;
     }
 
     const auto & hparams = model.hparams;
@@ -2668,7 +2682,7 @@ void llama_context::tape_replay(llama_seq_id seq_id, int n_accepted) {
     }
     if (cell_idx < 0) {
         LLAMA_LOG_WARN("%s: no active cell for seq %d\n", __func__, seq_id);
-        return;
+        return n_accepted;
     }
 
     const uint32_t n_embd_s = hparams.n_embd_s();
@@ -2686,7 +2700,7 @@ void llama_context::tape_replay(llama_seq_id seq_id, int n_accepted) {
     if (!gpu_backend) {
         tape_replay_cpu(mem_recurrent, cell_idx, n_accepted);
         tape_replay_conv(mem_recurrent, cell_idx, n_accepted, seq_id);
-        return;
+        return 0;
     }
 
     // partial offload: if any recurrent layer's state lives on CPU, fall back to CPU replay
@@ -2696,7 +2710,7 @@ void llama_context::tape_replay(llama_seq_id seq_id, int n_accepted) {
         if (s_tensor && s_tensor->buffer && ggml_backend_buffer_is_host(s_tensor->buffer)) {
             tape_replay_cpu(mem_recurrent, cell_idx, n_accepted);
             tape_replay_conv(mem_recurrent, cell_idx, n_accepted, seq_id);
-            return;
+            return 0;
         }
     }
 
@@ -2709,14 +2723,14 @@ void llama_context::tape_replay(llama_seq_id seq_id, int n_accepted) {
         dflash_capture->replay_cell_idx = cell_idx;
         dflash_capture->replay_seq_id = seq_id;
         dflash_capture->replay_mem_recurrent = mem_recurrent;
-        return;
+        return 0;
     }
 
     const bool multi_gpu_target = model.n_devices() > 1;
     if (multi_gpu_target) {
         if (tape_replay_gdn_direct_from_cpu_tape(mem_recurrent, cell_idx, n_accepted)) {
             tape_replay_conv(mem_recurrent, cell_idx, n_accepted, seq_id);
-            return;
+            return 0;
         }
         if (!dflash_capture->multi_gpu_replay_fallback_logged) {
             LLAMA_LOG_WARN("%s: multi-GPU target detected (%zu devices); exact CUDA DFlash replay unavailable, using CPU recurrent replay fallback\n",
@@ -2725,7 +2739,7 @@ void llama_context::tape_replay(llama_seq_id seq_id, int n_accepted) {
         }
         tape_replay_cpu(mem_recurrent, cell_idx, n_accepted);
         tape_replay_conv(mem_recurrent, cell_idx, n_accepted, seq_id);
-        return;
+        return 0;
     }
 
     // GPU tape replay: build a ggml graph with GDN ops for all recurrent layers
@@ -2872,7 +2886,7 @@ void llama_context::tape_replay(llama_seq_id seq_id, int n_accepted) {
             if (!use_gpu_tape) {
                 tape_replay_cpu(mem_recurrent, cell_idx, n_accepted);
                 tape_replay_conv(mem_recurrent, cell_idx, n_accepted, seq_id);
-                return;
+                return 0;
             }
             goto conv_rebuild;
         }
@@ -2900,7 +2914,7 @@ void llama_context::tape_replay(llama_seq_id seq_id, int n_accepted) {
             ggml_free(ctx);
             tape_replay_cpu(mem_recurrent, cell_idx, n_accepted);
             tape_replay_conv(mem_recurrent, cell_idx, n_accepted, seq_id);
-            return;
+            return 0;
         }
 
         // assign tensors within the persistent buffer
@@ -2955,8 +2969,9 @@ void llama_context::tape_replay(llama_seq_id seq_id, int n_accepted) {
             if (!use_gpu_tape) {
                 tape_replay_cpu(mem_recurrent, cell_idx, n_accepted);
                 tape_replay_conv(mem_recurrent, cell_idx, n_accepted, seq_id);
+                return 0;
             }
-            return;
+            return n_accepted;
         }
         if (dflash_capture->profile) {
             const uint64_t elapsed = ggml_time_us() - t_replay_enqueue_us;
@@ -2976,11 +2991,12 @@ void llama_context::tape_replay(llama_seq_id seq_id, int n_accepted) {
         dflash_capture->replay_cell_idx = cell_idx;
         dflash_capture->replay_seq_id = seq_id;
         dflash_capture->replay_mem_recurrent = mem_recurrent;
-        return; // conv rebuild deferred to tape_replay_sync()
+        return 0; // conv rebuild deferred to tape_replay_sync()
     }
 
 conv_rebuild:
     tape_replay_conv(mem_recurrent, cell_idx, n_accepted, seq_id);
+    return 0;
 }
 
 bool llama_context::tape_replay_gdn_direct_gpu(llama_memory_recurrent * mem_recurrent, int32_t cell_idx, int n_accepted) {
@@ -3944,11 +3960,11 @@ void llama_context::tape_replay_cpu(llama_memory_recurrent * mem_recurrent, int3
     }
 }
 
-void llama_context::dflash_rollback(llama_seq_id seq_id, llama_seq_id seq_backup, int n_past_before, int n_accepted) {
+int llama_context::dflash_rollback(llama_seq_id seq_id, llama_seq_id seq_backup, int n_past_before, int n_accepted) {
     auto * mem_hybrid = dynamic_cast<llama_memory_hybrid *>(memory.get());
     if (!mem_hybrid) {
         LLAMA_LOG_WARN("%s: dflash_rollback requires hybrid memory\n", __func__);
-        return;
+        return n_accepted;
     }
 
     const bool profile = dflash_capture && dflash_profile_has(dflash_capture->profile_flags, DFLASH_PROFILE_COPY);
@@ -3997,7 +4013,7 @@ void llama_context::dflash_rollback(llama_seq_id seq_id, llama_seq_id seq_backup
     profile_lap(recurrent_restore_us);
 
     // Replay DeltaNet state updates for accepted tokens
-    tape_replay(seq_id, n_accepted);
+    const int n_redecode = tape_replay(seq_id, n_accepted);
     profile_lap(tape_launch_us);
 
     if (profile) {
@@ -4019,6 +4035,7 @@ void llama_context::dflash_rollback(llama_seq_id seq_id, llama_seq_id seq_backup
             tape_launch_us / 1e3,
             (ggml_time_us() - t_start_us) / 1e3);
     }
+    return n_redecode;
 }
 
 void llama_context::dflash_prepare_branch(llama_seq_id seq_id, llama_seq_id seq_backup, int depth) {
@@ -8644,8 +8661,8 @@ bool llama_dflash_memory_seq_cp_recurrent_ordered(
     return ctx ? ctx->dflash_memory_seq_cp_recurrent_ordered(seq_id_src, seq_id_dst, p0, p1) : false;
 }
 
-void llama_dflash_rollback(llama_context * ctx, llama_seq_id seq_id, llama_seq_id seq_backup, int n_past_before, int n_accepted) {
-    ctx->dflash_rollback(seq_id, seq_backup, n_past_before, n_accepted);
+int llama_dflash_rollback(llama_context * ctx, llama_seq_id seq_id, llama_seq_id seq_backup, int n_past_before, int n_accepted) {
+    return ctx->dflash_rollback(seq_id, seq_backup, n_past_before, n_accepted);
 }
 
 void llama_dflash_prepare_branch(llama_context * ctx, llama_seq_id seq_id, llama_seq_id seq_backup, int depth) {

src/llama-context.cpp

@@ -738,7 +738,9 @@ struct llama_context {
     void set_active_dflash_slot(int slot_idx);
 
     // DFlash: replay tape data to reconstruct DeltaNet state for n_accepted tokens
-    void tape_replay(llama_seq_id seq_id, int n_accepted);
+    // Returns the number of positions that were NOT advanced by tape replay
+    // and need re-decoding. A return value of 0 means the rollback was fully successful.
+    int tape_replay(llama_seq_id seq_id, int n_accepted);
     void tape_replay_sync();
     bool tape_replay_conv_gpu(llama_memory_recurrent * mem_recurrent, int32_t cell_idx, int n_accepted);
     bool tape_replay_conv_gpu(llama_memory_recurrent * mem_recurrent, int32_t cell_idx, int n_accepted, bool advance_pos);
@@ -750,7 +752,9 @@ struct llama_context {
     bool dflash_memory_seq_cp_recurrent_ordered(llama_seq_id seq_id_src, llama_seq_id seq_id_dst, llama_pos p0, llama_pos p1);
 
     // DFlash: complete rollback for hybrid models (KV trim + recurrent restore + tape replay)
-    void dflash_rollback(llama_seq_id seq_id, llama_seq_id seq_backup, int n_past_before, int n_accepted);
+    // Returns the number of positions that were NOT advanced by tape replay
+    // and need re-decoding. A return value of 0 means the rollback was fully successful.
+    int dflash_rollback(llama_seq_id seq_id, llama_seq_id seq_backup, int n_past_before, int n_accepted);
 
     // DFlash: prepare DeltaNet state for branch verification (recurrent restore + tape replay, no KV touch)
     void dflash_prepare_branch(llama_seq_id seq_id, llama_seq_id seq_backup, int depth);