mamba : stop abusing attention metadata

This breaks existing converted-to-GGUF Mamba models,
but will allow supporting mixed architectures like MambaFormer
without needing to break Mamba models.

This will also allow changing the size of Mamba's states
without having to reconvert models in the future.
(e.g. using something else than d_conv - 1 columns for the conv_states
 will not require breaking existing converted Mamba models again)

* gguf-py : add new KV metadata key-value pairs for Mamba

* llama : add new metadata key-value pairs for Mamba

* llama : guard against divisions by zero when n_head is 0

* mamba : rename "unlimited" KV cache property to "recurrent"

Author: compilade

convert-hf-to-gguf.py

@@ -1862,21 +1862,28 @@ def set_vocab(self):
 
     def set_gguf_parameters(self):
         d_model = self.hparams["d_model"]
+        d_conv  = self.hparams.get("d_conv", 4)
         d_inner = self.hparams.get("d_inner", 2 * d_model)
+        d_state = self.hparams.get("d_state", 16)
+        # ceiling division
+        # ref: https://stackoverflow.com/a/17511341/22827863
+        # ref: https://github.com/state-spaces/mamba/blob/ce59daea3a090d011d6476c6e5b97f6d58ddad8b/mamba_ssm/modules/mamba_simple.py#L58
+        dt_rank = self.hparams.get("dt_rank", -(d_model // -16))
+
         # Fail early for models which don't have a block expansion factor of 2
         assert d_inner == 2 * d_model
 
         self.gguf_writer.add_name(self.dir_model.name)
         self.gguf_writer.add_context_length(2**20) # arbitrary value; for those who use the default
         self.gguf_writer.add_embedding_length(d_model)
         self.gguf_writer.add_feed_forward_length(0) # unused, but seemingly required when loading
-        self.gguf_writer.add_head_count(d_inner) # the number of rows in conv_state and ssm_state
+        self.gguf_writer.add_head_count(0) # unused, but seemingly required when loading
         self.gguf_writer.add_block_count(self.hparams["n_layer"])
+        self.gguf_writer.add_ssm_conv_kernel_size(d_conv)
+        self.gguf_writer.add_ssm_inner_length(d_inner)
+        self.gguf_writer.add_ssm_state_length(d_state)
+        self.gguf_writer.add_ssm_dt_rank(dt_rank)
         self.gguf_writer.add_layer_norm_rms_eps(self.hparams.get("rms_norm_eps", 1e-5))
-        # NOTE: (ab)using the KV cache metadata to store dimensions for conv_state and ssm_state
-        # Since the first column of the conv_state is shifted out each time, it's not actually needed
-        self.gguf_writer.add_key_length(self.hparams.get("d_conv", 4) - 1)
-        self.gguf_writer.add_value_length(self.hparams.get("d_state", 16))
         self.gguf_writer.add_file_type(self.ftype)
 
     def write_tensors(self):

gguf-py/gguf/constants.py

@@ -61,6 +61,12 @@ class Rope:
         SCALING_ORIG_CTX_LEN = "{arch}.rope.scaling.original_context_length"
         SCALING_FINETUNED    = "{arch}.rope.scaling.finetuned"
 
+    class SSM:
+        CONV_KERNEL_SIZE     = "{arch}.ssm.d_conv"
+        INNER_LENGTH         = "{arch}.ssm.d_inner"
+        STATE_LENGTH         = "{arch}.ssm.d_state"
+        DT_RANK              = "{arch}.ssm.dt_rank"
+
     class Tokenizer:
         MODEL            = "tokenizer.ggml.model"
         LIST             = "tokenizer.ggml.tokens"
@@ -747,6 +753,12 @@ def get_type(val: Any) -> GGUFValueType:
 KEY_ROPE_SCALING_ORIG_CTX_LEN = Keys.Rope.SCALING_ORIG_CTX_LEN
 KEY_ROPE_SCALING_FINETUNED    = Keys.Rope.SCALING_FINETUNED
 
+# SSM
+KEY_SSM_CONV_KERNEL_SIZE = Keys.SSM.CONV_KERNEL_SIZE
+KEY_SSM_INNER_LENGTH     = Keys.SSM.INNER_LENGTH
+KEY_SSM_STATE_LENGTH     = Keys.SSM.STATE_LENGTH
+KEY_SSM_DT_RANK          = Keys.SSM.DT_RANK
+
 # tokenization
 KEY_TOKENIZER_MODEL      = Keys.Tokenizer.MODEL
 KEY_TOKENIZER_LIST       = Keys.Tokenizer.LIST

gguf-py/gguf/gguf_writer.py

@@ -382,6 +382,18 @@ def add_rope_scaling_orig_ctx_len(self, value: int) -> None:
     def add_rope_scaling_finetuned(self, value: bool) -> None:
         self.add_bool(Keys.Rope.SCALING_FINETUNED.format(arch=self.arch), value)
 
+    def add_ssm_conv_kernel_size(self, value: int) -> None:
+        self.add_uint32(Keys.SSM.CONV_KERNEL_SIZE.format(arch=self.arch), value)
+
+    def add_ssm_inner_length(self, value: int) -> None:
+        self.add_uint32(Keys.SSM.INNER_LENGTH.format(arch=self.arch), value)
+
+    def add_ssm_state_length(self, value: int) -> None:
+        self.add_uint32(Keys.SSM.STATE_LENGTH.format(arch=self.arch), value)
+
+    def add_ssm_dt_rank(self, value: int) -> None:
+        self.add_uint32(Keys.SSM.DT_RANK.format(arch=self.arch), value)
+
     def add_tokenizer_model(self, model: str) -> None:
         self.add_string(Keys.Tokenizer.MODEL, model)