deepseek4: compacted sliding-window KV cache (--swa-compress) (#2266)

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

examples/server/server-context.cpp

@@ -2952,6 +2952,10 @@ void server_context::process_single_task(server_task&& task) {
             send_error(task, "slot save is unsupported for openPangu because per-sequence file state is not implemented", ERROR_TYPE_NOT_SUPPORTED);
             break;
         }
+        if (!llama_supports_full_state_io(ctx)) {
+            send_error(task, "slot save is unsupported with --swa-compress because file-session state is not implemented for compacted contexts", ERROR_TYPE_NOT_SUPPORTED);
+            break;
+        }
 
         const size_t token_count = slot->cache_tokens.size();
         const int64_t t_start = ggml_time_us();
@@ -2995,6 +2999,10 @@ void server_context::process_single_task(server_task&& task) {
             queue_tasks.defer(std::move(task));
             break;
         }
+        if (!llama_supports_full_state_io(ctx)) {
+            send_error(task, "slot restore is unsupported with --swa-compress because file-session state is not implemented for compacted contexts", ERROR_TYPE_NOT_SUPPORTED);
+            break;
+        }
         const int64_t t_start = ggml_time_us();
 
         std::string filename = task.data.at("filename");

include/llama.h

@@ -724,7 +724,7 @@ 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);
 
-    // true for every non-null context; none is excluded from full-state seq save/restore
+    // false when the context cannot serialize whole-context or file-session state (--swa-compress); per-sequence buffer state is unaffected
     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);

src/graphs/build_deepseek4.cpp

@@ -946,7 +946,7 @@ static void ds4_build_comp(ggml_tensor * cur, llm_build_context & llm, ggml_cont
 static ggml_tensor * ds4_attention(ggml_cgraph * gf, ggml_context * ctx0, llm_build_context & llm, ggml_tensor * inpL,
         ggml_tensor ** append_csa_state, ggml_tensor ** append_csa_score,
         ggml_tensor ** append_lid_state, ggml_tensor ** append_lid_score,
-        ggml_tensor * inp_pos, ggml_tensor * KQ_mask, int il) {
+        ggml_tensor * inp_pos, ggml_tensor * KQ_mask, ggml_tensor * KQ_mask_swa_win, int il) {
 
     ggml_tensor * residual = inpL;
     ggml_tensor * post = nullptr;
@@ -1072,8 +1072,19 @@ static ggml_tensor * ds4_attention(ggml_cgraph * gf, ggml_context * ctx0, llm_bu
         llm.llm_build_kv_store(lctx, ctx0, hparams, cparams, kv_self, gf, nullptr, kv, n_tokens, llm.kv_head, cb, il);
     }
 
+    const bool raw_compacted = kv_self.is_compacted((int) il);
+
     ggml_tensor * raw_k = nullptr;
-    if (hparams.n_head_kv(il) == 1 && lctx.dsv4.inputs.raw_k_read_idxs != nullptr) {
+    ggml_tensor * raw_mask = nullptr;
+    if (raw_compacted) {
+        // live window rows [win_off, win_off + w_view); KQ_mask_swa_win is that view's column-exact mask
+        GGML_ASSERT(hparams.n_head_kv(il) == 1 && KQ_mask_swa_win != nullptr && lctx.swa_window_view.active);
+        const size_t row_size = ggml_row_size(kv_self.k_l[il]->type, n_embd_head);
+        raw_k = ggml_view_3d(ctx0, kv_self.k_l[il],
+                n_embd_head, 1, lctx.swa_window_view.w_view,
+                row_size, row_size, row_size*(size_t) lctx.swa_window_view.win_off);
+        raw_mask = KQ_mask_swa_win;
+    } else if (hparams.n_head_kv(il) == 1 && lctx.dsv4.inputs.raw_k_read_idxs != nullptr) {
         raw_k = dsv4_raw_get_k(&lctx, ctx0, kv_self.k_l[il], lctx.dsv4.inputs.raw_k_read_idxs, n_embd_head, cb, il);
     }
     if (raw_k == nullptr) {
@@ -1085,22 +1096,26 @@ static ggml_tensor * ds4_attention(ggml_cgraph * gf, ggml_context * ctx0, llm_bu
     }
     cb(raw_k, "raw_k", il);
 
-    const int64_t raw_kq_n_kv = raw_k != nullptr && lctx.dsv4.raw.n_kv > 0
+    const int64_t raw_kq_n_kv = raw_compacted ? lctx.swa_window_view.w_view
+        : raw_k != nullptr && lctx.dsv4.raw.n_kv > 0
         ? lctx.dsv4.raw.n_kv
         : (raw_k != nullptr ? raw_k->ne[2] * raw_k->ne[3] : n_kv);
-    const int64_t raw_attn_n_kv = raw_kq_n_kv > 0 ? std::max<int64_t>(256, GGML_PAD(raw_kq_n_kv, 256)) : raw_kq_n_kv;
+    const int64_t raw_attn_n_kv = raw_compacted ? lctx.swa_window_view.w_view
-    if (raw_k != nullptr && raw_k->ne[3] == 1) {
+        : raw_kq_n_kv > 0 ? std::max<int64_t>(256, GGML_PAD(raw_kq_n_kv, 256)) : raw_kq_n_kv;
+    if (!raw_compacted && raw_k->ne[3] == 1) {
         raw_k = dsv4_pad_raw_k_to(ctx0, raw_k, raw_attn_n_kv);
     }
-    ggml_tensor * raw_mask = dsv4_build_raw_mask_view(ctx0, KQ_mask,
+    if (raw_mask == nullptr) {
-            lctx.dsv4.inputs.raw_k_read_idxs, raw_kq_n_kv, n_tokens, raw_k->ne[3], cb, il);
+        raw_mask = dsv4_build_raw_mask_view(ctx0, KQ_mask,
-    cb(raw_mask, "raw_mask_view", il);
+                lctx.dsv4.inputs.raw_k_read_idxs, raw_kq_n_kv, n_tokens, raw_k->ne[3], cb, il);
-    raw_mask = dsv4_pad_mask_tokens(ctx0, raw_mask, n_tokens);
+        cb(raw_mask, "raw_mask_view", il);
-    raw_mask = dsv4_pad_raw_mask_to(ctx0, raw_mask, raw_attn_n_kv, n_tokens);
+        raw_mask = dsv4_pad_mask_tokens(ctx0, raw_mask, n_tokens);
+        raw_mask = dsv4_pad_raw_mask_to(ctx0, raw_mask, raw_attn_n_kv, n_tokens);
+    }
     cb(raw_mask, "dsv4_raw_mask_padded", il);
     ggml_tensor * attn = nullptr;
 
-    if (hparams.n_swa > 0) {
+    if (hparams.n_swa > 0 && !raw_compacted) {
         constexpr int k_fa_chunk = 256;
         int n_swa = hparams.n_swa;
         int ntokens = std::max(k_fa_chunk, int(q->ne[2]));
@@ -1126,11 +1141,15 @@ static ggml_tensor * ds4_attention(ggml_cgraph * gf, ggml_context * ctx0, llm_bu
             extra_mask = dsv4_pad_mask_tokens(ctx0, extra_mask, n_tokens);
         }
         raw_k = dsv4_repeat_streams(ctx0, raw_k, extra_k->ne[3]);
-        if (!cparams.flash_attn) {
+        if (!cparams.flash_attn && !raw_compacted) {
             raw_mask = dsv4_build_raw_mask_view(ctx0, KQ_mask,
                     lctx.dsv4.inputs.raw_k_read_idxs, raw_kq_n_kv, n_tokens, extra_k->ne[3], cb, il);
             raw_mask = dsv4_pad_raw_mask_to(ctx0, raw_mask, raw_attn_n_kv, n_tokens);
         }
+        if (!cparams.flash_attn && raw_compacted && raw_mask->ne[1] != n_tokens) {
+            // non-FA masks are unpadded; drop the win mask's token padding before the concat
+            raw_mask = ggml_view_2d(ctx0, raw_mask, raw_mask->ne[0], n_tokens, raw_mask->nb[1], 0);
+        }
         if (cparams.flash_attn && extra_mask->type != GGML_TYPE_F16) {
             extra_mask = ggml_cast(ctx0, extra_mask, GGML_TYPE_F16);
         }
@@ -1257,7 +1276,16 @@ ggml_cgraph * llm_build_context::build_deepseek4() {
     dsv4_build_plan_inputs(ctx0, lctx.dsv4.inputs.lid, lctx.dsv4.lid_plan, "dsv4_lid", n_tokens, false, lctx.cparams.flash_attn);
 
     ggml_tensor * inp_pos = build_inp_pos();
-    ggml_tensor * KQ_mask = hparams.n_swa > 0 ? build_inp_KQ_mask_swa() : build_inp_KQ_mask();
+    // build only the mask the graph consumes; an input tensor without a consumer is never allocated
+    ggml_tensor * KQ_mask = nullptr;
+    ggml_tensor * KQ_mask_swa_win = nullptr;
+    if (kv_self.any_compacted()) {
+        bool KQ_mask_swa_windowed = false;
+        KQ_mask_swa_win = build_swa_mask_for_graph(hparams.n_swa, /* compacted = */ true, &KQ_mask_swa_windowed);
+        GGML_ASSERT(KQ_mask_swa_windowed && KQ_mask_swa_win != nullptr);
+    } else {
+        KQ_mask = hparams.n_swa > 0 ? build_inp_KQ_mask_swa() : build_inp_KQ_mask();
+    }
     ggml_tensor * inpL = nullptr;
 
     ggml_tensor * append_csa_state = nullptr;
@@ -1301,7 +1329,7 @@ ggml_cgraph * llm_build_context::build_deepseek4() {
         auto cur = ds4_attention(gf, ctx0, *this, inpL,
                              &append_csa_state, &append_csa_score,
                              &append_lid_state, &append_lid_score,
-                             inp_pos, KQ_mask, il);
+                             inp_pos, KQ_mask, KQ_mask_swa_win, il);
         inpL = cur;
 
         ggml_tensor *post, *comb;

src/graphs/build_openpangu.cpp

@@ -118,40 +118,6 @@ 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 compacted, bool * windowed) {
-    *windowed = false;
-    llm.lctx.swa_window_view = {};
-
-    if (window == 0) {
-        return nullptr;
-    }
-
-    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.swa_window_view = {
-        true,
-        compacted,
-        llm.n_kv,
-        llm.n_tokens,
-        window,
-        pad,
-        view.w_view,
-        view.win_off,
-    };
-    *windowed = true;
-    return llm.build_inp_KQ_mask_swa_win(view.w_view);
-}
-
 // openPangu-2.0-Flash graph.
 //
 // Attention runs absorbed MLA over a latent KV cache: per position the cache stores only
@@ -1034,7 +1000,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, /* compacted = */ false, &KQ_mask_swa_windowed)
+            ? build_swa_mask_for_graph(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);
@@ -1174,7 +1140,7 @@ ggml_cgraph * llm_build_context::build_openpangu() {
     // 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, kv_self.any_compacted(), &KQ_mask_swa_windowed)
+        ? build_swa_mask_for_graph(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);

src/llama-build-context.cpp

@@ -596,6 +596,39 @@ ggml_tensor * llm_build_context::build_inp_KQ_mask_swa_win(int64_t n_kv_win, boo
     return flash_attn ? ggml_cast(ctx0, lctx.inp_KQ_mask_swa_win, GGML_TYPE_F16) : lctx.inp_KQ_mask_swa_win;
 }
 
+ggml_tensor * llm_build_context::build_swa_mask_for_graph(uint32_t window, bool compacted, bool * windowed) {
+    *windowed = false;
+    lctx.swa_window_view = {};
+
+    if (window == 0) {
+        return nullptr;
+    }
+
+    const uint32_t pad = llama_kv_cache::get_padding(cparams.flash_attn);
+    const int64_t live = compacted
+        ? (int64_t) swa_head - (int64_t) kv_self.sink_rows + n_tokens : 0;
+    const llama_swa_window_view view = compacted
+        ? llama_swa_calc_window_view_compact(live, kv_self.sink_rows, n_tokens, window, pad)
+        : llama_swa_calc_window_view(n_kv, n_tokens, window, pad);
+
+    if (!view.engaged) {
+        return build_inp_KQ_mask_swa();
+    }
+
+    lctx.swa_window_view = {
+        true,
+        compacted,
+        n_kv,
+        n_tokens,
+        window,
+        pad,
+        view.w_view,
+        view.win_off,
+    };
+    *windowed = true;
+    return build_inp_KQ_mask_swa_win(view.w_view);
+}
+
 //build_mhc_post: x = 4096 x 4096 x 1 x 1, post = 4 x 4096 x 1 x 1, residual = 4096 x 4 x 4096 x 1, comb = 4 x 4 x 4096 x 1
 //build_mhc_post: x = 4096 x 1 x 1 x 1,    post = 4 x 1 x 1 x 1,    residual = 4096 x 4 x 1 x 1,    comb = 4 x 4 x 1 x 1
 // x        = n_embd x n_tokens           <--- y in Pangu

src/llama-build-context.h

@@ -153,6 +153,8 @@ struct llm_build_context {
 
     ggml_tensor * build_inp_KQ_mask_swa_win(int64_t n_kv_win, bool causal = true);
 
+    ggml_tensor * build_swa_mask_for_graph(uint32_t window, bool compacted, bool * windowed);
+
     ggml_tensor * build_inp_mean();
 
     ggml_tensor * build_inp_cls();

src/llama-dsv4.cpp

@@ -356,6 +356,22 @@ static bool dsv4_build_raw_context(
         return false;
     }
 
+    // compacted layers address raw K rows through [sinks | window] geometry rather than by cell
+    const bool compacted = kv.any_compacted();
+    if (compacted) {
+        if (kv.head_swa + (uint32_t) batch.n_tokens > kv.size_swa) {
+            LLAMA_LOG_ERROR("%s: DSV4 compacted raw write rows [%u, %u) are outside size_swa %u\n",
+                    __func__, kv.head_swa, kv.head_swa + (uint32_t) batch.n_tokens, kv.size_swa);
+            return false;
+        }
+        if (batch.pos != nullptr && batch.n_tokens > 0 &&
+            kv.pos_base_swa + (llama_pos) (kv.head_swa - kv.sink_rows) != batch.pos[0]) {
+            LLAMA_LOG_ERROR("%s: DSV4 compacted write row %u disagrees with batch position %d (base %d)\n",
+                    __func__, kv.head_swa, batch.pos[0], kv.pos_base_swa);
+            return false;
+        }
+    }
+
     raw.write_counts.push_back(batch.n_tokens);
     for (int32_t i = 0; i < batch.n_tokens; ++i) {
         const int32_t slot = kv.head + i;
@@ -368,7 +384,7 @@ static bool dsv4_build_raw_context(
         }
 
         raw.write_src_idxs.push_back(i);
-        raw.write_dst_idxs.push_back(slot);
+        raw.write_dst_idxs.push_back(compacted ? (int32_t) kv.head_swa + i : slot);
     }
 
     raw.n_kv = 0;
@@ -385,8 +401,14 @@ static bool dsv4_build_raw_context(
             if (!cell.has_seq_id(seq_id)) {
                 continue;
             }
-            raw.sinfo_read.idxs[s].push_back(slot);
+            if (compacted && cell.pos < kv.pos_base_swa) {
-            raw.read_dst_idxs.push_back((int32_t) slot);
+                // rows before the window base were overwritten by compaction
+                continue;
+            }
+            const uint32_t row = compacted
+                ? kv.sink_rows + (uint32_t) (cell.pos - kv.pos_base_swa) : slot;
+            raw.sinfo_read.idxs[s].push_back(row);
+            raw.read_dst_idxs.push_back((int32_t) row);
             ++count;
         }
         raw.read_counts.push_back(count);

src/llama-hparams.h

@@ -395,18 +395,14 @@ 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
+// sinks + retained window + one u-batch, padded as one sum so the total is pad-aligned for any
-// small u-batch from compacting every few tokens.
+// sink_rows; compaction then fires every C - W tokens, and 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);
+    const uint32_t min_slack = 256;
+    const uint32_t slack     = n_ubatch > min_slack ? n_ubatch : min_slack;
+    const uint32_t unpadded  = sink_rows + window + slack;
+    return pad > 1 ? ((unpadded + pad - 1)/pad)*pad : unpadded;
 }
 
 static inline uint32_t llama_kv_layer_rows(const llama_hparams & hparams, int il, uint32_t kv_size,

src/llama-model.cpp

@@ -2392,7 +2392,9 @@ size_t llama_model::cache_size(int il, ggml_type type_k, ggml_type type_v, ggml_
         const int64_t n_embd_head = hparams.n_embd_head_k(il);
         const int64_t n_indexer_head = hparams.indexer_head_size;
 
-        size_t size = ggml_row_size(type_k, n_embd_head) * hparams.n_head_kv(il) * kv_size;
+        const uint32_t raw_pad = llama_kv_cache::get_padding(flash_attn);
+        const uint32_t k_rows = llama_kv_layer_rows(hparams, il, kv_size, swa_compress, n_ubatch, raw_pad);
+        size_t size = ggml_row_size(type_k, n_embd_head) * hparams.n_head_kv(il) * k_rows;
         if (ratio == csa_ratio) {
             size += ggml_row_size(type_k, n_embd_head) * csa_kv * n_stream;
             size += ggml_row_size(idx_type_k, n_indexer_head) * csa_kv * n_stream;

src/llama-model.h

@@ -584,7 +584,7 @@ struct llama_model {
     // 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;
+        return arch == LLM_ARCH_OPENPANGU || arch == LLM_ARCH_DEEPSEEK4;
     }
 
     static inline int hadamard_size(int head_size) {

src/llama.cpp

@@ -1461,7 +1461,7 @@ static bool llama_kv_cache_init(
                 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, 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);
+                k = ggml_new_tensor_2d(ctx, this_type_k, n_embd_head_k, n_head_kv*cache.rows(i));
             } else {
                 k = ggml_new_tensor_2d(ctx, this_type_k, n_embd_head_k, n_head_kv*kv_size);
                 v = ggml_new_tensor_1d(ctx, this_type_v, v_ne);
@@ -1707,7 +1707,9 @@ static void llama_kv_cache_compact_swa(struct llama_context & lctx, uint32_t n_t
     }
 
     const uint32_t W = cache.window_swa;
-    const uint32_t C = cache.size_swa - cache.sink_rows;
+    // the graph views the window at pad granularity, so only the pad-aligned part of the allocation is usable capacity
+    const uint32_t pad = llama_kv_cache::get_padding(lctx.cparams.flash_attn);
+    const uint32_t C = pad > 1 ? ((cache.size_swa - cache.sink_rows)/pad)*pad : cache.size_swa - cache.sink_rows;
     GGML_ASSERT(n_tokens <= C);
 
     if (cache.live_swa() + n_tokens <= C) {
@@ -1728,7 +1730,9 @@ static void llama_kv_cache_compact_swa(struct llama_context & lctx, uint32_t n_t
             continue;
         }
         ggml_tensor * kl = cache.k_l[il];
-        const size_t stride = kl->nb[1];
+        // kl rows are position-major: kl->ne[1]/rows(il) rows per position (n_head_kv)
+        const size_t rows_per_pos = (size_t) kl->ne[1] / cache.rows((int) il);
+        const size_t stride = kl->nb[1] * rows_per_pos;
         const size_t nbytes = (size_t) W * stride;
         if (scratch.size() < nbytes) {
             scratch.resize(nbytes);
@@ -2519,6 +2523,11 @@ static void llama_kv_cache_defrag(struct llama_kv_cache & cache) {
         LLAMA_LOG_WARN("%s: defrag is not supported for this model's position-indexed KV cache - skipping\n", __func__);
         return;
     }
+    if (cache.any_compacted()) {
+        // defrag moves rows by cell index; a compacted layer has no per-cell rows
+        LLAMA_LOG_WARN("%s: defrag is not supported for a compacted KV cache (--swa-compress) - skipping\n", __func__);
+        return;
+    }
     cache.do_defrag = true;
 }
 
@@ -7893,6 +7902,13 @@ struct llama_context * llama_init_from_model(
         return nullptr;
     }
 
+    if (params.n_seq_max > 1 && params.swa_compress && model->arch == LLM_ARCH_DEEPSEEK4) {
+        // the compacted window is one position stream (head_swa/pos_base_swa are per-cache, not per-sequence)
+        LLAMA_LOG_ERROR("%s: --swa-compress supports a single sequence only (requested n_seq_max = %u); run with -np 1\n",
+                __func__, params.n_seq_max);
+        return nullptr;
+    }
+
     if (model->arch == LLM_ARCH_OPENPANGU) {
         std::string error_msg;
         if (!llama_openpangu_validate_latent_cache_types(params.type_k, params.type_v, &error_msg)) {
@@ -8635,7 +8651,7 @@ void llama_free(struct llama_context * ctx) {
 }
 
 bool llama_supports_full_state_io(const struct llama_context * ctx) {
-    return ctx != nullptr;
+    return ctx != nullptr && !ctx->kv_self.any_compacted();
 }
 
 const struct llama_vocab* llama_model_get_vocab(const struct llama_model* model) {
@@ -11081,6 +11097,10 @@ static bool llama_state_io_supported(
         LLAMA_LOG_ERROR("%s: only per-sequence state save/restore is supported for openPangu (whole-context and file-session state are not)\n", func);
         return false;
     }
+    if (ctx->kv_self.any_compacted() && seq_id < 0) {
+        LLAMA_LOG_ERROR("%s: whole-context state save/restore is not supported with --swa-compress; use per-sequence state\n", func);
+        return false;
+    }
     return true;
 }