Repository navigation
Performance: decode overlap, GQA-shared split attention, short-prompt GEMM tile (1.0.3) - #3
Merged
Conversation
…current one Encoding a forward pass depends only on (T, offset, computeLogits, allLogits); token ids are read from the ids buffer at execution time. After committing a final step, step() encodes the T=1 pass for offset+T before waiting, and commits it on the next call only on an exact match (otherwise it is dropped uncommitted). Same kernels, same dispatch order: outputs unchanged. The embedding lookup pipeline is resolved once at init instead of by name every step. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
…dding to 32 qgemm tiles 32 tokens, so a 5..15-token prefill paid the MMA work of 32 tokens (Mistral 7B: ~14.4 GFLOP/token x 32 / 2.45 TFLOPS = ~190 ms, the flat ~200 ms measured for T=5..15). qgemms uses a BM = 8 or 16 token tile with 4 simdgroups x 8 output columns, K-step 64 and register prefetch of the next step. Per output element the arithmetic is the same as qgemm (fp32 s*q+b dequant, zero-padded x, 8x8x8 MMAs with k ascending, 8-aligned tiles), so the accumulation order is unchanged. T=1, T=2..4 and T>=17 paths untouched. Kernel tests cover T = 8, 9, 15, 16, 17. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
…ernel shares K/V across GQA heads attention<E,ANSG> now loads K/V rows for 4 keys up front (clamped to L-1) and processes them in the original order, so more reads are in flight. For decode at context >= MLXR_ATTN_SPLIT_MIN (default 256), attention_split runs nKV x 8 threadgroups of 4 simdgroups; each simdgroup covers all query heads of its GQA group, so each K/V row is read once per group instead of once per head. It writes the same 32 per-simdgroup partials the fused kernel keeps in threadgroup memory, and attention_combine merges them in the same order: the arithmetic is unchanged and logits are bit-identical (new test splitDecodeAttentionIsBitIdentical). Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
…ter at short context Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
🚀 Deploying Preview to Cloudflare 🚀Preview Deployments by commit
|
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.
Changes
Transformer.swift): while the GPU runs token N, the CPU encodes the command buffer for token N+1. Encoding depends only on (T, offset), and token ids are read fromidsat execution, so the buffer is used only on an exact (T=1, offset) match and dropped otherwise.attention_split+attention_combine): each K/V row is read once per GQA group instead of once per query head. Partials are merged in the original order. The K/V prefetch (AU) stays at 1 because it measured no faster. New testsplitDecodeAttentionIsBitIdentical.qgemms): T = 5–16 uses an 8- or 16-token tile instead of padding to 32. The accumulation order is the same asqgemm. Kernel tests cover T = 8, 9, 15, 16, 17.Measurements (M2 Pro, 1.0.2 vs this branch, 4 alternating pairs, run-to-run spread ≤ 0.5%)
Split threshold check (split off vs on, same binary): @256 +0.3% / +1.0%, @512 +0.5% / +2.9% (Mistral / Qwen3), so 256 causes no regression.
Correctness
swift test: 95 tests pass.Rejected during the pass: a small-matrix matvec variant (
qmv4r), which made Qwen3 0.6B decode 5% slower.🤖 Generated with Claude Code