Refactor attention backend (#1381)

2024-09-11 11:44:26 -07:00
parent c03cece42f
commit fec185ce0c
16 changed files with 568 additions and 564 deletions
--- a/test/srt/test_triton_attention_kernels.py
+++ b/test/srt/test_triton_attention_kernels.py
@@ -96,23 +96,17 @@ class TestExtendAttention(unittest.TestCase):
            v_buffer,
            req_to_tokens,
            b_req_idx,
-            b_start_loc,
            b_seq_len,
-            b_seq_len_prefix,
-            b_start_loc_extend,
            b_seq_len_extend,
-            max_len_in_batch,
+            b_start_loc_extend,
            max_len_extend,
        )

        redundant_attention(
            q_extend,
-            k_extend,
-            v_extend,
            o_redundant,
            k_buffer,
            v_buffer,
-            req_to_tokens,
            b_req_idx,
            b_start_loc,
            b_seq_len,