Repository navigation
[AMD][Radeon] Add the Support of RDNA3/RDNA3.5(gfx11) WMMA - #2044
Conversation
|
👋 Hi! Thank you for contributing to the TileLang project. Please remember to run We appreciate you taking this step! Our team will review your contribution, and we look forward to your awesome work! 🚀 |
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
📝 WalkthroughWalkthroughAdds RDNA-generation awareness across HIP WMMA codegen and tilelang: targets are propagated into HIP codegen, RDNA generation is detected/validated (gfx11/gfx12), WMMA builtin naming, operand packing/typedefs, and WMMA layout/emitter logic are selected per generation. Changes
Sequence DiagramsequenceDiagram
participant Compiler as Compiler/Lowering
participant RtMod as rt_mod_hip.cc
participant CodeGen as CodeGenTileLangHIP
participant Utils as TargetGetRDNAGeneration
participant MacroGen as WMMAIntrinEmitter
participant LayoutLib as wmma_layout.py
Compiler->>RtMod: BuildTileLangHIP(target)
RtMod->>CodeGen: cg.Init(output_ssa)
RtMod->>CodeGen: cg.SetTarget(target)
RtMod->>CodeGen: cg.Finish()
CodeGen->>Utils: TargetGetRDNAGeneration(target)
Utils-->>CodeGen: rdna_gen (11/12/0)
CodeGen->>CodeGen: validate rdna_gen in {11,12}
CodeGen->>CodeGen: select builtin name, ab_half_elems, ab_vec_typedef
MacroGen->>Utils: detect rdna generation from target
Utils-->>MacroGen: rdna_gen
MacroGen->>LayoutLib: request per-gen layout/forward/replicate/store funcs
LayoutLib-->>MacroGen: gen-specific helpers
MacroGen->>CodeGen: emit WMMA call using selected names, typedefs, casts
Estimated code review effort🎯 4 (Complex) | ⏱️ ~45 minutes Possibly related PRs
Suggested reviewers
🚥 Pre-merge checks | ✅ 2 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (2 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
jiawei-real
changed the title
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 1
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
tilelang/tileop/gemm/__init__.py (1)
131-143:⚠️ Potential issue | 🔴 CriticalPass
mbar_phase_exprthrough to the implementation lowerer.Line 143 drops an argument that
tilelang/tileop/gemm/gemm_base.pystill defines andtilelang/tileop/gemm/gemm_tcgen05.pystill uses for barrier-phase selection. With the current call, TCGEN5 lowering will silently fall back to phase0, which can break pipelined mbarrier parity.Proposed fix
- return impl_class(self).lower(layout_map, target, thread_bounds, thread_var) + return impl_class(self).lower( + layout_map, target, thread_bounds, thread_var, mbar_phase_expr + )🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@tilelang/tileop/gemm/__init__.py` around lines 131 - 143, The lower method on the GEMM op drops the mbar_phase_expr argument when calling the implementation lowerer, causing implementations like gemm_tcgen05.py (and the base in gemm_base.py) that expect mbar_phase_expr to receive a default/0 phase; update the call in lower (and its callsite signature if needed) to pass mbar_phase_expr through to impl_class(self).lower so the selected implementation receives the barrier-phase expression; locate the lower method in tilelang/tileop/gemm/__init__.py and ensure its call to _select_gemm_instruction/_get_implementation_class then impl_class(self).lower(layout_map, target, thread_bounds, thread_var, mbar_phase_expr) matches the implementations in gemm_base.py and gemm_tcgen05.py.
🧹 Nitpick comments (2)
tilelang/intrinsics/wmma_macro_generator.py (1)
490-500: Maketargetmandatory instead of defaulting it toNone.
_detect_rdna_generation()now rejectsNoneunconditionally, butWMMAIntrinEmitter.__init__still advertisestarget: Target | None = None. That leaves a runtime trap for any caller that used the old default. Tightening the constructor signature would make the API match the new requirement.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@tilelang/intrinsics/wmma_macro_generator.py` around lines 490 - 500, The constructor signature of WMMAIntrinEmitter still allows target: Target | None = None while _detect_rdna_generation(Target | None) now rejects None; change WMMAIntrinEmitter.__init__ to require a non-optional Target (e.g., target: Target) and remove the default None, update any internal calls that passed None to pass a valid Target, and adjust all call sites/tests to provide a Target; ensure references to _detect_rdna_generation remain the same and that any callers that relied on the old default are updated to construct or pass a real Target instance.src/target/utils.cc (1)
291-299: ExportTargetGetRDNAGenerationthrough the existing FFI registry.Right now this helper is only callable from C++, so the Python side had to duplicate the same
mcpuparsing intilelang/intrinsics/wmma_macro_generator.py:490-500. Exposing it in theTVM_FFI_STATIC_INIT_BLOCK()below would keep RDNA target detection in one place.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@src/target/utils.cc` around lines 291 - 299, Add an FFI binding that exposes the existing C++ helper TargetGetRDNAGeneration to other languages: inside the existing TVM_FFI_STATIC_INIT_BLOCK() register a global/PackedFunc (e.g. "target.GetRDNAGeneration") whose body calls TargetGetRDNAGeneration(target) and returns the int result; ensure the registration signature accepts a Target (or ObjectRef) and returns an int so Python can call it instead of duplicating mcpu parsing. Use the exact symbol TargetGetRDNAGeneration in the registered body to locate the implementation.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@src/target/codegen_hip.cc`:
- Around line 1114-1160: The code emits __fp16 vector types unconditionally; add
a dtype guard that rejects any non-fp16 A/B before constructing
wmma_builtin/ab_vec_typedef/call_wmma_code: check the operand dtypes (e.g.
A_dtype and B_dtype) and require is_float16() for both, and if not,
ICHECK(false) or return an error describing unsupported non-fp16 WMMA operands;
keep the existing rdna_gen/wmma_builtin/ab_vec_typedef logic but only run it
after the dtype check so we fail fast (or clamp upstream) until bf16 codegen is
implemented.
---
Outside diff comments:
In `@tilelang/tileop/gemm/__init__.py`:
- Around line 131-143: The lower method on the GEMM op drops the mbar_phase_expr
argument when calling the implementation lowerer, causing implementations like
gemm_tcgen05.py (and the base in gemm_base.py) that expect mbar_phase_expr to
receive a default/0 phase; update the call in lower (and its callsite signature
if needed) to pass mbar_phase_expr through to impl_class(self).lower so the
selected implementation receives the barrier-phase expression; locate the lower
method in tilelang/tileop/gemm/__init__.py and ensure its call to
_select_gemm_instruction/_get_implementation_class then
impl_class(self).lower(layout_map, target, thread_bounds, thread_var,
mbar_phase_expr) matches the implementations in gemm_base.py and
gemm_tcgen05.py.
---
Nitpick comments:
In `@src/target/utils.cc`:
- Around line 291-299: Add an FFI binding that exposes the existing C++ helper
TargetGetRDNAGeneration to other languages: inside the existing
TVM_FFI_STATIC_INIT_BLOCK() register a global/PackedFunc (e.g.
"target.GetRDNAGeneration") whose body calls TargetGetRDNAGeneration(target) and
returns the int result; ensure the registration signature accepts a Target (or
ObjectRef) and returns an int so Python can call it instead of duplicating mcpu
parsing. Use the exact symbol TargetGetRDNAGeneration in the registered body to
locate the implementation.
In `@tilelang/intrinsics/wmma_macro_generator.py`:
- Around line 490-500: The constructor signature of WMMAIntrinEmitter still
allows target: Target | None = None while _detect_rdna_generation(Target | None)
now rejects None; change WMMAIntrinEmitter.__init__ to require a non-optional
Target (e.g., target: Target) and remove the default None, update any internal
calls that passed None to pass a valid Target, and adjust all call sites/tests
to provide a Target; ensure references to _detect_rdna_generation remain the
same and that any callers that relied on the old default are updated to
construct or pass a real Target instance.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro
Run ID: 607d155f-6357-41dc-9e58-3a5f9677ad4b
📒 Files selected for processing (8)
src/target/codegen_hip.ccsrc/target/codegen_hip.hsrc/target/rt_mod_hip.ccsrc/target/utils.ccsrc/target/utils.htilelang/intrinsics/wmma_layout.pytilelang/intrinsics/wmma_macro_generator.pytilelang/tileop/gemm/__init__.py
| // Get RDNA Generation | ||
| ICHECK(target_.defined()) << "CodeGenTileLangHIP target is not set"; | ||
| int rdna_gen = tvm::tl::TargetGetRDNAGeneration(target_); | ||
| ICHECK(rdna_gen == 11 || rdna_gen == 12) | ||
| << "Unsupported RDNA target for WMMA: gfx" << target_->str(); | ||
|
|
||
| // Determine wmma builtin name from shape | ||
| // shape = "f32_16x16x16_f16_w32" -> | ||
| // "__builtin_amdgcn_wmma_f32_16x16x16_f16_w32_gfx12" For gfx12 targets use | ||
| // the _gfx12 suffix variant. | ||
| std::string wmma_builtin = "__builtin_amdgcn_wmma_" + shape + "_gfx12"; | ||
| // the _gfx12 suffix variant, which gfx11 targets don't have. | ||
| std::string wmma_builtin = "__builtin_amdgcn_wmma_" + shape; | ||
|
|
||
| int ab_half_elems = 16; | ||
| std::string ab_vec_typedef = "tl_v16f16"; | ||
| if (rdna_gen == 12) { | ||
| wmma_builtin += "_gfx12"; | ||
| ab_half_elems = 8; | ||
| ab_vec_typedef = "tl_v8f16"; | ||
| } | ||
|
|
||
| // Emit the WMMA call. | ||
| // For gfx12: | ||
| // Signature: v8f32 = wmma_builtin(v8f16 a, v8f16 b, v8f32 c) | ||
| // where v8f16 = __fp16 x 8, v8f32 = float x 8. | ||
| // For gfx11: | ||
| // Signature: v8f32 = wmma_builtin(v16f16 a, v16f16 b, v8f32 c) | ||
| // where v16f16 = __fp16 x 16, v8f32 = float x 8. | ||
| // | ||
| // A/B buffers hold half_t (fp16), C/D buffers hold float. | ||
| // Each element index accesses a packed vector of 8 elements. | ||
| // Each element index accesses a packed vector of 8/16 elements. | ||
| // | ||
| // Using typedef'd vector types for the cast: | ||
| // typedef __attribute__((__vector_size__(8 * sizeof(__fp16)))) __fp16 | ||
| // tl_v8f16; typedef __attribute__((__vector_size__(8 * sizeof(float)))) | ||
| // typedef __attribute__((__vector_size__(8/16 * sizeof(__fp16)))) __fp16 | ||
| // tl_v8/16f16; typedef __attribute__((__vector_size__(8 * sizeof(float)))) | ||
| // float tl_v8f32; | ||
| std::string call_wmma_code = R"({ | ||
| typedef __attribute__((__vector_size__(8 * sizeof(__fp16)))) __fp16 tl_v8f16; | ||
| typedef __attribute__((__vector_size__({ab_half_elems} * sizeof(__fp16)))) __fp16 {ab_vec_typedef}; | ||
| typedef __attribute__((__vector_size__(8 * sizeof(float)))) float tl_v8f32; | ||
| *((tl_v8f32*){c_ref} + {c_bias}) = {wmma_builtin}( | ||
| *((tl_v8f16*){a_ref} + {a_bias}), | ||
| *((tl_v8f16*){b_ref} + {b_bias}), | ||
| *(({ab_vec_typedef}*){a_ref} + {a_bias}), | ||
| *(({ab_vec_typedef}*){b_ref} + {b_bias}), | ||
| *((tl_v8f32*){c_ref} + {c_bias})); | ||
| })"; | ||
| Replacer wmma_replacer; | ||
| wmma_replacer.register_rule("{wmma_builtin}", wmma_builtin); | ||
| wmma_replacer.register_rule("{ab_half_elems}", std::to_string(ab_half_elems)); | ||
| wmma_replacer.register_rule("{ab_vec_typedef}", ab_vec_typedef); |
There was a problem hiding this comment.
Reject non-float16 RDNA WMMA operands before emitting __fp16 vectors.
Lines 1126-1154 always lower A/B as tl_v8/16f16 backed by __fp16, regardless of the A_dtype / B_dtype passed into this intrinsic. If a non-fp16 configuration reaches this path, we'll generate the wrong builtin ABI instead of failing fast. Please add explicit dtype guards here, or clamp the upstream emitter to fp16-only until bf16 codegen exists.
Proposed guard
+ ICHECK(A_dtype == "float16x8" || A_dtype == "float16x16")
+ << "RDNA WMMA currently supports only float16 A operands";
+ ICHECK(B_dtype == "float16x8" || B_dtype == "float16x16")
+ << "RDNA WMMA currently supports only float16 B operands";
+
int ab_half_elems = 16;
std::string ab_vec_typedef = "tl_v16f16";🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@src/target/codegen_hip.cc` around lines 1114 - 1160, The code emits __fp16
vector types unconditionally; add a dtype guard that rejects any non-fp16 A/B
before constructing wmma_builtin/ab_vec_typedef/call_wmma_code: check the
operand dtypes (e.g. A_dtype and B_dtype) and require is_float16() for both, and
if not, ICHECK(false) or return an error describing unsupported non-fp16 WMMA
operands; keep the existing rdna_gen/wmma_builtin/ab_vec_typedef logic but only
run it after the dtype check so we fail fast (or clamp upstream) until bf16
codegen is implemented.
|
I am running this on AMD AI HX 370 with 890M on Ubuntu 24.04 LTS - linux 6.17.0-1017-oem (ROCM 7.2.1). Not successful. |
|
@petersktang Could you please give me more detailed information? For example, just run a SINGLE tuning config or the basic gemm kernel (examples/gemm/example_gemm.py)? |
And maybe your Linux version is not suitable: https://rocm.docs.amd.com/projects/radeon-ryzen/en/latest/docs/install/installryz/native_linux/install-ryzen.html "For ROCm on Ryzen, it is required to operate on the 6.14-1018 OEM kernel or newer." |
|
am already on linux 6.17.0 |
|
I ran python examples/gemm/example_gemm.py and got the below (gfx1150):
Traceback (most recent call last): ` File "...tilelang/tilelang/tileop/gemm/gemm_wmma.py", line 100, in lower File "...tilelang/tilelang/language/eager/builder.py", line 1286, in prim_func File "..i.tilelang/tilelang/language/eager/builder.py", line 1284, in impl File "...tilelang/tilelang/language/eager/builder.py", line 1277, in impl File "..i.tilelang/tilelang/tileop/gemm/gemm_wmma.py", line 107, in _gemm_ssr File "...tilelang/tilelang/intrinsics/wmma_macro_generator.py", line 212, in ldmatrix_a File "...tilelang/tilelang/language/eager/builder.py", line 781, in call File "...tilelang/tilelang/intrinsics/wmma_macro_generator.py", line 210, in _warp_ldmatrix_a File "...tilelang/3rdparty/tvm/python/tvm/tir/buffer.py", line 232, in getitem File "...tilelang/3rdparty/tvm/python/tvm/tir/expr.py", line 1120, in init File "python/tvm_ffi/cython/object.pxi", line 186, in tvm_ffi.core.Object.init_handle_by_constructor File "...tilelang/3rdparty/tvm/src/tir/ir/expr.cc", line 825, in tvm::tir::BufferLoad::BufferLoad(tvm::tir::Buffer, tvm::ffi::Array<tvm::PrimExpr, void>, tvm::ffi::Optional<tvm::PrimExpr, void>, tvm::Span) File "..tilelang/3rdparty/tvm/include/tvm/runtime/logging.h", line 321, in tvm::runtime::detail::LogFatal::~LogFatal() File "...tilelang/3rdparty/tvm/include/tvm/runtime/logging.h", line 337, in tvm::runtime::detail::LogFatal::Entry::Finalize() tvm.error.InternalError: Check failed: buffer->shape.size() == indices.size() (3 vs. 2) : Buffer A_shared is 3-dimensional, cannot be indexed with the 2-dimensional indices provided.` |
Thx for your reply! I've reproduced your problem, and it seems that something went wrong after merging updated main branch! Let me check if anything conflicts. |
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (1)
tilelang/intrinsics/wmma_macro_generator.py (1)
197-198: Potential staged-buffer misread if 3D buffers are passed unsliced: leading axes pinned to min=0.Line 197/243 extracts
A_other/B_otherfrom region mins. For a 2D buffer, this is empty; for rank > 2, it captures leading dimension mins (e.g.,[0]for stage axis). If a 3D staged buffer is passed directly without slicing,rkis effectively ignored and loads always read stage 0.Currently, all 46 callsites pass 2D buffers so this is not active. However, staged 3D buffers do exist in the codebase (e.g.,
example_warp_specialize_gemm_barrierpipe_stage2.py), and the method signature (in MFMA asBuffer | BufferRegion) permits them. Adding a runtime guard would prevent silent correctness errors if the interface is misused in the future.Proposed guard
@@ def ldmatrix_a(self, A_local_buf, A_shared_buf, ki, rk=0): A_region = self._legalize_to_buffer_region(A_shared_buf) + if isinstance(A_shared_buf, Buffer) and len(A_region.region) > 2: + raise ValueError( + "ldmatrix_a expects a 2D buffer or a sliced BufferRegion for rank > 2" + ) A_buf = A_region.buffer🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@tilelang/intrinsics/wmma_macro_generator.py` around lines 197 - 198, The code extracts A_other/B_other from A_region.region[:-2] and B_region.region[:-2] which silently captures leading axes for rank>2 and can mask staged-buffer misreads; add a runtime guard in the same scope (near where A_other, B_other and rk are used in wmma_macro_generator.py) that checks the region rank/leading dims and errors if a 3+D BufferRegion is passed unsliced: require len(A_region.region) <= 2 (and same for B_region) or verify all leading-region mins are 0 and document and raise a clear exception (ValueError) if not, so callers must explicitly slice staged 3D buffers before calling the function.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@tilelang/intrinsics/wmma_macro_generator.py`:
- Around line 206-207: The code uses a single-letter variable `l` in both A and
B buffer indexing which trips Ruff E741; rename these ambiguous locals to
descriptive names (e.g., replace `l` with `row_base` in the A-branch and with
`col_base` in the B-branch, or choose consistent names that match intent) and
update all uses in the A_local_buf and B_local_buf assignment expressions and
any subsequent calculations that reference `l` (look for symbols like
A_local_buf, B_local_buf, k_pack, chunk, ki, micro_size_k, warp_m,
warp_row_tiles, micro_size_x, local_size_a, local_size_b, local_id and replace
`l` in those expressions accordingly) so the indices and arithmetic remain
identical but no single-letter `l` remains.
---
Nitpick comments:
In `@tilelang/intrinsics/wmma_macro_generator.py`:
- Around line 197-198: The code extracts A_other/B_other from
A_region.region[:-2] and B_region.region[:-2] which silently captures leading
axes for rank>2 and can mask staged-buffer misreads; add a runtime guard in the
same scope (near where A_other, B_other and rk are used in
wmma_macro_generator.py) that checks the region rank/leading dims and errors if
a 3+D BufferRegion is passed unsliced: require len(A_region.region) <= 2 (and
same for B_region) or verify all leading-region mins are 0 and document and
raise a clear exception (ValueError) if not, so callers must explicitly slice
staged 3D buffers before calling the function.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro
Run ID: 23993cc1-c9f2-41ee-a64d-32f641bb9839
📒 Files selected for processing (1)
tilelang/intrinsics/wmma_macro_generator.py
| l, r = (rk * chunk + ki * (k_pack * micro_size_k), warp_m * warp_row_tiles + i * micro_size_x) | ||
| A_local_buf[i * k_pack * local_size_a + local_id] = A_buf[A_base0 + l + row, A_base1 + r + col] | ||
| A_local_buf[i * k_pack * local_size_a + local_id] = A_buf[ |
There was a problem hiding this comment.
Rename ambiguous l locals to satisfy Ruff E741.
Line 214 and Line 260 use l as a variable name; Ruff flags this as ambiguous (E741). Rename to explicit names (e.g., row_base, col_base) in both A/B branches.
Also applies to: 214-215, 252-253, 260-261
🧰 Tools
🪛 Ruff (0.15.9)
[error] 206-206: Ambiguous variable name: l
(E741)
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@tilelang/intrinsics/wmma_macro_generator.py` around lines 206 - 207, The code
uses a single-letter variable `l` in both A and B buffer indexing which trips
Ruff E741; rename these ambiguous locals to descriptive names (e.g., replace `l`
with `row_base` in the A-branch and with `col_base` in the B-branch, or choose
consistent names that match intent) and update all uses in the A_local_buf and
B_local_buf assignment expressions and any subsequent calculations that
reference `l` (look for symbols like A_local_buf, B_local_buf, k_pack, chunk,
ki, micro_size_k, warp_m, warp_row_tiles, micro_size_x, local_size_a,
local_size_b, local_id and replace `l` in those expressions accordingly) so the
indices and arithmetic remain identical but no single-letter `l` remains.
|
@petersktang I have found the conflict and already fixed it, could you plz 1. clear your local tilelang cache and 2. try again? I've tested locally and both of them passed. Root Cause
This issue was exposed by PR #2002 , which changed the pipeline flow so software-pipelined shared buffers could be multi-versioned earlier in lowering and gain an extra leading stage dimension, for example:
WMMA still assumed that those shared buffers were always 2-D, so it emitted 2-D indexing for a 3-D buffer. That mismatch caused errors such as: FixThe fix preserves any leading dimensions from the legalized In practice, the load path now:
This approach was implemented by following the corresponding MFMA shared-load logic in
This keeps WMMA shared loads compatible with both regular 2-D shared buffers and pipelined multi-versioned shared buffers. |
|
Results after running examples/gemm/example_gemm.py looks good: ` extern "C" global void launch_bounds(128) gemm_kernel(half_t* restrict A, half_t* restrict B, half_t* restrict C) { tilelang Latency: 0.543604384615385ms |
|
example_amd_flash_attn_fwd.py is successful.
Tuned Latency 21.765731811523438 with config {'block_M': 32, 'block_N': 32, 'num_split_q': 64, 'threads': 64, 'num_stages': 0, 'enable_rasterization': True, 'k_pack': 1, 'panel_size': 7, 'qk_coalesced_width': 8, 'v_coalesced_width': 4} at index 40 Tuned Latency 65.1008529663086 with config {'block_M': 32, 'block_N': 32, 'num_split_q': 32, 'threads': 32, 'num_stages': 0, 'enable_rasterization': True, 'k_pack': 1, 'panel_size': 8, 'qk_coalesced_width': 8, 'v_coalesced_width': 4} at index 41 |
That's it, bwd implementation does not perform special adaptation for RDNA (fwd does), which might be why bwd isn't working. |
|
@petersktang @zhangnju @LeiWang1999 BTW, I discovered a legacy issue from PR #1951 : layout inference in I made a quick-fixing version for this test, you can have a try: import sys
import torch
import torch.nn.functional as F
import tilelang
import tilelang.language as T
from tilelang.tileop.base import GemmWarpPolicy
import itertools
import argparse
from functools import partial
import numpy as np
import time
def IsRDNA():
if torch.cuda.is_available():
gpu_name = torch.cuda.get_device_name().strip()
return "Radeon" in gpu_name
else:
print("Error: GPU Device is not detected")
sys.exit(1)
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():
# Match the standalone forward example on RDNA. WMMA configs larger than
# 32x32 can trigger layout issues when bridging the softmax fragment into
# the second GEMM's A-layout.
if IsRDNA():
block_M = [16, 32, 64]
block_N = [16, 32, 64]
threads = [32, 64]
num_split_q = [16, 32, 64]
num_stages = [0]
k_pack = [1]
else:
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]
k_pack = [2]
enable_rasterization = [True]
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
):
if IsRDNA() and m == 16 and n == 16 and t == 64:
continue
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 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]
kv_shape = [batch, seq_len, head_kv, dim]
dtype = T.float16
accum_dtype = T.float32
vec_size = qk_coalesced_width
v_vec_size = v_coalesced_width
@T.prim_func
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(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(T.int32)
bx_loop_var = b_split
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)
# Bridge the WMMA D-layout softmax fragment into the A-layout
# expected by GEMM 2 on RDNA GPUs.
if IsRDNA():
P_shared = T.alloc_shared([block_M, block_N], 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):
m_i[i] = T.max(m_i[i], m_prev[i])
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):
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]
if IsRDNA():
for i, j in T.Parallel(block_M, block_N):
P_shared[i, j] = T.cast(acc_s[i, j], dtype)
T.copy(P_shared, acc_s_cast)
else:
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):
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):
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
def get_bwd_configs():
# Keep the RDNA search space aligned with the WMMA-friendly tile sizes
# verified above. Larger tiles and some warp/block combinations are either
# unsupported or known to trigger invalid lowering on RDNA.
if IsRDNA():
block_M = [16, 32]
block_N = [16, 32]
threads = [32, 64]
num_stages = [0]
panel_size = [7, 8]
else:
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]
panel_size = [7, 8, 9, 10]
enable_rasterization = [True]
configs = []
for m, n, stages, t, r, p in itertools.product(block_M, block_N, num_stages, threads, enable_rasterization, panel_size):
if IsRDNA() and m == 16 and n == 16 and t == 64:
continue
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):
dtype = T.float16
accum_dtype = T.float32
shape = [batch, seq_len, heads, dim]
blk = 32
@T.prim_func
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, 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):
acc[i, j] += o[i, j] * do[i, j]
T.reduce_sum(acc, delta, 1)
T.copy(delta, Delta[bz, bx, by * blk : (by + 1) * blk])
return flash_bwd_prep
@tilelang.autotune(configs=get_bwd_configs(), cache_input_tensors=True)
@tilelang.jit
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]
kv_shape = [batch, seq_len, head_kv, dim]
dtype = T.float16
accum_dtype = T.float32
@T.prim_func
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)
T.annotate_layout({dQ: make_dq_layout(dQ)})
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_shared = T.alloc_shared([block_N], accum_dtype)
ds_shared = T.alloc_shared([block_M, block_N], dtype)
if IsRDNA():
# Bridge the WMMA D-layout fragment produced by GEMM/elementwise
# ops into the A-layout expected by the following GEMM.
p_shared = T.alloc_shared([block_M, block_N], dtype)
ds_trans_shared = T.alloc_shared([block_N, block_M], 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=num_stages):
T.copy(Q[bz, k * block_N : (k + 1) * block_N, bx, :], q_shared)
T.clear(qkT)
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):
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):
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)
if IsRDNA():
for i, j in T.Parallel(block_M, block_N):
p_shared[i, j] = T.cast(P_acc[i, j], dtype)
T.copy(p_shared, p_cast)
else:
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_shared)
if IsRDNA():
for i, j in T.Parallel(block_M, block_N):
dP[i, j] = P_acc[i, j] * (dP[i, j] - delta_shared[j]) * sm_scale
for i, j in T.Parallel(block_M, block_N):
p_shared[i, j] = T.cast(dP[i, j], dtype)
T.copy(p_shared, p_cast)
T.gemm(p_cast, q_shared, dk, policy=T.GemmWarpPolicy.FullRow)
for i, j in T.Parallel(block_M, block_N):
ds_trans_shared[j, i] = T.cast(dP[i, j], dtype)
else:
for i, j in T.Parallel(block_M, block_N):
p_cast[i, j] = P_acc[i, j] * (dP[i, j] - delta_shared[j]) * sm_scale
T.gemm(p_cast, q_shared, dk, policy=T.GemmWarpPolicy.FullRow)
T.copy(p_cast, ds_shared)
T.clear(dq)
if IsRDNA():
T.gemm(ds_trans_shared, K_shared, dq, policy=T.GemmWarpPolicy.FullRow)
else:
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):
T.atomic_add(dV[bz, by * block_M + i, bx // groups, j], dv[i, j])
T.atomic_add(dK[bz, by * block_M + i, bx // groups, j], dk[i, j])
return flash_bwd_kernel
def make_dq_layout(dQ):
# atomicAdd can not be vectorized, so we need to reorder dq to match the
# 8x8 GEMM fragment layout before copying it back to a dense tensor.
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):
dtype = T.float16
accum_dtype = T.float32
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.annotate_layout({dQ_in: make_dq_layout(dQ_in)})
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=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.heads, args.seq_len, args.dim, args.is_causal, args.groups) |
|
Code result below (Target: hip -keys=hip,gpu -max_num_threads=256 -max_shared_memory_per_block=65536 -max_threads_per_block=256 -mcpu=gfx1150 -mtriple=amdgcn-amd-amdhsa-hcc -thread_warp_size=64): `Test configuration: batch=1, heads=8, seq_len=1024, dim=64, is_causal=False, groups=1 Starting FlashAttention-V2 backward pass autotuning... === dQ Comparison === === dK Comparison === === dV Comparison === === Performance Benchmarking === |
|
The issue of transpose_A=true is fixed: 
|
There was a problem hiding this comment.
Actionable comments posted: 2
🧹 Nitpick comments (1)
src/target/codegen_hip.cc (1)
1117-1118: Consider showing the mcpu value in the error message for clarity.The current error message uses
target_->str()which may produce verbose output. Showing the mcpu value directly would be more helpful for debugging.Suggested improvement
ICHECK(rdna_gen == 11 || rdna_gen == 12) - << "Unsupported RDNA target for WMMA: gfx" << target_->str(); + << "Unsupported RDNA target for WMMA: expected gfx11xx or gfx12xx, got mcpu=" + << Downcast<tvm::ffi::String>(target_->attrs.at("mcpu"));🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@src/target/codegen_hip.cc` around lines 1117 - 1118, The ICHECK that validates rdna_gen should include the specific mcpu value to make failures clearer: update the check message that currently prints target_->str() to also extract and show the target's mcpu (e.g., via target_->GetAttr<String>("mcpu") or the appropriate accessor) alongside the existing text so the logged message reads something like "Unsupported RDNA target for WMMA: gfx... (mcpu=...)" referencing the ICHECK, rdna_gen, and target_->str()/mcpu symbols.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@examples/amd/example_amd_flash_attn_bwd.py`:
- Around line 14-21: The IsRDNA() helper should not call sys.exit at import
time—replace the sys.exit(1) path with raising a descriptive exception (e.g.,
RuntimeError) so callers/tests can handle it, and make the function more robust
by expanding the GPU detection beyond a fragile "Radeon" substring (for example,
also check torch.version.hip or other ROCm/HIP indicators and accept "AMD" or
"Radeon" variants from torch.cuda.get_device_name()); also ensure any decorators
that call IsRDNA() run the check lazily (i.e., at runtime/validation time) or
handle the raised exception so imports/tests aren’t aborted.
In `@tilelang/intrinsics/wmma_layout.py`:
- Around line 268-271: The comment above the layout definition is mislabeled as
"gfx12 store index map" while the implementation and function name
(wmma_store_index_map_gfx11) implement the gfx11 layout; update that comment to
read "gfx11 store index map" (and any accompanying description that references
gfx12) so it correctly matches the function wmma_store_index_map_gfx11 and its
affine mapping description.
---
Nitpick comments:
In `@src/target/codegen_hip.cc`:
- Around line 1117-1118: The ICHECK that validates rdna_gen should include the
specific mcpu value to make failures clearer: update the check message that
currently prints target_->str() to also extract and show the target's mcpu
(e.g., via target_->GetAttr<String>("mcpu") or the appropriate accessor)
alongside the existing text so the logged message reads something like
"Unsupported RDNA target for WMMA: gfx... (mcpu=...)" referencing the ICHECK,
rdna_gen, and target_->str()/mcpu symbols.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro
Run ID: e2c3dfaa-8582-4583-89e2-f3b6abcccd07
📒 Files selected for processing (5)
examples/amd/example_amd_flash_attn_bwd.pysrc/target/codegen_hip.ccsrc/target/utils.cctilelang/intrinsics/wmma_layout.pytilelang/intrinsics/wmma_macro_generator.py
🚧 Files skipped from review as they are similar to previous changes (1)
- src/target/utils.cc
| def IsRDNA(): | ||
| if torch.cuda.is_available(): | ||
| gpu_name = torch.cuda.get_device_name().strip() | ||
| return "Radeon" in gpu_name | ||
| else: | ||
| print("Error: GPU Device is not detected") | ||
| sys.exit(1) | ||
|
|
There was a problem hiding this comment.
IsRDNA() may cause ungraceful exit at import time.
This function is called at module load time (via @tilelang.autotune(configs=get_fwd_configs(), ...) decorators). If CUDA is unavailable, the script exits with sys.exit(1) which could disrupt test harnesses or import-time checks. Additionally, the "Radeon" string match may be fragile for edge cases (e.g., future GPU naming).
Consider:
- Raising an exception instead of
sys.exit(1)to allow callers to handle gracefully. - Adding a fallback or more robust detection (e.g., checking for ROCm/HIP platform attributes).
Suggested improvement
def IsRDNA():
if torch.cuda.is_available():
gpu_name = torch.cuda.get_device_name().strip()
return "Radeon" in gpu_name
else:
- print("Error: GPU Device is not detected")
- sys.exit(1)
+ raise RuntimeError("IsRDNA() requires a CUDA/HIP device but none is available")🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@examples/amd/example_amd_flash_attn_bwd.py` around lines 14 - 21, The
IsRDNA() helper should not call sys.exit at import time—replace the sys.exit(1)
path with raising a descriptive exception (e.g., RuntimeError) so callers/tests
can handle it, and make the function more robust by expanding the GPU detection
beyond a fragile "Radeon" substring (for example, also check torch.version.hip
or other ROCm/HIP indicators and accept "AMD" or "Radeon" variants from
torch.cuda.get_device_name()); also ensure any decorators that call IsRDNA() run
the check lazily (i.e., at runtime/validation time) or handle the raised
exception so imports/tests aren’t aborted.
|
@jiawei-real, Thanks for your contributions on Radeon! Good job! have you tested your PR on other gfx11 GPUs (7900) or gfx12 GPU? |
Thanks to @zhangnju for the w7900 hardware support! Tested flashattn fwd/bwd, and both results are correct. The results have already been synced to the PR message. |
|
Tested example_amd_flash_attn_bwd.py on AMD 890M gfx1150 Target: hip -keys=hip,gpu -max_num_threads=256 -max_shared_memory_per_block=65536 -max_threads_per_block=256 -mcpu=gfx1150 -mtriple=amdgcn-amd-amdhsa-hcc -thread_warp_size=64 Starting FlashAttention-V2 backward pass autotuning... === dQ Comparison === === dK Comparison === === dV Comparison === === Performance Benchmarking === |
|
@petersktang Thanks for your validation, glad to know even 890M can run tilelang, which is good news for ROCm Tilelang ecosystem |
|
@zhangnju @petersktang Thanks again for your contributions to this work! |
|
@zhangnju @LeiWang1999 The correctness has already been verified on the 9070XT and W7900, and the verification results have been updated in the PR message. |
zhangnju
left a comment
zhangnju
left a comment
There was a problem hiding this comment.
TargetGetRDNAGeneration and _detect_rdna_generation() in wmma_macro_generator.py duplicates the same mcpu-string-parsing logic in Python
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (1)
tilelang/intrinsics/wmma_macro_generator.py (1)
380-391: Use#comments instead of a bare triple-quoted string.This triple-quoted block is an expression statement (not a docstring — the function’s first statements are
asserts above), so it allocates a string each call and isn't picked up by any tooling. Convert to regular#comments.♻️ Proposed change
- """ - gfx11 and gfx12 differ in how logical A/B fragment elements map to lanes. - - gfx11 duplicates each logical A/B element across the two half-waves - (lane t and lane t + 16). A single-owner forward_thread_fn cannot - faithfully represent this one-to-many ownership, so we model it with - T.Fragment(..., forward_fn=..., replicate=2), where `rep` selects the - lower/upper half-wave copy. - - gfx12 has a single unique owner for each logical element, so the - existing forward_thread_fn/forward_index_fn form is sufficient. - """ + # gfx11 and gfx12 differ in how logical A/B fragment elements map to lanes. + # + # gfx11 duplicates each logical A/B element across the two half-waves + # (lane t and lane t + 16). A single-owner forward_thread_fn cannot + # faithfully represent this one-to-many ownership, so we model it with + # T.Fragment(..., forward_fn=..., replicate=2), where `rep` selects the + # lower/upper half-wave copy. + # + # gfx12 has a single unique owner for each logical element, so the + # existing forward_thread_fn/forward_index_fn form is sufficient.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@tilelang/intrinsics/wmma_macro_generator.py` around lines 380 - 391, Replace the literal triple-quoted string block in wmma_macro_generator.py (the explanatory block about gfx11/gfx12 mapping near the assert lines) with regular line comments: delete the standalone """...""" expression and reformat its contents into one or more `#` comment lines immediately above the related code, preserving the exact text and indentation; ensure no runtime string is created and the explanatory text remains as comments describing the A/B fragment ownership behavior for gfx11/gfx12 in the same location (keep surrounding asserts and any variables unchanged).
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@src/target/codegen_hip.cc`:
- Around line 1196-1200: The error message currently prefixes a hardcoded "gfx"
to target_->str(), producing wrong output; update the ICHECK message in the code
around TargetGetRDNAGeneration(target_) to either remove the literal "gfx" (so
it prints the full target string via target_->str() only) or explicitly include
the mcpu value instead of "gfx" by querying the target's mcpu attribute (use the
Target/mcpu accessor used in your codebase) and print that; adjust the ICHECK
call that references target_, TargetGetRDNAGeneration, and target_->str()
accordingly.
---
Nitpick comments:
In `@tilelang/intrinsics/wmma_macro_generator.py`:
- Around line 380-391: Replace the literal triple-quoted string block in
wmma_macro_generator.py (the explanatory block about gfx11/gfx12 mapping near
the assert lines) with regular line comments: delete the standalone """..."""
expression and reformat its contents into one or more `#` comment lines
immediately above the related code, preserving the exact text and indentation;
ensure no runtime string is created and the explanatory text remains as comments
describing the A/B fragment ownership behavior for gfx11/gfx12 in the same
location (keep surrounding asserts and any variables unchanged).
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro
Run ID: 697326da-a1b3-44e3-bdaa-a328efaa75b3
📒 Files selected for processing (3)
src/target/codegen_hip.ccsrc/target/utils.cctilelang/intrinsics/wmma_macro_generator.py
🚧 Files skipped from review as they are similar to previous changes (1)
- src/target/utils.cc
| // Get RDNA Generation | ||
| ICHECK(target_.defined()) << "CodeGenTileLangHIP target is not set"; | ||
| int rdna_gen = tvm::tl::TargetGetRDNAGeneration(target_); | ||
| ICHECK(rdna_gen == 11 || rdna_gen == 12) | ||
| << "Unsupported RDNA target for WMMA: gfx" << target_->str(); |
There was a problem hiding this comment.
Error message concatenates a stray "gfx" in front of the full target string.
target_->str() returns the entire target description (e.g. "hip -keys=hip,gpu -mcpu=gfx1100 ..."), not just the arch suffix, so this will print Unsupported RDNA target for WMMA: gfxhip -keys=.... Drop the literal gfx, or look up mcpu explicitly.
🛠 Proposed fix
- ICHECK(rdna_gen == 11 || rdna_gen == 12)
- << "Unsupported RDNA target for WMMA: gfx" << target_->str();
+ ICHECK(rdna_gen == 11 || rdna_gen == 12)
+ << "Unsupported RDNA target for WMMA: " << target_->str();📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| // Get RDNA Generation | |
| ICHECK(target_.defined()) << "CodeGenTileLangHIP target is not set"; | |
| int rdna_gen = tvm::tl::TargetGetRDNAGeneration(target_); | |
| ICHECK(rdna_gen == 11 || rdna_gen == 12) | |
| << "Unsupported RDNA target for WMMA: gfx" << target_->str(); | |
| // Get RDNA Generation | |
| ICHECK(target_.defined()) << "CodeGenTileLangHIP target is not set"; | |
| int rdna_gen = tvm::tl::TargetGetRDNAGeneration(target_); | |
| ICHECK(rdna_gen == 11 || rdna_gen == 12) | |
| << "Unsupported RDNA target for WMMA: " << target_->str(); |
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@src/target/codegen_hip.cc` around lines 1196 - 1200, The error message
currently prefixes a hardcoded "gfx" to target_->str(), producing wrong output;
update the ICHECK message in the code around TargetGetRDNAGeneration(target_) to
either remove the literal "gfx" (so it prints the full target string via
target_->str() only) or explicitly include the mcpu value instead of "gfx" by
querying the target's mcpu attribute (use the Target/mcpu accessor used in your
codebase) and print that; adjust the ICHECK call that references target_,
TargetGetRDNAGeneration, and target_->str() accordingly.
LeiWang1999
left a comment
LeiWang1999
left a comment
There was a problem hiding this comment.
overall lgtm, but a minor question for the unexpected change.
|
@LeiWang1999 @zhangnju, please take a final review~ |
|
LGTM, Thanks! |
4 participants
Hi @LeiWang1999 @zhangnju, this PR extends the AMD RDNA WMMA path to better support
gfx11while keepinggfx12behavior intact. This implementation is based on PR #1951 and fixes Issue #1984 . Here are the main changes:gfx11andgfx12emit the correct builtin ABIgfx11-specific fragment modeling for duplicated half-wave ownership via replicatedT.FragmentDetailed Changes
src/target/codegen_hip.ccgfx11vsgfx12src/target/codegen_hip.h,src/target/rt_mod_hip.cc,src/target/utils.h,src/target/utils.ccTargetinto HIP codegen and add RDNA generation helperstilelang/intrinsics/wmma_layout.pygfx11fragment-forward helpers for duplicated half-wave ownershipA_colmajorsupport forgfx11&gfx12tilelang/intrinsics/wmma_macro_generator.pyWMMAIntrinEmitterlocal_size_a/local_size_bgfx11A/B fragment ownership withforward_fn + replicate=2Why
gfx11WMMA differs fromgfx12in two important ways:Validation
Summary by CodeRabbit
New Features
Improvements
Bug Fixes