github/ColossalAI
1
0
Fork 0
mirror of https://github.com/hpcaitech/ColossalAI.git synced 2025-09-02 01:28:31 +00:00
Code Issues Packages Projects Releases Wiki Activity

feat baichuan2 rmsnorm whose hidden size equals to 5120 (#5611)

Browse Source
This commit is contained in:
Steve Luo
2024-04-19 15:34:53 +08:00
committed by GitHub
parent e37ee2fb65
commit ccf72797e3
3 changed files with 9 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files

6
extensions/csrc/cuda/rms_layernorm_kernel.cu
View File

@@ -277,6 +277,9 @@ void rms_layernorm(
case 2:
RMSNORM_LAUNCHER(2, block);
break;
case 3:
RMSNORM_LAUNCHER(3, block);
break;
case 4:
RMSNORM_LAUNCHER(4, block);
break;
@@ -321,6 +324,9 @@ void fused_add_rms_layernorm(
case 2:
FUSED_ADD_RMSNORM_LAUNCHER(2, block);
break;
case 3:
FUSED_ADD_RMSNORM_LAUNCHER(3, block);
break;
case 4:
FUSED_ADD_RMSNORM_LAUNCHER(4, block);
break;