Skip to content

perf(moe): wire the two-phase parallel tile-plan builders #67

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
TheTom merged 1 commit into from
Jul 22, 2026
Merged

perf(moe): wire the two-phase parallel tile-plan builders #67

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
93 changes: 79 additions & 14 deletions Sources/FFAI/Models/MoELayer.swift
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -491,6 +491,10 @@ public final class MoELayer: Module, DecoderLayer {
/// gate). Raw env string cached at init; `sortPlanCountingEnabled`
/// below is the single place that interprets it.
let sortPlanCountingEnvRaw: String?
/// Parallelized tile-plan builder gate (F-85 small-M follow-up). Raw
/// env string cached at init; `parallelPlanEnabled` below is the
/// single place that interprets it.
let parallelPlanEnvRaw: String?

/// Lazy dequant of the gate weight + cached output. Skipped when
/// `gate.inner` is not a 8-bit QuantizedLinear or env
Expand Down Expand Up @@ -676,6 +680,7 @@ public final class MoELayer: Module, DecoderLayer {
self.pairedIndirectEnvRaw = env["FFAI_MOE_PAIRED_INDIRECT"]
self.coopIndirectEnvRaw = env["FFAI_MOE_COOP_INDIRECT"]
self.sortPlanCountingEnvRaw = env["FFAI_MOE_COUNTING_SORT"]
self.parallelPlanEnvRaw = env["FFAI_MOE_PARALLEL_PLAN"]
}

/// F-85 expert-aligned tile-plan gate. Default ON as of the wiring
Expand Down Expand Up @@ -712,6 +717,29 @@ public final class MoELayer: Module, DecoderLayer {
sortPlanCountingEnvRaw != "0"
}

/// F-85 parallelized tile-plan builder gate (small-M follow-up to the
/// idle-tile-cost/counting-sort round). `ffai_moe_build_tile_plan` /
/// `_bm32` / `_bm32_own` each dispatch ONE threadgroup for the whole
/// plan build - every expert's binary search, the prefix sum, and
/// every tile emission all run on a single GPU core, a per-layer fixed
/// cost that does not shrink with `mTotal` and was found
/// disproportionately expensive at low prefill `T`. The parallel
/// replacements (`Ops.moeBuildTilePlanParallel`/`Bm32Parallel`/
/// `Bm32OwnParallel`, see `moe_tile_plan_expert_counts.rs`'s module
/// doc) split the per-expert row-range search out to `nExperts`
/// threadgroups (one per expert, running concurrently across the
/// GPU's cores) followed by a small single-threadgroup prefix-sum +
/// emission pass - same output format/contract, byte-exact against the
/// original builders across the full fixture set (dual CPU-oracle +
/// live-original-kernel checks). Default ON. Escape hatch:
/// `FFAI_MOE_PARALLEL_PLAN=0` reverts to the single-threadgroup
/// builders. Applies wherever `useTilePlan` holds - the plain bm16
/// path, the default `useCoopTilePlan` path, and the paired lever's
/// "own" dispatch (see `decodeMany`).
public var parallelPlanEnabled: Bool {
parallelPlanEnvRaw != "0"
}

/// F-85 coop-core MoE gather GEMM gate (BM=32 tile plan over
/// `ffai_moe_gather_qmm_coop`, the gather counterpart of the dense
/// `ffai_qmm_coop` core). Default ON as of the wiring pass's quality +
Expand Down Expand Up @@ -1461,11 +1489,23 @@ public final class MoELayer: Module, DecoderLayer {
buffer: ownIndirectGateUpBuf, offset: 4, shape: [1], dtype: .u32)
let ownTileCountDown = Tensor(
buffer: ownIndirectDownBuf, offset: 4, shape: [1], dtype: .u32)
Ops.moeBuildTilePlanBm32Own(
sortedExperts: indices,
tileExpert: tileExpert, tileRowStart: tileRowStart, tileRowCount: tileRowCount,
tileCountGateUp: ownTileCountGateUp, tileCountDown: ownTileCountDown,
mTotal: mTotal, nExperts: nExperts, on: work)
// F-85 parallelized builder (see `parallelPlanEnabled`) -
// default ON, escape hatch FFAI_MOE_PARALLEL_PLAN=0.
if self.parallelPlanEnabled {
Ops.moeBuildTilePlanBm32OwnParallel(
sortedExperts: indices,
tileExpert: tileExpert, tileRowStart: tileRowStart,
tileRowCount: tileRowCount,
tileCountGateUp: ownTileCountGateUp, tileCountDown: ownTileCountDown,
mTotal: mTotal, nExperts: nExperts, on: work)
} else {
Ops.moeBuildTilePlanBm32Own(
sortedExperts: indices,
tileExpert: tileExpert, tileRowStart: tileRowStart,
tileRowCount: tileRowCount,
tileCountGateUp: ownTileCountGateUp, tileCountDown: ownTileCountDown,
mTotal: mTotal, nExperts: nExperts, on: work)
}
coopIndirectArgsGateUp = ownIndirectGateUpBuf
coopIndirectArgsDown = ownIndirectDownBuf

Expand Down Expand Up @@ -1527,18 +1567,43 @@ public final class MoELayer: Module, DecoderLayer {
buffer: indirectGateUpBuf, offset: 4, shape: [1], dtype: .u32)
let tileCountDown = Tensor(
buffer: indirectDownBuf, offset: 4, shape: [1], dtype: .u32)
Ops.moeBuildTilePlanBm32(
sortedExperts: indices,
tileExpert: tileExpert, tileRowStart: tileRowStart, tileRowCount: tileRowCount,
tileCountGateUp: tileCountGateUp, tileCountDown: tileCountDown,
mTotal: mTotal, nExperts: nExperts, on: work)
// F-85 parallelized builder (see `parallelPlanEnabled`) -
// default ON, escape hatch FFAI_MOE_PARALLEL_PLAN=0. This is
// the default production coop path's builder call.
if self.parallelPlanEnabled {
Ops.moeBuildTilePlanBm32Parallel(
sortedExperts: indices,
tileExpert: tileExpert, tileRowStart: tileRowStart,
tileRowCount: tileRowCount,
tileCountGateUp: tileCountGateUp, tileCountDown: tileCountDown,
mTotal: mTotal, nExperts: nExperts, on: work)
} else {
Ops.moeBuildTilePlanBm32(
sortedExperts: indices,
tileExpert: tileExpert, tileRowStart: tileRowStart,
tileRowCount: tileRowCount,
tileCountGateUp: tileCountGateUp, tileCountDown: tileCountDown,
mTotal: mTotal, nExperts: nExperts, on: work)
}
coopIndirectArgsGateUp = indirectGateUpBuf
coopIndirectArgsDown = indirectDownBuf
} else {
Ops.moeBuildTilePlan(
sortedExperts: indices,
tileExpert: tileExpert, tileRowStart: tileRowStart, tileRowCount: tileRowCount,
mTotal: mTotal, nExperts: nExperts, on: work)
// F-85 parallelized builder (see `parallelPlanEnabled`) -
// bm16 fallback path (coop core off or shape doesn't
// qualify - see `useCoopTilePlan`).
if self.parallelPlanEnabled {
Ops.moeBuildTilePlanParallel(
sortedExperts: indices,
tileExpert: tileExpert, tileRowStart: tileRowStart,
tileRowCount: tileRowCount,
mTotal: mTotal, nExperts: nExperts, on: work)
} else {
Ops.moeBuildTilePlan(
sortedExperts: indices,
tileExpert: tileExpert, tileRowStart: tileRowStart,
tileRowCount: tileRowCount,
mTotal: mTotal, nExperts: nExperts, on: work)
}
}
// xGathered is already pre-gathered (step 8 above) - identity
// row map, same helper the bm16 MPP path uses.
Expand Down
178 changes: 178 additions & 0 deletions Sources/FFAI/Ops/Ops.swift
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -6457,6 +6457,184 @@ public enum Ops {
gridSize: grid, threadgroupSize: tg, on: cmd)
}

/// Phase 1 of the parallelized tile-plan builders (F-85 small-M
/// follow-up, `ffai_moe_tile_plan_expert_counts`). Dispatches ONE
/// THREADGROUP PER EXPERT (`nExperts` groups, one thread each) so every
/// expert's binary search over `sortedExperts` runs concurrently across
/// the GPU's cores instead of serially on the single core the original
/// builders' one-threadgroup dispatch confines all `nExperts` lanes to.
/// Writes `expertRowBase[e]`/`expertCount[e]`, the same row-range values
/// any of the three single-threadgroup builders would compute locally
/// for expert `e` - shared by all three tile-height variants below,
/// since the row range a sorted-expert array implies for an expert does
/// not depend on the tile height it will be sliced into.
///
/// Callers (`moeBuildTilePlanParallel`/`Bm32Parallel`/`Bm32OwnParallel`
/// below) allocate `expertRowBase`/`expertCount` as fresh per-call
/// scratch, matching `moeSortPlanCounting`'s `blockCounts`/
/// `blockOffsets` convention - not cached, since `nExperts` is fixed
/// per model but this keeps every caller's allocation lifetime obvious.
public static func moeTilePlanExpertCounts(
sortedExperts: Tensor,
expertRowBase: Tensor, expertCount: Tensor,
mTotal: Int, nExperts: Int,
on cmd: MTLCommandBuffer
) {
precondition(
sortedExperts.dtype == .u32, "moeTilePlanExpertCounts: sortedExperts must be u32")
precondition(
expertRowBase.dtype == .u32 && expertCount.dtype == .u32,
"moeTilePlanExpertCounts: expertRowBase/expertCount must be u32")
precondition(
nExperts <= 256,
"moeTilePlanExpertCounts: nExperts (\(nExperts)) exceeds the kernel's 256-expert scope"
)
let grid = MTLSize(width: nExperts, height: 1, depth: 1)
let tg = MTLSize(width: 1, height: 1, depth: 1)
MetalTileKernels.ffai_moe_tile_plan_expert_counts_f32_threadgroups(
sorted_experts: sortedExperts.buffer, sorted_expertsOffset: sortedExperts.offset,
expert_row_base: expertRowBase.buffer, expert_row_baseOffset: expertRowBase.offset,
expert_count: expertCount.buffer, expert_countOffset: expertCount.offset,
m_total: UInt32(mTotal),
gridSize: grid, threadgroupSize: tg, on: cmd)
}

/// Parallelized sibling of `moeBuildTilePlan` (F-85 small-M follow-up):
/// dispatches `moeTilePlanExpertCounts` (phase 1, per-expert parallel)
/// then `ffai_moe_build_tile_plan_parallel` (phase 2 - prefix sum +
/// tile emission only, no per-lane binary search left). Identical
/// output format/contract to `moeBuildTilePlan` - same
/// `tileExpert`/`tileRowStart`/`tileRowCount`, same worst-case capacity
/// and zero-fill padding convention. Gated by `FFAI_MOE_PARALLEL_PLAN`
/// in `MoELayer.swift`; `moeBuildTilePlan` remains the escape hatch.
public static func moeBuildTilePlanParallel(
sortedExperts: Tensor,
tileExpert: Tensor, tileRowStart: Tensor, tileRowCount: Tensor,
mTotal: Int, nExperts: Int,
on cmd: MTLCommandBuffer
) {
precondition(
sortedExperts.dtype == .u32, "moeBuildTilePlanParallel: sortedExperts must be u32")
precondition(
tileExpert.dtype == .u32 && tileRowStart.dtype == .u32 && tileRowCount.dtype == .u32,
"moeBuildTilePlanParallel: tile_expert/tile_row_start/tile_row_count must be u32")
precondition(
nExperts <= 256,
"moeBuildTilePlanParallel: nExperts (\(nExperts)) exceeds the kernel's 256-lane threadgroup scratch capacity"
)
let rowBaseBuf = device.makeBuffer(length: nExperts * 4)
let countBuf = device.makeBuffer(length: nExperts * 4)
let expertRowBase = Tensor(buffer: rowBaseBuf, offset: 0, shape: [nExperts], dtype: .u32)
let expertCount = Tensor(buffer: countBuf, offset: 0, shape: [nExperts], dtype: .u32)
moeTilePlanExpertCounts(
sortedExperts: sortedExperts, expertRowBase: expertRowBase, expertCount: expertCount,
mTotal: mTotal, nExperts: nExperts, on: cmd)

let grid = MTLSize(width: 1, height: 1, depth: 1)
let tg = MTLSize(width: nExperts, height: 1, depth: 1)
MetalTileKernels.ffai_moe_build_tile_plan_parallel_f32_threadgroups(
expert_row_base: expertRowBase.buffer, expert_row_baseOffset: expertRowBase.offset,
expert_count: expertCount.buffer, expert_countOffset: expertCount.offset,
tile_expert: tileExpert.buffer, tile_expertOffset: tileExpert.offset,
tile_row_start: tileRowStart.buffer, tile_row_startOffset: tileRowStart.offset,
tile_row_count: tileRowCount.buffer, tile_row_countOffset: tileRowCount.offset,
n_experts: UInt32(nExperts),
gridSize: grid, threadgroupSize: tg, on: cmd)
}

/// Parallelized sibling of `moeBuildTilePlanBm32` (F-85 small-M
/// follow-up) - same two-phase split as `moeBuildTilePlanParallel`,
/// BM=32 tile height, including the `tileCountGateUp`/`tileCountDown`
/// indirect-dispatch outputs. Gated by `FFAI_MOE_PARALLEL_PLAN`;
/// `moeBuildTilePlanBm32` remains the escape hatch.
public static func moeBuildTilePlanBm32Parallel(
sortedExperts: Tensor,
tileExpert: Tensor, tileRowStart: Tensor, tileRowCount: Tensor,
tileCountGateUp: Tensor, tileCountDown: Tensor,
mTotal: Int, nExperts: Int,
on cmd: MTLCommandBuffer
) {
precondition(
sortedExperts.dtype == .u32, "moeBuildTilePlanBm32Parallel: sortedExperts must be u32")
precondition(
tileExpert.dtype == .u32 && tileRowStart.dtype == .u32 && tileRowCount.dtype == .u32,
"moeBuildTilePlanBm32Parallel: tile_expert/tile_row_start/tile_row_count must be u32")
precondition(
tileCountGateUp.dtype == .u32 && tileCountDown.dtype == .u32,
"moeBuildTilePlanBm32Parallel: tileCountGateUp/tileCountDown must be u32")
precondition(
nExperts <= 256,
"moeBuildTilePlanBm32Parallel: nExperts (\(nExperts)) exceeds the kernel's 256-lane threadgroup scratch capacity"
)
let rowBaseBuf = device.makeBuffer(length: nExperts * 4)
let countBuf = device.makeBuffer(length: nExperts * 4)
let expertRowBase = Tensor(buffer: rowBaseBuf, offset: 0, shape: [nExperts], dtype: .u32)
let expertCount = Tensor(buffer: countBuf, offset: 0, shape: [nExperts], dtype: .u32)
moeTilePlanExpertCounts(
sortedExperts: sortedExperts, expertRowBase: expertRowBase, expertCount: expertCount,
mTotal: mTotal, nExperts: nExperts, on: cmd)

let grid = MTLSize(width: 1, height: 1, depth: 1)
let tg = MTLSize(width: nExperts, height: 1, depth: 1)
MetalTileKernels.ffai_moe_build_tile_plan_bm32_parallel_f32_threadgroups(
expert_row_base: expertRowBase.buffer, expert_row_baseOffset: expertRowBase.offset,
expert_count: expertCount.buffer, expert_countOffset: expertCount.offset,
tile_expert: tileExpert.buffer, tile_expertOffset: tileExpert.offset,
tile_row_start: tileRowStart.buffer, tile_row_startOffset: tileRowStart.offset,
tile_row_count: tileRowCount.buffer, tile_row_countOffset: tileRowCount.offset,
tile_count_gateup: tileCountGateUp.buffer, tile_count_gateupOffset: tileCountGateUp.offset,
tile_count_down: tileCountDown.buffer, tile_count_downOffset: tileCountDown.offset,
n_experts: UInt32(nExperts),
gridSize: grid, threadgroupSize: tg, on: cmd)
}

/// Parallelized sibling of `moeBuildTilePlanBm32Own` (F-85 small-M
/// follow-up) - same two-phase split, excludes `1..=16`-row pairing
/// candidates like the original "own" builder. Gated by
/// `FFAI_MOE_PARALLEL_PLAN`; `moeBuildTilePlanBm32Own` remains the
/// escape hatch.
public static func moeBuildTilePlanBm32OwnParallel(
sortedExperts: Tensor,
tileExpert: Tensor, tileRowStart: Tensor, tileRowCount: Tensor,
tileCountGateUp: Tensor, tileCountDown: Tensor,
mTotal: Int, nExperts: Int,
on cmd: MTLCommandBuffer
) {
precondition(
sortedExperts.dtype == .u32, "moeBuildTilePlanBm32OwnParallel: sortedExperts must be u32"
)
precondition(
tileExpert.dtype == .u32 && tileRowStart.dtype == .u32 && tileRowCount.dtype == .u32,
"moeBuildTilePlanBm32OwnParallel: tile_expert/tile_row_start/tile_row_count must be u32")
precondition(
tileCountGateUp.dtype == .u32 && tileCountDown.dtype == .u32,
"moeBuildTilePlanBm32OwnParallel: tileCountGateUp/tileCountDown must be u32")
precondition(
nExperts <= 256,
"moeBuildTilePlanBm32OwnParallel: nExperts (\(nExperts)) exceeds the kernel's 256-lane threadgroup scratch capacity"
)
let rowBaseBuf = device.makeBuffer(length: nExperts * 4)
let countBuf = device.makeBuffer(length: nExperts * 4)
let expertRowBase = Tensor(buffer: rowBaseBuf, offset: 0, shape: [nExperts], dtype: .u32)
let expertCount = Tensor(buffer: countBuf, offset: 0, shape: [nExperts], dtype: .u32)
moeTilePlanExpertCounts(
sortedExperts: sortedExperts, expertRowBase: expertRowBase, expertCount: expertCount,
mTotal: mTotal, nExperts: nExperts, on: cmd)

let grid = MTLSize(width: 1, height: 1, depth: 1)
let tg = MTLSize(width: nExperts, height: 1, depth: 1)
MetalTileKernels.ffai_moe_build_tile_plan_bm32_own_parallel_f32_threadgroups(
expert_row_base: expertRowBase.buffer, expert_row_baseOffset: expertRowBase.offset,
expert_count: expertCount.buffer, expert_countOffset: expertCount.offset,
tile_expert: tileExpert.buffer, tile_expertOffset: tileExpert.offset,
tile_row_start: tileRowStart.buffer, tile_row_startOffset: tileRowStart.offset,
tile_row_count: tileRowCount.buffer, tile_row_countOffset: tileRowCount.offset,
tile_count_gateup: tileCountGateUp.buffer, tile_count_gateupOffset: tileCountGateUp.offset,
tile_count_down: tileCountDown.buffer, tile_count_downOffset: tileCountDown.offset,
n_experts: UInt32(nExperts),
gridSize: grid, threadgroupSize: tg, on: cmd)
}

/// Device-side PAIR-ONLY BM=32 tile-plan builder for the F-85
/// tile-pairing lever (`ffai_moe_build_tile_plan_bm32_paired`). Reads
/// the SAME `sortedExperts` array as `moeBuildTilePlanBm32Own` and
Expand Down
Loading
Loading