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.
Files changed (112) hide show
  1. {chad_code-1.0.7/src/chad_code.egg-info → chad_code-1.0.9}/PKG-INFO +21 -5
  2. {chad_code-1.0.7 → chad_code-1.0.9}/README.md +20 -4
  3. {chad_code-1.0.7 → chad_code-1.0.9}/pyproject.toml +13 -9
  4. {chad_code-1.0.7 → chad_code-1.0.9}/src/chad/__init__.py +1 -1
  5. {chad_code-1.0.7 → chad_code-1.0.9}/src/chad/mlx_fastpath.py +55 -4
  6. chad_code-1.0.9/src/chad/mlx_moe_fused.py +657 -0
  7. {chad_code-1.0.7 → chad_code-1.0.9}/src/chad/speech.py +22 -1
  8. {chad_code-1.0.7 → chad_code-1.0.9/src/chad_code.egg-info}/PKG-INFO +21 -5
  9. {chad_code-1.0.7 → chad_code-1.0.9}/src/chad_code.egg-info/SOURCES.txt +2 -0
  10. chad_code-1.0.9/tests/test_mlx_moe_fused.py +354 -0
  11. {chad_code-1.0.7 → chad_code-1.0.9}/tests/test_speech.py +18 -1
  12. {chad_code-1.0.7 → chad_code-1.0.9}/LICENSE +0 -0
  13. {chad_code-1.0.7 → chad_code-1.0.9}/setup.cfg +0 -0
  14. {chad_code-1.0.7 → chad_code-1.0.9}/src/chad/agent.py +0 -0
  15. {chad_code-1.0.7 → chad_code-1.0.9}/src/chad/atif.py +0 -0
  16. {chad_code-1.0.7 → chad_code-1.0.9}/src/chad/base_engine.py +0 -0
  17. {chad_code-1.0.7 → chad_code-1.0.9}/src/chad/bench.py +0 -0
  18. {chad_code-1.0.7 → chad_code-1.0.9}/src/chad/cli.py +0 -0
  19. {chad_code-1.0.7 → chad_code-1.0.9}/src/chad/compaction.py +0 -0
  20. {chad_code-1.0.7 → chad_code-1.0.9}/src/chad/completion_engine.py +0 -0
  21. {chad_code-1.0.7 → chad_code-1.0.9}/src/chad/config.py +0 -0
  22. {chad_code-1.0.7 → chad_code-1.0.9}/src/chad/diag.py +0 -0
  23. {chad_code-1.0.7 → chad_code-1.0.9}/src/chad/engine.py +0 -0
  24. {chad_code-1.0.7 → chad_code-1.0.9}/src/chad/guardrails.py +0 -0
  25. {chad_code-1.0.7 → chad_code-1.0.9}/src/chad/ignore.py +0 -0
  26. {chad_code-1.0.7 → chad_code-1.0.9}/src/chad/levers.py +0 -0
  27. {chad_code-1.0.7 → chad_code-1.0.9}/src/chad/lsp.py +0 -0
  28. {chad_code-1.0.7 → chad_code-1.0.9}/src/chad/lspclient.py +0 -0
  29. {chad_code-1.0.7 → chad_code-1.0.9}/src/chad/lspservers.py +0 -0
  30. {chad_code-1.0.7 → chad_code-1.0.9}/src/chad/mcp.py +0 -0
  31. {chad_code-1.0.7 → chad_code-1.0.9}/src/chad/mcp_oauth.py +0 -0
  32. {chad_code-1.0.7 → chad_code-1.0.9}/src/chad/mlx_qsdpa.py +0 -0
  33. {chad_code-1.0.7 → chad_code-1.0.9}/src/chad/parakeet/LICENSE +0 -0
  34. {chad_code-1.0.7 → chad_code-1.0.9}/src/chad/parakeet/__init__.py +0 -0
  35. {chad_code-1.0.7 → chad_code-1.0.9}/src/chad/parakeet/alignment.py +0 -0
  36. {chad_code-1.0.7 → chad_code-1.0.9}/src/chad/parakeet/attention.py +0 -0
  37. {chad_code-1.0.7 → chad_code-1.0.9}/src/chad/parakeet/audio.py +0 -0
  38. {chad_code-1.0.7 → chad_code-1.0.9}/src/chad/parakeet/cache.py +0 -0
  39. {chad_code-1.0.7 → chad_code-1.0.9}/src/chad/parakeet/conformer.py +0 -0
  40. {chad_code-1.0.7 → chad_code-1.0.9}/src/chad/parakeet/ctc.py +0 -0
  41. {chad_code-1.0.7 → chad_code-1.0.9}/src/chad/parakeet/parakeet.py +0 -0
  42. {chad_code-1.0.7 → chad_code-1.0.9}/src/chad/parakeet/rnnt.py +0 -0
  43. {chad_code-1.0.7 → chad_code-1.0.9}/src/chad/parakeet/tokenizer.py +0 -0
  44. {chad_code-1.0.7 → chad_code-1.0.9}/src/chad/parakeet/utils.py +0 -0
  45. {chad_code-1.0.7 → chad_code-1.0.9}/src/chad/profiles.py +0 -0
  46. {chad_code-1.0.7 → chad_code-1.0.9}/src/chad/prompt.py +0 -0
  47. {chad_code-1.0.7 → chad_code-1.0.9}/src/chad/prove.py +0 -0
  48. {chad_code-1.0.7 → chad_code-1.0.9}/src/chad/render.py +0 -0
  49. {chad_code-1.0.7 → chad_code-1.0.9}/src/chad/repomap.py +0 -0
  50. {chad_code-1.0.7 → chad_code-1.0.9}/src/chad/serve.py +0 -0
  51. {chad_code-1.0.7 → chad_code-1.0.9}/src/chad/session.py +0 -0
  52. {chad_code-1.0.7 → chad_code-1.0.9}/src/chad/skills.py +0 -0
  53. {chad_code-1.0.7 → chad_code-1.0.9}/src/chad/symbols.py +0 -0
  54. {chad_code-1.0.7 → chad_code-1.0.9}/src/chad/syntaxgate.py +0 -0
  55. {chad_code-1.0.7 → chad_code-1.0.9}/src/chad/toolcall_parse.py +0 -0
  56. {chad_code-1.0.7 → chad_code-1.0.9}/src/chad/tools.py +0 -0
  57. {chad_code-1.0.7 → chad_code-1.0.9}/src/chad/tui.py +0 -0
  58. {chad_code-1.0.7 → chad_code-1.0.9}/src/chad/validate.py +0 -0
  59. {chad_code-1.0.7 → chad_code-1.0.9}/src/chad_code.egg-info/dependency_links.txt +0 -0
  60. {chad_code-1.0.7 → chad_code-1.0.9}/src/chad_code.egg-info/entry_points.txt +0 -0
  61. {chad_code-1.0.7 → chad_code-1.0.9}/src/chad_code.egg-info/requires.txt +0 -0
  62. {chad_code-1.0.7 → chad_code-1.0.9}/src/chad_code.egg-info/top_level.txt +0 -0
  63. {chad_code-1.0.7 → chad_code-1.0.9}/tests/test_agent.py +0 -0
  64. {chad_code-1.0.7 → chad_code-1.0.9}/tests/test_agent_e2e.py +0 -0
  65. {chad_code-1.0.7 → chad_code-1.0.9}/tests/test_agent_guards.py +0 -0
  66. {chad_code-1.0.7 → chad_code-1.0.9}/tests/test_atif.py +0 -0
  67. {chad_code-1.0.7 → chad_code-1.0.9}/tests/test_bench.py +0 -0
  68. {chad_code-1.0.7 → chad_code-1.0.9}/tests/test_cli.py +0 -0
  69. {chad_code-1.0.7 → chad_code-1.0.9}/tests/test_compact_notice.py +0 -0
  70. {chad_code-1.0.7 → chad_code-1.0.9}/tests/test_compaction.py +0 -0
  71. {chad_code-1.0.7 → chad_code-1.0.9}/tests/test_completion_engine.py +0 -0
  72. {chad_code-1.0.7 → chad_code-1.0.9}/tests/test_config.py +0 -0
  73. {chad_code-1.0.7 → chad_code-1.0.9}/tests/test_confirm_preview.py +0 -0
  74. {chad_code-1.0.7 → chad_code-1.0.9}/tests/test_done_audit.py +0 -0
  75. {chad_code-1.0.7 → chad_code-1.0.9}/tests/test_drift_warn.py +0 -0
  76. {chad_code-1.0.7 → chad_code-1.0.9}/tests/test_edit.py +0 -0
  77. {chad_code-1.0.7 → chad_code-1.0.9}/tests/test_edit_corruption.py +0 -0
  78. {chad_code-1.0.7 → chad_code-1.0.9}/tests/test_engine.py +0 -0
  79. {chad_code-1.0.7 → chad_code-1.0.9}/tests/test_engine_kvquant.py +0 -0
  80. {chad_code-1.0.7 → chad_code-1.0.9}/tests/test_engine_pld_hybrid.py +0 -0
  81. {chad_code-1.0.7 → chad_code-1.0.9}/tests/test_feel_pack.py +0 -0
  82. {chad_code-1.0.7 → chad_code-1.0.9}/tests/test_garble_invariant.py +0 -0
  83. {chad_code-1.0.7 → chad_code-1.0.9}/tests/test_gate.py +0 -0
  84. {chad_code-1.0.7 → chad_code-1.0.9}/tests/test_ignore.py +0 -0
  85. {chad_code-1.0.7 → chad_code-1.0.9}/tests/test_intent.py +0 -0
  86. {chad_code-1.0.7 → chad_code-1.0.9}/tests/test_lever_bite.py +0 -0
  87. {chad_code-1.0.7 → chad_code-1.0.9}/tests/test_levers.py +0 -0
  88. {chad_code-1.0.7 → chad_code-1.0.9}/tests/test_log_redaction.py +0 -0
  89. {chad_code-1.0.7 → chad_code-1.0.9}/tests/test_lsp.py +0 -0
  90. {chad_code-1.0.7 → chad_code-1.0.9}/tests/test_lsp_live.py +0 -0
  91. {chad_code-1.0.7 → chad_code-1.0.9}/tests/test_lspclient.py +0 -0
  92. {chad_code-1.0.7 → chad_code-1.0.9}/tests/test_mcp.py +0 -0
  93. {chad_code-1.0.7 → chad_code-1.0.9}/tests/test_mcp_oauth.py +0 -0
  94. {chad_code-1.0.7 → chad_code-1.0.9}/tests/test_mlx_fastpath.py +0 -0
  95. {chad_code-1.0.7 → chad_code-1.0.9}/tests/test_mlx_qsdpa.py +0 -0
  96. {chad_code-1.0.7 → chad_code-1.0.9}/tests/test_plan_review.py +0 -0
  97. {chad_code-1.0.7 → chad_code-1.0.9}/tests/test_prove.py +0 -0
  98. {chad_code-1.0.7 → chad_code-1.0.9}/tests/test_render.py +0 -0
  99. {chad_code-1.0.7 → chad_code-1.0.9}/tests/test_replace_lines.py +0 -0
  100. {chad_code-1.0.7 → chad_code-1.0.9}/tests/test_repomap.py +0 -0
  101. {chad_code-1.0.7 → chad_code-1.0.9}/tests/test_repomap_polyglot.py +0 -0
  102. {chad_code-1.0.7 → chad_code-1.0.9}/tests/test_serve.py +0 -0
  103. {chad_code-1.0.7 → chad_code-1.0.9}/tests/test_session.py +0 -0
  104. {chad_code-1.0.7 → chad_code-1.0.9}/tests/test_skills.py +0 -0
  105. {chad_code-1.0.7 → chad_code-1.0.9}/tests/test_speech_tui.py +0 -0
  106. {chad_code-1.0.7 → chad_code-1.0.9}/tests/test_subagent.py +0 -0
  107. {chad_code-1.0.7 → chad_code-1.0.9}/tests/test_symbols.py +0 -0
  108. {chad_code-1.0.7 → chad_code-1.0.9}/tests/test_syntaxgate.py +0 -0
  109. {chad_code-1.0.7 → chad_code-1.0.9}/tests/test_toolcall_parse.py +0 -0
  110. {chad_code-1.0.7 → chad_code-1.0.9}/tests/test_tools.py +0 -0
  111. {chad_code-1.0.7 → chad_code-1.0.9}/tests/test_tui.py +0 -0
  112. {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.7
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 the `lsp` extra (`uv tool install 'chad-code[lsp]'`); without
168
- it chad uses the tree-sitter fallback automatically.
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 (`uv sync --extra speech` — just a mic
195
- library; no torch, no numba).
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 the `lsp` extra (`uv tool install 'chad-code[lsp]'`); without
132
- it chad uses the tree-sitter fallback automatically.
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 (`uv sync --extra speech` — just a mic
159
- library; no torch, no numba).
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"
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. Install with `uv sync --extra speech`.
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
- # mlx_fastpath patches mlx_lm classes that only exist on macOS: with mlx installed
174
- # (dev Macs, the macos CI jobs) its `type: ignore[method-assign]` comments are load-
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 this module.
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"
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
- layer._moe_fast = _compile_moe_step(layer)
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
- gates = mx.softmax(qmm(x, gate), axis=-1, precise=True)
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