Everything Bagel 2 - #10
Closed
TheTom wants to merge 167 commits into
Closed
Conversation
fix(release): dispatch notify-docs.yml explicitly after publishing
fix(ci): build tests so --skip-build doesn't fail on cold cache
orpheus-style tts on a llama 3.x backbone: reuses the existing LlamaModel engine for the acoustic decoder and adds the orpheus token protocol — prompt framing, the autoregressive snac-code decode loop, and snac plane de-interleaving. registered in AudioModelRegistry; the snac neural codec (waveform tail) is left behind a SNACDecoding protocol for a separate codec port. also passes the new has_sink/sink_logit params to the head_dim=128 ffai_sdpa_decode kernels (0/0 dense default — matches the regenerated kernel signature; gpt-oss continues its software sink fold).
sesame conversational-speech-model tts: a dual-transformer acoustic model — a backbone and a depth decoder, both built on FFAI's LlamaLayer blocks — plus text/audio-codebook embedding tables, codebook-0 + per- codebook audio heads, and the autoregressive frame-generation loop. generateFrames emits the [K, nFrames] mimi code matrix; the mimi neural codec (waveform tail) is left behind a MimiDecoding protocol for a separate codec port. registered in AudioModelRegistry.
qwen3tts is a four-part tts model (talker + code predictor + ecapa speaker encoder + intrinsic speech-tokenizer codec) — a staged port. this lands stage 1: decodes the nested talker_config / tokenizer_config / speaker_encoder_config blocks, detects the family, and registers it in AudioModelRegistry so a checkpoint reports textToSpeech. pins the typed Qwen3TTSModel / talker-config surface the later stages build on. synthesize throws Qwen3TTSError.synthesisNotWired until the talker (qwen3 stack + 3d mrope), code predictor and codec stages land.
unit tests for LlamaTTS / Marvis / Qwen3TTS — registry detection, config decoding, and the LlamaTTS SNAC de-interleave — so the families have fast parallel-safe coverage independent of the network-gated integration suites.
# Conflicts: # Sources/FFAI/Ops.swift
Ports the Qwen2.5-VL vision-language model: a dynamic-resolution windowed-attention ViT tower whose projected image tokens splice into the Qwen 2.x text backbone. The text backbone runs on the Llama dense engine, which gains an embedding-input forward path (forward over a supplied [hidden] row instead of an embedding-table gather) plus a single-token textEmbedding lookup for the text half of the VL splice. ModelRegistry routes Qwen2_5_VLForConditionalGeneration checkpoints.
VAD models have a distinct contract (audio waveform in, per-frame speech-probability stream out), so they load through a dedicated VADModelRegistry rather than ModelRegistry or AudioModelRegistry. SileroVAD is a streaming STFT + gated-conv detector with 16 kHz / 8 kHz branch configs; SmartTurn is a conversational endpoint detector. Both expose loadFromDirectory / fromPretrained and detect(audio:sampleRate:) returning a VADOutput. VADCompute holds the shared front-end helpers.
SNAC is a multi-scale residual-VQ codec — the waveform tail for Orpheus-style (LlamaTTS) synthesis. encode(waveform:) turns audio into discrete code planes; decode(codes:) renders them back. Lands under Sources/FFAI/Audio/ alongside the shared codec primitives (AudioMath, WeightNorm). Documents the new VAD, vision-language and codec sections in models.md.
Stack-interleaved hybrid of double-gated short-convolution layers and GQA attention layers. One family file covers the LFM2 and LFM2.5 collections (architecturally identical) plus the LFM2-MoE variant. Extends MoERouter with an optional per-expert additive bias for LFM2-MoE's softmax-then-bias routing.
Four unrelated compile errors blocked the whole module: - toFloatArray() declared in both Tensor.swift and VADCompute.swift — removed the duplicate, folded the integer-dtype cases into the canonical Tensor version. - Gemma4VL vision-linear loader: `??` whose RHS throws needed the operator expression marked `try`. - EncodecBlocks numQuantizers: mixed-literal expression overwhelmed the type-checker — split into explicitly-typed sub-expressions. - SenseVoice registry test: missing `await` on an async load call.
Cohere Command A+, the multimodal batch (Marlin / MiniCPM-V / SAM 3.1 / Falcon-OCR), and the LiquidAI LFM2 / LFM2.5 / LFM2.5-VL design docs.
`mt_*_mpp_*` kernels use Apple's `mpp::tensor_ops::matmul2d` cooperative tensors from `<MetalPerformancePrimitives/MetalPerformancePrimitives.h>`. The MPP header inlines cooperative-tensor type IDs + descriptor layouts that drift between SDK versions; the live OS runtime resolves them via its bundled header, while `xcrun metal -c` (run when `tile build --emit` produces `kernels.metallib`) resolves them via the SDK header at build time. When SDK macOS != runtime macOS (here SDK 26.5, runtime 26.4.1), the type IDs/descriptor layouts baked into the offline-compiled metallib disagree with the device runtime. Result: bit-deterministic wrong output on MPP kernels — bm64_mpp_bf16 at Qwen3.6-A3B down shape produced cos 0.816 vs m1 oracle (vs 0.999984 when live-compiled). forwardManyEquivalence with FFAI_MOE_BGEMM_BM64=1 then tripped argmax 279 vs ref 52290. Fix: in PSOCache, detect MPP kernels by name (`_mpp_` substring), locate the `.metal` source emitted alongside `kernels.metallib` (in Resources/kernels/<name>.metal), and compile via `device.makeLibrary(source:options:)` at PSO build time. The runtime header is then used → MPP type IDs match → kernel produces correct output. - Non-MPP kernels (the bulk of the metallib) still load from kernels.metallib — no perf change. - Opt-out via `FFAI_PSO_LIVE_COMPILE_MPP=0` for fallback to metallib loading if a future kernel name accidentally contains `_mpp_`. - PSO compilation is single-flight via compileLock so the extra cost (a few ms of source compile per kernel, paid once at first dispatch) doesn't multiply across concurrent callers. Verified: - forwardManyEquivalence Qwen3.6-A3B T=8 with FFAI_GDN_FUSED_PREP=1 FFAI_MOE_BGEMM=1 FFAI_MOE_BGEMM_BM64=1 → ref argmax 52290 == batched argmax 52290, top-5 logit drift < 0.15 - New MoEBgemmBm64MppTests.swift covers clean_tile, multi_tile, qwen3.6 gate/up, qwen3.6 down (production shape that triggered the bug), down-with-sin-inputs, raw-dispatchThreadgroups, and a live-compile cross-check. All ≥ 0.999 cosine vs m1 oracle. See ollama #15594, #14432, llama.cpp PR #16634 for the same class of bug across Metal 4 / MPP integrations. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
…t loader Merge waffles/ek/aura-port (0e68f05) into feat/metaltile-kernel-pack-integration. Conflicts resolved in MoELayer.swift, Qwen35.swift, Ops.swift: - MoELayer: take aura-port's LFM2-MoE expert_bias support (strict superset of HEAD; passthrough when expertBias=nil) - Qwen35: keep HEAD's activation-dtype-from-quantized-embed-scales path (Qwen3.6's quantized embed needs scales, not the u32 weight) and HEAD's lmHeadPrefix-driven lm_head loader; add aura-port's quantized tied-embed fallback - Qwen35.buildGDNMixer: drop the duplicate `quant` parameter the text-merge introduced (left only the `quantization quant:` one) - Qwen35: keep BOTH HEAD's batched-prefill forwardMany / _forwardManyPerTokenLegacy / _forwardManyBatched AND aura-port's VLM embedding-input forward(inputEmbedding:) + textEmbedding helpers - Ops.swift: take HEAD wholesale (all our batched-prefill MoE / SDPA / GDN / dynamic-M / scatter-sum wrappers stay), add Ops.layerNorm (mt_layer_norm kernel ships in our metallib) and stub the 6 aura-port-only VLM/audio ops (conv2d, patchEmbed, rope2D, melSpectrogram, audioConv1d, vocoderISTFT) with fatalError — their metaltile kernels live on origin/ek/aura-port, NOT on clandestine/dev, so the dependent VLM/audio Swift modules compile but fail loudly at first use. Rebase those metaltile kernels onto dev to drop the stubs. VL-gate fix in Model.swift: Qwen3.6 text-only checkpoints sometimes carry a vestigial `vision_config` block — the aura-port VL-detection gate then misroutes them to the visionModelNotIntegrated throw. Route `Qwen3_5MoeFor*` architectures through loadQwen35 inside the VL gate before the throw. Verified post-merge: - swift build clean - FFAI_GDN_FUSED_PREP=1 FFAI_MOE_BGEMM=1 forwardManyEquivalence T=8 ✓ (ref argmax=52290 == batched argmax=52290) Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
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.
metaltile now generates the dequant_gemv_int4_{f16,bf16}_indirect
Swift wrappers (metaltile feat(emit): _indirect wrapper generator), so
the hand-written fatalError shims are no longer needed. Delete
FFAIStubs.swift and re-enable IndirectDispatchTests — both indirect
cells pass bit-for-bit against the direct dispatch.
All 528 unit tests pass.
TheTom
added a commit
that referenced
this pull request
May 26, 2026
…ateKVMany, conv1dCausalStepSiluCastMany, gatedDeltaPrepChunk Eric's dev only has the single-token forms of these ops; the PR9/10 forwardMany / decodeMany prefill paths replace the per-token T-loop with a single batched dispatch. - `ropePartialMany` — single dispatch covering T rows of `qk` rotated by per-row `positions[T]`; saves T-1 encoder pairs and matches T sequential `ropePartial` calls bit-identically. - `ropePartialManyTwo` — Q + K share `positions` and constants on one encoder; row strides differ. - `kvCacheUpdateKVMany` — writes T tokens' K + V rows in one encoder, indexed by `positions[T]`. - `conv1dCausalStepSiluCastMany` — `T` depthwise conv + SiLU + fp32 cast collapsed into one dispatch with the conv state held in per-channel registers across the sweep. - `gatedDeltaPrepChunk` — chunked-prefill counterpart to `gatedDeltaPrepStep`; runs the fused GDN prep + recurrence over a token range in one dispatch, keeping per-head recurrent state in registers across the chunk. `tLen` arrives as a `[1]` u32 GPU buffer so chunked dispatch chains avoid CPU readback. All kernels (`ffai_rope_llama_many_*`, `kv_cache_update_many_*`, `ffai_conv1d_causal_step_silu_cast_many_*`, `mt_gated_delta_prep_chunk_*`) are already emitted on the dev metaltile. Tests: - `ropePartialManyMatchesSequential` — T=3 batched matches three sequential `ropePartial` outputs bit-identically. - `kvCacheUpdateKVManyF32` — verifies multi-position K + V writes land in the right cache rows and leave untouched slots zero. Migrated from PR #10 (ITERs 86–99). Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
TheTom
added a commit
that referenced
this pull request
May 26, 2026
…coder Mirror of `dequantGemvInt4Two/Three/Four` for the per-expert MoE down phase, where each chosen expert has its own activation tensor but all N projections share PSO, `inDim`, and `groupSize`. The wrapper sets the constexpr binding once and rotates `(weight, scales, biases, input, output)` per dispatch, saving N-1 encoder begin/end pairs versus N independent `dequantGemvInt4` calls. Test verifies the batched form matches three sequential `dequantGemvInt4` CPU-truth references across distinct (scale, bias) projections + distinct inputs. Migrated from PR #10. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
TheTom
added a commit
that referenced
this pull request
May 26, 2026
…code Same dispatch contract as `sdpaMulti(causal: false)`, but the scalar `causal` flag is replaced by an additive `[nQuery, nQuery]` mask tensor consulted only for in-block KV positions. Positions `< baseKV` (the cached prefix) are always attended. Mask is `0.0` to allow and `-inf` to block, added directly to pre-softmax scores. Used by speculative-decode tree-verify: one verifier call attends every leaf-to-root path of the draft tree at once, gated by a tree-causal mask the drafter produces from its node structure. Wraps `ffai_sdpa_multi_tree_mask_*` (headDim=128 only — no d256 variant for tree-verify on dev). Tests: - All-zero mask reproduces `sdpaMulti(causal: false)` output bit-equivalent (sanity check on the mask add path). - Diagonal-only -inf mask isolates each query to its own KV slot: with uniform Q/K and per-row distinct V, output collapses to a copy of the matching V row. Migrated from PR #10 ITER 70. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
TheTom
added a commit
that referenced
this pull request
May 26, 2026
Pass 1 partitions the KV cache into `blocks` even tiles and writes per-block partial `(O, max, lse)` triples; pass 2 merges them with the log-sum-exp trick. Total work matches single-pass `sdpaDecode` but the two passes parallelize across more GPU threadgroups, paying off at long KV (≥ 1K typical) where the single-pass kernel starves the GPU with one TG per Q-head. Block-count contract (the load-bearing detail PR10's port silently violated): the pass2 kernel hardcodes a 32-lane block reduction (`bn = 32`, `block_chunks = blocks / bn`), so `blocks` MUST be a multiple of 32 AND at least 32. Anything smaller makes `block_chunks` round to zero — pass2's outer loop iterates zero times and the output stays whatever pass1 left in scratch, typically all zeros. The previous port's "power of 2" check let `blocks=2` through and was the cause of the all-zero smoke failures during the earlier port attempt; the new precondition catches it at the wrapper. Scratch buffers (`partialO`, `partialM`, `partialL`) are caller-owned so a decode loop can amortise them across steps. Test: `blocks=32, nKV=64, nQHeads=2` — verifies the two-pass output matches single-pass `sdpaDecode` to within 1e-2 absolute. Migrated from PR #10 ITER 53. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
TheTom
added a commit
that referenced
this pull request
May 26, 2026
…ding Speculative-decode candidate-proposer interface used by the in-progress `SpecDecode` driver. Three implementations land here: - `NGramDrafter` — prompt-lookup n-gram. Scans `history` for the last `nMatch` tokens, falls back to shorter n-grams down to `minNMatch`, returns up to `gamma` follow-on tokens from the earlier occurrence. Zero ML cost; the workhorse for code and structured-chat workloads where prior context recurs. - `NeverDrafter` — stub that always returns `[]`. Used to bench pure spec-decode driver overhead vs raw single-token decode. - `NGramTreeDrafter` — real branching tree drafter. Conforms to both `Drafter` (linear γ fallback via `propose`) and `TreeDrafter` (`proposeTree` returning top-K continuations per depth) for tree-verify spec decode against `Ops.sdpaMultiTreeMask`. `DraftTreeNode` is the pure-Swift tree representation, with `flatten()`, `verify(oracleAtHistoryEnd:oracle:)`, and `treeCausalMask()` helpers — all pure functions, no model / cache dependencies, ready to wire into the SpecDecode driver once `forwardManyAllLogits` lands. `LinearTreeAdapter` lets any linear `Drafter` masquerade as a `TreeDrafter` (degenerate single-branch tree) so the tree-verify driver can A/B against the linear baseline using the same `NGramDrafter` instance. Tests in `Tests/FFAITests/Generation/DrafterTests.swift` cover all three drafters end-to-end: 20 host-side tests, no GPU dependencies. All pass. Migrated from PR #10 (ITERs 61, 68, 69, 71). Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
TheTom
added a commit
that referenced
this pull request
May 26, 2026
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.
TheTom
added a commit
that referenced
this pull request
May 26, 2026
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.
TheTom
added a commit
that referenced
this pull request
May 26, 2026
Both Q and K go through the same partial RoPE rotation with identical `(headDim, rotaryDim, position, thetaBase, scaling)` parameters every decode step on Qwen3-class models. `ropePartialTwo` sets the RoPE constants once and dispatches both rotations through a single shared compute encoder, saving one encoder begin/end pair per attention layer per token. The kernel is unchanged (`ffai_rope_llama_*`); the wrapper is a buffer-binding rotation around the existing in-place rotation. A test verifies bit-identical output against two sequential `ropePartial` calls on the same inputs. Migrated from PR #10.
TheTom
added a commit
that referenced
this pull request
May 26, 2026
Pairs two `ffai_gather_*` dispatches against the same `table` binding inside one compute encoder, saving an encoder begin/end pair when the caller needs to fetch two parallel id streams from the same embedding (e.g. the Qwen3 attention head-half split path). Migrated from PR #10.
TheTom
added a commit
that referenced
this pull request
May 26, 2026
Two `mt_rms_norm_*` dispatches inside one compute encoder, used by the Qwen3 attention mixer's pre-RoPE Q-norm + K-norm pair: both projections run the fast path (TPG = rowSize / 4) and have no data dependency between them. The wrapper shares a single eps buffer allocation when the two eps values are equal — the common case since both norms come from the same checkpoint config. The wrapper restricts callers to the fast `mt_rms_norm_*` path (rowSize ≤ 4096 and a multiple of 128). Wider rows must fall back to two separate `rmsNormRows` calls so the wide kernel is picked per call. Test verifies bit-equivalence against two sequential `rmsNormRows` calls on the same inputs. Migrated from PR #10.
TheTom
added a commit
that referenced
this pull request
May 26, 2026
Wraps two `kv_cache_update_*` dispatches in a single compute encoder. Cache layouts, dtypes, and writing thread are identical to `kvCacheUpdate`; this only saves the encoder begin/end pair between the K-cache write and the V-cache write. With 10 full attention layers on Qwen3.6-A3B, this drops 10 encoder pairs per decode token. Migrated from PR #10.
TheTom
added a commit
that referenced
this pull request
May 26, 2026
…n one encoder Used inside the Qwen3.5 GDN T-loop where, every token, one tensor goes through `siluCastToF32` (the gate input) while the two raw gate scalars `aRaw` / `bRaw` go through plain `castToF32`. The wrapper sets the silu PSO once, dispatches the silu, then switches to the plain-cast PSO and dispatches the two raw casts — all inside one compute encoder. Saves the encoder begin/end pairs the sequential 3-call path would create per GDN layer per token. Migrated from PR #10.
TheTom
added a commit
that referenced
this pull request
May 26, 2026
The Qwen3-MoE decode path historically had to commit + wait on a
single GPU layer just to read the router logits back to the host,
run a CPU top-K + softmax, and then dispatch the per-expert gate +
up + down GEMVs. Two new wrappers move both halves to the GPU:
- `moeRouterTopK` / `moeRouterTopKMany` — wrap the
`mt_moe_router_topk_*` kernel. Single-row form takes
`logits: [nExperts]` and writes `[k]` u32 indices plus `[k]`
weights; the T-batched form iterates `program_id<0>()` across
prefill rows in one dispatch.
- `dequantGemvInt4ExpertIndexed` /
`dequantGemvInt4ExpertIndexedMany` — wrap
`dequant_gemv_int4_expert_indexed_*`, which reads
`expertIndex[0]` at kernel entry to pick which slab of a
stacked `[nExperts, outDim, inDim/8]` u32 weight tensor to
dequantize. The `Many` variant runs N sequential dispatches
inside one compute encoder sharing the constexpr binding so
Qwen3.6-A3B's topK=8 gate + up phase collapses to one encoder
per MoE layer.
The router output indices buffer can be aliased into N single-u32
views (`buffer + offset = slot · 4`) and threaded straight into
the expert-indexed GEMV — no host readback per layer.
Tests:
- `dequantGemvInt4ExpertIndexedCorrectness` — stacks 4 experts'
weights, picks expert 2, and checks the output matches a
CPU-truth dequant + GEMV of that expert's slab.
- `moeRouterTopKDeterministic` — top-K of (1, 5, 3, 4) at k=2
picks indices {1, 3}; weights sum to 1.0 and the larger logit
dominates.
Migrated from PR #10 ITERs 54–60.
TheTom
added a commit
that referenced
this pull request
May 26, 2026
Eric's `sdpaMulti` on dev only switches on dtype — every call lands
on `ffai_sdpa_multi_*` (the head_dim=128-only kernel). Qwen3.6-A3B's
full-attention layers run at head_dim=256, and metaltile emits a
dedicated `ffai_sdpa_multi_d256_*` family with a 2-phase output
reduction that keeps the per-row threadgroup memory under Apple's
32 KB cap. Without the d256 routing the d128 kernel reads past
each row of K/V and produces silent garbage.
Switches `sdpaMulti` to a `(dtype, headDim)` tuple match so f32 /
f16 / bf16 each pick the correct kernel; relaxes
`OpsValidation.validateSdpaMulti` to accept `headDim ∈ {128, 256}`
and updates the matching validator test.
Adds a runtime test that mirrors the existing d=128 uniform-K case
at head_dim=256, exercising the new routing end-to-end.
Migrated from PR #10 ITER 93 (Qwen3.6-A3B prefill wiring needs it).
TheTom
added a commit
that referenced
this pull request
May 26, 2026
Eric's dev has the base versions (`castToF32`, `swiglu`, `gatedMixerNorm`) but not the batched variants our PR9/10 wirings rely on. Each variant reuses the same metaltile kernel and collapses N independent dispatches onto a single compute encoder: - `castToF32Two` / `castToF32Three` — Qwen3 GDN fused-prep path promotes 2-3 same-dtype tensors (convAct / aRaw / bRaw) to fp32 every token; one shared encoder replaces 2-3 separate ones. - `swigluMany` — MoE per-expert SwiGLU at decode runs topK independent dispatches per layer; one shared encoder per call. - `gatedMixerNormMany` — Qwen3.5 GDN mixer's per-token norm is evaluated `T·Hv` times during prefill; widening the grid X axis from `Hv` to `T·Hv` collapses the T-loop into a single dispatch. Tests in `Tests/FFAITests/Ops/OpsSpecialPathTests.swift` verify each batched form matches the corresponding sequential reference. Migrated from PR #10 (closes the Tier 2 paired-on-encoder gap identified during the Eric-overlap audit).
TheTom
added a commit
that referenced
this pull request
May 26, 2026
…ateKVMany, conv1dCausalStepSiluCastMany, gatedDeltaPrepChunk Eric's dev only has the single-token forms of these ops; the PR9/10 forwardMany / decodeMany prefill paths replace the per-token T-loop with a single batched dispatch. - `ropePartialMany` — single dispatch covering T rows of `qk` rotated by per-row `positions[T]`; saves T-1 encoder pairs and matches T sequential `ropePartial` calls bit-identically. - `ropePartialManyTwo` — Q + K share `positions` and constants on one encoder; row strides differ. - `kvCacheUpdateKVMany` — writes T tokens' K + V rows in one encoder, indexed by `positions[T]`. - `conv1dCausalStepSiluCastMany` — `T` depthwise conv + SiLU + fp32 cast collapsed into one dispatch with the conv state held in per-channel registers across the sweep. - `gatedDeltaPrepChunk` — chunked-prefill counterpart to `gatedDeltaPrepStep`; runs the fused GDN prep + recurrence over a token range in one dispatch, keeping per-head recurrent state in registers across the chunk. `tLen` arrives as a `[1]` u32 GPU buffer so chunked dispatch chains avoid CPU readback. All kernels (`ffai_rope_llama_many_*`, `kv_cache_update_many_*`, `ffai_conv1d_causal_step_silu_cast_many_*`, `mt_gated_delta_prep_chunk_*`) are already emitted on the dev metaltile. Tests: - `ropePartialManyMatchesSequential` — T=3 batched matches three sequential `ropePartial` outputs bit-identically. - `kvCacheUpdateKVManyF32` — verifies multi-position K + V writes land in the right cache rows and leave untouched slots zero. Migrated from PR #10 (ITERs 86–99).
TheTom
added a commit
that referenced
this pull request
May 26, 2026
…coder Mirror of `dequantGemvInt4Two/Three/Four` for the per-expert MoE down phase, where each chosen expert has its own activation tensor but all N projections share PSO, `inDim`, and `groupSize`. The wrapper sets the constexpr binding once and rotates `(weight, scales, biases, input, output)` per dispatch, saving N-1 encoder begin/end pairs versus N independent `dequantGemvInt4` calls. Test verifies the batched form matches three sequential `dequantGemvInt4` CPU-truth references across distinct (scale, bias) projections + distinct inputs. Migrated from PR #10.
TheTom
added a commit
that referenced
this pull request
May 26, 2026
…code Same dispatch contract as `sdpaMulti(causal: false)`, but the scalar `causal` flag is replaced by an additive `[nQuery, nQuery]` mask tensor consulted only for in-block KV positions. Positions `< baseKV` (the cached prefix) are always attended. Mask is `0.0` to allow and `-inf` to block, added directly to pre-softmax scores. Used by speculative-decode tree-verify: one verifier call attends every leaf-to-root path of the draft tree at once, gated by a tree-causal mask the drafter produces from its node structure. Wraps `ffai_sdpa_multi_tree_mask_*` (headDim=128 only — no d256 variant for tree-verify on dev). Tests: - All-zero mask reproduces `sdpaMulti(causal: false)` output bit-equivalent (sanity check on the mask add path). - Diagonal-only -inf mask isolates each query to its own KV slot: with uniform Q/K and per-row distinct V, output collapses to a copy of the matching V row. Migrated from PR #10 ITER 70.
TheTom
added a commit
that referenced
this pull request
May 26, 2026
Pass 1 partitions the KV cache into `blocks` even tiles and writes per-block partial `(O, max, lse)` triples; pass 2 merges them with the log-sum-exp trick. Total work matches single-pass `sdpaDecode` but the two passes parallelize across more GPU threadgroups, paying off at long KV (≥ 1K typical) where the single-pass kernel starves the GPU with one TG per Q-head. Block-count contract (the load-bearing detail PR10's port silently violated): the pass2 kernel hardcodes a 32-lane block reduction (`bn = 32`, `block_chunks = blocks / bn`), so `blocks` MUST be a multiple of 32 AND at least 32. Anything smaller makes `block_chunks` round to zero — pass2's outer loop iterates zero times and the output stays whatever pass1 left in scratch, typically all zeros. The previous port's "power of 2" check let `blocks=2` through and was the cause of the all-zero smoke failures during the earlier port attempt; the new precondition catches it at the wrapper. Scratch buffers (`partialO`, `partialM`, `partialL`) are caller-owned so a decode loop can amortise them across steps. Test: `blocks=32, nKV=64, nQHeads=2` — verifies the two-pass output matches single-pass `sdpaDecode` to within 1e-2 absolute. Migrated from PR #10 ITER 53.
TheTom
added a commit
that referenced
this pull request
May 26, 2026
…ding Speculative-decode candidate-proposer interface used by the in-progress `SpecDecode` driver. Three implementations land here: - `NGramDrafter` — prompt-lookup n-gram. Scans `history` for the last `nMatch` tokens, falls back to shorter n-grams down to `minNMatch`, returns up to `gamma` follow-on tokens from the earlier occurrence. Zero ML cost; the workhorse for code and structured-chat workloads where prior context recurs. - `NeverDrafter` — stub that always returns `[]`. Used to bench pure spec-decode driver overhead vs raw single-token decode. - `NGramTreeDrafter` — real branching tree drafter. Conforms to both `Drafter` (linear γ fallback via `propose`) and `TreeDrafter` (`proposeTree` returning top-K continuations per depth) for tree-verify spec decode against `Ops.sdpaMultiTreeMask`. `DraftTreeNode` is the pure-Swift tree representation, with `flatten()`, `verify(oracleAtHistoryEnd:oracle:)`, and `treeCausalMask()` helpers — all pure functions, no model / cache dependencies, ready to wire into the SpecDecode driver once `forwardManyAllLogits` lands. `LinearTreeAdapter` lets any linear `Drafter` masquerade as a `TreeDrafter` (degenerate single-branch tree) so the tree-verify driver can A/B against the linear baseline using the same `NGramDrafter` instance. Tests in `Tests/FFAITests/Generation/DrafterTests.swift` cover all three drafters end-to-end: 20 host-side tests, no GPU dependencies. All pass. Migrated from PR #10 (ITERs 61, 68, 69, 71).
TheTom
added a commit
that referenced
this pull request
May 26, 2026
…harness Same shape as Qwen3.6's forwardManyBench but pointed at the locally- available Qwen3.5-35B-A3B-4bit checkpoint. Engine is m.qwen35 in both cases, so this exercises every Bagel-migration wiring landed on tom/bagel-clean (GPU MoE router, batched QKV, GDN chunked recurrence, rmsNormQgemvInt4Fast finalNorm, 2-pass Flash-Decoding, MTLResidencySet weight pinning). Each test gates on the local checkpoint existing and prints a skip message when missing. M5 Max decode T=1 measurement: 96.87 tps (median of 5 × 32-step runs), matching the PR #10 Qwen3.6-A3B tip target of 96.3 tps.
TheTom
added a commit
that referenced
this pull request
May 26, 2026
…32 guard QuantizedLinear.callMany routes 4-bit weights to Ops.dequantGemmDynamicM (mt_qmm_mma backed), whose kernel requires outDim to be a multiple of 32 (the BN tile of the MMA kernel). Without an outDim guard the precondition tripped on Qwen3.5-35B-A3B's shared_expert_gate, which quantises to 4-bit at outDim=1 (a per-token gate scalar): 'Ops.dequantGemmDynamicM: nOut (1) must be multiple of 32 (BN tile)'. Add the multiple-of-32 check next to the 4-bit branch. Anything that doesn't fit (4-bit narrow projections + every non-4-bit) falls through to the per-row dequantGemv loop, bit-identical to the per-token path and irrelevant to perf at outDim=1. Verified on M5 Max with Qwen3.5-35B-A3B-4bit: forwardManyBench T=8 / T=128 / T=512 all complete cleanly. Headline perf: decode T=1 96.87 tps, prefill T=512 458.6 tps (4.75x speedup over per-token loop) — matches the PR #10 tip targets of 96.3 / 455 tps.
TheTom
added a commit
that referenced
this pull request
May 26, 2026
…e4/LFM2/Jamba dense MLP Two ITERs from PR #10 the wirings agent missed in the first pass. ITER 20: bump MetalTileLibrary.defaultMaxCommandBufferCount from 16 to 64. Lets the Metal driver pipeline more cmd-buffers in flight before applying backpressure; absorbs per-cmd-buffer encode latency across more concurrent submissions. PR #10 measured +2-3% prefill T=512 on Qwen3.6-A3B / M5 Max. FFAI_MAX_COMMAND_BUFFERS override preserved. ITER 17: replace Ops.mul(Ops.silu(g, cmd), u, cmd) with the fused Ops.swiglu in Granite4Text.MLP.forward, LFM2Text.MLP.forward, and JambaText.MLP.forward. Same semantic; one kernel dispatch instead of two, intermediate silu(g) stays in registers instead of round- tripping to DRAM. Brings these three cross-family models in line with Qwen3.5 (already wired via the agent's earlier ITER 16 commit on Qwen35DenseMLP + MoEFFN shared expert).
TheTom
added a commit
that referenced
this pull request
May 26, 2026
… bench cases Extends the bench harness with two long-context shapes that PR #10's work needs to validate but Eric's existing harness only covered up to T=512: - forwardManyT2K (T=2048): exercises chunked GDN recurrence over a longer window + larger SDPA prefill grid. - decodeAfterLongPrefill: prefill 1024 tokens via forwardMany, then measure decode-only tps over 32 steps. Catches any KV-cache / residency / sdpaDecode2Pass routing regression at non-trivial KV. Measured on M5 Max, Qwen3.5-35B-A3B-4bit: forwardManyBench T=2048: per_token=21499ms batched=4755ms speedup=4.52x batched_tps=430.7 decode T=1 after T=1024 prefill: runs=[0.335, 0.335, 0.336, 0.336, 0.337]s median=0.336s -> 95.37 tps Decode tps barely drops from empty-cache (96.51) to KV=1024 (95.37) - the kernel-compute floor PR #10 documented at 96 tps holds through the long-context regime.
Collaborator
Collaborator
|
Superseded by #14. |
ekryski
added a commit
that referenced
this pull request
May 27, 2026
PR #10's regression was the wrapper accepting blocks=2 (kernel pass-2 hardcodes a 32-lane block-merge reduction, so anything below 32 silently produces zero output). PR #14 closed it with a `blocks >= 32 && blocks % 32 == 0` precondition. The other half of the contract — that at the minimum legal blocks=32, two-pass output matches the canonical single-pass kernel within tolerance — had no host-side test. Add one. metaltile's per- kernel GPU correctness test would not have caught the dispatch- routing bug; the wrapper is what needs covering.
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
2 participants
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.
Summary
Continuation of the autonomous kernel-opt loop from PR #9 (merged at
98f82f2). Same playbook: iterate, bench, ship.Starting baseline (post-PR-9 merge): 55.17 tps decode median / 55.94 best on Qwen3.6-A3B M5 Max.
ITER 44 — cache xNorm + ffnNorm in Qwen35GDNLayer + Qwen35AttentionLayer
Per-call rmsNorm fresh allocations replaced with instance-cached
xNormScratch+ffnNormScratchon every decoder layer. 80 fresh allocations/token saved (2 per layer × 40 layers).Decode 55.17 → 56.09 tps median (best 56.28) = +1.7%.
3-run bench:
Cumulative session vs morning 16.29 tps: 56.09 = 3.44× median / 3.46× best.
forwardManyEquivalenceT=8 + T=128 argmax bit-identical (52290 / 760).Test plan
🤖 Generated with Claude Code