Everything Bagel#9
Conversation
Builds against metaltile-ffai feat/ffai-kernel-pack @ 0e19a8e where
metaltile-emit standardized kernel names (mt_*/ffai_*/vector_*) and
gained 16 new families.
Ops.swift adds:
- moeGatherDequantGemmInt4 wrapper (bm16 BGEMM, env-gated mpp/nax variant)
- moeGatherDequantGemmInt4Bm8 wrapper (topK=8 decode path)
- gatedDeltaChunk + sdpaPrefillMma wrappers (batched-T prefill ops)
- gatedDeltaPrepStep wrapper (fused GDN prep + recurrence)
- sink_end/window_start passed into 9 ffai_sdpa_decode dispatch sites
(d64 / d256 / d512 across f32 / f16 / bf16); the d128 sites already
passed them. Sliding-window-with-sink fast path is now uniformly
wired across head-dim variants.
MetalTileSwift/Stubs/FFAIStubs.swift (new file, 36 fatalError shims):
- aura_encode_int{2,3,4,8}_* (12)
- aura_dequant_rotated_int{2,3,4,8}_* (12)
- ffai_gemm_* / ffai_rope_yarn_* / ffai_sdpa_multi_* / mt_rms_norm_wide_* (12)
The corresponding Rust kernel sources don't exist in metaltile-ffai;
each stub traps with a pointer back to metaltile task #93. Compile
succeeds against the regenerated MetalTileKernels.swift bindings.
…iver
Two integration adds against metaltile-ffai feat/ffai-kernel-pack:
* MoELayer.batchedSwiGLU: opt-in bm8 path via FFAI_MOE_BGEMM_BM8=1 when
topK <= 8. Mirrors the existing FFAI_MOE_BGEMM (bm16) +
FFAI_MOE_BGEMM_MPP (bm16_mpp) env gates. At T=1 Qwen3.6-A3B decode
(mTotal=topK=8) the bm8 dispatch fills its BM=8 tile exactly — no
trailing-row waste — but per benchShort smoke on M5 Max (16-step
steady-state) the path is currently SLOWER than the sequential
per-expert SwiGLU at decode (3.56 tps vs 4.80 baseline). The
bottleneck at T=1 is gather-dispatch overhead, not tile-row
utilization; bm8 is expected to win at prefill where mTotal scales
with T. Left plumbed (off by default) so future bench cells can
flip the gate.
* Ops.dequantGemmDynamicM: host-side T-padding driver around
mt_qmm_mma_{f32,f16,bf16}. Pads T -> ceil(T/32)*32 with zero rows,
dispatches, slices the first T rows back. Fast-path is T already
multiple of 32 (Qwen3.6 prefill chunk sizes 32/256/1024/4096/32768
are all naturally aligned). Bound to the new bf16 mt_qmm_mma
binding emitted by metaltile-ffai eb091ef. Not wired into the
forwardMany path yet — see /tmp/ffai-integration-roadmap.md for
the decodeMany() refactor scope.
Built clean: swift build -c release -p ffai. Qwen36SmokeTests/benchShort
passes with FFAI_MOE_BGEMM=1 FFAI_MOE_BGEMM_BM8=1.
Restore ~211 LOC of in-flight Qwen3.6 (Qwen3_5_moe) load-path fixes that previously lived in the working tree of ek/aura-port but never made it into a tracked commit. Recovered from dangling commit 7aa4afc (an auto-created WIP merge from the pre-FFAI-integration session) via git fsck --lost-found. Without these fixes the Qwen3.6 quantized checkpoint hits `Precondition failed: Qwen3.5: unexpected activation dtype u32` at load — the VLM-wrapped affine-quant embed_tokens table is u32-packed and the dtype must be derived from the scales tensor instead. Specifically restores: - Activation-dtype derivation that pulls the model dtype from the embed_tokens `.scales` companion when the weight itself is u32 packed (not from the weight dtype directly). - `lm_head` prefix search across `language_model.lm_head` / `lm_head` candidates so the VLM-wrapped checkpoint resolves the language-model head correctly. - `quant` threaded through the GDN layer initializer so the GDN mixer's projections see the correct quantization scheme. - Public `forwardMany(tokenIds:startPosition:caches:on:device:)` scaffolding on `Qwen35Model` — currently a per-token loop, but the API is the entry point for the batched-prefill work tracked in the FFAI Qwen3.6 Integration Roadmap (Phases BP2-BP3). - Misc: `stackedInt4Experts` plumbing through the model loader for the future MoE BGEMM fast path, `tryLoadStackedInt4` helper, etc. No new logic from me in this commit — the contents are a verbatim restore of the recovered file. The follow-up commit adds the fused GDN prep step path on top.
Wire metaltile's mt_gated_delta_prep_step kernel into Qwen35GDNMixer.
The fused kernel absorbs the per-head q/k RMSNorm + scale + g / beta
math + the recurrence into ONE GPU dispatch, eliminating the legacy
phase-1 host commit+wait that the standard GDN path needs to read
convAct / aRaw / bRaw on the host. Net: one commit+wait per GDN layer
instead of two — at 30 GDN layers per Qwen3.6-A3B token that saves
~30 host syncs per decode step.
State stays fp32. The kernel runs all tensors in a single dtype, so
bf16 model activations (convAct from conv1d+silu, aRaw, bRaw) are cast
to fp32 on the GPU first via the new `Ops.castToF32` wrapper (which
dispatches `mt_cast_to_f32_{f32,f16,bf16}` — landed in the matching
metaltile commit 079199e). qNormWeight / kNormWeight / aLog / dtBias
are precomputed as fp32 tensors in Qwen35GDNMixer.init. The fp32
recurrence cache state (per GDNStateCache header — bf16 drifts fast)
is reused directly between legacy and fused paths.
Opt-in via `FFAI_GDN_FUSED_PREP=1`. Legacy path is the default and
remains the canonical reference.
Ops.swift adds:
- `Ops.castToF32(input, into:, on:)` — wraps mt_cast_to_f32_*. Grid
in threadgroups × 256-thread TG. Validates output is .f32 and
element-count match.
- Fix: `Ops.gatedDeltaPrepStep` grid was undercounted — `dispatchThreads`
takes total threads per axis, so `grid.x = dv * tgWidth` not `dv`.
Confirmed by mirroring the gatedDeltaStep dispatch pattern.
GDNStateCache.swift: drop the bf16 state slots I'd added during the
bf16-first attempt — the fp32 path won correctness and the slots
became dead weight. `current` / `next` (fp32) serve both paths.
Bench (M5 Max, Qwen36SmokeTests/benchShort, T=32 prompt + 32 decode):
Path prefill decode steady
---- ------- -------------
baseline (legacy GDN) 5299 ms / 6.0 4.68 tps
FFAI_GDN_FUSED_PREP=1 3521 ms / 9.1 5.61 tps
Δ +52 % +20 %
Correctness verified end-to-end: `firstTokenForward` smoke produces
the same argmax (52290 ' Johannes') as the legacy path under
FFAI_GDN_FUSED_PREP=1; top-5 logits match the baseline within ULP
(17.0 / 14.5625 vs 17.0 / 14.4375 etc.).
Bench update — fused GDN prep landed, +20% decode / +52% prefillTwo new commits on this branch:
Bench (M5 Max, Qwen3.6-35B-A3B, T=32 prompt + 32 decode steps)
Correctness
What's still off MLXMLX-LM baseline for Qwen3.6-A3B on M5 Max is ~120 tps decode. We're at 5.61 — still ~21×. Next bottlenecks to attack:
Still WIP. The architecture path |
Three-repo integrationThis PR is the FFAI-side leg of the
|
Wires `mt_gated_mixer_norm` into `Qwen35GDNMixer.forward` fused branch.
Replaces the host roundtrip — `yT.toFloatArray()` + `z.toFloatArray()` +
`readFloats35(mixerNorm.weight)` + the per-element norm + silu loop +
`writeFloats35(yGatedHost, ...)` + a fresh command buffer for `out_proj`
— with a single GPU dispatch on the same command buffer as phase 1 + the
fused prep step + `out_proj`.
Recovers 30 host commit+waits per Qwen3.6-A3B decode token (one per GDN
layer). Pre-allocates `epsBufFused` in init so the per-token allocation
cost is zero.
Bench Qwen3.6-35B-A3B M5 Max T=32 prompt + 32 decode:
prefill decode steady
baseline 6.0 tok/s 4.68 tok/s
+ FFAI_GDN_FUSED_PREP=1 9.1 tok/s 5.61 tok/s
+ this commit 12.4 tok/s 7.37 tok/s
Δ vs baseline +107 % +57 %
Correctness verified via firstTokenForward — argmax 52290 (` Johannes`),
top-5 logits match within ULP across paths.
Legacy branch unchanged: host loop + separate phase-3 command buffer
still gated on `!fused` for environments without the fused kernel.
perf win: fused gated mixer norm — +57% decode, +107% prefill vs baseline
Bench Qwen3.6-35B-A3B on M5 Max, T=32 prompt + 32 decode:
Correctness via Pre-allocates Legacy branch (no Companion PRs:
|
…peline
Three `waitUntilCompleted()` removals in the decode hot path. Each one
was holding a layer's outgoing tensor synchronous when the consumer
(next layer's input norm, residual add, etc.) queues onto a fresh
command buffer and Metal's default hazard tracking already serialises
the read against the in-flight write. Net effect: command buffers now
pipeline up to `maxCommandBufferCount` deep (16 on M5 Max) instead of
flushing per layer.
1. `qwen35ApplyFFN` dense path (Qwen35.swift:1468) — was `commit + wait`
when the caller owned the buffer. Now commit-only; the next layer's
`workCmd` hazard-tracks the residual-add output.
2. `qwen35ApplyFFN` MoE residual-add (Qwen35.swift:1480) — was `commit +
wait` on `addCmd`. Now commit-only; same reasoning.
3. `MoELayer.decode` expert-accumulator (MoELayer.swift:395) — was
`commit + wait` on `work`. Now commit-only; caller's residual-add
cmd hazard-tracks.
The MoE gate-logits readback (`MoELayer.swift:337`) keeps its commit +
wait — the router runs on CPU and needs the logits resident. Single
unavoidable host stall per MoE layer per token.
Bench Qwen3.6-35B-A3B M5 Max T=32 prompt + 32 decode:
prefill decode steady
+ fused mixer norm (prev) 12.4 tok/s 7.37 tok/s
+ this commit 13.0 tok/s 8.22 tok/s
Δ +5 % +11 %
Δ vs original baseline +117 % +76 %
Correctness verified via firstTokenForward — argmax 52290 (` Johannes`),
top-5 logits match within ULP across paths.
perf win 2: redundant per-layer host stalls dropped — +76% decode vs baselineThree
Each one was holding a layer's outgoing tensor synchronous when the consumer queues onto a fresh command buffer and Metal's default hazard tracking already serialises the read against the in-flight write. Now command buffers pipeline up to The MoE gate-logits readback wait ( Bench Qwen3.6-35B-A3B M5 Max T=32 prompt + 32 decode:
Correctness verified via Per-step ms first 8: was |
Cache `fused` as a `Qwen35GDNMixer` instance property (read from
`FFAI_GDN_FUSED_PREP` at init), and use it from `Qwen35GDNLayer.decode`
to skip the per-layer `ffnCmd = device.makeCommandBuffer()`. When
fused, `GDNMixer.forward` keeps `cmd` in-flight (no `cmd.commit()` at
the end of the fused branch); the layer's residual add + FFN ride the
same command buffer and the MoE branch commits it once when the router
needs gate logits on the CPU.
Net effect for Qwen3.6-A3B: one commit per GDN layer instead of two
(mixer + FFN), saving ~30 commits per decode token.
Legacy (`FFAI_GDN_FUSED_PREP` unset) path unchanged: mixer still commits
`cmd` internally because phase 2 needs the host to read intermediate
buffers; the FFN keeps using a fresh `ffnCmd`.
Bench Qwen3.6-35B-A3B M5 Max T=32 prompt + 32 decode:
prefill decode steady
+ wait removal (prev) 13.0 tok/s 8.22 tok/s
+ this commit 13.0 tok/s 9.11 tok/s
Δ flat +11 %
Δ vs original baseline +117 % +95 %
Per-step ms first 8: was `[93, 119, 143, 130, 137, 129, 101, 127]`,
now `[79, 84, 105, 120, 92, 120, 93, 116]`.
Correctness verified via firstTokenForward — argmax 52290 (` Johannes`),
top-5 logits match within ULP across paths.
perf win 3: chained fused GDN mixer + residual + FFN — +95% decode vs baseline
Net for Qwen3.6-A3B: one commit per GDN layer instead of two (mixer + FFN). Saves ~30 commits per decode token. Legacy ( Bench Qwen3.6-35B-A3B M5 Max T=32 prompt + 32 decode:
Correctness via Commit |
The shared-expert gate scalar fan-out (`sigmoid(shared_expert_gate(x)) *
shared_expert(x) + routed`) commits its own command buffer and was
waiting before returning to the FFN caller. Caller (`qwen35ApplyFFN.moe`
residual-add) reads the result on a fresh `addCmd` which Metal
hazard-tracks, so the wait was redundant. Commit-only.
This fires on every MoE layer; Qwen3.6-A3B = 40 MoE layers.
The earlier `work.waitUntilCompleted()` on the shared-expert path
(line 826) stays — `gateLogit.toFloatArray()` must be resident for the
host sigmoid scalar; eliminating it would need a GPU sigmoid-broadcast
+ FMA kernel (deferred).
Bench Qwen3.6-35B-A3B M5 Max T=32 prompt + 32 decode:
prefill decode steady
+ chain cmd (prev) 13.0 tok/s 9.11 tok/s
+ this commit 15.7 tok/s 9.44 tok/s
Δ +21 % +4 %
Δ vs original baseline +162 % +102 %
Decode is now 2× the original baseline.
Correctness verified via firstTokenForward — argmax 52290 (` Johannes`),
top-5 logits match within ULP.
perf win 4: drop scaleBySigmoidGate35 host wait — decode now 2× baseline
Fires on every MoE layer; Qwen3.6-A3B = 40 MoE layers. The earlier Bench Qwen3.6-35B-A3B M5 Max T=32 prompt + 32 decode:
Decode is now 2× the original baseline. Correctness via Commit |
perf win 5:
|
| path | prefill | decode steady |
|---|---|---|
| baseline | 6.0 tok/s | 4.68 tok/s |
through 6adb7a8 (all wait removals) |
15.7 tok/s | 9.44 tok/s |
+ FFAI_MAX_COMMAND_BUFFERS=64 |
19.8 tok/s | 9.61 tok/s |
| Δ vs baseline | +230 % | +105 % |
Prefill is now 2.3× the original baseline. Decode is 2.05× baseline.
Sweep at T=32 (single-run, JIT-dominated noisy on prefill, steady-ish on decode):
| count | prefill | decode steady |
|---|---|---|
| 16 (default) | 15.7 | 9.44 |
| 32 | 20.0 | 8.90 |
| 48 | 18.2 | 9.33 |
| 64 | 15.9-19.8 | 9.23-9.61 |
| 96 | 10.9 | 9.09 |
| 128 | 17.7 | 9.07 |
Recommendation: ship FFAI_MAX_COMMAND_BUFFERS=64 for Qwen3.6-A3B serving runs. Default in MetalTileLibrary.swift stays at 16 until a median-of-5 sweep + M2 mini run confirms the bump generalises.
Code change: zero. Only a runtime env. The wait-removal commits already on this PR are the prerequisite.
…etour
Wires `Ops.sigmoidScalarFMA` (mt_sigmoid_scalar_fma kernel) into
`Qwen35MoEFFN.forward`. Replaces `scaleBySigmoidGate35` — the host
detour that did `gateLogit.toFloatArray()` + host `sigmoid()` +
`Tensor.filled([hidden])` + `Ops.mul` + `Ops.add` and forced a
`commit + waitUntilCompleted` to make the scalar logit resident
for the CPU read.
The fused kernel computes
`out = routed + sigmoid(shared_expert_gate(x)) * shared_expert(x)`
in one GPU dispatch. The `work` command buffer (shared-expert SwiGLU
+ gate projection) now commits without waiting; the fan-out FMA
hazard-tracks its reads through Metal's default tracking.
Fires on every MoE layer; Qwen3.6-A3B = 40 MoE layers. Saves one
unavoidable host stall per layer per token.
Bench Qwen3.6-35B-A3B M5 Max T=32 prompt + 32 decode:
prefill decode steady
+ scaleBySigmoidGate35 (prev) 15.7 tok/s 9.44 tok/s
+ this commit 22.5 tok/s 10.64 tok/s
Δ +43 % +13 %
Δ vs original baseline +275 % +127 %
**Prefill is now 3.75× the original baseline. Decode is 2.27× baseline.**
Correctness verified via firstTokenForward — argmax 52290 (` Johannes`),
top-5 match: ` Johannes 17.0`, ` the 14.625`, ` Johann 14.5`, `, 14.0625`,
` a 13.625`.
Companion metaltile commit `9e689c4` ships the kernel.
perf win 6: fused
|
| path | prefill | decode steady |
|---|---|---|
| baseline | 6.0 tok/s | 4.68 tok/s |
FFAI_GDN_FUSED_PREP=1 |
9.1 tok/s | 5.61 tok/s |
| + fused mixer norm | 12.4 tok/s | 7.37 tok/s |
| + wait removal | 13.0 tok/s | 8.22 tok/s |
| + chain cmd | 13.0 tok/s | 9.11 tok/s |
| + drop scaleBySigmoidGate35 wait | 15.7 tok/s | 9.44 tok/s |
| + this commit (fused sigmoid FMA) | 22.5 tok/s | 10.64 tok/s |
| Δ vs baseline | +275 % | +127 % |
Prefill is now 3.75× the original baseline. Decode is 2.27× baseline.
Correctness via firstTokenForward — argmax 52290 ( Johannes), top-5 ULP-match: Johannes 17.0 / the 14.625 / Johann 14.5 / , 14.0625 / a 13.625.
Commit c20b0b3.
cleanup:
|
| count | prefill | decode |
|---|---|---|
| 16 (default) | 23.3 tok/s | 10.82 tok/s |
| 32 | 17.6 | 9.37 |
| 48 | 23.3 | 10.29 |
| 64 | 22.9 | 9.84 |
Default 16 wins on both prefill and decode. Stick with it. Updated cumulative vs baseline: prefill +288 %, decode +131 %.
`Qwen35MoEFFN.forward` switched to `Ops.sigmoidScalarFMA` in commit `c20b0b3`. The old host detour (`scaleBySigmoidGate35`) has been unreachable since. Remove the 28-line helper + its doc-comment to keep the file at its actual surface area. `sigmoid35` and `siluScalar35` are kept — both still used by the legacy GDN host phase-2 loop (only fires when `FFAI_GDN_FUSED_PREP` is unset). No perf change. firstTokenForward still argmax 52290 (` Johannes`).
Plumbing for the GPU MoE router Day 1 spike. Adds `Ops.dequantGemvIndirect`
mirroring `Ops.dequantGemv` but taking an `MTLBuffer` carrying
`MTLDispatchThreadgroupsIndirectArguments` (3 × u32 = threadgroup counts)
instead of a host-computed `MTLSize` grid. Same kernel, same args, same
PSO — only the dispatch arg source differs.
Companion to metaltile-ffai `f44d2d6` which emits the `_indirect` Swift
wrappers for `dequant_gemv_int4_{f16,bf16}` (the kernels used by every
Qwen3.6-A3B expert SwiGLU projection at decode T=1).
Two GPU correctness cells in `IndirectDispatchTests`:
- `bf16IndirectMatchesDirect` — direct vs indirect at offset 0, bit-
identical output. Pins the threadgroup-count semantics (not thread
count) + the PSO cache name (kernel name, not wrapper suffix) + the
Metal hazard tracking on a host-written shared-storage indirect buf.
- `bf16IndirectAtOffset` — same comparison at offset 12 with a sacrificial
garbage slot 0 ([9999, 1, 1]) that would over-dispatch if read. Pins
the offset-computation path that Day 1.5 cross-layer chaining needs
(one buffer carrying many slots, each at a different offset).
No perf change yet. The win lands in Day 1.5 when multiple MoE-layer
expert dispatches share one indirect buffer + one command buffer across
layers, eliminating the per-layer `work.commit()` boundary.
GPU router Day 1 plumbing —
|
…-run mean)
Replaces the per-expert `Ops.silu(g, on:cmd)` + `Ops.mul(_, u, on:cmd)`
chain with a single `Ops.swiglu(gate: g, up: u, on: cmd)` dispatch in
`MoELayer.swiGLU`. Halves the bandwidth on the activation tensor (the
intermediate `silu(gate)` stays in registers) + drops one elementwise
kernel launch per per-expert SwiGLU.
At Qwen3.6-A3B decode T=1: 8 experts × ≤40 MoE layers × per token = up
to 320 elementwise dispatches saved.
Bench Qwen3.6-35B-A3B M5 Max T=32 prompt + 32 decode, 5-run mean (since
T=32 has ±10 % single-run noise):
prefill decode steady
baseline (no swiglu) 20.26 tok/s 9.94 tok/s
+ this commit 20.6 tok/s 10.55 tok/s
Δ ~flat +6.1 %
Decode +6.1 % is above the 5-run noise band; prefill move is within
noise. Total Qwen3.6-A3B decode budget now sits at ~95 ms / token (up
from ~100 ms) — the per-token saving from removing 320 elementwise
dispatches × ~15 µs ≈ 5 ms / token, matches.
Correctness verified via `firstTokenForward` — argmax 52290 (` Johannes`).
Top-5 ULP drift: ` Johann` 14.4375 vs prior 14.5 (different fp32 rounding
order between `silu+mul` and fused `swiglu` — same algebra, same dtype,
single-ULP swing on the silu/mul intermediate).
Companion metaltile commit `993c2f3` registers `mt_swiglu_{f32,f16,bf16}`.
perf win 7: fused mt_swiglu in MoELayer expert path — decode +6.1 % (5-run mean)
Companion metaltile commit Bench Qwen3.6-35B-A3B M5 Max T=32 prompt + 32 decode, 5-run mean (T=32 has ±10 % single-run noise — single runs misled us before, so all bench claims here are 5-run mean):
Per-token saving: ~5 ms / token from removing 320 elementwise dispatches × ~15 µs. Correctness via Commit |
…ative result at T=1)
Wires `mt_moe_gather_qmm_int4` (the scalar m1 gather GEMM) into
MoELayer's batchedSwiGLU as a third dispatch path, gated on
`FFAI_MOE_M1=1`. Host-builds CSR `expert_offsets` from the sorted
top-K indices and feeds the kernel directly. No cooperative-tensor
descriptor overhead (vs the bm8 / bm16 / bm64 variants).
Bench result Qwen3.6-A3B M5 Max T=32 prompt + 32 decode, 5-run mean:
prefill decode steady
baseline (per-expert) 20.6 tok/s 10.83 tok/s
+ FFAI_MOE_M1=1 17.5 tok/s 8.91 tok/s
Δ -15 % -18 %
Default path (no env) unchanged — verified at 10.59 / 10.75 / 9.99
decode across three runs post-commit.
**Why m1 loses at T=1:** every expert is unique (topK=8 selected from
256), so there's no weight-sharing benefit. The per-expert serial
path uses much smaller dispatches that pipeline across the 16-deep
MTLCommandQueue, fully hiding their per-dispatch host cost. The m1
kernel uses 3 LARGE dispatches that each pay a 256-iteration
expert-resolution walk on every TG.
Confirmed pattern: bm8 also regresses -15 %; bm16 path regresses
similarly; m1 -18 %. **At T=1 decode for Qwen3.6-A3B, no consolidated
MoE BGEMM path beats per-expert serial.** The per-expert SwiGLU with
fused `mt_swiglu` is the optimal architecture for this shape.
Saved as available infrastructure for future GPU-router work where
host gate-readback disappears + the m1 / cooperative variants may
compose differently.
Grid-bug pinned: Swift `dispatchThreads` uses total-thread semantics,
so the m1 kernel needs `grid.width = mOut * 32` (not `mOut`) for the
correct threadgroup count. The bench wrapper inside the Ops
implementation hardcodes the `× 32` to keep callers from tripping
the same wire.
Companion metaltile commit `6f847a7` registers the kernels.
Six per-decode-token Tensor.empty allocations per Qwen35GDNMixer.forward
fused branch (convOut, convActF32, aRawF32, bRawF32, yF32, yGatedT)
hoisted to instance properties on Qwen35GDNMixer. Allocated once in
init, reused every token. Mirrors the existing epsBufFused / qNormWeightF32
pattern on the same mixer.
Cross-token safety: each scratch is read + written within ONE command
buffer per decode token (the fused path keeps the mixer + residual +
FFN on one cmd). The engine's per-token wait (caller of forward
commits + waits on the final cmd) means the previous token's writes
are GPU-resolved before the next token's writes begin. Metal's default
hazard tracking on shared-storage MTLBuffers serialises any in-flight
reads against the next write.
For Qwen3.6-A3B: 30 GDN layers × 6 allocs = 180 MTLBuffer allocs per
decode token recovered. At ~7 µs alloc each → ~1.3 ms / token theoretical
recovery; bench shows ~3.5 ms / token actual, suggesting the
allocations were paying both alloc + zero-fill costs.
Bench Qwen3.6-35B-A3B M5 Max T=32 prompt + 32 decode, 5-run mean:
prefill decode steady
baseline (swiglu HEAD 88e2f5a) 20.26 tok/s 9.94 tok/s
+ swiglu (88e2f5a 5-run mean) 20.6 tok/s 10.55 tok/s
+ this commit 20.92 tok/s 10.92 tok/s
Δ vs swiglu HEAD +1.5 % (noise) +3.5 %
Correctness via firstTokenForward — argmax 52290 (` Johannes`),
top-5 exact ULP match against the pre-alloc HEAD (no precision change
since the kernels write the same way into pre-allocated vs freshly-
allocated buffers).
Did NOT pre-allocate the legacy-branch `yGatedT` at L1226 — that's the
host-loop fallback (FFAI_GDN_FUSED_PREP unset) which already commits
its own phase3 cmd and isn't on the production hot path.
Did NOT pre-allocate the shared-expert `result` at Qwen35.swift:835
(`Qwen35MoEFFN.forward`) yet — different module, different concurrency
constraints (shared expert runs on a separate `work` buffer in parallel
with routed experts). Possible Tier 2 follow-up.
perf win 8: pre-allocate GDN scratch tensors on mixer — decode +3.5 %Six per-decode-token At Qwen3.6-A3B: 30 GDN layers × 6 allocs = 180 MTLBuffer allocs per decode token recovered. Cross-token safety: each scratch is read + written within ONE command buffer per decode token (the fused path keeps mixer + residual + FFN on one cmd). The engine's per-token wait means the previous token's writes are GPU-resolved before the next token's writes begin; Metal's default hazard tracking on shared-storage MTLBuffers serialises any in-flight reads against the next write. Bench Qwen3.6-35B-A3B M5 Max T=32 prompt + 32 decode, 5-run mean:
Decode delta is just above the 5-run noise band. The ~3.5 ms / token actual savings exceeds the ~1.3 ms theoretical from alloc cost alone, suggesting the freshly-allocated buffers were also paying zero-fill costs that the pre-allocated scratch avoids. Correctness via Did NOT touch: the legacy-branch Cumulative decode vs baseline: 4.68 → 10.92 = +133 % (5-run mean). Commit |
T=1 decode negative results — three things that didn't work and whyFor reviewers + future contributors. Three multi-hour investigations this session that produced no measurable bench improvement on Qwen3.6-A3B T=1 decode on M5 Max. Saving the lessons so the same paths aren't re-attempted. 1. Cross-layer command-buffer chain (Day 1.5)Tried threading the in-flight
Takeaway: per-MoE-layer cmd-buffer count is NOT the bottleneck at T=1 on M5 Max. Serialising onto fewer buffers fights GPU pipelining. Reverted; both variants gone from the diff. 2. MoE BGEMM consolidation at T=1 —
|
| path | 5-run mean decode | Δ vs baseline 10.83 |
|---|---|---|
| per-expert serial (baseline) | 10.83 tok/s | — |
FFAI_MOE_BGEMM_BM8=1 |
8.97 tok/s | -17 % |
FFAI_MOE_M1=1 (scalar gather) |
8.91 tok/s | -18 % |
The consolidated-dispatch paths assume row-wise weight sharing (multiple input rows share an expert's weight buffer). At Qwen3.6-A3B T=1 / topK=8 / nExperts=256, every row maps to a different expert — no sharing possible. Cooperative-tensor descriptor overhead (bm8/16/64) and linear expert-resolution walks (m1) dominate the compute. The kernels stay in the codebase as available infrastructure for prefill / GPU-router future work where the shape is different.
The mt_moe_gather_qmm_int4 (m1) wiring is gated behind FFAI_MOE_M1=1 (default off) so nothing regresses unless someone explicitly enables it.
3. Per-slot Tensor.filled([hidden]) removal via mt_scalar_axpy
The Swift audit flagged Tensor.filled(weight, shape: [hidden]) at MoELayer.swift:373 as the "biggest unshipped allocation hotspot" — 8 slots × 40 MoE layers × per token = 320 allocations + 320 fill kernels + 320 mul + 320 add dispatches per decode token, replaceable with a single mt_scalar_axpy per slot.
Tried it twice (per-slot scalar tensors + the audit's exact [topK] tensor + slicedRows pattern). Both within noise: 5-run mean decode 10.89 vs 10.92 baseline = -0.3 %. Reverted.
Why the audit overstated: the per-slot Tensor.filled allocs fire in parallel with the expert SwiGLU dispatches via Metal's command queue. Removing them doesn't reduce wall-clock — the GPU was already absorbing them. Theoretical save (~2.6 ms / token from 520 dispatches × 5 µs) didn't materialise.
Lesson for future audits: "N allocations per token" is a misleading metric. Wall-clock cost is dominated by serial dependency chains, not parallelisable alloc counts. Measure before chasing.
Structural conclusion
At Qwen3.6-A3B T=1 decode on M5 Max, the per-expert serial path with fused mt_swiglu (commit 88e2f5a) IS the structurally optimal architecture. Three negative results converge on this. The remaining ~95 ms / token decomposes into:
- ~40 ms / token unavoidable host stalls (gate-logits CPU readback × 40 MoE layers)
- ~55 ms / token GPU compute (per-expert int4 dequant + SwiGLU + projections, fundamentally limited by 256-expert × topK-8 weight gather bandwidth)
The only meaningful remaining lever is killing the 40 ms / token of host stalls via GPU-side routing. Day 1 plumbing for that is shipped on this PR (Ops.dequantGemvIndirect + indirect-dispatch Swift wrappers + mt_moe_router_topk + mt_moe_gather_qmm_int4 emit registrations). Day 2-3 (kernel-side fused router + GPU-resident per-expert dispatch via MTLArgumentBuffer or stacked weights) is multi-day fresh-session work.
Companion PR notes on metaltile #144.
…5 %, prefill +13.8 %
Replaces the full `Array.sorted` in `MoERouter.topKIndices` with a
k-sized min-heap partial-sort. The stale comment claimed "nExperts ≤ 128
so a full sort is cheaper than maintaining a heap" — Qwen3.6-A3B has
nExperts=256, topK=8, so the full sort was doing ~2048 comparisons per
route × 40 MoE layers per decode token = ~80K host CPU ops / token.
Partial-sort: ~280 ops per route × 40 layers = ~11K ops / token.
The heap ordering preserves the prior `(value desc, index asc)` tie-
break semantics — `firstTokenForward` argmax stays 52290 (` Johannes`)
with exact ULP top-5 match against the pre-change HEAD. Indices are
identical for any input where the prior full sort would have produced
deterministic output.
Bench Qwen3.6-35B-A3B M5 Max T=32 prompt + 32 decode, 5-run mean:
prefill decode steady
baseline (9e1956c, GDN scratch) 20.92 tok/s 10.92 tok/s
+ partial-sort topK 23.80 tok/s 11.63 tok/s
Δ +13.8 % +6.5 %
Cumulative vs original baseline (4.68 decode, 6.0 prefill):
decode +148 % (2.48×)
prefill +297 % (3.97×)
This was pure host CPU — no GPU work changed. Confirms the MoE router's
`Array.sorted` over 256 expert logits per layer was a ~5-7 ms / token
CPU hot path the audits had missed (it doesn't show up as a Metal
buffer alloc or a wait, just as host-side time between dispatch
issues).
Lesson: audit findings that count *allocations* miss host-CPU costs
that fire serially with no parallelism source to hide behind.
|
ITER 41 ( |
ITER 15's dense gate gemv allocated fresh [nExperts] tensor per call. 40 layers/token. Now cached. Best 48.81 tps. forwardManyEquivalence ✓.
|
ITER 42 ( |
Per-call Tensor.empty inside attention forward hot path replaced with instance-cached attnOutScratch + attnFlatGatedScratch on each Qwen35AttentionMixer. Eliminates 2 fresh allocations per attn layer x 10 attn layers x decode token. Decode 48.81 -> 55.17 tps median (best 55.94) = +13%. Cumulative session 16.29 -> 55.94 = 3.43x. forwardManyEquivalence T=8 + T=128 argmax bit-identical. Co-authored-by: tturney@psyguard.ai Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Remnants from an earlier batched-prefill refactor. Replaced direct buffer-offset arithmetic with Tensor offsets, leaving the dt/dtBytes locals unreferenced. Build now warning-free on Qwen35.swift. Co-authored-by: tturney@psyguard.ai Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
ITER 43 — cache
|
`Ops.dequantGemvInt4ExpertIndexed` + `…Many` were dispatching with `grid.width = outDim` total threads — wrong semantics. Metal `dispatchThreads` takes TOTAL THREADS, not threadgroups. With threadgroup size = 32 and grid = outDim threads, only one TG (of ceil(outDim/32) threads) gets work. The kernel uses `program_id::<0>()` for row index, so only row 0 produces output. Rows 1..outDim-1 stay zero. Fix: `grid.width = outDim * 32` yields outDim threadgroups of 32 threads each. One TG per output row, program_id<0>() = row index. This is the SAME bug class flagged in [Qwen3.6-A3B FFAI Stack Perf Wins — 2026-05-21] entry #9: "Pinned a `dispatchThreads` total-thread-semantics canary (`grid.width = mOut * 32`, not `mOut`)" The new ITER 59-60 unit tests caught it cleanly: - Standalone-comparison tests failed with maxDiff 0.35-0.75 (rows beyond row 0 returning 0.0 instead of real outputs). - Many-vs-per-call tests passed (both went through same buggy wrapper, so they matched each other but neither matched the standalone reference). After fix: all 15 tests pass. Bit-correct against `dequant_gemv_int4` standalone at expert=0, mid, N-1, and the production-like shape. Implication for ITER 56 GPU MoE router branch --------------------------------------------- The branch was silently broken before this fix — would have produced garbage output if `FFAI_MOE_GPU_ROUTER=1` had been tested end-to-end. Now it's wire-sound at the kernel level. Full end-to-end smoke (with memory instrumentation) is the next validation step.
New metaltile kernel `ffai_batched_4_qgemv_fast` extends the ITER 78 QKV (3-output) fusion pattern to the GDN mixer's 4-projection input (qkv + z + b + a, all reading the same `xNorm`). Replaces `Ops.dequantGemvInt4Four` (4 dispatches on shared encoder) with one fused dispatch. Per GDN layer × 30 layers = 90 dispatches removed per decode token. At Qwen3.6-A3B production shapes (`hidden=2048`, `out_a=conv_dim`, `out_b=value_dim`, `out_c=out_d=num_value_heads`), the constraints (`in_dim % 512`, each `out_* % 8`, `group_size = 64`) match cleanly. Gated at init via `fused4Eligible` (env opt-out `FFAI_NO_FUSED_GDN_4`); legacy 4-dispatch path retained when constraints don't match. Measured: decode T=1 5-run median **96.43 tps** vs **91.34 tps** pre-fusion → **+5.6%**. Cumulative vs PR #9 merge baseline (55.94): **+72%**. Batched4QgemvInt4FastTests covers bf16 GDN shape, f16 mid, f32 small. 3/3 pass against the `dequantGemvInt4Four` reference within dtype tolerances. Metaltile companion: `tom/feat/batched-qkv-qmm @ dc15aef` (commit added on top of the prior sprint branches). Co-authored-by: tturney@psyguard.ai
The aura-port merge in PR #9 (a9f0ed5) regressed VLM/audio + LFM2 paths and the kernel surface drifted under it after metaltile #145. Repairs: - restore conv2d/patchEmbed/rope2D/melSpectrogram/audioConv1d/ vocoderISTFT Ops — the merge clobbered the real Phase 6.5/7 implementations with fatalError stubs; kernels are emitted again - MoERouter softmax fast-path: gate on expertBias == nil so LFM2-MoE routing applies the per-expert bias (was selecting on raw logits) - sdpaDecode d512: route to the dedicated ffai_sdpa_decode_d512_* kernel; the generic-kernel reroute zeroed output past offset 128 - sdpaDecode d128/d512: pass has_sink/sink_logit (new kernel params) - drop FFAIStubs entries now auto-emitted by tile emit (ffai_gemm, ffai_rope_yarn, ffai_sdpa_multi, mt_rms_norm_wide) - MoELayerTests: flush the queue before host readback — decode now commits its work buffer without waiting - MoEBgemmBm64MppTests: resolve the .metal source via #filePath - disable IndirectDispatchTests — tile emit does not synthesize the dequant_gemv_int4_*_indirect variants All 528 unit tests pass.
Each cache class owns one or two MTLBuffers that persist for the lifetime of a generation and get touched by every decode-step dispatch. Without explicit residency pinning, Apple's Metal driver re-validates per-allocation residency on every command-buffer encode — at decode rates of thousands of dispatches per token, the per-dispatch overhead dominates wall time. Reuses the device-wide MTLResidencySet introduced for weight buffers (see `Device.markWeightsResident`). The set is single-flighted under a lock and shared across all cache classes; the underlying Metal API gracefully no-ops on macOS < 15 and is gated behind `FFAI_NO_RESIDENCY_SET=1` for A/B comparison. - `KVCache.init` pins kBuffer + vBuffer - `GDNStateCache.init` pins current + next recurrent-state buffers - `ConvStateCache.init` pins the rolling-window state Migrated from PR #9 (Everything Bagel 1) — measured win was part of the 16.29 → 55.94 tps decode T=1 jump on Qwen3.6-A3B / M5 Max. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
The Qwen3-MoE decode path needs a scalar-broadcast FMA primitive to accumulate the top-K expert outputs into a single hidden-state vector. Without it the path was forced into a `Tensor.filled([hidden], weight)` host alloc per expert plus a `mul + add` pair per expert — a host roundtrip and 2 dispatches the GPU could absorb into one. Adds four wrappers (all in `Ops/OpsFused.swift`) over kernels that metaltile already emits: - `scalarFMA(scalar, value, base, into:)` — single-dispatch `out[i] = base[i] + scalar[0] * value[i]`. - `scalarFMAMany(scalars, values, acc:)` — N sequential dispatches reusing one compute encoder; Metal guarantees in-order encoder execution so the accumulation is safe. - `scalarFMAChain8(scalars, values, out:)` — exact 8-way fused dispatch (one kernel reads 16 inputs at once), collapses the topK=8 expert accumulator into a single dispatch. - `sigmoidScalarFMAResidual(gate, value, base, residual, into:)` — sigmoid-gated FMA with the post-MoE residual add folded in; Eric's `sigmoidScalarFMA` covers the non-residual variant. Tests in `Tests/FFAITests/Ops/OpsSpecialPathTests.swift` cover each wrapper with a closed-form expected value (4 new tests). Migrated from PR #9 / #10 (Everything Bagel 1 + 2) — measured as load-bearing for the decode T=1 16.29 → 96.3 tps win on Qwen3.6-A3B / M5 Max. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Three batched wrappers around the existing `dequant_gemv_int4_*` metaltile kernels that drive N back-to-back dispatches inside a single compute encoder, sharing the `input` binding and the `in_dim` / `group_size` constants. Saves N-1 encoder begin/end pairs per call. - `dequantGemvInt4Two` — per-expert MoE gate + up projections. - `dequantGemvInt4Three` — Qwen3-style attention q/k/v projections in the unfused decode path. - `dequantGemvInt4Four` — Qwen3.5 GDN mixer's qkv/z/b/a input-projection quartet, all reading the same xNorm. CPU-correctness tests in `Tests/FFAITests/Ops/QuantizedOpsTests.swift` reuse a `makeInt4Projection` helper to build packed weights + scales + biases + expected GEMV output, then verify each batched wrapper matches CPU truth for every output tensor. Migrated from PR #9 / #10. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Each cache class owns one or two MTLBuffers that persist for the lifetime of a generation and get touched by every decode-step dispatch. Without explicit residency pinning, Apple's Metal driver re-validates per-allocation residency on every command-buffer encode — at decode rates of thousands of dispatches per token, the per-dispatch overhead dominates wall time. Reuses the device-wide MTLResidencySet introduced for weight buffers (see `Device.markWeightsResident`). The set is single-flighted under a lock and shared across all cache classes; the underlying Metal API gracefully no-ops on macOS < 15 and is gated behind `FFAI_NO_RESIDENCY_SET=1` for A/B comparison. - `KVCache.init` pins kBuffer + vBuffer - `GDNStateCache.init` pins current + next recurrent-state buffers - `ConvStateCache.init` pins the rolling-window state Migrated from PR #9 (Everything Bagel 1) — measured win was part of the 16.29 → 55.94 tps decode T=1 jump on Qwen3.6-A3B / M5 Max. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
The Qwen3-MoE decode path needs a scalar-broadcast FMA primitive to accumulate the top-K expert outputs into a single hidden-state vector. Without it the path was forced into a `Tensor.filled([hidden], weight)` host alloc per expert plus a `mul + add` pair per expert — a host roundtrip and 2 dispatches the GPU could absorb into one. Adds four wrappers (all in `Ops/OpsFused.swift`) over kernels that metaltile already emits: - `scalarFMA(scalar, value, base, into:)` — single-dispatch `out[i] = base[i] + scalar[0] * value[i]`. - `scalarFMAMany(scalars, values, acc:)` — N sequential dispatches reusing one compute encoder; Metal guarantees in-order encoder execution so the accumulation is safe. - `scalarFMAChain8(scalars, values, out:)` — exact 8-way fused dispatch (one kernel reads 16 inputs at once), collapses the topK=8 expert accumulator into a single dispatch. - `sigmoidScalarFMAResidual(gate, value, base, residual, into:)` — sigmoid-gated FMA with the post-MoE residual add folded in; Eric's `sigmoidScalarFMA` covers the non-residual variant. Tests in `Tests/FFAITests/Ops/OpsSpecialPathTests.swift` cover each wrapper with a closed-form expected value (4 new tests). Migrated from PR #9 / #10 (Everything Bagel 1 + 2) — measured as load-bearing for the decode T=1 16.29 → 96.3 tps win on Qwen3.6-A3B / M5 Max. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Three batched wrappers around the existing `dequant_gemv_int4_*` metaltile kernels that drive N back-to-back dispatches inside a single compute encoder, sharing the `input` binding and the `in_dim` / `group_size` constants. Saves N-1 encoder begin/end pairs per call. - `dequantGemvInt4Two` — per-expert MoE gate + up projections. - `dequantGemvInt4Three` — Qwen3-style attention q/k/v projections in the unfused decode path. - `dequantGemvInt4Four` — Qwen3.5 GDN mixer's qkv/z/b/a input-projection quartet, all reading the same xNorm. CPU-correctness tests in `Tests/FFAITests/Ops/QuantizedOpsTests.swift` reuse a `makeInt4Projection` helper to build packed weights + scales + biases + expected GEMV output, then verify each batched wrapper matches CPU truth for every output tensor. Migrated from PR #9 / #10. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Each cache class owns one or two MTLBuffers that persist for the lifetime of a generation and get touched by every decode-step dispatch. Without explicit residency pinning, Apple's Metal driver re-validates per-allocation residency on every command-buffer encode — at decode rates of thousands of dispatches per token, the per-dispatch overhead dominates wall time. Reuses the device-wide MTLResidencySet introduced for weight buffers (see `Device.markWeightsResident`). The set is single-flighted under a lock and shared across all cache classes; the underlying Metal API gracefully no-ops on macOS < 15 and is gated behind `FFAI_NO_RESIDENCY_SET=1` for A/B comparison. - `KVCache.init` pins kBuffer + vBuffer - `GDNStateCache.init` pins current + next recurrent-state buffers - `ConvStateCache.init` pins the rolling-window state Migrated from PR #9 (Everything Bagel 1) — measured win was part of the 16.29 → 55.94 tps decode T=1 jump on Qwen3.6-A3B / M5 Max.
The Qwen3-MoE decode path needs a scalar-broadcast FMA primitive to accumulate the top-K expert outputs into a single hidden-state vector. Without it the path was forced into a `Tensor.filled([hidden], weight)` host alloc per expert plus a `mul + add` pair per expert — a host roundtrip and 2 dispatches the GPU could absorb into one. Adds four wrappers (all in `Ops/OpsFused.swift`) over kernels that metaltile already emits: - `scalarFMA(scalar, value, base, into:)` — single-dispatch `out[i] = base[i] + scalar[0] * value[i]`. - `scalarFMAMany(scalars, values, acc:)` — N sequential dispatches reusing one compute encoder; Metal guarantees in-order encoder execution so the accumulation is safe. - `scalarFMAChain8(scalars, values, out:)` — exact 8-way fused dispatch (one kernel reads 16 inputs at once), collapses the topK=8 expert accumulator into a single dispatch. - `sigmoidScalarFMAResidual(gate, value, base, residual, into:)` — sigmoid-gated FMA with the post-MoE residual add folded in; Eric's `sigmoidScalarFMA` covers the non-residual variant. Tests in `Tests/FFAITests/Ops/OpsSpecialPathTests.swift` cover each wrapper with a closed-form expected value (4 new tests). Migrated from PR #9 / #10 (Everything Bagel 1 + 2) — measured as load-bearing for the decode T=1 16.29 → 96.3 tps win on Qwen3.6-A3B / M5 Max.
Three batched wrappers around the existing `dequant_gemv_int4_*` metaltile kernels that drive N back-to-back dispatches inside a single compute encoder, sharing the `input` binding and the `in_dim` / `group_size` constants. Saves N-1 encoder begin/end pairs per call. - `dequantGemvInt4Two` — per-expert MoE gate + up projections. - `dequantGemvInt4Three` — Qwen3-style attention q/k/v projections in the unfused decode path. - `dequantGemvInt4Four` — Qwen3.5 GDN mixer's qkv/z/b/a input-projection quartet, all reading the same xNorm. CPU-correctness tests in `Tests/FFAITests/Ops/QuantizedOpsTests.swift` reuse a `makeInt4Projection` helper to build packed weights + scales + biases + expected GEMV output, then verify each batched wrapper matches CPU truth for every output tensor. Migrated from PR #9 / #10.
…en3.6 Integration test from PR #9 that compares baseline greedy decode (forward loop) vs SpecDecode.generateGreedy + NGramDrafter at gamma in 1, 2, 4. Verifies generated sequences match baseline exactly (greedy + greedy drafter + correct snapshot/restore is deterministic) and reports decode-only tps so the cost of the verify forward is visible. Gated on the local Qwen3.6-A3B checkpoint at /Users/tom/models/Qwen3.6-35B-A3B-4bit; skips with a print when the path is missing.
|
Superseded by PR #14 |
…d greedy VisionModel.runGreedy -> runSampled: the multimodal decode (image/audio/video -> text) now samples every step via the shared Sampling.sample (temperature / top-p / top-k / min-p / repetition-penalty / seed), threading tokenHistory for the rep-penalty and an RNG from the seed — exactly like the text path. The prefill samples the final position's logits instead of argmax. generate(...) gains an optional 'parameters:' (nil -> deterministic greedy; real callers pass the model's defaultGenerationParameters). Gemma 4 image/audio/video integration tests updated to pass defaultGenerationParameters + seed 42 and all pass with sampled defaults (temp 1.0 / top-k 64): 'a beautiful Golden Retriever dog', 'sure i can help you with that', cat. Roadmap item #2 done. (Full streaming output via Generate.generateStream is still §7a / roadmap #9.)
…le harness Decode applied rope_theta=10000 to ALL 43 layers, but compressed layers (compress_ratio != 0 → layers 2-41) must use compress_rope_theta=160000 (DwarfStar ds4.c layer_rope_freq_base: ratio != 0 ⇒ 160000). Wired per-layer theta into the Q/KV/inverse rope calls in forwardFullAttnSubblock. Combined with advancing state.position per token (the validation loop was leaving every token at position 0 → identity rope), this removes the abnormal logit blow-up (was ~60 for a wrong token; now a normal ~17 distribution) on the oracle prompt. Added dsv4bench --prompt-ids for apples-to-apples oracle comparison against the working engines (llama.cpp fork + DwarfStar, both on the same GGUF, both give 'Tokyo'=30228 for token IDs 671,6102,294,6310,344). Remaining for correctness = task #9: wire the (already-loaded) CSA/HCA compressor + Lightning indexer + compressed-SDPA into the attention forward (all weights/Ops/kernels/config present; needs compressed-KV state in LayerKVState + the per-compress_ratio forward branch). Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Co-authored-by: tturney <tturney@psyguard.ai>
forwardFullAttnSubblock now runs the per-layer KV compressor on compress_ratio != 0 layers (DwarfStar compressor_decode_one ported to compressorStepCPU): projects attn_norm via attn_compressor_kv/gate, adds APE, rolls the per-token window, and on a ratio boundary emits one compressed KV row (per-dimension softmax pool — two-lane overlap for CSA, single-lane for HCA — RMSNorm(compressor_norm), compressed-RoPE). SDPA now attends [raw sliding window ∪ all compressed rows] (dense; the Lightning indexer top-512 is a no-op while n_comp <= 512). LayerKVState gained compCache + rolling window state. Also fixed: forwardFullAttnSubblock/makeDecodeState read cfg.layerCompressRatios (textConfig, loaded from the wrong short metadata key → all zeros) instead of the model's correctly-resolved layerCompressRatios (full key deepseek4.attention.compress_ratios + structural fallback). Verified structurally on the oracle prompt: ratios = [0,0,4,128,...,4,0] (matches official config); CSA layers emit 1 compressed row for a 5-token prompt, HCA emit 0. Numerical correctness of the compressed row + YaRN/mscale still being validated against DwarfStar. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Co-authored-by: tturney <tturney@psyguard.ai>
…patch
Adds the MetalTileSwift wiring that lets vllm-swift dispatch the highest-
impact metaltile kernels (mt_qmm_mma + mt_moe_gather_qmm_mma_int4_bm{16,64}_mpp)
in place of MLX's quantizedMatmul for Qwen3.6-class hybrid MoE serving.
Files:
- swift/Package.swift: depend on MetalTileSwift (pre-compiled metaltile
Metal library + Swift dispatch wrappers).
- swift/Sources/VLLMBridge/MetalTileShim.swift: env-gated bridging layer.
MLXArray ↔ MTLBuffer via MLXArray.asMTLBuffer; dispatch on the
MetalTileLibrary's command queue; results reconstructed via
mlx_array_new_data from MTLBuffer.contents(). Off by default
(VSM_USE_METALTILE=0); opt-in via VSM_USE_METALTILE=1.
- swift/Sources/VLLMBridge/Bridge.swift: connectivity proof at startup
(kernel count print + smoke dispatch of mt_qmm_mma at the Qwen3.6
prefill shape M=32 N=256 K=2048 fp16) plus per-call dispatch
integration in the hybrid sparse forward path.
Note: the shim does its own commit + wait per dispatch, which
caps the perf win at long contexts (see project notes — cross-queue
sync from Swift would unblock more, but is multi-day work on its own).
Numbers in flight; this PR exists so the diff is visible while the
companion changes land:
- metaltile-ffai PR for the kernel pack + emit + naming
standardization: 0xClandestine/metaltile#144
- FFAI PR for the matching Ops wrappers + Qwen35 fused-GDN-prep
wiring: thewafflehaus/FFAI#9
Together the three repos let us run Qwen3.6-A3B end-to-end via
vllm-swift → ffai → metaltile, with the integration story now bench-
validated on the FFAI leg (+52% prefill / +20% decode on M5 Max with
the new mt_gated_delta_prep_step kernel).
2 participants
Qwen3.6-A3B batched prefill on M5 Max: 260.76 tps at T=512 batched, 30x over per-token forward, 43.5x cumulative vs morning's 6.0 tps baseline. Best single run 1.449s = 353.3 tps. Pairs with metaltile-ffai PRs #144 (kernel-pack + emit overhaul) and #145 (follow-up MPP / steel-gemm).
Summary
End-to-end batched prefill path on Qwen3.5 / Qwen3.6-A3B + the metaltile-emit kernel-pack integration that unblocks it.
forwardMany/decodeManyplumbed throughQwen35Model,Qwen35AttentionLayer,Qwen35GDNLayer,Qwen35MoEFFN,Qwen35DenseMLP,MoELayer,AnyLinear,QuantizedLinear.forwardManyEquivalencetest passes at T=8 and T=128 (argmax matches reference within bf16 noise).MoELayer.decodeManyfans T*topK rows into one BGEMM per gate / up / down, sorted by expert id so BM=16 tiles fill 16-deep with shared expert weights. The very samemoeGatherDequantGemmInt4kernels that regress at T=1 decode now win at this shape.Qwen35GDNMixer.forwardManylifts 5 GDN projections out of the per-token loop; recurrence stays per-token (state crosses tokens) but rides one in-flight cmd buffer.FFAI_GDN_FUSED_PREP=1is load-bearing for the batched path.Ops.siluCastToF32(drops T30 dispatches per prefill),mt_moe_unpermutescatter-sum (collapses TtopK*2 small dispatches per stage), fusedmt_swigluin MoE expert path, fusedmt_sigmoid_scalar_fmafor shared expert,Ops.gatedDeltaPrepStep,Ops.dequantGemvIndirect,Ops.copyblit.forwardManyEquivalence T=128pins the path.Ops.dequantGemmDynamicMhost-side T-padding driver aroundmt_qmm_mma_{f32,f16,bf16}(fast path on T%32==0).topKIndicesO(n log n) -> O(n log k), skip full softmax whennormTopKProb=false, paralleliserouter.routeacross T tokens.castToF32lifted out of GDN T-loop.Device.markWeightsResidentpins allSafeTensorsBundleMTLBuffers in anMTLResidencySetattached to the command queue. Skips per-encode residency re-validation across thousands of dispatches per prefill.MTLResidencySettypedAny?so themacOS 15+cast sits inside@available; deployment target unchanged. Opt-out viaFFAI_NO_RESIDENCY_SET=1. KV / GDN / conv state buffers also pinned.scaleBySigmoidGate35host wait.mt_*/ffai_*/vector_*).FFAIStubs.swiftoriginally added 36 trap-on-call shims for kernels FFAI referenced but metaltile-ffai didn't yet emit; now dropped — indirect kernels are emitted by tile emit (PR #145).metaltilePSOCache live-compiles MPP kernels to dodge SDK skew on the macOS 26 host.Bench (M5 Max, Qwen3.6-A3B, 5-run median)
Cumulative session vs morning's 6.0 tps starting point: 43.5x.
Wins scale with T — at T=512,
mTotal = T * topK = 4096fills the BM=16 tile 256-deep with shared expert weights. Per-token comparison degrades with T because GDN's recurrent state walks T steps serially.forwardManyEquivalence(T=8, T=128) passes.Env gates
FFAI_GDN_FUSED_PREP=1— load-bearing for batched path.FFAI_MOE_BGEMM=1— enables MoE BGEMM at T*topK.FFAI_MOE_BGEMM_BM8=1— opt-in bm8 path whentopK <= 8(off by default; wins only at prefill shapes).FFAI_NO_RESIDENCY_SET=1— opt-out of weight residency pin.FFAI_MAX_COMMAND_BUFFERS— Metal queue depth tuning (default 16; 48-64 unlocks prefill pipelining).Files touched
Sources/FFAI/Models/Qwen35.swift,Sources/FFAI/Models/MoELayer.swift,Sources/FFAI/Models/MoERouter.swiftSources/FFAI/Models/GatedDelta.swift(GDN inner loop)Sources/FFAI/Ops.swift,Sources/FFAI/Layers/AnyLinear.swift,Sources/FFAI/Layers/QuantizedLinear.swiftSources/FFAI/Device.swift(markWeightsResident,MTLResidencySetplumbing)Sources/FFAI/Model.swift(residency hook inload)Sources/FFAI/Caches/*(KV / GDN / conv state pinning)Sources/MetalTileSwift/Stubs/FFAIStubs.swift(introduced, then dropped after metaltile #145 lands the indirect kernels)Tests/FFAITests/Qwen36SmokeTests.swift(forwardManyBench,forwardManyEquivalence)Branches
feat/metaltile-kernel-pack-integration-> baseek/aura-port