mirror of
https://github.com/tile-ai/tilelang.git
synced 2026-10-02 06:34:36 +08:00
[Misc] configure pymarkdown (#1505)
This commit is contained in:
Gabriel Wu
@@ -54,3 +54,8 @@ repos:
|
||||
^.+\.svg$|
|
||||
^.*\brequirements\b.*\.txt$
|
||||
)
|
||||
- repo: https://github.com/jackdewinter/pymarkdown
|
||||
rev: v0.9.33
|
||||
hooks:
|
||||
- id: pymarkdown
|
||||
args: ["--config", ".pymarkdown", "fix"]
|
||||
|
||||
+37
@@ -0,0 +1,37 @@
|
||||
{
|
||||
"plugins": {
|
||||
"md003": {
|
||||
"style": "atx"
|
||||
},
|
||||
"md004": {
|
||||
"style": "dash"
|
||||
},
|
||||
"md013": {
|
||||
"enabled": false
|
||||
},
|
||||
"md026": {
|
||||
"enabled": false
|
||||
},
|
||||
"md029": {
|
||||
"enabled": false
|
||||
},
|
||||
"md031": {
|
||||
"enabled": false
|
||||
},
|
||||
"md032": {
|
||||
"enabled": false
|
||||
},
|
||||
"md033": {
|
||||
"enabled": false
|
||||
},
|
||||
"md034": {
|
||||
"enabled": false
|
||||
},
|
||||
"md040": {
|
||||
"enabled": false
|
||||
},
|
||||
"md041": {
|
||||
"enabled": false
|
||||
}
|
||||
}
|
||||
}
|
||||
+10
-10
@@ -17,23 +17,23 @@ diverse, inclusive, and healthy community.
|
||||
Examples of behavior that contributes to a positive environment for our
|
||||
community include:
|
||||
|
||||
* Demonstrating empathy and kindness toward other people
|
||||
* Being respectful of differing opinions, viewpoints, and experiences
|
||||
* Giving and gracefully accepting constructive feedback
|
||||
* Accepting responsibility and apologizing to those affected by our mistakes,
|
||||
- Demonstrating empathy and kindness toward other people
|
||||
- Being respectful of differing opinions, viewpoints, and experiences
|
||||
- Giving and gracefully accepting constructive feedback
|
||||
- Accepting responsibility and apologizing to those affected by our mistakes,
|
||||
and learning from the experience
|
||||
* Focusing on what is best not just for us as individuals, but for the overall
|
||||
- Focusing on what is best not just for us as individuals, but for the overall
|
||||
community
|
||||
|
||||
Examples of unacceptable behavior include:
|
||||
|
||||
* The use of sexualized language or imagery, and sexual attention or advances of
|
||||
- The use of sexualized language or imagery, and sexual attention or advances of
|
||||
any kind
|
||||
* Trolling, insulting or derogatory comments, and personal or political attacks
|
||||
* Public or private harassment
|
||||
* Publishing others' private information, such as a physical or email address,
|
||||
- Trolling, insulting or derogatory comments, and personal or political attacks
|
||||
- Public or private harassment
|
||||
- Publishing others' private information, such as a physical or email address,
|
||||
without their explicit permission
|
||||
* Other conduct which could reasonably be considered inappropriate in a
|
||||
- Other conduct which could reasonably be considered inappropriate in a
|
||||
professional setting
|
||||
|
||||
## Enforcement Responsibilities
|
||||
|
||||
+1
-1
@@ -2,7 +2,7 @@
|
||||
|
||||
That would be awesome if you want to contribute something to TileLang!
|
||||
|
||||
### Table of Contents <!-- omit in toc --> <!-- markdownlint-disable heading-increment -->
|
||||
## Table of Contents <!-- omit in toc --> <!-- markdownlint-disable heading-increment -->
|
||||
|
||||
- [Report Bugs](#report-bugs)
|
||||
- [Ask Questions](#ask-questions)
|
||||
|
||||
@@ -24,7 +24,7 @@ Check out the preview here:
|
||||
This includes implementations across two branches:
|
||||
[ascendc_pto](https://github.com/tile-ai/tilelang-ascend) and
|
||||
[npuir](https://github.com/tile-ai/tilelang-ascend/tree/npuir).
|
||||
Feel free to explore and share your feedback!
|
||||
Feel free to explore and share your feedback!
|
||||
- 07/04/2025 🚀: Introduced `T.gemm_sp` for 2:4 sparse tensor core support, check out [Pull Request #526](https://github.com/tile-ai/tilelang/pull/526) for details.
|
||||
- 06/05/2025 ✨: Added [NVRTC Backend](https://github.com/tile-ai/tilelang/pull/461) to significantly reduce compilation time for cute templates!
|
||||
- 04/14/2025 🚀: Added high-performance FlashMLA implementation for AMD MI300X, achieving performance parity with hand-optimized assembly kernels of Aiter! See [example_mla_amd](./examples/deepseek_mla/amd/README.md) for details.
|
||||
@@ -49,7 +49,6 @@ Although tile-lang aims to be portable across a range of Devices, it has been sp
|
||||
|
||||
Within the `examples` directory, you will also find additional complex kernels—such as convolutions, forward/backward passes for FlashAttention, more operators will continuously be added.
|
||||
|
||||
|
||||
## Benchmark Summary
|
||||
|
||||
TileLang achieves exceptional performance across a variety of computational patterns. Comprehensive benchmark scripts and settings are available at [tilelang-benchmark](https://github.com/tile-ai/tilelang-benchmark). Below are selected results showcasing its capabilities:
|
||||
@@ -109,9 +108,9 @@ pip install -e . -v # remove -e option if you don't want to install in editable
|
||||
|
||||
### Method 2: Build from Source
|
||||
We currently provide three ways to install **tile-lang** from source:
|
||||
- [Install from Source (using your own TVM installation)](./docs/get_started/Installation.md#method-1-install-from-source-using-your-own-tvm-installation)
|
||||
- [Install from Source (using the bundled TVM submodule)](./docs/get_started/Installation.md#method-2-install-from-source-using-the-bundled-tvm-submodule)
|
||||
- [Install Using the Provided Script](./docs/get_started/Installation.md#method-3-install-using-the-provided-script)
|
||||
- [Install from Source (using your own TVM installation)](./docs/get_started/Installation.md#method-1-install-from-source-using-your-own-tvm-installation)
|
||||
- [Install from Source (using the bundled TVM submodule)](./docs/get_started/Installation.md#method-2-install-from-source-using-the-bundled-tvm-submodule)
|
||||
- [Install Using the Provided Script](./docs/get_started/Installation.md#method-3-install-using-the-provided-script)
|
||||
|
||||
### Method 3: Install with Nightly Version
|
||||
|
||||
|
||||
@@ -45,7 +45,6 @@ PY
|
||||
| 16384 | 2.531 | 135.711 |
|
||||
| 32768 | 5.076 | 135.379 |
|
||||
|
||||
|
||||
## Compare with Baselines
|
||||
|
||||
- Triton: v3.5.0, mamba-ssm: v2.2.6.post3
|
||||
@@ -56,4 +55,4 @@ PY
|
||||
<img src="mamba_benchmark_result.png" alt="Mamba2_chunk_scan Performance Comparison on H100">
|
||||
</a>
|
||||
<figcaption style="text-align: center;">Performance comparison across compilers on NVIDIA H100</figcaption>
|
||||
</figure>
|
||||
</figure>
|
||||
|
||||
+1
-1
@@ -27,4 +27,4 @@ cd _build/html
|
||||
python3 -m http.server
|
||||
```
|
||||
|
||||
Then you can view the documentation in your browser at `http://localhost:8000` (the port can be customized by appending ` -p PORT_NUMBER` in the python command above).
|
||||
Then you can view the documentation in your browser at `http://localhost:8000` (the port can be customized by appending `-p PORT_NUMBER` in the python command above).
|
||||
|
||||
@@ -384,4 +384,3 @@ Fix: pass correct scalar types, e.g., `scalar_check(1, True)`.
|
||||
## Closing Notes
|
||||
- Cross-check “shape / strides / device / dtype” against the kernel signature to localize issues efficiently.
|
||||
- For complex symbolic relations, print the host source to confirm binding/solving order, then adjust runtime shapes/layouts accordingly.
|
||||
|
||||
|
||||
@@ -1,8 +1,7 @@
|
||||
# 🚀 Write High Performance FlashMLA with TileLang on Hopper
|
||||
|
||||
|
||||
<div style="text-align: left;">
|
||||
<em>Author:</em> <a href="https://github.com/chengyupku">Yu Cheng</a>
|
||||
<em>Author:</em> <a href="https://github.com/chengyupku">Yu Cheng</a>
|
||||
<em>Author:</em> <a href="https://github.com/LeiWang1999">Lei Wang</a>
|
||||
</div>
|
||||
|
||||
@@ -32,7 +31,7 @@ Figure 1: Performance under batch size=64
|
||||
Figure 2: Performance under batch size=128
|
||||
```
|
||||
|
||||
As shown in the results, TileLang achieves performance comparable to FlashMLA in most cases, significantly outperforming both FlashInfer and Triton.
|
||||
As shown in the results, TileLang achieves performance comparable to FlashMLA in most cases, significantly outperforming both FlashInfer and Triton.
|
||||
Notably, **TileLang accomplishes this with just around 80 lines of Python code**, demonstrating its exceptional ease of use and efficiency. Let's dive in and see how TileLang achieves this.
|
||||
|
||||
## Implementation
|
||||
@@ -62,7 +61,7 @@ Compared to traditional attention operators like MHA (Multi-Headed Attention) or
|
||||
|
||||
This raises the question of how to partition the matrix multiplication operation. On the Hopper architecture, most computation kernels use [`wgmma.mma_async`](https://docs.nvidia.com/cuda/parallel-thread-execution/#asynchronous-warpgroup-level-matrix-instructions) instructions for optimal performance. The `wgmma.mma_async` instruction organizes 4 warps (128 threads) into a warpgroup for collective MMA operations. However, `wgmma.mma_async` instructions require a minimum M dimension of 64. This means each warpgroup's minimum M dimension can only be reduced to 64, but a tile size of 64*512 is too large for a single warpgroup, leading to register spilling.
|
||||
|
||||
Therefore, our only option is to partition `acc_o` along the `dim` dimension, with two warpgroups computing the left and right part of `acc_o` respectively. However, this introduces another challenge: both warpgroups require the complete `acc_s` result as input.
|
||||
Therefore, our only option is to partition `acc_o` along the `dim` dimension, with two warpgroups computing the left and right part of `acc_o` respectively. However, this introduces another challenge: both warpgroups require the complete `acc_s` result as input.
|
||||
|
||||
Our solution is to have each warpgroup compute half of `acc_s` during `Q @ K` computation, then obtain the other half computed by the other warpgroup through shared memory.
|
||||
|
||||
@@ -106,7 +105,6 @@ T.use_swizzle(panel_size: int, order: str = "row")
|
||||
|
||||
Here, `panel_size` specifies the width of the swizzled threadblock group, and `order` determines the swizzling pattern, which can be either "row" or "col".
|
||||
|
||||
|
||||
### Shared Memory Swizzling
|
||||
|
||||
In CUDA programming, shared memory is divided into multiple memory banks, with each bank capable of servicing one thread request per clock cycle in parallel. Bank conflicts occur when multiple threads simultaneously access different addresses mapped to the same bank, forcing these accesses to be serialized and degrading performance.
|
||||
@@ -123,17 +121,14 @@ T.annotate_layout({
|
||||
|
||||
Here, `T.annotate_layout` allows users to specify any desired layout for a buffer. For convenience, TileLang provides the `make_swizzled_layout` primitive to automatically generate a swizzled layout.
|
||||
|
||||
|
||||
### Warp-Specialization
|
||||
|
||||
The Hopper architecture commonly employs warp specialization for performance optimization. A typical approach is to designate one warpgroup as a producer that handles data movement using TMA (Tensor Memory Accelerator), while the remaining warpgroups serve as consumers performing computations. However, this programming pattern is complex, requiring developers to manually manage the execution logic for producers and consumers, including synchronization through the `mbarrier` objects.
|
||||
|
||||
In TileLang, users are completely shielded from these implementation details. The frontend script is automatically transformed into a warp-specialized form, where TileLang handles all producer-consumer synchronization automatically, enabling efficient computation.
|
||||
|
||||
|
||||
### Pipeline
|
||||
|
||||
|
||||
Pipeline is a technique used to improve memory access efficiency by overlapping memory access and computation. In TileLang, pipeline can be implemented through the `T.pipelined` annotation:
|
||||
|
||||
```python
|
||||
@@ -142,14 +137,12 @@ T.pipelined(range: int, stage: int)
|
||||
|
||||
Here, `range` specifies the range of the pipeline, and `stage` specifies the stage of the pipeline. Multi-stage pipelining enables overlapping of computation and memory access, which can significantly improve performance for memory-intensive operators. However, setting a higher number of stages consumes more shared memory resources, so the optimal configuration needs to be determined based on specific use cases.
|
||||
|
||||
|
||||
### Split-KV
|
||||
|
||||
We have also implemented Split-KV optimization similar to [FlashDecoding](https://pytorch.org/blog/flash-decoding/). Specifically, when the batch size is small, parallel SM resources cannot be fully utilized due to low parallelism. In such cases, we can split the kv_ctx dimension across multiple SMs for parallel computation and then merge the results.
|
||||
|
||||
In our implementation, we have developed both split and combine kernels, allowing users to control the split size through a `num_split` parameter.
|
||||
|
||||
|
||||
## 🚀 On AMD MI300X Accelerators
|
||||
|
||||
Following our previous demonstration of [high-performance FlashMLA implementation on NVIDIA Hopper architectures using TileLang](https://github.com/tile-ai/tilelang/blob/main/examples/deepseek_mla/README.md), this work presents an optimized implementation for AMD MI300X accelerators. We examine architectural differences and corresponding optimization strategies between these platforms.
|
||||
|
||||
@@ -280,8 +280,8 @@ To evaluate complexity, one could implement the same elementwise addition operat
|
||||
|
||||
```c++
|
||||
template<int NUM_ELE_PER_THREAD=8>
|
||||
__global__ void elementwise_add(nv_bfloat16* C,
|
||||
const nv_bfloat16* A,
|
||||
__global__ void elementwise_add(nv_bfloat16* C,
|
||||
const nv_bfloat16* A,
|
||||
const nv_bfloat16* B,
|
||||
int N) {
|
||||
using namespace cute;
|
||||
|
||||
@@ -206,7 +206,6 @@ def splitk_gemv(
|
||||
return main
|
||||
```
|
||||
|
||||
|
||||
## Vectorized Reads
|
||||
|
||||
GEMV is less computation intensive than GEMM as the computation intensity and memory throughput will be the optimization bottleneck. One effective strategy is to use vectorized load/store operations (e.g., `float2`, `float4`). In `TileLang`, you can specify vectorized operations via `T.vectorized`:
|
||||
@@ -254,7 +253,6 @@ def splitk_gemv_vectorized(
|
||||
|
||||
With vectorized read, now the kernel finishes in **~0.0084 ms**, which is getting close to cuBLAS performance.
|
||||
|
||||
|
||||
## `tvm_thread_allreduce` Instead of `atomicAdd`
|
||||
|
||||
[`tvm_thread_allreduce`](https://tvm.apache.org/docs/reference/api/python/tir/tir.html#tvm.tir.tvm_thread_allreduce) has implemented optimization when making an all-reduce across a number of threads, which should outperfrom out plain smem + `atomidAdd`:
|
||||
@@ -459,6 +457,5 @@ This corresponds closely to our `TileLang` program, with necessary synchronizati
|
||||
| splitk_gemv_vectorized | 0.00809 ms |
|
||||
| splitk_gemv_vectorized_tvm | 0.00675 ms |
|
||||
|
||||
|
||||
Triton Time: 0.0077344514429569244
|
||||
In this tutorial, we implemented a simple GEMV kernel and learn that `TileLang` exposes low level control to user such as thread-level programming and CUDA primitives.
|
||||
In this tutorial, we implemented a simple GEMV kernel and learn that `TileLang` exposes low level control to user such as thread-level programming and CUDA primitives.
|
||||
|
||||
@@ -14,11 +14,11 @@
|
||||
|
||||
TileLang is a domain-specific language (DSL) designed for writing high-performance GPU kernels. It provides three main levels of abstraction:
|
||||
|
||||
* **Level 1:** A user writes pure compute logic without knowledge of or concern for hardware details (e.g., GPU caches, tiling, etc.). The compiler or runtime performs automatic scheduling and optimization. This level is conceptually similar to the idea behind TVM.
|
||||
- **Level 1:** A user writes pure compute logic without knowledge of or concern for hardware details (e.g., GPU caches, tiling, etc.). The compiler or runtime performs automatic scheduling and optimization. This level is conceptually similar to the idea behind TVM.
|
||||
|
||||
* **Level 2:** A user is aware of GPU architecture concepts—such as shared memory, tiling, and thread blocks—but does not necessarily want to drop down to the lowest level of explicit thread control. This mode is somewhat comparable to Triton's programming model, where you can write tile-level operations and let the compiler do layout inference, pipelining, etc.
|
||||
- **Level 2:** A user is aware of GPU architecture concepts—such as shared memory, tiling, and thread blocks—but does not necessarily want to drop down to the lowest level of explicit thread control. This mode is somewhat comparable to Triton's programming model, where you can write tile-level operations and let the compiler do layout inference, pipelining, etc.
|
||||
|
||||
* **Level 3:** A user takes full control of thread-level primitives and can write code that is almost as explicit as a hand-written CUDA/HIP kernel. This is useful for performance experts who need to manage every detail, such as PTX inline assembly, explicit thread behavior, etc.
|
||||
- **Level 3:** A user takes full control of thread-level primitives and can write code that is almost as explicit as a hand-written CUDA/HIP kernel. This is useful for performance experts who need to manage every detail, such as PTX inline assembly, explicit thread behavior, etc.
|
||||
|
||||
```{figure} ../_static/img/overview.png
|
||||
:width: 50%
|
||||
@@ -52,12 +52,12 @@ While Level 1 in TileLang can be very comfortable for general users—since it r
|
||||
|
||||
Below is a simplified code snippet for a 1024 x 1024 x 1024 matrix multiplication. It uses:
|
||||
|
||||
* **`T.Kernel(...)`** to initialize the thread block configuration (grid dimensions, block size, etc.).
|
||||
* **`T.alloc_shared(...)`** to allocate GPU shared memory.
|
||||
* **`T.alloc_fragment(...)`** to allocate a register fragment for accumulation.
|
||||
* **`T.Pipelined(...)`** to express software pipelining across the K dimension.
|
||||
* **`T.Parallel(...)`** to parallelize data copy loops.
|
||||
* **`T.gemm(...)`** to perform tile-level GEMM operations (which map to the appropriate backends, such as MMA instructions on NVIDIA GPUs).
|
||||
- **`T.Kernel(...)`** to initialize the thread block configuration (grid dimensions, block size, etc.).
|
||||
- **`T.alloc_shared(...)`** to allocate GPU shared memory.
|
||||
- **`T.alloc_fragment(...)`** to allocate a register fragment for accumulation.
|
||||
- **`T.Pipelined(...)`** to express software pipelining across the K dimension.
|
||||
- **`T.Parallel(...)`** to parallelize data copy loops.
|
||||
- **`T.gemm(...)`** to perform tile-level GEMM operations (which map to the appropriate backends, such as MMA instructions on NVIDIA GPUs).
|
||||
|
||||
```python
|
||||
import tilelang
|
||||
@@ -147,14 +147,12 @@ with T.Kernel(T.ceildiv(N, block_N), T.ceildiv(M, block_M), threads=128) as (bx,
|
||||
- This sets up the block grid dimensions based on N/block_N and M/block_M.
|
||||
- `threads=128` specifies that each thread block uses 128 threads. The compiler will infer how loops map to these threads.
|
||||
|
||||
|
||||
```{figure} ../_static/img/Parallel.png
|
||||
:alt: Parallel
|
||||
:align: center
|
||||
|
||||
```
|
||||
|
||||
|
||||
2. **Shared & Fragment Memory**:
|
||||
|
||||
```python
|
||||
@@ -182,7 +180,6 @@ for ko in T.Pipelined(T.ceildiv(K, block_K), num_stages=3):
|
||||
|
||||
```
|
||||
|
||||
|
||||
4. **Parallel Copy**:
|
||||
|
||||
```python
|
||||
@@ -252,8 +249,8 @@ For more advanced usage—including partial lowering, explicitly controlling thr
|
||||
|
||||
## Further Resources
|
||||
|
||||
* [TileLang GitHub](https://github.com/tile-ai/tilelang)
|
||||
* [BitBLAS](https://github.com/tile-ai/bitblas)
|
||||
* [Triton](https://github.com/openai/triton)
|
||||
* [Cutlass](https://github.com/NVIDIA/cutlass)
|
||||
* [PyCUDA](https://documen.tician.de/pycuda/) <!-- codespell:ignore -->
|
||||
- [TileLang GitHub](https://github.com/tile-ai/tilelang)
|
||||
- [BitBLAS](https://github.com/tile-ai/bitblas)
|
||||
- [Triton](https://github.com/openai/triton)
|
||||
- [Cutlass](https://github.com/NVIDIA/cutlass)
|
||||
- [PyCUDA](https://documen.tician.de/pycuda/) <!-- codespell:ignore -->
|
||||
|
||||
@@ -51,7 +51,6 @@ Here, `A_sparse` contains all the non-zero elements of `A`, while `E` stores the
|
||||
The metadata is reordered internally to optimize memory access patterns (e.g., for ldsm instructions and vectorized loads).
|
||||
For more information, see **A note on `gemm_sp` and `gemm_sp_v2`**.
|
||||
|
||||
|
||||
## `T.gemm_sp` with CUTLASS's compressor
|
||||
|
||||
:::{warning}
|
||||
@@ -259,4 +258,4 @@ However, fixing a specific layout introduces several potential issues:
|
||||
|
||||
3. Alignment requirements: `CUTLASS` enforces strict alignment checks, and many hyperparameter configurations can lead to compilation errors. (For reference, sm8x was implemented in `CUTLASS 2`.)
|
||||
|
||||
`T.gemm_sp_v2` was designed to address these limitations, following the approach of `T.gemm_v2`. It lowers directly to PTX, removing the need for a fixed metadata layout.
|
||||
`T.gemm_sp_v2` was designed to address these limitations, following the approach of `T.gemm_v2`. It lowers directly to PTX, removing the need for a fixed metadata layout.
|
||||
|
||||
@@ -51,13 +51,13 @@ Figure 1: High-level overview of the TileLang compilation flow.
|
||||
|
||||
## Tile-based Programming Model
|
||||
|
||||
[Figure 2](#fig-overview-gemm) provides a concise matrix multiplication (GEMM) example in ``TileLang``,
|
||||
illustrating how developers can employ high-level constructs such as tiles, memory placement, pipelining,
|
||||
[Figure 2](#fig-overview-gemm) provides a concise matrix multiplication (GEMM) example in ``TileLang``,
|
||||
illustrating how developers can employ high-level constructs such as tiles, memory placement, pipelining,
|
||||
and operator calls to manage data movement and computation with fine-grained control.
|
||||
In particular, this snippet ([Figure 2](#fig-overview-gemm) (a)) demonstrates how multi-level tiling
|
||||
leverages different memory hierarchies (global, shared, and registers) to optimize bandwidth utilization
|
||||
In particular, this snippet ([Figure 2](#fig-overview-gemm) (a)) demonstrates how multi-level tiling
|
||||
leverages different memory hierarchies (global, shared, and registers) to optimize bandwidth utilization
|
||||
and reduce latency.
|
||||
Overall, [Figure 2](#fig-overview-gemm) (b) showcases how the Python-like syntax of ``TileLang``
|
||||
Overall, [Figure 2](#fig-overview-gemm) (b) showcases how the Python-like syntax of ``TileLang``
|
||||
allows developers to reason about performance-critical optimizations within a user-friendly programming model.
|
||||
|
||||
```{figure} ../_static/img/MatmulExample.png
|
||||
|
||||
@@ -17,7 +17,6 @@ get_started/overview
|
||||
get_started/targets
|
||||
:::
|
||||
|
||||
|
||||
:::{toctree}
|
||||
:maxdepth: 1
|
||||
:caption: TUTORIALS
|
||||
|
||||
@@ -171,8 +171,6 @@ Register control
|
||||
- `T.annotate_producer_reg_dealloc(n=24)`, `T.annotate_consumer_reg_alloc(n=240)`.
|
||||
- `T.no_set_max_nreg()`, `T.disable_warp_group_reg_alloc()`.
|
||||
|
||||
|
||||
|
||||
## Notes on Dtypes
|
||||
|
||||
Dtypes accept three equivalent forms:
|
||||
|
||||
@@ -39,4 +39,3 @@ Notes
|
||||
GEMM with `float16` inputs and `float32` accumulators).
|
||||
- The complete, up‑to‑date list is exposed in
|
||||
`tilelang.language.v2.dtypes` and rendered in the API Reference.
|
||||
|
||||
|
||||
@@ -14,7 +14,7 @@ Auto-tuning a Tile Language program involves three main steps:
|
||||
|
||||
## Matrix Multiplication Example
|
||||
|
||||
The following example demonstrates auto-tuning matrix multiplication. Code has been simplified for readability - see `examples/gemm/example_gemm.py` for complete implementation.
|
||||
The following example demonstrates auto-tuning matrix multiplication. Code has been simplified for readability - see `examples/gemm/example_gemm.py` for complete implementation.
|
||||
|
||||
### Step 1: Implement with Reserved Parameters
|
||||
Users can implement matrix multiplication in Tile Language while reserving parameters for optimization:
|
||||
@@ -145,4 +145,4 @@ for hint in roller_hints:
|
||||
config["thread_num"] = block_rows * block_cols * 32
|
||||
config["enable_rasteration"] = hint.rasterization_plan is not NoRasterization
|
||||
|
||||
```
|
||||
```
|
||||
|
||||
@@ -12,7 +12,6 @@ A Tile Language program (hereafter referred to as a *program*) is transformed in
|
||||
2. The program undergoes multiple *Passes* for transformation and optimization (the *lower* stage, see `tilelang/engine/lower.py`), finally producing an intermediate representation (e.g., LLVM or C for CPU, CUDA for NVIDIA GPUs, etc.).
|
||||
3. The generated code is compiled by the respective compiler (e.g., nvcc) into a hardware-executable file.
|
||||
|
||||
|
||||
```{figure} ../_static/img/overview.png
|
||||
:width: 300
|
||||
:alt: Overview of the compilation process
|
||||
@@ -22,9 +21,9 @@ A Tile Language program (hereafter referred to as a *program*) is transformed in
|
||||
|
||||
During this process, users may encounter roughly three categories of issues:
|
||||
|
||||
* **Generation issues**: The Tile Language program fails to generate a valid hardware-executable file (i.e., errors during the lowering process).
|
||||
* **Correctness issues**: The resulting executable runs, but produces incorrect results.
|
||||
* **Performance issues**: The executable runs with performance significantly below the expected theoretical hardware limits.
|
||||
- **Generation issues**: The Tile Language program fails to generate a valid hardware-executable file (i.e., errors during the lowering process).
|
||||
- **Correctness issues**: The resulting executable runs, but produces incorrect results.
|
||||
- **Performance issues**: The executable runs with performance significantly below the expected theoretical hardware limits.
|
||||
|
||||
This tutorial focuses on the first two issues—how to debug generation and correctness problems. Performance tuning often requires using vendor-provided profiling tools (e.g., **Nsight Compute**, **rocProf**, etc.) for further hardware-level analysis, which we will address in future materials.
|
||||
|
||||
@@ -52,7 +51,6 @@ func = matmul(1024, 1024, 1024, 128, 128, 32)
|
||||
|
||||
TileLang essentially performs *progressive lowering*. For example, a `T.copy` may first be expanded into `T.Parallel` (see the pass `LowerTileOP`), which is then expanded again, eventually resulting in lower-level statements that can be translated to CUDA C code.
|
||||
|
||||
|
||||
```{figure} ../_static/img/ir_transform_diagram.png
|
||||
:width: 400
|
||||
:alt: IR transformation diagram
|
||||
@@ -196,7 +194,6 @@ C_local inferenced layout:
|
||||
Index: [_j % 16 // 8 * 4 + _i % 16 // 8 * 2 + _j % 2]
|
||||
```
|
||||
|
||||
|
||||
## Conclusion
|
||||
|
||||
By carefully examining intermediate representations (IR) before final code generation—and by leveraging runtime printing through `T.print`—one can quickly diagnose where index calculations, copy logic, or other kernel operations deviate from the intended behavior. This two-pronged approach (inspecting IR transformations and using runtime prints) is often sufficient for resolving generation and correctness issues in TileLang programs.
|
||||
|
||||
@@ -25,10 +25,10 @@ VLOG(1) << "aaa";
|
||||
|
||||
- **LOG**: Standard logging preserved in code for displaying necessary information at different levels during runtime. Most Tilelang C++ error reporting is implemented via `LOG(FATAL) << "error msg"`.
|
||||
- **DLOG**: Debug logging for developer debugging output. DLOG is controlled at build time by the TVM_LOG_DEBUG environment variable and is **eliminated in Release builds through dead code elimination**.
|
||||
- The key difference between LOG(DEBUG) and DLOG is this build-time elimination. We recommend using DLOG over LOG(DEBUG), as the latter has overlapping functionality and gets compiled into the release runtime.
|
||||
- The key difference between LOG(DEBUG) and DLOG is this build-time elimination. We recommend using DLOG over LOG(DEBUG), as the latter has overlapping functionality and gets compiled into the release runtime.
|
||||
- **VLOG**: [Verbose logging](https://google.github.io/glog/stable/logging/#verbose-logging), primarily for debugging. Its main feature is customizable verbosity levels. For example, VLOG(n) where n can be 1, 2, 3, 4, 5, or 6, enabling complex tracing requirements. In contrast, LOG and DLOG typically use predefined verbose levels like INFO and DEBUG.
|
||||
- In practical Tilelang development, VLOG is used less frequently.
|
||||
- TVM's VLOG is implemented using DLOG, thus inheriting DLOG's characteristics.
|
||||
- In practical Tilelang development, VLOG is used less frequently.
|
||||
- TVM's VLOG is implemented using DLOG, thus inheriting DLOG's characteristics.
|
||||
|
||||
Additional useful macros include various **CHECK** variants:
|
||||
|
||||
@@ -114,5 +114,3 @@ Then you also need to specify the runtime environment variables. For example, to
|
||||
|
||||
These two should ideally have different names, but TVM uses the same name for both, which can cause confusion.
|
||||
:::
|
||||
|
||||
|
||||
|
||||
@@ -64,10 +64,10 @@ class AnalysisResult:
|
||||
```
|
||||
### `Analyzer` Class Methods
|
||||
#### `analysis(fn, device)`
|
||||
* Parameters:
|
||||
* fn: TVM IRModule or PrimFunc
|
||||
* device: Device configuration object
|
||||
* Returns: AnalysisResult
|
||||
- Parameters:
|
||||
- fn: TVM IRModule or PrimFunc
|
||||
- device: Device configuration object
|
||||
- Returns: AnalysisResult
|
||||
#### Supported Architectures
|
||||
```python
|
||||
# Extendable to custom hardware via: "compute_capability": (cores_per_SM, clock_GHz, flops_per_cycle, max_SM_count)
|
||||
|
||||
@@ -2,7 +2,6 @@
|
||||
|
||||
We compare with an optimized version of the official Triton implementation [here](https://github.com/openai/gpt-oss/blob/main/gpt_oss/triton/attention.py).
|
||||
|
||||
|
||||
## Algorithm
|
||||
### Forward
|
||||
The only change from vanilla FlashAttention is that `sinks` should be taken into consideration in the softmax, which requires an extra rescaling at the epilogue stage.
|
||||
@@ -43,4 +42,4 @@ where $P_{b, h, q}$ is the proportion of $sink_h$ in the softmax in the $b$-th b
|
||||
| 16384 | 64 | 309.46 | **400.62** | 1.29x |
|
||||
| 16384 | 128 | 418.99 | **549.11** | 1.31x |
|
||||
|
||||
> The backward performance will be further optimized in the future.
|
||||
> The backward performance will be further optimized in the future.
|
||||
|
||||
@@ -2,7 +2,6 @@
|
||||
license: mit
|
||||
---
|
||||
|
||||
|
||||
This is a Tilelang Implementation for the reproduced 1.58bit model from [1bitLLM/bitnet_b1_58-3B](https://huggingface.co/1bitLLM/bitnet_b1_58-3B). We replaced the original simulated Int8x3bit Quantized Inference Kernel with INT8xINT2 Kernel. We also evaluated the model's correctness and performance through `eval_correctness.py` and `benchmark_inference_latency.py`.
|
||||
|
||||
## Make Checkpoints for vLLM
|
||||
@@ -43,7 +42,6 @@ python3 inference_with_bitblas_format.py
|
||||
| bitnet-3b-1.58bits | vllm-tilelang | 379.25 | 117.43 | 752.55 |
|
||||
| bitnet-3b-1.58bits | vllm-tilelang-cuda-graph | 2543.58 | 1621.08 | 2731.79 |
|
||||
|
||||
|
||||
## BitBLAS Results
|
||||
|
||||
### Performance
|
||||
@@ -94,4 +92,4 @@ The differences between the reported numbers and the reproduced results are poss
|
||||
journal={arXiv preprint arXiv:2402.17764},
|
||||
year={2024}
|
||||
}
|
||||
```
|
||||
```
|
||||
|
||||
@@ -2,7 +2,6 @@
|
||||
license: mit
|
||||
---
|
||||
|
||||
|
||||
This is a BitBLAS Implementation for the reproduced 1.58bit model from [1bitLLM/bitnet_b1_58-3B](https://huggingface.co/1bitLLM/bitnet_b1_58-3B). We replaced the original simulated Int8x3bit Quantized Inference Kernel with BitBLAS INT8xINT2 Kernel. We also evaluated the model's correctness and performance through `eval_correctness.py` and `benchmark_inference_latency.py`.
|
||||
|
||||
## Latest News
|
||||
@@ -88,4 +87,4 @@ The differences between the reported numbers and the reproduced results are poss
|
||||
journal={arXiv preprint arXiv:2402.17764},
|
||||
year={2024}
|
||||
}
|
||||
```
|
||||
```
|
||||
|
||||
@@ -1,6 +1,5 @@
|
||||
# Block-Sparse Flash-Attention
|
||||
|
||||
Tilelang implementation of block-sparse flash-attention kernels.
|
||||
|
||||
The kernels have been used in [Rectified Sparse Attention](https://arxiv.org/abs/2506.04108) and [SeerAttention-R](https://arxiv.org/abs/2506.08889).
|
||||
Tilelang implementation of block-sparse flash-attention kernels.
|
||||
|
||||
The kernels have been used in [Rectified Sparse Attention](https://arxiv.org/abs/2506.04108) and [SeerAttention-R](https://arxiv.org/abs/2506.08889).
|
||||
|
||||
@@ -24,7 +24,7 @@ We benchmarked the performance of FlashMLA, TileLang, Torch, Triton, and FlashIn
|
||||
<figcaption style="text-align: center;">Figure 2：Performance under batch size=128</figcaption>
|
||||
</figure>
|
||||
|
||||
As shown in the results, TileLang achieves performance comparable to FlashMLA in most cases, significantly outperforming both FlashInfer and Triton.
|
||||
As shown in the results, TileLang achieves performance comparable to FlashMLA in most cases, significantly outperforming both FlashInfer and Triton.
|
||||
Notably, **TileLang accomplishes this with just around 80 lines of Python code**, demonstrating its exceptional ease of use and efficiency. Let's dive in and see how TileLang achieves this.
|
||||
|
||||
## Implementation
|
||||
@@ -54,7 +54,7 @@ Compared to traditional attention operators like MHA (Multi-Headed Attention) or
|
||||
|
||||
This raises the question of how to partition the matrix multiplication operation. On the Hopper architecture, most computation kernels use [`wgmma.mma_async`](https://docs.nvidia.com/cuda/parallel-thread-execution/#asynchronous-warpgroup-level-matrix-instructions) instructions for optimal performance. The `wgmma.mma_async` instruction organizes 4 warps (128 threads) into a warpgroup for collective MMA operations. However, `wgmma.mma_async` instructions require a minimum M dimension of 64. This means each warpgroup's minimum M dimension can only be reduced to 64, but a tile size of 64*512 is too large for a single warpgroup, leading to register spilling.
|
||||
|
||||
Therefore, our only option is to partition `acc_o` along the `dim` dimension, with two warpgroups computing the left and right part of `acc_o` respectively. However, this introduces another challenge: both warpgroups require the complete `acc_s` result as input.
|
||||
Therefore, our only option is to partition `acc_o` along the `dim` dimension, with two warpgroups computing the left and right part of `acc_o` respectively. However, this introduces another challenge: both warpgroups require the complete `acc_s` result as input.
|
||||
|
||||
Our solution is to have each warpgroup compute half of `acc_s` during `Q @ K` computation, then obtain the other half computed by the other warpgroup through shared memory.
|
||||
|
||||
@@ -96,7 +96,6 @@ T.use_swizzle(panel_size: int, order: str = "row")
|
||||
|
||||
Here, `panel_size` specifies the width of the swizzled threadblock group, and `order` determines the swizzling pattern, which can be either "row" or "col".
|
||||
|
||||
|
||||
### Shared Memory Swizzling
|
||||
|
||||
In CUDA programming, shared memory is divided into multiple memory banks, with each bank capable of servicing one thread request per clock cycle in parallel. Bank conflicts occur when multiple threads simultaneously access different addresses mapped to the same bank, forcing these accesses to be serialized and degrading performance.
|
||||
@@ -113,17 +112,14 @@ T.annotate_layout({
|
||||
|
||||
Here, `T.annotate_layout` allows users to specify any desired layout for a buffer. For convenience, TileLang provides the `make_swizzled_layout` primitive to automatically generate a swizzled layout.
|
||||
|
||||
|
||||
### Warp-Specialization
|
||||
|
||||
The Hopper architecture commonly employs warp specialization for performance optimization. A typical approach is to designate one warpgroup as a producer that handles data movement using TMA (Tensor Memory Accelerator), while the remaining warpgroups serve as consumers performing computations. However, this programming pattern is complex, requiring developers to manually manage the execution logic for producers and consumers, including synchronization through the `mbarrier` objects.
|
||||
|
||||
In TileLang, users are completely shielded from these implementation details. The frontend script is automatically transformed into a warp-specialized form, where TileLang handles all producer-consumer synchronization automatically, enabling efficient computation.
|
||||
|
||||
|
||||
### Pipeline
|
||||
|
||||
|
||||
Pipeline is a technique used to improve memory access efficiency by overlapping memory access and computation. In TileLang, pipeline can be implemented through the `T.pipelined` annotation:
|
||||
|
||||
```python
|
||||
@@ -132,9 +128,8 @@ T.pipelined(range: int, stage: int)
|
||||
|
||||
Here, `range` specifies the range of the pipeline, and `stage` specifies the stage of the pipeline. Multi-stage pipelining enables overlapping of computation and memory access, which can significantly improve performance for memory-intensive operators. However, setting a higher number of stages consumes more shared memory resources, so the optimal configuration needs to be determined based on specific use cases.
|
||||
|
||||
|
||||
### Split-KV
|
||||
|
||||
We have also implemented Split-KV optimization similar to [FlashDecoding](https://pytorch.org/blog/flash-decoding/). Specifically, when the batch size is small, parallel SM resources cannot be fully utilized due to low parallelism. In such cases, we can split the kv_ctx dimension across multiple SMs for parallel computation and then merge the results.
|
||||
|
||||
In our implementation, we have developed both split and combine kernels, allowing users to control the split size through a `num_split` parameter.
|
||||
In our implementation, we have developed both split and combine kernels, allowing users to control the split size through a `num_split` parameter.
|
||||
|
||||
@@ -121,7 +121,7 @@ for i_i in T.Pipelined(NI, num_stages=num_stages):
|
||||
# ... compute attention over selected tokens
|
||||
```
|
||||
|
||||
This reduces compute from O(seq_len * seq_len_kv) to O(seq_len * topk). The causal mask is enforced by checking whether each index position is valid:
|
||||
This reduces compute from O(seq_len *seq_len_kv) to O(seq_len* topk). The causal mask is enforced by checking whether each index position is valid:
|
||||
|
||||
```python
|
||||
for bi_i in T.Parallel(BI):
|
||||
@@ -204,7 +204,7 @@ for i_i in T.Pipelined(NI, num_stages=num_stages):
|
||||
|
||||
The key gradient computations are:
|
||||
- **dQ = dP @ K** (query gradients)
|
||||
- **dK = dP^T @ Q** (key gradients)
|
||||
- **dK = dP^T @ Q** (key gradients)
|
||||
- **dV = P^T @ dO** (value gradients)
|
||||
|
||||
**3. Atomic Sparse Updates**: Uses atomic operations for dKV accumulation:
|
||||
|
||||
@@ -11,4 +11,4 @@ Launch the interactive chat interface and start exploring DeepSeek's capabilitie
|
||||
```bash
|
||||
export CONFIG=config_671B_v3.2.json
|
||||
torchrun --nproc-per-node ${MP} generate.py --ckpt-path ${SAVE_PATH} --config ${CONFIG} --interactive
|
||||
```
|
||||
```
|
||||
|
||||
@@ -108,4 +108,4 @@ def flash_attention(
|
||||
|
||||
# Write back the final output block from acc_o to the Output buffer
|
||||
T.copy(acc_o, Output[bz, bx * block_M : (bx + 1) * block_M, by, :])
|
||||
```
|
||||
```
|
||||
|
||||
@@ -1 +1 @@
|
||||
**Notes**: Now we only support fp8 with mma instructions instead of `T.gemm`, because the cutlass version of tilelang is too old, we should update the cutlass version in future.
|
||||
**Notes**: Now we only support fp8 with mma instructions instead of `T.gemm`, because the cutlass version of tilelang is too old, we should update the cutlass version in future.
|
||||
|
||||
@@ -103,4 +103,3 @@ latency = profiler.do_bench()
|
||||
print(f"Latency: {latency} ms")
|
||||
print(f"Performance: {2 * M * N * K / (latency/1e3) / 1e12:.2f} TFLOPS")
|
||||
```
|
||||
|
||||
|
||||
@@ -1,12 +1,12 @@
|
||||
# Carver: A Tile-Structure Based Hint Recommend Framework for Machine Learning Compilers
|
||||
|
||||
**Carver** is a lightweight framework for generating and ranking tile configurations (also known as **tiling strategies**, **blocking schemes**, or **scheduling hints**) for common GPU, CPU, and accelerator backends. It helps you explore efficient mappings of loops for operations such as matrix multiplication, elementwise transforms, and other reduction-oriented kernels.
|
||||
**Carver** is a lightweight framework for generating and ranking tile configurations (also known as **tiling strategies**, **blocking schemes**, or **scheduling hints**) for common GPU, CPU, and accelerator backends. It helps you explore efficient mappings of loops for operations such as matrix multiplication, elementwise transforms, and other reduction-oriented kernels.
|
||||
|
||||
Carver combines hardware architecture information, user-defined tile structures, and built-in heuristics to recommend tiling strategies (or "hints"). The recommended hints are easily adaptable to multiple backends, including [TVM](https://tvm.apache.org/), [triton](https://github.com/openai/triton), [tilelang](https://github.com/tile-ai/tilelang) (or other domain-specific compilers).
|
||||
|
||||
---
|
||||
|
||||
### Key Features
|
||||
## Key Features
|
||||
- **Unified Tiling Framework**: Generate tile candidates for multiple backends under a unified API.
|
||||
- **Architecture-Specific Modeling**: Take into account architecture constraints (e.g., CUDA `smem_cap`, warp size, CPU cache structure, etc.) when generating hints.
|
||||
- **Flexible Templates**: High-level templates (like `MatmulTemplate`, `GeneralReductionTemplate`, `ElementwiseTemplate`) let you concisely specify kernel structures.
|
||||
@@ -72,7 +72,6 @@ A tile structure composed of S and R can simulate various cases. For example, st
|
||||
|
||||
We can specialize more advanced templates to provide finer-grained information, such as `MatmulTemplate`.
|
||||
|
||||
|
||||
### Matmul Template
|
||||
|
||||
Carver also provides a specialized `MatmulTemplate` for matrix multiplication (e.g., `C = A * B`), automatically inferring common tiling strategies (thread blocks, warps, use of tensor cores, etc.).
|
||||
@@ -190,8 +189,6 @@ You might interpret this in **Triton** as:
|
||||
|
||||
This helps quickly test multiple configurations without manually guessing.
|
||||
|
||||
|
||||
|
||||
## Supported Templates
|
||||
|
||||
Carver abstracts common loop patterns through templates:
|
||||
@@ -203,8 +200,6 @@ Carver abstracts common loop patterns through templates:
|
||||
|
||||
You can also create your own specialized templates if you have unique loop structures or constraints. For instance, you might define specialized templates for convolution, flash attention, etc.
|
||||
|
||||
|
||||
## TODO Items
|
||||
|
||||
- [ ] **Adapt to tile language**: Provide ready-made scheduling calls or wrappers for [tilelang](https://github.com/LeiYanggh/tilelang) to streamline end-to-end integration.
|
||||
|
||||
|
||||
Reference in New Issue
Block a user