diff --git a/common/common.cpp b/common/common.cpp index b67d80aa9..bc52ec338 100644 --- a/common/common.cpp +++ b/common/common.cpp @@ -1920,6 +1920,10 @@ bool gpt_params_find_arg(int argc, char ** argv, const std::string & arg, gpt_pa params.fused_idx_topk = true; return true; } + if (arg == "--swa-compress") { + params.swa_compress = true; + return true; + } if (arg == "-dsatk" || arg == "--dsa-top-k") { CHECK_ARG params.dsa_top_k = std::stoi(argv[i]); @@ -3054,6 +3058,7 @@ void gpt_params_print_usage(int /*argc*/, char ** argv, const gpt_params & param options.push_back({ "*", "-mla, --mla-use", "enable MLA (default: %d)", params.mla_attn }); options.push_back({ "*", "-dsa, --dsa", "enable GLM DSA sparse attention (GLM-DSA arch only; default: %s)", params.dsa ? "enabled" : "disabled" }); options.push_back({ "*", "-fidx, --fused-indexer-topk", "enable the fused indexer topk op (DSA only; default: %s)", params.fused_idx_topk ? "enabled" : "disabled" }); + options.push_back({ "*", " --swa-compress", "allocate sliding-window layers at window size instead of n_ctx (default: %s)", params.swa_compress ? "enabled" : "disabled" }); options.push_back({ "*", "-dsatk, --dsa-top-k", "DSA top-k override; <0 uses the model's configured indexer_top_k (default: %d)", params.dsa_top_k }); options.push_back({ "*", "-amb, --attention-max-batch", "max batch size for attention computations (default: %d)", params.attn_max_batch}); options.push_back({ "*", "-no-fmoe, --no-fused-moe", "disable fused MoE (default: %s)", params.fused_moe_up_gate ? "enabled" : "disabled" }); @@ -4246,6 +4251,7 @@ struct llama_model_params common_model_params_to_llama(const gpt_params & params mparams.mtp = params.speculative.has_stage_type(COMMON_SPECULATIVE_TYPE_MTP); mparams.flash_attn = params.flash_attn; mparams.defer_experts = params.defer_experts; + mparams.swa_compress = params.swa_compress; if (params.kv_overrides.empty()) { mparams.kv_overrides = NULL; } else { @@ -4326,6 +4332,7 @@ struct llama_context_params common_context_params_to_llama(const gpt_params & pa cparams.graph_reuse = params.graph_reuse; cparams.dsa = params.dsa; cparams.fused_idx_topk = params.fused_idx_topk; + cparams.swa_compress = params.swa_compress; cparams.dsa_top_k = params.dsa_top_k; cparams.k_cache_hadamard = params.k_cache_hadamard; cparams.v_cache_hadamard = params.v_cache_hadamard; diff --git a/common/common.h b/common/common.h index ada90425e..962c47fb4 100644 --- a/common/common.h +++ b/common/common.h @@ -422,6 +422,7 @@ struct gpt_params { bool graph_reuse = true; // if to reuse compute graphs bool dsa = false; // enable GLM DSA sparse attention (off by default; opt-in via --dsa) bool fused_idx_topk = true; // enable the fused indexer topk op (off by default; opt-in via -fidx or --fused-indexer-topk) + bool swa_compress = false; int dsa_top_k = -1; // DSA top-k override (<0 => use the model's configured indexer_top_k) int min_experts = -1; float thresh_experts = 0; diff --git a/examples/server/server-context.cpp b/examples/server/server-context.cpp index b28de4ca6..8abe97c12 100644 --- a/examples/server/server-context.cpp +++ b/examples/server/server-context.cpp @@ -391,6 +391,12 @@ void server_context::init() { reuse_forced_off = true; } + if (params_base.cache_ram_mib != 0 && !llama_supports_full_state_io(ctx)) { + LLAMA_LOG_WARN("prompt cache is disabled: this context cannot save full sequence state (--swa-compress)\n"); + params_base.cache_ram_mib = 0; + reuse_forced_off = true; + } + if (params_base.cache_ram_mib != 0 && llama_model_supports_partial_kv_reuse(model)) { if (params_base.cache_ram_mib < 0) { LLAMA_LOG_INFO("prompt cache is enabled, size limit: %s\n", "no limit"); @@ -3661,12 +3667,16 @@ void server_context::apply_checkpoint(server_slot & slot) { // restore the context checkpoint const int64_t t_start = ggml_time_us(); const size_t checkpoint_size = it->data.size(); - if (is_openpangu) { + const bool rewound = !is_openpangu || llama_kv_cache_seq_rm(slot.ctx, slot.id, it->pos_max + 1, -1); - } - const size_t n = llama_state_seq_set_data(ctx, it->data.data(), checkpoint_size, slot.id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY); + const size_t n = rewound + ? llama_state_seq_set_data(ctx, it->data.data(), checkpoint_size, slot.id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY) + : 0; - if (n != checkpoint_size) { + if (!rewound) { + SLT_ERR(slot, "checkpoint rewind to %d was refused; reprocessing from scratch\n", it->pos_max + 1); + do_reset = true; + } else if (n != checkpoint_size) { SLT_ERR(slot, "failed to restore context checkpoint (pos_min = %d, pos_max = %d, n_tokens = %" PRId64 ", size = %.3f MiB)\n", it->pos_min, it->pos_max, it->n_tokens, (float)checkpoint_size / 1024 / 1024); do_reset = true; //printf("[DEBUG] `do_reset` was set to `true` after failing to restore a checkpoint"); diff --git a/include/llama.h b/include/llama.h index 24e2eb9ec..eb7030c34 100644 --- a/include/llama.h +++ b/include/llama.h @@ -433,6 +433,7 @@ extern "C" { bool dry_run; // skip loading tensors bool flash_attn; bool defer_experts; // defer expert mmap residency to speed up model loading (Linux only) + bool swa_compress; // must match llama_context_params::swa_compress; the fit also assumes that context's n_ubatch }; // NOTE: changing the default values of parameters marked as [EXPERIMENTAL] may cause crashes or incorrect results in certain configurations @@ -494,6 +495,7 @@ extern "C" { bool graph_reuse; // whether to reuse graphs when possible [EXPERIMENTAL] bool dsa; // enable GLM DSA sparse attention (off by default) [EXPERIMENTAL] bool fused_idx_topk; // enable the fused indexer topk op (off by default) [EXPERIMENTAL] + bool swa_compress; // allocate sliding-window layers at window size instead of n_ctx (off by default) [EXPERIMENTAL] int dsa_top_k; // DSA top-k override (<0 => model's configured indexer_top_k) [EXPERIMENTAL] int min_experts; float thresh_experts; @@ -720,6 +722,10 @@ extern "C" { // Currently true for every model; no architecture is excluded from partial KV reuse. LLAMA_API bool llama_model_supports_partial_kv_reuse(const struct llama_model * model); + // False when full-state seq save/restore cannot work for this context (--swa-compress). + // PARTIAL_ONLY is unaffected. Context property, not a model one. + LLAMA_API bool llama_supports_full_state_io(const struct llama_context * ctx); + LLAMA_API const char * llama_model_arch_string(const struct llama_model * model); // Returns 0 on success diff --git a/src/graphs/build_dflash.cpp b/src/graphs/build_dflash.cpp index 66d9e6116..e7dfd5bba 100644 --- a/src/graphs/build_dflash.cpp +++ b/src/graphs/build_dflash.cpp @@ -207,7 +207,7 @@ ggml_cgraph * llm_build_context::build_dflash() { const int64_t ctx_len = lctx.dflash.visible_cross_ctx > 0 ? (int64_t) lctx.dflash.visible_cross_ctx : std::max(1, (int64_t) cparams.n_ctx - (int64_t) hparams.dflash_block_size); - const int64_t n_kv_total = GGML_PAD(ctx_len + n_tokens, flash_attn ? 256 : 32); + const int64_t n_kv_total = GGML_PAD(ctx_len + n_tokens, (int64_t) llama_kv_cache::get_padding(flash_attn)); GGML_ASSERT(n_embd_head_k == n_embd_head_v); GGML_ASSERT(n_target_features > 0); diff --git a/src/graphs/build_openpangu.cpp b/src/graphs/build_openpangu.cpp index c490e5fb2..e2596ac10 100644 --- a/src/graphs/build_openpangu.cpp +++ b/src/graphs/build_openpangu.cpp @@ -36,10 +36,6 @@ static void openpangu_register_cache_copy( copies[idx].step = step; } -static uint32_t openpangu_kv_cache_pad(const llama_cparams & cparams) { - return cparams.flash_attn ? 256u : 32u; -} - static bool openpangu_idx_score_should_chunk(int64_t n_tokens, int64_t chunk) { return chunk > 0 && n_tokens > 14 && n_tokens > chunk; } @@ -122,24 +118,29 @@ static ggml_tensor * openpangu_cast_gathered_latent_for_cache_type(ggml_context return !ggml_is_quantized(kl->type) && src->type != kl->type ? ggml_cast(ctx, src, kl->type) : src; } -static ggml_tensor * openpangu_build_swa_mask_for_graph(llm_build_context & llm, uint32_t window, bool * windowed) { +static ggml_tensor * openpangu_build_swa_mask_for_graph(llm_build_context & llm, uint32_t window, + bool compacted, bool * windowed) { *windowed = false; - llm.lctx.openpangu_swa_window_view = {}; + llm.lctx.swa_window_view = {}; if (window == 0) { return nullptr; } - const uint32_t pad = openpangu_kv_cache_pad(llm.cparams); - const llama_openpangu_swa_window_view view = - llama_openpangu_calc_swa_window_view(llm.n_kv, llm.n_tokens, window, pad); + const uint32_t pad = llama_kv_cache::get_padding(llm.cparams.flash_attn); + const int64_t live = compacted + ? (int64_t) llm.swa_head - (int64_t) llm.kv_self.sink_rows + llm.n_tokens : 0; + const llama_swa_window_view view = compacted + ? llama_swa_calc_window_view_compact(live, llm.kv_self.sink_rows, llm.n_tokens, window, pad) + : llama_swa_calc_window_view(llm.n_kv, llm.n_tokens, window, pad); if (!view.engaged) { return llm.build_inp_KQ_mask_swa(); } - llm.lctx.openpangu_swa_window_view = { + llm.lctx.swa_window_view = { true, + compacted, llm.n_kv, llm.n_tokens, window, @@ -337,8 +338,9 @@ ggml_tensor * llm_build_context::build_openpangu_attention( kpe_store = openpangu_cast_for_latent_cache_write(ctx0, kpe_store, kl); } ggml_tensor * k_latent = ggml_concat(ctx0, ckv_store, kpe_store, 0); + const int64_t store_row = kv_self.is_compacted(il) ? swa_head : kv_head; ggml_tensor * kl_full = ggml_view_2d(ctx0, kl, kv_lora_rank + n_embd_head_qk_rope, - n_tokens, kl->nb[1], kv_head*kl->nb[1]); + n_tokens, kl->nb[1], store_row*kl->nb[1]); ggml_tensor * cpy_kl = ggml_cpy(ctx0, k_latent, kl_full); openpangu_register_cache_copy(lctx, il, OPENPANGU_COPY_K_CKV, cpy_kl, kl->nb[1]); ggml_build_forward_expand(gf, cpy_kl); @@ -366,7 +368,7 @@ ggml_tensor * llm_build_context::build_openpangu_attention( GGML_ASSERT(topk > 0 && topk <= INT_MAX); dsa_topk = topk; - const uint32_t pad = openpangu_kv_cache_pad(cparams); + const uint32_t pad = llama_kv_cache::get_padding(cparams.flash_attn); const bool is_base_graph = cparams.mtp_op_type == MTP_OP_NONE; const bool dsa_gather_predicate = openpangu_dsa_gather_should_engage(n_kv, n_tokens, topk, pad); @@ -527,11 +529,11 @@ ggml_tensor * llm_build_context::build_openpangu_attention( // ---- latent attention over [sinks ++ cached tokens] (flash_attn is forced off) ---- // Keep sinks and cached tokens separate until after KQ so f16 latent caches do not need // unsupported non-f32 concat along dim1. - const bool use_swa_window = KQ_mask_swa_windowed && lctx.openpangu_swa_window_view.active; + const bool use_swa_window = KQ_mask_swa_windowed && lctx.swa_window_view.active; const bool use_dsa_gather = dsa_gather_engaged && sel_idx != nullptr; const int64_t n_kv_attn = use_dsa_gather ? dsa_topk : - use_swa_window ? lctx.openpangu_swa_window_view.w_view : n_kv; - const int64_t win_off = use_swa_window ? lctx.openpangu_swa_window_view.win_off : 0; + use_swa_window ? lctx.swa_window_view.w_view : n_kv; + const int64_t win_off = use_swa_window ? lctx.swa_window_view.win_off : 0; if (sel_mask) { GGML_ASSERT(!use_swa_window && "openPangu DSA/indexer layers must not use SWA window views"); } @@ -694,7 +696,7 @@ ggml_tensor * llm_build_context::build_openpangu_attention( const bool prefill_gather_chunk = can_prefill_gather && openpangu_dsa_prefill_gather_should_engage(n_kv, n_tokens, c0, tc, dsa_topk, - openpangu_kv_cache_pad(cparams)); + llama_kv_cache::get_padding(cparams.flash_attn)); ggml_tensor * kqv_c = nullptr; if (prefill_gather_chunk) { if (use_fused_attn) { @@ -997,6 +999,13 @@ ggml_cgraph * llm_build_context::build_openpangu() { if (batch.pos && batch.n_tokens > 0) { GGML_ASSERT((llama_pos) kv_head == batch.pos[0] && "openPangu KV cache is position-addressed; kv head must equal the first batch position"); + if (kv_self.any_compacted()) { + GGML_ASSERT(swa_head >= (int32_t) kv_self.sink_rows && + swa_head + n_tokens <= (int32_t) kv_self.size_swa && + "compacted store must fit inside the window region"); + GGML_ASSERT(kv_self.pos_base_swa + (llama_pos) (swa_head - (int32_t) kv_self.sink_rows) == batch.pos[0] && + "compacted row<->position map must agree with the batch position"); + } } const int64_t n_embd_head_k = hparams.n_embd_head_k(0); // 192 @@ -1025,7 +1034,7 @@ ggml_cgraph * llm_build_context::build_openpangu() { // uses hparams.n_swa_mtp when the graph is built with an MTP op type bool KQ_mask_swa_windowed = false; ggml_tensor * KQ_mask = hparams.n_swa_mtp > 0 && hparams.n_swa > 0 - ? openpangu_build_swa_mask_for_graph(*this, hparams.n_swa_mtp, &KQ_mask_swa_windowed) + ? openpangu_build_swa_mask_for_graph(*this, hparams.n_swa_mtp, /* compacted = */ false, &KQ_mask_swa_windowed) : build_inp_KQ_mask(); ggml_tensor * inp_out_ids = n_tokens > 1 ? build_inp_out_ids() : nullptr; lctx.inp_tokens = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, batch.n_tokens); @@ -1164,7 +1173,9 @@ ggml_cgraph * llm_build_context::build_openpangu() { // mask and add the indexer's top-k selection inside the attention builder. Absent // schedule keys (n_swa == 0) keep every layer dense (pre-DSA GGUF fallback). bool KQ_mask_swa_windowed = false; - ggml_tensor * KQ_mask_swa = hparams.n_swa > 0 ? openpangu_build_swa_mask_for_graph(*this, hparams.n_swa, &KQ_mask_swa_windowed) : nullptr; + ggml_tensor * KQ_mask_swa = hparams.n_swa > 0 + ? openpangu_build_swa_mask_for_graph(*this, hparams.n_swa, kv_self.any_compacted(), &KQ_mask_swa_windowed) + : nullptr; lctx.inp_s_seq_qnext = ggml_new_tensor_2d(ctx0, GGML_TYPE_I32, 1, n_tokens); cb(lctx.inp_s_seq_qnext, "inp_s_seq_qnext", -1); ggml_set_input(lctx.inp_s_seq_qnext); @@ -1219,7 +1230,7 @@ ggml_cgraph * llm_build_context::build_openpangu() { cur = llm_build_norm(ctx0, x, hparams, layer.attn_norm, NULL, LLM_NORM_RMS, cb, il); if (il == 0) ggml_set_name(cur, "opg0_attn_norm"); - const bool layer_swa = KQ_mask_swa && hparams.openpangu_window[il] > 0; + const bool layer_swa = KQ_mask_swa && hparams.swa_layers[il]; ggml_tensor * layer_mask = layer_swa ? KQ_mask_swa : KQ_mask; cur = build_openpangu_attention(gf, layer, il, cur, layer_mask, inp_pos, conv_state, seq_qnext, kq_scale, diff --git a/src/llama-build-context.cpp b/src/llama-build-context.cpp index 2f2aaf972..bb97687aa 100644 --- a/src/llama-build-context.cpp +++ b/src/llama-build-context.cpp @@ -69,6 +69,8 @@ llm_build_context::llm_build_context( n_outputs (worst_case ? n_outputs_ > 0 ? n_outputs_ : n_tokens : lctx.n_outputs), n_outputs_enc (worst_case ? n_tokens : lctx.embd_enc.size() / hparams.n_embd), kv_head (worst_case ? (kv_self.recurrent ? 0 : kv_self.size - n_tokens) : kv_self.head), + swa_head (!kv_self.any_compacted() ? kv_head : + worst_case ? (int32_t) (kv_self.size_swa - n_tokens) : (int32_t) kv_self.head_swa), n_ctx_orig (cparams.n_ctx_orig_yarn), flash_attn (cparams.flash_attn), mla_attn (cparams.mla_attn), @@ -107,7 +109,7 @@ void llm_build_context::init() { lctx.inp_KQ_mask = nullptr; lctx.inp_KQ_mask_swa = nullptr; lctx.inp_KQ_mask_swa_win = nullptr; - lctx.openpangu_swa_window_view = {}; + lctx.swa_window_view = {}; lctx.inp_K_shift = nullptr; lctx.inp_mean = nullptr; lctx.inp_cls = nullptr; diff --git a/src/llama-build-context.h b/src/llama-build-context.h index 0dc5fa0c3..87cb09da4 100644 --- a/src/llama-build-context.h +++ b/src/llama-build-context.h @@ -72,6 +72,7 @@ struct llm_build_context { const int32_t n_outputs; const int32_t n_outputs_enc; const int32_t kv_head; // index of where we store new KV data in the cache + const int32_t swa_head; // same, for --swa-compress layers; equals kv_head otherwise const int32_t n_ctx_orig; const bool flash_attn; diff --git a/src/llama-context.h b/src/llama-context.h index 391bd39e8..b3432ea17 100644 --- a/src/llama-context.h +++ b/src/llama-context.h @@ -13,15 +13,15 @@ struct llama_model; #include #include -struct llama_openpangu_swa_window_view { +struct llama_swa_window_view { int64_t w_view = 0; int64_t win_off = 0; bool engaged = false; }; -static inline llama_openpangu_swa_window_view llama_openpangu_calc_swa_window_view( +static inline llama_swa_window_view llama_swa_calc_window_view( int64_t n_kv, int64_t n_tokens, int64_t window, int64_t pad) { - llama_openpangu_swa_window_view result; + llama_swa_window_view result; if (window <= 0 || n_kv <= 0) { result.w_view = n_kv; return result; @@ -35,6 +35,19 @@ static inline llama_openpangu_swa_window_view llama_openpangu_calc_swa_window_vi return result; } +static inline llama_swa_window_view llama_swa_calc_window_view_compact( + int64_t live, int64_t sink_rows, int64_t n_tokens, int64_t window, int64_t pad) { + llama_swa_window_view result; + const int64_t live_padded = pad > 1 ? ((live + pad - 1) / pad) * pad : live; + + const int64_t unpadded = window + pad + n_tokens; + const int64_t overcovered = pad > 1 ? ((unpadded + pad - 1) / pad) * pad : unpadded; + result.w_view = overcovered < live_padded ? overcovered : live_padded; + result.win_off = sink_rows + live_padded - result.w_view; + result.engaged = true; + return result; +} + struct llama_kv_cell { llama_pos pos = -1; llama_pos delta = 0; @@ -57,6 +70,9 @@ struct llama_kv_cell { // ring-buffer of cached KV data struct llama_kv_cache { + // the FA kernels require padding to avoid extra runtime boundary checks + static uint32_t get_padding(bool flash_attn) { return flash_attn ? 256u : 32u; } + bool has_shift = false; bool do_defrag = false; bool do_copy = false; @@ -75,6 +91,28 @@ struct llama_kv_cache { uint32_t size = 0; uint32_t used = 0; // used cells (i.e. at least one seq_id) + std::vector row_count; + + bool any_compacted() const { return !row_count.empty(); } + + uint32_t rows(int il) const { + return row_count.empty() ? size : row_count[il]; + } + + bool is_compacted(int il) const { + return !row_count.empty() && row_count[il] < size; + } + + // rows [sinks|window]; row(pos) = sink_rows + pos - pos_base_swa + // sink_rows == hparams.param_sink_number, which only the openPangu loader sets (0 elsewhere) + uint32_t size_swa = 0; + uint32_t sink_rows = 0; + uint32_t window_swa = 0; + uint32_t head_swa = 0; + llama_pos pos_base_swa = 0; + + uint32_t live_swa() const { return head_swa - sink_rows; } + // computed before each graph build uint32_t n = 0; @@ -567,7 +605,7 @@ struct llama_context { struct ggml_tensor * inp_out_ids; // I32 [n_outputs] struct ggml_tensor * inp_KQ_mask; // F32 [kv_size, n_batch] struct ggml_tensor * inp_KQ_mask_swa; // F32 [kv_size, n_batch] - struct ggml_tensor * inp_KQ_mask_swa_win = nullptr; // F32 [openPangu SWA W_view, n_batch] + struct ggml_tensor * inp_KQ_mask_swa_win = nullptr; // F32 [SWA W_view, n_batch] struct ggml_tensor * inp_K_shift; // I32 [kv_size] struct ggml_tensor * inp_mean; // F32 [n_batch, n_batch] struct ggml_tensor * inp_cls; // I32 [n_batch] @@ -583,15 +621,16 @@ struct llama_context { struct ggml_tensor * inp_mtp_carry = nullptr; // F32 [n_embd, nextn-1] per-head hidden at the last committed position struct ggml_tensor * inp_dsa_sink = nullptr; // F32 [n_kv, n_tokens] per-sequence attention-sink boost for DSA indexer top-k - struct openpangu_swa_window_view_state { + struct swa_window_view_state { bool active = false; + bool compacted = false; int32_t n_kv = 0; int32_t n_tokens = 0; uint32_t window = 0; uint32_t pad = 0; int64_t w_view = 0; int64_t win_off = 0; - } openpangu_swa_window_view; + } swa_window_view; // multi-head MTP chaining state: head k's output row at the last committed position, // written back after each warmup/update decode and fed into the next MTP graph through @@ -601,6 +640,8 @@ struct llama_context { std::vector mtp_carry; bool mtp_carry_pending = false; + std::vector swa_compact_buf; + ggml_backend_t ggml_backend_by_name(const char * name); struct Prev; diff --git a/src/llama-cparams.h b/src/llama-cparams.h index 758ad69ba..5f0997286 100644 --- a/src/llama-cparams.h +++ b/src/llama-cparams.h @@ -45,6 +45,7 @@ struct llama_cparams { bool dsa_indexer_hadamard = true; // apply Walsh-Hadamard rotation to DSA indexer q/k (precision) bool dsa = false; // enable GLM DSA sparse attention (off by default; opt-in via --dsa) bool fused_idx_topk = false; // enable the fused indexer topk op (off by default; opt-in via -fidx or --fused-indexer-topk) + bool swa_compress = false; int dsa_top_k = -1; // DSA top-k override (<0 => use the model's configured indexer_top_k) bool split_mode_graph_scheduling; //bool split_mode_f16; diff --git a/src/llama-dflash.cpp b/src/llama-dflash.cpp index b5f51bf17..1d36dbef5 100644 --- a/src/llama-dflash.cpp +++ b/src/llama-dflash.cpp @@ -62,7 +62,7 @@ bool llama_context::ensure_dflash_kv_cache_tensors(int32_t cross_ctx) { const int32_t target_token_capacity = std::max( 1, std::max((int32_t) model.hparams.dflash_block_size, (int32_t) cparams.n_ubatch)); - const int32_t target_cache_n_kv_total = GGML_PAD(target_cross_ctx + target_token_capacity, cparams.flash_attn ? 256 : 32); + 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; const int64_t n_embd_head_k = model.hparams.n_embd_head_k(0); diff --git a/src/llama-hparams.cpp b/src/llama-hparams.cpp index 365747ea1..6d884d64b 100644 --- a/src/llama-hparams.cpp +++ b/src/llama-hparams.cpp @@ -1178,8 +1178,8 @@ void llm_load_hparams( // DSA/SWA schedule: openpangu.swa_layers lists the sliding-window layer ids and // openpangu.sliding_window_list the per-entry window; the remaining base layers // are DSA (indexer + top-k, no window). The NextN/MTP layers appear in the SWA - // list with their own (larger) window, used by the MTP graphs. Absent keys keep - // every window at 0 = dense fallback (pre-DSA GGUFs keep working). + // list with their own (larger) window, used by the MTP graphs. Absent keys leave + // swa_layers cleared = dense fallback (pre-DSA GGUFs keep working). { std::vector swa_ids, swa_windows; const bool have_ids = ml.get_arr("openpangu.swa_layers", swa_ids, false); @@ -1192,7 +1192,7 @@ void llm_load_hparams( if (il >= hparams.n_layer) { throw std::runtime_error(format("openpangu.swa_layers contains out-of-range layer %u", il)); } - hparams.openpangu_window[il] = swa_windows[i]; + hparams.swa_layers[il] = swa_windows[i] > 0 ? 1 : 0; if (il < n_base) { if (hparams.n_swa != 0 && hparams.n_swa != swa_windows[i]) { throw std::runtime_error("openpangu: non-uniform base sliding windows are not supported"); @@ -1208,6 +1208,13 @@ void llm_load_hparams( } else if (have_ids || have_win) { LLAMA_LOG_WARN("%s: openpangu SWA schedule keys are inconsistent - keeping dense fallback\n", __func__); } + // the graph derives head dims and the MoME conv slot width from layer 0 + if (hparams.n_swa > 0 && + (hparams.n_embd_head_k_swa != hparams.n_embd_head_k_full || + hparams.n_embd_head_v_swa != hparams.n_embd_head_v_full || + hparams.n_rot_swa != hparams.n_rot)) { + throw std::runtime_error("openpangu: per-layer SWA head dimensions are not supported"); + } } diff --git a/src/llama-hparams.h b/src/llama-hparams.h index 5fe94bde2..81b1e24f4 100644 --- a/src/llama-hparams.h +++ b/src/llama-hparams.h @@ -137,11 +137,8 @@ struct llama_hparams { uint32_t mhc_num_stream = 1; uint32_t mhc_recur_norm = 0; uint32_t param_sink_number = 0; - // openPangu DSA/SWA schedule: per-layer sliding window (0 = DSA layer, full causal - // attention over the indexer's top-k selection). The NextN/MTP layers carry their own - // (larger) window, applied when the graph is built with an MTP op type. + // window used by the NextN/MTP layers in place of n_swa uint32_t n_swa_mtp = 0; - std::array openpangu_window = {}; // DeepSeek-V4 uint32_t dsv4_o_group_count = 0; @@ -397,3 +394,32 @@ struct llama_hparams { }; static_assert(std::is_trivially_copyable::value, "llama_hparams must be trivially copyable"); + +// retained window + one u-batch; compaction then fires every C - W tokens. The slack floor keeps a +// small u-batch from compacting every few tokens. +static inline uint32_t llama_swa_compact_window_rows(uint32_t window, uint32_t pad, uint32_t n_ubatch) { + const uint32_t min_slack = 256; + const uint32_t slack = n_ubatch > min_slack ? n_ubatch : min_slack; + const uint32_t unpadded = window + slack; + return pad > 1 ? ((unpadded + pad - 1)/pad)*pad : unpadded; +} + +static inline uint32_t llama_swa_compact_rows(uint32_t window, uint32_t pad, uint32_t n_ubatch, + uint32_t sink_rows) { + return sink_rows + llama_swa_compact_window_rows(window, pad, n_ubatch); +} + +static inline uint32_t llama_kv_layer_rows(const llama_hparams & hparams, int il, uint32_t kv_size, + bool swa_compress, uint32_t n_ubatch, uint32_t pad) { + if (!swa_compress || il < 0 || il >= (int) hparams.n_layer) { + return kv_size; + } + if (il >= (int) (hparams.n_layer - hparams.nextn_predict_layers)) { + return kv_size; + } + if (!hparams.swa_layers[il]) { + return kv_size; + } + const uint32_t rows = llama_swa_compact_rows(hparams.n_swa, pad, n_ubatch, hparams.param_sink_number); + return rows < kv_size ? rows : kv_size; +} diff --git a/src/llama-model.cpp b/src/llama-model.cpp index fdcd75e2b..767da1106 100644 --- a/src/llama-model.cpp +++ b/src/llama-model.cpp @@ -1,4 +1,5 @@ #include "llama-model.h" +#include "llama-context.h" #include "llama-cparams.h" #include @@ -2253,7 +2254,8 @@ llm_tensor llm_tensor_type(llm_arch arch, const std::string & tensor_name, int i return LLM_TENSOR_UNKNOWN; } -size_t llama_model::cache_size(int il, ggml_type type_k, ggml_type type_v, ggml_type idx_type_k, uint32_t kv_size, int mla_attn, int n_seq_max, bool flash_attn) const { +size_t llama_model::cache_size(int il, ggml_type type_k, ggml_type type_v, ggml_type idx_type_k, uint32_t kv_size, int mla_attn, int n_seq_max, bool flash_attn, + bool swa_compress, uint32_t n_ubatch) const { if (il < 0 || il >= hparams.n_layer) return 0; if (hparams.recurrent_layer_arr[il]) { auto state_sots = std::min(std::max(1, n_seq_max), kv_size); @@ -2263,10 +2265,13 @@ size_t llama_model::cache_size(int il, ggml_type type_k, ggml_type type_v, ggml_ // MLA-latent cache: K row [ckv | roped k_pe]. The value-side latent is // rederived from K per graph. DSA layers also cache one indexer key per // position. The recurrent conv slot is constant-size and negligible here. - size_t size = ggml_row_size(type_k, hparams.n_lora_kv + hparams.n_rot) * kv_size; + // openPangu forces flash_attn off, so the pad is always the non-FA one + const uint32_t pad = llama_kv_cache::get_padding(/* flash_attn = */ false); + const uint32_t k_rows = llama_kv_layer_rows(hparams, il, kv_size, swa_compress, n_ubatch, pad); + size_t size = ggml_row_size(type_k, hparams.n_lora_kv + hparams.n_rot) * k_rows; if (hparams.indexer_head_size > 0 && hparams.n_swa > 0 && il < (int) hparams.n_layer - (int) hparams.nextn_predict_layers && - hparams.openpangu_window[il] == 0) { + !hparams.swa_layers[il]) { size += ggml_row_size(idx_type_k, hparams.indexer_head_size) * kv_size; } return size; diff --git a/src/llama-model.h b/src/llama-model.h index 73447913e..1ea835cf2 100644 --- a/src/llama-model.h +++ b/src/llama-model.h @@ -509,6 +509,7 @@ struct llama_model { int n_gpu_layers; bool mtp; // use mtp if is supported by the Model + bool swa_compress = false; // value the cache-size fit was computed with std::vector rpc_servers; std::vector devices; @@ -579,6 +580,13 @@ struct llama_model { return arch == LLM_ARCH_DEEPSEEK2 || arch == LLM_ARCH_GLM_DSA || arch == LLM_ARCH_MISTRAL4; } + // 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 also requires K-only cache + // rows and a single sequence + bool supports_swa_compress() const { + return arch == LLM_ARCH_OPENPANGU; + } + static inline int hadamard_size(int head_size) { if ((head_size & ~(head_size - 1)) == head_size) return head_size; // Note: we do not include 32 as an option because the CUDA Hadamard implementation @@ -599,7 +607,8 @@ struct llama_model { return hadamard_size(hparams.n_embd_head_v(il)); } - size_t cache_size(int il, ggml_type type_k, ggml_type type_v, ggml_type idx_type_k, uint32_t kv_size, int mla_attn, int n_seq_max, bool flash_attn) const; + size_t cache_size(int il, ggml_type type_k, ggml_type type_v, ggml_type idx_type_k, uint32_t kv_size, int mla_attn, int n_seq_max, bool flash_attn, + bool swa_compress = false, uint32_t n_ubatch = 0) const; void set_tensor_overrides(const llama_model_params& params); diff --git a/src/llama.cpp b/src/llama.cpp index b7030db73..e5e228b77 100644 --- a/src/llama.cpp +++ b/src/llama.cpp @@ -562,6 +562,8 @@ struct llama_context::Prev { llama_mtp_op_type mtp_op_type; int32_t mtp_step_idx; int32_t mtp_n_heads; + int64_t swa_w_view; + int64_t swa_win_off; ggml_cgraph * graph; }; @@ -599,6 +601,15 @@ bool llama_context::can_reuse_graph(const llama_batch & u_batch) { mtp_target_ctx != nullptr ? mtp_target_ctx->kv_self : kv_self; if (the_prev->save_per_step_ssm != kv_self_used.save_per_step_ssm || the_prev->per_step_max_allocated != kv_self_used.ckpt.per_step_max_allocated) return false; + if (kv_self_used.any_compacted()) { + const auto view = llama_swa_calc_window_view_compact( + (int64_t) kv_self_used.live_swa() + u_batch.n_tokens, kv_self_used.sink_rows, + u_batch.n_tokens, kv_self_used.window_swa, + llama_kv_cache::get_padding(cparams.flash_attn)); + if (view.w_view != the_prev->swa_w_view || view.win_off != the_prev->swa_win_off) { + return false; + } + } bool result = u_batch.all_seq_id == the_prev->all_seq_id && kv_self_used.head > 0 && kv_self_used.n == the_prev->n_kv && @@ -641,8 +652,10 @@ bool llama_context::update_cache_copies() { }; if (model.arch == LLM_ARCH_OPENPANGU) { auto & copies = cparams.mtp_op_type == MTP_OP_NONE ? openpangu_cache_copies : openpangu_cache_copies_mtp; + GGML_ASSERT(copies.size() == model.hparams.n_layer); bool any = false; - for (auto & c : copies) { + for (size_t idx = 0; idx < copies.size(); ++idx) { + auto & c = copies[idx]; if (!c.cpy) { continue; } @@ -650,7 +663,8 @@ bool llama_context::update_cache_copies() { if (c.cpy->op != GGML_OP_CPY || c.cpy->view_src == nullptr || c.cpy->src[1] == nullptr) { return false; } - c.cpy->view_offs = kv_self.head*c.step; + const uint32_t head = kv_self.is_compacted((int) idx) ? kv_self.head_swa : kv_self.head; + c.cpy->view_offs = head*c.step; c.cpy->src[1]->data = (char *)c.cpy->view_src->data + c.cpy->view_offs; c.cpy->data = c.cpy->src[1]->data; } @@ -896,15 +910,15 @@ static int llama_openpangu_chunked_graph_nodes(const llama_model & model, const hparams.n_layer - hparams.nextn_predict_layers : hparams.n_layer; const int64_t n_heads = hparams.n_head(0); const int64_t n_sinks = hparams.param_sink_number; - const int64_t pad = cparams.flash_attn ? 256 : 32; + const int64_t pad = llama_kv_cache::get_padding(cparams.flash_attn); const int64_t idx_chunks = idx_chunk > 0 && n_tokens > idx_chunk ? llama_div_ceil_i64(n_tokens, idx_chunk) : 0; int64_t extra_nodes = 0; for (int64_t il = 0; il < n_layer_base; ++il) { - const bool is_swa_layer = hparams.n_swa > 0 && hparams.openpangu_window[il] > 0; - const bool is_dsa_layer = has_dsa_indexer && hparams.openpangu_window[il] == 0; + const bool is_swa_layer = hparams.n_swa > 0 && hparams.swa_layers[il]; + const bool is_dsa_layer = has_dsa_indexer && !hparams.swa_layers[il]; const int64_t n_kv_eff = !is_swa_layer ? n_kv : - llama_openpangu_calc_swa_window_view(n_kv, n_tokens, hparams.openpangu_window[il], pad).w_view; + llama_swa_calc_window_view(n_kv, n_tokens, hparams.n_swa, pad).w_view; const double full_kq_bytes = (double) (n_kv_eff + n_sinks) * (double) n_heads * (double) n_tokens * (double) sizeof(float); const bool att_chunks = @@ -1099,6 +1113,38 @@ static bool llama_kv_cache_init( cache.size = kv_size; cache.used = 0; + cache.row_count.clear(); + if (cparams.swa_compress && !model.supports_swa_compress()) { + LLAMA_LOG_WARN("%s: --swa-compress is not implemented for this model; ignoring\n", __func__); + } else if (cparams.swa_compress) { + std::vector plan((size_t) hparams.n_layer, kv_size); + bool any = false; + for (int il = 0; il < (int) hparams.n_layer; ++il) { + plan[il] = llama_kv_layer_rows(hparams, il, kv_size, true, cparams.n_ubatch, + llama_kv_cache::get_padding(cparams.flash_attn)); + any = any || plan[il] < kv_size; + } + if (any) { + cache.row_count = std::move(plan); + uint32_t rows_compacted = 0; + for (int il = 0; il < (int) hparams.n_layer; ++il) { + if (cache.row_count[il] >= kv_size) { + continue; + } + GGML_ASSERT((rows_compacted == 0 || rows_compacted == cache.row_count[il]) && + "compacted sliding-window layers must share one row count"); + rows_compacted = cache.row_count[il]; + } + cache.size_swa = rows_compacted; + cache.sink_rows = hparams.param_sink_number; + cache.window_swa = hparams.n_swa; + 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__); + } + } + cache.type_k = type_k; cache.type_v = type_v; @@ -1130,6 +1176,12 @@ static bool llama_kv_cache_init( replicate_mla = true; } + if (cache.any_compacted() && (split_cache || replicate_mla)) { + LLAMA_LOG_ERROR("%s: --swa-compress is not supported with a split or replicated KV cache " + "(split mode graph/attn); run without --swa-compress or with a single device\n", __func__); + return false; + } + // count used buffer types std::map buft_layer_count; if (offload) { @@ -1375,7 +1427,7 @@ static bool llama_kv_cache_init( // MLA-latent cache: k_l holds [ckv_norm 512 | roped k_pe 64] per position. // The value-side latent is rederived from k_l per graph; no persistent V store. const int64_t n_lat = (int64_t) hparams.n_lora_kv + hparams.n_rot; // 576 - k = ggml_new_tensor_2d(ctx, this_type_k, n_lat, kv_size); + k = ggml_new_tensor_2d(ctx, this_type_k, n_lat, cache.rows(i)); } else if (is_dsv4_k_only) { k = ggml_new_tensor_2d(ctx, this_type_k, n_embd_head_k, n_head_kv*kv_size); } else { @@ -1401,10 +1453,10 @@ static bool llama_kv_cache_init( cache.s_l[i] = s_conv; cache.s_l_position_strict = true; - // DSA layers (window == 0, indexer present) also cache the per-position + // DSA layers (no window, indexer present) also cache the per-position // 128-d indexer key. Position-indexed like everything else in the cache, so // the same rollback invariant applies: committed columns never change. - if (has_openpangu_dsa_indexer && i < n_mtp_first_layer && hparams.openpangu_window[i] == 0) { + if (has_openpangu_dsa_indexer && i < n_mtp_first_layer && !hparams.swa_layers[i]) { ggml_tensor * idxk = ggml_new_tensor_2d(ctx, idx_type_k, hparams.indexer_head_size, kv_size); ggml_format_name(idxk, "cache_kr_l%d", i); cache.kr_l[i] = idxk; @@ -1614,6 +1666,48 @@ static bool llama_kv_cache_find_slot( return true; } +static void llama_kv_cache_compact_swa(struct llama_context & lctx, uint32_t n_tokens) { + llama_kv_cache & cache = lctx.kv_self; + std::vector & scratch = lctx.swa_compact_buf; + if (!cache.any_compacted()) { + return; + } + + const uint32_t W = cache.window_swa; + const uint32_t C = cache.size_swa - cache.sink_rows; + GGML_ASSERT(n_tokens <= C); + + if (cache.live_swa() + n_tokens <= C) { + return; + } + + const uint32_t live = cache.live_swa(); + GGML_ASSERT(live >= W && "the retained window must lie inside the live region"); + + // llama_graph_compute submits asynchronously and never waits + ggml_backend_sched_synchronize(lctx.sched); + + const uint32_t src_row = cache.sink_rows + live - W; + const uint32_t dst_row = cache.sink_rows; + + for (size_t il = 0; il < cache.k_l.size(); ++il) { + if (!cache.is_compacted((int) il) || cache.k_l[il] == nullptr) { + continue; + } + ggml_tensor * kl = cache.k_l[il]; + const size_t stride = kl->nb[1]; + const size_t nbytes = (size_t) W * stride; + if (scratch.size() < nbytes) { + scratch.resize(nbytes); + } + ggml_backend_tensor_get(kl, scratch.data(), (size_t) src_row*stride, nbytes); + ggml_backend_tensor_set(kl, scratch.data(), (size_t) dst_row*stride, nbytes); + } + + cache.pos_base_swa += (llama_pos) (live - W); + cache.head_swa = cache.sink_rows + W; +} + // find how many cells are currently in use static uint32_t llama_kv_cache_cell_max(const struct llama_kv_cache & cache) { for (uint32_t i = cache.size; i > 0; --i) { @@ -2085,6 +2179,8 @@ static void llama_kv_cache_clear(struct llama_kv_cache & cache) { } cache.head = 0; cache.used = 0; + cache.head_swa = cache.sink_rows; + cache.pos_base_swa = 0; for (auto & buf : cache.bufs) { ggml_backend_buffer_clear(buf, 0); @@ -2120,6 +2216,39 @@ static bool llama_kv_cache_seq_rm( } } + // a compacted layer holds one contiguous range, so a refusal must leave the cache untouched + bool compact_apply = false; + uint32_t compact_head = 0; + llama_pos compact_base = 0; + // seq_id > 0 owns no compacted rows here (single-sequence arch), so it must not move the map + if (cache.any_compacted() && seq_id <= 0 && p1 > p0) { + const llama_pos cur_end = cache.pos_base_swa + (llama_pos) cache.live_swa(); + GGML_ASSERT((cache.live_swa() > 0 || cache.pos_base_swa == 0) && + "compacted window cannot be empty above position 0"); + if (p0 < cur_end && p1 > cache.pos_base_swa) { + if (p1 < cur_end) { + LLAMA_LOG_ERROR("%s: --swa-compress cannot remove the interior range [%d, %d): " + "compacted layers hold one contiguous range\n", __func__, p0, p1); + return false; + } + if (p0 == 0) { + compact_apply = true; + compact_head = cache.sink_rows; + compact_base = 0; + } else if (cache.pos_base_swa == 0 || + p0 >= cache.pos_base_swa + (llama_pos) cache.window_swa) { + compact_apply = true; + compact_head = cache.sink_rows + (uint32_t) (p0 - cache.pos_base_swa); + compact_base = cache.pos_base_swa; + } else { + LLAMA_LOG_ERROR("%s: --swa-compress cannot rewind to position %d: the compacted " + "window starts at %d and would lose the %u positions before %d\n", + __func__, p0, cache.pos_base_swa, cache.window_swa, p0); + return false; + } + } + } + const bool has_qnext_state = llama_kv_has_qnext_state_storage(cache); for (uint32_t i = 0; i < cache.size; ++i) { @@ -2147,6 +2276,11 @@ static bool llama_kv_cache_seq_rm( // If we freed up a slot, set head to it so searching can start there. if (new_head != cache.size && new_head < cache.head) cache.head = new_head; + if (compact_apply) { + cache.head_swa = compact_head; + cache.pos_base_swa = compact_base; + } + return true; } @@ -2355,11 +2489,6 @@ static void llama_kv_cache_defrag(struct llama_kv_cache & cache) { cache.do_defrag = true; } -static uint32_t llama_kv_cache_get_padding(const struct llama_cparams & cparams) { - // the FA kernels require padding to avoid extra runtime boundary checks - return cparams.flash_attn ? 256u : 32u; -} - // // model loading and saving // @@ -3682,7 +3811,7 @@ struct expert_tensors { static std::pair, double> get_layer_sizes(const llama_model_loader & ml, const llama_model & model, ggml_type cache_type_k, ggml_type cache_type_v, ggml_type idx_type_k, uint32_t max_ctx_size, int mla_attn, int n_seq_max, int n_ubatch, - int amb, int worst_case_tokens, bool flash_attn, + int amb, int worst_case_tokens, bool flash_attn, bool swa_compress, std::vector & experts) { int n_layer = model.hparams.n_layer; std::vector result(n_layer+1, 0); @@ -3894,7 +4023,8 @@ static std::pair, double> get_layer_sizes(const llama_model_ LLAMA_LOG_INFO("------------------- Layer sizes:\n"); double tot_model = 0, tot_cache = 0, max_compute = 0; for (int il = 0; il < n_layer; ++il) { - auto kv_size = model.cache_size(il, cache_type_k, cache_type_v, idx_type_k, max_ctx_size, mla_attn, n_seq_max, flash_attn); + auto kv_size = model.cache_size(il, cache_type_k, cache_type_v, idx_type_k, max_ctx_size, mla_attn, n_seq_max, flash_attn, + swa_compress, (uint32_t) n_ubatch); LLAMA_LOG_INFO("Layer %2d: %9.2f, %9.2f, %9.2f %9.2f MiB\n", il, result[il]/1024./1024., kv_size/1024./1024., (result[il] + kv_size)/1024./1024., compute[il]/1024./1024.); max_compute = std::max(max_compute, compute[il]); tot_model += result[il]; @@ -3935,6 +4065,7 @@ static bool llm_load_tensors( const int * fit_margin_array, int worst_case_tokens, bool flash_attn, + bool swa_compress, bool use_mlock, bool validate_quants, bool mtp, @@ -4032,6 +4163,7 @@ static bool llm_load_tensors( model.max_gpu = max_gpu; model.n_gpu_layers = n_gpu_layers; model.mtp = mtp; + model.swa_compress = swa_compress; size_t mem_margin = fit_margin > 0 ? size_t(fit_margin)*1024*1024 : k_default_mem_margin; auto get_mem_margin = [mem_margin, fit_margin_array, n_gpu = int(model.devices.size()), func = __func__] (int gpu) { @@ -4133,7 +4265,7 @@ static bool llm_load_tensors( if (device_count > 0 && !model.devices.empty()) { std::vector experts; auto [layer_sizes, max_compute] = get_layer_sizes(ml, model, cache_type_k, cache_type_v, idx_type_k, max_ctx_size, mla_attn, n_seq_max, n_ubatch, - amb, worst_case_tokens, flash_attn, experts); + amb, worst_case_tokens, flash_attn, swa_compress, experts); size_t required_mem = 0; for (int i = 0; i <= n_layer; ++i) { required_mem += layer_sizes[i]; @@ -4760,7 +4892,7 @@ static int llama_model_load(const std::string & fname, llama_model & model, llam ml, model, params.n_gpu_layers, params.mla, params.split_mode, params.main_gpu, params.max_gpu, params.tensor_split, params.type_k, params.type_v, params.idx_type_k, params.extra_output_type, params.max_ctx_size, params.n_seq_max, params.n_ubatch, params.amb, params.fit_margin, params.fit_margin_array, - params.worst_graph_tokens, params.flash_attn, + params.worst_graph_tokens, params.flash_attn, params.swa_compress, params.use_mlock, params.validate_quants, params.mtp, params.fit, params.dry_run, params.progress_callback, params.progress_callback_user_data )) { @@ -5068,40 +5200,49 @@ static void llama_set_inputs(llama_context & lctx, const llama_batch & batch) { } if (data_swa_win || data_swa_win_f16) { - const auto & built = lctx.openpangu_swa_window_view; - const uint32_t pad = llama_kv_cache_get_padding(cparams); - const llama_openpangu_swa_window_view view = - llama_openpangu_calc_swa_window_view(n_kv, n_tokens, built.window, pad); - GGML_ASSERT(lctx.model.arch == LLM_ARCH_OPENPANGU && - "windowed SWA mask input is only valid for OpenPangu"); + const auto & built = lctx.swa_window_view; + const uint32_t pad = llama_kv_cache::get_padding(cparams.flash_attn); + const int64_t live = built.compacted + ? (int64_t) mask_kv_self.live_swa() + n_tokens : 0; + const llama_swa_window_view view = built.compacted + ? llama_swa_calc_window_view_compact(live, mask_kv_self.sink_rows, + n_tokens, built.window, pad) + : llama_swa_calc_window_view(n_kv, n_tokens, built.window, pad); GGML_ASSERT(built.active && view.engaged && - "openPangu SWA window view must be engaged when KQ_mask_swa_win is present"); + "SWA window view must be engaged when KQ_mask_swa_win is present"); GGML_ASSERT(built.n_kv == n_kv && built.n_tokens == n_tokens && built.pad == pad && built.w_view == view.w_view && built.win_off == view.win_off && - "openPangu SWA window view reuse-key mismatch"); + "SWA window view reuse-key mismatch"); const int64_t W_view = built.w_view; const int64_t win_off = built.win_off; + const bool compacted = built.compacted; + const int64_t row_base = mask_kv_self.sink_rows; + const llama_pos pos_base = mask_kv_self.pos_base_swa; for (int j = 0; j < n_tokens; ++j) { const llama_pos pos = batch.pos[j]; const llama_seq_id seq_id = batch.seq_id[j][0]; for (int64_t c = 0; c < W_view; ++c) { const int64_t i = win_off + c; + const llama_pos cell_pos = compacted + ? pos_base + (llama_pos) (i - row_base) : mask_kv_self.cells[i].pos; + const bool in_seq = compacted + ? cell_pos >= pos_base : mask_kv_self.cells[i].has_seq_id(seq_id); float f; - if (!mask_kv_self.cells[i].has_seq_id(seq_id) || mask_kv_self.cells[i].pos > pos) { + if (!in_seq || cell_pos > pos) { f = -INFINITY; } else { - f = hparams.use_alibi ? -std::abs(mask_kv_self.cells[i].pos - pos) : 0.0f; + f = hparams.use_alibi ? -std::abs(cell_pos - pos) : 0.0f; } if (f > -INFINITY) { if (hparams.n_attn_chunk) { llama_pos pos_chunk_start = (pos / hparams.n_attn_chunk) * hparams.n_attn_chunk; - if (mask_kv_self.cells[i].pos < pos_chunk_start || pos < pos_chunk_start) { + if (cell_pos < pos_chunk_start || pos < pos_chunk_start) { f = -INFINITY; } - } else if (pos - mask_kv_self.cells[i].pos >= (int32_t) built.window) { + } else if (pos - cell_pos >= (int32_t) built.window) { f = -INFINITY; } } @@ -6124,10 +6265,13 @@ static int llama_decode_internal( // a heuristic, to avoid attending the full cache if it is not yet utilized // after enough generations, the benefit from this heuristic disappears // if we start defragmenting the cache, the benefit from this will be more important - const uint32_t pad = llama_kv_cache_get_padding(cparams); + const uint32_t pad = llama_kv_cache::get_padding(cparams.flash_attn); auto max_cell = llama_kv_cache_cell_max(kv_self, pad); kv_self.n = std::min(kv_self.size, std::max(pad, GGML_PAD(max_cell, pad))); } + + // must run before can_reuse_graph() + llama_kv_cache_compact_swa(lctx, u_batch.n_tokens); } #if IK_PRINT_TIMING @@ -6175,7 +6319,9 @@ static int llama_decode_internal( (int)u_batch.all_seq_id, (int)lctx.n_outputs, (int)kv_self_used.n, (int)u_batch.n_tokens, kv_self_used.save_per_step_ssm, kv_self_used.ckpt.per_step_max_allocated, - cparams.mtp_op_type, lctx.mtp_step_idx, lctx.mtp_n_heads, gf}); + cparams.mtp_op_type, lctx.mtp_step_idx, lctx.mtp_n_heads, + lctx.swa_window_view.w_view, + lctx.swa_window_view.win_off, gf}); } } else { //printf("Reusing graph with type = %d, n_kv = %d, n_tokens = %d\n", cparams.mtp_op_type, (int)prev->n_kv, (int)prev->n_tokens); @@ -6288,9 +6434,14 @@ static int llama_decode_internal( reset_previous = true; } kv_self.head += n_tokens; + if (kv_self.any_compacted()) { + kv_self.head_swa += n_tokens; + GGML_ASSERT(kv_self.head_swa <= kv_self.size_swa); + } // Ensure kv cache head points to a valid index. if (kv_self.head >= kv_self.size) { + // only the generic head wraps; the compacted rows still hold the live window kv_self.head = 0; } } @@ -7263,6 +7414,7 @@ struct llama_model_params llama_model_default_params() { /*.dry_run =*/ false, /*.flash_attn =*/ true, /*.defer_experts =*/ false, + /*.swa_compress =*/ false, }; #ifdef GGML_USE_METAL @@ -7324,6 +7476,7 @@ struct llama_context_params llama_context_default_params() { /*.graph_reuse =*/ true, /*.dsa =*/ false, /*.fused_idx_topk =*/ true, + /*.swa_compress =*/ false, /*.dsa_top_k =*/ -1, /*.min_experts =*/ -1, /*.thtesh_experts =*/ 0.0f, @@ -7813,7 +7966,16 @@ struct llama_context * llama_init_from_model( cparams.graph_reuse = params.graph_reuse; cparams.dsa = params.dsa; cparams.fused_idx_topk = params.fused_idx_topk; + cparams.swa_compress = params.swa_compress; cparams.dsa_top_k = params.dsa_top_k; + + if (cparams.swa_compress != model->swa_compress) { + LLAMA_LOG_ERROR("%s: swa_compress differs between llama_model_params (%d) and llama_context_params (%d); " + "the cache-size fit would not match the allocation\n", + __func__, (int) model->swa_compress, (int) cparams.swa_compress); + llama_free(ctx); + return nullptr; + } // The DSA lightning indexer is built only in the layer-mode (non-TP) attention path. Under // -sm graph / -sm attn the model runs the tensor-parallel attention path, which has no indexer, // so --dsa would silently run dense MLA. Warn instead of degrading silently. @@ -7857,7 +8019,7 @@ struct llama_context * llama_init_from_model( cparams.rope_freq_scale = params.rope_freq_scale == 0.0f ? hparams.rope_freq_scale_train : params.rope_freq_scale; // 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)); + cparams.n_ctx = GGML_PAD(cparams.n_ctx, llama_kv_cache::get_padding(cparams.flash_attn)); // 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; @@ -8203,6 +8365,27 @@ struct llama_context * llama_init_from_model( } } + if (ctx->kv_self.any_compacted()) { + const auto & kv = ctx->kv_self; + uint32_t n_compacted = 0; + uint64_t rows_compacted = 0, rows_dense_equiv = 0; + uint32_t rows_min = UINT32_MAX, rows_max = 0; + for (int il = 0; il < (int) kv.row_count.size(); ++il) { + if (!kv.is_compacted(il)) continue; + const uint32_t r = kv.rows(il); + n_compacted++; + rows_compacted += r; + rows_dense_equiv += kv.size; + rows_min = std::min(rows_min, r); + rows_max = std::max(rows_max, r); + } + LLAMA_LOG_INFO("%s: SWA compress = %u of %d layers compacted, %u-%u rows each (dense %u), window %u, n_ubatch %u\n", + __func__, n_compacted, (int) kv.row_count.size(), rows_min, rows_max, + kv.size, hparams.n_swa, cparams.n_ubatch); + LLAMA_LOG_INFO("%s: SWA compress = %.2fx fewer rows on compacted layers\n", + __func__, rows_compacted ? (double) rows_dense_equiv / (double) rows_compacted : 0.0); + } + if (memory_size_k + memory_size_v) { if (cparams.mla_attn != 0 && !cparams.flash_attn) { LLAMA_LOG_INFO("%s: KV self size = %7.2f MiB, c^KV (%s): %7.2f MiB, kv^T (%s): %7.2f MiB\n", __func__, @@ -8230,6 +8413,19 @@ struct llama_context * llama_init_from_model( } } + if (ctx->kv_self.any_compacted() && cparams.mtp) { + LLAMA_LOG_ERROR("%s: --swa-compress is not supported together with MTP speculative decoding\n", __func__); + llama_free(ctx); + return nullptr; + } + + if (ctx->kv_self.any_compacted()) { + LLAMA_LOG_INFO("%s: --swa-compress: %u rows per compacted layer " + "(%u sink + %u window, window %u)\n", __func__, + ctx->kv_self.size_swa, ctx->kv_self.sink_rows, + ctx->kv_self.size_swa - ctx->kv_self.sink_rows, ctx->kv_self.window_swa); + } + // graph outputs buffer { // resized during inference when a batch uses more outputs @@ -8398,6 +8594,10 @@ void llama_free(struct llama_context * ctx) { delete ctx; } +bool llama_supports_full_state_io(const struct llama_context * ctx) { + return ctx != nullptr && !ctx->kv_self.any_compacted(); +} + const struct llama_vocab* llama_model_get_vocab(const struct llama_model* model) { return &model->vocab; } @@ -10767,6 +10967,13 @@ static bool llama_state_io_supported( const char * func, llama_state_seq_flags flags = 0, llama_seq_id seq_id = -1) { + if (ctx->kv_self.any_compacted() && flags == 0) { + LLAMA_LOG_ERROR("%s: full state save/restore is not supported with --swa-compress " + "(state I/O addresses cache rows by cell index; compacted layers hold a " + "translated subset)\n", func); + return false; + } + if (ctx->model.arch == LLM_ARCH_OPENPANGU) { if (seq_id >= 0 && llama_kv_qnext_seq_id_in_range(ctx->kv_self, seq_id) &&