• Lei Wang's avatar
    [Feat] Introduce new caching mechanism for compiled kernels (#176) · 7bde63d5
    Lei Wang authored
    * Add kernel caching mechanism to TileLang
    
    - Implement a new `cached` function in `tilelang/cache/__init__.py` to cache and reuse compiled kernels
    - Expose the `cached` function in the main `tilelang/__init__.py`
    - Add a test case for cached matrix multiplication in `testing/python/cache/test_tilelang_cache_matmul.py`
    - Provide a `clear_cache()` function to reset the kernel cache when needed
    
    * Refactor kernel caching test and implementation
    
    - Simplify the `cached` function in `tilelang/cache/__init__.py`
    - Update test script `test_tilelang_cache_matmul.py` to use `tilelang.testing.main()`
    - Remove unnecessary whitespace and improve code formatting
    
    * Update import for `cached` function in MHA examples
    
    - Modify import statement in `example_mha_bwd.py` and `test_tilelang_kernel_mha_bwd.py`
    - Change import from `tilelang.profiler import cached` to `tilelang import cached`
    - Align with recent refactoring of kernel caching mechanism
    
    * Refactor `cached` function signature in kernel caching
    
    - Update function signature to use keyword-only arguments for `target` and `target_host`
    - Improve parameter order and readability of the `cached` decorator
    - Maintain existing functionality while enhancing function definition
    7bde63d5
__init__.py 3.43 KB