diff --git a/crates/metaltile-std/src/kernels/sdpa/mod.rs b/crates/metaltile-std/src/kernels/sdpa/mod.rs index 49ad2d52..d3a7fa9a 100644 --- a/crates/metaltile-std/src/kernels/sdpa/mod.rs +++ b/crates/metaltile-std/src/kernels/sdpa/mod.rs @@ -34,3 +34,4 @@ pub mod sdpa_multi_d256; pub mod sdpa_prefill_d512_sink; pub mod sdpa_rel_pos_conformer; pub mod sdpa_vector; +pub mod steel_attn; diff --git a/crates/metaltile-std/src/mlx/steel/attn/mod.rs b/crates/metaltile-std/src/kernels/sdpa/steel_attn/mod.rs similarity index 100% rename from crates/metaltile-std/src/mlx/steel/attn/mod.rs rename to crates/metaltile-std/src/kernels/sdpa/steel_attn/mod.rs diff --git a/crates/metaltile-std/src/mlx/steel/attn/steel_attention.rs b/crates/metaltile-std/src/kernels/sdpa/steel_attn/steel_attention.rs similarity index 99% rename from crates/metaltile-std/src/mlx/steel/attn/steel_attention.rs rename to crates/metaltile-std/src/kernels/sdpa/steel_attn/steel_attention.rs index c795dc20..447c8c62 100644 --- a/crates/metaltile-std/src/mlx/steel/attn/steel_attention.rs +++ b/crates/metaltile-std/src/kernels/sdpa/steel_attn/steel_attention.rs @@ -344,7 +344,7 @@ pub mod kernel_benches { use super::mt_sdpa_prefill; use crate::{ - mlx::steel::attn::{ + kernels::sdpa::steel_attn::{ steel_attention_mma::mt_sdpa_prefill_mma, steel_attention_mma_bf16::mt_sdpa_prefill_mma_bf16, }, @@ -533,7 +533,7 @@ pub mod kernel_tests { use super::mt_sdpa_prefill; use crate::{ - mlx::steel::attn::{ + kernels::sdpa::steel_attn::{ steel_attention_mma::mt_sdpa_prefill_mma, steel_attention_mma_bf16::mt_sdpa_prefill_mma_bf16, }, diff --git a/crates/metaltile-std/src/mlx/steel/attn/steel_attention_mma.rs b/crates/metaltile-std/src/kernels/sdpa/steel_attn/steel_attention_mma.rs similarity index 99% rename from crates/metaltile-std/src/mlx/steel/attn/steel_attention_mma.rs rename to crates/metaltile-std/src/kernels/sdpa/steel_attn/steel_attention_mma.rs index cfcb7f96..63856764 100644 --- a/crates/metaltile-std/src/mlx/steel/attn/steel_attention_mma.rs +++ b/crates/metaltile-std/src/kernels/sdpa/steel_attn/steel_attention_mma.rs @@ -713,7 +713,7 @@ pub mod kernel_tests { use super::mt_sdpa_prefill_mma; use crate::{ - mlx::steel::attn::steel_attention_mma_bf16::mt_sdpa_prefill_mma_bf16, + kernels::sdpa::steel_attn::steel_attention_mma_bf16::mt_sdpa_prefill_mma_bf16, utils::{pack_f32, unpack_f32}, }; diff --git a/crates/metaltile-std/src/mlx/steel/attn/steel_attention_mma_bf16.rs b/crates/metaltile-std/src/kernels/sdpa/steel_attn/steel_attention_mma_bf16.rs similarity index 100% rename from crates/metaltile-std/src/mlx/steel/attn/steel_attention_mma_bf16.rs rename to crates/metaltile-std/src/kernels/sdpa/steel_attn/steel_attention_mma_bf16.rs diff --git a/crates/metaltile-std/src/mlx/steel/attn/steel_attention_nax.rs b/crates/metaltile-std/src/kernels/sdpa/steel_attn/steel_attention_nax.rs similarity index 100% rename from crates/metaltile-std/src/mlx/steel/attn/steel_attention_nax.rs rename to crates/metaltile-std/src/kernels/sdpa/steel_attn/steel_attention_nax.rs diff --git a/crates/metaltile-std/src/mlx/mod.rs b/crates/metaltile-std/src/mlx/mod.rs deleted file mode 100644 index 04e8c58a..00000000 --- a/crates/metaltile-std/src/mlx/mod.rs +++ /dev/null @@ -1,27 +0,0 @@ -//! Copyright 2026 0xClandestine, Ekryski, TheTom, Ambisphaeric -//! SPDX-License-Identifier: Apache-2.0 -//! MLX-compared kernels. -//! -//! Every kernel in this submodule has (or can have) a side-by-side -//! correctness/perf comparison against an MLX reference kernel — the -//! benches embed MLX's `.metal` source via `metal_file = "..."` and -//! dispatch the MLX kernel through `compile_with_bool_constants` / a -//! constructed kernel name. -//! -//! When a kernel can't be directly compared today (MLX template not -//! shipped at the pinned commit, or the comparison isn't wired yet) -//! but the implementation faithfully mirrors MLX semantics and is -//! expected to wire up eventually, it lives in `ffai/` until the -//! comparison lands. - -// block_scaled_dequant → quant family (migrated later); block_scaled_moe → -// kernels/moe/. The quantized matmuls moved to kernels/gemm/. -pub mod steel; - -// `conv.rs` and `shared.rs` are placeholder/stale stubs left over from -// the old `metaltile-bench` crate. They reference `crate::runner` which -// lives in `metaltile-cli`, so they don't compile — kept on disk for -// the kernel docs / future-work notes but intentionally not declared -// here. Delete or port when those kernels land in the #[kernel] DSL. -// `fft.rs` and `fence.rs` are now real `#[kernel]` ports (declared -// above / in `ffai/`). diff --git a/crates/metaltile-std/src/mlx/steel/conv/mod.rs b/crates/metaltile-std/src/mlx/steel/conv/mod.rs deleted file mode 100644 index b694da51..00000000 --- a/crates/metaltile-std/src/mlx/steel/conv/mod.rs +++ /dev/null @@ -1,5 +0,0 @@ -//! Copyright 2026 0xClandestine, Ekryski, TheTom, Ambisphaeric -//! SPDX-License-Identifier: Apache-2.0 -pub mod steel_conv; -pub mod steel_conv_3d; -pub mod steel_conv_general; diff --git a/crates/metaltile-std/src/mlx/steel/conv/steel_conv.rs b/crates/metaltile-std/src/mlx/steel/conv/steel_conv.rs deleted file mode 100644 index c90045cd..00000000 --- a/crates/metaltile-std/src/mlx/steel/conv/steel_conv.rs +++ /dev/null @@ -1,18 +0,0 @@ -//! Copyright 2026 0xClandestine, Ekryski, TheTom, Ambisphaeric -//! SPDX-License-Identifier: Apache-2.0 -//! Steel implicit-GEMM 2D conv — metal/steel/conv/kernels/steel_conv.metal -//! -//! 2D convolution via implicit GEMM (im2col × filter matrix): -//! implicit_gemm_conv_2d_{dtype}_bm{M}_bn{N}_bk{K}_wm{wm}_wn{wn}_channel_{c}_filter_{f} -//! Block shapes: 32×8, 64×8, 32×32, 32×64, 64×32, 64×64 (all ×16 K) -//! Channel variants: l (general), 1/2/3/4 (small fixed channel count) -//! Filter variants: s (small/separable), l (large) -//! Dtypes: float32, float16, bfloat16 -//! -//! NOT YET IMPLEMENTED in #[kernel] DSL: -//! Implicit GEMM convolution unfolds the input patch neighbourhood -//! into a virtual matrix via the `MLXConvParams` descriptor, then -//! runs tiled GEMM over the unfolded layout. This requires simdgroup -//! matrix ops (same blocker as steel_gemm_fused) plus the im2col -//! index arithmetic driven by `ImplicitGemmConv2DParams`. The DSL -//! has neither simdgroup matmul nor im2col primitives. diff --git a/crates/metaltile-std/src/mlx/steel/conv/steel_conv_3d.rs b/crates/metaltile-std/src/mlx/steel/conv/steel_conv_3d.rs deleted file mode 100644 index 9ad494da..00000000 --- a/crates/metaltile-std/src/mlx/steel/conv/steel_conv_3d.rs +++ /dev/null @@ -1,14 +0,0 @@ -//! Copyright 2026 0xClandestine, Ekryski, TheTom, Ambisphaeric -//! SPDX-License-Identifier: Apache-2.0 -//! Steel implicit-GEMM 3D conv — metal/steel/conv/kernels/steel_conv_3d.metal -//! -//! 3D convolution (D×H×W input) via implicit GEMM: -//! implicit_gemm_conv_3d_{dtype}_bm{M}_bn{N}_bk{K}_wm{wm}_wn{wn}_filter_{f} -//! Block shapes: 32×8, 64×8, 32×32, 32×64, 64×32, 64×64 (all ×16 K) -//! Filter variants: s (small), l (large) -//! Dtypes: float32, float16, bfloat16 -//! -//! NOT YET IMPLEMENTED in #[kernel] DSL: -//! Same blockers as steel_conv with the additional 3D volume indexing -//! over `MLXConvParams<3>`. No DSL support for 3D im2col or simdgroup -//! matrix ops. diff --git a/crates/metaltile-std/src/mlx/steel/conv/steel_conv_general.rs b/crates/metaltile-std/src/mlx/steel/conv/steel_conv_general.rs deleted file mode 100644 index 07042b4b..00000000 --- a/crates/metaltile-std/src/mlx/steel/conv/steel_conv_general.rs +++ /dev/null @@ -1,13 +0,0 @@ -//! Copyright 2026 0xClandestine, Ekryski, TheTom, Ambisphaeric -//! SPDX-License-Identifier: Apache-2.0 -//! Steel general implicit-GEMM 2D conv — metal/steel/conv/kernels/steel_conv_general.metal -//! -//! 2D convolution supporting arbitrary strides, dilation, padding, and groups: -//! implicit_gemm_conv_2d_general_{dtype}_bm{M}_bn{N}_bk{K}_wm{wm}_wn{wn} -//! Block shapes: 32×8, 64×8, 32×32, 32×64, 64×32, 64×64 (all ×16 K) -//! Dtypes: float32, float16, bfloat16 -//! -//! NOT YET IMPLEMENTED in #[kernel] DSL: -//! Extends steel_conv with `Conv2DGeneralJumpParams` and -//! `Conv2DGeneralBaseInfo` for non-unit strides/dilation and group -//! convolution. Same simdgroup matrix + im2col blockers as steel_conv. diff --git a/crates/metaltile-std/src/mlx/steel/mod.rs b/crates/metaltile-std/src/mlx/steel/mod.rs deleted file mode 100644 index efe6f405..00000000 --- a/crates/metaltile-std/src/mlx/steel/mod.rs +++ /dev/null @@ -1,4 +0,0 @@ -//! Copyright 2026 0xClandestine, Ekryski, TheTom, Ambisphaeric -//! SPDX-License-Identifier: Apache-2.0 -pub mod attn; -pub use crate::kernels::convolution::steel_conv as conv; diff --git a/crates/metaltile-std/tests/kernel_registry_consistency.rs b/crates/metaltile-std/tests/kernel_registry_consistency.rs index 5018295e..04815ef0 100644 --- a/crates/metaltile-std/tests/kernel_registry_consistency.rs +++ b/crates/metaltile-std/tests/kernel_registry_consistency.rs @@ -15,7 +15,7 @@ //! c. Defined locally in the same emit (template/inline function). //! //! 2. **Empty / unregistered kernel** — every `#[kernel] pub fn name` -//! in `crates/metaltile-std/src/{mlx,ffai}/**` should have a matching +//! in `crates/metaltile-std/src/kernels/**` should have a matching //! `inventory::submit!` for a BenchSpec that references its //! `kernel_ir_for` function. PR #19 silently emptied a kernel body //! via macro refactor; the inverse failure (kernel defined but @@ -41,7 +41,7 @@ use std::{ use metaltile::codegen::{MslGenerator, msl::MslConfig}; // Import the registries via `metaltile_std` (not `metaltile_core`) so the std // rlib — which carries the `#[kernel]` / `#[bench]` / `#[test_kernel]` inventory -// statics in its `ffai` / `mlx` modules — is force-linked into this test binary. +// statics in its `kernels` modules — is force-linked into this test binary. // Importing them from `metaltile_core` would yield empty registries here. use metaltile_std::{all_benches, all_kernels}; @@ -340,9 +340,9 @@ fn kernel_annotations_have_matching_inventory_submit() { // `CARGO_MANIFEST_DIR` points at `crates/metaltile-std`. let manifest_dir = env!("CARGO_MANIFEST_DIR"); let src_root = Path::new(manifest_dir).join("src"); - // Restrict the walk to mlx + ffai (where actual kernels live) to - // avoid scanning spec.rs / run_spec.rs glue files. - let scan_dirs = ["mlx", "ffai"]; + // Restrict the walk to the consolidated `kernels/` tree (where every + // kernel now lives) to avoid scanning spec.rs / run_spec.rs glue files. + let scan_dirs = ["kernels"]; // Track (kernel_name, source_file_path) for the second-pass check // that looks up `kernel_ir: ::kernel_ir_for` in inventory @@ -386,7 +386,7 @@ fn kernel_annotations_have_matching_inventory_submit() { assert!( !annotated.is_empty(), - "no `#[kernel]` annotations found under `src/{{mlx,ffai}}` — walker broken?" + "no `#[kernel]` annotations found under `src/kernels` — walker broken?" ); // Now: for each annotated kernel, it should match a registered diff --git a/crates/metaltile-std/tests/sdpa_prefill_mma_bf16_gpu.rs b/crates/metaltile-std/tests/sdpa_prefill_mma_bf16_gpu.rs index 46231de5..14aecfb1 100644 --- a/crates/metaltile-std/tests/sdpa_prefill_mma_bf16_gpu.rs +++ b/crates/metaltile-std/tests/sdpa_prefill_mma_bf16_gpu.rs @@ -37,7 +37,7 @@ mod common; use common::{Dt, gpu_lock, pack_bytes, unpack_bytes}; use metaltile::Context; -use metaltile_std::mlx::steel::attn::steel_attention_mma_bf16::mt_sdpa_prefill_mma_bf16; +use metaltile_std::kernels::sdpa::steel_attn::steel_attention_mma_bf16::mt_sdpa_prefill_mma_bf16; /// Causal-prefill SDPA reference for the (B, n_q_heads, q_len, head_dim) /// + (B, n_kv_heads, k_len, head_dim) layout the kernel reads. diff --git a/crates/metaltile-std/tests/sdpa_prefill_mma_long_t.rs b/crates/metaltile-std/tests/sdpa_prefill_mma_long_t.rs index cad8b40b..5c73736b 100644 --- a/crates/metaltile-std/tests/sdpa_prefill_mma_long_t.rs +++ b/crates/metaltile-std/tests/sdpa_prefill_mma_long_t.rs @@ -29,7 +29,7 @@ mod common; use common::{Dt, gpu_lock, pack_bytes, ramp, unpack_bytes}; use metaltile::Context; -use metaltile_std::mlx::steel::attn::steel_attention_mma::mt_sdpa_prefill_mma; +use metaltile_std::kernels::sdpa::steel_attn::steel_attention_mma::mt_sdpa_prefill_mma; /// Naive SDPA prefill (single batch, full causal). Q/K/V are /// `[n_heads_or_kv * T * D]` row-major; output is `[n_heads * T * D]`.