From b3d123e9aede6e29dc95aeb2656f817299bf205d Mon Sep 17 00:00:00 2001 From: Asher Feldman <59994+asher@users.noreply.github.com> Date: Sat, 17 Oct 2026 18:05:53 -0700 Subject: [PATCH 01/10] fix(chat): parse templates with the generation tag, and warn on --thinking off with no switch --- THIRD_PARTY_NOTICES.md | 2 +- gmlx/load/loader.py | 4 +- gmlx/load/template_parse.py | 37 + gmlx/tui/reasoning.py | 26 +- tests/fixtures/k2_horizon_template.jinja | 994 +++++++++++++++++++++++ tests/load/test_template_parse.py | 55 ++ tests/tui/test_reasoning.py | 43 + 7 files changed, 1153 insertions(+), 8 deletions(-) create mode 100644 gmlx/load/template_parse.py create mode 100644 tests/fixtures/k2_horizon_template.jinja create mode 100644 tests/load/test_template_parse.py diff --git a/THIRD_PARTY_NOTICES.md b/THIRD_PARTY_NOTICES.md index 43ebb5e5b..f5b68f397 100644 --- a/THIRD_PARTY_NOTICES.md +++ b/THIRD_PARTY_NOTICES.md @@ -10,7 +10,7 @@ in [`licenses/`](licenses/). | [mlx-lm](https://github.com/ml-explore/mlx-lm) | MIT | Release 0.31.3 vendored whole at `gmlx/_vendor/mlx_lm`, with its LICENSE, its module references rewritten and each gmlx edit marked `# gmlx:`. Model/cache classes from unmerged PRs: `gmlx/models/deepseek_v4/model.py` + `gmlx/models/deepseek_v4/hyper_connection.py` + `gmlx/models/deepseek_v4/cache.py` (PR #1192), `gmlx/models/minimax_m3.py` (PR #1401), `gmlx/models/hy_v3/model.py` (PR #1485); `gmlx/models/kimi_k3.py` builds on the released `kimi_linear` module. The KV cache classes in `gmlx/cache/kv_cache.py` and the mask and SDPA helpers in `gmlx/cache/attention.py` are mlx-lm's `models/cache.py` and `models/base.py` code, by way of mlx-vlm 0.7.3. The FunctionGemma tool parser in `gmlx/models/function_gemma.py` comes from release 0.31.3, by way of mlx-vlm 0.6.15. The Gemma 4 text weight sanitize in `gmlx/models/gemma4/lm_sanitize.py` comes from release 0.31.3. | | [mlx-vlm](https://github.com/Blaizzy/mlx-vlm) | MIT; Apache-2.0 for eight files | Release 0.7.3 vendored whole at `gmlx/_vendor/mlx_vlm`, with its LICENSE, its module references rewritten and each gmlx edit marked `# gmlx:`. Eight of its files keep an Apache-2.0 header, under `models/llada_image`, `models/mage_vl`, `models/nemotron_voicechat` and `models/sam3d_objects`. Also the qwen3.5 and qwen3.5 MoE text models of release 0.6.15, vendored in `gmlx/models/qwen35/` (`language.py`, `gated_delta.py`, `moe_language.py`); the Gemma 4 text model and unified container of release 0.7.3, vendored in `gmlx/models/gemma4/` (`language.py`, `unified.py`); the KV cache classes of release 0.7.3 (`models/cache.py`), vendored whole in `gmlx/cache/kv_cache.py`, and its mask and SDPA helpers (`models/base.py`) in `gmlx/cache/attention.py`; the chunked-eval `store_kv_blocks` override in `gmlx/cache/apc_manager.py`, the thinking-budget logits processor in `gmlx/gen/thinking_budget.py`, and `entropy_transfer_mask` in `gmlx/systemone/denoise.py`. | | [omlx](https://github.com/jundot/omlx) | Apache-2.0 | DeepSeek-V4-Flash MTP port (`gmlx/models/deepseek_v4/mtp.py`), the MTP undo log of the rotating cache classes in `gmlx/cache/kv_cache.py`, and the native-kernel dispatch pattern in `gmlx/models/deepseek_v4/model.py`. | -| [llama.cpp / ggml](https://github.com/ggml-org/llama.cpp) | MIT | Pre-tokenizer split regex patterns in `gmlx/load/tokenizer.py`; GGUF metadata and quantization-type conventions; the Muse Glimmer vision tower in `gmlx/models/muse_glimmer/vlm_model.py` (from `clip_graph_muse_glimmer` and `tools/mtmd/clip.cpp`) and its projector constants in `gmlx/load/vlm.py`; the Muse Spark audio encoder in `gmlx/models/muse_spark/audio.py` (from `clip_graph_muse_spark_audio`); the Kimi-K3 chat-template fixture `tests/fixtures/kimi_k3_template.jinja` (PR #26185). | +| [llama.cpp / ggml](https://github.com/ggml-org/llama.cpp) | MIT | Pre-tokenizer split regex patterns in `gmlx/load/tokenizer.py`; GGUF metadata and quantization-type conventions; the Muse Glimmer vision tower in `gmlx/models/muse_glimmer/vlm_model.py` (from `clip_graph_muse_glimmer` and `tools/mtmd/clip.cpp`) and its projector constants in `gmlx/load/vlm.py`; the Muse Spark audio encoder in `gmlx/models/muse_spark/audio.py` (from `clip_graph_muse_spark_audio`); the Kimi-K3 chat-template fixture `tests/fixtures/kimi_k3_template.jinja` (PR #26185); the K2 Horizon chat-template fixture `tests/fixtures/k2_horizon_template.jinja` (`models/templates/IFM-K2-Horizon.jinja`). | | [ds4.c (dwarfstar)](licenses/ds4.c-LICENSE) | MIT | QAT round-trip and sink-attention reference semantics reimplemented in `gmlx/models/deepseek_v4/model.py` (cited per function in comments); the DSpark drafter runtime reimplemented in `gmlx/models/deepseek_v4/dspark.py`; the `deepseek4-dspark` sidecar layout (`gguf-tools/deepseek4-quantize.c` and the published DSpark-support sidecar) emitted by `scripts/convert_dspark_sidecar.py`. | | [DeepSeek-V4-Flash-0731](https://huggingface.co/deepseek-ai/DeepSeek-V4-Flash-0731) | MIT | Reference DSparkBlock forward semantics (`inference/model.py`) reimplemented in `gmlx/models/deepseek_v4/dspark.py` (license text in [`licenses/deepseek-v4-LICENSE`](licenses/deepseek-v4-LICENSE)). | | [beellama.cpp](https://github.com/Anbeeld/beellama.cpp) | MIT | The KVarN KV-cache record format and cache semantics (128-token rotated records with three axis vectors, fp16 sink and precision tail, `kvarnN` widths, `--kv-tail-tokens`) reimplemented in `gmlx/cache/kvarn_cache.py` and `gmlx/cache/kvarn_sdpa.py`. The method is Huawei's KVarN ([arXiv:2606.03458](https://arxiv.org/abs/2606.03458), reference at [huawei-csl/KVarN](https://github.com/huawei-csl/KVarN), Apache-2.0; no code from it is used). The kernels live in mlx-kquant, whose notices cover its own beellama.cpp port. | diff --git a/gmlx/load/loader.py b/gmlx/load/loader.py index fe32c4915..500be70f0 100644 --- a/gmlx/load/loader.py +++ b/gmlx/load/loader.py @@ -1997,8 +1997,10 @@ def _resolve_chat_template(chat_template: str | None) -> str | None: import jinja2 except ImportError: # validated later by apply_chat_template instead return chat_template + from .template_parse import parse_chat_template + try: - jinja2.Environment().parse(chat_template) + parse_chat_template(chat_template) except jinja2.TemplateSyntaxError as e: raise ValueError( f"chat template is not valid Jinja (line {e.lineno}: {e.message})" diff --git a/gmlx/load/template_parse.py b/gmlx/load/template_parse.py new file mode 100644 index 000000000..baabe0934 --- /dev/null +++ b/gmlx/load/template_parse.py @@ -0,0 +1,37 @@ +# SPDX-License-Identifier: MIT +"""Parse a chat template with the tags its renderer accepts. + +transformers renders chat templates with ``loopcontrols`` and its own +``{% generation %}`` block (the assistant-mask tracker). A plain +``jinja2.Environment`` rejects both, so a static parse of such a template +fails although the template renders. This parser accepts the same tags. +The generation block parses as its body, so the variables inside it stay +visible to ``jinja2.meta``. +""" + +from __future__ import annotations + +import functools + + +@functools.cache +def _environment(): + import jinja2 + from jinja2.ext import Extension + + class _Generation(Extension): + tags = {"generation"} + + def parse(self, parser): + next(parser.stream) + return parser.parse_statements( + ("name:endgeneration",), drop_needle=True) + + return jinja2.Environment( + extensions=[_Generation, "jinja2.ext.loopcontrols"]) + + +def parse_chat_template(source: str): + """The Jinja AST of ``source``. Raises ``jinja2.TemplateSyntaxError`` + on a template the renderer would also reject.""" + return _environment().parse(source) diff --git a/gmlx/tui/reasoning.py b/gmlx/tui/reasoning.py index 681720ed1..7323d0a8d 100644 --- a/gmlx/tui/reasoning.py +++ b/gmlx/tui/reasoning.py @@ -636,12 +636,16 @@ def get_chat_template(processor, messages, add_generation_prompt, _BARE_THINKING_RE = re.compile(r"(? str: - """``template`` with its literal prose dropped. A name that only appears - in prose is text the model reads, not a variable the template reads.""" - return "\n".join(_JINJA_TAG_RE.findall(template)) + """``template`` with its literal prose and string literals dropped. A + name that only appears in prose, or in a string such as a + ``raise_exception`` message, is text, not a variable the template + reads.""" + return "\n".join(_JINJA_STRING_RE.sub("''", tag) + for tag in _JINJA_TAG_RE.findall(template)) def _inputs(template: str) -> frozenset[str] | None: @@ -650,11 +654,12 @@ def _inputs(template: str) -> frozenset[str] | None: its own ``thinking_mode`` from ``enable_thinking``) is not one of them, so passing it would be a no-op.""" try: - import jinja2 from jinja2 import meta + from gmlx.load.template_parse import parse_chat_template + return frozenset(meta.find_undeclared_variables( - jinja2.Environment().parse(template))) + parse_chat_template(template))) except Exception: return None @@ -667,6 +672,15 @@ def _reads(template: str, name: str) -> bool: return name in _jinja_code(template) +def _reads_bare_thinking(template: str) -> bool: + """True when ``template`` reads a bare ``thinking`` variable (Kimi + K2.x), not only ``enable_thinking`` or ``preserve_thinking``.""" + names = _inputs(template) + if names is not None: + return "thinking" in names + return bool(_BARE_THINKING_RE.search(_jinja_code(template))) + + # Templates that grade reasoning depth under a name of their own. The control # is spelled reasoning_effort throughout gmlx, so map it onto the template's # spelling rather than emitting a variable the template ignores (Muse Glimmer's @@ -734,7 +748,7 @@ def map_thinking_controls(base: dict, thinking=None, reasoning_effort=None, "chat template has no such variable - ignored") elif _reads(template, "enable_thinking") or not template: out["enable_thinking"] = mode == "on" - elif _BARE_THINKING_RE.search(_jinja_code(template)): + elif _reads_bare_thinking(template): # Kimi K2.x: a bare `thinking` variable gates the pre-filled # think-open (`thinking is defined and thinking is false` renders # the closed pair). The lookarounds skip preserve_thinking etc. diff --git a/tests/fixtures/k2_horizon_template.jinja b/tests/fixtures/k2_horizon_template.jinja new file mode 100644 index 000000000..41dde1f4b --- /dev/null +++ b/tests/fixtures/k2_horizon_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) -%} +{{- "" }} +{%- for tool in tools_list %} +{{- "\n" }} +{{- tool | tojson }} +{%- endfor %} +{{- "\n" }} +{%- 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) }}{{- "" }} +{%- else -%} +{{- "/>" }} +{%- endif -%} +{%- else -%} +{{- "<" + tag + ">" }}{{ render_xml_value(spec) }}{{- "" }} +{%- 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 -%} +{{- "" }}{{ render_xml_schema_children(spec, true, true) }}{{- "" }} +{%- else -%} +{{- "/>" }} +{%- endif -%} +{%- else -%} +{{- "" }}{{ render_xml_value(spec) }}{{- "" }} +{%- endif -%} +{%- endmacro -%} + +{%- macro render_xml_schema_children(spec, include_properties, include_description) -%} +{%- if include_description and spec.description is defined -%}{{- "" }}{{ spec.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 -%} +{{- "" }} +{%- for variant in spec.oneOf -%}{{ render_xml_schema_node("variant", variant, true) }}{%- endfor -%} +{{- "" }} +{%- endif -%} +{%- if spec.anyOf -%} +{{- "" }} +{%- for variant in spec.anyOf -%}{{ render_xml_schema_node("variant", variant, true) }}{%- endfor -%} +{{- "" }} +{%- endif -%} +{%- if spec.additionalProperties is mapping -%}{{ render_xml_schema_node("additionalProperties", spec.additionalProperties, true) }}{%- endif -%} +{%- if spec.patternProperties is mapping -%} +{{- "" }} +{%- for pattern, pattern_spec in spec.patternProperties | items -%}{{ render_xml_pattern_property(pattern, pattern_spec) }}{%- endfor -%} +{{- "" }} +{%- elif spec.patternProperties is defined -%}{{ render_xml_value(spec.patternProperties) }}{%- endif -%} +{%- if spec.returns is mapping -%}{{ render_xml_schema_node("returns", spec.returns, true) }}{%- elif spec.returns is defined -%}{{ render_xml_value(spec.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 -%} +{{- "" }} +{%- if spec.description -%}{{ spec.description }}{%- endif -%} +{{- render_xml_schema_children(spec, true, false) }} +{{- "" }} +{%- else -%} +{{- "/>" }} +{%- endif -%} +{%- endmacro -%} + +{%- macro render_tools_xml(tools_list) -%} +{{- "" }} +{%- 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" }} +{%- if fn.description -%} +{{- "" }}{{ fn.description }}{{- "" }} +{%- endif -%} +{{- "" }} +{%- 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 -%} +{{- "" }} +{%- 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 -%}{{ render_xml_value(fn_ret) }}{%- endif -%} +{{- "" }} +{%- endfor -%} +{{- "\n" }} +{%- 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) -%} +{{- "" }} +{%- 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" }} +{%- 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 block. For each call, emit one JSON object with the function name and arguments on the same line inside tags:\n\n\n{\"name\": , \"arguments\": }\n" }} +{%- elif fmt == 'xml' -%} +{{- "Wrap all tool calls in a single block. For each call, write the function name at the start of , followed by paired and tags for each argument:\n\n\n$FUNCTION_NAME\n$PARAMETER_NAME\n$PARAMETER_VALUE\n...\n\n\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 block. For each call, write the function name at the start of , followed by , , and tags for each argument:\n\n\n$FUNCTION_NAME\n$PARAMETER_NAME\n$ARGUMENT_TYPE\n$PARAMETER_VALUE\n...\n\n\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) -%} +{{- "" }} +{%- 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{\"name\": \"" + tool_call.name + "\", \"arguments\": " }}{{ tool_call.arguments | tojson }}{{- "}" }} + {%- elif fmt == 'xml' or fmt == 'xml_typed' -%} +{{- "\n" + tool_call.name + "\n" }} + {%- for key, value in tool_call.arguments | items -%} +{{- "" + key + "\n" }} +{%- if fmt == 'xml_typed' -%} +{{- "" + render_arg_type(tools_list, tool_call.name, key, value) + "\n" }} +{%- endif -%} +{{- "" }}{{ render_argument_value(value) }}{{- "\n" }} + {%- endfor -%} +{{- "" }} + {%- else -%} + {{- raise_exception("Unsupported tool_call_format: '" + fmt + "'. Supported formats: json, xml, xml_typed.") -}} + {%- endif -%} +{%- endfor -%} +{{- "\n" }} +{%- 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 + '' + content }} + {%- else -%} + {{- '<' + think_tag + '>\n' + 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\n' }} + {%- elif effort == 'medium' -%} + {{- '<|ifm|im_start|>assistant\n\n' }} + {%- elif effort == 'low' -%} + {{- '<|ifm|im_start|>assistant\n\n' }} + {%- else -%} + {{- raise_exception("Unsupported reasoning_effort: '" + effort + "'. Supported values: high, medium, low.") -}} + {%- endif -%} +{%- endif -%} diff --git a/tests/load/test_template_parse.py b/tests/load/test_template_parse.py new file mode 100644 index 000000000..d92c11b07 --- /dev/null +++ b/tests/load/test_template_parse.py @@ -0,0 +1,55 @@ +"""Static parse of chat templates that use the renderer's extension tags.""" + +from __future__ import annotations + +from pathlib import Path + +import pytest + +_K2 = Path(__file__).resolve().parents[1] / "fixtures" / "k2_horizon_template.jinja" + +_GENERATION = ( + "{% for m in messages %}{% if m.role == 'assistant' %}" + "{% generation %}{{ m.content }}{% endgeneration %}" + "{% else %}{{ m.content }}{% if loop.index > 9 %}{% break %}{% endif %}" + "{% endif %}{% endfor %}") + + +def test_plain_jinja_rejects_what_the_parser_accepts(): + import jinja2 + + from gmlx.load.template_parse import parse_chat_template + + with pytest.raises(jinja2.TemplateSyntaxError): + jinja2.Environment().parse(_GENERATION) + parse_chat_template(_GENERATION) + + +def test_generation_body_stays_visible_to_meta(): + from jinja2 import meta + + from gmlx.load.template_parse import parse_chat_template + + names = meta.find_undeclared_variables(parse_chat_template( + "{% generation %}{{ reasoning_effort }}{% endgeneration %}")) + assert names == {"reasoning_effort"} + + +def test_parser_still_rejects_bad_jinja(): + import jinja2 + + from gmlx.load.template_parse import parse_chat_template + + with pytest.raises(jinja2.TemplateSyntaxError): + parse_chat_template("{% generation %}unclosed") + + +def test_chat_template_override_accepts_the_generation_tag(tmp_path): + from gmlx.load.loader import _resolve_chat_template + + text = _K2.read_text() + assert "{% generation %}" in text + p = tmp_path / "k2.jinja" + p.write_text(text) + assert _resolve_chat_template(str(p)) == text + assert _resolve_chat_template(_GENERATION) == _GENERATION diff --git a/tests/tui/test_reasoning.py b/tests/tui/test_reasoning.py index 9d92330cb..802792149 100644 --- a/tests/tui/test_reasoning.py +++ b/tests/tui/test_reasoning.py @@ -8,6 +8,8 @@ from __future__ import annotations +from pathlib import Path + import pytest import gmlx.tui.chat as chat @@ -918,3 +920,44 @@ def test_bundled_template_wins_the_spelling_probe(monkeypatch, tmp_path): assert chat._template_text(args) == bundled_chat_template_for_arch( "deepseek41") assert chat.fold_thinking_flag(args, {}) == {"enable_thinking": False} + + +_K2_TEMPLATE = (Path(__file__).resolve().parents[1] / "fixtures" + / "k2_horizon_template.jinja") + + +def test_thinking_named_only_in_an_exception_message_is_no_switch(): + """K2 Horizon raises "...missing a thinking field..." inside a {{ }} + tag. That string must not read as Kimi's bare ``thinking`` switch, so + --thinking off warns instead of setting a no-op.""" + from gmlx.tui.reasoning import map_thinking_controls + + template = _K2_TEMPLATE.read_text() + warns: list = [] + assert map_thinking_controls({}, "off", None, template, + warn=warns.append) == {} + assert warns and "no thinking switch" in warns[0] + assert "reasoning_effort" in warns[0] + + +def test_unparsed_fallback_skips_string_literals(): + """With a template jinja cannot parse, the text probe still ignores a + word inside a string literal and still finds the bare variable.""" + from gmlx.tui.reasoning import map_thinking_controls + + broken = "{% if %}" + prose = broken + "{{ raise_exception('needs a thinking field') }}" + assert "thinking" not in map_thinking_controls( + {}, "off", None, prose, warn=lambda m: None) + assert map_thinking_controls({}, "off", None, broken + _KIMI_TAIL) == \ + {"thinking": False} + + +@pytest.mark.parametrize("name", ["kimi_k25_template.jinja", + "kimi_k3_template.jinja"]) +def test_kimi_templates_keep_the_bare_thinking_switch(name): + from gmlx.tui.reasoning import map_thinking_controls + + template = (_K2_TEMPLATE.parent / name).read_text() + assert map_thinking_controls({}, "off", None, template) == \ + {"thinking": False} From 59d6e18b5160a7a2fc07a351c4923db8c1352e17 Mon Sep 17 00:00:00 2001 From: Asher Feldman <59994+asher@users.noreply.github.com> Date: Sat, 17 Oct 2026 18:18:42 -0700 Subject: [PATCH 02/10] feat(k2-horizon): k2 horizon dense and mova model, with load wiring and the k2-horizon pre-tokenizer --- docs/arch-coverage.md | 3 +- gmlx/_vendor/manifest.json | 8 +- .../mlx_vlm/models/k2_horizon/__init__.py | 4 - .../mlx_vlm/models/k2_horizon/config.py | 40 -- .../mlx_vlm/models/k2_horizon/k2_horizon.py | 49 -- .../mlx_vlm/models/k2_horizon/language.py | 171 ------- gmlx/load/arch_table.py | 2 + gmlx/load/config_synth.py | 95 ++++ gmlx/load/loader.py | 7 + gmlx/load/remap.py | 62 +++ gmlx/load/rope_scaling.py | 3 +- gmlx/load/tokenizer.py | 17 +- gmlx/models/k2_horizon/__init__.py | 2 + gmlx/models/k2_horizon/model.py | 421 ++++++++++++++++ gmlx/models/owned.py | 3 + tests/load/test_adapter_save.py | 2 + tests/load/test_config_synth.py | 142 ++++++ tests/load/test_remap.py | 54 ++ tests/load/test_rope_swap.py | 5 +- tests/load/test_tokenizer.py | 61 +++ tests/models/test_k2_horizon.py | 465 ++++++++++++++++++ tests/tune/test_training_step.py | 3 +- tests/tune/tiny_train_archs.py | 9 + 23 files changed, 1354 insertions(+), 274 deletions(-) delete mode 100644 gmlx/_vendor/mlx_vlm/models/k2_horizon/__init__.py delete mode 100644 gmlx/_vendor/mlx_vlm/models/k2_horizon/config.py delete mode 100644 gmlx/_vendor/mlx_vlm/models/k2_horizon/k2_horizon.py delete mode 100644 gmlx/_vendor/mlx_vlm/models/k2_horizon/language.py create mode 100644 gmlx/models/k2_horizon/__init__.py create mode 100644 gmlx/models/k2_horizon/model.py create mode 100644 tests/models/test_k2_horizon.py diff --git a/docs/arch-coverage.md b/docs/arch-coverage.md index 9982f4c77..931a5a7af 100644 --- a/docs/arch-coverage.md +++ b/docs/arch-coverage.md @@ -1,6 +1,6 @@ # Supported architectures -gmlx maps 44 GGUF architectures to model types, and 43 of them load end-to-end +gmlx maps 45 GGUF architectures to model types, and 44 of them load end-to-end with no `hf_source`. `scripts/check-coverage.py` generates this page from the architecture table and the mlx-lm that gmlx carries, so do not edit it by hand. @@ -47,6 +47,7 @@ family, backend, status and caveat: | `hunyuan-moe` | hunyuan | mlx-lm | loadable | | | `hy_v3` | hunyuan | mlx-lm | loadable | Early GGUFs whose arch is hy-v3 with a dash are not mapped. Reconvert them. | | `hyv4` | hunyuan | mlx-lm | loadable | The GGUF has no MTP head, so speculative decoding is not available. `--kv-bits` is refused. | +| `k2-horizon` | k2-horizon | mlx-lm | loadable | The GGUF has no MTP head, so speculative decoding is not available. | | `kimi-k3` | kimi | mlx-lm | loadable | The hybrid cache cannot go back in place, so after `/retry` or `/undo`, chat reads the earlier conversation again with the next message. | | `llama` | llama | mlx-lm | loadable | | | `minimax-m2` | minimax | mlx-lm | loadable | | diff --git a/gmlx/_vendor/manifest.json b/gmlx/_vendor/manifest.json index 7f4dc3e39..3c52b47ae 100644 --- a/gmlx/_vendor/manifest.json +++ b/gmlx/_vendor/manifest.json @@ -714,10 +714,10 @@ "mlx_vlm/models/jina_vlm/language.py": {"status": "verbatim", "upstream": "65f112e5c77191c8e29cb01991ae21b5cf9d693cbf651664e27adbd25727b63c"}, "mlx_vlm/models/jina_vlm/processing_jinavlm.py": {"status": "verbatim", "upstream": "81fdae0d704f82f00081cb9c2dfd3ff4fe0f680f3275251eeab63ce6e946d875"}, "mlx_vlm/models/jina_vlm/vision.py": {"status": "verbatim", "upstream": "424382f0520c3ee7a0f6482c48a60f3d8498fc5567da39e395d87b42b88cfbf5"}, - "mlx_vlm/models/k2_horizon/__init__.py": {"status": "verbatim", "upstream": "ab9041063ffb321cc77fa4b03f2ef24c074af5de7d4f5ae832d05674a8da9081"}, - "mlx_vlm/models/k2_horizon/config.py": {"status": "verbatim", "upstream": "d2f6c6254fd7072612c1423c72846cd6c4f857ee5ab51eb3988438281b370ecb"}, - "mlx_vlm/models/k2_horizon/k2_horizon.py": {"status": "verbatim", "upstream": "ff47fd80ceb15716553a154029853abe080c5569d3fa8c9406d1bc85dbf8a180"}, - "mlx_vlm/models/k2_horizon/language.py": {"status": "verbatim", "upstream": "9b0df74e8226b437ad242fe488c8881737c4e75ee534b192136fd596cc63032c"}, + "mlx_vlm/models/k2_horizon/__init__.py": {"status": "excluded", "upstream": "ab9041063ffb321cc77fa4b03f2ef24c074af5de7d4f5ae832d05674a8da9081"}, + "mlx_vlm/models/k2_horizon/config.py": {"status": "excluded", "upstream": "d2f6c6254fd7072612c1423c72846cd6c4f857ee5ab51eb3988438281b370ecb"}, + "mlx_vlm/models/k2_horizon/k2_horizon.py": {"status": "excluded", "upstream": "ff47fd80ceb15716553a154029853abe080c5569d3fa8c9406d1bc85dbf8a180"}, + "mlx_vlm/models/k2_horizon/language.py": {"status": "excluded", "upstream": "9b0df74e8226b437ad242fe488c8881737c4e75ee534b192136fd596cc63032c"}, "mlx_vlm/models/kernels.py": {"status": "verbatim", "upstream": "7f18031e867b31eae7f57f38c53c653983bcd867e65080e66a32f60da8452360"}, "mlx_vlm/models/kimi_k25/__init__.py": {"status": "verbatim", "upstream": "647b1ea761370bbf2c7f0de5ca500bd185e837994abd04246035bb689b3f3e98"}, "mlx_vlm/models/kimi_k25/config.py": {"status": "verbatim", "upstream": "784f6bddddbaf0d38135e09e2596b5856573aafa551d3d628f246a5428e20034"}, diff --git a/gmlx/_vendor/mlx_vlm/models/k2_horizon/__init__.py b/gmlx/_vendor/mlx_vlm/models/k2_horizon/__init__.py deleted file mode 100644 index c2352e66c..000000000 --- a/gmlx/_vendor/mlx_vlm/models/k2_horizon/__init__.py +++ /dev/null @@ -1,4 +0,0 @@ -from .config import ModelConfig -from .k2_horizon import Model - -__all__ = ["Model", "ModelConfig"] diff --git a/gmlx/_vendor/mlx_vlm/models/k2_horizon/config.py b/gmlx/_vendor/mlx_vlm/models/k2_horizon/config.py deleted file mode 100644 index bf1dbc6bd..000000000 --- a/gmlx/_vendor/mlx_vlm/models/k2_horizon/config.py +++ /dev/null @@ -1,40 +0,0 @@ -from dataclasses import dataclass -from typing import Any, Dict, Optional - -from ..base import BaseModelConfig - - -@dataclass -class ModelConfig(BaseModelConfig): - model_type: str = "k2_horizon" - hidden_size: int = 4096 - num_hidden_layers: int = 36 - intermediate_size: int = 12288 - num_attention_heads: int = 32 - num_key_value_heads: int = 8 - head_dim: int = 128 - vocab_size: int = 250624 - rms_norm_eps: float = 1e-6 - max_position_embeddings: int = 524288 - rope_theta: float = 1e7 - rope_scaling: Optional[Dict[str, Any]] = None - attention_bias: bool = False - tie_word_embeddings: bool = False - - @classmethod - def from_dict(cls, params: Dict[str, Any]): - params = dict(params or {}) - # K2-Horizon carries RoPE config under `rope_parameters`; translate it to the - # (rope_theta, rope_scaling) pair that initialize_rope consumes. A "default" - # type collapses to plain RoPE (rope_scaling=None); "yarn" keeps the scaling - # dict (factor / beta_fast / beta_slow / original_max_position_embeddings). - rope = params.pop("rope_parameters", None) or params.get("rope_scaling") - if rope: - rope = dict(rope) - rope_type = rope.get("rope_type") or rope.get("type") or "default" - params["rope_theta"] = rope.pop("rope_theta", params.get("rope_theta")) - params["rope_scaling"] = None if rope_type == "default" else rope - return cls(**{k: v for k, v in params.items() if k in cls.__dataclass_fields__}) - - -K2HorizonConfig = ModelConfig diff --git a/gmlx/_vendor/mlx_vlm/models/k2_horizon/k2_horizon.py b/gmlx/_vendor/mlx_vlm/models/k2_horizon/k2_horizon.py deleted file mode 100644 index 30b4315b1..000000000 --- a/gmlx/_vendor/mlx_vlm/models/k2_horizon/k2_horizon.py +++ /dev/null @@ -1,49 +0,0 @@ -from typing import Optional - -import mlx.core as mx -import mlx.nn as nn - -from ..base import InputEmbeddingsFeatures, LanguageModelOutput -from .config import ModelConfig -from .language import LanguageModel - - -class Model(nn.Module): - def __init__(self, config: ModelConfig): - super().__init__() - self.config = config - self.model_type = config.model_type - self.language_model = LanguageModel(config) - - def get_input_embeddings( - self, - input_ids: Optional[mx.array] = None, - pixel_values: Optional[mx.array] = None, - **kwargs, - ) -> InputEmbeddingsFeatures: - return InputEmbeddingsFeatures( - inputs_embeds=self.language_model.model.embed_tokens(input_ids) - ) - - def __call__( - self, - input_ids: mx.array, - pixel_values: mx.array = None, - mask: mx.array = None, - cache=None, - **kwargs, - ) -> LanguageModelOutput: - return self.language_model(input_ids, cache=cache, **kwargs) - - def sanitize(self, weights): - if any(k.startswith("language_model.") for k in weights): - return weights - weights = self.language_model.sanitize(weights) - return {f"language_model.{k}": v for k, v in weights.items()} - - def make_cache(self): - return self.language_model.make_cache() - - @property - def layers(self): - return self.language_model.layers diff --git a/gmlx/_vendor/mlx_vlm/models/k2_horizon/language.py b/gmlx/_vendor/mlx_vlm/models/k2_horizon/language.py deleted file mode 100644 index f3e4d5ef4..000000000 --- a/gmlx/_vendor/mlx_vlm/models/k2_horizon/language.py +++ /dev/null @@ -1,171 +0,0 @@ -from typing import Any, Optional - -import mlx.core as mx -import mlx.nn as nn - -from ..base import ( - LanguageModelOutput, - create_attention_mask, - scaled_dot_product_attention, -) -from ..cache import KVCache -from ..mlp import SwiGLUMLP -from ..rope_utils import initialize_rope -from .config import ModelConfig - - -class Attention(nn.Module): - def __init__(self, args: ModelConfig): - 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 = head_dim = args.head_dim or dim // self.n_heads - self.scale = head_dim**-0.5 - - self.q_proj = nn.Linear(dim, self.n_heads * head_dim, bias=args.attention_bias) - self.k_proj = nn.Linear( - dim, self.n_kv_heads * head_dim, bias=args.attention_bias - ) - self.v_proj = nn.Linear( - dim, self.n_kv_heads * head_dim, bias=args.attention_bias - ) - self.o_proj = nn.Linear(self.n_heads * head_dim, dim, bias=args.attention_bias) - - self.rope = initialize_rope( - head_dim, - base=args.rope_theta, - traditional=False, - scaling_config=args.rope_scaling, - max_position_embeddings=args.max_position_embeddings, - ) - - def __call__( - self, - x: mx.array, - mask: Optional[mx.array] = None, - cache: Optional[Any] = None, - ) -> mx.array: - B, L, _ = x.shape - - queries, keys, values = self.q_proj(x), self.k_proj(x), self.v_proj(x) - - queries = queries.reshape(B, L, self.n_heads, -1).transpose(0, 2, 1, 3) - keys = keys.reshape(B, L, self.n_kv_heads, -1).transpose(0, 2, 1, 3) - values = values.reshape(B, L, self.n_kv_heads, -1).transpose(0, 2, 1, 3) - - if cache is not None: - queries = self.rope(queries, offset=cache.offset) - keys = self.rope(keys, offset=cache.offset) - keys, values = cache.update_and_fetch(keys, values) - else: - queries = self.rope(queries) - keys = self.rope(keys) - - output = scaled_dot_product_attention( - queries, keys, values, cache=cache, scale=self.scale, mask=mask - ) - output = output.transpose(0, 2, 1, 3).reshape(B, L, -1) - return self.o_proj(output) - - -class TransformerBlock(nn.Module): - def __init__(self, args: ModelConfig): - super().__init__() - self.self_attn = Attention(args) - self.mlp = SwiGLUMLP(args.hidden_size, args.intermediate_size) - self.input_layernorm = nn.RMSNorm(args.hidden_size, eps=args.rms_norm_eps) - self.post_attention_layernorm = nn.RMSNorm( - args.hidden_size, eps=args.rms_norm_eps - ) - - def __call__( - self, - x: mx.array, - mask: Optional[mx.array] = None, - cache: Optional[Any] = None, - ) -> mx.array: - 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: ModelConfig): - super().__init__() - self.args = args - self.embed_tokens = nn.Embedding(args.vocab_size, args.hidden_size) - self.layers = [TransformerBlock(args) for _ in range(args.num_hidden_layers)] - self.norm = nn.RMSNorm(args.hidden_size, eps=args.rms_norm_eps) - - def __call__( - self, - inputs: Optional[mx.array] = None, - cache=None, - input_embeddings: Optional[mx.array] = None, - ): - if input_embeddings is not None: - h = input_embeddings - else: - h = self.embed_tokens(inputs) - - if cache is None: - cache = [None] * len(self.layers) - mask = create_attention_mask(h, cache[0]) - - for layer, c in zip(self.layers, cache): - h = layer(h, mask, c) - - return self.norm(h) - - -class LanguageModel(nn.Module): - def __init__(self, args: ModelConfig): - super().__init__() - self.args = args - self.config = 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: Optional[mx.array] = None, - cache=None, - input_embeddings: Optional[mx.array] = None, - inputs_embeds: Optional[mx.array] = None, - **kwargs, - ) -> LanguageModelOutput: - if inputs is None: - inputs = kwargs.get("input_ids") - if inputs_embeds is None: - inputs_embeds = input_embeddings - - out = self.model(inputs, cache, inputs_embeds) - if self.args.tie_word_embeddings: - out = self.model.embed_tokens.as_linear(out) - else: - out = self.lm_head(out) - return LanguageModelOutput(logits=out) - - def sanitize(self, weights): - if self.args.tie_word_embeddings: - weights.pop("lm_head.weight", None) - return weights - - def make_cache(self): - return [KVCache() for _ in self.model.layers] - - @property - def layers(self): - return self.model.layers - - @property - def head_dim(self): - return ( - self.args.head_dim or self.args.hidden_size // self.args.num_attention_heads - ) - - @property - def n_kv_heads(self): - return self.args.num_key_value_heads diff --git a/gmlx/load/arch_table.py b/gmlx/load/arch_table.py index 3cf55cc4f..6b3a2d5ed 100644 --- a/gmlx/load/arch_table.py +++ b/gmlx/load/arch_table.py @@ -97,6 +97,7 @@ class ArchEntry: "kimi-k3": ("kimi", "Moonshot Kimi-K3 (2.8T-A50B, llama.cpp PR #26185); hybrid KDA linear attention (per-key-channel-decay delta rule via the fused gated_delta kernels, short conv1d Q/K/V, full-rank sigmoid output gate) + nope-only MLA every attn_res-th-ish layer (per-layer head_count_kv array, 0 = KDA; absorbed embed_q/unembed_out, sigmoid output gate, NO rope anywhere), cross-layer residual attention (softmax-scored convex mix over banked residual checkpoints, restart on bank layers), latent MoE (896 experts top-16 at routed_expert_hidden_size behind routed_down/routed_up, sigmoid gating + correction bias, shared experts at full width), and situ activation replacing SwiGLU everywhere. ssm_a arrives folded (-exp(A_log)) and stays folded. Native MXFP4 experts (gpt-oss codec). XTML chat format (<|open|>tag<|sep|>), tiktoken kimi-k2 pretokenizer, EOS 163586 <|end_of_msg|>. Hybrid cache is non-trimmable -> chat re-prefills on trim. Model class owned by gmlx and installed over any upstream class of the same name (kimi_linear is the nearest relative)"), "muse-glimmer": ("muse", "Meta Muse Glimmer 30B (llama.cpp 'muse-glimmer'); dense 52-layer sandwich-norm decoder - pre/post norms around both attention and FFN, at two epsilons (1e-5 pre/final, 1e-8 post) - with an afmoe-shaped attention output gate (sigmoid(x_norm @ W_gate) applied before o_proj), per-head qk-norm whose weights are synthesized at conversion to absorb qk_scale_factor (q_norm a uniform 3.87, k_norm ones), an unweighted RMSNorm on the token embeddings before layer 0, and a logit scale (output_multiplier 0.196) + gemma-style tanh softcap at 20. Attention runs a 3-of-4 sliding window (2048) with RoPE, and full attention with NoPE on every 4th layer - the inverse of the usual arrangement, and the reason the model card says '131072+': the largest positional offset ever resolved is the 2048 window, so there is no rope extrapolation ceiling and past 131k the only limit is KV memory at ~13 KB/token (just the 13 global layers grow; the 39 sliding ones are pinned by their rotating cache). The four per-layer norms carry a baked +1 that a plain RMSNorm consumes as-is (no gemma unbake). NORM rope with Q/K un-permuted at conversion => traditional=True, no qk_permute. Model class owned by gmlx and installed over any upstream class of the same name (afmoe is the nearest relative). Pairs with the Muse Glimmer mmproj (--mmproj) for vision and with the DFlash drafter GGUF for speculative decoding (autodetected next to the target, or --draft-gguf); the two compose, so a --mmproj load speculates on text turns"), "muse-spark": ("muse", "Meta Muse Spark Open 816B-A42B (llama.cpp 'muse-spark'); 62 layers at d=8192 with Muse Glimmer's attention (sigmoid output gate from the normed input, per-head qk-norm whose q weight carries the MetaP scale, RoPE theta 5e5 on the 2048-window sliding layers and NoPE on every 4th global layer, an unweighted RMSNorm on the token embeddings, logit scale 0.177 and a tanh softcap at 20) plus a weightless per-head RMSNorm on the attention output, weightless post-attention and post-FFN norms at post_norm_eps 1e-8, and norm-preserving residuals (x*alpha + branch*beta, per-channel alpha/beta resolved at conversion, kept fp32). Every layer routes a latent MoE: the router reads the full-width FFN input (fp32 sigmoid, selection-only exp_probs_b bias, top-8 of 256, renormalized), the experts run SwiGLU at moe_latent_size 4096 between latent_down and latent_up, and each expert output is scaled by post_exp_norm and normalized before the router weights apply. No shared expert, no MTP. The folded +1 norms load as plain gains. NORM rope with Q/K un-permuted at conversion => traditional=True, no qk_permute. Tool calls use Muse Glimmer's ATEM parser. Model class owned by gmlx (muse_glimmer is the nearest relative)"), + "k2-horizon": ("k2-horizon", "IFM K2 Horizon (llama.cpp 'k2-horizon'), dense 7B and MoVA 36B-A4B; GQA with NeoX rope at theta 1e7 and 512K native context. Every norm is a group RMSNorm: one RMS statistic per equal slice of the width (attention.group_norm_groups, 4 on the 7B, 2 on MoVA) and one full-width gain. MoVA adds an output gate on every layer (log2(1 + 2^g) of a projection of the normed input, applied before o_proj), 3 leading dense layers, then a sigmoid MoE (100 experts top-8, selection-only bias, renorm x 2.5, 1 shared expert) and mixture-of-value-attention: V is the weighted sum of silu(W_k x) over the top-4 of 64 value experts, routed with the FFN router's gating, renorm and scale and a selection-only bias. The value router weight is quantized on the wire. The value experts stay resident under expert streaming. Optional per-head Q/K norm before rope, absent in the public files. gpt2 byte-level BPE on the k2-horizon pre (llama3 split with marks, ZWNJ and ZWJ in letter runs); <|ifm|im_end|> ends a turn. Reasoning effort high/medium/low selects , or ; tool calls ride ."), "granitehybrid": ("granite", "IBM Granite 4.x hybrid (H-Micro/H-Tiny/H-Small); alternating Mamba2 + attention (layer_types from per-layer head_count_kv==0), softmax MoE + fused-input shared MLP (loader pre-fuses ffn_{gate,up}_shexp -> input_linear), granite runtime multipliers, NoPE via rope.scaling.finetuned=false. NORM rope => qk_permute on attention layers"), "falcon-h1": ("falcon", "TII Falcon-H1 (0.5B-34B); parallel attention + Mamba2 in every layer (one input_layernorm feeds both, outputs summed), dense gated MLP under feed_forward.*, explicit head_dim from key_length. NEOX rope (no qk-permute). The muP multiplier zoo is folded into the wire weights at convert => synth pins every multiplier neutral; ffn_norm/ssm_a/ssm_d stored with no .weight suffix"), "qwen3next": ("qwen", "Qwen3-Next-80B-A3B; gated-DeltaNet linear attention (3 of every 4 layers) + gated full attention (gate fused in attn_q), every-layer 512-expert MoE + shared expert. NEOX rope (no qk-permute, not qwen35's packing). Both GDN wire layouts load: legacy fused ssm_in -> in_proj_qkvz; the newer split attn_qkv/attn_gate via a load-time module split (loader swaps in_proj_qkvz for in_proj_qkv/in_proj_z, skipping the runtime de-interleave). V heads HF-grouped (kv_head_layout grouped, not qwen3.5's tiled order); +1 norm bake is what mlx-lm expects (passthrough)"), @@ -118,6 +119,7 @@ class ArchEntry: "hy_v3": "Early GGUFs whose arch is hy-v3 with a dash are not mapped. Reconvert them.", "hyv4": "The GGUF has no MTP head, so speculative decoding is not available. `--kv-bits` is refused.", "muse-spark": "The GGUF has no MTP head, so speculative decoding is not available.", + "k2-horizon": "The GGUF has no MTP head, so speculative decoding is not available.", "qwen4exp": "Speculative decoding needs a companion MTP GGUF beside the model or in an `MTP` folder, or named with `--draft-gguf`.", "kimi-k3": "The hybrid cache cannot go back in place, so after `/retry` or `/undo`, chat reads the earlier conversation again with the next message.", "glm5-next": "The hybrid cache cannot go back in place, so after `/retry` or `/undo`, chat reads the earlier conversation again with the next message.", diff --git a/gmlx/load/config_synth.py b/gmlx/load/config_synth.py index 02646347d..992014071 100644 --- a/gmlx/load/config_synth.py +++ b/gmlx/load/config_synth.py @@ -169,6 +169,11 @@ # residuals and a latent sigmoid MoE on every layer. Model class owned by # gmlx (no mlx-lm class). "muse-spark": "muse_spark", + # IFM K2 Horizon: GQA with NeoX rope and group RMSNorm, an optional + # softplus output gate, and on MoVA models a sigmoid MoE with a shared + # expert plus routed value experts in the attention. Model class owned + # by gmlx (the mlx-vlm class has none of the group norm, gate or MoVA). + "k2-horizon": "k2_horizon", # Moonshot Kimi-K3 (2.8T-A50B): hybrid KDA (linear, per-channel decay) + # nope-only MLA layers from the per-layer head_count_kv schedule, latent # 896-expert sigmoid MoE behind down/up projections, situ activation, @@ -1880,6 +1885,95 @@ def req(read, key): config["norm_topk_prob"] = req(_read_bool, "expert_weights_norm") +# k2-horizon (IFM K2 Horizon, dense and MoVA) + +# llama.cpp's expert_gating_func values that K2 Horizon accepts. +_K2_HORIZON_SCORING = {1: "softmax", 2: "sigmoid"} + + +def _synth_k2_horizon(meta, shapes, config: dict) -> None: + """Synthesize a k2_horizon config from a 'k2-horizon'-arch GGUF. + + The universal fields cover hidden/layers/heads/kv/ffn/ctx/eps/head_dim/ + tie/vocab. This adds the norm groups, the rope dims, the optional output + gate and Q/K norm (from tensor presence), and the MoE and MoVA routing. + The defaults follow llama.cpp's loader: 0 norm groups means 1, a shared + width of 0 means ``n_ff_exp * n_shared``, an unset gating function means + sigmoid, and the routing norm and scale are read only when the model + has FFN experts. A layer routes when it is at or after + ``leading_dense_block_count``; llama.cpp ignores ``moe_every_n_layers``. + """ + arch = "k2-horizon" + + def req(read, key): + return _require(read(meta, f"{arch}.{key}"), arch=arch, + gguf_field=f"{arch}.{key}") + + _require(config.get("rms_norm_eps"), arch=arch, + gguf_field=f"{arch}.attention.layer_norm_rms_epsilon") + head_dim = req(_read_int, "attention.key_length") + value_length = _read_int(meta, f"{arch}.attention.value_length") + if value_length is not None and value_length != head_dim: + raise NotImplementedError( + f"k2-horizon synth: value_length {value_length} differs from " + f"key_length {head_dim}; llama.cpp asserts they match.") + config["head_dim"] = head_dim + config["norm_groups"] = _read_int( + meta, f"{arch}.attention.group_norm_groups") or 1 + rope_dim = _read_int(meta, f"{arch}.rope.dimension_count") + if rope_dim is not None and rope_dim != head_dim: + config["rope_dim"] = rope_dim + config.setdefault("intermediate_size", 0) + + config["attn_output_gate"] = _has_tensor(shapes, "blk.0.attn_gate.weight") + config["qk_norm"] = _has_tensor(shapes, "blk.0.attn_q_norm.weight") + + n_experts = _read_int(meta, f"{arch}.expert_count") or 0 + gating = _read_int(meta, f"{arch}.expert_gating_func") or 0 + if n_experts > 0: + n_ff_exp = req(_first_or_scalar, "expert_feed_forward_length") + n_shared = _read_int(meta, f"{arch}.expert_shared_count") or 0 + n_ff_shexp = _read_int( + meta, f"{arch}.expert_shared_feed_forward_length") or 0 + if n_shared and not n_ff_shexp: + n_ff_shexp = n_ff_exp * n_shared + config["num_experts"] = n_experts + config["num_experts_per_tok"] = req(_read_int, "expert_used_count") + config["moe_intermediate_size"] = n_ff_exp + config["shared_expert_intermediate_size"] = ( + n_ff_shexp if n_shared else 0) + config["first_k_dense_replace"] = _read_int( + meta, f"{arch}.leading_dense_block_count") or 0 + config["routed_scaling_factor"] = _read_float( + meta, f"{arch}.expert_weights_scale") or 0.0 + config["norm_topk_prob"] = bool( + _read_bool(meta, f"{arch}.expert_weights_norm")) + gating = gating or 2 + + n_values = _read_int(meta, f"{arch}.attention.value_expert_count") or 0 + n_values_used = _read_int( + meta, f"{arch}.attention.value_expert_used_count") or 0 + if n_values > 0: + if not 0 < n_values_used <= n_values: + raise ValueError( + f"k2-horizon synth: {n_values_used} value experts used of " + f"{n_values}; llama.cpp needs 0 < used <= count.") + config["num_value_experts"] = n_values + config["num_value_experts_per_tok"] = n_values_used + elif n_values_used: + raise ValueError( + "k2-horizon synth: value_expert_used_count is set without " + "value_expert_count.") + + if n_experts > 0 or n_values > 0: + scoring = _K2_HORIZON_SCORING.get(gating) + if scoring is None: + raise NotImplementedError( + f"k2-horizon synth: {arch}.expert_gating_func {gating} is not " + "softmax (1) or sigmoid (2).") + config["scoring_func"] = scoring + + # granitehybrid (IBM Granite 4.x hybrid: H-Micro / H-Tiny / H-Small) def _synth_granite_hybrid(meta, shapes, config: dict) -> None: @@ -3815,6 +3909,7 @@ def _synth_qwen4exp(meta, shapes, config: dict) -> None: "hy_v3": _synth_hy_v3, "muse-glimmer": _synth_muse_glimmer, "muse-spark": _synth_muse_spark, + "k2-horizon": _synth_k2_horizon, "granitehybrid": _synth_granite_hybrid, "falcon-h1": _synth_falcon_h1, "qwen3next": _synth_qwen3next, diff --git a/gmlx/load/loader.py b/gmlx/load/loader.py index 500be70f0..4040cac7a 100644 --- a/gmlx/load/loader.py +++ b/gmlx/load/loader.py @@ -920,6 +920,13 @@ def _warm_touch_pass( # Kept in step with the model class's cast_predicate. "muse_spark": (".mlp.gate.weight", ".e_score_correction_bias", "_resid_alpha", "_resid_beta"), + # k2_horizon: sigmoid top-8-of-100 FFN routing and the MoVA value + # routing pick on scores plus a selection bias, near-tie-heavy, and + # llama.cpp routes fp32. The value router weight is quantized on the + # wire, so only its F32 bias qualifies. Kept in step with the model + # class's cast_predicate. + "k2_horizon": (".mlp.gate.weight", ".e_score_correction_bias", + ".selection_bias"), } # Params kept at their native f16 through the bf16 cast (no upcast). MLX diff --git a/gmlx/load/remap.py b/gmlx/load/remap.py index ca20f9597..66c3346d7 100644 --- a/gmlx/load/remap.py +++ b/gmlx/load/remap.py @@ -213,6 +213,11 @@ # traditional rope, as on Muse Glimmer, and the folded +1 norms load as # plain gains. "muse-spark": "MUSE_SPARK", + # IFM K2 Horizon (llama.cpp 'k2-horizon'): NeoX rope, so Q/K pass through + # un-permuted. The K2_HORIZON block claims every tensor, including the + # attention gate the canonical map homes on qwen3.5's linear_attn and + # the MoVA value router and experts gguf-py cannot name. + "k2-horizon": "K2_HORIZON", # DiffusionGemma (llama.cpp 'diffusion-gemma'): an encoder-decoder block- # diffusion model on the Gemma-4 MoE backbone. The decoder backbone uses the # exact Gemma-4 GGUF tensor names, but the mlx-vlm Model nests them under @@ -589,6 +594,63 @@ def _gemma_norm_transform(arch_string: str, hf_name: str | None, (re.compile(r"^blk\.(\d+)\.ffn_resid_beta\.weight$"), "model.layers.{bid}.ffn_resid_beta", "passthrough"), ], + "K2_HORIZON": [ + # Every tensor is claimed here. attn_v_gate.bias and exp_probs_b are + # selection-only router biases. A MoVA layer has the value router and + # experts in place of attn_v. + (re.compile(r"^token_embd\.weight$"), + "model.embed_tokens.weight", "passthrough"), + (re.compile(r"^output_norm\.weight$"), + "model.norm.weight", "passthrough"), + (re.compile(r"^output\.weight$"), + "lm_head.weight", "passthrough"), + (re.compile(r"^blk\.(\d+)\.attn_norm\.weight$"), + "model.layers.{bid}.input_layernorm.weight", "passthrough"), + (re.compile(r"^blk\.(\d+)\.attn_q\.weight$"), + "model.layers.{bid}.self_attn.q_proj.weight", "passthrough"), + (re.compile(r"^blk\.(\d+)\.attn_k\.weight$"), + "model.layers.{bid}.self_attn.k_proj.weight", "passthrough"), + (re.compile(r"^blk\.(\d+)\.attn_v\.weight$"), + "model.layers.{bid}.self_attn.v_proj.weight", "passthrough"), + (re.compile(r"^blk\.(\d+)\.attn_output\.weight$"), + "model.layers.{bid}.self_attn.o_proj.weight", "passthrough"), + (re.compile(r"^blk\.(\d+)\.attn_gate\.weight$"), + "model.layers.{bid}.self_attn.gate_proj.weight", "passthrough"), + (re.compile(r"^blk\.(\d+)\.attn_q_norm\.weight$"), + "model.layers.{bid}.self_attn.q_norm.weight", "passthrough"), + (re.compile(r"^blk\.(\d+)\.attn_k_norm\.weight$"), + "model.layers.{bid}.self_attn.k_norm.weight", "passthrough"), + (re.compile(r"^blk\.(\d+)\.attn_v_gate\.weight$"), + "model.layers.{bid}.self_attn.value_experts.v_router.weight", "passthrough"), + (re.compile(r"^blk\.(\d+)\.attn_v_gate\.bias$"), + "model.layers.{bid}.self_attn.value_experts.selection_bias", "passthrough"), + (re.compile(r"^blk\.(\d+)\.attn_v_exps\.weight$"), + "model.layers.{bid}.self_attn.value_experts.v_experts.weight", "passthrough"), + (re.compile(r"^blk\.(\d+)\.ffn_norm\.weight$"), + "model.layers.{bid}.post_attention_layernorm.weight", "passthrough"), + (re.compile(r"^blk\.(\d+)\.ffn_gate\.weight$"), + "model.layers.{bid}.mlp.gate_proj.weight", "passthrough"), + (re.compile(r"^blk\.(\d+)\.ffn_up\.weight$"), + "model.layers.{bid}.mlp.up_proj.weight", "passthrough"), + (re.compile(r"^blk\.(\d+)\.ffn_down\.weight$"), + "model.layers.{bid}.mlp.down_proj.weight", "passthrough"), + (re.compile(r"^blk\.(\d+)\.ffn_gate_inp\.weight$"), + "model.layers.{bid}.mlp.gate.weight", "passthrough"), + (re.compile(r"^blk\.(\d+)\.exp_probs_b\.bias$"), + "model.layers.{bid}.mlp.gate.e_score_correction_bias", "passthrough"), + (re.compile(r"^blk\.(\d+)\.ffn_gate_exps\.weight$"), + "model.layers.{bid}.mlp.switch_mlp.gate_proj.weight", "passthrough"), + (re.compile(r"^blk\.(\d+)\.ffn_up_exps\.weight$"), + "model.layers.{bid}.mlp.switch_mlp.up_proj.weight", "passthrough"), + (re.compile(r"^blk\.(\d+)\.ffn_down_exps\.weight$"), + "model.layers.{bid}.mlp.switch_mlp.down_proj.weight", "passthrough"), + (re.compile(r"^blk\.(\d+)\.ffn_gate_shexp\.weight$"), + "model.layers.{bid}.mlp.shared_experts.gate_proj.weight", "passthrough"), + (re.compile(r"^blk\.(\d+)\.ffn_up_shexp\.weight$"), + "model.layers.{bid}.mlp.shared_experts.up_proj.weight", "passthrough"), + (re.compile(r"^blk\.(\d+)\.ffn_down_shexp\.weight$"), + "model.layers.{bid}.mlp.shared_experts.down_proj.weight", "passthrough"), + ], "ERNIE4_5_MOE": [ # Baidu ERNIE-4.5-MoE: shared-expert fine-grained MoE with leading dense # layers. mlx-lm's ernie4_5_moe uses traditional=True rope, which consumes diff --git a/gmlx/load/rope_scaling.py b/gmlx/load/rope_scaling.py index 079008723..bb1225a88 100644 --- a/gmlx/load/rope_scaling.py +++ b/gmlx/load/rope_scaling.py @@ -33,7 +33,7 @@ USER_SCALING_TYPES = ("none", "linear", "yarn") # Header keys llama.cpp reads for some architectures only. -_BETA_ARCHES = frozenset({"mistral3"}) +_BETA_ARCHES = frozenset({"mistral3", "k2-horizon"}) _EXT_ATTN_ARCHES = frozenset({"grok"}) # log multiplier: the deepseek converters write 0.1 * mscale_all_dim. _LOG_MUL_DIV10_ARCHES = frozenset({"deepseek2", "deepseek32"}) @@ -452,6 +452,7 @@ def wrapper(*args, **kwargs): "mistral3", "gpt-oss", "seed_oss", "smollm3", "granite", "qwen3next", "minimax-m2", "qwen35", "qwen35moe", "qwen3vl", "qwen3vlmoe", "hy_v3", "hyv4", "muse-glimmer", "muse-spark", "minimax-m3", "falcon-h1", + "k2-horizon", "granitehybrid", "ernie4_5-moe", "gemma", "gemma2", "gemma3", "gemma4", }) diff --git a/gmlx/load/tokenizer.py b/gmlx/load/tokenizer.py index 7cc79799e..87430b6c3 100644 --- a/gmlx/load/tokenizer.py +++ b/gmlx/load/tokenizer.py @@ -136,6 +136,20 @@ def _bytelevel_pattern(digit_clause: str) -> str: r"\s+" ) +# K2 Horizon word-split, from the HF tokenizer.json (llama.cpp's +# LLAMA_VOCAB_PRE_TYPE_K2_HORIZON spells the same regex with explicit-case +# contractions). The llama3 split, with marks, ZWNJ and ZWJ inside letter runs. +# The HF tokenizer normalizes NFC, so this pre keeps the NFC step. +K2_HORIZON_PATTERN = ( + r"(?i:'s|'t|'re|'ve|'m|'ll|'d)|" + r"[^\r\n\p{L}\p{N}]?(?:\p{L}|\p{M}|\x{200C}|\x{200D})+|" + r"\p{N}{1,3}|" + r" ?[^\s\p{L}\p{N}]+[\r\n]*|" + r"\s*[\r\n]+|" + r"\s+(?!\S)|" + r"\s+" +) + # DeepSeek-V3 family word-split: llama.cpp applies these three regexes # sequentially (each further splits the previous pass's pieces), verbatim # from PRE_TYPE_DEEPSEEK3_LLM (llama-vocab.cpp:318-325; the same case also @@ -175,6 +189,7 @@ def _bytelevel_pattern(digit_clause: str) -> str: # pre fell back to the single-digit clause - over-splitting numbers and # gluing Han runs to Latin words. "kimi-k2": KIMI_K2_PATTERN, + "k2-horizon": K2_HORIZON_PATTERN, } @@ -733,7 +748,7 @@ def _metadata_stop_ids(meta, n_vocab: int | None = None) -> list[int]: "<|flush|>", "<|calls|>", "", "<|endoftext|>", "", "<|eom_id|>", "", "_", "[EOT]", "[EOS]", "<|end_of_text|>", "", "", "", "<|tool_response>", - "<｜end▁of▁sentence｜>", "[e~[", + "<｜end▁of▁sentence｜>", "[e~[", "<|ifm|im_end|>", ) diff --git a/gmlx/models/k2_horizon/__init__.py b/gmlx/models/k2_horizon/__init__.py new file mode 100644 index 000000000..937c9ca47 --- /dev/null +++ b/gmlx/models/k2_horizon/__init__.py @@ -0,0 +1,2 @@ +# SPDX-License-Identifier: MIT +"""Model package for IFM K2 Horizon (GGUF arch ``k2-horizon``).""" diff --git a/gmlx/models/k2_horizon/model.py b/gmlx/models/k2_horizon/model.py new file mode 100644 index 000000000..81866c53a --- /dev/null +++ b/gmlx/models/k2_horizon/model.py @@ -0,0 +1,421 @@ +# SPDX-License-Identifier: MIT +# Copyright (c) 2026 Asher Feldman +"""Model for IFM K2 Horizon (GGUF arch ``k2-horizon``). + +This is the runtime for llama.cpp's ``LLM_ARCH_K2_HORIZON`` conversions, +dense and MoVA. The attention is GQA with NeoX RoPE. These parts differ +from a Llama block: + + 1. Every norm is a group RMSNorm: one RMS statistic per equal slice of + the hidden width, then one full-width gain. + 2. An optional output gate scales the attention output before ``o_proj`` + by ``log2(1 + 2^g)``, with ``g`` projected from the normed input. + 3. An optional per-head Q/K RMSNorm runs before RoPE. + 4. MoVA layers route each token over value experts. V is the weighted + sum of ``silu(W_k x)`` over the selected experts. The router shares + the FFN router's gating, normalization and scale. + 5. MoE layers route the FFN over sigmoid-scored experts with a + selection-only bias, plus one shared expert. +""" + +import math +from dataclasses import dataclass +from typing import Any, Optional + +import mlx.core as mx +import mlx.nn as nn + +from gmlx._vendor.mlx_lm.models.base import BaseModelArgs +from gmlx._vendor.mlx_lm.models.rope_utils import initialize_rope +from gmlx._vendor.mlx_lm.models.switch_layers import ( + SwitchGLU, + SwitchLinear, + _gather_sort, + _scatter_unsort, +) + +from gmlx.cache.attention import ( + create_attention_mask, + scaled_dot_product_attention, +) +from gmlx.cache.kv_cache import KVCache + + +# llama.cpp clamps the routing weight sum to the smallest normal f16. +_WEIGHT_SUM_FLOOR = 6.103515625e-5 +_LN2 = math.log(2.0) + + +@dataclass +class ModelArgs(BaseModelArgs): + model_type: str + hidden_size: int + num_hidden_layers: int + intermediate_size: int + num_attention_heads: int + num_key_value_heads: int + head_dim: int + vocab_size: int + rms_norm_eps: float = 1e-6 + norm_groups: int = 1 + rope_theta: float = 1e7 + rope_dim: Optional[int] = None + rope_scaling: Optional[dict] = None + max_position_embeddings: int = 524288 + attn_output_gate: bool = False + qk_norm: bool = False + num_experts: int = 0 + num_experts_per_tok: int = 0 + moe_intermediate_size: int = 0 + shared_expert_intermediate_size: int = 0 + first_k_dense_replace: int = 0 + num_value_experts: int = 0 + num_value_experts_per_tok: int = 0 + scoring_func: str = "sigmoid" + norm_topk_prob: bool = False + routed_scaling_factor: float = 1.0 + tie_word_embeddings: bool = False + + +def _route_weights(scores, inds, norm, scale): + """Weights at ``inds`` from the unbiased scores, as llama.cpp computes + them: renormalized over a clamped sum, then scaled unless the scale is + 0 or 1.""" + w = mx.take_along_axis(scores, inds, axis=-1) + if norm: + w = w / mx.maximum(w.sum(axis=-1, keepdims=True), _WEIGHT_SUM_FLOOR) + if scale not in (0.0, 1.0): + w = w * scale + return w + + +def _route(logits, bias, top_k, scoring_func, norm, scale): + """``(inds, weights)``: top-k on the scores plus ``bias``, weights from + the unbiased scores.""" + logits = logits.astype(mx.float32) + if scoring_func == "softmax": + scores = mx.softmax(logits, axis=-1, precise=True) + else: + scores = mx.sigmoid(logits) + choice = scores if bias is None else scores + bias + inds = mx.stop_gradient( + mx.argpartition(-choice, kth=top_k - 1, axis=-1)[..., :top_k]) + return inds, _route_weights(scores, inds, norm, scale) + + +class GroupRMSNorm(nn.Module): + """RMS over ``groups`` equal slices of the last axis, then one + full-width gain.""" + + def __init__(self, dims: int, eps: float, groups: int = 1): + super().__init__() + if dims % groups: + raise ValueError(f"width {dims} is not divisible by {groups} groups") + self.weight = mx.ones((dims,)) + self.eps = float(eps) + self.groups = int(groups) + + def __call__(self, x: mx.array) -> mx.array: + if self.groups == 1: + return mx.fast.rms_norm(x, self.weight, self.eps) + shape = x.shape + y = x.reshape(*shape[:-1], self.groups, shape[-1] // self.groups) + y = mx.fast.rms_norm(y, None, self.eps).reshape(shape) + return y * self.weight + + +class HeadRMSNorm(nn.Module): + """RMSNorm over the head dim with a separate gain for every head. The + gains load flat, as llama.cpp stores them.""" + + def __init__(self, n_heads: int, head_dim: int, eps: float): + super().__init__() + self.weight = mx.ones((n_heads * head_dim,)) + self.n_heads = n_heads + self.head_dim = head_dim + self.eps = float(eps) + + def __call__(self, x: mx.array) -> mx.array: + # x: [B, L, H, D] + w = self.weight.reshape(self.n_heads, self.head_dim) + return mx.fast.rms_norm(x, None, self.eps) * w + + +class ValueExperts(nn.Module): + """MoVA: V = sum_k w_k * silu(W_k x) over the routed value experts.""" + + def __init__(self, args: ModelArgs): + super().__init__() + dim = args.hidden_size + out = args.num_key_value_heads * args.head_dim + n = args.num_value_experts + self.top_k = int(args.num_value_experts_per_tok) + self.scoring_func = args.scoring_func + self.norm_topk_prob = bool(args.norm_topk_prob) + self.routed_scaling_factor = float(args.routed_scaling_factor) + self.v_router = nn.Linear(dim, n, bias=False) + # The wire's attn_v_gate.bias. llama.cpp adds it to the scores for + # selection only, never to the logits. + self.selection_bias = mx.zeros((n,), dtype=mx.float32) + self.v_experts = SwitchLinear(dim, out, n, bias=False) + + def __call__(self, x: mx.array) -> mx.array: + inds, weights = _route( + self.v_router(x.astype(mx.float32)), self.selection_bias, + self.top_k, self.scoring_func, self.norm_topk_prob, + self.routed_scaling_factor) + + h = mx.expand_dims(x, (-2, -3)) + do_sort = inds.size >= 64 + idx, inv_order = inds, None + if do_sort: + h, idx, inv_order = _gather_sort(h, inds) + v = self.v_experts(h, idx, sorted_indices=do_sort) + if do_sort: + v = _scatter_unsort(v, inv_order, inds.shape) + v = nn.silu(v.squeeze(-2)) + return (v * weights[..., None].astype(v.dtype)).sum(axis=-2) + + +class Attention(nn.Module): + def __init__(self, args: ModelArgs, mova: bool): + 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 + + self.q_proj = nn.Linear(dim, self.n_heads * self.head_dim, bias=False) + self.k_proj = nn.Linear(dim, self.n_kv_heads * self.head_dim, bias=False) + if mova: + self.value_experts = ValueExperts(args) + else: + self.v_proj = nn.Linear( + dim, self.n_kv_heads * self.head_dim, bias=False) + self.o_proj = nn.Linear(self.n_heads * self.head_dim, dim, bias=False) + if args.attn_output_gate: + self.gate_proj = nn.Linear( + dim, self.n_heads * self.head_dim, bias=False) + if args.qk_norm: + self.q_norm = HeadRMSNorm( + self.n_heads, self.head_dim, args.rms_norm_eps) + self.k_norm = HeadRMSNorm( + self.n_kv_heads, self.head_dim, args.rms_norm_eps) + + self.rope = initialize_rope( + args.rope_dim or self.head_dim, + args.rope_theta, + False, + args.rope_scaling, + args.max_position_embeddings, + ) + + def __call__( + self, + x: mx.array, + mask: Optional[mx.array] = None, + cache: Optional[Any] = None, + ) -> mx.array: + B, L, _ = x.shape + + queries = self.q_proj(x).reshape(B, L, self.n_heads, -1) + keys = self.k_proj(x).reshape(B, L, self.n_kv_heads, -1) + if "value_experts" in self: + values = self.value_experts(x) + else: + values = self.v_proj(x) + values = values.reshape(B, L, self.n_kv_heads, -1) + + if "q_norm" in self: + queries = self.q_norm(queries) + keys = self.k_norm(keys) + + queries = queries.transpose(0, 2, 1, 3) + keys = keys.transpose(0, 2, 1, 3) + values = values.transpose(0, 2, 1, 3) + + offset = cache.offset if cache is not None else 0 + queries = self.rope(queries, offset=offset) + keys = self.rope(keys, offset=offset) + if cache is not None: + keys, values = cache.update_and_fetch(keys, values) + + output = scaled_dot_product_attention( + queries, keys, values, cache=cache, scale=self.scale, mask=mask + ) + output = output.transpose(0, 2, 1, 3).reshape(B, L, -1) + if "gate_proj" in self: + # log2(1 + 2^g), the softplus with beta ln 2. + g = self.gate_proj(x).astype(mx.float32) + gate = nn.softplus(g * _LN2) / _LN2 + output = output * gate.astype(output.dtype) + return self.o_proj(output) + + +class MLP(nn.Module): + """SwiGLU MLP: the dense layers and the shared expert.""" + + def __init__(self, dim: int, hidden: int): + super().__init__() + self.gate_proj = nn.Linear(dim, hidden, bias=False) + self.up_proj = nn.Linear(dim, hidden, bias=False) + self.down_proj = nn.Linear(hidden, dim, bias=False) + + def __call__(self, x: mx.array) -> mx.array: + return self.down_proj(nn.silu(self.gate_proj(x)) * self.up_proj(x)) + + +class K2HorizonMoEGate(nn.Module): + """FFN router: sigmoid scores, a selection-only correction bias, then + renormalized and scaled weights. + + A gate submodule returning ``(inds, weights)`` with an int ``top_k``, + the shape stream/moe_experts and stream/lookahead duck-type on. + """ + + def __init__(self, args: ModelArgs): + super().__init__() + self.top_k = int(args.num_experts_per_tok) + self.scoring_func = args.scoring_func + self.norm_topk_prob = bool(args.norm_topk_prob) + self.routed_scaling_factor = float(args.routed_scaling_factor) + # fp32 like the F32 wire tensors; llama.cpp routes in fp32. + self.weight = mx.zeros( + (args.num_experts, args.hidden_size), dtype=mx.float32) + self.e_score_correction_bias = mx.zeros( + (args.num_experts,), dtype=mx.float32) + + def _logits(self, x: mx.array) -> mx.array: + return x.astype(mx.float32) @ self.weight.T.astype(mx.float32) + + def __call__(self, x: mx.array): + return _route( + self._logits(x), self.e_score_correction_bias, self.top_k, + self.scoring_func, self.norm_topk_prob, self.routed_scaling_factor) + + def _kq_route_weights(self, x, inds): + # Route replay (stream/moe_routes): the weight branch of _route at + # caller-chosen ids. + logits = self._logits(x) + if self.scoring_func == "softmax": + scores = mx.softmax(logits, axis=-1, precise=True) + else: + scores = mx.sigmoid(logits) + return _route_weights( + scores, inds, self.norm_topk_prob, self.routed_scaling_factor) + + +class K2HorizonMoE(nn.Module): + def __init__(self, args: ModelArgs): + super().__init__() + self.gate = K2HorizonMoEGate(args) + self.switch_mlp = SwitchGLU( + args.hidden_size, args.moe_intermediate_size, args.num_experts) + if args.shared_expert_intermediate_size: + self.shared_experts = MLP( + args.hidden_size, args.shared_expert_intermediate_size) + + def __call__(self, x: mx.array) -> mx.array: + inds, scores = self.gate(x) + if getattr(self.switch_mlp, "_kq_mix_scores", False): + y = self.switch_mlp(x, inds, scores) + else: + y = self.switch_mlp(x, inds) + if y.ndim == scores.ndim + 1: + y = (y * scores[..., None].astype(y.dtype)).sum(-2) + if "shared_experts" in self: + y = y + self.shared_experts(x) + return y + + +class DecoderLayer(nn.Module): + def __init__(self, args: ModelArgs, layer_idx: int): + super().__init__() + dim = args.hidden_size + moe = args.num_experts > 0 and layer_idx >= args.first_k_dense_replace + mova = moe and args.num_value_experts > 0 + self.self_attn = Attention(args, mova) + self.mlp = K2HorizonMoE(args) if moe else MLP(dim, args.intermediate_size) + self.input_layernorm = GroupRMSNorm( + dim, args.rms_norm_eps, args.norm_groups) + self.post_attention_layernorm = GroupRMSNorm( + dim, args.rms_norm_eps, args.norm_groups) + + def __call__( + self, + x: mx.array, + mask: Optional[mx.array] = None, + cache: Optional[Any] = None, + ) -> mx.array: + 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.embed_tokens = nn.Embedding(args.vocab_size, args.hidden_size) + self.layers = [ + DecoderLayer(args, i) for i in range(args.num_hidden_layers)] + self.norm = GroupRMSNorm( + args.hidden_size, args.rms_norm_eps, args.norm_groups) + + def __call__( + self, + inputs: mx.array, + cache: Optional[Any] = None, + input_embeddings: Optional[mx.array] = None, + ): + h = self.embed_tokens(inputs) if input_embeddings is None \ + else input_embeddings + if cache is None: + cache = [None] * len(self.layers) + mask = create_attention_mask(h, cache[0]) + for layer, c in zip(self.layers, cache): + h = layer(h, 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: mx.array, + cache: Optional[Any] = None, + input_embeddings: Optional[mx.array] = None, + ): + out = self.model(inputs, cache, input_embeddings) + if self.args.tie_word_embeddings: + return self.model.embed_tokens.as_linear(out) + return self.lm_head(out) + + def sanitize(self, weights): + if self.args.tie_word_embeddings: + weights.pop("lm_head.weight", None) + return weights + + @property + def layers(self): + return self.model.layers + + def make_cache(self): + return [KVCache() for _ in self.layers] + + @property + def cast_predicate(self): + def predicate(path: str): + if path.endswith((".mlp.gate.weight", ".e_score_correction_bias", + ".selection_bias")): + return False + return True + + return predicate diff --git a/gmlx/models/owned.py b/gmlx/models/owned.py index 099220afe..29b0fb493 100644 --- a/gmlx/models/owned.py +++ b/gmlx/models/owned.py @@ -42,6 +42,9 @@ # relative and carries neither the independent hyper-connections nor # the sinks and the attention gate. "hy_v4": "gmlx.models.hy_v4.model", + # llama.cpp LLM_ARCH_K2_HORIZON; the mlx-vlm class is a plain Llama + # with none of the group norm, the output gate, MoE or MoVA. + "k2_horizon": "gmlx.models.k2_horizon.model", } # mlx-vlm model_type -> the gmlx module that stands for the model package diff --git a/tests/load/test_adapter_save.py b/tests/load/test_adapter_save.py index f931f4c12..7a9471137 100644 --- a/tests/load/test_adapter_save.py +++ b/tests/load/test_adapter_save.py @@ -108,6 +108,8 @@ def test_empty_modules_raises(tmp_path): ("model.layers.1.mlp.shared_mlp.gate_proj", "blk.1.ffn_gate_shexp.weight", IN)], "hyv4": [("model.layers.0.self_attn.q_a_proj", "blk.0.attn_q_a.weight", IN), ("model.layers.1.mlp.shared_experts.down_proj", "blk.1.ffn_down_shexp.weight", IN)], + "k2-horizon": [("model.layers.0.self_attn.gate_proj", "blk.0.attn_gate.weight", IN), + ("model.layers.1.mlp.shared_experts.down_proj", "blk.1.ffn_down_shexp.weight", IN)], "kimi-k3": [("model.layers.0.self_attn.b_proj", "blk.0.ssm_beta.weight", IN), ("model.layers.1.mlp.down_proj", "blk.1.ffn_down.weight", IN)], "minimax-m3": [("model.layers.0.self_attn.o_proj", "blk.0.attn_output.weight", IN), diff --git a/tests/load/test_config_synth.py b/tests/load/test_config_synth.py index 0872e4285..3c8f6dca1 100644 --- a/tests/load/test_config_synth.py +++ b/tests/load/test_config_synth.py @@ -2874,3 +2874,145 @@ def test_muse_spark_synth_refuses_softmax_gating_and_shared_experts(): m["muse-spark.expert_shared_count"] = 1 with pytest.raises(NotImplementedError, match="shared experts"): synthesize_config(m, tensor_shapes=_MUSE_SPARK_SHAPES) + + +# k2-horizon + + +def _k2_horizon_dense_meta() -> dict: + arch = "k2-horizon" + m = _base_meta(arch) + m[f"{arch}.attention.value_length"] = 16 + m[f"{arch}.attention.group_norm_groups"] = 4 + m[f"{arch}.rope.dimension_count"] = 16 + m[f"{arch}.rope.freq_base"] = 1e7 + # The dense 7B ships an expert count of 0 and a gating func anyway. + m[f"{arch}.expert_count"] = 0 + m[f"{arch}.expert_used_count"] = 0 + m[f"{arch}.expert_gating_func"] = 2 + return m + + +def _k2_horizon_mova_meta() -> dict: + arch = "k2-horizon" + m = _k2_horizon_dense_meta() + m[f"{arch}.block_count"] = 3 + m[f"{arch}.attention.group_norm_groups"] = 2 + m[f"{arch}.expert_count"] = 8 + m[f"{arch}.expert_used_count"] = 2 + m[f"{arch}.expert_feed_forward_length"] = 24 + m[f"{arch}.leading_dense_block_count"] = 1 + m[f"{arch}.moe_every_n_layers"] = 1 + m[f"{arch}.expert_shared_count"] = 1 + m[f"{arch}.expert_shared_feed_forward_length"] = 24 + m[f"{arch}.expert_weights_norm"] = True + m[f"{arch}.expert_weights_scale"] = 2.5 + m[f"{arch}.attention.value_expert_count"] = 6 + m[f"{arch}.attention.value_expert_used_count"] = 2 + return m + + +_K2_HORIZON_DENSE_SHAPES = {"output.weight": [64, VOCAB]} + +_K2_HORIZON_MOVA_SHAPES = { + "output.weight": [64, VOCAB], + "blk.0.attn_gate.weight": [64, 64], + "blk.1.attn_v_exps.weight": [64, 32, 6], + "blk.1.ffn_gate_exps.weight": [64, 24, 8], +} + + +def test_k2_horizon_dense_synth_fields(): + c = synthesize_config(_k2_horizon_dense_meta(), + tensor_shapes=_K2_HORIZON_DENSE_SHAPES) + assert c["model_type"] == "k2_horizon" + assert c["norm_groups"] == 4 + assert c["head_dim"] == 16 + assert c["rope_theta"] == 1e7 + assert "rope_dim" not in c + assert c["intermediate_size"] == 128 + assert c["attn_output_gate"] is False + assert c["qk_norm"] is False + assert "num_experts" not in c and "num_value_experts" not in c + assert "scoring_func" not in c + assert not c["tie_word_embeddings"] + + +def test_k2_horizon_mova_synth_fields(): + c = synthesize_config(_k2_horizon_mova_meta(), + tensor_shapes=_K2_HORIZON_MOVA_SHAPES) + assert c["norm_groups"] == 2 + assert c["attn_output_gate"] is True + assert (c["num_experts"], c["num_experts_per_tok"]) == (8, 2) + assert c["moe_intermediate_size"] == 24 + assert c["shared_expert_intermediate_size"] == 24 + assert c["first_k_dense_replace"] == 1 + assert c["routed_scaling_factor"] == 2.5 + assert c["norm_topk_prob"] is True + assert (c["num_value_experts"], c["num_value_experts_per_tok"]) == (6, 2) + assert c["scoring_func"] == "sigmoid" + + +def test_k2_horizon_synth_applies_llama_cpp_defaults(): + arch = "k2-horizon" + m = _k2_horizon_mova_meta() + for key in ("attention.group_norm_groups", "expert_gating_func", + "expert_shared_feed_forward_length", "expert_weights_scale", + "expert_weights_norm", "leading_dense_block_count"): + del m[f"{arch}.{key}"] + c = synthesize_config(m, tensor_shapes=_K2_HORIZON_MOVA_SHAPES) + assert c["norm_groups"] == 1 + assert c["scoring_func"] == "sigmoid" + # A shared width of 0 means n_ff_exp * n_shared. + assert c["shared_expert_intermediate_size"] == 24 + assert c["routed_scaling_factor"] == 0.0 + assert c["norm_topk_prob"] is False + assert c["first_k_dense_replace"] == 0 + + +def test_k2_horizon_synth_reads_partial_rope_and_qk_norm(): + m = _k2_horizon_dense_meta() + m["k2-horizon.rope.dimension_count"] = 8 + shapes = dict(_K2_HORIZON_DENSE_SHAPES) + shapes["blk.0.attn_q_norm.weight"] = [64] + c = synthesize_config(m, tensor_shapes=shapes) + assert c["rope_dim"] == 8 + assert c["qk_norm"] is True + + +@pytest.mark.parametrize("key", [ + "expert_feed_forward_length", "expert_used_count", + "attention.key_length", +]) +def test_k2_horizon_required_keys_fail_loud(key): + m = _k2_horizon_mova_meta() + del m[f"k2-horizon.{key}"] + with pytest.raises(ValueError, match=f"k2-horizon.{key}"): + synthesize_config(m, tensor_shapes=_K2_HORIZON_MOVA_SHAPES) + + +def test_k2_horizon_synth_refuses_bad_value_routing(): + m = _k2_horizon_mova_meta() + m["k2-horizon.attention.value_expert_used_count"] = 7 + with pytest.raises(ValueError, match="value experts used"): + synthesize_config(m, tensor_shapes=_K2_HORIZON_MOVA_SHAPES) + m = _k2_horizon_mova_meta() + m["k2-horizon.expert_gating_func"] = 3 + with pytest.raises(NotImplementedError, match="gating_func 3"): + synthesize_config(m, tensor_shapes=_K2_HORIZON_MOVA_SHAPES) + + +@pytest.mark.parametrize("meta_fn, shapes", [ + (_k2_horizon_dense_meta, _K2_HORIZON_DENSE_SHAPES), + (_k2_horizon_mova_meta, _K2_HORIZON_MOVA_SHAPES), +]) +def test_k2_horizon_synth_instantiates(meta_fn, shapes): + from gmlx.models.k2_horizon.model import Model, ModelArgs + + c = synthesize_config(meta_fn(), tensor_shapes=shapes) + model, args = _get_classes(c) + assert model is Model and args is ModelArgs + m = Model(ModelArgs.from_dict(c)) + mx.eval(m.parameters()) + logits = m(mx.array([[1, 2, 3]])) + assert logits.shape == (1, 3, VOCAB) diff --git a/tests/load/test_remap.py b/tests/load/test_remap.py index 2c8a24702..772975161 100644 --- a/tests/load/test_remap.py +++ b/tests/load/test_remap.py @@ -1622,3 +1622,57 @@ def test_muse_spark_claims_every_row_before_the_canonical_map(): def test_muse_spark_unknown_tensor_fails(): assert d("muse-spark", "blk.0.not_a_muse_spark_tensor.weight").kind == FAIL + + +# k2-horizon: every tensor is claimed by the K2_HORIZON block, all passthrough. +_K2_HORIZON_ROWS = ( + ("token_embd.weight", "model.embed_tokens.weight"), + ("output_norm.weight", "model.norm.weight"), + ("output.weight", "lm_head.weight"), + ("blk.0.attn_norm.weight", "model.layers.0.input_layernorm.weight"), + ("blk.0.attn_q.weight", "model.layers.0.self_attn.q_proj.weight"), + ("blk.0.attn_k.weight", "model.layers.0.self_attn.k_proj.weight"), + ("blk.0.attn_v.weight", "model.layers.0.self_attn.v_proj.weight"), + ("blk.0.attn_output.weight", "model.layers.0.self_attn.o_proj.weight"), + ("blk.0.attn_gate.weight", "model.layers.0.self_attn.gate_proj.weight"), + ("blk.0.attn_q_norm.weight", "model.layers.0.self_attn.q_norm.weight"), + ("blk.0.attn_k_norm.weight", "model.layers.0.self_attn.k_norm.weight"), + ("blk.0.ffn_norm.weight", "model.layers.0.post_attention_layernorm.weight"), + ("blk.0.ffn_gate.weight", "model.layers.0.mlp.gate_proj.weight"), + ("blk.0.ffn_up.weight", "model.layers.0.mlp.up_proj.weight"), + ("blk.0.ffn_down.weight", "model.layers.0.mlp.down_proj.weight"), + ("blk.3.attn_v_gate.weight", + "model.layers.3.self_attn.value_experts.v_router.weight"), + ("blk.3.attn_v_gate.bias", + "model.layers.3.self_attn.value_experts.selection_bias"), + ("blk.3.attn_v_exps.weight", + "model.layers.3.self_attn.value_experts.v_experts.weight"), + ("blk.3.ffn_gate_inp.weight", "model.layers.3.mlp.gate.weight"), + ("blk.3.exp_probs_b.bias", "model.layers.3.mlp.gate.e_score_correction_bias"), + ("blk.3.ffn_gate_exps.weight", "model.layers.3.mlp.switch_mlp.gate_proj.weight"), + ("blk.3.ffn_up_exps.weight", "model.layers.3.mlp.switch_mlp.up_proj.weight"), + ("blk.3.ffn_down_exps.weight", "model.layers.3.mlp.switch_mlp.down_proj.weight"), + ("blk.3.ffn_gate_shexp.weight", "model.layers.3.mlp.shared_experts.gate_proj.weight"), + ("blk.3.ffn_up_shexp.weight", "model.layers.3.mlp.shared_experts.up_proj.weight"), + ("blk.3.ffn_down_shexp.weight", "model.layers.3.mlp.shared_experts.down_proj.weight"), +) + + +@pytest.mark.parametrize("gguf, hf", _K2_HORIZON_ROWS) +def test_k2_horizon_rows(gguf, hf): + r = d("k2-horizon", gguf) + assert r.kind == MAP and r.hf_name == hf, (gguf, r.reason) + # NeoX rope: Q/K pass through un-permuted. + assert r.transform == "passthrough" + + +def test_k2_horizon_claims_every_row_before_the_canonical_map(): + from gmlx.load.remap import ARCH_PRIORITY_OVERRIDES + + claimed = ARCH_PRIORITY_OVERRIDES["K2_HORIZON"] + for gguf, _hf in _K2_HORIZON_ROWS: + assert any(p.match(gguf) for p, _t, _x in claimed), gguf + + +def test_k2_horizon_unknown_tensor_fails(): + assert d("k2-horizon", "blk.0.not_a_k2_horizon_tensor.weight").kind == FAIL diff --git a/tests/load/test_rope_swap.py b/tests/load/test_rope_swap.py index d9a97794d..a8d43eb4a 100644 --- a/tests/load/test_rope_swap.py +++ b/tests/load/test_rope_swap.py @@ -288,6 +288,7 @@ def _later_arches(): "hyv4": (t._hyv4_meta, t._HYV4_SHAPES), "muse-glimmer": (t._muse_glimmer_meta, t._MUSE_GLIMMER_SHAPES), "muse-spark": (t._muse_spark_meta, t._MUSE_SPARK_SHAPES), + "k2-horizon": (t._k2_horizon_mova_meta, t._K2_HORIZON_MOVA_SHAPES), "minimax-m3": (t._minimax_m3_meta, t._MINIMAX_M3_SHAPES), "falcon-h1": (t._falcon_h1_meta, t._FALCON_H1_SHAPES), "granitehybrid": (t._granitehybrid_meta, {}), @@ -296,8 +297,8 @@ def _later_arches(): @pytest.mark.parametrize("arch", ["hy_v3", "hyv4", "muse-glimmer", - "muse-spark", "minimax-m3", "falcon-h1", - "granitehybrid", "ernie4_5-moe"]) + "muse-spark", "k2-horizon", "minimax-m3", + "falcon-h1", "granitehybrid", "ernie4_5-moe"]) def test_later_class_built_arches_take_the_options(arch): assert arch in rs.SWAP_ARCHES and arch in rs.WIRED_ARCHES meta_fn, shapes = _later_arches()[arch] diff --git a/tests/load/test_tokenizer.py b/tests/load/test_tokenizer.py index 3178140d6..568adbc58 100644 --- a/tests/load/test_tokenizer.py +++ b/tests/load/test_tokenizer.py @@ -821,3 +821,64 @@ def test_ds4_marker_leaves_reachable_tokens_alone(): assert tok.encode("He", add_special_tokens=False) == [ meta["tokenizer.ggml.tokens"].index("He")] assert "He" not in tok.get_added_vocab() + + +def test_k2_horizon_pre_full_pattern_override(): + # The llama3 split with marks, ZWNJ and ZWJ inside letter runs, as the + # HF tokenizer.json and llama.cpp's K2_HORIZON pre spell it. + from tokenizers import Regex + from gmlx.load.tokenizer import K2_HORIZON_PATTERN, _bytelevel_split_patterns + + assert _bytelevel_split_patterns("k2-horizon") == (K2_HORIZON_PATTERN,) + split = pre_tokenizers.Split( + Regex(K2_HORIZON_PATTERN), behavior="isolated", invert=False) + zwnj = "Amy\u200ckhaham" + zwj = "a\u200db" + virama = "\u0915\u094d\u0937" # ka, virama (a mark), ssa + acute = "e\u0301clair" # decomposed e + combining acute + text = f"{zwnj} {zwj} {virama} {acute} 1234567 DON'T" + pieces = [p for p, _ in split.pre_tokenize_str(text)] + assert "".join(pieces) == text + assert zwnj in pieces + assert " " + zwj in pieces + assert " " + virama in pieces + assert " " + acute in pieces + assert ["123", "456", "7"] == [p for p in pieces if p.isdigit()] + assert "'T" in pieces # case-insensitive + + +def test_k2_horizon_pre_keeps_nfc(): + # The HF tokenizer.json normalizes NFC, so the pre keeps the default. + from gmlx.load.tokenizer import _build_bytelevel_bpe + + tok = _build_bytelevel_bpe(["a", "b"], [], pre_id="k2-horizon") + assert tok.normalizer is not None + assert tok.normalizer.normalize_str("e\u0301") == "\u00e9" + + +def test_k2_horizon_turn_end_is_a_stop_by_name(): + # The template raises on the turn-end heuristic's sentinel render (an + # assistant turn with no thinking field), so <|ifm|im_end|> reaches the + # stop set only through the name list, as in llama.cpp. + from pathlib import Path + + tmpl = (Path(__file__).parents[1] / "fixtures" + / "k2_horizon_template.jinja").read_text() + specials = ["<|ifm|begin_of_text|>", "<|ifm|endoftext|>", + "<|ifm|im_start|>", "<|ifm|im_end|>"] + toks = specials + _ALPHABET + _MERGED + meta = { + "general.architecture": "k2-horizon", + "tokenizer.ggml.model": "gpt2", + "tokenizer.ggml.pre": "k2-horizon", + "tokenizer.ggml.tokens": toks, + "tokenizer.ggml.merges": _MERGES, + "tokenizer.ggml.token_type": [3] * len(specials) + + [1] * (len(toks) - len(specials)), + "tokenizer.ggml.bos_token_id": 0, + "tokenizer.ggml.eos_token_id": 1, + "tokenizer.chat_template": tmpl, + } + tok = load_tokenizer_from_gguf(meta, "k2-horizon") + assert tok._gguf_eos_token_ids[0] == 1 + assert 3 in tok._gguf_eos_token_ids diff --git a/tests/models/test_k2_horizon.py b/tests/models/test_k2_horizon.py new file mode 100644 index 000000000..3bf18b540 --- /dev/null +++ b/tests/models/test_k2_horizon.py @@ -0,0 +1,465 @@ +# SPDX-License-Identifier: MIT +"""K2 Horizon model against a NumPy reference of llama.cpp's graph +(src/models/k2-horizon.cpp at 7182ef71e): group RMSNorm, NeoX rope, the +log2(1 + 2^g) output gate, MoVA value routing and the sigmoid MoE with a +shared expert.""" + +from __future__ import annotations + +import math + +import numpy as np +import pytest + +import mlx.core as mx +from mlx.utils import tree_flatten, tree_unflatten + +from gmlx.models.k2_horizon.model import ( + K2HorizonMoE, + K2HorizonMoEGate, + Model, + ModelArgs, + ValueExperts, +) + + +def _tiny_args(mova: bool = True, **over) -> ModelArgs: + kw = dict( + model_type="k2_horizon", + hidden_size=64, + num_hidden_layers=3, + intermediate_size=96, + num_attention_heads=4, + num_key_value_heads=2, + head_dim=16, + vocab_size=40, + rms_norm_eps=1e-6, + norm_groups=4, + rope_theta=1e7, + max_position_embeddings=4096, + ) + if mova: + kw.update( + norm_groups=2, + attn_output_gate=True, + num_experts=8, + num_experts_per_tok=3, + moe_intermediate_size=24, + shared_expert_intermediate_size=24, + first_k_dense_replace=1, + num_value_experts=6, + num_value_experts_per_tok=2, + scoring_func="sigmoid", + norm_topk_prob=True, + routed_scaling_factor=2.5, + ) + kw.update(over) + return ModelArgs(**kw) + + +def _randomize(module, seed: int): + rng = np.random.default_rng(seed) + upd = [] + for k, v in tree_flatten(module.parameters()): + shape = v.shape + if "norm" in k or k.endswith("layernorm.weight"): + a = 1.0 + 0.3 * rng.standard_normal(shape) + elif k.endswith(("e_score_correction_bias", "selection_bias")): + a = 0.3 * rng.standard_normal(shape) + else: + a = rng.standard_normal(shape) / np.sqrt(shape[-1]) + upd.append((k, mx.array(a.astype(np.float32)))) + module.update(tree_unflatten(upd)) + mx.eval(module.parameters()) + return module + + +def _np(x) -> np.ndarray: + return np.array(x.astype(mx.float32), dtype=np.float64) + + +def _silu(x): + return x / (1.0 + np.exp(-x)) + + +def _group_rms(x, w, groups, eps): + shape = x.shape + y = x.reshape(*shape[:-1], groups, shape[-1] // groups) + y = y / np.sqrt((y * y).mean(-1, keepdims=True) + eps) + return y.reshape(shape) * w + + +def _rope_neox(x, theta, dims=None): + # x: [T, H, D]; NeoX pairs dim i with dim i + dims/2. + T, _, D = x.shape + dims = dims or D + half = dims // 2 + freqs = theta ** (-np.arange(half) * 2.0 / dims) + ang = np.arange(T)[:, None] * freqs[None, :] + cos, sin = np.cos(ang)[:, None, :], np.sin(ang)[:, None, :] + x1, x2 = x[..., :half], x[..., half:dims] + out = x.copy() + out[..., :half] = x1 * cos - x2 * sin + out[..., half:dims] = x2 * cos + x1 * sin + return out + + +def _route(logits, bias, k, norm, scale): + scores = 1.0 / (1.0 + np.exp(-logits)) + choice = scores + bias + inds = np.argsort(-choice, axis=-1, kind="stable")[..., :k] + w = np.take_along_axis(scores, inds, axis=-1) + if norm: + w = w / np.maximum(w.sum(-1, keepdims=True), 6.103515625e-5) + if scale not in (0.0, 1.0): + w = w * scale + return inds, w + + +def _ref_values(ve, x, args): + W = _np(ve.v_router.weight) + inds, w = _route(x @ W.T, _np(ve.selection_bias), + args.num_value_experts_per_tok, args.norm_topk_prob, + args.routed_scaling_factor) + E = _np(ve.v_experts.weight) # [n, out, in] + out = np.zeros((x.shape[0], E.shape[1])) + for t in range(x.shape[0]): + for j in range(inds.shape[1]): + out[t] += w[t, j] * _silu(E[inds[t, j]] @ x[t]) + return out + + +def _ref_mlp(m, x): + g = x @ _np(m.gate_proj.weight).T + u = x @ _np(m.up_proj.weight).T + return (_silu(g) * u) @ _np(m.down_proj.weight).T + + +def _ref_moe(moe, x, args): + inds, w = _route(x @ _np(moe.gate.weight).T, + _np(moe.gate.e_score_correction_bias), + args.num_experts_per_tok, args.norm_topk_prob, + args.routed_scaling_factor) + sw = moe.switch_mlp + G, U, D = (_np(sw.gate_proj.weight), _np(sw.up_proj.weight), + _np(sw.down_proj.weight)) + out = np.zeros_like(x) + for t in range(x.shape[0]): + for j in range(inds.shape[1]): + e = inds[t, j] + h = _silu(G[e] @ x[t]) * (U[e] @ x[t]) + out[t] += w[t, j] * (D[e] @ h) + if "shared_experts" in moe: + out = out + _ref_mlp(moe.shared_experts, x) + return out + + +def _ref_attn(attn, x, args): + T = x.shape[0] + H, KV, D = args.num_attention_heads, args.num_key_value_heads, args.head_dim + q = (x @ _np(attn.q_proj.weight).T).reshape(T, H, D) + k = (x @ _np(attn.k_proj.weight).T).reshape(T, KV, D) + if "value_experts" in attn: + v = _ref_values(attn.value_experts, x, args).reshape(T, KV, D) + else: + v = (x @ _np(attn.v_proj.weight).T).reshape(T, KV, D) + if "q_norm" in attn: + qw = _np(attn.q_norm.weight).reshape(H, D) + kw = _np(attn.k_norm.weight).reshape(KV, D) + q = q / np.sqrt((q * q).mean(-1, keepdims=True) + args.rms_norm_eps) * qw + k = k / np.sqrt((k * k).mean(-1, keepdims=True) + args.rms_norm_eps) * kw + q = _rope_neox(q, args.rope_theta, args.rope_dim) + k = _rope_neox(k, args.rope_theta, args.rope_dim) + rep = H // KV + k, v = np.repeat(k, rep, axis=1), np.repeat(v, rep, axis=1) + s = np.einsum("thd,shd->hts", q, k) / math.sqrt(D) + s = s + np.triu(np.full((T, T), -np.inf), 1)[None] + p = np.exp(s - s.max(-1, keepdims=True)) + p = p / p.sum(-1, keepdims=True) + o = np.einsum("hts,shd->thd", p, v).reshape(T, H * D) + if "gate_proj" in attn: + g = x @ _np(attn.gate_proj.weight).T + o = o * np.log2(1.0 + np.exp2(g)) + return o @ _np(attn.o_proj.weight).T + + +def _ref_model(model, tokens, args): + eps, groups = args.rms_norm_eps, args.norm_groups + h = _np(model.model.embed_tokens.weight)[tokens] + for layer in model.layers: + a = _group_rms(h, _np(layer.input_layernorm.weight), groups, eps) + h = h + _ref_attn(layer.self_attn, a, args) + f = _group_rms(h, _np(layer.post_attention_layernorm.weight), groups, eps) + if isinstance(layer.mlp, K2HorizonMoE): + h = h + _ref_moe(layer.mlp, f, args) + else: + h = h + _ref_mlp(layer.mlp, f) + h = _group_rms(h, _np(model.model.norm.weight), groups, eps) + return h @ _np(model.lm_head.weight).T + + +def _rel(a, b): + return float(np.abs(a - b).max() / (np.abs(b).max() + 1e-12)) + + +@pytest.mark.parametrize("args", [ + _tiny_args(mova=False), + _tiny_args(mova=True), + _tiny_args(mova=True, qk_norm=True), + _tiny_args(mova=False, rope_dim=8), + _tiny_args(mova=True, norm_topk_prob=False, routed_scaling_factor=1.0), +], ids=["dense", "mova", "mova-qk-norm", "dense-partial-rope", "mova-no-norm"]) +def test_full_model_matches_reference(args): + model = _randomize(Model(args), seed=3) + tokens = np.array([1, 5, 9, 2, 33, 7, 11]) + got = _np(model(mx.array(tokens[None]))[0]) + ref = _ref_model(model, tokens, args) + assert _rel(got, ref) < 2e-4 + + +def test_layer_kinds_follow_the_dense_lead(): + model = Model(_tiny_args()) + kinds = [(type(ly.mlp).__name__, "value_experts" in ly.self_attn) + for ly in model.layers] + assert kinds == [("MLP", False), ("K2HorizonMoE", True), + ("K2HorizonMoE", True)] + assert all("gate_proj" in ly.self_attn for ly in model.layers) + + +def test_cached_decode_matches_full_forward(): + args = _tiny_args() + model = _randomize(Model(args), seed=5) + tokens = mx.array([[3, 1, 4, 1, 5, 9, 2, 6]]) + full = model(tokens) + cache = model.make_cache() + steps = [model(tokens[:, :5], cache=cache)[:, -1]] + for i in range(5, 8): + steps.append(model(tokens[:, i:i + 1], cache=cache)[:, -1]) + got = mx.stack(steps[1:], axis=1) + assert _rel(_np(got), _np(full[:, 5:])) < 1e-4 + + +def test_value_bias_moves_selection_but_not_weights(): + args = _tiny_args() + ve = _randomize(ValueExperts(args), seed=7) + x = mx.array(np.random.default_rng(0).standard_normal((1, 4, 64)), + dtype=mx.float32) + base = ve(x) + # A large bias on expert 0 forces it into every selection; its weight + # still comes from the unbiased score. + bias = np.zeros(args.num_value_experts, dtype=np.float32) + bias[0] = 10.0 + ve.selection_bias = mx.array(bias) + forced = _np(ve(x)) + ref = _ref_values(ve, _np(x[0]), args) + assert _rel(forced[0], ref) < 1e-5 + assert _rel(forced, _np(base)) > 1e-3 + + +def test_route_weight_sum_is_clamped_like_llama_cpp(): + args = _tiny_args(num_value_experts_per_tok=1) + ve = ValueExperts(args) + ve.v_router.weight = mx.full(ve.v_router.weight.shape, -30.0) + ve.selection_bias = mx.zeros((args.num_value_experts,)) + x = mx.ones((1, 1, 64)) + from gmlx.models.k2_horizon.model import _route + _, w = _route(ve.v_router(x), ve.selection_bias, 1, "sigmoid", True, 1.0) + # sigmoid(-1920) underflows to 0, and the clamped sum keeps the + # weight finite instead of 0/0. + assert np.isfinite(_np(w)).all() + + +def test_gate_route_weights_match_the_forward(): + args = _tiny_args() + gate = _randomize(K2HorizonMoEGate(args), seed=9) + x = mx.array(np.random.default_rng(1).standard_normal((2, 3, 64)), + dtype=mx.float32) + inds, w = gate(x) + assert _rel(_np(gate._kq_route_weights(x, inds)), _np(w)) < 1e-6 + + +def test_route_replay_supports_every_moe_layer(): + from gmlx.stream.moe_experts import _gate_submodule, _hook_target + from gmlx.stream.moe_routes import moe_layers, replay_unsupported + + model = Model(_tiny_args()) + assert replay_unsupported(model) == [] + assert moe_layers(model) == [1, 2] + for layer in model.layers[1:]: + assert _gate_submodule(layer.mlp) is layer.mlp.gate + assert _hook_target(layer.mlp) is layer.mlp.gate + assert _gate_submodule(layer.self_attn) is None + + +def test_cast_predicate_matches_the_loader_keep_list(): + from gmlx.load.loader import _FP32_KEEP_BY_MODEL_TYPE + + keep = _FP32_KEEP_BY_MODEL_TYPE["k2_horizon"] + model = Model(_tiny_args()) + pred = model.cast_predicate + pinned = set() + for path, _ in tree_flatten(model.parameters()): + assert pred(path) == (not any(s in path for s in keep)), path + if not pred(path): + pinned.add(path.split(".", 3)[-1]) + # The dense layer's mlp.gate_proj and the quantized value router stay + # on the activation dtype. + assert pinned == { + "mlp.gate.weight", + "mlp.gate.e_score_correction_bias", + "self_attn.value_experts.selection_bias", + } + + +_PER_LAYER_DENSE = [ + "attn_norm.weight", "attn_q.weight", "attn_k.weight", "attn_v.weight", + "attn_output.weight", "attn_gate.weight", "ffn_norm.weight", + "ffn_gate.weight", "ffn_up.weight", "ffn_down.weight", +] +_PER_LAYER_MOVA = [ + "attn_norm.weight", "attn_q.weight", "attn_k.weight", + "attn_v_gate.weight", "attn_v_gate.bias", "attn_v_exps.weight", + "attn_output.weight", "attn_gate.weight", "ffn_norm.weight", + "ffn_gate_inp.weight", "exp_probs_b.bias", "ffn_gate_exps.weight", + "ffn_up_exps.weight", "ffn_down_exps.weight", "ffn_gate_shexp.weight", + "ffn_up_shexp.weight", "ffn_down_shexp.weight", +] + + +@pytest.mark.parametrize("qk_norm", [False, True]) +def test_remap_covers_every_wire_tensor_onto_real_params(qk_norm): + """Every tensor llama.cpp's k2-horizon loader creates maps onto a + parameter of the model, and every parameter is loaded.""" + from gmlx.load.remap import RemapDecision, parse_gguf_name + + args = _tiny_args(qk_norm=qk_norm) + params = {k for k, _ in tree_flatten(Model(args).parameters())} + extra = ["attn_q_norm.weight", "attn_k_norm.weight"] if qk_norm else [] + names = ["token_embd.weight", "output_norm.weight", "output.weight"] + for i in range(args.num_hidden_layers): + per = _PER_LAYER_DENSE if i < args.first_k_dense_replace \ + else _PER_LAYER_MOVA + names += [f"blk.{i}.{t}" for t in per + extra] + mapped = set() + for name in names: + dec = parse_gguf_name("k2-horizon", name) + assert dec.kind == RemapDecision.KIND_MAP, (name, dec.reason) + assert dec.transform == "passthrough", name + assert dec.hf_name in params, (name, dec.hf_name) + mapped.add(dec.hf_name) + assert mapped == params + + +def test_unknown_tensor_is_not_claimed(): + from gmlx.load.remap import RemapDecision, parse_gguf_name + + dec = parse_gguf_name("k2-horizon", "blk.0.attn_sinks.weight") + assert dec.kind != RemapDecision.KIND_MAP + + +def test_synth_config_builds_this_model(): + from gmlx.load.config_synth import synthesize_config + from test_config_synth import _K2_HORIZON_MOVA_SHAPES, _k2_horizon_mova_meta + + c = synthesize_config(_k2_horizon_mova_meta(), + tensor_shapes=_K2_HORIZON_MOVA_SHAPES) + model = Model(ModelArgs.from_dict(c)) + ve = model.layers[1].self_attn.value_experts + assert ve.v_experts.weight.shape == (6, 32, 64) + assert ve.v_router.weight.shape == (6, 64) + assert model.layers[1].mlp.switch_mlp.gate_proj.weight.shape == (8, 24, 64) + assert model.layers[1].mlp.shared_experts.gate_proj.weight.shape == (24, 64) + + +def _loader_shaped(): + from tune.tiny_train_archs import build + + model, config, _ = build("k2_horizon") + return model, config + + +def test_loader_shaped_dtypes_and_fused_moe(monkeypatch): + """As the loader leaves it: the FFN router and both selection biases are + fp32, the value router is a K-quant linear (Q4_K on the wire), the + value experts gather K-quant, and every MoE layer's expert container + gets a bf16 input the fused decode gate accepts.""" + from mlx_kquant.nn import KQuantLinear, KQuantSwitchLinear + + import gmlx.load.modules as modules + from gmlx.load.dtypes import activation_dtype + + if activation_dtype() != mx.bfloat16: + pytest.skip("the loader picks float16 activations on this device") + model, _ = _loader_shaped() + moe = model.layers[1] + ve = moe.self_attn.value_experts + assert isinstance(ve.v_router, KQuantLinear) + assert isinstance(ve.v_experts, KQuantSwitchLinear) + assert ve.selection_bias.dtype == mx.float32 + assert moe.mlp.gate.weight.dtype == mx.float32 + assert moe.mlp.gate.e_score_correction_bias.dtype == mx.float32 + assert model.layers[0].mlp.gate_proj.weight.dtype != mx.float32 + + glu_cls = type(moe.mlp.switch_mlp) + if not hasattr(glu_cls, "_fused_ok"): + pytest.skip("fused MoE glu not installable in this build") + monkeypatch.setattr(modules, "_kq_fused_device_ok", lambda *m: True) + orig = glu_cls._fused_ok + verdicts = [] + + def record(self, x, indices): + verdicts.append((x.dtype, orig(self, x, indices))) + return False # the stock path runs on the CPU + + monkeypatch.setattr(glu_cls, "_fused_ok", record) + cache = model.make_cache() + model(mx.array([[1, 2, 3]]), cache=cache) + verdicts.clear() + out = model(mx.array([[4]]), cache=cache) + mx.eval(out) + assert out.dtype == mx.bfloat16 + assert len(verdicts) == 2 + assert all(v == (mx.bfloat16, True) for v in verdicts), verdicts + + +def test_expert_streaming_wraps_the_ffn_experts_only(monkeypatch): + """Streaming wraps each MoE layer's SwitchGLU; the MoVA value experts + stay resident, and the top-k override reaches the FFN gate only.""" + from mlx_kquant.nn import KQuantSwitchLinear + + from gmlx.stream.expert_streaming import ( + install_expert_streaming, + install_moe_experts_override, + ) + + model, _ = _loader_shaped() + monkeypatch.setattr( + mx, "device_info", lambda: {"max_recommended_working_set_size": 1024}) + tokens = mx.array([[1, 2, 3, 4, 5]]) + ref = _np(model(tokens)) + install_expert_streaming(model) + for layer in model.layers[1:]: + assert type(layer.mlp.switch_mlp).__name__.endswith("_CPUOffload") + assert type(layer.self_attn.value_experts.v_experts) is KQuantSwitchLinear + # The streamed experts run on the CPU stream, about one bf16 ulp off. + assert _rel(_np(model(tokens)), ref) < 3e-2 + + assert install_moe_experts_override(model, 1) == 2 + for layer in model.layers[1:]: + assert layer.mlp.gate.top_k == 1 + assert layer.self_attn.value_experts.top_k == 2 + + +def test_lookahead_router_replica_matches_the_gate(): + from gmlx.stream.lookahead import _router_fn_for + + args = _tiny_args() + moe = _randomize(K2HorizonMoE(args), seed=13) + select = _router_fn_for(moe) + assert select is not None + x = mx.array(np.random.default_rng(2).standard_normal((1, 1, 64)), + dtype=mx.float32) + inds, w = moe.gate(x) + ids, mass = select(x) + assert sorted(_np(ids).ravel().tolist()) == sorted(_np(inds).ravel().tolist()) + assert (np.diff(_np(mass).ravel()) <= 0).all() diff --git a/tests/tune/test_training_step.py b/tests/tune/test_training_step.py index 5a50bc54c..fae2752d0 100644 --- a/tests/tune/test_training_step.py +++ b/tests/tune/test_training_step.py @@ -64,7 +64,8 @@ # amplifies most: a one-ulp bump of its first layer's output moves the # logits by about 3%. PARITY_TOL = { - "qwen3_moe": 3e-2, "muse_glimmer": 3e-2, "muse_spark": 3e-2, "hy_v3": 3e-2, "glm5_next": 4e-2, + "qwen3_moe": 3e-2, "muse_glimmer": 3e-2, "muse_spark": 3e-2, "k2_horizon": 3e-2, + "hy_v3": 3e-2, "glm5_next": 4e-2, "kimi_k3": 1e-1, "hy_v4": 4e-2, "minimax_m3": 3e-2, "qwen4_exp": 2e-2, "qwen3_next": 4e-2, "qwen3_5": 3e-2, "qwen3_5_moe": 4e-2, "qwen3_5_hadamard": 4e-2, } diff --git a/tests/tune/tiny_train_archs.py b/tests/tune/tiny_train_archs.py index 7d4a5771f..ba8cc6879 100644 --- a/tests/tune/tiny_train_archs.py +++ b/tests/tune/tiny_train_archs.py @@ -60,6 +60,14 @@ def _qwen35(**over): output_multiplier=0.1767766922712326, post_norm_eps=1e-8, rms_norm_eps=1e-5, rope_parameters={"rope_theta": 5e5, "rope_type": "default"}, rope_theta=5e5, max_position_embeddings=4096, tie_word_embeddings=False), + "k2_horizon": dict( + model_type="k2_horizon", hidden_size=256, intermediate_size=512, num_hidden_layers=3, + num_attention_heads=4, num_key_value_heads=2, head_dim=64, vocab_size=512, + rms_norm_eps=1e-6, norm_groups=2, rope_theta=1e7, max_position_embeddings=4096, + attn_output_gate=True, num_experts=8, num_experts_per_tok=4, moe_intermediate_size=256, + shared_expert_intermediate_size=256, first_k_dense_replace=1, num_value_experts=6, + num_value_experts_per_tok=2, scoring_func="sigmoid", norm_topk_prob=True, + routed_scaling_factor=2.5, tie_word_embeddings=False), "hy_v3": dict( model_type="hy_v3", hidden_size=256, intermediate_size=512, expert_hidden_dim=256, num_hidden_layers=3, first_k_dense_replace=1, num_attention_heads=4, @@ -159,6 +167,7 @@ class Case: "qwen3_moe": Case("qwen3moe"), "muse_glimmer": Case("muse-glimmer"), "muse_spark": Case("muse-spark"), + "k2_horizon": Case("k2-horizon"), "hy_v3": Case("hy_v3", (r"router\.gate$",)), "glm5_next": Case("glm5next"), "kimi_k3": Case("kimi-k3"), From a0d3e8aee70ef7f682c2d0f607724cdafa1c4159 Mon Sep 17 00:00:00 2001 From: Asher Feldman <59994+asher@users.noreply.github.com> Date: Sat, 17 Oct 2026 18:08:23 -0700 Subject: [PATCH 03/10] feat(k2-horizon): tool-call parser and history normalizer for the ifm tag format --- gmlx/load/loader.py | 4 + gmlx/load/tokenizer.py | 4 + gmlx/models/k2_horizon/tools.py | 199 ++++++++++++++++++ tests/models/test_k2_horizon_tools.py | 291 ++++++++++++++++++++++++++ 4 files changed, 498 insertions(+) create mode 100644 gmlx/models/k2_horizon/tools.py create mode 100644 tests/models/test_k2_horizon_tools.py diff --git a/gmlx/load/loader.py b/gmlx/load/loader.py index 4040cac7a..822e20b03 100644 --- a/gmlx/load/loader.py +++ b/gmlx/load/loader.py @@ -191,6 +191,10 @@ def build_model(config_dict: dict, *, mtp: bool = False): import gmlx.models.hy_v3.tools as hy_v3_tools hy_v3_tools.ensure_registered() + if mt == "k2_horizon": + import gmlx.models.k2_horizon.tools as k2_horizon_tools + + k2_horizon_tools.ensure_registered() if mt == "glm5_next": # Its hybrid cache depends on the deepseek_v4 DS4PoolingCache. import gmlx.models.glm5_next.model as glm5_next_model diff --git a/gmlx/load/tokenizer.py b/gmlx/load/tokenizer.py index 87430b6c3..95162eb12 100644 --- a/gmlx/load/tokenizer.py +++ b/gmlx/load/tokenizer.py @@ -910,6 +910,10 @@ def finish_gguf_tokenizer(raw, model_type: str | None) -> None: import gmlx.models.deepseek_v41.tools as deepseek_v41_tools deepseek_v41_tools.install_message_normalizer(raw) + if model_type == "k2_horizon": + import gmlx.models.k2_horizon.tools as k2_horizon_tools + + k2_horizon_tools.install_message_normalizer(raw) def bundled_chat_template_for_arch(arch: str | None) -> str | None: diff --git a/gmlx/models/k2_horizon/tools.py b/gmlx/models/k2_horizon/tools.py new file mode 100644 index 000000000..86f002af9 --- /dev/null +++ b/gmlx/models/k2_horizon/tools.py @@ -0,0 +1,199 @@ +"""K2 Horizon tool-call parser (the ```` tag format). + +The K2 Horizon chat template wraps one or more calls in a section block. +Each call is the function name on its own line, then key and value pairs +(``tool_call_format`` ``xml``, the default):: + + + {name} + {key} + {value} + + + +``xml_typed`` adds ``{type}`` after each key, +which the models often leave out. ``json`` writes +``{"name": ..., "arguments": {...}}``. + +Parser contract (``mlx_vlm.tools.base.ToolParser``): ``tool_call_start`` and +``tool_call_end`` are the wrapper tags, so the streaming suppressor holds the +whole section, and ``parse_tool_call(text, tools)`` splits the per-call +segments of the inner text itself. + +Argument values follow llama.cpp's ``common/parsers/k2-horizon.cpp``. The +template writes a string raw and any other value with ``tojson``. A value +whose schema type is only string stays raw text. A union with string takes +the JSON reading only when it parses as one of the union's other types. Any +other value is JSON, and stays text when it does not parse. +""" + +from __future__ import annotations + +import json +import re +import sys +from typing import Any + +from gmlx.upstream.patching import link, wrapped_with + +tool_call_start = "" +tool_call_end = "" + +_CALL_RE = re.compile(r"(.*?)", re.DOTALL) +_ARG_KEY = "" +_ARG_RE = re.compile( + r"(.*?)\s*" + r"(?:.*?\s*)?" + r"(.*?)", + re.DOTALL, +) + +# The template's tool-call instructions spell both tags. Both are needed so +# no other template that mentions one of them is claimed. +_TEMPLATE_MARKERS = ((tool_call_start, _ARG_KEY),) + +# The reasoning fields the template reads from an assistant message. It +# raises when none is present. +_THINKING_FIELDS = ("think", "think_fast", "think_faster", + "reasoning_content", "reasoning") + +_JSON_TYPES = { + "object": (dict,), + "array": (list,), + "boolean": (bool,), + "null": (type(None),), +} + + +def _schema_types(schema: Any) -> set[str]: + """The JSON types a parameter schema allows. Empty means any type.""" + if not isinstance(schema, dict): + return set() + t = schema.get("type") + if isinstance(t, str): + return {t} + if isinstance(t, list): + return {x for x in t if isinstance(x, str)} + out: set[str] = set() + for key in ("anyOf", "oneOf"): + for sub in schema.get(key) or (): + sub_types = _schema_types(sub) + if not sub_types: + return set() + out |= sub_types + return out + + +def _param_types(tool_name: str, tools: list[Any] | None) -> dict[str, set[str]]: + for tool in tools or (): + func = tool.get("function") if isinstance(tool, dict) else None + if not func or func.get("name") != tool_name: + continue + properties = (func.get("parameters") or {}).get("properties") or {} + return {n: _schema_types(s) for n, s in properties.items()} + return {} + + +def _is_type(value: Any, types: set[str]) -> bool: + for t in types: + if t in ("number", "integer"): + if isinstance(value, (int, float)) and not isinstance(value, bool): + return True + elif isinstance(value, _JSON_TYPES.get(t, ())): + return True + return False + + +def _decode(value: str, types: set[str]) -> Any: + if types == {"string"}: + return value + try: + parsed = json.loads(value) + except ValueError: + return value + if "string" in types and not _is_type(parsed, types - {"string"}): + return value + return parsed + + +def _parse_json_call(body: str): + try: + call = json.loads(body) + except ValueError: + return None + if not isinstance(call, dict) or not isinstance(call.get("name"), str): + return None + args = call.get("arguments", {}) + if isinstance(args, str): + try: + args = json.loads(args) + except ValueError: + pass + return {"name": call["name"], "arguments": args if args is not None else {}} + + +def _parse_single(body: str, tools: list[Any] | None): + """One ```` body -> ``{"name", "arguments"}``.""" + if body.lstrip().startswith("{"): + call = _parse_json_call(body.strip()) + if call is not None: + return call + name, sep, rest = body.partition(_ARG_KEY) + name = name.strip() + if not sep: + if name and "<" not in name and not re.search(r"\s", name): + return {"name": name, "arguments": {}} + return {"name": "unknown", "arguments": {"raw": name}} + types = _param_types(name, tools) + arguments: dict[str, Any] = {} + for m in _ARG_RE.finditer(_ARG_KEY + rest): + key = m.group(1).strip() + arguments[key] = _decode(m.group(2), types.get(key, set())) + return {"name": name, "arguments": arguments} + + +def parse_tool_call(text: str, tools: list[Any] | None = None): + """Parse a section's inner text: a list of calls when ```` + segments are present, else the text as one bare call body.""" + bodies = _CALL_RE.findall(text) + if bodies: + return [_parse_single(b, tools) for b in bodies] + return _parse_single(text, tools) + + +def normalize_messages(messages: Any) -> Any: + """Give every assistant message a reasoning field. The template raises on + an assistant turn that has none, as a replayed client history often is.""" + if not isinstance(messages, list): + return messages + out, changed = [], False + for msg in messages: + if (isinstance(msg, dict) and msg.get("role") == "assistant" + and not any(f in msg for f in _THINKING_FIELDS)): + msg = {**msg, "reasoning_content": ""} + changed = True + out.append(msg) + return out if changed else messages + + +def install_message_normalizer(tokenizer) -> None: + """Run ``normalize_messages`` ahead of every chat-template render on + this tokenizer. Idempotent.""" + original = getattr(tokenizer, "apply_chat_template", None) + if original is None or wrapped_with(original, "_k2_horizon_normalized"): + return + + def apply_chat_template(conversation, *args, **kwargs): + return original(normalize_messages(conversation), *args, **kwargs) + + apply_chat_template._k2_horizon_normalized = True + tokenizer.apply_chat_template = link(apply_chat_template, original) + + +def ensure_registered() -> None: + """Register this module as the ``k2_horizon`` tool parser, ahead of any + vendored parser of that name, and route templates carrying the K2 + Horizon tags to it. Idempotent.""" + from gmlx.upstream.tool_parsers import register_tool_parser + + register_tool_parser("k2_horizon", sys.modules[__name__], _TEMPLATE_MARKERS) diff --git a/tests/models/test_k2_horizon_tools.py b/tests/models/test_k2_horizon_tools.py new file mode 100644 index 000000000..bdbfb14e3 --- /dev/null +++ b/tests/models/test_k2_horizon_tools.py @@ -0,0 +1,291 @@ +"""K2 Horizon tool-call parser (``gmlx.models.k2_horizon.tools``). CPU-only, +no model load. + +``parse_tool_call`` gets the inner text of an ```` section, +as mlx-vlm's ``process_tool_calls`` hands it over. The argument cases mirror +the K2 Horizon block of llama.cpp's ``tests/test-chat.cpp``. +""" + +from __future__ import annotations + +import json +from pathlib import Path + +import pytest + +import gmlx.models.k2_horizon.tools as k2_tools +from gmlx.models.k2_horizon.tools import ( + install_message_normalizer, + normalize_messages, + parse_tool_call, + tool_call_end, + tool_call_start, +) + +_FIXTURES = Path(__file__).resolve().parents[1] / "fixtures" +_TEMPLATE = (_FIXTURES / "k2_horizon_template.jinja").read_text() + + +def _tool(name, properties, required=()): + return {"type": "function", "function": { + "name": name, + "parameters": {"type": "object", "properties": properties, + "required": list(required)}}} + + +_SPECIAL = _tool("special_function", {"arg1": {"type": "integer"}}, ["arg1"]) +_GET_TIME = _tool("get_time", {"city": {"type": "string"}}, ["city"]) +_EMPTY = _tool("empty_args", {}) +_2REQ = _tool("tool_2req_4opt", { + "req1": {"type": "string"}, "req2": {"type": "integer"}, + "opt1": {"type": "string"}, "opt2": {"type": "integer"}, +}, ["req1", "req2"]) +_UNION = _tool("set_union", { + "value": {"type": ["string", "object"]}, + "amount": {"type": ["string", "integer"]}, +}, ["value", "amount"]) + + +def _call(name, args=(), typed=False): + body = f"{name}\n" + for key, value, *rest in args: + body += f"{key}\n" + if typed and rest: + body += f"{rest[0]}\n" + body += f"{value}\n" + return body + "" + + +def _inner(*calls): + return "\n" + "\n".join(calls) + "\n" + + +def _registry(): + # A plain import: a registry move must fail here, not skip. + from gmlx._vendor.mlx_vlm.tools import registry + + return registry + + +def _render(messages, tools=None, **kwargs): + from transformers.utils.chat_template_utils import _compile_jinja_template + + return _compile_jinja_template(_TEMPLATE).render( + messages=messages, tools=tools, add_generation_prompt=True, + bos_token="", **kwargs) + + +def test_delimiters_are_the_wrapper_tags(): + # The streaming suppressor keys on tool_call_start, so it must be the + # section tag, not the per-call tag it shares a prefix with. + assert tool_call_start == "" + assert tool_call_end == "" + + +def test_xml_call(): + out = parse_tool_call(_inner(_call("special_function", [("arg1", "1")])), + [_SPECIAL]) + assert out == [{"name": "special_function", "arguments": {"arg1": 1}}] + + +def test_xml_typed_call(): + text = _inner(_call("special_function", [("arg1", "1", "integer")], + typed=True)) + assert "" in text + assert parse_tool_call(text, [_SPECIAL]) == [ + {"name": "special_function", "arguments": {"arg1": 1}}] + + +def test_json_call(): + text = _inner('{"name": "special_function", ' + '"arguments": {"arg1": 1}}') + assert parse_tool_call(text, [_SPECIAL]) == [ + {"name": "special_function", "arguments": {"arg1": 1}}] + + +def test_json_call_with_string_arguments(): + text = _inner('{"name": "special_function", ' + '"arguments": "{\\"arg1\\": 1}"}') + assert parse_tool_call(text, [_SPECIAL]) == [ + {"name": "special_function", "arguments": {"arg1": 1}}] + + +def test_zero_arg_call(): + assert parse_tool_call(_inner(_call("empty_args")), [_EMPTY]) == [ + {"name": "empty_args", "arguments": {}}] + + +def test_parallel_calls(): + text = _inner(_call("get_time", [("city", "Paris")]), + _call("get_time", [("city", "Rome")])) + assert parse_tool_call(text, [_GET_TIME]) == [ + {"name": "get_time", "arguments": {"city": "Paris"}}, + {"name": "get_time", "arguments": {"city": "Rome"}}, + ] + + +def test_args_in_any_order_keep_schema_types(): + text = _inner(_call("tool_2req_4opt", [("req2", "7"), ("req1", "hello")])) + assert parse_tool_call(text, [_2REQ]) == [ + {"name": "tool_2req_4opt", "arguments": {"req2": 7, "req1": "hello"}}] + + +def test_string_schema_value_stays_raw(): + text = _inner(_call("tool_2req_4opt", [("req1", "123"), ("req2", "1"), + ("opt1", "true")])) + assert parse_tool_call(text, [_2REQ])[0]["arguments"] == { + "req1": "123", "req2": 1, "opt1": "true"} + + +@pytest.mark.parametrize( + "value", ["true", "42", "null", "[]", '"quoted"', "{not valid json"]) +def test_union_with_string_keeps_other_json_as_text(value): + text = _inner(_call("set_union", [("value", value), ("amount", "42")])) + assert parse_tool_call(text, [_UNION]) == [ + {"name": "set_union", "arguments": {"value": value, "amount": 42}}] + + +def test_union_with_string_takes_its_json_type(): + text = _inner(_call("set_union", [("value", '{"a": 1}'), + ("amount", "2 dollars")])) + assert parse_tool_call(text, [_UNION]) == [ + {"name": "set_union", + "arguments": {"value": {"a": 1}, "amount": "2 dollars"}}] + + +def test_multiline_string_value_kept_exactly(): + text = _inner(_call("get_time", [("city", "San\nFrancisco\n")])) + assert parse_tool_call(text, [_GET_TIME])[0]["arguments"] == { + "city": "San\nFrancisco\n"} + + +def test_untyped_value_falls_back_to_text(): + text = _inner(_call("unknown_tool", [("a", "Paris"), ("b", "[1, 2]")])) + assert parse_tool_call(text, None) == [ + {"name": "unknown_tool", "arguments": {"a": "Paris", "b": [1, 2]}}] + + +def test_freeform_text_is_not_a_call(): + out = parse_tool_call("just some text", None) + assert out == {"name": "unknown", "arguments": {"raw": "just some text"}} + + +@pytest.mark.parametrize("fmt", ["xml", "xml_typed", "json"]) +def test_template_render_round_trips(fmt): + tools = [_tool("get_weather", {"city": {"type": "string"}, + "days": {"type": "integer"}, + "opts": {"type": "object"}}, ["city"])] + calls = [{"city": "Paris\nFrance", "days": 3, "opts": {"a": [1, 2]}}, {}] + messages = [ + {"role": "user", "content": "hi"}, + {"role": "assistant", "content": "", "reasoning_content": "r", + "tool_calls": [{"type": "function", "function": { + "name": "get_weather", "arguments": a}} for a in calls]}, + {"role": "tool", "content": "sunny"}, + ] + out = _render(messages, tools, tool_call_format=fmt) + # The system turn spells the section tags in its instructions. + turn = out.split("<|ifm|im_start|>assistant", 1)[1] + inner = turn.split(tool_call_start, 1)[1].split(tool_call_end, 1)[0] + assert parse_tool_call(inner, tools) == [ + {"name": "get_weather", "arguments": a} for a in calls] + + +# --- message normalizer ----------------------------------------------------- + + +def test_normalizer_adds_empty_reasoning_to_bare_assistant_turns(): + msgs = [{"role": "user", "content": "hi"}, + {"role": "assistant", "content": "hello"}] + out = normalize_messages(msgs) + assert out[1] == {"role": "assistant", "content": "hello", + "reasoning_content": ""} + assert "reasoning_content" not in msgs[1] + assert out[0] is msgs[0] + + +@pytest.mark.parametrize("field", ["think", "think_fast", "think_faster", + "reasoning_content", "reasoning"]) +def test_normalizer_keeps_turns_that_carry_a_reasoning_field(field): + msgs = [{"role": "assistant", "content": "x", field: "r"}] + assert normalize_messages(msgs) is msgs + + +def test_normalizer_lets_the_template_render_client_history(): + from jinja2.exceptions import TemplateError + + msgs = [{"role": "user", "content": "hi"}, + {"role": "assistant", "content": "hello"}, + {"role": "user", "content": "again"}] + with pytest.raises(TemplateError, match="thinking field"): + _render(msgs) + out = _render(normalize_messages(msgs)) + assert "\nhello" in out + + +def test_install_message_normalizer_wraps_once(): + class _Tok: + def __init__(self): + self.seen = [] + + def apply_chat_template(self, conversation, *args, **kwargs): + self.seen.append(conversation) + return "OK" + + tok = _Tok() + install_message_normalizer(tok) + first = tok.apply_chat_template + install_message_normalizer(tok) + assert tok.apply_chat_template is first + assert tok.apply_chat_template( + [{"role": "assistant", "content": "x"}]) == "OK" + assert tok.seen == [[{"role": "assistant", "content": "x", + "reasoning_content": ""}]] + + +# --- registry --------------------------------------------------------------- + + +def _stock_choice(registry, text): + matches = [s for s in registry.SPECS + if s.matches(text) and s.name != "k2_horizon"] + return max(matches, key=lambda s: s.priority).name if matches else None + + +def test_stock_registry_would_misroute_the_template(): + # The template's XML tool listing carries both minicpm5 markers. + assert _stock_choice(_registry(), _TEMPLATE) == "minicpm5" + + +def test_ensure_registered_infers_k2_horizon_from_template(): + registry = _registry() + k2_tools.ensure_registered() + assert registry._infer_tool_parser(_TEMPLATE) == "k2_horizon" + + +def test_k2_markers_claim_no_other_fixture_template(): + registry = _registry() + k2_tools.ensure_registered() + spec = next(s for s in registry.SPECS if s.name == "k2_horizon") + others = [p for p in _FIXTURES.glob("*template*.jinja") + if p.name != "k2_horizon_template.jinja"] + assert others + claimed = [p.name for p in others if spec.matches(p.read_text())] + assert claimed == [] + + +def test_ensure_registered_is_idempotent_and_loadable(): + registry = _registry() + k2_tools.ensure_registered() + before = registry.SPECS + k2_tools.ensure_registered() + assert registry.SPECS == before + assert sum(s.name == "k2_horizon" for s in registry.SPECS) == 1 + assert registry.load_tool_module("k2_horizon") is k2_tools + + +def test_parser_output_is_json_serializable(): + out = parse_tool_call(_inner(_call("set_union", [("value", '{"a": 1}'), + ("amount", "3")])), + [_UNION]) + json.dumps(out) From 50cea92a55f7a7b3ae42bef9ed9979b677b5b092 Mon Sep 17 00:00:00 2001 From: Asher Feldman <59994+asher@users.noreply.github.com> Date: Sat, 17 Oct 2026 18:08:23 -0700 Subject: [PATCH 04/10] fix(serve): decode string tool-call arguments for k2 horizon templates --- gmlx/serve/patches/render.py | 11 +++++++---- tests/serve/test_faithful_history.py | 27 +++++++++++++++++++++++++++ 2 files changed, 34 insertions(+), 4 deletions(-) diff --git a/gmlx/serve/patches/render.py b/gmlx/serve/patches/render.py index 9bf3f65c3..0217d237c 100644 --- a/gmlx/serve/patches/render.py +++ b/gmlx/serve/patches/render.py @@ -86,6 +86,9 @@ "qwen3_5", "qwen3_5_moe", "kimi_k25", "muse_glimmer", "muse_spark", "glm5_next", "gemma4_unified", )) +# The model types whose chat template raises on a JSON-string +# ``tool_calls[*].function.arguments``. +_DICT_TOOL_ARGUMENTS = frozenset(("muse_glimmer", "muse_spark", "k2_horizon")) def _tool_images(message) -> int: @@ -265,9 +268,9 @@ def _decode_tool_arguments(msgs: list) -> None: """Decode ``tool_calls[*].function.arguments`` from the OpenAI wire form (a JSON string) into a mapping, in place. - The Muse Glimmer and Muse Spark ATEM templates raise outright on a - string, since the jinja sandbox cannot parse JSON, so a client replaying - its own tool turn would 500 on render. Undecodable arguments become an empty + The ``_DICT_TOOL_ARGUMENTS`` templates (the Muse ATEM ones and K2 + Horizon) raise outright on a string, since the jinja sandbox cannot parse + JSON, so a client replaying its own tool turn would 500 on render. Undecodable arguments become an empty mapping, the same fallback the retirement mirror uses (``retire_key.assistant_message_from_reply``). """ @@ -356,7 +359,7 @@ def apply_chat_template(processor, config, prompt, if (isinstance(m, dict) and "thinking" not in m and m.get("reasoning_content")): msgs[i] = {**m, "thinking": m["reasoning_content"]} - elif cfg.get("model_type") in ("muse_glimmer", "muse_spark"): + elif cfg.get("model_type") in _DICT_TOOL_ARGUMENTS: _decode_tool_arguments(msgs) if return_messages: return msgs diff --git a/tests/serve/test_faithful_history.py b/tests/serve/test_faithful_history.py index 0b9cba150..f5873eaad 100644 --- a/tests/serve/test_faithful_history.py +++ b/tests/serve/test_faithful_history.py @@ -187,3 +187,30 @@ def test_retire_capture_wraps_faithful_render(): inner = [c.cell_contents for c in outer.__closure__ if callable(getattr(c, "cell_contents", None))] assert any(getattr(f, "_kq_faithful_history", False) for f in inner) + + +@pytest.mark.parametrize("model_type", ["muse_spark", "k2_horizon"]) +def test_string_tool_arguments_decode_for_dict_templates(model_type): + # These templates raise on the OpenAI JSON-string form. + install_faithful_history() + history = _history() + history[1]["tool_calls"] = [{ + "id": "c1", "type": "function", + "function": {"name": "f", "arguments": '{"city": "Paris"}'}}] + msgs = openai_mod.apply_chat_template( + _Proc(), {"model_type": model_type}, history, return_messages=True) + assert msgs[1]["tool_calls"][0]["function"]["arguments"] == { + "city": "Paris"} + assert msgs[1]["tool_calls"][0]["id"] == "c1" + + +def test_string_tool_arguments_stay_strings_elsewhere(): + install_faithful_history() + history = _history() + history[1]["tool_calls"] = [{ + "id": "c1", "type": "function", + "function": {"name": "f", "arguments": '{"city": "Paris"}'}}] + msgs = openai_mod.apply_chat_template( + _Proc(), PLAIN_CFG, history, return_messages=True) + assert msgs[1]["tool_calls"][0]["function"]["arguments"] == ( + '{"city": "Paris"}') From ee08ed601c9bb78abc197f5197477ef8056afc18 Mon Sep 17 00:00:00 2001 From: Asher Feldman <59994+asher@users.noreply.github.com> Date: Sat, 17 Oct 2026 18:17:00 -0700 Subject: [PATCH 05/10] feat(reasoning): k2 horizon think pairs picked per reasoning effort, and its family defaults --- docs/family-defaults.md | 1 + gmlx/commands/cli.py | 1 + gmlx/gen/generation.py | 3 + gmlx/gen/profiles.py | 14 +++ gmlx/gen/thinking_budget.py | 111 ++++++++++++++++------ gmlx/serve/patches/chat_behavior.py | 28 ++++++ gmlx/tui/chat.py | 2 + gmlx/tui/reasoning.py | 8 ++ tests/gen/test_thinking_budget.py | 78 +++++++++++++++ tests/k2_think_tokens.py | 72 ++++++++++++++ tests/serve/test_patches_chat_behavior.py | 84 ++++++++++++++++ tests/tui/test_reasoning.py | 21 ++++ 12 files changed, 396 insertions(+), 27 deletions(-) create mode 100644 tests/k2_think_tokens.py diff --git a/docs/family-defaults.md b/docs/family-defaults.md index 1dbbedab7..28ea6e5ae 100644 --- a/docs/family-defaults.md +++ b/docs/family-defaults.md @@ -39,6 +39,7 @@ a model ends up with. | `kimi-k2` | `deepseek2 named (?i)\bkimi` | temperature=1.0 top_p=0.95 | - | | `muse` | `muse-glimmer` | temperature=1.0 top_p=0.95 top_k=64 thinking_start_token=<\|start\|>assistant to=self<\|message\|> thinking_end_token=<\|eom\|> | `@reasoning-high`: reasoning_strength=high; `@reasoning-low`: reasoning_strength=low; `@reasoning-medium`: reasoning_strength=medium; `@reasoning-xhigh`: reasoning_strength=xhigh | | `muse-spark` | `muse-spark` | temperature=1.0 top_p=1.0 top_k=64 thinking_start_token=<\|start\|>assistant to=self<\|message\|> thinking_end_token=<\|eom\|> | `@reasoning-high`: reasoning_effort=high; `@reasoning-low`: reasoning_effort=low; `@reasoning-medium`: reasoning_effort=medium; `@reasoning-xhigh`: reasoning_effort=xhigh | +| `k2-horizon` | `k2-horizon` | temperature=1.0 top_p=0.95 | `@reasoning-high`: reasoning_effort=high; `@reasoning-low`: reasoning_effort=low; `@reasoning-medium`: reasoning_effort=medium | | `llama` | `llama`, `smollm3` | temperature=0.6 top_p=0.9 | - | | `mistral` | `mistral3` | temperature=0.15 | - | | `default` | (anything else) | temperature=0.7 top_p=0.95 | `@coding`: temperature=0.3; `@creative`: temperature=1.0 min_p=0.05 | diff --git a/gmlx/commands/cli.py b/gmlx/commands/cli.py index 54046b70a..6f2e51fa9 100644 --- a/gmlx/commands/cli.py +++ b/gmlx/commands/cli.py @@ -1859,6 +1859,7 @@ def _run_vlm(args) -> int: think_tokenizer_for(processor), None, interruptible=True, start_token=args.thinking_start_token, end_token=args.thinking_end_token, + prompt=prompt, ) if tbp is not None: extra["logits_processors"] = [tbp] diff --git a/gmlx/gen/generation.py b/gmlx/gen/generation.py index 3cc6ed7e3..67bb2fe67 100644 --- a/gmlx/gen/generation.py +++ b/gmlx/gen/generation.py @@ -452,6 +452,7 @@ def generate( interruptible=True, start_token=thinking_start_token, end_token=thinking_end_token, + prompt=prompt, ) if tbp is not None: logits_processors = list(logits_processors) + [tbp] @@ -1401,6 +1402,7 @@ def generate_speculative_owned( start_token=thinking_start_token, end_token=thinking_end_token, prompt_ids=prompt_ids, + prompt=prompt, ) set_finish_key_target(hook) tic = time.perf_counter() @@ -1586,6 +1588,7 @@ def _stream_generate_speculative_owned( start_token=thinking_start_token, end_token=thinking_end_token, prompt_ids=prompt_ids, + prompt=prompt, ) set_finish_key_target(hook) tic = time.perf_counter() diff --git a/gmlx/gen/profiles.py b/gmlx/gen/profiles.py index 9684b1077..99ff11a82 100644 --- a/gmlx/gen/profiles.py +++ b/gmlx/gen/profiles.py @@ -277,6 +277,20 @@ "reasoning-xhigh": {"chat_template_kwargs": {"reasoning_effort": "xhigh"}}, }, }, + # IFM/K2-Horizon-36B model card, 2026-10 (generation_config carries no + # sampling): t=1.0, top_p=0.95. The template takes reasoning_effort + # high/medium/low, defaults to high, and pre-opens a different think + # pair for each, so the family sets no markers: the prompt picks them. + "k2-horizon": { + "label": "K2 Horizon", + "arches": ("k2-horizon",), + "base": {"sampling": {"temperature": 1.0, "top_p": 0.95}}, + "intents": { + "reasoning-low": {"chat_template_kwargs": {"reasoning_effort": "low"}}, + "reasoning-medium": {"chat_template_kwargs": {"reasoning_effort": "medium"}}, + "reasoning-high": {"chat_template_kwargs": {"reasoning_effort": "high"}}, + }, + }, # Llama 3.x generation_config (t=0.6/top_p=0.9); SmolLM3 card matches # closely enough to share. "llama": { diff --git a/gmlx/gen/thinking_budget.py b/gmlx/gen/thinking_budget.py index a90f1e428..04c059945 100644 --- a/gmlx/gen/thinking_budget.py +++ b/gmlx/gen/thinking_budget.py @@ -54,7 +54,13 @@ # prompt stops one marker short (at "<|start|>assistant"), so # the block is not prompt-opened and the model emits the whole # header itself. - ("<|start|>assistant to=self<|message|>", "<|eom|>")) + ("<|start|>assistant to=self<|message|>", "<|eom|>"), + # K2 Horizon: one pair per reasoning_effort (high, medium, + # low). The template names all three and pre-opens the one + # the effort picks, so the prompt decides the pair. + ("", ""), + ("", ""), + ("", "")) # Forced into the thinking block ahead of the first budget-triggered close. # The model must see itself DECIDE to answer: a bare close tag cuts the @@ -75,27 +81,68 @@ ) +def think_pairs_in(template) -> list: + """Every ``_THINK_PAIRS`` spelling ``template`` names, in table order. + Most templates name one; K2 Horizon names one per reasoning effort.""" + if not isinstance(template, str) or not template: + return [] + pairs = [p for p in _THINK_PAIRS if p[0] in template or p[1] in template] + # Kimi-K3 builds its markers by macro concatenation ('<|open|>' + tag + + # '<|sep|>'), so the literal pair never appears in the template source; + # detect the macro call on the 'think' tag instead. + if not pairs and "<|open|>" in template and "open_tag('think" in template: + pairs = [("<|open|>think<|sep|>", "<|close|>think<|sep|>")] + return pairs + + +def template_think_pairs(tokenizer) -> list: + """:func:`think_pairs_in` for the tokenizer's chat template.""" + return think_pairs_in(getattr(tokenizer, "chat_template", None)) + + def _template_think_pair(tokenizer): """The ``_THINK_PAIRS`` spelling the model's chat template actually uses, or None. The template is authoritative over vocab probing: a vocab can carry legacy ```` entries while the template thinks in another spelling (MiniMax-M3 ships both ```` and ```` as real single tokens but only ever emits the latter), and probing then resolves - - and forces - a tag the model treats as ordinary text.""" - tmpl = getattr(tokenizer, "chat_template", None) - if not isinstance(tmpl, str) or not tmpl: + - and forces - a tag the model treats as ordinary text. A template that + names several pairs gives the first; :func:`prompt_think_pair` picks the + one a rendered prompt opens.""" + pairs = template_think_pairs(tokenizer) + return pairs[0] if pairs else None + + +def open_think_pair(text, pairs): + """The pair in ``pairs`` whose block ``text`` ends inside (its last start + marker not yet closed by a later end marker), or None.""" + if not isinstance(text, str): return None - for pair in _THINK_PAIRS: - if pair[0] in tmpl or pair[1] in tmpl: - return pair - # Kimi-K3 builds its markers by macro concatenation ('<|open|>' + tag + - # '<|sep|>'), so the literal pair never appears in the template source; - # detect the macro call on the 'think' tag instead. - if "<|open|>" in tmpl and "open_tag('think" in tmpl: - return ("<|open|>think<|sep|>", "<|close|>think<|sep|>") + for start, end in pairs: + i = text.rfind(start) + if i != -1 and i > text.rfind(end): + return (start, end) return None +def prompt_think_pair(prompt, tokenizer): + """The template pair a rendered ``prompt`` ends inside, or None. + ``prompt`` is the rendered text or its token ids, whose tail is decoded. + With K2 Horizon this is the pair the request's reasoning effort chose.""" + pairs = template_think_pairs(tokenizer) + if not pairs or prompt is None: + return None + if not isinstance(prompt, str): + try: + ids = prompt.tolist() if hasattr(prompt, "tolist") else list(prompt) + while ids and isinstance(ids[0], list): + ids = ids[-1] + prompt = tokenizer.decode([int(t) for t in ids[-64:]]) + except Exception: # noqa: BLE001 - no decodable tail: no pair + return None + return open_think_pair(prompt, pairs) + + def prompt_open_think_tag(prompt, start_token: str | None = None, end_token: str | None = None, tokenizer=None): """The start marker of the thinking block a rendered ``prompt`` ends @@ -113,20 +160,17 @@ def prompt_open_think_tag(prompt, start_token: str | None = None, else: pairs = [(start_token or "", end_token or "")] if tokenizer is not None: - tpair = _template_think_pair(tokenizer) + tpairs = template_think_pairs(tokenizer) s = getattr(tokenizer, "think_start", None) e = getattr(tokenizer, "think_end", None) - if tpair is not None: - pairs = [tpair] + if tpairs: + pairs = tpairs elif s and e: pairs = [(s, e)] if pairs == [("", "")]: pairs = list(_THINK_PAIRS) - for start, end in pairs: - i = prompt.rfind(start) - if i != -1 and i > prompt.rfind(end): - return start - return None + pair = open_think_pair(prompt, pairs) + return pair[0] if pair else None def prompt_opens_thinking(prompt, start_token: str | None = None, @@ -369,7 +413,8 @@ def _tag_pair_seqs(tokenizer, start_tag, end_tag): tuple(int(t) for t in end_ids)) -def _thinking_token_seqs(tokenizer, start_token=None, end_token=None): +def _thinking_token_seqs(tokenizer, start_token=None, end_token=None, + prompt=None): """``(start_seq, end_seq)`` token-id tuples for the model's thinking delimiters, from mlx-lm's ``TokenizerWrapper`` inference (covers ````, ````, and the ``<|channel>thought`` / ```` channel @@ -383,10 +428,17 @@ def _thinking_token_seqs(tokenizer, start_token=None, end_token=None): names one of the ``_THINK_PAIRS`` spellings wins over everything except a configured ``end_token``: the template is the ground truth for which tags the model actually emits (see ``_template_think_pair``), and configured - markers are the user's word for a model neither one recognizes.""" + markers are the user's word for a model neither one recognizes. + ``prompt`` (the rendered prompt, text or token ids) picks among the pairs + a template names: the one the prompt opens wins.""" seqs = _tag_pair_seqs(tokenizer, start_token, end_token) if seqs is not None: return seqs + pair = prompt_think_pair(prompt, tokenizer) if end_token is None else None + if pair is not None: + seqs = _tag_pair_seqs(tokenizer, *pair) + if seqs is not None: + return seqs pair = _template_think_pair(tokenizer) if pair is not None: seqs = _tag_pair_seqs(tokenizer, *pair) @@ -704,6 +756,7 @@ def _announce_budget(tokenizer, budget, end_seq) -> None: def make_thinking_budget_processor( tokenizer, budget, *, start_in_thinking=True, verbose=False, eos_floor=True, interruptible=False, start_token=None, end_token=None, + prompt=None, ): """Build a thinking-budget logits processor, or ``None`` if unsupported. @@ -717,6 +770,8 @@ def make_thinking_budget_processor( ```` channel format), so the cap is not tied to one spelling. ``start_token``/``end_token`` (a family card's or the user's configured markers) override that inference for a model it does not recognize. + ``prompt`` (the rendered prompt) picks the pair it opens when the + template names several, so the close matches the open. Returns ``None`` (warning when ``verbose``) when no thinking-end token is detected; the caller then generates uncapped. ``budget=0`` force-closes as soon as a thinking block opens. @@ -743,12 +798,12 @@ def make_thinking_budget_processor( # probe can't handle must not break plain generation. try: start_seq, end_seq = _thinking_token_seqs( - tokenizer, start_token, end_token) + tokenizer, start_token, end_token, prompt) except Exception: # noqa: BLE001 return None else: start_seq, end_seq = _thinking_token_seqs( - tokenizer, start_token, end_token) + tokenizer, start_token, end_token, prompt) if not end_seq: if verbose and budget is not None: print( @@ -891,7 +946,7 @@ def take_forced(self): def make_mtp_finish_hook( tokenizer, *, budget=None, start_in_thinking=False, verbose=False, - start_token=None, end_token=None, prompt_ids=None, + start_token=None, end_token=None, prompt_ids=None, prompt=None, ): """Build the thinking-budget/^T hook for the owned MTP round loop, or ``None`` when the model's thinking delimiters can't be resolved (the @@ -900,12 +955,14 @@ def make_mtp_finish_hook( Best-effort by design: a tokenizer the probe can't handle must not break MTP decode. ``prompt_ids`` (the prompt's token ids, a list or a 1-D array) lets a start marker that the prompt begins complete on the - generated tokens.""" + generated tokens. ``prompt`` (the rendered text, else ``prompt_ids``) + picks the pair it opens when the template names several.""" if budget is not None and budget < 0: return None try: start_seq, end_seq = _thinking_token_seqs( - tokenizer, start_token, end_token) + tokenizer, start_token, end_token, + prompt if prompt is not None else prompt_ids) except Exception: # noqa: BLE001 return None if not end_seq: diff --git a/gmlx/serve/patches/chat_behavior.py b/gmlx/serve/patches/chat_behavior.py index 026ccbdbc..27ff0a2b9 100644 --- a/gmlx/serve/patches/chat_behavior.py +++ b/gmlx/serve/patches/chat_behavior.py @@ -13,9 +13,11 @@ import gmlx.serve.bridge_vlm as serving from gmlx.gen.thinking_budget import ( marker_ids, + open_think_pair, prompt_opens_block, prompt_tail_ids, seq_thinking_budget_criteria_cls, + think_pairs_in, ) from gmlx.upstream.patching import link, wrapped_with @@ -166,6 +168,7 @@ def _stash_template_kwargs(args, request, _processor): if merged: args._kq_template_kwargs = merged args._kq_thinking_explicit = thinking_explicit + _LAST_GEN_ARGS.set(args) return args @@ -447,6 +450,29 @@ def _split_thinking_keep_indent(rs, cls, text, start_token, end_token): _LAST_RENDERED_PROMPT: contextvars.ContextVar = contextvars.ContextVar( "kq_last_rendered_prompt", default=None) +# The request's GenerationArguments, so the render seam can set the think +# pair its prompt opens (see _adopt_prompt_think_pair). +_LAST_GEN_ARGS: contextvars.ContextVar = contextvars.ContextVar( + "kq_last_gen_args", default=None) + + +def _adopt_prompt_think_pair(args, processor, rendered) -> None: + """Set the request's think markers to the pair its rendered prompt opens, + when the chat template names several and none is configured. K2 Horizon + names one pair per reasoning_effort. The budget criteria, the MTP hook, + the stream splitter and the reply split read the markers from ``args``, + so each request closes with its own pair.""" + if args is None or not isinstance(rendered, str): + return + if (getattr(args, "thinking_start_token", None) + or getattr(args, "thinking_end_token", None)): + return + pairs = think_pairs_in(_model_template_text(processor)) + if len(pairs) < 2: + return + pair = open_think_pair(rendered, pairs) + if pair is not None: + args.thinking_start_token, args.thinking_end_token = pair def _prompt_tail_opens_thinking(text, open_close_markers) -> bool: @@ -518,6 +544,8 @@ def _wrap(fn): def apply_chat_template(*a, **kw): out = fn(*a, **kw) _LAST_RENDERED_PROMPT.set(out if isinstance(out, str) else None) + _adopt_prompt_think_pair( + _LAST_GEN_ARGS.get(), a[0] if a else kw.get("processor"), out) return out apply_chat_template.__dict__[_STREAM_SEED_FLAG] = True return apply_chat_template diff --git a/gmlx/tui/chat.py b/gmlx/tui/chat.py index 639028d30..584d42c7b 100644 --- a/gmlx/tui/chat.py +++ b/gmlx/tui/chat.py @@ -4033,6 +4033,7 @@ def _apply_session(doc: dict, name: str) -> None: interruptible=True, start_token=state.thinking_start_token, end_token=state.thinking_end_token, + prompt=prompt if isinstance(prompt, str) else None, ) set_finish_key_target(tbp) try: @@ -4230,6 +4231,7 @@ def _apply_session(doc: dict, name: str) -> None: interruptible=True, start_token=state.thinking_start_token, end_token=state.thinking_end_token, + prompt=prompt_text, ) if tbp is not None: logits_processors = list(logits_processors) + [tbp] diff --git a/gmlx/tui/reasoning.py b/gmlx/tui/reasoning.py index 7323d0a8d..61556fc03 100644 --- a/gmlx/tui/reasoning.py +++ b/gmlx/tui/reasoning.py @@ -102,6 +102,14 @@ ("<|close|>response<|sep|>", _DROP), ("<|close|>message<|sep|>", _DROP), ("<|end_of_msg|>", _DROP), + # K2 Horizon: one pair per reasoning_effort (high, medium, low). The + # template pre-opens the chosen one, so generation streams its close. + ("", _REASON), + ("", _ANSWER), + ("", _REASON), + ("", _ANSWER), + ("", _REASON), + ("", _ANSWER), ) # Markers that open a message header - text from here to the closing diff --git a/tests/gen/test_thinking_budget.py b/tests/gen/test_thinking_budget.py index 846bed7aa..7b4f6aa45 100644 --- a/tests/gen/test_thinking_budget.py +++ b/tests/gen/test_thinking_budget.py @@ -1038,3 +1038,81 @@ def test_atem_mtp_hook_commentary_close_keeps_it_armed(): h.observe(t) assert not h.done and h.in_thinking assert h.take_forced() is not None + + +# --- K2 Horizon: one think pair per reasoning effort ------------------------ + +_K2_EFFORTS = ("high", "medium", "low") + + +@pytest.mark.parametrize("effort", _K2_EFFORTS) +def test_k2_prompt_picks_its_effort_pair(effort): + import k2_think_tokens as k2 + + from gmlx.gen.thinking_budget import prompt_open_think_tag + + tok = k2.K2Tokenizer() + prompt = k2.generation_prompt(effort) + start, end = k2.PAIRS[effort] + assert prompt_open_think_tag(prompt, tokenizer=tok) == start + want = ((k2.PIECES[start],), (k2.PIECES[end],)) + assert _thinking_token_seqs(tok, prompt=prompt) == want + assert _thinking_token_seqs(tok, prompt=tok.encode(prompt)) == want + + +def test_k2_without_a_prompt_takes_the_first_pair(): + import k2_think_tokens as k2 + + _, end = _thinking_token_seqs(k2.K2Tokenizer()) + assert end == (k2.PIECES[""],) + + +def _forced_run(processor, prompt_ids, steps, close_at=None): + """Decode filler ids; collect every forced id. ``close_at`` requests ^T + before that step.""" + tokens, forced = list(prompt_ids), [] + for i in range(steps): + if i == close_at: + processor.request_close() + f = _one_hot_id(processor(mx.array(tokens), mx.zeros((1, 260000)))) + forced.append(f) + tokens.append(f if f is not None else 1000 + ord("x")) + return [f for f in forced if f is not None] + + +@pytest.mark.parametrize("effort", _K2_EFFORTS) +def test_k2_budget_close_uses_its_effort_tag(effort): + import k2_think_tokens as k2 + + tok = k2.K2Tokenizer() + prompt = k2.generation_prompt(effort) + p = make_thinking_budget_processor( + tok, 2, start_in_thinking=prompt_opens_thinking(prompt, tokenizer=tok), + prompt=prompt) + forced = _forced_run(p, tok.encode(prompt), 200) + assert forced[-1] == k2.PIECES[k2.PAIRS[effort][1]] + + +@pytest.mark.parametrize("effort", _K2_EFFORTS) +def test_k2_finish_key_closes_with_its_effort_tag(effort): + import k2_think_tokens as k2 + + tok = k2.K2Tokenizer() + prompt = k2.generation_prompt(effort) + p = make_thinking_budget_processor( + tok, None, start_in_thinking=True, interruptible=True, prompt=prompt) + forced = _forced_run(p, tok.encode(prompt), 200, close_at=3) + assert forced[-1] == k2.PIECES[k2.PAIRS[effort][1]] + + +@pytest.mark.parametrize("effort", _K2_EFFORTS) +def test_k2_mtp_finish_key_closes_with_its_effort_tag(effort): + import k2_think_tokens as k2 + + tok = k2.K2Tokenizer() + prompt = k2.generation_prompt(effort) + close = k2.PIECES[k2.PAIRS[effort][1]] + for kw in ({"prompt_ids": tok.encode(prompt)}, {"prompt": prompt}): + hook = make_mtp_finish_hook(tok, start_in_thinking=True, **kw) + hook.request_close() + assert hook.take_forced()[-1] == close diff --git a/tests/k2_think_tokens.py b/tests/k2_think_tokens.py new file mode 100644 index 000000000..a682b6b62 --- /dev/null +++ b/tests/k2_think_tokens.py @@ -0,0 +1,72 @@ +"""A K2 Horizon tokenizer stand-in for the think-pair tests. No model is +needed. + +The control pieces carry the ids the real GGUF vocab gives them. Other text +falls back to one id per character. The chat template is the real one, so a +rendered generation prompt opens the pair its reasoning effort picks. +""" + +from __future__ import annotations + +from pathlib import Path + +TEMPLATE = (Path(__file__).parent / "fixtures" + / "k2_horizon_template.jinja").read_text() + +PAIRS = { + "high": ("", ""), + "medium": ("", ""), + "low": ("", ""), +} + +PIECES: dict[str, int] = { + "<|ifm|begin_of_text|>": 0, + "<|ifm|endoftext|>": 1, + "<|ifm|im_start|>": 250018, + "<|ifm|im_end|>": 250019, + "": 250029, + "": 250030, + "": 250050, + "": 250051, + "": 250052, + "": 250053, +} + +_CHAR_BASE = 1000 + + +class K2Tokenizer: + chat_template = TEMPLATE + eos_token_ids = [1, 250019] + eos_token_id = 1 + + def __init__(self): + self._by_id = {v: k for k, v in PIECES.items()} + self._longest = max(len(k) for k in PIECES) + + def encode(self, text, add_special_tokens=False): + out, i = [], 0 + while i < len(text): + for n in range(min(self._longest, len(text) - i), 0, -1): + if text[i:i + n] in PIECES: + out.append(PIECES[text[i:i + n]]) + i += n + break + else: + out.append(_CHAR_BASE + ord(text[i])) + i += 1 + return out + + def decode(self, ids): + return "".join(self._by_id.get(int(t)) or chr(int(t) - _CHAR_BASE) + for t in ids) + + +def generation_prompt(effort: str) -> str: + """The real template's render of one user turn at ``effort``.""" + from transformers.utils.chat_template_utils import _compile_jinja_template + + return _compile_jinja_template(TEMPLATE).render( + messages=[{"role": "user", "content": "hi"}], + add_generation_prompt=True, bos_token="<|ifm|begin_of_text|>", + reasoning_effort=effort) diff --git a/tests/serve/test_patches_chat_behavior.py b/tests/serve/test_patches_chat_behavior.py index 2a1012248..2fc652ad4 100644 --- a/tests/serve/test_patches_chat_behavior.py +++ b/tests/serve/test_patches_chat_behavior.py @@ -998,3 +998,87 @@ def fake(processor, config, prompt, add_generation_prompt=True, "reasoning_content": "r"}], return_messages=True) assert "thinking" not in msgs[0] + + +# K2 Horizon names one think pair per reasoning_effort. Each request takes +# the pair its rendered prompt opens, so every marker reader agrees. +def _k2_request_args(effort, *, proc, render_out, **args_kw): + gen = importlib.import_module("gmlx._vendor.mlx_vlm.server.generation") + openai_mod = importlib.import_module("gmlx._vendor.mlx_vlm.server.openai") + original_fn = openai_mod.apply_chat_template + rs = importlib.import_module("gmlx._vendor.mlx_vlm.server.responses_state") + original_init = rs.ThinkingStreamState.__init__ + req = types.SimpleNamespace(model_fields_set=set(), + chat_template_kwargs=None, + reasoning_effort=effort) + tok = sp_chat._LAST_GEN_ARGS.set(None) + try: + openai_mod.apply_chat_template = lambda *a, **kw: render_out + sp.install_stream_thinking_seed() + args = sp_chat._stash_template_kwargs( + gen.GenerationArguments(**args_kw), req, proc) + out = openai_mod.apply_chat_template(proc, "config", []) + return args, out + finally: + openai_mod.apply_chat_template = original_fn + rs.ThinkingStreamState.__init__ = original_init + sp_chat._LAST_GEN_ARGS.reset(tok) + sp_chat._LAST_RENDERED_PROMPT.set(None) + + +@pytest.mark.parametrize("effort", ["high", "medium", "low"]) +def test_k2_request_effort_picks_the_server_markers(effort): + import k2_think_tokens as k2 + + tok = k2.K2Tokenizer() + proc = types.SimpleNamespace(chat_template=k2.TEMPLATE, tokenizer=tok) + prompt = k2.generation_prompt(effort) + args, _ = _k2_request_args(effort, proc=proc, render_out=prompt) + start, end = k2.PAIRS[effort] + assert (args.thinking_start_token, args.thinking_end_token) == (start, end) + + rg = types.SimpleNamespace(tokenizer=tok) + ids = tok.encode(prompt) + assert sp_chat._prompt_opens_thinking_tokens(rg, args, mx.array([ids])) + + split = _APP._split_thinking_text + assert split(f"why{end}Hi", start, end) == ("why", "Hi") + + rs = importlib.import_module("gmlx._vendor.mlx_vlm.server.responses_state") + st = rs.ThinkingStreamState(True, start, end) + d1, d2 = st.feed("why"), st.feed(f"{end}Hi", last=True) + assert (d1.reasoning, d2.content) == ("why", "Hi") + + gen = importlib.import_module("gmlx._vendor.mlx_vlm.server.generation") + cls = gen.ResponseGenerator + original = cls._make_thinking_budget_criteria + try: + sp.install_thinking_budget_fix() + args.thinking_budget = 3 + crit = cls._make_thinking_budget_criteria(rg, args, mx.array([ids])) + finally: + cls._make_thinking_budget_criteria = original + out = _atem_serve_run(crit, "weighing it all up", "Hi", tok) + assert out.split(end)[1] == "Hi" + for other in k2.PAIRS.values(): + if other[1] != end: + assert other[1] not in out + + +def test_k2_configured_markers_are_left_alone(): + import k2_think_tokens as k2 + + proc = types.SimpleNamespace(chat_template=k2.TEMPLATE, + tokenizer=k2.K2Tokenizer()) + args, _ = _k2_request_args( + "medium", proc=proc, render_out=k2.generation_prompt("medium"), + thinking_start_token="", thinking_end_token="") + assert (args.thinking_start_token, args.thinking_end_token) == ("", "") + + +def test_single_pair_template_markers_are_left_alone(): + proc = types.SimpleNamespace(chat_template="{{ '' }}") + args, _ = _k2_request_args( + None, proc=proc, render_out="<|im_start|>assistant\n\n") + assert args.thinking_start_token is None + assert args.thinking_end_token is None diff --git a/tests/tui/test_reasoning.py b/tests/tui/test_reasoning.py index 802792149..9b3f0ef0c 100644 --- a/tests/tui/test_reasoning.py +++ b/tests/tui/test_reasoning.py @@ -961,3 +961,24 @@ def test_kimi_templates_keep_the_bare_thinking_switch(name): template = (_K2_TEMPLATE.parent / name).read_text() assert map_thinking_controls({}, "off", None, template) == \ {"thinking": False} + + +@pytest.mark.parametrize("effort", ["high", "medium", "low"]) +def test_k2_reply_splits_at_its_effort_close_tag(effort): + """The template pre-opens the effort's pair, so generation streams only + its close. A tag split across chunks still matches.""" + import k2_think_tokens as k2 + + from gmlx.tui.sessions import split_thinking + + prompt = k2.generation_prompt(effort) + assert chat._opens_thinking(prompt) + close = k2.PAIRS[effort][1] + reason, answer = split_thinking(f"weighing it{close}Hello!", think_open=True) + assert (reason, answer) == ("weighing it", "Hello!") + rf = ReasoningFilter(start_in_thinking=True) + cut = len(close) // 2 + spans = (rf.feed("weighing it" + close[:cut]) + rf.feed(close[cut:] + "Hi") + + rf.flush()) + assert "".join(t for t, m in spans if m == "reason") == "weighing it" + assert "".join(t for t, m in spans if m == "answer") == "Hi" From cdcfef24cfe8ca0084600fb647cad446a8994085 Mon Sep 17 00:00:00 2001 From: Asher Feldman <59994+asher@users.noreply.github.com> Date: Sat, 17 Oct 2026 18:21:21 -0700 Subject: [PATCH 06/10] test(long-context): k2-horizon in the 16k parity and integrity sweep --- tests/gen/test_long_context.py | 1 + 1 file changed, 1 insertion(+) diff --git a/tests/gen/test_long_context.py b/tests/gen/test_long_context.py index aa10be456..47dbbde93 100644 --- a/tests/gen/test_long_context.py +++ b/tests/gen/test_long_context.py @@ -46,6 +46,7 @@ "seed_oss", "smollm3", "granite", "ernie4_5-moe", "minimax-m2", "minimax-m3", "hunyuan-moe", "granitehybrid", "falcon-h1", "qwen3next", "hy_v3", "kimi-k3", "muse-glimmer", "muse-spark", "qwen4exp", "glm5-next", "glm5next", "hyv4", + "k2-horizon", "qwen3vl", ] From b8e1cd800b3291f69408b66c89a0b81d12e9aefc Mon Sep 17 00:00:00 2001 From: Asher Feldman <59994+asher@users.noreply.github.com> Date: Sat, 17 Oct 2026 18:21:21 -0700 Subject: [PATCH 07/10] docs(k2-horizon): changelog, readme and rope options rows for k2 horizon --- CHANGELOG.md | 6 ++++++ README.md | 1 + docs/long-context.md | 1 + 3 files changed, 8 insertions(+) diff --git a/CHANGELOG.md b/CHANGELOG.md index c19c36a6e..ebcc0c3a1 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -23,6 +23,8 @@ adhere to [Semantic Versioning](https://semver.org/spec/v2.0.0.html). and image and audio input through its `--mmproj` companion. - An image LoRA whose PEFT header sets `use_rslora` loads, scaled as PEFT scales it, where it was refused. +- K2 Horizon 7B and K2 Horizon MoVA 36B-A4B (GGUF arch `k2-horizon`), with + `reasoning-low`, `-medium` and `-high` intents and tool calls. ### Changed @@ -39,6 +41,10 @@ adhere to [Semantic Versioning](https://semver.org/spec/v2.0.0.html). ### Fixed +- A chat template with a `{% generation %}` block was refused by + `--chat-template` as invalid Jinja. +- `--thinking off` was silently ignored on a model with no thinking switch + when its template has the word thinking in its text. It now warns. - A streamed image request refused while it ran, such as one too large for memory, got a server error frame where the same request unstreamed got a 400. - The `gmlx pull` plan of an image model could say its parts fit in stages diff --git a/README.md b/README.md index f42f37437..9406dc1cb 100644 --- a/README.md +++ b/README.md @@ -139,6 +139,7 @@ back. - GLM 4 through 5.3, Kimi-K3, MiniMax M2 and M3, and gpt-oss - Hunyuan, Hy3, HY4, Muse Glimmer and Muse Spark - Granite, Nemotron-H and Falcon-H1 +- K2 Horizon, dense and MoVA Each family's output is checked against llama.cpp at 16k context. The [coverage table](https://asher.github.io/gmlx/arch-coverage.html) lists the diff --git a/docs/long-context.md b/docs/long-context.md index f1d913d9d..ed44e14d2 100644 --- a/docs/long-context.md +++ b/docs/long-context.md @@ -79,6 +79,7 @@ flags refuses. | `gpt-oss` | Sliding window layers keep the header's factor 32. | | `granite`, `granitehybrid` | | | `hy_v3`, `hyv4` | | +| `k2-horizon` | | | `llama` | | | `minimax-m2`, `minimax-m3` | | | `mistral3` | | From 9d48e38a9e7cdd194a41bff9e2decc4bfb7a3c46 Mon Sep 17 00:00:00 2001 From: Asher Feldman <59994+asher@users.noreply.github.com> Date: Sat, 17 Oct 2026 18:25:01 -0700 Subject: [PATCH 08/10] perf(k2-horizon): kq router kernel at decode and the fused shared-expert fold --- gmlx/models/k2_horizon/model.py | 55 ++++++++++++++++++++++++++++--- tests/models/test_k2_horizon.py | 58 +++++++++++++++++++++++++++++++++ 2 files changed, 108 insertions(+), 5 deletions(-) diff --git a/gmlx/models/k2_horizon/model.py b/gmlx/models/k2_horizon/model.py index 81866c53a..2a6d20f60 100644 --- a/gmlx/models/k2_horizon/model.py +++ b/gmlx/models/k2_horizon/model.py @@ -45,6 +45,30 @@ _WEIGHT_SUM_FLOOR = 6.103515625e-5 _LN2 = math.log(2.0) +_KQ_ROUTER_STATE = {"ok": None} +# A/B seam: False routes every gate through the eager select. +_KQ_ROUTER_ENABLED = True + + +def _kq_router_ok(n_tokens: int, top_k: int) -> bool: + """Whether the kq router kernel takes this decode-scale call. The + kernel is Metal-only, and its sigmoid arm is docstring-sniffed so an + older kquant build keeps the eager select.""" + if (not _KQ_ROUTER_ENABLED or n_tokens * top_k >= 64 + or mx.default_device() != mx.gpu): + return False + ok = _KQ_ROUTER_STATE["ok"] + if ok is None: + try: + import mlx_kquant as kq + + ok = mx.metal.is_available() and "sigmoid" in ( + getattr(kq.moe_router_topk, "__doc__", "") or "") + except Exception: # noqa: BLE001 - optional dependency + ok = False + _KQ_ROUTER_STATE["ok"] = ok + return ok + @dataclass class ModelArgs(BaseModelArgs): @@ -89,10 +113,23 @@ def _route_weights(scores, inds, norm, scale): return w -def _route(logits, bias, top_k, scoring_func, norm, scale): +def _route(logits, bias, top_k, scoring_func, norm, scale, fused=True): """``(inds, weights)``: top-k on the scores plus ``bias``, weights from the unbiased scores.""" logits = logits.astype(mx.float32) + n = logits.shape[-1] + t = logits.size // n + if fused and scoring_func == "sigmoid" and _kq_router_ok(t, top_k): + # One kq dispatch for the sigmoid, the biased top-k, the renorm and + # the scale. Its renorm guard differs from llama.cpp's clamp only + # when every picked score is below about 1e-5. + import mlx_kquant as kq + + inds, w = kq.moe_router_topk( + logits.reshape(t, n), top_k, norm, shared_gate=False, bias=bias, + scoring="sigmoid", scale=1.0 if scale in (0.0, 1.0) else scale) + shp = logits.shape[:-1] + return inds.reshape(*shp, top_k), w.reshape(*shp, top_k) if scoring_func == "softmax": scores = mx.softmax(logits, axis=-1, precise=True) else: @@ -163,7 +200,7 @@ def __call__(self, x: mx.array) -> mx.array: inds, weights = _route( self.v_router(x.astype(mx.float32)), self.selection_bias, self.top_k, self.scoring_func, self.norm_topk_prob, - self.routed_scaling_factor) + self.routed_scaling_factor, fused=not self.training) h = mx.expand_dims(x, (-2, -3)) do_sort = inds.size >= 64 @@ -256,6 +293,10 @@ def __call__( class MLP(nn.Module): """SwiGLU MLP: the dense layers and the shared expert.""" + # Plain SiLU. The fused shared-expert fold reads this to accept the + # shared expert as one more slot of the expert gathers. + swiglu_limit = 0.0 + def __init__(self, dim: int, hidden: int): super().__init__() self.gate_proj = nn.Linear(dim, hidden, bias=False) @@ -292,7 +333,8 @@ def _logits(self, x: mx.array) -> mx.array: def __call__(self, x: mx.array): return _route( self._logits(x), self.e_score_correction_bias, self.top_k, - self.scoring_func, self.norm_topk_prob, self.routed_scaling_factor) + self.scoring_func, self.norm_topk_prob, self.routed_scaling_factor, + fused=not self.training) def _kq_route_weights(self, x, inds): # Route replay (stream/moe_routes): the weight branch of _route at @@ -322,8 +364,11 @@ def __call__(self, x: mx.array) -> mx.array: y = self.switch_mlp(x, inds, scores) else: y = self.switch_mlp(x, inds) - if y.ndim == scores.ndim + 1: - y = (y * scores[..., None].astype(y.dtype)).sum(-2) + if y.ndim == scores.ndim + 1: + y = (y * scores[..., None].astype(y.dtype)).sum(-2) + elif getattr(self.switch_mlp, "_kq_shexp_mod", None) is not None: + # A mixed return from the stamped shared-expert fold has it. + return y if "shared_experts" in self: y = y + self.shared_experts(x) return y diff --git a/tests/models/test_k2_horizon.py b/tests/models/test_k2_horizon.py index 3bf18b540..bdf4dd167 100644 --- a/tests/models/test_k2_horizon.py +++ b/tests/models/test_k2_horizon.py @@ -463,3 +463,61 @@ def test_lookahead_router_replica_matches_the_gate(): ids, mass = select(x) assert sorted(_np(ids).ravel().tolist()) == sorted(_np(inds).ravel().tolist()) assert (np.diff(_np(mass).ravel()) <= 0).all() + + +_ON_GPU = mx.default_device() == mx.gpu + + +@pytest.mark.skipif(not _ON_GPU, reason="the kq router kernel is Metal-only") +@pytest.mark.parametrize("n, k, norm, scale", [ + (100, 8, True, 2.5), (64, 4, True, 2.5), (64, 4, False, 1.0)]) +def test_fused_router_matches_the_eager_select(n, k, norm, scale): + from gmlx.models.k2_horizon.model import _kq_router_ok, _route + + if not _kq_router_ok(1, k): + pytest.skip("this kquant build has no sigmoid router arm") + rng = np.random.default_rng(n + k) + logits = mx.array(rng.standard_normal((1, 1, n)) * 3, dtype=mx.float32) + bias = mx.array(rng.standard_normal(n) * 0.3, dtype=mx.float32) + fi, fw = _route(logits, bias, k, "sigmoid", norm, scale, fused=True) + ei, ew = _route(logits, bias, k, "sigmoid", norm, scale, fused=False) + fo, eo = np.argsort(_np(fi).ravel()), np.argsort(_np(ei).ravel()) + assert (_np(fi).ravel()[fo] == _np(ei).ravel()[eo]).all() + np.testing.assert_allclose(_np(fw).ravel()[fo], _np(ew).ravel()[eo], + rtol=1e-5) + + +@pytest.mark.skipif(not _ON_GPU, reason="the fused MoE kernels are Metal-only") +def test_shared_expert_fold_matches_the_unfolded_block(): + model, _ = _loader_shaped() + moe = model.layers[1].mlp + se = getattr(moe.switch_mlp, "_kq_shexp_mod", None) + if se is None: + pytest.skip("this kquant build does not fold the shared expert") + assert se is moe.shared_experts + x = mx.array(np.random.default_rng(4).standard_normal((1, 1, 256)), + dtype=mx.bfloat16) + folded = _np(moe(x)) + object.__setattr__(moe.switch_mlp, "_kq_shexp_mod", None) + plain = _np(moe(x)) + assert _rel(folded, plain) < 2e-2 + + +def test_router_kernel_is_never_called_off_the_gpu(monkeypatch): + import mlx_kquant as kq + + from gmlx.models.k2_horizon import model as k2 + + def boom(*a, **k): + raise AssertionError("kq.moe_router_topk called off the GPU") + + monkeypatch.setattr(kq, "moe_router_topk", boom, raising=False) + monkeypatch.setitem(k2._KQ_ROUTER_STATE, "ok", True) + prev = mx.default_device() + mx.set_default_device(mx.cpu) + try: + logits = mx.zeros((1, 1, 64)) + inds, _ = k2._route(logits, None, 4, "sigmoid", True, 2.5) + mx.eval(inds) + finally: + mx.set_default_device(prev) From a4f366f149d7e5f55026288c2269084e00dd5997 Mon Sep 17 00:00:00 2001 From: Asher Feldman <59994+asher@users.noreply.github.com> Date: Sat, 17 Oct 2026 20:38:30 -0700 Subject: [PATCH 09/10] fix(reasoning): end k2 medium and low blocks on the 7B's close --- gmlx/_vendor/manifest.json | 2 +- .../_vendor/mlx_vlm/server/responses_state.py | 9 +++ gmlx/gen/thinking_budget.py | 67 ++++++++++++++++--- gmlx/serve/patches/chat_behavior.py | 2 + tests/gen/test_thinking_budget.py | 43 ++++++++++++ tests/serve/test_patches_chat_behavior.py | 35 ++++++++++ 6 files changed, 149 insertions(+), 9 deletions(-) diff --git a/gmlx/_vendor/manifest.json b/gmlx/_vendor/manifest.json index 3c52b47ae..032e79880 100644 --- a/gmlx/_vendor/manifest.json +++ b/gmlx/_vendor/manifest.json @@ -1525,7 +1525,7 @@ "mlx_vlm/server/realtime.py": {"sha256": "942671328088936df3145eb72ba7e6650d0171c22ca17321d79ad113cd0482b9", "status": "codemod", "upstream": "815998c835bd348649779ce480152c48cf8ff770ca53503f1a1c31649668b416"}, "mlx_vlm/server/request_normalization.py": {"sha256": "38c69ac6dd59534f5b6ad93adf15be145bf299c1aae9af0c9d0216338c006c59", "status": "edited", "upstream": "37eb0727a61eb8a647cf010af20098f388688336d0e9df9f217c2409fc3a00e5"}, "mlx_vlm/server/reranking.py": {"status": "verbatim", "upstream": "36b03fb077741261696795a3b32bfcae64339d0d50d091ee39f2df4b97343179"}, - "mlx_vlm/server/responses_state.py": {"sha256": "58919f35ec6bed9a40390ce45e9af28d65dbdf52e166ace2854afd9267167862", "status": "codemod", "upstream": "138184cb972c43c6f48fc104642608add2605180c1682d3cc1dfb0bd408f3657"}, + "mlx_vlm/server/responses_state.py": {"sha256": "0a195e66a156624c7fc8710c54fb3a782f6a14f7bfa2cd75b59e099e060eebf4", "status": "edited", "upstream": "138184cb972c43c6f48fc104642608add2605180c1682d3cc1dfb0bd408f3657"}, "mlx_vlm/server/runtime.py": {"status": "verbatim", "upstream": "de80f06d985f98270d7f5d16d2c4e41f1f8ecd69c15c9a819ec8e891ff1f83d2"}, "mlx_vlm/server/runtime_config.py": {"sha256": "ae05cb5e399e51700340a7ed31e24f7385f5373325e37fe15425ceee5e570b8a", "status": "codemod", "upstream": "cccded87d36ee60bb302be4898b2a674335448e9eb6838584ca797095e1b4c07"}, "mlx_vlm/server/schemas.py": {"status": "verbatim", "upstream": "80c863e32dd543ca878876e3b183559adfa1a27848d44599ae0852efe74fea68"}, diff --git a/gmlx/_vendor/mlx_vlm/server/responses_state.py b/gmlx/_vendor/mlx_vlm/server/responses_state.py index 6e1f748d2..2880253d1 100644 --- a/gmlx/_vendor/mlx_vlm/server/responses_state.py +++ b/gmlx/_vendor/mlx_vlm/server/responses_state.py @@ -134,9 +134,14 @@ def _build_open_close_markers( thinking_start_token: Optional[str], thinking_end_token: Optional[str], ) -> Tuple[Tuple[str, str], ...]: + from gmlx.gen.thinking_budget import end_aliases + markers = [] if thinking_start_token and thinking_end_token: markers.append((thinking_start_token, thinking_end_token)) + # gmlx: the pair's alias closes end its block too. + for alias in end_aliases(thinking_end_token): + markers.append((thinking_start_token, alias)) for marker_pair in cls._DEFAULT_OPEN_CLOSE_MARKERS: if marker_pair not in markers: markers.append(marker_pair) @@ -391,6 +396,10 @@ def _split_thinking( "Falling back from tokenizer response-template parser", exc_info=True ) + # gmlx: a pair's alias close (thinking_budget._END_ALIASES) splits too. + from gmlx.gen.thinking_budget import normalize_end_aliases + + text = normalize_end_aliases(text, thinking_end_token) for start_marker, end_marker in ThinkingStreamState._build_open_close_markers( thinking_start_token, thinking_end_token ): diff --git a/gmlx/gen/thinking_budget.py b/gmlx/gen/thinking_budget.py index 04c059945..f34e4e0fc 100644 --- a/gmlx/gen/thinking_budget.py +++ b/gmlx/gen/thinking_budget.py @@ -62,6 +62,41 @@ ("", ""), ("", "")) +# Other close tags a model ends a pair's block with. The K2 Horizon 7B ends +# its medium and low blocks with ; MoVA uses each pair's own +# tag. A close is detected on either, and a forced close writes the pair's +# own tag. +_END_ALIASES = { + "": ("",), + "": ("",), +} + + +def end_aliases(end_tag) -> tuple: + """The other close tags that end ``end_tag``'s block.""" + return _END_ALIASES.get(end_tag, ()) if isinstance(end_tag, str) else () + + +def normalize_end_aliases(text, end_tag): + """``text`` with each alias close of ``end_tag`` spelled as ``end_tag``.""" + if not isinstance(text, str): + return text + for alias in end_aliases(end_tag): + text = text.replace(alias, end_tag) + return text + + +def alias_end_seqs(tokenizer, end_seq) -> tuple: + """Token ids of each alias close of the pair ``end_seq`` closes.""" + if not end_seq: + return () + try: + text = tokenizer.decode(list(end_seq)) + except Exception: # noqa: BLE001 - undecodable: no aliases + return () + return tuple(s for s in (marker_ids(tokenizer, a) + for a in end_aliases(text)) if s) + # Forced into the thinking block ahead of the first budget-triggered close. # The model must see itself DECIDE to answer: a bare close tag cuts the # thought mid-sentence, and the unresolved reasoning tends to resume untagged @@ -274,9 +309,13 @@ def __init__(self, tokenizer, thinking_budget, thinking_start_token=thinking_start_token, **kw) self._start_seq = marker_ids(tokenizer, thinking_start_token) self._end_seq = marker_ids(tokenizer, thinking_end_token) + self._end_seqs = (self._end_seq,) + tuple( + s for s in (marker_ids(tokenizer, a) + for a in end_aliases(thinking_end_token)) if s) self._forced_sequence = (self._forced_sequence[:-1] + list(self._end_seq)) - self._tail_max = max(len(self._start_seq), len(self._end_seq)) + self._tail_max = max(len(self._start_seq), + *(len(s) for s in self._end_seqs)) keep = self._tail_max - 1 self._prompt_tail = ([int(t) for t in prompt_ids[-keep:]] if prompt_ids is not None and keep else []) @@ -295,7 +334,7 @@ def __call__(self, token_id): if self.enable_thinking and ends_with_seq(self._tail, self._start_seq): self.in_thinking = True return None - if ends_with_seq(self._tail, self._end_seq): + if any(ends_with_seq(self._tail, s) for s in self._end_seqs): # Also a commentary <|eom|> outside reasoning, where it # changes nothing. self.in_thinking = False @@ -520,8 +559,10 @@ def __init__( eos_ids=None, reclose_ids=None, skip_ids=None, + alt_end_seqs=(), ): self.end_seq = tuple(end_seq) + self.alt_end_seqs = tuple(tuple(s) for s in alt_end_seqs if s) self.start_seq = tuple(start_seq) if start_seq else None self.forced_ids = list(forced_ids) self.reclose_ids = list(reclose_ids) if reclose_ids else self.forced_ids @@ -532,7 +573,8 @@ def __init__( self._close_requested = False self._baseline = None # len(tokens) at first call -> skip the prompt self._tail: list[int] = [] # last tokens seen, for marker matching - self._tail_max = max(len(self.end_seq), len(self.start_seq or ()), 1) + self._tail_max = max(len(self.end_seq), len(self.start_seq or ()), 1, + *(len(s) for s in self.alt_end_seqs)) self._forcing = False self._forced_idx = 0 self._force_seq = self.forced_ids @@ -564,6 +606,10 @@ def _tail_is(self, seq) -> bool: k = len(seq) return k > 0 and len(self._tail) >= k and self._tail[-k:] == list(seq) + def _tail_closes(self) -> bool: + return self._tail_is(self.end_seq) or any( + self._tail_is(s) for s in self.alt_end_seqs) + def _seed_tail(self, tokens) -> None: """Seed the marker tail from the prompt tokens of the first call. mlx-lm hands over only the prompt's last token, so when those tokens @@ -609,8 +655,7 @@ def __call__(self, tokens, logits): self._tail.append(int(tokens[-1].item())) del self._tail[:-self._tail_max] if not self._forcing: - if (self._tail_is(self.end_seq) - and not self._outside_block()): + if self._tail_closes() and not self._outside_block(): self.in_thinking = False if self._landing: self._landing = False # our own forced close arriving @@ -824,6 +869,7 @@ def make_thinking_budget_processor( eos_ids=_eos_id_list(tokenizer) if eos_floor else [], reclose_ids=reclose_ids, skip_ids=skip_ids, + alt_end_seqs=alias_end_seqs(tokenizer, end_seq), ) @@ -851,8 +897,9 @@ class MTPFinishThinking: def __init__(self, *, end_seq, skip_ids, reclose_ids, forced_ids=None, budget=None, start_seq=None, start_in_thinking=False, - prompt_ids=None): + prompt_ids=None, alt_end_seqs=()): self.end_seq = tuple(end_seq) + self.alt_end_seqs = tuple(tuple(s) for s in alt_end_seqs if s) self.start_seq = tuple(start_seq) if start_seq else None self.skip_ids = list(skip_ids) self.reclose_ids = list(reclose_ids) @@ -865,7 +912,8 @@ def __init__(self, *, end_seq, skip_ids, reclose_ids, forced_ids=None, self._spent = False self._answer_pending = False self._strikes = 0 - self._tail_max = max(len(self.end_seq), len(self.start_seq or ()), 1) + self._tail_max = max(len(self.end_seq), len(self.start_seq or ()), 1, + *(len(s) for s in self.alt_end_seqs)) # The prompt's last tokens, so a start marker that the generation # prompt begins (ATEM's "<|start|>assistant") completes on the # generated tokens. Never a whole marker, so the prompt itself @@ -897,7 +945,9 @@ def observe(self, tok) -> None: self._tail.append(int(tok)) if len(self._tail) > self._tail_max: del self._tail[: len(self._tail) - self._tail_max] - if self._tail_is(self.end_seq) and not self._outside_block(): + closes = self._tail_is(self.end_seq) or any( + self._tail_is(s) for s in self.alt_end_seqs) + if closes and not self._outside_block(): self.in_thinking = False if not self._spent: self.done = True # natural close - disarm, as on the plain path @@ -985,6 +1035,7 @@ def make_mtp_finish_hook( start_seq=start_seq, start_in_thinking=start_in_thinking, prompt_ids=prompt_ids, + alt_end_seqs=alias_end_seqs(tokenizer, end_seq), ) diff --git a/gmlx/serve/patches/chat_behavior.py b/gmlx/serve/patches/chat_behavior.py index 27ff0a2b9..a15c15a27 100644 --- a/gmlx/serve/patches/chat_behavior.py +++ b/gmlx/serve/patches/chat_behavior.py @@ -13,6 +13,7 @@ import gmlx.serve.bridge_vlm as serving from gmlx.gen.thinking_budget import ( marker_ids, + normalize_end_aliases, open_think_pair, prompt_opens_block, prompt_tail_ids, @@ -429,6 +430,7 @@ def _split_thinking_keep_indent(rs, cls, text, start_token, end_token): (upstream body seam-fingerprinted via app._split_thinking_text).""" if not text: return None, text + text = normalize_end_aliases(text, end_token) for start_marker, end_marker in cls._build_open_close_markers( start_token, end_token): start = text.find(start_marker) diff --git a/tests/gen/test_thinking_budget.py b/tests/gen/test_thinking_budget.py index 7b4f6aa45..ba379b016 100644 --- a/tests/gen/test_thinking_budget.py +++ b/tests/gen/test_thinking_budget.py @@ -1116,3 +1116,46 @@ def test_k2_mtp_finish_key_closes_with_its_effort_tag(effort): hook = make_mtp_finish_hook(tok, start_in_thinking=True, **kw) hook.request_close() assert hook.take_forced()[-1] == close + + +# The 7B closes medium and low blocks with , not the pair's own +# close tag. +@pytest.mark.parametrize("effort", ["medium", "low"]) +def test_k2_alias_close_disarms_the_budget(effort): + import k2_think_tokens as k2 + + tok = k2.K2Tokenizer() + prompt = k2.generation_prompt(effort) + p = make_thinking_budget_processor( + tok, 4, start_in_thinking=True, prompt=prompt) + filler = 1000 + ord("x") + tokens = tok.encode(prompt) + script = [filler, filler, k2.PIECES[""]] + [filler] * 20 + for t in script: + assert _one_hot_id(p(mx.array(tokens), mx.zeros((1, 260000)))) is None + tokens.append(t) + assert p.done and not p.in_thinking + + +@pytest.mark.parametrize("effort", ["medium", "low"]) +def test_k2_alias_close_disarms_the_mtp_hook(effort): + import k2_think_tokens as k2 + + tok = k2.K2Tokenizer() + prompt = k2.generation_prompt(effort) + hook = make_mtp_finish_hook(tok, budget=4, start_in_thinking=True, + prompt=prompt) + for t in (1000, 1001, k2.PIECES[""]): + hook.observe(t) + assert hook.done and not hook.in_thinking + hook.request_close() + assert hook.take_forced() is None + + +def test_k2_high_close_has_no_alias(): + import k2_think_tokens as k2 + + tok = k2.K2Tokenizer() + p = make_thinking_budget_processor( + tok, 4, start_in_thinking=True, prompt=k2.generation_prompt("high")) + assert p.alt_end_seqs == () diff --git a/tests/serve/test_patches_chat_behavior.py b/tests/serve/test_patches_chat_behavior.py index 2fc652ad4..c0a72ec3f 100644 --- a/tests/serve/test_patches_chat_behavior.py +++ b/tests/serve/test_patches_chat_behavior.py @@ -1082,3 +1082,38 @@ def test_single_pair_template_markers_are_left_alone(): None, proc=proc, render_out="<|im_start|>assistant\n\n") assert args.thinking_start_token is None assert args.thinking_end_token is None + + +# The 7B closes medium and low blocks with , not the pair's own +# close tag. +@pytest.mark.parametrize("effort", ["medium", "low"]) +def test_k2_alias_close_splits_on_the_server(effort): + import k2_think_tokens as k2 + + tok = k2.K2Tokenizer() + proc = types.SimpleNamespace(chat_template=k2.TEMPLATE, tokenizer=tok) + prompt = k2.generation_prompt(effort) + args, _ = _k2_request_args(effort, proc=proc, render_out=prompt) + start, end = k2.PAIRS[effort] + alias = "" + + assert _APP._split_thinking_text(f"why{alias}Hi", start, end) == ("why", "Hi") + rs = importlib.import_module("gmlx._vendor.mlx_vlm.server.responses_state") + assert rs._split_thinking(f"why{alias}Hi", start, end) == ("why", "Hi") + st = rs.ThinkingStreamState(True, start, end) + d1, d2 = st.feed("whyHi", last=True) + assert (d1.reasoning, d2.reasoning, d2.content) == ("why", None, "Hi") + + rg = types.SimpleNamespace(tokenizer=tok) + gen = importlib.import_module("gmlx._vendor.mlx_vlm.server.generation") + cls = gen.ResponseGenerator + original = cls._make_thinking_budget_criteria + try: + sp.install_thinking_budget_fix() + args.thinking_budget = 6 + crit = cls._make_thinking_budget_criteria( + rg, args, mx.array([tok.encode(prompt)])) + finally: + cls._make_thinking_budget_criteria = original + script = f"ok{alias}" + "the answer is long " * 4 + assert _atem_serve_run(crit, script, "", tok) == script From 063293a06f2d4b7518868d0ddd3f6b5a279422a1 Mon Sep 17 00:00:00 2001 From: Asher Feldman <59994+asher@users.noreply.github.com> Date: Sat, 17 Oct 2026 22:59:43 -0700 Subject: [PATCH 10/10] test(hf-cache): run the fetch deadline tests on the fake hub's clock --- gmlx/serve/hf_cache.py | 10 ++++++---- tests/serve/fake_hub.py | 15 ++++++++++----- tests/serve/test_hf_cache.py | 23 +++++++++++------------ 3 files changed, 27 insertions(+), 21 deletions(-) diff --git a/gmlx/serve/hf_cache.py b/gmlx/serve/hf_cache.py index d8038685e..433101ca7 100644 --- a/gmlx/serve/hf_cache.py +++ b/gmlx/serve/hf_cache.py @@ -38,6 +38,8 @@ # Seconds that one call of fetch_pinned may take in all: to connect, and to # read each file it requests. FETCH_TIMEOUT = 5.0 +# The clock the fetch deadline runs on. Tests swap in a fake one. +_clock = time.monotonic # A file larger than this is not fetched. Model cards and module configs # are a few kB. _FETCH_MAX_BYTES = 4 << 20 @@ -103,7 +105,7 @@ def fetch_pinned(hf, model_ref: str, snapshot: str, wanted, *, import httpx from huggingface_hub.utils import build_hf_headers - deadline = time.monotonic() + FETCH_TIMEOUT + deadline = _clock() + FETCH_TIMEOUT tried: set = set() outcome = _GOT answered = False @@ -143,7 +145,7 @@ def _inside(folder: str, name: str) -> str | None: def _fetch_one(hf, client, model_ref, commit, store, snapshot, name, deadline, mark, answered) -> str: """One GET of ``name`` at ``commit``, given up at ``deadline`` (a - ``time.monotonic`` value). ``answered``: an earlier request of the fetch + ``_clock`` value). ``answered``: an earlier request of the fetch was answered. Returns what it ended in: _GOT, _FINAL or _RETRY.""" import httpx @@ -154,7 +156,7 @@ def _fetch_one(hf, client, model_ref, commit, store, snapshot, name, f"{commit}/{quote(name)}") outcome = _FINAL try: - left = deadline - time.monotonic() + left = deadline - _clock() if left <= 0: raise TimeoutError("the fetch ran out of time") with client.stream("GET", url, timeout=left) as r: @@ -169,7 +171,7 @@ def _fetch_one(hf, client, model_ref, commit, store, snapshot, name, body += chunk if len(body) > _FETCH_MAX_BYTES: raise OSError("the file is too large") - if time.monotonic() > deadline: + if _clock() > deadline: raise TimeoutError("the read timed out") etag = _etag(r) outcome = _RETRY # a failure to store it diff --git a/tests/serve/fake_hub.py b/tests/serve/fake_hub.py index e0aaeffc4..776dc1105 100644 --- a/tests/serve/fake_hub.py +++ b/tests/serve/fake_hub.py @@ -9,7 +9,9 @@ request is recorded. :class:`BlackHole` accepts connections and never answers, and :func:`refused_port` gives a port nothing listens on. ``fail_status`` makes the Hub answer every request with that status, and -``delay`` makes it wait that many seconds before each answer. +``delay`` moves the Hub's clock (:meth:`FakeHub.clock`) that many seconds +on before each answer. A test that sets it runs the fetch deadline on that +clock, so each answer lands at a fixed time on any machine. :func:`use_hub` points huggingface_hub at an endpoint and a scratch cache for one test, by the environment and by the constants it read at import, @@ -22,7 +24,6 @@ import os import socket import threading -import time from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer from urllib.parse import unquote, urlparse @@ -37,7 +38,8 @@ def __init__(self): self.requests: list = [] # (method, path) self.tokens: list = [] # the Authorization headers received self.fail_status = None # answer every request with this status - self.delay = 0.0 # seconds to wait before each answer + self.delay = 0.0 # clock seconds before each answer + self.now = 0.0 # the Hub's clock self._server = ThreadingHTTPServer(("127.0.0.1", 0), self._handler()) self._server.daemon_threads = True self._thread = threading.Thread(target=self._server.serve_forever, @@ -50,6 +52,10 @@ def endpoint(self) -> str: host, port = self._server.server_address[:2] return f"http://{host}:{port}" + def clock(self) -> float: + """The Hub's clock, a stand-in for ``time.monotonic``.""" + return self.now + def commit(self, repo: str, sha: str, files: dict, head: bool = True): """Add commit ``sha`` of ``repo`` holding ``files`` (name -> bytes or str), and make it the head when ``head``.""" @@ -100,8 +106,7 @@ def _route(self): hub.requests.append((self.command, path)) if self.headers.get("Authorization"): hub.tokens.append(self.headers["Authorization"]) - if hub.delay: - time.sleep(hub.delay) + hub.now += hub.delay if hub.fail_status: return self._send(hub.fail_status, {}) parts = path.strip("/").split("/") diff --git a/tests/serve/test_hf_cache.py b/tests/serve/test_hf_cache.py index 040f3fc57..96c05992f 100644 --- a/tests/serve/test_hf_cache.py +++ b/tests/serve/test_hf_cache.py @@ -341,23 +341,21 @@ def test_a_failed_connection_leaves_no_mark_and_the_next_start_fetches( def test_one_deadline_bounds_the_whole_fetch(hub, tmp_path, monkeypatch, capsys): - # A Hub that answers each request after 2 s. The first file arrives, + # A Hub that answers each request after 20 s. The first file arrives, # and the second request gets what is left of one FETCH_TIMEOUT, so the - # call ends at about 2.5 s, not 3 x 2 s. The Hub answered a request, so + # call ends there, not after 3 x 20 s. The Hub answered a request, so # the timeout is no answer about the rest: no mark, and the next start # asks for the files still missing. import os - import time - monkeypatch.setattr(hf_cache, "FETCH_TIMEOUT", 2.5) - hub.delay = 2.0 - started = time.monotonic() + monkeypatch.setattr(hf_cache, "FETCH_TIMEOUT", 25.0) + monkeypatch.setattr(hf_cache, "_clock", hub.clock) + hub.delay = 20.0 with hf_cache.offline_resolve(REPO, wanted=_card): pass - assert time.monotonic() - started < 4.0 assert len(hub.requests) == 2 - assert ("modules.json not fetched (ReadTimeout: timed out); the next " - "start tries again") in capsys.readouterr().err + assert ("modules.json not fetched (TimeoutError: the read timed out); " + "the next start tries again") in capsys.readouterr().err assert not os.path.exists(_mark(tmp_path)) hub.requests.clear() with hf_cache.offline_resolve(REPO, wanted=_card): @@ -369,7 +367,7 @@ def test_a_slow_hub_finishes_over_two_starts_and_the_copy_pools_as_upstream( hub, tmp_path, monkeypatch, capsys): # An unnamed bert snapshot holding config.json only, whose upstream # Pooling config takes the first token. A Hub that answers each request - # correctly after 0.8 s gets two of the three files in the 2 s deadline. + # correctly after 8 s gets two of the three files in the 20 s deadline. # The first start leaves no mark, so the second fetches the third file, # then marks the snapshot, and the copy pools as upstream does. import json @@ -388,8 +386,9 @@ def test_a_slow_hub_finishes_over_two_starts_and_the_copy_pools_as_upstream( "README.md": "# acme encoder\n", "1_Pooling/config.json": json.dumps({"pooling_mode_cls_token": True}), }, head=False) - monkeypatch.setattr(hf_cache, "FETCH_TIMEOUT", 2.0) - hub.delay = 0.8 + monkeypatch.setattr(hf_cache, "FETCH_TIMEOUT", 20.0) + monkeypatch.setattr(hf_cache, "_clock", hub.clock) + hub.delay = 8.0 with hf_cache.offline_resolve(REPO, wanted=card_files_to_fetch): snapshot = _window_snapshot(REPO) assert len(hub.requests) == 3