openpangu: opt-in compacted sliding-window KV cache (--swa-compress) (#2253)

* openpangu: opt-in compacted sliding-window KV cache (--swa-compress)

* openpangu: shrink the compacted window and drop the zero fill

* openpangu: correct the --swa-compress state I/O refusal message

---------

Co-authored-by: Joel Farthing <262452229+joelfarthing@users.noreply.github.com>

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;

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;

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 = rewound
-                const size_t n = llama_state_seq_set_data(ctx, it->data.data(), checkpoint_size, slot.id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY);
+                    ? 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");

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

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<int64_t>(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);

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 uint32_t pad = llama_kv_cache::get_padding(llm.cparams.flash_attn);
-    const llama_openpangu_swa_window_view view =
+    const int64_t live = compacted
-        llama_openpangu_calc_swa_window_view(llm.n_kv, llm.n_tokens, window, pad);
+        ? (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;
+                                use_swa_window ? lctx.swa_window_view.w_view  : n_kv;
-    const int64_t win_off     = use_swa_window ? lctx.openpangu_swa_window_view.win_off : 0;
+    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,

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;

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;

src/llama-context.h

@@ -13,15 +13,15 @@ struct llama_model;
 #include <set>
 #include <memory>
 
-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<uint32_t> 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<float> mtp_carry;
     bool mtp_carry_pending = false;
 
+    std::vector<uint8_t> swa_compact_buf;
+
     ggml_backend_t ggml_backend_by_name(const char * name);
 
     struct Prev;

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;

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<int32_t>(
             1,
             std::max<int32_t>((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);

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
+                // list with their own (larger) window, used by the MTP graphs. Absent keys leave
-                // every window at 0 = dense fallback (pre-DSA GGUFs keep working).
+                // swa_layers cleared = dense fallback (pre-DSA GGUFs keep working).
                 {
                     std::vector<uint32_t> 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");
+                    }
                 }
 
 

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
+    // window used by the NextN/MTP layers in place of n_swa
-    // 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.
     uint32_t n_swa_mtp = 0;
-    std::array<uint32_t, LLAMA_MAX_LAYERS> openpangu_window = {};
 
     // DeepSeek-V4
     uint32_t dsv4_o_group_count     = 0;
@@ -397,3 +394,32 @@ struct llama_hparams {
 };
 
 static_assert(std::is_trivially_copyable<llama_hparams>::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;
+}

src/llama-model.cpp

@@ -1,4 +1,5 @@
 #include "llama-model.h"
+#include "llama-context.h"
 #include "llama-cparams.h"
 
 #include <map>
@@ -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<uint32_t>(std::max<uint32_t>(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;

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_device> rpc_servers;
     std::vector<int32_t> 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);
 

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_swa_layer = hparams.n_swa > 0 && hparams.swa_layers[il];
-        const bool is_dsa_layer = has_dsa_indexer && hparams.openpangu_window[il] == 0;
+        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<uint32_t> 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<ggml_backend_buffer_type_t, int> 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<uint8_t> & 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<std::vector<double>, 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<expert_tensors> & experts) {
     int n_layer = model.hparams.n_layer;
     std::vector<double> result(n_layer+1, 0);
@@ -3894,7 +4023,8 @@ static std::pair<std::vector<double>, 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<expert_tensors> 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 auto & built = lctx.swa_window_view;
-                const uint32_t pad = llama_kv_cache_get_padding(cparams);
+                const uint32_t pad = llama_kv_cache::get_padding(cparams.flash_attn);
-                const llama_openpangu_swa_window_view view =
+                const int64_t live = built.compacted
-                    llama_openpangu_calc_swa_window_view(n_kv, n_tokens, built.window, pad);
+                    ? (int64_t) mask_kv_self.live_swa() + n_tokens : 0;
-                GGML_ASSERT(lctx.model.arch == LLM_ARCH_OPENPANGU &&
+                const llama_swa_window_view view = built.compacted
-                        "windowed SWA mask input is only valid for OpenPangu");
+                    ? 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) &&