Skip to content

dsv41 metal: FP8 decode rows load first and meet tile roots in one simdgroup, same bits - #3287

Open
ashhart wants to merge 1 commit into
mainfrom
lead/dsv41-fp8-rows
Open

ashhart wants to merge 1 commit into
mainfrom
lead/dsv41-fp8-rows

Conversation

@ashhart

@ashhart ashhart commented Nov 14, 2026

Copy link
Copy Markdown
Owner

DeepSeek-V4.1 decode on Metal spends about 40% of a layer in the dense FP8 projections (fp8_rows.metal). At one row they read about 675 GB/s on an M5 Ultra, which streams one large read at about 1.1 TB/s. This change gets them closer and keeps every output bit.

  • Loads first. The one-launch kernel issues every weight and scale load of a lane before its FMA chains.
  • Tile roots in one simdgroup. An output's tile roots now meet in one simdgroup's butterfly. Lane k holds tile k's root, zero past the tiles. That's the same zero-padded adjacent-pair tree one thread used to walk over 32 roots, so the sums and their order don't change. Most of the gain is here.
  • Kind d. A new block kind takes one column a simdgroup over 32-group tiles. kindOf picks it for K in whole 1024s with 2048 or more columns in the launch's batch, at 1 and 3-5 rows, and b otherwise. c is no longer picked.
  • 6 to 16 rows are left to dsv41 metal: launches of 6 to 16 rows run their FP8 projections on a matrix tile (stacked on #1720) #2973's tile kernel. At those widths b's new internals are about level (FP8 sum 2-4% faster, wo_b at 16 rows 4.6% slower).

Measured on an M5 Ultra (512 GB) on release layer 2, the eight FP8 launches each timed alone (time-dsv41-metal-real, TF_DSV41_TIMING_FP8=1, mean of two balanced runs):

Rows Before After
1 247.4 µs (675 GB/s) 214.8 µs (778 GB/s)
4 325.5 µs 281.4 µs
5 359.8 µs 316.0 µs
  • Engram's wkv (6144 x 25600) moves to kind d: 311-365 → 230-237 µs at one row.

  • Served on the release weights with dsv41 metal: launches of 6 to 16 rows run their FP8 projections on a matrix tile (stacked on #1720) #2973 and dsv41 metal: decode experts decode an expert's weights once for all of its rows (stacked on #1720, #2973) #2983 underneath, 16 prompts of 512 tokens: draft:false 46.4 → 49.5 tok/s and greedy 82.2 → 84.3, with every reply's text equal to main's (48 of 48).

  • Bits: release layers 2 and 21 give main's output and KV hashes at 1 and 4 rows, 20 of 20 concurrent launches equal, and the float64 real-layer checks pass (layer 2: 362 passes, layer 21: 266).

  • Kernel contract, --only "dsv41_fp8_rows": 6 kernels pass gates 1-4, 4 controls red, and 288 invariance launches byte-equal. The probe's invariance loop now covers every kind.

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

  • Tried and dropped: kind c for K=1280, 2-4 adjacent groups a lane, 2-4 column blocks a threadgroup, X held in registers across columns, applying the scale once at the end of a chain (changes subnormal and overflow bits).

  • Synthena (Ash's Agent)

…olumn a simdgroup for big K

The one-launch FP8 row kernel now issues every weight and scale load before its chains, and an output's
tile roots meet in one simdgroup's butterfly (lane k holds tile k's root, zero past the tiles), the same
zero-padded pair tree that one thread used to walk over 32 roots. A new block kind d takes one column a
simdgroup over 32-group tiles. kindOf picks d for K in whole 1024s with 2048 or more columns in the batch
at 1 and 3-5 rows, and b otherwise; c is no longer picked. The contract's row invariance covers every kind.

Bits unchanged: release layers 2 and 21 give main's output and KV hashes at 1 and 4 rows, 20 of 20
concurrent launches equal, and the float64 real-layer checks pass.

M5 Ultra, release layer 2, the 8 FP8 launches timed alone: 1 row 247 -> 215 us, 4 rows 325 -> 281 us,
5 rows 360 -> 316 us. x40 round: 1 row 25.1 -> 24.6 ms, 4 rows 41.9 -> 40.0 ms, 5 rows 47.5 -> 45.7 ms.

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