Skip to content

[Engine] support Qwen3.8-Flash-Next Coder (256 experts) on Apple Silicon Metal - #3662

Open
chuaaron wants to merge 2 commits into
Niko1221:mainfrom
chuaaron:feat/metal-coder
Open

chuaaron wants to merge 2 commits into
Niko1221:mainfrom
chuaaron:feat/metal-coder

Conversation

@chuaaron

@chuaaron chuaaron commented Oct 21, 2026 •

Copy link
Copy Markdown

Problem and result

Strata previously supported Qwen3.8-Flash-Next 512-expert models on Apple Silicon (Metal), but attempting to run the Qwen3.8-Flash-Next Coder variant (256 experts) failed at multiple levels:

  1. Setup scripts and profile selectors rejected it ("Select one of the three Qwen3.8-Flash-Next sizes (Q2_0, IQ3_XXS, IQ3_S); coder is not one of them").
  2. The Metal kernels (router.metal, moe.metal, prompt.metal) hardcoded 512 experts with 16-unroll routing and MoE dispatches, causing out-of-bounds indexing or shape mismatches when given a 256-expert model.
  3. The C++ engine (window.cpp, stream.cpp, stream.h, prompt_chain.inc) hardcoded 512 expert slots for routing logits, buffers, and offsets.
  4. On macOS 15 SDK, runtime shader compilation failed when prompt tokens triggered prefill because #include <MetalPerformancePrimitives/MetalPerformancePrimitives.h> was unconditionally included in dense_mm.metal without feature detection.

With this change:

  • coder family and IQ1_M are supported on Apple Silicon (Metal).
  • Shaders and engine dynamically handle 256 experts alongside existing 512-expert models (100% byte-identical default path for 512-expert models).
  • macOS 15 without Metal 4 / MPP safely falls back to universal decode window prompt prefill.
  • Verified end-to-end inference in CLI and HTTP server on Apple M2 Max (96 GB Unified Memory) with speculative decoding (MTP) yielding ~29-32 tok/s decode and ~42 tok/s prefill.

Related discussion

none

Approach

  1. Shaders:
  • In router.metal: added n_expert == 256 routing branch (unroll 8) alongside 512 (unroll 16).
  • In moe.metal: added n_expert == 256 MoE calculation branch.
  • In prompt.metal: added route_topk_256 kernel specialization.
  • In dense_mm.metal & mm_bench.metal: guarded MetalPerformancePrimitives include and gemm_f16 with #if __has_include(<MetalPerformancePrimitives/MetalPerformancePrimitives.h>).
  1. C++ Engine:
  • In stream.h, stream.cpp, window.cpp, prompt_chain.inc: replaced static 512 constants with n_expert_ detected dynamically from model metadata.
  • In window.cpp: added supports_metal4() check to fall back cleanly to decode window prefill on M1-M4 / macOS 15 devices.
  1. Setup & Control:
  • In setup.py, metal/setup_mac.py, strata_control/mac_fit.py: enabled coder family and IQ1_M on macOS, configured shard file naming and profile defaults.
  • In metal/tools/stream_experts.py: read n_expert from GGUF metadata instead of assuming 512.

Validation

Passed:

  • macOS Setup test suite on feat/metal-coder:
    python3 tools/test_setup_mac.py -> Ran 24 tests in 0.090s, OK.
  • Metal build:
    ninja -C build-metal strata -> clean build without errors.
  • Direct CLI Inference (101 tokens prompt prefill + decode):
    ./build-metal/strata --tokens "100,101...200" --pack <pack> --native <shard1> --ple-gguf <shard2> --expert-profile data/expert-profile-coder.bin --expert-cache 64 --prefill auto --spec 2 --mtp <mtp_dir> --max-context 65536 --mmap-experts --max-new 5
    Result:
    • Prefill 100 tokens in 2344.2 ms (42.66 tok/s).
    • Decode 5 tokens in 169.7 ms (29.46 tok/s).
    • Speculation: 3 rounds, drafts accepted 2 of 2 (1.000).
  • HTTP Server multi-turn chat (/v1/chat/completions):
    • Request 1: 53 prompt tokens, 20 generated tokens at 29.4 tok/s.
    • Request 2 (follow-up): 95 prompt tokens (48 reused + 47 read), 50 generated tokens at 32.2 tok/s, draft acceptance 27 of 29 (93.1%).
  • Hardware environment: Apple M2 Max (12 cores, 38-core GPU), 96 GB Unified Memory, macOS 15.7.5 (Build 24G624).

Failed:
none observed

Not run:
CUDA / ROCm / SYCL execution (change is isolated to Apple Metal backend and macOS setup scripts; non-Mac code paths unaffected).

Limits

On Apple Silicon without Metal 4 (M1-M4, macOS 15), prompt prefill uses the universal decode window batched prefill path; Metal 4 matrix-unit prefill (matmul2d) requires macOS 26+ and supported hardware.

@github-actions github-actions Bot added the needs-format Title or template does not follow CONTRIBUTING.md yet label Oct 21, 2026
@github-actions

Copy link
Copy Markdown

Thanks for the pull request. To keep the tracker readable, every item needs an area prefix in the title and the parts of the pull request template filled in. This one is missing:

  • Fill in "Validation": the exact commands, the revision and environment tested, what you saw. Keep Passed, Failed and Not run apart.

Edit the pull request (title and description) and this note goes away on its own. If it is still missing parts in 5 days it will be closed, which is not a rejection: fix it and comment /reopen.

Details: CONTRIBUTING.md

@github-actions github-actions Bot removed the needs-format Title or template does not follow CONTRIBUTING.md yet label Oct 21, 2026
@github-actions

Copy link
Copy Markdown

Thanks, this follows the format now.

…con Metal

Support Qwen3.8-Flash-Next Coder variants with 256 experts (including IQ1_M)
on Apple Silicon Mac (Metal backend).

- metal/kernels/router.metal, moe.metal, prompt.metal: add 256-expert routing
  and MoE dispatch paths while preserving 512-expert paths for standard models.
- metal/src/engine/window.cpp, stream.cpp, stream.h, prompt_chain.inc: dynamically
  read expert count from model metadata instead of assuming static 512.
- metal/src/engine/window.cpp, kernels/dense_mm.metal: guard MetalPerformancePrimitives
  and add supports_metal4() check to fall back cleanly to decode window prefill
  on macOS 15 / M1-M4 hardware.
- setup.py, metal/setup_mac.py, strata_control/mac_fit.py: add support for family 'coder'
  and model 'IQ1_M' on macOS.
@salty2011

Copy link
Copy Markdown

Tested this PR on an M5 Pro (18-core CPU, 20-core GPU, 64 GB, macOS 27.0.1). Hardware that's both M5 and 64 GB isn't covered by docs/MAC.md, so posting the numbers here in case they help with review.

Revision: v0.2.4.1 (c33fa7ba) merged with this PR's head e145db30. Built with setup.py --family coder --model IQ1_M --context 131072 --vision no, then --max-context raised to 262144 in the generated config. iogpu.wired_limit_mb=56000, --mmap-experts, --spec 2 with the shared MTP layer.

Passed

No regression on Q2_0 (same machine, same pack). I alternated four builds: base, PR, base, PR. On each I ran 3 greedy prompts (temperature 0, seed 1, 400 tokens) and the 23 coding-agent turns from tools/agent_turns_bench.py.

  • Greedy output text was identical on all four builds.
  • Agent turns:
Build Median TTFT Median decode
base 1 9.63 s 66.8 tok/s
PR 1 9.57 s 66.4 tok/s
base 2 9.73 s 65.4 tok/s
PR 2 9.64 s 66.3 tok/s

Coder IQ1_M loads and runs on M5:

  • Sanity prompts were correct.
  • MTP draft acceptance was 84-99%.
  • 23/23 agent turns completed, with a median TTFT of 10.96 s and a median decode of 59.6 tok/s. That's about 10% slower than Q2_0 on the same box.

Long-context recall (tools/needle_bench.py, depths 10/50/90%):

Context Found Prompt time (cold) Peak wired Lowest free
32K / 64K / 128K 9/9 37 s / 60-80 s / 76-156 s 40.8 GB 38%
192K / 250K 6/6 215-235 s / 297 s 43.1 GB 35%

There was no swap at any size, and the engine footprint was about 22 GB.

Coding tasks: I ran 2 multi-file Go tasks in an agent loop, graded by hidden tests. Coder passed both (7/7 and 11/11) in about the same time as Q2_0 (11.9 and 5.6 min, against 9.2 and 5.7 min for Q2_0). Tool calls were well formed.

The M5 kernels are still used with the new __has_include guards.

  • On this Mac, MTLDevice reports supportsFamily(.metal4) and .apple10 as true.
  • A runtime makeLibrary compile of the same #if __has_include(<MetalPerformancePrimitives/...>) guard takes the MPP branch.
  • So the gemm_f16_* and coop_layout_* kernels are compiled, and the Metal 4 prompt path is taken rather than the fallback.
  • Prompt speed agrees: Coder read a 31.6K-token prompt in 37 s, about 850 tok/s, against about 60 tok/s decode.

Failed

  • None that I could trace to the engine.
  • In one harness with strict shell rules, Coder chained commands more often than Q2_0 did. After some were refused it stopped running tests, and it left a compile error. Q2_0 didn't hit this under the same rules. It's a harness and model-behaviour issue, not an engine one, but worth knowing.

Not run

  • Any other M5 or 64 GB variant, or 128K+ contexts with real agent sessions. Recall at 256K was only tested with the needle bench.
  • A build on an older SDK without MetalPerformancePrimitives, to check the #else side of the guards. Note that when Metal 4 is missing, the engine quietly uses the slower prompt path. A one-line log of which path was chosen would make reports like this easier.

Happy to rerun anything specific on this hardware.

@Thomas-Zenkel

Copy link
Copy Markdown

Tested this PR on a Mac where the Coder streams its experts from the SSD: M4 Pro, 48 GB, macOS 27.0.1 (GPU working set 40.2 GB). The two earlier reports were on Macs that keep the model resident (M2 Max 96 GB, M5 Pro 64 GB).

Revision: main 7dd5cf18 merged with this PR's head e145db30 (no conflicts), engine built with ninja -C build-metal strata. Coder IQ1_M at the pinned revision 5348543e (SHA-256 checked); its shard 2 is byte-identical to the Q2_0 shard 2 already installed (SHA-256 316b46f3…). Pack from tools/iq_pack.py, experts' copy from this PR's metal/tools/stream_experts.py --verify ("48 layers x 256 experts identical to the GGUF", 23.42 GiB). Engine arguments as mac_fit.engine_args writes them for a streaming 48 GB Mac: --expert-cache 64 --prefill auto --spec 2 --suffix-draft 0 --mmap-experts, STRATA_METAL_STREAM=1, no MTP. Requests through serve/server.py, greedy, thinking off, 200 new tokens; the server restarted for every run, arms alternated.

Passed

Q2_0, main vs main+PR (2 rounds, 4 server starts): the answers are byte-identical on all four prompts in every run (3 short prompts and a 3,723-token prompt). Medians, decode tok/s (main / PR): 22.6 / 23.5, 25.1 / 27.2, 26.3 / 29.0; the 3.7K prompt 247 / 295 tok/s. Run-to-run spread on this Mac was larger than the difference (one main run was 13-15% slower on decode and 24% on the long prompt; a 6 MB/s download was running during this A/B). Nothing was swapped out.

Coder IQ1_M streams and answers (2 rounds each, the same greedy answers in every run, with and without the copy):

with the experts' copy without (STRATA_METAL_CONTIG=0)
Decode, 3 short prompts (medians) 21.0 / 23.5 / 24.6 tok/s 21.3 / 24.0 / 24.9 tok/s
Prompt, 3,723 tokens 287 tok/s 299 tok/s

The arena was 13.2-16.8 GB (6,000-7,600 of 12,288 experts), depending on the memory free at start. The copy gave no measurable gain here. Swap: 0 in these runs, 7,332 pages (115 MB) in one later run.

Needles (tools/needle_bench.py, 8K/16K/32K at 10/50/90%, with the copy): 9 of 9 found; 32K prompts in 112-114 s (about 280 tok/s); nothing swapped.

The 256-expert routing kernels against a CPU reference (a scratch program after metal/tests/fuse_parity.cpp, BF16 router 2560 -> 256, 1-8 tokens): moe_route with 256 threads picks the same ids as a double-precision softmax/top-10 (weights within 2.2e-7 relative); route_topk_256 is bitwise moe_route. fuse_parity (512) passes on the PR build. tools/test_setup_mac.py: 26 tests OK.

Failed

  1. STRATA_METAL_ROUTE=1 (the opt-in fused router) is broken for 256 experts. window.cpp launches router_route with grid = n_expert_ / 16 and tg = n_expert_ (256 threads = 8 simdgroups), but the kernel maps row = tgx * 16 + sg, so rows tgx*16 + 8..15 are never written. Measured with the scratch program: 128 of 256 logits unwritten per token, ids, weights and plan differ from gemv + moe_route. With the Coder through the server: every answer is empty (1 token). The same kernel with 512 threads (grid 16) is bitwise equal to gemv + moe_route for 256 experts too, so keeping r.tg[0] = 512 in that branch (or row = tgx * (nth / 32) + sg) fixes it. Default is off, so the normal path is not affected.
  2. The streaming arena ignores --expert-profile. window.cpp (around line 1308) loads data/expert-profile.bin unless STRATA_METAL_EXPERT_PROFILE is set. For the Coder its header (48 x 512) does not match, the log says experts read from the profile's top (0 ranked), and the arena is filled in the default order. With STRATA_METAL_EXPERT_PROFILE=data/expert-profile-coder.bin it reads 12288 ranked; on these short prompts that was not measurably faster (2 rounds; the arenas were smaller in those runs), so the cost is unmeasured, but the fallback is not what the config asks for.
  3. Setup would keep the Coder resident on this Mac. With main's mac_fit rule, streams("IQ1_M", 48) is False and setup writes STRATA_METAL_STREAM=0; the engine's own rule (29.6 GB + 12 GiB > 40.2 GB working set) would stream. That is the Gyro-S case of [Engine] Metal: Gyro-S resident mode wraps the whole 58.5 GB file and OOMs on a 64 GB Mac (docs: GPU maps 29.7 GB) #3784, fixed by [Other] mac_fit: decide resident vs streaming with the engine's rule (Gyro-S looped on a 48 GB Mac) #3983 (with it, streams("IQ1_M", 48) is True). Not run resident here: Gyro-S did not finish a 53-token prompt in 60 s that way on a 48 GB M4 Pro.
  4. Setting up by hand: --ple-gguf must name the shard 2 beside shard 1 (iq_pack.py also wants it there). Pointing it at the same file in another folder stops the engine with native dense: duplicate split shard number. Setup does it right; just noting it for anyone testing by hand.

Not run

  • metal/setup_mac.py --family coder end to end (it creates a venv and installs packages); I ran the steps by hand with the PR's tools.
  • MTP with the Coder on a streaming Mac (setup turns it off below 96 GB), images, the macOS 15 / no-MetalPerformancePrimitives side of the new guards.
  • A comparison with the CUDA engine on the same file.

Small notes: mac_fit.engine_args now picks the profile by looking for "coder" in the paths, and with profile=None it returns the bare name expert-profile.bin instead of data/expert-profile.bin (all callers pass a path today). The macOS 15 MPP guard is a second change in this PR; CONTRIBUTING asks for one problem per PR.

@chuaaron

Copy link
Copy Markdown
Author

@Thomas-Zenkel @salty2011

Thank you both for the incredibly thorough and detailed tests on both resident (M5 Pro 64 GB) and streaming (M4 Pro 48 GB) configurations!

I have addressed the findings and pushed the fixes in commit bdf92628:

1. STRATA_METAL_ROUTE=1 (Fused Router for 256 experts)

  • Root cause confirmed: window.cpp launched router_route with r.tg[0] = n_expert_ (256 threads = 8 simdgroups). Since the kernel maps row = tgx * 16 + sg, each threadgroup only computed rows 0..7 and left rows 8..15 unwritten.
  • Fix: Kept r.tg[0] = 512 (16 simdgroups) in the router_route dispatch so all 16 rows are computed per threadgroup for 256 experts as well.
  • Verification: Ran ./build-metal/fuse_parity across BF16 and F32 (1 to 4 tokens) — all tests report bitwise PASS (ALL PASS).

2. Streaming Arena & --expert-profile

  • Wired strata::metal_engine::set_expert_profile() from generate.cpp to window.cpp, so any --expert-profile passed via CLI / config is now properly respected by the streaming arena rather than being ignored.
  • When no profile is explicitly passed and n_expert_ == 256, window.cpp now automatically falls back to data/expert-profile-coder.bin (reading 12,288 ranked experts) rather than attempting to load the 512-expert profile.
  • Fixed strata_control/mac_fit.py engine_args so that default profile path transformations retain the directory prefix (data/expert-profile-coder.bin).

3. Prompt Prefill Path Logging

  • Added a one-line stderr message on first prefill indicating the active path:
    • strata metal prompt: Metal 4 resident prefill path
    • strata metal prompt: Metal 4 streaming prefill path
    • strata metal prompt: universal decode window prefill path (no Metal 4 on this OS/device)

Note regarding macOS 15 / MPP Guard

Regarding the MetalPerformancePrimitives guard: on macOS 15.x SDKs (M1–M4), <MetalPerformancePrimitives/MetalPerformancePrimitives.h> does not exist in the Metal compiler toolchain. Without #if __has_include(...), runtime compilation of dense_mm.metal fails immediately with exit code 3 as soon as a prompt triggers prefill. Guarding it ensures the Coder model is runnable on both macOS 15 and macOS 26+.

All 26 tests in python3 tools/test_setup_mac.py pass cleanly.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants