Sync server : context checkpointing for hybrid and recurrent models

Author: JamePeng

llama_cpp/llama_cpp.py

@@ -1694,6 +1694,14 @@ def llama_model_is_recurrent(model: llama_model_p, /) -> bool:
     ...
 
 
+# // Returns true if the model is hybrid (like Jamba, Granite, etc.)
+# LLAMA_API bool llama_model_is_hybrid(const struct llama_model * model);
+@ctypes_function("llama_model_is_hybrid", [llama_model_p_ctypes], ctypes.c_bool)
+def llama_model_is_hybrid(model: llama_model_p, /) -> bool:
+    """Returns true if the model is hybrid (like Jamba, Granite, etc.)"""
+    ...
+
+
 # // Returns true if the model is diffusion-based (like LLaDA, Dream, etc.)
 # LLAMA_API bool llama_model_is_diffusion(const struct llama_model * model);
 @ctypes_function("llama_model_is_diffusion", [llama_model_p_ctypes], ctypes.c_bool)
@@ -2539,6 +2547,92 @@ def llama_state_seq_load_file(
     ...
 
 
+# // for backwards-compat
+LLAMA_STATE_SEQ_FLAGS_SWA_ONLY = 1
+
+# // work only with partial states, such as SWA KV cache or recurrent cache (e.g. Mamba)
+LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY = 1
+
+llama_state_seq_flags = ctypes.c_uint32
+
+# LLAMA_API size_t llama_state_seq_get_size_ext(
+#         struct llama_context * ctx,
+#                 llama_seq_id   seq_id,
+#         llama_state_seq_flags   flags);
+@ctypes_function(
+    "llama_state_seq_get_size_ext",
+    [
+        llama_context_p_ctypes,
+        llama_seq_id,
+        llama_state_seq_flags,
+    ],
+    ctypes.c_size_t,
+)
+def llama_state_seq_get_size_ext(
+    ctx: llama_context_p,
+    seq_id: llama_seq_id,
+    flags: llama_state_seq_flags,
+    /,
+) -> int:
+    ...
+
+
+# LLAMA_API size_t llama_state_seq_get_data_ext(
+#         struct llama_context * ctx,
+#                         uint8_t * dst,
+#                         size_t   size,
+#                 llama_seq_id   seq_id,
+#         llama_state_seq_flags   flags);
+@ctypes_function(
+    "llama_state_seq_get_data_ext",
+    [
+        llama_context_p_ctypes,
+        ctypes.POINTER(ctypes.c_uint8),
+        ctypes.c_size_t,
+        llama_seq_id,
+        llama_state_seq_flags,
+    ],
+    ctypes.c_size_t,
+)
+def llama_state_seq_get_data_ext(
+    ctx: llama_context_p,
+    dst: ctypes.POINTER(ctypes.c_uint8),
+    size: Union[int, ctypes.c_size_t],
+    seq_id: llama_seq_id,
+    flags: llama_state_seq_flags,
+    /,
+) -> int:
+    ...
+
+
+# LLAMA_API size_t llama_state_seq_set_data_ext(
+#         struct llama_context * ctx,
+#                 const uint8_t * src,
+#                         size_t   size,
+#                 llama_seq_id   dest_seq_id,
+#         llama_state_seq_flags   flags);
+@ctypes_function(
+    "llama_state_seq_set_data_ext",
+    [
+        llama_context_p_ctypes,
+        ctypes.POINTER(ctypes.c_uint8),
+        ctypes.c_size_t,
+        llama_seq_id,
+        llama_state_seq_flags,
+    ],
+    ctypes.c_size_t,
+)
+def llama_state_seq_set_data_ext(
+    ctx: llama_context_p,
+    src: ctypes.POINTER(ctypes.c_uint8),
+    size: Union[int, ctypes.c_size_t],
+    dest_seq_id: llama_seq_id,
+    flags: llama_state_seq_flags,
+    /,
+) -> int:
+    ...
+
+
 # //
 # // Decoding
 # //