liger-kernel-nightly 0.6.0.dev20250721100112__tar.gz → 0.6.0.dev20250722224228__tar.gz

{liger_kernel_nightly-0.6.0.dev20250721100112 → liger_kernel_nightly-0.6.0.dev20250722224228}/PKG-INFO

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

{liger_kernel_nightly-0.6.0.dev20250721100112 → liger_kernel_nightly-0.6.0.dev20250722224228}/benchmark/scripts/benchmark_cpo_loss.py

@@ -36,17 +36,20 @@ def bench_memory_fused_linear_cpo_loss(
     dtype = input.extra_benchmark_config["dtype"]
     provider = input.kernel_provider
 
-    torch_lm_head_cpo = lambda x, target: TorchLMHeadCPO(H=H, V=V, dtype=dtype).to(device)(x, target)[0]
-    liger_lm_head_cpo = lambda x, target: LigerLMHeadCPO(H=H, V=V, dtype=dtype).to(device)(x, target)[0]
+    # Instantiate once and retrieve the first output only
+    torch_lm_head_cpo = TorchLMHeadCPO(H=H, V=V, dtype=dtype).to(device)
+    liger_lm_head_cpo = LigerLMHeadCPO(H=H, V=V, dtype=dtype).to(device)
+    torch_fwd = lambda x, target: torch_lm_head_cpo(x, target)[0]
+    liger_fwd = lambda x, target: liger_lm_head_cpo(x, target)[0]
 
     _input = torch.randn(B, T, H, requires_grad=True, dtype=dtype, device=device)
     target = torch.randint(V, (B, T), dtype=torch.long, device=device)
 
     def fwd():
         if provider == "liger":
-            return liger_lm_head_cpo(_input, target)
+            return liger_fwd(_input, target)
         elif provider == "huggingface":
-            return torch_lm_head_cpo(_input, target)
+            return torch_fwd(_input, target)
 
     def full():
         y = fwd()
@@ -79,17 +82,20 @@ def bench_speed_fused_linear_cpo_loss(
     provider = input.kernel_provider
     mode = input.kernel_operation_mode
 
-    torch_lm_head_cpo = lambda x, target: TorchLMHeadCPO(H=H, V=V, dtype=dtype).to(device)(x, target)[0]
-    liger_lm_head_cpo = lambda x, target: LigerLMHeadCPO(H=H, V=V, dtype=dtype).to(device)(x, target)[0]
+    # Instantiate once and retrieve the first output only
+    torch_lm_head_cpo = TorchLMHeadCPO(H=H, V=V, dtype=dtype).to(device)
+    liger_lm_head_cpo = LigerLMHeadCPO(H=H, V=V, dtype=dtype).to(device)
+    torch_fwd = lambda x, target: torch_lm_head_cpo(x, target)[0]
+    liger_fwd = lambda x, target: liger_lm_head_cpo(x, target)[0]
 
     _input = torch.randn(B, T, H, requires_grad=True, dtype=dtype, device=device)
     target = torch.randint(V, (B, T), dtype=torch.long, device=device)
 
     def fwd():
         if provider == "liger":
-            return liger_lm_head_cpo(_input, target)
+            return liger_fwd(_input, target)
         elif provider == "huggingface":
-            return torch_lm_head_cpo(_input, target)
+            return torch_fwd(_input, target)
 
     if mode == "forward":
         ms_50, ms_20, ms_80 = triton.testing.do_bench(

{liger_kernel_nightly-0.6.0.dev20250721100112 → liger_kernel_nightly-0.6.0.dev20250722224228}/benchmark/scripts/benchmark_dpo_loss.py

@@ -32,12 +32,11 @@ def bench_memory_dpo_loss(input: SingleBenchmarkRunInput) -> SingleBenchmarkRunO
     ignore_index = input.extra_benchmark_config["ignore_index"]
     provider = input.kernel_provider
 
-    torch_dpo_loss = lambda x, ref_x, target: TorchLMHeadDPO(
-        H=H, V=V, dtype=dtype, beta=beta, ignore_index=ignore_index, bias=bias
-    ).to(device)(x, ref_x, target)[0]
-    liger_dpo_loss = lambda x, ref_x, target: LigerLMHeadDPO(
-        H=H, V=V, dtype=dtype, beta=beta, ignore_index=ignore_index, bias=bias
-    ).to(device)(x, ref_x, target)[0]
+    # Instantiate once and retrieve the first output only
+    torch_dpo_loss = TorchLMHeadDPO(H=H, V=V, dtype=dtype, beta=beta, ignore_index=ignore_index, bias=bias).to(device)
+    liger_dpo_loss = LigerLMHeadDPO(H=H, V=V, dtype=dtype, beta=beta, ignore_index=ignore_index, bias=bias).to(device)
+    torch_fwd = lambda x, ref_x, target: torch_dpo_loss(x, ref_x, target)[0]
+    liger_fwd = lambda x, ref_x, target: liger_dpo_loss(x, ref_x, target)[0]
 
     # Input shape: [B, T, H]
     _input = torch.randn(B, T, H, device=device, dtype=dtype)
@@ -52,9 +51,9 @@ def bench_memory_dpo_loss(input: SingleBenchmarkRunInput) -> SingleBenchmarkRunO
 
     def fwd():
         if provider == "liger":
-            return liger_dpo_loss(_input, ref_input, target)
+            return liger_fwd(_input, ref_input, target)
         elif provider == "huggingface":
-            return torch_dpo_loss(_input, ref_input, target)
+            return torch_fwd(_input, ref_input, target)
 
     def full():
         y = fwd()
@@ -83,12 +82,11 @@ def bench_speed_dpo_loss(input: SingleBenchmarkRunInput) -> SingleBenchmarkRunOu
     provider = input.kernel_provider
     mode = input.kernel_operation_mode
 
-    torch_dpo_loss = lambda x, ref_x, target: TorchLMHeadDPO(
-        H=H, V=V, dtype=dtype, beta=beta, ignore_index=ignore_index, bias=bias
-    ).to(device)(x, ref_x, target)[0]
-    liger_dpo_loss = lambda x, ref_x, target: LigerLMHeadDPO(
-        H=H, V=V, dtype=dtype, beta=beta, ignore_index=ignore_index, bias=bias
-    ).to(device)(x, ref_x, target)[0]
+    # Instantiate once and retrieve the first output only
+    torch_dpo_loss = TorchLMHeadDPO(H=H, V=V, dtype=dtype, beta=beta, ignore_index=ignore_index, bias=bias).to(device)
+    liger_dpo_loss = LigerLMHeadDPO(H=H, V=V, dtype=dtype, beta=beta, ignore_index=ignore_index, bias=bias).to(device)
+    torch_fwd = lambda x, ref_x, target: torch_dpo_loss(x, ref_x, target)[0]
+    liger_fwd = lambda x, ref_x, target: liger_dpo_loss(x, ref_x, target)[0]
 
     # Input shape: [B, T, H]
     _input = torch.randn(B, T, H, device=device, dtype=dtype)
@@ -103,9 +101,9 @@ def bench_speed_dpo_loss(input: SingleBenchmarkRunInput) -> SingleBenchmarkRunOu
 
     def fwd():
         if provider == "liger":
-            return liger_dpo_loss(_input, ref_input, target)
+            return liger_fwd(_input, ref_input, target)
         elif provider == "huggingface":
-            return torch_dpo_loss(_input, ref_input, target)
+            return torch_fwd(_input, ref_input, target)
 
     if mode == "forward":
         ms_50, ms_20, ms_80 = triton.testing.do_bench(

{liger_kernel_nightly-0.6.0.dev20250721100112 → liger_kernel_nightly-0.6.0.dev20250722224228}/benchmark/scripts/benchmark_orpo_loss.py

@@ -36,8 +36,11 @@ def bench_memory_fused_linear_orpo_loss(
     dtype = input.extra_benchmark_config["dtype"]
     provider = input.kernel_provider
 
-    torch_lm_head_orpo = lambda x, target: TorchLMHeadORPO(H=H, V=V, dtype=dtype).to(device)(x, target)[0]
-    liger_lm_head_orpo = lambda x, target: LigerLMHeadORPO(H=H, V=V, dtype=dtype).to(device)(x, target)[0]
+    # Instantiate once and retrieve the first output only
+    torch_lm_head_orpo = TorchLMHeadORPO(H=H, V=V, dtype=dtype).to(device)
+    liger_lm_head_orpo = LigerLMHeadORPO(H=H, V=V, dtype=dtype).to(device)
+    torch_fwd = lambda x, target, nll_target: torch_lm_head_orpo(x, target, nll_target)[0]
+    liger_fwd = lambda x, target, nll_target: liger_lm_head_orpo(x, target, nll_target)[0]
 
     _input = torch.randn(B, T, H, requires_grad=True, dtype=dtype, device=device)
     target = torch.randint(V, (B, T), dtype=torch.long, device=device)
@@ -45,9 +48,9 @@ def bench_memory_fused_linear_orpo_loss(
 
     def fwd():
         if provider == "liger":
-            return liger_lm_head_orpo(_input, target, nll_target)
+            return liger_fwd(_input, target, nll_target)
         elif provider == "huggingface":
-            return torch_lm_head_orpo(_input, target, nll_target)
+            return torch_fwd(_input, target, nll_target)
 
     def full():
         y = fwd()
@@ -80,8 +83,11 @@ def bench_speed_fused_linear_orpo_loss(
     provider = input.kernel_provider
     mode = input.kernel_operation_mode
 
-    torch_lm_head_orpo = lambda x, target: TorchLMHeadORPO(H=H, V=V, dtype=dtype).to(device)(x, target)[0]
-    liger_lm_head_orpo = lambda x, target: LigerLMHeadORPO(H=H, V=V, dtype=dtype).to(device)(x, target)[0]
+    # Instantiate once and retrieve the first output only
+    torch_lm_head_orpo = TorchLMHeadORPO(H=H, V=V, dtype=dtype).to(device)
+    liger_lm_head_orpo = LigerLMHeadORPO(H=H, V=V, dtype=dtype).to(device)
+    torch_fwd = lambda x, target, nll_target: torch_lm_head_orpo(x, target, nll_target)[0]
+    liger_fwd = lambda x, target, nll_target: liger_lm_head_orpo(x, target, nll_target)[0]
 
     _input = torch.randn(B, T, H, requires_grad=True, dtype=dtype, device=device)
     target = torch.randint(V, (B, T), dtype=torch.long, device=device)
@@ -89,9 +95,9 @@ def bench_speed_fused_linear_orpo_loss(
 
     def fwd():
         if provider == "liger":
-            return liger_lm_head_orpo(_input, target, nll_target)
+            return liger_fwd(_input, target, nll_target)
         elif provider == "huggingface":
-            return torch_lm_head_orpo(_input, target, nll_target)
+            return torch_fwd(_input, target, nll_target)
 
     if mode == "forward":
         ms_50, ms_20, ms_80 = triton.testing.do_bench(

{liger_kernel_nightly-0.6.0.dev20250721100112 → liger_kernel_nightly-0.6.0.dev20250722224228}/benchmark/scripts/benchmark_simpo_loss.py

@@ -36,17 +36,20 @@ def bench_memory_fused_linear_simpo_loss(
     dtype = input.extra_benchmark_config["dtype"]
     provider = input.kernel_provider
 
-    torch_lm_head_simpo = lambda x, target: TorchLMHeadCPO(H=H, V=V, dtype=dtype).to(device)(x, target)[0]
-    liger_lm_head_simpo = lambda x, target: LigerLMHeadSimPO(H=H, V=V, dtype=dtype).to(device)(x, target)[0]
+    # Instantiate once and retrieve the first output only
+    torch_lm_head_simpo = TorchLMHeadCPO(H=H, V=V, dtype=dtype).to(device)
+    liger_lm_head_simpo = LigerLMHeadSimPO(H=H, V=V, dtype=dtype).to(device)
+    torch_fwd = lambda x, target: torch_lm_head_simpo(x, target)[0]
+    liger_fwd = lambda x, target: liger_lm_head_simpo(x, target)[0]
 
     _input = torch.randn(B, T, H, requires_grad=True, dtype=dtype, device=device)
     target = torch.randint(V, (B, T), dtype=torch.long, device=device)
 
     def fwd():
         if provider == "liger":
-            return liger_lm_head_simpo(_input, target)
+            return liger_fwd(_input, target)
         elif provider == "huggingface":
-            return torch_lm_head_simpo(_input, target)
+            return torch_fwd(_input, target)
 
     def full():
         y = fwd()
@@ -79,17 +82,20 @@ def bench_speed_fused_linear_simpo_loss(
     provider = input.kernel_provider
     mode = input.kernel_operation_mode
 
-    torch_lm_head_simpo = lambda x, target: TorchLMHeadCPO(H=H, V=V, dtype=dtype).to(device)(x, target)[0]
-    liger_lm_head_simpo = lambda x, target: LigerLMHeadSimPO(H=H, V=V, dtype=dtype).to(device)(x, target)[0]
+    # Instantiate once and retrieve the first output only
+    torch_lm_head_simpo = TorchLMHeadCPO(H=H, V=V, dtype=dtype).to(device)
+    liger_lm_head_simpo = LigerLMHeadSimPO(H=H, V=V, dtype=dtype).to(device)
+    torch_fwd = lambda x, target: torch_lm_head_simpo(x, target)[0]
+    liger_fwd = lambda x, target: liger_lm_head_simpo(x, target)[0]
 
     _input = torch.randn(B, T, H, requires_grad=True, dtype=dtype, device=device)
     target = torch.randint(V, (B, T), dtype=torch.long, device=device)
 
     def fwd():
         if provider == "liger":
-            return liger_lm_head_simpo(_input, target)
+            return liger_fwd(_input, target)
         elif provider == "huggingface":
-            return torch_lm_head_simpo(_input, target)
+            return torch_fwd(_input, target)
 
     if mode == "forward":
         ms_50, ms_20, ms_80 = triton.testing.do_bench(

{liger_kernel_nightly-0.6.0.dev20250721100112 → liger_kernel_nightly-0.6.0.dev20250722224228}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "liger_kernel_nightly"
-version = "0.6.0.dev20250721100112"
+version = "0.6.0.dev20250722224228"
 description = "Efficient Triton kernels for LLM Training"
 urls = { "Homepage" = "https://github.com/linkedin/Liger-Kernel" }
 readme = { file = "README.md", content-type = "text/markdown" }

{liger_kernel_nightly-0.6.0.dev20250721100112 → liger_kernel_nightly-0.6.0.dev20250722224228}/src/liger_kernel_nightly.egg-info/PKG-INFO

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