build/torch-cuda/__init__.py

"""Triton ops vendored from `flash-attn` (`flash_attn/ops/triton` and
`flash_attn/losses`), packaged as a self-contained, `kernels`-compliant kernel.

Upstream: https://github.com/Dao-AILab/flash-attention
Pinned commit: b02b07e1a10238fe12831b80a8937ed59b1353a5

These ops are pure Triton + PyTorch: there is nothing to compile ahead of time,
and the package does not import `flash_attn` at runtime.
"""

from .cross_entropy import cross_entropy_loss
from .losses import CrossEntropyLoss
from .rotary import apply_rotary
from .layer_norm import (
    LayerNormFn,
    RMSNorm,
    layer_norm_fn,
    layer_norm_linear_fn,
    rms_norm_fn,
)

__all__ = [
    # cross entropy
    "cross_entropy_loss",
    "CrossEntropyLoss",
    # rotary
    "apply_rotary",
    # layer / rms norm
    "layer_norm_fn",
    "rms_norm_fn",
    "layer_norm_linear_fn",
    "LayerNormFn",
    "RMSNorm",
]