Add files using upload-large-folder tool

.gitattributes

@@ -33,3 +33,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+tokenizer.json filter=lfs diff=lfs merge=lfs -text

added_tokens.json

@@ -0,0 +1,28 @@
+{
+  "</think>": 151668,
+  "</tool_call>": 151658,
+  "</tool_response>": 151666,
+  "<think>": 151667,
+  "<tool_call>": 151657,
+  "<tool_response>": 151665,
+  "<|box_end|>": 151649,
+  "<|box_start|>": 151648,
+  "<|endoftext|>": 151643,
+  "<|file_sep|>": 151664,
+  "<|fim_middle|>": 151660,
+  "<|fim_pad|>": 151662,
+  "<|fim_prefix|>": 151659,
+  "<|fim_suffix|>": 151661,
+  "<|im_end|>": 151645,
+  "<|im_start|>": 151644,
+  "<|image_pad|>": 151655,
+  "<|object_ref_end|>": 151647,
+  "<|object_ref_start|>": 151646,
+  "<|quad_end|>": 151651,
+  "<|quad_start|>": 151650,
+  "<|repo_name|>": 151663,
+  "<|video_pad|>": 151656,
+  "<|vision_end|>": 151653,
+  "<|vision_pad|>": 151654,
+  "<|vision_start|>": 151652
+}

chat_template.jinja

@@ -0,0 +1,85 @@
+{%- if tools %}
+    {{- '<|im_start|>system\n' }}
+    {%- if messages[0].role == 'system' %}
+        {{- messages[0].content + '\n\n' }}
+    {%- endif %}
+    {{- "# Tools\n\nYou may call one or more functions to assist with the user query.\n\nYou are provided with function signatures within <tools></tools> XML tags:\n<tools>" }}
+    {%- for tool in tools %}
+        {{- "\n" }}
+        {{- tool | tojson }}
+    {%- endfor %}
+    {{- "\n</tools>\n\nFor each function call, return a json object with function name and arguments within <tool_call></tool_call> XML tags:\n<tool_call>\n{\"name\": <function-name>, \"arguments\": <args-json-object>}\n</tool_call><|im_end|>\n" }}
+{%- else %}
+    {%- if messages[0].role == 'system' %}
+        {{- '<|im_start|>system\n' + messages[0].content + '<|im_end|>\n' }}
+    {%- endif %}
+{%- endif %}
+{%- set ns = namespace(multi_step_tool=true, last_query_index=messages|length - 1) %}
+{%- for message in messages[::-1] %}
+    {%- set index = (messages|length - 1) - loop.index0 %}
+    {%- if ns.multi_step_tool and message.role == "user" and not(message.content.startswith('<tool_response>') and message.content.endswith('</tool_response>')) %}
+        {%- set ns.multi_step_tool = false %}
+        {%- set ns.last_query_index = index %}
+    {%- endif %}
+{%- endfor %}
+{%- for message in messages %}
+    {%- if (message.role == "user") or (message.role == "system" and not loop.first) %}
+        {{- '<|im_start|>' + message.role + '\n' + message.content + '<|im_end|>' + '\n' }}
+    {%- elif message.role == "assistant" %}
+        {%- set content = message.content %}
+        {%- set reasoning_content = '' %}
+        {%- if message.reasoning_content is defined and message.reasoning_content is not none %}
+            {%- set reasoning_content = message.reasoning_content %}
+        {%- else %}
+            {%- if '</think>' in message.content %}
+                {%- set content = message.content.split('</think>')[-1].lstrip('\n') %}
+                {%- set reasoning_content = message.content.split('</think>')[0].rstrip('\n').split('<think>')[-1].lstrip('\n') %}
+            {%- endif %}
+        {%- endif %}
+        {%- if loop.index0 > ns.last_query_index %}
+            {%- if loop.last or (not loop.last and reasoning_content) %}
+                {{- '<|im_start|>' + message.role + '\n<think>\n' + reasoning_content.strip('\n') + '\n</think>\n\n' + content.lstrip('\n') }}
+            {%- else %}
+                {{- '<|im_start|>' + message.role + '\n' + content }}
+            {%- endif %}
+        {%- else %}
+            {{- '<|im_start|>' + message.role + '\n' + content }}
+        {%- endif %}
+        {%- if message.tool_calls %}
+            {%- for tool_call in message.tool_calls %}
+                {%- if (loop.first and content) or (not loop.first) %}
+                    {{- '\n' }}
+                {%- endif %}
+                {%- if tool_call.function %}
+                    {%- set tool_call = tool_call.function %}
+                {%- endif %}
+                {{- '<tool_call>\n{"name": "' }}
+                {{- tool_call.name }}
+                {{- '", "arguments": ' }}
+                {%- if tool_call.arguments is string %}
+                    {{- tool_call.arguments }}
+                {%- else %}
+                    {{- tool_call.arguments | tojson }}
+                {%- endif %}
+                {{- '}\n</tool_call>' }}
+            {%- endfor %}
+        {%- endif %}
+        {{- '<|im_end|>\n' }}
+    {%- elif message.role == "tool" %}
+        {%- if loop.first or (messages[loop.index0 - 1].role != "tool") %}
+            {{- '<|im_start|>user' }}
+        {%- endif %}
+        {{- '\n<tool_response>\n' }}
+        {{- message.content }}
+        {{- '\n</tool_response>' }}
+        {%- if loop.last or (messages[loop.index0 + 1].role != "tool") %}
+            {{- '<|im_end|>\n' }}
+        {%- endif %}
+    {%- endif %}
+{%- endfor %}
+{%- if add_generation_prompt %}
+    {{- '<|im_start|>assistant\n' }}
+    {%- if enable_thinking is defined and enable_thinking is false %}
+        {{- '<think>\n\n</think>\n\n' }}
+    {%- endif %}
+{%- endif %}

config.json

@@ -0,0 +1,107 @@
+{
+  "architectures": [
+    "BrumbyForCausalLM"
+  ],
+  "attention_bias": false,
+  "attention_dropout": 0.0,
+  "auto_map": {
+    "AutoConfig": "configuration_brumby.BrumbyConfig",
+    "AutoModelForCausalLM": "modeling_brumby.BrumbyForCausalLM"
+  },
+  "bos_token_id": 151643,
+  "chunk_size": 64,
+  "dtype": "bfloat16",
+  "eos_token_id": 151643,
+  "head_dim": 128,
+  "hidden_act": "silu",
+  "hidden_size": 5120,
+  "initializer_range": 0.02,
+  "intermediate_size": 17408,
+  "layer_types": [
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention"
+  ],
+  "max_position_embeddings": 32768,
+  "max_window_layers": 40,
+  "model_type": "brumby",
+  "num_attention_heads": 40,
+  "num_hidden_layers": 40,
+  "num_key_value_heads": 8,
+  "p": 2,
+  "pad_token_id": 151643,
+  "prefill_chunk_size": 1024,
+  "quantization_config": {
+    "bits": 4,
+    "checkpoint_format": "gptq",
+    "desc_act": false,
+    "group_size": 32,
+    "lm_head": false,
+    "meta": {
+      "act_group_aware": true,
+      "damp_auto_increment": 0.01,
+      "damp_percent": 0.05,
+      "mse": 0.0,
+      "quantizer": [
+        "gptqmodel:5.1.0-dev"
+      ],
+      "static_groups": false,
+      "true_sequential": true,
+      "uri": "https://github.com/modelcloud/gptqmodel",
+      "v2": false,
+      "v2_alpha": 0.25
+    },
+    "pack_dtype": "int32",
+    "pack_impl": "cpu",
+    "quant_method": "gptq",
+    "sym": true
+  },
+  "rms_norm_eps": 1e-06,
+  "rope_scaling": null,
+  "rope_theta": 1000000,
+  "sliding_window": null,
+  "switch_over_seq_len": 8192,
+  "tie_word_embeddings": false,
+  "transformers_version": "4.57.1",
+  "use_cache": true,
+  "use_exp": false,
+  "use_sliding_window": false,
+  "vocab_size": 151936
+}

configuration_brumby.py

@@ -0,0 +1,235 @@
+# coding=utf-8
+# Copyright 2025 Manifest AI.
+# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Brumby model configuration"""
+
+from transformers.configuration_utils import PretrainedConfig, layer_type_validation
+from transformers.modeling_rope_utils import rope_config_validation
+from transformers.utils import logging
+
+
+logger = logging.get_logger(__name__)
+
+
+class BrumbyConfig(PretrainedConfig):
+    r"""
+    This is the configuration class to store the configuration of a [`BrumbyModel`]. It is used to instantiate a
+    Brumby model according to the specified arguments, defining the model architecture. Instantiating a configuration
+    with the defaults will yield a similar configuration to that of
+    Brumby-14B-base [Brumby/Brumby-14B-base](https://huggingface.co/manifestai/Brumby-14B-base).
+
+    Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the
+    documentation from [`PretrainedConfig`] for more information.
+
+
+    Args:
+        vocab_size (`int`, *optional*, defaults to 151936):
+            Vocabulary size of the Brumby model. Defines the number of different tokens that can be represented by the
+            `inputs_ids` passed when calling [`BrumbyModel`]
+        hidden_size (`int`, *optional*, defaults to 4096):
+            Dimension of the hidden representations.
+        intermediate_size (`int`, *optional*, defaults to 22016):
+            Dimension of the MLP representations.
+        num_hidden_layers (`int`, *optional*, defaults to 32):
+            Number of hidden layers in the Transformer encoder.
+        num_attention_heads (`int`, *optional*, defaults to 32):
+            Number of attention heads for each attention layer in the Transformer encoder.
+        num_key_value_heads (`int`, *optional*, defaults to 32):
+            This is the number of key_value heads that should be used to implement Grouped Query Attention. If
+            `num_key_value_heads=num_attention_heads`, the model will use Multi Head Attention (MHA), if
+            `num_key_value_heads=1` the model will use Multi Query Attention (MQA) otherwise GQA is used. When
+            converting a multi-head checkpoint to a GQA checkpoint, each group key and value head should be constructed
+            by meanpooling all the original heads within that group. For more details, check out [this
+            paper](https://huggingface.co/papers/2305.13245). If it is not specified, will default to `32`.
+        head_dim (`int`, *optional*, defaults to 128):
+            The attention head dimension.
+        hidden_act (`str` or `function`, *optional*, defaults to `"silu"`):
+            The non-linear activation function (function or string) in the decoder.
+        max_position_embeddings (`int`, *optional*, defaults to 32768):
+            The maximum sequence length that this model might ever be used with.
+        initializer_range (`float`, *optional*, defaults to 0.02):
+            The standard deviation of the truncated_normal_initializer for initializing all weight matrices.
+        rms_norm_eps (`float`, *optional*, defaults to 1e-06):
+            The epsilon used by the rms normalization layers.
+        use_cache (`bool`, *optional*, defaults to `True`):
+            Whether or not the model should return the last key/values attentions (not used by all models). Only
+            relevant if `config.is_decoder=True`.
+        tie_word_embeddings (`bool`, *optional*, defaults to `False`):
+            Whether the model's input and output word embeddings should be tied.
+        rope_theta (`float`, *optional*, defaults to 10000.0):
+            The base period of the RoPE embeddings.
+        rope_scaling (`Dict`, *optional*):
+            Dictionary containing the scaling configuration for the RoPE embeddings. NOTE: if you apply new rope type
+            and you expect the model to work on longer `max_position_embeddings`, we recommend you to update this value
+            accordingly.
+            Expected contents:
+                `rope_type` (`str`):
+                    The sub-variant of RoPE to use. Can be one of ['default', 'linear', 'dynamic', 'yarn', 'longrope',
+                    'llama3'], with 'default' being the original RoPE implementation.
+                `factor` (`float`, *optional*):
+                    Used with all rope types except 'default'. The scaling factor to apply to the RoPE embeddings. In
+                    most scaling types, a `factor` of x will enable the model to handle sequences of length x *
+                    original maximum pre-trained length.
+                `original_max_position_embeddings` (`int`, *optional*):
+                    Used with 'dynamic', 'longrope' and 'llama3'. The original max position embeddings used during
+                    pretraining.
+                `attention_factor` (`float`, *optional*):
+                    Used with 'yarn' and 'longrope'. The scaling factor to be applied on the attention
+                    computation. If unspecified, it defaults to value recommended by the implementation, using the
+                    `factor` field to infer the suggested value.
+                `beta_fast` (`float`, *optional*):
+                    Only used with 'yarn'. Parameter to set the boundary for extrapolation (only) in the linear
+                    ramp function. If unspecified, it defaults to 32.
+                `beta_slow` (`float`, *optional*):
+                    Only used with 'yarn'. Parameter to set the boundary for interpolation (only) in the linear
+                    ramp function. If unspecified, it defaults to 1.
+                `short_factor` (`list[float]`, *optional*):
+                    Only used with 'longrope'. The scaling factor to be applied to short contexts (<
+                    `original_max_position_embeddings`). Must be a list of numbers with the same length as the hidden
+                    size divided by the number of attention heads divided by 2
+                `long_factor` (`list[float]`, *optional*):
+                    Only used with 'longrope'. The scaling factor to be applied to long contexts (<
+                    `original_max_position_embeddings`). Must be a list of numbers with the same length as the hidden
+                    size divided by the number of attention heads divided by 2
+                `low_freq_factor` (`float`, *optional*):
+                    Only used with 'llama3'. Scaling factor applied to low frequency components of the RoPE
+                `high_freq_factor` (`float`, *optional*):
+                    Only used with 'llama3'. Scaling factor applied to high frequency components of the RoPE
+        attention_bias (`bool`, defaults to `False`, *optional*, defaults to `False`):
+            Whether to use a bias in the query, key, value and output projection layers during self-attention.
+        use_sliding_window (`bool`, *optional*, defaults to `False`):
+            Whether to use sliding window attention.
+        sliding_window (`int`, *optional*, defaults to 4096):
+            Sliding window attention (SWA) window size. If not specified, will default to `4096`.
+        max_window_layers (`int`, *optional*, defaults to 28):
+            The number of layers using full attention. The first `max_window_layers` layers will use full attention, while any
+            additional layer afterwards will use SWA (Sliding Window Attention).
+        layer_types (`list`, *optional*):
+            Attention pattern for each layer.
+        attention_dropout (`float`, *optional*, defaults to 0.0):
+            The dropout ratio for the attention probabilities.
+
+    ```python
+    >>> from transformers import BrumbyModel, BrumbyConfig
+
+    >>> # Initializing a Brumby style configuration
+    >>> configuration = BrumbyConfig()
+
+    >>> # Initializing a model from the Brumby-14B-base style configuration
+    >>> model = BrumbyModel(configuration)
+
+    >>> # Accessing the model configuration
+    >>> configuration = model.config
+    ```"""
+
+    model_type = "brumby"
+    keys_to_ignore_at_inference = ["past_key_values"]
+
+    # Default tensor parallel plan for base model `Brumby`
+    base_model_tp_plan = {
+        "layers.*.self_attn.q_proj": "colwise",
+        "layers.*.self_attn.k_proj": "colwise",
+        "layers.*.self_attn.v_proj": "colwise",
+        "layers.*.self_attn.o_proj": "rowwise",
+        "layers.*.mlp.gate_proj": "colwise",
+        "layers.*.mlp.up_proj": "colwise",
+        "layers.*.mlp.down_proj": "rowwise",
+    }
+    base_model_pp_plan = {
+        "embed_tokens": (["input_ids"], ["inputs_embeds"]),
+        "layers": (["hidden_states", "attention_mask"], ["hidden_states"]),
+        "norm": (["hidden_states"], ["hidden_states"]),
+    }
+
+    def __init__(
+        self,
+        vocab_size=151936,
+        hidden_size=5120,
+        intermediate_size=17408,
+        num_hidden_layers=40,
+        num_attention_heads=40,
+        num_key_value_heads=8,
+        head_dim=128,
+        hidden_act="silu",
+        max_position_embeddings=32768,
+        initializer_range=0.02,
+        rms_norm_eps=1e-6,
+        use_cache=True,
+        tie_word_embeddings=False,
+        rope_theta=10000.0,
+        rope_scaling=None,
+        attention_bias=False,
+        use_sliding_window=False,
+        sliding_window=4096,
+        max_window_layers=40,
+        layer_types=None,
+        attention_dropout=0.0,
+        chunk_size=64,
+        switch_over_seq_len=8192,
+        prefill_chunk_size=1024,
+        use_exp=False,
+        **kwargs,
+    ):
+        self.vocab_size = vocab_size
+        self.max_position_embeddings = max_position_embeddings
+        self.hidden_size = hidden_size
+        self.intermediate_size = intermediate_size
+        self.num_hidden_layers = num_hidden_layers
+        self.num_attention_heads = num_attention_heads
+        self.use_sliding_window = use_sliding_window
+        self.sliding_window = sliding_window if self.use_sliding_window else None
+        self.max_window_layers = max_window_layers
+
+        # for backward compatibility
+        if num_key_value_heads is None:
+            num_key_value_heads = num_attention_heads
+
+        self.num_key_value_heads = num_key_value_heads
+        self.head_dim = head_dim
+        self.hidden_act = hidden_act
+        self.initializer_range = initializer_range
+        self.rms_norm_eps = rms_norm_eps
+        self.use_cache = use_cache
+        self.rope_theta = rope_theta
+        self.rope_scaling = rope_scaling
+        self.attention_bias = attention_bias
+        self.attention_dropout = attention_dropout
+        # Validate the correctness of rotary position embeddings parameters
+        # BC: if there is a 'type' field, move it to 'rope_type'.
+        if self.rope_scaling is not None and "type" in self.rope_scaling:
+            self.rope_scaling["rope_type"] = self.rope_scaling["type"]
+        rope_config_validation(self)
+
+        self.layer_types = layer_types
+        if self.layer_types is None:
+            self.layer_types = [
+                "sliding_attention"
+                if self.sliding_window is not None and i >= self.max_window_layers
+                else "full_attention"
+                for i in range(self.num_hidden_layers)
+            ]
+        layer_type_validation(self.layer_types, self.num_hidden_layers)
+        self.chunk_size = chunk_size
+        self.switch_over_seq_len = switch_over_seq_len
+        self.prefill_chunk_size = prefill_chunk_size
+        self.use_exp = use_exp
+        self.p = 2 # power of expansion for power attention, only 2 is supported for now
+        super().__init__(
+            tie_word_embeddings=tie_word_embeddings,
+            **kwargs,
+        )
+
+
+__all__ = ["BrumbyConfig"]

generation_config.json

@@ -0,0 +1,7 @@
+{
+  "bos_token_id": 151643,
+  "do_sample": true,
+  "eos_token_id": 151643,
+  "max_new_tokens": 2048,
+  "transformers_version": "4.57.1"
+}

model-00001-of-00003.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:b7a5fb6d7c5f0fbe8881269a0d64950655c2116fd54d57c0faa2ac418455cc8f
+size 4285645429

model-00002-of-00003.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:edebaa9a0222e8d49729b0765b38ee700a5df083da31a710aa50d6bb601b69cb
+size 4275983127

model-00003-of-00003.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:89ff91b78f9922f4309be337441185574508c9ef5e5ff6a4381da637c9579455
+size 2200218491

modeling_brumby.py

@@ -0,0 +1,857 @@
+# coding=utf-8
+# Copyright 2024 HuggingFace Inc. team.
+# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""PyTorch Brumby model."""
+
+from typing import Callable, Optional, Union, Any, Dict, Tuple
+
+import torch
+from torch import nn
+
+from transformers.activations import ACT2FN
+from transformers.cache_utils import Cache, DynamicCache, StaticCache, SlidingWindowCache
+from transformers.generation import GenerationMixin
+from transformers.integrations import use_kernel_forward_from_hub
+from transformers.modeling_attn_mask_utils import AttentionMaskConverter
+from transformers.modeling_flash_attention_utils import FlashAttentionKwargs
+from transformers.modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast
+from transformers.modeling_rope_utils import ROPE_INIT_FUNCTIONS, dynamic_rope_update
+from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel
+from transformers.processing_utils import Unpack
+from transformers.utils import TransformersKwargs, auto_docstring, can_return_tuple
+from transformers.utils.deprecation import deprecate_kwarg
+from transformers.utils.generic import check_model_inputs
+from .configuration_brumby import BrumbyConfig
+
+
+try:
+    from retention.triton import power_retention, power_retention_inference
+    from retention._utils import compute_expanded_dim
+except ImportError:
+    raise ImportError("Retention is required by the Brumby model. Please install it with `pip install retention`.")
+
+
+class PowerAttentionDynamicCache(DynamicCache):
+    """
+    A dynamic cache that encompasses 2 sets of caches:
+    1. attention cache with a short seq_len (determined by the chunk_size parameter):
+       - key_cache: [batch_size, num_heads, chunk_size, head_dim]
+       - value_cache: [batch_size, num_heads, chunk_size, head_dim]
+       - gating_cache: [batch_size, num_heads, chunk_size]
+    2. fixed-size state-based cache:
+       - state: [batch_size, num_heads, state_dim, head_dim]
+       - sum_of_keys: [batch_size, num_heads, state_dim]
+
+    where state_dim is determined by the power of expansion for power attention.
+    """
+    def __init__(self, config: BrumbyConfig, batch_size: int, dtype=torch.bfloat16, device=None):
+        super().__init__()
+        self.config = config
+        self.batch_size = batch_size
+        self.device = device
+        self.dtype = dtype
+        self.chunk_size = config.chunk_size
+        self.head_dim = config.hidden_size // config.num_attention_heads
+        self.p = config.p
+        self.state_dim = compute_expanded_dim(self.head_dim, deg=self.p)
+        
+        self.states = [None for _ in range(config.num_hidden_layers)]
+        self.sum_of_keys = [None for _ in range(config.num_hidden_layers)]
+        self.key_cache = [torch.tensor([[]] * batch_size, device=device, dtype=dtype) for _ in range(config.num_hidden_layers)]
+        self.value_cache = [torch.tensor([[]] * batch_size, device=device, dtype=dtype) for _ in range(config.num_hidden_layers)]
+        self.gate_cache = [torch.tensor([[]] * batch_size, device=device, dtype=torch.float32) for _ in range(config.num_hidden_layers)]
+
+    def clean_cache(self, layer_idx: int) -> None:
+        self.key_cache[layer_idx] = torch.tensor([[]] * self.batch_size, device=self.device, dtype=self.dtype)
+        self.value_cache[layer_idx] = torch.tensor([[]] * self.batch_size, device=self.device, dtype=self.dtype)
+        self.gate_cache[layer_idx] = torch.tensor([[]] * self.batch_size, device=self.device, dtype=torch.float32)
+
+    def update_cache(
+        self,
+        key_states: torch.Tensor,
+        value_states: torch.Tensor,
+        gate_states: torch.Tensor,
+        layer_idx: int,
+        cache_kwargs: Optional[Dict[str, Any]] = None,
+    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        if self.key_cache[layer_idx].shape[-1] == 0:
+            self.key_cache[layer_idx] = key_states
+            self.value_cache[layer_idx] = value_states
+            self.gate_cache[layer_idx] = gate_states
+        else:
+            self.key_cache[layer_idx] = torch.cat([self.key_cache[layer_idx], key_states], dim=2)
+            self.value_cache[layer_idx] = torch.cat([self.value_cache[layer_idx], value_states], dim=2)
+            self.gate_cache[layer_idx] = torch.cat([self.gate_cache[layer_idx], gate_states], dim=2)
+        return self.key_cache[layer_idx], self.value_cache[layer_idx], self.gate_cache[layer_idx], self.states[layer_idx], self.sum_of_keys[layer_idx]
+
+    def reorder_cache(self, beam_idx: torch.LongTensor):
+        """Reorders the cache for beam search, given the selected beam indices."""
+        for layer_idx in range(len(self.key_cache)):
+            device = self.key_cache[layer_idx].device
+            self.key_cache[layer_idx] = self.key_cache[layer_idx].index_select(0, beam_idx.to(device))
+            device = self.value_cache[layer_idx].device
+            self.value_cache[layer_idx] = self.value_cache[layer_idx].index_select(0, beam_idx.to(device))
+            device = self.gate_cache[layer_idx].device
+            self.gate_cache[layer_idx] = self.gate_cache[layer_idx].index_select(0, beam_idx.to(device))
+            device = self.states[layer_idx].device
+            self.states[layer_idx] = self.states[layer_idx].index_select(0, beam_idx.to(device))
+            device = self.sum_of_keys[layer_idx].device
+            self.sum_of_keys[layer_idx] = self.sum_of_keys[layer_idx].index_select(0, beam_idx.to(device))
+
+    def get_seq_length(self, layer_idx: Optional[int] = 0) -> int:
+        """Returns the sequence length of the cached states. A layer index can be optionally passed."""
+        # take any layer that contains cache and not empty tensor
+        if layer_idx is None:
+            layer_idx = 0
+        if layer_idx >= len(self.key_cache):
+            return 0
+        # Check if the cache for this layer is empty
+        if self.key_cache[layer_idx].numel() == 0:
+            return 0
+        return self.key_cache[layer_idx].shape[-2]
+
+    def to_legacy_cache(self) -> tuple[tuple[torch.Tensor, torch.Tensor]]:
+        raise NotImplementedError("PowerAttentionDynamicCache does not have a legacy cache equivalent.")
+
+    @classmethod
+    def from_legacy_cache(cls, past_key_values: Optional[Tuple[Tuple[torch.FloatTensor]]] = None) -> "PowerAttentionDynamicCache":
+        raise NotImplementedError("PowerAttentionDynamicCache does not have a legacy cache equivalent.")
+
+    def update_state(self, layer_idx: int, new_state: torch.Tensor, new_sum_of_keys: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
+        self.states[layer_idx] = new_state
+        self.sum_of_keys[layer_idx] = new_sum_of_keys
+        return self.states[layer_idx], self.sum_of_keys[layer_idx]
+
+    def reset(self):
+        self.states.zero_()
+        self.sum_of_keys.zero_()
+        self.key_cache.zero_()
+        self.value_cache.zero_()
+        self.gate_cache.zero_()
+
+
+
+@use_kernel_forward_from_hub("RMSNorm")
+class BrumbyRMSNorm(nn.Module):
+    def __init__(self, hidden_size, eps: float = 1e-6) -> None:
+        """
+        BrumbyRMSNorm is equivalent to T5LayerNorm
+        """
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(hidden_size))
+        self.variance_epsilon = eps
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        input_dtype = hidden_states.dtype
+        hidden_states = hidden_states.to(torch.float32)
+        variance = hidden_states.pow(2).mean(-1, keepdim=True)
+        hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)
+        return self.weight * hidden_states.to(input_dtype)
+
+    def extra_repr(self):
+        return f"{tuple(self.weight.shape)}, eps={self.variance_epsilon}"
+
+
+class BrumbyMLP(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.config = config
+        self.hidden_size = config.hidden_size
+        self.intermediate_size = config.intermediate_size
+        self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
+        self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
+        self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=False)
+        self.act_fn = ACT2FN[config.hidden_act]
+
+    def forward(self, x):
+        down_proj = self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x))
+        return down_proj
+
+
+def rotate_half(x):
+    """Rotates half the hidden dims of the input."""
+    x1 = x[..., : x.shape[-1] // 2]
+    x2 = x[..., x.shape[-1] // 2 :]
+    return torch.cat((-x2, x1), dim=-1)
+
+
+def apply_rotary_pos_emb(q, k, cos, sin, position_ids=None, unsqueeze_dim=1):
+    """Applies Rotary Position Embedding to the query and key tensors.
+
+    Args:
+        q (`torch.Tensor`): The query tensor.
+        k (`torch.Tensor`): The key tensor.
+        cos (`torch.Tensor`): The cosine part of the rotary embedding.
+        sin (`torch.Tensor`): The sine part of the rotary embedding.
+        position_ids (`torch.Tensor`, *optional*):
+            Deprecated and unused.
+        unsqueeze_dim (`int`, *optional*, defaults to 1):
+            The 'unsqueeze_dim' argument specifies the dimension along which to unsqueeze cos[position_ids] and
+            sin[position_ids] so that they can be properly broadcasted to the dimensions of q and k. For example, note
+            that cos[position_ids] and sin[position_ids] have the shape [batch_size, seq_len, head_dim]. Then, if q and
+            k have the shape [batch_size, heads, seq_len, head_dim], then setting unsqueeze_dim=1 makes
+            cos[position_ids] and sin[position_ids] broadcastable to the shapes of q and k. Similarly, if q and k have
+            the shape [batch_size, seq_len, heads, head_dim], then set unsqueeze_dim=2.
+    Returns:
+        `tuple(torch.Tensor)` comprising of the query and key tensors rotated using the Rotary Position Embedding.
+    """
+    cos = cos.unsqueeze(unsqueeze_dim)
+    sin = sin.unsqueeze(unsqueeze_dim)
+    q_embed = (q * cos) + (rotate_half(q) * sin)
+    k_embed = (k * cos) + (rotate_half(k) * sin)
+    return q_embed, k_embed
+
+
+def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor:
+    """
+    This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch,
+    num_key_value_heads, seqlen, head_dim) to (batch, num_attention_heads, seqlen, head_dim)
+    """
+    batch, num_key_value_heads, slen, head_dim = hidden_states.shape
+    if n_rep == 1:
+        return hidden_states
+    hidden_states = hidden_states[:, :, None, :, :].expand(batch, num_key_value_heads, n_rep, slen, head_dim)
+    return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim)
+
+
+def eager_attention_forward(
+    module: nn.Module,
+    query: torch.Tensor,
+    key: torch.Tensor,
+    value: torch.Tensor,
+    attention_mask: Optional[torch.Tensor],
+    scaling: float,
+    dropout: float = 0.0,
+    **kwargs: Unpack[TransformersKwargs],
+):
+    key_states = repeat_kv(key, module.num_key_value_groups)
+    value_states = repeat_kv(value, module.num_key_value_groups)
+
+    attn_weights = torch.matmul(query, key_states.transpose(2, 3)) * scaling
+    if attention_mask is not None:
+        causal_mask = attention_mask[:, :, :, : key_states.shape[-2]]
+        attn_weights = attn_weights + causal_mask
+
+    attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query.dtype)
+    attn_weights = nn.functional.dropout(attn_weights, p=dropout, training=module.training)
+    attn_output = torch.matmul(attn_weights, value_states)
+    attn_output = attn_output.transpose(1, 2).contiguous()
+
+    return attn_output, attn_weights
+
+
+class BrumbyAttention(nn.Module):
+    """Multi-headed attention from 'Attention Is All You Need' paper"""
+
+    def __init__(self, config: BrumbyConfig, layer_idx: int):
+        super().__init__()
+        self.config = config
+        self.layer_idx = layer_idx
+        self.head_dim = getattr(config, "head_dim", config.hidden_size // config.num_attention_heads)
+        self.num_key_value_groups = config.num_attention_heads // config.num_key_value_heads
+        self.scaling = self.head_dim**-0.5
+        self.attention_dropout = config.attention_dropout
+        self.is_causal = True
+        self.use_exp = config.use_exp
+        self.prefill_chunk_size = config.prefill_chunk_size
+        self.chunk_size = config.chunk_size
+        self.switch_over_seq_len = config.switch_over_seq_len
+
+        self.q_proj = nn.Linear(
+            config.hidden_size, config.num_attention_heads * self.head_dim, bias=config.attention_bias
+        )
+        self.k_proj = nn.Linear(
+            config.hidden_size, config.num_key_value_heads * self.head_dim, bias=config.attention_bias
+        )
+        self.v_proj = nn.Linear(
+            config.hidden_size, config.num_key_value_heads * self.head_dim, bias=config.attention_bias
+        )
+        self.g_proj = nn.Linear(
+            config.hidden_size, config.num_key_value_heads, bias=config.attention_bias
+        )
+        self.o_proj = nn.Linear(
+            config.num_attention_heads * self.head_dim, config.hidden_size, bias=config.attention_bias
+        )
+        self.q_norm = BrumbyRMSNorm(self.head_dim, eps=config.rms_norm_eps)  # unlike olmo, only on the head dim!
+        self.k_norm = BrumbyRMSNorm(self.head_dim, eps=config.rms_norm_eps)  # thus post q_norm does not need reshape
+        self.sliding_window = config.sliding_window if config.layer_types[layer_idx] == "sliding_attention" else None
+
+    @deprecate_kwarg("past_key_value", new_name="past_key_values", version="4.58")
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        position_embeddings: tuple[torch.Tensor, torch.Tensor],
+        attention_mask: Optional[torch.Tensor],
+        past_key_values: Optional[PowerAttentionDynamicCache] = None,
+        cache_position: Optional[torch.LongTensor] = None,
+        **kwargs: Unpack[FlashAttentionKwargs],
+    ) -> tuple[torch.Tensor, Optional[torch.Tensor]]:
+        input_shape = hidden_states.shape[:-1]
+        hidden_shape = (*input_shape, -1, self.head_dim)
+
+        query_states = self.q_norm(self.q_proj(hidden_states).view(hidden_shape)).transpose(1, 2)
+        key_states = self.k_norm(self.k_proj(hidden_states).view(hidden_shape)).transpose(1, 2)
+        value_states = self.v_proj(hidden_states).view(hidden_shape).transpose(1, 2)
+        gate_states = self.g_proj(hidden_states).view(hidden_shape[:-1]).transpose(1, 2)
+        gate_states = nn.functional.logsigmoid(gate_states.to(torch.float32))
+
+        cos, sin = position_embeddings
+        query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin)
+
+        if past_key_values is not None:
+            # sin and cos are specific to RoPE models; cache_position needed for the static cache
+            cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position}
+            key_states, value_states, gate_states, state, sum_of_keys = past_key_values.update_cache(key_states, value_states, gate_states, self.layer_idx, cache_kwargs)
+
+        if self.use_exp:
+            attention_interface: Callable = eager_attention_forward
+            if self.config._attn_implementation != "eager":
+                attention_interface = ALL_ATTENTION_FUNCTIONS[self.config._attn_implementation]
+
+            attn_output, attn_weights = attention_interface(
+                self,
+                query_states,
+                key_states,
+                value_states,
+                attention_mask,
+                dropout=0.0 if not self.training else self.attention_dropout,
+                scaling=self.scaling,
+                sliding_window=self.sliding_window,  # diff with Llama
+                **kwargs,
+            )
+        
+        elif query_states.shape[2] == 1:
+            key_len = key_states.shape[2]
+            attn_output, state, sum_of_keys = power_retention_inference(
+                query_states.transpose(1, 2),
+                key_states.transpose(1, 2),
+                value_states.transpose(1, 2),
+                gate_states.transpose(1, 2),
+                initial_state=state,
+                sum_of_keys=sum_of_keys,
+                deg=2,
+                scale=self.scaling,
+                switch_over_seq_len=self.chunk_size,
+            )
+            if self.chunk_size is not None and key_len >= self.chunk_size:
+                past_key_values.clean_cache(self.layer_idx)
+                past_key_values.update_state(self.layer_idx, state, sum_of_keys)
+
+            attn_weights = None
+
+        else:
+            key_len = key_states.shape[2]
+            attn_output = power_retention(
+                query_states.transpose(1, 2),
+                key_states.transpose(1, 2),
+                value_states.transpose(1, 2),
+                gate_states.transpose(1, 2),
+                deg=2,
+                scale=self.scaling,
+                chunk_size=self.prefill_chunk_size, # enable chunked prefilling by default
+                switch_over_seq_len=self.switch_over_seq_len,
+            )
+            attn_weights = None
+
+
+        attn_output = attn_output.reshape(*input_shape, -1).contiguous()
+        attn_output = self.o_proj(attn_output)
+        return attn_output, attn_weights
+
+
+class BrumbyDecoderLayer(nn.Module):
+    def __init__(self, config: BrumbyConfig, layer_idx: int):
+        super().__init__()
+        self.hidden_size = config.hidden_size
+
+        self.self_attn = BrumbyAttention(config=config, layer_idx=layer_idx)
+
+        self.mlp = BrumbyMLP(config)
+        self.input_layernorm = BrumbyRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.post_attention_layernorm = BrumbyRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.attention_type = config.layer_types[layer_idx]
+
+    @deprecate_kwarg("past_key_value", new_name="past_key_values", version="4.58")
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_values: Optional[Cache] = None,
+        use_cache: Optional[bool] = False,
+        cache_position: Optional[torch.LongTensor] = None,
+        position_embeddings: Optional[tuple[torch.Tensor, torch.Tensor]] = None,  # necessary, but kept here for BC
+        **kwargs: Unpack[TransformersKwargs],
+    ) -> torch.Tensor:
+        residual = hidden_states
+        hidden_states = self.input_layernorm(hidden_states)
+        # Self Attention
+        hidden_states, _ = self.self_attn(
+            hidden_states=hidden_states,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            past_key_values=past_key_values,
+            use_cache=use_cache,
+            cache_position=cache_position,
+            position_embeddings=position_embeddings,
+            **kwargs,
+        )
+        hidden_states = residual + hidden_states
+
+        # Fully Connected
+        residual = hidden_states
+        hidden_states = self.post_attention_layernorm(hidden_states)
+        hidden_states = self.mlp(hidden_states)
+        hidden_states = residual + hidden_states
+        return hidden_states
+
+
+@auto_docstring
+class BrumbyPreTrainedModel(PreTrainedModel):
+    config: BrumbyConfig
+    base_model_prefix = "model"
+    supports_gradient_checkpointing = True
+    _no_split_modules = ["BrumbyDecoderLayer"]
+    _skip_keys_device_placement = ["past_key_values"]
+    _supports_flash_attn = True
+    _supports_sdpa = True
+    _supports_flex_attn = True
+
+    _can_compile_fullgraph = True
+    _supports_attention_backend = True
+    _can_record_outputs = {
+        "hidden_states": BrumbyDecoderLayer,
+        "attentions": BrumbyAttention,
+    }
+
+
+class BrumbyRotaryEmbedding(nn.Module):
+    inv_freq: torch.Tensor  # fix linting for `register_buffer`
+
+    def __init__(self, config: BrumbyConfig, device=None):
+        super().__init__()
+        # BC: "rope_type" was originally "type"
+        if hasattr(config, "rope_scaling") and isinstance(config.rope_scaling, dict):
+            self.rope_type = config.rope_scaling.get("rope_type", config.rope_scaling.get("type"))
+        else:
+            self.rope_type = "default"
+        self.max_seq_len_cached = config.max_position_embeddings
+        self.original_max_seq_len = config.max_position_embeddings
+
+        self.config = config
+        self.rope_init_fn = ROPE_INIT_FUNCTIONS[self.rope_type]
+
+        inv_freq, self.attention_scaling = self.rope_init_fn(self.config, device)
+        self.register_buffer("inv_freq", inv_freq, persistent=False)
+        self.original_inv_freq = self.inv_freq
+
+    @torch.no_grad()
+    @dynamic_rope_update  # power user: used with advanced RoPE types (e.g. dynamic rope)
+    def forward(self, x, position_ids):
+        inv_freq_expanded = self.inv_freq[None, :, None].float().expand(position_ids.shape[0], -1, 1).to(x.device)
+        position_ids_expanded = position_ids[:, None, :].float()
+
+        device_type = x.device.type if isinstance(x.device.type, str) and x.device.type != "mps" else "cpu"
+        with torch.autocast(device_type=device_type, enabled=False):  # Force float32
+            freqs = (inv_freq_expanded.float() @ position_ids_expanded.float()).transpose(1, 2)
+            emb = torch.cat((freqs, freqs), dim=-1)
+            cos = emb.cos() * self.attention_scaling
+            sin = emb.sin() * self.attention_scaling
+
+        return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype)
+
+
+@auto_docstring
+class BrumbyModel(BrumbyPreTrainedModel):
+    def __init__(self, config: BrumbyConfig):
+        super().__init__(config)
+        self.padding_idx = config.pad_token_id
+        self.vocab_size = config.vocab_size
+
+        self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx)
+        self.layers = nn.ModuleList(
+            [BrumbyDecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)]
+        )
+        self.norm = BrumbyRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.rotary_emb = BrumbyRotaryEmbedding(config=config)
+        self.gradient_checkpointing = False
+        self.has_sliding_layers = "sliding_attention" in self.config.layer_types
+
+        # Initialize weights and apply final processing
+        self.post_init()
+
+    @check_model_inputs
+    @auto_docstring
+    def forward(
+        self,
+        input_ids: Optional[torch.LongTensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_values: Optional[PowerAttentionDynamicCache] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        use_cache: Optional[bool] = None,
+        cache_position: Optional[torch.LongTensor] = None,
+        **kwargs: Unpack[TransformersKwargs],
+    ) -> BaseModelOutputWithPast:
+        if (input_ids is None) ^ (inputs_embeds is not None):
+            raise ValueError("You must specify exactly one of input_ids or inputs_embeds")
+
+        if inputs_embeds is None:
+            inputs_embeds = self.embed_tokens(input_ids)
+
+        if use_cache and past_key_values is None:
+            raise ValueError("Brumby requires an initialized `PowerAttentionDynamicCache` to return a cache. None was provided, so no cache will be returned.")
+
+        if cache_position is None:
+            past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0
+            cache_position = torch.arange(
+                past_seen_tokens, past_seen_tokens + inputs_embeds.shape[1], device=inputs_embeds.device
+            )
+
+        if position_ids is None:
+            position_ids = cache_position.unsqueeze(0)
+
+        causal_mask = self._update_causal_mask(
+            attention_mask, inputs_embeds, cache_position, past_key_values, output_attentions=False
+        )
+
+        hidden_states = inputs_embeds
+
+        # create position embeddings to be shared across the decoder layers
+        position_embeddings = self.rotary_emb(hidden_states, position_ids)
+
+        for decoder_layer in self.layers[: self.config.num_hidden_layers]:
+            hidden_states = decoder_layer(
+                hidden_states,
+                attention_mask=causal_mask,
+                position_ids=position_ids,
+                past_key_values=past_key_values,
+                use_cache=use_cache,
+                cache_position=cache_position,
+                position_embeddings=position_embeddings,
+                **kwargs,
+            )
+
+        hidden_states = self.norm(hidden_states)
+        return BaseModelOutputWithPast(
+            last_hidden_state=hidden_states,
+            past_key_values=past_key_values if use_cache else None,
+        )
+
+    def _update_causal_mask(
+        self,
+        attention_mask: torch.Tensor,
+        input_tensor: torch.Tensor,
+        cache_position: torch.Tensor,
+        past_key_values: Cache,
+        output_attentions: bool = False,
+    ):
+        if self.config._attn_implementation == "flash_attention_2":
+            if attention_mask is not None and past_key_values is not None:
+                is_padding_right = attention_mask[:, -1].sum().item() != input_tensor.size()[0]
+                if is_padding_right:
+                    raise ValueError(
+                        "You are attempting to perform batched generation with padding_side='right'"
+                        " this may lead to unexpected behaviour for Flash Attention version of Qwen3. Make sure to "
+                        " call `tokenizer.padding_side  = 'left'` before tokenizing the input. "
+                    )
+            if attention_mask is not None and 0.0 in attention_mask:
+                return attention_mask
+            return None
+
+        # For SDPA, when possible, we will rely on its `is_causal` argument instead of its `attn_mask` argument, in
+        # order to dispatch on Flash Attention 2. This feature is not compatible with static cache, as SDPA will fail
+        # to infer the attention mask.
+        past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0
+        using_static_cache = isinstance(past_key_values, StaticCache)
+        using_sliding_window_cache = isinstance(past_key_values, SlidingWindowCache)
+
+        # When output attentions is True, sdpa implementation's forward method calls the eager implementation's forward
+        if (
+            self.config._attn_implementation == "sdpa"
+            and not (using_static_cache or using_sliding_window_cache)
+            and not output_attentions
+        ):
+            if AttentionMaskConverter._ignore_causal_mask_sdpa(
+                attention_mask,
+                inputs_embeds=input_tensor,
+                past_key_values_length=past_seen_tokens,
+                sliding_window=self.config.sliding_window,
+                is_training=self.training,
+            ):
+                return None
+
+        dtype, device = input_tensor.dtype, input_tensor.device
+        min_dtype = torch.finfo(dtype).min
+        sequence_length = input_tensor.shape[1]
+        # SlidingWindowCache or StaticCache
+        if using_sliding_window_cache or using_static_cache:
+            target_length = past_key_values.get_max_cache_shape()
+        # DynamicCache or no cache
+        else:
+            target_length = (
+                attention_mask.shape[-1]
+                if isinstance(attention_mask, torch.Tensor)
+                else past_seen_tokens + sequence_length + 1
+            )
+
+        # In case the provided `attention` mask is 2D, we generate a causal mask here (4D).
+        causal_mask = self._prepare_4d_causal_attention_mask_with_cache_position(
+            attention_mask,
+            sequence_length=sequence_length,
+            target_length=target_length,
+            dtype=dtype,
+            device=device,
+            cache_position=cache_position,
+            batch_size=input_tensor.shape[0],
+            config=self.config,
+            past_key_values=past_key_values,
+        )
+
+        if (
+            self.config._attn_implementation == "sdpa"
+            and attention_mask is not None
+            and attention_mask.device.type in ["cuda", "xpu"]
+            and not output_attentions
+        ):
+            # Attend to all tokens in fully masked rows in the causal_mask, for example the relevant first rows when
+            # using left padding. This is required by F.scaled_dot_product_attention memory-efficient attention path.
+            # Details: https://github.com/pytorch/pytorch/issues/110213
+            causal_mask = AttentionMaskConverter._unmask_unattended(causal_mask, min_dtype)
+
+        return causal_mask
+
+    @staticmethod
+    def _prepare_4d_causal_attention_mask_with_cache_position(
+        attention_mask: torch.Tensor,
+        sequence_length: int,
+        target_length: int,
+        dtype: torch.dtype,
+        device: torch.device,
+        cache_position: torch.Tensor,
+        batch_size: int,
+        config: BrumbyConfig,
+        past_key_values: Cache,
+    ):
+        """
+        Creates a causal 4D mask of shape `(batch_size, 1, query_length, key_value_length)` from a 2D mask of shape
+        `(batch_size, key_value_length)`, or if the input `attention_mask` is already 4D, do nothing.
+
+        Args:
+            attention_mask (`torch.Tensor`):
+                A 2D attention mask of shape `(batch_size, key_value_length)` or a 4D attention mask of shape `(batch_size, 1, query_length, key_value_length)`.
+            sequence_length (`int`):
+                The sequence length being processed.
+            target_length (`int`):
+                The target length: when generating with static cache, the mask should be as long as the static cache, to account for the 0 padding, the part of the cache that is not filled yet.
+            dtype (`torch.dtype`):
+                The dtype to use for the 4D attention mask.
+            device (`torch.device`):
+                The device to place the 4D attention mask on.
+            cache_position (`torch.Tensor`):
+                Indices depicting the position of the input sequence tokens in the sequence.
+            batch_size (`torch.Tensor`):
+                Batch size.
+            config (`Qwen3Config`):
+                The model's configuration class
+            past_key_values (`Cache`):
+                The cache class that is being used currently to generate
+        """
+        if attention_mask is not None and attention_mask.dim() == 4:
+            # In this case we assume that the mask comes already in inverted form and requires no inversion or slicing.
+            causal_mask = attention_mask
+        else:
+            min_dtype = torch.finfo(dtype).min
+            causal_mask = torch.full(
+                (sequence_length, target_length), fill_value=min_dtype, dtype=dtype, device=device
+            )
+            diagonal_attend_mask = torch.arange(target_length, device=device) > cache_position.reshape(-1, 1)
+            if config.sliding_window is not None:
+                # if we have sliding window, we should not attend to tokens beyond sliding window length, so we mask them out also
+                # the check is needed to verify is current checkpoint was trained with sliding window or not
+                if not isinstance(past_key_values, SlidingWindowCache) or sequence_length > target_length:
+                    sliding_attend_mask = torch.arange(target_length, device=device) <= (
+                        cache_position.reshape(-1, 1) - config.sliding_window
+                    )
+                    diagonal_attend_mask.bitwise_or_(sliding_attend_mask)
+            causal_mask *= diagonal_attend_mask
+            causal_mask = causal_mask[None, None, :, :].expand(batch_size, 1, -1, -1)
+            if attention_mask is not None:
+                causal_mask = causal_mask.clone()  # copy to contiguous memory for in-place edit
+                if attention_mask.shape[-1] > target_length:
+                    attention_mask = attention_mask[:, :target_length]
+                mask_length = attention_mask.shape[-1]
+                padding_mask = causal_mask[:, :, :, :mask_length] + attention_mask[:, None, None, :].to(
+                    causal_mask.device
+                )
+                padding_mask = padding_mask == 0
+                causal_mask[:, :, :, :mask_length] = causal_mask[:, :, :, :mask_length].masked_fill(
+                    padding_mask, min_dtype
+                )
+        return causal_mask
+
+
+@auto_docstring
+class BrumbyForCausalLM(BrumbyPreTrainedModel, GenerationMixin):
+    _tied_weights_keys = ["lm_head.weight"]
+    _tp_plan = {"lm_head": "colwise_rep"}
+    _pp_plan = {"lm_head": (["hidden_states"], ["logits"])}
+
+    def __init__(self, config):
+        super().__init__(config)
+        self.model = BrumbyModel(config)
+        self.vocab_size = config.vocab_size
+        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
+
+        # Initialize weights and apply final processing
+        self.post_init()
+
+    def prepare_inputs_for_generation(
+        self,
+        input_ids,
+        past_key_values=None,
+        attention_mask=None,
+        inputs_embeds=None,
+        cache_position=None,
+        position_ids=None,
+        use_cache=True,
+        **kwargs,
+    ):
+        # Copy from https://github.com/huggingface/transformers/blob/main/src/transformers/models/jamba/modeling_jamba.py
+        # Overwitten -- uses `past_key_values` as opposed to `past_key_values`
+        empty_past_kv = past_key_values is None
+        
+        # First pass
+        if not isinstance(past_key_values, PowerAttentionDynamicCache):
+            past_key_values = PowerAttentionDynamicCache(self.config, input_ids.shape[0], self.dtype, device=self.device)
+
+        # If we have cache: let's slice `input_ids` through `cache_position`, to keep only the unprocessed tokens
+        # Exception 1: when passing input_embeds, input_ids may be missing entries
+        # Exception 2: some generation methods do special slicing of input_ids, so we don't need to do it here
+        # Exception 3: with synced GPUs cache_position may go out of bounds, but we only want dummy token in that case.
+        #              (we can't check exception 3 while compiling)
+        if not empty_past_kv:
+            if (
+                inputs_embeds is not None  # Exception 1
+                or cache_position[-1] >= input_ids.shape[1]  # Exception 3
+            ):
+                input_ids = input_ids[:, -cache_position.shape[0] :]
+            elif input_ids.shape[1] != cache_position.shape[0]:  # Default case (the "else", a no op, is Exception 2)
+                input_ids = input_ids[:, cache_position]
+        else:
+            past_key_values = PowerAttentionDynamicCache(
+                self.config, input_ids.shape[0], self.dtype, device=self.device
+            )
+
+        if attention_mask is not None and position_ids is None:
+            # create position_ids on the fly for batch generation
+            position_ids = attention_mask.long().cumsum(-1) - 1
+            position_ids.masked_fill_(attention_mask == 0, 1)
+            if not empty_past_kv:
+                position_ids = position_ids[:, -input_ids.shape[1] :]
+
+        # if `inputs_embeds` are passed, we only want to use them in the 1st generation step
+        if inputs_embeds is not None and empty_past_kv:
+            # TODO(pjin): workaround fix for properly extending inputs_embeds;
+            # longer term, may be better handled elsewhere in .generate().
+            if input_ids is not None and inputs_embeds.shape[1] < input_ids.shape[1]:
+                new_token_embeds = self.get_input_embeddings()(input_ids[:,inputs_embeds.shape[1]:])
+                inputs_embeds = torch.cat([inputs_embeds, new_token_embeds], dim=1)
+            model_inputs = {"inputs_embeds": inputs_embeds}
+        else:
+            model_inputs = {"input_ids": input_ids.contiguous()}  # `contiguous()` needed for compilation use cases
+
+        model_inputs.update(
+            {
+                "position_ids": position_ids,
+                "past_key_values": past_key_values,
+                "use_cache": use_cache,
+                "attention_mask": attention_mask,
+                "cache_position": cache_position,
+            }
+        )
+        return model_inputs
+
+    @can_return_tuple
+    @auto_docstring
+    def forward(
+        self,
+        input_ids: Optional[torch.LongTensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_values: Optional[PowerAttentionDynamicCache] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        labels: Optional[torch.LongTensor] = None,
+        use_cache: Optional[bool] = None,
+        cache_position: Optional[torch.LongTensor] = None,
+        logits_to_keep: Union[int, torch.Tensor] = 0,
+        **kwargs: Unpack[TransformersKwargs],
+    ) -> CausalLMOutputWithPast:
+        r"""
+        labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
+            Labels for computing the masked language modeling loss. Indices should either be in `[0, ...,
+            config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored
+            (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`.
+
+        Example:
+
+        ```python
+        >>> from transformers import AutoTokenizer, BrumbyForCausalLM
+
+        >>> model = BrumbyForCausalLM.from_pretrained("Qwen/Brumby-8B")
+        >>> tokenizer = AutoTokenizer.from_pretrained("Qwen/Brumby-8B")
+
+        >>> prompt = "Hey, are you conscious? Can you talk to me?"
+        >>> inputs = tokenizer(prompt, return_tensors="pt")
+
+        >>> # Generate
+        >>> generate_ids = model.generate(inputs.input_ids, max_length=30)
+        >>> tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]
+        "Hey, are you conscious? Can you talk to me?\nI'm not conscious, but I can talk to you."
+        ```"""
+        outputs: BaseModelOutputWithPast = self.model(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            past_key_values=past_key_values,
+            inputs_embeds=inputs_embeds,
+            use_cache=use_cache,
+            cache_position=cache_position,
+            **kwargs,
+        )
+
+        hidden_states = outputs.last_hidden_state
+        # Only compute necessary logits, and do not upcast them to float if we are not computing the loss
+        slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep
+        logits = self.lm_head(hidden_states[:, slice_indices, :])
+
+        loss = None
+        if labels is not None:
+            loss = self.loss_function(logits=logits, labels=labels, vocab_size=self.config.vocab_size, **kwargs)
+
+        return CausalLMOutputWithPast(
+            loss=loss,
+            logits=logits,
+            past_key_values=outputs.past_key_values,
+            hidden_states=outputs.hidden_states,
+            attentions=outputs.attentions,
+        )
+
+
+__all__ = [
+    "BrumbyForCausalLM",
+    "BrumbyPreTrainedModel",
+    "BrumbyModel",
+]

quant_log.csv

@@ -0,0 +1,281 @@
+layer,module,loss,samples,damp,time
+0,self_attn.v_proj,0.0000000055,0.05000,4.851
+0,self_attn.k_proj,0.0000000067,0.05000,4.868
+0,self_attn.q_proj,0.0000000245,0.05000,4.875
+0,self_attn.o_proj,0.0000001023,0.05000,1.463
+0,mlp.gate_proj,0.0000014832,0.05000,3.275
+0,mlp.up_proj,0.0000012088,0.05000,3.334
+0,mlp.down_proj,0.0000007021,0.05000,5.384
+1,self_attn.q_proj,0.0000000516,0.05000,4.584
+1,self_attn.v_proj,0.0000000140,0.05000,4.683
+1,self_attn.k_proj,0.0000000128,0.05000,4.723
+1,self_attn.o_proj,0.0000001024,0.05000,1.334
+1,mlp.up_proj,0.0000235446,0.05000,3.642
+1,mlp.gate_proj,0.0000679487,0.05000,3.677
+1,mlp.down_proj,0.0000010552,0.05000,5.334
+2,self_attn.k_proj,0.0000000361,0.05000,4.120
+2,self_attn.v_proj,0.0000000386,0.05000,4.189
+2,self_attn.q_proj,0.0000001411,0.05000,4.244
+2,self_attn.o_proj,0.0000002282,0.05000,1.332
+2,mlp.gate_proj,0.0000977210,0.05000,3.152
+2,mlp.up_proj,0.0000492783,0.05000,3.153
+2,mlp.down_proj,0.0000008231,0.05000,5.410
+3,self_attn.q_proj,0.0000003171,0.05000,4.748
+3,self_attn.k_proj,0.0000000801,0.05000,4.758
+3,self_attn.v_proj,0.0000000896,0.05000,4.779
+3,self_attn.o_proj,0.0000002193,0.05000,1.414
+3,mlp.gate_proj,0.0001480885,0.05000,3.153
+3,mlp.up_proj,0.0000745903,0.05000,3.224
+3,mlp.down_proj,0.0000013200,0.05000,5.582
+4,self_attn.v_proj,0.0000001349,0.05000,5.075
+4,self_attn.q_proj,0.0000004968,0.05000,5.101
+4,self_attn.k_proj,0.0000001217,0.05000,5.105
+4,self_attn.o_proj,0.0000003706,0.05000,1.499
+4,mlp.up_proj,0.0001282589,0.05000,3.353
+4,mlp.gate_proj,0.0002110812,0.05000,3.373
+4,mlp.down_proj,0.0000014401,0.05000,5.379
+5,self_attn.q_proj,0.0000005028,0.05000,4.952
+5,self_attn.v_proj,0.0000001383,0.05000,4.969
+5,self_attn.k_proj,0.0000001241,0.05000,4.978
+5,self_attn.o_proj,0.0000003825,0.05000,1.373
+5,mlp.up_proj,0.0000865598,0.05000,3.568
+5,mlp.gate_proj,0.0001970606,0.05000,3.592
+5,mlp.down_proj,0.0000036706,0.05000,5.314
+6,self_attn.v_proj,0.0000002469,0.05000,4.978
+6,self_attn.q_proj,0.0000008867,0.05000,5.061
+6,self_attn.k_proj,0.0000002098,0.05000,5.067
+6,self_attn.o_proj,0.0000005215,0.05000,1.377
+6,mlp.up_proj,0.0001740953,0.05000,2.960
+6,mlp.gate_proj,0.0003093244,0.05000,3.049
+6,mlp.down_proj,0.0000451711,0.05000,5.501
+7,self_attn.q_proj,0.0000032589,0.05000,4.941
+7,self_attn.k_proj,0.0000007494,0.05000,4.953
+7,self_attn.v_proj,0.0000008548,0.05000,4.969
+7,self_attn.o_proj,0.0000011350,0.05000,1.371
+7,mlp.gate_proj,0.0002542938,0.05000,3.356
+7,mlp.up_proj,0.0001292021,0.05000,3.370
+7,mlp.down_proj,0.0000053761,0.05000,5.684
+8,self_attn.v_proj,0.0000008264,0.05000,4.899
+8,self_attn.q_proj,0.0000030318,0.05000,4.924
+8,self_attn.k_proj,0.0000007100,0.05000,4.975
+8,self_attn.o_proj,0.0000013708,0.05000,1.409
+8,mlp.up_proj,0.0000463701,0.05000,3.645
+8,mlp.gate_proj,0.0000779276,0.05000,3.699
+8,mlp.down_proj,0.0000070508,0.05000,5.465
+9,self_attn.v_proj,0.0000006791,0.05000,5.006
+9,self_attn.q_proj,0.0000025509,0.05000,5.023
+9,self_attn.k_proj,0.0000006144,0.05000,5.032
+9,self_attn.o_proj,0.0000013103,0.05000,1.424
+9,mlp.gate_proj,0.0000426071,0.05000,3.332
+9,mlp.up_proj,0.0000384853,0.05000,3.368
+9,mlp.down_proj,0.0000094974,0.05000,5.585
+10,self_attn.k_proj,0.0000008871,0.05000,5.058
+10,self_attn.v_proj,0.0000010250,0.05000,5.070
+10,self_attn.q_proj,0.0000039080,0.05000,5.100
+10,self_attn.o_proj,0.0000023923,0.05000,1.394
+10,mlp.gate_proj,0.0000476698,0.05000,2.910
+10,mlp.up_proj,0.0000438941,0.05000,2.918
+10,mlp.down_proj,0.0000099568,0.05000,5.422
+11,self_attn.v_proj,0.0000016619,0.05000,5.104
+11,self_attn.q_proj,0.0000061273,0.05000,5.127
+11,self_attn.k_proj,0.0000014110,0.05000,5.151
+11,self_attn.o_proj,0.0000034863,0.05000,1.411
+11,mlp.up_proj,0.0000549275,0.05000,3.648
+11,mlp.gate_proj,0.0000592392,0.05000,3.697
+11,mlp.down_proj,0.0000136331,0.05000,5.316
+12,self_attn.k_proj,0.0000015385,0.05000,4.950
+12,self_attn.v_proj,0.0000017779,0.05000,4.997
+12,self_attn.q_proj,0.0000064274,0.05000,5.019
+12,self_attn.o_proj,0.0000026011,0.05000,1.423
+12,mlp.gate_proj,0.0000721614,0.05000,3.257
+12,mlp.up_proj,0.0000664742,0.05000,3.284
+12,mlp.down_proj,0.0000165714,0.05000,5.587
+13,self_attn.v_proj,0.0000035299,0.05000,4.510
+13,self_attn.q_proj,0.0000127921,0.05000,4.531
+13,self_attn.k_proj,0.0000029702,0.05000,4.564
+13,self_attn.o_proj,0.0000040227,0.05000,1.394
+13,mlp.up_proj,0.0000723846,0.05000,2.925
+13,mlp.gate_proj,0.0000838850,0.05000,2.940
+13,mlp.down_proj,0.0000187251,0.05000,5.496
+14,self_attn.v_proj,0.0000024263,0.05000,4.790
+14,self_attn.q_proj,0.0000085204,0.05000,4.817
+14,self_attn.k_proj,0.0000020249,0.05000,4.848
+14,self_attn.o_proj,0.0000040741,0.05000,1.353
+14,mlp.gate_proj,0.0000798361,0.05000,2.961
+14,mlp.up_proj,0.0000716579,0.05000,2.974
+14,mlp.down_proj,0.0000192025,0.05000,5.455
+15,self_attn.v_proj,0.0000024046,0.05000,4.509
+15,self_attn.k_proj,0.0000021407,0.05000,4.558
+15,self_attn.q_proj,0.0000092248,0.05000,4.569
+15,self_attn.o_proj,0.0000045808,0.05000,1.348
+15,mlp.up_proj,0.0000751190,0.05000,3.045
+15,mlp.gate_proj,0.0000782121,0.05000,3.079
+15,mlp.down_proj,0.0000209370,0.05000,5.471
+16,self_attn.q_proj,0.0000100945,0.05000,5.142
+16,self_attn.v_proj,0.0000027845,0.05000,5.142
+16,self_attn.k_proj,0.0000023761,0.05000,5.241
+16,self_attn.o_proj,0.0000062379,0.05000,1.380
+16,mlp.up_proj,0.0000748809,0.05000,2.870
+16,mlp.gate_proj,0.0000727667,0.05000,2.884
+16,mlp.down_proj,0.0000219730,0.05000,5.353
+17,self_attn.k_proj,0.0000028254,0.05000,4.584
+17,self_attn.q_proj,0.0000126655,0.05000,4.615
+17,self_attn.v_proj,0.0000032815,0.05000,4.642
+17,self_attn.o_proj,0.0000061028,0.05000,1.371
+17,mlp.gate_proj,0.0000801209,0.05000,3.312
+17,mlp.up_proj,0.0000837841,0.05000,3.340
+17,mlp.down_proj,0.0000245772,0.05000,5.519
+18,self_attn.v_proj,0.0000047651,0.05000,4.578
+18,self_attn.k_proj,0.0000041778,0.05000,4.587
+18,self_attn.q_proj,0.0000184517,0.05000,4.621
+18,self_attn.o_proj,0.0000063895,0.05000,1.366
+18,mlp.up_proj,0.0000915147,0.05000,3.025
+18,mlp.gate_proj,0.0000842209,0.05000,3.032
+18,mlp.down_proj,0.0000275443,0.05000,5.559
+19,self_attn.k_proj,0.0000049127,0.05000,4.855
+19,self_attn.v_proj,0.0000055556,0.05000,4.940
+19,self_attn.q_proj,0.0000216752,0.05000,4.962
+19,self_attn.o_proj,0.0000083279,0.05000,1.398
+19,mlp.up_proj,0.0001011631,0.05000,2.898
+19,mlp.gate_proj,0.0000937019,0.05000,2.930
+19,mlp.down_proj,0.0000327259,0.05000,5.399
+20,self_attn.v_proj,0.0000109982,0.05000,4.754
+20,self_attn.k_proj,0.0000086623,0.05000,4.805
+20,self_attn.q_proj,0.0000405637,0.05000,4.808
+20,self_attn.o_proj,0.0000105337,0.05000,1.392
+20,mlp.gate_proj,0.0000998617,0.05000,3.189
+20,mlp.up_proj,0.0001093770,0.05000,3.217
+20,mlp.down_proj,0.0000415528,0.05000,5.438
+21,self_attn.k_proj,0.0000117670,0.05000,4.439
+21,self_attn.q_proj,0.0000514724,0.05000,4.453
+21,self_attn.v_proj,0.0000140596,0.05000,4.516
+21,self_attn.o_proj,0.0000151081,0.05000,1.453
+21,mlp.gate_proj,0.0001088591,0.05000,2.926
+21,mlp.up_proj,0.0001167877,0.05000,2.938
+21,mlp.down_proj,0.0000477426,0.05000,5.551
+22,self_attn.v_proj,0.0000117180,0.05000,4.538
+22,self_attn.k_proj,0.0000094837,0.05000,4.573
+22,self_attn.q_proj,0.0000441959,0.05000,4.617
+22,self_attn.o_proj,0.0000190037,0.05000,1.410
+22,mlp.up_proj,0.0001366769,0.05000,3.587
+22,mlp.gate_proj,0.0001259929,0.05000,3.622
+22,mlp.down_proj,0.0000632469,0.05000,5.323
+23,self_attn.v_proj,0.0000206423,0.05000,5.066
+23,self_attn.q_proj,0.0000767674,0.05000,5.099
+23,self_attn.k_proj,0.0000157870,0.05000,5.111
+23,self_attn.o_proj,0.0000160790,0.05000,1.398
+23,mlp.up_proj,0.0001518091,0.05000,3.378
+23,mlp.gate_proj,0.0001383609,0.05000,3.395
+23,mlp.down_proj,0.0000856998,0.05000,5.398
+24,self_attn.v_proj,0.0000329967,0.05000,5.030
+24,self_attn.k_proj,0.0000237478,0.05000,5.054
+24,self_attn.q_proj,0.0001217306,0.05000,5.056
+24,self_attn.o_proj,0.0000220392,0.05000,1.360
+24,mlp.gate_proj,0.0001632924,0.05000,3.625
+24,mlp.up_proj,0.0001751961,0.05000,3.657
+24,mlp.down_proj,0.0000941419,0.05000,5.674
+25,self_attn.k_proj,0.0000201400,0.05000,4.789
+25,self_attn.q_proj,0.0000949212,0.05000,4.936
+25,self_attn.v_proj,0.0000253956,0.05000,4.953
+25,self_attn.o_proj,0.0000235002,0.05000,1.350
+25,mlp.up_proj,0.0002028768,0.05000,3.754
+25,mlp.gate_proj,0.0001922596,0.05000,3.767
+25,mlp.down_proj,0.0001267621,0.05000,5.595
+26,self_attn.k_proj,0.0000225750,0.05000,4.881
+26,self_attn.v_proj,0.0000295418,0.05000,4.921
+26,self_attn.q_proj,0.0001108223,0.05000,4.949
+26,self_attn.o_proj,0.0000265471,0.05000,1.385
+26,mlp.up_proj,0.0002334800,0.05000,3.388
+26,mlp.gate_proj,0.0002268478,0.05000,3.422
+26,mlp.down_proj,0.0001893703,0.05000,5.634
+27,self_attn.q_proj,0.0001638606,0.05000,5.063
+27,self_attn.v_proj,0.0000431965,0.05000,5.074
+27,self_attn.k_proj,0.0000341267,0.05000,5.090
+27,self_attn.o_proj,0.0000498931,0.05000,1.395
+27,mlp.up_proj,0.0002877421,0.05000,3.456
+27,mlp.gate_proj,0.0002798097,0.05000,3.483
+27,mlp.down_proj,0.0002758613,0.05000,5.413
+28,self_attn.q_proj,0.0002983511,0.05000,5.019
+28,self_attn.k_proj,0.0000565795,0.05000,5.021
+28,self_attn.v_proj,0.0000813798,0.05000,5.037
+28,self_attn.o_proj,0.0000459752,0.05000,1.567
+28,mlp.gate_proj,0.0003511385,0.05000,3.554
+28,mlp.up_proj,0.0003660403,0.05000,3.605
+28,mlp.down_proj,0.0003470357,0.05000,5.451
+29,self_attn.v_proj,0.0000988175,0.05000,4.698
+29,self_attn.q_proj,0.0003714425,0.05000,4.765
+29,self_attn.k_proj,0.0000788545,0.05000,4.812
+29,self_attn.o_proj,0.0000517854,0.05000,1.332
+29,mlp.up_proj,0.0004072312,0.05000,3.735
+29,mlp.gate_proj,0.0003876266,0.05000,3.765
+29,mlp.down_proj,0.0004905408,0.05000,5.507
+30,self_attn.q_proj,0.0005447937,0.05000,4.696
+30,self_attn.v_proj,0.0001534902,0.05000,4.769
+30,self_attn.k_proj,0.0001148308,0.05000,4.786
+30,self_attn.o_proj,0.0000775836,0.05000,1.395
+30,mlp.gate_proj,0.0004827204,0.05000,3.138
+30,mlp.up_proj,0.0005053384,0.05000,3.150
+30,mlp.down_proj,0.0005774746,0.05000,5.607
+31,self_attn.k_proj,0.0001142838,0.05000,4.573
+31,self_attn.v_proj,0.0001417302,0.05000,4.617
+31,self_attn.q_proj,0.0005521921,0.05000,4.641
+31,self_attn.o_proj,0.0000683407,0.05000,1.471
+31,mlp.gate_proj,0.0005245574,0.05000,3.371
+31,mlp.up_proj,0.0005602224,0.05000,3.391
+31,mlp.down_proj,0.0007208518,0.05000,5.466
+32,self_attn.k_proj,0.0002052982,0.05000,4.842
+32,self_attn.q_proj,0.0009621031,0.05000,4.916
+32,self_attn.v_proj,0.0002829138,0.05000,4.930
+32,self_attn.o_proj,0.0000929007,0.05000,1.420
+32,mlp.gate_proj,0.0005781620,0.05000,3.086
+32,mlp.up_proj,0.0006268456,0.05000,3.146
+32,mlp.down_proj,0.0007982043,0.05000,5.426
+33,self_attn.v_proj,0.0003873876,0.05000,4.661
+33,self_attn.q_proj,0.0013329935,0.05000,4.679
+33,self_attn.k_proj,0.0002684569,0.05000,4.699
+33,self_attn.o_proj,0.0001004743,0.05000,1.397
+33,mlp.gate_proj,0.0006186912,0.05000,3.072
+33,mlp.up_proj,0.0006873148,0.05000,3.080
+33,mlp.down_proj,0.0008831521,0.05000,5.923
+34,self_attn.k_proj,0.0004681831,0.05000,4.314
+34,self_attn.q_proj,0.0022188602,0.05000,4.330
+34,self_attn.v_proj,0.0006818156,0.05000,4.369
+34,self_attn.o_proj,0.0001390811,0.05000,1.400
+34,mlp.gate_proj,0.0006815079,0.05000,2.850
+34,mlp.up_proj,0.0007714760,0.05000,2.856
+34,mlp.down_proj,0.0010444433,0.05000,5.671
+35,self_attn.v_proj,0.0009464067,0.05000,4.743
+35,self_attn.q_proj,0.0028209405,0.05000,4.756
+35,self_attn.k_proj,0.0005987322,0.05000,4.770
+35,self_attn.o_proj,0.0001368956,0.05000,1.364
+35,mlp.up_proj,0.0008114978,0.05000,3.085
+35,mlp.gate_proj,0.0007047556,0.05000,3.087
+35,mlp.down_proj,0.0012339309,0.05000,5.406
+36,self_attn.k_proj,0.0005470300,0.05000,4.830
+36,self_attn.q_proj,0.0024462599,0.05000,4.882
+36,self_attn.v_proj,0.0008094013,0.05000,4.888
+36,self_attn.o_proj,0.0002597887,0.05000,1.454
+36,mlp.up_proj,0.0008480489,0.05000,3.371
+36,mlp.gate_proj,0.0007133146,0.05000,3.388
+36,mlp.down_proj,0.0015971187,0.05000,5.952
+37,self_attn.q_proj,0.0035302981,0.05000,4.597
+37,self_attn.k_proj,0.0007176153,0.05000,4.631
+37,self_attn.v_proj,0.0012357151,0.05000,4.644
+37,self_attn.o_proj,0.0002945466,0.05000,1.404
+37,mlp.gate_proj,0.0007103522,0.05000,3.077
+37,mlp.up_proj,0.0008592182,0.05000,3.084
+37,mlp.down_proj,0.0022288549,0.05000,5.345
+38,self_attn.v_proj,0.0011880093,0.05000,4.358
+38,self_attn.q_proj,0.0031908930,0.05000,4.372
+38,self_attn.k_proj,0.0006751707,0.05000,4.399
+38,self_attn.o_proj,0.0003943942,0.05000,1.340
+38,mlp.gate_proj,0.0008062866,0.05000,3.148
+38,mlp.up_proj,0.0009217460,0.05000,3.174
+38,mlp.down_proj,0.0034524150,0.05000,5.362
+39,self_attn.k_proj,0.0002554454,0.05000,4.257
+39,self_attn.v_proj,0.0003735272,0.05000,4.294
+39,self_attn.q_proj,0.0012166777,0.05000,4.295
+39,self_attn.o_proj,0.0003847522,0.05000,1.344
+39,mlp.gate_proj,0.0009165438,0.05000,3.561
+39,mlp.up_proj,0.0010139998,0.05000,3.578
+39,mlp.down_proj,0.0066875947,0.05000,5.467

quantize_config.json

@@ -0,0 +1,25 @@
+{
+  "bits": 4,
+  "group_size": 32,
+  "desc_act": false,
+  "sym": true,
+  "lm_head": false,
+  "quant_method": "gptq",
+  "checkpoint_format": "gptq",
+  "pack_dtype": "int32",
+  "meta": {
+    "quantizer": [
+      "gptqmodel:5.1.0-dev"
+    ],
+    "uri": "https://github.com/modelcloud/gptqmodel",
+    "damp_percent": 0.05,
+    "damp_auto_increment": 0.01,
+    "static_groups": false,
+    "true_sequential": true,
+    "mse": 0.0,
+    "v2": false,
+    "v2_alpha": 0.25,
+    "act_group_aware": true
+  },
+  "pack_impl": "cpu"
+}

special_tokens_map.json

@@ -0,0 +1,25 @@
+{
+  "additional_special_tokens": [
+    "<|im_start|>",
+    "<|im_end|>",
+    "<|object_ref_start|>",
+    "<|object_ref_end|>",
+    "<|box_start|>",
+    "<|box_end|>",
+    "<|quad_start|>",
+    "<|quad_end|>",
+    "<|vision_start|>",
+    "<|vision_end|>",
+    "<|vision_pad|>",
+    "<|image_pad|>",
+    "<|video_pad|>"
+  ],
+  "eos_token": {
+    "content": "<|endoftext|>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "pad_token": "<|endoftext|>"
+}

tokenizer.json

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:aeb13307a71acd8fe81861d94ad54ab689df773318809eed3cbe794b4492dae4
+size 11422654

tokenizer_config.json

@@ -0,0 +1,239 @@
+{
+  "add_bos_token": false,
+  "add_prefix_space": false,
+  "added_tokens_decoder": {
+    "151643": {
+      "content": "<|endoftext|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151644": {
+      "content": "<|im_start|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151645": {
+      "content": "<|im_end|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151646": {
+      "content": "<|object_ref_start|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151647": {
+      "content": "<|object_ref_end|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151648": {
+      "content": "<|box_start|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151649": {
+      "content": "<|box_end|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151650": {
+      "content": "<|quad_start|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151651": {
+      "content": "<|quad_end|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151652": {
+      "content": "<|vision_start|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151653": {
+      "content": "<|vision_end|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151654": {
+      "content": "<|vision_pad|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151655": {
+      "content": "<|image_pad|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151656": {
+      "content": "<|video_pad|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151657": {
+      "content": "<tool_call>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151658": {
+      "content": "</tool_call>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151659": {
+      "content": "<|fim_prefix|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151660": {
+      "content": "<|fim_middle|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151661": {
+      "content": "<|fim_suffix|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151662": {
+      "content": "<|fim_pad|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151663": {
+      "content": "<|repo_name|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151664": {
+      "content": "<|file_sep|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151665": {
+      "content": "<tool_response>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151666": {
+      "content": "</tool_response>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151667": {
+      "content": "<think>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151668": {
+      "content": "</think>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    }
+  },
+  "additional_special_tokens": [
+    "<|im_start|>",
+    "<|im_end|>",
+    "<|object_ref_start|>",
+    "<|object_ref_end|>",
+    "<|box_start|>",
+    "<|box_end|>",
+    "<|quad_start|>",
+    "<|quad_end|>",
+    "<|vision_start|>",
+    "<|vision_end|>",
+    "<|vision_pad|>",
+    "<|image_pad|>",
+    "<|video_pad|>"
+  ],
+  "bos_token": null,
+  "clean_up_tokenization_spaces": false,
+  "eos_token": "<|endoftext|>",
+  "errors": "replace",
+  "extra_special_tokens": {},
+  "model_max_length": 131072,
+  "pad_token": "<|endoftext|>",
+  "split_special_tokens": false,
+  "tokenizer_class": "Qwen2Tokenizer",
+  "unk_token": null
+}