diff --git a/VULKAN_BUILD_VERIFICATION.md b/VULKAN_BUILD_VERIFICATION.md new file mode 100644 index 000000000000..ae01dceaa078 --- /dev/null +++ b/VULKAN_BUILD_VERIFICATION.md @@ -0,0 +1,153 @@ +# Vulkan Build Verification for PR #79 + +Verifies that PR #79 builds cleanly for the Vulkan backend on independent host environments. + +## PR head history + +The PR began as a single Qwen3-Coder-Next Vulkan enablement (`b788b4af1`). It has +since advanced with layered DFlash fixes, all of which must continue to build +cleanly for Vulkan: + +| Commit | Change | +|---|---| +| `b788b4af1` | `dflash: enable Qwen3-Coder-Next on Vulkan` (original PR head) | +| `4efe54303` | updated PR head at the time of Verification 2 | +| `bcd42973f` | sync the latest working DFlash code from `vulkan-dflash-cross-ring` (fixes Qwen3.6 divergence + Coder-Next crash on Vulkan) | +| `ab2be17fd` | keep `n_rs_seq` for QWEN35MOE (fix 122B GPU cross-ring corruption) | +| `efb07137c` | rotate DeltaNet RS snapshots for `n_seq_tokens llama_decode returns 1). The + // DFlash drafter's self-attention is block-local (non-causal over the + // query block) plus cross-attention to target hiddens, so evicting the + // oldest accepted-prefix cells is safe: that history is never attended. + // This previously ran only with the GPU ring enabled, which left the + // GGML_DFLASH_GPU_RING=0 (CPU cross) path filling its small KV and + // intermittently failing; the slide has no ring dependency. trim_drafter_prefix_window(); + if (!gpu_ring_handle) { + return; + } + const int n_update = std::min(n_written, cross_ctx); const int n_full_kv_update = std::min(n_update, drafter_prefix_window()); const int gpu_write_pos = ring_write_pos % cross_ctx; @@ -2574,6 +2606,19 @@ struct common_speculative_impl_dflash : public common_speculative_impl { llama_set_dflash_gpu_capture(ctx_tgt, false); LOG_WRN("dflash: GPU cross ring unavailable; using CPU hidden capture\n"); } + // Tell the drafter graph whether its normal KV cache will be populated with + // TARGET context K/V from the ring. This needs the backend's KV D2D helpers + // (dflash_kv_cache_write_d2d / _append_d2d / _interleave), which are + // CUDA-only today; Vulkan does not register them, so llama_dflash_kv_cache_init + // returns false there (the function lookup fails before any allocation). + // Enabling the target-KV path without a populated cache makes the drafter + // self-attend with empty target KV -> ~0% acceptance (the ring-ON regression: + // Coder-Next ring-ON went from ~1.2% to ~58% once this was gated on the actual + // KV-cache availability instead of just the ring handle). Gate the flag on + // the real KV-cache availability, not merely on the ring handle existing. + llama_dflash_kv_cache_set_active_seq(ctx_dft, 0); + const bool drafter_kv_cache_ok = llama_dflash_kv_cache_init(ctx_dft, cross_ctx); + llama_set_dflash_target_kv_available(ctx_dft, gpu_ring_handle != nullptr && drafter_kv_cache_ok); } ~common_speculative_impl_dflash() override { @@ -2656,7 +2701,17 @@ struct common_speculative_impl_dflash : public common_speculative_impl { if (!validate_target_hiddens("begin")) { return; } + if (common_dflash_rx_diag_enabled()) { + LOG_INF("DFLASH_RX begin: pre_capture seq=%d ring_filled=%d committed_len=%d prefill_flushed=%d prefill_flush_written=%d prefill_suffix_seen=%d\n", + seq_id, ring_filled, committed_len, + prefill_flushed ? 1 : 0, prefill_flush_written, + prefill_suffix_seen ? 1 : 0); + } capture_target_hiddens(); + if (common_dflash_rx_diag_enabled()) { + LOG_INF("DFLASH_RX begin: post_capture seq=%d ring_filled=%d committed_len=%d\n", + seq_id, ring_filled, committed_len); + } } bool process(const llama_batch & /*batch*/) override { @@ -2982,10 +3037,25 @@ struct common_speculative_impl_dflash : public common_speculative_impl { const llama_token id_last = dp.id_last; const int32_t n_max_eff = dp.n_max > 0 ? dp.n_max : (block_size - 1); const int n_draft = std::min(block_size - 1, n_max_eff); + + if (common_dflash_rx_diag_enabled()) { + LOG_INF("DFLASH_RX draft: seq=%d precheck committed_len=%d ring_filled=%d ring_write_pos=%d cpu_ring_valid=%d gpu_ring=%d n_draft=%d\n", + seq_id, committed_len, ring_filled, ring_write_pos, + cpu_ring_valid ? 1 : 0, gpu_ring_handle ? 1 : 0, n_draft); + } + if (committed_len == 0) { + if (common_dflash_rx_diag_enabled()) { + LOG_INF("DFLASH_RX draft: seq=%d SKIP reason=committed_len_zero ring_filled=%d\n", + seq_id, ring_filled); + } continue; } if (invalid_reduced_logits_fail_closed) { + if (common_dflash_rx_diag_enabled()) { + LOG_INF("DFLASH_RX draft: seq=%d SKIP reason=invalid_reduced_logits_fail_closed committed_len=%d ring_filled=%d\n", + seq_id, committed_len, ring_filled); + } continue; } @@ -2996,7 +3066,16 @@ struct common_speculative_impl_dflash : public common_speculative_impl { common_dflash_align_drafter_seq_or_clear(ctx_dft, seq_id, committed_len, "flat draft"); int cross_len = build_cross_data(ctx_dft); + if (common_dflash_rx_diag_enabled()) { + LOG_INF("DFLASH_RX draft: seq=%d after_build_cross_data cross_len=%d committed_len=%d ring_filled=%d cpu_ring_valid=%d gpu_ring=%d\n", + seq_id, cross_len, committed_len, ring_filled, + cpu_ring_valid ? 1 : 0, gpu_ring_handle ? 1 : 0); + } if (cross_len <= 0) { + if (common_dflash_rx_diag_enabled()) { + LOG_INF("DFLASH_RX draft: seq=%d SKIP reason=cross_len_nonpositive cross_len=%d committed_len=%d ring_filled=%d\n", + seq_id, cross_len, committed_len, ring_filled); + } continue; } @@ -3005,13 +3084,18 @@ struct common_speculative_impl_dflash : public common_speculative_impl { // build drafter batch: [id_last, mask, mask, ..., mask] // positions stay on the target model's absolute timeline so RoPE // tracks the accepted suffix instead of restarting at the window. - // batch size adapts to n_draft+1 (saves compute when n_max < block_size-1) - const int batch_len = n_draft + 1; + // IMPORTANT: DFlash attention is non-causal over the query block. + // The drafter was trained/inferred against a fixed full block, so + // truncating the batch to n_draft+1 changes the attention graph and + // can change even the first draft token. Keep the full block here + // and only consume the first n_draft predictions below. + const int batch_len = block_size; + const int output_len = n_draft + 1; // root + consumed draft positions const int draft_pos_base = committed_len; common_batch_clear(batch_dft); common_batch_add(batch_dft, id_last, draft_pos_base, { seq_id }, true); for (int i = 1; i < batch_len; ++i) { - common_batch_add(batch_dft, mask_token_id, draft_pos_base + i, { seq_id }, true); + common_batch_add(batch_dft, mask_token_id, draft_pos_base + i, { seq_id }, i < output_len); } const int64_t t2 = ggml_time_us(); @@ -3033,7 +3117,7 @@ struct common_speculative_impl_dflash : public common_speculative_impl { const int argmax_rows = llama_get_logits_argmax_n(ctx_dft); if (argmax) { const int n_vocab = llama_vocab_n_tokens(llama_model_get_vocab(model_dft)); - if (!common_dflash_argmax_shape_valid(__func__, argmax_rows, batch_len, K_flat)) { + if (!common_dflash_argmax_shape_valid(__func__, argmax_rows, output_len, K_flat)) { if (dp.draft_log_probs) { dp.draft_log_probs->clear(); } @@ -3041,8 +3125,30 @@ struct common_speculative_impl_dflash : public common_speculative_impl { continue; } + if (common_dflash_topk_trace_enabled() && output_len > 1) { + const int row = 1; + std::string ids = "["; + std::string probs = "["; + for (int k = 0; k < K_flat; ++k) { + if (k) { + ids += ","; + probs += ","; + } + ids += std::to_string(argmax[row * K_flat + k]); + if (argmax_probs) { + probs += std::to_string(argmax_probs[row * K_flat + k]); + } else { + probs += "nan"; + } + } + ids += "]"; + probs += "]"; + LOG_INF("DFLASH_TOPK_TRACE flat seq=%d committed=%d id_last=%d row=%d K=%d ids=%s logp=%s\n", + (int) seq_id, committed_len, (int) id_last, row, K_flat, ids.c_str(), probs.c_str()); + } + // GPU argmax path - only top-k ids/probs are transferred. - for (int i = 1; i < batch_len && (int) result.size() < n_draft; ++i) { + for (int i = 1; i < output_len && (int) result.size() < n_draft; ++i) { const auto params = dp; if (argmax_probs && p_min > 0.0f && (int) result.size() >= params.n_min) { float log_prob = argmax_probs[i * K_flat]; @@ -3073,7 +3179,7 @@ struct common_speculative_impl_dflash : public common_speculative_impl { } else { // fallback: CPU argmax over full vocab const int n_vocab_dft = llama_vocab_n_tokens(llama_model_get_vocab(model_dft)); - for (int i = 1; i < batch_len && (int) result.size() < n_draft; ++i) { + for (int i = 1; i < output_len && (int) result.size() < n_draft; ++i) { float * logits = llama_get_logits_ith(ctx_dft, i); if (!logits) { break; @@ -3616,6 +3722,10 @@ struct common_speculative_impl_dflash : public common_speculative_impl { // called after initial prefill — grab all hidden states void capture_target_hiddens() { if (!common_dflash_target_capture_ready_or_skip(ctx_tgt)) { + if (common_dflash_rx_diag_enabled()) { + LOG_INF("DFLASH_RX capture_target_hiddens: seq=%d EARLY_EXIT reason=capture_not_ready committed_len=%d ring_filled=%d\n", + seq_id, committed_len, ring_filled); + } ring_write_pos = 0; ring_filled = 0; committed_len = 0; @@ -3626,10 +3736,22 @@ struct common_speculative_impl_dflash : public common_speculative_impl { llama_dflash_set_active_slot(ctx_tgt, seq_id); int32_t n_slots = llama_get_n_layer_hiddens(ctx_tgt); - if (n_slots == 0) return; + if (n_slots == 0) { + if (common_dflash_rx_diag_enabled()) { + LOG_INF("DFLASH_RX capture_target_hiddens: seq=%d EARLY_EXIT reason=n_slots_zero committed_len=%d ring_filled=%d\n", + seq_id, committed_len, ring_filled); + } + return; + } int64_t n_tokens = llama_get_layer_hidden_n_tokens(ctx_tgt, 0); - if (n_tokens <= 0) return; + if (n_tokens <= 0) { + if (common_dflash_rx_diag_enabled()) { + LOG_INF("DFLASH_RX capture_target_hiddens: seq=%d EARLY_EXIT reason=n_tokens_nonpositive n_tokens=%lld committed_len=%d ring_filled=%d\n", + seq_id, (long long) n_tokens, committed_len, ring_filled); + } + return; + } LOG_DBG("DFLASH_DBG capture_target_hiddens: n_tokens=%lld ring_write_pos=%d ring_filled=%d committed_len=%d\n", (long long) n_tokens, ring_write_pos, ring_filled, committed_len); @@ -3643,10 +3765,19 @@ struct common_speculative_impl_dflash : public common_speculative_impl { reset_invalid_reduced_logits_state(); common_dflash_reset_drafter_seq_and_kv_cache(ctx_dft, seq_id, "capture target hiddens"); const int actual_written = ring_write(to_store, start_offset, true); + if (common_dflash_rx_diag_enabled()) { + LOG_INF("DFLASH_RX capture_target_hiddens: seq=%d n_tokens=%lld start_offset=%d to_store=%d actual_written=%d ring_filled=%d ring_write_pos=%d discarded=%d\n", + seq_id, (long long) n_tokens, start_offset, to_store, actual_written, + ring_filled, ring_write_pos, ring_write_discarded ? 1 : 0); + } if (ring_write_discarded) { return; } committed_len = start_offset + actual_written; + if (common_dflash_rx_diag_enabled()) { + LOG_INF("DFLASH_RX capture_target_hiddens: seq=%d committed_after=%d ring_filled=%d\n", + seq_id, committed_len, ring_filled); + } update_drafter_kv_cache(actual_written); } @@ -3659,7 +3790,18 @@ struct common_speculative_impl_dflash : public common_speculative_impl { llama_dflash_set_active_slot(ctx_tgt, seq_id); int32_t n_slots = llama_get_n_layer_hiddens(ctx_tgt); - if (n_slots == 0 || n_accepted <= 0) { + if (n_slots == 0) { + if (common_dflash_rx_diag_enabled()) { + LOG_INF("DFLASH_RX append: seq=%d EARLY_EXIT reason=n_slots_zero requested=%d committed_len=%d ring_filled=%d\n", + seq_id, n_accepted, committed_len, ring_filled); + } + return; + } + if (n_accepted <= 0) { + if (common_dflash_rx_diag_enabled()) { + LOG_INF("DFLASH_RX append: seq=%d EARLY_EXIT reason=request_zero requested=%d committed_len=%d ring_filled=%d\n", + seq_id, n_accepted, committed_len, ring_filled); + } return; } @@ -3678,7 +3820,19 @@ struct common_speculative_impl_dflash : public common_speculative_impl { last5[0], last5[1], last5[2], last5[3], last5[4]); } + const int ring_filled_before = ring_filled; + const int committed_before = committed_len; + const int write_pos_before = ring_write_pos; + const int actual_written = ring_write(n_accepted); + if (common_dflash_rx_diag_enabled()) { + LOG_INF("DFLASH_RX append: seq=%d requested=%d actual_written=%d ring_filled_before=%d ring_filled_after=%d committed_before=%d write_pos_before=%d write_pos_after=%d discarded=%d defer_kv=%d\n", + seq_id, n_accepted, actual_written, + ring_filled_before, ring_filled, + committed_before, write_pos_before, ring_write_pos, + ring_write_discarded ? 1 : 0, + defer_drafter_kv_cache ? 1 : 0); + } if (ring_write_discarded) { return; } @@ -3689,6 +3843,10 @@ struct common_speculative_impl_dflash : public common_speculative_impl { return; } committed_len += actual_written; + if (common_dflash_rx_diag_enabled()) { + LOG_INF("DFLASH_RX append: seq=%d committed_after=%d ring_filled=%d deferred_kv_tokens=%d\n", + seq_id, committed_len, ring_filled, deferred_drafter_kv_tokens); + } if (defer_drafter_kv_cache) { defer_drafter_kv_cache_update(actual_written); } else { @@ -4561,7 +4719,11 @@ void common_speculative_draft_batch( const llama_model * model_dft = llama_get_model(ctx_dft); const int block_size = llama_model_dflash_block_size(model_dft); const int n_draft = std::min(block_size - 1, params.n_max); - const int batch_len = n_draft + 1; + // Keep the full query block for flat/batched DFlash too. Tree drafting + // already does this. Non-causal attention over query tokens means shorter + // batches are not semantically equivalent to the trained full block. + const int batch_len = block_size; + const int output_len = n_draft + 1; // root + consumed draft positions const llama_token mask_tok = (llama_token) llama_model_dflash_mask_token_id(model_dft); const int64_t t0 = ggml_time_us(); @@ -4641,7 +4803,7 @@ void common_speculative_draft_batch( for (const auto & rs : ready) { common_batch_add(batch, id_last_per_spec[rs.spec_idx], rs.draft_pos_base, { rs.seq_id }, true); for (int i = 1; i < batch_len; i++) { - common_batch_add(batch, mask_tok, rs.draft_pos_base + i, { rs.seq_id }, true); + common_batch_add(batch, mask_tok, rs.draft_pos_base + i, { rs.seq_id }, i < output_len); } } @@ -4664,11 +4826,11 @@ void common_speculative_draft_batch( auto & rs = ready[r]; auto & result = result_per_spec[rs.spec_idx]; std::vector * log_probs = log_probs_per_spec ? &(*log_probs_per_spec)[rs.spec_idx] : nullptr; - const int offset = r * batch_len; + const int offset = r * output_len; if (argmax) { const int n_vocab = llama_vocab_n_tokens(llama_model_get_vocab(model_dft)); - if (!common_dflash_argmax_shape_valid(__func__, argmax_rows, n_ready * batch_len, K_flat)) { + if (!common_dflash_argmax_shape_valid(__func__, argmax_rows, n_ready * output_len, K_flat)) { if (log_probs) { log_probs->clear(); } @@ -4676,7 +4838,7 @@ void common_speculative_draft_batch( return; } - for (int i = 1; i < batch_len && (int) result.size() < n_draft; i++) { + for (int i = 1; i < output_len && (int) result.size() < n_draft; i++) { if (argmax_probs && params.p_min > 0.0f && (int) result.size() >= params.n_min) { float log_prob = argmax_probs[(offset + i) * K_flat]; if (log_prob < logf(params.p_min)) { @@ -4701,14 +4863,12 @@ void common_speculative_draft_batch( log_probs->push_back(argmax_probs[(offset + i) * K_flat]); } } - } else { - const int n_vocab = llama_vocab_n_tokens(llama_model_get_vocab(model_dft)); - for (int i = 1; i < batch_len && (int) result.size() < n_draft; i++) { + } else { const int n_vocab = llama_vocab_n_tokens(llama_model_get_vocab(model_dft)); + for (int i = 1; i < output_len && (int) result.size() < n_draft; i++) { float * logits = llama_get_logits_ith(ctx_dft, offset + i); if (!logits) { break; - } - llama_token best = (llama_token)(std::max_element(logits, logits + n_vocab) - logits); + } llama_token best = (llama_token)(std::max_element(logits, logits + n_vocab) - logits); result.push_back(best); if (log_probs) { log_probs->push_back(0.0f); diff --git a/docs/enable-and-fix-dflash-on-vulkan-various-models.md b/docs/enable-and-fix-dflash-on-vulkan-various-models.md new file mode 100644 index 000000000000..462499341de2 --- /dev/null +++ b/docs/enable-and-fix-dflash-on-vulkan-various-models.md @@ -0,0 +1,717 @@ +# Enable and Fix DFlash on Vulkan — Various Models + +This document is the canonical rationale for **enabling and fixing DFlash on the +Vulkan backend across the model set** tested on AMD Strix Halo (Radeon 8060S / +RADV GFX1151, Mesa 25.0.7, Vulkan 1.4.309, 96 GB UMA). It consolidates the +original Qwen3-Coder-Next enablement work with the later multi-model fixes +(Qwen3.5-122B-A10B GPU cross-ring, Qwen3.6-35B-A3B stochastic-sampling +corruption) and the per-model status/performance for Qwen3.6-27B, Gemma 4, and +the 256K stress test results. + +It replaces the scattered investigation notes that were created while debugging +the original 0% draft-token acceptance problem and the later qwen35moe +recurrent-state regressions. + +**Current PR head**: `d6531e021` (`enable-dflash-qwen3-coder-next-on-vulkan`) +**PR**: [Anbeeld/beellama.cpp#79](https://github.com/Anbeeld/beellama.cpp/pull/79) + +## Summary + +DFlash on Vulkan was enabled and stabilized across a range of target/drafter +pairs. Each model exposed a distinct issue; the fixes are layered and do not +conflict. + +| Target model | Arch | Drafter | Original Vulkan status | Fix | +|---|---|---|---|---| +| Qwen3-Coder-Next | `qwen3next` (MoE) | qwen3-coder-next-dflash | 0% acceptance | Drafter graph target-KV branch gate + Qwen3Next recurrent-state tape capture + full-block flat drafting | +| Qwen3.6-27B | `qwen35` (dense) | Qwen3.6-27B-DFlash | worked | (no fix needed; performance characterized) | +| Qwen3.5-122B-A10B | `qwen35moe` | qwen35-122b-a10b-dflash | GPU cross-ring recurrent-state desync under `n_rs_seq=0` | Keep `n_rs_seq` for QWEN35MOE (RS partial rollback) instead of clamping to 0 | +| Qwen3.6-35B-A3B | `qwen35moe` | qwen36-35b-a3b-dflash | gibberish under stochastic sampling (0% acceptance) | Rotate DeltaNet RS snapshots for `n_seq_tokens < K` | +| Gemma 4 31B | `gemma4` | gemma4-31b-it-dflash | works (FA off); FA on → "own" loop (Vulkan FA kernel ISWA bug) | `--flash-attn off` workaround; stochastic less sensitive | + +The two decisive runtime fixes layered on top of the coder-next enablement are: + +1. **Keep `n_rs_seq` for QWEN35MOE** (122B): under the GPU cross-ring, + clamping `n_rs_seq=0` forced checkpoint-restore for draft rejection and the + target's DeltaNet recurrent-state cells desynchronized from committed token + positions (`non-consecutive token position` warnings) → wrong drafts accepted + → repetition. Keeping `n_rs_seq = draft.n_max` (RS partial rollback) makes both + ring=0 and ring=1 coherent. + +2. **Rotate DeltaNet RS snapshots for `n_seq_tokens < K`** (35B-A3B): the + `n_rs_seq>0` write-back loop in `build_recurrent_attn` wrote all `K` + `gdn_out` snapshot slots into the recurrent-state tensor every step, but for a + decode / short batch the gated_delta_net kernel only populates the last + `min(n_seq_tokens, K)` slots; the leading slots are uninitialised. Under + stochastic sampling the adaptive draft-max controller drives spec depth to 0 + once acceptance drops, so every step becomes a 1-token decode → snapshot slots + `1..K-1` stay garbage → repeated draft-rejection rollbacks read stale + recurrent state → output degrades to multilingual gibberish. The fix writes + only the populated slots and rotates the older snapshots down. + +## Test setup + +- Hardware: AMD Ryzen AI MAX+ 395 / Radeon 8060S, 96 GB UMA, Vulkan 1.4.309 (RADV). +- Build: Vulkan Release, `GGML_VULKAN=ON`. Vulkan has no DFlash GPU cross-ring + and no TurboQuant/TCQ cache types, so tests use `--cache-type-k q4_0 + --cache-type-v q4_1` and `GGML_DFLASH_GPU_RING` to select ring mode. +- Server: `llama-server`, single-sequence DFlash (`-np 1`), `-ngl all`, + `--flash-attn on`, `--no-cache-prompt`, `--device Vulkan0`. +- DFlash args: `--spec-type dflash --spec-draft-n-max 4 --spec-dflash-cross-ctx + 1024 --spec-branch-budget 0 --spec-draft-ngl all --spec-draft-device Vulkan0`. +- `GGML_DFLASH_GPU_RING=0` → CPU hidden capture (ring=0); `=1` → GPU cross-ring + (ring=1), which on main falls back to CPU because main has no GPU ring code. + +## Post-fix performance (greedy, 256 tokens, PR fixed, all rep=0) + +| Target model | baseline tok/s | DFlash ring=0 (acc) | DFlash ring=1 (acc) | +|---|---:|---:|---:| +| Qwen3.6-35B-A3B | 52.04 | 75.60 (0.827) | 76.88 (0.828) | +| Qwen3.6-27B (dense) | 12.47 | 29.80 (0.835) | 29.94 (0.816) | +| Qwen3-Coder-Next (MoE) | 53.26 | 34.07 (0.754) | 50.09 (0.621) | +| Qwen3.5-122B-A10B (MoE) | 13.92 | 8.20 (0.189) | 14.92 (0.256) | +| Gemma 4 31B | 11.34 | 17.57 (0.231) | 16.93 (0.185) | + +Interpretation: + +- **Qwen3.6-35B-A3B**: after the snapshot-rotation fix, DFlash is both coherent + **and 1.45–1.48× faster** than baseline (52 → 76 tok/s) with 0.83 acceptance. + ring=0 ≈ ring=1 on Vulkan for this model. +- **Qwen3.6-27B (dense)**: strongest DFlash win — **2.4×** (12.5 → 30 tok/s), + acceptance 0.82. ring=0 ≈ ring=1. +- **MoE targets (Coder-Next, 122B)**: DFlash is net-slower to ~neutral, matching + the known Vulkan/MoE pattern — cheap active decode means drafter overhead + dominates. ring=1 is at or slightly above baseline; ring=0 is slower (CPU + capture overhead). These were already coherent; the fixes preserve that and + make their stochastic-sampling paths coherent too. +- **Gemma 4 31B** (`gemma4`, no DeltaNet recurrent state): DFlash is coherent + at ~1.5× speedup (11.34 → 17.57 tok/s) with `--flash-attn off`. ring=0 (CPU + capture) beats ring=1 for Gemma 4's large `n_embd` (5376). See the Gemma 4 + section for the Vulkan DFlash + FA corruption issue. +- **Gemma 4 31B** (`gemma4`, no DeltaNet recurrent state) is unaffected by the + recurrent-state fixes. On Vulkan, ring=0 (CPU capture) beats ring=1 for Gemma 4's + large `n_embd` (5376). See the Gemma 4 section below for the Vulkan DFlash + FA + corruption issue and its workaround (`--flash-attn off`). + +## Qwen3-Coder-Next enablement (`qwen3next` MoE) + +Qwen3-Coder-Next DFlash originally produced near-0% accepted draft tokens on +Vulkan. The failure was not caused by a generic DFlash verifier bug: Qwen3.6 +DFlash worked, and Qwen3-Coder-Next could work under other implementations. The +issue was a chain of Qwen3Next-specific runtime and drafter-input mismatches. + +The decisive Vulkan runtime fix is: + +- Vulkan does not have the DFlash GPU cross ring. +- Without the GPU cross ring, the drafter's normal full-attention KV cache is not + populated with target-context K/V. +- The Qwen3-Coder-Next drafter was still taking the full-attention KV-cache + branch, reading empty/stale context. +- The DFlash drafter graph now only builds that branch when target KV is actually + available. On Vulkan it falls back to projecting K/V freshly from CPU-captured + `target_hidden`. + +With the runtime fixes and a drafter trained for the live C++ position semantics, +in-distribution acceptance was observed at roughly 68-72% instead of 0%. + +### User-facing command shape + +The working Vulkan launch shape is: + +```bash +build-vulkan/bin/llama-server \ + -m /crypt/models/Qwen3-Coder-Next-Q8_0-00001-of-00003.gguf \ + --spec-draft-model /crypt/tmp/zlab_finetuned_pos1_f16.gguf \ + --spec-type dflash \ + --spec-dflash-cross-ctx 512 \ + --spec-draft-n-max 4 \ + --spec-branch-budget 0 \ + -np 1 \ + --kv-unified \ + -ngl all \ + --spec-draft-ngl all \ + -b 2048 \ + --ctx-size 8192 \ + --flash-attn on \ + --jinja \ + --temp 0.0 \ + --top-k 1 +``` + +The exact drafter file matters. The runtime fix enables the correct Vulkan path, +but a drafter trained against the wrong input/position semantics can still show +poor acceptance. + +### Root-cause chain + +#### 1. Qwen3Next conv state was not advanced during DFlash rollback + +Qwen3-Coder-Next has recurrent / convolutional state. During DFlash verification +and rollback, accepted tokens must advance the target recurrent state. The tape +replay path had to reconstruct both DeltaNet state and convolution state. + +The Qwen3Next graph emits the pre-conv QKV tensor as: + +```text +linear_attn_qkv_mixed- +``` + +with a legacy path named: + +```text +qkv_mixed- +``` + +The DFlash tape map did not capture these names, so `tape_replay_conv()` skipped +every Qwen3Next recurrent layer. The conv state `r_l` remained frozen at the +pre-draft backup state. That could produce garbled target output even when no +draft tokens were accepted. + +Fix: + +```cpp +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}; +``` + +This made the conv replay path see the Qwen3Next pre-conv data and advance `r_l` +correctly. + +#### 2. Vulkan lacked a populated drafter full-attention KV cache + +DFlash keeps recent target hidden states in a cross-attention ring. On CUDA, the +GPU cross ring can also populate the drafter-side full-attention KV cache with +target-context K/V. + +On Vulkan, the GPU cross ring is unavailable, so DFlash uses CPU hidden capture: + +```text +gpu_ring_handle == nullptr +``` + +In that case, `update_drafter_kv_cache()` returns early and the drafter +full-attention KV cache is not populated with target context. + +However, the DFlash drafter graph still built `inp_attn_kv_full` whenever the +drafter had a memory context. For all-full-attention drafters, this selected the +branch that reads target context from the normal KV cache: + +```cpp +if (inp_attn_kv_full && !hparams.is_swa(il)) { + // read target context from drafter KV cache +} +``` + +On Vulkan that cache was empty or stale. The drafter therefore generated from +little or no real target context, producing repeated bad predictions and 0% +acceptance. + +Fix: add an explicit context parameter: + +```cpp +bool dflash_target_kv_available = false; +``` + +and a public setter: + +```cpp +llama_set_dflash_target_kv_available(ctx_dft, gpu_ring_handle != nullptr); +``` + +Then build the full-attention KV input only when the target KV cache is actually +populated: + +```cpp +llm_graph_input_attn_kv * inp_attn_kv_full = nullptr; +if (mctx && cparams.dflash_target_kv_available) { + inp_attn_kv_full = dflash_build_base_attn_input(...); +} +``` + +When the flag is false, the graph uses the existing fresh-projection path: + +```cpp +Kcur_ctx = wk(fused_target) +Vcur_ctx = wv(fused_target) +``` + +That path works with CPU-captured `target_hidden`, which is the correct Vulkan +fallback. + +#### 3. Flat DFlash must keep the full query block + +DFlash query-token attention is non-causal over the draft block. Shortening the +query block can change logits, including early positions. + +The flat path previously used: + +```cpp +batch_len = n_draft + 1; +``` + +Tree DFlash already used the full block. The flat and batched flat paths now +keep the trained/inferred block shape: + +```cpp +batch_len = block_size; +output_len = n_draft + 1; +``` + +Only `output_len` rows are marked for output and consumed by the sampler. This +preserves full-block drafter semantics while still limiting active draft horizon. + +#### 4. Runtime parity required comparing exact live C++ inputs + +The debugging path added several diagnostics that made it possible to prove where +the mismatch lived (see the diagnostics table below). The cross dump showed that +after the Vulkan KV branch fix, C++ and PyTorch forward passes matched for the +same live input. That narrowed the remaining 0% behavior to drafter +training/input semantics rather than a C++ graph mismatch. + +#### 5. Drafter position semantics matter + +C++ consumes DFlash predictions starting at row / position 1. Position 0 is the +known root / bonus token. A drafter trained to predict from position 0 can +appear valid in a standalone training loop while still producing poor live C++ +acceptance, because C++ reads position 1. Training the live-consumed position +fixed the in-distribution acceptance collapse once the runtime graph was correct. + +This document does not cover broader drafter generalization work; that is a +separate training/data problem. + +## Qwen3.5-122B-A10B GPU cross-ring fix (`qwen35moe`) + +### Symptom + +122B DFlash under the GPU cross-ring (`GGML_DFLASH_GPU_RING=1`) produced +repetition (`rep=12`) with dozens of `non-consecutive token position N after N` +warnings in `llama-memory-recurrent.cpp`. ring=0 (CPU capture) was coherent. + +### Root cause + +`need_n_rs_seq() == draft.n_max` (e.g. 4) requests per-token recurrent-state +snapshots so the target can roll back after a partial draft rejection. An earlier +guard clamped `n_rs_seq=0` for QWEN35MOE ("DeltaNet RS-rollback snapshot bug +under split_equal"). With `n_rs_seq=0`, draft rejection falls back to +checkpoint-restore, but under the GPU cross-ring the target's DeltaNet +recurrent-state cells desynchronize from committed token positions → wrong drafts +accepted → repetition. + +### Fix + +In `src/llama-context.cpp`, do not clamp `n_rs_seq` to 0 for QWEN35MOE — keep it +as requested by DFlash (`need_n_rs_seq() == draft.n_max`). The QWEN35 dense path +(Qwen3.6-27B, `n_rs_seq=8`) was never clamped and already works. + +### Caveat resolved by the 35B-A3B fix + +The 122B fix's comment originally said the `n_rs_seq>0 → garbage` note was "not +reproduced in the single-sequence DFlash path". That was true for 122B under +greedy sampling but **false for Qwen3.6-35B-A3B under stochastic sampling**, +which exposed a separate DeltaNet snapshot write-back bug (see next section). +That bug is now fixed, so `n_rs_seq>0` is correct for QWEN35MOE under both greedy +and stochastic sampling. + +## Qwen3.6-35B-A3B stochastic-sampling fix (`qwen35moe`) + +### Symptom + +35B-A3B DFlash produced multilingual gibberish (`"Here's a thinking +fertig是否经济的 whether puppermodele..."`) with **0.000 acceptance** and an HTTP +500 parse error under stochastic sampling (`--temp 1.0 --top-k 64 --top-p 0.95`, +the quickstart setting), for both `--reasoning on` and `--reasoning off`, and +for both ring=0 and ring=1. Greedy (`--temp 0`) was coherent. + +### Root-cause investigation (evidence-based, gdb) + +The 122B fix removed the QWEN35MOE `n_rs_seq` clamp, so 35B-A3B ran with +`n_rs_seq=4`. gdb on `llama_memory_recurrent::seq_rm` / `set_rs_idx` showed: + +- Greedy (coherent): 10 rollbacks over 24 tokens, snapshot slot indices + scattered (1, 3, 4, 2). +- Stochastic (garbage): 17 rollbacks over 32 tokens, **`set_rs_idx idx=2` on + nearly every step** after the first rejection. + +The differentiator is **stochastic sampling**, not reasoning mode, length, or +ring wrap (confirmed by isolation tests): greedy → coherent at 24/1024/2048 +tokens; stochastic → garbage at 128 tokens. + +### The bug + +In `build_recurrent_attn` (`src/models/delta-net-base.cpp`), the `keep` +(`n_rs_seq>0`) write-back loop copies all `K` `gdn_out` snapshot slots into +`ssm_states_all` every step: + +```cpp +for (int64_t k_i = 0; k_i < K; ++k_i) { + const uint32_t cache_slot = (uint32_t) (K - 1 - k_i); + ggml_tensor * src = ggml_view_4d(ctx0, gdn_out, ..., attn_score_elems + k_i * state_size_per_snap); + ggml_tensor * dst = ggml_view_2d(ctx0, ssm_states_all, ..., cache_slot * ...); + ggml_build_forward_expand(gf, ggml_cpy(ctx0, src, dst)); +} +``` + +The gated_delta_net kernel only writes the last `min(n_seq_tokens, K)` snapshot +slots of `gdn_out` (slots `(K-n_seq_tokens) .. (K-1)`); the leading slots are +**uninitialised** (caller-owned). For a **decode / short batch** +(`n_seq_tokens < K`) the loop writes that uninitialised data into state slots +`1..K-1` on every step. + +Under DFlash with stochastic sampling, acceptance drops → the adaptive draft-max +controller drives spec depth to 0 → **every step becomes a 1-token decode** +(`n_seq_tokens=1`) → snapshot slots `1..K-1` stay garbage → repeated +draft-rejection rollbacks (`idx=2`) read stale recurrent state → output degrades +to gibberish. Greedy / high-acceptance runs were unaffected because verify +batches (`n_seq_tokens = K`) overwrite the corruption with valid data. + +### The fix + +Write only the kernel-populated `gdn_out` slots (newest `n_pop` states → state +slots `0..n_pop-1`) and **rotate** the older snapshots down by `n_pop` slots +(old slot `s` → slot `s + n_pop`) so the recurrent-state history stays correct +for rollback: + +```cpp +const int64_t n_pop = (n_seq_tokens < K) ? n_seq_tokens : K; +for (int64_t k_i = 0; k_i < K; ++k_i) { + const uint32_t cache_slot = (uint32_t) (K - 1 - k_i); + ggml_tensor * dst = ggml_view_2d(ctx0, ssm_states_all, ..., cache_slot * ...); + if (k_i >= K - n_pop) { + // gdn_out slot k_i holds the state after one of this batch's tokens. + ggml_tensor * src = ggml_view_4d(ctx0, gdn_out, ..., attn_score_elems + k_i * state_size_per_snap); + ggml_build_forward_expand(gf, ggml_cpy(ctx0, src, dst)); + } else { + // gdn_out slot k_i is uninitialised (n_seq_tokens < K): rotate the older + // snapshot down so slot history is preserved for rollback. + const uint32_t old_slot = cache_slot - (uint32_t) n_pop; + ggml_tensor * src = ggml_view_2d(ctx0, ssm_states_all, ..., old_slot * ...); + ggml_build_forward_expand(gf, ggml_cpy(ctx0, src, dst)); + } +} +``` + +Iterating `k_i` ascending reads each old slot before it is overwritten, so the +in-place shift is safe (copies to the same persistent buffer are serialised by +the scheduler). For `n_seq_tokens >= K` the behaviour is unchanged (full verify +batch). The `!keep` path (`n_rs_seq=0`, e.g. Qwen3-Coder-Next) is untouched. + +### Verification + +Post-fix (PR `efb07137c`), all coherent (`rep=0`): + +| model | arch | n_rs_seq | greedy r1 | stochastic r1 | +|---|---|---|---|---| +| 35B-A3B | qwen35moe | 4 | 0.788 acc ✓ | **0.000→0.711 acc FIXED** ✓ | +| 122B | qwen35moe | 4 | 0.229 acc ✓ (no reg) | **0.209 acc (bonus)** ✓ | +| qwen3.6-27B | qwen35 dense | 8 | 0.771 acc ✓ | 0.519 acc ✓ | +| Coder-Next | qwen3next | 0 | 0.621 acc ✓ | 0.579 acc ✓ (unaffected) | + +The fix turns 35B-A3B from broken (gibberish under stochastic) into a 1.48× +speedup, and makes 122B's stochastic path coherent as a bonus, with **zero +regressions** across all models. + +## Gemma 4 31B (`gemma4`) + +Gemma 4 31B DFlash works on the PR and crashes on main (DFlash +`std::out_of_range` vocab index). It has no DeltaNet recurrent state, so it is +unaffected by the recurrent-state fixes. On Vulkan, ring=0 (CPU capture) beats +ring=1 (GPU ring) for Gemma 4's large `n_embd` (5376), consistent with the +quickstart's "Vulkan: not recommended for DFlash, falls back to CPU ring" +guidance. + +### Vulkan DFlash + FA corruption ("own own own" loop) + +Under greedy sampling (`--temp 0`) with `--flash-attn on`, Gemma 4 DFlash on +Vulkan produces "own own own" repetition with extremely low vocabulary diversity +(3% unique words over 512 tokens). The same prompt with `--flash-attn off` is +fully coherent (64% unique words, correct C++ code). + +**Evidence table** (hard prompt, ~665 tokens, 512 generated, greedy `--temp 0`): + +| Configuration | Result | Vocabulary diversity | +|---|---|---| +| G4 baseline + FA on | ✅ Coherent C++ code | 66% unique words | +| G4 + DFlash + FA **off** | ✅ Coherent C++ code | 64% unique words | +| G4 + DFlash + FA **on** | ❌ "own" repetition | 3% unique words | + +**Root cause**: The Vulkan flash attention kernel produces incorrect attention +output during DFlash draft-verification for Gemma 4's ISWA (Interleaved Sliding +Window Attention) layers. The precise mechanism is likely a combination of: +(a) FA mask stride/tiling issues on ISWA's mixed SWA/FULL attention pattern, +(b) fp16 precision accumulation during the verification batch, and +(c) the unique batch shapes of speculative verification. + +Gemma 4 has a mix of SWA and FULL attention layers (`swa=4 full=2`). The bug +is consistent with known upstream issues: +- Vulkan FA mask stride fixes (#12720, #12853) addressed mask loads but the + underlying shader may still miscompute for ISWA + draft-verification pattern +- Vulkan + Gemma 4 fp16 MMQ precision bug (ollama#15261, fixed in ollama + v0.30.0-rc31) — Vulkan explicitly produces wrong answers for Gemma 4 under + certain conditions +- The bug is specific to the FA kernel path during DFlash verification + (normal decoding with FA is correct) + +This is a **Vulkan backend limitation**, not a DFlash logic bug. On CUDA the +issue does not reproduce. The bug may also affect other ISWA models (Gemma 3, +Ministral 3) under similar conditions. + +**Workaround**: Use `--flash-attn off` for Gemma 4 DFlash on Vulkan. With FA +off, DFlash produces coherent output at ~17-19 tok/s (ring=0) — still ~50% +faster than baseline. Stochastic sampling (`--temp 1.0 --top-k 64 --top-p +0.95`) is also less sensitive to the precision issue. + +### GLM second-opinion review + +An independent review (glm-5.2:cloud via Ollama) confirmed: +- The A/B evidence **strongly supports** the interaction between DFlash and FA + as the trigger +- The workaround is **highly sound** (FA is a perf optimization, not required + for correctness) +- The root cause statement was broadened from "ISWA mask corruption" to + "Vulkan FA precision/mask handling during DFlash verification batches" to + account for fp16 accumulation and batch-shape irregularities +- The bug may affect other ISWA models (Gemma 3, Ministral 3) and should be + tested if those models are used with Vulkan DFlash + +### 256K context stress test + +At `-c 262144` (256K context) with the hard prompt (~665 tokens, 512 generated): +- All 5 models fit at 256K, zero crashes, zero OOM +- Gemma 4 ring=0 (FA off): coherent, 10.49 tok/s, acc 0.215 +- Gemma 4 ring=1 stochastic: 12.37 tok/s, acc 0.564 (above baseline) +- Correctness (FA off): 131/200 unique words (66%), matching baseline quality + +| Model (256K) | Baseline gen | DFlash ring=0 | DFlash ring=1 | ring=1 stochastic | Result | +|---|---|---|---|---|---| +| Qwen3.6-27B | 12.36 tok/s | 21.72 (acc .746) | 21.42 (acc .743) | 14.55 (acc .434) | ✅ Best DFlash win (1.73×) | +| 35B-A3B | 42.80 tok/s | 54.39 (acc .705) | 54.97 (acc .710) | 38.52 (acc .444) | ✅ Strong DFlash win (1.28×) | +| Coder-Next | 53.22 tok/s | 30.25 (acc .328) | 46.68 (acc .478) | 45.29 (acc .284) | ✅ Stable, slower than baseline | +| Gemma 4 31B | 11.33 tok/s | 10.49 (acc .215) | 10.83 (acc .185) | 12.37 (acc .564) | ✅ Stable; stochastic above baseline | +| 122B-A10B | 19.04 tok/s | 9.67 (acc .144) | 13.22 (acc .146) | 16.23 (acc .164) | ✅ Stable, slower than baseline | + +## Ruled-out or deprioritized hypotheses + +### (coder-next) k_norm ordering + +Suspected mismatch around whether target K normalization happened before or +after concatenating context/noise K. Not the cause of 0% acceptance; the +attempted k_norm fix was reverted. + +### (coder-next) Basic cross-ring emptiness + +Diagnostics showed the cross window was not generally empty. The real problem +was full-attention layers reading target context from an unpopulated drafter KV +cache on Vulkan instead of using the fresh `target_hidden` projection path. + +### (coder-next) Generic DFlash verifier failure + +Qwen3.6 DFlash worked, and the server verification path could reject all drafts +while still producing coherent target output when recurrent state was handled +correctly. A global verifier failure was unlikely. + +### (coder-next) C++ vs PyTorch graph mismatch after KV-branch fix + +After dumping the exact live hidden-state window and replaying it in PyTorch, C++ +and PyTorch forward outputs matched. The C++ graph was not the remaining source +of the live acceptance issue after the Vulkan KV fix. + +### (35B-A3B) 122B-style RS desync + +The 122B smoking gun (`non-consecutive token position` warnings + `n_rs_seq=0`) +was **absent** in the broken 35B-A3B run (0 warnings, `n_rs_seq=4`) and **present** +in the working coder-next run (22 warnings, `n_rs_seq=0`). The RS rollback path +is clean in 35B-A3B; the corruption is in the snapshot write-back, not the +rollback/restore. Confirmed by gdb: no `RS-ROLLBACK-OVERFLOW`, no non-consec +warnings, but repeated `set_rs_idx idx=2` reading stale slots. + +### (35B-A3B) Reasoning mode / generation length / ring wrap as the trigger + +Isolation tests disproved all three: reasoning ON + greedy = coherent; +reasoning OFF + stochastic = garbage; 2048-token greedy (wraps past +`cross_ctx=1024`) = coherent. The sole trigger is **stochastic sampling + +`n_rs_seq>0`**, because it drives the adaptive controller to spec depth 0 +(all-decode steps). + +## Relevant code changes + +### Public API + +`include/llama.h`: + +```cpp +LLAMA_API void llama_set_dflash_target_kv_available(struct llama_context * ctx, bool avail); +``` + +Debug-only recurrent state dump API: + +```cpp +LLAMA_API void llama_dflash_dump_recurrent_state_dbg( + struct llama_context * ctx, + llama_seq_id seq_id, + const char * tag); +``` + +### Context parameter + +`src/llama-cparams.h`: + +```cpp +bool dflash_target_kv_available = false; +``` + +### Flag propagation + +`common/speculative.cpp`: + +```cpp +llama_set_dflash_target_kv_available(ctx_dft, gpu_ring_handle != nullptr); +``` + +### Drafter graph branch gate + +`src/models/dflash_draft.cpp`: + +```cpp +if (mctx && cparams.dflash_target_kv_available) { + inp_attn_kv_full = dflash_build_base_attn_input(...); +} +``` + +### Qwen3Next tape capture + +`src/llama-context.cpp` registers the Qwen3Next pre-conv QKV names: + +```cpp +linear_attn_qkv_mixed- +qkv_mixed- +``` + +### Keep n_rs_seq for QWEN35MOE (122B) + +`src/llama-context.cpp` no longer clamps `n_rs_seq` to 0 for QWEN35MOE; it keeps +the DFlash-requested value (`need_n_rs_seq() == draft.n_max`). + +### Rotate DeltaNet RS snapshots for n_seq_tokens < K (35B-A3B) + +`src/models/delta-net-base.cpp`, `build_recurrent_attn` keep-path write-back loop +(see the 35B-A3B fix section above for the full code). + +### Full-block flat drafting + +`common/speculative.cpp` uses: + +```cpp +const int batch_len = block_size; +const int output_len = n_draft + 1; +``` + +and consumes only rows `1..output_len-1`. + +### Drafter KV slide without GPU cross-ring (commit `5b0e7533c`) + +Drafter KV cache is now slid correctly even without the GPU cross-ring, fixing +the "drafter decode failed with 1" errors at 256K context. This is verified by +the 256K regression test: all 20 runs (5 models × 4 configs) pass with zero +crashes, zero OOM, and zero "failed to slide" warnings. + +### Regression unit tests (commit `da8806748`) + +Five new test functions in `tests/test-dflash-plumbing.cpp` guard all PR +behavior changes. Each test was confirmed to FAIL when its corresponding fix is +reverted. Production refactor: extracted `build_recurrent_attn`'s RS write-back +slot decision into `llama_dflash_rs_writeback_slot_for_test()` in +`llama-context.h` (minimal, behavior-identical). + +## Diagnostics retained + +The following environment variables are useful for future Vulkan / Qwen3Next / +qwen35moe DFlash debugging: + +| Variable | Purpose | +|---|---| +| `GGML_DFLASH_RX_DIAG=1` | Logs cross-window, ring, and hidden-row availability. | +| `GGML_DFLASH_TOKEN_TRACE=1` | Logs sampled/drafted/verified token IDs. | +| `GGML_DFLASH_TOPK_TRACE=1` | Logs drafter top-k rows for C++ vs PyTorch comparison. | +| `GGML_DFLASH_CROSS_DUMP=/path/file.txt` | Dumps a live cross-attention hidden window for replay. | +| `GGML_DFLASH_FORCE_REDECODE=1` | Forces re-decode instead of tape replay, useful to isolate recurrent-state replay issues. | +| `GGML_DFLASH_DISABLE_KV_CACHE=1` | Disables the DFlash projection cache for isolation. | +| `GGML_DFLASH_GPU_RING` | `0` = CPU hidden capture; `1` = GPU cross-ring (Vulkan falls back to CPU on main). | +| `GGML_DFLASH_GPU_RING_DEBUG=1` | Logs GPU cross-ring D2D allocation/transfer traces. | + +For recurrent-state rollback inspection, `gdb` on +`llama_memory_recurrent::seq_rm` and `llama_memory_recurrent::set_rs_idx` (both +in `src/llama-memory-recurrent.cpp`, symbols in `libllama.so`) shows the rollback +requests and the snapshot slot index selected. + +## Final status + +The Vulkan path enables DFlash across the model set by: + +- **Coder-Next**: making the drafter graph choose the correct target-context + source (GPU-ring → drafter KV cache; Vulkan/CPU-hidden → fresh + `target_hidden` projection), advancing Qwen3Next recurrent state during + rollback, and keeping the full query block. +- **122B**: keeping `n_rs_seq` for QWEN35MOE so the GPU cross-ring can roll + back DeltaNet recurrent state without desyncing. +- **35B-A3B**: rotating DeltaNet RS snapshots for `n_seq_tokens < K` so + stochastic sampling (which drives spec depth to 0 and turns every step into a + decode) does not clobber the snapshot history. +- **Gemma 4**: Vulkan FA kernel corruption during DFlash verification for ISWA + layers documented; `--flash-attn off` workaround produces coherent output at + ~50% speedup over baseline. + +### PR head history + +| Commit | Description | +|---|---| +| `b788b4af1` | Original PR: dflash enable Qwen3-Coder-Next on Vulkan | +| `648c0e9d6` | Original head: validate reduced-verify argmax on Vulkan | +| `bcd42973f` | Sync Vulkan DFlash code from `vulkan-dflash-cross-ring` | +| `ab2be17fd` | Keep `n_rs_seq` for QWEN35MOE (fix 122B GPU cross-ring corruption) | +| `efb07137c` | Rotate DeltaNet RS snapshots for `n_seq_tokens < K` (fix 35B stochastic) | +| `da8806748` | Regression unit tests + RS write-back slot helper extraction | +| `5b0e7533c` | Slide drafter KV without GPU cross-ring (fix 256K decode fails) | +| `d6531e021` | Document Gemma 4 Vulkan DFlash FA corruption root cause and workaround | + +### Verified across + +- **5 models** × **4 configs** (baseline, r0, r1, r1 stochastic) at both 4K and + 256K context: all coherent, zero crashes, zero OOM +- **Unit tests**: `test-dflash-plumbing.cpp` passes on PR head; each new test + confirmed to FAIL when its fix is reverted +- **Build**: Vulkan Release, 0 errors, 50 benign warnings (no -Werror), all + binaries produced +- **256K stress**: all 20 runs pass, zero "drafter decode failed with 1" errors + +Remaining low acceptance on out-of-distribution prompts is a drafter +training/generalization limitation, tracked separately from the Vulkan +enablement/fix rationale. + +## Archived investigation documents + +The following detailed investigation notes were consolidated into this document +and moved under: + +```text +docs/archive/dflash-q3cn-0acceptance-investigation/ +``` + +- `dflash-q3cn-vulkan-working.md` +- `dflash-q3cn-state-corruption-confirmed.md` +- `qwen3next-drafter-runtime-investigation.md` +- `dflash-drafter-debug-findings.md` +- `dflash-drafter-retrain-status.md` +- `dflash-cpp-vs-pytorch-graph-comparison.md` +- `dflash-rx-diag-results.md` +- `dflash-k-norm-fix-analysis.md` +- `dflash-k-norm-order-mismatch.md` +- `dflash-k-norm-reverted.md` +- `dflash-glm-state-review.md` +- `dflash-acceptance-diagnostic-report.md` +- `dflash-acceptance-diagnostic-report-v2.md` +- `dflash-acceptance-intermediate-report.md` +- `dflash-drafter-rootcause.md` +- `dflash-runtime-trace-results.md` +- `dflash-next-steps.md` +- `dflash-qwen3-coder-next-diagnostic-report.md` +- `vulkan-dflash-status.md` \ No newline at end of file diff --git a/ggml/src/ggml-vulkan/ggml-vulkan.cpp b/ggml/src/ggml-vulkan/ggml-vulkan.cpp index 6735afbd7f6d..6bd4630b9305 100644 --- a/ggml/src/ggml-vulkan/ggml-vulkan.cpp +++ b/ggml/src/ggml-vulkan/ggml-vulkan.cpp @@ -7682,6 +7682,7 @@ static void ggml_vk_buffer_copy(vk_buffer& dst, size_t dst_offset, vk_buffer& sr } } + static void ggml_vk_buffer_memset_async(vk_context& ctx, vk_buffer& dst, size_t offset, uint32_t c, size_t size) { VK_LOG_DEBUG("ggml_vk_buffer_memset_async(" << offset << ", " << c << ", " << size << ")"); @@ -17112,11 +17113,268 @@ static ggml_backend_dev_t ggml_backend_vk_reg_get_device(ggml_backend_reg_t reg, return devices[device]; } +// ---- DFlash GPU cross-ring (Vulkan) ---- +// Mirrors the CUDA dflash_cross_ring_gpu_* C-ABI. set_tensor/write_d2d use the +// tensor-variant signatures (ggml_tensor*) because Vulkan tensor t->data is a shared +// sentinel (vk_ptr_base + offset), not a real device address, so the drafter/target +// vk_buffer must be resolved via tensor->buffer->context. CUDA never registers these +// tensor-variant names (it uses the raw-ptr variants); see +// docs/superpowers/specs/2026-06-21-vulkan-dflash-cross-ring-design.md §7. +struct dflash_cross_ring_vk { + vk_device device; + int n_layers = 0; + int n_embd = 0; + int ring_size = 0; + vk_buffer rings; // device-local, size = n_layers*ring_size*n_embd*4 + vk_buffer staging; // device-local interleaved output, size = ring_size*n_layers*n_embd*4 + bool uma_host_visible = false; +}; + +// Runtime debug tracing for the Vulkan cross-ring. Enable with GGML_DFLASH_GPU_RING_DEBUG=1. +static bool dflash_vk_ring_debug_enabled() { + static int v = -1; + if (v == -1) { + const char * e = std::getenv("GGML_DFLASH_GPU_RING_DEBUG"); + v = (e && (e[0]=='1' || e[0]=='y' || (e[0]=='o' && e[1]=='n'))) ? 1 : 0; + } + return v == 1; +} +#define DFLASH_VK_RING_DBG(fmt, ...) do { if (dflash_vk_ring_debug_enabled()) fprintf(stderr, "[dflash-vk-ring] " fmt "\n", ##__VA_ARGS__); } while (0) + +// Tiny registry of our own staging buffers (keyed on bda_addr) so set_tensor_tensor can +// resolve a raw src ptr (a real device address) back to the staging vk_buffer. +static std::mutex g_dflash_vk_staging_mtx; +static std::vector> g_dflash_vk_staging; +static void dflash_vk_staging_register(vk_buffer & b) { + if (!b || b->bda_addr == 0) return; + std::lock_guard lk(g_dflash_vk_staging_mtx); + g_dflash_vk_staging.emplace_back(b->bda_addr, b); +} +static void dflash_vk_staging_unregister(vk::DeviceAddress addr) { + if (addr == 0) return; + std::lock_guard lk(g_dflash_vk_staging_mtx); + for (auto it = g_dflash_vk_staging.begin(); it != g_dflash_vk_staging.end(); ++it) { + if (it->first == addr) { g_dflash_vk_staging.erase(it); return; } + } +} +static vk_buffer dflash_vk_staging_resolve(vk::DeviceAddress addr) { + std::lock_guard lk(g_dflash_vk_staging_mtx); + for (auto & e : g_dflash_vk_staging) if (e.first == addr) return e.second; + return nullptr; +} + +extern "C" void * dflash_cross_ring_gpu_alloc_device(int dev_idx, int n_layers, int n_embd, int ring_size) { + DFLASH_VK_RING_DBG("alloc_device dev_idx=%d n_layers=%d n_embd=%d ring_size=%d", dev_idx, n_layers, n_embd, ring_size); + const char * env = std::getenv("GGML_DFLASH_GPU_RING"); + if (env && env[0] == '0') { DFLASH_VK_RING_DBG("alloc: disabled by GGML_DFLASH_GPU_RING=0"); return nullptr; } + if (n_layers <= 0 || n_embd <= 0 || ring_size <= 0) return nullptr; + vk_device dev = ggml_vk_get_device(dev_idx); + if (!dev) { DFLASH_VK_RING_DBG("alloc: dev_idx=%d not found", dev_idx); return nullptr; } + // set_tensor_tensor resolves our staging by its real device address, so the backend + // must support bufferDeviceAddress. If not, the cross-ring is unavailable (CPU fallback). + if (!dev->buffer_device_address) { DFLASH_VK_RING_DBG("alloc: buffer_device_address unsupported on dev=%s — CPU fallback", dev->name.c_str()); return nullptr; } + auto * ring = new dflash_cross_ring_vk(); + ring->device = dev; + ring->n_layers = n_layers; + ring->n_embd = n_embd; + ring->ring_size = ring_size; + size_t ring_bytes = (size_t)n_layers * ring_size * n_embd * sizeof(float); + size_t staging_bytes = (size_t)ring_size * n_layers * n_embd * sizeof(float); + try { + ring->rings = ggml_vk_create_buffer_device(dev, ring_bytes); + ring->staging = ggml_vk_create_buffer_device(dev, staging_bytes); + } catch (...) { + delete ring; + return nullptr; + } + ring->uma_host_visible = (ring->rings->ptr != nullptr); + dflash_vk_staging_register(ring->staging); + DFLASH_VK_RING_DBG("alloc: OK dev=%s uma=%d staging_bda=0x%llx", dev->name.c_str(), (int)ring->uma_host_visible, (unsigned long long)ring->staging->bda_addr); + return (void *)ring; +} +extern "C" void * dflash_cross_ring_gpu_alloc(int n_layers, int n_embd, int ring_size) { + return dflash_cross_ring_gpu_alloc_device(0, 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; + vk::DeviceAddress addr = ring->staging ? ring->staging->bda_addr : (vk::DeviceAddress)0; + dflash_vk_staging_unregister(addr); + delete ring; // vk_buffer shared_ptrs release on drop +} +extern "C" void dflash_cross_ring_gpu_write(void * handle, int layer, int ring_pos, const float * data, int n_tokens, int n_embd) { + if (!handle || !data || n_tokens <= 0 || n_embd <= 0) return; + auto * ring = (dflash_cross_ring_vk *)handle; + if (layer < 0 || layer >= ring->n_layers || n_embd != ring->n_embd) { DFLASH_VK_RING_DBG("write: reject layer=%d n_embd=%d (ring n_layers=%d n_embd=%d)", layer, n_embd, ring->n_layers, ring->n_embd); return; } + DFLASH_VK_RING_DBG("write layer=%d ring_pos=%d n_tokens=%d n_embd=%d uma=%d", layer, ring_pos, n_tokens, n_embd, (int)ring->uma_host_visible); + size_t off = ((size_t)layer * ring->ring_size + ring_pos) * ring->n_embd * sizeof(float); + size_t bytes = (size_t)n_tokens * n_embd * sizeof(float); + if (ring->uma_host_visible && ring->rings->ptr) { + memcpy((uint8_t*)ring->rings->ptr + off, data, bytes); + } else { + ggml_vk_buffer_write(ring->rings, off, data, bytes); + } +} +extern "C" void dflash_cross_ring_gpu_synchronize(void * handle) { + if (!handle) return; + auto * ring = (dflash_cross_ring_vk *)handle; + ring->device->transfer_queue.queue.waitIdle(); +} +extern "C" bool dflash_cross_ring_gpu_snapshot(void * handle, int write_pos, int filled, int ctx_window, float * host_data, int n_tokens, int n_layers, int n_embd) { + if (!handle || !host_data) return false; + auto * ring = (dflash_cross_ring_vk *)handle; + if (n_layers != ring->n_layers || n_embd != ring->n_embd) return false; + if (ctx_window <= 0 || n_tokens < 0) return false; + int cross_len = filled < ctx_window ? filled : ctx_window; + if (cross_len > ring->ring_size) cross_len = ring->ring_size; + if (n_tokens != cross_len) return false; + if (cross_len == 0) return true; + int read_start = ((write_pos - cross_len) % ring->ring_size + ring->ring_size) % ring->ring_size; + DFLASH_VK_RING_DBG("snapshot write_pos=%d filled=%d ctx_window=%d cross_len=%d read_start=%d", write_pos, filled, ctx_window, cross_len, read_start); + // D2H the whole rings buffer once, then layer-major copy with wrap on the host. + size_t ring_total = (size_t)ring->ring_size * ring->n_layers * ring->n_embd; + std::vector hrings(ring_total); + if (ring->uma_host_visible && ring->rings->ptr) { + memcpy(hrings.data(), ring->rings->ptr, ring_total * sizeof(float)); + } else { + ggml_vk_buffer_read(ring->rings, 0, hrings.data(), ring_total * sizeof(float)); + } + for (int layer = 0; layer < ring->n_layers; ++layer) { + const float * src_base = hrings.data() + (size_t)layer * ring->ring_size * ring->n_embd; + float * dst = host_data + (size_t)layer * (size_t)cross_len * (size_t)n_embd; + for (int t = 0; t < cross_len; ++t) { + int slot = (read_start + t) % ring->ring_size; + memcpy(dst + (size_t)t * n_embd, src_base + (size_t)slot * n_embd, (size_t)n_embd * sizeof(float)); + } + } + return true; +} +extern "C" const float * dflash_cross_ring_gpu_interleave(void * handle, int write_pos, int filled, int ctx_window) { + if (!handle) return nullptr; + auto * ring = (dflash_cross_ring_vk *)handle; + int cross_len = filled < ctx_window ? filled : ctx_window; + if (cross_len <= 0) return nullptr; + int read_start = ((write_pos - cross_len) % ring->ring_size + ring->ring_size) % ring->ring_size; + DFLASH_VK_RING_DBG("interleave write_pos=%d filled=%d ctx_window=%d cross_len=%d read_start=%d n_layers=%d", write_pos, filled, ctx_window, cross_len, read_start, ring->n_layers); + // GPU-side interleave: record one vkCmdCopyBuffer per (token, layer) slice into the + // interleaved staging buffer, all in a single transfer command buffer, then submit+wait. + // Keeps data on-GPU (no GPU->CPU->GPU round-trip). A dedicated compute shader could + // fuse these into one dispatch later. + std::lock_guard guard(ring->device->mutex); + vk_context subctx = ggml_vk_create_temporary_context(ring->device->transfer_queue.cmd_pool); + ggml_vk_ctx_begin(ring->device, subctx); + const size_t bytes = (size_t)ring->n_embd * sizeof(float); + for (int t = 0; t < cross_len; ++t) { + int slot = (read_start + t) % ring->ring_size; + for (int l = 0; l < ring->n_layers; ++l) { + size_t src_off = ((size_t)l * ring->ring_size + slot) * ring->n_embd * sizeof(float); + size_t dst_off = ((size_t)t * ring->n_layers + l) * ring->n_embd * sizeof(float); + ggml_vk_buffer_copy_async(subctx, ring->staging, dst_off, ring->rings, src_off, bytes); + } + } + ggml_vk_ctx_end(subctx); + ggml_vk_submit(subctx, ring->device->fence); + VK_CHECK(ring->device->device.waitForFences({ ring->device->fence }, true, UINT64_MAX), "dflash_cross_ring_gpu_interleave waitForFences"); + ring->device->device.resetFences({ ring->device->fence }); + ggml_vk_queue_command_pools_cleanup(ring->device); + return (const float *)ring->staging->bda_addr; +} +// Vulkan-only tensor variants. CUDA does not register these names. +// Vulkan equivalent of CUDA's dflash_cuda_backend_wait_for_stream: ensure the target +// graph's GPU hidden capture (ggml_cpy into hgpu->layers, run on the compute queue) +// is complete and visible to the cross-ring's transfer-queue D2D reads +// (write_d2d_tensor). Without this, write_d2d_tensor reads stale/incomplete +// hgpu->layers -> wrong drafter input -> low acceptance / target corruption. +// Registered under the CUDA name so dflash_capture_add_wait_backend finds it via +// the Vulkan backend registry. +extern "C" bool dflash_vk_backend_wait_for_stream(ggml_backend_t backend) { + if (!backend || !backend->context) return false; + // Only handle Vulkan backends (guard against a non-Vulkan backend being passed). + if (backend->iface.get_name != ggml_backend_vk_name) return false; + auto * ctx = (ggml_backend_vk_context *) backend->context; + if (!ctx || !ctx->device) return false; + // Full device wait: guarantees all compute-queue graph copies (the hidden + // capture) are complete and visible to subsequent transfer-queue reads. + // Heavier than CUDA's event-stream-wait but correct for backends without a + // cheap cross-queue event/semaphore exposed to the DFlash layer; a + // semaphore-based fast path can follow once the data-correctness is proven. + ctx->device->device.waitIdle(); + return true; +} + +extern "C" void dflash_cross_ring_gpu_set_tensor_tensor(ggml_tensor * dst, const void * src, size_t offset, size_t size) { + if (!dst || !src || size == 0 || !dst->buffer) return; + vk::DeviceAddress src_addr = (vk::DeviceAddress)(uintptr_t)src; + vk_buffer src_buf = dflash_vk_staging_resolve(src_addr); + if (!src_buf) return; + // If dst is NOT a Vulkan tensor (e.g. the drafter placed a target_hidden graph + // input on the CPU), the D2D vkCmdCopyBuffer path would compute a garbage + // vk_tensor_offset(dst) (dst->data is a real CPU pointer, not the 0x1000 + // sentinel) and write out of bounds -> SIGBUS. Fall back to a GPU->CPU + // readback of the staging region + ggml_backend_tensor_set (CPU memcpy into + // the CPU dst). This keeps the cross-ring correct for mixed CPU/GPU drafter + // input placements at the cost of one readback for the CPU-resident tensor. + if (!ggml_backend_buffer_is_vk(dst->buffer)) { + DFLASH_VK_RING_DBG("set_tensor_tensor CPU-dst fallback dst_off=%zu src_bda=0x%llx size=%zu", offset, (unsigned long long)src_addr, size); + std::vector cpu_tmp(size); + ggml_vk_buffer_read(src_buf, 0, cpu_tmp.data(), size); + ggml_backend_tensor_set(dst, cpu_tmp.data(), offset, size); + return; + } + if (!dst->buffer->context) return; + ggml_backend_vk_buffer_context * buf_ctx = (ggml_backend_vk_buffer_context *)dst->buffer->context; + vk_buffer dst_buf = buf_ctx->dev_buffer; + size_t dst_off = vk_tensor_offset(dst) + dst->view_offs + offset; + DFLASH_VK_RING_DBG("set_tensor_tensor dst_off=%zu src_bda=0x%llx size=%zu resolved=%d", dst_off, (unsigned long long)src_addr, size, (int)(src_buf != nullptr)); + if (!dst_buf) return; + // NOTE: cross-queue memory-visibility bug -- the drafter's graph reads + // target_hidden on the COMPUTE queue but this D2D copy writes it on the + // TRANSFER queue. Per the Vulkan memory model, cross-queue visibility needs a + // transfer->compute semaphore wired into the graph scheduler. vkDeviceWaitIdle + // and pipeline barriers were tested and only lift Coder-Next ring-ON acceptance + // from ~1.2% to ~5% (vs ~44% CPU-capture) and do NOT fix Qwen3.6 (hybrid + // arch: the GPU hidden capture also disrupts the target's DeltaNet compute). + // Until the semaphore is wired in, the cross-ring D2D produces wrong drafts; + // use GGML_DFLASH_GPU_RING=0 (CPU hidden capture) for correct DFlash. + ggml_vk_buffer_copy(dst_buf, dst_off, src_buf, 0, size); +} +extern "C" bool dflash_cross_ring_gpu_write_d2d_tensor(void * handle, int layer, int ring_pos, ggml_tensor * src, int src_offset, int n_tokens, int n_embd) { + // Phase 3: D2D copy target hidden tensor -> ring (no CPU readback). src is the + // target's device-resident hidden tensor (qwen3next GPU hidden capture); resolve its + // vk_buffer via tensor->buffer->context + vk_tensor_offset(). dst is our ring (handle). + if (!handle || !src || n_tokens <= 0 || n_embd <= 0 || !src->buffer || !src->buffer->context) return false; + auto * ring = (dflash_cross_ring_vk *)handle; + if (layer < 0 || layer >= ring->n_layers || n_embd != ring->n_embd) return false; + ggml_backend_vk_buffer_context * buf_ctx = (ggml_backend_vk_buffer_context *)src->buffer->context; + vk_buffer src_buf = buf_ctx->dev_buffer; + size_t src_off = vk_tensor_offset(src) + src->view_offs + (size_t)src_offset * (size_t)n_embd * sizeof(float); + size_t dst_off = ((size_t)layer * ring->ring_size + ring_pos) * ring->n_embd * sizeof(float); + size_t bytes = (size_t)n_tokens * n_embd * sizeof(float); + DFLASH_VK_RING_DBG("write_d2d_tensor layer=%d ring_pos=%d src_offset=%d n_tokens=%d n_embd=%d src_off=%zu dst_off=%zu bytes=%zu resolved=%d", layer, ring_pos, src_offset, n_tokens, n_embd, src_off, dst_off, bytes, (int)(src_buf != nullptr)); + if (!src_buf) return false; + ggml_vk_buffer_copy(ring->rings, dst_off, src_buf, src_off, bytes); + return true; +} + +static void * ggml_backend_vk_reg_get_proc_address(ggml_backend_reg_t /*reg*/, const char * name) { + if (std::strcmp(name, "dflash_cross_ring_gpu_alloc") == 0) return (void *)dflash_cross_ring_gpu_alloc; + if (std::strcmp(name, "dflash_cross_ring_gpu_alloc_device") == 0) return (void *)dflash_cross_ring_gpu_alloc_device; + if (std::strcmp(name, "dflash_cross_ring_gpu_free") == 0) return (void *)dflash_cross_ring_gpu_free; + if (std::strcmp(name, "dflash_cross_ring_gpu_write") == 0) return (void *)dflash_cross_ring_gpu_write; + if (std::strcmp(name, "dflash_cross_ring_gpu_synchronize") == 0) return (void *)dflash_cross_ring_gpu_synchronize; + if (std::strcmp(name, "dflash_cross_ring_gpu_snapshot") == 0) return (void *)dflash_cross_ring_gpu_snapshot; + if (std::strcmp(name, "dflash_cross_ring_gpu_interleave") == 0) return (void *)dflash_cross_ring_gpu_interleave; + if (std::strcmp(name, "dflash_cross_ring_gpu_set_tensor_tensor") == 0)return (void *)dflash_cross_ring_gpu_set_tensor_tensor; + if (std::strcmp(name, "dflash_cross_ring_gpu_write_d2d_tensor") == 0) return (void *)dflash_cross_ring_gpu_write_d2d_tensor; + if (std::strcmp(name, "dflash_cuda_backend_wait_for_stream") == 0) return (void *)dflash_vk_backend_wait_for_stream; + 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() { diff --git a/include/llama.h b/include/llama.h index ba01f215c0b3..06b8d1bf5d88 100644 --- a/include/llama.h +++ b/include/llama.h @@ -1159,6 +1159,8 @@ extern "C" { // cache and forces a reserve on next decode. LLAMA_API void llama_set_dflash_n_slots(struct llama_context * ctx, int n); + LLAMA_API void llama_set_dflash_target_kv_available(struct llama_context * ctx, bool avail); + // DFlash: enable/disable tape recording for DeltaNet rollback // When enabled, the eval callback records per-token DeltaNet inputs (k, v, gate, beta) // during verification decode for efficient state replay instead of full re-evaluation @@ -1200,13 +1202,24 @@ 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, int n_past_before, int n_accepted); + // Debug-only: dump a checksum of the committed recurrent state (s_l for the + // first recurrent layer) for seq_id. Used to compare tape-replay vs re-decode + // state advancement. No-op effect on the state. + LLAMA_API void llama_dflash_dump_recurrent_state_dbg( + struct llama_context * ctx, + llama_seq_id seq_id, + const char * tag); + // DFlash: wait for async tape replay to complete (must be called before next verify) LLAMA_API void llama_tape_replay_sync(struct llama_context * ctx); @@ -1265,6 +1278,19 @@ extern "C" { LLAMA_API void llama_dflash_cross_ring_gpu_synchronize(void * handle); LLAMA_API bool llama_dflash_cross_ring_gpu_snapshot(void * handle, int ring_write_pos, int ring_filled, int ctx_window, float * data, int n_tokens, int n_layers, int n_embd); + // DFlash: dump captured hidden states to a file for drafter training. + // Format: one line per token, tab-separated: + // slot_layer_idx_token_posn_embdf1_f2_..._fn_embd + // Returns number of tokens written. + LLAMA_API int64_t llama_dflash_dump_hidden_states(struct llama_context * ctx, const char * path); + + // DFlash: query captured hidden state counts. + // Returns total tokens captured across all layers for a given slot. + LLAMA_API int64_t llama_dflash_hidden_state_count(const struct llama_context * ctx, int slot); + + // DFlash: layer IDs being captured (returns count, fills array if provided) + LLAMA_API int llama_dflash_capture_layer_ids(const struct llama_context * ctx, int32_t * out, int out_size); + LLAMA_API void llama_dflash_cross_ring_gpu_set_cross(struct llama_context * ctx, void * handle, llama_seq_id seq_id, int ring_write_pos, int ring_filled, int n_layers, int n_embd, int ctx_window); LLAMA_API bool llama_dflash_kv_cache_init(struct llama_context * ctx, int ctx_size); LLAMA_API void llama_dflash_kv_cache_reset(struct llama_context * ctx); diff --git a/src/llama-context.cpp b/src/llama-context.cpp index 2dfa35538b2f..5471d0454c6f 100644 --- a/src/llama-context.cpp +++ b/src/llama-context.cpp @@ -22,6 +22,7 @@ #include "ggml-alloc.h" #include +#include #include #include #include @@ -68,6 +69,9 @@ static ggml_backend_reg_t dflash_gpu_backend_reg() { if (!reg) { reg = ggml_backend_reg_by_name("ROCm"); } + if (!reg) { + reg = ggml_backend_reg_by_name("Vulkan"); + } return reg; } @@ -244,6 +248,28 @@ llama_context::llama_context( __func__, cparams.n_rs_seq); cparams.n_rs_seq = 0; } + // QWEN35MOE (Qwen3.5-122B-A10B): keep n_rs_seq as requested by DFlash + // (common.cpp need_n_rs_seq() == draft.n_max, e.g. 4). Do NOT clamp to 0 here. + // + // Evidence (AMD Strix Halo / Vulkan, single-sequence DFlash, -np 1): + // clamping n_rs_seq=0 forces checkpoint-restore for draft rejection, but under + // the GPU cross-ring (GGML_DFLASH_GPU_RING=1) the target's DeltaNet recurrent-state + // cells desync from committed token positions (dozens of "non-consecutive token + // position N after N" warnings in llama-memory-recurrent.cpp) -> wrong drafts + // accepted -> repetition (rep=12). Keeping n_rs_seq=4 (RS partial rollback) makes + // BOTH ring=0 (CPU capture) and ring=1 (GPU ring) coherent (rep=0, zero desync + // warnings) and lets the GPU ring beat ring=0 (~16.0 vs ~14.6 tok/s). The earlier + // "n_rs_seq>0 -> garbage" note (commit 9d09310cc, DeltaNet RS-rollback snapshot bug + // under split_equal) is not reproduced in the single-sequence DFlash path; if + // split_equal multi-sequence regressions appear, fix the DeltaNet snapshot logic + // there, not by clamping n_rs_seq=0 here (which breaks the GPU cross-ring). QWEN35 + // dense (Qwen3.6, n_rs_seq=8) was never clamped and its RS path already works. + // + // Update (2026-06): the DeltaNet RS-snapshot write-back bug that produced + // "n_rs_seq>0 -> garbage" is now fixed in build_recurrent_attn (delta-net-base.cpp): + // for n_seq_tokens < K the older snapshot slots are rotated down instead of being + // overwritten with uninitialised kernel output. This makes n_rs_seq>0 correct under + // stochastic sampling too (Qwen3.6-35B-A3B and Qwen3.5-122B-A10B, qwen35moe, rep=0). cparams.ctx_type = params.ctx_type; cparams.yarn_ext_factor = params.yarn_ext_factor >= 0.0f ? params.yarn_ext_factor : hparams.yarn_ext_factor; cparams.yarn_attn_factor = params.yarn_attn_factor >= 0.0f ? params.yarn_attn_factor : hparams.yarn_attn_factor; @@ -1568,6 +1594,14 @@ static bool dflash_eval_callback(struct ggml_tensor * t, bool ask, void * user_d } return true; } + if (cap->tape_enabled) { + if (dflash_diagnostic_debug_enabled()) { + const char * nm = t->name ? t->name : ""; + if (strstr(nm, "qkv") || strstr(nm, "mixed")) { + fprintf(stderr, "[dflash-tape-cb] check name='%s' in_map=%d\n", nm, (int)cap->tape_name_map.count(t->name)); + } + } + } if (cap->tape_enabled && cap->tape_name_map.count(t->name)) { if (cap->profile) { cap->profile_cb_tape_ask++; @@ -1622,6 +1656,14 @@ static bool dflash_eval_callback(struct ggml_tensor * t, bool ask, void * user_d const size_t add_elems = (size_t) new_n * (size_t) new_embd; buf.data.resize(old_elems + add_elems); dflash_read_tensor_to(t, buf.data.data() + old_elems, add_elems); + // Capture token IDs for training data extraction + if (ub) { + const size_t old_n = buf.n_tokens; + buf.token_ids.resize(old_n + new_n); + for (int64_t ti = 0; ti < new_n; ++ti) { + buf.token_ids[old_n + ti] = ub->token[old_n + ti]; + } + } buf.n_tokens += new_n; return true; } @@ -1658,6 +1700,10 @@ static bool dflash_eval_callback(struct ggml_tensor * t, bool ask, void * user_d std::memcpy(buf.data.data() + old_elems, cap->scatter_buf.data() + (size_t) i * (size_t) new_embd, (size_t) new_embd * sizeof(float)); + // Capture token ID for training data extraction + if (ub) { + buf.token_ids.push_back(ub->token[i]); + } buf.n_tokens += 1; } return true; @@ -1767,6 +1813,16 @@ void llama_context::set_dflash_n_slots(int n) { gf_res_prev->reset(); } +void llama_context::set_dflash_target_kv_available(bool avail) { + if (cparams.dflash_target_kv_available == avail) { + return; + } + cparams.dflash_target_kv_available = avail; + // drafter graph attention branch depends on this → force a fresh reserve + sched_need_reserve = true; + gf_res_prev->reset(); +} + void llama_context::set_dflash_capture(const int32_t * layer_ids, int32_t n_layers) { if (layer_ids == nullptr || n_layers <= 0) { // Permanent deconfiguration: clear layer config and remove callback. @@ -1948,6 +2004,7 @@ void llama_context::dflash_reset_hidden_capture() { for (auto & buf : slot_bufs) { buf.n_tokens = 0; std::vector().swap(buf.data); + std::vector().swap(buf.token_ids); } } for (auto & hidden : dflash_capture->hidden_gpu) { @@ -1993,7 +2050,20 @@ 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 (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}; } } dflash_capture->tape_layers.resize(dflash_capture->recurrent_layer_ids.size()); @@ -2609,9 +2679,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 @@ -2624,7 +2698,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(memory.get()); @@ -2636,7 +2710,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 0; } const auto & hparams = model.hparams; @@ -2653,7 +2727,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 0; } const uint32_t n_embd_s = hparams.n_embd_s(); @@ -2667,11 +2741,10 @@ void llama_context::tape_replay(llama_seq_id seq_id, int n_accepted) { break; } } - if (!gpu_backend) { - tape_replay_cpu(mem_recurrent, cell_idx, n_accepted); + int n_fallback = tape_replay_cpu(mem_recurrent, cell_idx, n_accepted); tape_replay_conv(mem_recurrent, cell_idx, n_accepted, seq_id); - return; + return n_fallback; } // partial offload: if any recurrent layer's state lives on CPU, fall back to CPU replay @@ -2679,9 +2752,24 @@ void llama_context::tape_replay(llama_seq_id seq_id, int n_accepted) { for (int li = 0; li < (int) rec_ids.size(); ++li) { ggml_tensor * s_tensor = mem_recurrent->s_l[rec_ids[li]]; if (s_tensor && s_tensor->buffer && ggml_backend_buffer_is_host(s_tensor->buffer)) { - tape_replay_cpu(mem_recurrent, cell_idx, n_accepted); + int n_fallback = tape_replay_cpu(mem_recurrent, cell_idx, n_accepted); tape_replay_conv(mem_recurrent, cell_idx, n_accepted, seq_id); - return; + return n_fallback; + } + } + + // Vulkan GPU tape replay crashes (GDN ops not properly supported on Vulkan backends). + // Use CPU fallback with NaN guard instead. + // Note: Vulkan reports as GGML_BACKEND_DEVICE_TYPE_GPU, so we detect by driver name. + { + const char * dev_name = ggml_backend_dev_name(ggml_backend_get_device(gpu_backend)); + bool is_vulkan = dev_name && (strstr(dev_name, "vulkan") || strstr(dev_name, "Vulkan") || + strstr(dev_name, "RADV") || strstr(dev_name, "ANV") || + strstr(dev_name, "radeon")); + if (is_vulkan) { + int n_fallback = tape_replay_cpu(mem_recurrent, cell_idx, n_accepted); + tape_replay_conv(mem_recurrent, cell_idx, n_accepted, seq_id); + return n_fallback; } } @@ -2694,23 +2782,23 @@ 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", __func__, model.n_devices()); dflash_capture->multi_gpu_replay_fallback_logged = true; } - tape_replay_cpu(mem_recurrent, cell_idx, n_accepted); + int n_fallback = tape_replay_cpu(mem_recurrent, cell_idx, n_accepted); tape_replay_conv(mem_recurrent, cell_idx, n_accepted, seq_id); - return; + return n_fallback; } // GPU tape replay: build a ggml graph with GDN ops for all recurrent layers @@ -2855,9 +2943,9 @@ void llama_context::tape_replay(llama_seq_id seq_id, int n_accepted) { if (replay_graph_unsafe) { ggml_free(ctx); if (!use_gpu_tape) { - tape_replay_cpu(mem_recurrent, cell_idx, n_accepted); + int n_fallback = tape_replay_cpu(mem_recurrent, cell_idx, n_accepted); tape_replay_conv(mem_recurrent, cell_idx, n_accepted, seq_id); - return; + return n_fallback; } goto conv_rebuild; } @@ -2871,6 +2959,7 @@ void llama_context::tape_replay(llama_seq_id seq_id, int n_accepted) { ggml_backend_buffer_type_t gpu_buft = ggml_backend_get_default_buffer_type(gpu_backend); size_t needed = ggml_backend_alloc_ctx_tensors_from_buft_size(ctx, gpu_buft); + if (needed > dflash_capture->replay_buf_size) { if (dflash_capture->replay_buf) { ggml_backend_buffer_free(dflash_capture->replay_buf); @@ -2883,9 +2972,9 @@ void llama_context::tape_replay(llama_seq_id seq_id, int n_accepted) { if (!dflash_capture->replay_buf) { LLAMA_LOG_WARN("%s: failed to allocate GPU buffer for tape replay, falling back to CPU\n", __func__); ggml_free(ctx); - tape_replay_cpu(mem_recurrent, cell_idx, n_accepted); + int n_fallback = tape_replay_cpu(mem_recurrent, cell_idx, n_accepted); tape_replay_conv(mem_recurrent, cell_idx, n_accepted, seq_id); - return; + return n_fallback; } // assign tensors within the persistent buffer @@ -2902,8 +2991,12 @@ void llama_context::tape_replay(llama_seq_id seq_id, int n_accepted) { } } + // upload data for tensors that need it - for (auto & inp : inputs) { + + for (size_t ii = 0; ii < inputs.size(); ++ii) { + auto & inp = inputs[ii]; + // Q: always needs zeros { const int64_t S = inp.q->ne[0]; @@ -2930,18 +3023,21 @@ void llama_context::tape_replay(llama_seq_id seq_id, int n_accepted) { } // compute: launch GDN ops + state copies on GPU (async — overlap with next draft) + const int64_t t_replay_enqueue_us = dflash_capture->profile ? ggml_time_us() : 0; const ggml_status replay_status = ggml_backend_graph_compute_async(gpu_backend, graph); + if (replay_status != GGML_STATUS_SUCCESS) { LLAMA_LOG_WARN("%s: GPU DFlash recurrent replay graph failed with status %d; %s\n", __func__, (int) replay_status, use_gpu_tape ? "CPU fallback unavailable for GPU tape" : "falling back to CPU"); ggml_free(ctx); if (!use_gpu_tape) { - tape_replay_cpu(mem_recurrent, cell_idx, n_accepted); + int n_fallback = tape_replay_cpu(mem_recurrent, cell_idx, n_accepted); tape_replay_conv(mem_recurrent, cell_idx, n_accepted, seq_id); + return n_fallback; } - return; + return 0; } if (dflash_capture->profile) { const uint64_t elapsed = ggml_time_us() - t_replay_enqueue_us; @@ -2961,11 +3057,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) { @@ -3656,10 +3753,16 @@ void llama_context::tape_replay_conv(llama_memory_recurrent * mem_recurrent, int if (!mem_recurrent->r_l[il]) continue; const bool use_gpu_qkv = gpu_backend && gpu_layer && gpu_layer->qkv; + // tape.qkv_mixed layout is [conv_ch, n_seq_tokens, n_seqs] (ggml row-major); the + // rebuild below indexes qkv_mixed[(src_pos-conv_window)*conv_ch + ch] which is only + // valid for n_seqs==1 (seq 0). The CPU tape path returns false for n_seqs_unq>1 + // upstream, so this is safe today. If multi-seq CPU tape is ever enabled, the seq + // stride (n_seq_tokens*conv_ch) must be added to the index. if (!use_gpu_qkv) { - if (tape.n_tokens <= 0 || n_accepted > tape.n_tokens) continue; - if (tape.qkv_mixed.empty()) continue; + if (tape.n_tokens <= 0 || n_accepted > tape.n_tokens) { if (dflash_diagnostic_debug_enabled()) fprintf(stderr, "[dflash-tr-conv] layer=%d SKIP n_tokens=%d n_accepted=%d\n", il, tape.n_tokens, n_accepted); continue; } + if (tape.qkv_mixed.empty()) { if (dflash_diagnostic_debug_enabled()) fprintf(stderr, "[dflash-tr-conv] layer=%d SKIP qkv_mixed EMPTY\n", il); continue; } } + if (dflash_diagnostic_debug_enabled()) fprintf(stderr, "[dflash-tr-conv] layer=%d OK use_gpu_qkv=%d qkv_mixed.size=%zu n_tokens=%d n_accepted=%d\n", il, (int)use_gpu_qkv, tape.qkv_mixed.size(), tape.n_tokens, n_accepted); size_t qkv_seq_offset = 0; if (!use_gpu_qkv && tape.n_seqs > 1) { @@ -3859,7 +3962,8 @@ void llama_context::tape_replay_sync() { } // CPU fallback for tape replay (used when no GPU backend available) -void llama_context::tape_replay_cpu(llama_memory_recurrent * mem_recurrent, int32_t cell_idx, int n_accepted) { +// Returns 0 on success, n_accepted on failure (caller should re-decode instead). +int llama_context::tape_replay_cpu(llama_memory_recurrent * mem_recurrent, int32_t cell_idx, int n_accepted) { const auto & hparams = model.hparams; const auto & rec_ids = dflash_capture->recurrent_layer_ids; auto & tape_layers = dflash_capture->tape_layers; @@ -3869,25 +3973,56 @@ void llama_context::tape_replay_cpu(llama_memory_recurrent * mem_recurrent, int3 dflash_capture->profile_replay_layers += rec_ids.size(); } + // Helper to check for NaN/Inf in a float range + auto tape_data_ok = [](const float * data, size_t n) { + for (size_t i = 0; i < n; ++i) { + if (std::isnan(data[i]) || std::isinf(data[i])) { + return false; + } + } + return true; + }; + + // Accumulate results for all layers before writing back. + // This prevents partial write-back: if layer 2 fails, layers 0-1 are not + // committed to GPU. On failure, no state is written and the caller + // re-decodes from the backup position, advancing all layers consistently. + struct layer_result { + int layer_id; + ggml_tensor * s_tensor; + size_t s_offset; + size_t state_bytes; + std::vector state; + }; + std::vector results; + results.reserve(rec_ids.size()); + + int n_layers_skipped_dim = 0; + int n_layers_processed = 0; + const bool dbg = dflash_diagnostic_debug_enabled(); for (size_t li = 0; li < rec_ids.size(); ++li) { int il = rec_ids[li]; auto & tape = tape_layers[li]; - if (tape.n_tokens <= 0 || n_accepted > tape.n_tokens) continue; + if (tape.n_tokens <= 0 || n_accepted > tape.n_tokens) { ++n_layers_skipped_dim; continue; } const int64_t S = tape.S_k; const int64_t H_k = tape.H_k; const int64_t H_v = tape.H_v; if (S <= 0 || H_k <= 0 || H_v <= 0) { - continue; + if (dbg) fprintf(stderr, "[dflash-tr-cpu] layer=%d SKIP zero-dim S=%lld H_k=%lld H_v=%lld\n", il, (long long)S, (long long)H_k, (long long)H_v); + ++n_layers_skipped_dim; continue; } if ((size_t) S * (size_t) S * (size_t) H_v != (size_t) n_embd_s) { - continue; + if (dbg) fprintf(stderr, "[dflash-tr-cpu] layer=%d SKIP dim-check S*S*H_v=%zu != n_embd_s=%u (S=%lld H_v=%lld)\n", + il, (size_t)(S*S*H_v), n_embd_s, (long long)S, (long long)H_v); + ++n_layers_skipped_dim; continue; } ggml_tensor * s_tensor = mem_recurrent->s_l[il]; if (!s_tensor) { - continue; + if (dbg) fprintf(stderr, "[dflash-tr-cpu] layer=%d SKIP no s_tensor\n", il); + ++n_layers_skipped_dim; continue; } const size_t s_offset = (size_t)cell_idx * n_embd_s * ggml_element_size(s_tensor); const size_t state_bytes = (size_t) n_embd_s * ggml_element_size(s_tensor); @@ -3901,6 +4036,27 @@ void llama_context::tape_replay_cpu(llama_memory_recurrent * mem_recurrent, int3 std::vector state(n_embd_s); ggml_backend_tensor_get(s_tensor, state.data(), s_offset, n_embd_s * sizeof(float)); + // Guard: check tape data is available and valid before using it + size_t tape_k_size = (size_t)n_accepted * S * H_k; + size_t tape_v_size = (size_t)n_accepted * S * H_v; + size_t tape_gate_size = (size_t)n_accepted * H_v; + size_t tape_beta_size = (size_t)n_accepted * H_v; + if (tape.k.empty() || tape.v.empty() || tape.gate.empty() || tape.beta.empty() || + tape.k.size() < tape_k_size || tape.v.size() < tape_v_size || + tape.gate.size() < tape_gate_size || tape.beta.size() < tape_beta_size || + !tape_data_ok(tape.k.data(), tape_k_size) || + !tape_data_ok(tape.v.data(), tape_v_size) || + !tape_data_ok(tape.gate.data(), tape_gate_size) || + !tape_data_ok(tape.beta.data(), tape_beta_size) || + !tape_data_ok(state.data(), n_embd_s)) { + LLAMA_LOG_WARN("%s: DFlash CPU tape replay invalid tape data " + "(layer=%d n_accepted=%d k_size=%zu/%zu v_size=%zu/%zu) -> falling back to re-decode\n", + __func__, il, n_accepted, + tape_k_size, tape.k.size(), + tape_v_size, tape.v.size()); + return n_accepted; + } + for (int tok = 0; tok < n_accepted; ++tok) { for (int64_t hv = 0; hv < H_v; ++hv) { const int64_t hk = hv % H_k; @@ -3925,15 +4081,49 @@ void llama_context::tape_replay_cpu(llama_memory_recurrent * mem_recurrent, int3 } } - ggml_backend_tensor_set(s_tensor, state.data(), s_offset, n_embd_s * sizeof(float)); + // Guard: check result for NaN/Inf before accumulating + if (!tape_data_ok(state.data(), n_embd_s)) { + LLAMA_LOG_WARN("%s: DFlash CPU tape replay produced NaN/Inf in result " + "(layer=%d n_accepted=%d) -> falling back to re-decode\n", + __func__, il, n_accepted); + return n_accepted; + } + + // Accumulate result; write back only after all layers succeed + if (dbg) { + double sum = 0.0, abs = 0.0; + for (size_t i = 0; i < n_embd_s; ++i) { sum += state[i]; abs += std::fabs(state[i]); } + fprintf(stderr, "[dflash-tr-cpu] layer=%d OK n_accepted=%d S=%lld H_k=%lld H_v=%lld cell=%d " + "state[0..3]=%.6g,%.6g,%.6g,%.6g sum=%.4g abs_sum=%.4g gate0=%.6g beta0=%.6g k0=%.6g v0=%.6g\n", + il, n_accepted, (long long)S, (long long)H_k, (long long)H_v, cell_idx, + state[0], state[1], state[2], state[3], sum, abs, + tape.gate.empty()?0.0f:tape.gate[0], tape.beta.empty()?0.0f:tape.beta[0], + tape.k.empty()?0.0f:tape.k[0], tape.v.empty()?0.0f:tape.v[0]); + } + results.push_back({il, s_tensor, s_offset, state_bytes, std::move(state)}); + ++n_layers_processed; } + + if (dbg) { + fprintf(stderr, "[dflash-tr-cpu] DONE n_accepted=%d cell=%d n_rec=%zu processed=%d skipped=%d results=%zu -> returning %d\n", + n_accepted, cell_idx, rec_ids.size(), n_layers_processed, n_layers_skipped_dim, results.size(), + results.empty() ? n_accepted : 0); + } + + // All layers passed — write back results in one pass + for (auto & r : results) { + ggml_backend_tensor_set(r.s_tensor, r.state.data(), r.s_offset, r.state_bytes); + } + + return results.empty() ? n_accepted : 0; } -void llama_context::dflash_rollback(llama_seq_id seq_id, llama_seq_id seq_backup, int n_past_before, int n_accepted) { +__attribute__((noinline)) +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(memory.get()); if (!mem_hybrid) { LLAMA_LOG_WARN("%s: dflash_rollback requires hybrid memory\n", __func__); - return; + return 0; } const bool profile = dflash_capture && dflash_profile_has(dflash_capture->profile_flags, DFLASH_PROFILE_COPY); @@ -3981,8 +4171,20 @@ void llama_context::dflash_rollback(llama_seq_id seq_id, llama_seq_id seq_backup mem_recr->seq_rm(seq_backup, -1, -1); profile_lap(recurrent_restore_us); - // Replay DeltaNet state updates for accepted tokens - tape_replay(seq_id, n_accepted); + // Replay DeltaNet state updates for accepted tokens. + // On Vulkan, tape_replay_cpu is used (GPU tape crashes due to missing GDN ops). + // If tape_replay_cpu fails (NaN guard or empty tape), return n_accepted to signal + // the server to re-decode the accepted positions instead. + int n_positions_needing_reeval = 0; + { + n_positions_needing_reeval = tape_replay(seq_id, n_accepted); + if (n_positions_needing_reeval > 0) { + // Tape replay was skipped (e.g., Vulkan). The recurrent state is behind + // by n_positions_needing_reeval positions. Trim the attention memory + // to match the recurrent state, so seq_pos_max is consistent. + mem_attn->seq_rm(seq_id, n_past_before, -1); + } + } profile_lap(tape_launch_us); if (profile) { @@ -4004,6 +4206,59 @@ 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_positions_needing_reeval; +} + +void llama_context::dflash_dump_recurrent_state_dbg(llama_seq_id seq_id, const char * tag) { + auto * cap = dflash_capture.get(); + if (!cap || cap->recurrent_layer_ids.empty()) { fprintf(stderr, "[dflash-rs-dump] %s seq=%d no recurrent layers\n", tag ? tag : "?", (int)seq_id); return; } + auto * mem_hybrid = dynamic_cast(memory.get()); + if (!mem_hybrid) return; + auto * mem_recr = mem_hybrid->get_mem_recr(); + if (!mem_recr || (uint32_t)seq_id >= mem_recr->size) return; + int32_t cell_idx = mem_recr->cells[seq_id].tail; + const uint32_t n_embd_s = model.hparams.n_embd_s(); + int il0 = (int) cap->recurrent_layer_ids[0]; + ggml_tensor * s_tensor = (il0 >= 0 && il0 < (int)mem_recr->s_l.size()) ? mem_recr->s_l[il0] : nullptr; + if (!s_tensor || cell_idx < 0 || n_embd_s == 0) { + fprintf(stderr, "[dflash-rs-dump] %s seq=%d cell=%d il0=%d n_embd_s=%u (no state)\n", tag ? tag : "?", (int)seq_id, cell_idx, il0, n_embd_s); + return; + } + const size_t off = (size_t)cell_idx * n_embd_s * ggml_element_size(s_tensor); + std::vector st(n_embd_s); + ggml_backend_tensor_get(s_tensor, st.data(), off, n_embd_s * sizeof(float)); + double sum = 0, ab = 0; for (size_t i=0;irecurrent_layer_ids.back(); + ggml_tensor * s_last = (il_last >= 0 && il_last < (int)mem_recr->s_l.size()) ? mem_recr->s_l[il_last] : nullptr; + double sumL=0, abL=0, sL0=0; + if (s_last) { + std::vector sl(n_embd_s); + ggml_backend_tensor_get(s_last, sl.data(), off, n_embd_s * sizeof(float)); + for (size_t i=0;i= 0 && il0 < (int)mem_recr->r_l.size()) ? mem_recr->r_l[il0] : nullptr; + const uint32_t n_embd_r = model.hparams.n_embd_r(); + double rsum=0, rab=0, r0=0; size_t rn=0; + if (r_tensor && n_embd_r > 0) { + rn = n_embd_r; + const size_t roff = (size_t)cell_idx * n_embd_r * ggml_element_size(r_tensor); + const size_t rbytes = (size_t)n_embd_r * ggml_element_size(r_tensor); + if (dflash_tensor_span_in_bounds(r_tensor, roff, rbytes)) { + std::vector rv(n_embd_r); + ggml_backend_tensor_get(r_tensor, rv.data(), roff, n_embd_r * sizeof(float)); + for (size_t i=0;icells[seq_id].pos, + st[0], sum, ab, sL0, sumL, abL, r0, rsum, rab, rn); } void llama_context::dflash_prepare_branch(llama_seq_id seq_id, llama_seq_id seq_backup, int depth) { @@ -4104,7 +4359,8 @@ std::unique_ptr & llama_context::dflash_kv_cache_active_re void llama_context::set_cross_data_gpu( llama_seq_id seq_id, const void * d_staging, int cross_len, - int n_layers, int n_embd_layer, set_tensor_d2d_fn_t fn_d2d) { + int n_layers, int n_embd_layer, set_tensor_d2d_fn_t fn_d2d, + set_tensor_d2d_tensor_fn_t fn_d2d_tensor) { int64_t n_target_features = (int64_t)n_layers * n_embd_layer; const int64_t max_ctx = dflash_max_cross_ctx(); @@ -4120,12 +4376,13 @@ void llama_context::set_cross_data_gpu( cross.v_embd_gpu = d_staging; cross.v_embd_gpu_n_enc_real = cross_len; cross.fn_set_tensor_d2d = fn_d2d; + cross.fn_set_tensor_d2d_tensor = fn_d2d_tensor; cross.dflash_kv_cache = nullptr; if (seq_id >= 0) { dflash_kv_cache_active_seq = seq_id; } - const bool use_gpu_only = d_staging != nullptr && fn_d2d != nullptr && cparams.dflash_n_slots <= 1; + const bool use_gpu_only = d_staging != nullptr && (fn_d2d != nullptr || fn_d2d_tensor != nullptr) && cparams.dflash_n_slots <= 1; if (use_gpu_only) { std::vector().swap(cross.v_embd); } else { @@ -4721,8 +4978,9 @@ bool llama_context::dflash_kv_cache_update_gpu( int n_tokens, int n_layers, int n_embd_layer, - set_tensor_d2d_fn_t fn_d2d) { - if (!d_hidden || n_tokens <= 0 || n_layers <= 0 || n_embd_layer <= 0 || !fn_d2d) { + set_tensor_d2d_fn_t fn_d2d, + set_tensor_d2d_tensor_fn_t fn_d2d_tensor) { + if (!d_hidden || n_tokens <= 0 || n_layers <= 0 || n_embd_layer <= 0 || (!fn_d2d && !fn_d2d_tensor)) { return false; } @@ -4735,6 +4993,7 @@ bool llama_context::dflash_kv_cache_update_gpu( cross.dflash_kv_update_n_embd = n_target_features; cross.dflash_kv_update_n_enc_real = n_tokens; cross.dflash_kv_update_fn_set_tensor_d2d = fn_d2d; + cross.dflash_kv_update_fn_set_tensor_d2d_tensor = fn_d2d_tensor; const bool ok = dflash_kv_cache_update(n_tokens); @@ -4742,6 +5001,7 @@ bool llama_context::dflash_kv_cache_update_gpu( cross.dflash_kv_update_n_embd = 0; cross.dflash_kv_update_n_enc_real = 0; cross.dflash_kv_update_fn_set_tensor_d2d = nullptr; + cross.dflash_kv_update_fn_set_tensor_d2d_tensor = nullptr; return ok; } @@ -4895,8 +5155,9 @@ bool llama_context::dflash_target_kv_cache_update_gpu( int n_tokens, int n_layers, int n_embd_layer, - set_tensor_d2d_fn_t fn_d2d) { - if (!d_hidden || n_tokens <= 0 || n_layers <= 0 || n_embd_layer <= 0 || !fn_d2d) { + set_tensor_d2d_fn_t fn_d2d, + set_tensor_d2d_tensor_fn_t fn_d2d_tensor) { + if (!d_hidden || n_tokens <= 0 || n_layers <= 0 || n_embd_layer <= 0 || (!fn_d2d && !fn_d2d_tensor)) { return false; } if (!llm_arch_is_dflash_drafter(model.arch) || !memory) { @@ -4997,6 +5258,7 @@ bool llama_context::dflash_target_kv_cache_update_gpu( cross.dflash_kv_update_n_embd = n_target_features; cross.dflash_kv_update_n_enc_real = n_tokens; cross.dflash_kv_update_fn_set_tensor_d2d = fn_d2d; + cross.dflash_kv_update_fn_set_tensor_d2d_tensor = fn_d2d_tensor; const ggml_status status = [&] { res->set_inputs(&ubatch); @@ -5007,6 +5269,7 @@ bool llama_context::dflash_target_kv_cache_update_gpu( cross.dflash_kv_update_n_embd = 0; cross.dflash_kv_update_n_enc_real = 0; cross.dflash_kv_update_fn_set_tensor_d2d = nullptr; + cross.dflash_kv_update_fn_set_tensor_d2d_tensor = nullptr; // Reused commit scratch lives on cudaStreamPerThread, while the graph runs on // the backend stream. Order those streams before this function returns so the @@ -6230,6 +6493,32 @@ int llama_context::decode(const llama_batch & batch_inp) { const uint32_t n_tokens_all = balloc->get_n_tokens(); const uint32_t n_outputs_all = balloc->get_n_outputs(); + // DFlash reduced-verify safety: when logits are expected by the caller + // (batch.logits has at least one true entry), the output buffer must + // contain logits. If dflash_reduced_consumer_active suppresses logits + // allocation but the batch still requests them, the subsequent + // llama_get_logits_ith() call will crash. Catch this early. + if (n_outputs_all == 0 && !output_all) { + bool any_logits_requested = false; + for (int32_t i = 0; i < batch_inp.n_tokens; ++i) { + if (batch_inp.logits && batch_inp.logits[i]) { + any_logits_requested = true; + break; + } + } + if (any_logits_requested) { + const bool dflash_reduced = cparams.dflash_reduced_consumer_active && cparams.dflash_verify_logits; + LLAMA_LOG_ERROR("%s: batch requests logits but n_outputs_all=0 " + "(dflash_reduced_consumer=%d dflash_verify_logits=%d dflash_reduced=%d)\n", + __func__, + cparams.dflash_reduced_consumer_active ? 1 : 0, + cparams.dflash_verify_logits ? 1 : 0, + dflash_reduced ? 1 : 0); + GGML_ASSERT(n_outputs_all > 0 && "batch requests logits but n_outputs_all=0; " + "if DFlash reduced verify is active, logits buffer will be null and sampling will crash"); + } + } + if (output_all) { // require that all tokens are output if (n_outputs_all != n_tokens_all) { @@ -6757,7 +7046,7 @@ int llama_context::decode(const llama_batch & batch_inp) { if (dflash_capture && dflash_crash_trace_enabled()) { LLAMA_LOG_INFO( - "%s: dflash crash breadcrumb: before process_ubatch tokens=%u seq_tokens=%u seqs=%u seqs_unq=%u " + "%s: dflash crash breadcrumb: before process_ubatch tokens=%u seq_tokens=%u seqs=%u seqs_unq=%u " "capture_active=%d capture_n_tokens=%d capture_n_seqs=%d gpu_ready=%d tree_active=%d " "prefill_plan_active=%d prefill_max=%d use_prefill=%d prefill_active=%d prefill_n_seqs=%d prefill_n_tokens=%d " "hidden_ready=%d hidden_n_seqs=%d tape_ready=%d tape_n_seqs=%d skip_cb=%d suppress_view=%d seqs_changed=%d " @@ -8589,6 +8878,10 @@ void llama_set_dflash_n_slots(llama_context * ctx, int n) { ctx->set_dflash_n_slots(n); } +void llama_set_dflash_target_kv_available(llama_context * ctx, bool avail) { + ctx->set_dflash_target_kv_available(avail); +} + void llama_set_tape_recording(llama_context * ctx, bool enable) { ctx->set_tape_recording(enable); } @@ -8625,8 +8918,12 @@ 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_dump_recurrent_state_dbg(llama_context * ctx, llama_seq_id seq_id, const char * tag) { + if (ctx) ctx->dflash_dump_recurrent_state_dbg(seq_id, tag); } void llama_dflash_prepare_branch(llama_context * ctx, llama_seq_id seq_id, llama_seq_id seq_backup, int depth) { @@ -8652,13 +8949,17 @@ struct dflash_cross_ring_handle { bool (*fn_snapshot)(void *, int, int, int, float *, int, int, int); const float * (*fn_interleave)(void *, int, int, int); void (*fn_set_tensor)(void *, const void *, size_t, size_t); + // Tensor-variant fn pointers (Vulkan): resolve vk_buffer from ggml_tensor*. + // Null on CUDA (CUDA uses the raw-ptr variants above); the raw variants are null on Vulkan. + void (*fn_set_tensor_tensor)(ggml_tensor *, const void *, size_t, size_t); + bool (*fn_write_d2d_tensor)(void *, int, int, ggml_tensor *, int, int, int); }; void * llama_context::init_cross_ring_gpu(int n_layers, int n_embd, int ring_size) { ggml_backend_reg_t cuda_reg = dflash_gpu_backend_reg(); if (!cuda_reg) { if (dflash_multi_gpu_debug_enabled() || dflash_diagnostic_debug_enabled()) { - LLAMA_LOG_WARN("%s: dflash GPU ring unavailable: CUDA/ROCm backend registry not found\n", __func__); + LLAMA_LOG_WARN("%s: dflash GPU ring unavailable: CUDA/ROCm/Vulkan backend registry not found\n", __func__); } return nullptr; } @@ -8678,6 +8979,8 @@ void * llama_context::init_cross_ring_gpu(int n_layers, int n_embd, int ring_siz using snapshot_fn_t = bool (*)(void *, int, int, int, float *, int, int, int); using interleave_fn_t = const float * (*)(void *, int, int, int); using set_tensor_fn_t = void (*)(void *, const void *, size_t, size_t); + using set_tensor_tensor_fn_t = void (*)(ggml_tensor *, const void *, size_t, size_t); + using write_d2d_tensor_fn_t = bool (*)(void *, int, int, ggml_tensor *, int, int, int); auto fn_alloc_device = (alloc_device_fn_t) ggml_backend_reg_get_proc_address(cuda_reg, "dflash_cross_ring_gpu_alloc_device"); @@ -8689,10 +8992,16 @@ void * llama_context::init_cross_ring_gpu(int n_layers, int n_embd, int ring_siz auto fn_snapshot = (snapshot_fn_t) ggml_backend_reg_get_proc_address(cuda_reg, "dflash_cross_ring_gpu_snapshot"); auto fn_interleave = (interleave_fn_t) ggml_backend_reg_get_proc_address(cuda_reg, "dflash_cross_ring_gpu_interleave"); auto fn_set_tensor = (set_tensor_fn_t) ggml_backend_reg_get_proc_address(cuda_reg, "dflash_cross_ring_gpu_set_tensor"); + // Tensor variants (Vulkan-only). Null on CUDA, which uses the raw variants above. + auto fn_set_tensor_tensor = (set_tensor_tensor_fn_t) ggml_backend_reg_get_proc_address(cuda_reg, "dflash_cross_ring_gpu_set_tensor_tensor"); + auto fn_write_d2d_tensor = (write_d2d_tensor_fn_t) ggml_backend_reg_get_proc_address(cuda_reg, "dflash_cross_ring_gpu_write_d2d_tensor"); - if (!fn_alloc || !fn_free || !fn_write || !fn_write_d2d || !fn_sync || !fn_snapshot || !fn_interleave || !fn_set_tensor) { + // Need the 7 shared fns, plus at least one of each {set_tensor, set_tensor_tensor} / {write_d2d, write_d2d_tensor} pair. + if (!fn_alloc || !fn_free || !fn_write || !fn_sync || !fn_snapshot || !fn_interleave) { return nullptr; } + if (!fn_write_d2d && !fn_write_d2d_tensor) return nullptr; + if (!fn_set_tensor && !fn_set_tensor_tensor) return nullptr; const int ring_device = dflash_gpu_ring_device_override(); @@ -8713,6 +9022,8 @@ void * llama_context::init_cross_ring_gpu(int n_layers, int n_embd, int ring_siz handle->fn_snapshot = fn_snapshot; handle->fn_interleave = fn_interleave; handle->fn_set_tensor = fn_set_tensor; + handle->fn_set_tensor_tensor = fn_set_tensor_tensor; + handle->fn_write_d2d_tensor = fn_write_d2d_tensor; return handle; } @@ -8800,9 +9111,17 @@ bool llama_context::cross_ring_gpu_write_hidden(void * handle, int layer, int ri } auto * h = (dflash_cross_ring_handle *)handle; + // Vulkan: tensor-variant D2D resolves the target hidden vk_buffer from ggml_tensor* + // and takes src_offset explicitly (the raw variant bakes the offset into the pointer, + // which is a Vulkan sentinel and cannot be used). CUDA leaves fn_write_d2d_tensor null. + if (h->fn_write_d2d_tensor && + h->fn_write_d2d_tensor(h->gpu_ring, layer, ring_pos, tensor, src_offset, n_tokens, n_embd)) { + return true; + } + // CUDA: raw-ptr D2D (offset baked into the pointer; tensor->data is a real device address). 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)) { + if (h->fn_write_d2d && h->fn_write_d2d(h->gpu_ring, layer, ring_pos, src, n_tokens, n_embd)) { return true; } @@ -8851,9 +9170,14 @@ bool llama_context::prefill_gpu_write_hidden(void * handle, int slot, int layer, } auto * h = (dflash_cross_ring_handle *)handle; + // Vulkan: tensor-variant D2D (see cross_ring_gpu_write_hidden). CUDA leaves it null. + if (h->fn_write_d2d_tensor && + h->fn_write_d2d_tensor(h->gpu_ring, layer, ring_pos, tensor, src_offset, n_tokens, n_embd)) { + return true; + } 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)) { + if (h->fn_write_d2d && h->fn_write_d2d(h->gpu_ring, layer, ring_pos, src, n_tokens, n_embd)) { return true; } @@ -8928,7 +9252,103 @@ void llama_dflash_cross_ring_gpu_set_cross( if (!d_staging) return; int cross_len = ring_filled < ctx_window ? ring_filled : ctx_window; - ctx->set_cross_data_gpu(seq_id, d_staging, cross_len, n_layers, n_embd, h->fn_set_tensor); + ctx->set_cross_data_gpu(seq_id, d_staging, cross_len, n_layers, n_embd, h->fn_set_tensor, h->fn_set_tensor_tensor); +} + +// DFlash: dump captured hidden states to file for drafter training +int64_t llama_context::dflash_dump_hidden_states(const char * path) { + auto * cap = dflash_capture.get(); + if (!cap || cap->hiddens == nullptr) + return 0; + + // Append mode ("a") so repeated calls across speculative cycles accumulate data. + // The header is written only on a new file; subsequent appends skip it. + FILE * f = fopen(path, "a"); + if (!f) return 0; + + // Check if file is empty (first write) to decide whether to write header + fseek(f, 0, SEEK_END); + bool empty = ftell(f) == 0; + rewind(f); + + if (empty) { + // Header: layer_ids + const auto & layer_ids = cap->layer_ids; + fprintf(f, "# layers: "); + for (size_t i = 0; i < layer_ids.size(); i++) { + if (i) fprintf(f, " "); + fprintf(f, "%d", layer_ids[i]); + } + fprintf(f, "\n"); + } + + int64_t total_tokens = 0; + const auto & hiddens = *cap->hiddens; + for (int slot = 0; slot < (int) hiddens.size(); slot++) { + const auto & slot_bufs = hiddens[slot]; + for (size_t h = 0; h < slot_bufs.size(); h++) { + const auto & buf = slot_bufs[h]; + if (buf.n_tokens == 0) continue; + for (int64_t t = 0; t < buf.n_tokens; t++) { + const float * row = buf.data.data() + (size_t) t * buf.n_embd; + int32_t token_id = (t < (int64_t) buf.token_ids.size()) ? buf.token_ids[t] : -1; + fprintf(f, "%d %zu %ld %d ", slot, h, t, token_id); + for (int64_t e = 0; e < buf.n_embd; e++) { + if (e) fprintf(f, " "); + fprintf(f, "%.8g", row[e]); + } + fprintf(f, "\n"); + } + total_tokens += buf.n_tokens; + } + } + + fclose(f); + return total_tokens; +} + +// DFlash: query captured hidden state count for a slot +int64_t llama_context::dflash_hidden_state_count(int slot) const { + auto * cap = dflash_capture.get(); + if (!cap || cap->hiddens == nullptr) + return 0; + const auto & hiddens = *cap->hiddens; + if (slot < 0 || slot >= (int) hiddens.size()) + return 0; + int64_t total = 0; + for (const auto & buf : hiddens[slot]) { + total += buf.n_tokens; + } + return total; +} + +// DFlash: get capture layer IDs +int llama_context::dflash_capture_layer_ids(int32_t * out, int out_size) const { + auto * cap = dflash_capture.get(); + if (!cap) + return 0; + const auto & ids = cap->layer_ids; + int n = (int) ids.size(); + if (out && out_size > 0) { + int copy = std::min(n, out_size); + for (int i = 0; i < copy; i++) { + out[i] = ids[i]; + } + } + return n; +} + +// External API wrappers +int64_t llama_dflash_dump_hidden_states(llama_context * ctx, const char * path) { + return ctx ? ctx->dflash_dump_hidden_states(path) : 0; +} + +int64_t llama_dflash_hidden_state_count(const llama_context * ctx, int slot) { + return ctx ? ctx->dflash_hidden_state_count(slot) : 0; +} + +int llama_dflash_capture_layer_ids(const llama_context * ctx, int32_t * out, int out_size) { + return ctx ? ctx->dflash_capture_layer_ids(out, out_size) : 0; } bool llama_dflash_kv_cache_init(llama_context * ctx, int ctx_size) { @@ -8964,7 +9384,7 @@ bool llama_dflash_kv_cache_update_from_ring( } const int n_update = std::min(ring_filled, n_tokens); - return ctx->dflash_kv_cache_update_gpu(d_staging, n_update, n_layers, n_embd, h->fn_set_tensor); + return ctx->dflash_kv_cache_update_gpu(d_staging, n_update, n_layers, n_embd, h->fn_set_tensor, h->fn_set_tensor_tensor); } bool llama_dflash_kv_cache_update_from_ring_seq( @@ -9002,7 +9422,7 @@ bool llama_dflash_target_kv_cache_update_from_ring( const int n_update = std::min(ring_filled, n_tokens); return ctx->dflash_target_kv_cache_update_gpu( - seq_id, start_pos, d_staging, n_update, n_layers, n_embd, h->fn_set_tensor); + seq_id, start_pos, d_staging, n_update, n_layers, n_embd, h->fn_set_tensor, h->fn_set_tensor_tensor); } void llama_set_tree_mask(llama_context * ctx, const uint8_t * visibility, int n_tree_tokens) { diff --git a/src/llama-context.h b/src/llama-context.h index 0c9e80297aaa..f33320724840 100644 --- a/src/llama-context.h +++ b/src/llama-context.h @@ -21,6 +21,7 @@ static inline bool llama_dflash_gpu_hidden_supported_arch(llm_arch arch) { switch (arch) { case LLM_ARCH_QWEN35: case LLM_ARCH_QWEN35MOE: + case LLM_ARCH_QWEN3NEXT: case LLM_ARCH_GEMMA4: return true; default: @@ -94,6 +95,59 @@ static inline bool llama_dflash_view_span_in_bounds_for_test( return offset_bytes <= total_bytes && n_bytes <= total_bytes - offset_bytes; } +// DFlash DeltaNet RS-snapshot write-back slot decision. +// +// build_recurrent_attn() (src/models/delta-net-base.cpp) writes the K recurrent-state +// snapshot slots back into ssm_states_all after a gated_delta_net step. The kernel +// only populates the last min(n_seq_tokens, K) gdn_out snapshot slots; for a decode / +// short batch (n_seq_tokens < K) the leading gdn_out slots are uninitialised, so the +// older state snapshots must be rotated down by n_pop slots (old slot s -> s + n_pop) +// instead of being clobbered with garbage. This helper captures that decision as a +// pure function so it is unit-testable (see tests/test-dflash-plumbing.cpp) and so the +// production write-back loop and the test share one source of truth. +// +// For each output slot index k_i in [0, K): +// cache_slot = K - 1 - k_i (destination state slot) +// n_pop = min(n_seq_tokens, K) (kernel-populated gdn_out slots) +// if k_i >= K - n_pop: from_gdn=true, gdn_slot = k_i (new state from this batch) +// else: from_gdn=false, old_slot = cache_slot - n_pop (rotated older snapshot) +// Iterating k_i ascending visits cache_slot descending, so each old_slot is read before +// it is overwritten by a later iteration (in-place shift is safe). +// +// Returns false (and leaves outputs untouched) if k_i is out of range [0, K) or the +// inputs are invalid (K < 1, n_seq_tokens < 0, n_pop == 0). +struct llama_dflash_rs_slot_src { + bool from_gdn; // true: source is gdn_out snapshot slot gdn_slot; false: rotate from old state slot old_slot + int64_t gdn_slot; // valid when from_gdn == true + uint32_t old_slot; // valid when from_gdn == false; the older state slot to rotate FROM +}; + +static inline bool llama_dflash_rs_writeback_slot_for_test( + int64_t K, + int64_t n_seq_tokens, + int64_t k_i, + uint32_t & cache_slot, + llama_dflash_rs_slot_src & src) { + if (K < 1 || n_seq_tokens < 0 || k_i < 0 || k_i >= K) { + return false; + } + const int64_t n_pop = (n_seq_tokens < K) ? n_seq_tokens : K; // kernel-populated slots + if (n_pop == 0) { + return false; + } + cache_slot = (uint32_t) (K - 1 - k_i); + if (k_i >= K - n_pop) { + src.from_gdn = true; + src.gdn_slot = k_i; + src.old_slot = 0; + } else { + src.from_gdn = false; + src.gdn_slot = 0; + src.old_slot = cache_slot - (uint32_t) n_pop; // rotate older snapshot down by n_pop + } + return true; +} + class llama_memory_recurrent; struct llama_model; @@ -109,6 +163,7 @@ struct llama_memory_context_i; // DFlash: hidden state buffer for captured layer activations struct dflash_layer_hidden_buf { std::vector data; + std::vector token_ids; // token ID at each position (for training data) int64_t n_embd = 0; int64_t n_tokens = 0; }; @@ -700,6 +755,11 @@ struct llama_context { void set_dflash_consume_reduced(bool enabled); void set_dflash_n_slots(int n); + // DFlash: mark whether the drafter's normal KV cache is populated with TARGET + // context K/V (GPU cross ring available). When false, the drafter graph must + // fall back to fresh Kcur_ctx = wk(fused_target) instead of reading the cache. + void set_dflash_target_kv_available(bool avail); + // DFlash: reset hidden-state capture for a fresh decode() call so the // eval callback accumulates across this call's ubatches void dflash_reset_hidden_capture(); @@ -709,6 +769,11 @@ struct llama_context { bool dflash_hidden_capture_available() const; const char * dflash_hidden_capture_unavailable_reason() const; + // DFlash: dump captured hidden states to file for drafter training + int64_t dflash_dump_hidden_states(const char * path); + int64_t dflash_hidden_state_count(int slot) const; + int dflash_capture_layer_ids(int32_t * out, int out_size) const; + // DFlash: enable/disable tape recording for DeltaNet state rollback void set_tape_recording(bool enable); void dflash_ensure_recurrent_setup(); @@ -733,7 +798,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 0 if tape replay succeeded, or n_accepted if it was skipped (e.g., Vulkan) + // and the server needs to re-decode the accepted positions. + 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); @@ -741,11 +808,16 @@ struct llama_context { bool tape_replay_gdn_direct_from_cpu_tape(llama_memory_recurrent * mem_recurrent, int32_t cell_idx, int n_accepted); bool tape_replay_conv_gpu_from_cpu_tape(llama_memory_recurrent * mem_recurrent, int32_t cell_idx, int n_accepted, llama_seq_id seq_id); void tape_replay_conv(llama_memory_recurrent * mem_recurrent, int32_t cell_idx, int n_accepted, llama_seq_id seq_id = 0); - void tape_replay_cpu(llama_memory_recurrent * mem_recurrent, int32_t cell_idx, int n_accepted); + int tape_replay_cpu(llama_memory_recurrent * mem_recurrent, int32_t cell_idx, int n_accepted); 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 need re-decoding after rollback. + // 0 means tape replay succeeded; >0 means the server must re-decode. + int dflash_rollback(llama_seq_id seq_id, llama_seq_id seq_backup, int n_past_before, int n_accepted); + + // DFlash debug: dump committed recurrent state checksum for seq_id (no-op on state) + void dflash_dump_recurrent_state_dbg(llama_seq_id seq_id, const char * tag); // 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); @@ -766,15 +838,17 @@ struct llama_context { bool prefill_gpu_write_hidden(void * handle, int slot, int layer, int ring_pos, int src_offset, int n_tokens, int n_embd); // DFlash GPU ring: set GPU device pointer as cross data source (D2D path) - using set_tensor_d2d_fn_t = void (*)(void *, const void *, size_t, size_t); + using set_tensor_d2d_fn_t = void (*)(void *, const void *, size_t, size_t); + using set_tensor_d2d_tensor_fn_t = void (*)(ggml_tensor *, const void *, size_t, size_t); void set_cross_data_gpu(llama_seq_id seq_id, const void * d_staging, int cross_len, - int n_layers, int n_embd, set_tensor_d2d_fn_t fn_d2d); + int n_layers, int n_embd, set_tensor_d2d_fn_t fn_d2d, + set_tensor_d2d_tensor_fn_t fn_d2d_tensor = nullptr); bool dflash_kv_cache_init(int ctx_size); void dflash_kv_cache_reset(); bool dflash_kv_cache_update(int n_tokens); - bool dflash_kv_cache_update_gpu(const void * d_hidden, int n_tokens, int n_layers, int n_embd_layer, set_tensor_d2d_fn_t fn_d2d); - bool dflash_target_kv_cache_update_gpu(llama_seq_id seq_id, llama_pos start_pos, const void * d_hidden, int n_tokens, int n_layers, int n_embd_layer, set_tensor_d2d_fn_t fn_d2d); + bool dflash_kv_cache_update_gpu(const void * d_hidden, int n_tokens, int n_layers, int n_embd_layer, set_tensor_d2d_fn_t fn_d2d, set_tensor_d2d_tensor_fn_t fn_d2d_tensor = nullptr); + bool dflash_target_kv_cache_update_gpu(llama_seq_id seq_id, llama_pos start_pos, const void * d_hidden, int n_tokens, int n_layers, int n_embd_layer, set_tensor_d2d_fn_t fn_d2d, set_tensor_d2d_tensor_fn_t fn_d2d_tensor = nullptr); bool dflash_kv_cache_prepare(int ctx_window); bool dflash_kv_cache_prepare_batch(const llama_seq_id * seq_ids, int n_seq, int ctx_window); dflash_kv_cache_data * dflash_kv_cache_active(); diff --git a/src/llama-cparams.h b/src/llama-cparams.h index 1643b2b7376c..fc5d64526a35 100644 --- a/src/llama-cparams.h +++ b/src/llama-cparams.h @@ -71,6 +71,16 @@ struct llama_cparams { // DFlash: cross-attention window in tokens (how many target hidden states the drafter sees) int dflash_cross_ctx = 512; + // DFlash drafter: true when the drafter's normal KV cache is actually populated + // with TARGET context K/V (via the GPU cross ring + update_drafter_kv_cache). + // When false (e.g. Vulkan where the GPU cross ring is unavailable and CPU + // hidden capture is used), the full-attention if-branch would read an + // empty/stale KV cache and get NO target context -> 0%% acceptance for + // all-full-attention cross-attention drafters (e.g. Z-Lab Qwen3-Coder-Next). + // In that case the graph MUST fall back to the else branch (fresh + // Kcur_ctx = wk(fused_target) from the CPU-captured target_hidden). + bool dflash_target_kv_available = false; + // DFlash drafter: number of concurrent slots the batched drafter graph is reserved // for. ctx_len in the drafter graph = dflash_n_slots * dflash_cross_ctx, // and drafter n_tokens reservation = dflash_n_slots * block_size. Set on the diff --git a/src/llama-graph.h b/src/llama-graph.h index f1cc11096504..793d97bf9a67 100644 --- a/src/llama-graph.h +++ b/src/llama-graph.h @@ -93,6 +93,8 @@ struct llama_cross { const void * v_embd_gpu = nullptr; int64_t v_embd_gpu_n_enc_real = 0; void (*fn_set_tensor_d2d)(void * d_dst, const void * d_src, size_t offset, size_t size) = nullptr; + // Vulkan tensor-variant: resolves vk_buffer from ggml_tensor* (CUDA uses the raw variant above). + void (*fn_set_tensor_d2d_tensor)(ggml_tensor * d_dst, const void * d_src, size_t offset, size_t size) = nullptr; // Temporary DFlash K/V-update input. This lets the drafter-side projection // cache read newly committed hidden states without mutating the main cross @@ -102,6 +104,7 @@ struct llama_cross { int64_t dflash_kv_update_n_embd = 0; int64_t dflash_kv_update_n_enc_real = 0; void (*dflash_kv_update_fn_set_tensor_d2d)(void * d_dst, const void * d_src, size_t offset, size_t size) = nullptr; + void (*dflash_kv_update_fn_set_tensor_d2d_tensor)(ggml_tensor * d_dst, const void * d_src, size_t offset, size_t size) = nullptr; // Per-seq cross buffers for DFlash multi-slot. // When non-empty, graph builders should pack these into target_hidden per slot diff --git a/src/llama-model.cpp b/src/llama-model.cpp index 7fefd6ba38ee..f9faedd4dc54 100644 --- a/src/llama-model.cpp +++ b/src/llama-model.cpp @@ -2492,8 +2492,24 @@ float llama_model_rope_freq_scale_train_swa(const llama_model * model) { } void llama_model_share_tensors(llama_model * dst, const llama_model * src) { - dst->tok_embd = src->tok_embd; - dst->output = src->output; + // DFlash: share tok_embd and output tensors from src (target) to dst (drafter). + // Only share when the drafter does NOT already have its own tensor loaded from + // its GGUF. A drafter fine-tuned with its own lm_head (output.weight present in + // the GGUF) must keep its own output tensor; overwriting it with the target's + // lm_head would destroy the fine-tuned readout and yield 0%% acceptance. + // The original Z-Lab drafter (no output.weight in GGUF) still shares the target's. + if (!dst->tok_embd) { + dst->tok_embd = src->tok_embd; + fprintf(stderr, "[dflash share] shared tok_embd from target (drafter had none)\n"); fflush(stderr); + } else { + fprintf(stderr, "[dflash share] kept drafter's own tok_embd\n"); fflush(stderr); + } + if (!dst->output) { + dst->output = src->output; + fprintf(stderr, "[dflash share] shared output from target (drafter had none)\n"); fflush(stderr); + } else { + fprintf(stderr, "[dflash share] kept drafter's own output (fine-tuned lm_head)\n"); fflush(stderr); + } } int32_t llama_model_dflash_block_size(const llama_model * model) { diff --git a/src/models/delta-net-base.cpp b/src/models/delta-net-base.cpp index ae5f4d1b3dcb..a627ec6c141c 100644 --- a/src/models/delta-net-base.cpp +++ b/src/models/delta-net-base.cpp @@ -1,6 +1,7 @@ #include "models.h" #include "llama-impl.h" +#include "llama-context.h" // llama_dflash_rs_writeback_slot_for_test (RS snapshot write-back decision) #include "llama-memory-recurrent.h" #include "ggml.h" @@ -629,19 +630,54 @@ ggml_tensor * llm_build_delta_net_base::build_recurrent_attn( cb(output, "attn_output", il); const size_t row_size = hparams.n_embd_s() * ggml_element_size(ssm_states_all); + // The gated_delta_net kernel only writes the last min(n_seq_tokens, K) snapshot + // slots of gdn_out; the leading slots are left untouched (caller-owned). For a + // full verify batch (n_seq_tokens >= K) every slot is populated, but for a + // decode / short batch (n_seq_tokens < K) the leading gdn_out slots are + // uninitialised. Writing them into ssm_states_all would clobber the older + // snapshot slots with garbage on every decode step. Under DFlash with stochastic + // sampling the adaptive controller drives spec depth to 0 once acceptance drops, + // so every step becomes a 1-token decode -> slots 1..K-1 stay garbage -> repeated + // draft-rejection rollbacks read stale recurrent state -> output degrades to + // gibberish (reproduced on Qwen3.6-35B-A3B, qwen35moe, n_rs_seq=4). + // + // Fix: write only the populated gdn_out slots (the newest n_pop states -> state + // slots 0..n_pop-1) and rotate the older snapshots down by n_pop slots + // (old slot s -> slot s + n_pop) so the recurrent-state history stays correct + // for rollback. Iterating k_i ascending reads each old slot before it is + // overwritten, so the in-place shift is safe (copies to the same persistent + // buffer are serialised by the scheduler). + // The write-back slot decision (destination cache_slot + whether to read a + // kernel-populated gdn_out slot or rotate an older state snapshot) is centralised + // in llama_dflash_rs_writeback_slot_for_test() so the production loop and the unit + // test (tests/test-dflash-plumbing.cpp) share one source of truth. See the helper + // doc in src/llama-context.h for the full invariant. for (int64_t k_i = 0; k_i < K; ++k_i) { - const uint32_t cache_slot = (uint32_t) (K - 1 - k_i); - ggml_tensor * src = ggml_view_4d(ctx0, gdn_out, - S_v, S_v, H_v, n_seqs, - ggml_row_size(gdn_out->type, S_v), - ggml_row_size(gdn_out->type, S_v * S_v), - ggml_row_size(gdn_out->type, S_v * S_v * H_v), - ggml_row_size(gdn_out->type, attn_score_elems + k_i * state_size_per_snap)); - + uint32_t cache_slot = 0; + llama_dflash_rs_slot_src slot_src{}; + if (!llama_dflash_rs_writeback_slot_for_test(K, n_seq_tokens, k_i, cache_slot, slot_src)) { + continue; // unreachable for k_i in [0,K) with valid K / n_seq_tokens + } ggml_tensor * dst = ggml_view_2d(ctx0, ssm_states_all, hparams.n_embd_s(), n_seqs, ssm_states_all->nb[1], ((size_t) cache_slot * mem_size + kv_head) * row_size); + ggml_tensor * src = nullptr; + if (slot_src.from_gdn) { + // gdn_out slot slot_src.gdn_slot holds the state after one of this batch's tokens. + src = ggml_view_4d(ctx0, gdn_out, + S_v, S_v, H_v, n_seqs, + ggml_row_size(gdn_out->type, S_v), + ggml_row_size(gdn_out->type, S_v * S_v), + ggml_row_size(gdn_out->type, S_v * S_v * H_v), + ggml_row_size(gdn_out->type, attn_score_elems + slot_src.gdn_slot * state_size_per_snap)); + } else { + // gdn_out slot k_i is uninitialised (n_seq_tokens < K): rotate the older + // snapshot (slot_src.old_slot) down so slot history is preserved for rollback. + src = ggml_view_2d(ctx0, ssm_states_all, + hparams.n_embd_s(), n_seqs, ssm_states_all->nb[1], + ((size_t) slot_src.old_slot * mem_size + kv_head) * row_size); + } ggml_build_forward_expand(gf, ggml_cpy(ctx0, src, dst)); } diff --git a/src/models/dflash_draft.cpp b/src/models/dflash_draft.cpp index fbf34cb485f4..732742ebf4ab 100644 --- a/src/models/dflash_draft.cpp +++ b/src/models/dflash_draft.cpp @@ -6,6 +6,7 @@ #include #include #include +#include #include #include @@ -308,6 +309,14 @@ static dflash_kv_cache_mode dflash_kv_cache_mode_env() { return mode; } +static bool dflash_rx_diag_enabled() { + static const bool enabled = [] { + const char * env = getenv("GGML_DFLASH_RX_DIAG"); + return env && std::atoi(env) != 0; + }(); + return enabled; +} + static bool dflash_input_shape_debug() { static const bool v = [] { const char * e = getenv("GGML_DFLASH_INPUT_DEBUG"); @@ -422,7 +431,8 @@ void llm_graph_input_dflash_update::set_input(const llama_ubatch * ubatch) { const float * src_data = nullptr; const void * src_gpu = nullptr; - auto fn_d2d = cross ? cross->fn_set_tensor_d2d : nullptr; + auto fn_d2d = cross ? cross->fn_set_tensor_d2d : nullptr; + auto fn_d2d_tensor = cross ? cross->fn_set_tensor_d2d_tensor : nullptr; int64_t src_real = 0; int64_t n_feat = 0; @@ -432,6 +442,7 @@ void llm_graph_input_dflash_update::set_input(const llama_ubatch * ubatch) { src_gpu = cross->dflash_kv_update_gpu; src_real = cross->dflash_kv_update_n_enc_real; fn_d2d = cross->dflash_kv_update_fn_set_tensor_d2d; + fn_d2d_tensor = cross->dflash_kv_update_fn_set_tensor_d2d_tensor; } else if (cross->v_embd_gpu) { n_feat = cross->n_embd; src_gpu = cross->v_embd_gpu; @@ -449,8 +460,9 @@ void llm_graph_input_dflash_update::set_input(const llama_ubatch * ubatch) { if (n_copy > 0 && n_feat == target_hidden->ne[0] && (src_gpu || src_data)) { const size_t actual_bytes = std::min(copy_bytes, tensor_bytes); - if (src_gpu && fn_d2d) { - fn_d2d(target_hidden->data, src_gpu, 0, actual_bytes); + if (src_gpu && (fn_d2d || fn_d2d_tensor)) { + if (fn_d2d_tensor) { fn_d2d_tensor(target_hidden, src_gpu, 0, actual_bytes); } + else { fn_d2d(target_hidden->data, src_gpu, 0, actual_bytes); } } else { ggml_backend_tensor_set(target_hidden, src_data, 0, actual_bytes); } @@ -504,6 +516,21 @@ void llm_graph_input_dflash::set_input(const llama_ubatch * ubatch) { const int64_t n_copy = std::min(src_real, ctx_len); const int64_t win_off = (src_real > ctx_len) ? (src_real - ctx_len) : 0; + if (dflash_rx_diag_enabled()) { + fprintf(stderr, + "DFLASH_RX set_input single: n_seqs=%d src_data=%s src_gpu=%s src_n_real=%lld src_real=%lld ctx_len=%lld n_copy=%lld win_off=%lld cross_n_embd=%lld use_kv_cache=%d\n", + n_seqs, + src_data ? "nonnull" : "null", + src_gpu ? "nonnull" : "null", + (long long) src_n_real, + (long long) src_real, + (long long) ctx_len, + (long long) n_copy, + (long long) win_off, + (long long) (cross ? cross->n_embd : 0), + use_kv_cache ? 1 : 0); + } + const bool skip_target_hidden_upload = use_kv_cache && dflash_kv_cache_mode_env() == DFLASH_KV_CACHE_BOTH; @@ -526,10 +553,14 @@ void llm_graph_input_dflash::set_input(const llama_ubatch * ubatch) { const size_t copy_bytes = (size_t) n_feat * (size_t) n_copy * sizeof(float); const size_t actual_bytes = std::min(copy_bytes, tensor_bytes); - if (src_gpu && cross->fn_set_tensor_d2d) { + if (src_gpu && (cross->fn_set_tensor_d2d || cross->fn_set_tensor_d2d_tensor)) { // GPU D2D path const void * gpu_src = (const char *)src_gpu + (size_t)win_off * n_feat * sizeof(float); - cross->fn_set_tensor_d2d(target_hidden->data, gpu_src, 0, actual_bytes); + if (cross->fn_set_tensor_d2d_tensor) { + cross->fn_set_tensor_d2d_tensor(target_hidden, gpu_src, 0, actual_bytes); + } else { + cross->fn_set_tensor_d2d(target_hidden->data, gpu_src, 0, actual_bytes); + } } else { // CPU H2D path (fallback) const float * src = src_data + win_off * n_feat; @@ -545,6 +576,43 @@ void llm_graph_input_dflash::set_input(const llama_ubatch * ubatch) { const int64_t n_real = n_copy; + if (const char * dump_path = getenv("GGML_DFLASH_CROSS_DUMP"); dump_path && dump_path[0] != '\0') { + const int64_t n_feat_dump = cross ? cross->n_embd : 0; + if (src_data && n_copy > 0 && n_feat_dump > 0) { + static std::atomic dump_count{0}; + // Dump the first ALL-REAL cycle (n_copy == ctx_len, no padding) so the + // PyTorch replay matches the W=ctx_len all-real training regime. + const bool all_real = (n_copy >= ctx_len) && (ctx_len >= 512); + if (all_real && dump_count.fetch_add(1) == 0) { + FILE * f = fopen(dump_path, "w"); + if (f) { + const int token0 = (ubatch && ubatch->token && ubatch->n_tokens > 0) ? (int) ubatch->token[0] : -1; + const int pos0 = (ubatch && ubatch->pos && ubatch->n_tokens > 0) ? (int) ubatch->pos[0] : -1; + fprintf(f, "# token0=%d pos0=%d n_copy=%lld n_feat=%lld ctx_len=%lld n_block=%lld win_off=%lld src_n_real=%lld\n", + token0, pos0, (long long) n_copy, (long long) n_feat_dump, + (long long) ctx_len, (long long) n_block, (long long) win_off, + (long long) src_n_real); + for (int64_t r = 0; r < n_copy; ++r) { + const float * row = src_data + (win_off + r) * n_feat_dump; + fprintf(f, "%lld", (long long) r); + for (int64_t j = 0; j < n_feat_dump; ++j) { + fprintf(f, " %.9g", (double) row[j]); + } + fprintf(f, "\n"); + } + fclose(f); + } + } + } + } + + // RX diagnostic: log n_real at set_input time (before mask fill) + if (getenv("GGML_DFLASH_RX_DIAG")) { + fprintf(stderr, "[dflash-rx] set_input: n_real=%lld n_copy=%lld ctx_len=%lld n_block=%lld\n", + (long long)n_real, (long long)n_copy, (long long)ctx_len, (long long)n_block); + fflush(stderr); + } + const bool have_pos = (ubatch != nullptr) && (ubatch->pos != nullptr) && ((int64_t) ubatch->n_tokens >= n_block); const int32_t q_pos_base = have_pos && n_block > 0 @@ -569,6 +637,13 @@ void llm_graph_input_dflash::set_input(const llama_ubatch * ubatch) { } if (kq_mask && kq_mask->buffer) { + // RX diagnostic: log n_real at mask fill time + if (getenv("GGML_DFLASH_RX_DIAG")) { + fprintf(stderr, "[dflash-rx] mask fill: n_real=%lld ctx_len=%lld n_block=%lld (cross visible=%lld/%lld)\n", + (long long)n_real, (long long)ctx_len, (long long)n_block, + (long long)std::min(n_real, (int64_t)ctx_len), (long long)ctx_len); + fflush(stderr); + } dflash_fill_kq_mask(kq_mask, [&](auto * data) { const int64_t n_kv = ctx_len + n_block; for (int64_t q = 0; q < n_block; ++q) { @@ -651,6 +726,23 @@ void llm_graph_input_dflash::set_input(const llama_ubatch * ubatch) { slot_win_off[s] = (nr > per_slot_ctx) ? (nr - per_slot_ctx) : 0; } + if (dflash_rx_diag_enabled()) { + fprintf(stderr, + "DFLASH_RX set_input multi: n_seqs=%d ctx_len=%lld per_slot_ctx=%d n_feat=%zu", + n_seqs, (long long) ctx_len, per_slot_ctx, n_feat); + for (int s = 0; s < n_seqs && s < LLAMA_DFLASH_MAX_SLOTS; s++) { + fprintf(stderr, + " slot%d{gpu=%d,cpu=%d,n_real=%lld,n_copy=%lld,win_off=%lld}", + s, + slot_info[s].gpu ? 1 : 0, + slot_info[s].data ? 1 : 0, + (long long) slot_info[s].n_real, + (long long) slot_n_copy[s], + (long long) slot_win_off[s]); + } + fprintf(stderr, "\n"); + } + if (target_hidden && target_hidden->buffer && target_hidden->data && n_feat > 0) { ggml_backend_tensor_memset(target_hidden, 0, 0, ggml_nbytes(target_hidden)); @@ -665,16 +757,24 @@ void llm_graph_input_dflash::set_input(const llama_ubatch * ubatch) { const size_t dst_offset = (size_t) s * (size_t) per_slot_ctx * n_feat * sizeof(float); - if (slot_info[s].gpu && cross && cross->fn_set_tensor_d2d) { + if (slot_info[s].gpu && cross && (cross->fn_set_tensor_d2d || cross->fn_set_tensor_d2d_tensor)) { const void * gpu_src = (const char *) slot_info[s].gpu + (size_t) slot_win_off[s] * n_feat * sizeof(float); - cross->fn_set_tensor_d2d( - target_hidden->data, - gpu_src, - dst_offset, - copy_bytes); + if (cross->fn_set_tensor_d2d_tensor) { + cross->fn_set_tensor_d2d_tensor( + target_hidden, + gpu_src, + dst_offset, + copy_bytes); + } else { + cross->fn_set_tensor_d2d( + target_hidden->data, + gpu_src, + dst_offset, + copy_bytes); + } } else if (slot_info[s].data) { const float * src = slot_info[s].data + slot_win_off[s] * n_feat; @@ -880,8 +980,13 @@ llm_build_dflash_draft::llm_build_dflash_draft( // Full-attention draft layers consume accepted-prefix K/V from the normal // drafter KV cache. Sliding layers still read the current target-hidden // window directly because their context is already bounded by SWA. + // ONLY build this KV-cache input when the cache will actually be populated + // with TARGET context K/V (GPU cross ring available). Otherwise (e.g. Vulkan, + // CPU hidden capture) the cache would be empty and the if-branch below would + // self-attend with NO target context -> 0%% acceptance for all-full-attention + // cross-attention drafters. Fall back to the else branch (fresh Kcur_ctx). llm_graph_input_attn_kv * inp_attn_kv_full = nullptr; - if (mctx) { + if (mctx && cparams.dflash_target_kv_available) { const bool rebind_base_from_iswa = hparams.swa_type != LLAMA_SWA_TYPE_NONE; const llama_kv_cache_context * base_ctx = rebind_base_from_iswa @@ -892,6 +997,7 @@ llm_build_dflash_draft::llm_build_dflash_draft( // --- Fusion layer: project concatenated target hidden states --- ggml_tensor * fused_target = build_lora_mm(model.dflash_fc, target_hidden); + fused_target = build_norm(fused_target, model.dflash_hidden_norm, nullptr, LLM_NORM_RMS, -1); cb(fused_target, "fused_target", -1); @@ -1132,7 +1238,8 @@ llm_build_dflash_kv_update::llm_build_dflash_kv_update( for (int il = 0; il < n_layer; ++il) { ggml_tensor * Kcur_ctx = build_lora_mm(model.layers[il].wk, fused_target); Kcur_ctx = ggml_reshape_3d(ctx0, Kcur_ctx, n_embd_head, n_head_kv, n_tokens); - Kcur_ctx = build_norm(Kcur_ctx, model.layers[il].attn_k_norm, nullptr, LLM_NORM_RMS, il); + // Store raw K (k_norm applied after concat in the attention block) + cb(Kcur_ctx, "dflash_kv_update_k_raw", il); ggml_tensor * Vcur_ctx = build_lora_mm(model.layers[il].wv, fused_target); Vcur_ctx = ggml_reshape_3d(ctx0, Vcur_ctx, n_embd_head, n_head_kv, n_tokens); diff --git a/src/models/qwen35.cpp b/src/models/qwen35.cpp index bda246b7fc38..6ef548ce7b80 100644 --- a/src/models/qwen35.cpp +++ b/src/models/qwen35.cpp @@ -211,8 +211,21 @@ llama_model_qwen35::graph::graph(const llama_model & model, const llm_graph_para cur = build_cvec(cur, il); cb(cur, "l_out", il); +#ifdef GGML_DFLASH_DEBUG + // [DIAG] GGML_DFLASH_CAPTURE_SKIP: 0=off, 1=skip recurrent layers, 2=skip all layers + static int dflash_capture_skip_mode = ([]() { + const char * e = std::getenv("GGML_DFLASH_CAPTURE_SKIP"); + return e ? std::atoi(e) : 0; }()); + const bool dflash_capture_skip_this = + (dflash_capture_skip_mode == 2) || + (dflash_capture_skip_mode == 1 && hparams.is_recurrent(il)); if (cparams.hidden_gpu_n_seqs > 0 && + !dflash_capture_skip_this && cur->ne[1] == dflash_capture_n_tokens * dflash_capture_n_seqs) { +#else + if (cparams.hidden_gpu_n_seqs > 0 && + cur->ne[1] == dflash_capture_n_tokens * dflash_capture_n_seqs) { +#endif for (int s = 0; s < (int) dflash_capture_n_seqs && s < cparams.hidden_gpu_n_seqs; ++s) { auto * hgpu = cparams.hidden_gpu_seqs[s]; if (!hgpu) { @@ -234,11 +247,12 @@ llama_model_qwen35::graph::graph(const llama_model & model, const llm_graph_para cur->ne[0], dflash_capture_n_tokens, cur->nb[1], (size_t) s * (size_t) dflash_capture_n_tokens * cur->nb[1]); - ggml_tensor * h_cont = ggml_cont(ctx0, h_slice); ggml_tensor * h_dst = ggml_view_2d(ctx0, hgpu->layers[hi], hgpu->layers[hi]->ne[0], (int64_t) dflash_capture_n_tokens, hgpu->layers[hi]->nb[1], 0); - ggml_build_forward_expand(gf, ggml_cpy(ctx0, h_cont, h_dst)); + // [FIX-TRY] drop ggml_cont: cpy the strided view directly so the scheduler + // does not allocate a cont temp that can alias the recurrent RS snapshot buffer. + ggml_build_forward_expand(gf, ggml_cpy(ctx0, h_slice, h_dst)); } } diff --git a/src/models/qwen3next.cpp b/src/models/qwen3next.cpp index 12b8d904ba50..efb115730002 100644 --- a/src/models/qwen3next.cpp +++ b/src/models/qwen3next.cpp @@ -1,4 +1,5 @@ #include "models.h" +#include "llama-context.h" #include "llama-memory-recurrent.h" void llama_model_qwen3next::load_arch_hparams(llama_model_loader & ml) { @@ -120,6 +121,13 @@ llama_model_qwen3next::graph::graph(const llama_model & model, const llm_graph_p ggml_tensor * inp_pos = build_inp_pos(); ggml_tensor * inp_out_ids = build_inp_out_ids(); + // DFlash GPU hidden-capture window sizing (mirrors qwen35). + const int64_t n_seq_tokens = ubatch.n_seq_tokens; + const int64_t dflash_capture_n_seqs = + ubatch.n_seqs_unq > 1 ? (int64_t) ubatch.n_seqs_unq : 1; + const int64_t dflash_capture_n_tokens = + ubatch.n_seqs_unq > 1 ? n_seq_tokens : (int64_t) ubatch.n_tokens; + for (int il = 0; il < n_layer; ++il) { ggml_tensor * inpSA = inpL; @@ -164,6 +172,77 @@ llama_model_qwen3next::graph::graph(const llama_model & model, const llm_graph_p cur = build_cvec(cur, il); cb(cur, "l_out", il); + // DFlash GPU hidden capture (ported from qwen35): copy this layer's output into the + // per-slot device-resident hidden_gpu tensors so the drafter cross-ring can D2D-copy + // them into the ring (no CPU readback). Gated by hidden_gpu_n_seqs; only layers listed + // in hgpu->layer_ids are captured. Runs on the target backend (Vulkan/CUDA). + if (cparams.hidden_gpu_n_seqs > 0 && + cur->ne[1] == dflash_capture_n_tokens * dflash_capture_n_seqs) { + for (int s = 0; s < (int) dflash_capture_n_seqs && s < cparams.hidden_gpu_n_seqs; ++s) { + auto * hgpu = cparams.hidden_gpu_seqs[s]; + if (!hgpu) { + continue; + } + int hi = -1; + for (int i = 0; i < (int) hgpu->layer_ids.size(); ++i) { + if (hgpu->layer_ids[i] == il) { + hi = i; + break; + } + } + if (hi < 0 || dflash_capture_n_tokens > hgpu->max_tokens) { + continue; + } + ggml_tensor * h_slice = ggml_view_2d(ctx0, cur, + cur->ne[0], dflash_capture_n_tokens, + cur->nb[1], + (size_t) s * (size_t) dflash_capture_n_tokens * cur->nb[1]); + ggml_tensor * h_cont = ggml_cont(ctx0, h_slice); + ggml_tensor * h_dst = ggml_view_2d(ctx0, hgpu->layers[hi], + hgpu->layers[hi]->ne[0], (int64_t) dflash_capture_n_tokens, + hgpu->layers[hi]->nb[1], 0); + ggml_build_forward_expand(gf, ggml_cpy(ctx0, h_cont, h_dst)); + } + } + if (cparams.prefill_gpu_n_seqs > 0 && + cparams.dflash_prefill_capture_active && + cparams.dflash_prefill_n_tokens > 0 && + cur->ne[1] == dflash_capture_n_tokens * dflash_capture_n_seqs) { + for (int s = 0; s < (int) dflash_capture_n_seqs && s < cparams.prefill_gpu_n_seqs; ++s) { + const int n_copy = cparams.dflash_prefill_n_tokens_seqs[s]; + const int src_off = cparams.dflash_prefill_src_offsets[s]; + const int dst_off = cparams.dflash_prefill_dst_offsets[s]; + if (n_copy <= 0 || src_off < 0 || src_off + n_copy > dflash_capture_n_tokens) { + continue; + } + auto * pgpu = cparams.prefill_gpu_seqs[s]; + if (!pgpu) { + continue; + } + int hi = -1; + for (int i = 0; i < (int) pgpu->layer_ids.size(); ++i) { + if (pgpu->layer_ids[i] == il) { + hi = i; + break; + } + } + if (hi < 0 || dst_off < 0 || dst_off + n_copy > pgpu->max_tokens) { + continue; + } + ggml_tensor * h_slice = ggml_view_2d(ctx0, cur, + cur->ne[0], (int64_t) n_copy, + cur->nb[1], + (size_t) s * (size_t) dflash_capture_n_tokens * cur->nb[1] + + (size_t) src_off * cur->nb[1]); + ggml_tensor * h_cont = ggml_cont(ctx0, h_slice); + ggml_tensor * h_dst = ggml_view_2d(ctx0, pgpu->layers[hi], + pgpu->layers[hi]->ne[0], (int64_t) n_copy, + pgpu->layers[hi]->nb[1], + (size_t) dst_off * pgpu->layers[hi]->nb[1]); + ggml_build_forward_expand(gf, ggml_cpy(ctx0, h_cont, h_dst)); + } + } + // Input for next layer inpL = cur; } @@ -231,7 +310,7 @@ ggml_tensor * llama_model_qwen3next::graph::build_layer_attn( ggml_tensor * gate = ggml_view_4d(ctx0, Qcur_full, n_embd_head, n_head, n_tokens, 1, Qcur_full->nb[1], Qcur_full->nb[2], Qcur_full->nb[3], n_embd_head * ggml_element_size(Qcur_full)); - cb(gate, "gate", il); + cb(gate, "attn_gate", il); ggml_tensor * Kcur = build_lora_mm(model.layers[il].wk, cur); cb(Kcur, "Kcur", il); diff --git a/tests/test-dflash-plumbing.cpp b/tests/test-dflash-plumbing.cpp index 63d609452f9f..19eacecf5d48 100644 --- a/tests/test-dflash-plumbing.cpp +++ b/tests/test-dflash-plumbing.cpp @@ -2,11 +2,16 @@ #include "dflash-profile.h" #include "speculative.h" +#include "ggml.h" +#include "ggml-backend.h" + #include #include +#include #include #include #include +#include static std::string read_file(const std::string & path) { std::ifstream file(path); @@ -60,6 +65,312 @@ static std::string slice_between(const std::string & text, const std::string & b return text.substr(b, e - b); } +static bool test_vk_dflash_cross_ring_roundtrip() { + ggml_backend_reg_t vk_reg = ggml_backend_reg_by_name("Vulkan"); + if (!vk_reg) { return true; } // no Vulkan backend in this build — skip + using alloc_fn_t = void *(*)(int,int,int); + using free_fn_t = void (*)(void*); + using write_fn_t = void (*)(void*,int,int,const float*,int,int); + using sync_fn_t = void (*)(void*); + using snap_fn_t = bool (*)(void*,int,int,int,float*,int,int,int); + using inter_fn_t = const float* (*)(void*,int,int,int); + using setT_fn_t = void (*)(ggml_tensor*, const void*, size_t, size_t); + auto falloc = (alloc_fn_t)ggml_backend_reg_get_proc_address(vk_reg, "dflash_cross_ring_gpu_alloc"); + auto ffree = (free_fn_t) ggml_backend_reg_get_proc_address(vk_reg, "dflash_cross_ring_gpu_free"); + auto fwrite = (write_fn_t)ggml_backend_reg_get_proc_address(vk_reg, "dflash_cross_ring_gpu_write"); + auto fsync = (sync_fn_t) ggml_backend_reg_get_proc_address(vk_reg, "dflash_cross_ring_gpu_synchronize"); + auto fsnap = (snap_fn_t) ggml_backend_reg_get_proc_address(vk_reg, "dflash_cross_ring_gpu_snapshot"); + auto finter = (inter_fn_t)ggml_backend_reg_get_proc_address(vk_reg, "dflash_cross_ring_gpu_interleave"); + auto fsetT = (setT_fn_t) ggml_backend_reg_get_proc_address(vk_reg, "dflash_cross_ring_gpu_set_tensor_tensor"); + if (!falloc || !ffree || !fwrite || !fsync || !fsnap || !finter || !fsetT) return false; + + void * h = falloc(/*n_layers*/2, /*n_embd*/3, /*ring_size*/4); + if (!h) { return true; } // Vulkan present but cross-ring alloc failed (e.g. no buffer_device_address) — skip + + float l0[3] = {1,2,3}, l1[3] = {4,5,6}; + fwrite(h, 0, 1, l0, 1, 3); + fwrite(h, 1, 1, l1, 1, 3); + fsync(h); + + float snap[12] = {0}; + bool ok = fsnap(h, /*write_pos*/2, /*filled*/2, /*ctx_window*/2, snap, /*n_tokens*/2, /*n_layers*/2, /*n_embd*/3); + if (ok) { + const float exp_snap[12] = {0,0,0,1,2,3, 0,0,0,4,5,6}; + for (int i = 0; i < 12; ++i) { if (snap[i] != exp_snap[i]) { ok = false; break; } } + } + + const float * staging = finter(h, /*write_pos*/2, /*filled*/2, /*ctx_window*/2); + ok = ok && staging != nullptr; + + // Drafter-side Vulkan tensor to receive the set_tensor_tensor D2D copy. + ggml_backend_dev_t dev = ggml_backend_reg_dev_get(vk_reg, 0); + ggml_backend_t backend = ggml_backend_dev_init(dev, nullptr); + ggml_backend_buffer_type_t buft = ggml_backend_dev_buffer_type(dev); + size_t mem_size = ggml_tensor_overhead()*4; + void * mem = malloc(mem_size); + struct ggml_init_params ip = {}; + ip.mem_buffer = mem; + ip.mem_size = mem_size; + ip.no_alloc = true; + ggml_context * gctx = ggml_init(ip); + ggml_tensor * t = ggml_new_tensor_1d(gctx, GGML_TYPE_F32, 12); + ggml_backend_buffer_t buf = ggml_backend_alloc_ctx_tensors(gctx, backend); + ok = ok && buf != nullptr && t != nullptr && t->buffer != nullptr; + if (ok) { + fsetT(t, staging, 0, 12 * sizeof(float)); + float got[12] = {0}; + ggml_backend_tensor_get(t, got, 0, 12 * sizeof(float)); + const float exp[12] = {0,0,0, 0,0,0, 1,2,3, 4,5,6}; + for (int i = 0; i < 12; ++i) { if (got[i] != exp[i]) { ok = false; break; } } + } + if (buf) ggml_backend_buffer_free(buf); + if (gctx) ggml_free(gctx); + free(mem); + ggml_backend_free(backend); + ffree(h); + return ok; +} + +// ============================================================================ +// Regression guards for PR #79 DFlash changes (added beyond the original +// cross-ring plumbing test). These cover the behaviour changes that are +// unit-testable without a full model: +// - 122B fix: llm_arch_supports_rs_rollback gate keeps QWEN35MOE on. +// - DFlash enablement: common_params_speculative::need_n_rs_seq() returns +// draft.n_max for DFlash. +// - 35B-A3B fix: RS snapshot write-back slot mapping (rotation for +// n_seq_tokens < K) + the gated_delta_net kernel snapshot +// contract (only the last min(n_seq_tokens,K) slots are +// written). +// - Coverage gap: the existing pure _for_test helpers in llama-context.h +// had no direct unit tests. +// Integration-only paths (Coder-Next KV fallback, hidden-state dump APIs, the +// rollback int return) are guarded by the rebenchmark, not here. +// ============================================================================ + +static bool test_llm_arch_supports_rs_rollback() { + bool ok = true; + ok &= expect( llm_arch_supports_rs_rollback(LLM_ARCH_QWEN35), "QWEN35 must support RS rollback"); + ok &= expect( llm_arch_supports_rs_rollback(LLM_ARCH_QWEN35MOE), "QWEN35MOE must support RS rollback (122B fix: keep n_rs_seq)"); + ok &= expect(!llm_arch_supports_rs_rollback(LLM_ARCH_QWEN3NEXT), "QWEN3NEXT must NOT support RS rollback (Coder-Next uses checkpoint-restore)"); + ok &= expect(!llm_arch_supports_rs_rollback(LLM_ARCH_KIMI_LINEAR),"Kimi-linear must NOT support RS rollback"); + ok &= expect(!llm_arch_supports_rs_rollback(LLM_ARCH_UNKNOWN), "unknown arch must NOT support RS rollback"); + return ok; +} + +static bool test_need_n_rs_seq() { + bool ok = true; + auto make = [](std::vector types, int32_t n_max) -> uint32_t { + common_params_speculative s{}; + s.types = std::move(types); + s.draft.n_max = n_max; + return s.need_n_rs_seq(); + }; + ok &= expect(make({COMMON_SPECULATIVE_TYPE_DFLASH}, 4) == 4u, "DFlash -> need_n_rs_seq == draft.n_max"); + ok &= expect(make({COMMON_SPECULATIVE_TYPE_DRAFT_MTP}, 3) == 3u, "MTP -> need_n_rs_seq == draft.n_max"); + ok &= expect(make({}, 4) == 0u, "no speculative types -> need_n_rs_seq == 0"); + ok &= expect(make({COMMON_SPECULATIVE_TYPE_DFLASH, COMMON_SPECULATIVE_TYPE_DRAFT_MTP}, 5) == 5u, + "DFlash+MTP -> need_n_rs_seq == draft.n_max"); + ok &= expect(make({COMMON_SPECULATIVE_TYPE_DFLASH}, 0) == 0u, "DFlash with draft.n_max=0 -> 0"); + return ok; +} + +// Pure-logic guard for the 35B-A3B fix: the RS snapshot write-back slot mapping +// produced by llama_dflash_rs_writeback_slot_for_test() (the helper build_recurrent_attn +// now calls). Verifies the rotation invariant for decode (n_seq_tokens < K) and the +// full-batch case, plus the explicit 35B-A3B decode scenario (n_rs_seq=4 -> K=5). +static bool test_rs_writeback_slot_mapping() { + bool ok = true; + const int64_t cases_K[] = {2, 5}; + const int64_t cases_nst[] = {0, 1, 2, 3, 5, 8}; + for (int64_t K : cases_K) { + for (int64_t nst : cases_nst) { + const int64_t n_pop = (nst < K) ? nst : K; + int64_t n_from_gdn = 0; + for (int64_t k_i = 0; k_i < K; ++k_i) { + uint32_t cache_slot = 0xDEADu; + llama_dflash_rs_slot_src src{}; + const bool valid = llama_dflash_rs_writeback_slot_for_test(K, nst, k_i, cache_slot, src); + if (n_pop == 0) { + ok &= expect(!valid, "n_seq_tokens=0 must yield invalid (n_pop==0)"); + continue; + } + ok &= expect(valid, "k_i in [0,K) with n_pop>0 must be valid"); + ok &= expect(cache_slot == (uint32_t)(K - 1 - k_i), "cache_slot == K-1-k_i"); + if (src.from_gdn) { + ++n_from_gdn; + ok &= expect(src.gdn_slot == k_i, "from_gdn: gdn_slot == k_i"); + ok &= expect(k_i >= K - n_pop, "from_gdn only for k_i >= K-n_pop"); + } else { + ok &= expect(k_i < K - n_pop, "rotate only for k_i < K-n_pop"); + ok &= expect(src.old_slot == cache_slot - (uint32_t)n_pop, "rotate: old_slot == cache_slot - n_pop"); + ok &= expect(src.old_slot < cache_slot, "rotate: old_slot < cache_slot (read-before-overwrite)"); + } + } + ok &= expect(n_from_gdn == n_pop, "exactly n_pop slots read from gdn_out"); + } + } + // 35B-A3B decode scenario: n_rs_seq=4 -> K=5, n_seq_tokens=1. + // After the step: state slot 0 <- newest gdn_out; state slot s <- old slot s-1 (s=1..4). + { + const int64_t K = 5, nst = 1; + for (int64_t k_i = 0; k_i < K; ++k_i) { + uint32_t cache_slot = 0; + llama_dflash_rs_slot_src src{}; + llama_dflash_rs_writeback_slot_for_test(K, nst, k_i, cache_slot, src); + if (cache_slot == 0) { + ok &= expect(src.from_gdn && src.gdn_slot == K - 1, "decode: slot 0 <- newest gdn_out slot K-1"); + } else { + ok &= expect(!src.from_gdn && src.old_slot == cache_slot - 1u, "decode: slot s <- old slot s-1 (rotation)"); + } + } + } + return ok; +} + +// Kernel-contract guard for the 35B-A3B fix: ggml_gated_delta_net, for a decode +// batch (n_seq_tokens=1 < K=4), must write ONLY the last min(n_seq_tokens,K)=1 +// snapshot slot of its output; the leading K-1 state slots are left untouched +// (caller-owned). The fix's rotation in build_recurrent_attn relies on exactly +// this contract. Runs a single-op graph on the CPU backend (always present). +static bool test_gated_delta_net_snapshot_contract() { + ggml_backend_reg_t cpu_reg = ggml_backend_reg_by_name("CPU"); + if (!cpu_reg) { return true; } // CPU backend missing (never expected) - skip + ggml_backend_dev_t dev = ggml_backend_reg_dev_get(cpu_reg, 0); + ggml_backend_t backend = ggml_backend_dev_init(dev, nullptr); + if (!backend) { return true; } + + bool ok = true; + const int64_t S_v = 2, H = 1, n_tokens = 1, n_seqs = 1, K = 4; + const int64_t D = S_v * S_v * H; // = 4 floats per snapshot + const int64_t attn_elems = S_v * H * n_tokens * n_seqs; // = 2 + const int64_t state_per_snap = S_v * S_v * H * n_seqs; // = 4 + const float sentinel = -12345.0f; + + size_t tensors_mem = ggml_tensor_overhead() * 16 + ggml_graph_overhead() + (1 << 16); + void * mem = malloc(tensors_mem); + struct ggml_init_params ip = {}; + ip.mem_buffer = mem; + ip.mem_size = tensors_mem; + ip.no_alloc = true; + ggml_context * gctx = ggml_init(ip); + ok &= expect(gctx != nullptr, "gated_delta_net test: ggml_init"); + if (!ok) { free(mem); ggml_backend_free(backend); return ok; } + + auto mk4 = [&](int64_t n0, int64_t n1, int64_t n2, int64_t n3) { + return ggml_new_tensor_4d(gctx, GGML_TYPE_F32, n0, n1, n2, n3); + }; + ggml_tensor * q = mk4(S_v, H, n_tokens, n_seqs); + ggml_tensor * k = mk4(S_v, H, n_tokens, n_seqs); + ggml_tensor * v = mk4(S_v, H, n_tokens, n_seqs); + ggml_tensor * g = mk4(1, H, n_tokens, n_seqs); // scalar gate + ggml_tensor * beta = mk4(1, H, n_tokens, n_seqs); + // 3D snapshot state (D, K, n_seqs) - mirrors build_recurrent_attn's s_3d_pad. + ggml_tensor * state = ggml_new_tensor_3d(gctx, GGML_TYPE_F32, D, K, n_seqs); + ggml_tensor * result = ggml_gated_delta_net(gctx, q, k, v, g, beta, state); + ok &= expect(q && k && v && g && beta && state && result, "gated_delta_net test: tensor creation"); + if (!ok) { ggml_free(gctx); free(mem); ggml_backend_free(backend); return ok; } + + ggml_backend_buffer_t buf = ggml_backend_alloc_ctx_tensors(gctx, backend); + ok &= expect(buf != nullptr, "gated_delta_net test: alloc_ctx_tensors"); + if (!ok) { ggml_free(gctx); free(mem); ggml_backend_free(backend); return ok; } + + // Inputs: simple finite values. State slot 0 = 2x2 identity, pad slots = 0. + float qv[2] = {1.0f, 1.0f}; + float gv[1] = {0.0f}; + float betav[1] = {0.0f}; + float statev[(size_t)(D * K * n_seqs)]; + for (int64_t i = 0; i < D * K * n_seqs; ++i) statev[i] = 0.0f; + statev[0] = 1.0f; statev[3] = 1.0f; // identity in slot 0 + ggml_backend_tensor_set(q, qv, 0, sizeof(qv)); + ggml_backend_tensor_set(k, qv, 0, sizeof(qv)); + ggml_backend_tensor_set(v, qv, 0, sizeof(qv)); + ggml_backend_tensor_set(g, gv, 0, sizeof(gv)); + ggml_backend_tensor_set(beta, betav, 0, sizeof(betav)); + ggml_backend_tensor_set(state, statev, 0, sizeof(statev)); + + // Pre-fill the whole result with a sentinel so untouched slots keep it. + const int64_t result_elems = attn_elems + K * state_per_snap; // = 2 + 16 = 18 + std::vector sentinel_buf((size_t)result_elems, sentinel); + ggml_backend_tensor_set(result, sentinel_buf.data(), 0, sizeof(float) * (size_t)result_elems); + + ggml_cgraph * gf = ggml_new_graph(gctx); + ggml_build_forward_expand(gf, result); + const ggml_status status = ggml_backend_graph_compute(backend, gf); + ggml_backend_synchronize(backend); + ok &= expect(status == GGML_STATUS_SUCCESS, "gated_delta_net test: graph compute succeeded"); + if (!ok) { ggml_backend_buffer_free(buf); ggml_free(gctx); free(mem); ggml_backend_free(backend); return ok; } + + std::vector out((size_t)result_elems, 0.0f); + ggml_backend_tensor_get(result, out.data(), 0, sizeof(float) * (size_t)result_elems); + + // n_tokens=1, K=4 -> n_pop=1, shift = n_tokens-K = -3, target_slot = 0-(-3) = 3 = K-1. + // Only slot 3 is written; slots 0,1,2 must retain the sentinel. + auto slot_start = [&](int64_t s) { return (size_t)(attn_elems + s * state_per_snap); }; + for (int64_t s = 0; s < K - 1; ++s) { // untouched slots 0..K-2 + const size_t off = slot_start(s); + bool untouched = true; + for (int64_t j = 0; j < state_per_snap; ++j) { + if (out[off + (size_t)j] != sentinel) { untouched = false; break; } + } + ok &= expect(untouched, "gated_delta_net decode: leading snapshot slot untouched (kernel writes only last n_pop)"); + } + { // written slot K-1 must differ from the sentinel + const size_t off = slot_start(K - 1); + bool written = false; + for (int64_t j = 0; j < state_per_snap; ++j) { + if (out[off + (size_t)j] != sentinel) { written = true; break; } + } + ok &= expect(written, "gated_delta_net decode: last snapshot slot is written by the kernel"); + } + + ggml_backend_buffer_free(buf); + ggml_free(gctx); + free(mem); + ggml_backend_free(backend); + return ok; +} + +// Coverage gap: the existing pure _for_test helpers in llama-context.h are called by +// production code but had no direct unit tests. +static bool test_existing_for_test_helpers() { + bool ok = true; + // llama_dflash_capture_tokens_per_seq(n_tokens, n_seq_tokens, n_seqs_unq) + ok &= expect(llama_dflash_capture_tokens_per_seq(10, 1, 1) == 10, "capture: single unique seq -> all ubatch tokens"); + ok &= expect(llama_dflash_capture_tokens_per_seq(10, 4, 3) == 4, "capture: multi unique seq -> n_seq_tokens"); + ok &= expect(llama_dflash_capture_tokens_per_seq(7, 7, 1) == 7, "capture: single seq, n_seq_tokens==n_tokens"); + // llama_dflash_replay_gdn_supported_s_for_test + ok &= expect( llama_dflash_replay_gdn_supported_s_for_test(16), "gdn replay s=16"); + ok &= expect( llama_dflash_replay_gdn_supported_s_for_test(32), "gdn replay s=32"); + ok &= expect( llama_dflash_replay_gdn_supported_s_for_test(64), "gdn replay s=64"); + ok &= expect( llama_dflash_replay_gdn_supported_s_for_test(128), "gdn replay s=128"); + ok &= expect(!llama_dflash_replay_gdn_supported_s_for_test(8), "gdn replay s=8 not supported"); + ok &= expect(!llama_dflash_replay_gdn_supported_s_for_test(256), "gdn replay s=256 not supported"); + ok &= expect(!llama_dflash_replay_gdn_supported_s_for_test(0), "gdn replay s=0 not supported"); + // llama_dflash_replay_state_shape_valid_for_test(s, h_v, n_embd_s): valid iff s*s*h_v == n_embd_s + ok &= expect( llama_dflash_replay_state_shape_valid_for_test(16, 1, 16*16), "state shape: 16*16*1 == 256"); + ok &= expect( llama_dflash_replay_state_shape_valid_for_test(8, 2, 8*8*2), "state shape: 8*8*2 == 128"); + ok &= expect(!llama_dflash_replay_state_shape_valid_for_test(8, 2, 127), "state shape: 8*8*2 != 127"); + ok &= expect(!llama_dflash_replay_state_shape_valid_for_test(0, 1, 0), "state shape: s<=0 invalid"); + ok &= expect(!llama_dflash_replay_state_shape_valid_for_test(-1, 1, 0), "state shape: negative s invalid"); + // llama_dflash_view_span_in_bounds_for_test(total, offset, n) + ok &= expect( llama_dflash_view_span_in_bounds_for_test(100, 0, 100), "span: full"); + ok &= expect( llama_dflash_view_span_in_bounds_for_test(100, 30, 40), "span: interior"); + ok &= expect(!llama_dflash_view_span_in_bounds_for_test(100, 101, 0), "span: offset>total"); + ok &= expect(!llama_dflash_view_span_in_bounds_for_test(100, 90, 20), "span: overflow"); + ok &= expect( llama_dflash_view_span_in_bounds_for_test(100, 100, 0), "span: zero-length at end"); + // llama_dflash_suppress_callback_for_prefill_ubatch_for_test(active, staging, intersection) + ok &= expect( llama_dflash_suppress_callback_for_prefill_ubatch_for_test(true, true, false), "suppress: active+staging+no-intersection"); + ok &= expect(!llama_dflash_suppress_callback_for_prefill_ubatch_for_test(true, true, true), "no suppress: has intersection"); + ok &= expect(!llama_dflash_suppress_callback_for_prefill_ubatch_for_test(true, false, false), "no suppress: not staging"); + ok &= expect(!llama_dflash_suppress_callback_for_prefill_ubatch_for_test(false, true, false), "no suppress: no prefill plan"); + // llama_dflash_prefill_plan_needs_staging_for_test(planned, current): staging iff planned > LLAMA_DFLASH_MAX_VERIFY_TOKENS (25) + ok &= expect( llama_dflash_prefill_plan_needs_staging_for_test((int)LLAMA_DFLASH_MAX_VERIFY_TOKENS + 1, 0), "staging: planned > MAX_VERIFY"); + ok &= expect(!llama_dflash_prefill_plan_needs_staging_for_test((int)LLAMA_DFLASH_MAX_VERIFY_TOKENS, 0), "no staging: planned == MAX_VERIFY"); + ok &= expect(!llama_dflash_prefill_plan_needs_staging_for_test((int)LLAMA_DFLASH_MAX_VERIFY_TOKENS - 1, 0), "no staging: planned < MAX_VERIFY"); + return ok; +} + int main(int argc, char ** argv) { bool ok = true; @@ -68,6 +379,12 @@ int main(int argc, char ** argv) { ok &= expect(!llama_dflash_gpu_tape_supported_arch(LLM_ARCH_QWEN3NEXT), "Qwen3Next must stay on fallback"); ok &= expect(!llama_dflash_gpu_tape_supported_arch(LLM_ARCH_KIMI_LINEAR), "Kimi-linear must stay on fallback"); ok &= expect(!llama_dflash_gpu_tape_supported_arch(LLM_ARCH_UNKNOWN), "unknown arch must stay on fallback"); + ok &= expect(test_vk_dflash_cross_ring_roundtrip(), "Vulkan DFlash cross-ring alloc/write/snapshot/interleave/set_tensor_tensor round-trip"); + ok &= expect(test_llm_arch_supports_rs_rollback(), "llm_arch_supports_rs_rollback gate (122B fix: keep n_rs_seq for QWEN35MOE)"); + ok &= expect(test_need_n_rs_seq(), "common_params_speculative::need_n_rs_seq (DFlash RS-seq enablement)"); + ok &= expect(test_rs_writeback_slot_mapping(), "DFlash RS snapshot write-back slot mapping (35B-A3B fix: rotation for n_seq_tokens(1200, 42), SERVER_LOOP_REGION_VISIBLE); + const auto result = guard.check(SERVER_LOOP_REGION_VISIBLE); + assert_triggered(result, "periodic_tail"); + assert(result.period == 1); + } + + // Visible region: independent from reasoning region + { + server_loop_guard guard(test_params()); + accept_many(guard, std::vector(1200, 42), SERVER_LOOP_REGION_REASONING); + const auto r1 = guard.check(SERVER_LOOP_REGION_REASONING); + assert_triggered(r1, "periodic_tail"); + const auto r2 = guard.check(SERVER_LOOP_REGION_VISIBLE); + assert_not_triggered(r2); + } + + // Visible region: no loop with varied output + { + server_loop_guard guard(test_params()); + std::vector tokens; + tokens.reserve(1200); + std::mt19937 rng(5678); + std::uniform_int_distribution dist(0, 500); + for (int i = 0; i < 1200; ++i) { + tokens.push_back((llama_token) dist(rng)); + } + accept_many(guard, tokens, SERVER_LOOP_REGION_VISIBLE); + assert_not_triggered(guard.check(SERVER_LOOP_REGION_VISIBLE)); + } + return 0; } diff --git a/tools/server/server-context.cpp b/tools/server/server-context.cpp index df782c76729b..f2b127d16e8b 100644 --- a/tools/server/server-context.cpp +++ b/tools/server/server-context.cpp @@ -60,6 +60,42 @@ static bool dflash_server_crash_trace_enabled() { return enabled; } +static bool dflash_rx_diag_enabled() { + /* Per-cycle ring/cross/append diagnostics for DFlash investigation. + * Logs committed_len, ring_filled, cross_len, n_enc_real, n_real, + * n_accepted_draft, n_hidden_keep, and append success on every cycle. + * Zero cost when disabled (static const guard). */ + static const bool enabled = [] { + const char * env = std::getenv("GGML_DFLASH_RX_DIAG"); + return env && std::atoi(env) != 0; + }(); + return enabled; +} + +static bool dflash_token_trace_enabled() { + static const bool enabled = [] { + const char * env = std::getenv("GGML_DFLASH_TOKEN_TRACE"); + return env && std::atoi(env) != 0; + }(); + return enabled; +} + +static std::string dflash_format_tokens(const llama_tokens & tokens, size_t limit = 16) { + std::string s = "["; + const size_t n = std::min(tokens.size(), limit); + for (size_t i = 0; i < n; ++i) { + if (i) { + s += ","; + } + s += std::to_string(tokens[i]); + } + if (tokens.size() > limit) { + s += ",..."; + } + s += "]"; + return s; +} + static bool dflash_verify_padding_enabled() { static const bool enabled = [] { const char * env = getenv("GGML_DFLASH_VERIFY_PAD"); @@ -1809,6 +1845,12 @@ struct server_context_impl { llama_context * ctx_tgt = nullptr; + // DFlash: if reduced verify (TOPK-based argmax) failed on first attempt, + // the backend likely lacks GGML_OP_TOPK support (e.g., Vulkan). Set to true + // to disable reduced verify permanently for this model session. Resets when + // the model is reloaded (since server_context_impl is destroyed/recreated). + bool dflash_reduced_verify_broken = false; + // DFlash: one drafter context shared across all slots' // common_speculative states (non-owning refs). Must outlive all specs — the // destroy() order below (specs first, then this) enforces that; when destroy() @@ -3274,8 +3316,7 @@ struct server_context_impl { bool loop_guard_accept_enabled(const server_slot & slot) const { return slot.task && slot.smpl && - slot.task->params.reasoning_loop_guard.mode != COMMON_REASONING_LOOP_GUARD_OFF && - slot.task->params.sampling.reasoning_budget_tracking; + slot.task->params.reasoning_loop_guard.mode != COMMON_REASONING_LOOP_GUARD_OFF; } bool handle_loop_guard_accept(server_slot & slot, const common_sampler_accept_info & info) { @@ -3288,7 +3329,6 @@ struct server_context_impl { slot.reasoning_output_tokens++; } else { slot.visible_output_tokens++; - return true; } const bool forcing_reasoning_end = info.reasoning_state_before == REASONING_BUDGET_FORCING || @@ -3297,14 +3337,18 @@ struct server_context_impl { return true; } - slot.loop_guard.accept(info.token, SERVER_LOOP_REGION_REASONING); + const server_loop_guard_region region = is_reasoning + ? SERVER_LOOP_REGION_REASONING + : SERVER_LOOP_REGION_VISIBLE; + + slot.loop_guard.accept(info.token, region); const bool token_is_eog = llama_vocab_is_eog(vocab, info.token); - if (!slot.loop_guard.should_check(SERVER_LOOP_REGION_REASONING, token_is_eog, forcing_reasoning_end)) { + if (!slot.loop_guard.should_check(region, token_is_eog, forcing_reasoning_end)) { return true; } - const auto check = slot.loop_guard.check(SERVER_LOOP_REGION_REASONING); + const auto check = slot.loop_guard.check(region); if (!check.triggered) { return true; } @@ -3313,7 +3357,8 @@ struct server_context_impl { slot.loop_guard_reason = server_loop_guard_reason_to_string(check); const auto & params = slot.task->params.reasoning_loop_guard; - if (params.mode == COMMON_REASONING_LOOP_GUARD_FORCE_CLOSE && + if (region == SERVER_LOOP_REGION_REASONING && + params.mode == COMMON_REASONING_LOOP_GUARD_FORCE_CLOSE && slot.loop_guard_interventions < params.interventions_max && common_sampler_force_reasoning_end(slot.smpl.get())) { slot.loop_guard_interventions++; @@ -4252,6 +4297,18 @@ struct server_context_impl { } void update_slots() { + // DFlash: check for hidden state dump trigger file + { + const char * dump_path_env = getenv("LLAMA_DFLASH_HIDDEN_DUMP"); + if (dump_path_env && strlen(dump_path_env) > 0) { + // Dump once per cycle if the path exists (idempotent) + int64_t n = llama_dflash_dump_hidden_states(ctx_tgt, dump_path_env); + if (n > 0) { + SRV_INF("dflash: dumped %ld hidden state tokens to %s\n", (long)n, dump_path_env); + } + } + } + // check if all slots are idle { bool all_idle = true; @@ -4394,27 +4451,31 @@ struct server_context_impl { }; auto dflash_backup_recurrent_state = [&](llama_seq_id seq_id_src, llama_seq_id seq_id_dst) { auto * mem = llama_get_memory(ctx_tgt); + if (const char * e = std::getenv("GGML_DFLASH_DEBUG"); e && std::atoi(e) != 0) { + // dump committed recurrent state (s_l for first recurrent layer) at cycle start + llama_dflash_dump_recurrent_state_dbg(ctx_tgt, seq_id_src, "backup_pre"); + } dflash_recurrent_profile_reset(mem); const int64_t t_backup_start = dflash_profile_start(); if (dflash_server_crash_trace_enabled()) { - SRV_INF("dflash crash breadcrumb: recurrent backup enter src=%d dst=%d\n", + SRV_DBG("dflash crash breadcrumb: recurrent backup enter src=%d dst=%d\n", (int) seq_id_src, (int) seq_id_dst); } const bool ordered = llama_dflash_memory_seq_cp_recurrent_ordered(ctx_tgt, seq_id_src, seq_id_dst, -1, -1); if (!ordered) { if (dflash_server_crash_trace_enabled()) { - SRV_INF("dflash crash breadcrumb: recurrent backup fallback copy src=%d dst=%d\n", + SRV_DBG("dflash crash breadcrumb: recurrent backup fallback copy src=%d dst=%d\n", (int) seq_id_src, (int) seq_id_dst); } llama_memory_seq_cp_recurrent(mem, seq_id_src, seq_id_dst, -1, -1); } else if (dflash_server_crash_trace_enabled()) { - SRV_INF("dflash crash breadcrumb: recurrent backup ordered copy complete src=%d dst=%d\n", + SRV_DBG("dflash crash breadcrumb: recurrent backup ordered copy complete src=%d dst=%d\n", (int) seq_id_src, (int) seq_id_dst); } dflash_profile_add(t_recurrent_backup_total, t_backup_start); dflash_recurrent_profile_collect(mem); if (dflash_server_crash_trace_enabled()) { - SRV_INF("dflash crash breadcrumb: recurrent backup exit src=%d dst=%d\n", + SRV_DBG("dflash crash breadcrumb: recurrent backup exit src=%d dst=%d\n", (int) seq_id_src, (int) seq_id_dst); } }; @@ -4802,6 +4863,15 @@ struct server_context_impl { } slot.spec_pad_i_batch.clear(); + if (dflash_server_crash_trace_enabled() && !slot.spec_draft.empty()) { + SRV_DBG("dflash draft produced: slot=%d n_draft=%zu adaptive_n_max=%d n_past=%d\n", + slot.id, slot.spec_draft.size(), slot.adaptive_n_max, draft_n_past); + } + if (dflash_token_trace_enabled() && !slot.spec_draft.empty()) { + SRV_INF("dflash token trace: slot=%d sampled=%d draft=%s n_past=%d\n", + slot.id, (int) slot.sampled, dflash_format_tokens(slot.spec_draft).c_str(), draft_n_past); + } + if (slot.spec_draft.size() > (size_t) n_draft_max) { SLT_WRN(slot, "draft size %d exceeds max %d, truncating\n", (int) slot.spec_draft.size(), n_draft_max); slot.spec_draft.resize(n_draft_max); @@ -5750,7 +5820,18 @@ struct server_context_impl { int32_t dflash_reduced_verify_view_start = 0; bool dflash_verify_graph_enabled = false; - if (dflash_verify_plan.enabled) { + // If reduced verify has been attempted but never succeeded, the backend + // likely doesn't support the TOPK compute op (e.g. Vulkan). Disable + // it permanently so we fall back to full logits. + const bool dflash_use_reduced_verify = dflash_verify_plan.enabled && !dflash_reduced_verify_broken; + + // Set to true when reduced verify fails on the current cycle. Skips all + // speculative token processing (KV rollback, accept, output) because we + // have no logits or argmax to work with. The draft state is cleared and + // the slot continues generating with a fresh single-token decode. + bool dflash_reduced_verify_recovery = false; + + if (dflash_use_reduced_verify) { for (int32_t j = 0; j < batch.n_tokens; j += std::min(n_batch, batch.n_tokens - j)) { const int32_t n_tokens_probe = std::min(n_batch, batch.n_tokens - j); const char * dflash_reduce_reason_probe = dflash_verify_plan.reason; @@ -5763,7 +5844,7 @@ struct server_context_impl { } } if (dflash_server_crash_trace_enabled() && params_base.speculative.type() == COMMON_SPECULATIVE_TYPE_DFLASH) { - SRV_INF("dflash crash breadcrumb: before verify logits config graph=%d enabled=%d top_k=%d batch_tokens=%d\n", + SRV_DBG("dflash crash breadcrumb: before verify logits config graph=%d enabled=%d top_k=%d batch_tokens=%d\n", dflash_verify_graph_enabled ? 1 : 0, dflash_verify_plan.enabled ? 1 : 0, dflash_verify_plan.top_k, @@ -5771,7 +5852,7 @@ struct server_context_impl { } llama_set_dflash_verify_logits(ctx_tgt, dflash_verify_graph_enabled, dflash_verify_plan.top_k); if (dflash_server_crash_trace_enabled() && params_base.speculative.type() == COMMON_SPECULATIVE_TYPE_DFLASH) { - SRV_INF("dflash crash breadcrumb: after verify logits config graph=%d enabled=%d top_k=%d batch_tokens=%d\n", + SRV_DBG("dflash crash breadcrumb: after verify logits config graph=%d enabled=%d top_k=%d batch_tokens=%d\n", dflash_verify_graph_enabled ? 1 : 0, dflash_verify_plan.enabled ? 1 : 0, dflash_verify_plan.top_k, @@ -5891,7 +5972,7 @@ struct server_context_impl { dflash_reduced_verify_ready = false; const char * dflash_reduce_reason = dflash_verify_plan.reason; bool dflash_reduce_this_view = false; - if (dflash_verify_plan.enabled) { + if (dflash_use_reduced_verify) { dflash_reduce_this_view = dflash_batch_view_is_reduced_verify(slots, params_base.sampling, params_base.speculative, use_rejection_sampling, ddtree_batch_active, i, n_tokens, @@ -5978,7 +6059,7 @@ struct server_context_impl { } if (dflash_server_profile_enabled(DFLASH_PROFILE_PREFILL)) { - SRV_INF("dflash prefill capture: view_start=%d n_tokens=%d enabled=%d\n", + SRV_DBG("dflash prefill capture: view_start=%d n_tokens=%d enabled=%d\n", i, n_tokens, dflash_capture_needed_for_view ? 1 : 0); } } @@ -5999,8 +6080,25 @@ struct server_context_impl { dflash_profile_add(t_replay_sync_total, t_replay_sync_start); } + if (dflash_server_crash_trace_enabled() && params_base.speculative.type() == COMMON_SPECULATIVE_TYPE_DFLASH) { + int n_logits_requested = 0; + for (int32_t j = 0; j < batch_view.n_tokens; ++j) { + if (batch_view.logits[j]) { n_logits_requested++; } + } + SRV_DBG("dflash crash breadcrumb: before decode ret=? batch_tokens=%d logits_requested=%d dflash_verify_graph=%d dflash_reduce_view=%d adaptive_n_max=%d\n", + batch_view.n_tokens, n_logits_requested, + dflash_verify_graph_enabled ? 1 : 0, + dflash_reduce_this_view ? 1 : 0, + slots.empty() ? -1 : slots[0].adaptive_n_max); + } const int64_t t_verify_start = ggml_time_us(); const int ret = llama_decode(ctx_tgt, batch_view); + if (dflash_server_crash_trace_enabled() && params_base.speculative.type() == COMMON_SPECULATIVE_TYPE_DFLASH) { + const float * logits_ptr = llama_get_logits(ctx_tgt); + SRV_DBG("dflash crash breadcrumb: after decode ret=%d batch_tokens=%d logits_ptr=%p reduced_verify_ready=%d\n", + ret, batch_view.n_tokens, (const void *) logits_ptr, + dflash_reduced_verify_ready ? 1 : 0); + } if (params_base.speculative.type() == COMMON_SPECULATIVE_TYPE_DFLASH && dflash_server_crash_trace_enabled() && !pending_prefill_flushes.empty()) { @@ -6011,7 +6109,7 @@ struct server_context_impl { const int64_t prefill_gpu_n = llama_dflash_prefill_gpu_n_tokens(ctx_tgt, pf.slot_id); llama_dflash_set_active_slot(ctx_tgt, pf.slot_id); const int64_t cpu_hidden_n = llama_get_layer_hidden_n_tokens(ctx_tgt, 0); - SRV_INF("dflash prefill capture result: slot=%d requested=%d src_offset=%d ret=%d plan=%d planned=%d written=%d prefill_gpu_active=%d prefill_gpu_n=%lld cpu_hidden_n=%lld\n", + SRV_DBG("dflash prefill capture result: slot=%d requested=%d src_offset=%d ret=%d plan=%d planned=%d written=%d prefill_gpu_active=%d prefill_gpu_n=%lld cpu_hidden_n=%lld\n", pf.slot_id, pf.span.n_tokens, pf.span.src_offset, @@ -6028,16 +6126,46 @@ struct server_context_impl { int32_t * compact_argmax = llama_get_logits_argmax(ctx_tgt); const int32_t compact_n = llama_get_logits_argmax_n(ctx_tgt); const int compact_k = llama_get_logits_argmax_k(ctx_tgt); - if (compact_argmax != nullptr && + // Value-sanity check: on backends that allocate the argmax buffer but + // don't implement GGML_OP_TOPK (e.g. Vulkan), the buffer is non-null + // with the right n/k but contains UNINITIALIZED garbage ids. The + // structural check below would pass and dflash_sample_reduced_verify + // would read garbage -> out-of-range bonus token. Validate that the + // first n_tokens*top_k ids are in [0, vocab_size); if not, treat as + // broken so the cycle falls back to full logits. + const int32_t vocab_size = llama_n_vocab(llama_model_get_vocab(llama_get_model(ctx_tgt))); + bool compact_argmax_valid = (compact_argmax != nullptr); + if (compact_argmax_valid) { + const size_t n_check = (size_t) std::min(compact_n, n_tokens) * (size_t) std::max(1, compact_k); + for (size_t v = 0; v < n_check; ++v) { + const int32_t id = compact_argmax[v]; + if (id < 0 || id >= vocab_size) { + compact_argmax_valid = false; + break; + } + } + } + if (compact_argmax_valid && compact_n >= n_tokens && compact_k == dflash_verify_plan.top_k) { dflash_reduced_verify_ready = true; dflash_reduced_verify_top_k = dflash_verify_plan.top_k; dflash_reduced_verify_view_start = i; - } else if (dflash_server_profile_enabled(DFLASH_PROFILE_VERIFY)) { - SRV_INF("dflash compact-output mismatch: view_start=%d n_tokens=%d expected_top_k=%d got_ptr=%d got_n=%d got_k=%d\n", - i, n_tokens, dflash_verify_plan.top_k, - compact_argmax != nullptr ? 1 : 0, compact_n, compact_k); + } else { + // Backend doesn't produce argmax/TOPK output (e.g. Vulkan + // doesn't support GGML_OP_TOPK). Disable reduced verify + // permanently so subsequent cycles fall back to full logits. + if (!dflash_reduced_verify_broken) { + SRV_WRN("DFlash reduced verify failed (no argmax output, top_k=%d); " + "disabling reduced verify for this session (backend may not support TOPK)\n", + dflash_verify_plan.top_k); + } + dflash_reduced_verify_broken = true; + if (dflash_server_profile_enabled(DFLASH_PROFILE_VERIFY)) { + SRV_INF("dflash compact-output mismatch: view_start=%d n_tokens=%d expected_top_k=%d got_ptr=%d got_n=%d got_k=%d\n", + i, n_tokens, dflash_verify_plan.top_k, + compact_argmax != nullptr ? 1 : 0, compact_n, compact_k); + } } } const int64_t t_verify_elapsed = ggml_time_us() - t_verify_start; @@ -6397,7 +6525,40 @@ struct server_context_impl { GGML_ABORT("DFlash reduced verifier output missing; falling back is unsafe because raw logits were not copied\n"); } } else { - prefetched.ids = common_sampler_sample_and_accept_n(slot.smpl.get(), ctx_tgt, slot.spec_i_batch, slot.spec_draft, false, on_accept); + if (dflash_reduce_this_view && !dflash_reduced_verify_ready) { + // Reduced verify was active but failed to produce argmax results, + // and the logits buffer was not allocated. Skip all speculative + // token processing; clear the draft state and let the slot + // continue with a fresh single-token decode. + SLT_WRN(slot, "DFlash reduced verify failed and logits unavailable; " + "discarding %d draft tokens and recovering\n", + (int) slot.spec_draft.size()); + dflash_reduced_verify_recovery = true; + } else { +#ifdef GGML_DFLASH_DEBUG + { + static int fv_probe_count = 0; + if (fv_probe_count < 2) { + fv_probe_count++; + const int n_vocab = llama_n_vocab(llama_model_get_vocab(llama_get_model(ctx_tgt))); + std::string dbg; + for (size_t p = 0; p < slot.spec_i_batch.size() && p < 4; ++p) { + const float * l = llama_get_logits_ith(ctx_tgt, slot.spec_i_batch[p]); + if (l) { + int amax = 0; float vmax = l[0]; bool has_nan = false; + for (int v = 0; v < n_vocab; ++v) { + if (std::isnan(l[v]) || std::isinf(l[v])) { has_nan = true; break; } + if (l[v] > vmax) { vmax = l[v]; amax = v; } + } + dbg += "p" + std::to_string(p) + ":amax=" + std::to_string(amax) + ":vmax=" + std::to_string(vmax) + ":nan=" + std::to_string((int)has_nan) + " "; + } else { dbg += "p" + std::to_string(p) + ":NULL "; } + } + SLT_WRN(slot, "FV_PROBE full_verify raw logits: %s\n", dbg.c_str()); + } + } +#endif + prefetched.ids = common_sampler_sample_and_accept_n(slot.smpl.get(), ctx_tgt, slot.spec_i_batch, slot.spec_draft, false, on_accept); + } } } prefetched.sample_us = ggml_time_us() - t_sample_start; @@ -6417,6 +6578,11 @@ struct server_context_impl { std::max(0, (int) prefetched.ids.size() - (prefetched.speculative_has_bonus ? 1 : 0)); prefetched.n_hidden_keep = prefetched.ids.empty() ? 0 : prefetched.n_accepted_draft + 1; prefetched.ready = true; + if (dflash_server_crash_trace_enabled() && params_base.speculative.type() == COMMON_SPECULATIVE_TYPE_DFLASH) { + SRV_INF("dflash prefetch verify: slot=%d ids=%zu n_accepted=%d n_draft=%zu reduced_ready=%d\n", + slot.id, prefetched.ids.size(), prefetched.n_accepted_draft, slot.spec_draft.size(), + dflash_reduced_verify_ready ? 1 : 0); + } } for (auto & slot : slots) { @@ -6765,7 +6931,26 @@ struct server_context_impl { GGML_ABORT("DFlash reduced verifier output missing; falling back is unsafe because raw logits were not copied\n"); } } else { - ids = common_sampler_sample_and_accept_n(slot.smpl.get(), ctx_tgt, slot.spec_i_batch, slot.spec_draft, false, on_accept); + // When DFlash reduced verify was active (dflash_reduce_this_view=true), + // output_reserve skips the logits buffer allocation. If the reduced + // verify path then fails to produce argmax output + // (reduced_verify_ready=false), we have no logits and no argmax. + // Accept 0 draft tokens and recover on the next cycle. + if (dflash_reduce_this_view && !dflash_reduced_verify_ready) { + // Reduced verify was active but failed to produce argmax results, + // and the logits buffer was not allocated because output_reserve + // saw dflash_reduced_consumer_active. We have no logits and no + // argmax output. Skip all speculative token processing; clear + // the draft state and let the slot continue with a fresh + // single-token decode on the next cycle (which will use full + // logits since dflash_reduced_verify_broken is now true). + SLT_WRN(slot, "DFlash reduced verify failed and logits unavailable; " + "discarding %d draft tokens and recovering\n", + (int) slot.spec_draft.size()); + dflash_reduced_verify_recovery = true; + } else { + ids = common_sampler_sample_and_accept_n(slot.smpl.get(), ctx_tgt, slot.spec_i_batch, slot.spec_draft, false, on_accept); + } } } profile_accept_lap(profile_accept_sample_us); @@ -6783,6 +6968,16 @@ struct server_context_impl { n_accepted_draft = std::max(0, (int) ids.size() - (speculative_has_bonus ? 1 : 0)); n_hidden_keep = ids.empty() ? 0 : n_accepted_draft + 1; + if (dflash_rx_diag_enabled()) { + SRV_INF("dflash rx: slot=%d verify_result ids_size=%zu n_accepted_draft=%d n_hidden_keep=%d has_bonus=%d\n", + slot.id, ids.size(), n_accepted_draft, n_hidden_keep, + speculative_has_bonus ? 1 : 0); + } + if (dflash_token_trace_enabled()) { + SRV_INF("dflash token trace: slot=%d verify ids=%s accepted=%d has_bonus=%d\n", + slot.id, dflash_format_tokens(ids).c_str(), n_accepted_draft, speculative_has_bonus ? 1 : 0); + } + if (params_base.speculative.type() == COMMON_SPECULATIVE_TYPE_DFLASH) { llama_dflash_set_active_slot(ctx_tgt, slot.id); } @@ -6790,12 +6985,54 @@ struct server_context_impl { batch_tokens.push_back(slot.sampled); batch_tokens.insert(batch_tokens.end(), slot.spec_draft.begin(), slot.spec_draft.end()); common_speculative_update_logits_deferred_dflash_kv(slot.get_spec(), ctx_tgt, batch_tokens, n_hidden_keep); + if (dflash_rx_diag_enabled()) { + SRV_INF("dflash rx: slot=%d update_logits_deferred called n_hidden_keep=%d batch_tokens=%zu\n", + slot.id, n_hidden_keep, batch_tokens.size()); + } profile_accept_lap(profile_accept_update_us); } + if (dflash_server_crash_trace_enabled() && params_base.speculative.type() == COMMON_SPECULATIVE_TYPE_DFLASH) { + SRV_DBG("dflash verify result: slot=%d ids=%zu n_accepted=%d n_draft=%zu reduced_ready=%d\n", + slot.id, ids.size(), n_accepted_draft, n_draft, + dflash_reduced_verify_ready ? 1 : 0); + } GGML_ASSERT(slot.remaining_generation_budget(params_base) == -1 || (int32_t) ids.size() <= slot.remaining_generation_budget(params_base)); + // When reduced verify failed, we have no logits or argmax output. + // Skip all speculative token processing (KV rollback, accept, output) + // but still reset the draft model's state so the next cycle is consistent. + if (dflash_reduced_verify_recovery) { + if (params_base.speculative.type() == COMMON_SPECULATIVE_TYPE_DFLASH) { + llama_dflash_set_active_slot(ctx_tgt, slot.id); + common_speculative_accept(slot.get_spec(), 0); + if (dflash_rx_diag_enabled()) { + SRV_INF("dflash rx: slot=%d REDUCED_VERIFY_RECOVERY append_skipped=1 accept_reset=1 draft_size=%zu\n", + slot.id, slot.spec_draft.size()); + } + } + // Clean up draft tokens from KV cache and prompt so the next cycle + // starts from a consistent position. The draft was already decoded + // into the KV cache, but we have no logits to verify against. + { + // Remove from n_pos_before_draft onwards (target + draft tokens) + // Use llama_memory_seq_rm directly (not common_context_seq_rm) to + // avoid abort on failure — recurrent memory may not support partial rm + auto * mem = llama_get_memory(ctx_tgt); + llama_memory_seq_rm(mem, slot.id, slot.n_pos_before_draft, -1); + if (ctx_dft) { + auto * mem_dft = llama_get_memory(ctx_dft.get()); + llama_memory_seq_rm(mem_dft, slot.id, slot.n_pos_before_draft, -1); + } + // Remove draft tokens + target token from prompt + slot.prompt.tokens.keep_first(slot.prompt.n_tokens() - (size_t) slot.spec_draft.size() - 1); + } + dflash_skip_tg_slot_for_next_cycle(slot); + slot.has_draft_backup = false; + continue; + } + const int64_t t_current = ggml_time_us(); if (ids.empty()) { @@ -6943,9 +7180,21 @@ struct server_context_impl { // Plain MTP is handled by the dedicated block above (use_mtp_spec_accept). // This path is only reached for DFlash, tree, or other speculative types. + if (dflash_server_crash_trace_enabled() && params_base.speculative.type() == COMMON_SPECULATIVE_TYPE_DFLASH) { + SRV_INF("dflash before speculative_accept: slot=%d n_accepted_draft=%d\n", + slot.id, n_accepted_draft); + } { common_speculative_accept(slot.get_spec(), n_accepted_draft); } + if (dflash_rx_diag_enabled()) { + SRV_INF("dflash rx: slot=%d speculative_accept called n_accepted_draft=%d\n", + slot.id, n_accepted_draft); + } + if (dflash_server_crash_trace_enabled() && params_base.speculative.type() == COMMON_SPECULATIVE_TYPE_DFLASH) { + SRV_INF("dflash after speculative_accept: slot=%d\n", + slot.id); + } if (is_draft_tree) { slot.prompt.tokens.keep_first(slot.n_tokens_before_draft + 1); @@ -6962,6 +7211,13 @@ struct server_context_impl { GGML_ASSERT(seq_backup >= 0); const bool all_accepted_flat = (n_accepted_draft == (int) n_draft) && !had_dflash_padding; + if (dflash_server_crash_trace_enabled() && params_base.speculative.type() == COMMON_SPECULATIVE_TYPE_DFLASH) { + SRV_DBG("dflash rollback enter: slot=%d all_accepted_flat=%d n_accepted_draft=%d n_draft=%d " + "n_hidden_keep=%d seq_backup=%d n_pos_before_draft=%d is_draft_tree=%d\n", + slot.id, all_accepted_flat, n_accepted_draft, n_draft, + n_hidden_keep, seq_backup, slot.n_pos_before_draft, is_draft_tree); + } + if (params_base.speculative.type() == COMMON_SPECULATIVE_TYPE_DFLASH && is_draft_tree) { llama_dflash_set_active_slot(ctx_tgt, slot.id); @@ -7044,12 +7300,37 @@ struct server_context_impl { llama_memory_seq_rm(mem, slot.id, slot.prompt.tokens.pos_next(), -1); } else { llama_clear_tree_parent_ids(ctx_tgt); - llama_dflash_rollback(ctx_tgt, slot.id, seq_backup, slot.n_pos_before_draft, n_hidden_keep); + if (dflash_server_crash_trace_enabled()) { + SRV_DBG("dflash before dflash_rollback: slot=%d n_hidden_keep=%d n_pos_before_draft=%d\n", + slot.id, n_hidden_keep, slot.n_pos_before_draft); + } + const int n_reeval = llama_dflash_rollback(ctx_tgt, slot.id, seq_backup, slot.n_pos_before_draft, n_hidden_keep); + if (dflash_server_crash_trace_enabled()) { + SRV_DBG("dflash after dflash_rollback: slot=%d n_reeval=%d\n", + slot.id, n_reeval); + } if (n_slots_drafted > 1) { // Multi-slot accept can immediately roll back another seq; make this // seq's async recurrent replay visible before the next mutation. llama_tape_replay_sync(ctx_tgt); } + // When tape replay was unavailable (e.g., Vulkan), re-decode + // the accepted positions to advance the recurrent state. + // logits=false: we only need the state advance, not the output. + if (n_reeval > 0) { + llama_batch batch_reeval = llama_batch_init(n_reeval, 0, 1); + for (int j = 0; j < n_reeval; ++j) { + const llama_pos pos = slot.n_pos_before_draft + j; + common_batch_add(batch_reeval, slot.prompt.tokens[slot.n_tokens_before_draft + j], pos, { slot.id }, false); + } + const int ret_reeval = llama_decode(ctx_tgt, batch_reeval); + llama_batch_free(batch_reeval); + if (ret_reeval != 0) { + SLT_WRN(slot, "re-decode of %d accepted tokens failed (ret=%d), " + "recurrent state may be inconsistent\n", + n_reeval, ret_reeval); + } + } } } else { auto * mem = llama_get_memory(ctx_tgt); diff --git a/tools/server/server-loop-guard.cpp b/tools/server/server-loop-guard.cpp index 0fcf10fd7529..e84074028ff0 100644 --- a/tools/server/server-loop-guard.cpp +++ b/tools/server/server-loop-guard.cpp @@ -54,7 +54,9 @@ bool server_loop_guard::should_check(server_loop_guard_region region, bool token return false; } const int32_t n_seen = seen(region); - if (region == SERVER_LOOP_REGION_REASONING && n_seen < params.min_reasoning_tokens) { + const int32_t min_tokens = (region == SERVER_LOOP_REGION_REASONING) + ? params.min_reasoning_tokens : params.min_reasoning_tokens; + if (n_seen < min_tokens) { return false; } return n_seen > 0 && n_seen % params.check_interval == 0; @@ -64,7 +66,10 @@ server_loop_guard_result server_loop_guard::check(server_loop_guard_region regio if (params.mode == COMMON_REASONING_LOOP_GUARD_OFF) { return {}; } - if (region == SERVER_LOOP_REGION_REASONING && reasoning_seen < params.min_reasoning_tokens) { + const int32_t n_seen = seen(region); + const int32_t min_tokens = (region == SERVER_LOOP_REGION_REASONING) + ? params.min_reasoning_tokens : params.min_reasoning_tokens; + if (n_seen < min_tokens) { return {}; } diff --git a/tools/server/tests/unit/test_reasoning_loop_guard.py b/tools/server/tests/unit/test_reasoning_loop_guard.py index deb9d823c12a..7f9ddef2b3fd 100644 --- a/tools/server/tests/unit/test_reasoning_loop_guard.py +++ b/tools/server/tests/unit/test_reasoning_loop_guard.py @@ -136,3 +136,32 @@ def test_reasoning_loop_guard_force_close_intervenes_then_continues(): assert res.body["visible_completion_tokens"] > 0 assert res.body["loop_guard"]["triggered"] is True assert res.body["loop_guard"]["action"] == "force-close" + + +def test_reasoning_loop_guard_stop_triggers_on_visible_output_loop(): + """Regression test: visible output loops (e.g. repeated Assistant: snippets) + should be detected by the loop guard even when no reasoning tags are used.""" + global server + server.start() + + res = server.make_request("POST", "/completion", data={ + "prompt": "Assistant:\n\n\n\n\n\nAssistant:\n\n\n\n\n\nAssistant:\n\n\n\n\n\nAssistant:\n\n", + "n_predict": 16, + "temperature": 0.0, + "logit_bias": _single_token_bias(" loop"), + "reasoning_loop_guard": "stop", + "reasoning_loop_min_tokens": 4, + "reasoning_loop_window": 8, + "reasoning_loop_max_period": 2, + "reasoning_loop_min_coverage": 3, + "reasoning_loop_check_interval": 1, + "reasoning_loop_interventions": 0, + }) + + assert res.status_code == 200 + assert res.body["stop_type"] == "limit" + assert res.body["stop_detail"] == "reasoning_loop_guard" + assert res.body["tokens_predicted"] < 16 + assert res.body["loop_guard"]["triggered"] is True + assert res.body["loop_guard"]["action"] == "stop" + assert res.body["loop_guard"]["reason"] == "periodic_tail"