pr-include-rev-in-flake

README.md

@@ -1,80 +1,9 @@
----
-license: bsd-3-clause
-tags:
-- kernel
----
-# Triton layer normalization kernels.
-
-This kernel implements layers normalization using Triton. This kernel is from
-the [flash-attention](https://github.com/Dao-AILab/flash-attention) project.
-
-## Functions
-
-### Function `layer_norm`
-
-`(x: torch.Tensor, weight: torch.Tensor, bias: torch.Tensor, residual: Optional[torch.Tensor] = None, x1: Optional[torch.Tensor] = None, weight1: Optional[torch.Tensor] = None, bias1: Optional[torch.Tensor] = None, eps: float = 1e-06, dropout_p: float = 0.0, rowscale=None, prenorm: bool = False, residual_in_fp32: bool = False, is_rms_norm: bool = False, return_dropout_mask: bool = False, out: Optional[torch.Tensor] = None, residual_out: Optional[torch.Tensor] = None)`
-
-Apply layer normalization to the input tensor with Triton acceleration.
-
-### Parameters
-
-- **x** (*torch.Tensor*) --
-  Input tensor to normalize.
-- **weight** (*torch.Tensor*) --
-  Scale parameter for normalization.
-- **bias** (*torch.Tensor*) --
-  Shift parameter for normalization.
-- **residual** (*torch.Tensor*, *optional*) --
-  Optional residual tensor to add to the input before normalization.
-- **x1** (*torch.Tensor*, *optional*) --
-  Optional second input tensor to combine with *x*. When provided, the function
-  first adds *x1* to *x* and then applies normalization.
-- **weight1** (*torch.Tensor*, *optional*) --
-  Scale parameter for the second normalization.
-- **bias1** (*torch.Tensor*, *optional*) --
-  Shift parameter for the second normalization.
-- **eps** (*float*, *optional*, defaults to 1e-6) --
-  Small constant added for numerical stability in normalization.
-- **dropout_p** (*float*, *optional*, defaults to 0.0) --
-  Dropout probability. If greater than 0, applies dropout to the input before
-  normalization and residual addition.
-- **rowscale** (*torch.Tensor*, *optional*) --
-  Optional scaling factor applied to each row of the input tensor.
-  Not compatible with the use of *x1*.
-- **prenorm** (*bool*, *optional*, defaults to False) --
-  If True, returns both the normalized output and the unnormalized input+residual.
-- **residual_in_fp32** (*bool*, *optional*, defaults to False) --
-  If True, performs the residual connection in FP32 precision.
-- **is_rms_norm** (*bool*, *optional*, defaults to False) --
-  If True, uses RMS normalization instead of layer normalization.
-- **return_dropout_mask** (*bool*, *optional*, defaults to False) --
-  If True, returns the dropout mask used for the computation.
-- **out** (*torch.Tensor*, *optional*) --
-  Output tensor for the normalized result. If *None*, a new tensor is allocated.
-- **residual_out** (*torch.Tensor*, *optional*) --
-  Output tensor for the residual result when using prenorm. If *None*, a new tensor
-  is allocated when needed.
-
-### Returns
-
-**Type**: *torch.Tensor* or tuple of *torch.Tensor*
-
-- The normalized input.
-- The second normalization of the input if *weight1* is provided.
-- The residual tensor if *prenorm* is set.
-- The dropout mask if *return_dropout_mask* is set.
-- The dropout mask for *x1* if *x1* is provided and *return_dropout_mask* is set.
-
-## Layers
-
-### Class `LlamaRMSNorm`
-
-No documentation available.
-
-#### Methods
-
-##### Method `forward`
-
-`(self, hidden_states: torch.Tensor) -> torch.Tensor`
-
-No documentation available.
+---
+license: bsd-3-clause
+tags:
+  - kernel
+---
+
+## triton-layer-norm
+
+Triton layer norm [from flash-attention](https://github.com/Dao-AILab/flash-attention).

build.toml

@@ -1,3 +1,5 @@
 [general]
 name = "triton_layer_norm"
+
+[torch]
 universal = true

build/torch-universal/triton_layer_norm/__init__.py

@@ -1,114 +1,5 @@
-"""Triton layer normalization kernels
-
-This kernel implements layers normalization using Triton. This kernel is from
-the `flash-attention <https://github.com/Dao-AILab/flash-attention>`_ project.
-"""
-
-from typing import Optional
-
-import torch
-
-from . import layers
 from .layer_norm import layer_norm_fn, layer_norm_linear_fn, rms_norm_fn
 
+from . import layers
 
-def layer_norm(
-    x: torch.Tensor,
-    weight: torch.Tensor,
-    bias: torch.Tensor,
-    residual: Optional[torch.Tensor] = None,
-    x1: Optional[torch.Tensor] = None,
-    weight1: Optional[torch.Tensor] = None,
-    bias1: Optional[torch.Tensor] = None,
-    eps: float = 1e-6,
-    dropout_p: float = 0.0,
-    rowscale=None,
-    prenorm: bool = False,
-    residual_in_fp32: bool = False,
-    is_rms_norm: bool = False,
-    return_dropout_mask: bool = False,
-    out: Optional[torch.Tensor] = None,
-    residual_out: Optional[torch.Tensor] = None,
-):
-    """
-    Apply layer normalization to the input tensor with Triton acceleration.
-
-    Args:
-        x (`torch.Tensor`):
-            Input tensor to normalize.
-        weight (`torch.Tensor`):
-            Scale parameter for normalization.
-        bias (`torch.Tensor`):
-            Shift parameter for normalization.
-        residual (`torch.Tensor`, *optional*):
-            Optional residual tensor to add to the input before normalization.
-        x1 (`torch.Tensor`, *optional*):
-            Optional second input tensor to combine with `x`. When provided, the function
-            first adds `x1` to `x` and then applies normalization.
-        weight1 (`torch.Tensor`, *optional*):
-            Scale parameter for the second normalization.
-        bias1 (`torch.Tensor`, *optional*):
-            Shift parameter for the second normalization.
-        eps (`float`, *optional*, defaults to 1e-6):
-            Small constant added for numerical stability in normalization.
-        dropout_p (`float`, *optional*, defaults to 0.0):
-            Dropout probability. If greater than 0, applies dropout to the input before
-            normalization and residual addition.
-        rowscale (`torch.Tensor`, *optional*):
-            Optional scaling factor applied to each row of the input tensor.
-            Not compatible with the use of `x1`.
-        prenorm (`bool`, *optional*, defaults to False):
-            If True, returns both the normalized output and the unnormalized input+residual.
-        residual_in_fp32 (`bool`, *optional*, defaults to False):
-            If True, performs the residual connection in FP32 precision.
-        is_rms_norm (`bool`, *optional*, defaults to False):
-            If True, uses RMS normalization instead of layer normalization.
-        return_dropout_mask (`bool`, *optional*, defaults to False):
-            If True, returns the dropout mask used for the computation.
-        out (`torch.Tensor`, *optional*):
-            Output tensor for the normalized result. If `None`, a new tensor is allocated.
-        residual_out (`torch.Tensor`, *optional*):
-            Output tensor for the residual result when using prenorm. If `None`, a new tensor
-            is allocated when needed.
-
-    Returns:
-        `torch.Tensor` or tuple of `torch.Tensor`:
-            - The normalized input.
-            - The second normalization of the input if `weight1` is provided.
-            - The residual tensor if `prenorm` is set.
-            - The dropout mask if `return_dropout_mask` is set.
-            - The dropout mask for `x1` if `x1` is provided and `return_dropout_mask` is set.
-    """
-    return layer_norm_fn(
-        x,
-        weight,
-        bias,
-        residual,
-        x1,
-        weight1,
-        bias1,
-        eps,
-        dropout_p,
-        rowscale,
-        prenorm,
-        residual_in_fp32,
-        is_rms_norm,
-        return_dropout_mask,
-        out=out,
-        residual_out=residual_out,
-    )
-
-
-__kernel_metadata__ = {
-    "license": "bsd-3-clause",
-}
-
-
-__all__ = [
-    "__kernel_metadata__",
-    "layers",
-    "layer_norm",
-    "layer_norm_fn",
-    "layer_norm_linear_fn",
-    "rms_norm_fn",
-]
+__all__ = ["layers", "layer_norm_fn", "layer_norm_linear_fn", "rms_norm_fn"]

build/torch-universal/triton_layer_norm/_ops.py

@@ -1,8 +0,0 @@
-import torch
-ops = torch.ops._triton_layer_norm_4dc3a9b_dirty
-
-def add_op_namespace_prefix(op_name: str):
-    """
-    Prefix op by namespace.
-    """
-    return f"_triton_layer_norm_4dc3a9b_dirty::{op_name}"

build/torch-universal/triton_layer_norm/layers.py

@@ -1,46 +1,4 @@
-import torch
-from torch import nn
+from .layer_norm import RMSNorm
 
-from .layer_norm import rms_norm_fn
 
-
-class LlamaRMSNorm(nn.Module):
-    """
-    RMS Layer Norm for Llama models.
-
-    Triton-optimized RMS layer norm. The interface is compatible with `LLamaRMSNorm` in
-    `transformers`.
-
-    Attributes:
-        weight (`torch.Tensor`): The learnable scaling parameter.
-        variance_epsilon (`float`): The epsilon value for numerical stability.
-    """
-    weight: torch.Tensor
-    variance_epsilon: float
-
-    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
-        """
-        Apply RMS normalization to the input hidden states.
-
-        Args:
-            hidden_states (`torch.Tensor`):
-                Input tensor of shape `(batch_size, sequence_length, hidden_size)` or any shape
-                where the last dimension is the feature dimension to be normalized.
-
-        Returns:
-            `torch.Tensor`:
-                The normalized tensor with the same shape as the input `hidden_states`.
-        """
-        return rms_norm_fn(
-            hidden_states,
-            self.weight,
-            bias=None,
-            residual=None,
-            eps=self.variance_epsilon,
-            dropout_p=0.0,
-            prenorm=False,
-            residual_in_fp32=False,
-        )
-
-
-__all__ = ["LlamaRMSNorm"]
+__all__ = ["RMSNorm"]

flake.lock

@@ -1,168 +0,0 @@
-{
-  "nodes": {
-    "flake-compat": {
-      "locked": {
-        "lastModified": 1747046372,
-        "narHash": "sha256-CIVLLkVgvHYbgI2UpXvIIBJ12HWgX+fjA8Xf8PUmqCY=",
-        "owner": "edolstra",
-        "repo": "flake-compat",
-        "rev": "9100a0f413b0c601e0533d1d94ffd501ce2e7885",
-        "type": "github"
-      },
-      "original": {
-        "owner": "edolstra",
-        "repo": "flake-compat",
-        "type": "github"
-      }
-    },
-    "flake-compat_2": {
-      "locked": {
-        "lastModified": 1733328505,
-        "narHash": "sha256-NeCCThCEP3eCl2l/+27kNNK7QrwZB1IJCrXfrbv5oqU=",
-        "owner": "edolstra",
-        "repo": "flake-compat",
-        "rev": "ff81ac966bb2cae68946d5ed5fc4994f96d0ffec",
-        "type": "github"
-      },
-      "original": {
-        "owner": "edolstra",
-        "repo": "flake-compat",
-        "type": "github"
-      }
-    },
-    "flake-utils": {
-      "inputs": {
-        "systems": "systems"
-      },
-      "locked": {
-        "lastModified": 1731533236,
-        "narHash": "sha256-l0KFg5HjrsfsO/JpG+r7fRrqm12kzFHyUHqHCVpMMbI=",
-        "owner": "numtide",
-        "repo": "flake-utils",
-        "rev": "11707dc2f618dd54ca8739b309ec4fc024de578b",
-        "type": "github"
-      },
-      "original": {
-        "owner": "numtide",
-        "repo": "flake-utils",
-        "type": "github"
-      }
-    },
-    "flake-utils_2": {
-      "inputs": {
-        "systems": "systems_2"
-      },
-      "locked": {
-        "lastModified": 1731533236,
-        "narHash": "sha256-l0KFg5HjrsfsO/JpG+r7fRrqm12kzFHyUHqHCVpMMbI=",
-        "owner": "numtide",
-        "repo": "flake-utils",
-        "rev": "11707dc2f618dd54ca8739b309ec4fc024de578b",
-        "type": "github"
-      },
-      "original": {
-        "owner": "numtide",
-        "repo": "flake-utils",
-        "type": "github"
-      }
-    },
-    "hf-nix": {
-      "inputs": {
-        "flake-compat": "flake-compat_2",
-        "flake-utils": "flake-utils_2",
-        "nixpkgs": "nixpkgs"
-      },
-      "locked": {
-        "lastModified": 1750234878,
-        "narHash": "sha256-q9DRC9zdpzUf88qqg1qbhP1qgJbE2cMtn8oUmosuyT8=",
-        "owner": "huggingface",
-        "repo": "hf-nix",
-        "rev": "c7132f90763d756da3e77da62e01be0a4546dc57",
-        "type": "github"
-      },
-      "original": {
-        "owner": "huggingface",
-        "repo": "hf-nix",
-        "type": "github"
-      }
-    },
-    "kernel-builder": {
-      "inputs": {
-        "flake-compat": "flake-compat",
-        "flake-utils": "flake-utils",
-        "hf-nix": "hf-nix",
-        "nixpkgs": [
-          "kernel-builder",
-          "hf-nix",
-          "nixpkgs"
-        ]
-      },
-      "locked": {
-        "lastModified": 1750409351,
-        "narHash": "sha256-xkzrwee77LrBDtwNNihBkYbY7yUwdOv0/4+J3B5xCZE=",
-        "owner": "huggingface",
-        "repo": "kernel-builder",
-        "rev": "9e61fba877153bffa6eaff023243fd81220c0eea",
-        "type": "github"
-      },
-      "original": {
-        "owner": "huggingface",
-        "repo": "kernel-builder",
-        "type": "github"
-      }
-    },
-    "nixpkgs": {
-      "locked": {
-        "lastModified": 1747820358,
-        "narHash": "sha256-fTqsZsUX6M3yeEvgyQvXcbGmT2CaRVyVwsi8eK29Oj4=",
-        "owner": "danieldk",
-        "repo": "nixpkgs",
-        "rev": "d3c1681180717528068082103bf323147de6ab0b",
-        "type": "github"
-      },
-      "original": {
-        "owner": "danieldk",
-        "ref": "cudatoolkit-12.9-kernel-builder",
-        "repo": "nixpkgs",
-        "type": "github"
-      }
-    },
-    "root": {
-      "inputs": {
-        "kernel-builder": "kernel-builder"
-      }
-    },
-    "systems": {
-      "locked": {
-        "lastModified": 1681028828,
-        "narHash": "sha256-Vy1rq5AaRuLzOxct8nz4T6wlgyUR7zLU309k9mBC768=",
-        "owner": "nix-systems",
-        "repo": "default",
-        "rev": "da67096a3b9bf56a91d16901293e51ba5b49a27e",
-        "type": "github"
-      },
-      "original": {
-        "owner": "nix-systems",
-        "repo": "default",
-        "type": "github"
-      }
-    },
-    "systems_2": {
-      "locked": {
-        "lastModified": 1681028828,
-        "narHash": "sha256-Vy1rq5AaRuLzOxct8nz4T6wlgyUR7zLU309k9mBC768=",
-        "owner": "nix-systems",
-        "repo": "default",
-        "rev": "da67096a3b9bf56a91d16901293e51ba5b49a27e",
-        "type": "github"
-      },
-      "original": {
-        "owner": "nix-systems",
-        "repo": "default",
-        "type": "github"
-      }
-    }
-  },
-  "root": "root",
-  "version": 7
-}

flake.nix

@@ -10,8 +10,5 @@
       self,
       kernel-builder,
     }:
-    kernel-builder.lib.genFlakeOutputs {
-      path = ./.;
-      rev = self.shortRev or self.dirtyShortRev or self.lastModifiedDate;
-    };
+    kernel-builder.lib.genFlakeOutputs ./.;
 }

torch-ext/triton_layer_norm/__init__.py

@@ -1,114 +1,5 @@
-"""Triton layer normalization kernels
-
-This kernel implements layers normalization using Triton. This kernel is from
-the `flash-attention <https://github.com/Dao-AILab/flash-attention>`_ project.
-"""
-
-from typing import Optional
-
-import torch
-
-from . import layers
 from .layer_norm import layer_norm_fn, layer_norm_linear_fn, rms_norm_fn
 
+from . import layers
 
-def layer_norm(
-    x: torch.Tensor,
-    weight: torch.Tensor,
-    bias: torch.Tensor,
-    residual: Optional[torch.Tensor] = None,
-    x1: Optional[torch.Tensor] = None,
-    weight1: Optional[torch.Tensor] = None,
-    bias1: Optional[torch.Tensor] = None,
-    eps: float = 1e-6,
-    dropout_p: float = 0.0,
-    rowscale=None,
-    prenorm: bool = False,
-    residual_in_fp32: bool = False,
-    is_rms_norm: bool = False,
-    return_dropout_mask: bool = False,
-    out: Optional[torch.Tensor] = None,
-    residual_out: Optional[torch.Tensor] = None,
-):
-    """
-    Apply layer normalization to the input tensor with Triton acceleration.
-
-    Args:
-        x (`torch.Tensor`):
-            Input tensor to normalize.
-        weight (`torch.Tensor`):
-            Scale parameter for normalization.
-        bias (`torch.Tensor`):
-            Shift parameter for normalization.
-        residual (`torch.Tensor`, *optional*):
-            Optional residual tensor to add to the input before normalization.
-        x1 (`torch.Tensor`, *optional*):
-            Optional second input tensor to combine with `x`. When provided, the function
-            first adds `x1` to `x` and then applies normalization.
-        weight1 (`torch.Tensor`, *optional*):
-            Scale parameter for the second normalization.
-        bias1 (`torch.Tensor`, *optional*):
-            Shift parameter for the second normalization.
-        eps (`float`, *optional*, defaults to 1e-6):
-            Small constant added for numerical stability in normalization.
-        dropout_p (`float`, *optional*, defaults to 0.0):
-            Dropout probability. If greater than 0, applies dropout to the input before
-            normalization and residual addition.
-        rowscale (`torch.Tensor`, *optional*):
-            Optional scaling factor applied to each row of the input tensor.
-            Not compatible with the use of `x1`.
-        prenorm (`bool`, *optional*, defaults to False):
-            If True, returns both the normalized output and the unnormalized input+residual.
-        residual_in_fp32 (`bool`, *optional*, defaults to False):
-            If True, performs the residual connection in FP32 precision.
-        is_rms_norm (`bool`, *optional*, defaults to False):
-            If True, uses RMS normalization instead of layer normalization.
-        return_dropout_mask (`bool`, *optional*, defaults to False):
-            If True, returns the dropout mask used for the computation.
-        out (`torch.Tensor`, *optional*):
-            Output tensor for the normalized result. If `None`, a new tensor is allocated.
-        residual_out (`torch.Tensor`, *optional*):
-            Output tensor for the residual result when using prenorm. If `None`, a new tensor
-            is allocated when needed.
-
-    Returns:
-        `torch.Tensor` or tuple of `torch.Tensor`:
-            - The normalized input.
-            - The second normalization of the input if `weight1` is provided.
-            - The residual tensor if `prenorm` is set.
-            - The dropout mask if `return_dropout_mask` is set.
-            - The dropout mask for `x1` if `x1` is provided and `return_dropout_mask` is set.
-    """
-    return layer_norm_fn(
-        x,
-        weight,
-        bias,
-        residual,
-        x1,
-        weight1,
-        bias1,
-        eps,
-        dropout_p,
-        rowscale,
-        prenorm,
-        residual_in_fp32,
-        is_rms_norm,
-        return_dropout_mask,
-        out=out,
-        residual_out=residual_out,
-    )
-
-
-__kernel_metadata__ = {
-    "license": "bsd-3-clause",
-}
-
-
-__all__ = [
-    "__kernel_metadata__",
-    "layers",
-    "layer_norm",
-    "layer_norm_fn",
-    "layer_norm_linear_fn",
-    "rms_norm_fn",
-]
+__all__ = ["layers", "layer_norm_fn", "layer_norm_linear_fn", "rms_norm_fn"]

torch-ext/triton_layer_norm/layers.py

@@ -1,46 +1,4 @@
-import torch
-from torch import nn
+from .layer_norm import RMSNorm
 
-from .layer_norm import rms_norm_fn
 
-
-class LlamaRMSNorm(nn.Module):
-    """
-    RMS Layer Norm for Llama models.
-
-    Triton-optimized RMS layer norm. The interface is compatible with `LLamaRMSNorm` in
-    `transformers`.
-
-    Attributes:
-        weight (`torch.Tensor`): The learnable scaling parameter.
-        variance_epsilon (`float`): The epsilon value for numerical stability.
-    """
-    weight: torch.Tensor
-    variance_epsilon: float
-
-    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
-        """
-        Apply RMS normalization to the input hidden states.
-
-        Args:
-            hidden_states (`torch.Tensor`):
-                Input tensor of shape `(batch_size, sequence_length, hidden_size)` or any shape
-                where the last dimension is the feature dimension to be normalized.
-
-        Returns:
-            `torch.Tensor`:
-                The normalized tensor with the same shape as the input `hidden_states`.
-        """
-        return rms_norm_fn(
-            hidden_states,
-            self.weight,
-            bias=None,
-            residual=None,
-            eps=self.variance_epsilon,
-            dropout_p=0.0,
-            prenorm=False,
-            residual_in_fp32=False,
-        )
-
-
-__all__ = ["LlamaRMSNorm"]
+__all__ = ["RMSNorm"]