3d tensor for RMSNorm (#111)

Signed-off-by: Mayank Mishra <[email protected]>
mayank31398 · Dec 23, 2024 · 333d10e · 333d10e
1 parent df6c243
commit 333d10e
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 3 deletions.
diff --git a/cute_kernels/kernels/rmsnorm/triton_implementation/kernels_backward.py b/cute_kernels/kernels/rmsnorm/triton_implementation/kernels_backward.py
@@ -107,7 +107,8 @@ def _rmsnorm_backward_no_weight_triton(
     BLOCK_SIZE_B: int,
     BLOCK_SIZE_H: int,
 ) -> None:
-    num_elements, hidden_size = x.size()
+    hidden_size = x.size(-1)
+    num_elements = x.numel() // hidden_size
 
     if BLOCK_SIZE_H < hidden_size:
         raise ValueError(f"hidden_size should be more than the BLOCK_SIZE_H")
@@ -158,7 +159,8 @@ def _rmsnorm_backward_triton(
     BLOCK_SIZE_B: int,
     BLOCK_SIZE_H: int,
 ) -> torch.Tensor:
-    num_elements, hidden_size = x.size()
+    hidden_size = x.size(-1)
+    num_elements = x.numel() // hidden_size
 
     if BLOCK_SIZE_H < hidden_size:
         raise ValueError(f"hidden_size should be more than the BLOCK_SIZE_H")

diff --git a/cute_kernels/kernels/rmsnorm/triton_implementation/kernels_forward.py b/cute_kernels/kernels/rmsnorm/triton_implementation/kernels_forward.py
@@ -84,7 +84,8 @@ def rmsnorm_forward_triton(
     BLOCK_SIZE_B: int,
     BLOCK_SIZE_H: int,
 ) -> None:
-    num_elements, hidden_size = x.size()
+    hidden_size = x.size(-1)
+    num_elements = x.numel() // hidden_size
 
     if BLOCK_SIZE_H < hidden_size:
         raise ValueError(f"hidden_size should be more than the BLOCK_SIZE_H")