chad-code 1.0.7__tar.gz → 1.0.9__tar.gz
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- {chad_code-1.0.7/src/chad_code.egg-info → chad_code-1.0.9}/PKG-INFO +21 -5
- {chad_code-1.0.7 → chad_code-1.0.9}/README.md +20 -4
- {chad_code-1.0.7 → chad_code-1.0.9}/pyproject.toml +13 -9
- {chad_code-1.0.7 → chad_code-1.0.9}/src/chad/__init__.py +1 -1
- {chad_code-1.0.7 → chad_code-1.0.9}/src/chad/mlx_fastpath.py +55 -4
- chad_code-1.0.9/src/chad/mlx_moe_fused.py +657 -0
- {chad_code-1.0.7 → chad_code-1.0.9}/src/chad/speech.py +22 -1
- {chad_code-1.0.7 → chad_code-1.0.9/src/chad_code.egg-info}/PKG-INFO +21 -5
- {chad_code-1.0.7 → chad_code-1.0.9}/src/chad_code.egg-info/SOURCES.txt +2 -0
- chad_code-1.0.9/tests/test_mlx_moe_fused.py +354 -0
- {chad_code-1.0.7 → chad_code-1.0.9}/tests/test_speech.py +18 -1
- {chad_code-1.0.7 → chad_code-1.0.9}/LICENSE +0 -0
- {chad_code-1.0.7 → chad_code-1.0.9}/setup.cfg +0 -0
- {chad_code-1.0.7 → chad_code-1.0.9}/src/chad/agent.py +0 -0
- {chad_code-1.0.7 → chad_code-1.0.9}/src/chad/atif.py +0 -0
- {chad_code-1.0.7 → chad_code-1.0.9}/src/chad/base_engine.py +0 -0
- {chad_code-1.0.7 → chad_code-1.0.9}/src/chad/bench.py +0 -0
- {chad_code-1.0.7 → chad_code-1.0.9}/src/chad/cli.py +0 -0
- {chad_code-1.0.7 → chad_code-1.0.9}/src/chad/compaction.py +0 -0
- {chad_code-1.0.7 → chad_code-1.0.9}/src/chad/completion_engine.py +0 -0
- {chad_code-1.0.7 → chad_code-1.0.9}/src/chad/config.py +0 -0
- {chad_code-1.0.7 → chad_code-1.0.9}/src/chad/diag.py +0 -0
- {chad_code-1.0.7 → chad_code-1.0.9}/src/chad/engine.py +0 -0
- {chad_code-1.0.7 → chad_code-1.0.9}/src/chad/guardrails.py +0 -0
- {chad_code-1.0.7 → chad_code-1.0.9}/src/chad/ignore.py +0 -0
- {chad_code-1.0.7 → chad_code-1.0.9}/src/chad/levers.py +0 -0
- {chad_code-1.0.7 → chad_code-1.0.9}/src/chad/lsp.py +0 -0
- {chad_code-1.0.7 → chad_code-1.0.9}/src/chad/lspclient.py +0 -0
- {chad_code-1.0.7 → chad_code-1.0.9}/src/chad/lspservers.py +0 -0
- {chad_code-1.0.7 → chad_code-1.0.9}/src/chad/mcp.py +0 -0
- {chad_code-1.0.7 → chad_code-1.0.9}/src/chad/mcp_oauth.py +0 -0
- {chad_code-1.0.7 → chad_code-1.0.9}/src/chad/mlx_qsdpa.py +0 -0
- {chad_code-1.0.7 → chad_code-1.0.9}/src/chad/parakeet/LICENSE +0 -0
- {chad_code-1.0.7 → chad_code-1.0.9}/src/chad/parakeet/__init__.py +0 -0
- {chad_code-1.0.7 → chad_code-1.0.9}/src/chad/parakeet/alignment.py +0 -0
- {chad_code-1.0.7 → chad_code-1.0.9}/src/chad/parakeet/attention.py +0 -0
- {chad_code-1.0.7 → chad_code-1.0.9}/src/chad/parakeet/audio.py +0 -0
- {chad_code-1.0.7 → chad_code-1.0.9}/src/chad/parakeet/cache.py +0 -0
- {chad_code-1.0.7 → chad_code-1.0.9}/src/chad/parakeet/conformer.py +0 -0
- {chad_code-1.0.7 → chad_code-1.0.9}/src/chad/parakeet/ctc.py +0 -0
- {chad_code-1.0.7 → chad_code-1.0.9}/src/chad/parakeet/parakeet.py +0 -0
- {chad_code-1.0.7 → chad_code-1.0.9}/src/chad/parakeet/rnnt.py +0 -0
- {chad_code-1.0.7 → chad_code-1.0.9}/src/chad/parakeet/tokenizer.py +0 -0
- {chad_code-1.0.7 → chad_code-1.0.9}/src/chad/parakeet/utils.py +0 -0
- {chad_code-1.0.7 → chad_code-1.0.9}/src/chad/profiles.py +0 -0
- {chad_code-1.0.7 → chad_code-1.0.9}/src/chad/prompt.py +0 -0
- {chad_code-1.0.7 → chad_code-1.0.9}/src/chad/prove.py +0 -0
- {chad_code-1.0.7 → chad_code-1.0.9}/src/chad/render.py +0 -0
- {chad_code-1.0.7 → chad_code-1.0.9}/src/chad/repomap.py +0 -0
- {chad_code-1.0.7 → chad_code-1.0.9}/src/chad/serve.py +0 -0
- {chad_code-1.0.7 → chad_code-1.0.9}/src/chad/session.py +0 -0
- {chad_code-1.0.7 → chad_code-1.0.9}/src/chad/skills.py +0 -0
- {chad_code-1.0.7 → chad_code-1.0.9}/src/chad/symbols.py +0 -0
- {chad_code-1.0.7 → chad_code-1.0.9}/src/chad/syntaxgate.py +0 -0
- {chad_code-1.0.7 → chad_code-1.0.9}/src/chad/toolcall_parse.py +0 -0
- {chad_code-1.0.7 → chad_code-1.0.9}/src/chad/tools.py +0 -0
- {chad_code-1.0.7 → chad_code-1.0.9}/src/chad/tui.py +0 -0
- {chad_code-1.0.7 → chad_code-1.0.9}/src/chad/validate.py +0 -0
- {chad_code-1.0.7 → chad_code-1.0.9}/src/chad_code.egg-info/dependency_links.txt +0 -0
- {chad_code-1.0.7 → chad_code-1.0.9}/src/chad_code.egg-info/entry_points.txt +0 -0
- {chad_code-1.0.7 → chad_code-1.0.9}/src/chad_code.egg-info/requires.txt +0 -0
- {chad_code-1.0.7 → chad_code-1.0.9}/src/chad_code.egg-info/top_level.txt +0 -0
- {chad_code-1.0.7 → chad_code-1.0.9}/tests/test_agent.py +0 -0
- {chad_code-1.0.7 → chad_code-1.0.9}/tests/test_agent_e2e.py +0 -0
- {chad_code-1.0.7 → chad_code-1.0.9}/tests/test_agent_guards.py +0 -0
- {chad_code-1.0.7 → chad_code-1.0.9}/tests/test_atif.py +0 -0
- {chad_code-1.0.7 → chad_code-1.0.9}/tests/test_bench.py +0 -0
- {chad_code-1.0.7 → chad_code-1.0.9}/tests/test_cli.py +0 -0
- {chad_code-1.0.7 → chad_code-1.0.9}/tests/test_compact_notice.py +0 -0
- {chad_code-1.0.7 → chad_code-1.0.9}/tests/test_compaction.py +0 -0
- {chad_code-1.0.7 → chad_code-1.0.9}/tests/test_completion_engine.py +0 -0
- {chad_code-1.0.7 → chad_code-1.0.9}/tests/test_config.py +0 -0
- {chad_code-1.0.7 → chad_code-1.0.9}/tests/test_confirm_preview.py +0 -0
- {chad_code-1.0.7 → chad_code-1.0.9}/tests/test_done_audit.py +0 -0
- {chad_code-1.0.7 → chad_code-1.0.9}/tests/test_drift_warn.py +0 -0
- {chad_code-1.0.7 → chad_code-1.0.9}/tests/test_edit.py +0 -0
- {chad_code-1.0.7 → chad_code-1.0.9}/tests/test_edit_corruption.py +0 -0
- {chad_code-1.0.7 → chad_code-1.0.9}/tests/test_engine.py +0 -0
- {chad_code-1.0.7 → chad_code-1.0.9}/tests/test_engine_kvquant.py +0 -0
- {chad_code-1.0.7 → chad_code-1.0.9}/tests/test_engine_pld_hybrid.py +0 -0
- {chad_code-1.0.7 → chad_code-1.0.9}/tests/test_feel_pack.py +0 -0
- {chad_code-1.0.7 → chad_code-1.0.9}/tests/test_garble_invariant.py +0 -0
- {chad_code-1.0.7 → chad_code-1.0.9}/tests/test_gate.py +0 -0
- {chad_code-1.0.7 → chad_code-1.0.9}/tests/test_ignore.py +0 -0
- {chad_code-1.0.7 → chad_code-1.0.9}/tests/test_intent.py +0 -0
- {chad_code-1.0.7 → chad_code-1.0.9}/tests/test_lever_bite.py +0 -0
- {chad_code-1.0.7 → chad_code-1.0.9}/tests/test_levers.py +0 -0
- {chad_code-1.0.7 → chad_code-1.0.9}/tests/test_log_redaction.py +0 -0
- {chad_code-1.0.7 → chad_code-1.0.9}/tests/test_lsp.py +0 -0
- {chad_code-1.0.7 → chad_code-1.0.9}/tests/test_lsp_live.py +0 -0
- {chad_code-1.0.7 → chad_code-1.0.9}/tests/test_lspclient.py +0 -0
- {chad_code-1.0.7 → chad_code-1.0.9}/tests/test_mcp.py +0 -0
- {chad_code-1.0.7 → chad_code-1.0.9}/tests/test_mcp_oauth.py +0 -0
- {chad_code-1.0.7 → chad_code-1.0.9}/tests/test_mlx_fastpath.py +0 -0
- {chad_code-1.0.7 → chad_code-1.0.9}/tests/test_mlx_qsdpa.py +0 -0
- {chad_code-1.0.7 → chad_code-1.0.9}/tests/test_plan_review.py +0 -0
- {chad_code-1.0.7 → chad_code-1.0.9}/tests/test_prove.py +0 -0
- {chad_code-1.0.7 → chad_code-1.0.9}/tests/test_render.py +0 -0
- {chad_code-1.0.7 → chad_code-1.0.9}/tests/test_replace_lines.py +0 -0
- {chad_code-1.0.7 → chad_code-1.0.9}/tests/test_repomap.py +0 -0
- {chad_code-1.0.7 → chad_code-1.0.9}/tests/test_repomap_polyglot.py +0 -0
- {chad_code-1.0.7 → chad_code-1.0.9}/tests/test_serve.py +0 -0
- {chad_code-1.0.7 → chad_code-1.0.9}/tests/test_session.py +0 -0
- {chad_code-1.0.7 → chad_code-1.0.9}/tests/test_skills.py +0 -0
- {chad_code-1.0.7 → chad_code-1.0.9}/tests/test_speech_tui.py +0 -0
- {chad_code-1.0.7 → chad_code-1.0.9}/tests/test_subagent.py +0 -0
- {chad_code-1.0.7 → chad_code-1.0.9}/tests/test_symbols.py +0 -0
- {chad_code-1.0.7 → chad_code-1.0.9}/tests/test_syntaxgate.py +0 -0
- {chad_code-1.0.7 → chad_code-1.0.9}/tests/test_toolcall_parse.py +0 -0
- {chad_code-1.0.7 → chad_code-1.0.9}/tests/test_tools.py +0 -0
- {chad_code-1.0.7 → chad_code-1.0.9}/tests/test_tui.py +0 -0
- {chad_code-1.0.7 → chad_code-1.0.9}/tests/test_validate.py +0 -0
|
@@ -1,6 +1,6 @@
|
|
|
1
1
|
Metadata-Version: 2.4
|
|
2
2
|
Name: chad-code
|
|
3
|
-
Version: 1.0.
|
|
3
|
+
Version: 1.0.9
|
|
4
4
|
Summary: Local MLX-backed, Claude-Code-style coding agent (Apple Silicon, Ornith 35B/9B)
|
|
5
5
|
License-Expression: MIT
|
|
6
6
|
Project-URL: Repository, https://github.com/nathansutton/chad
|
|
@@ -139,6 +139,21 @@ uv run chad "add a --json flag to main.py and update the tests" # one-shot, he
|
|
|
139
139
|
uv run chad -c # resume this directory's last conversation
|
|
140
140
|
```
|
|
141
141
|
|
|
142
|
+
**Optional extras.** Everything core is in the base install; two features are opt-in
|
|
143
|
+
because they pull deps not every install wants — `speech` (voice mode: a mic library,
|
|
144
|
+
no torch) and `highlight` (syntax colour in diffs/previews). An extra rides on the
|
|
145
|
+
**install spec**, not on a separate command, so how you add it depends on how you
|
|
146
|
+
installed chad:
|
|
147
|
+
|
|
148
|
+
```bash
|
|
149
|
+
uv tool install --force 'chad-code[speech]' # add to an existing `uv tool` install
|
|
150
|
+
uvx --from 'chad-code[speech]' chad # one-off run, nothing installed
|
|
151
|
+
uv sync --extra speech # from a clone
|
|
152
|
+
```
|
|
153
|
+
|
|
154
|
+
`/speech` in the TUI prints whichever of those matches your install, so you never have
|
|
155
|
+
to work it out from here.
|
|
156
|
+
|
|
142
157
|
**The model.** chad picks one for you by RAM and downloads it once into the shared Hugging
|
|
143
158
|
Face cache (`~/.cache/huggingface`, reused across every project). Override with
|
|
144
159
|
`--model 9b` / `--model 35b`, or `--model <repo or local dir>` for anything else.
|
|
@@ -164,8 +179,8 @@ re-downloads the model.
|
|
|
164
179
|
**Development.** `uv sync` once, then `uv run pytest -q` — the fast unit gate loads **no
|
|
165
180
|
model weights**, runs in seconds, and is what CI runs. Throughput on your own machine:
|
|
166
181
|
`uv run chad-bench` (see [Throughput & performance](docs/benchmarks.md)). LSP-precise
|
|
167
|
-
find-references / rename need
|
|
168
|
-
|
|
182
|
+
find-references / rename need no extra — chad fetches pyright via `uvx` on first use and
|
|
183
|
+
falls back to tree-sitter when a language server can't start.
|
|
169
184
|
|
|
170
185
|
## Interactive UX
|
|
171
186
|
|
|
@@ -191,8 +206,9 @@ it chad uses the tree-sitter fallback automatically.
|
|
|
191
206
|
isn't clipped, and a personal word table (`~/.chad/speech_words.json`) teaches it your
|
|
192
207
|
identifiers — `{"pie test": "pytest"}`. Dictation cost is linear in take length, so a
|
|
193
208
|
long thought is fine; `/speech` off releases both the mic and the weights. Nothing
|
|
194
|
-
leaves the machine. Needs the `speech` extra
|
|
195
|
-
|
|
209
|
+
leaves the machine. Needs the `speech` extra — just a mic library; no torch, no numba
|
|
210
|
+
(see [Installing & upgrading](#installing--upgrading); on a `uv tool` install that's
|
|
211
|
+
`uv tool install --force 'chad-code[speech]'`).
|
|
196
212
|
|
|
197
213
|
**Usage.** `uv run chad --help` is the source of truth:
|
|
198
214
|
|
|
@@ -103,6 +103,21 @@ uv run chad "add a --json flag to main.py and update the tests" # one-shot, he
|
|
|
103
103
|
uv run chad -c # resume this directory's last conversation
|
|
104
104
|
```
|
|
105
105
|
|
|
106
|
+
**Optional extras.** Everything core is in the base install; two features are opt-in
|
|
107
|
+
because they pull deps not every install wants — `speech` (voice mode: a mic library,
|
|
108
|
+
no torch) and `highlight` (syntax colour in diffs/previews). An extra rides on the
|
|
109
|
+
**install spec**, not on a separate command, so how you add it depends on how you
|
|
110
|
+
installed chad:
|
|
111
|
+
|
|
112
|
+
```bash
|
|
113
|
+
uv tool install --force 'chad-code[speech]' # add to an existing `uv tool` install
|
|
114
|
+
uvx --from 'chad-code[speech]' chad # one-off run, nothing installed
|
|
115
|
+
uv sync --extra speech # from a clone
|
|
116
|
+
```
|
|
117
|
+
|
|
118
|
+
`/speech` in the TUI prints whichever of those matches your install, so you never have
|
|
119
|
+
to work it out from here.
|
|
120
|
+
|
|
106
121
|
**The model.** chad picks one for you by RAM and downloads it once into the shared Hugging
|
|
107
122
|
Face cache (`~/.cache/huggingface`, reused across every project). Override with
|
|
108
123
|
`--model 9b` / `--model 35b`, or `--model <repo or local dir>` for anything else.
|
|
@@ -128,8 +143,8 @@ re-downloads the model.
|
|
|
128
143
|
**Development.** `uv sync` once, then `uv run pytest -q` — the fast unit gate loads **no
|
|
129
144
|
model weights**, runs in seconds, and is what CI runs. Throughput on your own machine:
|
|
130
145
|
`uv run chad-bench` (see [Throughput & performance](docs/benchmarks.md)). LSP-precise
|
|
131
|
-
find-references / rename need
|
|
132
|
-
|
|
146
|
+
find-references / rename need no extra — chad fetches pyright via `uvx` on first use and
|
|
147
|
+
falls back to tree-sitter when a language server can't start.
|
|
133
148
|
|
|
134
149
|
## Interactive UX
|
|
135
150
|
|
|
@@ -155,8 +170,9 @@ it chad uses the tree-sitter fallback automatically.
|
|
|
155
170
|
isn't clipped, and a personal word table (`~/.chad/speech_words.json`) teaches it your
|
|
156
171
|
identifiers — `{"pie test": "pytest"}`. Dictation cost is linear in take length, so a
|
|
157
172
|
long thought is fine; `/speech` off releases both the mic and the weights. Nothing
|
|
158
|
-
leaves the machine. Needs the `speech` extra
|
|
159
|
-
|
|
173
|
+
leaves the machine. Needs the `speech` extra — just a mic library; no torch, no numba
|
|
174
|
+
(see [Installing & upgrading](#installing--upgrading); on a `uv tool` install that's
|
|
175
|
+
`uv tool install --force 'chad-code[speech]'`).
|
|
160
176
|
|
|
161
177
|
**Usage.** `uv run chad --help` is the source of truth:
|
|
162
178
|
|
|
@@ -4,7 +4,7 @@
|
|
|
4
4
|
# import name, and command name are independent. `uvx chad-code` runs the alias
|
|
5
5
|
# script added under [project.scripts].
|
|
6
6
|
name = "chad-code"
|
|
7
|
-
version = "1.0.
|
|
7
|
+
version = "1.0.9"
|
|
8
8
|
description = "Local MLX-backed, Claude-Code-style coding agent (Apple Silicon, Ornith 35B/9B)"
|
|
9
9
|
readme = "README.md"
|
|
10
10
|
license = "MIT"
|
|
@@ -72,11 +72,15 @@ dependencies = [
|
|
|
72
72
|
"mcp>=2,<3",
|
|
73
73
|
]
|
|
74
74
|
|
|
75
|
+
# Extras ride on the INSTALL SPEC, so the command depends on how chad was installed:
|
|
76
|
+
# `uv tool install --force 'chad-code[<extra>]'` (or `uvx --from 'chad-code[<extra>]' chad`)
|
|
77
|
+
# for a PyPI install, `uv sync --extra <extra>` only in a clone. speech.install_hint()
|
|
78
|
+
# derives the right one at runtime — keep it in step with any extra added here.
|
|
79
|
+
# NOTE: the PyPI distribution is `chad-code`; 'pip install chad' installs an unrelated
|
|
80
|
+
# squatted package.
|
|
81
|
+
#
|
|
75
82
|
# Optional syntax highlighting in diffs / confirm previews. Pure-Python,
|
|
76
83
|
# import-guarded in render.py: absent, output is byte-identical to the un-highlighted path.
|
|
77
|
-
# Install with `uv sync --extra highlight`.
|
|
78
|
-
# NOTE: no PyPI release exists; 'pip install chad' installs an unrelated squatted
|
|
79
|
-
# package. Document uv/git installs only.
|
|
80
84
|
[project.optional-dependencies]
|
|
81
85
|
highlight = ["pygments>=2.17"]
|
|
82
86
|
# Voice mode (/speech in the TUI): push-to-talk dictation via chad's vendored
|
|
@@ -85,7 +89,7 @@ highlight = ["pygments>=2.17"]
|
|
|
85
89
|
# `say`. All on-device, nothing leaves the machine. mlx/numpy/huggingface-hub
|
|
86
90
|
# are already core deps, so the extra adds only the mic (sounddevice/PortAudio).
|
|
87
91
|
# Optional and import-guarded in speech.py: absent, /speech prints the install
|
|
88
|
-
# hint and everything else is untouched.
|
|
92
|
+
# hint and everything else is untouched.
|
|
89
93
|
speech = [
|
|
90
94
|
"sounddevice>=0.5",
|
|
91
95
|
]
|
|
@@ -170,10 +174,10 @@ module = "chad.parakeet.*"
|
|
|
170
174
|
ignore_errors = true
|
|
171
175
|
|
|
172
176
|
[[tool.mypy.overrides]]
|
|
173
|
-
#
|
|
174
|
-
# (dev Macs, the macos CI jobs)
|
|
177
|
+
# These modules patch mlx_lm classes that only exist on macOS: with mlx installed
|
|
178
|
+
# (dev Macs, the macos CI jobs) their `type: ignore[method-assign]` comments are load-
|
|
175
179
|
# bearing; on the Linux lint runner mlx_lm is absent, the patched classes type as Any,
|
|
176
180
|
# and warn_unused_ignores flags those same comments — unfixable in the code for both
|
|
177
|
-
# platforms at once, so the unused-ignore warning alone is relaxed for
|
|
178
|
-
module = "chad.mlx_fastpath"
|
|
181
|
+
# platforms at once, so the unused-ignore warning alone is relaxed for them.
|
|
182
|
+
module = ["chad.mlx_fastpath", "chad.mlx_moe_fused"]
|
|
179
183
|
warn_unused_ignores = false
|
|
@@ -4,7 +4,7 @@ A flat collection of cooperating modules behind one console script (``chad``):
|
|
|
4
4
|
the inference engine, the tool layer, the agent loop, and the terminal UI.
|
|
5
5
|
"""
|
|
6
6
|
|
|
7
|
-
__version__ = "1.0.
|
|
7
|
+
__version__ = "1.0.9"
|
|
8
8
|
|
|
9
9
|
# chad sets no MLX_* runtime vars. MLX_METAL_FAST_SYNCH, MLX_MAX_OPS_PER_BUFFER
|
|
10
10
|
# and MLX_MAX_MB_PER_BUFFER were each measured end-to-end on the 35B and every
|
|
@@ -61,6 +61,10 @@ def install(model: Any) -> bool:
|
|
|
61
61
|
return False
|
|
62
62
|
_concat_expert_gate_up(model)
|
|
63
63
|
_concat_gdn_in_projs(model)
|
|
64
|
+
# Fused MoE decode kernels + router|seg / shared gate|up concats;
|
|
65
|
+
# engages only on the exact 35B geometry, silent no-op elsewhere.
|
|
66
|
+
from . import mlx_moe_fused
|
|
67
|
+
mlx_moe_fused.install(model)
|
|
64
68
|
_install_layer_fastpath(model)
|
|
65
69
|
log.info("FASTPATH installed: fused expert/GDN projections + compiled "
|
|
66
70
|
"S=1 layer step")
|
|
@@ -129,6 +133,8 @@ def _patch_switch_glu() -> None:
|
|
|
129
133
|
import mlx.core as mx
|
|
130
134
|
from mlx_lm.models import switch_layers as sl
|
|
131
135
|
|
|
136
|
+
if getattr(sl.SwitchGLU.__call__, "_chad_fastpath", False):
|
|
137
|
+
return # already patched; re-wrapping would stack guards
|
|
132
138
|
stock_call = sl.SwitchGLU.__call__
|
|
133
139
|
|
|
134
140
|
def fused_call(self, x, indices):
|
|
@@ -151,6 +157,7 @@ def _patch_switch_glu() -> None:
|
|
|
151
157
|
y = sl._scatter_unsort(y, inv_order, indices.shape)
|
|
152
158
|
return y.squeeze(-2)
|
|
153
159
|
|
|
160
|
+
fused_call._chad_fastpath = True # type: ignore[attr-defined]
|
|
154
161
|
sl.SwitchGLU.__call__ = fused_call # type: ignore[method-assign]
|
|
155
162
|
|
|
156
163
|
|
|
@@ -187,6 +194,8 @@ def _patch_gdn_call() -> None:
|
|
|
187
194
|
import mlx.nn as nn
|
|
188
195
|
from mlx_lm.models import qwen3_5 as q35
|
|
189
196
|
|
|
197
|
+
if getattr(q35.GatedDeltaNet.__call__, "_chad_fastpath", False):
|
|
198
|
+
return
|
|
190
199
|
stock_call = q35.GatedDeltaNet.__call__
|
|
191
200
|
|
|
192
201
|
def call(self, inputs, mask=None, cache=None):
|
|
@@ -236,6 +245,7 @@ def _patch_gdn_call() -> None:
|
|
|
236
245
|
out = self.norm(out, z)
|
|
237
246
|
return self.out_proj(out.reshape(B, S, -1))
|
|
238
247
|
|
|
248
|
+
call._chad_fastpath = True # type: ignore[attr-defined]
|
|
239
249
|
q35.GatedDeltaNet.__call__ = call # type: ignore[method-assign]
|
|
240
250
|
|
|
241
251
|
|
|
@@ -246,13 +256,32 @@ def _install_layer_fastpath(model) -> None:
|
|
|
246
256
|
from mlx_lm.models import qwen3_5 as q35
|
|
247
257
|
|
|
248
258
|
for layer in model.language_model.model.layers:
|
|
249
|
-
|
|
259
|
+
body = _compile_moe_step(layer)
|
|
260
|
+
fused = getattr(layer.mlp, "_fused_step", None)
|
|
261
|
+
if fused is not None:
|
|
262
|
+
# mlx_moe_fused kernels handle the single-token (1,1,2048) decode
|
|
263
|
+
# step; anything else (batched decode, foreign width) takes the
|
|
264
|
+
# compiled stock-graph body. Pinning the exact element count keeps
|
|
265
|
+
# a future geometry change from reaching kernels that hardcode it.
|
|
266
|
+
from .mlx_moe_fused import HID as moe_hid
|
|
267
|
+
layer._moe_body_fast = body
|
|
268
|
+
layer._moe_fast = (lambda h, f=fused, b=body, n=moe_hid:
|
|
269
|
+
f(h) if h.size == n else b(h))
|
|
270
|
+
else:
|
|
271
|
+
layer._moe_fast = body
|
|
250
272
|
if layer.is_linear and hasattr(layer.linear_attn, "_fused_w"):
|
|
251
273
|
layer._gdn_fast = _compile_gdn_step(layer)
|
|
252
274
|
|
|
275
|
+
if getattr(q35.DecoderLayer.__call__, "_chad_fastpath", False):
|
|
276
|
+
return
|
|
253
277
|
stock_layer_call = q35.DecoderLayer.__call__
|
|
254
278
|
|
|
255
279
|
def layer_call(self, x, mask=None, cache=None):
|
|
280
|
+
# S==1 cannot distinguish decode from a stray 1-token prefill chunk
|
|
281
|
+
# (mask is None for both); such a chunk takes this path and pays one
|
|
282
|
+
# token's worth of the same rounding class decode itself injects into
|
|
283
|
+
# the cache every step — bounded, unlike S>1 compile-fusion drift,
|
|
284
|
+
# which is why prefill proper stays on the stock graph.
|
|
256
285
|
if mask is None and x.shape[1] == 1 and cache is not None \
|
|
257
286
|
and getattr(self, "_moe_fast", None) is not None:
|
|
258
287
|
if self.is_linear:
|
|
@@ -270,6 +299,7 @@ def _install_layer_fastpath(model) -> None:
|
|
|
270
299
|
return self._moe_fast(x + r)
|
|
271
300
|
return stock_layer_call(self, x, mask=mask, cache=cache)
|
|
272
301
|
|
|
302
|
+
layer_call._chad_fastpath = True # type: ignore[attr-defined]
|
|
273
303
|
q35.DecoderLayer.__call__ = layer_call # type: ignore[method-assign]
|
|
274
304
|
|
|
275
305
|
|
|
@@ -300,14 +330,22 @@ def _moe_body(layer):
|
|
|
300
330
|
sw = mlp.switch_mlp
|
|
301
331
|
ln_w = layer.post_attention_layernorm.weight
|
|
302
332
|
ln_eps = layer.post_attention_layernorm.eps
|
|
303
|
-
gate = mlp.gate
|
|
304
333
|
fw, fs, fb = sw._fused_w, sw._fused_s, sw._fused_b
|
|
305
334
|
sgs, sbits = sw._fused_gs, sw._fused_bits
|
|
306
335
|
dp = sw.down_proj
|
|
307
336
|
se = mlp.shared_expert
|
|
308
|
-
seg = mlp.shared_expert_gate
|
|
309
337
|
k = mlp.top_k
|
|
310
338
|
norm_topk = mlp.norm_topk_prob
|
|
339
|
+
E = mlp.num_experts
|
|
340
|
+
# mlx_moe_fused may have concatenated router|seg and shared gate|up; this
|
|
341
|
+
# body must then read the fused copies (originals are placeholders). Both
|
|
342
|
+
# variants are the same math per row.
|
|
343
|
+
moek = hasattr(mlp, "_rt_w")
|
|
344
|
+
if moek:
|
|
345
|
+
rt = (mlp._rt_w, mlp._rt_s, mlp._rt_b)
|
|
346
|
+
shg = (mlp._sh_w, mlp._sh_s, mlp._sh_b)
|
|
347
|
+
else:
|
|
348
|
+
gate, seg = mlp.gate, mlp.shared_expert_gate
|
|
311
349
|
|
|
312
350
|
def qmm(x, m):
|
|
313
351
|
return mx.quantized_matmul(x, m.weight, scales=m.scales, biases=m.biases,
|
|
@@ -316,7 +354,13 @@ def _moe_body(layer):
|
|
|
316
354
|
|
|
317
355
|
def fwd(h):
|
|
318
356
|
x = mx.fast.rms_norm(h, ln_w, ln_eps)
|
|
319
|
-
|
|
357
|
+
if moek:
|
|
358
|
+
rg = mx.quantized_matmul(x, rt[0], scales=rt[1], biases=rt[2],
|
|
359
|
+
transpose=True, group_size=mlp._rt_gs,
|
|
360
|
+
bits=mlp._rt_bits)
|
|
361
|
+
gates = mx.softmax(rg[..., :E], axis=-1, precise=True)
|
|
362
|
+
else:
|
|
363
|
+
gates = mx.softmax(qmm(x, gate), axis=-1, precise=True)
|
|
320
364
|
inds = mx.argpartition(gates, kth=-k, axis=-1)[..., -k:]
|
|
321
365
|
scores = mx.take_along_axis(gates, inds, axis=-1)
|
|
322
366
|
if norm_topk:
|
|
@@ -329,6 +373,13 @@ def _moe_body(layer):
|
|
|
329
373
|
y = mx.gather_qmm(hh, dp.weight, dp.scales, dp.biases, rhs_indices=inds,
|
|
330
374
|
transpose=True, group_size=dp.group_size, bits=dp.bits)
|
|
331
375
|
y = (y.squeeze(-2) * scores[..., None]).sum(axis=-2)
|
|
376
|
+
if moek:
|
|
377
|
+
sgu = mx.quantized_matmul(x, shg[0], scales=shg[1], biases=shg[2],
|
|
378
|
+
transpose=True, group_size=mlp._sh_gs,
|
|
379
|
+
bits=mlp._sh_bits)
|
|
380
|
+
g, u = mx.split(sgu, 2, axis=-1)
|
|
381
|
+
sh = qmm(nn.silu(g) * u, se.down_proj)
|
|
382
|
+
return h + y + mx.sigmoid(rg[..., E:E + 1]) * sh
|
|
332
383
|
sh = qmm(nn.silu(qmm(x, se.gate_proj)) * qmm(x, se.up_proj), se.down_proj)
|
|
333
384
|
return h + y + mx.sigmoid(qmm(x, seg)) * sh
|
|
334
385
|
|