liger-kernel-nightly 0.6.2.dev20250921193116__py3-none-any.whl → 0.6.2.dev20250923161350__py3-none-any.whl

liger_kernel/ops/layer_norm.py

@@ -63,12 +63,11 @@ def _layer_norm_forward_kernel(
     X_f32 = X_row.to(tl.float32)
 
     # Compute statistics in fp32 for numerical stability
-    n_cols_f32 = n_cols.to(tl.float32)
-    mean = tl.sum(X_f32, axis=0) / n_cols_f32
+    mean = tl.sum(X_f32, axis=0) / n_cols
     X_centered = X_f32 - mean
     # Apply mask to variance calculation to exclude contributions from masked elements
     X_centered_masked = tl.where(mask, X_centered, 0.0)
-    var = tl.sum(X_centered_masked * X_centered_masked, axis=0) / n_cols_f32
+    var = tl.sum(X_centered_masked * X_centered_masked, axis=0) / n_cols
     rstd = rsqrt(var + eps)
 
     # Store statistics (convert back to original dtype only once)
@@ -113,7 +112,6 @@ def _layer_norm_backward_kernel(
     # Pre-load weights once (same optimization as forward pass)
     w = tl.load(W_ptr + cols, mask=mask, other=0.0)
     w_f32 = w.to(tl.float32)
-    n_cols_f32 = n_cols.to(tl.float32)
 
     # Calculate pointers for this specific row
     row_X_ptr = X_ptr + row_idx * stride_x
@@ -137,8 +135,8 @@ def _layer_norm_backward_kernel(
     # Compute backward pass for this row
     x_hat = (x_f32 - mean_f32) * rstd_f32
     wdy = w_f32 * dy_f32
-    c1 = tl.sum(x_hat * wdy, axis=0) / n_cols_f32
-    c2 = tl.sum(wdy, axis=0) / n_cols_f32
+    c1 = tl.sum(x_hat * wdy, axis=0) / n_cols
+    c2 = tl.sum(wdy, axis=0) / n_cols
     dx = (wdy - (x_hat * c1 + c2)) * rstd_f32
 
     # Store input gradient

{liger_kernel_nightly-0.6.2.dev20250921193116.dist-info → liger_kernel_nightly-0.6.2.dev20250923161350.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: liger_kernel_nightly
-Version: 0.6.2.dev20250921193116
+Version: 0.6.2.dev20250923161350
 Summary: Efficient Triton kernels for LLM Training
 License: BSD 2-CLAUSE LICENSE
         Copyright 2024 LinkedIn Corporation 

{liger_kernel_nightly-0.6.2.dev20250921193116.dist-info → liger_kernel_nightly-0.6.2.dev20250923161350.dist-info}/RECORD

@@ -28,7 +28,7 @@ liger_kernel/ops/group_norm.py,sha256=qD4D4lSjSgVtO52EBNLC2iTseALRgPgqXE50U2wogg
 liger_kernel/ops/grpo_loss.py,sha256=anRnv7k1-AV3pCC6_TqP0GMg78YYUfRAJrbpx6PVhl0,9448
 liger_kernel/ops/jsd.py,sha256=onHp5T3MbvJaVz5Vup7Ww6EQp_HTaZeayTjJk6FgQMY,7042
 liger_kernel/ops/kl_div.py,sha256=ZjGdDLKWksHT9dZ0xF_TDgAkj5cuMTwwT5tr9E-_24o,8734
-liger_kernel/ops/layer_norm.py,sha256=BHPDuaogMTfIJkBJdqLZbOQouNWTf3fJVyOQOD7blCE,9901
+liger_kernel/ops/layer_norm.py,sha256=WmiORsIyufOhazmYZTPjeSc5Z-xTAYwXAKqUcCv_dlY,9807
 liger_kernel/ops/llama4_rope.py,sha256=-aqdZzllklTN8b9--e-TsWY_ntGCN8-tyseT4x0bd8s,8223
 liger_kernel/ops/multi_token_attention.py,sha256=Oz_RXDp-OSS_R_HuGmaETHdAJ7Toda_70OfE7TXMUlY,7645
 liger_kernel/ops/qwen2vl_mrope.py,sha256=3GExhYpLgB4VUtyZyjRk8XjEur3W4EWF6HQ67ML5vBU,8481
@@ -97,9 +97,9 @@ liger_kernel/transformers/trainer/__init__.py,sha256=p7yQfklV8-467qSz_ZMimkbDF7H
 liger_kernel/transformers/trainer/orpo_trainer.py,sha256=tX0h63aOFe3rNqTmk6JpMf75UPo981yzEa6TghnjS0Q,5370
 liger_kernel/triton/__init__.py,sha256=qCiCamzCRv6lpV8IqpAc9YMdNKC7GKurClWceQPnlis,92
 liger_kernel/triton/monkey_patch.py,sha256=Rd0hUHAzDkFfHvnX7-PBaNK5EKnZhtfM_h-fgQH9HPY,1568
-liger_kernel_nightly-0.6.2.dev20250921193116.dist-info/LICENSE,sha256=OhzLDHJ0to4a8sodVLELZiCFylZ1NAAYLs-HrjPy0ag,1312
-liger_kernel_nightly-0.6.2.dev20250921193116.dist-info/METADATA,sha256=J9dqNzPCm0jzT8w6An9KVhxn2B_vrIXEw8nhMqT1UKs,24605
-liger_kernel_nightly-0.6.2.dev20250921193116.dist-info/NOTICE,sha256=njwnoPZLh9AN8SJQzxvCGLHi-8X__AvWRze6joNXIY8,2066
-liger_kernel_nightly-0.6.2.dev20250921193116.dist-info/WHEEL,sha256=iAkIy5fosb7FzIOwONchHf19Qu7_1wCWyFNR5gu9nU0,91
-liger_kernel_nightly-0.6.2.dev20250921193116.dist-info/top_level.txt,sha256=2eghu4hA3LnkM7ElW92tQ8zegWKgSbeo-k-aGe1YnvY,13
-liger_kernel_nightly-0.6.2.dev20250921193116.dist-info/RECORD,,
+liger_kernel_nightly-0.6.2.dev20250923161350.dist-info/LICENSE,sha256=OhzLDHJ0to4a8sodVLELZiCFylZ1NAAYLs-HrjPy0ag,1312
+liger_kernel_nightly-0.6.2.dev20250923161350.dist-info/METADATA,sha256=WqiloVBMdQO2pUusvcTNwjrib3CUnO9Q_iBffY-JaM8,24605
+liger_kernel_nightly-0.6.2.dev20250923161350.dist-info/NOTICE,sha256=njwnoPZLh9AN8SJQzxvCGLHi-8X__AvWRze6joNXIY8,2066
+liger_kernel_nightly-0.6.2.dev20250923161350.dist-info/WHEEL,sha256=iAkIy5fosb7FzIOwONchHf19Qu7_1wCWyFNR5gu9nU0,91
+liger_kernel_nightly-0.6.2.dev20250923161350.dist-info/top_level.txt,sha256=2eghu4hA3LnkM7ElW92tQ8zegWKgSbeo-k-aGe1YnvY,13
+liger_kernel_nightly-0.6.2.dev20250923161350.dist-info/RECORD,,