docs: kernel organization spec - #259
Merged
ekryski merged 4 commits intoJun 3, 2026
Merged
Conversation
github-actions
Bot
added
the
documentation
| - One **cohesive file per kernel** = the op + all its shape / precision / dtype | ||
| / bit-width variants, with its kernel spec, bench spec, and test spec | ||
| together. | ||
| - Make the MLX reference a purely **optional per-bench attribute** that any |
Owner
There was a problem hiding this comment.
It's currently fully optional, we are no longer hardcoding the name mlx or mlx_ref etc, just refer to them as metal refs moving foward.
Collaborator
Author
There was a problem hiding this comment.
ya there may be a couple hallucinations or things that are now stale (since yesterday?!) in the doc. More for a general idea of how to clean up the files, the why, and maybe a rubric for how we decide whether something deserves it's own file or not.
Comment on lines
+143
to
+144
| - a different IR / algorithm with little shared body (decode SDPA vs prefill | ||
| SDPA vs bidirectional SDPA → three files under `sdpa/`). |
Owner
There was a problem hiding this comment.
One thing I'd add here:
Would expect folder structure/naming to look like this:
sdpa/
bidirectional.rs
decode.rs
prefill.rs
Collaborator
Author
There was a problem hiding this comment.
agreed. I prefer this structure too. At least for now. Otherwise I think one file would be way too big.
Collaborator
Author
There was a problem hiding this comment.
Added in 148ecdd — the explicit per-family sub-folder shape (sdpa/{bidirectional,decode,prefill}.rs + mod.rs) right where the family splits along genuine algorithm boundaries, with a note that dtype/head-dim/shape variants of one op still collapse into that op's single file.
Collaborator
Author
|
Making minor modifications and will merge. Will move specs to a |
Extract the metaltile-std kernel-layout proposal (file granularity, canonical per-kernel file shape, family-by-family migration plan) into a standalone doc so it can be reviewed and landed independently of the kernel-implementation work it shipped alongside. Linked from the docs index under Design & planning.
…hape
Reference comparators are a generic optional metal ref (`ref = MetalRef
{ … }` on #[bench], default none), not hardcoded mlx/mlx_ref — verified
against the macro (valid #[bench] keys: name, dtypes, bytes, flops, ref)
and the RefKernel / with_reference API. Reworded the goals, §1 why, §5
canonical-file example (the stale op=/subop=/class=/mlx= form is gone),
§7 policy, and §8 to match. Added the per-family sub-folder layout
(sdpa/{bidirectional,decode,prefill}.rs) per review.
PR #250 landed comprehensive precision support — ~30 QFormat block-scaled formats with shared format.rs/codec.rs infra, but materialized as 15 per-op *_block_scaled.rs files (~240 #[kernel] fns, ~39 kLOC) where the op body is copy-pasted once per format. Rewrote §6 around consolidating that: a `dequant_block_scaled(...)` DSL op that codegen lowers per QFormat (element unpack + decode + block-scale), plus promoting format to a first-class macro axis (the int family already does this via int_conv2d_*! macros; the float formats never got it). That makes op and format orthogonal, collapsing the matrix toward ~one body per op. Reconciled the §3 quant subtree, goals, and sequencing notes with the now-merged state.
ekryski
force-pushed
the
ek/docs-kernel-organization
branch
from
June 3, 2026 16:29
148ecdd to
8c58786
Compare
Cross-references STYLE_GUIDE.md as the authority on per-kernel conventions (file skeleton, naming, oracle, bench, DSL) and reconciles the overlap: - §5/§6 macro axis now points at the real #[kernel(variants(...))] from #261 (bit-width, head-dim/lane-packing, and quant-format axes) instead of a hypothetical "macro we'd need to build" — the remaining work is generalising hand-written bodies over the variant param, not building the macro. - reference comparator API matches the guide: .with_reference(RefKernel:: new(...)) everywhere (dropped the stale MetalRef / `ref =` framing). - naming (§8) defers to the guide's mt_<op> / ffai_<op> / mt_<family>_< variant> table + path-style bench names; dropped the dead op=/subop= bench-spec attrs. - §1 acknowledges the guide documents the current mlx/ + ffai/ convention that this spec proposes evolving past.
ekryski
enabled auto-merge (squash)
June 3, 2026 16:43
3 tasks
0xClandestine
pushed a commit
that referenced
this pull request
Jun 3, 2026
## Summary Lands `docs/KERNEL_ORGANIZATION.md` as a standalone, **docs-only** change so it can be reviewed and discussed without wading through kernel-implementation diffs. The doc is a **proposal** for `metaltile-std`'s kernel source layout — file granularity, the canonical per-kernel file shape (kernel → `kernel_tests` → `kernel_benches`), and a family-by-family migration plan that intentionally avoids a big-bang pass. It coordinates with (but is orthogonal to) the CLI subprocess-rewrite spec: that spec governs the crate/CLI architecture; this one governs the kernel files inside `metaltile-std`. ## Why split it out It was originally written alongside the vision/audio/TTS/video kernel batch in #252. Pulling it into its own PR lets reviewers evaluate the proposed organization on its own merits and land it independently — #252 drops the file in the same change. ## Changes - Add `docs/KERNEL_ORGANIZATION.md` (extracted verbatim from #252). - Index it under **Design & planning** in `docs/README.md`. No code changes.
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.
Summary
Lands
docs/KERNEL_ORGANIZATION.mdas a standalone, docs-only change so it can be reviewed and discussed without wading through kernel-implementation diffs.The doc is a proposal for
metaltile-std's kernel source layout — file granularity, the canonical per-kernel file shape (kernel →kernel_tests→kernel_benches), and a family-by-family migration plan that intentionally avoids a big-bang pass. It coordinates with (but is orthogonal to) the CLI subprocess-rewrite spec: that spec governs the crate/CLI architecture; this one governs the kernel files insidemetaltile-std.Why split it out
It was originally written alongside the vision/audio/TTS/video kernel batch in #252. Pulling it into its own PR lets reviewers evaluate the proposed organization on its own merits and land it independently — #252 drops the file in the same change.
Changes
docs/KERNEL_ORGANIZATION.md(extracted verbatim from feat(ffai): vision/audio/TTS/video kernel batch + kernel_tests_harness dead-strip fix #252).docs/README.md.No code changes.