From 5e561560eaab135a51e7e4e4d39b2f5c1c5375ed Mon Sep 17 00:00:00 2001
From: AS <1280924+alexsarmiento@users.noreply.github.com>
Date: Thu, 8 May 2025 19:32:51 -0700
Subject: [PATCH] Update fwd_prefill.py

Remove cudagraph from triton autotune. Buggy under zluda
---
 modules/flash_attn_triton_amd/fwd_prefill.py | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/modules/flash_attn_triton_amd/fwd_prefill.py b/modules/flash_attn_triton_amd/fwd_prefill.py
index cf78cd0ee..123242694 100644
--- a/modules/flash_attn_triton_amd/fwd_prefill.py
+++ b/modules/flash_attn_triton_amd/fwd_prefill.py
@@ -218,8 +218,7 @@ autotune_configs, autotune_keys = get_autotune_configs()
 
 @triton.autotune(
     configs=autotune_configs,
-    key=autotune_keys,
-    use_cuda_graph=True,
+    key=autotune_keys
 )
 @triton.jit
 def attn_fwd(Q, K, V, bias, Cache_seqlens, Cache_batch_idx, # pylint: disable=unused-argument