Developer documentation

docs/developer/architecture_overview.rst

@@ -0,0 +1,178 @@
+..
+    Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+
+    See LICENSE for license information.
+
+.. _architecture-overview:
+
+Architecture Overview
+=====================
+
+Transformer Engine is structured as a layered system: framework-specific Python frontends
+sit on top of a shared C++ core that dispatches to optimized CUDA kernels.
+
+.. figure:: ./img/architecture_layers.svg
+   :align: center
+   :width: 80%
+
+   The four layers of Transformer Engine.
+
+..
+   Diagram description for ``architecture_layers.svg``:
+   Four horizontal bands stacked vertically.
+   Top band: "Python API" with two sub-boxes "PyTorch Frontend" and "JAX Frontend".
+   Second band: "Binding Layer" with sub-boxes "pybind11 (PyTorch)" and "XLA FFI (JAX)".
+   Third band: "C++ Core Library" (single box spanning full width).
+   Bottom band: "CUDA Kernels" with sub-boxes for major areas: GEMM, Normalization,
+   Quantize, Activation, Attention, Other.
+   Arrows flow downward between bands. A side annotation shows "transformer_engine/common/"
+   pointing at the bottom two bands.
+
+Layer Diagram
+-------------
+
+**Python API** (``transformer_engine/pytorch/``, ``transformer_engine/jax/``)
+   Framework-specific modules, autograd functions, quantized tensor types, and distributed
+   utilities. This is what users import and interact with.
+
+**Binding Layer** (``transformer_engine/pytorch/csrc/``, ``transformer_engine/jax/csrc/``)
+   pybind11 extensions (PyTorch) and XLA custom-call registrations (JAX) that bridge
+   Python to C++. These translate framework tensor types into the C API's ``NVTETensor``.
+
+**C++ Core** (``transformer_engine/common/``)
+   Framework-agnostic library implementing the actual computation logic. Exposes a C API
+   (``transformer_engine.h``, etc.) for stability across framework bindings.
+
+**CUDA Kernels** (within ``transformer_engine/common/``)
+   Hand-optimized CUDA kernels and cuDNN/cuBLASLt integrations organized by functional
+   area.
+
+Directory Map
+-------------
+
+.. code-block:: text
+
+   transformer_engine/
+   ├── common/                        # C++ core (framework-agnostic)
+   │   ├── include/transformer_engine/ # Public C API headers
+   │   │   ├── transformer_engine.h    #   Base types: NVTEDType, NVTETensor, NVTEScalingMode
+   │   │   ├── fused_attn.h           #   Attention C API
+   │   │   ├── gemm.h                 #   GEMM C API
+   │   │   └── ...                    #   Other kernel APIs
+   │   ├── common.h                   #   C++ types: DType, SimpleTensor, Tensor
+   │   ├── gemm/                      #   cuBLASLt GEMM kernels
+   │   ├── fused_attn/                #   Fused attention (cuDNN + custom)
+   │   ├── normalization/             #   LayerNorm, RMSNorm kernels
+   │   ├── activation/                #   GeLU, SiLU, etc.
+   │   ├── cast/                      #   Quantization cast kernels
+   │   ├── transpose/                 #   Transpose + cast fusion
+   │   ├── fused_rope/                #   Rotary positional embeddings
+   │   ├── fused_softmax/             #   Scaled softmax variants
+   │   ├── fused_router/              #   MoE router kernels
+   │   └── comm_gemm/                 #   Communication-GEMM overlap
+   │
+   ├── pytorch/                       # PyTorch frontend
+   │   ├── module/                    #   nn.Module subclasses (Linear, LN, Attention, etc.)
+   │   ├── tensor/                    #   Quantized tensor implementations
+   │   ├── attention/                 #   Attention backends and dispatch
+   │   ├── ops/                       #   Op fusion framework
+   │   ├── cpp_extensions/            #   Python wrappers for C++ extensions
+   │   ├── csrc/                      #   pybind11 C++ source
+   │   ├── distributed.py             #   TP/SP/FSDP utilities
+   │   ├── quantized_tensor.py        #   Base quantization classes
+   │   └── transformer.py             #   TransformerLayer composition
+   │
+   └── jax/                           # JAX frontend
+       ├── flax/                      #   Flax nn.Module wrappers
+       ├── cpp_extensions/            #   XLA FFI primitive registration
+       ├── quantize/                  #   JAX quantizer and ScaledTensor
+       ├── attention.py               #   JAX attention entry points
+       ├── sharding.py                #   Mesh and sharding utilities
+       └── dense.py                   #   JAX dense (linear) layer
+
+Data Flow: Forward Pass
+-----------------------
+
+A typical forward pass through a quantized linear layer follows this path (see
+:doc:`linear_walkthrough` for a detailed end-to-end trace):
+
+.. figure:: ./img/data_flow_forward.svg
+   :align: center
+   :width: 90%
+
+   Data flow through a quantized linear forward pass.
+
+..
+   Diagram description for ``data_flow_forward.svg``:
+   Left-to-right flow diagram.
+   1. Box "Input (BF16/FP16)" with arrow to
+   2. Box "Quantize (cast kernel)" which outputs "FP8 data + scale" arrow to
+   3. Box "GEMM (cuBLASLt)" which also receives "Weight (FP8 + scale)" from above, arrow to
+   4. Box "Output (BF16/FP16)".
+   A dashed feedback arrow from GEMM back to Quantize labeled "amax history (delayed scaling)".
+   Below the main flow, a note: "Scale computation depends on NVTEScalingMode".
+
+1. **Input arrives** in high precision (BF16 or FP16).
+2. **Quantization** casts the input to FP8 (or other low-precision format), producing
+   quantized data and associated scaling factors. The specific cast kernel and scale
+   granularity depend on the :ref:`scaling mode <scaling-modes>`.
+3. **GEMM** executes in low precision via cuBLASLt, consuming quantized activations and
+   weights along with their scale inverses.
+4. **Output** is produced in high precision.
+
+For delayed tensor scaling, an amax (absolute maximum) feedback loop maintains scaling
+factor history across iterations.
+
+The backward pass follows the same pattern twice: once for the activation gradient (dgrad)
+and once for the weight gradient (wgrad). Both backward GEMMs use inputs already quantized
+during the forward pass (saved via autograd). The wgrad path uses *columnwise* quantized
+data to satisfy cuBLASLt's transposed layout requirements — see
+:doc:`quantization/rowwise_columnwise` for details.
+
+Guiding Principles
+-------------------
+
+These principles apply across the entire codebase and inform design decisions at every
+layer.
+
+**Three API surfaces, one core**
+   Transformer Engine exposes three distinct API surfaces — a C API
+   (``transformer_engine.h``), PyTorch ``nn.Module`` subclasses, and JAX XLA primitives —
+   all backed by the same framework-agnostic C++ core. CUDA kernels live in ``common/``
+   and know nothing about PyTorch or JAX; framework frontends are responsible for
+   converting their tensor types before calling into the C API.
+
+**Feel native to each framework**
+   Each frontend should feel like a natural extension of its framework, not a foreign
+   library. PyTorch modules follow ``nn.Module`` conventions (autograd, ``state_dict``,
+   ``torch.compile``). JAX primitives register with XLA's custom-call and sharding
+   machinery. Users should be able to mix and match TE components with native framework
+   code freely — a ``te.Linear`` should be a drop-in replacement for ``torch.nn.Linear``.
+
+**Leverage the ecosystem, don't reinvent it**
+   TE provides custom CUDA kernels where they deliver clear performance wins (fused
+   attention, quantized GEMM, fused normalization + cast). For everything else, we prefer
+   the framework's native capabilities. A custom kernel must justify its
+   existence against the ecosystem alternative.
+
+**No memory allocation in the C++ layer**
+   The C++ core never allocates GPU memory directly. All memory is allocated by the
+   framework frontend and passed down to the C API as pointers within ``NVTETensor``
+   handles. This keeps memory management in the framework's control (important for
+   memory pools, garbage collection, and distributed training) and avoids hidden
+   allocation surprises.
+
+**Hide complexity, but stay extensible**
+   Low-precision training involves many moving parts (scaling modes, amax history, block
+   sizes, format selection). TE hides this complexity behind simple APIs — a user can
+   pass a recipe to ``autocast`` and get FP8 training without understanding the
+   internals. But we do not want to be a black box: the quantization system is designed
+   to be extensible (see :doc:`quantization/adding_new_type`), scaling recipes are
+   configurable, and advanced users can provide custom quantizers.
+
+Next Steps
+----------
+
+- :doc:`linear_walkthrough` — End-to-end trace through a Linear forward and backward pass
+- :doc:`cpp_core/index` — C++ core type system, kernels, and build system
+- :doc:`quantization/index` — Quantization architecture and adding new types

docs/developer/attention/backend_selection.rst

@@ -0,0 +1,159 @@
+..
+    Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+
+    See LICENSE for license information.
+
+.. _backend-selection:
+
+Backend Selection
+=================
+
+Transformer Engine automatically selects the best attention backend based on the
+hardware, input configuration, and user preferences. This page documents the selection
+logic and how to override it.
+
+Selection Flow
+--------------
+
+The backend is selected at the start of each ``DotProductAttention.forward()`` call.
+The selection result is **cached** — if the attention parameters haven't changed since the
+last call, the cached backend is reused without re-running the selection logic (see
+``_attention_backends`` in ``dot_product_attention.py``).
+
+The selection logic lives in ``get_attention_backend()`` in
+``transformer_engine/pytorch/attention/dot_product_attention/utils.py``. It filters
+backends by compatibility (head dimension, sequence length, mask type, dropout, GQA
+groups, etc.) and then applies a priority order:
+
+.. code-block:: text
+
+   On Hopper+ (sm90):  FusedAttention > FlashAttention > Unfused
+   On pre-Hopper:      FlashAttention > FusedAttention > Unfused
+
+FusedAttention is preferred on Hopper+ for performance reasons. On pre-Hopper hardware,
+FlashAttention takes priority when both are available.
+
+SM120 (Blackwell) Architecture Notes
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+
+SM120 GPUs have additional restrictions that affect backend selection:
+
+- **KV caching**: Both FusedAttention and FlashAttention are disabled when KV caching
+  (inference mode) is active. Only the Unfused backend is available for KV-cached
+  inference on SM120.
+- **cuDNN version**: FusedAttention requires cuDNN >= 9.18.1 for THD layout on SM120.
+  With older cuDNN versions, FusedAttention is disabled for THD.
+- **Layout restrictions**: Even with cuDNN >= 9.18.1, the T3HD and TH3D layouts
+  (interleaved Q/K/V) are not supported on SM120.
+- **Deterministic mode**: Deterministic attention is not supported on SM120.
+- **Softmax LSE packed format**: The packed softmax LSE format for THD is excluded on
+  SM120 (requires cuDNN >= 9.6.0 and a non-SM120 architecture).
+
+Environment Variables
+---------------------
+
+.. list-table::
+   :header-rows: 1
+   :widths: 35 65
+
+   * - Variable
+     - Description
+   * - ``NVTE_FUSED_ATTN``
+     - ``0`` to disable cuDNN fused attention, ``1`` to enable (default: ``1``)
+   * - ``NVTE_FLASH_ATTN``
+     - ``0`` to disable FlashAttention, ``1`` to enable (default: ``1``)
+   * - ``NVTE_UNFUSED_ATTN``
+     - ``0`` to disable native PyTorch attention, ``1`` to enable (default: ``1``)
+
+Selecting a Backend
+-------------------
+
+``DotProductAttention`` does not expose a constructor argument that forces one backend.
+To select a backend for debugging or testing, enable the desired backend and disable the
+other two before the module is first called. For example, to select FlashAttention:
+
+.. code-block:: python
+
+   import os
+
+   os.environ["NVTE_FLASH_ATTN"] = "1"
+   os.environ["NVTE_FUSED_ATTN"] = "0"
+   os.environ["NVTE_UNFUSED_ATTN"] = "0"
+
+The environment variables are read during backend selection. Set them before the first
+``DotProductAttention.forward()`` call because the selection result is cached.
+
+Feature Compatibility Matrix
+-----------------------------
+
+Not all backends support all features. Key restrictions:
+
+.. list-table::
+   :header-rows: 1
+   :widths: 30 15 15 15 15
+
+   * - Feature
+     - Unfused
+     - Flash
+     - cuDNN F16
+     - cuDNN FP8
+   * - Causal mask
+     - Yes
+     - Yes
+     - Yes
+     - Yes
+   * - Padding mask
+     - Yes
+     - Yes
+     - Yes
+     - Yes
+   * - Sliding window
+     - Yes (via arbitrary mask)
+     - Yes
+     - Yes
+     - Yes
+   * - GQA/MQA
+     - Yes
+     - Yes
+     - Yes
+     - Yes
+   * - Dropout
+     - Yes
+     - Yes
+     - Yes
+     - Yes
+   * - FP8 Q/K/V
+     - No
+     - No
+     - No
+     - Yes
+   * - Context parallel
+     - No
+     - Limited
+     - Yes
+     - Yes
+   * - Head dim > 256
+     - Yes
+     - No
+     - Yes
+     - Limited
+
+Debugging Selection
+-------------------
+
+To see which backend was selected, set:
+
+.. code-block:: bash
+
+   NVTE_DEBUG=1 NVTE_DEBUG_LEVEL=1
+
+``NVTE_DEBUG_LEVEL`` controls verbosity: ``1`` logs the selected backend (INFO level),
+``2`` logs the full selection process including each filter step and why backends were
+disabled (DEBUG level). Output goes through Python's ``logging`` module under the
+``DotProductAttention`` logger name.
+
+See Also
+--------
+
+- :doc:`backends` — Detailed description of each backend
+- :doc:`context_parallel` — How CP affects backend selection