Feat/support encoder model (like bert) (#4887)

2025-04-17 16:50:48 +08:00
parent 6fb29ffd9e
commit 3bface15e6
8 changed files with 593 additions and 3 deletions
--- a/test/srt/test_triton_attention_kernels.py
+++ b/test/srt/test_triton_attention_kernels.py
@@ -116,6 +116,7 @@ class TestTritonAttention(CustomTestCase):
            kv_indptr,
            kv_indices,
            custom_mask,
+            True,
            mask_indptr,
            max_len_extend,
        )
@@ -150,6 +151,7 @@ class TestTritonAttention(CustomTestCase):
            kv_indptr,
            kv_indices,
            custom_mask,
+            True,
            mask_indptr,
            max_len_extend,
        )