Support SWA compression with DFlash and DSpark (#2384)

Co-authored-by: Kawrakow <iwankawrakow@gmail.com>

common/speculative-dflash-impl.h

@@ -128,6 +128,7 @@ struct common_speculative_state_dflash : public common_speculative_state {
     uint64_t rebuild_rows = 0;
     uint64_t rebuild_decode_time_us = 0;
     uint64_t generated_tokens = 0;
+    uint64_t target_window_ring_wraps = 0;
 
     common_speculative_state_dflash(
             enum common_speculative_type type,
@@ -265,18 +266,20 @@ struct common_speculative_state_dflash : public common_speculative_state {
 
         llama_set_dflash_visible_cross_ctx(ctx_dft, this->cross_ctx);
         llama_set_dflash_dspark(ctx_dft, is_dspark);
-        LOG_INF("%s: DFlash context ready (n_ctx=%d, block_size=%d, query_capacity=%d, active_width=%d, cross_ctx=%d, n_target_features=%d, n_target_layers=%d)\n",
+        LOG_INF("%s: DFlash context ready (n_ctx=%d, block_size=%d, query_capacity=%d, active_width=%d, cross_ctx=%d, n_target_features=%d, n_target_layers=%d, host_ring=%.2f MiB)\n",
                 __func__, llama_n_ctx(ctx_dft), block_size, query_capacity, active_width, this->cross_ctx,
-                n_target_features, n_target_layers);
+                n_target_features, n_target_layers,
+                (double) target_window_ring.size() * sizeof(float) / (1024.0 * 1024.0));
     }
     ~common_speculative_state_dflash() override {
         if (rebuild_count > 0) {
-            LOG_INF("%s: DFlash cache rebuilds=%llu rows=%llu rebuild+decode=%.3f ms rows/token=%.3f\n",
+            LOG_INF("%s: DFlash cache rebuilds=%llu rows=%llu rebuild+decode=%.3f ms rows/token=%.3f host_ring_wraps=%llu\n",
                     __func__,
                     (unsigned long long) rebuild_count,
                     (unsigned long long) rebuild_rows,
                     (double) rebuild_decode_time_us / 1000.0,
-                    generated_tokens > 0 ? (double) rebuild_rows / (double) generated_tokens : 0.0);
+                    generated_tokens > 0 ? (double) rebuild_rows / (double) generated_tokens : 0.0,
+                    (unsigned long long) target_window_ring_wraps);
         }
         llama_clear_dflash_capture(ctx_tgt);
         if (ctx_dft) {
@@ -297,6 +300,7 @@ struct common_speculative_state_dflash : public common_speculative_state {
         rebuild_rows = 0;
         rebuild_decode_time_us = 0;
         generated_tokens = 0;
+        target_window_ring_wraps = 0;
     }
 
     void draft(
@@ -518,6 +522,7 @@ static void dflash_ring_append_rows(
     int32_t write_pos = state.target_window_ring_write_pos;
     int32_t remaining = n_rows;
     const float * src = rows;
+    state.target_window_ring_wraps += ((uint64_t) write_pos + (uint64_t) n_rows) / (uint64_t) state.cross_ctx;
     while (remaining > 0) {
         const int32_t chunk_rows = std::min<int32_t>(remaining, state.cross_ctx - write_pos);
         std::memcpy(

common/speculative.cpp

@@ -3,6 +3,7 @@
 #include "common.h"
 #include "ggml.h"
 #include "llama.h"
+#include "llama-spec-features-dflash.h"
 #include "log.h"
 #include "ngram-cache.h"
 #include "ngram-map.h"
@@ -1355,17 +1356,15 @@ common_speculative * common_speculative_init(
                 LOG_INF("%s: DFlash draft context %d exceeds target context %d, clamping to target capacity\n",
                         __func__, requested_draft_ctx, target_ctx);
             }
-            const int32_t cross_ctx = effective_draft_ctx - query_capacity;
+            const int32_t logical_cross_ctx = effective_draft_ctx - query_capacity;
-            if (cross_ctx <= 0) {
+            if (logical_cross_ctx <= 0) {
                 LOG_ERR("%s: invalid DFlash draft context size draft=%d target=%d query_capacity=%d, draft context must exceed the query block\n",
                         __func__, requested_draft_ctx, target_ctx, query_capacity);
                 return nullptr;
             }
 
             cparams_dft.n_ctx = (uint32_t) effective_draft_ctx;
-            dflash_cross_ctx = cross_ctx;
+            cparams_dft.dflash_query_capacity = query_capacity;
-            LOG_INF("%s: DFlash context target/slot=%d logical=%d cross_ctx=%d query_block=%d\n",
-                    __func__, target_ctx, effective_draft_ctx, cross_ctx, query_capacity);
         }
 
         ctx_dft = llama_init_from_model(params.model_dft, cparams_dft);
@@ -1373,6 +1372,14 @@ common_speculative * common_speculative_init(
             LOG_ERR("%s", "failed to create draft context\n");
             return nullptr;
         }
+        if (has_dflash_stage) {
+            dflash_cross_ctx = llama_get_dflash_visible_cross_ctx(ctx_dft);
+            if (dflash_cross_ctx <= 0) {
+                LOG_ERR("%s: DFlash context did not expose a valid resolved physical cross-context\n", __func__);
+                llama_free(ctx_dft);
+                return nullptr;
+            }
+        }
     }
 
     // Compute the implementations to use based on the resolved stage chain.

include/llama.h

@@ -518,6 +518,7 @@ extern "C" {
         void *              abort_callback_data;
         void *              offload_policy;
         void *              cuda_params;
+        int32_t             dflash_query_capacity; // internal DFlash query capacity override
     };
 
     // model quantization parameters

src/llama-cparams.h

@@ -54,6 +54,7 @@ struct llama_cparams {
     float thresh_experts;
     bool mtp;
     int  worst_graph_tokens;
+    int  dflash_query_capacity = 0; // internal DFlash query capacity override
 
     enum ggml_type reduce_type;
     enum ggml_type graph_attn_precision;

src/llama-dflash.cpp

@@ -59,7 +59,9 @@ static ggml_backend_t llama_backend_for_tensor(const llama_context & lctx, const
 
 bool llama_context::ensure_dflash_kv_cache_tensors(int32_t cross_ctx) {
     const int32_t target_cross_ctx = std::max<int32_t>(1, cross_ctx);
-    const int32_t target_token_capacity = std::max<int32_t>(1, (int32_t) model.hparams.dflash_block_size);
+    const int32_t target_token_capacity = std::max<int32_t>(
+            std::max<int32_t>(1, (int32_t) model.hparams.dflash_block_size),
+            cparams.dflash_query_capacity);
     const int32_t target_cache_n_kv_total = GGML_PAD(target_cross_ctx + target_token_capacity, (int32_t) llama_kv_cache::get_padding(cparams.flash_attn));
     const ggml_type target_cache_type = cparams.flash_attn ? GGML_TYPE_F16 : GGML_TYPE_F32;
     const int32_t n_layer = model.hparams.n_layer;
@@ -153,6 +155,25 @@ bool llama_context::ensure_dflash_kv_cache_tensors(int32_t cross_ctx) {
 
     llama_reset_dflash_kv_cache_state(this);
 
+    size_t k_bytes = 0;
+    size_t v_bytes = 0;
+    for (size_t i = 0; i < dflash.kv.cache_bufs.size(); ++i) {
+        if (dflash.kv.cache_bufs[i] == nullptr) {
+            continue;
+        }
+        if ((i & 1) == 0) {
+            k_bytes += ggml_backend_buffer_get_size(dflash.kv.cache_bufs[i]);
+        } else {
+            v_bytes += ggml_backend_buffer_get_size(dflash.kv.cache_bufs[i]);
+        }
+    }
+    LLAMA_LOG_INFO("%s: DFlash custom K/V cache bytes = %.2f MiB (K %.2f MiB, V %.2f MiB, cross_ctx %d, query_capacity %d, tail_capacity %d)\n",
+            __func__,
+            (double) (k_bytes + v_bytes) / (1024.0 * 1024.0),
+            (double) k_bytes / (1024.0 * 1024.0),
+            (double) v_bytes / (1024.0 * 1024.0),
+            target_cross_ctx, cparams.dflash_query_capacity, target_token_capacity);
+
     return true;
 }
 

src/llama-model.h

@@ -8,6 +8,7 @@
 
 #include "ggml-backend.h"
 
+#include <algorithm>
 #include <vector>
 #include <unordered_map>
 #include <set>
@@ -647,9 +648,29 @@ struct llama_model {
 
     // a compacted sliding-window cache needs the graph to build its KQ mask over the compacted
     // layout, and the compacted mask keys on position alone, so it requires a single sequence
+    bool supports_dflash_swa_compress() const {
+        if (!llm_arch_is_dflash_family(arch) || hparams.n_swa == 0 || hparams.n_layer == 0) {
+            return false;
+        }
+        for (uint32_t il = 0; il < hparams.n_layer; ++il) {
+            if (!hparams.swa_layers[il]) {
+                return false;
+            }
+        }
+        return true;
+    }
+
+    int32_t dflash_swa_compress_cross_ctx(int32_t logical_cross_ctx, bool enabled) const {
+        const int32_t safe_cross_ctx = std::max<int32_t>(1, logical_cross_ctx);
+        return enabled && supports_dflash_swa_compress()
+                ? std::min<int32_t>(safe_cross_ctx, (int32_t) hparams.n_swa)
+                : safe_cross_ctx;
+    }
+
     bool supports_swa_compress() const {
         return arch == LLM_ARCH_OPENPANGU || arch == LLM_ARCH_DEEPSEEK4
-            || arch == LLM_ARCH_LAGUNA   || arch == LLM_ARCH_GEMMA4;
+            || arch == LLM_ARCH_LAGUNA    || arch == LLM_ARCH_GEMMA4
+            || supports_dflash_swa_compress() ;
     }
 
     static inline int hadamard_size(int head_size) {

src/llama-spec-features-dflash.cpp

@@ -67,7 +67,10 @@ void llama_set_dflash_visible_cross_ctx(
         return;
     }
 
-    ctx->dflash.visible_cross_ctx = std::max<int32_t>(0, cross_ctx);
+    const int32_t safe_cross_ctx = std::max<int32_t>(0, cross_ctx);
+    ctx->dflash.visible_cross_ctx = ctx->cparams.swa_compress
+            ? ctx->model.dflash_swa_compress_cross_ctx(safe_cross_ctx, true)
+            : safe_cross_ctx;
 }
 
 int32_t llama_get_dflash_visible_cross_ctx(

src/llama.cpp

@@ -1225,7 +1225,11 @@ static bool llama_kv_cache_init(
             cache.head_swa     = cache.sink_rows;
             cache.pos_base_swa = 0;
         } else {
-            LLAMA_LOG_WARN("%s: --swa-compress had no effect: no compactable sliding-window layers\n", __func__);
+            if (model.supports_dflash_swa_compress()) {
+                LLAMA_LOG_INFO("%s: --swa-compress uses the DFlash custom cache, ordinary KV bookkeeping has no compactable layers\n", __func__);
+            } else {
+                LLAMA_LOG_WARN("%s: --swa-compress had no effect: no compactable sliding-window layers\n", __func__);
+            }
         }
     }
 
@@ -8053,6 +8057,7 @@ struct llama_context_params llama_context_default_params() {
         /*.abort_callback_data         =*/ nullptr,
         /*.offload_policy              =*/ nullptr,
         /*.cuda_params                 =*/ nullptr,
+        /*.dflash_query_capacity       =*/ 0,
     };
 
     return result;
@@ -8563,6 +8568,7 @@ struct llama_context * llama_init_from_model(
     cparams.cuda_params      = params.cuda_params;
     cparams.mtp              = params.mtp;
     cparams.worst_graph_tokens = params.worst_case_tokens;
+    cparams.dflash_query_capacity = params.dflash_query_capacity;
 
     cparams.reduce_type      = params.type_reduce;
     cparams.graph_attn_precision = params.type_graph_attn;
@@ -8580,6 +8586,19 @@ struct llama_context * llama_init_from_model(
     // this is necessary due to kv_self.n being padded later during inference
     cparams.n_ctx            = GGML_PAD(cparams.n_ctx, llama_kv_cache::get_padding(cparams.flash_attn));
 
+    if (llm_arch_is_dflash_family(model->arch)) {
+        const int32_t query_capacity = cparams.dflash_query_capacity > 0
+                ? cparams.dflash_query_capacity
+                : std::max<int32_t>(1, (int32_t) model->hparams.dflash_block_size);
+        const int32_t logical_cross_ctx = std::max<int32_t>(1,
+                (int32_t) cparams.n_ctx - query_capacity);
+        const int32_t physical_cross_ctx = model->dflash_swa_compress_cross_ctx(
+                logical_cross_ctx, cparams.swa_compress);
+        ctx->dflash.visible_cross_ctx = physical_cross_ctx;
+        LLAMA_LOG_INFO("%s: DFlash context logical_cross_ctx=%d physical_cross_ctx=%d query_capacity=%d window=%u\n",
+                __func__, logical_cross_ctx, physical_cross_ctx, query_capacity, model->hparams.n_swa);
+    }
+
     // with causal attention, the batch size is limited by the context size
     cparams.n_batch          = hparams.causal_attn ? std::min(cparams.n_ctx, params.n_batch) : params.n_batch;