Skip to content

refactor(std): restructure conv + rope kernels into kernels/<family>/ - #286

Closed
ekryski wants to merge 14 commits into
0xClandestine:devfrom
thewafflehaus:ek/kernel-reorg-conv-rope
Closed

refactor(std): restructure conv + rope kernels into kernels/<family>/#286
ekryski wants to merge 14 commits into
0xClandestine:devfrom
thewafflehaus:ek/kernel-reorg-conv-rope

Conversation

@ekryski

@ekryski ekryski commented Jun 12, 2026 •
edited
Loading

Copy link
Copy Markdown
Collaborator

What

Restructures metaltile-std's kernels toward the kernels/<family>/ layout
from docs/specs/KERNEL_CONSOLIDATION_PLAN.md, completing the first two
families: conv and rope. No kernel math changes — this is a relocation,
naming, and redundancy-removal pass.

Net −23,494 lines, almost all of it redundant per-format conv test/bench
boilerplate collapsed into macro_rules! generators.

Conv

  • Delete stale ffai/ conv + kokoro duplicates that convolution/ already
    superseded.
  • DRY the conv2d_mma block-scaled benches and the conv2d/conv3d/depthwise
    block-scaled per-format tests via local *_fmt! macros (1 macro + N
    invocations replacing N hand-written copies).
  • git mv convolution/ -> kernels/conv/ and rewire the two references
    (lib.rs, the mlx/steel re-export). Pure rename — kernel pub fn names are
    preserved, so the inventory and FFAI emit are unaffected.
  • Drop the duplicate PATTERNS.md / PLAN.md (folded into the consolidation
    plan).

Rope

  • Migrate the rope family out of the legacy mlx/ + ffai/ split into
    kernels/rope/.
  • Drop the ffai_ / dsv4_ prefixes — every kernel now carries mt_
    (mt_rope, mt_rope_llama, mt_rope_yarn, mt_rope_2d,
    mt_partial_rope). Names change, so the FFAI emit consumer is regenerated
    against mt_*.
  • Collapse the single-token kernels into the batched form. mt_rope_llama
    and mt_partial_rope now take a per-row/per-token position plus 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 scalar-position kernels were
    redundant: #[constexpr] is a runtime buffer read (one PSO per dtype), so a
    scalar-position variant bought nothing over a single-row dispatch. Tests cover
    decode, prefill, and banding; benches keep both a decode and a prefill shape.
    Decode callers change to pass positions=[pos] / base_position=pos, T=1.

Verification

  • cargo clippy --workspace --all-targets clean.
  • tile test --filter rope — 33/33 pass, including the merged
    mt_rope_llama_many.
  • tile test --filter conv2d — 312/312 pass.
  • dsv4_partial_rope_rows_correctness integration test passes.

Notes

The KERNEL_AUDIT.md conv rows already drifted (they point at ffai/conv2d.rs
paths that predate the convolution/ canonicalization); reconciling the full
audit against the live inventory is a separate follow-up. Only the rope rows —
which this PR's rename made stale-from-accurate — were corrected here.

ekryski added 6 commits June 12, 2026 05:59
@ekryski
docs(conv): drop duplicate PATTERNS.md / PLAN.md from convolution/
8010efe 
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.
@ekryski
docs(arch): update ARCHITECTURE.md to the current codebase
bc9ced4 
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).
@ekryski
moving temporary hip and cuda markdown files to specs directory
4505ee9
@ekryski
docs: consolidate kernel roadmap + move STYLE_GUIDE to docs/
e84d473 
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.
@ekryski
docs: bring each crate README up to date with the current codebase
a41eb9a 
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.
@ekryski
docs: update BENCH_METRICS_SPEC + resolve KERNEL_CONSOLIDATION_PLAN q…
be4929d 
…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).
@github-actions github-actions Bot added the ignore-for-release Hidden from release notes label Jun 12, 2026
ekryski added 8 commits June 12, 2026 07:38
@ekryski
refactor(std): migrate rope family into kernels/rope/ (first family)
213b658 
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.
@ekryski
refactor(ffai): delete stale conv/kokoro duplicates superseded by con…
f1d0799 
…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.
@ekryski
refactor(conv): DRY conv2d_mma block-scaled benches via macro
f513877 
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.
@ekryski
refactor(conv): consolidate block-scaled per-format tests via macro
98809c0 
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).
@ekryski
refactor(conv): move convolution/ to kernels/conv/ per consolidation …
09a1aa5 
…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)
@ekryski
refactor(rope): rename to mt_ prefix and merge decode/prefill forms
5cd0e2e 
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).
@ekryski
refactor(rope): collapse single-token kernels into the batched form
8509edd 
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.
@ekryski
refactor(kernels): name by operation not model; conv -> convolution
bc1b5a4 
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.
@ekryski
ekryski force-pushed the ek/kernel-reorg-conv-rope branch from ac584ad to bc1b5a4 Compare June 12, 2026 13:50
@ekryski

ekryski commented Jun 12, 2026

Copy link
Copy Markdown
Collaborator Author

Updated: renamed kernels/conv/ → kernels/convolution/ (spell out abbreviated family names) and de-model-ified the rope kernel (mt_rope_llama → mt_rope_banded; Llama-3/Qwen noted in a comment). The branch is now based on the docs branch (#284) so the consolidation plan + style guide travel with the refactor and are updated in place. Build + clippy clean; 36 rope + conv tests pass.

This was referenced Jun 12, 2026
Closed
Closed
Closed
@ekryski

ekryski commented Jun 12, 2026

Copy link
Copy Markdown
Collaborator Author

As per @0xClandestine's request, as relayed to me by @TheTom, closing in favour of moving rapid iteration to forks

@ekryski ekryski closed this Jun 12, 2026
This was referenced Jun 12, 2026
Merged
Merged
Merged
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

No reviews

Assignees

No one assigned

Labels

ignore-for-release Hidden from release notes

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

1 participant

@ekryski