nicolas
/
ik_llama_opt
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Support SWA compression with DFlash and DSpark (#2384) 

Co-authored-by: Kawrakow <iwankawrakow@gmail.com>
This commit is contained in:
Samuel Oliveira Alves 2026-09-03 13:23:23 -03:00 committed by GitHub
parent caf7eae528
commit 6d6fe936db
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
8 changed files with 91 additions and 13 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

13
common/speculative-dflash-impl.h
View File

@ -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(

17
common/speculative.cpp
View File

@ -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.

1
include/llama.h
View File

@ -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

1
src/llama-cparams.h
View File

@ -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;

23
src/llama-dflash.cpp
View File

@ -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;
}

23
src/llama-model.h
View File

@ -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) {

5
src/llama-spec-features-dflash.cpp
View File

@ -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(

19
src/llama.cpp
View File

@ -1224,10 +1224,14 @@ static bool llama_kv_cache_init(
cache.window_swa = hparams.n_swa;
cache.head_swa = cache.sink_rows;
cache.pos_base_swa = 0;
} else {
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__);
}
}
}
cache.type_k = type_k;
cache.type_v = type_v;
@ -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;