K2-Horizon-3.7B MLX 4-bit release

.gitattributes

@@ -0,0 +1,37 @@
+*.7z filter=lfs diff=lfs merge=lfs -text
+*.arrow filter=lfs diff=lfs merge=lfs -text
+*.bin filter=lfs diff=lfs merge=lfs -text
+*.bz2 filter=lfs diff=lfs merge=lfs -text
+*.ckpt filter=lfs diff=lfs merge=lfs -text
+*.ftz filter=lfs diff=lfs merge=lfs -text
+*.gz filter=lfs diff=lfs merge=lfs -text
+*.h5 filter=lfs diff=lfs merge=lfs -text
+*.joblib filter=lfs diff=lfs merge=lfs -text
+*.lfs.* filter=lfs diff=lfs merge=lfs -text
+*.mlmodel filter=lfs diff=lfs merge=lfs -text
+*.model filter=lfs diff=lfs merge=lfs -text
+*.msgpack filter=lfs diff=lfs merge=lfs -text
+*.npy filter=lfs diff=lfs merge=lfs -text
+*.npz filter=lfs diff=lfs merge=lfs -text
+*.onnx filter=lfs diff=lfs merge=lfs -text
+*.ot filter=lfs diff=lfs merge=lfs -text
+*.parquet filter=lfs diff=lfs merge=lfs -text
+*.pb filter=lfs diff=lfs merge=lfs -text
+*.pickle filter=lfs diff=lfs merge=lfs -text
+*.pkl filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text
+*.pth filter=lfs diff=lfs merge=lfs -text
+*.rar filter=lfs diff=lfs merge=lfs -text
+*.safetensors filter=lfs diff=lfs merge=lfs -text
+saved_model/**/* filter=lfs diff=lfs merge=lfs -text
+*.tar.* filter=lfs diff=lfs merge=lfs -text
+*.tar filter=lfs diff=lfs merge=lfs -text
+*.tflite filter=lfs diff=lfs merge=lfs -text
+*.tgz filter=lfs diff=lfs merge=lfs -text
+*.wasm filter=lfs diff=lfs merge=lfs -text
+*.xz 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
+assets/k2-horizon-3.7b-benchmarks.png filter=lfs diff=lfs merge=lfs -text
+tokenizer.json filter=lfs diff=lfs merge=lfs -text

README.md

@@ -0,0 +1,53 @@
+---
+base_model: IFM/K2-Horizon-3.7B
+base_model_relation: quantized
+license: apache-2.0
+language:
+  - en
+pipeline_tag: text-generation
+library_name: mlx
+tags:
+  - mlx
+  - mlx-lm
+  - 4-bit
+  - k2-horizon
+  - long-context
+  - 512k-context
+  - dense
+---
+
+# K2-Horizon-3.7B MLX — 4-bit
+
+MLX 4-bit conversion of [IFM/K2-Horizon-3.7B](https://huggingface.co/IFM/K2-Horizon-3.7B), a 3.7B dense decoder-only model for reasoning, coding, long-context work, and tool use. The source checkpoint supports a native context length of **524,288 tokens (512K)**.
+
+## Benchmarks
+
+![IFM/K2-Horizon-3.7B benchmark results](assets/k2-horizon-3.7b-benchmarks.png)
+
+*Benchmark results reported by IFM for the original IFM/K2-Horizon-3.7B checkpoint.*
+
+## Release
+
+| Format | Quantization | Size |
+| --- | --- | ---: |
+| MLX safetensors | Affine 4-bit, group size 64 | 2.87 GB |
+
+The included `k2_horizon_mlx.py` adapter preserves K2 Horizon's grouped RMSNorm. Use it with MLX-LM and `--trust-remote-code`. The model is text-only; no vision projector or MTP files are included.
+
+## Usage
+
+```bash
+pip install -U mlx-lm
+mlx_lm.generate \
+  --model . \
+  --trust-remote-code \
+  --prompt "Explain why reproducible builds matter." \
+  --max-tokens 512
+```
+
+## Source
+
+- Model: [IFM/K2-Horizon-3.7B](https://huggingface.co/IFM/K2-Horizon-3.7B)
+- Source revision: [`633f52a`](https://huggingface.co/IFM/K2-Horizon-3.7B/tree/633f52ad28b17edeabd82afc61d2d13b4c59a561)
+- License: [Apache-2.0](https://www.apache.org/licenses/LICENSE-2.0)
+- Checksums: [SHA256SUMS.txt](SHA256SUMS.txt)

SHA256SUMS.txt

@@ -0,0 +1,11 @@
+0d8e479c380d5d9c8b28ea753205dd9736e53682aaf0038962088cd27c08a0d0  README.md
+a892cd0b0195599f283a8c706787520d9a6747640efb2f4dec4144b0abb62590  chat_template.jinja
+fec604c00546023bc73b9c00f7aa594831839e9db47ff2f25c288825cbeac1b1  config.json
+5c2f993c1053d9462ebea6dea416c897fddfbb4a5edd904e486936b20d4badc5  configuration_k2_horizon.py
+2da7d47641f4509da4ae47711e31d8b5f0f3f801ee08d87e7e9f07f814bdc4a3  generation_config.json
+56c88b1f4ece81dc946a0b8a2431822aaa67f0ff45c04496d683697c0cb5fef9  k2_horizon_mlx.py
+f2155c94d950cbc6d55929b0e46b114a912e211d006ae89bc7ed987b4b1cd2e8  model.safetensors
+0328a4fa1d70bdea50b7e747286e52ada17c145a142fa2cdfe5c944abe91c521  model.safetensors.index.json
+fb09e010956bd51cfa7d4055b4381cff34c9e06164066b49e3546f38b2e6242f  modeling_k2_horizon.py
+838d767b7c9925ff257feb20eaa4299a8e3cc35bb3d805589c373f51d2cc3cb6  tokenizer.json
+068cfdcf2bcef44fd77f935a9fb41b4d45af547fd41b95079817bd40b24fe518  tokenizer_config.json

chat_template.jinja

@@ -0,0 +1,994 @@
+{{- bos_token }}
+{%- if tool_presentation is defined -%}
+    {{- raise_exception("Unsupported argument: tool_presentation. Use tool_presentation_format with one of: json, xml, markdown.") -}}
+{%- endif -%}
+{%- if tool_calling_format is defined -%}
+    {{- raise_exception("Unsupported argument: tool_calling_format. Use tool_call_format with one of: json, xml, xml_typed.") -}}
+{%- endif -%}
+{%- if tool_format is defined -%}
+    {{- raise_exception("Unsupported argument: tool_format. Use tool_call_format with one of: json, xml, xml_typed.") -}}
+{%- endif -%}
+{%- set tool_presentation_fmt = tool_presentation_format | default('markdown') -%}
+{%- set tool_call_fmt = tool_call_format | default('xml') -%}
+{%- if tool_presentation_fmt != 'json' and tool_presentation_fmt != 'xml' and tool_presentation_fmt != 'markdown' -%}
+    {{- raise_exception("Unsupported tool_presentation_format: '" ~ tool_presentation_fmt ~ "'. Supported formats: json, xml, markdown.") -}}
+{%- endif -%}
+{%- if tool_call_fmt != 'json' and tool_call_fmt != 'xml' and tool_call_fmt != 'xml_typed' -%}
+    {{- raise_exception("Unsupported tool_call_format: '" ~ tool_call_fmt ~ "'. Supported formats: json, xml, xml_typed.") -}}
+{%- endif -%}
+
+{#- Renderability state, computed during validate_tools (single walk, no extra -#}
+{#- traversal at render time): ok = working flag for the tool being validated; -#}
+{#- bad = pipe-delimited indices of tools that must render as verbatim JSON. -#}
+{%- set RB = namespace(ok=true, bad='|') -%}
+
+{%- macro value_contains_mapping(v) -%}
+{%- if v is mapping -%}
+true
+{%- elif v is sequence and v is not string -%}
+{%- set f = namespace(x='false') -%}
+{%- for c in v -%}{%- if value_contains_mapping(c) == 'true' -%}{%- set f.x = 'true' -%}{%- endif -%}{%- endfor -%}
+{{- f.x -}}
+{%- else -%}
+false
+{%- endif -%}
+{%- endmacro -%}
+
+{#- $ref inlining state: defs = local $defs of the tool being rendered; seen = -#}
+{#- pipe-delimited names already expanded for this tool (each def inlines at most -#}
+{#- once; later references render by def name; cycles terminate immediately). -#}
+{#- $ref-sibling annotations (description/default/...) merge OVER the def at -#}
+{#- the inline site, so use-site annotations win and are never dropped. -#}
+{%- set REFS = namespace(defs={}, seen='|') -%}
+
+{%- macro render_compact_type_name(type_name, spec) -%}
+{%- if type_name == "array" -%}
+array[{%- if 'items' in spec -%}{{ render_compact_type(spec['items']) }}{%- else -%}any{%- endif -%}]
+{%- elif type_name -%}
+{{- type_name -}}
+{%- else -%}
+any
+{%- endif -%}
+{%- endmacro -%}
+
+{%- macro render_compact_type(spec) -%}
+{%- if spec is not mapping -%}
+any
+{%- elif spec.type is defined and spec.type is sequence and spec.type is not string and spec.type | length > 0 -%}
+{%- for type_name in spec.type -%}{{ render_compact_type_name(type_name, spec) }}{%- if not loop.last -%}|{%- endif -%}{%- endfor -%}
+{%- elif spec.type is defined and spec.type is sequence and spec.type is not string -%}
+any
+{%- elif spec.type -%}
+{{- render_compact_type_name(spec.type, spec) -}}
+{%- elif spec['$ref'] is string -%}
+{{- spec['$ref'].split('/') | last -}}
+{%- elif spec.oneOf -%}
+oneOf[{%- for variant in spec.oneOf -%}{{ render_compact_type(variant) }}{%- if not loop.last -%}|{%- endif -%}{%- endfor -%}]
+{%- elif spec.anyOf -%}
+anyOf[{%- for variant in spec.anyOf -%}{{ render_compact_type(variant) }}{%- if not loop.last -%}|{%- endif -%}{%- endfor -%}]
+{%- elif spec.properties -%}
+object
+{%- elif 'items' in spec -%}
+array[{{ render_compact_type(spec['items']) }}]
+{%- else -%}
+any
+{%- endif -%}
+{%- endmacro -%}
+
+{%- macro render_markdown_type_name(type_name, spec) -%}
+{%- if type_name == "array" -%}
+array of {% if 'items' in spec %}{{ render_markdown_type(spec['items']) }}{% else %}any{% endif %}
+{%- elif type_name -%}
+{{- type_name -}}
+{%- else -%}
+any
+{%- endif -%}
+{%- endmacro -%}
+
+{%- macro render_markdown_type(spec) -%}
+{%- if spec is sameas true -%}
+True
+{%- elif spec is sameas false -%}
+False
+{%- elif spec is not mapping -%}
+any
+{%- elif spec.type is defined and spec.type is sequence and spec.type is not string and spec.type | length > 0 -%}
+{%- for type_name in spec.type -%}{{ render_markdown_type_name(type_name, spec) }}{% if not loop.last %} or {% endif %}{%- endfor -%}
+{%- elif spec.type is defined and spec.type is sequence and spec.type is not string -%}
+any
+{%- elif spec.type -%}
+{{- render_markdown_type_name(spec.type, spec) -}}
+{%- elif spec['$ref'] is string -%}
+{{- spec['$ref'].split('/') | last -}}
+{%- elif spec.oneOf -%}
+oneOf[{%- for variant in spec.oneOf -%}{{ render_markdown_type(variant) }}{% if not loop.last %} or {% endif %}{%- endfor -%}]
+{%- elif spec.anyOf -%}
+anyOf[{%- for variant in spec.anyOf -%}{{ render_markdown_type(variant) }}{% if not loop.last %} or {% endif %}{%- endfor -%}]
+{%- elif spec.properties -%}
+object
+{%- elif 'items' in spec -%}
+array of {{ render_markdown_type(spec['items']) }}
+{%- else -%}
+any
+{%- endif -%}
+{%- endmacro -%}
+
+{%- macro render_xml_text(value) -%}
+{{- value.split() | join(" ") -}}
+{%- endmacro -%}
+
+{%- macro render_python_string(value) -%}
+'{{- value.split() | join(" ") | replace("\\", "\\\\") | replace("'", "\\'") -}}'
+{%- endmacro -%}
+
+{%- macro render_python_repr(value) -%}
+{%- if value is string -%}
+{{ render_python_string(value) }}
+{%- elif value is sameas true -%}
+True
+{%- elif value is sameas false -%}
+False
+{%- elif value is none -%}
+None
+{%- elif value is mapping -%}
+{{- "{" -}}
+{%- for key, child in value | items -%}
+{{ render_python_repr(key) }}: {{ render_python_repr(child) }}{%- if not loop.last -%}, {% endif -%}
+{%- endfor -%}
+{{- "}" -}}
+{%- elif value is sequence -%}
+{{- "[" -}}
+{%- for child in value -%}
+{{ render_python_repr(child) }}{%- if not loop.last -%}, {% endif -%}
+{%- endfor -%}
+{{- "]" -}}
+{%- else -%}
+{{- value -}}
+{%- endif -%}
+{%- endmacro -%}
+
+{%- macro render_xml_value(value) -%}
+{%- if value is string -%}{{ render_xml_text(value) }}{%- else -%}{{ render_python_repr(value) }}{%- endif -%}
+{%- endmacro -%}
+
+{%- macro render_xml_enum_value(value) -%}
+{%- if value is string -%}"{{- value | replace("\\", "\\\\") | replace("\"", "\\\"") -}}"{%- else -%}"{{- render_python_repr(value) | replace("\\", "\\\\") | replace("\"", "\\\"") -}}"{%- endif -%}
+{%- endmacro -%}
+
+{%- macro render_xml_enum(values) -%}
+{%- for value in values -%}{{ render_xml_enum_value(value) }}{%- if not loop.last -%}|{%- endif -%}{%- endfor -%}
+{%- endmacro -%}
+
+{%- macro render_xml_default_attr(value) -%}
+{{- " default=" }}{%- if value is string -%}"{{- value | replace("\\", "\\\\") | replace("\"", "\\\"") -}}"{%- else -%}{{ render_xml_value(value) }}{%- endif -%}
+{%- endmacro -%}
+
+{%- macro render_xml_attr(name, value) -%}
+{{- " " + name + "=" }}{%- if value == "" -%}""{%- else -%}{{ render_xml_value(value) }}{%- endif -%}
+{%- endmacro -%}
+
+{%- macro validate_schema(spec, path, lenient=false, classify=true, in_variant=false) -%}
+{%- if spec is mapping -%}
+    {%- if not lenient -%}
+    {%- if spec.required is defined -%}
+        {%- if spec.required is string or spec.required is not sequence -%}
+            {{- raise_exception("Schema '" + path + "' has 'required' but it is not a list.") -}}
+        {%- endif -%}
+        {%- if spec.required | length > 0 and not spec.properties and not in_variant -%}
+            {{- raise_exception("Schema '" + path + "' has required fields but no properties object to define them.") -}}
+        {%- endif -%}
+        {%- if spec.properties -%}
+            {%- for required_name in spec.required -%}
+                {%- if required_name not in spec.properties -%}
+                    {{- raise_exception("Schema '" + path + "' marks '" + required_name + "' as required, but that property is not defined in properties.") -}}
+                {%- endif -%}
+            {%- endfor -%}
+        {%- endif -%}
+    {%- endif -%}
+    {%- endif -%}
+    {#- renderability classification, piggybacking on this walk (no raises here): -#}
+    {#- constructs the pretty renderer does not fully handle flip RB.ok so the -#}
+    {#- tool falls back to verbatim JSON. Skipped entirely for json presentation. -#}
+    {%- if classify -%}
+    {%- for key, value in spec | items -%}
+        {%- if key == '$ref' -%}
+            {%- if value is not string -%}{%- set RB.ok = false -%}
+            {%- elif not (value.startswith('#/$defs/') or value.startswith('#/definitions/')) -%}{%- set RB.ok = false -%}{%- endif -%}
+        {%- elif key == '$defs' or key == 'definitions' -%}
+            {%- if value is mapping -%}
+                {%- for dk, dv in value | items -%}
+                    {{- validate_schema(dv, path + ".$defs." + dk, true) -}}
+                {%- endfor -%}
+            {%- else -%}{%- set RB.ok = false -%}{%- endif -%}
+        {%- elif key == 'type' -%}
+            {%- if value is mapping -%}{%- set RB.ok = false -%}{%- endif -%}
+        {%- elif key == 'enum' -%}
+            {%- if value is string or value is mapping or value is not sequence -%}{%- set RB.ok = false -%}{%- endif -%}
+        {%- elif key == 'items' -%}
+            {#- any items shape renders: mapping structurally, others via repr detail -#}
+        {%- elif key == 'oneOf' or key == 'anyOf' -%}
+            {%- if value is mapping or value is string or value is not sequence -%}{%- set RB.ok = false -%}{%- endif -%}
+        {%- elif key == 'required' -%}
+            {%- if value and not spec.properties -%}{%- set RB.ok = false -%}{%- endif -%}
+        {%- elif ('|' ~ key ~ '|') in '|description|default|title|examples|properties|patternProperties|additionalProperties|returns|' -%}
+        {%- elif value is mapping -%}
+            {%- for uk, uv in value | items -%}
+                {%- if value_contains_mapping(uv) == 'true' -%}{%- set RB.ok = false -%}{%- endif -%}
+            {%- endfor -%}
+        {%- elif value is sequence and value is not string -%}
+            {%- if value_contains_mapping(value) == 'true' -%}{%- set RB.ok = false -%}{%- endif -%}
+        {%- endif -%}
+    {%- endfor -%}
+    {%- endif -%}
+    {%- if spec.properties -%}
+        {%- for child_name, child_spec in spec.properties | items -%}
+            {{- validate_schema(child_spec, path + "." + child_name, lenient, classify) -}}
+        {%- endfor -%}
+    {%- endif -%}
+    {%- if 'items' in spec -%}{{- validate_schema(spec['items'], path + "[]", lenient, classify) -}}{%- endif -%}
+    {%- if spec.oneOf -%}
+        {%- for variant in spec.oneOf -%}{{- validate_schema(variant, path + ".oneOf[" + (loop.index0 | string) + "]", lenient, classify, true) -}}{%- endfor -%}
+    {%- endif -%}
+    {%- if spec.anyOf -%}
+        {%- for variant in spec.anyOf -%}{{- validate_schema(variant, path + ".anyOf[" + (loop.index0 | string) + "]", lenient, classify, true) -}}{%- endfor -%}
+    {%- endif -%}
+    {%- if spec.additionalProperties is mapping -%}{{- validate_schema(spec.additionalProperties, path + ".additionalProperties", lenient, classify) -}}{%- endif -%}
+    {%- if spec.patternProperties is mapping -%}
+        {%- for pattern, pattern_spec in spec.patternProperties | items -%}
+            {{- validate_schema(pattern_spec, path + ".patternProperties[" + pattern + "]", lenient, classify) -}}
+        {%- endfor -%}
+    {%- endif -%}
+    {%- if spec.returns is mapping -%}{{- validate_schema(spec.returns, path + ".returns", lenient, classify) -}}{%- endif -%}
+{%- endif -%}
+{%- endmacro -%}
+
+{%- macro validate_tools(tools_list, classify=true) -%}
+{%- set RB.bad = '|' -%}
+{%- for tool in tools_list -%}
+    {%- set fn = tool.function if tool.function is defined else tool -%}
+    {%- set RB.ok = true -%}
+    {%- if fn.parameters is defined and fn.parameters is string -%}
+        {{- raise_exception("tool.function.parameters must be a dict, not a JSON string. Parse it before passing to the template.") -}}
+    {%- endif -%}
+    {%- if fn.parameters is not defined or fn.parameters is none -%}
+        {%- if fn.arguments is defined -%}
+            {{- raise_exception("Tool '" + fn.name + "' has 'arguments' instead of 'parameters'. Rename 'arguments' to 'parameters'.") -}}
+        {%- else -%}
+            {{- raise_exception("Tool '" + fn.name + "' is missing required 'parameters' field. Each tool must have a 'parameters' dict with 'type', 'properties', and 'required' keys.") -}}
+        {%- endif -%}
+    {%- endif -%}
+    {{- validate_schema(fn.parameters, "tool." + fn.name + ".parameters", false, classify) -}}
+    {%- if classify -%}
+    {%- if fn.parameters is mapping -%}
+        {#- unknown container-valued keys at the parameters ROOT are never rendered -#}
+        {#- by the pretty path (root extras are dropped) -> verbatim fallback. -#}
+        {%- for rk, rv in fn.parameters | items -%}
+            {%- if rk not in ['type', 'description', 'enum', 'default', 'properties', 'required', 'optional', 'title', 'items', 'oneOf', 'anyOf', 'additionalProperties', 'patternProperties', 'returns', 'examples', '$defs', 'definitions', '$ref'] -%}
+                {%- if rv is mapping or (rv is sequence and rv is not string) -%}{%- set RB.ok = false -%}{%- endif -%}
+            {%- endif -%}
+        {%- endfor -%}
+    {%- else -%}
+        {%- set RB.ok = false -%}
+    {%- endif -%}
+    {%- endif -%}
+    {%- if fn.returns is mapping -%}{{- validate_schema(fn.returns, "tool." + fn.name + ".returns", false, classify) -}}{%- endif -%}
+    {%- if classify and fn.returns is not defined and fn.response is mapping -%}{{- validate_schema(fn.response, "tool." + fn.name + ".response", true) -}}{%- endif -%}
+    {#- unknown container-valued keys at the FUNCTION level are never rendered -> fallback. -#}
+    {%- if classify -%}
+    {%- for fk, fv in fn | items -%}
+        {%- if fk not in ['name', 'description', 'parameters', 'returns', 'response', 'type', 'function'] -%}
+            {%- if fv is mapping or (fv is sequence and fv is not string) -%}{%- set RB.ok = false -%}{%- endif -%}
+        {%- endif -%}
+    {%- endfor -%}
+    {%- endif -%}
+    {%- if not RB.ok -%}{%- set RB.bad = RB.bad ~ loop.index0 ~ '|' -%}{%- endif -%}
+{%- endfor -%}
+{%- endmacro -%}
+
+{%- macro render_tools_json(tools_list) -%}
+{{- "<ifm|tools>" }}
+{%- for tool in tools_list %}
+{{- "\n" }}
+{{- tool | tojson }}
+{%- endfor %}
+{{- "\n</ifm|tools>" }}
+{%- endmacro -%}
+
+{%- macro render_xml_schema_attrs(spec, include_value_attrs) -%}
+{%- if spec is mapping -%}
+{%- set structural_keys = ["type", "description", "enum", "default", "properties", "required", "items", "oneOf", "anyOf", "additionalProperties", "patternProperties", "returns"] -%}
+{%- if include_value_attrs and spec.enum -%}{{- " enum=" }}{{ render_xml_enum(spec.enum) }}{%- endif -%}
+{%- if include_value_attrs and spec.default is defined -%}{{ render_xml_default_attr(spec.default) }}{%- endif -%}
+{%- if spec.additionalProperties is defined and spec.additionalProperties is not mapping -%}{{ render_xml_attr("additionalProperties", spec.additionalProperties) }}{%- endif -%}
+{%- if spec.patternProperties is defined and spec.patternProperties is not mapping -%}{{ render_xml_attr("patternProperties", spec.patternProperties) }}{%- endif -%}
+{%- for key, value in spec | items -%}
+    {%- if key not in structural_keys -%}
+{{ render_xml_attr(key, value) }}
+    {%- endif -%}
+{%- endfor -%}
+{%- endif -%}
+{%- endmacro -%}
+
+{%- macro xml_schema_has_children(spec, include_properties, include_description) -%}
+{%- if spec is not mapping -%}
+false
+{%- elif (include_description and spec.description is defined) or (include_properties and spec.properties) or 'items' in spec or spec.oneOf or spec.anyOf or spec.additionalProperties is mapping or spec.patternProperties is mapping or spec.returns is defined -%}
+true
+{%- else -%}
+false
+{%- endif -%}
+{%- endmacro -%}
+
+{%- macro render_xml_schema_node(tag, spec, include_properties) -%}
+{%- if spec is mapping and spec['$ref'] is string -%}
+    {%- set _r = spec['$ref'] -%}
+    {%- set _k = _r[8:] if _r.startswith('#/$defs/') else (_r[14:] if _r.startswith('#/definitions/') else none) -%}
+    {%- if _k is not none and ('|' + _k + '|') not in REFS.seen and REFS.defs[_k] is mapping -%}
+        {%- set spec = dict((REFS.defs[_k] | items | list) + (spec | items | rejectattr('0', 'equalto', '$ref') | list)) -%}
+        {%- set REFS.seen = REFS.seen + _k + '|' -%}
+        {%- if spec['$ref'] is string -%}
+            {%- set _r2 = spec['$ref'] -%}
+            {%- set _k2 = _r2[8:] if _r2.startswith('#/$defs/') else (_r2[14:] if _r2.startswith('#/definitions/') else none) -%}
+            {%- if _k2 is not none and REFS.defs[_k2] is mapping -%}
+                {%- set spec = dict((REFS.defs[_k2] | items | list) + (spec | items | rejectattr('0', 'equalto', '$ref') | list)) -%}
+                {%- set REFS.seen = REFS.seen + _k2 + '|' -%}
+            {%- endif -%}
+        {%- endif -%}
+    {%- endif -%}
+{%- endif -%}
+{%- if spec is mapping -%}
+{{- "<" + tag + " type=" + render_compact_type(spec) }}{{ render_xml_schema_attrs(spec, true) }}
+{%- if xml_schema_has_children(spec, include_properties, true) == 'true' -%}
+{{- ">" }}{{ render_xml_schema_children(spec, include_properties, true) }}{{- "</" + tag + ">" }}
+{%- else -%}
+{{- "/>" }}
+{%- endif -%}
+{%- else -%}
+{{- "<" + tag + ">" }}{{ render_xml_value(spec) }}{{- "</" + tag + ">" }}
+{%- endif -%}
+{%- endmacro -%}
+
+{%- macro render_xml_pattern_property(pattern, spec) -%}
+{%- if spec is mapping and spec['$ref'] is string -%}
+    {%- set _r = spec['$ref'] -%}
+    {%- set _k = _r[8:] if _r.startswith('#/$defs/') else (_r[14:] if _r.startswith('#/definitions/') else none) -%}
+    {%- if _k is not none and ('|' + _k + '|') not in REFS.seen and REFS.defs[_k] is mapping -%}
+        {%- set spec = dict((REFS.defs[_k] | items | list) + (spec | items | rejectattr('0', 'equalto', '$ref') | list)) -%}
+        {%- set REFS.seen = REFS.seen + _k + '|' -%}
+        {%- if spec['$ref'] is string -%}
+            {%- set _r2 = spec['$ref'] -%}
+            {%- set _k2 = _r2[8:] if _r2.startswith('#/$defs/') else (_r2[14:] if _r2.startswith('#/definitions/') else none) -%}
+            {%- if _k2 is not none and REFS.defs[_k2] is mapping -%}
+                {%- set spec = dict((REFS.defs[_k2] | items | list) + (spec | items | rejectattr('0', 'equalto', '$ref') | list)) -%}
+                {%- set REFS.seen = REFS.seen + _k2 + '|' -%}
+            {%- endif -%}
+        {%- endif -%}
+    {%- endif -%}
+{%- endif -%}
+{%- if spec is mapping -%}
+{{- "<patternProperty" }}{{ render_xml_attr("pattern", pattern) }}{{- " type=" + render_compact_type(spec) }}{{ render_xml_schema_attrs(spec, true) }}
+{%- if xml_schema_has_children(spec, true, true) == 'true' -%}
+{{- ">" }}{{ render_xml_schema_children(spec, true, true) }}{{- "</patternProperty>" }}
+{%- else -%}
+{{- "/>" }}
+{%- endif -%}
+{%- else -%}
+{{- "<patternProperty" }}{{ render_xml_attr("pattern", pattern) }}{{- ">" }}{{ render_xml_value(spec) }}{{- "</patternProperty>" }}
+{%- endif -%}
+{%- endmacro -%}
+
+{%- macro render_xml_schema_children(spec, include_properties, include_description) -%}
+{%- if include_description and spec.description is defined -%}{{- "<description>" }}{{ spec.description }}{{- "</description>" }}{%- endif -%}
+{%- if include_properties and spec.properties -%}
+{%- for child_name, child_spec in spec.properties | items -%}
+{{- render_xml_param(child_name, child_spec, spec.required or []) }}
+{%- endfor -%}
+{%- endif -%}
+{%- if 'items' in spec -%}{{ render_xml_schema_node("items", spec['items'], true) }}{%- endif -%}
+{%- if spec.oneOf -%}
+{{- "<oneOf>" }}
+{%- for variant in spec.oneOf -%}{{ render_xml_schema_node("variant", variant, true) }}{%- endfor -%}
+{{- "</oneOf>" }}
+{%- endif -%}
+{%- if spec.anyOf -%}
+{{- "<anyOf>" }}
+{%- for variant in spec.anyOf -%}{{ render_xml_schema_node("variant", variant, true) }}{%- endfor -%}
+{{- "</anyOf>" }}
+{%- endif -%}
+{%- if spec.additionalProperties is mapping -%}{{ render_xml_schema_node("additionalProperties", spec.additionalProperties, true) }}{%- endif -%}
+{%- if spec.patternProperties is mapping -%}
+{{- "<patternProperties>" }}
+{%- for pattern, pattern_spec in spec.patternProperties | items -%}{{ render_xml_pattern_property(pattern, pattern_spec) }}{%- endfor -%}
+{{- "</patternProperties>" }}
+{%- elif spec.patternProperties is defined -%}<patternProperties>{{ render_xml_value(spec.patternProperties) }}</patternProperties>{%- endif -%}
+{%- if spec.returns is mapping -%}{{ render_xml_schema_node("returns", spec.returns, true) }}{%- elif spec.returns is defined -%}<returns>{{ render_xml_value(spec.returns) }}</returns>{%- endif -%}
+{%- endmacro -%}
+
+{%- macro render_xml_param(name, spec, required_list) -%}
+{%- if spec is mapping and spec['$ref'] is string -%}
+    {%- set _r = spec['$ref'] -%}
+    {%- set _k = _r[8:] if _r.startswith('#/$defs/') else (_r[14:] if _r.startswith('#/definitions/') else none) -%}
+    {%- if _k is not none and ('|' + _k + '|') not in REFS.seen and REFS.defs[_k] is mapping -%}
+        {%- set spec = dict((REFS.defs[_k] | items | list) + (spec | items | rejectattr('0', 'equalto', '$ref') | list)) -%}
+        {%- set REFS.seen = REFS.seen + _k + '|' -%}
+        {%- if spec['$ref'] is string -%}
+            {%- set _r2 = spec['$ref'] -%}
+            {%- set _k2 = _r2[8:] if _r2.startswith('#/$defs/') else (_r2[14:] if _r2.startswith('#/definitions/') else none) -%}
+            {%- if _k2 is not none and REFS.defs[_k2] is mapping -%}
+                {%- set spec = dict((REFS.defs[_k2] | items | list) + (spec | items | rejectattr('0', 'equalto', '$ref') | list)) -%}
+                {%- set REFS.seen = REFS.seen + _k2 + '|' -%}
+            {%- endif -%}
+        {%- endif -%}
+    {%- endif -%}
+{%- endif -%}
+{{- "<param name=" + name + " type=" + render_compact_type(spec) }}
+{%- if name in (required_list or []) -%}{{- " required=true" }}{%- endif -%}
+{%- if spec.enum -%}{{- " enum=" }}{{ render_xml_enum(spec.enum) }}{%- endif -%}
+{%- if spec.default is defined -%}{{ render_xml_default_attr(spec.default) }}{%- endif -%}
+{{- render_xml_schema_attrs(spec, false) }}
+{%- if spec.description or xml_schema_has_children(spec, true, false) == 'true' -%}
+{{- ">" }}
+{%- if spec.description -%}{{ spec.description }}{%- endif -%}
+{{- render_xml_schema_children(spec, true, false) }}
+{{- "</param>" }}
+{%- else -%}
+{{- "/>" }}
+{%- endif -%}
+{%- endmacro -%}
+
+{%- macro render_tools_xml(tools_list) -%}
+{{- "<ifm|tools>" }}
+{%- for tool in tools_list -%}
+    {%- set fn = tool.function if tool.function is defined else tool -%}
+    {%- set REFS.defs = fn.parameters['$defs'] if (fn.parameters is mapping and fn.parameters['$defs'] is mapping) else (fn.parameters['definitions'] if (fn.parameters is mapping and fn.parameters['definitions'] is mapping) else {}) -%}
+    {%- set REFS.seen = '|' -%}
+    {%- set fnp = namespace(p=fn.parameters) -%}
+    {%- if fnp.p is mapping and fnp.p['$ref'] is string -%}
+        {%- set _r = fnp.p['$ref'] -%}
+        {%- set _k = _r[8:] if _r.startswith('#/$defs/') else (_r[14:] if _r.startswith('#/definitions/') else none) -%}
+        {%- if _k is not none and REFS.defs[_k] is mapping -%}
+            {%- set fnp.p = dict((REFS.defs[_k] | items | list) + (fnp.p | items | rejectattr('0', 'equalto', '$ref') | list)) -%}
+            {%- set REFS.seen = REFS.seen + _k + '|' -%}
+        {%- endif -%}
+    {%- endif -%}
+{{- "\n<function name=" + fn.name + ">" }}
+{%- if fn.description -%}
+{{- "<description>" }}{{ fn.description }}{{- "</description>" }}
+{%- endif -%}
+{{- "<parameters>" }}
+{%- if fnp.p and fnp.p.properties -%}
+    {%- for pname, pspec in fnp.p.properties | items -%}
+{{- render_xml_param(pname, pspec, fnp.p.required or []) }}
+    {%- endfor -%}
+{%- elif fnp.p is mapping and (fnp.p.oneOf or fnp.p.anyOf or 'items' in fnp.p) -%}
+{{- render_xml_schema_children(fnp.p, true, false) }}
+{%- endif -%}
+{{- "</parameters>" }}
+{%- set fn_ret = fn.returns if fn.returns is defined else fn.response -%}
+{%- if fn_ret is mapping -%}{{ render_xml_schema_node("returns", fn_ret, true) }}{%- elif fn_ret is defined -%}<returns>{{ render_xml_value(fn_ret) }}</returns>{%- endif -%}
+{{- "</function>" }}
+{%- endfor -%}
+{{- "\n</ifm|tools>" }}
+{%- endmacro -%}
+
+{%- macro render_markdown_literal(value) -%}
+{%- if value is string and value == "" -%}""
+{%- elif value is string -%}`{{ value | replace("\n", "\\n") }}`
+{%- else -%}`{{ render_python_repr(value) }}`
+{%- endif -%}
+{%- endmacro -%}
+
+{%- macro render_allowed_values(values) -%}
+{%- for value in values -%}{{ render_markdown_literal(value) }}{% if not loop.last %}, {% endif %}{%- endfor -%}
+{%- endmacro -%}
+
+{%- macro render_markdown_value(value) -%}
+{%- if value is string and value == "" -%}""{%- elif value is string -%}{{ value }}{%- else -%}{{ render_python_repr(value) }}{%- endif -%}
+{%- endmacro -%}
+
+{%- macro render_markdown_detail(indent, label, value) -%}
+{{- "\n" + indent + "  - " + label + ": " }}{{ render_markdown_value(value) }}
+{%- endmacro -%}
+
+{%- macro render_markdown_metadata_detail(label, value) -%}
+{{- "\n- " + label + ": " }}{{ render_markdown_value(value) }}
+{%- endmacro -%}
+
+{%- macro render_markdown_schema_annotations(spec, indent, include_value_details) -%}
+{%- if include_value_details and spec.description is defined -%}{{ render_markdown_detail(indent, "Description", spec.description | replace("\n", "\n" + indent + "    ")) }}{%- endif -%}
+{%- if include_value_details and spec.enum is defined -%}{{- "\n" + indent + "  - Allowed values: " }}{{ render_allowed_values(spec.enum) }}{%- endif -%}
+{%- if include_value_details and spec.default is defined -%}{{- "\n" + indent + "  - Default: " }}{{ render_markdown_literal(spec.default) }}{%- endif -%}
+{%- if spec.additionalProperties is defined -%}
+    {%- if spec.additionalProperties is mapping -%}
+{{- "\n" + indent + "  - Additional properties *(" + render_markdown_type(spec.additionalProperties) + ")*" }}
+{{- render_markdown_schema_details(spec.additionalProperties, indent + "  ", true) }}
+    {%- else -%}
+{{ render_markdown_detail(indent, "Additional properties", spec.additionalProperties) }}
+    {%- endif -%}
+{%- endif -%}
+{%- endmacro -%}
+
+{%- macro render_markdown_metadata_annotations(spec) -%}
+{%- if spec.description is defined -%}{{ render_markdown_metadata_detail("Description", spec.description | replace("\n", "\n    ")) }}{%- endif -%}
+{%- if spec.enum is defined -%}{{- "\n- Allowed values: " }}{{ render_allowed_values(spec.enum) }}{%- endif -%}
+{%- if spec.default is defined -%}{{- "\n- Default: " }}{{ render_markdown_literal(spec.default) }}{%- endif -%}
+{%- if spec.additionalProperties is defined -%}
+    {%- if spec.additionalProperties is mapping -%}
+{{- "\n- Additional properties *(" + render_markdown_type(spec.additionalProperties) + ")*" }}
+{{- render_markdown_schema_details(spec.additionalProperties, "", true) }}
+    {%- else -%}
+{{ render_markdown_metadata_detail("Additional properties", spec.additionalProperties) }}
+    {%- endif -%}
+{%- endif -%}
+{%- endmacro -%}
+
+{%- macro render_markdown_schema_extras(spec, indent) -%}
+{%- set rendered_keys = ["type", "description", "enum", "default", "properties", "required", "items", "oneOf", "anyOf", "additionalProperties", "patternProperties", "returns"] -%}
+{%- for key, value in spec | items -%}
+    {%- if key not in rendered_keys -%}
+{{- "\n" + indent + "  - " + key + ": " }}{{ render_markdown_value(value) }}
+    {%- endif -%}
+{%- endfor -%}
+{%- endmacro -%}
+
+{%- macro render_markdown_metadata_extras(spec) -%}
+{%- set rendered_keys = ["type", "description", "enum", "default", "properties", "required", "items", "oneOf", "anyOf", "additionalProperties", "patternProperties", "returns"] -%}
+{%- for key, value in spec | items -%}
+    {%- if key not in rendered_keys -%}
+{{- "\n- " + key + ": " }}{{ render_markdown_value(value) }}
+    {%- endif -%}
+{%- endfor -%}
+{%- endmacro -%}
+
+{%- macro markdown_schema_has_extra(spec) -%}
+{%- set rendered_keys = ["type", "description", "enum", "default", "properties", "required", "items", "oneOf", "anyOf", "additionalProperties", "patternProperties", "returns"] -%}
+{%- set found = namespace(value='false') -%}
+{%- for key, value in spec | items -%}
+    {%- if key not in rendered_keys -%}{%- set found.value = 'true' -%}{%- endif -%}
+{%- endfor -%}
+{{- found.value -}}
+{%- endmacro -%}
+
+{%- macro markdown_parameter_schema_has_details(spec) -%}
+{%- if spec.description is defined or spec.enum is defined or spec.default is defined or spec.additionalProperties is defined or spec.patternProperties is defined or 'items' in spec or spec.oneOf or spec.anyOf or spec.returns is defined or markdown_schema_has_extra(spec) == 'true' -%}
+true
+{%- else -%}
+false
+{%- endif -%}
+{%- endmacro -%}
+
+{%- macro render_markdown_schema_structure(spec, indent, include_properties) -%}
+{%- if include_properties and spec.properties -%}
+    {%- for child_name, child_spec in spec.properties | items -%}
+{{- render_markdown_param(child_name, child_spec, spec.required or [], indent + "  ") }}
+    {%- endfor -%}
+{%- endif -%}
+{%- if 'items' in spec and spec['items'] is mapping -%}
+{{- "\n" + indent + "  - Items *(" + render_markdown_type(spec['items']) + ")*" }}
+{{- render_markdown_schema_details(spec['items'], indent + "  ", true) }}
+{%- elif 'items' in spec -%}
+{{ render_markdown_detail(indent, "Items", spec['items']) }}
+{%- endif -%}
+{%- if spec.oneOf -%}
+{{- "\n" + indent + "  - oneOf:" }}
+    {%- for variant in spec.oneOf -%}
+{{- "\n" + indent + "    - Variant " }}{{ loop.index }}{{- " *(" + render_markdown_type(variant) + ")*" }}
+{{- render_markdown_schema_details(variant, indent + "    ", true) }}
+    {%- endfor -%}
+{%- endif -%}
+{%- if spec.anyOf -%}
+{{- "\n" + indent + "  - anyOf:" }}
+    {%- for variant in spec.anyOf -%}
+{{- "\n" + indent + "    - Variant " }}{{ loop.index }}{{- " *(" + render_markdown_type(variant) + ")*" }}
+{{- render_markdown_schema_details(variant, indent + "    ", true) }}
+    {%- endfor -%}
+{%- endif -%}
+{%- if spec.patternProperties is mapping -%}
+{{- "\n" + indent + "  - Pattern properties:" }}
+    {%- for pattern, pattern_spec in spec.patternProperties | items -%}
+        {%- if pattern_spec is mapping -%}
+{{- "\n" + indent + "    - `" + pattern + "` *(" + render_markdown_type(pattern_spec) + ")*" }}
+{{- render_markdown_schema_details(pattern_spec, indent + "    ", true) }}
+        {%- else -%}
+{{- "\n" + indent + "    - `" + pattern + "`: " }}{{ render_markdown_value(pattern_spec) }}
+        {%- endif -%}
+    {%- endfor -%}
+{%- elif spec.patternProperties is defined -%}
+{{ render_markdown_detail(indent, "Pattern properties", spec.patternProperties) }}
+{%- endif -%}
+{%- if spec.returns is mapping -%}
+{{- "\n" + indent + "  - Returns *(" + render_markdown_type(spec.returns) + ")*" }}
+{{- render_markdown_schema_details(spec.returns, indent + "  ", true) }}
+{%- elif spec.returns is defined -%}
+{{ render_markdown_detail(indent, "Returns", spec.returns) }}
+{%- endif -%}
+{%- endmacro -%}
+
+{%- macro render_markdown_schema_details(spec, indent, include_value_details) -%}
+{%- if spec is mapping and spec['$ref'] is string -%}
+    {%- set _r = spec['$ref'] -%}
+    {%- set _k = _r[8:] if _r.startswith('#/$defs/') else (_r[14:] if _r.startswith('#/definitions/') else none) -%}
+    {%- if _k is not none and ('|' + _k + '|') not in REFS.seen and REFS.defs[_k] is mapping -%}
+        {%- set spec = dict((REFS.defs[_k] | items | list) + (spec | items | rejectattr('0', 'equalto', '$ref') | list)) -%}
+        {%- set REFS.seen = REFS.seen + _k + '|' -%}
+        {%- if spec['$ref'] is string -%}
+            {%- set _r2 = spec['$ref'] -%}
+            {%- set _k2 = _r2[8:] if _r2.startswith('#/$defs/') else (_r2[14:] if _r2.startswith('#/definitions/') else none) -%}
+            {%- if _k2 is not none and REFS.defs[_k2] is mapping -%}
+                {%- set spec = dict((REFS.defs[_k2] | items | list) + (spec | items | rejectattr('0', 'equalto', '$ref') | list)) -%}
+                {%- set REFS.seen = REFS.seen + _k2 + '|' -%}
+            {%- endif -%}
+        {%- endif -%}
+    {%- endif -%}
+{%- endif -%}
+{%- if spec is mapping -%}
+{{- render_markdown_schema_annotations(spec, indent, include_value_details) }}
+{{- render_markdown_schema_structure(spec, indent, true) }}
+{{- render_markdown_schema_extras(spec, indent) }}
+{%- elif spec is not sameas true and spec is not sameas false -%}
+{{- "\n" + indent + "  - Value: " }}{{ render_markdown_literal(spec) }}
+{%- endif -%}
+{%- endmacro -%}
+
+{%- macro render_markdown_parameter_schema(spec) -%}
+{%- if spec is mapping and spec['$ref'] is string -%}
+    {%- set _r = spec['$ref'] -%}
+    {%- set _k = _r[8:] if _r.startswith('#/$defs/') else (_r[14:] if _r.startswith('#/definitions/') else none) -%}
+    {%- if _k is not none and ('|' + _k + '|') not in REFS.seen and REFS.defs[_k] is mapping -%}
+        {%- set spec = dict((REFS.defs[_k] | items | list) + (spec | items | rejectattr('0', 'equalto', '$ref') | list)) -%}
+        {%- set REFS.seen = REFS.seen + _k + '|' -%}
+        {%- if spec['$ref'] is string -%}
+            {%- set _r2 = spec['$ref'] -%}
+            {%- set _k2 = _r2[8:] if _r2.startswith('#/$defs/') else (_r2[14:] if _r2.startswith('#/definitions/') else none) -%}
+            {%- if _k2 is not none and REFS.defs[_k2] is mapping -%}
+                {%- set spec = dict((REFS.defs[_k2] | items | list) + (spec | items | rejectattr('0', 'equalto', '$ref') | list)) -%}
+                {%- set REFS.seen = REFS.seen + _k2 + '|' -%}
+            {%- endif -%}
+        {%- endif -%}
+    {%- endif -%}
+{%- endif -%}
+{%- if spec is mapping -%}
+{{- render_markdown_metadata_annotations(spec) }}
+{%- if 'items' in spec and spec['items'] is mapping -%}
+{{- "\n- Items *(" + render_markdown_type(spec['items']) + ")*" }}
+{{- render_markdown_schema_details(spec['items'], "", true) }}
+{%- elif 'items' in spec -%}
+{{ render_markdown_metadata_detail("Items", spec['items']) }}
+{%- endif -%}
+{%- if spec.oneOf -%}
+{{- "\n- oneOf:" }}
+    {%- for variant in spec.oneOf -%}
+{{- "\n  - Variant " }}{{ loop.index }}{{- " *(" + render_markdown_type(variant) + ")*" }}
+{{- render_markdown_schema_details(variant, "  ", true) }}
+    {%- endfor -%}
+{%- endif -%}
+{%- if spec.anyOf -%}
+{{- "\n- anyOf:" }}
+    {%- for variant in spec.anyOf -%}
+{{- "\n  - Variant " }}{{ loop.index }}{{- " *(" + render_markdown_type(variant) + ")*" }}
+{{- render_markdown_schema_details(variant, "  ", true) }}
+    {%- endfor -%}
+{%- endif -%}
+{%- if spec.patternProperties is mapping -%}
+{{- "\n- Pattern properties:" }}
+    {%- for pattern, pattern_spec in spec.patternProperties | items -%}
+        {%- if pattern_spec is mapping -%}
+{{- "\n  - `" + pattern + "` *(" + render_markdown_type(pattern_spec) + ")*" }}
+{{- render_markdown_schema_details(pattern_spec, "  ", true) }}
+        {%- else -%}
+{{- "\n  - `" + pattern + "`: " }}{{ render_markdown_value(pattern_spec) }}
+        {%- endif -%}
+    {%- endfor -%}
+{%- elif spec.patternProperties is defined -%}
+{{ render_markdown_metadata_detail("Pattern properties", spec.patternProperties) }}
+{%- endif -%}
+{%- if spec.returns is mapping -%}
+{{- "\n- Returns *(" + render_markdown_type(spec.returns) + ")*" }}
+{{- render_markdown_schema_details(spec.returns, "", true) }}
+{%- elif spec.returns is defined -%}
+{{ render_markdown_metadata_detail("Returns", spec.returns) }}
+{%- endif -%}
+{{- render_markdown_metadata_extras(spec) }}
+{%- endif -%}
+{%- endmacro -%}
+
+{%- macro render_markdown_param(name, spec, required_list, indent) -%}
+{%- if spec is mapping and spec['$ref'] is string -%}
+    {%- set _r = spec['$ref'] -%}
+    {%- set _k = _r[8:] if _r.startswith('#/$defs/') else (_r[14:] if _r.startswith('#/definitions/') else none) -%}
+    {%- if _k is not none and ('|' + _k + '|') not in REFS.seen and REFS.defs[_k] is mapping -%}
+        {%- set spec = dict((REFS.defs[_k] | items | list) + (spec | items | rejectattr('0', 'equalto', '$ref') | list)) -%}
+        {%- set REFS.seen = REFS.seen + _k + '|' -%}
+        {%- if spec['$ref'] is string -%}
+            {%- set _r2 = spec['$ref'] -%}
+            {%- set _k2 = _r2[8:] if _r2.startswith('#/$defs/') else (_r2[14:] if _r2.startswith('#/definitions/') else none) -%}
+            {%- if _k2 is not none and REFS.defs[_k2] is mapping -%}
+                {%- set spec = dict((REFS.defs[_k2] | items | list) + (spec | items | rejectattr('0', 'equalto', '$ref') | list)) -%}
+                {%- set REFS.seen = REFS.seen + _k2 + '|' -%}
+            {%- endif -%}
+        {%- endif -%}
+    {%- endif -%}
+{%- endif -%}
+{{- "\n" + indent + "- `" + name + "` *(" + render_markdown_type(spec) }}
+{%- if name in (required_list or []) -%}{{- ", required" }}{%- endif -%}
+{{- ")*" }}
+{%- if spec.description -%}{{- " - " + spec.description | replace("\n", "\n" + indent + "  ") }}{%- endif -%}
+{%- if spec.enum -%}
+{{- "\n" + indent + "  - Allowed values: " }}{{ render_allowed_values(spec.enum) }}
+{%- endif -%}
+{%- if spec.default is defined -%}
+{{- "\n" + indent + "  - Default: " }}{{ render_markdown_literal(spec.default) }}
+{%- endif -%}
+{{- render_markdown_schema_details(spec, indent, false) }}
+{%- endmacro -%}
+
+{%- macro render_tools_markdown(tools_list) -%}
+{{- "<ifm|tools>" }}
+{%- for tool in tools_list -%}
+    {%- set fn = tool.function if tool.function is defined else tool -%}
+    {%- set REFS.defs = fn.parameters['$defs'] if (fn.parameters is mapping and fn.parameters['$defs'] is mapping) else (fn.parameters['definitions'] if (fn.parameters is mapping and fn.parameters['definitions'] is mapping) else {}) -%}
+    {%- set REFS.seen = '|' -%}
+    {%- set fnp = namespace(p=fn.parameters) -%}
+    {%- if fnp.p is mapping and fnp.p['$ref'] is string -%}
+        {%- set _r = fnp.p['$ref'] -%}
+        {%- set _k = _r[8:] if _r.startswith('#/$defs/') else (_r[14:] if _r.startswith('#/definitions/') else none) -%}
+        {%- if _k is not none and REFS.defs[_k] is mapping -%}
+            {%- set fnp.p = dict((REFS.defs[_k] | items | list) + (fnp.p | items | rejectattr('0', 'equalto', '$ref') | list)) -%}
+            {%- set REFS.seen = REFS.seen + _k + '|' -%}
+        {%- endif -%}
+    {%- endif -%}
+{{- "\n## " + fn.name }}
+{%- if fn.description -%}
+{{- "\n" + fn.description }}
+{%- endif -%}
+{{- "\n\n**Parameters**" }}
+{%- if fnp.p and fnp.p.properties -%}
+    {%- for pname, pspec in fnp.p.properties | items -%}
+{{- render_markdown_param(pname, pspec, fnp.p.required or [], "") }}
+    {%- endfor -%}
+{%- elif fnp.p is mapping and (fnp.p.oneOf or fnp.p.anyOf or 'items' in fnp.p) -%}
+{{- render_markdown_parameter_schema(fnp.p) }}
+{%- else -%}
+{{- "\n- None" }}
+{%- endif -%}
+{%- set fn_ret = fn.returns if fn.returns is defined else fn.response -%}
+{%- if fn_ret is mapping -%}
+{{- "\n\n**Returns**" }}
+{{- "\n- Return *(" + render_markdown_type(fn_ret) + ")*" }}
+{{- render_markdown_schema_details(fn_ret, "", true) }}
+{%- elif fn_ret is defined -%}
+{{- "\n\n**Returns**\n- " }}{{ render_markdown_value(fn_ret) }}
+{%- endif -%}
+{%- if not loop.last -%}{{- "\n" }}{%- endif -%}
+{%- endfor -%}
+{{- "\n</ifm|tools>" }}
+{%- endmacro -%}
+
+{%- macro render_tool_presentation(tools_list, fmt) -%}
+{%- if fmt == 'json' -%}
+{{- render_tools_json(tools_list) }}
+{%- elif RB.bad != '|' -%}
+{#- some tool uses constructs the pretty renderers cannot represent (verdicts -#}
+{#- computed during validate_tools): render the WHOLE toolset exactly as the -#}
+{#- json presentation would, so the block stays uniform and model-familiar. -#}
+{{- render_tools_json(tools_list) }}
+{%- elif fmt == 'xml' -%}
+{{- render_tools_xml(tools_list) }}
+{%- elif fmt == 'markdown' -%}
+{{- render_tools_markdown(tools_list) }}
+{%- else -%}
+{{- raise_exception("Unsupported tool_presentation_format: '" + fmt + "'. Supported formats: json, xml, markdown.") }}
+{%- endif -%}
+{%- endmacro -%}
+
+{%- macro render_call_instructions(fmt) -%}
+{%- if fmt == 'json' -%}
+{{- "Wrap all tool calls in a single <ifm|tool_calls></ifm|tool_calls> block. For each call, emit one JSON object with the function name and arguments on the same line inside <ifm|tool_call></ifm|tool_call> tags:\n\n<ifm|tool_calls>\n<ifm|tool_call>{\"name\": <function-name>, \"arguments\": <args-json-object>}</ifm|tool_call>\n</ifm|tool_calls>" }}
+{%- elif fmt == 'xml' -%}
+{{- "Wrap all tool calls in a single <ifm|tool_calls></ifm|tool_calls> block. For each call, write the function name at the start of <ifm|tool_call>, followed by paired <ifm|arg_key> and <ifm|arg_value> tags for each argument:\n\n<ifm|tool_calls>\n<ifm|tool_call>$FUNCTION_NAME\n<ifm|arg_key>$PARAMETER_NAME</ifm|arg_key>\n<ifm|arg_value>$PARAMETER_VALUE</ifm|arg_value>\n...\n</ifm|tool_call>\n</ifm|tool_calls>\n\nString and scalar parameters should be written as plain text. Array and object parameters should be written as JSON literals." }}
+{%- elif fmt == 'xml_typed' -%}
+{{- "Wrap all tool calls in a single <ifm|tool_calls></ifm|tool_calls> block. For each call, write the function name at the start of <ifm|tool_call>, followed by <ifm|arg_key>, <ifm|arg_type>, and <ifm|arg_value> tags for each argument:\n\n<ifm|tool_calls>\n<ifm|tool_call>$FUNCTION_NAME\n<ifm|arg_key>$PARAMETER_NAME</ifm|arg_key>\n<ifm|arg_type>$ARGUMENT_TYPE</ifm|arg_type>\n<ifm|arg_value>$PARAMETER_VALUE</ifm|arg_value>\n...\n</ifm|tool_call>\n</ifm|tool_calls>\n\nUse the parameter type shown in the tool definition. If that type contains anyOf or oneOf, use the actual argument value type instead. String and scalar parameters should be written as plain text. Array and object parameters should be written as JSON literals." }}
+{%- else -%}
+{{- raise_exception("Unsupported tool_call_format: '" + fmt + "'. Supported formats: json, xml, xml_typed.") }}
+{%- endif -%}
+{%- endmacro -%}
+
+{%- macro render_system_with_tools(tools_list, system_content, presentation_fmt, call_fmt) -%}
+{{- "<|ifm|im_start|>system\n# Tools\nYou may call one or more tools to assist with the user query.\n\nAvailable tools are:\n\n" }}
+{{- render_tool_presentation(tools_list, presentation_fmt) }}
+{{- "\n\nWhen calling tools, you MUST follow the tool-call format below:\n\n" }}
+{{- render_call_instructions(call_fmt) }}
+{%- if system_content -%}
+{{- "\n\n" + system_content }}
+{%- endif -%}
+{{- "<|ifm|im_end|>" }}
+{%- endmacro -%}
+
+{%- macro render_argument_value(value) -%}
+{%- if value is string -%}{{- value -}}{%- else -%}{{- value | tojson -}}{%- endif -%}
+{%- endmacro -%}
+
+{%- macro render_value_type(value) -%}
+{%- if value is none -%}null
+{%- elif value is boolean -%}boolean
+{%- elif value is integer -%}integer
+{%- elif value is number -%}number
+{%- elif value is string -%}string
+{%- elif value is mapping -%}object
+{%- elif value is sequence -%}array
+{%- else -%}any
+{%- endif -%}
+{%- endmacro -%}
+
+{%- macro schema_has_combinator(spec) -%}
+{%- if spec.oneOf or spec.anyOf -%}
+true
+{%- elif spec.type is defined and spec.type is sequence and spec.type is not string and spec.type | length > 1 -%}
+true
+{%- elif spec.type == "array" and 'items' in spec -%}
+{{- schema_has_combinator(spec['items']) -}}
+{%- elif spec.properties -%}
+    {%- set found = namespace(value='false') -%}
+    {%- for child_name, child_spec in spec.properties | items -%}
+        {%- if schema_has_combinator(child_spec) == 'true' -%}
+            {%- set found.value = 'true' -%}
+        {%- endif -%}
+    {%- endfor -%}
+{{- found.value -}}
+{%- else -%}
+false
+{%- endif -%}
+{%- endmacro -%}
+
+{%- macro render_arg_type(tools_list, tool_name, arg_name, value) -%}
+{%- set found = namespace(type='any') -%}
+{%- for tool in tools_list -%}
+    {%- set fn = tool.function if tool.function is defined else tool -%}
+    {%- if fn.name == tool_name and fn.parameters and fn.parameters.properties and arg_name in fn.parameters.properties -%}
+        {%- set spec = fn.parameters.properties[arg_name] -%}
+        {%- if spec is mapping and spec['$ref'] is string -%}
+            {%- set _r = spec['$ref'] -%}
+            {%- set _k = _r[8:] if _r.startswith('#/$defs/') else (_r[14:] if _r.startswith('#/definitions/') else none) -%}
+            {%- set _d = fn.parameters['$defs'] if fn.parameters['$defs'] is mapping else fn.parameters['definitions'] -%}
+            {%- set spec = dict((_d[_k] | items | list) + (spec | items | rejectattr('0', 'equalto', '$ref') | list)) if (_k is not none and _d is mapping and _d[_k] is mapping) else spec -%}
+        {%- endif -%}
+        {%- if schema_has_combinator(spec) == 'true' -%}
+            {%- set found.type = render_value_type(value) -%}
+        {%- else -%}
+            {%- set found.type = render_compact_type(spec) -%}
+        {%- endif -%}
+    {%- endif -%}
+{%- endfor -%}
+{{- found.type -}}
+{%- endmacro -%}
+
+{%- macro render_tool_calls_block(tool_calls, fmt, tools_list) -%}
+{{- "<ifm|tool_calls>" }}
+{%- for raw_tool_call in tool_calls -%}
+    {%- set tool_call = raw_tool_call.function if raw_tool_call.function else raw_tool_call -%}
+    {%- if tool_call.arguments is string -%}
+        {{- raise_exception("tool_call.arguments must be a dict, not a JSON string. Parse it before passing to the template.") -}}
+    {%- endif -%}
+    {%- if fmt == 'json' -%}
+{{- "\n<ifm|tool_call>{\"name\": \"" + tool_call.name + "\", \"arguments\": " }}{{ tool_call.arguments | tojson }}{{- "}</ifm|tool_call>" }}
+    {%- elif fmt == 'xml' or fmt == 'xml_typed' -%}
+{{- "\n<ifm|tool_call>" + tool_call.name + "\n" }}
+        {%- for key, value in tool_call.arguments | items -%}
+{{- "<ifm|arg_key>" + key + "</ifm|arg_key>\n" }}
+{%- if fmt == 'xml_typed' -%}
+{{- "<ifm|arg_type>" + render_arg_type(tools_list, tool_call.name, key, value) + "</ifm|arg_type>\n" }}
+{%- endif -%}
+{{- "<ifm|arg_value>" }}{{ render_argument_value(value) }}{{- "</ifm|arg_value>\n" }}
+        {%- endfor -%}
+{{- "</ifm|tool_call>" }}
+    {%- else -%}
+        {{- raise_exception("Unsupported tool_call_format: '" + fmt + "'. Supported formats: json, xml, xml_typed.") -}}
+    {%- endif -%}
+{%- endfor -%}
+{{- "\n</ifm|tool_calls>" }}
+{%- endmacro -%}
+
+{%- macro render_tool_response_messages(raw_content) -%}
+{%- if raw_content is string -%}
+{{- '<|ifm|im_start|>tool\n' + raw_content + '<|ifm|im_end|>' }}
+{%- elif raw_content is sequence and raw_content is not string and raw_content is not mapping -%}
+    {%- if raw_content | length == 0 -%}
+        {{- raise_exception("tool message content list must not be empty.") -}}
+    {%- endif -%}
+{{- '<|ifm|im_start|>tool\n' -}}
+    {%- for item in raw_content -%}
+        {%- if not loop.first -%}{{- '\n' -}}{%- endif -%}
+        {%- if item is string -%}
+{{- item -}}
+        {%- elif item is mapping and item.text is string -%}
+{{- item.text -}}
+        {%- else -%}
+{{- (item | tojson) -}}
+        {%- endif -%}
+    {%- endfor -%}
+{{- '<|ifm|im_end|>' -}}
+{%- else -%}
+{{- '<|ifm|im_start|>tool\n' }}{{ raw_content | tojson }}{{- '<|ifm|im_end|>' }}
+{%- endif -%}
+{%- endmacro -%}
+
+{%- set available_tools = tools if tools else [] -%}
+{%- if (not available_tools) and messages[0].role == 'system' and messages[0].get('tools') -%}
+    {%- set available_tools = messages[0]['tools'] -%}
+{%- endif -%}
+{%- if available_tools -%}
+    {{- validate_tools(available_tools, tool_presentation_fmt != 'json') }}
+    {%- set system_content = '' -%}
+    {%- if messages[0].role == 'system' and messages[0].content -%}
+        {%- set system_content = messages[0].content -%}
+    {%- endif -%}
+    {{- render_system_with_tools(available_tools, system_content, tool_presentation_fmt, tool_call_fmt) }}
+{%- else -%}
+    {%- if messages[0].role == 'system' -%}
+        {{- '<|ifm|im_start|>system\n' + messages[0].content + '<|ifm|im_end|>' }}
+    {%- endif -%}
+{%- endif -%}
+
+{%- for message in messages -%}
+    {%- if message.content is string -%}
+        {%- set content = message.content -%}
+    {%- else -%}
+        {%- set content = '' -%}
+    {%- endif -%}
+    {%- if (message.role == "user") or (message.role == "system" and not loop.first) -%}
+        {{- '<|ifm|im_start|>' + message.role + '\n' + content + '<|ifm|im_end|>' }}
+    {%- elif message.role == "assistant" -%}
+        {%- set thinking_content = '' -%}
+        {%- set think_tag = 'ifm|think' -%}
+        {%- if message.think is defined and message.think is string -%}
+            {%- set thinking_content = message.think -%}
+            {%- set think_tag = 'ifm|think' -%}
+        {%- elif message.think_fast is defined and message.think_fast is string -%}
+            {%- set thinking_content = message.think_fast -%}
+            {%- set think_tag = 'ifm|think_fast' -%}
+        {%- elif message.think_faster is defined and message.think_faster is string -%}
+            {%- set thinking_content = message.think_faster -%}
+            {%- set think_tag = 'ifm|think_faster' -%}
+        {%- elif message.reasoning_content is defined and message.reasoning_content is string -%}
+            {%- set thinking_content = message.reasoning_content -%}
+            {%- set think_tag = 'ifm|think' -%}
+        {%- elif message.reasoning is defined and message.reasoning is string -%}
+            {%- set thinking_content = message.reasoning -%}
+            {%- set think_tag = 'ifm|think' -%}
+        {%- elif message.think is not defined and message.reasoning is not defined and message.reasoning_content is not defined and message.think_fast is not defined and message.think_faster is not defined -%}
+            {{- raise_exception("Assistant message is missing a thinking field. Provide one of: think, reasoning, reasoning_content, think_fast, think_faster.") -}}
+        {%- else -%}
+            {{- raise_exception("Assistant thinking fields must be strings. Provide one of: think, reasoning, reasoning_content, think_fast, think_faster as a string.") -}}
+        {%- endif -%}
+        {{- '<|ifm|im_start|>' + message.role }}
+        {% generation %}
+        {%- if think_tag -%}
+            {%- if thinking_content -%}
+                {{- '<' + think_tag + '>\n' + thinking_content + '</' + think_tag + '>' + content }}
+            {%- else -%}
+                {{- '<' + think_tag + '>\n</' + think_tag + '>' + content }}
+            {%- endif -%}
+        {%- else -%}
+            {{- content }}
+        {%- endif -%}
+        {%- if message.tool_calls -%}
+            {{- render_tool_calls_block(message.tool_calls, tool_call_fmt, available_tools) }}
+        {%- endif -%}
+        {{- '<|ifm|im_end|>' -}}
+        {%- endgeneration -%}
+    {%- elif message.role == "tool" -%}
+        {{- render_tool_response_messages(message.content) }}
+    {%- endif -%}
+{%- endfor -%}
+{%- if add_generation_prompt -%}
+    {%- set effort = reasoning_effort | default('high') -%}
+    {%- if effort == 'high' -%}
+        {{- '<|ifm|im_start|>assistant\n<ifm|think>\n' }}
+    {%- elif effort == 'medium' -%}
+        {{- '<|ifm|im_start|>assistant\n<ifm|think_fast>\n' }}
+    {%- elif effort == 'low' -%}
+        {{- '<|ifm|im_start|>assistant\n<ifm|think_faster>\n' }}
+    {%- else -%}
+        {{- raise_exception("Unsupported reasoning_effort: '" + effort + "'. Supported values: high, medium, low.") -}}
+    {%- endif -%}
+{%- endif -%}

config.json

@@ -0,0 +1,104 @@
+{
+  "architectures": [
+    "K2HorizonForCausalLM"
+  ],
+  "attention_bias": false,
+  "attention_dropout": 0.0,
+  "attention_gate_func": null,
+  "auto_map": {
+    "AutoConfig": "configuration_k2_horizon.K2HorizonConfig",
+    "AutoModel": "modeling_k2_horizon.K2HorizonModel",
+    "AutoModelForCausalLM": "modeling_k2_horizon.K2HorizonForCausalLM"
+  },
+  "bos_token_id": 0,
+  "decoder_sparse_step": 1,
+  "dtype": "bfloat16",
+  "eos_token_id": [
+    1,
+    250019
+  ],
+  "head_dim": 128,
+  "hidden_act": "silu",
+  "hidden_size": 2560,
+  "initializer_range": 0.02,
+  "intermediate_size": 10240,
+  "layernorm_num_groups": 2,
+  "max_position_embeddings": 524288,
+  "mlp_only_layers": [
+    0,
+    1,
+    2,
+    3,
+    4,
+    5,
+    6,
+    7,
+    8,
+    9,
+    10,
+    11,
+    12,
+    13,
+    14,
+    15,
+    16,
+    17,
+    18,
+    19,
+    20,
+    21,
+    22,
+    23,
+    24,
+    25,
+    26,
+    27,
+    28,
+    29,
+    30,
+    31,
+    32,
+    33,
+    34,
+    35
+  ],
+  "model_type": "k2_horizon",
+  "moe_gate_bias": false,
+  "moe_intermediate_size": 0,
+  "mova_num_experts": 0,
+  "mova_num_experts_per_tok": 0,
+  "norm_topk_prob": true,
+  "num_attention_heads": 32,
+  "num_experts": 0,
+  "num_experts_per_tok": 0,
+  "num_hidden_layers": 36,
+  "num_key_value_heads": 8,
+  "num_shared_experts": 0,
+  "output_router_logits": false,
+  "pad_token_id": null,
+  "quantization": {
+    "group_size": 64,
+    "bits": 4,
+    "mode": "affine"
+  },
+  "quantization_config": {
+    "group_size": 64,
+    "bits": 4,
+    "mode": "affine"
+  },
+  "query_key_norm": false,
+  "rms_norm_eps": 1e-06,
+  "rope_head_dim": 128,
+  "rope_scaling": null,
+  "rope_theta": 10000000.0,
+  "router_aux_loss_coef": 0.001,
+  "router_scaling_factor": 1.0,
+  "router_score_func": "sigmoid",
+  "sliding_window": null,
+  "tie_word_embeddings": false,
+  "transformers_version": "5.13.0",
+  "use_cache": true,
+  "use_sliding_window": false,
+  "vocab_size": 250624,
+  "model_file": "k2_horizon_mlx.py"
+}

configuration_k2_horizon.py

@@ -0,0 +1,96 @@
+# Copyright 2024 The Qwen team, Alibaba Group and the HuggingFace Inc. team. 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.
+"""K2Horizon model configuration"""
+
+from huggingface_hub.dataclasses import strict
+
+from transformers.configuration_utils import PreTrainedConfig
+from transformers.modeling_rope_utils import RopeParameters
+
+
+@strict
+class K2HorizonConfig(PreTrainedConfig):
+    r"""
+    decoder_sparse_step (`int`, *optional*, defaults to 1):
+        The frequency of the MoE layer.
+    mlp_only_layers (`list[int]`, *optional*, defaults to `[]`):
+        Indicate which layers use K2HorizonMLP rather than K2HorizonSparseMoeBlock
+        The list contains layer index, from 0 to num_layers-1 if we have num_layers layers
+        If `mlp_only_layers` is empty, `decoder_sparse_step` is used to determine the sparsity.
+
+    ```python
+    >>> from transformers import K2HorizonModel, K2HorizonConfig
+
+    >>> # Initializing a K2Horizon style configuration
+    >>> configuration = K2HorizonConfig()
+    >>> model = K2HorizonModel(configuration)
+
+    >>> # Accessing the model configuration
+    >>> configuration = model.config
+    ```
+    """
+
+    model_type = "k2_horizon"
+    keys_to_ignore_at_inference = ["past_key_values"]
+
+    vocab_size: int = 151936
+    hidden_size: int = 2048
+    intermediate_size: int = 6144
+    num_hidden_layers: int = 24
+    num_attention_heads: int = 32
+    num_key_value_heads: int = 4
+    hidden_act: str = "silu"
+    max_position_embeddings: int = 32768
+    initializer_range: float = 0.02
+    rms_norm_eps: float = 1e-6
+    use_cache: bool = True
+    tie_word_embeddings: bool = False
+    rope_parameters: RopeParameters | dict | None = None
+    attention_bias: bool = False
+    use_sliding_window: bool = False
+    sliding_window: int | None = 4096
+    attention_dropout: float | int = 0.0
+    decoder_sparse_step: int = 1
+    moe_intermediate_size: int = 768
+    num_experts_per_tok: int = 8
+    num_experts: int = 128
+    norm_topk_prob: bool = False
+    output_router_logits: bool = False
+    router_aux_loss_coef: float = 0.001
+    mlp_only_layers: list[int] | None = None
+    pad_token_id: int | None = None
+    bos_token_id: int | None = None
+    eos_token_id: int | list[int] | None = None
+
+    head_dim: int = 128
+    query_key_norm: bool = True
+    moe_gate_bias: bool = False
+    layernorm_num_groups: int = 1
+    num_shared_experts: int = 0
+    router_score_func: str = "softmax"
+    router_scaling_factor: float | None = 1.0
+    rope_head_dim: int | None = None
+    attention_gate_func: str | None = None
+    mova_num_experts: int = 0
+    mova_num_experts_per_tok: int = 0
+
+    def __post_init__(self, **kwargs):
+        self.sliding_window = self.sliding_window if self.use_sliding_window else None
+        self.mlp_only_layers = [] if self.mlp_only_layers is None else self.mlp_only_layers
+        if self.router_scaling_factor is None:
+            self.router_scaling_factor = 1.0
+        super().__post_init__(**kwargs)
+
+
+__all__ = ["K2HorizonConfig"]

generation_config.json

@@ -0,0 +1,7 @@
+{
+    "bos_token_id": 0,
+    "eos_token_id": [
+        1,
+        250019
+    ]
+}

k2_horizon_mlx.py

@@ -0,0 +1,198 @@
+"""MLX-LM adapter for the dense K2 Horizon configuration.
+
+K2 Horizon uses the Llama-style decoder/attention layout, but its RMSNorm
+normalizes two groups of the hidden dimension independently.  This module
+keeps that detail instead of silently treating the checkpoint as Llama.
+"""
+
+from dataclasses import dataclass
+from typing import Any, Dict, List, Optional, Union
+
+import mlx.core as mx
+import mlx.nn as nn
+
+from mlx_lm.models.activations import swiglu
+from mlx_lm.models.base import BaseModelArgs, create_attention_mask, scaled_dot_product_attention
+from mlx_lm.models.cache import KVCache, RotatingKVCache
+from mlx_lm.models.rope_utils import initialize_rope
+
+
+@dataclass
+class ModelArgs(BaseModelArgs):
+    model_type: str
+    hidden_size: int
+    num_hidden_layers: int
+    intermediate_size: int
+    num_attention_heads: int
+    rms_norm_eps: float
+    vocab_size: int
+    num_key_value_heads: Optional[int] = None
+    head_dim: Optional[int] = None
+    max_position_embeddings: Optional[int] = None
+    # K2 stores RoPE settings under rope_parameters; this is the model default.
+    rope_theta: float = 10_000_000.0
+    rope_traditional: bool = False
+    rope_scaling: Optional[Dict[str, Union[float, str]]] = None
+    tie_word_embeddings: bool = False
+    attention_bias: bool = False
+    mlp_bias: bool = False
+    layernorm_num_groups: int = 2
+    sliding_window: Optional[int] = None
+    layer_types: Optional[List[str]] = None
+
+    def __post_init__(self):
+        if self.num_key_value_heads is None:
+            self.num_key_value_heads = self.num_attention_heads
+        if self.head_dim is None:
+            self.head_dim = self.hidden_size // self.num_attention_heads
+        if self.layer_types is None:
+            self.layer_types = ["full_attention"] * self.num_hidden_layers
+
+
+class GroupRMSNorm(nn.Module):
+    """K2's T5-style grouped RMS normalization."""
+
+    def __init__(self, dims: int, eps: float, groups: int):
+        super().__init__()
+        if dims % groups:
+            raise ValueError(f"hidden size {dims} is not divisible by {groups} groups")
+        self.weight = mx.ones((dims,))
+        self.groups = groups
+        self.eps = eps
+
+    def __call__(self, x: mx.array) -> mx.array:
+        x = mx.unflatten(x, axis=-1, shape=(self.groups, -1))
+        x = mx.fast.rms_norm(x, weight=None, eps=self.eps)
+        return self.weight * mx.flatten(x, -2)
+
+
+class Attention(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+        dim = args.hidden_size
+        self.n_heads = args.num_attention_heads
+        self.n_kv_heads = args.num_key_value_heads
+        self.head_dim = args.head_dim
+        self.scale = self.head_dim**-0.5
+        bias = args.attention_bias
+        self.q_proj = nn.Linear(dim, self.n_heads * self.head_dim, bias=bias)
+        self.k_proj = nn.Linear(dim, self.n_kv_heads * self.head_dim, bias=bias)
+        self.v_proj = nn.Linear(dim, self.n_kv_heads * self.head_dim, bias=bias)
+        self.o_proj = nn.Linear(self.n_heads * self.head_dim, dim, bias=bias)
+        self.rope = initialize_rope(
+            self.head_dim,
+            args.rope_theta,
+            args.rope_traditional,
+            args.rope_scaling,
+            args.max_position_embeddings,
+        )
+
+    def __call__(self, x: mx.array, mask=None, cache=None) -> mx.array:
+        B, L, _ = x.shape
+        q = self.q_proj(x).reshape(B, L, self.n_heads, self.head_dim).transpose(0, 2, 1, 3)
+        k = self.k_proj(x).reshape(B, L, self.n_kv_heads, self.head_dim).transpose(0, 2, 1, 3)
+        v = self.v_proj(x).reshape(B, L, self.n_kv_heads, self.head_dim).transpose(0, 2, 1, 3)
+        offset = cache.offset if cache is not None else 0
+        q = self.rope(q, offset=offset)
+        k = self.rope(k, offset=offset)
+        if cache is not None:
+            k, v = cache.update_and_fetch(k, v)
+        out = scaled_dot_product_attention(q, k, v, cache=cache, scale=self.scale, mask=mask)
+        out = out.transpose(0, 2, 1, 3).reshape(B, L, -1)
+        return self.o_proj(out)
+
+
+class MLP(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+        self.gate_proj = nn.Linear(args.hidden_size, args.intermediate_size, bias=args.mlp_bias)
+        self.down_proj = nn.Linear(args.intermediate_size, args.hidden_size, bias=args.mlp_bias)
+        self.up_proj = nn.Linear(args.hidden_size, args.intermediate_size, bias=args.mlp_bias)
+
+    def __call__(self, x):
+        return self.down_proj(swiglu(self.gate_proj(x), self.up_proj(x)))
+
+
+class TransformerBlock(nn.Module):
+    def __init__(self, args: ModelArgs, use_sliding: bool = False):
+        super().__init__()
+        self.self_attn = Attention(args)
+        self.mlp = MLP(args)
+        self.input_layernorm = GroupRMSNorm(
+            args.hidden_size, args.rms_norm_eps, args.layernorm_num_groups
+        )
+        self.post_attention_layernorm = GroupRMSNorm(
+            args.hidden_size, args.rms_norm_eps, args.layernorm_num_groups
+        )
+        self.use_sliding = use_sliding
+
+    def __call__(self, x, mask=None, cache=None):
+        h = x + self.self_attn(self.input_layernorm(x), mask, cache)
+        return h + self.mlp(self.post_attention_layernorm(h))
+
+
+class K2HorizonModel(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+        self.args = args
+        self.vocab_size = args.vocab_size
+        self.embed_tokens = nn.Embedding(args.vocab_size, args.hidden_size)
+        self.layers = [
+            TransformerBlock(args, use_sliding=t == "sliding_attention")
+            for t in args.layer_types
+        ]
+        self.norm = GroupRMSNorm(args.hidden_size, args.rms_norm_eps, args.layernorm_num_groups)
+        self.sliding_window = args.sliding_window
+        self.fa_idx = 0
+        self.swa_idx = None
+        for i, layer in enumerate(self.layers):
+            if layer.use_sliding:
+                self.swa_idx = i
+                break
+
+    def __call__(self, inputs, cache=None, input_embeddings=None):
+        h = self.embed_tokens(inputs) if input_embeddings is None else input_embeddings
+        if cache is None:
+            cache = [None] * len(self.layers)
+        fa_mask = create_attention_mask(h, cache[self.fa_idx])
+        swa_mask = None
+        if self.swa_idx is not None:
+            swa_mask = create_attention_mask(
+                h, cache[self.swa_idx], window_size=self.sliding_window
+            )
+        for layer, c in zip(self.layers, cache):
+            h = layer(h, swa_mask if layer.use_sliding else fa_mask, c)
+        return self.norm(h)
+
+
+class Model(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+        self.args = args
+        self.model_type = args.model_type
+        self.model = K2HorizonModel(args)
+        if not args.tie_word_embeddings:
+            self.lm_head = nn.Linear(args.hidden_size, args.vocab_size, bias=False)
+
+    def __call__(self, inputs, cache=None, input_embeddings=None):
+        h = self.model(inputs, cache, input_embeddings)
+        return self.lm_head(h)
+
+    def sanitize(self, weights):
+        return {
+            k: v
+            for k, v in weights.items()
+            if "rotary_emb.inv_freq" not in k
+        }
+
+    @property
+    def layers(self):
+        return self.model.layers
+
+    def make_cache(self):
+        return [
+            RotatingKVCache(max_size=self.model.sliding_window)
+            if layer.use_sliding
+            else KVCache()
+            for layer in self.model.layers
+        ]

model.safetensors

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

model.safetensors.index.json

@@ -0,0 +1,843 @@
+{
+    "metadata": {
+        "total_size": 2845537280,
+        "total_parameters": 5058255360
+    },
+    "weight_map": {
+        "lm_head.biases": "model.safetensors",
+        "lm_head.scales": "model.safetensors",
+        "lm_head.weight": "model.safetensors",
+        "model.embed_tokens.biases": "model.safetensors",
+        "model.embed_tokens.scales": "model.safetensors",
+        "model.embed_tokens.weight": "model.safetensors",
+        "model.layers.0.input_layernorm.weight": "model.safetensors",
+        "model.layers.0.mlp.down_proj.biases": "model.safetensors",
+        "model.layers.0.mlp.down_proj.scales": "model.safetensors",
+        "model.layers.0.mlp.down_proj.weight": "model.safetensors",
+        "model.layers.0.mlp.gate_proj.biases": "model.safetensors",
+        "model.layers.0.mlp.gate_proj.scales": "model.safetensors",
+        "model.layers.0.mlp.gate_proj.weight": "model.safetensors",
+        "model.layers.0.mlp.up_proj.biases": "model.safetensors",
+        "model.layers.0.mlp.up_proj.scales": "model.safetensors",
+        "model.layers.0.mlp.up_proj.weight": "model.safetensors",
+        "model.layers.0.post_attention_layernorm.weight": "model.safetensors",
+        "model.layers.0.self_attn.k_proj.biases": "model.safetensors",
+        "model.layers.0.self_attn.k_proj.scales": "model.safetensors",
+        "model.layers.0.self_attn.k_proj.weight": "model.safetensors",
+        "model.layers.0.self_attn.o_proj.biases": "model.safetensors",
+        "model.layers.0.self_attn.o_proj.scales": "model.safetensors",
+        "model.layers.0.self_attn.o_proj.weight": "model.safetensors",
+        "model.layers.0.self_attn.q_proj.biases": "model.safetensors",
+        "model.layers.0.self_attn.q_proj.scales": "model.safetensors",
+        "model.layers.0.self_attn.q_proj.weight": "model.safetensors",
+        "model.layers.0.self_attn.v_proj.biases": "model.safetensors",
+        "model.layers.0.self_attn.v_proj.scales": "model.safetensors",
+        "model.layers.0.self_attn.v_proj.weight": "model.safetensors",
+        "model.layers.1.input_layernorm.weight": "model.safetensors",
+        "model.layers.1.mlp.down_proj.biases": "model.safetensors",
+        "model.layers.1.mlp.down_proj.scales": "model.safetensors",
+        "model.layers.1.mlp.down_proj.weight": "model.safetensors",
+        "model.layers.1.mlp.gate_proj.biases": "model.safetensors",
+        "model.layers.1.mlp.gate_proj.scales": "model.safetensors",
+        "model.layers.1.mlp.gate_proj.weight": "model.safetensors",
+        "model.layers.1.mlp.up_proj.biases": "model.safetensors",
+        "model.layers.1.mlp.up_proj.scales": "model.safetensors",
+        "model.layers.1.mlp.up_proj.weight": "model.safetensors",
+        "model.layers.1.post_attention_layernorm.weight": "model.safetensors",
+        "model.layers.1.self_attn.k_proj.biases": "model.safetensors",
+        "model.layers.1.self_attn.k_proj.scales": "model.safetensors",
+        "model.layers.1.self_attn.k_proj.weight": "model.safetensors",
+        "model.layers.1.self_attn.o_proj.biases": "model.safetensors",
+        "model.layers.1.self_attn.o_proj.scales": "model.safetensors",
+        "model.layers.1.self_attn.o_proj.weight": "model.safetensors",
+        "model.layers.1.self_attn.q_proj.biases": "model.safetensors",
+        "model.layers.1.self_attn.q_proj.scales": "model.safetensors",
+        "model.layers.1.self_attn.q_proj.weight": "model.safetensors",
+        "model.layers.1.self_attn.v_proj.biases": "model.safetensors",
+        "model.layers.1.self_attn.v_proj.scales": "model.safetensors",
+        "model.layers.1.self_attn.v_proj.weight": "model.safetensors",
+        "model.layers.10.input_layernorm.weight": "model.safetensors",
+        "model.layers.10.mlp.down_proj.biases": "model.safetensors",
+        "model.layers.10.mlp.down_proj.scales": "model.safetensors",
+        "model.layers.10.mlp.down_proj.weight": "model.safetensors",
+        "model.layers.10.mlp.gate_proj.biases": "model.safetensors",
+        "model.layers.10.mlp.gate_proj.scales": "model.safetensors",
+        "model.layers.10.mlp.gate_proj.weight": "model.safetensors",
+        "model.layers.10.mlp.up_proj.biases": "model.safetensors",
+        "model.layers.10.mlp.up_proj.scales": "model.safetensors",
+        "model.layers.10.mlp.up_proj.weight": "model.safetensors",
+        "model.layers.10.post_attention_layernorm.weight": "model.safetensors",
+        "model.layers.10.self_attn.k_proj.biases": "model.safetensors",
+        "model.layers.10.self_attn.k_proj.scales": "model.safetensors",
+        "model.layers.10.self_attn.k_proj.weight": "model.safetensors",
+        "model.layers.10.self_attn.o_proj.biases": "model.safetensors",
+        "model.layers.10.self_attn.o_proj.scales": "model.safetensors",
+        "model.layers.10.self_attn.o_proj.weight": "model.safetensors",
+        "model.layers.10.self_attn.q_proj.biases": "model.safetensors",
+        "model.layers.10.self_attn.q_proj.scales": "model.safetensors",
+        "model.layers.10.self_attn.q_proj.weight": "model.safetensors",
+        "model.layers.10.self_attn.v_proj.biases": "model.safetensors",
+        "model.layers.10.self_attn.v_proj.scales": "model.safetensors",
+        "model.layers.10.self_attn.v_proj.weight": "model.safetensors",
+        "model.layers.11.input_layernorm.weight": "model.safetensors",
+        "model.layers.11.mlp.down_proj.biases": "model.safetensors",
+        "model.layers.11.mlp.down_proj.scales": "model.safetensors",
+        "model.layers.11.mlp.down_proj.weight": "model.safetensors",
+        "model.layers.11.mlp.gate_proj.biases": "model.safetensors",
+        "model.layers.11.mlp.gate_proj.scales": "model.safetensors",
+        "model.layers.11.mlp.gate_proj.weight": "model.safetensors",
+        "model.layers.11.mlp.up_proj.biases": "model.safetensors",
+        "model.layers.11.mlp.up_proj.scales": "model.safetensors",
+        "model.layers.11.mlp.up_proj.weight": "model.safetensors",
+        "model.layers.11.post_attention_layernorm.weight": "model.safetensors",
+        "model.layers.11.self_attn.k_proj.biases": "model.safetensors",
+        "model.layers.11.self_attn.k_proj.scales": "model.safetensors",
+        "model.layers.11.self_attn.k_proj.weight": "model.safetensors",
+        "model.layers.11.self_attn.o_proj.biases": "model.safetensors",
+        "model.layers.11.self_attn.o_proj.scales": "model.safetensors",
+        "model.layers.11.self_attn.o_proj.weight": "model.safetensors",
+        "model.layers.11.self_attn.q_proj.biases": "model.safetensors",
+        "model.layers.11.self_attn.q_proj.scales": "model.safetensors",
+        "model.layers.11.self_attn.q_proj.weight": "model.safetensors",
+        "model.layers.11.self_attn.v_proj.biases": "model.safetensors",
+        "model.layers.11.self_attn.v_proj.scales": "model.safetensors",
+        "model.layers.11.self_attn.v_proj.weight": "model.safetensors",
+        "model.layers.12.input_layernorm.weight": "model.safetensors",
+        "model.layers.12.mlp.down_proj.biases": "model.safetensors",
+        "model.layers.12.mlp.down_proj.scales": "model.safetensors",
+        "model.layers.12.mlp.down_proj.weight": "model.safetensors",
+        "model.layers.12.mlp.gate_proj.biases": "model.safetensors",
+        "model.layers.12.mlp.gate_proj.scales": "model.safetensors",
+        "model.layers.12.mlp.gate_proj.weight": "model.safetensors",
+        "model.layers.12.mlp.up_proj.biases": "model.safetensors",
+        "model.layers.12.mlp.up_proj.scales": "model.safetensors",
+        "model.layers.12.mlp.up_proj.weight": "model.safetensors",
+        "model.layers.12.post_attention_layernorm.weight": "model.safetensors",
+        "model.layers.12.self_attn.k_proj.biases": "model.safetensors",
+        "model.layers.12.self_attn.k_proj.scales": "model.safetensors",
+        "model.layers.12.self_attn.k_proj.weight": "model.safetensors",
+        "model.layers.12.self_attn.o_proj.biases": "model.safetensors",
+        "model.layers.12.self_attn.o_proj.scales": "model.safetensors",
+        "model.layers.12.self_attn.o_proj.weight": "model.safetensors",
+        "model.layers.12.self_attn.q_proj.biases": "model.safetensors",
+        "model.layers.12.self_attn.q_proj.scales": "model.safetensors",
+        "model.layers.12.self_attn.q_proj.weight": "model.safetensors",
+        "model.layers.12.self_attn.v_proj.biases": "model.safetensors",
+        "model.layers.12.self_attn.v_proj.scales": "model.safetensors",
+        "model.layers.12.self_attn.v_proj.weight": "model.safetensors",
+        "model.layers.13.input_layernorm.weight": "model.safetensors",
+        "model.layers.13.mlp.down_proj.biases": "model.safetensors",
+        "model.layers.13.mlp.down_proj.scales": "model.safetensors",
+        "model.layers.13.mlp.down_proj.weight": "model.safetensors",
+        "model.layers.13.mlp.gate_proj.biases": "model.safetensors",
+        "model.layers.13.mlp.gate_proj.scales": "model.safetensors",
+        "model.layers.13.mlp.gate_proj.weight": "model.safetensors",
+        "model.layers.13.mlp.up_proj.biases": "model.safetensors",
+        "model.layers.13.mlp.up_proj.scales": "model.safetensors",
+        "model.layers.13.mlp.up_proj.weight": "model.safetensors",
+        "model.layers.13.post_attention_layernorm.weight": "model.safetensors",
+        "model.layers.13.self_attn.k_proj.biases": "model.safetensors",
+        "model.layers.13.self_attn.k_proj.scales": "model.safetensors",
+        "model.layers.13.self_attn.k_proj.weight": "model.safetensors",
+        "model.layers.13.self_attn.o_proj.biases": "model.safetensors",
+        "model.layers.13.self_attn.o_proj.scales": "model.safetensors",
+        "model.layers.13.self_attn.o_proj.weight": "model.safetensors",
+        "model.layers.13.self_attn.q_proj.biases": "model.safetensors",
+        "model.layers.13.self_attn.q_proj.scales": "model.safetensors",
+        "model.layers.13.self_attn.q_proj.weight": "model.safetensors",
+        "model.layers.13.self_attn.v_proj.biases": "model.safetensors",
+        "model.layers.13.self_attn.v_proj.scales": "model.safetensors",
+        "model.layers.13.self_attn.v_proj.weight": "model.safetensors",
+        "model.layers.14.input_layernorm.weight": "model.safetensors",
+        "model.layers.14.mlp.down_proj.biases": "model.safetensors",
+        "model.layers.14.mlp.down_proj.scales": "model.safetensors",
+        "model.layers.14.mlp.down_proj.weight": "model.safetensors",
+        "model.layers.14.mlp.gate_proj.biases": "model.safetensors",
+        "model.layers.14.mlp.gate_proj.scales": "model.safetensors",
+        "model.layers.14.mlp.gate_proj.weight": "model.safetensors",
+        "model.layers.14.mlp.up_proj.biases": "model.safetensors",
+        "model.layers.14.mlp.up_proj.scales": "model.safetensors",
+        "model.layers.14.mlp.up_proj.weight": "model.safetensors",
+        "model.layers.14.post_attention_layernorm.weight": "model.safetensors",
+        "model.layers.14.self_attn.k_proj.biases": "model.safetensors",
+        "model.layers.14.self_attn.k_proj.scales": "model.safetensors",
+        "model.layers.14.self_attn.k_proj.weight": "model.safetensors",
+        "model.layers.14.self_attn.o_proj.biases": "model.safetensors",
+        "model.layers.14.self_attn.o_proj.scales": "model.safetensors",
+        "model.layers.14.self_attn.o_proj.weight": "model.safetensors",
+        "model.layers.14.self_attn.q_proj.biases": "model.safetensors",
+        "model.layers.14.self_attn.q_proj.scales": "model.safetensors",
+        "model.layers.14.self_attn.q_proj.weight": "model.safetensors",
+        "model.layers.14.self_attn.v_proj.biases": "model.safetensors",
+        "model.layers.14.self_attn.v_proj.scales": "model.safetensors",
+        "model.layers.14.self_attn.v_proj.weight": "model.safetensors",
+        "model.layers.15.input_layernorm.weight": "model.safetensors",
+        "model.layers.15.mlp.down_proj.biases": "model.safetensors",
+        "model.layers.15.mlp.down_proj.scales": "model.safetensors",
+        "model.layers.15.mlp.down_proj.weight": "model.safetensors",
+        "model.layers.15.mlp.gate_proj.biases": "model.safetensors",
+        "model.layers.15.mlp.gate_proj.scales": "model.safetensors",
+        "model.layers.15.mlp.gate_proj.weight": "model.safetensors",
+        "model.layers.15.mlp.up_proj.biases": "model.safetensors",
+        "model.layers.15.mlp.up_proj.scales": "model.safetensors",
+        "model.layers.15.mlp.up_proj.weight": "model.safetensors",
+        "model.layers.15.post_attention_layernorm.weight": "model.safetensors",
+        "model.layers.15.self_attn.k_proj.biases": "model.safetensors",
+        "model.layers.15.self_attn.k_proj.scales": "model.safetensors",
+        "model.layers.15.self_attn.k_proj.weight": "model.safetensors",
+        "model.layers.15.self_attn.o_proj.biases": "model.safetensors",
+        "model.layers.15.self_attn.o_proj.scales": "model.safetensors",
+        "model.layers.15.self_attn.o_proj.weight": "model.safetensors",
+        "model.layers.15.self_attn.q_proj.biases": "model.safetensors",
+        "model.layers.15.self_attn.q_proj.scales": "model.safetensors",
+        "model.layers.15.self_attn.q_proj.weight": "model.safetensors",
+        "model.layers.15.self_attn.v_proj.biases": "model.safetensors",
+        "model.layers.15.self_attn.v_proj.scales": "model.safetensors",
+        "model.layers.15.self_attn.v_proj.weight": "model.safetensors",
+        "model.layers.16.input_layernorm.weight": "model.safetensors",
+        "model.layers.16.mlp.down_proj.biases": "model.safetensors",
+        "model.layers.16.mlp.down_proj.scales": "model.safetensors",
+        "model.layers.16.mlp.down_proj.weight": "model.safetensors",
+        "model.layers.16.mlp.gate_proj.biases": "model.safetensors",
+        "model.layers.16.mlp.gate_proj.scales": "model.safetensors",
+        "model.layers.16.mlp.gate_proj.weight": "model.safetensors",
+        "model.layers.16.mlp.up_proj.biases": "model.safetensors",
+        "model.layers.16.mlp.up_proj.scales": "model.safetensors",
+        "model.layers.16.mlp.up_proj.weight": "model.safetensors",
+        "model.layers.16.post_attention_layernorm.weight": "model.safetensors",
+        "model.layers.16.self_attn.k_proj.biases": "model.safetensors",
+        "model.layers.16.self_attn.k_proj.scales": "model.safetensors",
+        "model.layers.16.self_attn.k_proj.weight": "model.safetensors",
+        "model.layers.16.self_attn.o_proj.biases": "model.safetensors",
+        "model.layers.16.self_attn.o_proj.scales": "model.safetensors",
+        "model.layers.16.self_attn.o_proj.weight": "model.safetensors",
+        "model.layers.16.self_attn.q_proj.biases": "model.safetensors",
+        "model.layers.16.self_attn.q_proj.scales": "model.safetensors",
+        "model.layers.16.self_attn.q_proj.weight": "model.safetensors",
+        "model.layers.16.self_attn.v_proj.biases": "model.safetensors",
+        "model.layers.16.self_attn.v_proj.scales": "model.safetensors",
+        "model.layers.16.self_attn.v_proj.weight": "model.safetensors",
+        "model.layers.17.input_layernorm.weight": "model.safetensors",
+        "model.layers.17.mlp.down_proj.biases": "model.safetensors",
+        "model.layers.17.mlp.down_proj.scales": "model.safetensors",
+        "model.layers.17.mlp.down_proj.weight": "model.safetensors",
+        "model.layers.17.mlp.gate_proj.biases": "model.safetensors",
+        "model.layers.17.mlp.gate_proj.scales": "model.safetensors",
+        "model.layers.17.mlp.gate_proj.weight": "model.safetensors",
+        "model.layers.17.mlp.up_proj.biases": "model.safetensors",
+        "model.layers.17.mlp.up_proj.scales": "model.safetensors",
+        "model.layers.17.mlp.up_proj.weight": "model.safetensors",
+        "model.layers.17.post_attention_layernorm.weight": "model.safetensors",
+        "model.layers.17.self_attn.k_proj.biases": "model.safetensors",
+        "model.layers.17.self_attn.k_proj.scales": "model.safetensors",
+        "model.layers.17.self_attn.k_proj.weight": "model.safetensors",
+        "model.layers.17.self_attn.o_proj.biases": "model.safetensors",
+        "model.layers.17.self_attn.o_proj.scales": "model.safetensors",
+        "model.layers.17.self_attn.o_proj.weight": "model.safetensors",
+        "model.layers.17.self_attn.q_proj.biases": "model.safetensors",
+        "model.layers.17.self_attn.q_proj.scales": "model.safetensors",
+        "model.layers.17.self_attn.q_proj.weight": "model.safetensors",
+        "model.layers.17.self_attn.v_proj.biases": "model.safetensors",
+        "model.layers.17.self_attn.v_proj.scales": "model.safetensors",
+        "model.layers.17.self_attn.v_proj.weight": "model.safetensors",
+        "model.layers.18.input_layernorm.weight": "model.safetensors",
+        "model.layers.18.mlp.down_proj.biases": "model.safetensors",
+        "model.layers.18.mlp.down_proj.scales": "model.safetensors",
+        "model.layers.18.mlp.down_proj.weight": "model.safetensors",
+        "model.layers.18.mlp.gate_proj.biases": "model.safetensors",
+        "model.layers.18.mlp.gate_proj.scales": "model.safetensors",
+        "model.layers.18.mlp.gate_proj.weight": "model.safetensors",
+        "model.layers.18.mlp.up_proj.biases": "model.safetensors",
+        "model.layers.18.mlp.up_proj.scales": "model.safetensors",
+        "model.layers.18.mlp.up_proj.weight": "model.safetensors",
+        "model.layers.18.post_attention_layernorm.weight": "model.safetensors",
+        "model.layers.18.self_attn.k_proj.biases": "model.safetensors",
+        "model.layers.18.self_attn.k_proj.scales": "model.safetensors",
+        "model.layers.18.self_attn.k_proj.weight": "model.safetensors",
+        "model.layers.18.self_attn.o_proj.biases": "model.safetensors",
+        "model.layers.18.self_attn.o_proj.scales": "model.safetensors",
+        "model.layers.18.self_attn.o_proj.weight": "model.safetensors",
+        "model.layers.18.self_attn.q_proj.biases": "model.safetensors",
+        "model.layers.18.self_attn.q_proj.scales": "model.safetensors",
+        "model.layers.18.self_attn.q_proj.weight": "model.safetensors",
+        "model.layers.18.self_attn.v_proj.biases": "model.safetensors",
+        "model.layers.18.self_attn.v_proj.scales": "model.safetensors",
+        "model.layers.18.self_attn.v_proj.weight": "model.safetensors",
+        "model.layers.19.input_layernorm.weight": "model.safetensors",
+        "model.layers.19.mlp.down_proj.biases": "model.safetensors",
+        "model.layers.19.mlp.down_proj.scales": "model.safetensors",
+        "model.layers.19.mlp.down_proj.weight": "model.safetensors",
+        "model.layers.19.mlp.gate_proj.biases": "model.safetensors",
+        "model.layers.19.mlp.gate_proj.scales": "model.safetensors",
+        "model.layers.19.mlp.gate_proj.weight": "model.safetensors",
+        "model.layers.19.mlp.up_proj.biases": "model.safetensors",
+        "model.layers.19.mlp.up_proj.scales": "model.safetensors",
+        "model.layers.19.mlp.up_proj.weight": "model.safetensors",
+        "model.layers.19.post_attention_layernorm.weight": "model.safetensors",
+        "model.layers.19.self_attn.k_proj.biases": "model.safetensors",
+        "model.layers.19.self_attn.k_proj.scales": "model.safetensors",
+        "model.layers.19.self_attn.k_proj.weight": "model.safetensors",
+        "model.layers.19.self_attn.o_proj.biases": "model.safetensors",
+        "model.layers.19.self_attn.o_proj.scales": "model.safetensors",
+        "model.layers.19.self_attn.o_proj.weight": "model.safetensors",
+        "model.layers.19.self_attn.q_proj.biases": "model.safetensors",
+        "model.layers.19.self_attn.q_proj.scales": "model.safetensors",
+        "model.layers.19.self_attn.q_proj.weight": "model.safetensors",
+        "model.layers.19.self_attn.v_proj.biases": "model.safetensors",
+        "model.layers.19.self_attn.v_proj.scales": "model.safetensors",
+        "model.layers.19.self_attn.v_proj.weight": "model.safetensors",
+        "model.layers.2.input_layernorm.weight": "model.safetensors",
+        "model.layers.2.mlp.down_proj.biases": "model.safetensors",
+        "model.layers.2.mlp.down_proj.scales": "model.safetensors",
+        "model.layers.2.mlp.down_proj.weight": "model.safetensors",
+        "model.layers.2.mlp.gate_proj.biases": "model.safetensors",
+        "model.layers.2.mlp.gate_proj.scales": "model.safetensors",
+        "model.layers.2.mlp.gate_proj.weight": "model.safetensors",
+        "model.layers.2.mlp.up_proj.biases": "model.safetensors",
+        "model.layers.2.mlp.up_proj.scales": "model.safetensors",
+        "model.layers.2.mlp.up_proj.weight": "model.safetensors",
+        "model.layers.2.post_attention_layernorm.weight": "model.safetensors",
+        "model.layers.2.self_attn.k_proj.biases": "model.safetensors",
+        "model.layers.2.self_attn.k_proj.scales": "model.safetensors",
+        "model.layers.2.self_attn.k_proj.weight": "model.safetensors",
+        "model.layers.2.self_attn.o_proj.biases": "model.safetensors",
+        "model.layers.2.self_attn.o_proj.scales": "model.safetensors",
+        "model.layers.2.self_attn.o_proj.weight": "model.safetensors",
+        "model.layers.2.self_attn.q_proj.biases": "model.safetensors",
+        "model.layers.2.self_attn.q_proj.scales": "model.safetensors",
+        "model.layers.2.self_attn.q_proj.weight": "model.safetensors",
+        "model.layers.2.self_attn.v_proj.biases": "model.safetensors",
+        "model.layers.2.self_attn.v_proj.scales": "model.safetensors",
+        "model.layers.2.self_attn.v_proj.weight": "model.safetensors",
+        "model.layers.20.input_layernorm.weight": "model.safetensors",
+        "model.layers.20.mlp.down_proj.biases": "model.safetensors",
+        "model.layers.20.mlp.down_proj.scales": "model.safetensors",
+        "model.layers.20.mlp.down_proj.weight": "model.safetensors",
+        "model.layers.20.mlp.gate_proj.biases": "model.safetensors",
+        "model.layers.20.mlp.gate_proj.scales": "model.safetensors",
+        "model.layers.20.mlp.gate_proj.weight": "model.safetensors",
+        "model.layers.20.mlp.up_proj.biases": "model.safetensors",
+        "model.layers.20.mlp.up_proj.scales": "model.safetensors",
+        "model.layers.20.mlp.up_proj.weight": "model.safetensors",
+        "model.layers.20.post_attention_layernorm.weight": "model.safetensors",
+        "model.layers.20.self_attn.k_proj.biases": "model.safetensors",
+        "model.layers.20.self_attn.k_proj.scales": "model.safetensors",
+        "model.layers.20.self_attn.k_proj.weight": "model.safetensors",
+        "model.layers.20.self_attn.o_proj.biases": "model.safetensors",
+        "model.layers.20.self_attn.o_proj.scales": "model.safetensors",
+        "model.layers.20.self_attn.o_proj.weight": "model.safetensors",
+        "model.layers.20.self_attn.q_proj.biases": "model.safetensors",
+        "model.layers.20.self_attn.q_proj.scales": "model.safetensors",
+        "model.layers.20.self_attn.q_proj.weight": "model.safetensors",
+        "model.layers.20.self_attn.v_proj.biases": "model.safetensors",
+        "model.layers.20.self_attn.v_proj.scales": "model.safetensors",
+        "model.layers.20.self_attn.v_proj.weight": "model.safetensors",
+        "model.layers.21.input_layernorm.weight": "model.safetensors",
+        "model.layers.21.mlp.down_proj.biases": "model.safetensors",
+        "model.layers.21.mlp.down_proj.scales": "model.safetensors",
+        "model.layers.21.mlp.down_proj.weight": "model.safetensors",
+        "model.layers.21.mlp.gate_proj.biases": "model.safetensors",
+        "model.layers.21.mlp.gate_proj.scales": "model.safetensors",
+        "model.layers.21.mlp.gate_proj.weight": "model.safetensors",
+        "model.layers.21.mlp.up_proj.biases": "model.safetensors",
+        "model.layers.21.mlp.up_proj.scales": "model.safetensors",
+        "model.layers.21.mlp.up_proj.weight": "model.safetensors",
+        "model.layers.21.post_attention_layernorm.weight": "model.safetensors",
+        "model.layers.21.self_attn.k_proj.biases": "model.safetensors",
+        "model.layers.21.self_attn.k_proj.scales": "model.safetensors",
+        "model.layers.21.self_attn.k_proj.weight": "model.safetensors",
+        "model.layers.21.self_attn.o_proj.biases": "model.safetensors",
+        "model.layers.21.self_attn.o_proj.scales": "model.safetensors",
+        "model.layers.21.self_attn.o_proj.weight": "model.safetensors",
+        "model.layers.21.self_attn.q_proj.biases": "model.safetensors",
+        "model.layers.21.self_attn.q_proj.scales": "model.safetensors",
+        "model.layers.21.self_attn.q_proj.weight": "model.safetensors",
+        "model.layers.21.self_attn.v_proj.biases": "model.safetensors",
+        "model.layers.21.self_attn.v_proj.scales": "model.safetensors",
+        "model.layers.21.self_attn.v_proj.weight": "model.safetensors",
+        "model.layers.22.input_layernorm.weight": "model.safetensors",
+        "model.layers.22.mlp.down_proj.biases": "model.safetensors",
+        "model.layers.22.mlp.down_proj.scales": "model.safetensors",
+        "model.layers.22.mlp.down_proj.weight": "model.safetensors",
+        "model.layers.22.mlp.gate_proj.biases": "model.safetensors",
+        "model.layers.22.mlp.gate_proj.scales": "model.safetensors",
+        "model.layers.22.mlp.gate_proj.weight": "model.safetensors",
+        "model.layers.22.mlp.up_proj.biases": "model.safetensors",
+        "model.layers.22.mlp.up_proj.scales": "model.safetensors",
+        "model.layers.22.mlp.up_proj.weight": "model.safetensors",
+        "model.layers.22.post_attention_layernorm.weight": "model.safetensors",
+        "model.layers.22.self_attn.k_proj.biases": "model.safetensors",
+        "model.layers.22.self_attn.k_proj.scales": "model.safetensors",
+        "model.layers.22.self_attn.k_proj.weight": "model.safetensors",
+        "model.layers.22.self_attn.o_proj.biases": "model.safetensors",
+        "model.layers.22.self_attn.o_proj.scales": "model.safetensors",
+        "model.layers.22.self_attn.o_proj.weight": "model.safetensors",
+        "model.layers.22.self_attn.q_proj.biases": "model.safetensors",
+        "model.layers.22.self_attn.q_proj.scales": "model.safetensors",
+        "model.layers.22.self_attn.q_proj.weight": "model.safetensors",
+        "model.layers.22.self_attn.v_proj.biases": "model.safetensors",
+        "model.layers.22.self_attn.v_proj.scales": "model.safetensors",
+        "model.layers.22.self_attn.v_proj.weight": "model.safetensors",
+        "model.layers.23.input_layernorm.weight": "model.safetensors",
+        "model.layers.23.mlp.down_proj.biases": "model.safetensors",
+        "model.layers.23.mlp.down_proj.scales": "model.safetensors",
+        "model.layers.23.mlp.down_proj.weight": "model.safetensors",
+        "model.layers.23.mlp.gate_proj.biases": "model.safetensors",
+        "model.layers.23.mlp.gate_proj.scales": "model.safetensors",
+        "model.layers.23.mlp.gate_proj.weight": "model.safetensors",
+        "model.layers.23.mlp.up_proj.biases": "model.safetensors",
+        "model.layers.23.mlp.up_proj.scales": "model.safetensors",
+        "model.layers.23.mlp.up_proj.weight": "model.safetensors",
+        "model.layers.23.post_attention_layernorm.weight": "model.safetensors",
+        "model.layers.23.self_attn.k_proj.biases": "model.safetensors",
+        "model.layers.23.self_attn.k_proj.scales": "model.safetensors",
+        "model.layers.23.self_attn.k_proj.weight": "model.safetensors",
+        "model.layers.23.self_attn.o_proj.biases": "model.safetensors",
+        "model.layers.23.self_attn.o_proj.scales": "model.safetensors",
+        "model.layers.23.self_attn.o_proj.weight": "model.safetensors",
+        "model.layers.23.self_attn.q_proj.biases": "model.safetensors",
+        "model.layers.23.self_attn.q_proj.scales": "model.safetensors",
+        "model.layers.23.self_attn.q_proj.weight": "model.safetensors",
+        "model.layers.23.self_attn.v_proj.biases": "model.safetensors",
+        "model.layers.23.self_attn.v_proj.scales": "model.safetensors",
+        "model.layers.23.self_attn.v_proj.weight": "model.safetensors",
+        "model.layers.24.input_layernorm.weight": "model.safetensors",
+        "model.layers.24.mlp.down_proj.biases": "model.safetensors",
+        "model.layers.24.mlp.down_proj.scales": "model.safetensors",
+        "model.layers.24.mlp.down_proj.weight": "model.safetensors",
+        "model.layers.24.mlp.gate_proj.biases": "model.safetensors",
+        "model.layers.24.mlp.gate_proj.scales": "model.safetensors",
+        "model.layers.24.mlp.gate_proj.weight": "model.safetensors",
+        "model.layers.24.mlp.up_proj.biases": "model.safetensors",
+        "model.layers.24.mlp.up_proj.scales": "model.safetensors",
+        "model.layers.24.mlp.up_proj.weight": "model.safetensors",
+        "model.layers.24.post_attention_layernorm.weight": "model.safetensors",
+        "model.layers.24.self_attn.k_proj.biases": "model.safetensors",
+        "model.layers.24.self_attn.k_proj.scales": "model.safetensors",
+        "model.layers.24.self_attn.k_proj.weight": "model.safetensors",
+        "model.layers.24.self_attn.o_proj.biases": "model.safetensors",
+        "model.layers.24.self_attn.o_proj.scales": "model.safetensors",
+        "model.layers.24.self_attn.o_proj.weight": "model.safetensors",
+        "model.layers.24.self_attn.q_proj.biases": "model.safetensors",
+        "model.layers.24.self_attn.q_proj.scales": "model.safetensors",
+        "model.layers.24.self_attn.q_proj.weight": "model.safetensors",
+        "model.layers.24.self_attn.v_proj.biases": "model.safetensors",
+        "model.layers.24.self_attn.v_proj.scales": "model.safetensors",
+        "model.layers.24.self_attn.v_proj.weight": "model.safetensors",
+        "model.layers.25.input_layernorm.weight": "model.safetensors",
+        "model.layers.25.mlp.down_proj.biases": "model.safetensors",
+        "model.layers.25.mlp.down_proj.scales": "model.safetensors",
+        "model.layers.25.mlp.down_proj.weight": "model.safetensors",
+        "model.layers.25.mlp.gate_proj.biases": "model.safetensors",
+        "model.layers.25.mlp.gate_proj.scales": "model.safetensors",
+        "model.layers.25.mlp.gate_proj.weight": "model.safetensors",
+        "model.layers.25.mlp.up_proj.biases": "model.safetensors",
+        "model.layers.25.mlp.up_proj.scales": "model.safetensors",
+        "model.layers.25.mlp.up_proj.weight": "model.safetensors",
+        "model.layers.25.post_attention_layernorm.weight": "model.safetensors",
+        "model.layers.25.self_attn.k_proj.biases": "model.safetensors",
+        "model.layers.25.self_attn.k_proj.scales": "model.safetensors",
+        "model.layers.25.self_attn.k_proj.weight": "model.safetensors",
+        "model.layers.25.self_attn.o_proj.biases": "model.safetensors",
+        "model.layers.25.self_attn.o_proj.scales": "model.safetensors",
+        "model.layers.25.self_attn.o_proj.weight": "model.safetensors",
+        "model.layers.25.self_attn.q_proj.biases": "model.safetensors",
+        "model.layers.25.self_attn.q_proj.scales": "model.safetensors",
+        "model.layers.25.self_attn.q_proj.weight": "model.safetensors",
+        "model.layers.25.self_attn.v_proj.biases": "model.safetensors",
+        "model.layers.25.self_attn.v_proj.scales": "model.safetensors",
+        "model.layers.25.self_attn.v_proj.weight": "model.safetensors",
+        "model.layers.26.input_layernorm.weight": "model.safetensors",
+        "model.layers.26.mlp.down_proj.biases": "model.safetensors",
+        "model.layers.26.mlp.down_proj.scales": "model.safetensors",
+        "model.layers.26.mlp.down_proj.weight": "model.safetensors",
+        "model.layers.26.mlp.gate_proj.biases": "model.safetensors",
+        "model.layers.26.mlp.gate_proj.scales": "model.safetensors",
+        "model.layers.26.mlp.gate_proj.weight": "model.safetensors",
+        "model.layers.26.mlp.up_proj.biases": "model.safetensors",
+        "model.layers.26.mlp.up_proj.scales": "model.safetensors",
+        "model.layers.26.mlp.up_proj.weight": "model.safetensors",
+        "model.layers.26.post_attention_layernorm.weight": "model.safetensors",
+        "model.layers.26.self_attn.k_proj.biases": "model.safetensors",
+        "model.layers.26.self_attn.k_proj.scales": "model.safetensors",
+        "model.layers.26.self_attn.k_proj.weight": "model.safetensors",
+        "model.layers.26.self_attn.o_proj.biases": "model.safetensors",
+        "model.layers.26.self_attn.o_proj.scales": "model.safetensors",
+        "model.layers.26.self_attn.o_proj.weight": "model.safetensors",
+        "model.layers.26.self_attn.q_proj.biases": "model.safetensors",
+        "model.layers.26.self_attn.q_proj.scales": "model.safetensors",
+        "model.layers.26.self_attn.q_proj.weight": "model.safetensors",
+        "model.layers.26.self_attn.v_proj.biases": "model.safetensors",
+        "model.layers.26.self_attn.v_proj.scales": "model.safetensors",
+        "model.layers.26.self_attn.v_proj.weight": "model.safetensors",
+        "model.layers.27.input_layernorm.weight": "model.safetensors",
+        "model.layers.27.mlp.down_proj.biases": "model.safetensors",
+        "model.layers.27.mlp.down_proj.scales": "model.safetensors",
+        "model.layers.27.mlp.down_proj.weight": "model.safetensors",
+        "model.layers.27.mlp.gate_proj.biases": "model.safetensors",
+        "model.layers.27.mlp.gate_proj.scales": "model.safetensors",
+        "model.layers.27.mlp.gate_proj.weight": "model.safetensors",
+        "model.layers.27.mlp.up_proj.biases": "model.safetensors",
+        "model.layers.27.mlp.up_proj.scales": "model.safetensors",
+        "model.layers.27.mlp.up_proj.weight": "model.safetensors",
+        "model.layers.27.post_attention_layernorm.weight": "model.safetensors",
+        "model.layers.27.self_attn.k_proj.biases": "model.safetensors",
+        "model.layers.27.self_attn.k_proj.scales": "model.safetensors",
+        "model.layers.27.self_attn.k_proj.weight": "model.safetensors",
+        "model.layers.27.self_attn.o_proj.biases": "model.safetensors",
+        "model.layers.27.self_attn.o_proj.scales": "model.safetensors",
+        "model.layers.27.self_attn.o_proj.weight": "model.safetensors",
+        "model.layers.27.self_attn.q_proj.biases": "model.safetensors",
+        "model.layers.27.self_attn.q_proj.scales": "model.safetensors",
+        "model.layers.27.self_attn.q_proj.weight": "model.safetensors",
+        "model.layers.27.self_attn.v_proj.biases": "model.safetensors",
+        "model.layers.27.self_attn.v_proj.scales": "model.safetensors",
+        "model.layers.27.self_attn.v_proj.weight": "model.safetensors",
+        "model.layers.28.input_layernorm.weight": "model.safetensors",
+        "model.layers.28.mlp.down_proj.biases": "model.safetensors",
+        "model.layers.28.mlp.down_proj.scales": "model.safetensors",
+        "model.layers.28.mlp.down_proj.weight": "model.safetensors",
+        "model.layers.28.mlp.gate_proj.biases": "model.safetensors",
+        "model.layers.28.mlp.gate_proj.scales": "model.safetensors",
+        "model.layers.28.mlp.gate_proj.weight": "model.safetensors",
+        "model.layers.28.mlp.up_proj.biases": "model.safetensors",
+        "model.layers.28.mlp.up_proj.scales": "model.safetensors",
+        "model.layers.28.mlp.up_proj.weight": "model.safetensors",
+        "model.layers.28.post_attention_layernorm.weight": "model.safetensors",
+        "model.layers.28.self_attn.k_proj.biases": "model.safetensors",
+        "model.layers.28.self_attn.k_proj.scales": "model.safetensors",
+        "model.layers.28.self_attn.k_proj.weight": "model.safetensors",
+        "model.layers.28.self_attn.o_proj.biases": "model.safetensors",
+        "model.layers.28.self_attn.o_proj.scales": "model.safetensors",
+        "model.layers.28.self_attn.o_proj.weight": "model.safetensors",
+        "model.layers.28.self_attn.q_proj.biases": "model.safetensors",
+        "model.layers.28.self_attn.q_proj.scales": "model.safetensors",
+        "model.layers.28.self_attn.q_proj.weight": "model.safetensors",
+        "model.layers.28.self_attn.v_proj.biases": "model.safetensors",
+        "model.layers.28.self_attn.v_proj.scales": "model.safetensors",
+        "model.layers.28.self_attn.v_proj.weight": "model.safetensors",
+        "model.layers.29.input_layernorm.weight": "model.safetensors",
+        "model.layers.29.mlp.down_proj.biases": "model.safetensors",
+        "model.layers.29.mlp.down_proj.scales": "model.safetensors",
+        "model.layers.29.mlp.down_proj.weight": "model.safetensors",
+        "model.layers.29.mlp.gate_proj.biases": "model.safetensors",
+        "model.layers.29.mlp.gate_proj.scales": "model.safetensors",
+        "model.layers.29.mlp.gate_proj.weight": "model.safetensors",
+        "model.layers.29.mlp.up_proj.biases": "model.safetensors",
+        "model.layers.29.mlp.up_proj.scales": "model.safetensors",
+        "model.layers.29.mlp.up_proj.weight": "model.safetensors",
+        "model.layers.29.post_attention_layernorm.weight": "model.safetensors",
+        "model.layers.29.self_attn.k_proj.biases": "model.safetensors",
+        "model.layers.29.self_attn.k_proj.scales": "model.safetensors",
+        "model.layers.29.self_attn.k_proj.weight": "model.safetensors",
+        "model.layers.29.self_attn.o_proj.biases": "model.safetensors",
+        "model.layers.29.self_attn.o_proj.scales": "model.safetensors",
+        "model.layers.29.self_attn.o_proj.weight": "model.safetensors",
+        "model.layers.29.self_attn.q_proj.biases": "model.safetensors",
+        "model.layers.29.self_attn.q_proj.scales": "model.safetensors",
+        "model.layers.29.self_attn.q_proj.weight": "model.safetensors",
+        "model.layers.29.self_attn.v_proj.biases": "model.safetensors",
+        "model.layers.29.self_attn.v_proj.scales": "model.safetensors",
+        "model.layers.29.self_attn.v_proj.weight": "model.safetensors",
+        "model.layers.3.input_layernorm.weight": "model.safetensors",
+        "model.layers.3.mlp.down_proj.biases": "model.safetensors",
+        "model.layers.3.mlp.down_proj.scales": "model.safetensors",
+        "model.layers.3.mlp.down_proj.weight": "model.safetensors",
+        "model.layers.3.mlp.gate_proj.biases": "model.safetensors",
+        "model.layers.3.mlp.gate_proj.scales": "model.safetensors",
+        "model.layers.3.mlp.gate_proj.weight": "model.safetensors",
+        "model.layers.3.mlp.up_proj.biases": "model.safetensors",
+        "model.layers.3.mlp.up_proj.scales": "model.safetensors",
+        "model.layers.3.mlp.up_proj.weight": "model.safetensors",
+        "model.layers.3.post_attention_layernorm.weight": "model.safetensors",
+        "model.layers.3.self_attn.k_proj.biases": "model.safetensors",
+        "model.layers.3.self_attn.k_proj.scales": "model.safetensors",
+        "model.layers.3.self_attn.k_proj.weight": "model.safetensors",
+        "model.layers.3.self_attn.o_proj.biases": "model.safetensors",
+        "model.layers.3.self_attn.o_proj.scales": "model.safetensors",
+        "model.layers.3.self_attn.o_proj.weight": "model.safetensors",
+        "model.layers.3.self_attn.q_proj.biases": "model.safetensors",
+        "model.layers.3.self_attn.q_proj.scales": "model.safetensors",
+        "model.layers.3.self_attn.q_proj.weight": "model.safetensors",
+        "model.layers.3.self_attn.v_proj.biases": "model.safetensors",
+        "model.layers.3.self_attn.v_proj.scales": "model.safetensors",
+        "model.layers.3.self_attn.v_proj.weight": "model.safetensors",
+        "model.layers.30.input_layernorm.weight": "model.safetensors",
+        "model.layers.30.mlp.down_proj.biases": "model.safetensors",
+        "model.layers.30.mlp.down_proj.scales": "model.safetensors",
+        "model.layers.30.mlp.down_proj.weight": "model.safetensors",
+        "model.layers.30.mlp.gate_proj.biases": "model.safetensors",
+        "model.layers.30.mlp.gate_proj.scales": "model.safetensors",
+        "model.layers.30.mlp.gate_proj.weight": "model.safetensors",
+        "model.layers.30.mlp.up_proj.biases": "model.safetensors",
+        "model.layers.30.mlp.up_proj.scales": "model.safetensors",
+        "model.layers.30.mlp.up_proj.weight": "model.safetensors",
+        "model.layers.30.post_attention_layernorm.weight": "model.safetensors",
+        "model.layers.30.self_attn.k_proj.biases": "model.safetensors",
+        "model.layers.30.self_attn.k_proj.scales": "model.safetensors",
+        "model.layers.30.self_attn.k_proj.weight": "model.safetensors",
+        "model.layers.30.self_attn.o_proj.biases": "model.safetensors",
+        "model.layers.30.self_attn.o_proj.scales": "model.safetensors",
+        "model.layers.30.self_attn.o_proj.weight": "model.safetensors",
+        "model.layers.30.self_attn.q_proj.biases": "model.safetensors",
+        "model.layers.30.self_attn.q_proj.scales": "model.safetensors",
+        "model.layers.30.self_attn.q_proj.weight": "model.safetensors",
+        "model.layers.30.self_attn.v_proj.biases": "model.safetensors",
+        "model.layers.30.self_attn.v_proj.scales": "model.safetensors",
+        "model.layers.30.self_attn.v_proj.weight": "model.safetensors",
+        "model.layers.31.input_layernorm.weight": "model.safetensors",
+        "model.layers.31.mlp.down_proj.biases": "model.safetensors",
+        "model.layers.31.mlp.down_proj.scales": "model.safetensors",
+        "model.layers.31.mlp.down_proj.weight": "model.safetensors",
+        "model.layers.31.mlp.gate_proj.biases": "model.safetensors",
+        "model.layers.31.mlp.gate_proj.scales": "model.safetensors",
+        "model.layers.31.mlp.gate_proj.weight": "model.safetensors",
+        "model.layers.31.mlp.up_proj.biases": "model.safetensors",
+        "model.layers.31.mlp.up_proj.scales": "model.safetensors",
+        "model.layers.31.mlp.up_proj.weight": "model.safetensors",
+        "model.layers.31.post_attention_layernorm.weight": "model.safetensors",
+        "model.layers.31.self_attn.k_proj.biases": "model.safetensors",
+        "model.layers.31.self_attn.k_proj.scales": "model.safetensors",
+        "model.layers.31.self_attn.k_proj.weight": "model.safetensors",
+        "model.layers.31.self_attn.o_proj.biases": "model.safetensors",
+        "model.layers.31.self_attn.o_proj.scales": "model.safetensors",
+        "model.layers.31.self_attn.o_proj.weight": "model.safetensors",
+        "model.layers.31.self_attn.q_proj.biases": "model.safetensors",
+        "model.layers.31.self_attn.q_proj.scales": "model.safetensors",
+        "model.layers.31.self_attn.q_proj.weight": "model.safetensors",
+        "model.layers.31.self_attn.v_proj.biases": "model.safetensors",
+        "model.layers.31.self_attn.v_proj.scales": "model.safetensors",
+        "model.layers.31.self_attn.v_proj.weight": "model.safetensors",
+        "model.layers.32.input_layernorm.weight": "model.safetensors",
+        "model.layers.32.mlp.down_proj.biases": "model.safetensors",
+        "model.layers.32.mlp.down_proj.scales": "model.safetensors",
+        "model.layers.32.mlp.down_proj.weight": "model.safetensors",
+        "model.layers.32.mlp.gate_proj.biases": "model.safetensors",
+        "model.layers.32.mlp.gate_proj.scales": "model.safetensors",
+        "model.layers.32.mlp.gate_proj.weight": "model.safetensors",
+        "model.layers.32.mlp.up_proj.biases": "model.safetensors",
+        "model.layers.32.mlp.up_proj.scales": "model.safetensors",
+        "model.layers.32.mlp.up_proj.weight": "model.safetensors",
+        "model.layers.32.post_attention_layernorm.weight": "model.safetensors",
+        "model.layers.32.self_attn.k_proj.biases": "model.safetensors",
+        "model.layers.32.self_attn.k_proj.scales": "model.safetensors",
+        "model.layers.32.self_attn.k_proj.weight": "model.safetensors",
+        "model.layers.32.self_attn.o_proj.biases": "model.safetensors",
+        "model.layers.32.self_attn.o_proj.scales": "model.safetensors",
+        "model.layers.32.self_attn.o_proj.weight": "model.safetensors",
+        "model.layers.32.self_attn.q_proj.biases": "model.safetensors",
+        "model.layers.32.self_attn.q_proj.scales": "model.safetensors",
+        "model.layers.32.self_attn.q_proj.weight": "model.safetensors",
+        "model.layers.32.self_attn.v_proj.biases": "model.safetensors",
+        "model.layers.32.self_attn.v_proj.scales": "model.safetensors",
+        "model.layers.32.self_attn.v_proj.weight": "model.safetensors",
+        "model.layers.33.input_layernorm.weight": "model.safetensors",
+        "model.layers.33.mlp.down_proj.biases": "model.safetensors",
+        "model.layers.33.mlp.down_proj.scales": "model.safetensors",
+        "model.layers.33.mlp.down_proj.weight": "model.safetensors",
+        "model.layers.33.mlp.gate_proj.biases": "model.safetensors",
+        "model.layers.33.mlp.gate_proj.scales": "model.safetensors",
+        "model.layers.33.mlp.gate_proj.weight": "model.safetensors",
+        "model.layers.33.mlp.up_proj.biases": "model.safetensors",
+        "model.layers.33.mlp.up_proj.scales": "model.safetensors",
+        "model.layers.33.mlp.up_proj.weight": "model.safetensors",
+        "model.layers.33.post_attention_layernorm.weight": "model.safetensors",
+        "model.layers.33.self_attn.k_proj.biases": "model.safetensors",
+        "model.layers.33.self_attn.k_proj.scales": "model.safetensors",
+        "model.layers.33.self_attn.k_proj.weight": "model.safetensors",
+        "model.layers.33.self_attn.o_proj.biases": "model.safetensors",
+        "model.layers.33.self_attn.o_proj.scales": "model.safetensors",
+        "model.layers.33.self_attn.o_proj.weight": "model.safetensors",
+        "model.layers.33.self_attn.q_proj.biases": "model.safetensors",
+        "model.layers.33.self_attn.q_proj.scales": "model.safetensors",
+        "model.layers.33.self_attn.q_proj.weight": "model.safetensors",
+        "model.layers.33.self_attn.v_proj.biases": "model.safetensors",
+        "model.layers.33.self_attn.v_proj.scales": "model.safetensors",
+        "model.layers.33.self_attn.v_proj.weight": "model.safetensors",
+        "model.layers.34.input_layernorm.weight": "model.safetensors",
+        "model.layers.34.mlp.down_proj.biases": "model.safetensors",
+        "model.layers.34.mlp.down_proj.scales": "model.safetensors",
+        "model.layers.34.mlp.down_proj.weight": "model.safetensors",
+        "model.layers.34.mlp.gate_proj.biases": "model.safetensors",
+        "model.layers.34.mlp.gate_proj.scales": "model.safetensors",
+        "model.layers.34.mlp.gate_proj.weight": "model.safetensors",
+        "model.layers.34.mlp.up_proj.biases": "model.safetensors",
+        "model.layers.34.mlp.up_proj.scales": "model.safetensors",
+        "model.layers.34.mlp.up_proj.weight": "model.safetensors",
+        "model.layers.34.post_attention_layernorm.weight": "model.safetensors",
+        "model.layers.34.self_attn.k_proj.biases": "model.safetensors",
+        "model.layers.34.self_attn.k_proj.scales": "model.safetensors",
+        "model.layers.34.self_attn.k_proj.weight": "model.safetensors",
+        "model.layers.34.self_attn.o_proj.biases": "model.safetensors",
+        "model.layers.34.self_attn.o_proj.scales": "model.safetensors",
+        "model.layers.34.self_attn.o_proj.weight": "model.safetensors",
+        "model.layers.34.self_attn.q_proj.biases": "model.safetensors",
+        "model.layers.34.self_attn.q_proj.scales": "model.safetensors",
+        "model.layers.34.self_attn.q_proj.weight": "model.safetensors",
+        "model.layers.34.self_attn.v_proj.biases": "model.safetensors",
+        "model.layers.34.self_attn.v_proj.scales": "model.safetensors",
+        "model.layers.34.self_attn.v_proj.weight": "model.safetensors",
+        "model.layers.35.input_layernorm.weight": "model.safetensors",
+        "model.layers.35.mlp.down_proj.biases": "model.safetensors",
+        "model.layers.35.mlp.down_proj.scales": "model.safetensors",
+        "model.layers.35.mlp.down_proj.weight": "model.safetensors",
+        "model.layers.35.mlp.gate_proj.biases": "model.safetensors",
+        "model.layers.35.mlp.gate_proj.scales": "model.safetensors",
+        "model.layers.35.mlp.gate_proj.weight": "model.safetensors",
+        "model.layers.35.mlp.up_proj.biases": "model.safetensors",
+        "model.layers.35.mlp.up_proj.scales": "model.safetensors",
+        "model.layers.35.mlp.up_proj.weight": "model.safetensors",
+        "model.layers.35.post_attention_layernorm.weight": "model.safetensors",
+        "model.layers.35.self_attn.k_proj.biases": "model.safetensors",
+        "model.layers.35.self_attn.k_proj.scales": "model.safetensors",
+        "model.layers.35.self_attn.k_proj.weight": "model.safetensors",
+        "model.layers.35.self_attn.o_proj.biases": "model.safetensors",
+        "model.layers.35.self_attn.o_proj.scales": "model.safetensors",
+        "model.layers.35.self_attn.o_proj.weight": "model.safetensors",
+        "model.layers.35.self_attn.q_proj.biases": "model.safetensors",
+        "model.layers.35.self_attn.q_proj.scales": "model.safetensors",
+        "model.layers.35.self_attn.q_proj.weight": "model.safetensors",
+        "model.layers.35.self_attn.v_proj.biases": "model.safetensors",
+        "model.layers.35.self_attn.v_proj.scales": "model.safetensors",
+        "model.layers.35.self_attn.v_proj.weight": "model.safetensors",
+        "model.layers.4.input_layernorm.weight": "model.safetensors",
+        "model.layers.4.mlp.down_proj.biases": "model.safetensors",
+        "model.layers.4.mlp.down_proj.scales": "model.safetensors",
+        "model.layers.4.mlp.down_proj.weight": "model.safetensors",
+        "model.layers.4.mlp.gate_proj.biases": "model.safetensors",
+        "model.layers.4.mlp.gate_proj.scales": "model.safetensors",
+        "model.layers.4.mlp.gate_proj.weight": "model.safetensors",
+        "model.layers.4.mlp.up_proj.biases": "model.safetensors",
+        "model.layers.4.mlp.up_proj.scales": "model.safetensors",
+        "model.layers.4.mlp.up_proj.weight": "model.safetensors",
+        "model.layers.4.post_attention_layernorm.weight": "model.safetensors",
+        "model.layers.4.self_attn.k_proj.biases": "model.safetensors",
+        "model.layers.4.self_attn.k_proj.scales": "model.safetensors",
+        "model.layers.4.self_attn.k_proj.weight": "model.safetensors",
+        "model.layers.4.self_attn.o_proj.biases": "model.safetensors",
+        "model.layers.4.self_attn.o_proj.scales": "model.safetensors",
+        "model.layers.4.self_attn.o_proj.weight": "model.safetensors",
+        "model.layers.4.self_attn.q_proj.biases": "model.safetensors",
+        "model.layers.4.self_attn.q_proj.scales": "model.safetensors",
+        "model.layers.4.self_attn.q_proj.weight": "model.safetensors",
+        "model.layers.4.self_attn.v_proj.biases": "model.safetensors",
+        "model.layers.4.self_attn.v_proj.scales": "model.safetensors",
+        "model.layers.4.self_attn.v_proj.weight": "model.safetensors",
+        "model.layers.5.input_layernorm.weight": "model.safetensors",
+        "model.layers.5.mlp.down_proj.biases": "model.safetensors",
+        "model.layers.5.mlp.down_proj.scales": "model.safetensors",
+        "model.layers.5.mlp.down_proj.weight": "model.safetensors",
+        "model.layers.5.mlp.gate_proj.biases": "model.safetensors",
+        "model.layers.5.mlp.gate_proj.scales": "model.safetensors",
+        "model.layers.5.mlp.gate_proj.weight": "model.safetensors",
+        "model.layers.5.mlp.up_proj.biases": "model.safetensors",
+        "model.layers.5.mlp.up_proj.scales": "model.safetensors",
+        "model.layers.5.mlp.up_proj.weight": "model.safetensors",
+        "model.layers.5.post_attention_layernorm.weight": "model.safetensors",
+        "model.layers.5.self_attn.k_proj.biases": "model.safetensors",
+        "model.layers.5.self_attn.k_proj.scales": "model.safetensors",
+        "model.layers.5.self_attn.k_proj.weight": "model.safetensors",
+        "model.layers.5.self_attn.o_proj.biases": "model.safetensors",
+        "model.layers.5.self_attn.o_proj.scales": "model.safetensors",
+        "model.layers.5.self_attn.o_proj.weight": "model.safetensors",
+        "model.layers.5.self_attn.q_proj.biases": "model.safetensors",
+        "model.layers.5.self_attn.q_proj.scales": "model.safetensors",
+        "model.layers.5.self_attn.q_proj.weight": "model.safetensors",
+        "model.layers.5.self_attn.v_proj.biases": "model.safetensors",
+        "model.layers.5.self_attn.v_proj.scales": "model.safetensors",
+        "model.layers.5.self_attn.v_proj.weight": "model.safetensors",
+        "model.layers.6.input_layernorm.weight": "model.safetensors",
+        "model.layers.6.mlp.down_proj.biases": "model.safetensors",
+        "model.layers.6.mlp.down_proj.scales": "model.safetensors",
+        "model.layers.6.mlp.down_proj.weight": "model.safetensors",
+        "model.layers.6.mlp.gate_proj.biases": "model.safetensors",
+        "model.layers.6.mlp.gate_proj.scales": "model.safetensors",
+        "model.layers.6.mlp.gate_proj.weight": "model.safetensors",
+        "model.layers.6.mlp.up_proj.biases": "model.safetensors",
+        "model.layers.6.mlp.up_proj.scales": "model.safetensors",
+        "model.layers.6.mlp.up_proj.weight": "model.safetensors",
+        "model.layers.6.post_attention_layernorm.weight": "model.safetensors",
+        "model.layers.6.self_attn.k_proj.biases": "model.safetensors",
+        "model.layers.6.self_attn.k_proj.scales": "model.safetensors",
+        "model.layers.6.self_attn.k_proj.weight": "model.safetensors",
+        "model.layers.6.self_attn.o_proj.biases": "model.safetensors",
+        "model.layers.6.self_attn.o_proj.scales": "model.safetensors",
+        "model.layers.6.self_attn.o_proj.weight": "model.safetensors",
+        "model.layers.6.self_attn.q_proj.biases": "model.safetensors",
+        "model.layers.6.self_attn.q_proj.scales": "model.safetensors",
+        "model.layers.6.self_attn.q_proj.weight": "model.safetensors",
+        "model.layers.6.self_attn.v_proj.biases": "model.safetensors",
+        "model.layers.6.self_attn.v_proj.scales": "model.safetensors",
+        "model.layers.6.self_attn.v_proj.weight": "model.safetensors",
+        "model.layers.7.input_layernorm.weight": "model.safetensors",
+        "model.layers.7.mlp.down_proj.biases": "model.safetensors",
+        "model.layers.7.mlp.down_proj.scales": "model.safetensors",
+        "model.layers.7.mlp.down_proj.weight": "model.safetensors",
+        "model.layers.7.mlp.gate_proj.biases": "model.safetensors",
+        "model.layers.7.mlp.gate_proj.scales": "model.safetensors",
+        "model.layers.7.mlp.gate_proj.weight": "model.safetensors",
+        "model.layers.7.mlp.up_proj.biases": "model.safetensors",
+        "model.layers.7.mlp.up_proj.scales": "model.safetensors",
+        "model.layers.7.mlp.up_proj.weight": "model.safetensors",
+        "model.layers.7.post_attention_layernorm.weight": "model.safetensors",
+        "model.layers.7.self_attn.k_proj.biases": "model.safetensors",
+        "model.layers.7.self_attn.k_proj.scales": "model.safetensors",
+        "model.layers.7.self_attn.k_proj.weight": "model.safetensors",
+        "model.layers.7.self_attn.o_proj.biases": "model.safetensors",
+        "model.layers.7.self_attn.o_proj.scales": "model.safetensors",
+        "model.layers.7.self_attn.o_proj.weight": "model.safetensors",
+        "model.layers.7.self_attn.q_proj.biases": "model.safetensors",
+        "model.layers.7.self_attn.q_proj.scales": "model.safetensors",
+        "model.layers.7.self_attn.q_proj.weight": "model.safetensors",
+        "model.layers.7.self_attn.v_proj.biases": "model.safetensors",
+        "model.layers.7.self_attn.v_proj.scales": "model.safetensors",
+        "model.layers.7.self_attn.v_proj.weight": "model.safetensors",
+        "model.layers.8.input_layernorm.weight": "model.safetensors",
+        "model.layers.8.mlp.down_proj.biases": "model.safetensors",
+        "model.layers.8.mlp.down_proj.scales": "model.safetensors",
+        "model.layers.8.mlp.down_proj.weight": "model.safetensors",
+        "model.layers.8.mlp.gate_proj.biases": "model.safetensors",
+        "model.layers.8.mlp.gate_proj.scales": "model.safetensors",
+        "model.layers.8.mlp.gate_proj.weight": "model.safetensors",
+        "model.layers.8.mlp.up_proj.biases": "model.safetensors",
+        "model.layers.8.mlp.up_proj.scales": "model.safetensors",
+        "model.layers.8.mlp.up_proj.weight": "model.safetensors",
+        "model.layers.8.post_attention_layernorm.weight": "model.safetensors",
+        "model.layers.8.self_attn.k_proj.biases": "model.safetensors",
+        "model.layers.8.self_attn.k_proj.scales": "model.safetensors",
+        "model.layers.8.self_attn.k_proj.weight": "model.safetensors",
+        "model.layers.8.self_attn.o_proj.biases": "model.safetensors",
+        "model.layers.8.self_attn.o_proj.scales": "model.safetensors",
+        "model.layers.8.self_attn.o_proj.weight": "model.safetensors",
+        "model.layers.8.self_attn.q_proj.biases": "model.safetensors",
+        "model.layers.8.self_attn.q_proj.scales": "model.safetensors",
+        "model.layers.8.self_attn.q_proj.weight": "model.safetensors",
+        "model.layers.8.self_attn.v_proj.biases": "model.safetensors",
+        "model.layers.8.self_attn.v_proj.scales": "model.safetensors",
+        "model.layers.8.self_attn.v_proj.weight": "model.safetensors",
+        "model.layers.9.input_layernorm.weight": "model.safetensors",
+        "model.layers.9.mlp.down_proj.biases": "model.safetensors",
+        "model.layers.9.mlp.down_proj.scales": "model.safetensors",
+        "model.layers.9.mlp.down_proj.weight": "model.safetensors",
+        "model.layers.9.mlp.gate_proj.biases": "model.safetensors",
+        "model.layers.9.mlp.gate_proj.scales": "model.safetensors",
+        "model.layers.9.mlp.gate_proj.weight": "model.safetensors",
+        "model.layers.9.mlp.up_proj.biases": "model.safetensors",
+        "model.layers.9.mlp.up_proj.scales": "model.safetensors",
+        "model.layers.9.mlp.up_proj.weight": "model.safetensors",
+        "model.layers.9.post_attention_layernorm.weight": "model.safetensors",
+        "model.layers.9.self_attn.k_proj.biases": "model.safetensors",
+        "model.layers.9.self_attn.k_proj.scales": "model.safetensors",
+        "model.layers.9.self_attn.k_proj.weight": "model.safetensors",
+        "model.layers.9.self_attn.o_proj.biases": "model.safetensors",
+        "model.layers.9.self_attn.o_proj.scales": "model.safetensors",
+        "model.layers.9.self_attn.o_proj.weight": "model.safetensors",
+        "model.layers.9.self_attn.q_proj.biases": "model.safetensors",
+        "model.layers.9.self_attn.q_proj.scales": "model.safetensors",
+        "model.layers.9.self_attn.q_proj.weight": "model.safetensors",
+        "model.layers.9.self_attn.v_proj.biases": "model.safetensors",
+        "model.layers.9.self_attn.v_proj.scales": "model.safetensors",
+        "model.layers.9.self_attn.v_proj.weight": "model.safetensors",
+        "model.norm.weight": "model.safetensors"
+    }
+}

modeling_k2_horizon.py

@@ -0,0 +1,1116 @@
+#                🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨
+#           This file was automatically generated from src/transformers/models/qwen3_moe/modular_qwen3_moe.py.
+#               Do NOT edit this file manually as any edits will be overwritten by the generation of
+#             the file from the modular. If any change should be done, please apply the change to the
+#                          modular_qwen3_moe.py file directly. One of our CI enforces this.
+#                🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨
+# coding=utf-8
+# Copyright 2025 The Qwen team, Alibaba Group and the HuggingFace Inc. team. 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.
+
+from typing import Callable, Optional, Union
+
+import math
+import torch
+import torch.nn.functional as F
+from torch import nn
+
+from transformers.activations import ACT2FN
+from transformers.cache_utils import Cache, DynamicCache
+from transformers.generation import GenerationMixin
+from transformers.masking_utils import create_causal_mask, create_sliding_window_causal_mask
+from transformers.modeling_flash_attention_utils import FlashAttentionKwargs
+from transformers.modeling_layers import (
+    GenericForQuestionAnswering,
+    GenericForSequenceClassification,
+    GenericForTokenClassification,
+    GradientCheckpointingLayer,
+)
+from transformers.modeling_outputs import MoeCausalLMOutputWithPast, MoeModelOutputWithPast
+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.generic import maybe_autocast
+from transformers.utils.deprecation import deprecate_kwarg
+from transformers.utils.output_capturing import OutputRecorder
+
+from .configuration_k2_horizon import K2HorizonConfig
+
+
+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 split_to_interleaved(x):
+    # Split halves:  x0 x1 x2 x3 ... y0 y1 y2 y3 ...
+    # Interleaved:   x0 y0 x1 y1 x2 y2 x3 y3 ...
+    return x.reshape(*x.shape[:-1], 2, -1).transpose(-1, -2).reshape(*x.shape[:-1], -1)
+
+
+def interleaved_to_split(x):
+    # Interleaved:   x0 y0 x1 y1 x2 y2 x3 y3 ...
+    # Split halves:  x0 x1 x2 x3 ... y0 y1 y2 y3 ...
+    return x.reshape(*x.shape[:-1], -1, 2).transpose(-1, -2).reshape(*x.shape[:-1], -1)
+
+
+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
+
+
+def calc_router_weights(
+    router_logits: torch.Tensor,
+    router_bias: Optional[torch.Tensor],
+    score_func: str,
+    top_k: int,
+    scaling_factor: Optional[float],
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Return native-XLLM-compatible routing weights and selected experts.
+
+    XLLM applies router bias only to the values used for top-k selection. The
+    selected routes are still weighted by the original router probabilities,
+    then optionally normalized and scaled.
+    """
+    if score_func == "softmax":
+        routing_scores = F.softmax(router_logits, dim=-1, dtype=torch.float32)
+    elif score_func == "sigmoid":
+        routing_scores = torch.sigmoid(router_logits.to(torch.float32))
+    else:
+        raise ValueError(f"Unsupported router score function: {score_func}")
+
+    selection_scores = routing_scores
+    if router_bias is not None:
+        selection_scores = selection_scores + router_bias.to(selection_scores)
+
+    selected_indices = torch.topk(selection_scores, top_k, dim=-1).indices
+    routing_weights = torch.gather(routing_scores, dim=-1, index=selected_indices)
+    if top_k > 1:
+        routing_weights = routing_weights / routing_weights.sum(dim=-1, keepdim=True)
+    if scaling_factor is not None:
+        routing_weights = routing_weights * scaling_factor
+    return routing_weights, selected_indices
+
+
+def combine_routed_experts(
+    hidden_states: torch.Tensor,
+    routing_weights: torch.Tensor,
+    selected_indices: torch.Tensor,
+    experts: nn.ModuleList,
+    activation: Optional[Callable[[torch.Tensor], torch.Tensor]] = None,
+) -> torch.Tensor:
+    num_tokens, hidden_dim = hidden_states.shape
+    final_hidden_states = torch.zeros(
+        (num_tokens, experts[0].out_features),
+        dtype=hidden_states.dtype,
+        device=hidden_states.device,
+    )
+    expert_mask = torch.nn.functional.one_hot(
+        selected_indices, num_classes=len(experts)
+    ).permute(2, 1, 0)
+
+    for expert_idx in torch.nonzero(expert_mask.sum(dim=(-1, -2)), as_tuple=False).flatten():
+        topk_positions, token_positions = torch.where(expert_mask[int(expert_idx)])
+        expert_states = experts[int(expert_idx)](hidden_states[token_positions])
+        if activation is not None:
+            expert_states = activation(expert_states)
+        expert_states = expert_states * routing_weights[token_positions, topk_positions, None].to(expert_states.dtype)
+        final_hidden_states.index_add_(0, token_positions, expert_states.to(hidden_states.dtype))
+
+    return final_hidden_states
+
+
+class K2HorizonAttention(nn.Module):
+    """Multi-headed attention from 'Attention Is All You Need' paper"""
+
+    def __init__(self, config: K2HorizonConfig, 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.rope_head_dim = self.head_dim if config.rope_head_dim is None else config.rope_head_dim
+
+        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.o_proj = nn.Linear(
+            config.num_attention_heads * self.head_dim, config.hidden_size, bias=config.attention_bias
+        )
+
+        self.gate_func = config.attention_gate_func
+        if self.gate_func is not None:
+            self.gate_proj = nn.Linear(
+                config.hidden_size,
+                config.num_attention_heads * self.head_dim,
+                bias=False)
+
+        if config.query_key_norm:
+            self.q_norm = K2HorizonRMSNorm(
+                hidden_size=config.num_attention_heads * self.head_dim,
+                n_groups=config.num_attention_heads,
+                eps=config.rms_norm_eps)
+            self.k_norm = K2HorizonRMSNorm(
+                hidden_size=config.num_key_value_heads * self.head_dim,
+                n_groups=config.num_key_value_heads,
+                eps=config.rms_norm_eps)
+
+        self.sliding_window = getattr(config, "sliding_window", 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[Cache] = 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)
+
+        if self.config.query_key_norm:
+            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)
+        else:
+            query_states = self.q_proj(hidden_states).view(hidden_shape).transpose(1, 2)
+            key_states = self.k_proj(hidden_states).view(hidden_shape).transpose(1, 2)
+
+        value_states = self.v_proj(hidden_states).view(hidden_shape).transpose(1, 2)
+
+        cos, sin = position_embeddings
+        if self.rope_head_dim == self.head_dim:
+            query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin)
+        else:
+            query_states, query_states_ = torch.split(
+                split_to_interleaved(query_states),
+                split_size_or_sections=[self.rope_head_dim, self.head_dim - self.rope_head_dim],
+                dim=-1)
+
+            key_states, key_states_ = torch.split(
+                split_to_interleaved(key_states),
+                split_size_or_sections=[self.rope_head_dim, self.head_dim - self.rope_head_dim],
+                dim=-1)
+
+            query_states, key_states = apply_rotary_pos_emb(
+                interleaved_to_split(query_states),
+                interleaved_to_split(key_states),
+                cos,
+                sin)
+
+            query_states = interleaved_to_split(torch.cat(
+                [split_to_interleaved(query_states), query_states_], dim=-1))
+            key_states = interleaved_to_split(torch.cat(
+                [split_to_interleaved(key_states), key_states_], dim=-1))
+
+        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 = past_key_values.update(key_states, value_states, self.layer_idx, cache_kwargs)
+
+        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,
+        )
+
+        if self.gate_func is not None:
+            gate = self.gate_proj(hidden_states).view(
+                input_shape + (-1, self.head_dim))
+            if self.gate_func == 'silu':
+                gate = F.silu(gate)
+            else:
+                assert self.gate_func == 'softplus'
+                gate = F.softplus(gate, beta=math.log(2))
+
+            attn_output = attn_output * gate
+
+        attn_output = attn_output.reshape(*input_shape, -1).contiguous()
+        attn_output = self.o_proj(attn_output)
+        return attn_output, attn_weights
+
+
+def apply_rotary_pos_emb_xllm(q, k, freqs_cis):
+    if q.shape[-1] % 2 != 0 or k.shape[-1] % 2 != 0:
+        raise ValueError(f"RoPE dimensions must be even, got q={q.shape[-1]} and k={k.shape[-1]}")
+    q_ = torch.view_as_complex(q.float().reshape(*q.shape[:-1], -1, 2))
+    k_ = torch.view_as_complex(k.float().reshape(*k.shape[:-1], -1, 2))
+    if freqs_cis.ndim == 2:
+        freqs_cis = freqs_cis.unsqueeze(1)
+    elif freqs_cis.ndim == 3:
+        freqs_cis = freqs_cis.unsqueeze(2)
+    else:
+        raise ValueError(f"Unsupported freqs_cis shape: {tuple(freqs_cis.shape)}")
+    if freqs_cis.shape[-1] != q_.shape[-1] or freqs_cis.shape[-1] != k_.shape[-1]:
+        raise ValueError(
+            "RoPE frequency dimension mismatch: "
+            f"q_rope_dim={q.shape[-1]}, k_rope_dim={k.shape[-1]}, "
+            f"q_complex_dim={q_.shape[-1]}, k_complex_dim={k_.shape[-1]}, "
+            f"freqs_complex_dim={freqs_cis.shape[-1]}, freqs_shape={tuple(freqs_cis.shape)}"
+        )
+    q_embed = torch.view_as_real(q_ * freqs_cis).flatten(-2).to(q.dtype)
+    k_embed = torch.view_as_real(k_ * freqs_cis).flatten(-2).to(k.dtype)
+    return q_embed, k_embed
+
+
+class K2HorizonMoVAAttention(nn.Module):
+    """MoVA attention with routed value experts and optional post-attention gate."""
+
+    def __init__(self, config: K2HorizonConfig, 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.num_experts_per_tok = config.mova_num_experts_per_tok
+        self.router_score_func = config.router_score_func
+        self.router_scaling_factor = config.router_scaling_factor
+        self.gate_func = config.attention_gate_func
+
+        self.rope_head_dim = self.head_dim if config.rope_head_dim is None else config.rope_head_dim
+
+        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.o_proj = nn.Linear(
+            config.num_attention_heads * self.head_dim, config.hidden_size, bias=config.attention_bias
+        )
+        self.v_router = nn.Linear(
+            config.hidden_size,
+            config.mova_num_experts,
+            bias=config.moe_gate_bias)
+        self.v_experts = nn.ModuleList([
+            nn.Linear(
+                config.hidden_size,
+                config.num_key_value_heads * self.head_dim,
+                bias=False
+            ) for _ in range(config.mova_num_experts)
+        ])
+
+        if self.gate_func is not None:
+            self.gate_proj = nn.Linear(
+                config.hidden_size,
+                config.num_attention_heads * self.head_dim,
+                bias=False)
+
+        if config.query_key_norm:
+            self.q_norm = K2HorizonRMSNorm(
+                hidden_size=config.num_attention_heads * self.head_dim,
+                n_groups=config.num_attention_heads,
+                eps=config.rms_norm_eps,
+            )
+            self.k_norm = K2HorizonRMSNorm(
+                hidden_size=config.num_key_value_heads * self.head_dim,
+                n_groups=config.num_key_value_heads,
+                eps=config.rms_norm_eps,
+            )
+
+        self.sliding_window = getattr(config, "sliding_window", None)
+
+    @deprecate_kwarg("past_key_value", new_name="past_key_values", version="4.58")
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        position_embeddings: torch.Tensor,
+        attention_mask: Optional[torch.Tensor],
+        past_key_values: Optional[Cache] = 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)
+        flat_hidden_states = hidden_states.reshape(-1, hidden_states.shape[-1])
+
+        # Match native MOVAttention router semantics exactly: compute logits with
+        # the weight-only linear and apply router bias only to selection scores.
+        router_logits = F.linear(flat_hidden_states, self.v_router.weight)
+        routing_weights, selected_values = calc_router_weights(
+            router_logits=router_logits,
+            router_bias=self.v_router.bias,
+            score_func=self.router_score_func,
+            top_k=self.num_experts_per_tok,
+            scaling_factor=self.router_scaling_factor,
+        )
+
+        mixed_value_states = combine_routed_experts(
+            hidden_states=flat_hidden_states,
+            routing_weights=routing_weights,
+            selected_indices=selected_values,
+            experts=self.v_experts,
+            activation=F.silu)
+
+        value_states = mixed_value_states.view(hidden_shape).transpose(1, 2)
+
+        if self.config.query_key_norm:
+            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)
+        else:
+            query_states = self.q_proj(hidden_states).view(hidden_shape).transpose(1, 2)
+            key_states = self.k_proj(hidden_states).view(hidden_shape).transpose(1, 2)
+
+        cos, sin = position_embeddings
+        if self.rope_head_dim == self.head_dim:
+            query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin)
+        else:
+            query_states, query_states_ = torch.split(
+                split_to_interleaved(query_states),
+                split_size_or_sections=[self.rope_head_dim, self.head_dim - self.rope_head_dim],
+                dim=-1)
+
+            key_states, key_states_ = torch.split(
+                split_to_interleaved(key_states),
+                split_size_or_sections=[self.rope_head_dim, self.head_dim - self.rope_head_dim],
+                dim=-1)
+
+            query_states, key_states = apply_rotary_pos_emb(
+                interleaved_to_split(query_states),
+                interleaved_to_split(key_states),
+                cos,
+                sin)
+
+            query_states = interleaved_to_split(torch.cat(
+                [split_to_interleaved(query_states), query_states_], dim=-1))
+            key_states = interleaved_to_split(torch.cat(
+                [split_to_interleaved(key_states), key_states_], dim=-1))
+
+        if past_key_values is not None:
+            cache_kwargs = {"cache_position": cache_position}
+            key_states, value_states = past_key_values.update(
+                key_states, value_states, self.layer_idx, cache_kwargs
+            )
+
+        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,
+            **kwargs,
+        )
+
+        if self.gate_func is not None:
+            gate = self.gate_proj(hidden_states).view(input_shape + (-1, self.head_dim))
+            if self.gate_func == 'silu':
+                gate = F.silu(gate)
+            else:
+                assert self.gate_func == 'softplus'
+                gate = F.softplus(gate, beta=math.log(2))
+
+            attn_output = attn_output * gate
+
+        attn_output = attn_output.reshape(*input_shape, -1).contiguous()
+        attn_output = self.o_proj(attn_output)
+        return attn_output, attn_weights
+
+
+class K2HorizonMLP(nn.Module):
+    def __init__(self, config, intermediate_size=None):
+        super().__init__()
+        self.config = config
+        self.hidden_size = config.hidden_size
+        self.intermediate_size = intermediate_size if intermediate_size is not None else 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
+
+
+class K2HorizonSparseMoeBlock(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.num_experts = config.num_experts
+        self.top_k = config.num_experts_per_tok
+        self.norm_topk_prob = config.norm_topk_prob
+        self.num_shared_experts = config.num_shared_experts
+        self.router_score_func = config.router_score_func
+        self.router_scaling_factor = config.router_scaling_factor
+
+        # gating
+        self.gate = nn.Linear(config.hidden_size, config.num_experts, bias=config.moe_gate_bias)
+        self.experts = nn.ModuleList(
+            [K2HorizonMLP(config, intermediate_size=config.moe_intermediate_size) for _ in range(self.num_experts)]
+        )
+
+        if config.num_shared_experts > 0:
+            self.shared_experts = K2HorizonMLP(
+                config=config,
+                intermediate_size=config.moe_intermediate_size * config.num_shared_experts)
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        """ """
+        residuals = hidden_states
+
+        batch_size, sequence_length, hidden_dim = hidden_states.shape
+        hidden_states = hidden_states.view(-1, hidden_dim)
+        # router_logits: (batch * sequence_length, n_experts)
+        # router_logits = self.gate(hidden_states)
+        router_logits = F.linear(hidden_states, self.gate.weight)
+
+        if self.router_score_func == "softmax":
+            routing_weights = F.softmax(router_logits, dim=1, dtype=torch.float)
+        else:
+            assert self.router_score_func == "sigmoid"
+            routing_weights = F.sigmoid(router_logits.to(torch.float32))
+
+        routing_weights_for_choice = routing_weights
+        if self.gate.bias is not None:
+            routing_weights_for_choice = routing_weights + self.gate.bias.to(routing_weights.dtype)
+
+        _, selected_experts = torch.topk(routing_weights_for_choice, self.top_k, dim=-1)
+        routing_weights = torch.gather(routing_weights, dim=-1, index=selected_experts)
+
+        if self.norm_topk_prob:  # only diff with mixtral sparse moe block!
+            routing_weights /= routing_weights.sum(dim=-1, keepdim=True)
+        routing_weights = routing_weights * self.router_scaling_factor
+        # we cast back to the input dtype
+        routing_weights = routing_weights.to(hidden_states.dtype)
+
+        final_hidden_states = torch.zeros(
+            (batch_size * sequence_length, hidden_dim), dtype=hidden_states.dtype, device=hidden_states.device
+        )
+
+        # One hot encode the selected experts to create an expert mask
+        # this will be used to easily index which expert is going to be sollicitated
+        expert_mask = torch.nn.functional.one_hot(selected_experts, num_classes=self.num_experts).permute(2, 1, 0)
+
+        # Loop over all available experts in the model and perform the computation on each expert
+        expert_hit = torch.greater(expert_mask.sum(dim=(-1, -2)), 0).nonzero()
+        for expert_idx in expert_hit:
+            expert_layer = self.experts[expert_idx]
+            idx, top_x = torch.where(expert_mask[expert_idx].squeeze(0))
+
+            # Index the correct hidden states and compute the expert hidden state for
+            # the current expert. We need to make sure to multiply the output hidden
+            # states by `routing_weights` on the corresponding tokens (top-1 and top-2)
+            current_state = hidden_states[None, top_x].reshape(-1, hidden_dim)
+            current_hidden_states = expert_layer(current_state) * routing_weights[top_x, idx, None]
+
+            # However `index_add_` only support torch tensors for indexing so we'll use
+            # the `top_x` tensor here.
+            final_hidden_states.index_add_(0, top_x, current_hidden_states.to(hidden_states.dtype))
+        final_hidden_states = final_hidden_states.reshape(batch_size, sequence_length, hidden_dim)
+
+        if self.num_shared_experts > 0:
+            final_hidden_states = final_hidden_states + self.shared_experts(residuals)
+
+        return final_hidden_states, router_logits
+
+
+# @use_kernel_forward_from_hub("RMSNorm")
+class K2HorizonRMSNorm(nn.Module):
+    def __init__(self, hidden_size: int, n_groups: int, eps=1e-6):
+        """
+        K2HorizonRMSNorm is equivalent to T5LayerNorm
+        """
+        super().__init__()
+        self.n_groups = n_groups
+        self.hidden_size = hidden_size
+        assert hidden_size % n_groups == 0
+        self.weight = nn.Parameter(torch.ones(hidden_size))
+        self.variance_epsilon = eps
+
+    def forward(self, hidden_states):
+        input_dtype = hidden_states.dtype
+        hidden_states = hidden_states.to(torch.float32)
+
+        hidden_states = hidden_states.reshape(*hidden_states.shape[:-1], self.n_groups, -1)
+        variance = hidden_states.pow(2).mean(-1, keepdim=True)
+
+        hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)
+        hidden_states = hidden_states.reshape(*hidden_states.shape[:-2], -1)
+        hidden_states = self.weight * hidden_states
+
+        return hidden_states.to(input_dtype)
+
+    def extra_repr(self):
+        return f"{tuple(self.weight.shape)}, eps={self.variance_epsilon}"
+
+
+class K2HorizonDecoderLayer(GradientCheckpointingLayer):
+    def __init__(self, config: K2HorizonConfig, layer_idx: int):
+        super().__init__()
+        self.hidden_size = config.hidden_size
+
+        is_sparse_layer = (layer_idx not in config.mlp_only_layers) and (
+                config.num_experts > 0 and (layer_idx + 1) % config.decoder_sparse_step == 0)
+
+        if is_sparse_layer and config.mova_num_experts > 0:
+            self.self_attn = K2HorizonMoVAAttention(config=config, layer_idx=layer_idx)
+        else:
+            self.self_attn = K2HorizonAttention(config, layer_idx)
+
+        if is_sparse_layer:
+            self.mlp = K2HorizonSparseMoeBlock(config)
+        else:
+            self.mlp = K2HorizonMLP(config, intermediate_size=config.intermediate_size)
+
+        assert config.hidden_size % config.layernorm_num_groups == 0
+        self.input_layernorm = K2HorizonRMSNorm(
+            hidden_size=config.hidden_size,
+            n_groups=config.layernorm_num_groups,
+            eps=config.rms_norm_eps)
+        self.post_attention_layernorm = K2HorizonRMSNorm(
+            hidden_size=config.hidden_size,
+            n_groups=config.layernorm_num_groups,
+            eps=config.rms_norm_eps)
+
+    @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] = None,
+            position_ids: Optional[torch.LongTensor] = None,
+            past_key_values: Optional[Cache] = None,
+            cache_position: Optional[torch.LongTensor] = None,
+            **kwargs: Unpack[FlashAttentionKwargs],
+    ) -> torch.FloatTensor:
+        """
+        Args:
+            hidden_states (`torch.FloatTensor`): input to the layer of shape `(batch, seq_len, embed_dim)`
+            attention_mask (`torch.FloatTensor`, *optional*): attention mask of size
+                `(batch, sequence_length)` where padding elements are indicated by 0.
+            output_attentions (`bool`, *optional*):
+                Whether or not to return the attentions tensors of all attention layers. See `attentions` under
+                returned tensors for more detail.
+            output_router_logits (`bool`, *optional*):
+                Whether or not to return the logits of all the routers. They are useful for computing the router loss,
+                and should not be returned during inference.
+            use_cache (`bool`, *optional*):
+                If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding
+                (see `past_key_values`).
+            past_key_values (`Cache`, *optional*): cached past key and value projection states
+            cache_position (`torch.LongTensor` of shape `(sequence_length)`, *optional*):
+                Indices depicting the position of the input sequence tokens in the sequence.
+            position_embeddings (`tuple[torch.FloatTensor, torch.FloatTensor]`, *optional*):
+                Tuple containing the cosine and sine positional embeddings of shape `(batch_size, seq_len, head_dim)`,
+                with `head_dim` being the embedding dimension of each attention head.
+            kwargs (`dict`, *optional*):
+                Arbitrary kwargs to be ignored, used for FSDP and other methods that injects code
+                into the model
+        """
+        residual = hidden_states
+        hidden_states = self.input_layernorm(hidden_states)
+
+        # Self Attention
+        hidden_states, _ = self.self_attn(
+            hidden_states=hidden_states,
+            position_embeddings=position_embeddings,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            past_key_values=past_key_values,
+            cache_position=cache_position,
+            **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)
+        # For the MoE layers, we need to unpack
+        if isinstance(hidden_states, tuple):
+            hidden_states, _ = hidden_states
+        hidden_states = residual + hidden_states
+
+        return hidden_states
+
+
+class K2HorizonRotaryEmbedding(nn.Module):
+    inv_freq: torch.Tensor  # fix linting for `register_buffer`
+
+    def __init__(self, config: K2HorizonConfig, device=None):
+        super().__init__()
+        self.max_seq_len_cached = config.max_position_embeddings
+        self.original_max_seq_len = config.max_position_embeddings
+
+        self.config = config
+
+        self.rope_type = self.config.rope_parameters["rope_type"]
+        rope_init_fn: Callable = self.compute_default_rope_parameters
+        if self.rope_type != "default":
+            rope_init_fn = ROPE_INIT_FUNCTIONS[self.rope_type]
+        inv_freq, self.attention_scaling = rope_init_fn(self.config, device)
+
+        self.register_buffer("inv_freq", inv_freq, persistent=False)
+        self.register_buffer("original_inv_freq", inv_freq.clone(), persistent=False)
+
+    @staticmethod
+    def compute_default_rope_parameters(
+        config: K2HorizonConfig | None = None,
+        device: Optional["torch.device"] = None,
+        seq_len: int | None = None,
+    ) -> tuple["torch.Tensor", float]:
+        """
+        Computes the inverse frequencies according to the original RoPE implementation
+        Args:
+            config ([`~transformers.PreTrainedConfig`]):
+                The model configuration.
+            device (`torch.device`):
+                The device to use for initialization of the inverse frequencies.
+            seq_len (`int`, *optional*):
+                The current sequence length. Unused for this type of RoPE.
+        Returns:
+            Tuple of (`torch.Tensor`, `float`), containing the inverse frequencies for the RoPE embeddings and the
+            post-processing scaling factor applied to the computed cos/sin (unused in this type of RoPE).
+        """
+        base = config.rope_parameters["rope_theta"]
+        # dim = getattr(config, "head_dim", None) or config.hidden_size // config.num_attention_heads
+        dim = (
+            config.rope_head_dim
+            if config.rope_head_dim is not None
+            else getattr(config, "head_dim", None) or config.hidden_size // config.num_attention_heads
+        )
+
+        attention_factor = 1.0  # Unused in this type of RoPE
+
+        # Compute the inverse frequencies
+        inv_freq = 1.0 / (
+            base ** (torch.arange(0, dim, 2, dtype=torch.int64).to(device=device, dtype=torch.float) / dim)
+        )
+        return inv_freq, attention_factor
+
+    @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 maybe_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 K2HorizonPreTrainedModel(PreTrainedModel):
+    config: K2HorizonConfig
+    base_model_prefix = "model"
+    supports_gradient_checkpointing = True
+    _no_split_modules = ["K2HorizonDecoderLayer"]
+    _skip_keys_device_placement = ["past_key_values"]
+    _supports_flash_attn = True
+    _supports_sdpa = True
+    _supports_flex_attn = True
+    _can_compile_fullgraph = False  # MoE models don't work with torch.compile (`torch.where(condition)` not supported)
+    _supports_attention_backend = True
+    _can_record_outputs = {
+        "router_logits": OutputRecorder(K2HorizonSparseMoeBlock, index=1),
+        "hidden_states": K2HorizonDecoderLayer,
+        "attentions": K2HorizonAttention,
+    }
+
+
+@auto_docstring
+class K2HorizonModel(K2HorizonPreTrainedModel):
+    def __init__(self, config: K2HorizonConfig):
+        super().__init__(config)
+        # self.padding_idx = config.pad_token_id
+        self.padding_idx = getattr(config, "padding_idx", None)
+        self.vocab_size = config.vocab_size
+
+        self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx)
+        self.layers = nn.ModuleList(
+            [K2HorizonDecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)]
+        )
+
+        assert config.hidden_size % config.layernorm_num_groups == 0
+        self.norm = K2HorizonRMSNorm(
+            hidden_size=config.hidden_size,
+            n_groups=config.layernorm_num_groups,
+            eps=config.rms_norm_eps)
+
+        self.rotary_emb = K2HorizonRotaryEmbedding(config=config)
+        self.gradient_checkpointing = False
+
+        # Initialize weights and apply final processing
+        self.post_init()
+
+    @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[Cache] = None,
+            inputs_embeds: Optional[torch.FloatTensor] = None,
+            use_cache: Optional[bool] = None,
+            cache_position: Optional[torch.LongTensor] = None,
+            **kwargs: Unpack[TransformersKwargs],
+    ) -> MoeModelOutputWithPast:
+        if (input_ids is None) ^ (inputs_embeds is not None):
+            raise ValueError("You must specify exactly one of input_ids or inputs_embeds")
+
+        if use_cache and past_key_values is None:
+            past_key_values = DynamicCache(config=self.config)
+
+        if inputs_embeds is None:
+            inputs_embeds = self.embed_tokens(input_ids)
+
+        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)
+
+        mask_function = create_causal_mask if self.config.sliding_window is None else create_sliding_window_causal_mask
+        causal_mask = mask_function(
+            config=self.config,
+            inputs_embeds=inputs_embeds,
+            attention_mask=attention_mask,
+            past_key_values=past_key_values,
+            position_ids=position_ids,
+        )
+
+        hidden_states = inputs_embeds
+
+        # create position embeddings to be shared across the decoder layers
+        position_embeddings = self.rotary_emb(hidden_states, position_ids)
+
+        for layer_idx, decoder_layer in enumerate(self.layers[: self.config.num_hidden_layers]):
+            hidden_states = decoder_layer(
+                hidden_states,
+                position_embeddings=position_embeddings,
+                attention_mask=causal_mask,
+                position_ids=position_ids,
+                past_key_values=past_key_values,
+                use_cache=use_cache,
+                cache_position=cache_position,
+                **kwargs,
+            )
+
+        hidden_states = self.norm(hidden_states)
+
+        return MoeModelOutputWithPast(  # only diff with Mistral is the output type, we need MoE
+            last_hidden_state=hidden_states,
+            past_key_values=past_key_values,
+        )
+
+
+def load_balancing_loss_func(
+        gate_logits: Union[torch.Tensor, tuple[torch.Tensor], None],
+        num_experts: Optional[int] = None,
+        top_k=2,
+        attention_mask: Optional[torch.Tensor] = None,
+) -> Union[torch.Tensor, int]:
+    r"""
+    Computes auxiliary load balancing loss as in Switch Transformer - implemented in Pytorch.
+
+    See Switch Transformer (https://huggingface.co/papers/2101.03961) for more details. This function implements the loss
+    function presented in equations (4) - (6) of the paper. It aims at penalizing cases where the routing between
+    experts is too unbalanced.
+
+    Args:
+        gate_logits:
+            Logits from the `gate`, should be a tuple of model.config.num_hidden_layers tensors of
+            shape [batch_size X sequence_length, num_experts].
+        num_experts:
+            Number of experts
+        top_k:
+            The number of experts to route per-token, can be also interpreted as the `top-k` routing
+            parameter.
+        attention_mask (`torch.Tensor`, *optional*):
+            The attention_mask used in forward function
+            shape [batch_size X sequence_length] if not None.
+
+    Returns:
+        The auxiliary loss.
+    """
+    if gate_logits is None or not isinstance(gate_logits, tuple):
+        return 0
+
+    if isinstance(gate_logits, tuple):
+        compute_device = gate_logits[0].device
+        concatenated_gate_logits = torch.cat([layer_gate.to(compute_device) for layer_gate in gate_logits], dim=0)
+
+    routing_weights = torch.nn.functional.softmax(concatenated_gate_logits, dim=-1)
+
+    _, selected_experts = torch.topk(routing_weights, top_k, dim=-1)
+
+    expert_mask = torch.nn.functional.one_hot(selected_experts, num_experts)
+
+    if attention_mask is None:
+        # Compute the percentage of tokens routed to each experts
+        tokens_per_expert = torch.mean(expert_mask.float(), dim=0)
+
+        # Compute the average probability of routing to these experts
+        router_prob_per_expert = torch.mean(routing_weights, dim=0)
+    else:
+        batch_size, sequence_length = attention_mask.shape
+        num_hidden_layers = concatenated_gate_logits.shape[0] // (batch_size * sequence_length)
+
+        # Compute the mask that masks all padding tokens as 0 with the same shape of expert_mask
+        expert_attention_mask = (
+            attention_mask[None, :, :, None, None]
+            .expand((num_hidden_layers, batch_size, sequence_length, top_k, num_experts))
+            .reshape(-1, top_k, num_experts)
+            .to(compute_device)
+        )
+
+        # Compute the percentage of tokens routed to each experts
+        tokens_per_expert = torch.sum(expert_mask.float() * expert_attention_mask, dim=0) / torch.sum(
+            expert_attention_mask, dim=0
+        )
+
+        # Compute the mask that masks all padding tokens as 0 with the same shape of tokens_per_expert
+        router_per_expert_attention_mask = (
+            attention_mask[None, :, :, None]
+            .expand((num_hidden_layers, batch_size, sequence_length, num_experts))
+            .reshape(-1, num_experts)
+            .to(compute_device)
+        )
+
+        # Compute the average probability of routing to these experts
+        router_prob_per_expert = torch.sum(routing_weights * router_per_expert_attention_mask, dim=0) / torch.sum(
+            router_per_expert_attention_mask, dim=0
+        )
+
+    overall_loss = torch.sum(tokens_per_expert * router_prob_per_expert.unsqueeze(0))
+    return overall_loss * num_experts
+
+
+@auto_docstring
+class K2HorizonForCausalLM(K2HorizonPreTrainedModel, 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 = K2HorizonModel(config)
+        self.vocab_size = config.vocab_size
+        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
+        self.router_aux_loss_coef = config.router_aux_loss_coef
+        self.num_experts = config.num_experts
+        self.num_experts_per_tok = config.num_experts_per_tok
+
+        # Initialize weights and apply final processing
+        self.post_init()
+
+    @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[Cache] = None,
+            inputs_embeds: Optional[torch.FloatTensor] = None,
+            labels: Optional[torch.LongTensor] = None,
+            use_cache: Optional[bool] = None,
+            output_router_logits: Optional[bool] = None,
+            cache_position: Optional[torch.LongTensor] = None,
+            logits_to_keep: Union[int, torch.Tensor] = 0,
+            **kwargs: Unpack[TransformersKwargs],
+    ) -> MoeCausalLMOutputWithPast:
+        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, K2HorizonForCausalLM
+
+        >>> model = K2HorizonForCausalLM.from_pretrained("Qwen/Qwen3-MoE-15B-A2B")
+        >>> tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen3-MoE-15B-A2B")
+
+        >>> 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."
+        ```"""
+
+        output_router_logits = (
+            output_router_logits if output_router_logits is not None else self.config.output_router_logits
+        )
+
+        # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn)
+        outputs: MoeModelOutputWithPast = 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,
+            output_router_logits=output_router_logits,
+            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, labels, self.vocab_size, **kwargs)
+
+        aux_loss = None
+        if output_router_logits:
+            aux_loss = load_balancing_loss_func(
+                outputs.router_logits,
+                self.num_experts,
+                self.num_experts_per_tok,
+                attention_mask,
+            )
+            if labels is not None:
+                loss += self.router_aux_loss_coef * aux_loss.to(loss.device)  # make sure to reside in the same device
+
+        return MoeCausalLMOutputWithPast(
+            loss=loss,
+            aux_loss=aux_loss,
+            logits=logits,
+            past_key_values=outputs.past_key_values,
+            hidden_states=outputs.hidden_states,
+            attentions=outputs.attentions,
+            router_logits=outputs.router_logits,
+        )
+
+
+class K2HorizonForSequenceClassification(GenericForSequenceClassification, K2HorizonPreTrainedModel):
+    pass
+
+
+class K2HorizonForTokenClassification(GenericForTokenClassification, K2HorizonPreTrainedModel):
+    pass
+
+
+class K2HorizonForQuestionAnswering(GenericForQuestionAnswering, K2HorizonPreTrainedModel):
+    base_model_prefix = "transformer"  # For BC, where `transformer` was used instead of `model`
+
+
+__all__ = [
+    "K2HorizonForCausalLM",
+    "K2HorizonForQuestionAnswering",
+    "K2HorizonModel",
+    "K2HorizonPreTrainedModel",
+    "K2HorizonForSequenceClassification",
+    "K2HorizonForTokenClassification",
+]

tokenizer.json

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

tokenizer_config.json

@@ -0,0 +1,10 @@
+{
+  "backend": "tokenizers",
+  "bos_token": "<|ifm|begin_of_text|>",
+  "clean_up_tokenization_spaces": false,
+  "eos_token": "<|ifm|endoftext|>",
+  "is_local": true,
+  "local_files_only": false,
+  "model_max_length": 1000000000000000019884624838656,
+  "tokenizer_class": "TokenizersBackend"
+}