[Model] Support MAP-NEO model (#5081)

Co-authored-by: Zhuohan Li <zhuohan123@gmail.com>
2024-05-31 10:24:41 +08:00
parent 533c217792
commit a22dea54d3
8 changed files with 18 additions and 6 deletions
--- a/tests/kernels/test_cache.py
+++ b/tests/kernels/test_cache.py
@@ -11,7 +11,7 @@ DTYPES = [torch.half, torch.bfloat16, torch.float]
 NUM_TOKENS = [42]  # Arbitrary values for testing
 NUM_LAYERS = [1]  # Arbitrary values for testing
 NUM_HEADS = [8]  # Arbitrary values for testing
-HEAD_SIZES = [64, 80, 96, 112, 128, 256]
+HEAD_SIZES = [64, 80, 96, 112, 128, 192, 256]
 BLOCK_SIZES = [8, 16, 32]

 # Arbitrary values for testing