flagos-ai/FlagGems

FlagGems is an operator library for large language models implemented in the Triton Language.

What it solves

FlagGems is designed to eliminate fragmentation among AI chipset-specific software stacks. It solves the problem of having to rewrite or maintain separate codebases for different AI accelerators, enabling a "develop once, run anywhere" workflow for LLM training and inference.

How it works

It is a high-performance, generic operator library implemented in the Triton language. By registering with the PyTorch ATen backend, it allows developers to use familiar PyTorch APIs while the underlying operations are executed via Triton kernels. This provides a backend-neutral approach that supports over 10 different hardware platforms without requiring the user to change low-level APIs.

Who it’s for

  • Model developers who want to hardware acceleration without changing their PyTorch code.
  • Kernel developers who want to write high-performance kernels using a more readable and user-friendly language than CUDA.

Highlights

  • Broad Hardware Support: Supports over 10 backends for diverse AI accelerators.
  • PyTorch Compatibility: Seamlessly integrates with PyTorch APIs, remaining independent of torch.compile and ready for eager-mode.
  • Triton-based: Leverages Triton for performance comparable to CUDA with better readability.
  • Automatic Codegen: Includes automatic pointwise operator codegen supporting arbitrary input types and layouts.
  • Efficient Dispatching: Features fast per-function runtime kernel dispatching.

Related

  • Project
    flagos-ai/FlagTree

    An open-source unified compiler for multiple AI chips that leverages the Triton ecosystem to enable a "develop once, run anywhere" workflow across diverse hardware accelerators.

  • Project
    ServiceNow/Fast-LLM

    A high-performance library for training large language models that uses 3D parallelism and optimized kernels to reduce training time and cost.

  • Project
    meta-pytorch/tritonbench

    A collection of PyTorch operators used to evaluate the performance of Triton and its integration with PyTorch.

  • Project
    linkedin/Liger-Kernel

    A collection of efficient Triton kernels for LLM training that increases throughput by 20% and reduces memory usage by 60% through kernel fusion and chunking.

  • Project
    ROCm/aiter

    AITER is AMD's high-performance AI operator library providing optimized GPU kernels for inference and training workloads on ROCm-compatible GPUs.