fix bug&add amd examples (#966)

* [Enhancement] Refactor buffer index handling for improved precision and clarity (#668)

- Enhanced buffer index handling to address precision issues by removing redundant operations.
- Streamlined the logic for determining buffer overlaps, ensuring more accurate conflict detection.
- Updated related documentation to reflect changes in buffer management practices.

* Remove obsolete test script for AMD example, streamlining the examples directory.

* Remove unused dtype_size variable in AMD example script to streamline code.

* Add input configuration file and update AMD example script for enhanced flexibility

- Introduced a new input.txt file for configurable parameters.
- Modified the example_amd_flash_attn_fwd.py script to allow for a wider range of configurations, including additional options for num_stages, enable_rasterization, and k_pack.
- Streamlined the main function for better clarity and organization.
- Added a new test script to facilitate running the example with specified parameters.

* Remove input configuration file and obsolete test script; enhance AMD example with swizzle layout annotations

- Deleted input.txt and test.sh files as they are no longer needed.
- Updated example_amd_flash_attn_fwd.py to include swizzle layout annotations for shared memory, improving bank conflict avoidance.
- Reintroduced swizzle usage in the kernel for better performance.

* Refactor AMD example script for FlashAttention-2

- Updated function names for clarity, changing `get_v2_configs` to `get_configs` and `fast_flashattn_v2` to `fast_flashattn`.
- Streamlined the main function by renaming `main_v2` to `main` and adjusting the corresponding calls.
- Removed outdated comments and improved code organization for better readability.

* Refactor formatting in AMD FlashAttention example script

- Improved code readability by adjusting line breaks and indentation in the `fast_flashattn` function.
- Streamlined the `main` function parameter formatting for consistency.
- Removed unnecessary blank lines to enhance overall code organization.

* Update example_amd_flash_attn_fwd.py

* Enhance AMD example script and update CI workflows

- Improved the `example_amd_flash_attn_fwd.py` script for better clarity and organization.
- Added new CI workflows for AMD and documentation publishing.
- Updated various requirements files to include necessary dependencies.
- Introduced new test cases and examples for better coverage and functionality.
- Refactored existing code for improved readability and maintainability.

* Remove redundant tool cache cleanup step in AMD CI workflow

* Remove `torch` dependency from `requirements-rocm.txt` to streamline requirements.

* Add new AMD FlashAttention example and test script

- Introduced `example_amd_flash_attn_bwd.py` for backward attention computation using TileLang.
- Added `test.sh` script to facilitate running the new example with specified parameters.
- Enhanced the overall structure and organization of the example for better clarity and usability.

* Update configurations in `example_amd_flash_attn_fwd.py` for autotuner

- Reduced the number of threads and `num_split_q` options for improved performance.
- Adjusted `panel_size` options to streamline configuration settings.

* Update submodule 'tvm' to commit 6ccc74f622c7ec4ac25d430d0f6546e7b9edb217

* Update submodule 'tvm' to commit 14ff70ab142b9e5a31bbf9c7923c8a697d41e86c

* Add example for AMD Flash Attention backward pass implementation

- Introduced a new example script `example_amd_flash_attn_bwd.py` demonstrating the forward and backward operations of Flash Attention using TileLang.
- Implemented JIT-compiled functions for both forward and backward passes, including preprocessing and postprocessing steps.
- Added a main function to facilitate testing and benchmarking of the attention mechanism with configurable parameters.
- Included reference implementation for validation against PyTorch's attention mechanism.

This addition enhances the examples directory by providing a comprehensive guide for users to understand and utilize Flash Attention in their applications.

* Enhance AMD Flash Attention example with additional testing capabilities

- Updated `example_amd_flash_attn_bwd.py` to include more comprehensive testing features for the Flash Attention implementation.
- Improved the main function to allow for better parameter configuration and benchmarking.
- Added validation checks against PyTorch's attention mechanism to ensure accuracy and reliability of the example.

This update aims to provide users with a more robust tool for understanding and utilizing Flash Attention in their applications.

* Update submodule TVM to commit a64a5926a6e59f5417ef2501f9d88b467337cf6a

* Refactor HIP intrinsic rules to CUDA

- Updated file name from `intrin_rule_hip.cc` to `intrin_rule_cuda.cc` to reflect the change in focus from HIP to CUDA intrinsic rules.
- Adjusted include paths for better organization and clarity in the code structure.

* Update AMD CI workflow to uninstall specific PyTorch packages before installation

- Removed the installation of `flash_attn==2.5.8` to streamline the CI process.
- Added a step to uninstall `torch`, `torchvision`, and `torchaudio` prior to installing pre-release versions, ensuring compatibility and reducing potential conflicts.

* Remove unused shared memory allocations in AMD Flash Attention backward example

- Eliminated the allocation of shared memory for `dv_shared` and `dk_shared` in `example_amd_flash_attn_bwd.py` to streamline memory usage and improve performance.
- This change focuses on optimizing the backward pass implementation by reducing unnecessary memory overhead.

* Remove unnecessary pip uninstall command from AMD CI workflow

- Eliminated the step to uninstall `torch`, `torchvision`, and `torchaudio` in the AMD CI workflow, as it is no longer required for the installation of pre-release versions.
- This change simplifies the CI process and reduces potential overhead during package management.

* Refactor DispatchHIPWarpActiveMask function in HIP intrinsic rules

- Updated the return statement to use std::string for concatenation in the case of 16-bit types, improving code clarity.
- Added a null check for the CallNode pointer in DispatchHIPWarpActiveMask to enhance robustness and prevent potential dereferencing issues.

* Refactor formatting of HIP intrinsic rule registrations

- Adjusted the formatting of TVM_REGISTER_OP calls for better readability by aligning method chaining.
- No functional changes were made; this update focuses on code style improvements to enhance maintainability.

* Update file name and documentation for HIP intrinsic rules

- Renamed the file from `intrin_rule_cuda.cc` to `intrin_rule_hip.cc` to accurately reflect the focus on HIP intrinsic rules.
- Updated the file documentation to clarify its purpose as related to HIP rather than CUDA.

* Enhance DispatchHIPShuffle function with clang-analyzer comments

- Added NOLINTBEGIN and NOLINTEND comments to the DispatchHIPShuffle function to suppress clang-analyzer warnings related to inner pointer usage.
- This change improves code clarity and maintains compliance with static analysis tools.

* lint fix

* fix

* Enhance autotuner configurations in example_amd_flash_attn_fwd.py by adding new block sizes, stages, and panel sizes. Update test script to use relative Python path and adjust parameters for consistency.

* Add backward attention example to test script

- Extended the test.sh script to include a new backward attention example using example_amd_flash_attn_bwd.py.
- Added parameters for batch size, context length, and head dimensions to ensure consistency with the forward example.
- Updated the command for the backward tile example to match the new configuration.

* Refactor FlashAttention implementation in example_amd_flash_attn_bwd.py and example_amd_flash_attn_fwd.py

- Introduced new functions for forward and backward configurations to enhance autotuning capabilities.
- Updated the FlashAttention forward and backward functions to improve performance and maintainability.
- Adjusted test script parameters for consistency and clarity, including the addition of group handling.
- Enhanced the autotuner configurations by refining block sizes and stages for better performance tuning.
- Updated the main function to reflect changes in parameter names and types for better usability.

* Enhance FlashAttention backward implementation in example_amd_flash_attn_bwd.py

- Updated the backward function to return additional outputs, including log-sum-exp (LSE) values for improved gradient calculations.
- Refined autotuner configurations by adding new block sizes and adjusting parameters for better performance tuning.
- Improved shared memory usage in the backward pass to optimize memory access patterns and enhance computational efficiency.
- Updated the main function to reflect changes in parameter handling and ensure consistency with the forward pass.
- Enhanced correctness checks in the main function to include LSE validation alongside gradient checks.

* Enhance FlashAttention backward implementation in example_amd_flash_attn_bwd.py

- Introduced a scaling factor for improved numerical stability in gradient calculations.
- Optimized shared memory usage by adding new shared buffers for intermediate calculations.
- Refined the handling of tensor fragments to improve performance and maintainability.
- Updated the main function to ensure compatibility with the new output parameters for backward operations.
- Removed unnecessary parameters from the test script to streamline execution.

* Refactor FlashAttention implementation in example_amd_flash_attn_bwd.py and example_mha_bwd.py

- Updated the forward and backward functions to improve numerical stability and performance.
- Enhanced shared memory usage by optimizing buffer allocations and reducing unnecessary parameters.
- Adjusted autotuner configurations for better performance tuning and compatibility with new output parameters.
- Added debugging and benchmarking functions for improved correctness verification and performance analysis.
- Updated the main function to reflect changes in parameter handling and ensure consistency across examples.

* Enhance FlashAttention backward implementation in example_amd_flash_attn_bwd.py

- Updated scaling factor application for improved numerical stability in gradient calculations.
- Refined tensor handling to ensure consistency with forward pass operations.
- Optimized atomic operations for writing gradients to dK and dV using fp32 for better precision.
- Adjusted comments for clarity and alignment with standard implementation practices.

* Expand autotuner configurations in example_amd_flash_attn_bwd.py and update test.sh

- Increased the range of block sizes and stages for forward and backward configurations to enhance performance tuning.
- Adjusted the test script to include additional parameters for batch size and head dimensions, ensuring consistency with the forward example.
- Improved comments for clarity and alignment with the updated configurations.

* Enhance performance calculations and benchmarking in example_amd_flash_attn_bwd.py

- Updated FLOPs calculation to account for both forward and backward passes, clarifying the total computational cost.
- Modified benchmarking functions to evaluate the complete forward and backward performance of both reference and Tile-lang implementations.
- Improved comments for better understanding of the performance metrics and implementation details.
- Removed unnecessary parameter from test.sh to streamline execution.

* Remove forward attention test commands from test.sh and retain backward attention execution for streamlined testing.

* Refactor FlashAttention forward and backward implementations in example_amd_flash_attn_bwd.py and example_amd_flash_attn_fwd.py

- Updated the forward function to return both output and log-sum-exp (LSE) values for improved gradient calculations.
- Enhanced autotuner configurations for forward pass, including new parameters for better performance tuning.
- Refined scaling factor calculations for numerical stability in both forward and backward passes.
- Improved comments and documentation for clarity and consistency across implementations.
- Adjusted main function to reflect changes in parameter handling and ensure compatibility with new output requirements.

* Refactor FlashAttention implementation in example_amd_flash_attn_bwd.py

- Removed outdated comments and improved clarity in the code.
- Enhanced the forward function to consistently return output and log-sum-exp (LSE) values.
- Updated autotuner configurations to include new parameters for better performance tuning.
- Refined tensor handling and scaling factor calculations for improved numerical stability.
- Adjusted the main function to ensure compatibility with updated output requirements and parameter handling.

* Enhance FlashAttention backward implementation in example_amd_flash_attn_bwd.py

- Updated configuration parameters for backward calculations, including new options for block sizes, threads, and rasterization.
- Added new parameters (k_pack, qk_coalesced_width, v_coalesced_width) to improve performance tuning and memory access patterns.
- Modified tensor copy operations to utilize coalesced widths for optimized memory loads.
- Enhanced GEMM operations with k_pack for improved computational efficiency.
- Refined the configuration generation logic to accommodate the new parameters, ensuring comprehensive coverage for backward pass scenarios.

* Refactor configuration and tensor operations in example_amd_flash_attn_bwd.py

- Updated backward configuration parameters to include larger block sizes and a wider range of threads for enhanced performance tuning.
- Removed unnecessary parameters (k_pack, qk_coalesced_width, v_coalesced_width) from function signatures and tensor operations to simplify the implementation.
- Optimized tensor copy operations by eliminating coalesced width specifications, streamlining memory access patterns.
- Adjusted GEMM operations to improve computational efficiency without the use of k_pack.

* Enhance HIP code generation and FP8 type support

- Added support for additional FP8 types (e4m3, e4m3b11fnuz, e5m2fnuz, e8m0) in codegen_hip.cc to improve compatibility.
- Updated error logging to include unsupported FP8 type details for better debugging.
- Implemented handling for loop break and no-op register management in HIP within VisitExpr_ method.
- Introduced new FP8 vector types (e5 and e8) in hip_fp8.h for enhanced functionality.
- Added overloads for AtomicAdd in common.h to support both pointer and value arguments.

* Enhance FP8 type support and clarify accumulator handling in HIP

- Expanded FP8 type support in codegen_hip.cc to include additional float8 formats.
- Updated gemm.h to clarify the handling of the accumulator when clear_accum is true.
- Added comments in hip_fp8.h to indicate that E8M0 types are not supported in the current HIP version.

* Remove deprecated files and update print statements for clarity in example_amd_flash_attn_bwd.py

* Update print statement formatting for clarity in example_amd_flash_attn_bwd.py

* Remove redundant verification results summary print statement in example_amd_flash_attn_bwd.py for cleaner output.

* Fix formatting inconsistencies in example_amd_flash_attn_bwd.py and example_amd_flash_attn_fwd.py by adding spaces for improved readability in configuration parameters and print statements.

* Refactor and enhance HIP code generation for improved FP8 support

- Reorganized and cleaned up code in codegen_hip.cc for better readability and maintainability.
- Enhanced handling of FP8 types, including additional formats and improved error logging for unsupported types.
- Updated AtomicAdd function in common.h to streamline its implementation.
- Refined the PrintVecElemLoadExpr method to handle volatile loads more effectively.
- Added function to manage the addition of new functions in the code generation process.

* Fix formatting issue in HIP code generation for MFMA call

- Adjusted the indentation of the MFMA call code block in codegen_hip.cc for improved readability and consistency.

* Refactor HIP code generation and enhance FP8 type handling

- Reintroduced necessary includes and reorganized code in codegen_hip.cc for improved structure and readability.
- Enhanced the GetFP8Type function to support additional FP8 formats and improved error handling for unsupported types.
- Updated PrintType and PrintVecElemLoadExpr methods to better manage type conversions and vector element loading.
- Refined the AddFunction method to streamline function addition in the code generation process.

* Remove unnecessary blank line in example_amd_flash_attn_bwd.py for improved code cleanliness.

* Refactor backward attention implementation in example_amd_flash_attn_bwd.py

- Updated the GEMM operation to use shared memory for improved performance.
- Adjusted parallelization parameters to enhance efficiency in the backward pass.

* Fix formatting by removing an unnecessary blank line in example_amd_flash_attn_bwd.py for improved code cleanliness.

* Add additional test cases for `assert_tl_matmul_correctness` with `float8_e4m3fnuz` and various configurations

* Refactor test case formatting for `assert_tl_matmul_correctness` in `test_tilelang_gemm_mfma_intrinsic.py`

---------
Co-authored-by: xinxyxiao <xinyxiao@amd.com>
Co-authored-by: Lei Wang <34334180+LeiWang1999@users.noreply.github.com>
Co-authored-by: LeiWang1999 <leiwang1999@outlook.com>

examples/amd/example_amd_flash_attn_bwd.py

 import torch
 import torch.nn.functional as F
 import tilelang
-from tilelang.autotuner import *
 import tilelang.language as T
+from tilelang.primitives.gemm.base import GemmWarpPolicy
+import itertools
 import argparse
+from functools import partial
+import numpy as np
+import time
 
 
+def ref_program(Q, K, V, is_causal, groups=1):
+    assert Q.size(
+        2) == K.size(2) * groups, f"Q heads {Q.size(2)} K heads {K.size(2)} groups {groups}"
+    assert Q.size(
+        2) == V.size(2) * groups, f"Q heads {Q.size(2)} V heads {V.size(2)} groups {groups}"
+    dim = Q.size(-1)
+    K_ref = K.repeat_interleave(groups, dim=2)
+    V_ref = V.repeat_interleave(groups, dim=2)
+    scores = torch.einsum('bqhd,bkhd->bhqk', Q, K_ref)
+    scores = scores / torch.sqrt(torch.tensor(dim, dtype=scores.dtype))
+    if is_causal:
+        seq_len = Q.size(1)
+        mask = torch.tril(torch.ones(seq_len, seq_len, device=scores.device))
+        mask = mask.unsqueeze(0).unsqueeze(0)
+        scores = scores.masked_fill(mask == 0, float('-inf'))
+    attention_weights = F.softmax(scores, dim=-1)
+    output = torch.einsum('bhqk,bkhd->bqhd', attention_weights, V_ref)
+    lse = torch.logsumexp(scores, dim=-1).float()
+    return output, lse
+
+
+def get_fwd_configs():
+    block_M = [32, 64, 128, 256]
+    block_N = [32, 64, 128, 256]
+    threads = [128, 256, 512]
+    num_split_q = [64, 128, 256]
+    num_stages = [0, 1]
+    enable_rasterization = [True]
+    k_pack = [2]
+    panel_size = [7, 8, 9, 10]
+    qk_coalesced_width = [8]
+    v_coalesced_width = [4]
+
+    valid_configs = []
+
+    for m, n, s, t, stages, r, k, p, qkw, vw in itertools.product(block_M, block_N, num_split_q,
+                                                                  threads, num_stages,
+                                                                  enable_rasterization, k_pack,
+                                                                  panel_size, qk_coalesced_width,
+                                                                  v_coalesced_width):
+        valid_configs.append({
+            "block_M": m,
+            "block_N": n,
+            "num_split_q": s,
+            "threads": t,
+            "num_stages": stages,
+            "enable_rasterization": r,
+            "k_pack": k,
+            "panel_size": p,
+            "qk_coalesced_width": qkw,
+            "v_coalesced_width": vw,
+        })
+    return valid_configs
+
+
+@tilelang.autotune(configs=get_fwd_configs(), cache_input_tensors=True)
 @tilelang.jit(out_idx=[3, 4])
-def flashattn_fwd(batch, heads, seq_len, dim_qk, dim_v, is_causal, block_M, block_N, groups=1):
-    scale = (1.0 / dim_qk)**0.5 * 1.44269504  # log2(e)
+def fast_flashattn(
+    batch,
+    heads,
+    seq_len,
+    dim,
+    is_causal,
+    groups,
+    block_M: int,
+    block_N: int,
+    num_split_q: int,
+    threads: int,
+    num_stages: int,
+    enable_rasterization: bool,
+    k_pack: int,
+    panel_size: int,
+    qk_coalesced_width: int,
+    v_coalesced_width: int,
+):
+    scale = (1.0 / dim)**0.5
     head_kv = heads // groups
-    q_shape = [batch, seq_len, heads, dim_qk]
-    k_shape = [batch, seq_len, head_kv, dim_qk]
-    v_shape = [batch, seq_len, head_kv, dim_v]
+    q_shape = [batch, seq_len, heads, dim]
+    kv_shape = [batch, seq_len, head_kv, dim]
     dtype = "float16"
     accum_dtype = "float"
 
+    vec_size = qk_coalesced_width
+    v_vec_size = v_coalesced_width
+
     @T.prim_func
-    def flash_fwd(
-            Q: T.Tensor(q_shape, dtype),  # type: ignore
-            K: T.Tensor(k_shape, dtype),  # type: ignore
-            V: T.Tensor(v_shape, dtype),  # type: ignore
-            Output: T.Tensor([batch, seq_len, heads, dim_v], dtype),  # type: ignore
-            lse: T.Tensor([batch, heads, seq_len], accum_dtype),  # type: ignore
+    def main(
+            Q: T.Tensor(q_shape, dtype),
+            K: T.Tensor(kv_shape, dtype),
+            V: T.Tensor(kv_shape, dtype),
+            Output: T.Tensor(q_shape, dtype),
+            LSE: T.Tensor([batch, heads, seq_len], accum_dtype),
     ):
-        with T.Kernel(T.ceildiv(seq_len, block_M), heads, batch, threads=256) as (bx, by, bz):
-            Q_shared = T.alloc_shared([block_M, dim_qk], dtype)
-            K_shared = T.alloc_shared([block_N, dim_qk], dtype)
-            V_shared = T.alloc_shared([block_N, dim_v], dtype)
-            acc_s = T.alloc_fragment([block_M, block_N], accum_dtype)
-            acc_s_cast = T.alloc_fragment([block_M, block_N], dtype)
-            acc_o = T.alloc_fragment([block_M, dim_v], accum_dtype)
-            scores_max = T.alloc_fragment([block_M], accum_dtype)
-            scores_max_prev = T.alloc_fragment([block_M], accum_dtype)
-            scores_scale = T.alloc_fragment([block_M], accum_dtype)
-            scores_sum = T.alloc_fragment([block_M], accum_dtype)
-            logsum = T.alloc_fragment([block_M], accum_dtype)
-
-            T.copy(Q[bz, bx * block_M:(bx + 1) * block_M, by, :], Q_shared)
-            T.fill(acc_o, 0)
-            T.fill(logsum, 0)
-            T.fill(scores_max, -T.infinity(accum_dtype))
-            loop_range = (
-                T.ceildiv(
-                    (bx + 1) * block_M, block_N) if is_causal else T.ceildiv(seq_len, block_N))
-            for k in T.Pipelined(loop_range, num_stages=1):
-                T.copy(K[bz, k * block_N:(k + 1) * block_N, by // groups, :], K_shared)
-                if is_causal:
+        with T.Kernel(num_split_q, batch * heads, threads=threads) as (b_split, byz_combined):
+            T.use_swizzle(panel_size, enable=enable_rasterization)
+
+            bz = byz_combined // heads
+            by = byz_combined % heads
+
+            num_q_blocks = T.ceildiv(seq_len, block_M)
+
+            bx_loop_var = T.alloc_var("int32")
+            bx_loop_var = b_split
+
+            with T.While(bx_loop_var < num_q_blocks):
+                acc_o = T.alloc_fragment([block_M, dim], accum_dtype)
+                m_i = T.alloc_fragment([block_M], accum_dtype)
+                l_i = T.alloc_fragment([block_M], accum_dtype)
+
+                T.fill(acc_o, 0)
+                T.fill(m_i, -T.infinity(accum_dtype))
+                T.fill(l_i, 0)
+
+                current_bx = bx_loop_var
+                q_block_offset = current_bx * block_M
+
+                Q_shared = T.alloc_shared([block_M, dim], dtype)
+                K_shared = T.alloc_shared([block_N, dim], dtype)
+                V_shared = T.alloc_shared([block_N, dim], dtype)
+                acc_s_cast = T.alloc_fragment([block_M, block_N], dtype)
+
+                acc_s = T.alloc_fragment([block_M, block_N], accum_dtype)
+                m_prev = T.alloc_fragment([block_M], accum_dtype)
+                scale_factor = T.alloc_fragment([block_M], accum_dtype)
+
+                T.copy(
+                    Q[bz, q_block_offset:q_block_offset + block_M, by, :],
+                    Q_shared,
+                    coalesced_width=vec_size)
+
+                loop_end_k = (
+                    T.ceildiv(q_block_offset +
+                              block_M, block_N) if is_causal else T.ceildiv(seq_len, block_N))
+
+                row_sum = T.alloc_fragment([block_M], accum_dtype)
+
+                for k in T.Pipelined(loop_end_k, num_stages=num_stages):
+                    kv_idx = k * block_N
+
+                    T.copy(
+                        K[bz, kv_idx:kv_idx + block_N, by // groups, :],
+                        K_shared,
+                        coalesced_width=vec_size)
+                    T.copy(
+                        V[bz, kv_idx:kv_idx + block_N, by // groups, :],
+                        V_shared,
+                        coalesced_width=v_vec_size)
+
+                    if is_causal:
+                        for i, j in T.Parallel(block_M, block_N):
+                            acc_s[i, j] = T.if_then_else(q_block_offset + i >= kv_idx + j, 0,
+                                                         -T.infinity(acc_s.dtype))
+                    else:
+                        T.clear(acc_s)
+                    T.gemm(
+                        Q_shared,
+                        K_shared,
+                        acc_s,
+                        transpose_B=True,
+                        k_pack=k_pack,
+                        policy=GemmWarpPolicy.FullRow,
+                    )
+
+                    for i, j in T.Parallel(block_M, block_N):
+                        acc_s[i, j] = acc_s[i, j] * scale
+
+                    T.copy(m_i, m_prev)
+                    T.reduce_max(acc_s, m_i, dim=1, clear=False)
+
+                    for i in T.Parallel(block_M):
+                        if m_prev[i] == -T.infinity(accum_dtype):
+                            scale_factor[i] = 0.0
+                        else:
+                            scale_factor[i] = T.exp(m_prev[i] - m_i[i])
+
+                        l_i[i] *= scale_factor[i]
+
+                    for i, j in T.Parallel(block_M, dim):
+                        acc_o[i, j] *= scale_factor[i]
+
                     for i, j in T.Parallel(block_M, block_N):
-                        acc_s[i, j] = T.if_then_else(bx * block_M + i >= k * block_N + j, 0,
-                                                     -T.infinity(acc_s.dtype))
-                else:
-                    T.clear(acc_s)
-                T.gemm(Q_shared, K_shared, acc_s, transpose_B=True, policy=T.GemmWarpPolicy.FullRow)
-                T.copy(V[bz, k * block_N:(k + 1) * block_N, by // groups, :], V_shared)
-                T.copy(scores_max, scores_max_prev)
-                T.reduce_max(acc_s, scores_max, dim=1, clear=False)
+                        if acc_s[i, j] == -T.infinity(acc_s.dtype):
+                            acc_s[i, j] = 0.0
+                        else:
+                            acc_s[i, j] = T.exp(acc_s[i, j] - m_i[i])
+
+                    T.reduce_sum(acc_s, row_sum, dim=1)
+                    for i in T.Parallel(block_M):
+                        l_i[i] += row_sum[i]
+
+                    T.copy(acc_s, acc_s_cast)
+
+                    T.gemm(acc_s_cast, V_shared, acc_o, policy=GemmWarpPolicy.FullRow)
+
+                l_inv = T.alloc_fragment([block_M], accum_dtype)
                 for i in T.Parallel(block_M):
-                    scores_scale[i] = T.exp2(scores_max_prev[i] * scale - scores_max[i] * scale)
-                for i, j in T.Parallel(block_M, dim_v):
-                    acc_o[i, j] *= scores_scale[i]
-                for i, j in T.Parallel(block_M, block_N):
-                    acc_s[i, j] = T.exp2(acc_s[i, j] * scale - scores_max[i] * scale)
-                T.copy(acc_s, acc_s_cast)
-                T.gemm(acc_s_cast, V_shared, acc_o, policy=T.GemmWarpPolicy.FullRow)
-                T.reduce_sum(acc_s, scores_sum, dim=1)
+                    safe_l = T.if_then_else(l_i[i] > 1e-6, l_i[i], 1.0)
+                    l_inv[i] = 1.0 / safe_l
+
+                for i, j in T.Parallel(block_M, dim):
+                    Output[bz, q_block_offset + i, by, j] = acc_o[i, j] * l_inv[i]
+
                 for i in T.Parallel(block_M):
-                    logsum[i] = logsum[i] * scores_scale[i] + scores_sum[i]
-            for i, j in T.Parallel(block_M, dim_v):
-                acc_o[i, j] /= logsum[i]
-            T.copy(acc_o, Output[bz, bx * block_M:(bx + 1) * block_M, by, :])
-            for i in T.Parallel(block_M):
-                logsum[i] = T.log2(logsum[i]) + scores_max[i] * scale
-            T.copy(logsum, lse[bz, by, bx * block_M:(bx + 1) * block_M])
+                    if q_block_offset + i < seq_len:
+                        lse_val = T.if_then_else(l_i[i] > 0,
+                                                 T.log(l_i[i]) + m_i[i], -T.infinity(accum_dtype))
+                        LSE[bz, by, q_block_offset + i] = lse_val
+
+                bx_loop_var = current_bx + num_split_q
+
+    return main
+
 
-    return flash_fwd
+def get_bwd_configs():
+    block_M = [16, 32, 64, 128, 256]
+    block_N = [16, 32, 64, 128, 256]
+    threads = [64, 128, 256, 512, 1024]
+    num_stages = [0, 1, 2]
+    enable_rasterization = [True]
+    panel_size = [7, 8, 9, 10]
+
+    configs = []
+    for m, n, stages, t, r, p in itertools.product(block_M, block_N, num_stages, threads,
+                                                   enable_rasterization, panel_size):
+        configs.append({
+            "block_M": m,
+            "block_N": n,
+            "num_stages": stages,
+            "threads": t,
+            "enable_rasterization": r,
+            "panel_size": p,
+        })
+
+    return configs
 
 
 @tilelang.jit(out_idx=[2])
-def flashattn_bwd_preprocess(batch, heads, seq_len, dim_v):
+def flashattn_bwd_preprocess(batch, heads, seq_len, dim):
     dtype = "float16"
     accum_dtype = "float"
-    shape = [batch, seq_len, heads, dim_v]
+    shape = [batch, seq_len, heads, dim]
     blk = 32
 
     @T.prim_func
-    def flash_bwd_prep(
-            O: T.Tensor(shape, dtype),  # type: ignore
-            dO: T.Tensor(shape, dtype),  # type: ignore
-            Delta: T.Tensor([batch, heads, seq_len], accum_dtype),  # type: ignore
-    ):
-        with T.Kernel(heads, T.ceildiv(seq_len, blk), batch) as (bx, by, bz):
+    def flash_bwd_prep(O: T.Tensor(shape, dtype), dO: T.Tensor(shape, dtype),
+                       Delta: T.Tensor([batch, heads, seq_len], accum_dtype)):
+        with T.Kernel(batch, heads, T.ceildiv(seq_len, blk)) as (bz, bx, by):
             o = T.alloc_fragment([blk, blk], dtype)
             do = T.alloc_fragment([blk, blk], dtype)
             acc = T.alloc_fragment([blk, blk], accum_dtype)
             delta = T.alloc_fragment([blk], accum_dtype)
             T.clear(acc)
-            for k in range(T.ceildiv(dim_v, blk)):
+            for k in range(T.ceildiv(dim, blk)):
                 T.copy(O[bz, by * blk:(by + 1) * blk, bx, k * blk:(k + 1) * blk], o)
                 T.copy(dO[bz, by * blk:(by + 1) * blk, bx, k * blk:(k + 1) * blk], do)
                 for i, j in T.Parallel(blk, blk):
@@ -107,256 +273,330 @@ def flashattn_bwd_preprocess(batch, heads, seq_len, dim_v):
     return flash_bwd_prep
 
 
-def make_dq_layout(dQ):
-    # atomicAdd can not be vectorized, so we need to reorder dq to match the 8x8 gemm fragment
-    return T.Layout(dQ.shape,
-                    lambda b, l, h, d: [b, l // 8, h, d // 8, (d % 2), 4 * (l % 8) + (d % 8) // 2])
-
-
-@tilelang.jit(out_idx=[1])
-def flashattn_bwd_postprocess(batch, heads, seq_len, dim_qk):
-    dtype = "float16"
-    accum_dtype = "float"
-    shape = [batch, seq_len, heads, dim_qk]
-    blk = 64
-
-    @T.prim_func
-    def flash_bwd_post(
-            dQ: T.Tensor(shape, accum_dtype),  # type: ignore
-            dQ_out: T.Tensor(shape, dtype),  # type: ignore
-    ):
-        with T.Kernel(T.ceildiv(seq_len, blk), heads, batch, threads=128) as (bx, by, bz):
-            T.copy(
-                dQ[bz, bx * blk:(bx + 1) * blk, by, :],
-                dQ_out[bz, bx * blk:(bx + 1) * blk, by, :],
-            )
-
-    return flash_bwd_post
-
-
+@tilelang.autotune(configs=get_bwd_configs(), cache_input_tensors=True)
 @tilelang.jit
-def flashattn_bwd(batch, heads, seq_len, dim_qk, dim_v, is_causal, block_M, block_N, groups=1):
-    sm_scale = (1.0 / dim_qk)**0.5
-    scale = (1.0 / dim_qk)**0.5 * 1.44269504  # log2(e)
+def flashattn_bwd(batch, heads, seq_len, dim, is_causal, groups, block_M: int, block_N: int,
+                  num_stages: int, threads: int, enable_rasterization: bool, panel_size: int):
+    sm_scale = (1.0 / dim)**0.5
     head_kv = heads // groups
-    q_shape = [batch, seq_len, heads, dim_qk]
-    k_shape = [batch, seq_len, head_kv, dim_qk]
-    v_shape = [batch, seq_len, head_kv, dim_v]
+    q_shape = [batch, seq_len, heads, dim]
+    kv_shape = [batch, seq_len, head_kv, dim]
     dtype = "float16"
     accum_dtype = "float"
 
     @T.prim_func
-    def flash_bwd(
-            Q: T.Tensor(q_shape, dtype),  # type: ignore
-            K: T.Tensor(k_shape, dtype),  # type: ignore
-            V: T.Tensor(v_shape, dtype),  # type: ignore
-            dO: T.Tensor([batch, seq_len, heads, dim_v], dtype),  # type: ignore
-            lse: T.Tensor([batch, heads, seq_len], accum_dtype),  # type: ignore
-            Delta: T.Tensor([batch, heads, seq_len], accum_dtype),  # type: ignore
-            dQ: T.Tensor(q_shape, accum_dtype),  # type: ignore
-            dK: T.Tensor(k_shape, accum_dtype),  # type: ignore
-            dV: T.Tensor(v_shape, accum_dtype),  # type: ignore
-    ):
-        with T.Kernel(heads, T.ceildiv(seq_len, block_M), batch, threads=128) as (bx, by, bz):
-            K_shared = T.alloc_shared([block_M, dim_qk], dtype)
-            dsT_shared = T.alloc_shared([block_M, block_N], dtype)
-            q = T.alloc_shared([block_N, dim_qk], dtype)
-            V_shared = T.alloc_shared([block_M, dim_v], dtype)
-            qkT = T.alloc_fragment([block_M, block_N], accum_dtype)
-            dsT = T.alloc_fragment([block_M, block_N], accum_dtype)
-            qkT_cast = T.alloc_fragment([block_M, block_N], dtype)
-            dsT_cast = T.alloc_fragment([block_M, block_N], dtype)
+    def flash_bwd_kernel(Q: T.Tensor(q_shape,
+                                     dtype), K: T.Tensor(kv_shape,
+                                                         dtype), V: T.Tensor(kv_shape, dtype),
+                         dO: T.Tensor(q_shape, dtype), lse: T.Tensor([batch, heads, seq_len],
+                                                                     accum_dtype),
+                         Delta: T.Tensor([batch, heads, seq_len],
+                                         accum_dtype), dQ: T.Tensor(q_shape, accum_dtype),
+                         dK: T.Tensor(kv_shape, accum_dtype), dV: T.Tensor(kv_shape, accum_dtype)):
+        with T.Kernel(heads, T.ceildiv(seq_len, block_M), batch, threads=threads) as (bx, by, bz):
+            T.use_swizzle(panel_size, enable=enable_rasterization)
+
+            K_shared = T.alloc_shared([block_M, dim], dtype)
+            V_shared = T.alloc_shared([block_M, dim], dtype)
+            q_shared = T.alloc_shared([block_N, dim], dtype)
+            do_shared = T.alloc_shared([block_N, dim], dtype)
             lse_shared = T.alloc_shared([block_N], accum_dtype)
-            delta = T.alloc_shared([block_N], accum_dtype)
-            do = T.alloc_shared([block_N, dim_v], dtype)
-            dv = T.alloc_fragment([block_M, dim_v], accum_dtype)
-            dk = T.alloc_fragment([block_M, dim_qk], accum_dtype)
-            dq = T.alloc_fragment([block_N, dim_qk], accum_dtype)
+            delta_shared = T.alloc_shared([block_N], accum_dtype)
+            ds_shared = T.alloc_shared([block_M, block_N], dtype)
+
+            p_cast = T.alloc_fragment([block_M, block_N], dtype)
+            qkT = T.alloc_fragment([block_M, block_N], accum_dtype)
+            P_acc = T.alloc_fragment([block_M, block_N], accum_dtype)
+            dP = T.alloc_fragment([block_M, block_N], accum_dtype)
+
+            dv = T.alloc_fragment([block_M, dim], accum_dtype)
+            dk = T.alloc_fragment([block_M, dim], accum_dtype)
+            dq = T.alloc_fragment([block_N, dim], accum_dtype)
 
             T.copy(K[bz, by * block_M:(by + 1) * block_M, bx // groups, :], K_shared)
             T.copy(V[bz, by * block_M:(by + 1) * block_M, bx // groups, :], V_shared)
             T.clear(dv)
             T.clear(dk)
+
             loop_st = T.floordiv(by * block_M, block_N) if is_causal else 0
             loop_ed = T.ceildiv(seq_len, block_N)
-            for k in T.Pipelined(loop_st, loop_ed, num_stages=1):
-                T.copy(Q[bz, k * block_N:(k + 1) * block_N, bx, :], q)
+
+            for k in T.Pipelined(loop_st, loop_ed, num_stages=num_stages):
+                T.copy(Q[bz, k * block_N:(k + 1) * block_N, bx, :], q_shared)
                 T.clear(qkT)
-                T.gemm(K_shared, q, qkT, transpose_B=True, policy=T.GemmWarpPolicy.FullRow)
+
+                T.gemm(K_shared, q_shared, qkT, transpose_B=True, policy=T.GemmWarpPolicy.FullRow)
+
                 T.copy(lse[bz, bx, k * block_N:(k + 1) * block_N], lse_shared)
+
                 for i, j in T.Parallel(block_M, block_N):
-                    qkT[i, j] = T.exp2(qkT[i, j] * scale - lse_shared[j])
+                    P_acc[i, j] = T.exp(qkT[i, j] * sm_scale - lse_shared[j])
+
                 if is_causal:
                     for i, j in T.Parallel(block_M, block_N):
-                        qkT[i, j] = T.if_then_else(by * block_M + i <= k * block_N + j, qkT[i, j],
-                                                   0)
-                T.copy(dO[bz, k * block_N:(k + 1) * block_N, bx, :], do)
-                T.clear(dsT)
-                T.gemm(V_shared, do, dsT, transpose_B=True, policy=T.GemmWarpPolicy.FullRow)
-                T.copy(qkT, qkT_cast)
-                T.gemm(qkT_cast, do, dv, policy=T.GemmWarpPolicy.FullRow)
+                        P_acc[i, j] = T.if_then_else(by * block_M + i <= k * block_N + j,
+                                                     P_acc[i, j], 0.0)
+
+                T.copy(dO[bz, k * block_N:(k + 1) * block_N, bx, :], do_shared)
+                T.clear(dP)
+
+                T.gemm(V_shared, do_shared, dP, transpose_B=True, policy=T.GemmWarpPolicy.FullRow)
+
+                T.copy(P_acc, p_cast)
+                T.gemm(p_cast, do_shared, dv, policy=T.GemmWarpPolicy.FullRow)
 
-                T.copy(Delta[bz, bx, k * block_N:(k + 1) * block_N], delta)
+                T.copy(Delta[bz, bx, k * block_N:(k + 1) * block_N], delta_shared)
 
                 for i, j in T.Parallel(block_M, block_N):
-                    dsT_cast[i, j] = qkT[i, j] * (dsT[i, j] - delta[j]) * sm_scale
-                T.gemm(dsT_cast, q, dk, policy=T.GemmWarpPolicy.FullRow)
+                    p_cast[i, j] = P_acc[i, j] * (dP[i, j] - delta_shared[j]) * sm_scale
 
-                T.copy(dsT_cast, dsT_shared)
+                T.gemm(p_cast, q_shared, dk, policy=T.GemmWarpPolicy.FullRow)
+
+                T.copy(p_cast, ds_shared)
                 T.clear(dq)
-                T.gemm(dsT_shared, K_shared, dq, transpose_A=True)
-                for i, j in T.Parallel(block_N, dim_qk):
+                T.gemm(ds_shared, K_shared, dq, transpose_A=True)
+                for i, j in T.Parallel(block_N, dim):
                     T.atomic_add(dQ[bz, k * block_N + i, bx, j], dq[i, j])
 
-            for i, j in T.Parallel(block_M, dim_v):
+            for i, j in T.Parallel(block_M, dim):
                 T.atomic_add(dV[bz, by * block_M + i, bx // groups, j], dv[i, j])
-            for i, j in T.Parallel(block_M, dim_qk):
                 T.atomic_add(dK[bz, by * block_M + i, bx // groups, j], dk[i, j])
 
-    return flash_bwd
-
-
-@torch.compile
-class _attention(torch.autograd.Function):
-
-    @staticmethod
-    def forward(ctx, q, k, v, causal, groups=1):
-        BATCH, N_CTX, H, D_HEAD_QK = q.shape
-        D_HEAD_V = v.shape[-1]
-        block_M = 128
-        block_N = 64
-        mod = flashattn_fwd(BATCH, H, N_CTX, D_HEAD_QK, D_HEAD_V, causal, block_M, block_N, groups)
-        o, lse = mod(q, k, v)
-        ctx.save_for_backward(q, k, v, o, lse)
-        ctx.causal = causal
-        return o
-
-    @staticmethod
-    def backward(ctx, do):
-        q, k, v, o, lse = ctx.saved_tensors
-        BATCH, N_CTX, H, D_HEAD_QK = q.shape
-        HEAD_KV, D_HEAD_V, = v.shape[-2], v.shape[-1]
-        groups = H // HEAD_KV
-
-        def maybe_contiguous(x):
-            if x.stride(-1) != 1:
-                return x.contiguous()
-            return x
-
-        do, q, k, v, o = [maybe_contiguous(x) for x in (do, q, k, v, o)]
-        block_M = 64
-        block_N = 32
-        mod_prep = flashattn_bwd_preprocess(BATCH, H, N_CTX, D_HEAD_V)
-        mod_post = flashattn_bwd_postprocess(BATCH, H, N_CTX, D_HEAD_QK)
-        delta = mod_prep(o, do)
-        kernel = flashattn_bwd(BATCH, H, N_CTX, D_HEAD_QK, D_HEAD_V, ctx.causal, block_M, block_N,
-                               groups)
-        shape_q = [BATCH, N_CTX, H, D_HEAD_QK]
-        shape_k = [BATCH, N_CTX, HEAD_KV, D_HEAD_QK]
-        shape_v = [BATCH, N_CTX, HEAD_KV, D_HEAD_V]
-        dq = torch.zeros(shape_q, dtype=torch.float32, device=q.device)
-        dk = torch.zeros(shape_k, dtype=torch.float32, device=q.device)
-        dv = torch.zeros(shape_v, dtype=torch.float32, device=q.device)
-        kernel(q, k, v, do, lse, delta, dq, dk, dv)
-        dq = mod_post(dq)
-        return dq, dk, dv, None, None
-
-
-attention = _attention.apply
+    return flash_bwd_kernel
 
 
-def ref_program(Q, K, V, is_causal, groups=1):
-    # Q: [B, T, HQ, D_QK]
-    # K: [B, T, HK, D_QK]
-    # V: [B, T, HV, D_V]
-    # HQ = HKV * groups
-    assert Q.size(2) == K.size(
-        2) * groups, f"Q.size(2): {Q.size(2)}, K.size(2): {K.size(2)}, groups: {groups}"
-    assert Q.size(2) == V.size(
-        2) * groups, f"Q.size(2): {Q.size(2)}, V.size(2): {V.size(2)}, groups: {groups}"
-
-    dim_qk = Q.size(-1)
-    K = K.repeat_interleave(groups, dim=2)
-    V = V.repeat_interleave(groups, dim=2)
-    scores = torch.einsum('bqhd,bkhd->bhqk', Q, K)
-    scores = scores / torch.sqrt(torch.tensor(dim_qk, dtype=scores.dtype))
-    if is_causal:
-        seq_len = Q.size(1)
-        mask = torch.tril(torch.ones(seq_len, seq_len, device=scores.device))
-        mask = mask.unsqueeze(0).unsqueeze(0)
-        scores = scores.masked_fill(mask == 0, float('-inf'))
-    attention_weights = F.softmax(scores, dim=-1)
-    output = torch.einsum('bhqk,bkhd->bqhd', attention_weights, V)
-    return output
-
-
-def main(BATCH: int = 1,
-         H: int = 32,
-         N_CTX: int = 256,
-         D_HEAD_QK: int = 192,
-         D_HEAD_V: int = 128,
-         groups: int = 16,
-         causal: bool = False):
-    flops_per_qk = 2.0 * BATCH * H * N_CTX * N_CTX * D_HEAD_QK
-    flops_per_v = 2.0 * BATCH * H * N_CTX * N_CTX * D_HEAD_V
-    total_flops = 3 * flops_per_qk + 2 * flops_per_v
-    if causal:
-        total_flops *= 0.5
-    Q = (
-        torch.empty(BATCH, N_CTX, H, D_HEAD_QK, dtype=torch.half,
-                    device="cuda").normal_().requires_grad_())
-
-    head_kv = H // groups
-    K = (
-        torch.empty(BATCH, N_CTX, head_kv, D_HEAD_QK, dtype=torch.half,
-                    device="cuda").normal_().requires_grad_())
-    V = (
-        torch.empty(BATCH, N_CTX, head_kv, D_HEAD_V, dtype=torch.half,
-                    device="cuda").normal_().requires_grad_())
-    dO = (
-        torch.empty(BATCH, N_CTX, H, D_HEAD_V, dtype=torch.half,
-                    device="cuda").normal_().requires_grad_())
-    O = attention(Q, K, V, causal, groups)
-    O.backward(dO, retain_graph=True)
-    dQ, Q.grad = Q.grad.clone(), None
-    dK, K.grad = K.grad.clone(), None
-    dV, V.grad = V.grad.clone(), None
-
-    O_ref = ref_program(Q, K, V, causal, groups)
-    O_ref.backward(dO, retain_graph=True)
-    dQ_ref, Q.grad = Q.grad.clone(), None
-    dK_ref, K.grad = K.grad.clone(), None
-    dV_ref, V.grad = V.grad.clone(), None
-
-    torch.testing.assert_close(O, O_ref, rtol=1e-2, atol=1e-2)
-    torch.testing.assert_close(dV, dV_ref, rtol=1e-2, atol=1e-2)
-    torch.testing.assert_close(dK, dK_ref, rtol=1e-2, atol=1e-2)
-    torch.testing.assert_close(dQ, dQ_ref, rtol=1e-2, atol=1e-2)
-
-    def run():
-        O_ref.backward(dO, retain_graph=True)
-
-    def run1():
-        O.backward(dO, retain_graph=True)
-
-    from tilelang.profiler import do_bench
-
-    latency = do_bench(run, warmup=500)
-    print("torch: {:.2f} ms".format(latency))
-    print("torch: {:.2f} TFlops".format(total_flops / latency * 1e-9))
-    latency = do_bench(run1, warmup=500)
-    print("tilelang: {:.2f} ms".format(latency))
-    print("tilelang: {:.2f} TFlops".format(total_flops / latency * 1e-9))
+@tilelang.jit(out_idx=[1])
+def flashattn_bwd_postprocess(batch, heads, seq_len, dim):
+    dtype = "float16"
+    accum_dtype = "float"
+    shape = [batch, seq_len, heads, dim]
+    blk = 64
+
+    @T.prim_func
+    def flash_bwd_post(dQ_in: T.Tensor(shape, accum_dtype), dQ_out: T.Tensor(shape, dtype)):
+        with T.Kernel(T.ceildiv(seq_len, blk), heads, batch, threads=128) as (bx, by, bz):
+            T.copy(
+                dQ_in[bz, bx * blk:(bx + 1) * blk, by, :],
+                dQ_out[bz, bx * blk:(bx + 1) * blk, by, :],
+            )
+
+    return flash_bwd_post
+
+
+def debug_tensor_comparison(tensor1, tensor2, name, rtol=1e-3, atol=1e-3):
+    print(f"\n=== {name} Comparison ===")
+    print(f"Shape: {tensor1.shape} vs {tensor2.shape}")
+    print(f"Data type: {tensor1.dtype} vs {tensor2.dtype}")
+    print(f"Device: {tensor1.device} vs {tensor2.device}")
+
+    diff = torch.abs(tensor1 - tensor2)
+    max_diff = diff.max().item()
+    mean_diff = diff.mean().item()
+    std_diff = diff.std().item()
+
+    print(f"Max difference: {max_diff:.6f}")
+    print(f"Mean difference: {mean_diff:.6f}")
+    print(f"Difference std: {std_diff:.6f}")
+
+    if max_diff > atol:
+        max_idx = torch.argmax(diff)
+        max_idx = np.unravel_index(max_idx.cpu().numpy(), tensor1.shape)
+        print(f"Max difference position: {max_idx}")
+        print(f"Value1: {tensor1[max_idx].item():.6f}, Value2: {tensor2[max_idx].item():.6f}")
+
+    nan_count1 = torch.isnan(tensor1).sum().item()
+    nan_count2 = torch.isnan(tensor2).sum().item()
+    inf_count1 = torch.isinf(tensor1).sum().item()
+    inf_count2 = torch.isinf(tensor2).sum().item()
+
+    print(f"NaN count: {nan_count1} vs {nan_count2}")
+    print(f"Inf count: {inf_count1} vs {inf_count2}")
+
+    relative_diff = diff / (torch.abs(tensor2) + 1e-8)
+    max_relative_diff = relative_diff.max().item()
+    mean_relative_diff = relative_diff.mean().item()
+
+    print(f"Max relative difference: {max_relative_diff:.6f}")
+    print(f"Mean relative difference: {mean_relative_diff:.6f}")
+
+    close = torch.allclose(tensor1, tensor2, rtol=rtol, atol=atol)
+    print(f"Within tolerance (rtol={rtol}, atol={atol}): {close}")
+
+    return close, max_diff, mean_diff
+
+
+def benchmark_function(func, *args, warmup=10, repeat=100):
+    for _ in range(warmup):
+        func(*args)
+
+    if torch.cuda.is_available():
+        torch.cuda.synchronize()
+
+    times = []
+    for _ in range(repeat):
+        start = time.time()
+        func(*args)
+        if torch.cuda.is_available():
+            torch.cuda.synchronize()
+        end = time.time()
+        times.append((end - start) * 1000)
+
+    return np.median(times)
+
+
+def main(batch: int = 1,
+         heads: int = 8,
+         seq_len: int = 4096,
+         dim: int = 128,
+         is_causal: bool = False,
+         groups: int = 1):
+
+    device = "cuda"
+    dtype = torch.float16
+
+    torch.manual_seed(42)
+    torch.cuda.manual_seed(42)
+
+    print(
+        f"Test configuration: batch={batch}, heads={heads}, seq_len={seq_len}, dim={dim}, is_causal={is_causal}, groups={groups}"
+    )
+
+    flops_per_gemm = 2.0 * batch * heads * seq_len * seq_len * dim
+    total_flops = 5 * flops_per_gemm
+
+    print(f"Total FLOPs: {total_flops / 1e12:.2f} TFlops")
+
+    q = torch.randn(batch, seq_len, heads, dim, device=device, dtype=dtype)
+    k = torch.randn(batch, seq_len, heads // groups, dim, device=device, dtype=dtype)
+    v = torch.randn(batch, seq_len, heads // groups, dim, device=device, dtype=dtype)
+    dO = torch.randn_like(q)
+
+    print("Starting autotuning for Fast FlashAttention-V2 Forward Pass...")
+    fwd_kernel = fast_flashattn(batch, heads, seq_len, dim, is_causal, groups)
+    if fwd_kernel is None or fwd_kernel.config is None:
+        print("Forward pass auto-tuning failed.")
+        return
+    print(f"Autotuning finished. Best Forward Configuration: {fwd_kernel.config}")
+
+    ref_program_processed = partial(ref_program, is_causal=is_causal, groups=groups)
+
+    profiler = fwd_kernel.get_profiler(tensor_supply_type=tilelang.TensorSupplyType.Normal)
+
+    print("Verifying correctness...")
+    profiler.assert_allclose(ref_program_processed, rtol=0.01, atol=0.01)
+    print("Forward pass is correct.")
+
+    o_tl, lse_tl = fwd_kernel(q, k, v)
+
+    bwd_prep = flashattn_bwd_preprocess(batch, heads, seq_len, dim)
+    delta_tl = bwd_prep(o_tl, dO)
+
+    print("\nStarting FlashAttention-V2 backward pass autotuning...")
+    bwd_kernel = flashattn_bwd(batch, heads, seq_len, dim, is_causal, groups)
+    if bwd_kernel is None or bwd_kernel.config is None:
+        print("Backward pass autotuning failed.")
+        return
+    print(f"Autotuning completed. Best backward pass configuration: {bwd_kernel.config}")
+
+    dQ_accum = torch.zeros_like(q, dtype=torch.float32)
+    dK_tl = torch.zeros_like(k, dtype=torch.float32)
+    dV_tl = torch.zeros_like(v, dtype=torch.float32)
+
+    bwd_kernel(q, k, v, dO, lse_tl, delta_tl, dQ_accum, dK_tl, dV_tl)
+
+    post_kernel = flashattn_bwd_postprocess(batch, heads, seq_len, dim)
+    dQ_tl = post_kernel(dQ_accum)
+
+    q_ref = q.clone().detach().requires_grad_()
+    k_ref = k.clone().detach().requires_grad_()
+    v_ref = v.clone().detach().requires_grad_()
+
+    o_ref, _ = ref_program(q_ref, k_ref, v_ref, is_causal, groups)
+    o_ref.backward(dO)
+
+    print("Verifying backward pass correctness...")
+    dq_close, dq_max_diff, dq_mean_diff = debug_tensor_comparison(
+        dQ_tl, q_ref.grad, "dQ", rtol=0.05, atol=0.05)
+    if dq_close:
+        print("dQ is correct.")
+    else:
+        print("dQ mismatch detected.")
+
+    dk_close, dk_max_diff, dk_mean_diff = debug_tensor_comparison(
+        dK_tl.to(torch.float16), k_ref.grad, "dK", rtol=0.05, atol=0.05)
+    if dk_close:
+        print("dK is correct.")
+    else:
+        print("dK mismatch detected.")
+
+    dv_close, dv_max_diff, dv_mean_diff = debug_tensor_comparison(
+        dV_tl.to(torch.float16), v_ref.grad, "dV", rtol=0.05, atol=0.05)
+    if dv_close:
+        print("dV is correct.")
+    else:
+        print("dV mismatch detected.")
+
+    print("\n=== Performance Benchmarking ===")
+
+    def run_reference_fwd_bwd():
+        q_ref_bench = q.clone().detach().requires_grad_()
+        k_ref_bench = k.clone().detach().requires_grad_()
+        v_ref_bench = v.clone().detach().requires_grad_()
+
+        o_ref_bench, _ = ref_program(q_ref_bench, k_ref_bench, v_ref_bench, is_causal, groups)
+
+        o_ref_bench.backward(dO)
+
+        if torch.cuda.is_available():
+            torch.cuda.synchronize()
+
+    ref_latency = benchmark_function(run_reference_fwd_bwd, warmup=10, repeat=100)
+    print(
+        f"Reference PyTorch Forward+Backward: {ref_latency:.2f} ms | {total_flops / ref_latency * 1e-9:.2f} TFlops"
+    )
+
+    def run_complete_fwd_bwd():
+        o_tl_bench, lse_tl_bench = fwd_kernel(q, k, v)
+
+        delta_tl_bench = bwd_prep(o_tl_bench, dO)
+
+        dQ_bench = torch.zeros_like(q, dtype=torch.float32)
+        dK_bench = torch.zeros_like(k, dtype=torch.float32)
+        dV_bench = torch.zeros_like(v, dtype=torch.float32)
+        bwd_kernel(q, k, v, dO, lse_tl_bench, delta_tl_bench, dQ_bench, dK_bench, dV_bench)
+
+        post_kernel(dQ_bench)
+
+        if torch.cuda.is_available():
+            torch.cuda.synchronize()
+
+    tile_latency = benchmark_function(run_complete_fwd_bwd, warmup=10, repeat=100)
+    print(
+        f"Complete Flash Attention V2 Forward+Backward (Tile-lang): {tile_latency:.2f} ms | {total_flops / tile_latency * 1e-9:.2f} TFlops"
+    )
+
+    speedup = ref_latency / tile_latency
+    print(f"Speedup: {speedup:.2f}x")
+
+    print("Forward output: Passed")
+    print(f"dQ: {'Passed' if dq_close else 'Failed'} (Max diff: {dq_max_diff:.6f})")
+    print(f"dK: {'Passed' if dk_close else 'Failed'} (Max diff: {dk_max_diff:.6f})")
+    print(f"dV: {'Passed' if dv_close else 'Failed'} (Max diff: {dv_max_diff:.6f})")
+
+    if all([dq_close, dk_close, dv_close]):
+        print("All checks passed!")
+    else:
+        print("Some checks failed, may need further debugging.")
 
 
 if __name__ == "__main__":
     parser = argparse.ArgumentParser()
-    parser.add_argument('--batch', type=int, default=8, help='Batch size')
-    parser.add_argument('--h', type=int, default=32, help='Number of heads')
-    parser.add_argument('--n_ctx', type=int, default=1024, help='Context size')
-    parser.add_argument('--d_head_qk', type=int, default=192, help='Head dimension for Q/K')
-    parser.add_argument('--d_head_v', type=int, default=128, help='Head dimension for V')
-    parser.add_argument('--causal', type=bool, default=False, help='Causal flag')
-    parser.add_argument('--groups', type=int, default=16, help='groups')
+    parser.add_argument('--batch', type=int, default=1, help='batch size')
+    parser.add_argument('--heads', type=int, default=8, help='heads')
+    parser.add_argument('--seq_len', type=int, default=1024, help='sequence length')
+    parser.add_argument('--dim', type=int, default=64, help='dim')
+    parser.add_argument('--is_causal', action='store_true', help='causal')
+    parser.add_argument('--groups', type=int, default=1, help='groups')
     args = parser.parse_args()
-    main(args.batch, args.h, args.n_ctx, args.d_head_qk, args.d_head_v, args.groups, args.causal)
+
+    main(args.batch, args.heads, args.seq_len, args.dim, args.is_causal, args.groups)

examples/amd/example_amd_flash_attn_fwd.py

@@ -34,7 +34,7 @@ def get_configs():
     block_N = [32, 64, 128, 256]
     threads = [128, 256, 512]
     num_split_q = [64, 128, 256]
-    num_stages = [0]
+    num_stages = [0, 1]
     enable_rasterization = [True]
     k_pack = [2]
     panel_size = [7, 8]
@@ -60,18 +60,6 @@ def get_configs():
             "qk_coalesced_width": qkw,
             "v_coalesced_width": vw,
         })
-    valid_configs.append({
-        'block_M': 64,
-        'block_N': 64,
-        'num_split_q': 64,
-        'threads': 256,
-        'num_stages': 1,
-        'enable_rasterization': True,
-        'k_pack': 2,
-        'panel_size': 64,
-        'qk_coalesced_width': 8,
-        'v_coalesced_width': 8,
-    })
     return valid_configs
 
 
@@ -95,7 +83,7 @@ def fast_flashattn(
     qk_coalesced_width: int,
     v_coalesced_width: int,
 ):
-    scale = (1.0 / dim)**0.5 * 1.44269504
+    scale = (1.0 / dim)**0.5
     head_kv = heads // groups
     q_shape = [batch, seq_len, heads, dim]
     kv_shape = [batch, seq_len, head_kv, dim]
@@ -185,7 +173,7 @@ def fast_flashattn(
                     T.reduce_max(acc_s, m_i, dim=1, clear=False)
 
                     for i in T.Parallel(block_M):
-                        sf = T.exp2(m_prev[i] * scale - m_i[i] * scale)
+                        sf = T.exp(m_prev[i] * scale - m_i[i] * scale)
                         l_i[i] *= sf
                         scale_factor[i] = sf
 
@@ -193,7 +181,7 @@ def fast_flashattn(
                         acc_o[i, j] *= scale_factor[i]
 
                     for i, j in T.Parallel(block_M, block_N):
-                        acc_s[i, j] = T.exp2(acc_s[i, j] * scale - m_i[i] * scale)
+                        acc_s[i, j] = T.exp(acc_s[i, j] * scale - m_i[i] * scale)
 
                     T.reduce_sum(acc_s, row_sum, dim=1)
                     for i in T.Parallel(block_M):

examples/amd/test.sh

-/root/miniconda3/envs/py312/bin/python3 examples/amd/example_amd_flash_attn_fwd.py \
-    --batch 2 \
-    --heads 16 \
-    --seq_len 4096 \
-    --dim 128 \
-    --is_causal \
-    --groups 2
-
-/root/composable_kernel/build/bin/tile_example_fmha_fwd  \
--b=2 -h=16 -s=4096 -d=128 -mask=t -v=1 -warmup=5 -repeat=20

examples/flash_attention/example_mha_bwd.py

@@ -38,14 +38,10 @@ def flashattn_fwd(batch, heads, seq_len, dim, is_causal, block_M, block_N):
             scores_sum = T.alloc_fragment([block_M], accum_dtype)
             logsum = T.alloc_fragment([block_M], accum_dtype)
 
-            T.annotate_layout({Q_shared: tilelang.layout.make_swizzled_layout(Q_shared)})
             T.copy(Q[bz, bx * block_M:(bx + 1) * block_M, by, :], Q_shared)
             T.fill(acc_o, 0)
             T.fill(logsum, 0)
             T.fill(scores_max, -T.infinity(accum_dtype))
-            # T.copy(Q_shared, Q_local)
-            # for i, j in T.Parallel(block_M, dim):
-            #     Q_local[i, j] *= scale
             loop_range = (
                 T.ceildiv(
                     (bx + 1) * block_M, block_N) if is_causal else T.ceildiv(seq_len, block_N))
@@ -192,9 +188,6 @@ def flashattn_bwd(batch, heads, seq_len, dim, is_causal, block_M, block_N):
 
             T.annotate_layout({
                 dQ: make_dq_layout(dQ),
-                K_shared: tilelang.layout.make_swizzled_layout(K_shared),
-                dv_shared: tilelang.layout.make_swizzled_layout(dv_shared),
-                dk_shared: tilelang.layout.make_swizzled_layout(dk_shared),
             })
             T.copy(K[bz, by * block_M:(by + 1) * block_M, bx, :], K_shared)
             T.copy(V[bz, by * block_M:(by + 1) * block_M, bx, :], V_shared)

src/target/codegen_hip.cc

@@ -41,10 +41,18 @@ static std::string GetFP8Type(DataType type) {
     stream << "fp8_e4" << vec << "_t";
   } else if (type.code() == DataType::kFloat8_e4m3fnuz) {
     stream << "fp8_e4" << vec << "_t";
+  } else if (type.code() == DataType::kFloat8_e4m3) {
+    stream << "fp8_e4" << vec << "_t";
+  } else if (type.code() == DataType::kFloat8_e4m3b11fnuz) {
+    stream << "fp8_e4" << vec << "_t";
   } else if (type.code() == DataType::kFloat8_e5m2) {
     stream << "fp8_e5" << vec << "_t";
+  } else if (type.code() == DataType::kFloat8_e5m2fnuz) {
+    stream << "fp8_e5" << vec << "_t";
+  } else if (type.code() == DataType::kFloat8_e8m0fnu) {
+    stream << "fp8_e8" << vec << "_t";
   } else {
-    LOG(FATAL) << "Unsupported FP8 type in HIP codegen";
+    LOG(FATAL) << "Unsupported FP8 type in HIP codegen: " << type;
   }
   return stream.str();
 }
@@ -926,10 +934,10 @@ void CodeGenTileLangHIP::VisitExpr_(const CallNode *op, std::ostream &os) {
         {"float8_e4m3fnuzx8", "long"},
         {"float32x16", "float32x16"}};
     std::string call_mfma_code = R"({
-    *((({C_dtype}*){c_ref}) + {c_bias}) = {mfma_buildin}(*((({A_dtype}*){a_ref}) + {a_bias}),
-                  *((({B_dtype}*){b_ref}) + {b_bias}),
-                  *((({C_dtype}*){c_ref}) + {c_bias}), 0, 0, 0);
-  })";
+      *((({C_dtype}*){c_ref}) + {c_bias}) = {mfma_buildin}(*((({A_dtype}*){a_ref}) + {a_bias}),
+                    *((({B_dtype}*){b_ref}) + {b_bias}),
+                    *((({C_dtype}*){c_ref}) + {c_bias}), 0, 0, 0);
+    })";
     std::string mfma_buildin = "__builtin_amdgcn_mfma_" + prefix;
     Replacer replacer;
 
@@ -955,6 +963,13 @@ void CodeGenTileLangHIP::VisitExpr_(const CallNode *op, std::ostream &os) {
                           op->args, true, os);
   } else if (op->op.same_as(tl::tl_gemm_sp())) {
     LOG(FATAL) << "tl_gemm_sp is not supported on HIP";
+  } else if (op->op.same_as(tl::loop_break())) {
+    this->PrintIndent();
+    this->stream << "break;\n";
+  } else if (op->op.same_as(tl::no_set_max_nreg())) {
+    // HIP doesn't need explicit register management like CUDA
+    // This is a no-op for HIP
+    return;
   } else {
     CodeGenC::VisitExpr_(op, os);
   }
@@ -1160,7 +1175,8 @@ inline void PrintConst(const FloatImmNode *op, std::ostream &os,
     os << "bfloat16_t";
     os << '(' << std::scientific << op->value << 'f' << ')';
     return;
-  } else if (op->dtype.is_float8_e4m3fnuz()) {
+  } else if (op->dtype.is_float8_e4m3fnuz() || op->dtype.is_float8_e4m3() ||
+             op->dtype.is_float8_e4m3fn()) {
     os << "fp8_e4_t";
     os << '(' << std::scientific << op->value << 'f' << ')';
     return;

src/tl_templates/hip/common.h

@@ -109,3 +109,13 @@ template <typename T1, typename T2>
 TL_DEVICE void AtomicAdd(T1 *address, T2 val) {
   atomicAdd(reinterpret_cast<T1 *>(address), static_cast<T1>(val));
 }
+
+// Overload for when the first argument is a value instead of a pointer
+template <typename T1, typename T2>
+TL_DEVICE void AtomicAdd(T1 address, T2 val) {
+  atomicAdd(reinterpret_cast<T1 *>(&address), static_cast<T1>(val));
+}
+
+template <typename T1, typename T2> TL_DEVICE T1 AtomicAddRet(T1 &ref, T2 val) {
+  return atomicAdd(&ref, static_cast<T1>(val));
+}

src/tl_templates/hip/gemm.h

@@ -70,7 +70,9 @@ template <int M, int N, int K, int num_warp_m, int num_warp_n, bool TransposeA,
           typename B_type, typename C_type, typename AccDataType = float>
 class GemmTensorOp {
 public:
-  static_assert(!clear_accum, "clear_accum=true is not supported yet");
+  // Note: clear_accum=true is not fully supported in HIP implementation
+  // but we'll handle it by manually clearing the accumulator
+  // static_assert(!clear_accum, "clear_accum=true is not supported yet");
 
   static constexpr int micro_size_x = 16;
   static constexpr int micro_size_y = 16;

src/tl_templates/hip/hip_fp8.h

@@ -5,6 +5,13 @@
 using fp8_e4_t = __hip_fp8_e4m3_fnuz;
 using fp8_e4_2_t = __hip_fp8x2_e4m3_fnuz;
 
+// Additional FP8 types for compatibility
+using fp8_e5_t = __hip_fp8_e5m2_fnuz;
+using fp8_e5_2_t = __hip_fp8x2_e5m2_fnuz;
+// Note: E8M0 types are not supported in current HIP version
+// using fp8_e8_t = __hip_fp8_e8m0_fnuz;
+// using fp8_e8_2_t = __hip_fp8x2_e8m0_fnuz;
+
 // Simple wrapper that provides member access for generated code
 struct fp8_e4_4_t {
   union {
@@ -43,6 +50,54 @@ struct __align__(16) fp8_e4_16_t {
   fp8_e4_8_t y;
 };
 
+// FP8 E5M2 vector types
+struct fp8_e5_4_t {
+  union {
+    __hip_fp8x4_e5m2_fnuz data;
+    struct {
+      fp8_e5_t x, y, z, w;
+    };
+  };
+  __device__ fp8_e5_4_t() = default;
+  __device__ fp8_e5_4_t(const __hip_fp8x4_e5m2_fnuz &val) : data(val) {}
+  __device__ operator __hip_fp8x4_e5m2_fnuz() const { return data; }
+};
+
+struct __align__(8) fp8_e5_8_t {
+  fp8_e5_4_t x;
+  fp8_e5_4_t y;
+};
+
+struct __align__(16) fp8_e5_16_t {
+  fp8_e5_8_t x;
+  fp8_e5_8_t y;
+};
+
+// FP8 E8M0 vector types - not supported in current HIP version
+/*
+struct fp8_e8_4_t {
+  union {
+    __hip_fp8x4_e8m0_fnuz data;
+    struct {
+      fp8_e8_t x, y, z, w;
+    };
+  };
+  __device__ fp8_e8_4_t() = default;
+  __device__ fp8_e8_4_t(const __hip_fp8x4_e8m0_fnuz &val) : data(val) {}
+  __device__ operator __hip_fp8x4_e8m0_fnuz() const { return data; }
+};
+
+struct __align__(8) fp8_e8_8_t {
+  fp8_e8_4_t x;
+  fp8_e8_4_t y;
+};
+
+struct __align__(16) fp8_e8_16_t {
+  fp8_e8_8_t x;
+  fp8_e8_8_t y;
+};
+*/
+
 __device__ fp8_e4_4_t make_fp8_e4_4_t(fp8_e4_t x, fp8_e4_t y, fp8_e4_t z,
                                       fp8_e4_t w) {
   // reinterpret the 4 fp8_e4_t values to signed char value and shift

testing/python/amd/test_tilelang_gemm_mfma_intrinsic.py

@@ -238,6 +238,12 @@ def test_assert_tl_matmul():
         128, 256, 256, "int8", "int32", b_transposed=False, accum_dtype="int32")
     assert_tl_matmul_correctness(
         128, 256, 256, "int8", "int32", b_transposed=False, accum_dtype="int32", k_pack=2)
+    assert_tl_matmul_correctness(128, 128, 128, "float8_e4m3fnuz", "float16")
+    assert_tl_matmul_correctness(128, 256, 256, "float8_e4m3fnuz", "float32")
+    assert_tl_matmul_correctness(128, 256, 256, "float8_e4m3fnuz", "float32", k_pack=2)
+    assert_tl_matmul_correctness(128, 256, 256, "float8_e4m3fnuz", "float32", b_transposed=False)
+    assert_tl_matmul_correctness(
+        128, 256, 256, "float8_e4m3fnuz", "float32", b_transposed=False, k_pack=2)
 
 
 if __name__ == "__main__":