[BugFix][Metal] Resolve auto backend to torch and skip disk cache for torch (#2722)
execution_backend="auto" previously resolved to tvm_ffi for metal targets
because it picks the first registered spec and tvm_ffi was registered before
torch, contradicting the documented default (metal->torch). Register torch
first so bare @tilelang.jit runs on Metal via torch.mps.compile_shader; tvm_ffi
remains available as an explicit opt-in.
KernelCache also crashed on the torch backend: MetalKernelAdapter has no libpath
artifact (the MSL source is compiled in-process via torch.mps.compile_shader),
but _save_so_cubin_to_disk unconditionally read kernel.adapter.libpath, logging
"Error during atomic cache save" with an AttributeError on every run and never
persisting a usable entry (the database-load path has no torch branch either).
The torch backend now disables disk caching when it compiles, so the disk cache
is never written for it, and libpath access is guarded so lib-less adapters can
never crash the save path again.
Adds regression tests in testing/python/metal/test_metal_cache.py.
* [Metal] Add cooperative tensor intrinsics
Expose TileLang-owned cooperative tensor builtins so Metal MPP lowering does not depend on extra TVM fork APIs.
* [Metal] Select cooperative tensor GEMM lowering
Add a shape-aware MPP instruction choice for shared-output Metal GEMM while preserving simdgroup fallback for fragments and unsupported tiles.
* [Metal] Emit MPP cooperative tensor shaders
Generate Metal 4 MPP matmul2d code for cooperative tensor intrinsics and keep source-only codegen separate from runtime compilation.
* [Metal] Lower GEMM through MPP tensor ops
Split Metal GEMM lowering into simdgroup and cooperative tensor emitters so M5 tiles use MPP while fragment accumulators keep the existing path.
* [Metal] Guard generic passes for cooperative tensors
Keep generic allocation and storage rewrites away from opaque Metal cooperative tensor scopes to avoid invalid scope analysis.
* [Metal] Test cooperative tensor GEMM coverage
Add runtime and source-only coverage for non-square MPP GEMM so the new cooperative tensor path is reproducible in CI and on M5.
* [Metal] Update TVM Metal 4 runtime guard
Point the submodule at the macOS SDK guarded Metal 4 runtime update used by cooperative tensor shaders.
* [Metal] Document cooperative tensor GEMM
Add a reference page covering the two Metal GEMM paths, selection rules, current limitations, and planned follow-up work.
* lint
* Improve Metal cooperative tensor GEMM
* Optimize Metal cooperative tensor GEMM
* Update TVM Metal 4 support
* Add Metal backend internals documentation
* Harden Metal cooperative tensor lowering
* test older impl with new framework
* Clean up Metal transform exports
* Use macos-latest for Metal CI
* resolve comments and lint
---------
Co-authored-by: SiriusNEO <chaofan@deepseek.com>
* [Metal] Add Metal GEMM support with simdgroup_matrix MMA
Add T.gemm support for Apple Metal using simdgroup_matrix 8x8 operations
(simdgroup_load/store/multiply_accumulate). Works on all Apple Silicon
(M1-M5) without requiring a TVM fork.
Key changes:
- codegen_metal.cc/h: Fork TVM Metal codegen to tilelang with
simdgroup intrinsic emission and 128-bit vectorized copy
- gemm_metal.py: GemmMetal tile operator for sharedxshared GEMM
- metal_macro_generator.py: MPSIntrinEmitter for simdgroup MMA macros
- metal_fragment_to_simdgroup.py: Pass rewrites local.fragment GEMM
accumulators to metal.simdgroup scope before layout inference
- LowerSIMDGroupCopy in copy.cc for fragment->device simdgroup_store
24 Metal tests (codegen cross-platform + correctness on device).
* Move Metal buffer helpers to Metal backend
* fix(metal): remove layout_inference bypass, fix build & clarify docs
- Remove duplicate BuildTileLangMetal in rt_mod_metal.cc (fixes linker error)
- Remove TargetIsMetal bypass in layout_inference.cc. The bypass is
unnecessary because MetalFragmentToSimdgroup already converts all
GEMM accumulator fragment buffers to metal.simdgroup scope before
LayoutInference runs, so IsFragmentBuffer never matches them.
- Clarify docstrings for MetalFragmentToSimdgroup pass explaining why
it runs before LayoutInference (simdgroup matrices are opaque and
have no explicit thread-level layout).
- Restore is_metal_simdgroup check in GemmMetal.lower() so the GEMM
lowering correctly handles buffers already converted to simdgroup
scope by the pass.
* fix: tir->tirx migration build fixes for metal backend
- Replace tir::Buffer with Buffer, tir:: with tirx:: in metal backend
- Update codegen_metal to use AllocBufferNode instead of AllocateNode
- Create runtime/metal/metal_module.h with MetalModuleCreate factory
- Fix ICHECK -> TVM_FFI_ICHECK macros
- Fix 'from tvm import tir' -> 'from tvm import tirx as tir' in Python
- Restore tilelang/backend/__init__.py for metal gemm registration
- Add metal backend import to tilelang/__init__.py
- Fix metal_fragment_to_simdgroup.py: Block->SBlock, Allocate->AllocBuffer
* fix: improve metal_fragment_to_simdgroup buffer remapping
- Add AllocBuffer node handling for standalone allocations
- Remap buffer references in SBlock reads/writes/match_buffers
- Add _remap_buffer_region and _remap_match_buffer helpers
* lint
* fix: add full-body var subst in MetalFragmentToSimdgroup to reach gemm call_intrin args
The per-SBlock substitute() in _rewrite_scope missed Var references
buried inside BufferRegion objects nested in call_intrin arguments (e.g.,
the GEMM op's C buffer). This leaked the old local.fragment-typed C_local
var into the subsequent gemm lowering, producing a duplicate variable that
MakePackedAPI flagged as undefined.
Fix: apply tir.stmt_functor.substitute() to the entire body before the
ir_transform pass, so all Var references — including those inside opaque
Call args — are remapped to the metal.simdgroup-typed replacements.
Also remove the redundant per-SBlock substitute() and the duplicate
_remap_buffer function.
* fix: add missing <cmath> include for std::isinf/isnan in codegen_metal
* [Metal] Fix review findings: dedup warp partition, kernel_only, barrier comment, immutable gemm ops
- Extract shared ComputeSquareWarpPartition to src/backend/metal/op/utils.h,
used by both gemm.cc (Square policy) and copy.cc (simdgroup store lowering).
- Implement kernel_only parameter in MetalKernelAdapter.get_kernel_source.
- Add comment documenting that Metal's simdgroup_barrier synchronizes the
full threadgroup (no per-simdgroup barrier exists).
- Replace module-level mutable global in metal_fragment_to_simdgroup.py
with @lru_cache and remove CUDA-only gemm ops from the set.
* lint
* Refactor MetalFragmentToSimdgroup import path in LowerAndLegalize function
* Remove metal_fragment_to_simdgroup.py file and update import in metal/__init__.py
---------
Co-authored-by: LeiWang1999 <leiwang1999@outlook.com>
* [Enhancement] Update examples and tests for improved type handling and functionality
- Enhanced various example scripts to support new data types and improve compatibility with PyTorch.
- Updated tests across multiple modules to ensure correct functionality with the latest changes in type handling.
- Refactored code in examples to streamline operations and improve clarity, particularly in tensor operations and memory management.
- Added comprehensive tests for new features and fixed existing issues related to type conversions and buffer handling.
* [Refactor] Update accumulation data type to float32 across examples
- Changed accumulation data type from "float" to T.float32 in multiple example scripts to ensure consistency and improve numerical stability.
- This update affects various modules including flash attention, GEMM analysis, convolution, and deepseek MLA examples, enhancing type handling across the board.
* [Refactor] Standardize data type usage across benchmark scripts
- Updated data type definitions in benchmark scripts to use T.float16 and T.float32 consistently, enhancing clarity and type handling.
- Adjusted dtype assignments in matmul functions and configuration setups to align with the new standard.
- Improved overall code consistency and maintainability by ensuring uniform data type usage across various modules.
* [Refactor] Standardize data type usage in templates and scripts
- Updated data type definitions in various templates and scripts to use string representations (e.g., "float16", "int32") instead of T.float16 and T.int32 for improved consistency and clarity.
- Enhanced overall code maintainability by ensuring uniform data type usage across multiple modules, including convolution, elementwise operations, and matrix multiplication templates.
- This change aims to streamline type handling and improve compatibility with existing workflows.
* [Refactor] Standardize data type usage in examples and benchmarks
- Updated data type definitions in various example and benchmark scripts to use T.float16 and T.int32 consistently, enhancing clarity and maintainability.
- Adjusted dtype assignments in kernel functions and configuration setups to align with the new standard.
- Improved overall code consistency by ensuring uniform data type usage across multiple modules, including attention mechanisms, matrix multiplication, and GEMM examples.
* [Refactor] Import dtypes from language.v2 module
- Added import statement for dtypes from the language.v2 module to enhance type handling and maintain consistency across the codebase.
- This change aims to streamline data type management and improve overall code clarity.
* fix
* [Refactor] Standardize data type usage across scripts
- Updated data type definitions in various scripts to use string representations (e.g., "float16", "int8") instead of T.float16 and T.int8 for improved consistency and clarity.
- Adjusted dtype assignments in functions and configuration setups to align with the new standard, enhancing overall code maintainability.
- This change affects multiple modules, including benchmark and attention mechanisms, ensuring uniform data type usage throughout the codebase.
* [Refactor] Update data type handling for consistency and clarity
- Changed string representations of data types in the Hint class to use T.float32 and T.int32 for improved consistency.
- Added new data types "int4" and "int16" to the dtypes module, enhancing type support across the codebase.
- Updated function signatures and assertions in the lop3 and mxfp modules to utilize the new data types, ensuring uniformity in type handling.
- This refactor aims to streamline data type management and improve overall code clarity and maintainability.
* [Enhancement] Improve data type handling and error messaging
- Introduced a mapping for canonical data types to their display strings, enhancing clarity in type representation.
- Updated the dtype creation logic to utilize the new mapping, ensuring more intuitive handling of string inputs.
- Refined error messages in the lop3 module to provide clearer feedback on invalid source formats, improving debugging and user experience.
* [Fix] Correct boolean flag in GEMM SP test case
- Updated the boolean flag in the test_gemm_sp_sm90 function to ensure proper functionality in the test case.
- This change enhances the accuracy of the test and aligns it with expected behavior for the GEMM SP implementation.
* [Refactor] Standardize data type usage across scripts
- Updated data type definitions in various scripts to use T.float16 and T.bfloat16 consistently, enhancing clarity and maintainability.
- Adjusted dtype assignments in function signatures and argument parsing to align with the new standard, ensuring uniform data type usage throughout the codebase.
- This change affects multiple modules, including benchmarks and examples, improving overall code consistency and readability.
* [Refactor] Standardize data type usage in various modules
- Updated data type assignments in multiple scripts to utilize T.float32, T.int8, and T.int32 consistently, enhancing clarity and maintainability.
- Adjusted function signatures and parameter types across benchmarks, examples, and tests to align with the new standard, ensuring uniform data type usage throughout the codebase.
- This change improves overall code consistency and readability, impacting modules related to matrix multiplication, GEMM, and tensor operations.
* [Refactor] Update argument parsing for data types in benchmarks
- Changed argument parsing for data types in benchmark_matmul_intrinsic.py and benchmark_matmul_sp.py to use string representations ("float16", "int8", "float") instead of T.float16 and T.float.
- This update enhances consistency in data type handling across benchmark scripts, improving clarity and maintainability.
* [Refactor] Update data type handling in benchmark and example scripts
- Changed data type arguments in benchmark and example scripts to use string representations ("float16") instead of T.float16 for improved consistency.
- Updated function signatures and argument parsing to align with the new standard, enhancing clarity and maintainability across the codebase.
- This change affects multiple modules related to attention mechanisms and tensor operations, ensuring uniform data type usage throughout the examples.
* [Refactor] Fix data type conversion in multiple scripts
- Corrected the usage of the data type conversion method from dtype..as_torch() to dtype.as_torch() across various benchmark and example scripts.
- This change enhances consistency in data type handling and improves code readability, impacting modules related to attention mechanisms and tensor operations.
* [Refactor] Update float8 data type usage across multiple scripts
- Changed instances of T.float8_e4m3 to T.float8_e4m3fn in various benchmark, example, and test scripts to ensure consistency in data type handling.
- This update enhances clarity and maintainability across the codebase, particularly in modules related to matrix multiplication and tensor operations.
* [Refactor] Enhance float8 data type handling in CUDA code generation
- Updated the handling of float8 data types in the CUDA code generation to include additional float8 variants, improving type conversion logic.
- Adjusted conditions to ensure proper type checks for float8 conversions, enhancing clarity and maintainability in the codebase.
- Modified layout inference to streamline float8 type checks, ensuring consistency across the implementation.
- This change impacts modules related to matrix operations and CUDA code generation, improving overall type handling and conversion accuracy.
* [Refactor] Streamline float8 data type handling in CUDA and related modules
- Enhanced float8 data type handling in CUDA code generation by refining type conversion logic and ensuring consistent type checks.
- Updated layout inference for float8 types to improve clarity and maintainability across the implementation.
- This change impacts modules related to matrix operations and CUDA code generation, improving overall type handling and conversion accuracy.
* [Refactor] Remove unnecessary cache disabling in float8 example script
- Eliminated the call to tilelang.disable_cache() in example_group_per_split_token_cast_to_fp8.py to streamline the code.
- This change enhances clarity and maintainability of the example script without affecting its functionality.
* [Refactor] Update data type usage in debug print tests
- Changed the argument for dtype in the test_debug_print_buffer function from a string representation to the corresponding T.bool type.
- This update enhances consistency in data type handling within the test suite, improving clarity and maintainability.
* lint fix
* Update function parameter types from `str` to `T.dtype` for improved type safety in attention sink and related examples
* Refactor `gemv_alloc_reducer` function signature for improved readability by formatting parameters across multiple lines.
Yichen Yan