Repository navigation
perf(model): exact Metal speed-ups of the forward pass - #45
Merged
Merged
Conversation
The student scorer spends 157 ms per hook at 12 candidates, and the next speed-up has to come from the kernels without changing a single ranking. To choose which kernels to touch, profile_forward times the hook's scoring path and every op of an encoder and a head layer at the model's real shapes, pipelined as in a real forward, and sets each op's share against the measured batch. At b=8, s=187 the ops sum to 97% of the batch, so the table explains where the time goes. The parity test also reads its fixtures from LAYA_CODEX_PARITY_FIXTURES, so that the student export and r1 can be checked against fixtures made from their own model dirs. author: Tin Dang
…rkers Only the [MASK] option positions reach the scorer, yet the head's last layer computed queries, the output projection and its 4x-wide FFN for every token. Keys and values still need every position, but everything after them is row-wise and attention is independent per query row, so the last layer now projects its queries, attends, and runs its FFN on the b * kmax marker rows alone. The scored rows are the same. On the student at W=128, K=12 (Metal, f16) the hook's scoring call drops from 156.7 ms to 148 ms (p50, alternating runs). Parity against the Python fixtures is unchanged for the student and r1 on CPU and Metal. Attention over a subset of query rows gets its own entry point. With a single query row it takes the explicit path on Metal, because the fused single-query kernel ignores the mask; a unit test covers one, two and three rows against the matching rows of full attention, and a layer-by-layer f32 reference checks the head's logits on CPU and Metal. author: Tin Dang
candle's rope_thd kernel computes every element's position with 64-bit integer division, which the GPU emulates: at b=8, s=187 it took 0.41 ms per layer, about 9% of the student's forward pass, for a memory-bound elementwise op. The crate now compiles its own kernel with the same arithmetic on a 3-D grid (half-dim, head, row) with 32-bit indices: 0.034 ms for the same tensor, 12x faster. The output is bit-identical to candle's (a unit test compares f16 and f32 on Metal with zero tolerance), the kernel is compiled once per device with Metal's default fast-math options, the ones candle uses for its own kernels, and the CPU path is unchanged. Student and r1 parity against the Python fixtures is unchanged on CPU and Metal. candle-metal-kernels and objc2-metal come in under the metal feature at the versions candle already uses. author: Tin Dang
The fused attention kernel wrote (b, h, s, hd) and every layer then assembled the (b, s, h * hd) rows the output projection reads from b * h blits plus a zero fill: 0.21 ms per layer at b=8, about 5% of the student's forward pass. The kernel already takes output strides, so a small custom op calls it with the strides of a (b, s, h, hd) buffer and the transpose disappears. The kernel, its tiling and its inputs are unchanged, so the output is bit-identical (unit test with zero tolerance, full and two-row queries); the mask must start at offset 0 because the kernel binds it without one, so an offset view is made contiguous first. Student and r1 parity against the Python fixtures is unchanged on CPU and Metal. author: Tin Dang
candle adds a linear layer's bias through its strided broadcast kernel, which computes every element's offset with 64-bit '%' per dimension: on Metal the decision head's biased projections took about twice as long as their matmuls (e.g. 1.49 ms instead of 0.67 ms for the packed q|k projection at b=8, s=187). The encoder has no biases; the head's first layer and its last layer's key/value projections still run over every token. Dense keeps candle's matmul and adds the bias with the crate's 2-D kernel, the same x + b in the same dtype: bit-identical to candle_nn::Linear on Metal and identical code on the CPU (unit tests with zero tolerance on both). Student and r1 parity against the Python fixtures is unchanged. With the three previous changes, the student's scoring call at W=128, K=12 goes from 157-167 ms to 132-138 ms (p50 in alternating runs on a loaded machine). author: Tin Dang
profile_forward's op table described candle's ops before the kernel work. It now times the ops the model runs (the crate's rope and row-major attention kernels, Dense projections, a head whose last layer works at the marker rows) with their real call counts, and keeps the replaced candle ops as zero-count reference rows. At b=8, s=187 the ops sum to 96% of the measured batch; the encoder's matmuls are 77% of it. author: Tin Dang
…uild The crate's kernel library declared its bf16 variants unconditionally, and `bfloat` needs MSL 3.1 (macOS 14). On macOS 12-13 the library would fail to compile, every forward pass would error, the hook would fall back to lexical ranking, and the compile was retried under the global cache lock on every call. The bf16 kernels are now behind `#if defined(__HAVE_BFLOAT__)`, as candle guards its own. A library or kernel that cannot be built is remembered as unavailable per (device, source, kernel) and logged once, and rope_thd and add_bias then use candle's op: older macOS keeps the speed it has today instead of losing the model. Tests force a library that does not compile: both ops match candle bit for bit, the compile is attempted once, and a missing kernel is unavailable. The bit-equality tests now also assert that the crate's kernel, not the fallback, produced the result, and bf16 is checked wherever candle itself supports it. author: Tin Dang
…ds it The fused attention kernel binds the mask buffer without its start offset, so sdpa_rows meant to move an offset mask into a fresh buffer first. It used contiguous(), which returns an already-contiguous view unchanged, offset included; such a mask then failed the op's own offset check, contrary to what the commit that added sdpa_rows said. force_contiguous() always copies into a new buffer at offset 0. A test with a contiguous mask view one batch row into its buffer failed before and now matches the copied mask bit for bit. The model's own masks start at offset 0, so its forward pass is unchanged. author: Tin Dang
…them candle's fused attention kernel takes only the (batch, head, row) strides of q, k, v, the mask and the output and reads every row with unit stride, but sdpa_rows never checked that and its doc promised "any strides". A q or mask whose last dim was not unit-stride (e.g. a transposed view) gave silently wrong attention. The op now refuses such inputs, and sdpa_rows copies them to contiguous tensors first (the same values, so the same result; candle's own sdpa has the same assumption, so it is not a safer fallback). The doc states what is read in place and what is copied. The model's q, k, v and masks are unit-stride already, so its forward pass makes no extra copy. A test with transposed q, k, v and mask views failed before (wrong values) and now matches the contiguous inputs bit for bit. author: Tin Dang
metal_ops calls candle-metal-kernels' call_sdpa_full, an internal function whose argument list and output-stride handling can change in any release, and its rope and bias kernels must stay bit-identical to candle's. A caret requirement would let `cargo update` pull a 0.11.x that silently changes either. candle-core, candle-nn and candle-metal-kernels are now `=0.11.0`, with a comment saying why and what to re-run on an upgrade. Cargo.lock already resolved 0.11.0 and is unchanged (the --locked builds pass with and without the metal feature). author: Tin Dang
…n paths Every Metal test returned early when Metal was missing, so a build without the metal feature, or a runner without a GPU, passed them all. With LAYA_REQUIRE_METAL=1 the shared Metal test device and the Metal parity test now fail instead of skipping (documented in the crate docs). Before: `LAYA_REQUIRE_METAL=1 cargo test -p laya-model --lib` without the feature passed 16 of 16; now 3 tests fail and say why. The fused-attention equality test now runs f16 and f32, a batch of one and of three, key lengths 37, 64 and 704 (the kernel's aligned paths up to the student's max_len), with full and two-row queries. New tests: a single query row is refused rather than computed with the mask ignored, and the bias add on a non-contiguous input takes candle's path and matches it. The bf16, offset-mask and failed-build cases are covered by the earlier fixes. author: Tin Dang
The crate's kernels index with u32 and dispatched whatever grid they were given. Their docs now state the limit, and a tensor with more than u32::MAX elements takes candle's op instead of erroring (the student at max_len 704 and 64 rows has about 3.5e7, so this never triggers in practice). Empty inputs return an empty result without a zero-size dispatch. The empty-input test passed before the change too (candle tolerated the zero-size launches here); it now guards the early returns. author: Tin Dang
The head's last layer took the first kmax query rows of the square (b, 1, s, s) padding mask as the mask of its marker rows. That is right only because every row of that mask is the same, an assumption nothing enforced. AttnMasks now keeps the (b, 1, 1, s) key-padding row it is built from, the head takes that row (and checks its shape), broadcasts it over every query for its full layer, and passes it as is to the marker-row attention. The values the kernels read are unchanged, so the logits are too. The subset-attention test now checks both the square mask's rows and the broadcast key row, on CPU and Metal. author: Tin Dang
Since the parity tests can read fixtures from LAYA_CODEX_PARITY_FIXTURES, a missing or malformed file still failed with "fixtures/laya_parity.json", pointing at the wrong file. The panic now names the path it tried and the error. author: Tin Dang
This was referenced Oct 9, 2026
# Conflicts: # crates/laya-model/tests/parity.rs
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
2 participants
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.
Exact speed-ups of the Metal forward pass. Outputs are unchanged.
Measured: the student's model stage goes from 156 to about 121 ms at 12 candidates, and the hook p50 from about 163 to about 130 ms on a quiet machine. r1 also gets faster (503 → 449 ms).
Four changes, found by profiling one hook-sized forward pass:
Exactness:
Review fixes:
__HAVE_BFLOAT__. If a kernel can't be built, that's cached and candle's op is used instead, so macOS 12–13 keeps today's speed.=0.11.0, because one change calls a candle-internal function.LAYA_REQUIRE_METAL=1makes the Metal tests fail instead of skipping. Hosted runners may have no Metal device, so run it locally before merging:LAYA_REQUIRE_METAL=1 cargo test --workspace --release --locked.