refactor(ops): migrate core/elementwise primitives into kernels/ops/ (completes wave 1) - #289
Closed
ekryski wants to merge 17 commits into
Closed
refactor(ops): migrate core/elementwise primitives into kernels/ops/ (completes wave 1)#289ekryski wants to merge 17 commits into
ekryski wants to merge 17 commits into
Conversation
These two planning docs under crates/metaltile-std/src/convolution/ are duplicates of the canonical specs in docs/specs/: convolution/PATTERNS.md -> docs/specs/CONV_KERNEL_PATTERNS.md convolution/PLAN.md -> docs/specs/CONV_CONSOLIDATION_PLAN.md Keep the docs/specs/ copies as the single source of truth; nothing references the in-tree duplicates.
The doc described the pre-#275 single-backend, in-process-runner design. Bring it current: - Multi-backend codegen: the CodegenBackend seam (Target = Metal/Cuda/Hip/Spirv), TargetProfile (lane width, MmaStrategy), and the msl/cuda/hip/spirv generators. Metal is the default/mature backend; CUDA/HIP/Vulkan device execution is feature-gated and validated by the GPU-vs-GPU reference corpora. - Subprocess execution: the runner is no longer in-process — `tile` spawns `__tile_runner` (project-linked, inventory populated) and streams ProtocolMessage JSON lines. Rewrote the "subprocess readiness" section accordingly. - metaltile-std modules: mlx / ffai / convolution / quant (was mlx + ffai only). - Commands: added clean + config; clarified `emit` is `build --emit msl|metallib|swift|ir|all` and that --backend applies to bench/test only. - Fixed companion-doc links (../cli.md, ../developing.md) and the now-renamed metric path (to_gflops / estimate_profile / classify_bottleneck / device_specs::lookup → ProfileInfo; was derive_metrics / DerivedMetrics).
Three overlapping docs (KERNEL_ORGANIZATION, CONV_CONSOLIDATION_PLAN,
CONV_KERNEL_PATTERNS) collapse into one singular reference,
docs/specs/KERNEL_CONSOLIDATION_PLAN.md — the kernels/<family>/ target layout
(no more mlx/ + ffai/), the three LOC-reduction tools (shared primitives,
#[kernel(variants)], op×format merge), the proven conv/ exemplar, and the
family-by-family migration order.
STYLE_GUIDE.md moves crates/metaltile-std/ → docs/ and becomes the authoring
authority ("how we want kernels written post-consolidation"): de-mlx/ffai'd,
the removed #[bench] name key fixed, a shared-primitives section added, family
mod.rs paths. developing.md's kernel-writing-philosophy now links to it; the
authoring patterns from CONV_KERNEL_PATTERNS fold in here rather than the plan.
Also fixes the stale ../specs/ links in docs/README.md (specs live at
docs/specs/), repoints ARCHITECTURE.md, and indexes both new docs.
Roadmap (where files go) vs authoring (how to write one kernel) are now one doc
each, with the redundancy across the old four removed.
The crate READMEs had drifted across several big changes. Fixes: - metaltile-codegen / -runtime / metaltile (facade): no longer "Metal only". Codegen is multi-backend (CodegenBackend, Target = Metal/Cuda/Hip/Spirv, TargetProfile, the cuda/hip/spirv generators); runtime owns device/ (Metal + feature-gated cuda/hip/vulkan), dispatch/, cache/; the facade hosts harness/ + runner/ and runs GPU work in the __tile_runner subprocess. - metaltile-core: GpuFamily lives in src/protocol.rs (not src/gpu_family.rs); documents the protocol wire types. - metaltile-macros: corrected module paths (kernel/body.rs, kernel/sig.rs, bench.rs, derive/mod.rs), added kernel/variants.rs, and replaced the dead #[kernel(bench(op=,class=,mlx=))] system with the current #[bench] / #[test_kernel] / #[kernel(variants(...))]. - metaltile-std: rewritten — it's the kernel stdlib (mlx/ ffai/ convolution/ quant/ probe/ utils), not "benchmark metadata"; dropped the removed BenchSpec/bench_types/run_spec framing. - metaltile-cli: added the missing subcommands (test, clean, config, init, update, completions) and the --backend flag. Cross-doc links point at the relocated style guide / consolidation plan.
…uestions - BENCH_METRICS_SPEC.md: mark the metrics work + Appendix B precision support as implemented/complete; note M5 fp8/fp4 acceleration arrives with Metal 4.1; trim the stale capture context. - KERNEL_CONSOLIDATION_PLAN.md: turn §9 open questions into decisions — (1) rename ffai_/model-prefixed kernels to mt_<op> (the FFAI emit consumer is regenerated, so names are not pinned); (2) keep vision/ + audio/ folders; (3) f16-scale twins fold into a ScaleKind format axis, not separate kernels; (4) kernels that mirror an upstream reference keep their .with_reference() comparator indefinitely. Reconciled §6 (rename step), §7 (scale axis), and §8 (inventory names do change; emit is regenerated).
First family moved into the new kernels/<family>/ umbrella from the legacy mlx/ + ffai/ split (the consolidation target — see KERNEL_CONSOLIDATION_PLAN.md). git mv of all 7 rope files: mlx/rope.rs -> kernels/rope/base.rs (mt_rope) ffai/rope_2d/_llama/_llama_many/_yarn, ffai/dsv4_partial_rope[_rows] -> kernels/rope/ Adds src/kernels/mod.rs + kernels/rope/mod.rs, wires lib.rs (pub mod kernels;), removes the 7 decls from mlx/mod.rs + ffai/mod.rs, and repoints the two intra-family/external references (rope_llama_many's test helper import and the dsv4_partial_rope_rows correctness test). base.rs avoids the rope::rope module-inception lint. Pure structural move — pub fn names are unchanged (mt_rope, ffai_rope_2d, …), so the kernel inventory and FFAI emit are byte-identical. Build + workspace clippy clean; kernel_tests_harness green. LOC-reducing merges (rope_llama + _many, partial_rope folding, ffai_/dsv4_ prefix drop) are the documented next phase.
…volution/
The #268 conv consolidation moved every 2D/3D/depthwise/winograd kernel into the
new convolution/ module (declared in lib.rs, fully built/tested/benched/emitted)
but left byte-identical copies behind in ffai/, un-declared in ffai/mod.rs —
13 orphaned files that compiled nowhere and duplicated the live kernels:
conv2d{,_block_scaled,_mma,_mma_block_scaled}, conv3d{ -"- },
depthwise_conv2d{,_block_scaled,_nhwc}, winograd_conv, kokoro
10 are bit-identical to their convolution/ twin; conv2d/conv3d_block_scaled
differ only by a one-line comment (the int8_decode rename). kokoro's two kernels
(lstm_cell, adain1d) are likewise superseded by the dedicated ffai/lstm.rs and
ffai/adain1d.rs modules. Nothing imports any of them.
Delete all 13, and fix the now-stale `ffai::conv2d` doc-ref in convolution/
conv3d.rs to `super::conv2d`. No kernel lost; the live ones stay in convolution/.
Harness green (0 failures), workspace clippy clean.
conv2d_mma_block_scaled was the lone block-scaled conv file still spelling out all 29 per-format benches as explicit fns (~384 lines), while every sibling — conv2d/conv3d/depthwise block_scaled and conv3d_mma — generates them from a `*_bench_fmt!` macro. Replace the 29 fns with a `conv2d_mma_bench_fmt!` macro (mirroring conv3d_mma's) + 29 one-line invocations, and adopt the siblings' short bench names (bench_mxfp4, …) for naming consistency. Behaviour-preserving: same kernels, same representative MMA shape (1×64×32×32 → out_ch 1024, 1×1), same 29 formats incl. fp8_e4m3_f16 reusing the nvfp8_f16 kernel. Build clean; kernel_registry_consistency green (all benches codegen, no undefined symbols). −384 lines.
The four block-scaled conv files each spelled out 30 per-format correctness tests as explicit fns (~600 lines/file), all following the identical pattern of calling a shared setup helper with only (kernel, QFormat) differing — while the sibling benches already collapse to a `*_bench_fmt!` macro. Mirror that on the test side with a `*_test_fmt!` macro per file (conv2d/conv3d/depthwise + conv3d_mma), replacing the explicit fns with one macro + 30 one-line invocations. Behaviour-preserving: same 120 tests, same names, same uniform shape/tol per file, same kernels (incl. fp8_e4m3[_f16] reusing the nvfp8[_f16] kernel). Each file's kernel_ir_for count verified at 30. Full kernel_tests_harness green; workspace clippy clean. ~2100 lines removed with zero coverage change. This is option (a) of the test consolidation — a macro_rules! mirror of the benches; migrating both to #[test_kernel(variants(...))] is deferred to the broader kernel cleanup (see CONV_CONSOLIDATION_PLAN.md Phase 2).
…plan Relocates the convolution family under the kernels/<family>/ layout (docs/specs/KERNEL_CONSOLIDATION_PLAN.md). Pure git-mv — kernel pub fn names are preserved so the inventory and FFAI emit are unaffected. - crate::convolution -> crate::kernels::conv (lib.rs, mlx/steel re-export) - drop the stale duplicate PATTERNS.md / PLAN.md (folded into the plan)
Wraps up the rope family per docs/specs/KERNEL_CONSOLIDATION_PLAN.md: - drop the legacy ffai_/dsv4_ name prefixes; every kernel is now mt_* (mt_rope_llama, mt_rope_yarn, mt_rope_2d, mt_partial_rope[_rows]) - merge rope_llama_many into rope_llama.rs (decode + batched-prefill forms share the banded inv_freq oracle in one file) - fold dsv4_partial_rope_rows into partial_rope.rs (same single + rows split). band_inv_freq is no longer imported across modules. Kernel bodies are unchanged; names change so the FFAI emit consumer is regenerated against mt_*. Updated the one integration test import, the cross-referencing doc comments, and the KERNEL_AUDIT rope rows. All 33 rope kernel tests pass (incl. merged mt_rope_llama_many).
mt_rope_llama and mt_partial_rope now take a per-row/per-token position and an outer grid axis, so decode is just the single-row case (T=1 / n_tokens=1) and prefill ropes every token in one dispatch. The former single-token mt_rope_llama / mt_partial_rope kernels were redundant: #[constexpr] is a runtime buffer read (not a baked literal, one PSO per dtype), so a scalar-position variant bought nothing over the single-row dispatch. - delete the two scalar-position kernels; rename the batched forms (mt_rope_llama_many -> mt_rope_llama, mt_partial_rope_rows -> mt_partial_rope) - tests now cover decode (single-row), prefill (multi-row), and banding; benches keep both a decode and a prefill shape - rename the integration test file/fn to drop the now-meaningless _rows - decode callers change: pass positions=[pos] / base_position=pos with T=1 All 36 rope kernel tests pass; partial_rope integration test passes.
Applies the naming directive: kernel names describe the operation / layout, never a model — many models share an op in different permutations, so the differentiator is the layout. Model-specific usage notes move to comments. - rope: mt_rope_llama -> mt_rope_banded (rope_llama.rs -> rope_banded.rs); the frequency-band scaling is the operation, Llama-3 / Qwen are noted as users in the header comment. test_rope_llama_banding -> _scaled. - folder: kernels/conv/ -> kernels/convolution/ (spell out abbreviated single-word family names; keep standard acronyms gemm/sdpa/moe/rope/ssm). - docs: KERNEL_CONSOLIDATION_PLAN (no-model-name + folder-naming rules, mark convolution + rope done), STYLE_GUIDE (naming example, §1 layout), KERNEL_AUDIT (rope row). Cross-ref doc comments updated. cargo build + clippy clean; 36 rope + 198 conv3d tests pass.
Moves the normalization kernels out of the legacy mlx/ + ffai/ split into
kernels/norm/ and drops the ffai_ / bare-name prefixes per the consolidation
plan (names describe the operation, not the source/model):
- mlx/rms_norm, mlx/layer_norm -> kernels/norm/ (mt_rms_norm*, mt_layer_norm)
- ffai/{rms_norm_residual,rms_norm_rope,gated_rmsnorm,rms_norm_qgemv,
gated_rms_norm_qgemv} -> mt_* (drop ffai_)
- ffai/adain1d: adain1d -> mt_adain1d
- the *_block_scaled_qgemv format matrices move as-is (already mt_<format>_*);
the format-axis fold (plan §7) is deferred to a later pass.
Generated MSL per kernel is unchanged; only paths + inventory names change.
Updated the two integration-test imports, cross-ref doc comments, the plan
(norm marked done) and KERNEL_AUDIT rows.
cargo build + clippy clean; all norm kernel tests + the rms_norm integration
tests pass.
Moves the logits->token pipeline out of the legacy mlx/ + ffai/ split into
kernels/sampling/ and adds the mt_ prefix to the unprefixed ffai kernels:
- mlx/{softmax,sort} -> kernels/sampling/
- ffai/sampling.rs -> kernels/sampling/categorical_sample.rs
(softmax_categorical_sample -> mt_softmax_categorical_sample; renamed the
file off the family name to avoid module inception)
- ffai/logits_{topk,top_p,min_p,processors} -> kernels/sampling/, with mt_
prefixes (mt_logits_topk_mask, mt_logits_top_p_mask, mt_logits_min_p_mask,
mt_logits_temperature, mt_logits_repetition_penalty)
Generated MSL per kernel is unchanged; only paths + inventory names change.
Updated the softmax_categorical_sample integration-test imports, cross-ref
doc comments (dsv4_indexer_topk, spirv codegen), the plan (sampling done),
and KERNEL_AUDIT rows.
cargo build + clippy clean; all sampling kernel tests + the categorical-sample
integration matrix pass.
Completes wave 1: moves the core/elementwise primitive kernels out of the
legacy mlx/ + ffai/ split into kernels/ops/ per the consolidation plan.
- mlx/{arange,arg_reduce,binary,binary_two,copy,gather_axis,hadamard,
hadamard_m,indexing,logsumexp,random,reduce,scan,scatter_axis,strided,
ternary,unary} -> kernels/ops/ (already mt_*; inventory-registered, so no
import sites to update beyond one snapshot test)
- ffai/{axpy_scalar_inplace,clamp_scalar,gather} + mt_vector_add.rs (renamed
vector_add.rs) -> kernels/ops/ with ffai_ -> mt_ prefixes
- delete ffai/arg_reduce.rs: ffai_argmax was a byte-for-byte duplicate of
mt_argmax (verified) -> real dedup
- fence.rs moves but stays intentionally undeclared (device/system fences have
no DSL representation — the '1 out of scope' op)
Updated the steel_msl_snapshots import, the op_group/name-filter test fixtures,
the plan (ops done, wave 1 complete) and KERNEL_AUDIT rows.
cargo build + workspace clippy clean; all ops kernel tests + steel_msl_snapshots
pass.
Collaborator
Author
|
As per @0xClandestine's request, as relayed to me by @TheTom, closing in favour of moving rapid iteration to forks |
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
1 participant
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.
What
Final wave-1 kernel-family port: moves the core / elementwise primitive
kernels out of the legacy
mlx/+ffai/split intokernels/ops/perdocs/specs/KERNEL_CONSOLIDATION_PLAN.md. With this, wave 1 is complete(
rope·norm·sampling·ops).Changes
mlx/{arange,arg_reduce,binary,binary_two,copy,gather_axis,hadamard,hadamard_m, indexing,logsumexp,random,reduce,scan,scatter_axis,strided,ternary,unary}→kernels/ops/(alreadymt_*; these are inventory-registered, so there wereno path-import sites to fix beyond one snapshot test).
ffai/{axpy_scalar_inplace,clamp_scalar,gather}+mt_vector_add.rs(→vector_add.rs) →kernels/ops/withffai_→mt_prefixes.ffai/arg_reduce.rs—ffai_argmaxwas a byte-for-byteduplicate of
mt_argmax(verified line-by-line), so it's dropped, not moved.fence.rsmoves but stays intentionally undeclared — device/system memoryfences and atomics have no
#[kernel]DSL representation (the "1 out of scope"op in the audit); the file is kept for its implementation notes.
steel_msl_snapshotsimport, theop_group/ name-filter unit-testfixtures, the plan (ops done, wave 1 complete), and the
KERNEL_AUDITrows.Generated MSL per kernel is unchanged — only module paths and inventory names.
Verification
cargo build -p metaltile-std+cargo clippy --workspace --all-targetsclean.tile test(binary, unary, ternary/select, copy,strided, arange, random, reduce, argmax/argmin, scan, indexing, gather(_axis),
scatter_axis, hadamard(_m), logsumexp, clamp, axpy, vector_add).
steel_msl_snapshotsintegration test passes.Note on stacking
Stacked on sampling (#288) → norm (#287) → conv+rope (#286) → docs (#284); the
diff includes those until they merge. Each family lands as its own PR.