From d01487db82c79df188b78b3f35b8aa67fcd864fe Mon Sep 17 00:00:00 2001 From: Linfeng Zheng Date: Thu, 16 Jul 2026 16:28:55 +0800 Subject: [PATCH] quick fix for fmha regression (#3387) * quick fix for the v4.7 fmha regression * enhance rules --- .../cute/blackwell/kernel/attention/fmha/fmha.py | 12 +++++++++++- 1 file changed, 11 insertions(+), 1 deletion(-) diff --git a/examples/python/CuTeDSL/cute/blackwell/kernel/attention/fmha/fmha.py b/examples/python/CuTeDSL/cute/blackwell/kernel/attention/fmha/fmha.py index a717df31b..58bddfef7 100644 --- a/examples/python/CuTeDSL/cute/blackwell/kernel/attention/fmha/fmha.py +++ b/examples/python/CuTeDSL/cute/blackwell/kernel/attention/fmha/fmha.py @@ -3795,6 +3795,16 @@ def run( ) print("Compiling kernel with cute.compile ...") + cutlass_major_minor = tuple( + map(int, cutlass.__version__.split(".")[:2]) + ) + cuda_version = cutlass.CUDA_VERSION + compile_options = ( + "--opt-level=2" + if cutlass_major_minor >= (4, 6) + and (cuda_version.major, cuda_version.minor) == (12, 9) + else "" + ) start_time = time.time() # compile fmha kernel compiled_fmha = cute.compile( @@ -3818,7 +3828,7 @@ def run( total_softmax_count, current_stream, use_pdl, - options=f"--opt-level=2" if cutlass.__version__[0:3]=="4.6" else "", + options=compile_options, ) compilation_time = time.time() - start_time print(f"Compilation time: {compilation_time:.4f} seconds")