DOCSEARCHbyAlgolia
Submit Docs

FBGEMM

repository·main·Indexed 23 days ago

https://github.com/pytorch/fbgemm

A project providing highly-optimized kernels for deep learning, including FBGEMM for CPU-based low-precision inference and FBGEMM_GPU for recommendation systems and Generative AI. FBGEMM GenAI supports FP8 and INT4 low-precision inference, MetaShuffling MoE kernels, and optimized operators for Blackwell GPUs, including fused multi-head attention (FMHA) and multi-head latent attention (MLA). It also includes HSTU high-performance attention implementations for NVIDIA Ampere and Hopper architectures.

Tokens
29.4K
Snippets
72
Records
124
Agent score
77%

What's inside pytorch-fbgemm

  1. Overview of FBGEMM GenAI

    main

    FBGEMM GenAI (FBGEMM Generative AI Kernels Library) is a collection of PyTorch GPU operator libraries specifically designed for generative AI applications. Its primary capabilities include:

    • FP8 row-wise quantization: Optimized kernels for low-precision quantization.
    • Collective communications: Specialized operators for distributed generative AI workloads.

  2. Overview of FBGEMM_GPU

    main

    FBGEMM_GPU (FBGEMM GPU Kernels Library) is a collection of high-performance PyTorch GPU operator libraries designed for both training and inference. It specializes in providing efficient implementations for:

    • Table batched embedding bag: Optimized operations for large-scale embedding lookups.
    • Data layout transformation: Efficiently changing data formats for GPU processing.
    • Quantization support: Tools and kernels for quantized operations to improve performance and reduce memory footprint.

  3. Overview of FBGEMM GenAI capabilities

    main

    FBGEMM GenAI provides optimized kernels for generative AI workloads, specifically focusing on low-precision inference (FP8 and INT4).

    Supported FP8 Configurations:

    • GEMM Operators: Supports {CUTLASS, CK, Triton} backends with {BF16, FP8} precision and {tensor-wise, row-wise, block-wise} scaling across Nvidia H100 and AMD MI300x hardware.
    • Precision Conversion: High/low precision conversion (FP32 / BF16 <-> FP8) with scaling options {tensor-wise, row-wise, block-wise} using {Triton, CUDA/HIP} programming models.

    Other Supported Features:

    • INT4 Support: On-the-fly quantization kernels (used in Llama 4).
    • Communication: Customized AllReduce communications optimized for small message sizes.
    • Attention & Cache: Optimized GQA (Grouped Query Attention) for decoding, and KV cache quantization.
    • Embeddings: Rotary Positional Embedding (RoPE).
    • MoE: MetaShuffling MoE (Mixture of Experts) operators.

  5. Overview of MetaShuffling MoE support in FBGEMM GenAI

    main
    FBGEMM GenAI provides MetaShuffling MoE (Mixture-of-Experts) kernel support. MetaShuffling is an inference solution designed to address the high memory and communication pressure, as well as the dynamism and sparsity, inherent in MoE architectures like Llama 4. It enables efficient deployment for real-world inference scenarios by optimizing how experts are accessed and computed.

  6. Overview of the FBGEMM Project packages

    main

    The FBGEMM Project consists of three distinct packages tailored for different deep learning workloads:

    1. FBGEMM: A low-precision, high-performance matrix-matrix multiplication and convolution library optimized for server-side inference. It serves as the backend for PyTorch quantized operators on x86 machines.
    2. FBGEMM_GPU: A collection of high-performance PyTorch GPU operator libraries built on top of FBGEMM. It focuses on recommendation systems, providing efficient table batched embedding (TBE) bag, data layout transformation, and quantization support.
    3. FBGEMM GenAI: A collection of PyTorch GPU operator libraries specifically designed for generative AI applications, including support for FP8 row-wise quantization and collective communications.

  7. HSTU-Ampere Features and Capabilities

    main

    HSTU-Ampere is designed for Ampere and Ada GPUs (e.g., A100, L20).

    Capabilities:

    • Data types: FP16, BF16.
    • Head dimensions: 32, 64, 128, 256.
    • Paged attention:
      • Supports one specific mask configuration.
      • The sequence length of k and v must match q.
      • The first few items of k and v (representing new history) are invalid; actual values reside in the paged kv_cache.
      • Supported page sizes: 32 and 64.
    • Attention masks supported:
      • No mask
      • Local mask (0 <= window_size_left < max_seqlen_k or 0 <= window_size_right < max_seqlen_k)
      • Causal mask (window_size_left = -1 and window_size_right = 0)
      • Context mask + causal mask
      • Target mask + causal mask
      • Context mask + causal mask + target mask
      • Delta_q (automatically used if seqlen_q < seqlen_k)
      • Delta_q + local mask
      • Delta_q + causal mask
      • Arbitrary mask (using an array to determine mask situation for each row of S).

  8. Understand the FBGEMM Project packages

    main

    The FBGEMM Project is organized into three distinct packages, each serving different hardware and application needs:

    • FBGEMM: A low-precision, high-performance matrix multiplication and convolution library optimized for server-side inference on x86 machines. It is used as a backend for PyTorch quantized operators.
    • FBGEMM_GPU: A collection of PyTorch GPU operator libraries built on top of FBGEMM, specifically focused on training and inference for recommendation systems.
    • FBGEMM_GPU GenAI: A collection of PyTorch GPU operator libraries designed for generative AI applications, featuring capabilities like FP8 row-wise quantization and collective communications.

  9. Understand FBGEMM_GPU Stable API guarantees

    main

    Starting from FBGEMM_GPU v1.0, certain APIs are designated as 'Stable'. These APIs come with the following guarantees:

    • Backward Compatibility: Stable APIs are guaranteed to be backward compatible within a major version. For example, stable APIs in v1.0.0 will remain compatible with future releases unless a breaking change is explicitly announced.
    • Documentation: Every stable API is accompanied by comprehensive and up-to-date documentation.
    • Functionality: Functionality is only guaranteed for features that are explicitly documented and covered by the unit testing framework.

    Note on Performance: There are no explicit performance guarantees for stable APIs, though the project provides support on a best-effort basis.

  10. Convert between Jagged and Dense tensors

    main

    FBGEMM_GPU allows converting between jagged and dense tensor formats.

    Jagged to Dense

    When converting a jagged tensor J to a dense tensor D:

    • The rank of the dense tensor is: rank(D) = len(MaxLengths) + 2.
    • For dimensions 0 < i < D.rank-1, the size is: dim(i) = MaxLengths[i-1].
    • Elements from Values are mapped to the dense tensor based on the ranges in Offsets.
    • Areas in the dense tensor not covered by the jagged partitions are padded with a pad value (typically 0).

    Dense to Jagged

    When converting from a dense tensor to a jagged tensor, values are loaded into the jagged tensor's Values. If the dense tensor's dimensions are smaller than what the Offsets expect, the corresponding Values are filled with the padded value.