[Kernel] Add punica dimension for Baichuan-13B (#4053)

vllm-project · Apr 13, 2024 · 989ae25 · 989ae25
1 parent 0a430b4
commit 989ae25
Show file tree

Hide file tree

Showing 3 changed files with 3 additions and 1 deletion.
diff --git a/csrc/punica/bgmv/bgmv_config.h b/csrc/punica/bgmv/bgmv_config.h
@@ -47,6 +47,7 @@ void bgmv_kernel(out_T *__restrict__ Y, const in_T *__restrict__ X,
     f(in_T, out_T, W_T, narrow, 13696) \
     f(in_T, out_T, W_T, narrow, 13824) \
     f(in_T, out_T, W_T, narrow, 14336) \
+    f(in_T, out_T, W_T, narrow, 15360) \
     f(in_T, out_T, W_T, narrow, 16384) \
     f(in_T, out_T, W_T, narrow, 20480) \
     f(in_T, out_T, W_T, narrow, 22016) \

diff --git a/tests/lora/test_baichuan.py b/tests/lora/test_baichuan.py
@@ -62,7 +62,7 @@ def test_baichuan_lora(baichuan_lora_files):
 
 
 @pytest.mark.skip("Requires multiple GPUs")
-def test_llama_tensor_parallel_equality(baichuan_lora_files):
+def test_baichuan_tensor_parallel_equality(baichuan_lora_files):
     # Cannot use as it will initialize torch.cuda too early...
     # if torch.cuda.device_count() < 4:
     #     pytest.skip(f"Not enough GPUs for tensor parallelism {4}")

diff --git a/tests/lora/test_punica.py b/tests/lora/test_punica.py
@@ -72,6 +72,7 @@ def _lora_ref_impl(
     11008,
     13824,
     14336,
+    15360,
     22016,
     24576,
     27392,