Add support for BLOOM (#331)

2023-07-03 13:12:35 -07:00
parent d6fa1be3a8
commit e41f06702c
11 changed files with 479 additions and 18 deletions
--- a/tests/kernels/test_attention.py
+++ b/tests/kernels/test_attention.py
@@ -216,6 +216,7 @@ def run_single_query_cached_kv_attention(
        context_lens,
        block_size,
        max_context_len,
+        None,  # ALiBi slopes.
    )

    ref_output = torch.empty_like(query)