[Kernel] Add punica dimension for Swallow-MS-7B LoRA (vllm-project#4134)

ywang96 · Apr 17, 2024 · a532225 · a532225
1 parent fe3b5bb
commit a532225
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 0 deletions.
diff --git a/csrc/punica/bgmv/bgmv_config.h b/csrc/punica/bgmv/bgmv_config.h
@@ -60,6 +60,7 @@ void bgmv_kernel(out_T *__restrict__ Y, const in_T *__restrict__ X,
     f(in_T, out_T, W_T, narrow, 32768) \
     f(in_T, out_T, W_T, narrow, 33024) \
     f(in_T, out_T, W_T, narrow, 36864) \
+    f(in_T, out_T, W_T, narrow, 43264) \
     f(in_T, out_T, W_T, narrow, 49152) \
     f(in_T, out_T, W_T, narrow, 64000) \
     f(in_T, out_T, W_T, narrow, 64256) \

diff --git a/tests/lora/test_punica.py b/tests/lora/test_punica.py
@@ -82,6 +82,7 @@ def _lora_ref_impl(
     32768,
     33024,
     36864,
+    43264,
     49152,
     64000,
     64256,