[fix] fix illegal mem access and clean up triton attention backend (#4571)

2025-03-20 17:01:52 +08:00
parent fad86a6863
commit 9e93ef3f8e
7 changed files with 124 additions and 125 deletions
--- a/test/srt/test_triton_attention_kernels.py
+++ b/test/srt/test_triton_attention_kernels.py
@@ -265,7 +265,8 @@ class TestTritonAttention(unittest.TestCase):
            o,
            kv_indptr,
            kv_indices,
-            (attn_logits, attn_lse),
+            attn_logits,
+            attn_lse,
            num_kv_splits,
            max_kv_splits,
            sm_scale,
@@ -329,7 +330,8 @@ class TestTritonAttention(unittest.TestCase):
            o,
            kv_indptr,
            kv_indices,
-            (attn_logits, attn_lse),
+            attn_logits,
+            attn_lse,
            num_kv_splits,
            max_kv_splits,
            sm_scale,
@@ -353,7 +355,8 @@ class TestTritonAttention(unittest.TestCase):
            o_grouped,
            kv_indptr,
            kv_indices,
-            (attn_logits1, attn_lse1),
+            attn_logits1,
+            attn_lse1,
            num_kv_splits,
            max_kv_splits,
            sm_scale,