Repository navigation
Metal M5: prompt lane projections at every width take a chunk's 16-row blocks in one dispatch - #2851
Merged
Merged
Conversation
This was referenced Nov 11, 2026
Closed
8 tasks done
Give tf_lane the reg path's 16-row grid axis and runtime sums stride. Each block keeps the same group products, FMAs and K reduction; cooperative and older-chip encodes retain 16 rows. Qwen27 selects a whole-prompt key for non-reg widths behind the existing Ultra gate. Decode and mixed stacks stay short. Compare each width with main's frozen short kernel, including ragged tails, producer sums, strided outputs and three controls. Pin real quant dispatch counts with main-red host regressions. Refs #2818
The full lane contract compiles all four tf_lane_sums sources for M1 and M2 at a 1024-thread pipeline limit. Record those measured values and keep the actual 256-thread launch. Tensor-op sources remain explicitly unmeasured on these chips. Pin the distinction in a host regression; no shader changes. Refs #2818
Only whole-prompt layouts compile the independent row-block header. Short layouts retain a scalar group coordinate and their original row extent; guarded sums remain bounds-safe. Keep older-chip sources on that short variant as well. Compare both short and long programs against frozen main at every maintained byte-check shape and row count. Refs #2818
Leave other affine lane callers on their scalar 16-row source and capacity. Only the Qwen prompt key and byte fixture request independent row blocks. Pin the shared shader bytes, describe the row layout accurately and cover the 16/17-row boundary. Refs #2818
ashhart
force-pushed
the
lead/m5-lane-prompt-widths
branch
from
November 13, 2026 13:43
54b0bd8 to
3668b1a
Compare
ashhart
added a commit
that referenced
this pull request
Nov 13, 2026
Figures and scope corrected to the PRs (#2948, #2863, #2934, #2933, #2921, #2906, #2878, #2892, #2851, #2728, #3072, #2915, #2864, #2854 and others), promises of later work removed (#2903, #2875, #1630, #1720), #3074 and #3047/#3046 said as their PRs scope them, the duplicate #2861, #2585, #2321 and #2719 entries and two copies of released 1.2.5 lines removed, #1953's line named for its PR, and a line for #3255. Released sections unchanged.
scchumbletech
added a commit
to scchumbletech/TensorFold
that referenced
this pull request
Nov 14, 2026
… Mac ashhart#2851 gates the 16-row-block prompt path on Ultras; an M5 Max falls back to one 16-row window a dispatch and streams the whole weight slab per dispatch. Open the path on tensor units at large, Ultras keeping their long keys. Receipt (Apple M5 Max 40-core, 128 GB, Qwen3.8-27B-MLX-8bit, 8k salted unique-prefix prompt, streaming TTFT, 3-run median, --drafter z-lab/Qwen3.8-27B-DFlash2, --parallel 4, --context 262144): - main 09a9eb3: 336 tok/s cold 8-bit prefill - this change: 531 tok/s (+58%), replies byte-identical on the measured prompts - 4-bit unchanged (842-858 tok/s stock-level across runs)
scchumbletech
added a commit
to scchumbletech/TensorFold
that referenced
this pull request
Nov 14, 2026
The dispatch test held every non-Ultra Mac to 16 rows a dispatch; on a tensor-unit Mac the change takes the row blocks ashhart#2851 gave Ultras. The non-tensor-unit expectation stays.
8 tasks done
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
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.
On M5, the Metal lane projection took a prompt chunk's 16-row blocks one dispatch at a time for every width except the 4-bit path #1738 widened, so Qwen3.8-27B MLX 8-bit and the oQ checkpoints' other widths made ~128 dispatches a projection a chunk.
What changes: whole non-reg prompt chunks put their independent 16-row blocks on the grid's y, advancing X, Y and the scales per block and keeping each block's extent at min(remaining, 16). Products, group FMAs and split-K sums keep main's order. Layouts admit 16 to 2,048 rows and scratch follows capacity; short layouts compile the scalar grid without row-offset work. Qwen27 selects the long key only above 16 prompt rows through the unchanged Ultra gate. The reg and cooperative paths keep their limits; the Ultra-only gate (P1a) and the mixed stack are untouched.
Exactness (on an M5 Ultra): a byte suite against main's frozen kernel (kept under
zig/tests/fixtures/lane_prompt/) matches in all 352 reg/cooperative, 54 long-reg, 192 short and 192 long cases: widths 8/g64 (inline and producer sums), 6/g32, 5/g64 (inline and prepared), 3/g32, 2/g64 and 8/g128 at 17, 1,065 and 2,048 rows, with partial and strided columns and output padding. All three wrong-block and sums controls are caught. Lane contract: 10/10 kernels pass gates 1-4, 8/8 controls red, no limit problems. Main fails the new host test (InvalidSchedule at 17 rows).Measured on an M5 Ultra, Qwen3.8-27B MLX 8-bit, warm/base/new/new/base fresh processes, warm excluded, prompt cache off; all 12 prompt keys have equal full token hashes in every process:
One-stream decode is level (+0.2% chat, +0.4% code). Host, fmt and lean (no new findings) pass.
For #2818.