WIP: metaltile shim — Qwen3.6 mt_qmm_mma / mt_moe_gather dispatch from VLLMBridge#22
Draft
TheTom wants to merge 8 commits into
Draft
WIP: metaltile shim — Qwen3.6 mt_qmm_mma / mt_moe_gather dispatch from VLLMBridge#22TheTom wants to merge 8 commits into
TheTom wants to merge 8 commits into
Conversation
Hybrid Qwen3.5 / Qwen3.6 prefill > 2048 crashed with broadcast mismatch `(1, n_kv, T+1, head_dim)` vs `(1, n_kv, 2048, head_dim)`. Root cause: both batched-cache construction paths (`initBatchedHybrid` + `prefillBatchedUniformHybrid`) passed `engine.generateParams` straight through to `model.newBatchedHybridCache(parameters:)`, which falls back to `maxSeq = parameters?.maxKVSize ?? 2048`. Mirror the sizing logic the non-hybrid path already uses. Three changes in `prefillBatchedUniformHybrid` and one in `initBatchedHybrid`: 1. Inject `maxKVSize = T + 2 + decodeMargin` (or `maxPrefillOffset + decodeMargin` for `initBatchedHybrid`) when the caller didn't override it, so the BatchedKVCache fits the actual prefill. 2. Cap `maxBatch = B` for `prefill_batched_uniform`. Previous default `max(B, engine.maxConcurrentRequests=64)` allocated 64 attn KV slots even at B=1 — 40+ GB at ctx=32K. The batched-uniform path knows the batch up front; new concurrent requests trigger a fresh init. 3. Chunk the prefill via `lmModel.defaultPrefillStepSize` + asyncEval per chunk + final `Memory.clearCache()`. Mirrors `TokenIterator.prepare` (mlx-swift-lm `LLMModel.swift`). Previous single-call `lmModel(prefillChunk, ...)` materialized all T-1 token activations in one graph, ballooning MLX's buffer pool to 30+ GB at ctx=32K. Chunked path bounds the pool to one chunk's activations. 4. Memory-snapshot instrumentation (`memLog`) at engine-create, prefill-alloc-before/after, periodic decode (every 50 steps), and engine_destroy — for diagnosing future long-ctx memory issues. Validated against Qwen3.6-35B-A3B-4bit B=1: ctx prefill-tok/s decode-tok/s cache-pool 2048 2,172 113.8 ~1 GB 4096 2,288 114.8 ~1.3 GB 8192 884 107.4 ~2 GB 16384 440 113.2 ~3 GB 32768 751 104.3 ~3.5 GB Coherency check at ctx=32768: model answers a question planted at the end of a 32,728-token prompt with a faithful one-sentence summary. Prefill rate degrades >4K because GDN's MLX `gated_delta_step` kernel runs sequentially per-token within each threadgroup. Separate work (spec 028 chunked-WY) tracks that. Local-only; no push to origin per Tom's per-push-confirmation rule. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com> Co-Authored-By: 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).
Owner
Author
Three-repo integrationThis PR is the serving-layer leg of the
|
Production path is moving to `vllm-swift → FFAI → metaltile`; kernels flow through FFAI's inner `Sources/MetalTileSwift/` target once FFAI is wired in. The standalone `/Users/tom/dev/MetalTileSwift` Swift package has been stuck on the pre-PR-#144 emit (no `mt_qmm_mma_mpp_bf16` etc.) and is no longer a viable dep route. Three coordinated removals: - `swift/Package.swift`: drop the `MetalTileSwift` package dep + the `.product(name: "MetalTileSwift", ...)` from VLLMBridge target deps. - `swift/Sources/VLLMBridge/MetalTileShim.swift`: delete (was the hand-written wrapper layer over `MetalTileKernels.mt_qmm_mma_f16` etc.; only used by the connectivity-proof smoke dispatch in `Bridge.swift`). - `swift/Sources/VLLMBridge/Bridge.swift`: drop the `import MetalTileSwift`, the `MetalTileLibrary.shared` connectivity proof, and the `MetalTileShim.qmmMma(...)` smoke dispatch. Replaced with a one-line comment pointing at the new path. `~/dev/MetalTileSwift` is intentionally left on disk — no deletion, just no longer referenced from this branch. Re-add once vllm-swift takes its FFAI dep. Validated: `swift build` clean (210/211 → libVLLMBridge.dylib), no references to MetalTileKernels / MetalTileLibrary / MetalTileShim remain in `swift/Sources/`.
Adds .package(path: ~/dev/FFAI) + the FFAI product to the VLLMBridge target, per the pre-planned architecture (the MetalTileShim was removed in anticipation). FFAI carries the GGUF DSv4-Flash engine (Swift) and the metaltile-generated Metal kernels (its inner MetalTileSwift target). Verified VLLMBridge builds clean with both FFAI and mlx-swift-lm as deps — the two Metal kernel bundles + dependency graphs coexist without conflict. Bridge→FFAI is native Swift→Swift; the only C-ABI edge stays at the Python↔dylib boundary (@_cdecl). Engine-routing (an FFAI path in vsm_engine_create/decode_step for DSv4-GGUF) lands next. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Co-Authored-By: tturney@psyguard.ai
Adds FFAIDsv4Engine (Swift): detects a DeepSeek-V4 GGUF dir, loads via FFAI.DeepSeekV4Flash.loadFlashFromGGUF, holds the DecodeState. Registered in a parallel ffaiEngines handle map so the MLX path is entirely unaffected — vsm_engine_create early-returns an FFAI handle for DSv4-GGUF; destroy/vocab_size/num_layers branch to it. Builds clean with both FFAI and mlx-swift-lm deps. Decode routing (prefill/decode_step/get_logits → forwardAllLayers) lands next; batched paths stay MLX (FFAI is single-stream). Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Co-Authored-By: tturney@psyguard.ai
… vLLM-swift Wires the FFAI engine through the single-stream vsm_engine_* C-ABI: - prefill_req: feed the prompt sequentially (FFAI has no batched prefill), return the first greedy token. - decode_step_req: feed the carried lastToken, return next greedy token. - get_logits: hand Python a pointer into FFAI's stable logits buffer (previously a no-op stub even for MLX — FFAI leads here). - reset: fresh DecodeState. FFAIDsv4Engine holds DecodeState + lastToken + a stable logits buffer; forwardAndArgmax copies logits + returns greedy argmax. All native Swift→Swift; the @_cdecl C exports are unchanged. Only the `_default` single-stream request routes to FFAI — batched/multi-session and the real long-context (CSA/HCA, Track B) stay on their own paths. Builds clean with FFAI + mlx-swift-lm. End-to-end: a DSv4-GGUF dir now loads + decodes through vLLM-swift at ~32 tps single-stream. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Co-Authored-By: tturney@psyguard.ai
…metaltile Mirrors engine_bridge.py's exact call sequence (create/prefill/decode_step/ get_logits/destroy) against libVLLMBridge.dylib + the GGUF DSv4-Flash dir. VERIFIED working: engine loads (vocab=129280, layers=43), prefill returns a token, get_logits hands back real logits, decode sustains ~31.9 tps through the full vllm-swift → FFAI → metaltile chain. Confirms the bridge serves the GGUF model at the same ~32 tps the standalone FFAI bench measures — Swift engine, Rust kernels, C only at the Python boundary. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Co-Authored-By: tturney@psyguard.ai
Owner
Author
Integration milestone: bridged FFAI hits DeepSeek-V4 prefill parity 🎯No code change here — flagging that the FFAI side this shim bridges now clears prefill parity, so the integration is serving a parity-class DSv4 backend:
Relevant for this shim: the metaltile |
… prefill Extract storeLogitsAndArgmax from forwardAndArgmax so both the per-token decode path and a prefill-chunk path reuse the logits → stable-buffer → greedy-argmax tail; forwardAndArgmax now also advances state.position. Add an opt-in batched-prefill path gated on VSM_FFAI_EXPERIMENTAL_BATCHED_PREFILL=1 that calls FFAI forwardPrefillChunk. Default remains the validated sequential per-token prefill — the batched path still diverges on some short prompts.
TheTom
force-pushed
the
integration/metaltile-ffai
branch
from
June 3, 2026 20:05
718c1d0 to
cc3a4ed
Compare
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Draft / WIP. Surfacing the vllm-swift leg of the three-repo
vllm-swift → ffai → metaltileintegration so the diff is visible while the companion PRs settle.Two commits
eedf3f8(pre-existing):fix(bridge): unblock Qwen3.6 long-ctx serving — break 2048 barrier—initBatchedHybridandprefillBatchedUniformHybridwere passing rawgenerateParamstonewBatchedHybridCache(parameters:), which defaultedmaxKVSize = 2048. Fixed by injectingmaxKVSize = T + 2 + decodeMarginwhen the caller doesn't override. The non-hybrid path already had this logic; this brings hybrid (Qwen3.5 / Qwen3.6) into parity.c85cd41(this slice): wireMetalTileSwiftintoVLLMBridgeso the highest-impact metaltile kernels (mt_qmm_mma+mt_moe_gather_qmm_mma_int4_bm{16,64}_mpp) can substitute for MLX'squantizedMatmulin the Qwen3.6 forward path.What
c85cd41addsswift/Package.swift— depend onMetalTileSwift(the pre-compiled metaltile Metal library + Swift dispatch wrappers — produced frommetaltile-ffaiPR #144).swift/Sources/VLLMBridge/MetalTileShim.swift(NEW, 114 LOC) — env-gated bridging layer.MLXArray ↔ MTLBufferviaMLXArray.asMTLBuffer(mlx-swift public API). Dispatch on theMetalTileLibrary.shared.commandQueue. Result MLXArrays reconstructed viamlx_array_new_dataagainstMTLBuffer.contents(). Off by default (VSM_USE_METALTILEunset); opt-in viaVSM_USE_METALTILE=1.swift/Sources/VLLMBridge/Bridge.swift— startup connectivity proof (kernel count print + smoke dispatch ofmt_qmm_mmaat the Qwen3.6 prefill cellM=32 N=256 K=2048 fp16) + per-call dispatch integration on the hybrid sparse forward path.Companion PRs
0xClandestine/metaltilethewafflehaus/FFAIBench (M5 Max, FFAI side)
Bench from the FFAI integration — proves the kernel pack is making the combined product faster, not just compiling:
Correctness preserved (Johannes / Gutenberg prompt produces argmax 52290
Johannesunder both paths, top-5 logits match within ULP).Known caveats
quantizedMatmuldirectly in most of its hybrid forward paths; this PR's shim is the entry point for replacing those calls. Per-call wiring is incremental — the connectivity proof + the smoke dispatch are the first step.Stream.commandBufferfrom mlx-swift.mlx-swift-lmend-to-end withFFAIundervllm-swift— the north star path that lets us drive Qwen3.6 throughvllm-swift → ffai → metaltilewithout MLX in the model dispatch — is multi-week scope (Bridge.swift is 4000+ LOC deeply integrated withmlx-swift-lm). Tracked separately; the shim shipped here is the smaller, immediate-value step.WIP. Not blocking; landing in pieces so the diff is reviewable.