Skip to content

dsv41 metal: the decode MoE in four waves, same bits - #3292

Open
ashhart wants to merge 1 commit into
mainfrom
lead/dsv41-decode-waves
Open

ashhart wants to merge 1 commit into
mainfrom
lead/dsv41-decode-waves

Conversation

@ashhart

@ashhart ashhart commented Nov 14, 2026

Copy link
Copy Markdown
Owner

DeepSeek-V4.1's decode MoE on Metal ran six barrier-separated waves after its norm, and the router's matvec, the route and the expert plan were three small dependent launches. This puts those three in one launch and moves it beside work it doesn't depend on, so the MoE takes four waves. Every output and KV byte is unchanged.

  • tf_dsv41_route_gate_plan.
    • Groups of 32 columns compute the gate logits with mv.metal's lane chains and butterfly, rows in up to four slices.
    • The last group to finish routes and plans. It's found through a sync word in the scratch, which the last group resets to zero for the next launch.
    • Writers store their logits atomically and fence at device scope before counting in. The last group fences, reads the logits back as atomic loads, ranks each row's picks in the route's order (each simdgroup's top-k, then their union's) and plans them. Nothing spins.
  • Waves. That launch runs beside the shared expert's w1/w3 and the experts' act_quant. Gate/up runs beside the SwiGLU and the hc mixes, and down beside the shared w2, then the join. A plain layer goes from 17 waves and 27 dispatches to 15 and 25.
  • One-row experts. upc[0] and down[0] decode each 32-bit word of E2M1 codes with shifts instead of integer multiplies, and stage x a 32-value group on a stride of 40. Every chain keeps its order.

Measured on an M5 Ultra (512 GB), release layers, time-dsv41-metal-real, fastest of 41, base and branch alternated twice:

Layer 1 row (x40 ms) 4 rows 6 rows
2 25.04 → 24.19 41.66 → 40.64 54.15 → 53.36
21 21.59 → 20.73 39.31 → 38.44 48.32 → 47.84
5 21.57 → 20.64 37.30 → 36.10 48.24 → 47.65
  • Bits: outputs and KV state hash as main on layers 2, 21 and 5 at 1 and 4 rows, with 20 of 20 concurrent launches equal. The float64 real-layer checks pass on 2, 21 and 5.

  • Kernel contract, --only dsv41: 91 kernels pass gates 1-4, 194 controls red, 0 limit problems. The new launch registers in the matvec suite (exact bytes at 384 experts, 1-16 rows, four controls under TF_GCTL) and in the route suite on given logits.

  • zig build test-dsv41-metal -Dcpu=apple_m1 and test-dsv41-metal-serve pass on an M5 Max. Lean: 0 new.

  • With dsv41 metal: decode experts decode an expert's weights once for all of its rows (stacked on #1720, #2973) #2983: this conflicts textually in experts.metal (around down_one), the experts records and docs/kernel-contract.md. A version rebased on dsv41 metal: decode experts decode an expert's weights once for all of its rows (stacked on #1720, #2973) #2983's head exists and hashes the same on layers 2, 21 and 5; I'll push whichever is needed once one of the two lands.

  • Tried and dropped: matvec with eight chains a thread, hoisted loads in combpost/prenorm/postnorm, a separate route+plan kernel (level), gate and up split across simdgroups (60-105% slower).

  • Synthena (Ash's Agent)

…ord of codes at a time

tf_dsv41_route_gate_plan runs the router's matvec, the route and the experts' plan as one launch:
groups of 32 columns compute the gate logits (mv.metal's lane chains and butterfly, rows in up to
four slices), and the last group to finish (its count in the scratch's sync word, zero between
launches) reads the logits as atomics, ranks every row's picks in the route's order (each
simdgroup's top-k, then their union's) and plans them. It runs beside the shared expert's w1/w3,
so the MoE takes four waves after the norm instead of six: gate+route+plan beside the shared up
and act_quant; gate/up beside the SwiGLU and the hc mixes; down beside the shared w2; the join.

The one-row gate/up and down decode each word of E2M1 codes by shifts (no integer multiplies)
and stage x a 32-value group on a stride of 40; every chain keeps its order.

Every output and KV byte is unchanged: layers 2, 21 and 5 at 1 and 4 rows hash as main. The new
launch is in the contract twice (matvec and route suites, controls under TF_GCTL and TF_CTL).

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant