liger-kernel 0.5.10__tar.gz → 0.6.0__tar.gz

liger_kernel-0.6.0/.github/workflows/benchmark.yml

@@ -0,0 +1,93 @@
+name: Benchmarks
+
+on:
+  schedule:
+    # Runs at 00:00 UTC every Friday
+    - cron: '0 0 * * 5'
+  workflow_dispatch:  # Enables manual trigger
+
+permissions:
+  contents: write
+  
+concurrency:
+  # This causes it to cancel previous in-progress actions on the same PR / branch,
+  group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.ref }}
+  cancel-in-progress: true
+  
+jobs:
+  benchmarks:
+    runs-on: ubuntu-latest
+    env:
+      MODAL_TOKEN_ID: ${{ secrets.MODAL_TOKEN_ID }}
+      MODAL_TOKEN_SECRET: ${{ secrets.MODAL_TOKEN_SECRET }}
+      GITHUB_USERNAME: linkedin
+      REPO_NAME: Liger-Kernel
+      OUTPUT_DIR: benchmarks
+      OUTPUT_FILENAME: benchmark.csv
+      GENERATED_CSV: benchmark/data/all_benchmark_data.csv
+
+
+    steps:
+      - name: Checkout code
+        uses: actions/checkout@v3
+
+      # Get the latest commit hash from main branch
+      - name: Get commit hash
+        id: get_hash
+        run: echo "hash=$(git rev-parse --short HEAD)" >> $GITHUB_OUTPUT
+
+      - name: Set up Python
+        uses: actions/setup-python@v3
+        with:
+          python-version: '3.10'
+
+      # Install dependencies
+      - name: Install dependencies
+        run: |
+          python -m pip install --upgrade pip
+          pip install modal
+          pip install pandas
+
+      # Delete previous benchmark results.
+      - name: Remove previous benchmark data
+        run: |
+          rm -f benchmark/data/all_benchmark_data.csv
+          
+      - name: Run benchmarks on GPU
+        run: |
+          modal run dev.modal.benchmarks
+        
+       # Step 5: Checkout gh-pages branch in a subfolderAdd commentMore actions
+      - name: Checkout gh-pages
+        uses: actions/checkout@v3
+        with:
+          ref: gh-pages
+          path: gh-pages
+      
+      # Step 6: Copy benchmark CSV to gh-pages directory
+      - name: Copy generated benchmark to gh-pages
+        run: |
+          mkdir -p gh-pages/${OUTPUT_DIR}/${{ steps.get_hash.outputs.hash }}
+          cp ${GENERATED_CSV} gh-pages/${OUTPUT_DIR}/${{ steps.get_hash.outputs.hash }}/${OUTPUT_FILENAME}
+      # Step 7: Append commit hash to commits.txt if not already present
+      - name: Update commits.txt
+        run: |
+          cd gh-pages
+          echo "commits.txt file path: ${OUTPUT_DIR}/commits.txt"
+          # Create file if it doesn't exist
+          mkdir -p ${OUTPUT_DIR}
+          touch ${OUTPUT_DIR}/commits.txt
+          # Append only if not already present
+          if ! grep -q "${{ steps.get_hash.outputs.hash }}" ${OUTPUT_DIR}/commits.txt; then
+            echo "${{ steps.get_hash.outputs.hash }}" >> ${OUTPUT_DIR}/commits.txt
+          fi
+      # Step 7: Commit and push
+      - name: Commit and push to gh-pages
+        run: |
+          cd gh-pages
+          git config user.name github-actions[bot]
+          git config user.email 41898282+github-actions[bot]@users.noreply.github.com
+          git add .
+          git commit -m "Add benchmark for commit ${{ steps.get_hash.outputs.hash }}" || echo "No changes to commit"
+          git push origin gh-pages
+

{liger_kernel-0.5.10 → liger_kernel-0.6.0}/.github/workflows/docs.yml

@@ -3,10 +3,12 @@ on:
   push:
     branches:
       - main
+
 permissions:
   contents: write
 jobs:
   deploy:
+    if: False
     runs-on: ubuntu-latest
     steps:
       - uses: actions/checkout@v4
@@ -25,4 +27,4 @@ jobs:
           restore-keys: |
             mkdocs-material-
       - run: pip install mkdocs-material 
-      - run: mkdocs gh-deploy --force
+      - run: mkdocs gh-deploy --force

{liger_kernel-0.5.10 → liger_kernel-0.6.0}/.gitignore

@@ -6,6 +6,7 @@ site/
 venv/
 .ipynb_checkpoints/
 .vscode/
+.idea/
 
 # Misc
 .DS_Store
@@ -14,6 +15,9 @@ venv/
 build/
 dist/
 
+# Doc Build
+site/
+
 # Lockfiles
 uv.lock
 

{liger_kernel-0.5.10 → liger_kernel-0.6.0}/Makefile

@@ -48,13 +48,19 @@ run-benchmarks:
 # MkDocs Configuration
 MKDOCS = mkdocs
 CONFIG_FILE = mkdocs.yml
+SITE_DIR = site
 
 # MkDocs targets
+
+# Serve the documentation 
 serve:
 	$(MKDOCS) serve -f $(CONFIG_FILE)
 
+# Build the documentation into the specified site directory
 build:
-	$(MKDOCS) build -f $(CONFIG_FILE)
+	$(MKDOCS) build -f $(CONFIG_FILE) --site-dir $(SITE_DIR)
 
+# Clean the output directory
 clean:
-	rm -rf site/
+	rm -rf $(SITE_DIR)/
+

{liger_kernel-0.5.10 → liger_kernel-0.6.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: liger_kernel
-Version: 0.5.10
+Version: 0.6.0
 Summary: Efficient Triton kernels for LLM Training
 License: BSD 2-CLAUSE LICENSE
         Copyright 2024 LinkedIn Corporation 
@@ -33,7 +33,7 @@ License-File: NOTICE
 Requires-Dist: torch>=2.1.2
 Requires-Dist: triton>=2.3.1
 Provides-Extra: dev
-Requires-Dist: transformers>=4.44.2; extra == "dev"
+Requires-Dist: transformers>=4.49.0; extra == "dev"
 Requires-Dist: matplotlib>=3.7.2; extra == "dev"
 Requires-Dist: flake8>=4.0.1.1; extra == "dev"
 Requires-Dist: black>=24.4.2; extra == "dev"
@@ -45,6 +45,7 @@ Requires-Dist: datasets>=2.19.2; extra == "dev"
 Requires-Dist: seaborn; extra == "dev"
 Requires-Dist: mkdocs; extra == "dev"
 Requires-Dist: mkdocs-material; extra == "dev"
+Requires-Dist: torchvision>=0.20; extra == "dev"
 Dynamic: license-file
 Dynamic: provides-extra
 Dynamic: requires-dist
@@ -114,6 +115,8 @@ Dynamic: requires-dist
 
 We've also added optimized Post-Training kernels that deliver **up to 80% memory savings** for alignment and distillation tasks. We support losses like DPO, CPO, ORPO, SimPO, KTO, JSD, and many more. Check out [how we optimize the memory](https://x.com/hsu_byron/status/1866577403918917655).
 
+You can view the documentation site for additional installation, usage examples, and API references:https://linkedin.github.io/Liger-Kernel/
+
 ## Supercharge Your Model with Liger Kernel
 
 ![Banner](https://raw.githubusercontent.com/linkedin/Liger-Kernel/main/docs/images/banner.GIF)
@@ -290,6 +293,7 @@ loss.backward()
 
 | **Model**   | **API**                                                      | **Supported Operations**                                                |
 |-------------|--------------------------------------------------------------|-------------------------------------------------------------------------|
+| Llama4 (Text) & (Multimodal)      | `liger_kernel.transformers.apply_liger_kernel_to_llama4`   | RMSNorm, LayerNorm, GeGLU, CrossEntropyLoss, FusedLinearCrossEntropy         |
 | LLaMA 2 & 3 | `liger_kernel.transformers.apply_liger_kernel_to_llama`   | RoPE, RMSNorm, SwiGLU, CrossEntropyLoss, FusedLinearCrossEntropy        |
 | LLaMA 3.2-Vision | `liger_kernel.transformers.apply_liger_kernel_to_mllama`   | RoPE, RMSNorm, SwiGLU, CrossEntropyLoss, FusedLinearCrossEntropy        |
 | Mistral     | `liger_kernel.transformers.apply_liger_kernel_to_mistral`  | RoPE, RMSNorm, SwiGLU, CrossEntropyLoss, FusedLinearCrossEntropy        |
@@ -326,6 +330,8 @@ loss.backward()
 | GeGLU                           | `liger_kernel.transformers.LigerGEGLUMLP`                   |
 | CrossEntropy                    | `liger_kernel.transformers.LigerCrossEntropyLoss`           |
 | Fused Linear CrossEntropy       | `liger_kernel.transformers.LigerFusedLinearCrossEntropyLoss`|
+| Multi Token Attention           | `liger_kernel.transformers.LigerMultiTokenAttention`        |
+| Softmax                         | `liger_kernel.transformers.LigerSoftmax`                    |
 | Sparsemax                       | `liger_kernel.transformers.LigerSparsemax`                  |
 
 

{liger_kernel-0.5.10 → liger_kernel-0.6.0}/README.md

@@ -63,6 +63,8 @@
 
 We've also added optimized Post-Training kernels that deliver **up to 80% memory savings** for alignment and distillation tasks. We support losses like DPO, CPO, ORPO, SimPO, KTO, JSD, and many more. Check out [how we optimize the memory](https://x.com/hsu_byron/status/1866577403918917655).
 
+You can view the documentation site for additional installation, usage examples, and API references:https://linkedin.github.io/Liger-Kernel/
+
 ## Supercharge Your Model with Liger Kernel
 
 ![Banner](https://raw.githubusercontent.com/linkedin/Liger-Kernel/main/docs/images/banner.GIF)
@@ -239,6 +241,7 @@ loss.backward()
 
 | **Model**   | **API**                                                      | **Supported Operations**                                                |
 |-------------|--------------------------------------------------------------|-------------------------------------------------------------------------|
+| Llama4 (Text) & (Multimodal)      | `liger_kernel.transformers.apply_liger_kernel_to_llama4`   | RMSNorm, LayerNorm, GeGLU, CrossEntropyLoss, FusedLinearCrossEntropy         |
 | LLaMA 2 & 3 | `liger_kernel.transformers.apply_liger_kernel_to_llama`   | RoPE, RMSNorm, SwiGLU, CrossEntropyLoss, FusedLinearCrossEntropy        |
 | LLaMA 3.2-Vision | `liger_kernel.transformers.apply_liger_kernel_to_mllama`   | RoPE, RMSNorm, SwiGLU, CrossEntropyLoss, FusedLinearCrossEntropy        |
 | Mistral     | `liger_kernel.transformers.apply_liger_kernel_to_mistral`  | RoPE, RMSNorm, SwiGLU, CrossEntropyLoss, FusedLinearCrossEntropy        |
@@ -275,6 +278,8 @@ loss.backward()
 | GeGLU                           | `liger_kernel.transformers.LigerGEGLUMLP`                   |
 | CrossEntropy                    | `liger_kernel.transformers.LigerCrossEntropyLoss`           |
 | Fused Linear CrossEntropy       | `liger_kernel.transformers.LigerFusedLinearCrossEntropyLoss`|
+| Multi Token Attention           | `liger_kernel.transformers.LigerMultiTokenAttention`        |
+| Softmax                         | `liger_kernel.transformers.LigerSoftmax`                    |
 | Sparsemax                       | `liger_kernel.transformers.LigerSparsemax`                  |
 
 

{liger_kernel-0.5.10 → liger_kernel-0.6.0}/benchmark/benchmarks_visualizer.py

@@ -22,6 +22,9 @@ class VisualizationsConfig:
         kernel_name (str): Kernel name to benchmark. (Will run `scripts/benchmark_{kernel_name}.py`)
         metric_name (str): Metric name to visualize (speed/memory)
         kernel_operation_mode (str): Kernel operation mode to visualize (forward/backward/full). Defaults to "full"
+        extra_config_filter (str, optional): A string to filter extra_benchmark_config.
+                                            Can be a substring to match or a 'key=value' pair (e.g., "'H': 4096").
+                                            Defaults to None, which means the first available config will be used if multiple exist.
         display (bool): Display the visualization. Defaults to False
         overwrite (bool): Overwrite existing visualization, if none exist this flag has no effect as ones are always created and saved. Defaults to False
 
@@ -30,6 +33,7 @@ class VisualizationsConfig:
     kernel_name: str
     metric_name: str
     kernel_operation_mode: str = "full"
+    extra_config_filter: str | None = None
     display: bool = False
     overwrite: bool = False
 
@@ -55,6 +59,14 @@ def parse_args() -> VisualizationsConfig:
         default=None,
         help="Kernel operation modes to visualize (forward/backward/full). If not provided, generate for all available modes.",
     )
+    parser.add_argument(
+        "--extra-config-filter",
+        type=str,
+        default=None,
+        help="A string to filter extra_benchmark_config. "
+        "Can be a substring to match or a JSON-like 'key=value' pair (e.g., \"'H': 4096\" or \"H=4096\" for simple cases). "
+        "Defaults to None (first available config if multiple exist).",
+    )
     parser.add_argument("--display", action="store_true", help="Display the visualization")
     parser.add_argument(
         "--overwrite",
@@ -81,19 +93,101 @@ def load_data(config: VisualizationsConfig) -> pd.DataFrame:
     df = pd.read_csv(DATA_PATH)
     df["extra_benchmark_config"] = df["extra_benchmark_config_str"].apply(json.loads)
 
-    filtered_df = df[
+    base_filtered_df = df[
         (df["kernel_name"] == config.kernel_name)
         & (df["metric_name"] == config.metric_name)
         & (df["kernel_operation_mode"] == config.kernel_operation_mode)
-        # Use this to filter by extra benchmark configuration property
-        # & (data['extra_benchmark_config'].apply(lambda x: x.get('H') == 4096))
-        # FIXME: maybe add a way to filter using some configuration, except of hardcoding it
     ]
 
-    if filtered_df.empty:
-        raise ValueError("No data found for the given filters")
+    if base_filtered_df.empty:
+        raise ValueError(
+            f"No data found for kernel_name='{config.kernel_name}', "
+            f"metric_name='{config.metric_name}', "
+            f"kernel_operation_mode='{config.kernel_operation_mode}'."
+        )
 
-    return filtered_df
+    unique_extra_configs_str = base_filtered_df["extra_benchmark_config_str"].unique()
+    selected_extra_config_str = None
+
+    if len(unique_extra_configs_str) == 0:
+        print(
+            "Warning: No extra_benchmark_config found for the initial filters. "
+            "Proceeding with all data from initial filter."
+        )
+        return base_filtered_df
+
+    if config.extra_config_filter:
+        matched_configs = []
+        try:
+            if "=" in config.extra_config_filter:
+                key_filter, value_filter = config.extra_config_filter.split("=", 1)
+                for cfg_str in unique_extra_configs_str:
+                    cfg_json = json.loads(cfg_str)
+                    if str(cfg_json.get(key_filter.strip("'\" "))) == value_filter.strip("'\" "):
+                        matched_configs.append(cfg_str)
+            if not matched_configs:
+                matched_configs = [
+                    cfg_str for cfg_str in unique_extra_configs_str if config.extra_config_filter in cfg_str
+                ]
+        except Exception as e:
+            print(
+                f"Note: Could not parse extra_config_filter '{config.extra_config_filter}' as key=value ({e}), using substring match."
+            )
+            matched_configs = [cfg_str for cfg_str in unique_extra_configs_str if config.extra_config_filter in cfg_str]
+
+        if matched_configs:
+            if len(matched_configs) > 1:
+                print(
+                    f"Warning: Multiple extra_benchmark_configs match filter '{config.extra_config_filter}': {matched_configs}. "
+                    f"Using the first one: {matched_configs[0]}"
+                )
+            selected_extra_config_str = matched_configs[0]
+        else:
+            print(
+                f"Warning: No extra_benchmark_config matches filter '{config.extra_config_filter}'. "
+                f"Available configs for {config.kernel_name} ({config.metric_name}, {config.kernel_operation_mode}): {list(unique_extra_configs_str)}"
+            )
+            if len(unique_extra_configs_str) > 0:
+                selected_extra_config_str = unique_extra_configs_str[0]
+                print(f"Defaulting to the first available extra_benchmark_config: {selected_extra_config_str}")
+            else:
+                raise ValueError("No extra_benchmark_config available to select after failed filter attempt.")
+
+    elif len(unique_extra_configs_str) > 1:
+        selected_extra_config_str = unique_extra_configs_str[0]
+        print(
+            f"Warning: Multiple extra_benchmark_configs found for {config.kernel_name} ({config.metric_name}, {config.kernel_operation_mode})."
+        )
+        print(f"Defaulting to use: {selected_extra_config_str}")
+        print(f"Available configs: {list(unique_extra_configs_str)}")
+        print(
+            "Use the --extra-config-filter argument to select a specific one "
+            "(e.g., --extra-config-filter \"'H': 4096\" or a substring like \"'seq_len': 512\")."
+        )
+    elif len(unique_extra_configs_str) == 1:
+        selected_extra_config_str = unique_extra_configs_str[0]
+        print(f"Using unique extra_benchmark_config: {selected_extra_config_str}")
+
+    if selected_extra_config_str:
+        final_filtered_df = base_filtered_df[
+            base_filtered_df["extra_benchmark_config_str"] == selected_extra_config_str
+        ]
+    else:
+        print("Warning: Could not select an extra_benchmark_config. Using data from initial filter if any.")
+        final_filtered_df = base_filtered_df
+
+    if final_filtered_df.empty:
+        raise ValueError(
+            f"No data found after attempting to filter by extra_benchmark_config. "
+            f"Selected/Defaulted extra_config_str: {selected_extra_config_str}"
+            if selected_extra_config_str
+            else "No specific extra_config was selected."
+        )
+
+    print(
+        f"Plotting data for extra_benchmark_config: {json.loads(selected_extra_config_str if selected_extra_config_str else '{}')}"
+    )
+    return final_filtered_df
 
 
 def plot_data(df: pd.DataFrame, config: VisualizationsConfig):
@@ -103,6 +197,10 @@ def plot_data(df: pd.DataFrame, config: VisualizationsConfig):
         df (pd.DataFrame): Filtered benchmark dataframe.
         config (VisualizationsConfig): Configuration object for the visualizations script.
     """
+    for col in ["y_value_20", "y_value_50", "y_value_80"]:
+        if col in df.columns:
+            df[col] = pd.to_numeric(df[col], errors="coerce")
+
     xlabel = df["x_label"].iloc[0]
     ylabel = f"{config.metric_name} ({df['metric_unit'].iloc[0]})"
     # Sort by "kernel_provider" to ensure consistent color assignment
@@ -110,15 +208,26 @@ def plot_data(df: pd.DataFrame, config: VisualizationsConfig):
 
     plt.figure(figsize=(10, 6))
     sns.set(style="whitegrid")
-    ax = sns.lineplot(
-        data=df,
-        x="x_value",
-        y="y_value_50",
-        hue="kernel_provider",
-        marker="o",
-        palette="tab10",
-        errorbar=("ci", None),
-    )
+    try:
+        ax = sns.lineplot(
+            data=df,
+            x="x_value",
+            y="y_value_50",
+            hue="kernel_provider",
+            marker="o",
+            palette="tab10",
+            errorbar=("ci", None),
+        )
+    except Exception:
+        ax = sns.lineplot(
+            data=df,
+            x="x_value",
+            y="y_value_50",
+            hue="kernel_provider",
+            marker="o",
+            palette="tab10",
+            errorbar=None,
+        )
 
     # Seaborn can't plot pre-computed error bars, so we need to do it manually
     lines = ax.get_lines()