[Enhancement] Refactor buffer index handling for improved precision a… #671

Alex4210987 · 2025-07-29T03:27:47Z

…nd 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.

…nd clarity (tile-ai#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.

github-actions · 2025-07-29T03:27:55Z

👋 Hi! Thank you for contributing to the TileLang project.

Please remember to run bash format.sh in the root directory of the project to ensure your changes are properly linted and formatted. This will help ensure your contribution passes the format check.

We appreciate you taking this step! Our team will review your contribution, and we look forward to your awesome work!

🚀

gemini-code-assist

Summary of Changes

Hello @Alex4210987, I'm Gemini Code Assist¹! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

This pull request adds a new example for FlashAttention-2 forward pass, specifically tailored for AMD platforms using tilelang. The implementation focuses on refining buffer index handling and overlap detection for greater precision and clarity, while also setting up robust performance tuning and correctness verification.

Highlights

New FlashAttention-2 Example: Introduced a new example (example_amd_flash_attn_fwd.py) demonstrating a FlashAttention-2 forward pass implementation using tilelang for AMD platforms. This example incorporates enhanced buffer index handling and streamlined logic for determining buffer overlaps, aligning with the PR's goal of improved precision and clarity.
Macro Refactoring for Clarity: Within the new FlashAttention-2 implementation, internal tilelang macros (compute_block and scale_and_write_back) were refactored and moved outside the main kernel definition. This change aims to improve code structure, clarity, and compiler understanding of the macro calls.
Performance and Correctness Verification: The new FlashAttention-2 example includes a PyTorch reference implementation for correctness verification and utilizes tilelang.autotune to find optimal configurations, providing benchmarking capabilities for performance comparison.

Using Gemini Code Assist

The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.

Invoking Gemini

You can request assistance from Gemini at any point in your pull request via creating an issue comment (i.e. comment on the pull request page) using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands.

Feature	Command	Description
Code Review	`/gemini review`	Performs a code review for the current pull request in its current state.
Pull Request Summary	`/gemini summary`	Provides a summary of the current pull request in its current state.
Comment	@gemini-code-assist	Responds in comments when explicitly tagged, both in issue comments and review comments.
Help	`/gemini help`	Displays a list of available commands.

Customization

To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.

Limitations & Feedback

Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments or fill out our survey to provide feedback.

You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.

Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩

gemini-code-assist

Code Review

This pull request adds a new example implementation of FlashAttention-V2 for AMD GPUs using tilelang, along with a shell script for running it. The new Python example is well-structured and demonstrates the capabilities of tilelang for complex kernel implementation.

My review focuses on improving code clarity, maintainability, and adherence to common best practices. I've pointed out a few areas for improvement:

Removing dead/unreachable code.
Improving readability by replacing a magic number with a named constant.
Ensuring consistency in comment language (English).
Making the test script more portable and compliant with POSIX standards.

gemini-code-assist · 2025-07-29T03:30:14Z

examples/amd/example_amd_flash_attn_fwd.py

+from functools import partial
+
+
+# PyTorch 参考实现保持不变


The comments in this file are in Chinese, which is inconsistent with the English used in the rest of the project. For better maintainability and accessibility for all contributors, please translate the comments to English.

Suggested change

# PyTorch 参考实现保持不变

# PyTorch reference implementation remains unchanged

gemini-code-assist · 2025-07-29T03:30:14Z

examples/amd/example_amd_flash_attn_fwd.py

+    if not valid_configs:
+        valid_configs.append({
+            'block_M': 64,
+            'block_N': 64,
+            'num_split_q': 64,
+            'threads': 256,
+            'num_stages': 1,
+            'enable_rasterization': True,
+            'k_pack': 2
+        })


This if not valid_configs: block appears to be unreachable. The itertools.product call on line 49 will only produce an empty iterator if one of the input iterables is empty. In this function, all lists used in the product are non-empty, so valid_configs will always contain at least one configuration. This block is dead code and should be removed.

gemini-code-assist · 2025-07-29T03:30:14Z

examples/amd/example_amd_flash_attn_fwd.py

+    enable_rasterization: bool,
+    k_pack: int,
+):
+    scale = (1.0 / dim)**0.5 * 1.44269504


The magic number 1.44269504 is 1/log(2), used for converting the base of the exponent from e to 2. To improve readability and maintainability, it's better to define it as a named constant.

LOG2_E = 1.44269504 # 1.0 / math.log(2) scale = (1.0 / dim)**0.5 * LOG2_E

gemini-code-assist · 2025-07-29T03:30:14Z

examples/amd/example_amd_flash_attn_fwd.py

+    dtype = "float16"
+    accum_dtype = "float"
+
+    dtype_size = 2


The variable dtype_size is defined but never used. It should be removed to avoid dead code.

gemini-code-assist · 2025-07-29T03:30:14Z

examples/amd/test.sh

@@ -0,0 +1,16 @@
+/bin/python /workspace/tilelang/examples/amd/example_amd_flash_attn_fwd.py     --batch 2     --seq_len 4096     --dim 128     --heads 16     --groups 8     --is_causal


This script is missing a shebang (e.g., #!/bin/bash or #!/usr/bin/env bash) which prevents it from being directly executable. Additionally, it uses a hardcoded, non-standard path for the python interpreter (/bin/python), which harms portability.

Please add a shebang and use a more standard way to invoke python, like python3.

Suggested change

/bin/python /workspace/tilelang/examples/amd/example_amd_flash_attn_fwd.py --batch 2 --seq_len 4096 --dim 128 --heads 16 --groups 8 --is_causal

#!/usr/bin/env bash

python3 /workspace/tilelang/examples/amd/example_amd_flash_attn_fwd.py --batch 2 --seq_len 4096 --dim 128 --heads 16 --groups 8 --is_causal

gemini-code-assist · 2025-07-29T03:30:15Z

examples/amd/test.sh

+#   ~/.tilelang/cache/0dad5010ab3d8e01b2a86e56208af334ad62865b92cb03038f9df84fda8d8a99/kernel.cu \
+#   -o ./kernel.s \
+#   -I/workspace/tilelang/3rdparty/composable_kernel/include \
+#   -I/workspace/tilelang/src


The file should end with a newline character. It's a POSIX standard for text files, and some tools might not process the last line correctly without it.

Suggested change

# -I/workspace/tilelang/src

# -I/workspace/tilelang/src

…s directory.

… code.

This reverts commit e8cc372.

…" (#672) This reverts commit e8cc372.

gemini-code-assist bot reviewed Jul 29, 2025

View reviewed changes

xinxyxiao added 2 commits July 29, 2025 03:30

Remove obsolete test script for AMD example, streamlining the example…

499daa3

…s directory.

Remove unused dtype_size variable in AMD example script to streamline…

555537a

… code.

Alex4210987 merged commit e8cc372 into tile-ai:main Jul 29, 2025
3 checks passed

LeiWang1999 added a commit that referenced this pull request Jul 29, 2025

Revert "[Enhancement] Add flash attn example for AMD MI300 series(#671)"

c07c458

This reverts commit e8cc372.

gemini-code-assist bot mentioned this pull request Jul 29, 2025

Reverts tile-ai/tilelang#671 #672

Merged

LeiWang1999 added a commit that referenced this pull request Jul 29, 2025

Revert "[Enhancement] Add flash attn example for AMD MI300 series(#671)…

56a8a64

…" (#672) This reverts commit e8cc372.

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

[Enhancement] Refactor buffer index handling for improved precision a… #671

[Enhancement] Refactor buffer index handling for improved precision a… #671

Uh oh!

Alex4210987 commented Jul 29, 2025

Uh oh!

github-actions bot commented Jul 29, 2025

Uh oh!

gemini-code-assist bot left a comment

Uh oh!

gemini-code-assist bot left a comment

Uh oh!

gemini-code-assist bot Jul 29, 2025

Uh oh!

gemini-code-assist bot Jul 29, 2025

Uh oh!

gemini-code-assist bot Jul 29, 2025

Uh oh!

gemini-code-assist bot Jul 29, 2025

Uh oh!

gemini-code-assist bot Jul 29, 2025

Uh oh!

gemini-code-assist bot Jul 29, 2025

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant

		from functools import partial


		# PyTorch 参考实现保持不变

	# PyTorch 参考实现保持不变
	# PyTorch reference implementation remains unchanged

		@@ -0,0 +1,16 @@
		/bin/python /workspace/tilelang/examples/amd/example_amd_flash_attn_fwd.py --batch 2 --seq_len 4096 --dim 128 --heads 16 --groups 8 --is_causal

	/bin/python /workspace/tilelang/examples/amd/example_amd_flash_attn_fwd.py --batch 2 --seq_len 4096 --dim 128 --heads 16 --groups 8 --is_causal
	#!/usr/bin/env bash
	python3 /workspace/tilelang/examples/amd/example_amd_flash_attn_fwd.py --batch 2 --seq_len 4096 --dim 128 --heads 16 --groups 8 --is_causal

[Enhancement] Refactor buffer index handling for improved precision a… #671

[Enhancement] Refactor buffer index handling for improved precision a… #671

Uh oh!

Conversation

Alex4210987 commented Jul 29, 2025

Uh oh!

github-actions bot commented Jul 29, 2025

Uh oh!

gemini-code-assist bot left a comment

Choose a reason for hiding this comment

Summary of Changes

Highlights

Footnotes

Uh oh!

gemini-code-assist bot left a comment

Choose a reason for hiding this comment

Code Review

Uh oh!

gemini-code-assist bot Jul 29, 2025

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist bot Jul 29, 2025

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist bot Jul 29, 2025

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist bot Jul 29, 2025

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist bot Jul 29, 2025

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist bot Jul 29, 2025

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant