[Example] Update GEMM FP8 Example (#123)

* Add DeepSeek MLA decode example with Flash Attention implementation

* Add GEMM SplitK and StreamK example implementations

This commit introduces two new example scripts demonstrating advanced GEMM (matrix multiplication) techniques:
- `example_tilelang_gemm_splitk.py`: Implements a Split-K GEMM kernel using TileLang
- `example_tilelang_gemm_streamk.py`: Implements a Stream-K GEMM kernel using TileLang

Both examples showcase different parallel computation strategies for matrix multiplication, with comprehensive testing using PyTorch reference implementations.

* Refactor GEMM SplitK and StreamK example implementations

Clean up and improve code formatting for the SplitK and StreamK GEMM example scripts:
- Remove unused import (Profiler) in splitk example
- Simplify line breaks and improve code readability
- Standardize indentation and remove unnecessary whitespace
- Optimize atomic add and copy operations for better clarity

* Add block sparse attention benchmarks for multiple libraries

This commit introduces comprehensive block sparse attention benchmarks for different libraries:
- TileLang block sparse FMHA implementation
- Triton block sparse FMHA implementation
- PyTorch reference block sparse FMHA implementation
- FlashAttention dense FMHA reference implementation

The benchmarks include:
- Configurable benchmark parameters (batch size, heads, sequence length, etc.)
- Sparse mask generation using top-k and threshold methods
- Performance measurement for different sparse attention configurations
- Utility functions for mask generation and benchmarking

* Refactor block sparse attention benchmarks with code style improvements

- Add Ruff linter ignore comments to benchmark files
- Improve code formatting and line breaks
- Remove unused imports
- Standardize print statement formatting
- Enhance code readability across multiple library benchmarks

* lint fix

* Add CUDA atomic operations for BFLOAT16 and update function naming

- Implement AtomicAdd functions for BFLOAT16 and BFLOAT16x2 in CUDA common header
- Rename existing atomic add functions to use PascalCase (atomicAdd -> AtomicAdd)
- Add a new __pack_nv_bfloat162 function for packing BFLOAT16 values
- Update kernel and language customization to use new function names
- Add return type annotations in profiler module

* lint fix

* Add example for Group Query Attention (GQA) forward pass using Flash Attention in TileLang

This commit introduces a new example script `example_gqa_fwd_bshd.py` that demonstrates:
- Group Query Attention (GQA) implementation
- Flash Attention forward pass
- Performance benchmarking
- Configurable parameters for batch, heads, sequence length, and dimension
- Autotuning support
- Reference implementation comparison

* Refactor IR lowering pipeline into modular phases

This commit introduces a new module `phase.py` to modularize the IR lowering process by splitting the complex lowering pipeline into two distinct phases:
- `LowerAndLegalize`: Handles initial IR legalization and transformation
- `OptimizeForTarget`: Applies target-specific optimizations

The changes simplify the lowering logic in multiple files by extracting the transformation steps into reusable functions, improving code readability and maintainability.

* lintfix

* nas kernel

* Enhance Native Sparse Attention Examples with Code Improvements and Parameter Updates

- Updated example_tilelang_nsa.py and example_triton_nsa.py with code formatting and style improvements
- Increased default number of heads and selected blocks in TileLang NSA example
- Added Ruff linter ignore comments to reference.py
- Standardized function signatures and improved code readability across NSA implementations

* Add utility math functions for integer operations

- Implement `next_power_of_2()` to calculate the next power of 2 for an integer
- Add `cdiv()` function for ceiling division of integers

* Add utility math functions for integer operations

- Implement `next_power_of_2()` to calculate the next power of 2 for an integer
- Add `cdiv()` function for ceiling division of integers

* Refactor DeepSeek MLA Decode Example with Enhanced Flash Attention Implementation

- Update flash attention kernel to support positional embeddings (PE)
- Modify reference implementation to handle PE and group query attention
- Increase default batch size and adjust benchmarking parameters
- Improve kernel performance and readability
- Add einops and torch operations for more flexible tensor manipulation

* Update README.md with corrected Flash MLA Decoding example path

- Modify the example link for Flash MLA Decoding to point to the correct directory
- Ensure accurate navigation to the DeepSeek MLA decoding example

Author: LeiWang1999

README.md

@@ -26,7 +26,7 @@ Although tile-lang aims to be portable across a range of Devices, it has been sp
 - [Dequantization GEMM](./examples/dequantize_gemm/)
 - [Flash Attention](./examples/flash_attention/)
 - [Flash Linear Attention](./examples/linear_attention/)
-- [Flash MLA Decoding](./examples/flash_decoding/example_mla_decode.py)
+- [Flash MLA Decoding](./examples/deepseek_mla/)
 - [Native Sparse Attention](./examples/native_sparse_attention/)
 
 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.

examples/deepseek_mla/example_mla_decode.py

@@ -3,17 +3,13 @@
 import tilelang
 from tilelang.autotuner import *
 import tilelang.language as T
+from einops import rearrange, einsum
 
-num_split = 4
+num_split = 1
 
 
 def flashattn(batch, heads, kv_head_num, seqlen_kv, dim, pe_dim, block_N, block_H):
     scale = (1.0 / (dim + pe_dim))**0.5 * 1.44269504  # log2(e)
-    shape_q = [batch, heads, (dim + pe_dim)]
-    shape_k = [batch, seqlen_kv, kv_head_num, (dim + pe_dim)]
-    shape_v = [batch, seqlen_kv, kv_head_num, dim]
-    shape_o = [batch, heads, dim]
-    part_shape = [batch, heads, num_split, dim]
     dtype = "float16"
     accum_dtype = "float"
     kv_group_num = heads // kv_head_num
@@ -22,19 +18,23 @@ def flashattn(batch, heads, kv_head_num, seqlen_kv, dim, pe_dim, block_N, block_
 
     @T.macro
     def flash_attn_split(
-            Q: T.Buffer(shape_q, dtype),
-            K: T.Buffer(shape_k, dtype),
-            V: T.Buffer(shape_v, dtype),
+            Q: T.Buffer([batch, heads, dim], dtype),
+            Q_pe: T.Buffer([batch, heads, pe_dim], dtype),
+            KV: T.Buffer([batch, seqlen_kv, kv_head_num, dim], dtype),
+            K_pe: T.Buffer([batch, seqlen_kv, kv_head_num, pe_dim], dtype),
             glse: T.Buffer([batch, heads, num_split], dtype),
-            Output_partial: T.Buffer(part_shape, dtype),
+            Output_partial: T.Buffer([batch, heads, num_split, dim], dtype),
     ):
         with T.Kernel(
-                batch, heads // min(block_H, kv_group_num), num_split, threads=128) as (bx, by, bz):
-            Q_shared = T.alloc_shared([block_H, (dim + pe_dim)], dtype)
-            K_shared = T.alloc_shared([block_N, (dim + pe_dim)], dtype)
-            V_shared = T.alloc_shared([block_N, dim], dtype)
+                batch, heads // min(block_H, kv_group_num), num_split, threads=256) as (bx, by, bz):
+            Q_shared = T.alloc_shared([block_H, dim], dtype)
+            S_shared = T.alloc_shared([block_H, block_N], dtype)
+            Q_pe_shared = T.alloc_shared([block_H, pe_dim], dtype)
+            KV_shared = T.alloc_shared([block_N, dim], dtype)
+            K_pe_shared = T.alloc_shared([block_N, pe_dim], dtype)
             O_shared = T.alloc_shared([block_H, dim], dtype)
             acc_s = T.alloc_fragment([block_H, block_N], accum_dtype)
+            acc_s_0 = T.alloc_fragment([block_H, block_N], accum_dtype)
             acc_s_cast = T.alloc_fragment([block_H, block_N], dtype)
             acc_o = T.alloc_fragment([block_H, dim], accum_dtype)
             scores_max = T.alloc_fragment([block_H], accum_dtype)
@@ -53,20 +53,32 @@ def flash_attn_split(
             })
 
             T.copy(Q[bid, hid * VALID_BLOCK_H:(hid + 1) * VALID_BLOCK_H, :], Q_shared)
+            T.copy(Q_pe[bid, hid * VALID_BLOCK_H:(hid + 1) * VALID_BLOCK_H, :], Q_pe_shared)
             T.fill(acc_o, 0)
             T.fill(logsum, 0)
             T.fill(scores_max, -T.infinity(accum_dtype))
 
             loop_range = T.ceildiv((seqlen_kv // num_split), block_N)
-            for k in T.Pipelined(loop_range, num_stages=1):
-                T.copy(
-                    K[bid, (seqlen_kv // num_split) * sid +
-                      k * block_N:(seqlen_kv // num_split) * sid + (k + 1) * block_N,
-                      cur_kv_head, :], K_shared)
-                T.clear(acc_s)
-                T.gemm(Q_shared, K_shared, acc_s, transpose_B=True, policy=T.GemmWarpPolicy.FullRow)
+            for k in T.Pipelined(loop_range, num_stages=2):
+                kv_start = (seqlen_kv // num_split) * sid + k * block_N
+                kv_end = (seqlen_kv // num_split) * sid + (k + 1) * block_N
+
+                T.copy(KV[bid, kv_start:kv_end, cur_kv_head, :], KV_shared)
+                T.copy(K_pe[bid, kv_start:kv_end, cur_kv_head, :], K_pe_shared)
+
+                T.clear(acc_s_0)
+                T.gemm(
+                    Q_shared, KV_shared, acc_s_0, transpose_B=True, policy=T.GemmWarpPolicy.FullCol)
+                T.gemm(
+                    Q_pe_shared,
+                    K_pe_shared,
+                    acc_s_0,
+                    transpose_B=True,
+                    policy=T.GemmWarpPolicy.FullCol)
                 T.copy(scores_max, scores_max_prev)
                 T.fill(scores_max, -T.infinity(accum_dtype))
+                T.copy(acc_s_0, S_shared)
+                T.copy(S_shared, acc_s)
                 T.reduce_max(acc_s, scores_max, dim=1, clear=False)
                 for i in T.Parallel(block_H):
                     scores_scale[i] = T.exp2(scores_max_prev[i] * scale - scores_max[i] * scale)
@@ -78,11 +90,7 @@ def flash_attn_split(
                 T.copy(acc_s, acc_s_cast)
                 for i, j in T.Parallel(block_H, dim):
                     acc_o[i, j] *= scores_scale[i]
-                T.copy(
-                    V[bid, (seqlen_kv // num_split) * sid +
-                      k * block_N:(seqlen_kv // num_split) * sid + (k + 1) * block_N,
-                      cur_kv_head, :], V_shared)
-                T.gemm(acc_s_cast, V_shared, acc_o, policy=T.GemmWarpPolicy.FullRow)
+                T.gemm(acc_s_cast, KV_shared, acc_o, policy=T.GemmWarpPolicy.FullCol)
             for i, j in T.Parallel(block_H, dim):
                 acc_o[i, j] /= logsum[i]
             for i in T.Parallel(block_H):
@@ -96,8 +104,8 @@ def flash_attn_split(
     @T.macro
     def combine(
             glse: T.Buffer([batch, heads, num_split], dtype),
-            Output_partial: T.Buffer(part_shape, dtype),
-            Output: T.Buffer(shape_o, dtype),
+            Output_partial: T.Buffer([batch, heads, num_split, dim], dtype),
+            Output: T.Buffer([batch, heads, dim], dtype),
     ):
         with T.Kernel(heads, batch, threads=128) as (by, bz):
             po_local = T.alloc_fragment([dim], dtype)
@@ -133,50 +141,61 @@ def combine(
 
     @T.prim_func
     def main(
-            Q: T.Buffer(shape_q, dtype),
-            K: T.Buffer(shape_k, dtype),
-            V: T.Buffer(shape_v, dtype),
+            Q: T.Buffer([batch, heads, dim], dtype),
+            Q_pe: T.Buffer([batch, heads, pe_dim], dtype),
+            KV: T.Buffer([batch, seqlen_kv, kv_head_num, dim], dtype),
+            K_pe: T.Buffer([batch, seqlen_kv, kv_head_num, pe_dim], dtype),
             glse: T.Buffer([batch, heads, num_split], dtype),
-            Output_partial: T.Buffer(part_shape, dtype),  # [batch, heads, num_split, dim]
-            Output: T.Buffer(shape_o, dtype),
+            Output_partial: T.Buffer([batch, heads, num_split, dim], dtype),
+            Output: T.Buffer([batch, heads, dim], dtype),
     ):
-        flash_attn_split(Q, K, V, glse, Output_partial)
+        flash_attn_split(Q, Q_pe, KV, K_pe, glse, Output_partial)
         combine(glse, Output_partial, Output)
 
     return main
 
 
-def ref_program(query, key, value, glse, Output_partial):
+def ref_program(q, q_pe, kv, k_pe, glse, Output_partial):
     #     """
     #     Inputs:
-    #     - query (Tensor): [batch, heads, dim]
-    #     - key (Tensor): [batch, seqlen_kv, kv_head_num, dim]
-    #     - value (Tensor): [batch, seqlen_kv, kv_head_num, dim]
-
+    #     - q (Tensor): [batch, heads, dim]
+    #     - q_pe (Tensor): [batch, heads, pe_dim]
+    #     - kv (Tensor): [batch, seqlen_kv, kv_head_num, dim]
+    #     - k_pe (Tensor): [batch, seqlen_kv, kv_head_num, pe_dim]
+    #     - glse (Tensor): [batch, heads, num_split]
+    #     - Output_partial (Tensor): [batch, heads, num_split, dim]
     #     Outputs:
     #     - output (Tensor): [batch, heads, dim]
     #     """
-    from einops import rearrange
-    batch_size, query_heads, dim = query.shape  # [batch_size, query_heads, dim]
-    _, seqlen_kv, kv_heads, _ = key.shape  # [batch_size, seqlen_kv, kv_heads, kv_dim]
-    dim_v = value.shape[-1]
-    assert kv_heads == 1, "kv_heads must be 1"
-
-    query_expanded = rearrange(query, 'b h d -> b h 1 d')  # [batch_size, query_heads, 1, dim]
-    key_expanded = key.expand(-1, -1, query_heads, -1)  # [batch_size, query_heads, seqlen_kv, dim]
-    value_expanded = value.expand(-1, -1, query_heads,
-                                  -1)  # [batch_size, query_heads, seqlen_kv, dim]
-    key_expanded = rearrange(key_expanded,
-                             'b n h d -> b h n d')  # [batch_size, kv_head_num, seqlen_kv, dim]
-    value_expanded = rearrange(value_expanded,
-                               'b n h d -> b h n d')  # [batch_size, query_heads, seqlen_kv, dim]
-
-    scores = torch.matmul(query_expanded,
-                          key_expanded.transpose(-1, -2))  # [batch_size, query_heads, 1, seqlen_kv]
-    scores = scores / torch.sqrt(torch.tensor(dim, dtype=scores.dtype))
-    attention_weights = F.softmax(scores, dim=-1)  # [batch_size, query_heads, 1, seqlen_kv]
-    output = torch.matmul(attention_weights, value_expanded)  # [batch_size, query_heads, 1, dim]
-    return output.view(batch_size, query_heads, dim_v)
+    dim = q.shape[-1]
+    pe_dim = q_pe.shape[-1]
+    num_head_groups = q.shape[1] // kv.shape[2]
+    scale = (dim + pe_dim)**0.5
+    q = rearrange(
+        q, 'b (h g) d -> b g h d', g=num_head_groups)  # [batch_size, num_head_groups, groups, dim]
+
+    q_pe = rearrange(
+        q_pe, 'b (h g) d -> b g h d',
+        g=num_head_groups)  # [batch_size, num_head_groups, groups, pe_dim]
+
+    kv = rearrange(kv, 'b n h d -> b h n d')  # [batch_size, groups, seqlen_kv, dim]
+
+    k_pe = rearrange(k_pe, 'b n h d -> b h n d')  # [batch_size, num_head_groups, groups, pe_dim]
+
+    query = torch.concat([q, q_pe], dim=-1)
+    key = torch.concat([kv, k_pe], dim=-1)
+
+    scores = einsum(
+        query, key,
+        'b g h d, b h s d -> b g h s')  # [batch_size, num_head_groups, groups, seqlen_kv]
+
+    attention = F.softmax(
+        scores / scale, dim=-1)  # [batch_size, num_head_groups, groups, seqlen_kv]
+
+    out = einsum(attention, kv,
+                 'b g h s, b h s d -> b g h d')  # [batch_size, num_head_groups, groups, dim]
+    out = rearrange(out, 'b g h d -> b (h g) d')  # [batch_size, heads, dim]
+    return out
 
 
 def flash_split_ref(Q, K, V):
@@ -251,7 +270,7 @@ def reduce_ref(Q, K, V, glse, Output_partial):
 
 
 if __name__ == "__main__":
-    BATCH, H_Q, KV_H, KV_CTX, D_HEAD, DPE = 64, 128, 1, 8192, 512, 64
+    BATCH, H_Q, KV_H, KV_CTX, D_HEAD, DPE = 128, 128, 1, 8192, 512, 64
     qk_flops = 2 * BATCH * H_Q * KV_CTX * (D_HEAD + DPE)
     pv_flops = 2 * BATCH * H_Q * KV_CTX * D_HEAD
     total_flops = qk_flops + pv_flops
@@ -260,8 +279,9 @@ def reduce_ref(Q, K, V, glse, Output_partial):
 
     program = flashattn(BATCH, H_Q, KV_H, KV_CTX, D_HEAD, DPE, BLOCK_N, BLOCK_H)
     mod, params = tilelang.lower(program)
-    mod = tilelang.Profiler(mod, params, [5], tilelang.TensorSupplyType.Normal)
+    mod = tilelang.Profiler(mod, params, [6], tilelang.TensorSupplyType.Normal)
     mod.assert_allclose(ref_program, rtol=0.01, atol=0.01)
-    latency = mod.do_bench(mod.func, warmup=500)
+    print("All close")
+    latency = mod.do_bench(mod.func, n_warmup=10, n_repeat=10, profiler="torch")
     print("Tile-lang: {:.2f} ms".format(latency))
-    print("Tile-lang: {:.2f} TFlops".format(total_flops / latency * 1e-9))
+    print("Tile-lang: {:.2f} TFlops".format(total_flops / latency * 1e-9))