perf(moe): wire the two-phase parallel tile-plan builders#67
Merged
Conversation
Ops.moeBuildTilePlanParallel / Bm32Parallel / Bm32OwnParallel dispatch ffai_moe_tile_plan_expert_counts (phase 1, one threadgroup per expert) chained into the matching phase-2 emission kernel, in place of the single-threadgroup ffai_moe_build_tile_plan / _bm32 / _bm32_own. Same output format/contract as the originals (byte-exact per the kernels-repo dual-oracle correctness suite). Wired into all three useTilePlan call sites in MoELayer.decodeMany: the plain bm16 path, the default useCoopTilePlan path, and the paired lever's "own" dispatch. Gate: FFAI_MOE_PARALLEL_PLAN, default ON, escape hatch =0 reverts to the single-threadgroup builders. e2e prefill ladder on Qwen3.6-35B-A3B-4bit (TilePlanParallelGateTests, one (T, config) pair per process - a within-process sweep across T briefly held two configs' KV caches/activations resident and pushed this shared M5 Max under 2GB free, so each point here is an isolated process run): T=512 (7 trials): parallel ~3115-3195 tps, single-threadgroup 3206 tps T=1024 (7 trials): parallel 3689 tps, single-threadgroup 3695 tps T=2048 (3 trials): parallel 3878 tps, single-threadgroup 3883 tps T=4096 (3 trials): parallel 3708 tps, single-threadgroup 3727 tps T=16384(3 trials): parallel 2267 tps, single-threadgroup 2019 tps (+12.3%) T=32768(1 trial, reduced for memory safety): parallel 1449 tps, single-threadgroup 1501 tps T=512/1024/2048/4096 are a wash within run-to-run noise (repeat measurement at T=512 spanned 3116-3195 tps on the identical config); T=16384 shows a real win; T=32768's single-trial gap is within the noise band seen everywhere else but wasn't repeated to confirm, given the memory risk of repeated large-T runs on this shared box. All six points clear their prior no-regression baselines. --dump-prefill-logits is byte-identical between FFAI_MOE_PARALLEL_PLAN default and =0 at T=512/1024/4096/16384, on a deterministic repeat, and on an adversarial multi-topic multi-chunk prompt - pure plumbing, same routing plan in means same GEMM output out.
TheTom
added a commit
that referenced
this pull request
Jul 22, 2026
…ctions (#68) * docs(campaign): fill in missing F-85 env-var entries, fix polarity, drop em dashes The env-var reference block (added #60) documents every prefill-hot-path toggle in one place; five vars added by later commits in the campaign (#61/#64/#65/#66/#67) had no entry: FFAI_MOE_COOP, FFAI_MOE_COOP_INDIRECT, FFAI_MOE_PAIRED_INDIRECT, FFAI_MOE_COUNTING_SORT, FFAI_MOE_PARALLEL_PLAN. FFAI_MOE_PAIRED's existing entry also had the polarity backwards: it read "=0 opts out, default ON" but MoELayer.pairedEnabled actually returns true only when the raw value is exactly "1" (default OFF, opt-in). Fixed to match the code and MoELayer.swift's own (correct) doc comment. Added a caution note on FFAI_QMM_COOP=1: forcing the coop kernel on below the MPP/NAX capability floor is untested territory, not a bypass of the requirement. Also fixes 4 em-dash violations in a relocated `gatedDeltaChunk` doc comment (moved, not reworded, when the new tile-plan section was inserted above it, but the moved text still counts as added-since-#52 for the campaign's own no-em-dash rule). Pure documentation and comment changes, no code behavior touched. Green release build. * perf(moe): cache the default tile-plan buffers per instance instead of per call MoELayer.decodeMany allocated three fresh MTLBuffers (tile_expert / tile_row_start / tile_row_count) via device.makeBuffer every layer, every forward call, on the useTilePlan/useCoopTilePlan path - the DEFAULT prefill path since tilePlanEnabled and coopEnabled both default ON. The paired lever's own six buffers already got the reuse treatment in the same round that shipped useCoopPaired; these three, used by every prefill call regardless of the paired flag, were missed. Cached grow-only per instance (tilePlanBuffers(capacity:device:)), mirroring the existing pairedTilePlanBuffers pattern in this same file and Ops.identityRowsBuffer's grow-only reuse elsewhere. Capacity varies with mTotal (per-chunk row count), unlike the paired buffers' fixed per-instance capacity, so growth uses a max-so-far buffer with a per-call capacity-sized view rather than exact-match reallocation; the zero-fill every call still only touches the current call's requested capacity, preserving the same stale-tail-data safety argument the paired buffers already document. Quality gate: --dump-prefill-logits byte-identical vs unmodified dev at T=512 and T=4096 (FFAI_PREFILL_STEP pinned). Interleaved e2e sanity (2 order-balanced pairs each at T=512/T=4096) shows no regression beyond run-to-run noise. * perf(qwen3x): cache per-instance env-var tiers in the attention/GDN mixers Qwen35AttentionMixer.forwardMany re-derived the d=256 SDPA kernel tier (FFAI_SDPA_QTILED_D256 / FFAI_SDPA_LEGACY_D256) from ProcessInfo.environment on every call, and Qwen35GDNMixer.forwardManyChunked did the same for the chunked-WY eligibility (FFAI_GDN_WY) and the fast prep_chunk kernel eligibility (FFAI_GDN_NO_PREP_CHUNK_FAST) - once per layer per prefill chunk, and once per layer per decode token on the batched decode path. Both mixers already had a precedent for caching env-derived, instance- constant values at init instead (Qwen35AttentionMixer.sdpa2PassThreshold, Qwen35GDNMixer.fused) that this campaign's own new gates didn't follow. All four gates depend only on values fixed for the mixer's lifetime (headDim / keyHeadDim / valueHeadDim / numKeyHeads / numValueHeads and the env), so they're now computed once in each init and stored as instance lets (useMmaD256Cached/useQtiledD256Cached, wyShapeEligible/fastPrepChunkEligible). The per-call `t >= 64` / `!useWY` halves that genuinely vary per call stay as cheap Int/Bool comparisons at the call site. Quality gate: --dump-prefill-logits byte-identical vs unmodified dev at T=512 and T=4096 (FFAI_PREFILL_STEP pinned). Interleaved e2e sanity (2 order-balanced pairs each at T=512/T=4096) shows no regression beyond run-to-run noise. * perf(ops): cache FFAI_SDPA_BK16 / FFAI_QMM_MPP / FFAI_QMM_COOP as static lets sdpaPrefillMmaD256 (the default d=256 prefill attention kernel) and dispatchQmmMma (the shared dispatch entry for every dense int4 projection at prefill - GDN in-projections, attention QKV/out, every other quantized linear) both re-read ProcessInfo.environment on every call to resolve their kernel tier. dispatchQmmMma in particular fires many times per layer per forward call, so this was the hottest of the uncached env reads found in this campaign's additions. Cached as static lets (sdpaBK16Enabled, qmmMppEnvRaw, qmmCoopEnvRaw), mirroring the existing Ops.mppAutoCapable pattern already used for the hardware-capability check right next to these gates. FFAI_QMM_MPP's own read was pre-existing (predates this campaign) but shared the same `let env = ProcessInfo.processInfo.environment` line the campaign's new FFAI_QMM_COOP read was added next to, so fixing both together avoided leaving a still-uncached sibling read one line away from a cached one. Quality gate: --dump-prefill-logits byte-identical vs unmodified dev at T=512 and T=4096 (FFAI_PREFILL_STEP pinned). Interleaved e2e sanity (2 order-balanced pairs each at T=512/T=4096) shows no regression beyond run-to-run noise. * fix(moe): gate the coop-core MoE gather GEMM on MPP/NAX hardware capability ffai_moe_gather_qmm_coop (the F-85 coop-core MoE gather GEMM, default ON via MoELayer.coopEnabled) wraps the same matmul2d cooperative-tensor core as the dense ffai_qmm_coop kernel, which requires Apple GPU Family 9+ / macOS 26+. dispatchQmmMma's dense coop tier already auto-detects this via Ops.mppAutoCapable and falls back to a cheaper tier when incapable; coopEnabled had no equivalent check at all - it was unconditionally `coopEnvRaw != "0"`, meaning the DEFAULT prefill path for any Qwen3.5/3.6 MoE checkpoint would unconditionally dispatch this cooperative-tensor kernel even on hardware that cannot run it, with no lower-level fallback inside the Ops wrapper to catch it (unlike the dense path). Fixed by gating the auto-default (unset / not "0"/"1") on Ops.mppAutoCapable, same as dispatchQmmMma's coop/mpp tiers. Explicit FFAI_MOE_COOP=1/=0 still force the gate either way for A/B measurement, unchanged from before. Ops.mppAutoCapable bumped from `private` to module-internal visibility so MoELayer can reuse the single cached check instead of duplicating the OS-version + GPU-family query. No effect on Apple Family 9+ / macOS 26+ hardware (the auto-default already resolved to true there) - --dump-prefill-logits is byte- identical vs unmodified dev at T=512 and T=4096 (FFAI_PREFILL_STEP pinned), and an interleaved e2e sanity (3 order-balanced pairs each at T=512/T=4096, covering this commit plus every other change in this review pass) shows no regression beyond run-to-run noise. Effect is confined to hardware below the capability floor, which this box cannot exercise - flagged in the review report as untestable here rather than assumed safe.
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.
Pairs with kernels PR #66 (merged to kernels dev). Routes the three tile-plan build call sites in MoELayer through the two-dispatch parallel builders (shared per-expert counts pass + prefix/emission pass). FFAI_MOE_PARALLEL_PLAN default ON, =0 reverts to the single-threadgroup originals.