diff --git a/common/speculative-dflash-impl.h b/common/speculative-dflash-impl.h index fe17f0e1..1fa3c7fa 100644 --- a/common/speculative-dflash-impl.h +++ b/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(remaining, state.cross_ctx - write_pos); std::memcpy( diff --git a/common/speculative.cpp b/common/speculative.cpp index 317f165e..a2aa6580 100644 --- a/common/speculative.cpp +++ b/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; - if (cross_ctx <= 0) { + const int32_t logical_cross_ctx = effective_draft_ctx - query_capacity; + 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; - 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); + cparams_dft.dflash_query_capacity = 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. diff --git a/include/llama.h b/include/llama.h index da39924b..a7fb80c7 100644 --- a/include/llama.h +++ b/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 diff --git a/src/llama-cparams.h b/src/llama-cparams.h index 5f099728..6422e57e 100644 --- a/src/llama-cparams.h +++ b/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; diff --git a/src/llama-dflash.cpp b/src/llama-dflash.cpp index 60cf38bb..b1d2d058 100644 --- a/src/llama-dflash.cpp +++ b/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(1, cross_ctx); - const int32_t target_token_capacity = std::max(1, (int32_t) model.hparams.dflash_block_size); + const int32_t target_token_capacity = std::max( + std::max(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; } diff --git a/src/llama-model.h b/src/llama-model.h index 4c28e8f9..7301f070 100644 --- a/src/llama-model.h +++ b/src/llama-model.h @@ -8,6 +8,7 @@ #include "ggml-backend.h" +#include #include #include #include @@ -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(1, logical_cross_ctx); + return enabled && supports_dflash_swa_compress() + ? std::min(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) { diff --git a/src/llama-spec-features-dflash.cpp b/src/llama-spec-features-dflash.cpp index 451bc769..cbd72270 100644 --- a/src/llama-spec-features-dflash.cpp +++ b/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(0, cross_ctx); + const int32_t safe_cross_ctx = std::max(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( diff --git a/src/llama.cpp b/src/llama.cpp index cc665fa3..1fc46380 100644 --- a/src/llama.cpp +++ b/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(1, (int32_t) model->hparams.dflash_block_size); + const int32_t logical_cross_ctx = std::max(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;