Upload folder using huggingface_hub

.meta/create_model.py

@@ -107,6 +107,25 @@ def create_model(source_model_id: str, source_revision: str, save_folder: str) -
                 past_key_values=past_key_values,
                 position_ids=position_ids,
             )"""
+    cache_api_code = '''    def get_mask_sizes(self, cache_position: torch.Tensor, layer_idx: int) -> tuple[int, int]:
+        """
+        Return a tuple (kv_length, kv_offset) corresponding to the length and offset that will be returned for
+        the given layer at `layer_idx`.
+        The masks are then prepared according to the given lengths (kv_length, kv_offset) and patterns for each layer.
+        """
+        kv_offset = 0
+        query_length = cache_position.shape[0]'''
+    cache_api_compat_code = '''    def get_query_offset(self, layer_idx: int) -> int:
+        return self.get_seq_length(layer_idx)
+
+    def get_mask_sizes(self, cache_position: torch.Tensor | int, layer_idx: int) -> tuple[int, int]:
+        """
+        Return a tuple (kv_length, kv_offset) corresponding to the length and offset that will be returned for
+        the given layer at `layer_idx`.
+        The masks are then prepared according to the given lengths (kv_length, kv_offset) and patterns for each layer.
+        """
+        kv_offset = 0
+        query_length = cache_position if isinstance(cache_position, int) else cache_position.shape[0]'''
 
     replace_file(
         f"{folder}/modeling_kimi_linear.py",
@@ -128,6 +147,7 @@ def create_model(source_model_id: str, source_revision: str, save_folder: str) -
     )""",
             ),
             (causal_mask_code, causal_mask_compat_code),
+            (cache_api_code, cache_api_compat_code),
             (
                 force_flash_code,
                 '        config._attn_implementation = getattr(config, "_attn_implementation", "eager")',
@@ -163,6 +183,26 @@ def create_model(source_model_id: str, source_revision: str, save_folder: str) -
                 else:
                     first_layer_past_key_value = past_key_values[0][0][:, :, :, 0]""",
             ),
+            # Transformers 5.x drops an all-ones attention mask and passes None; the
+            # multimodal merge and the cached generation branch still expect a concrete mask.
+            (
+                "                # Get the target length\n"
+                "                target_length = input_ids.shape[1]\n"
+                "                past_length = first_layer_past_key_value.shape[-1]\n",
+                "                # Get the target length\n"
+                "                target_length = input_ids.shape[1]\n"
+                "                past_length = first_layer_past_key_value.shape[-1]\n\n"
+                "                if attention_mask is None:\n"
+                "                    attention_mask = torch.ones_like(input_ids)\n",
+            ),
+            (
+                "        _, embed_dim = image_features[0].shape\n"
+                "        feature_lengths = [x.shape[0] for x in image_features]",
+                "        if attention_mask is None:\n"
+                "            attention_mask = torch.ones_like(input_ids)\n\n"
+                "        _, embed_dim = image_features[0].shape\n"
+                "        feature_lengths = [x.shape[0] for x in image_features]",
+            ),
         ],
     )
 

.meta/env.json

@@ -18,7 +18,28 @@
         "torch": "2.13.0+cu126",
         "transformers": "5.22.0.dev0"
       },
-      "cuda": "12.6"
+      "cuda": "12.6",
+      "hardware": {
+        "gpu_model": "NVIDIA H20"
+      }
+    },
+    "transformers": {
+      "python": "3.12.13",
+      "packages": {
+        "einops": "0.9.0.dev0",
+        "fla-core": "0.5.2",
+        "huggingface_hub": "2.3.0",
+        "numpy": "2.5.2",
+        "pillow": "12.3.0",
+        "safetensors": "0.8.0",
+        "tokenizers": "0.23.3",
+        "torch": "2.13.0+cu126",
+        "transformers": "5.22.0.dev0"
+      },
+      "cuda": "12.6",
+      "hardware": {
+        "gpu_model": "NVIDIA H20"
+      }
     }
   }
 }

README.md

@@ -35,7 +35,7 @@ K3.1-Pro ships no modeling code; the included pinned K3 `modeling_kimi_k3.py`/`m
 
 ### Example usage
 
-#### Transformers (unverified)
+#### Transformers (✅ verified)
 
 Loads with the pinned K3 reference modeling on GPU (KDA delta attention uses Triton via fla-core). The K3.1 gated-DSA layers run their MLA attention core; the LSA indexer tensors are present but not consumed because released Transformers lacks the gated-DSA operator.
 

modeling_kimi_k3.py

@@ -976,6 +976,9 @@ class KimiK3ForConditionalGeneration(KimiK3PreTrainedModel):
             labels (:obj:`torch.Tensor` of shape :obj:`(batch_size, sequence_length)`, *optional*):
                 The labels.
         """
+        if attention_mask is None:
+            attention_mask = torch.ones_like(input_ids)
+
         _, embed_dim = image_features[0].shape
         feature_lengths = [x.shape[0] for x in image_features]
         image_features = torch.cat(image_features, dim=0)
@@ -1187,6 +1190,9 @@ class KimiK3ForConditionalGeneration(KimiK3PreTrainedModel):
                 target_length = input_ids.shape[1]
                 past_length = first_layer_past_key_value.shape[-1]
 
+                if attention_mask is None:
+                    attention_mask = torch.ones_like(input_ids)
+
                 extended_attention_mask = torch.ones(
                     (attention_mask.shape[0], past_length),
                     dtype=attention_mask.dtype,

modeling_kimi_linear.py

@@ -203,14 +203,17 @@ class KimiDynamicCache:
             return 0
         return self.key_cache[layer_idx].shape[-2]
 
-    def get_mask_sizes(self, cache_position: torch.Tensor, layer_idx: int) -> tuple[int, int]:
+    def get_query_offset(self, layer_idx: int) -> int:
+        return self.get_seq_length(layer_idx)
+
+    def get_mask_sizes(self, cache_position: torch.Tensor | int, layer_idx: int) -> tuple[int, int]:
         """
         Return a tuple (kv_length, kv_offset) corresponding to the length and offset that will be returned for
         the given layer at `layer_idx`.
         The masks are then prepared according to the given lengths (kv_length, kv_offset) and patterns for each layer.
         """
         kv_offset = 0
-        query_length = cache_position.shape[0]
+        query_length = cache_position if isinstance(cache_position, int) else cache_position.shape[0]
         past_seen_tokens = self.get_seq_length(layer_idx)
         kv_length = query_length + past_seen_tokens
         return kv_length, kv_offset