diff --git a/examples/tp-numerics/.gitignore b/examples/tp-numerics/.gitignore
new file mode 100644
index 00000000000..540ffde29e7
--- /dev/null
+++ b/examples/tp-numerics/.gitignore
@@ -0,0 +1,12 @@
+# Result logs (kept local, not in upstream PR)
+results/
+output/
+*.log
+
+# Bridge-based legacy scripts (replaced by submit_qwen3_*_tp_invariant.sh)
+validate_e2e_*.py
+submit_e2e_*.sh
+test_tp_numerics.py
+
+# Local-only notes (never push)
+*.local.md
diff --git a/examples/tp-numerics/README.md b/examples/tp-numerics/README.md
new file mode 100644
index 00000000000..04f36aa9dcd
--- /dev/null
+++ b/examples/tp-numerics/README.md
@@ -0,0 +1,167 @@
+# TP-Invariant Numerics: Bitwise Identical Training Across TP Degrees
+
+Bitwise identical forward, backward, and E2E training for Megatron-Core TransformerBlocks
+regardless of Tensor Parallelism (TP) degree — TP=1, 2, 4, 8 produce the same result.
+
+**Source branches** (view diffs):
+- [Megatron-LM](https://github.com/jinzex/Megatron-LM/tree/jinzex/tp-invariant-numerics) — MCore changes (clip_grads, batch_invariant_kernels, etc.)
+- [TransformerEngine](https://github.com/jinzex/TransformerEngine/tree/jinzex/tp-invariant-numerics) — TE changes (layernorm_linear, linear, etc.)
+
+## Status
+
+| Model | Unit Test (single fwd+bwd) | E2E Training |
+|-------|-----------|--------------|
+| **Dense** | TP=1/2/4/8 bitwise identical | TP=1/2/4 bitwise identical loss+grad_norm, **100 iters** |
+| **MoE** | TP=1/2/4/8 bitwise identical (BIK) | Pending |
+
+Config: `BIK=1 NVTE_TP_INVARIANT_MODE=1` | TE 2.9 | Qwen3-0.6B / 8B | H100, B300
+
+**Cross-architecture:** B300 ≡ H100 bitwise (Qwen3-0.6B TP=1, Qwen3-8B TP=4, 100 iters). See [Cross-Architecture Validation](#cross-architecture-validation-b300-vs-h100).
+
+## Baseline vs TP-Invariant (Qwen3-0.6B)
+
+Without TP-invariant mode, loss diverges from iteration 1:
+
+| Iter | Baseline TP=1 | Baseline TP=2 (diverges) | Baseline TP=4 (diverges) | TP-Inv TP=1/2/4 (all identical) |
+|------|--------------|--------------------------|--------------------------|-------------------------------|
+| 1 | 1.213320E+01 | 1.213298E+01 | 1.213321E+01 | **1.213320E+01** |
+| 2 | 1.187939E+01 | 1.188345E+01 | 1.188395E+01 | **1.187872E+01** |
+| 5 | 1.252811E+01 | 1.253201E+01 | 1.253240E+01 | **1.257282E+01** |
+| 10 | 1.165137E+01 | 1.164608E+01 | 1.164373E+01 | **1.143649E+01** |
+| 100 | 8.277067E+00 | - | - | **8.277067E+00** |
+
+## Components & Patches
+
+All fixes are required together. TE patches must be copied into the container at runtime
+(TE is a site-package). MCore changes are committed directly to this branch.
+
+TE-side patches live in the companion TE branch
+[`jinzex/tp-invariant-numerics`](https://github.com/jinzex/TransformerEngine/tree/jinzex/tp-invariant-numerics);
+MCore changes are committed in this branch.
+
+| Component | Fix | Location |
+|-----------|-----|----------|
+| **TP-Invariant GEMM** (fwd+bwd) | All-gather sharded weight, full-K GEMM | TE `transformer_engine/pytorch/module/{layernorm_linear,linear}.py` |
+| **Gated deinterleave** (bwd) | Reorder via `partition_stride` after all-gather | TE `transformer_engine/pytorch/module/layernorm_linear.py` |
+| **Cross-entropy** (fwd) | All-gather exp_logits, local sum | `megatron/core/tensor_parallel/cross_entropy.py` |
+| **Output projection** (bwd) | All-gather weight+grad, full dgrad GEMM | `megatron/core/tensor_parallel/layers.py` |
+| **Gradient clipping** | Float64 norm + pow2 clip_coeff rounding | `megatron/core/optimizer/clip_grads.py` |
+| **RMSNorm dgamma** (bwd) | All-gather tokens + rank-0-only reduction | `megatron/core/transformer/custom_layers/batch_invariant_kernels.py` |
+| **BIK** (fwd+bwd) | M-invariant Triton matmul_persistent | `megatron/core/transformer/custom_layers/batch_invariant_kernels.py` |
+
+### Gradient Clipping Fix
+
+`multi_tensor_l2norm` (Apex CUDA) reduces over different-shaped TP shards → different
+FP32 partial sums → different clip_factor. Fix: float64 norm computation + pow2 rounding
+of clip_coeff. Mismatch rate: 40.3% (float32) → 0.000% (pow2).
+
+## Quick Start
+
+**Prerequisites:** TE 2.9, 1+ GPU (1 for 0.6B-TP1, 4 for 8B-TP4, 8 for MoE-toy), this Megatron-LM branch on PYTHONPATH. **No Bridge dependency** — dense scripts use raw MLM `pretrain_gpt.py` directly.
+
+### TE patches
+
+Submit scripts install the matching TE patches into the container at launch
+(squashfs containers are read-only per-srun, so patches reapply each run).
+Version-pinned snapshots live under `patches/v<TE_VERSION>/` mirroring the TE
+source tree — see [`patches/README.md`](patches/README.md) for the supported
+versions and how to refresh from the companion
+[TE branch](https://github.com/jinzex/TransformerEngine/tree/jinzex/tp-invariant-numerics).
+No matching version → submit script fails fast.
+
+### Run tests
+
+```bash
+cd $PROJ  # this Megatron-LM branch
+
+# Unit test: TP=1 ≡ TP=2 ≡ TP=4 bitwise on a small TransformerBlock
+NVTE_TP_INVARIANT_MODE=1 \
+torchrun --nproc_per_node=4 -m pytest \
+    tests/unit_tests/transformer/test_tp_invariant.py -v -s
+
+# E2E: Qwen3-0.6B (100 iters, TP=1; 1 GPU)
+bash examples/tp-numerics/submit_qwen3_0.6b_tp_invariant.sh
+
+# E2E: Qwen3-8B (100 iters, TP=4; 4 GPUs)
+bash examples/tp-numerics/submit_qwen3_8b_tp_invariant.sh
+
+# Smoke test: toy MoE (Qwen3-30B-A3B 4L 8E top-2, 10 iters, 1 GPU)
+bash examples/tp-numerics/submit_qwen3_moe_toy_tp_invariant.sh
+
+# Override TP / iters via env:
+# TP_SIZE=2 TRAIN_ITERS=10 bash submit_qwen3_0.6b_tp_invariant.sh
+# Each script also runs as `sbatch <script>` for SLURM batch.
+```
+
+**Expected iter 1** (Qwen3-0.6B TP=1): `lm loss: 1.213320E+01 | grad norm: 18.259`.
+
+## Validation Results
+
+| Model | Backend | TP | Iters | Result |
+|-------|---------|-----|-------|--------|
+| Qwen3-0.6B | unfused | 1/2/4 | 100 | Bitwise identical |
+| Qwen3-0.6B | auto (cuDNN) | 1/2 | 100 | Bitwise identical |
+| Qwen3-0.6B | flash (FA3) | 1/2 | 100 | Bitwise identical |
+| Qwen3-8B | unfused | 4/8 | 10 | Bitwise identical |
+| Qwen3-8B | auto (cuDNN) | 4/8 | 100 | Bitwise identical |
+| Qwen3-8B | — | 1/2 | — | OOM (all-gather doubles peak memory) |
+
+All three attention backends (unfused, auto/cuDNN, flash/FA3) achieve TP-invariance —
+attention is intra-rank with no cross-rank reduction. Recommended: **unfused** (no extra setup).
+
+### Qwen3-8B: TP=4 vs TP=8 (bitwise identical, 100 iters)
+
+![TP=4 vs TP=8 loss curves](assets/tp4_vs_tp8_invariant.png)
+
+### Qwen3-8B: TP-invariant vs baseline (TP=4)
+
+![TP-invariant vs baseline](assets/invariant_vs_baseline_tp4.png)
+
+## Cross-Architecture Validation (B300 vs H100)
+
+Bitwise identical training across NVIDIA SM_90 (H100) and SM_100 (B300) with
+the same patch stack. Loss + grad_norm + validation match every iter.
+
+| Model      | TP | Iters | Result |
+|------------|----|-------|--------|
+| Qwen3-0.6B | 1  | 100   | Bitwise identical |
+| Qwen3-8B   | 4  | 100   | Bitwise identical |
+
+![B300 vs H100 (Qwen3-0.6B, TP=1)](assets/b300_vs_h100_tp1_invariant.png)
+
+Reproduce by running the Quick Start commands on each GPU and diffing
+the per-iter lm loss + grad norm streams:
+
+```bash
+diff <(grep -oP "lm loss: \S+|grad norm: \S+" b300_run.log) \
+     <(grep -oP "lm loss: \S+|grad norm: \S+" h100_run.log)
+# Empty output = bitwise identical.
+```
+
+## Performance Overhead (Qwen3-8B, 4x H100, TP=4, auto/cuDNN)
+
+| Mode | Avg step (ms) | Overhead | What changes |
+|------|-------------|----------|-------------|
+| Baseline | 227.5 | — | Standard training |
+| **TP-invariant fwd** | 386.7 | **+70%** | 2/4 fwd GEMMs all-gather + cross-entropy |
+| TP-invariant E2E | 598.2 | +163% | All 7 components (fwd+bwd) |
+
+Fwd-only overhead (+70%) is the relevant cost for RL training where only forward logits need TP-invariance. Full E2E (+163%) is for debugging/validation only.
+
+## MoE E2E
+
+Unit tests: **bitwise identical** fwd+bwd across TP=1/2/4/8 (with BIK=1).
+E2E training: **pending**.
+
+Forward is bitwise identical (iter 1 loss matches across TP). Backward diverges
+because expert wgrad accumulates over different token subsets per TP rank (partial-K).
+At DP=1: expert wgrads are unsynchronized. At DP>1: synced via all-reduce, but
+sum(partial wgrads) ≠ full wgrad due to FP32 accumulation order. Same class as
+the Dense partial-K issue, but expert weights are replicated (not TP-sharded),
+so the Dense all-gather-weight fix doesn't apply directly.
+
+## Known Limitations
+
+- **PP**: Megatron seeds each PP stage differently → can't compare across PP configs. All validation uses PP=1.
+- **DP>1**: Different DP means different gradient accumulation order. DP-invariance is a separate problem.
+- **8B OOM at TP≤2**: TP-invariant all-gather doubles peak weight memory. Validation-mode tradeoff.
diff --git a/examples/tp-numerics/assets/b300_vs_h100_tp1_invariant.png b/examples/tp-numerics/assets/b300_vs_h100_tp1_invariant.png
new file mode 100644
index 00000000000..706bd908f4d
Binary files /dev/null and b/examples/tp-numerics/assets/b300_vs_h100_tp1_invariant.png differ
diff --git a/examples/tp-numerics/assets/invariant_vs_baseline_tp4.png b/examples/tp-numerics/assets/invariant_vs_baseline_tp4.png
new file mode 100644
index 00000000000..57945647072
Binary files /dev/null and b/examples/tp-numerics/assets/invariant_vs_baseline_tp4.png differ
diff --git a/examples/tp-numerics/assets/tp4_vs_tp8_invariant.png b/examples/tp-numerics/assets/tp4_vs_tp8_invariant.png
new file mode 100644
index 00000000000..123a1e41527
Binary files /dev/null and b/examples/tp-numerics/assets/tp4_vs_tp8_invariant.png differ
diff --git a/examples/tp-numerics/diff_dumps.py b/examples/tp-numerics/diff_dumps.py
new file mode 100644
index 00000000000..177f4159973
--- /dev/null
+++ b/examples/tp-numerics/diff_dumps.py
@@ -0,0 +1,45 @@
+"""Diff cross-GPU tensor dumps to localize where B300 vs H100 first disagree.
+
+Usage:
+    python diff_dumps.py dump_b300.pt dump_h100.pt
+"""
+import sys
+import torch
+
+
+def main():
+    if len(sys.argv) != 3:
+        print("usage: python diff_dumps.py <a.pt> <b.pt>", file=sys.stderr)
+        sys.exit(1)
+
+    a = torch.load(sys.argv[1], map_location="cpu", weights_only=False)
+    b = torch.load(sys.argv[2], map_location="cpu", weights_only=False)
+
+    common = sorted(set(a.keys()) & set(b.keys()))
+    only_a = sorted(set(a.keys()) - set(b.keys()))
+    only_b = sorted(set(b.keys()) - set(a.keys()))
+
+    if only_a:
+        print(f"only in {sys.argv[1]}: {only_a}")
+    if only_b:
+        print(f"only in {sys.argv[2]}: {only_b}")
+
+    print(f"\n{'key':<40s} {'shape':<25s} {'bitwise':>10s} {'max_abs':>14s} {'mean_abs':>14s} {'nonzero%':>10s}")
+    print("-" * 115)
+    for k in common:
+        ta, tb = a[k], b[k]
+        if not isinstance(ta, torch.Tensor) or not isinstance(tb, torch.Tensor):
+            continue
+        if ta.shape != tb.shape:
+            print(f"{k:<40s} shape mismatch: {tuple(ta.shape)} vs {tuple(tb.shape)}")
+            continue
+        # Exact bitwise check: torch.equal compares dtype + every element exactly.
+        bitwise = "✅" if torch.equal(ta, tb) else "❌"
+        d = (ta.float() - tb.float()).abs()
+        print(f"{k:<40s} {str(tuple(ta.shape)):<25s} {bitwise:>10s} "
+              f"{d.max().item():>14.6e} {d.mean().item():>14.6e} "
+              f"{100.0 * (d > 0).float().mean().item():>9.4f}%")
+
+
+if __name__ == "__main__":
+    main()
diff --git a/examples/tp-numerics/patches/README.md b/examples/tp-numerics/patches/README.md
new file mode 100644
index 00000000000..ebe59aed521
--- /dev/null
+++ b/examples/tp-numerics/patches/README.md
@@ -0,0 +1,35 @@
+# TE patches
+
+Snapshots of TE files the submit scripts copy into the container's TE
+site-packages before training. Squashfs containers are read-only per-srun, so
+patches must re-apply on every launch.
+
+Submit scripts detect the container's TE version and pick `v<version>/`. No
+matching subdirectory → fail-fast.
+
+## Layout
+
+```
+v<version>/transformer_engine/pytorch/  # mirrors TE source tree
+  module/{layernorm_linear,linear}.py
+  attention/dot_product_attention/{backends,dot_product_attention}.py
+```
+
+## Supported versions
+
+| TE version | Base commit | Source |
+|---|---|---|
+| `v2.9.0` | `70f53666` | [jinzex/TransformerEngine@d1bcd9a6](https://github.com/jinzex/TransformerEngine/tree/jinzex/tp-invariant-numerics) |
+
+Gated on `NVTE_TP_INVARIANT_MODE=1` — env-var unset = stock TE behavior.
+
+## Refresh
+
+```bash
+TE=<TransformerEngine checkout>; VER=2.9.0
+DEST=v${VER}/transformer_engine/pytorch
+mkdir -p $DEST/module $DEST/attention/dot_product_attention
+cp $TE/transformer_engine/pytorch/module/{layernorm_linear,linear}.py $DEST/module/
+cp $TE/transformer_engine/pytorch/attention/dot_product_attention/{backends,dot_product_attention}.py \
+   $DEST/attention/dot_product_attention/
+```
diff --git a/examples/tp-numerics/patches/v2.9.0/transformer_engine/pytorch/attention/dot_product_attention/backends.py b/examples/tp-numerics/patches/v2.9.0/transformer_engine/pytorch/attention/dot_product_attention/backends.py
new file mode 100644
index 00000000000..d9444ba05f6
--- /dev/null
+++ b/examples/tp-numerics/patches/v2.9.0/transformer_engine/pytorch/attention/dot_product_attention/backends.py
@@ -0,0 +1,1948 @@
+# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+#
+# See LICENSE for license information.
+
+"""Attention Backends."""
+from contextlib import nullcontext
+from importlib.metadata import version as get_pkg_version
+from importlib.metadata import PackageNotFoundError
+import os
+from typing import Any, Callable, Dict, List, Optional, Tuple, Union
+import warnings
+import logging
+from packaging.version import Version as PkgVersion
+
+import torch
+import torch.nn.functional as F
+import transformer_engine_torch as tex
+from transformer_engine.pytorch.utils import (
+    get_device_compute_capability,
+    split_tensor_along_dim,
+)
+from transformer_engine.pytorch.utils import attention_mask_func, nvtx_range_push, nvtx_range_pop
+from transformer_engine.pytorch.tensor.float8_tensor import (
+    Float8Quantizer,
+    Float8CurrentScalingQuantizer,
+)
+from transformer_engine.pytorch.tensor.quantized_tensor import (
+    QuantizedTensorStorage,
+    prepare_for_saving,
+    restore_from_saved,
+)
+from transformer_engine.pytorch.tensor.float8_tensor import Float8Tensor
+from transformer_engine.pytorch.constants import (
+    TE_DType,
+    QKVLayouts,
+    dist_group_type,
+)
+from transformer_engine.pytorch.cpp_extensions.fused_attn import (
+    fused_attn_fwd,
+    fused_attn_bwd,
+    FusedAttnBackend,
+    META_O,
+    META_QKV,
+)
+from transformer_engine.pytorch.quantization import get_fp8_torch_dtype, FP8GlobalStateManager
+from transformer_engine.pytorch.distributed import get_distributed_world_size
+from transformer_engine.pytorch.jit import no_torch_dynamo
+from transformer_engine.pytorch.attention.dot_product_attention.context_parallel import (
+    attn_forward_func_with_cp,
+)
+from transformer_engine.pytorch.attention.dot_product_attention.softmax import FusedScaleMaskSoftmax
+from transformer_engine.pytorch.attention.inference import InferenceParams
+
+# Import attention utils
+import transformer_engine.pytorch.attention.dot_product_attention.utils as dpa_utils
+from transformer_engine.pytorch.attention.dot_product_attention.utils import (
+    FlashAttentionUtils as fa_utils,
+    combine_and_quantize,
+    combine_and_dequantize,
+    print_quantizers,
+    ConvertTHDtoBSHD,
+    ConvertBSHDtoTHD,
+)
+from transformer_engine.pytorch.attention.dot_product_attention.utils import (
+    AttentionLogging as attn_log,
+)
+from transformer_engine.pytorch import export
+from transformer_engine.pytorch.export import is_in_onnx_export_mode
+
+# Global vars for flash attn v2 and v3 imports
+flash_attn_cuda_bwd = None
+flash_attn_func = None
+flash_attn_varlen_func = None
+_flash_attn_fwd = None
+_flash_attn_bwd = None
+_flash_attn_varlen_fwd = None
+_flash_attn_varlen_bwd = None
+try:
+    fa_utils.version = PkgVersion(get_pkg_version("flash-attn"))
+except PackageNotFoundError:
+    pass  # only print warning if use_flash_attention_2 = True in get_attention_backend
+else:
+    if torch.cuda.is_available() and get_device_compute_capability() >= (10, 0):
+        if fa_utils.version_required_blackwell <= fa_utils.version <= fa_utils.max_version:
+            fa_utils.is_installed = True
+    elif fa_utils.version_required <= fa_utils.version <= fa_utils.max_version:
+        fa_utils.is_installed = True
+
+    if fa_utils.is_installed:
+        from flash_attn_2_cuda import varlen_bwd as flash_attn_cuda_bwd
+        from flash_attn.flash_attn_interface import flash_attn_func, flash_attn_varlen_func
+        from flash_attn.flash_attn_interface import _flash_attn_forward as _flash_attn_fwd
+        from flash_attn.flash_attn_interface import _flash_attn_backward as _flash_attn_bwd
+        from flash_attn.flash_attn_interface import (
+            _flash_attn_varlen_forward as _flash_attn_varlen_fwd,
+        )
+        from flash_attn.flash_attn_interface import (
+            _flash_attn_varlen_backward as _flash_attn_varlen_bwd,
+        )
+
+        # Setup Flash attention utils
+        fa_utils.set_flash_attention_version()
+    elif (
+        torch.cuda.is_available()
+        and get_device_compute_capability() >= (8, 0)
+        and dpa_utils._NVTE_FLASH_ATTN
+    ):
+        attn_log.fa_logger.warning(
+            "Supported flash-attn versions are %s. Found flash-attn %s.",
+            dpa_utils._get_supported_versions(
+                (
+                    fa_utils.version_required
+                    if get_device_compute_capability() < (10, 0)
+                    else fa_utils.version_required_blackwell
+                ),
+                fa_utils.max_version,
+            ),
+            fa_utils.version,
+        )
+try:
+    fa_utils.fa3_version = PkgVersion(get_pkg_version("flash-attn-3"))
+except PackageNotFoundError:
+    flash_attn_func_v3 = None
+    flash_attn_varlen_func_v3 = None
+    flash_attn_with_kvcache_v3 = None
+    # pass  # only print warning if use_flash_attention_3 = True in get_attention_backend
+else:
+    from flash_attn_3.flash_attn_interface import flash_attn_func as flash_attn_func_v3
+    from flash_attn_3.flash_attn_interface import (
+        flash_attn_varlen_func as flash_attn_varlen_func_v3,
+    )
+    from flash_attn_3.flash_attn_interface import (
+        flash_attn_with_kvcache as flash_attn_with_kvcache_v3,
+    )
+    from flash_attn_3.flash_attn_interface import _flash_attn_forward as _flash_attn_fwd_v3
+    from flash_attn_3.flash_attn_interface import _flash_attn_backward as _flash_attn_bwd_v3
+
+    fa_utils.set_flash_attention_3_params()
+
+# Float8CurrentScaling: fused_attn_bwd takes O in FP8 by default, this flag allows it in F16
+_dpa_fp8_cs_o_in_f16 = os.getenv("NVTE_DPA_FP8CS_O_in_F16", "1") == "1"
+
+
+class FP8EmulationFunc(torch.autograd.Function):
+    """
+    Emulate the effects of FP8 quantization on tensors. Used in UnfusedDotProductAttention as follows:
+    - forward : QKV (quantize+dequantize),  P (pass-through),  S (quantize+dequantize),    O (pass-through)
+    - backward:  dO (quantize+dequantize), dS (pass-through), dP (quantize+dequantize), dQKV (pass-through)
+    """
+
+    @staticmethod
+    def forward(ctx, tensor1, tensor2, tensor3, quantizer, quantizer_name, qkv_layout):
+        # pylint: disable=missing-function-docstring
+        if quantizer_name == "QKV_quantizer":
+            query_layer, key_layer, value_layer = [
+                x.contiguous() for x in [tensor1, tensor2, tensor3]
+            ]
+            q_fp8, k_fp8, v_fp8 = combine_and_quantize(
+                qkv_layout, query_layer, key_layer, value_layer, quantizer
+            )
+            tensors = combine_and_dequantize(
+                qkv_layout, q_fp8, k_fp8, v_fp8, src_nominal_dtype=query_layer.dtype
+            )
+        elif quantizer_name in ["S_quantizer", "O_quantizer"]:
+            t_fp8 = quantizer(tensor1)
+            tensors = (t_fp8.dequantize(dtype=tensor1.dtype), tensor2, tensor3)
+        else:
+            tensors = (tensor1, tensor2, tensor3)
+        ctx.quantizer = quantizer
+        ctx.quantizer_name = quantizer_name
+        ctx.qkv_layout = qkv_layout
+        return tensors[0], tensors[1], tensors[2]
+
+    @staticmethod
+    def backward(ctx, grad1, grad2, grad3):
+        # pylint: disable=missing-function-docstring
+        if ctx.quantizer_name in ["dO_quantizer", "dP_quantizer"]:
+            dt_fp8 = ctx.quantizer(grad1)
+            tensors = dt_fp8.dequantize(dtype=grad1.dtype), grad2, grad3
+        elif ctx.quantizer_name == "dQKV_quantizer":
+            query_grad, key_grad, value_grad = [x.contiguous() for x in [grad1, grad2, grad3]]
+            dq_fp8, dk_fp8, dv_fp8 = combine_and_quantize(
+                ctx.qkv_layout, query_grad, key_grad, value_grad, ctx.quantizer
+            )
+            tensors = combine_and_dequantize(
+                ctx.qkv_layout, dq_fp8, dk_fp8, dv_fp8, src_nominal_dtype=query_grad.dtype
+            )
+        else:
+            tensors = grad1, grad2, grad3
+        return tensors[0], tensors[1], tensors[2], None, None, None
+
+
+class UnfusedDotProductAttention(torch.nn.Module):
+    """Parallel attention w/o QKV and Proj Gemms
+    BMM1 -> softmax + dropout -> BMM2
+    """
+
+    def __init__(
+        self,
+        softmax_scale: float,
+        attention_type: str = "self",
+        attention_dropout: float = 0.0,
+        attention_dropout_ctx: Optional[Callable] = nullcontext,
+        layer_number: Optional[int] = None,
+        softmax_type: str = "vanilla",
+        return_max_logit: Optional[bool] = False,
+    ) -> None:
+        super().__init__()
+
+        self.softmax_scale = softmax_scale
+        self.attention_type = attention_type
+        self.attention_dropout_ctx = attention_dropout_ctx
+        self.layer_number = layer_number
+        self.softmax_type = softmax_type
+        self.return_max_logit = return_max_logit
+
+        def mask_func(x, y):
+            return (
+                export.onnx_attention_mask_func(x, y)
+                if is_in_onnx_export_mode()
+                else attention_mask_func(x, y)
+            )
+
+        self.mask_func = mask_func
+        self.scale_mask_softmax = FusedScaleMaskSoftmax(mask_func)
+
+        # Dropout. Note that for a single iteration, this layer will generate
+        # different outputs on different number of parallel partitions but
+        # on average it should not be partition dependent.
+        self.attention_dropout = torch.nn.Dropout(attention_dropout)
+
+        # An FP16 training trick required for certain GPT-like models.
+        self.apply_qk_layer_scaling = (
+            bool(int(os.getenv("NVTE_APPLY_QK_LAYER_SCALING", "0"))) and layer_number is not None
+        )
+
+    def forward(
+        self,
+        _alibi_cache: Dict[str, Any],
+        query_layer: torch.Tensor,
+        key_layer: torch.Tensor,
+        value_layer: torch.Tensor,
+        qkv_layout: str = "sbh3d",
+        cu_seqlens_q: Optional[torch.Tensor] = None,  # pylint: disable=unused-argument
+        cu_seqlens_kv: Optional[torch.Tensor] = None,  # pylint: disable=unused-argument
+        max_seqlen_q: Optional[torch.Tensor] = None,  # pylint: disable=unused-argument
+        max_seqlen_kv: Optional[torch.Tensor] = None,  # pylint: disable=unused-argument
+        attn_mask_type: str = "causal",
+        attention_mask: Optional[Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]] = None,
+        window_size: Optional[Tuple[int, int]] = None,
+        core_attention_bias_type: str = "no_bias",
+        core_attention_bias: Optional[torch.Tensor] = None,
+        alibi_slopes: Optional[torch.Tensor] = None,
+        inference_params: Optional[InferenceParams] = None,
+        softmax_offset: torch.Tensor = None,
+        fp8: bool = False,
+        fp8_meta: Optional[Dict[str, Any]] = None,
+        quantizers=None,
+        fp8_output: bool = False,
+    ) -> torch.Tensor:
+        """Unfused attention fprop"""
+        assert (
+            qkv_layout in QKVLayouts
+        ), f"UnfusedDotProductAttention does not support qkv_layout = {qkv_layout}!"
+
+        # get q_format and kv_format for training and inference
+        qkv_format, q_format, _ = dpa_utils.get_qkv_format(qkv_layout, inference_params)
+        if inference_params is not None and inference_params.is_paged:
+            key_layer, value_layer = inference_params.convert_paged_to_nonpaged(self.layer_number)
+
+        # convert to sbhd
+        # training: bshd, thd
+        # inference: bshd, sbhd_2bshd, thd_2bshd
+        if qkv_format == "bshd":
+            # convert to sbhd and use sbhd implementation for now
+            query_layer, key_layer, value_layer = [
+                x.transpose(0, 1) for x in [query_layer, key_layer, value_layer]
+            ]
+        if qkv_format == "sbhd_2bshd":
+            key_layer, value_layer = [x.transpose(0, 1) for x in [key_layer, value_layer]]
+
+        if qkv_format == "thd_2bshd":
+            batch_size = key_layer.shape[0]
+            query_layer = tex.convert_thd_to_bshd(
+                query_layer,
+                cu_seqlens_q,
+                batch_size,
+                inference_params.max_ctx_len,
+            )
+            query_layer, key_layer, value_layer = [
+                x.transpose(0, 1) for x in [query_layer, key_layer, value_layer]
+            ]
+        if qkv_format == "thd":
+            assert cu_seqlens_q is not None and cu_seqlens_kv is not None
+            assert max_seqlen_q is not None and max_seqlen_kv is not None
+            query_layer = ConvertTHDtoBSHD.apply(
+                query_layer,
+                cu_seqlens_q,
+                max_seqlen_q,
+            )
+            key_layer, value_layer = [
+                ConvertTHDtoBSHD.apply(
+                    x,
+                    cu_seqlens_kv,
+                    max_seqlen_kv,
+                )
+                for x in [key_layer, value_layer]
+            ]
+            query_layer, key_layer, value_layer = [
+                x.transpose(0, 1).contiguous() for x in [query_layer, key_layer, value_layer]
+            ]
+
+        batch_size, max_seqlen_q, max_seqlen_kv = (
+            query_layer.shape[1],
+            query_layer.shape[0],
+            key_layer.shape[0],
+        )
+
+        if "padding" in attn_mask_type and attention_mask is None:
+            attention_mask = dpa_utils.get_padding_mask(
+                batch_size,
+                cu_seqlens_q,
+                cu_seqlens_kv,
+                max_seqlen_q,
+                max_seqlen_kv,
+                self.attention_type,
+            )
+        attn_mask_type, attention_mask, actual_seqlens_q, actual_seqlens_kv = (
+            dpa_utils.get_full_mask(
+                max_seqlen_q,
+                max_seqlen_kv,
+                attn_mask_type=attn_mask_type,
+                attention_mask=attention_mask,
+                window_size=window_size,
+                attention_type=self.attention_type,
+            )
+        )
+
+        batch_size, seqlen = query_layer.shape[1], query_layer.shape[0]
+        apply_qk_layer_scaling = self.apply_qk_layer_scaling and key_layer.dtype == torch.float16
+
+        # [b, np, sq, sk]
+        output_size = (
+            query_layer.size(1),
+            query_layer.size(2),
+            query_layer.size(0),
+            key_layer.size(0),
+        )
+
+        if key_layer.shape[2] != query_layer.shape[2]:
+            assert (
+                query_layer.shape[2] % key_layer.shape[2] == 0
+            ), "The number of attention heads must be divisible by the number of GQA groups!"
+            key_layer = key_layer.repeat_interleave(
+                int(query_layer.shape[2] / key_layer.shape[2]), dim=2
+            )
+            value_layer = value_layer.repeat_interleave(
+                int(query_layer.shape[2] / value_layer.shape[2]), dim=2
+            )
+
+        # [sq, b, np, hn] -> [sq, b * np, hn]
+        query_layer = query_layer.reshape(output_size[2], output_size[0] * output_size[1], -1)
+        # [sk, b, np, hn] -> [sk, b * np, hn]
+        key_layer = key_layer.reshape(output_size[3], output_size[0] * output_size[1], -1)
+
+        # preallocting result tensor: [b * np, sq, sk]
+        matmul_result = torch.empty(
+            output_size[0] * output_size[1],
+            output_size[2],
+            output_size[3],
+            dtype=query_layer.dtype,
+            device=torch.cuda.current_device(),
+        )
+
+        scale = self.softmax_scale
+        if apply_qk_layer_scaling:
+            scale /= self.layer_number
+
+        if fp8:
+            # get quantizers from DPA; all Nones if not fp8
+            QKV_quantizer, O_quantizer, S_quantizer, dQKV_quantizer, dO_quantizer, dP_quantizer = (
+                dpa_utils.get_attention_quantizers(fp8, quantizers)
+            )
+            # S/dP are forced to use DS quantizers in DPA.init_fp8_metadata; revert them here for true CS emulation
+            fp8_recipe = FP8GlobalStateManager.get_fp8_recipe()
+            if fp8_meta is not None and fp8_meta.get("local_recipes", None) is not None:
+                fp8_recipe = fp8_meta["local_recipes"][0]
+            if fp8_recipe.float8_current_scaling():
+                S_quantizer = Float8CurrentScalingQuantizer(
+                    fp8_dtype=S_quantizer.dtype, device="cuda"
+                )
+                dP_quantizer = Float8CurrentScalingQuantizer(
+                    fp8_dtype=dP_quantizer.dtype, device="cuda"
+                )
+
+            if "2" in qkv_layout or "3" in qkv_layout:
+                qkv_format, *_ = dpa_utils.get_qkv_format(qkv_layout)
+                qkv_layout = "_".join([qkv_format] * 3)
+            # quantize and dequantize QKV to emulate FP8
+            query_layer, key_layer, value_layer = FP8EmulationFunc.apply(
+                query_layer, key_layer, value_layer, QKV_quantizer, "QKV_quantizer", qkv_layout
+            )
+            # quantize and dequantize dQKV to emulate FP8
+            query_layer, key_layer, value_layer = FP8EmulationFunc.apply(
+                query_layer, key_layer, value_layer, dQKV_quantizer, "dQKV_quantizer", qkv_layout
+            )
+
+        # Raw attention scores. [b * np, sq, sk]
+        if core_attention_bias_type == "no_bias":
+            matmul_result = torch.baddbmm(
+                matmul_result,
+                query_layer.transpose(0, 1),  # [b * np, sq, hn]
+                key_layer.transpose(0, 1).transpose(1, 2),  # [b * np, hn, sk]
+                beta=0.0,
+                alpha=scale,
+            ).view(*output_size)
+
+        elif core_attention_bias_type == "pre_scale_bias":
+            assert core_attention_bias is not None, "core_attention_bias should not be None!"
+            matmul_result = torch.bmm(
+                query_layer.transpose(0, 1),  # [b * np, sq, hn]
+                key_layer.transpose(0, 1).transpose(1, 2),  # [b * np, hn, sk]
+            )
+            matmul_result = matmul_result.view(*output_size) + core_attention_bias
+            matmul_result *= scale
+
+        elif core_attention_bias_type in ["post_scale_bias", "alibi"]:
+            if core_attention_bias_type == "post_scale_bias":
+                assert core_attention_bias is not None, "core_attention_bias should not be None!"
+            if core_attention_bias_type == "alibi":
+                _, core_attention_bias = dpa_utils.get_alibi(
+                    _alibi_cache,
+                    output_size[1],
+                    output_size[2],
+                    output_size[3],
+                    actual_seqlens_q=actual_seqlens_q if "padding" in attn_mask_type else None,
+                    actual_seqlens_kv=actual_seqlens_kv if "padding" in attn_mask_type else None,
+                    alibi_slopes=alibi_slopes,
+                    bottom_right_alignment=attn_mask_type not in ["causal", "padding_causal"],
+                )
+            matmul_result = torch.baddbmm(
+                matmul_result,
+                query_layer.transpose(0, 1),  # [b * np, sq, hn]
+                key_layer.transpose(0, 1).transpose(1, 2),  # [b * np, hn, sk]
+                beta=0.0,
+                alpha=scale,
+            )
+            matmul_result = (matmul_result.view(*output_size) + core_attention_bias).to(
+                dtype=query_layer.dtype
+            )
+
+        if fp8:
+            # quantize and dequantize dP to emulate FP8
+            matmul_result, *_ = FP8EmulationFunc.apply(
+                matmul_result, None, None, dP_quantizer, "dP_quantizer", None
+            )
+
+        # max attention score
+        max_logit = None
+        if self.return_max_logit:
+            # matmul_result [b, np, sq, dk], max_logit [np]
+            max_logit = matmul_result
+            if attn_mask_type != "no_mask":
+                max_logit = self.mask_func(matmul_result, attention_mask)
+            max_logit = torch.amax(max_logit, dim=(0, 2, 3))
+
+        # add attention sink to the last column: [b, np, sq, sk+1]
+        if self.softmax_type != "vanilla":
+            matmul_result = torch.cat(
+                [
+                    matmul_result,
+                    softmax_offset.to(dtype=matmul_result.dtype).expand(
+                        matmul_result.size(0), -1, matmul_result.size(2), -1
+                    ),
+                ],
+                dim=-1,
+            )
+            attention_mask = F.pad(attention_mask, (0, 1), mode="constant", value=False)
+            attn_mask_type = "arbitrary"
+
+        # attention scores and attention mask
+        softmax_scale = self.layer_number if apply_qk_layer_scaling else None
+        attention_probs = self.scale_mask_softmax(
+            matmul_result, attention_mask, attn_mask_type, softmax_scale
+        )
+
+        # mask out the pad positions in softmax results, mostly for the rows (pad tokens from q)
+        # the columns (pad tokens from k) are already zeroed out during softmax
+        if "padding" in attn_mask_type:
+            attention_probs = attention_probs.masked_fill(attention_mask, 0)
+
+        # remove attention sink: [b, np, sq, sk]
+        if self.softmax_type != "vanilla":
+            attention_probs = attention_probs[..., :-1]
+
+        # This is actually dropping out entire tokens to attend to, which might
+        # seem a bit unusual, but is taken from the original Transformer paper.
+        with self.attention_dropout_ctx():
+            attention_probs = self.attention_dropout(attention_probs)
+
+        # value_layer -> context layer.
+        # [sk, b, np, hn] --> [b, np, sq, hn]
+        output_size = (
+            value_layer.size(1),
+            value_layer.size(2),
+            query_layer.size(0),
+            value_layer.size(3),
+        )
+
+        # change view [sk, b * np, hn]
+        value_layer = value_layer.reshape(value_layer.size(0), output_size[0] * output_size[1], -1)
+
+        # change view [b * np, sq, sk]
+        attention_probs = attention_probs.view(output_size[0] * output_size[1], output_size[2], -1)
+
+        if fp8:
+            # quantize and dequantize S to emulate FP8
+            attention_probs, *_ = FP8EmulationFunc.apply(
+                attention_probs, None, None, S_quantizer, "S_quantizer", None
+            )
+
+        # matmul: [b * np, sq, hn]
+        context_layer = torch.bmm(attention_probs, value_layer.transpose(0, 1))
+
+        # change view [b, np, sq, hn]
+        context_layer = context_layer.view(*output_size)
+
+        if q_format == "sbhd":
+            # [b, np, sq, hn] --> [sq, b, np, hn]
+            context_layer = context_layer.permute(2, 0, 1, 3).contiguous()
+
+            # [sq, b, np, hn] --> [sq, b, hp]
+            context_layer = context_layer.view(seqlen, batch_size, -1)
+
+        if q_format == "bshd":
+            # [b, np, sq, hn] --> [b, sq, np, hn]
+            context_layer = context_layer.permute(0, 2, 1, 3).contiguous()
+
+            # [b, sq, np, hn] --> [b, sq, hp]
+            context_layer = context_layer.view(batch_size, seqlen, -1)
+
+        if q_format == "thd":
+            # [b, np, sq, hn] --> [b, sq, np, hn]
+            context_layer = context_layer.permute(0, 2, 1, 3).contiguous()
+
+            # [b, sq, np, hn] --> [tq, np, hn]
+            context_layer = ConvertBSHDtoTHD.apply(
+                context_layer,
+                cu_seqlens_q,
+            )
+
+            # [tq, np, hn] --> [tq, hp]
+            context_layer = context_layer.view(context_layer.shape[0], -1)
+
+        if fp8:
+            # quantize and dequantize O to emulate FP8
+            context_layer, *_ = FP8EmulationFunc.apply(
+                context_layer, None, None, O_quantizer, "O_quantizer", None
+            )
+            # quantize and dequantize dO to emulate FP8
+            context_layer, *_ = FP8EmulationFunc.apply(
+                context_layer, None, None, dO_quantizer, "dO_quantizer", None
+            )
+
+            # quantize O
+            if fp8_output:
+                context_layer = O_quantizer(context_layer)
+
+        if self.return_max_logit:
+            return context_layer, max_logit
+
+        return context_layer
+
+
+class _PrepareQKVForFA(torch.autograd.Function):
+    """This class converts QKV from interleaved (s, b, ...) layout
+    to separate contiguous q, k, v tensors in (b, s, ...) layout."""
+
+    @staticmethod
+    def forward(
+        _ctx: torch.autograd.function.FunctionCtx,  # unused
+        query_layer: torch.Tensor,
+        key_layer: torch.Tensor,
+        value_layer: torch.Tensor,
+    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        # pylint: disable=missing-function-docstring
+        # All inputs received are non-contiguous tensors.
+        # The `query_layer` tensor is used to access the
+        # full memory region of the QKV tensor.
+        qkv = tex.fa_prepare_fwd(query_layer)
+        q, k, v = split_tensor_along_dim(qkv, 0, 3)
+        query_layer = torch.squeeze(q, 0)
+        key_layer = torch.squeeze(k, 0)
+        value_layer = torch.squeeze(v, 0)
+        return query_layer, key_layer, value_layer
+
+    @staticmethod
+    def backward(
+        _ctx: torch.autograd.function.FunctionCtx,  # unused
+        dq: torch.Tensor,
+        dk: torch.Tensor,
+        dv: torch.Tensor,
+    ) -> Tuple[Union[torch.Tensor, None], ...]:
+        # pylint: disable=missing-function-docstring
+        dqkv = tex.fa_prepare_bwd(dq, dk, dv)
+        dq, dk, dv = split_tensor_along_dim(dqkv, -1, 3)
+        return dq, dk, dv
+
+
+class FlashAttention(torch.nn.Module):
+    """Dot product attention, using HazyResearch flash-attn package:
+    https://github.com/Dao-AILab/flash-attention
+    """
+
+    def __init__(
+        self,
+        softmax_scale: float,
+        attention_dropout: float = 0.0,
+        attention_dropout_ctx: Optional[Callable] = nullcontext,
+        attention_type: str = "self",
+        layer_number: Optional[int] = None,
+        deterministic: bool = False,
+    ) -> None:
+        super().__init__()
+
+        if fa_utils.is_installed:
+            assert (
+                fa_utils.version >= fa_utils.version_required
+            ), f"FlashAttention minimum version {fa_utils.version_required} is required."
+            assert (
+                fa_utils.version <= fa_utils.max_version
+            ), f"FlashAttention maximum version {fa_utils.max_version} is supported."
+
+        self.softmax_scale = softmax_scale
+        self.attention_dropout_ctx = attention_dropout_ctx
+        self.attention_dropout = attention_dropout
+        self.attention_type = attention_type
+        self.layer_number = 1 if layer_number is None else layer_number
+        self.deterministic = deterministic
+        self.logger = logging.getLogger("FlashAttention")
+        if attn_log._is_logging_setup is False:
+            attn_log.setup_logging()
+        self.logger.setLevel(attn_log._log_level)
+        if not self.logger.hasHandlers():
+            self.logger.addHandler(attn_log._stream_handler)
+
+    def forward(
+        self,
+        query_layer: torch.Tensor,
+        key_layer: torch.Tensor,
+        value_layer: torch.Tensor,
+        attention_mask: Optional[Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]] = None,
+        qkv_layout: str = "sbh3d",
+        cu_seqlens_q: Optional[torch.Tensor] = None,
+        cu_seqlens_kv: Optional[torch.Tensor] = None,
+        max_seqlen_q: Optional[int] = None,
+        max_seqlen_kv: Optional[int] = None,
+        attn_mask_type: str = "causal",
+        window_size: Optional[Tuple[int, int]] = None,
+        alibi_slopes: Optional[torch.Tensor] = None,
+        cp_group: Optional[Union[dist_group_type, List[dist_group_type]]] = None,
+        cp_global_ranks: List[int] = None,
+        cp_stream: torch.cuda.Stream = None,
+        cp_comm_type: str = "p2p",
+        fp8: bool = False,
+        fp8_meta: Optional[Dict[str, Any]] = None,
+        quantizers=None,
+        inference_params: Optional[InferenceParams] = None,
+        flash_attention_backend: Optional[PkgVersion] = PkgVersion("0"),
+        fp8_output: bool = False,
+        num_splits: Optional[int] = 1,
+    ) -> torch.Tensor:
+        """flash-attn fprop"""
+
+        assert all(
+            x.dtype in [torch.float16, torch.bfloat16] or isinstance(x, Float8Tensor)
+            for x in [query_layer, key_layer, value_layer]
+        ), "FlashAttention only supports FP16 and BF16 data types, or Float8Tensors."
+        assert (
+            query_layer.is_cuda and key_layer.is_cuda and value_layer.is_cuda
+        ), "FlashAttention currently only supports CUDA tensors."
+        assert (
+            qkv_layout in QKVLayouts
+        ), f"FlashAttention does not support qkv_layout = {qkv_layout}!"
+
+        cp_size = 1
+        if isinstance(cp_group, dist_group_type):
+            cp_size = get_distributed_world_size(cp_group)
+        elif isinstance(cp_group, list):
+            for group in cp_group:
+                cp_size *= get_distributed_world_size(group)
+        context_parallel = cp_size > 1
+
+        # get q_format and kv_format for training and inference
+        qkv_format, q_format, kv_format = dpa_utils.get_qkv_format(qkv_layout, inference_params)
+
+        # convert q, k, v to bshd if they are in sbhd; qkv_format doesn't change
+        if all(not isinstance(x, Float8Tensor) for x in [query_layer, key_layer, value_layer]):
+            if qkv_format == "sbhd":
+                # For now just 128, will make it more general in the future
+                if (
+                    query_layer.shape[-1] == 128
+                    and query_layer.shape[0] * query_layer.shape[1] >= 512
+                    and qkv_layout == "sbh3d"
+                ):
+                    query_layer, key_layer, value_layer = _PrepareQKVForFA.apply(
+                        query_layer, key_layer, value_layer
+                    )
+                else:
+                    query_layer, key_layer, value_layer = [
+                        x.transpose(0, 1).contiguous()
+                        for x in (query_layer, key_layer, value_layer)
+                    ]
+            elif q_format == "sbhd" and kv_format == "bshd":
+                query_layer = query_layer.transpose(0, 1).contiguous()
+            if context_parallel:
+                query_layer, key_layer, value_layer = [
+                    x.contiguous() for x in (query_layer, key_layer, value_layer)
+                ]
+        else:
+            if qkv_format == "sbhd":
+                query_layer._data, key_layer._data, value_layer._data = [
+                    x.transpose(0, 1).contiguous()
+                    for x in (query_layer._data, key_layer._data, value_layer._data)
+                ]
+                query_layer, key_layer, value_layer = [
+                    Float8Tensor.make_like(x, data=x._data, shape=x._data.shape)
+                    for x in (query_layer, key_layer, value_layer)
+                ]
+            elif q_format == "sbhd" and kv_format == "bshd":
+                query_layer._data = query_layer._data.transpose(0, 1).contiguous()
+                query_layer = Float8Tensor.make_like(
+                    query_layer, data=query_layer._data, shape=query_layer._data.shape
+                )
+            if context_parallel:
+                query_layer._data, key_layer._data, value_layer._data = [
+                    x.contiguous() for x in (query_layer._data, key_layer._data, value_layer._data)
+                ]
+
+        # get batch_size, max_seqlen and cu_seqlens
+        batch_size, context_len = None, None
+        if inference_params is None:
+            if qkv_format in ["sbhd", "bshd"]:
+                batch_size = query_layer.shape[0]
+                max_seqlen_q, max_seqlen_kv = query_layer.shape[1], key_layer.shape[1]
+                max_seqlen_q *= cp_size
+                max_seqlen_kv *= cp_size
+
+                if "padding" in attn_mask_type:
+                    assert (
+                        not context_parallel
+                    ), "Padding mask not supported with context parallelism!"
+
+                    # [b * s, h, d]
+                    query_layer, key_layer, value_layer = [
+                        x.reshape(x.shape[0] * x.shape[1], *x.shape[2:])
+                        for x in [query_layer, key_layer, value_layer]
+                    ]
+
+                    if self.attention_type == "self":
+                        assert (
+                            max_seqlen_q == max_seqlen_kv
+                        ), "Maximum sequence length for Q and KV should be the same."
+                        if cu_seqlens_q is None:
+                            assert (
+                                attention_mask is not None
+                            ), "Please provide attention_mask for padding!"
+                            cu_seqlens_q, indices_q = dpa_utils.get_cu_seqlens_and_indices(
+                                attention_mask
+                            )
+                        else:
+                            indices_q = dpa_utils.get_indices(max_seqlen_q, cu_seqlens_q)
+                        cu_seqlens_kv = cu_seqlens_q
+                        query_layer, key_layer, value_layer = dpa_utils.PackTensors.apply(
+                            indices_q, query_layer, key_layer, value_layer
+                        )
+                    else:
+                        if cu_seqlens_q is None or cu_seqlens_kv is None:
+                            assert (
+                                attention_mask is not None
+                            ), "Please provide attention_mask for padding!"
+                            cu_seqlens_q, indices_q = dpa_utils.get_cu_seqlens_and_indices(
+                                attention_mask[0]
+                            )
+                            cu_seqlens_kv, indices_kv = dpa_utils.get_cu_seqlens_and_indices(
+                                attention_mask[1]
+                            )
+                        else:
+                            indices_q = dpa_utils.get_indices(max_seqlen_q, cu_seqlens_q)
+                            indices_kv = dpa_utils.get_indices(max_seqlen_kv, cu_seqlens_kv)
+                        query_layer = dpa_utils.PackTensors.apply(indices_q, query_layer)
+                        key_layer, value_layer = dpa_utils.PackTensors.apply(
+                            indices_kv, key_layer, value_layer
+                        )
+                else:
+                    # Cumulative sequence lengths for unpadded data
+                    if cu_seqlens_q is None:
+                        cu_seqlens_q = dpa_utils.get_full_cu_seqlens(
+                            batch_size,
+                            max_seqlen_q,
+                            query_layer.device,
+                        )
+                    if cu_seqlens_kv is None:
+                        cu_seqlens_kv = dpa_utils.get_full_cu_seqlens(
+                            batch_size,
+                            max_seqlen_kv,
+                            key_layer.device,
+                        )
+            elif qkv_format == "thd":
+                assert (
+                    cu_seqlens_q is not None and cu_seqlens_kv is not None
+                ), "cu_seqlens_q and cu_seqlens_kv can not be None when qkv_format = thd!"
+                if max_seqlen_q is None:
+                    seqlens_q = cu_seqlens_q[1:] - cu_seqlens_q[:-1]
+                    max_seqlen_q = seqlens_q.max().item()
+                if max_seqlen_kv is None:
+                    seqlens_kv = cu_seqlens_kv[1:] - cu_seqlens_kv[:-1]
+                    max_seqlen_kv = seqlens_kv.max().item()
+        else:
+            if qkv_format in ["sbhd_2bshd", "bshd"]:
+                # q is in bshd in both cases from conversion above or the original input
+                batch_size, context_len = query_layer.shape[:2]
+                cu_seqlens_q = cu_seqlens_q[: batch_size + 1]
+                cu_seqlens_kv = cu_seqlens_kv[: batch_size + 1]
+                # convert from bshd to thd_2bshd for flash_attn_varlen_func/_with_kvcache;
+                # kernel assumes tensor is contiguous
+                if isinstance(query_layer, Float8Tensor):
+                    query_layer._data = tex.convert_bshd_to_thd(
+                        query_layer._data,
+                        cu_seqlens_q,
+                        batch_size * context_len,
+                    )
+                    query_layer = Float8Tensor.make_like(
+                        query_layer, data=query_layer._data, shape=query_layer._data.shape
+                    )
+                else:
+                    query_layer = tex.convert_bshd_to_thd(
+                        query_layer,
+                        cu_seqlens_q,
+                        batch_size * context_len,
+                    )
+
+        use_flash_attn_3 = False
+        if flash_attention_backend is not None and flash_attention_backend > PkgVersion("3.0.0b"):
+            use_flash_attn_3 = True
+        if context_parallel and all(
+            not isinstance(x, Float8Tensor) for x in [query_layer, key_layer, value_layer]
+        ):
+            assert (
+                alibi_slopes is None
+            ), "Alibi slope bias addition is not supported with context parallelism."
+            with self.attention_dropout_ctx():
+                output = attn_forward_func_with_cp(
+                    self.training,
+                    query_layer,
+                    key_layer,
+                    value_layer,
+                    cu_seqlens_q,
+                    cu_seqlens_kv,
+                    max_seqlen_q,
+                    max_seqlen_kv,
+                    cu_seqlens_q if qkv_format == "thd" else None,
+                    cu_seqlens_kv if qkv_format == "thd" else None,
+                    self.attention_dropout if self.training else 0.0,
+                    cp_group,
+                    cp_global_ranks,
+                    cp_stream,
+                    cp_comm_type,
+                    softmax_scale=self.softmax_scale,
+                    qkv_format="bshd" if qkv_format == "sbhd" else qkv_format,
+                    attn_mask_type=attn_mask_type,
+                    deterministic=self.deterministic,
+                    window_size=window_size,
+                    quantizers=quantizers,
+                    pad_between_seqs=False,
+                    use_flash_attn_3=use_flash_attn_3,
+                    fp8_output=fp8_output,
+                )
+        else:
+            from transformer_engine.pytorch.cpu_offload import (
+                CPUOffloadEnabled,
+                mark_activation_offload,
+            )
+
+            if CPUOffloadEnabled:
+                mark_activation_offload(
+                    query_layer, key_layer, value_layer, cu_seqlens_q, cu_seqlens_kv
+                )
+
+            with self.attention_dropout_ctx():
+                #       | API                     | use cases
+                # ----------------------------------------------------------------------
+                # FA v2 | flash_attn_func         | bshd/sbhd + not padding
+                #       | flash_attn_varlen_func  | bshd/sbhd + padding
+                #       |                         | thd + padding
+                #       |                         | KV cache (not-paged/paged), i.e.
+                #       |                         |     bshd/sbhd/thd + padding
+                # FA v3 | flash_attn_func         | bshd/sbhd + not padding
+                #       | flash_attn_varlen_func  | bshd/sbhd + padding
+                #       |                         | thd + padding
+                #       | flash_attn_with_kvcache | KV cache (not-paged/paged), i.e.
+                #       |                         |     bshd/sbhd/thd + padding
+                fa_optional_forward_args_thd = []
+                if qkv_format in ["bshd", "sbhd"] and "padding" not in attn_mask_type:
+                    func = (
+                        flash_attn_func if not use_flash_attn_3 else flash_attn_func_v3
+                    )  # pylint: disable=possibly-used-before-assignment
+                else:
+                    if not use_flash_attn_3:
+                        func = flash_attn_varlen_func
+                    elif inference_params is None:
+                        func = flash_attn_varlen_func_v3  # pylint: disable=possibly-used-before-assignment
+                    else:
+                        func = flash_attn_with_kvcache_v3  # pylint: disable=possibly-used-before-assignment
+                    if not use_flash_attn_3 or inference_params is None:
+                        fa_optional_forward_args_thd.append(cu_seqlens_q)
+                        fa_optional_forward_args_thd.append(cu_seqlens_kv)
+                        fa_optional_forward_args_thd.append(max_seqlen_q)
+                        fa_optional_forward_args_thd.append(max_seqlen_kv)
+                if not use_flash_attn_3:
+                    fa_optional_forward_kwargs = {}
+                    if fa_utils.v2_3_plus:
+                        fa_optional_forward_kwargs["window_size"] = window_size
+                    if fa_utils.v2_4_plus:
+                        fa_optional_forward_kwargs["alibi_slopes"] = alibi_slopes
+                    if fa_utils.v2_4_1_plus:
+                        fa_optional_forward_kwargs["deterministic"] = self.deterministic
+                    if inference_params is not None:
+                        # use block_table kwarg to support thd_2bshd for non-paged
+                        fa_optional_forward_kwargs["block_table"] = (
+                            inference_params.cache_manager.page_table[:batch_size]
+                            if inference_params.is_paged
+                            else inference_params.cache_manager.batch_indices_post_step.unsqueeze(
+                                1
+                            )[:batch_size]
+                        )
+                    output = func(
+                        query_layer,
+                        key_layer,
+                        value_layer,
+                        *fa_optional_forward_args_thd,
+                        self.attention_dropout if self.training else 0.0,
+                        softmax_scale=self.softmax_scale,
+                        causal="causal" in attn_mask_type,
+                        **fa_optional_forward_kwargs,
+                    )
+                else:
+                    fa_3_optional_forward_kwargs = {}
+                    fa_3_optional_forward_kwargs["window_size"] = window_size
+                    if num_splits is not None and num_splits != 1:
+                        if not use_flash_attn_3:
+                            raise ValueError(
+                                "num_splits != 1 is only supported with FlashAttention-3."
+                            )
+                    if num_splits is not None:
+                        fa_3_optional_forward_kwargs["num_splits"] = num_splits
+                    if inference_params is None:
+                        fa_3_optional_forward_kwargs["deterministic"] = self.deterministic
+                    else:
+                        fa_3_optional_forward_kwargs["cu_seqlens_q"] = cu_seqlens_q
+                        fa_3_optional_forward_kwargs["max_seqlen_q"] = max_seqlen_q
+                        cache_seqlens = cu_seqlens_kv[1:] - cu_seqlens_kv[:-1]
+                        fa_3_optional_forward_kwargs["cache_seqlens"] = cache_seqlens
+                        # flash_attn_with_kvcache accepts thd_2bshd for non-paged
+                        if inference_params.is_paged:
+                            fa_3_optional_forward_kwargs["page_table"] = (
+                                inference_params.cache_manager.page_table[:batch_size]
+                            )
+                    if fp8:
+                        QKV_quantizer = quantizers["scaling_fwd"][META_QKV]
+                        torch_dtype = get_fp8_torch_dtype(fp8_meta["recipe"], fprop_tensor=True)
+                        torch_orig_dtype = query_layer.dtype
+
+                        def convert_to_torch_float8(tensor, dtype):
+                            out = torch.Tensor().to(device=tensor.device, dtype=dtype)
+                            out.set_(
+                                tensor._data.untyped_storage(),
+                                tensor._data.storage_offset(),
+                                tensor._data.shape,
+                                tensor._data.stride(),
+                            )
+                            return out
+
+                        assert isinstance(key_layer, query_layer.__class__) and isinstance(
+                            value_layer, query_layer.__class__
+                        ), "q, k, and v must have the same type."
+                        if not isinstance(query_layer, Float8Tensor):
+                            query_layer, key_layer, value_layer = (
+                                QKV_quantizer(x) for x in [query_layer, key_layer, value_layer]
+                            )
+                        batch_size = cu_seqlens_q.shape[0] - 1
+                        num_heads_k = key_layer.shape[-2]
+                        fa_3_optional_forward_kwargs["q_descale"] = (
+                            query_layer._scale_inv.unsqueeze(0).repeat(batch_size, num_heads_k)
+                        )
+                        fa_3_optional_forward_kwargs["k_descale"] = key_layer._scale_inv.unsqueeze(
+                            0
+                        ).repeat(batch_size, num_heads_k)
+                        fa_3_optional_forward_kwargs["v_descale"] = (
+                            value_layer._scale_inv.unsqueeze(0).repeat(batch_size, num_heads_k)
+                        )
+                        query_layer, key_layer, value_layer = (
+                            convert_to_torch_float8(x, torch_dtype)
+                            for x in [query_layer, key_layer, value_layer]
+                        )
+                    try:
+                        output = func(
+                            query_layer,
+                            key_layer,
+                            value_layer,
+                            *fa_optional_forward_args_thd,
+                            softmax_scale=self.softmax_scale,
+                            causal="causal" in attn_mask_type,
+                            **fa_3_optional_forward_kwargs,
+                        )
+                        if isinstance(output, (List, Tuple)):
+                            output = output[0]
+                    except TypeError as e:
+                        if fa_utils.v3_0_0_beta:
+                            e.args = (
+                                e.args[0]
+                                + ". Please update your flash-attn v3 (beta) installation as it "
+                                + "may have added more supported arguments to its API. \n"
+                                + fa_utils.v3_installation_steps,
+                            ) + e.args[1:]
+                        raise
+
+                    if fp8:
+                        output = output.to(dtype=torch_orig_dtype)
+                    if fp8 and fp8_output:
+                        O_quantizer = quantizers["scaling_fwd"][META_O]
+                        output = O_quantizer(output)
+
+        if inference_params is None:
+            if qkv_format in ["sbhd", "bshd"] and "padding" in attn_mask_type:
+                output = dpa_utils.UnpackTensor.apply(indices_q, batch_size * max_seqlen_q, output)
+        elif qkv_format in ["bshd", "sbhd_2bshd"]:
+            # all KV caching cases use thd_2bshd for calculation
+            # convert results back to bshd from thd_2bshd
+            if isinstance(query_layer, Float8Tensor):
+                output._data = tex.convert_thd_to_bshd(
+                    output._data,
+                    cu_seqlens_q,
+                    batch_size,
+                    context_len,
+                )
+                output = Float8Tensor.make_like(output, data=output._data, shape=output._data.shape)
+            else:
+                output = tex.convert_thd_to_bshd(
+                    output,
+                    cu_seqlens_q,
+                    batch_size,
+                    context_len,
+                )
+
+        if q_format == "sbhd":
+            # (bs)hd -> bs(hd) -> sb(hd)
+            if fp8 and fp8_output:
+                output_data = (
+                    output._data.reshape(batch_size, max_seqlen_q // cp_size, -1)
+                    .transpose(0, 1)
+                    .contiguous()
+                )
+                output = Float8Tensor.make_like(
+                    output,
+                    data=output_data,
+                    shape=output_data.shape,
+                )
+            else:
+                output = output.view(batch_size, max_seqlen_q // cp_size, -1).transpose(0, 1)
+        elif q_format == "bshd":
+            # (bs)hd -> bs(hd)
+            output = output.reshape(batch_size, max_seqlen_q // cp_size, -1)
+        elif q_format == "thd":
+            # thd -> t(hd)
+            output = output.reshape(output.shape[0], -1)
+
+        return output.contiguous()
+
+
+class FusedAttnFunc(torch.autograd.Function):
+    """FusedAttention forward and backward implementation"""
+
+    @staticmethod
+    def forward(
+        ctx,
+        is_training,
+        max_seqlen_q,
+        max_seqlen_kv,
+        cu_seqlens_q,
+        cu_seqlens_kv,
+        cu_seqlens_q_padded,
+        cu_seqlens_kv_padded,
+        page_table_k,
+        page_table_v,
+        q,
+        k,
+        v,
+        attn_bias,
+        attn_scale,
+        dropout_p,
+        fast_zero_fill,
+        qkv_layout,
+        attn_bias_type,
+        attn_mask_type,
+        softmax_type,
+        window_size,
+        rng_gen,
+        fused_attention_backend,
+        use_FAv2_bwd,
+        fp8,
+        fp8_meta,
+        quantizers,
+        deterministic,
+        softmax_offset,
+        fp8_output,
+        layer_number,
+        return_max_logit,
+    ):
+        # pylint: disable=missing-function-docstring
+
+        # add NVTX range
+        nvtx_label = "transformer_engine.FusedAttnFunc.forward"
+        nvtx_range_push(f"{nvtx_label}")
+
+        # recipe passed in through autocast or set by NVTE_DPA_FP8_RECIPE;
+        # may be different from fp8_meta["recipe"]
+        fp8_recipe = FP8GlobalStateManager.get_fp8_recipe()
+        if fp8_meta is not None and fp8_meta.get("local_recipes", None) is not None:
+            fp8_recipe = fp8_meta["local_recipes"][0]
+
+        # input types are inferred from the real data while output types are controlled by fp8_output
+        # fp8_output should be set upstream as (DPA.fp8 and DPA.fp8_meta["recipe"].fp8_mha)
+        assert isinstance(k, q.__class__) and isinstance(
+            v, q.__class__
+        ), "q, k, v must be of the same class, e.g. torch.Tensor or Float8Tensor."
+        is_input_fp8 = isinstance(q, Float8Tensor)
+        is_output_fp8 = fp8_output
+
+        # whether fwd kernel in FP8: fp8 = (DPA.fp8 and DPA.fp8_meta["recipe"].fp8_dpa)
+        # whether bwd kernel in FP8:
+        is_bwd_fp8 = fp8 and int(os.getenv("NVTE_FP8_DPA_BWD", "1"))
+
+        # get quantizers from DPA; all Nones if not fp8
+        QKV_quantizer, O_quantizer, S_quantizer, dQKV_quantizer, dO_quantizer, dP_quantizer = (
+            dpa_utils.get_attention_quantizers(fp8, quantizers)
+        )
+
+        # get nominal data type for out
+        # FP16/BF16 attention: torch.float16 or torch.bfloat16
+        # FP8 attention:       torch.float16 or torch.bfloat16
+        out_nominal_dtype = q.dtype
+
+        max_logit = None
+        if fp8:
+            fused_attention_backend = FusedAttnBackend["FP8"]
+
+            # q, k, v:             torch.Tensor; dtype = torch.float16 or torch.bfloat16
+            # q_fp8, k_fp8, v_fp8: Float8Tensor; dtype = torch.float16 or torch.bfloat16
+            #                                    fp8_dtype = tex.DType.kFloat8E4M3
+            if is_input_fp8:
+                q_fp8, k_fp8, v_fp8 = q, k, v
+            else:
+                q_fp8, k_fp8, v_fp8 = combine_and_quantize(qkv_layout, q, k, v, QKV_quantizer)
+
+            # print quantizers
+            print_quantizers(
+                "FusedAttnFunc.forward >> before: ",
+                layer_number,
+                QKV_quantizer,
+                O_quantizer,
+                S_quantizer,
+                dQKV_quantizer,
+                dO_quantizer,
+                dP_quantizer,
+            )
+
+            # out_:
+            # DelayedScaling:       Float8Tensor; dtype = torch.float16 or torch.bfloat16
+            #                                     fp8_dtype = tex.DType.kFloat8E4M3
+            # Float8CurrentScaling: torch.Tensor; dtype = torch.float16 or torch.bfloat16
+            out_, aux_ctx_tensors, *_ = fused_attn_fwd(
+                is_training,
+                max_seqlen_q,
+                max_seqlen_kv,
+                cu_seqlens_q,
+                cu_seqlens_kv,
+                q_fp8,
+                k_fp8,
+                v_fp8,
+                out_nominal_dtype,
+                fused_attention_backend,
+                attn_bias,
+                cu_seqlens_q_padded,
+                cu_seqlens_kv_padded,
+                None,
+                None,
+                S_quantizer,
+                O_quantizer,
+                attn_scale,
+                dropout_p,
+                fast_zero_fill,
+                qkv_layout,
+                attn_bias_type,
+                attn_mask_type,
+                softmax_type,
+                window_size,
+                rng_gen,
+                softmax_offset,
+            )
+
+            # out_fp8: Float8Tensor; dtype = torch.float16 or torch.bfloat16
+            #                        fp8_dtype = tex.DType.kFloat8E4M3
+            # out:     torch.Tensor; dtype = torch.float16 or torch.bfloat16
+            out_fp8 = out_
+            out = out_
+
+            if isinstance(out_, Float8Tensor):
+                if not is_output_fp8 or not is_bwd_fp8:
+                    out = out_.dequantize().view(out_.shape)
+            else:
+                if is_output_fp8 or (
+                    is_bwd_fp8
+                    and not (fp8_recipe.float8_current_scaling() and _dpa_fp8_cs_o_in_f16)
+                ):
+                    out_fp8 = O_quantizer(out_)
+
+            # print quantizers
+            print_quantizers(
+                "FusedAttnFunc.forward >> after:  ",
+                layer_number,
+                QKV_quantizer,
+                O_quantizer,
+                S_quantizer,
+                dQKV_quantizer,
+                dO_quantizer,
+                dP_quantizer,
+            )
+
+            # return appropriate tensors
+            out_ret = out_fp8 if is_output_fp8 else out
+
+            # save appropriate tensors
+            fp8_tensors = (None, None, None, None)
+            qkvo_tensors = (None, None, None, None)
+            if is_bwd_fp8:
+                if fp8_recipe.float8_current_scaling() and _dpa_fp8_cs_o_in_f16:
+                    fp8_tensors = (q_fp8, k_fp8, v_fp8, None)
+                    qkvo_tensors = (None, None, None, out)
+                else:
+                    fp8_tensors = (q_fp8, k_fp8, v_fp8, out_fp8)
+            else:
+                if is_input_fp8:
+                    q, k, v = combine_and_dequantize(qkv_layout, q_fp8, k_fp8, v_fp8)
+                qkvo_tensors = (q, k, v, out)
+        else:
+            # q, k, v, out_: torch.Tensor; dtype = torch.float16 or torch.bfloat16
+            out_, aux_ctx_tensors, *max_logit = fused_attn_fwd(
+                is_training,
+                max_seqlen_q,
+                max_seqlen_kv,
+                cu_seqlens_q,
+                cu_seqlens_kv,
+                q,
+                k,
+                v,
+                out_nominal_dtype,
+                fused_attention_backend,
+                attn_bias,
+                cu_seqlens_q_padded,
+                cu_seqlens_kv_padded,
+                page_table_k,
+                page_table_v,
+                None,  # s_quantizer
+                None,  # o_quantizer
+                attn_scale,
+                dropout_p,
+                fast_zero_fill,
+                qkv_layout,
+                attn_bias_type,
+                attn_mask_type,
+                softmax_type,
+                window_size,
+                rng_gen,
+                softmax_offset,
+                return_max_logit,
+            )
+            out = out_
+            out_ret = out_
+            fp8_tensors = (None, None, None, None)
+            qkvo_tensors = (q, k, v, out)
+
+        nvtx_range_pop(f"{nvtx_label}")
+
+        ctx.fp8_recipe = fp8_recipe
+        ctx.fp8 = is_bwd_fp8
+        # assume fwd and bwd always use the same high precision, i.e. torch.float16 or torch.bfloat16
+        # used when some tensors are base tensors and loose the "dtype" attribute
+        ctx.nominal_dtype = out_nominal_dtype
+
+        from transformer_engine.pytorch.cpu_offload import (
+            CPUOffloadEnabled,
+            mark_activation_offload,
+        )
+
+        if CPUOffloadEnabled:
+            if ctx.fp8:
+                tensor_list = fp8_tensors
+            else:
+                tensor_list = [q, k, v, out]
+
+            mark_activation_offload(*tensor_list)
+            mark_activation_offload(*aux_ctx_tensors)
+
+        ctx.is_input_fp8 = is_input_fp8
+        ctx.is_output_fp8 = is_output_fp8
+        tensors_to_save, tensor_objects = prepare_for_saving(
+            *fp8_tensors,
+            *qkvo_tensors,
+            cu_seqlens_q,
+            cu_seqlens_kv,
+            cu_seqlens_q_padded,
+            cu_seqlens_kv_padded,
+            *aux_ctx_tensors,
+        )
+        ctx.save_for_backward(*tensors_to_save)
+        ctx.tensor_objects = tensor_objects
+        ctx.fp8_meta = fp8_meta
+
+        ctx.layer_number = layer_number
+        ctx.QKV_quantizer = QKV_quantizer
+        ctx.O_quantizer = O_quantizer
+        ctx.dQKV_quantizer = dQKV_quantizer
+        ctx.dO_quantizer = dO_quantizer
+        ctx.dP_quantizer = dP_quantizer
+        ctx.S_quantizer = S_quantizer
+        if ctx.fp8 and isinstance(ctx.S_quantizer, Float8Quantizer):
+            ctx.S_quantizer = S_quantizer.copy()
+            ctx.S_quantizer.scale = S_quantizer.scale.clone()
+
+        ctx.max_seqlen_q = max_seqlen_q
+        ctx.max_seqlen_kv = max_seqlen_kv
+        ctx.attn_scale = attn_scale
+        ctx.dropout_p = dropout_p
+        ctx.fast_zero_fill = fast_zero_fill
+
+        from transformer_engine.pytorch.cpu_offload import (
+            CPUOffloadedLayer,
+        )
+
+        # If interleaved tensor is offloaded, reloaded tensor will be
+        # non-interleaved, so we need to modify the QKV layout
+        # for backward
+        if CPUOffloadedLayer and CPUOffloadEnabled:
+            reload_layout = ""
+            split_list = qkv_layout.split("_")
+            for split in split_list:
+                temp_layout = ""
+                rep_count = 1
+                for s in split:
+                    if s.isalpha():
+                        temp_layout = temp_layout + s
+                    else:
+                        rep_count = int(s)
+                for _ in range(rep_count):
+                    reload_layout = reload_layout + temp_layout + "_"
+            ctx.qkv_layout = reload_layout[:-1]
+        else:
+            ctx.qkv_layout = qkv_layout
+
+        ctx.attn_bias_type = attn_bias_type
+        ctx.attn_mask_type = attn_mask_type
+        ctx.softmax_type = softmax_type
+        ctx.window_size = window_size
+        ctx.fused_attention_backend = (
+            fused_attention_backend if ctx.fp8 else FusedAttnBackend["F16_arbitrary_seqlen"]
+        )
+        ctx.use_FAv2_bwd = use_FAv2_bwd
+        ctx.deterministic = deterministic
+
+        if return_max_logit:
+            return out_ret, *max_logit
+        return out_ret
+
+    @staticmethod
+    def backward(ctx, d_out, *_args):
+        # pylint: disable=missing-function-docstring
+
+        # d_out is expected to be in FP8 if is_output_fp8=True,
+        # but in the case it's not, convert it to FP8 before any operation
+        if ctx.fp8 and ctx.is_output_fp8 and not isinstance(d_out, QuantizedTensorStorage):
+            d_out = ctx.dO_quantizer(d_out)
+            if not ctx.use_FAv2_bwd:
+                d_out._data = d_out._data.contiguous()
+        elif not ctx.use_FAv2_bwd:
+            d_out = d_out.contiguous()
+        (
+            q_fp8,
+            k_fp8,
+            v_fp8,
+            out_fp8,
+            q,
+            k,
+            v,
+            out,
+            cu_seqlens_q,
+            cu_seqlens_kv,
+            cu_seqlens_q_padded,
+            cu_seqlens_kv_padded,
+            *other_tensors,
+        ) = restore_from_saved(ctx.tensor_objects, ctx.saved_tensors)
+
+        aux_ctx_tensors = other_tensors
+
+        if not aux_ctx_tensors[0].is_contiguous():
+            aux_ctx_tensors[0] = aux_ctx_tensors[0].contiguous()
+        rest = [None]
+        if ctx.use_FAv2_bwd:
+            softmax_lse, rng_state = aux_ctx_tensors
+            dq = torch.empty_like(q)
+            dk = torch.empty_like(k)
+            dv = torch.empty_like(v)
+            d_out, q, k, v, out = [dpa_utils.maybe_contiguous(x) for x in (d_out, q, k, v, out)]
+            # from transformer_engine.pytorch.attention.dot_product_attention import flash_attn_cuda_bwd
+            flash_attn_cuda_bwd(
+                d_out,
+                q,
+                k,
+                v,
+                out,
+                softmax_lse,
+                dq,
+                dk,
+                dv,
+                cu_seqlens_q,
+                cu_seqlens_kv,
+                ctx.max_seqlen_q,
+                ctx.max_seqlen_kv,
+                ctx.dropout_p,
+                ctx.attn_scale,
+                False,
+                "causal" in ctx.attn_mask_type,
+                None,
+                rng_state,
+            )
+            dq = dq[..., : d_out.shape[-1]]
+            dk = dk[..., : d_out.shape[-1]]
+            dv = dv[..., : d_out.shape[-1]]
+        else:
+            with torch.cuda.nvtx.range("FusedAttnFunc.backward"):
+                # get nominal data type of dq, dk, dv
+                # FP16/BF16 attention: torch.float16 or torch.bfloat16
+                # FP8 attention:       torch.float16 or torch.bfloat16
+                dqkv_nominal_dtype = ctx.nominal_dtype
+
+                if ctx.fp8:
+                    # d_out:     torch.Tensor; dtype = torch.float16 or torch.bfloat16
+                    # d_out_fp8: Float8Tensor; dtype = torch.float16 or torch.bfloat16
+                    #                          fp8_dtype = tex.DType.kFloat8E5M2
+                    if ctx.is_output_fp8:
+                        d_out_fp8 = d_out
+                    else:
+                        d_out_fp8 = ctx.dO_quantizer(d_out)
+
+                    # print quantizers
+                    print_quantizers(
+                        "FusedAttnFunc.backward >> before: ",
+                        ctx.layer_number,
+                        ctx.QKV_quantizer,
+                        ctx.O_quantizer,
+                        ctx.S_quantizer,
+                        ctx.dQKV_quantizer,
+                        ctx.dO_quantizer,
+                        ctx.dP_quantizer,
+                    )
+
+                    # get tex.DType for dq, dk, dv data
+                    dqkv_te_dtype = d_out_fp8._fp8_dtype
+
+                    # q_fp8, k_fp8, v_fp8, out_fp8: Float8Tensor; dtype = torch.float16 or torch.bfloat16,
+                    #                               fp8_dtype = tex.DType.kFloat8E4M3
+                    # d_out_fp8:                    Float8Tensor; dtype = torch.float16 or torch.bfloat16
+                    #                               fp8_dtype = tex.DType.kFloat8E5M2
+                    # out_:
+                    # DelayedScaling:               Float8Tensor; dtype = torch.float16 or torch.bfloat16
+                    #                               fp8_dtype = tex.DType.kFloat8E4M3
+                    # Float8CurrentScaling:         torch.Tensor; dtype = torch.float16 or torch.bfloat16
+                    #
+                    # dq_, dk_, dv_:
+                    # DelayedScaling:               Float8Tensor; dtype = torch.float16 or torch.bfloat16
+                    #                               fp8_dtype = tex.DType.kFloat8E5M2
+                    # Float8CurrentScaling:         torch.Tensor; dtype = torch.float16 or torch.bfloat16
+                    out_ = (
+                        out
+                        if ctx.fp8_recipe.float8_current_scaling() and _dpa_fp8_cs_o_in_f16
+                        else out_fp8
+                    )
+                    dq_, dk_, dv_, *rest = fused_attn_bwd(
+                        ctx.max_seqlen_q,
+                        ctx.max_seqlen_kv,
+                        cu_seqlens_q,
+                        cu_seqlens_kv,
+                        q_fp8,
+                        k_fp8,
+                        v_fp8,
+                        out_,
+                        d_out_fp8,
+                        dqkv_nominal_dtype,
+                        dqkv_te_dtype,
+                        aux_ctx_tensors,
+                        ctx.fused_attention_backend,
+                        cu_seqlens_q_padded,
+                        cu_seqlens_kv_padded,
+                        ctx.S_quantizer,
+                        ctx.dP_quantizer,
+                        ctx.dQKV_quantizer,
+                        ctx.attn_scale,
+                        ctx.dropout_p,
+                        ctx.fast_zero_fill,
+                        ctx.qkv_layout,
+                        ctx.attn_bias_type,
+                        ctx.attn_mask_type,
+                        ctx.softmax_type,
+                        ctx.window_size,
+                        ctx.deterministic,
+                    )
+
+                    # dq, dk, dv:             torch.Tensor; dtype = torch.float16 or torch.bfloat16
+                    dq, dk, dv = dq_, dk_, dv_
+                    is_float8tensor = isinstance(dq_, Float8Tensor)
+                    if is_float8tensor and not ctx.is_input_fp8:
+                        # return in F16
+                        dq, dk, dv = combine_and_dequantize(
+                            ctx.qkv_layout,
+                            dq_,
+                            dk_,
+                            dv_,
+                            src_nominal_dtype=dq_.dtype,
+                        )
+                    if not is_float8tensor and ctx.is_input_fp8:
+                        # return in FP8
+                        dq, dk, dv = combine_and_quantize(
+                            ctx.qkv_layout, dq_, dk_, dv_, ctx.dQKV_quantizer
+                        )
+
+                    # print quantizers
+                    print_quantizers(
+                        "FusedAttnFunc.backward >> after:  ",
+                        ctx.layer_number,
+                        ctx.QKV_quantizer,
+                        ctx.O_quantizer,
+                        ctx.S_quantizer,
+                        ctx.dQKV_quantizer,
+                        ctx.dO_quantizer,
+                        ctx.dP_quantizer,
+                    )
+                else:
+                    if isinstance(d_out, QuantizedTensorStorage):
+                        d_out = d_out.dequantize(dtype=ctx.nominal_dtype)
+                    dqkv_te_dtype = TE_DType[d_out.dtype]
+                    # q, k, v, out, d_out, dq, dk, dv: torch.Tensor; torch.float16 or torch.bfloat16
+                    dq, dk, dv, *rest = fused_attn_bwd(
+                        ctx.max_seqlen_q,
+                        ctx.max_seqlen_kv,
+                        cu_seqlens_q,
+                        cu_seqlens_kv,
+                        q,
+                        k,
+                        v,
+                        out,
+                        d_out,
+                        dqkv_nominal_dtype,
+                        dqkv_te_dtype,
+                        aux_ctx_tensors,
+                        ctx.fused_attention_backend,
+                        cu_seqlens_q_padded,
+                        cu_seqlens_kv_padded,
+                        None,
+                        None,
+                        None,
+                        ctx.attn_scale,
+                        ctx.dropout_p,
+                        ctx.fast_zero_fill,
+                        ctx.qkv_layout,
+                        ctx.attn_bias_type,
+                        ctx.attn_mask_type,
+                        ctx.softmax_type,
+                        ctx.window_size,
+                        ctx.deterministic,
+                    )
+
+        d_bias = None
+        if ctx.attn_bias_type not in ["no_bias", "alibi"]:
+            d_bias = rest[0]
+        d_softmax_offset = None
+        if ctx.softmax_type != "vanilla":
+            d_softmax_offset = rest[1]
+        return (
+            None,
+            None,
+            None,
+            None,
+            None,
+            None,
+            None,
+            None,
+            None,
+            dq,
+            dk,
+            dv,
+            d_bias,
+            None,
+            None,
+            None,
+            None,
+            None,
+            None,
+            None,
+            None,
+            None,
+            None,
+            None,
+            None,
+            None,
+            None,
+            None,
+            d_softmax_offset,
+            None,
+            None,
+            None,
+        )
+
+
+class FusedAttention(torch.nn.Module):
+    """Dot product attention, with multiple backends:
+
+    1. FusedAttnBackend["F16_max512_seqlen"]
+       cuDNN based fused attention for FP16/BF16 and <=512 sequence length.
+    2. FusedAttnBackend["F16_arbitrary_seqlen"]
+       cuDNN based fused attention for FP16/BF16 and any sequence length.
+
+    Support matrix:
+
+    | backend       | 1                       | 2                              |
+    | flash based   | no                      | yes                            |
+    | cuDNN based   | yes                     | yes                            |
+    | qkv dtype     | fp16/bf16               | fp16/bf16                      |
+    | attn_type     | self/cross              | self/cross                     |
+    | qkv_layout    |                         |                                |
+    |  - (q,k,v)    | sb3hd, bs3hd            | sb3hd, bs3hd, sbh3d, bsh3d     |
+    |               | sbhd_sb2hd, bshd_bs2hd  | sbhd_sb2hd, bshd_bs2hd         |
+    |               | bshd_bshd_bshd          | sbhd_sbh2d, bshd_bsh2d         |
+    |               |                         | sbhd_sbhd_sbhd, bshd_bshd_bshd |
+    | mask_type     | causal/padding/no_mask  | causal/padding/no_mask         |
+    | bias_type     | post_scale_bias/no_bias | post_scale_bias/alibi/no_bias  |
+    | dropout       | yes                     | yes                            |
+    | max_seqlen    | <=512, multiple of 64   | any, multiple of 64            |
+    | head_dim      | 64                      | <=128, multiple of 8           |
+    | output dtype  | fp16/bf16               | fp16/bf16                      |
+    """
+
+    def __init__(
+        self,
+        softmax_scale: float,
+        attention_dropout: float = 0.0,
+        attention_dropout_ctx: Optional[Callable] = nullcontext,
+        attention_type: str = "self",
+        layer_number: Optional[int] = None,
+        deterministic: bool = False,
+        softmax_type: str = "vanilla",
+        return_max_logit: Optional[bool] = False,
+    ) -> None:
+        super().__init__()
+
+        self.softmax_scale = softmax_scale
+        self.attention_dropout = attention_dropout
+        self.attention_dropout_ctx = attention_dropout_ctx
+        self.attention_type = attention_type
+        self.use_FAv2_bwd = os.getenv(
+            "NVTE_FUSED_ATTN_USE_FAv2_BWD", "0"
+        ) == "1" and get_device_compute_capability() == (9, 0)
+        self.layer_number = 1 if layer_number is None else layer_number
+        self.deterministic = deterministic
+        self.softmax_type = softmax_type
+        self.return_max_logit = return_max_logit
+
+        def remove_extra_states_check(self, incompatible_keys):  # pylint: disable=unused-argument
+            """
+            Temporarily remove fused_attention._extra_state as a missing key
+            or an unexpected key when loading Transformer Engine checkpoints.
+            Please store FP8 metadata as DotProductAttention's _extra_state,
+            rather than FusedAttention's _extra_state. This hook will be
+            phased out in Transformer Engine 2.0.
+            """
+            for key in incompatible_keys.missing_keys:
+                if "fused_attention._extra_state" in key:
+                    incompatible_keys.missing_keys.remove(key)
+            for key in incompatible_keys.unexpected_keys:
+                if "fused_attention._extra_state" in key:
+                    incompatible_keys.unexpected_keys.remove(key)
+                    warnings.warn(
+                        "fused_attention._extra_state is not loaded from checkpoint. Please map "
+                        "FusedAttention's _extra_state to DotProductAttention's _extra_state."
+                    )
+
+        self.register_load_state_dict_post_hook(remove_extra_states_check)
+
+    @no_torch_dynamo()
+    def forward(
+        self,
+        query_layer: torch.Tensor,
+        key_layer: torch.Tensor,
+        value_layer: torch.Tensor,
+        qkv_layout: str = "sbh3d",
+        cu_seqlens_q: Optional[torch.Tensor] = None,
+        cu_seqlens_kv: Optional[torch.Tensor] = None,
+        cu_seqlens_q_padded: Optional[torch.Tensor] = None,
+        cu_seqlens_kv_padded: Optional[torch.Tensor] = None,
+        max_seqlen_q: Optional[int] = None,
+        max_seqlen_kv: Optional[int] = None,
+        attn_mask_type: str = "causal",
+        attention_mask: Optional[Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]] = None,
+        window_size: Optional[Tuple[int, int]] = None,
+        fused_attention_backend: tex.NVTE_Fused_Attn_Backend = tex.NVTE_Fused_Attn_Backend.NVTE_No_Backend,
+        core_attention_bias_type: str = "no_bias",
+        core_attention_bias: Optional[torch.Tensor] = None,
+        fast_zero_fill: bool = True,
+        cp_group: Optional[Union[dist_group_type, List[dist_group_type]]] = None,
+        cp_global_ranks: List[int] = None,
+        cp_stream: torch.cuda.Stream = None,
+        cp_comm_type: str = "p2p",
+        fp8: bool = False,
+        fp8_meta: Optional[Dict[str, Any]] = None,
+        quantizers=None,
+        pad_between_seqs: bool = False,
+        inference_params: Optional[InferenceParams] = None,
+        softmax_offset: torch.Tensor = None,
+        fp8_output: bool = False,
+    ) -> torch.Tensor:
+        """fused attention fprop"""
+        assert (
+            fused_attention_backend != tex.NVTE_Fused_Attn_Backend.NVTE_No_Backend
+        ), "No fused attention backend supports this input combination!"
+        assert all(
+            x.dtype in [torch.float16, torch.bfloat16] or isinstance(x, Float8Tensor)
+            for x in [query_layer, key_layer, value_layer]
+        ), "FusedAttention only supports FP16 and BF16 data types, or Float8Tensors."
+        assert (
+            query_layer.is_cuda and key_layer.is_cuda and value_layer.is_cuda
+        ), "FusedAttention only supports CUDA tensors."
+        assert (
+            qkv_layout in QKVLayouts
+        ), f"FusedAttention does not support qkv_layout = {qkv_layout}!"
+
+        cp_size = 1
+        if isinstance(cp_group, dist_group_type):
+            cp_size = get_distributed_world_size(cp_group)
+        elif isinstance(cp_group, list):
+            for group in cp_group:
+                cp_size *= get_distributed_world_size(group)
+        context_parallel = cp_size > 1
+
+        # get q_format and kv_format for training and inference
+        qkv_format, q_format, kv_format = dpa_utils.get_qkv_format(qkv_layout, inference_params)
+
+        # cuDNN can work with 0-length sequences in the batch for both bshd/sbhd and thd formats
+        # however, for bshd/sbhd, q/k/v tensors need to have the same batch size as indicated by
+        # cu_seqlens, whereas thd does not have this requirement
+        # e.g. if q_format = bshd, and q.shape = [3, 1, 16, 64], we should have k.shape[0] =
+        # v.shape[0] = q.shape[0], and cu_seqlens_q.shape = cu_seqlens_kv.shape = [4]
+        if q_format in ["bshd", "sbhd"] or kv_format in ["bshd", "sbhd"]:
+            batch_size = query_layer.shape[0] if q_format == "bshd" else query_layer.shape[1]
+            cu_seqlens_q = cu_seqlens_q[: batch_size + 1]
+            cu_seqlens_kv = cu_seqlens_kv[: batch_size + 1]
+
+        page_table = None
+        if inference_params is None:
+            if qkv_format in ["sbhd", "bshd"]:
+                if qkv_format == "sbhd":
+                    batch_size = query_layer.shape[1]
+                    max_seqlen_q = query_layer.shape[0]
+                    max_seqlen_kv = key_layer.shape[0]
+                if qkv_format == "bshd":
+                    batch_size = query_layer.shape[0]
+                    max_seqlen_q = query_layer.shape[1]
+                    max_seqlen_kv = key_layer.shape[1]
+                max_seqlen_q *= cp_size
+                max_seqlen_kv *= cp_size
+                if "padding" in attn_mask_type:
+                    assert (
+                        not context_parallel
+                    ), "Padding mask not supported with context parallelism!"
+                    if cu_seqlens_q is None or cu_seqlens_kv is None:
+                        if attention_mask is None:
+                            raise RuntimeError(
+                                "Please provide attention_mask or cu_seqlens for padding!"
+                            )
+                        if self.attention_type == "self":
+                            cu_seqlens_q = dpa_utils.get_cu_seqlens(attention_mask)
+                            cu_seqlens_kv = cu_seqlens_q
+                        else:
+                            cu_seqlens_q = dpa_utils.get_cu_seqlens(attention_mask[0])
+                            cu_seqlens_kv = dpa_utils.get_cu_seqlens(attention_mask[1])
+                else:
+                    if cu_seqlens_q is None:
+                        cu_seqlens_q = dpa_utils.get_full_cu_seqlens(
+                            batch_size,
+                            max_seqlen_q,
+                            query_layer.device,
+                        )
+                    if cu_seqlens_kv is None:
+                        cu_seqlens_kv = dpa_utils.get_full_cu_seqlens(
+                            batch_size,
+                            max_seqlen_kv,
+                            key_layer.device,
+                        )
+            if qkv_format == "thd":
+                assert (
+                    max_seqlen_q is not None
+                    and max_seqlen_kv is not None
+                    and cu_seqlens_q is not None
+                    and cu_seqlens_kv is not None
+                ), "max_seqlen_q/kv and cu_seqlens_q/kv can not be None when qkv_format is thd!"
+        elif inference_params.is_paged:
+            page_table = inference_params.cache_manager.page_table
+
+        if (q_format == "thd" or "padding" in attn_mask_type) and cu_seqlens_q_padded is None:
+            cu_seqlens_q_padded = cu_seqlens_q
+        if (kv_format == "thd" or "padding" in attn_mask_type) and cu_seqlens_kv_padded is None:
+            cu_seqlens_kv_padded = cu_seqlens_kv
+
+        use_FAv2_bwd = (
+            self.use_FAv2_bwd
+            and (core_attention_bias_type == "no_bias")
+            and (fused_attention_backend == tex.NVTE_Fused_Attn_Backend.NVTE_F16_arbitrary_seqlen)
+        )
+
+        if fp8:
+            fp8_recipe = FP8GlobalStateManager.get_fp8_recipe()
+            if fp8_meta is not None and fp8_meta.get("local_recipes", None) is not None:
+                fp8_recipe = fp8_meta["local_recipes"][0]
+            assert fused_attention_backend == tex.NVTE_Fused_Attn_Backend.NVTE_FP8, (
+                f"cuDNN attention sub-backend {int(tex.NVTE_Fused_Attn_Backend.NVTE_FP8)}"
+                " is required for FP8 attention!"
+            )
+            assert fp8_meta is not None, "FP8 metadata fp8_meta is required for FP8 attention!"
+            if fp8_recipe.delayed():
+                assert not context_parallel or fp8_recipe.reduce_amax, (
+                    "Amax reduction across TP+CP group is necessary when using context parallelism"
+                    " with FP8!"
+                )
+            if fp8_recipe.float8_current_scaling() and context_parallel:
+                all_quantizers = dpa_utils.get_attention_quantizers(fp8, quantizers)
+                for q in all_quantizers:
+                    if isinstance(q, Float8CurrentScalingQuantizer):
+                        q.with_amax_reduction = True
+                        q.amax_reduction_group = (
+                            cp_group[0] if cp_comm_type == "a2a+p2p" else cp_group
+                        )
+
+        if context_parallel:
+            assert (
+                fp8
+                or fused_attention_backend == tex.NVTE_Fused_Attn_Backend.NVTE_F16_arbitrary_seqlen
+            ), f"{fused_attention_backend} does not work with context parallelism!"
+            assert core_attention_bias_type not in [
+                "alibi"
+            ], f"{core_attention_bias_type} is not supported with context parallelism!"
+            query_layer, key_layer, value_layer = [
+                x.contiguous() for x in (query_layer, key_layer, value_layer)
+            ]
+            with self.attention_dropout_ctx():
+                output = attn_forward_func_with_cp(
+                    self.training,
+                    query_layer,
+                    key_layer,
+                    value_layer,
+                    cu_seqlens_q,
+                    cu_seqlens_kv,
+                    max_seqlen_q,
+                    max_seqlen_kv,
+                    cu_seqlens_q_padded,
+                    cu_seqlens_kv_padded,
+                    self.attention_dropout if self.training else 0.0,
+                    cp_group,
+                    cp_global_ranks,
+                    cp_stream,
+                    cp_comm_type,
+                    softmax_scale=self.softmax_scale,
+                    qkv_format=qkv_format,
+                    attn_mask_type=attn_mask_type,
+                    attn_bias_type=core_attention_bias_type,
+                    attn_bias=core_attention_bias,
+                    deterministic=self.deterministic,
+                    use_fused_attention=True,
+                    window_size=window_size,
+                    fp8=fp8,
+                    fp8_meta=fp8_meta,
+                    quantizers=quantizers,
+                    pad_between_seqs=pad_between_seqs,
+                    softmax_type=self.softmax_type,
+                    softmax_offset=softmax_offset,
+                    fp8_output=fp8_output,
+                    layer_number=self.layer_number,
+                    return_max_logit=self.return_max_logit,
+                )
+        else:
+            with self.attention_dropout_ctx():
+                output = FusedAttnFunc.apply(
+                    self.training,
+                    max_seqlen_q,
+                    max_seqlen_kv,
+                    cu_seqlens_q,
+                    cu_seqlens_kv,
+                    cu_seqlens_q_padded,
+                    cu_seqlens_kv_padded,
+                    page_table,
+                    page_table,
+                    query_layer,
+                    key_layer,
+                    value_layer,
+                    core_attention_bias,
+                    self.softmax_scale,
+                    self.attention_dropout if self.training else 0.0,
+                    fast_zero_fill,
+                    qkv_layout,
+                    core_attention_bias_type,
+                    attn_mask_type,
+                    self.softmax_type,
+                    window_size,
+                    None,  # rng_gen
+                    fused_attention_backend,
+                    use_FAv2_bwd,
+                    fp8,
+                    fp8_meta,
+                    quantizers,
+                    self.deterministic,
+                    softmax_offset,
+                    fp8_output,
+                    self.layer_number,
+                    self.return_max_logit,
+                )
+
+        if self.return_max_logit:
+            # ...hd -> ...(hd)
+            return output[0].view(*output[0].shape[:-2], -1), output[1]
+        # ...hd -> ...(hd)
+        return output.view(*output.shape[:-2], -1)
diff --git a/examples/tp-numerics/patches/v2.9.0/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py b/examples/tp-numerics/patches/v2.9.0/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py
new file mode 100644
index 00000000000..dc33e02da3b
--- /dev/null
+++ b/examples/tp-numerics/patches/v2.9.0/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py
@@ -0,0 +1,1560 @@
+# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+#
+# See LICENSE for license information.
+
+"""Attention."""
+from contextlib import nullcontext
+import math
+import os
+from typing import Any, Callable, Dict, List, Optional, Tuple, Union
+import warnings
+import logging
+
+import torch
+from torch.nn.parameter import Parameter
+
+import transformer_engine_torch as tex
+from transformer_engine.common.recipe import (
+    Format,
+    Recipe,
+    DelayedScaling,
+    Float8CurrentScaling,
+)
+from transformer_engine.pytorch.utils import get_cudnn_version
+from transformer_engine.pytorch.quantization import (
+    get_fp8_te_dtype,
+    FP8GlobalStateManager,
+    RecipeState,
+    DelayedScalingRecipeState,
+    MXFP8BlockScalingRecipeState,
+    Float8CurrentScalingRecipeState,
+    Float8BlockScalingRecipeState,
+)
+from transformer_engine.pytorch.tensor.float8_tensor import Float8Tensor
+from transformer_engine.pytorch.module.base import TransformerEngineBaseModule
+from transformer_engine.pytorch.export import is_in_onnx_export_mode
+from transformer_engine.pytorch.constants import (
+    AttnMaskTypes,
+    AttnTypes,
+    dist_group_type,
+)
+from transformer_engine.pytorch.distributed import (
+    get_distributed_world_size,
+    checkpoint,
+    set_all_rng_states,
+    CudaRNGStatesTracker,
+    graph_safe_rng_available,
+)
+from transformer_engine.pytorch.jit import no_torch_dynamo
+from transformer_engine.pytorch.graph import is_graph_capturing
+from transformer_engine.pytorch.attention.inference import InferenceParams
+
+# Import attention utils
+import transformer_engine.pytorch.attention.dot_product_attention.utils as dpa_utils
+from transformer_engine.pytorch.attention.dot_product_attention.utils import (
+    AttentionLogging as attn_log,
+)
+
+from transformer_engine.pytorch.attention.dot_product_attention.backends import (
+    UnfusedDotProductAttention,
+    FusedAttention,
+    FlashAttention,
+)
+
+
+# Setup Attention Logging
+attn_log.setup_logging()
+
+# Global vars for available attention backends and ALiBi cache
+_attention_backends = {
+    "attention_params": None,
+    "use_flash_attention": None,
+    "flash_attention_backend": None,
+    "use_fused_attention": None,
+    "fused_attention_backend": None,
+    "use_unfused_attention": None,
+    "backend_selection_requires_update": False,
+}
+
+_alibi_cache = {
+    "_num_heads": None,
+    "_alibi_slopes": None,
+    "_max_seqlen_q": None,
+    "_max_seqlen_kv": None,
+    "_bottom_right_alignment": True,
+    "_alibi_bias": None,
+    "_alibi_slopes_require_update": False,
+    "_alibi_bias_require_update": False,
+}
+
+"""
+This feature is **experimental** and subject to change.
+
+Some models may use different FP8 recipes for their linear layers and attention layers. To support this,
+users can either use multiple, nested autocast() contexts to assign a distinct recipe for each layer,
+or use a single autocast() for the non-attention layers and configure the recipe for the attention
+layers as follows.
+
++-------------------+-----------+-----------------------------------------------------------------------------------+
+| Linear            | Attention | Configuration                                                                     |
++===================+===========+===================================================================================+
+| FP8DS/FP8CS/NVFP4 | FP16/BF16 | Pass FP8DS, FP8CS or NVFP4 to autocast();                                     |
+|                   |           | export NVTE_DPA_FP8_RECIPE="F16"                                                  |
++-------------------+-----------+-----------------------------------------------------------------------------------+
+| FP8DS             | FP8DS     | Pass FP8DS to autocast();                                                     |
++-------------------+-----------+-----------------------------------------------------------------------------------+
+| FP8CS             | FP8DS     | Pass FP8CS to autocast();                                                     |
+|                   |           | Attention FP8DS reuses the fp8_format, fp8_dpa, fp8_mha values from linear FP8CS; |
+|                   |           | export NVTE_DPA_FP8_RECIPE="DelayedScaling"       # switch to DS                  |
+|                   |           | export NVTE_DPA_FP8DS_AMAX_ALGO="most_recent"     # or "max"                      |
+|                   |           | export NVTE_DPA_FP8DS_AMAX_HISTLEN=1              # or any other integer          |
+|                   |           | export NVTE_DPA_FP8DS_REDUCE_AMAX=1               # or 0                          |
++-------------------+-----------+-----------------------------------------------------------------------------------+
+| NVFP4             | FP8DS     | Pass NVFP4 to autocast();                                                     |
+|                   |           | Attention FP8DS reuses the fp8_dpa, fp8_mha values from linear NVFP4;             |
+|                   |           | export NVTE_DPA_FP8_RECIPE="DelayedScaling"       # switch to DS                  |
+|                   |           | export NVTE_DPA_FP8_FORMAT="HYBRID"               # or "E4M3", "E5M2"             |
+|                   |           | export NVTE_DPA_FP8DS_AMAX_ALGO="most_recent"     # or "max"                      |
+|                   |           | export NVTE_DPA_FP8DS_AMAX_HISTLEN=1              # or any other integer          |
+|                   |           | export NVTE_DPA_FP8DS_REDUCE_AMAX=1               # or 0                          |
++-------------------+-----------+-----------------------------------------------------------------------------------+
+| FP8DS             | FP8CS     | Pass FP8DS to autocast();                                                     |
+|                   |           | Attention uses FP8DS for S, dP tensors, and creates a new FP8CS recipe for QKV, O,|
+|                   |           | dO, dQKV tensors based on fp8_format, fp8_dpa, fp8_mha from linear FP8DS;         |
+|                   |           | export NVTE_DPA_FP8_RECIPE="Float8CurrentScaling" # switch to CS                  |
++-------------------+-----------+-----------------------------------------------------------------------------------+
+| FP8CS             | FP8CS     | Pass FP8CS to autocast();                                                     |
+|                   |           | Attention uses FP8CS for QKV, O, dO, dQKV tensors, and creates a new FP8DS recipe |
+|                   |           | for S, dP tensors based on fp8_format, fp8_dpa, fp8_mha from linear FP8CS and:    |
+|                   |           | export NVTE_DPA_FP8DS_AMAX_ALGO="most_recent"     # or "max"                      |
+|                   |           | export NVTE_DPA_FP8DS_AMAX_HISTLEN=1              # or any other integer          |
+|                   |           | export NVTE_DPA_FP8DS_REDUCE_AMAX=1               # or 0                          |
++-------------------+-----------+-----------------------------------------------------------------------------------+
+| NVFP4             | FP8CS     | Pass NVFP4 to autocast();                                                     |
+|                   |           | Attention creates a new FP8CS recipe for QKV, O, dO, dQKV, and a new FP8DS recipe |
+|                   |           | for S, dP, based on the fp8_dpa, fp8_mha values from linear NVFP4 and:            |
+|                   |           | export NVTE_DPA_FP8_RECIPE="Float8CurrentScaling" # switch to CS                  |
+|                   |           | export NVTE_DPA_FP8_FORMAT="HYBRID"               # or "E4M3", "E5M2"             |
+|                   |           | export NVTE_DPA_FP8DS_AMAX_ALGO="most_recent"     # or "max"                      |
+|                   |           | export NVTE_DPA_FP8DS_AMAX_HISTLEN=1              # or any other integer          |
+|                   |           | export NVTE_DPA_FP8DS_REDUCE_AMAX=1               # or 0                          |
++-------------------+-----------+-----------------------------------------------------------------------------------+
+"""
+_dpa_fp8_recipe = os.getenv("NVTE_DPA_FP8_RECIPE", "")
+formats = {"HYBRID": Format.HYBRID, "E4M3": Format.E4M3, "E5M2": Format.E5M2}
+_dpa_fp8_format = formats[os.getenv("NVTE_DPA_FP8_FORMAT", "HYBRID")]
+_dpa_fp8ds_amax_algo = os.getenv("NVTE_DPA_FP8DS_AMAX_ALGO", "most_recent")
+_dpa_fp8ds_amax_histlen = int(os.getenv("NVTE_DPA_FP8DS_AMAX_HISTLEN", "1"))
+_dpa_fp8ds_reduce_amax = os.getenv("NVTE_DPA_FP8DS_REDUCE_AMAX", "1") == "1"
+
+
+__all__ = ["DotProductAttention"]
+
+
+class DotProductAttention(TransformerEngineBaseModule):
+    """Allows the model to jointly attend to information from different
+    representation subspaces as described in the paper:
+    `Attention Is All You Need <https://arxiv.org/abs/1706.03762>`_.
+
+    .. note::
+
+        Argument :attr:`attention_mask` in the `forward` call is only used when
+        :attr:`attn_mask_type` includes '"padding"' or `"arbitrary"`.
+
+    .. warning::
+
+        FlashAttention uses a non-deterministic algorithm for optimal performance. To observe
+        deterministic behavior at the cost of performance, use FlashAttention version >= `2.4.1`
+        and set the environment variable :attr:`NVTE_ALLOW_NONDETERMINISTIC_ALGO=0`. In order
+        to disable`flash-attn` entirely, set :attr:`NVTE_FLASH_ATTN=0`.
+
+    .. note::
+
+        Transformer Engine stores the FP8 metadata under a `._extra_state` key when checkpointing.
+        As the FP8 attention support expands from one backend to multiple backends, the location
+        of that key has also shifted (see `FP8 checkpoint compatibility <https://docs.nvidia.com/deeplearning/transformer-engine/user-guide/faq.html#fp8-checkpoint-compatibility>`_).
+
+
+    Parameters
+    ----------
+    num_attention_heads : int
+                         number of attention heads in the transformer layer.
+    kv_channels : Union[int, Tuple[int, int]]
+                the head size in key and value tensors. If the same, :attr:`kv_channels` can be
+                an integer; if not, :attr:`kv_channels` should be a tuple of two integers.
+    num_gqa_groups : Optional[int] = None
+                    number of GQA groups in the transformer layer.
+                    Grouped Query Attention is described in
+                    `this paper <https://arxiv.org/pdf/2305.13245.pdf>`_.
+                    This only affects the keys and values, not the queries.
+                    GQA-1 is equivalent to Multi-Query Attention
+                    (`MQA <https://arxiv.org/pdf/1911.02150.pdf>`_), while GQA-H
+                    is equivalent to MHA, i.e. `num_gqa_groups = num_attention_heads`.
+    attention_dropout: float, default = 0.0
+                      dropout probability for the dropout op during multi-head attention.
+    attn_mask_type: str, default = `causal`
+                   type of attention mask passed into softmax operation, options are "`no_mask`",
+                   "`padding`", "`causal`", "`padding,causal`", "`causal,padding`",
+                   "`padding_causal`", "`causal_bottom_right`", "`padding_causal_bottom_right`", and
+                   "`arbitrary`", where "`padding,causal`", "`causal,padding`" and "`padding_causal`"
+                   are equivalent. This arg can be overridden by :attr:`attn_mask_type` in the
+                   `forward` method. It is useful for cases involving compilation/tracing, e.g.
+                   ONNX export, and the forward arg is useful for dynamically changing mask types,
+                   e.g. a different mask for training and inference.
+                   1. For "`no_mask`", no attention mask is applied.
+                   2. For "`causal`", "`causal_bottom_right`", or the causal mask in
+                   "`padding_causal`" and "`padding_causal_bottom_right`", Transformer Engine
+                   calculates and applies an upper triangular mask to the softmax input.
+                   No user input is needed. Causal masks without the "`bottom_right`" appendix align
+                   the diagonal line to the top left corner of the softmax matrix. With
+                   "`bottom_right`", the causal mask is aligned to the bottom right corner, which is
+                   often used in inference/KV caching.
+                   3. For "`padding`", or the padding mask in "`padding_causal`" and
+                   "`padding_causal_bottom_right`", users need to provide the locations of padded
+                   tokens, either via :attr:`cu_seqlens_q` and :attr:`cu_seqlens_kv` (both in shape
+                   [batch_size + 1]), or via :attr:`attention_mask` (one tensor for self-attention
+                   in shape [batch_size, 1, 1, max_seqlen_q], or two tensors in a tuple for
+                   cross-attention in shapes [batch_size, 1, 1, max_seqlen_q] and
+                   [batch_size, 1, 1, max_seqlen_kv]).
+                   4. For "`arbitrary`", users need to provide a mask that is broadcastable to
+                   the shape of softmax input [batch_size, num_heads, max_seqlen_q, max_seqlen_kv].
+    window_size: Optional[Tuple[int, int]], default = `None`
+                sliding window size for local attention, where query at position i attends to keys
+                in [i + seqlen_k - seqlen_q - window_size[0], i + seqlen_k - seqlen_q
+                + window_size[1]] inclusive. Special cases (-1, -1) and (-1, 0) mean no sliding
+                window and causal mask specifically. Both `causal` and `causal_bottom_right` masks
+                map to `window_size = (-1, 0)` and Transformer Engine distinguishes them based on
+                `attn_mask_type`. Similar to :attr:`attn_mask_type`, `window_size` can
+                be overridden by :attr:`window_size` in `forward` as well.
+    attention_type: str, default = `self`
+                   type of attention, either "`self`" and "`cross`".
+    layer_number: int, default = `None`
+                 layer number of the current `DotProductAttention` when multiple such modules
+                 are concatenated, for instance in consecutive transformer blocks.
+    qkv_format: str, default = `sbhd`
+               dimension format for `query_layer`, `key_layer` and `value_layer`,
+               {`sbhd`, `bshd`, `thd`}. `s` stands for the sequence length, `b` batch size,
+               `h` the number of heads, `d` head size, and `t` the total number of tokens
+               in a batch, with `t = sum(s_i), for i = 0...b-1`. `sbhd` and `bshd` formats
+               are used for when sequences in a batch are of equal length or padded to
+               equal length, and the `thd` format is used for when sequences in a batch
+               have different lengths. Please note that these formats do not reflect how
+               tensors `query_layer`, `key_layer`, `value_layer` are laid out in memory.
+               For that, please use `get_qkv_layout` to gain the layout information.
+    softmax_scale: Optional[float], default = `None`
+                softmax scale for the attention scores. If `None`, defaults to
+                `1.0/math.sqrt(kv_channels if isinstance(kv_channels, int) else kv_channels[0])`.
+    softmax_type: str = {'vanilla', 'off-by-one', 'learnable'}, default = 'vanilla'
+                 softmax type as described in this paper:
+                 `Efficient Streaming Language Models with Attention Sinks
+                 <https://arxiv.org/pdf/2309.17453v3>`_.
+                 For a given attention score S = Q*K^T, of shape [b, h, s_q, s_kv],
+                 'vanilla': S[:,:,:,i] = exp(S[:,:,:,i])/sum(exp(S[:,:,:,:]), dim=-1),
+                 'off-by-one': S[:,:,:,i] = exp(S[:,:,:,i])/(1 + sum(exp(S[:,:,:,:]), dim=-1)), and
+                 'learnable': S[:,j,:,i] = exp(S[:,j,:,i])/(exp(alpha[j]) + sum(exp(S[:,j,:,:]), dim=-1)),
+                 where alpha is a learnable parameter in shape [h].
+                 'off-by-one' and 'learnable' softmax types are also called sink attention
+                 ('zero sink' and 'learnable sink').
+    return_max_logit: Optional[bool], default = `False`
+                     If true, returns the maximum attention score that can be used in a Muon optimizer to
+                     rescale the Q and K projection weights (see `Muon is Scalable for LLM Training
+                     <https://arxiv.org/pdf/2502.16982>`_).
+                     max_logit = max(S), where S = mask(Q*K^T*softmax_scale + bias) in shape [b, h, s_q, s_kv],
+                     and max_logit is in shape [h].
+
+    Parallelism parameters
+    ----------------------
+    sequence_parallel : bool, default = `False`
+                       if set to `True`, uses sequence parallelism.
+    tp_size : int, default = 1
+             tensor parallel world size.
+    tp_group : ProcessGroup, default = `None`
+              tensor parallel process group.
+    cp_group : Union[ProcessGroup, List[ProcessGroup]], default = `None`
+              context parallel process group.
+              ProcessGroup is for cp_comm_type of "p2p", "all_gather", and "a2a".
+              List[ProcessGroup] is for cp_comm_type of "a2a+p2p", where cp_group[0]
+              and cp_group[1] are for a2a and p2p communications respectively.
+    cp_global_ranks : list of global rank IDs, default = `None`
+                     global rank IDs of GPUs that are in cp_group.
+    cp_stream : CUDA stream, default = `None`
+               context parallelism splits flash attention into multiple steps for
+               compute and communication overlapping. To address the wave quantization
+               issue of each split step, we add an additional CUDA stream so that we
+               can overlap two flash attention kernels.
+    cp_comm_type : str, default = `p2p`
+                  inter-gpu communication type for context parallelism.
+                  Can be "p2p" or "all_gather" or "a2a" or "a2a+p2p".
+                  "p2p": Exchange KV chunks with P2P communications in ring topology.
+                         P2P is async and can be overlapped with attention compute.
+                  "all_gather": All-gather to get full sequence of KV before attention.
+                                The all-gather is not async, and cannot be overlapped.
+                  "a2a": Like DeepSpeed Ulysses, scatter attention heads across the CP
+                         group, and gather to get full sequence of QKV.
+                  "a2a+p2p": hierarchical CP implementation. First applying a2a to QKV
+                  across each CP sub-group (e.g., via NVLink), then exchanging KV with
+                  p2p between sub-groups (e.g., via IBLink).
+    """
+
+    def __init__(
+        self,
+        num_attention_heads: int,
+        kv_channels: Union[int, Tuple[int, int]],
+        num_gqa_groups: Optional[int] = None,
+        attention_dropout: float = 0.0,
+        qkv_format: str = "sbhd",
+        attn_mask_type: str = "causal",
+        window_size: Optional[Tuple[int, int]] = None,
+        sequence_parallel: bool = False,
+        tp_size: int = 1,
+        get_rng_state_tracker: Optional[Callable] = None,
+        tp_group: Optional[dist_group_type] = None,
+        layer_number: Optional[int] = None,
+        attention_type: str = "self",
+        cp_group: Optional[Union[dist_group_type, List[dist_group_type]]] = None,
+        cp_global_ranks: List[int] = None,
+        cp_stream: torch.cuda.Stream = None,
+        cp_comm_type: str = "p2p",
+        softmax_scale: Optional[float] = None,
+        softmax_type: str = "vanilla",
+        return_max_logit: Optional[bool] = False,
+    ) -> None:
+        super().__init__()
+
+        self.logger = logging.getLogger("DotProductAttention")
+        self.logger.setLevel(attn_log._log_level)
+        if not self.logger.hasHandlers():
+            self.logger.addHandler(attn_log._stream_handler)
+        self.qkv_format = qkv_format
+        attn_mask_type = attn_mask_type.replace(",", "_")
+        if attn_mask_type == "causal_padding":
+            attn_mask_type = "padding_causal"
+        self.attn_mask_type = attn_mask_type
+        self.window_size = dpa_utils.check_set_window_size(attn_mask_type, window_size)
+        if tp_group is None:
+            self.tp_size = tp_size
+            if tp_size == 1:
+                self.set_tensor_parallel_group(tp_group)
+        else:
+            self.tp_size = get_distributed_world_size(tp_group)
+            self.set_tensor_parallel_group(tp_group)
+        self.get_rng_state_tracker = get_rng_state_tracker
+        self.num_attention_heads = num_attention_heads
+        self.layer_number = 1 if layer_number is None else layer_number
+        self.cp_group = cp_group
+        self.cp_global_ranks = cp_global_ranks
+        self.cp_stream = cp_stream
+        self.cp_comm_type = cp_comm_type
+
+        self.hidden_size_per_attention_head_k = (
+            kv_channels if isinstance(kv_channels, int) else kv_channels[0]
+        )
+        self.hidden_size_per_attention_head_v = (
+            kv_channels if isinstance(kv_channels, int) else kv_channels[1]
+        )
+
+        self.num_gqa_groups = num_attention_heads if num_gqa_groups is None else num_gqa_groups
+        self.num_gqa_groups_per_partition = int(self.num_gqa_groups // self.tp_size)
+
+        assert (
+            num_attention_heads % self.num_gqa_groups == 0
+        ), "The number of attention heads must be divisible by the number of GQA groups!"
+
+        self.rng_states_tracker = None
+        if sequence_parallel or get_rng_state_tracker is None:
+            attention_dropout_ctx = nullcontext
+        else:
+            self.rng_states_tracker = get_rng_state_tracker()
+            set_all_rng_states(self.rng_states_tracker.get_states())
+            attention_dropout_ctx = self.rng_states_tracker.fork
+
+        if softmax_scale is None:
+            softmax_scale = 1.0 / math.sqrt(
+                kv_channels if isinstance(kv_channels, int) else kv_channels[0]
+            )
+
+        self.deterministic = (
+            not bool(int(os.getenv("NVTE_ALLOW_NONDETERMINISTIC_ALGO", "1")))
+            or torch.are_deterministic_algorithms_enabled()
+        )
+        # To use the workspace optimization path for determinism, please
+        # set NVTE_FUSED_ATTN_FORCE_WORKSPACE_OPT=1 for cuDNN >=8.9.5 and <9.0.0,
+        # and set NVTE_ALLOW_NONDETERMINISTIC_ALGO=0 for cuDNN >=9.0.0.
+        cudnn_version = get_cudnn_version()
+        if (8, 9, 5) <= cudnn_version < (9, 0, 0):
+            if self.deterministic:
+                os.environ["NVTE_FUSED_ATTN_FORCE_WORKSPACE_OPT"] = "1"
+
+            # CUDNN_FRONTEND_ATTN_DP_WORKSPACE_LIMIT
+            # - unset:       enables workspace optimization when required workspace is <= 256MB
+            #                or when bias gradient needs to be computed
+            # - n:           enables workspace optimization when required workspace is <= n bytes
+            # - -1:          enables workspace optimization always
+            # - 0:           disables workspace optimization always
+            if "NVTE_FUSED_ATTN_FORCE_WORKSPACE_OPT" in os.environ:
+                if os.environ["NVTE_FUSED_ATTN_FORCE_WORKSPACE_OPT"] == "0":
+                    os.environ["CUDNN_FRONTEND_ATTN_DP_WORKSPACE_LIMIT"] = "0"
+                if os.environ["NVTE_FUSED_ATTN_FORCE_WORKSPACE_OPT"] == "1":
+                    os.environ["CUDNN_FRONTEND_ATTN_DP_WORKSPACE_LIMIT"] = "-1"
+
+        assert attention_type in AttnTypes, f"attention_type {attention_type} not supported"
+
+        self.attention_type = attention_type
+        self.attention_dropout = attention_dropout
+        self.return_max_logit = return_max_logit
+
+        self.softmax_type = softmax_type
+        if self.softmax_type == "vanilla":
+            self.softmax_offset = None
+        if self.softmax_type == "off-by-one":
+            self.softmax_offset = torch.zeros(
+                self.num_attention_heads // self.tp_size, device="cuda"
+            )
+        if self.softmax_type == "learnable":
+            self.register_parameter(
+                "softmax_offset",
+                Parameter(torch.empty(self.num_attention_heads // self.tp_size, device="cuda")),
+                get_rng_state_tracker=get_rng_state_tracker,
+            )
+
+        attn_kwargs = {
+            "attention_dropout": attention_dropout,
+            "attention_dropout_ctx": attention_dropout_ctx,
+        }
+
+        self.flash_attention = FlashAttention(
+            softmax_scale,
+            attention_type=attention_type,
+            layer_number=layer_number,
+            deterministic=self.deterministic,
+            **attn_kwargs,
+        )
+
+        # Instantiating three types since use of flash-attn and FusedAttention
+        # might be ruled out due to forward inputs.
+        self.fused_attention = FusedAttention(
+            softmax_scale,
+            attention_type=attention_type,
+            layer_number=layer_number,
+            deterministic=self.deterministic,
+            **attn_kwargs,
+            softmax_type=self.softmax_type,
+            return_max_logit=self.return_max_logit,
+        )
+
+        self.unfused_attention = UnfusedDotProductAttention(
+            softmax_scale,
+            attention_type=attention_type,
+            **attn_kwargs,
+            layer_number=layer_number,
+            softmax_type=self.softmax_type,
+            return_max_logit=self.return_max_logit,
+        )
+
+        def remove_extra_states_check(self, incompatible_keys):  # pylint: disable=unused-argument
+            """
+            Temporarily remove core_attention._extra_state as a missing key
+            when loading older Transformer Engine checkpoints. Will phase out
+            this hook in Transformer Engine 2.0.
+            """
+            for key in incompatible_keys.missing_keys:
+                if "core_attention._extra_state" in key:
+                    incompatible_keys.missing_keys.remove(key)
+
+        self.register_load_state_dict_post_hook(remove_extra_states_check)
+
+    def _load_from_state_dict(
+        self, state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs
+    ):
+        """
+        This function helps to load Transformer Engine 1.6 and 1.7 checkpoints, where FP8 attention
+        metadata is stored under the `core_attention.fused_attention._extra_state` key and not the
+        `core_attention._extra_state` key. Please see `FP8 checkpoint compatibility
+        <https://docs.nvidia.com/deeplearning/transformer-engine/user-guide/faq.html#fp8-checkpoint-compatibility>`_ for more details.
+        """
+        fused_attn_key = False
+        dot_product_attn_key = False
+        for k in state_dict.keys():
+            if "core_attention.fused_attention._extra_state" in k:
+                fused_attn_key = True
+            if "core_attention._extra_state" in k:
+                dot_product_attn_key = True
+        if fused_attn_key and not dot_product_attn_key:
+            prefix = prefix + "fused_attention."
+        super()._load_from_state_dict(
+            state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs
+        )
+
+    def _checkpointed_attention_forward(
+        self,
+        attention_func: Callable,
+        *forward_args: Tuple[torch.Tensor, ...],
+        **forward_kwargs: Dict[str, Any],
+    ) -> torch.Tensor:
+        """Forward method with activation checkpointing."""
+
+        def custom_forward(*input_args, **input_kwargs):
+            return attention_func(*input_args, **input_kwargs)
+
+        hidden_states = checkpoint(
+            custom_forward,
+            distribute_saved_activations=False,
+            get_rng_state_tracker=self.get_rng_state_tracker,
+            tp_group=self.tp_group,
+            *forward_args,
+            **forward_kwargs,
+        )
+
+        return hidden_states
+
+    def set_context_parallel_group(
+        self,
+        cp_group: Union[dist_group_type, List[dist_group_type], None],
+        cp_global_ranks: List[int],
+        cp_stream: torch.cuda.Stream,
+        cp_comm_type: str = "p2p",
+    ) -> None:
+        """
+        Set the context parallel attributes for the given
+        module before executing the forward pass.
+
+        Parameters
+        ----------
+        cp_group : Union[ProcessGroup, List[ProcessGroup]]
+                  context parallel process group.
+                  ProcessGroup is for cp_comm_type of "p2p", "all_gather", and "a2a".
+                  List[ProcessGroup] is for cp_comm_type of "a2a+p2p", where cp_group[0]
+                  and cp_group[1] are for a2a and p2p communications respectively.
+        cp_global_ranks : List[int]
+                         list of global ranks in the context group.
+        cp_stream : torch.cuda.Stream
+                   cuda stream for context parallel execution.
+        cp_comm_type : str, default = `p2p`
+                      inter-gpu communication type for context parallelism.
+                      Can be "p2p" or "all_gather" or "a2a" or "a2a+p2p".
+                      "p2p": Exchange KV chunks with P2P communications in ring topology.
+                             P2P is async and can be overlapped with attention compute.
+                      "all_gather": All-gather to get full sequence of KV before attention.
+                                    The all-gather is not async, and cannot be overlapped.
+                      "a2a": Like DeepSpeed Ulysses, scatter attention heads across the CP
+                             group, and gather to get full sequence of QKV.
+                      "a2a+p2p": hierarchical CP implementation. First applying a2a to QKV
+                      across each CP sub-group (e.g., via NVLink), then exchanging KV with
+                      p2p between sub-groups (e.g., via IBLink).
+        """
+        self.cp_group = cp_group
+        self.cp_global_ranks = cp_global_ranks
+        self.cp_stream = cp_stream
+        self.cp_comm_type = cp_comm_type
+
+    def init_fp8_metadata(self, num_gemms: int = 1) -> None:
+        """
+        Override TransformerEngineBaseModule.init_fp8_metadata to allow for more flexible recipe support.
+        Initialize fp8 related metadata and tensors during fprop.
+        """
+        _original_recipe = self.fp8_meta.get("recipe", None)
+
+        # global recipe set in autocast()
+        fp8_recipe = FP8GlobalStateManager.get_fp8_recipe()
+        if fp8_recipe.custom():
+            return
+
+        # switch/append recipe: fp8_recipe stays unchanged, but DPA.fp8_meta["recipe"] may be set to
+        # a different recipe than fp8_recipe. DPA.quantizers may be a mix of different quantizers as well.
+        #
+        # fp8_recipe                | NVTE_DPA_FP8_RECIPE | self.fp8_meta["recipe"] | self.quantizers
+        # --------------------------------------------------------------------------------------------
+        # DelayedScaling (DS)       | unset               | DS                      | all DS
+        # Float8CurrentScaling (CS) | unset               | DS                      | CS for QKV, O, dO, dQKV; DS for S, dP
+        # x={DS, CS}                | y                   | refer to row x=y        | refer to row x=y
+        fp8_recipe_dpa = fp8_recipe
+        fp8_recipes = fp8_recipe
+        if _dpa_fp8_recipe == "F16":
+            # ignore the recipe from autocast, set fp8_dpa = False, fp8_mha = False
+            fp8_recipe.fp8_dpa = False
+            fp8_recipe.fp8_mha = False
+        elif fp8_recipe.float8_current_scaling() and _dpa_fp8_recipe == "DelayedScaling":
+            # reuse fp8_format, fp8_dpa, fp8_mha from fp8_recipe, and construct a DS recipe
+            fake_recipe = DelayedScaling(
+                fp8_format=fp8_recipe.fp8_format,
+                amax_history_len=_dpa_fp8ds_amax_histlen,
+                amax_compute_algo=_dpa_fp8ds_amax_algo,
+                fp8_dpa=fp8_recipe.fp8_dpa,
+                fp8_mha=fp8_recipe.fp8_mha,
+                reduce_amax=_dpa_fp8ds_reduce_amax,
+            )
+            fp8_recipe_dpa = fake_recipe
+            fp8_recipes = fp8_recipe_dpa
+        elif fp8_recipe.nvfp4() and _dpa_fp8_recipe == "DelayedScaling":
+            # reuse fp8_dpa, fp8_mha from fp8_recipe but not fp8_format; construct a DS recipe
+            fake_recipe = DelayedScaling(
+                fp8_format=_dpa_fp8_format,
+                amax_history_len=_dpa_fp8ds_amax_histlen,
+                amax_compute_algo=_dpa_fp8ds_amax_algo,
+                fp8_dpa=fp8_recipe.fp8_dpa,
+                fp8_mha=fp8_recipe.fp8_mha,
+                reduce_amax=_dpa_fp8ds_reduce_amax,
+            )
+            fp8_recipe_dpa = fake_recipe
+            fp8_recipes = fp8_recipe_dpa
+        elif fp8_recipe.delayed() and _dpa_fp8_recipe == "Float8CurrentScaling":
+            # reuse fp8_format, fp8_dpa, fp8_mha from fp8_recipe, and construct a CS+DS recipe
+            fake_recipes = [
+                Float8CurrentScaling(
+                    fp8_format=fp8_recipe.fp8_format,
+                    fp8_dpa=fp8_recipe.fp8_dpa,
+                    fp8_mha=fp8_recipe.fp8_mha,
+                ),
+                fp8_recipe,
+            ]
+            fp8_recipe_dpa = fake_recipes[1]
+            fp8_recipes = fake_recipes
+        elif (
+            fp8_recipe.float8_current_scaling()
+            and _dpa_fp8_recipe in ("", "Float8CurrentScaling")
+            and (fp8_recipe.fp8_dpa or fp8_recipe.fp8_mha)
+        ):
+            # use fp8_recipe for QKV, O, dO, dQKV, and construct a DS recipe for S, dP
+            # reuse fp8_format, fp8_dpa, fp8_mha from fp8_recipe
+            fake_recipe = DelayedScaling(
+                fp8_format=fp8_recipe.fp8_format,
+                amax_history_len=_dpa_fp8ds_amax_histlen,
+                amax_compute_algo=_dpa_fp8ds_amax_algo,
+                fp8_dpa=fp8_recipe.fp8_dpa,
+                fp8_mha=fp8_recipe.fp8_mha,
+                reduce_amax=_dpa_fp8ds_reduce_amax,
+            )
+            fp8_recipe_dpa = fake_recipe
+            fp8_recipes = [fp8_recipe, fp8_recipe_dpa]
+        elif fp8_recipe.nvfp4() and _dpa_fp8_recipe == "Float8CurrentScaling":
+            # reuse fp8_dpa, fp8_mha from fp8_recipe but not fp8_format
+            # construct a CS recipe for QKV, O, dO, dQKV and a DS recipe for S, dP
+            fake_recipes = [
+                Float8CurrentScaling(
+                    fp8_format=_dpa_fp8_format,
+                    fp8_dpa=fp8_recipe.fp8_dpa,
+                    fp8_mha=fp8_recipe.fp8_mha,
+                ),
+                DelayedScaling(
+                    fp8_format=_dpa_fp8_format,
+                    amax_history_len=_dpa_fp8ds_amax_histlen,
+                    amax_compute_algo=_dpa_fp8ds_amax_algo,
+                    fp8_dpa=fp8_recipe.fp8_dpa,
+                    fp8_mha=fp8_recipe.fp8_mha,
+                    reduce_amax=_dpa_fp8ds_reduce_amax,
+                ),
+            ]
+            fp8_recipe_dpa = fake_recipes[1]
+            fp8_recipes = fake_recipes
+        # DPA only support DS and CS; other recipes should have fp8_dpa=False, fp8_mha=False
+        if not fp8_recipe_dpa.float8_per_tensor_scaling():
+            assert not (
+                fp8_recipe_dpa.fp8_dpa or fp8_recipe_dpa.fp8_mha
+            ), f"DotProductAttention does not support {fp8_recipe_dpa.__class__.__name__} recipe"
+
+        # reduce over TP+CP groups; expect fp8_group to be set up so
+        # assume attention uses the same fp8_group as GEMMs
+        fp8_group = FP8GlobalStateManager.get_fp8_group()
+
+        self.fp8_parameters = FP8GlobalStateManager.with_fp8_parameters()
+        self.fp8 = FP8GlobalStateManager.is_fp8_enabled()
+        self.fp8_calibration = FP8GlobalStateManager.is_fp8_calibration()
+        fp8_enabled = self.fp8 or self.fp8_calibration
+        self.fp8_meta["fp8_checkpoint"] = self.fp8 or self.fp8_calibration
+        if self.fp8_parameters or fp8_enabled:
+            self.fp8_meta["global_recipe"] = fp8_recipe
+            self.fp8_meta["local_recipes"] = (
+                fp8_recipes if isinstance(fp8_recipes, List) else [fp8_recipes]
+            )
+
+        if self.fp8_parameters or fp8_enabled:
+            if self.fp8_initialized and fp8_recipe_dpa == self.fp8_meta["recipe"]:
+                # FP8 init has already been run and recipe is the same, don't do anything.
+                return
+            self.fp8_meta["recipe"] = fp8_recipe_dpa
+            if fp8_recipe != fp8_recipe_dpa:
+                # fp8_recipe has changed, rehash the key.
+                autocast_key = FP8GlobalStateManager.get_unique_autocast_key(
+                    fp8_recipe_dpa, fp8_group
+                )
+                FP8GlobalStateManager.autocast_arguments[autocast_key] = (
+                    fp8_recipe_dpa,
+                    fp8_group,
+                )
+        else:
+            # If fp8 isn't enabled, turn off and return.
+            self.fp8_initialized = False
+            return
+
+        if self.fp8_parameters and not self.fp8_initialized:
+            self.fp8_meta["num_gemms"] = num_gemms
+            self.init_fp8_meta_tensors(fp8_recipes)
+
+        if fp8_enabled:
+            # Set FP8 and other FP8 metadata
+            self.fp8_meta["num_gemms"] = num_gemms
+            self.fp8_meta["fp8_group"] = fp8_group
+
+            # Set FP8_MAX per tensor according to recipe
+            self.fp8_meta["fp8_max_fwd"] = self.fp8_meta["recipe"].fp8_format.value.max_fwd
+            self.fp8_meta["fp8_max_bwd"] = self.fp8_meta["recipe"].fp8_format.value.max_bwd
+
+            # Allocate scales and amaxes
+            self.init_fp8_meta_tensors(fp8_recipes)
+            self.fp8_initialized = True
+
+            self.fp8_meta["recipe"] = fp8_recipe_dpa
+            if fp8_recipe != fp8_recipe_dpa:
+                # fp8_recipe has changed, rehash the key.
+                autocast_key = FP8GlobalStateManager.get_unique_autocast_key(
+                    fp8_recipe_dpa, fp8_group
+                )
+                FP8GlobalStateManager.autocast_arguments[autocast_key] = (
+                    fp8_recipe_dpa,
+                    fp8_group,
+                )
+
+        _current_recipe = self.fp8_meta["recipe"]
+        if _original_recipe is not None and not (
+            issubclass(_current_recipe.__class__, _original_recipe.__class__)
+            or issubclass(_original_recipe.__class__, _current_recipe.__class__)
+        ):
+            warnings.warn(
+                f"Recipe type changed from {_original_recipe.__class__.__name__} "
+                f"to {_current_recipe.__class__.__name__}. "
+                "This may affect model behavior."
+            )
+            # Clear cached workspaces as they were created with the old recipe/quantizer type
+            self._fp8_workspaces.clear()
+
+    def set_meta_tensor(self, fwd: bool, recipe: Union[Recipe, List[Recipe]]) -> None:
+        """Override to allow multiple recipes. Init scales and amaxes for fwd | bwd."""
+        if isinstance(recipe, Recipe):
+            recipe = [recipe]
+        fp8_recipe_dpa = recipe[-1]
+        fp8_meta_tensor_key = "scaling_fwd" if fwd else "scaling_bwd"
+
+        # Return early if recipe state matches recipe
+        if self.fp8_meta_tensors_initialized:
+            recipe_state = self.fp8_meta[fp8_meta_tensor_key]
+            if fp8_recipe_dpa.delayed() and isinstance(recipe_state, DelayedScalingRecipeState):
+                self.adjust_amax_history_length(fp8_recipe_dpa.amax_history_len, fwd=fwd)
+                return
+            if fp8_recipe_dpa.mxfp8() and isinstance(recipe_state, MXFP8BlockScalingRecipeState):
+                return
+            if fp8_recipe_dpa.float8_current_scaling() and isinstance(
+                recipe_state, Float8CurrentScalingRecipeState
+            ):
+                return
+            if fp8_recipe_dpa.float8_block_scaling() and isinstance(
+                recipe_state, Float8BlockScalingRecipeState
+            ):
+                return
+
+        # When fp8_recipe=Float8CurrentScaling, recipe=[CS, DS], and QKV/dQKV, O/dO use CS quantizers, S/dP use DS quantizers.
+        # See table above in init_fp8_metadata for more detail.
+        num_gemms = [2, 1] if len(recipe) == 2 else [3]
+        # Max. number of fp8 tensors per GEMM = 3 (input, weight, output) for fwd and
+        # 2 (grad_output and grad_input) for bwd
+        num_fp8_tensors = [x * 3 if fwd else x * 2 for x in num_gemms]
+
+        # Initialize recipe state and quantizers
+        recipe_states = [
+            RecipeState.create(
+                recipe[i],
+                mode=("forward" if fwd else "backward"),
+                num_quantizers=num_fp8_tensors[i],
+            )
+            for i in range(len(recipe))
+        ]
+
+        self.fp8_meta[fp8_meta_tensor_key] = (
+            recipe_states[-1] if len(recipe) == 2 else recipe_states[0]
+        )
+        self.quantizers[fp8_meta_tensor_key] = []
+        for recipe_state in recipe_states:
+            self.quantizers[fp8_meta_tensor_key].extend(recipe_state.make_quantizers())
+
+    @no_torch_dynamo(recursive=False)
+    def forward(
+        self,
+        query_layer: torch.Tensor,
+        key_layer: torch.Tensor,
+        value_layer: torch.Tensor,
+        attention_mask: Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]] = None,
+        qkv_format: str = None,
+        cu_seqlens_q: torch.Tensor = None,
+        cu_seqlens_kv: torch.Tensor = None,
+        cu_seqlens_q_padded: torch.Tensor = None,
+        cu_seqlens_kv_padded: torch.Tensor = None,
+        max_seqlen_q: int = None,
+        max_seqlen_kv: int = None,
+        attn_mask_type: Optional[str] = None,
+        window_size: Optional[Tuple[int, int]] = None,
+        checkpoint_core_attention: bool = False,
+        core_attention_bias_type: str = "no_bias",
+        core_attention_bias: Optional[torch.Tensor] = None,
+        alibi_slopes: Optional[torch.Tensor] = None,
+        fast_zero_fill: bool = True,
+        inference_params: Optional[InferenceParams] = None,
+        pad_between_seqs: Optional[bool] = None,
+        fp8_output: Optional[bool] = False,
+        num_splits: Optional[int] = 1,
+    ) -> torch.Tensor:
+        """
+        Dot Product Attention Layer.
+
+        .. note::
+
+            Argument :attr:`attention_mask` is only used when :attr:`attn_mask_type`
+            includes '"padding"' or `"arbitrary"`.
+
+        .. note::
+
+            DotProductAttention supports three backends: 1) FlashAttention which calls
+            HazyResearch/Dao-AILab's `flash-attn <https://arxiv.org/pdf/2305.13245.pdf>`_
+            PyTorch API, 2) FusedAttention which has multiple fused attention implementations
+            based on `cuDNN Graph API
+            <https://docs.nvidia.com/deeplearning/cudnn/developer-guide/index.html#op-fusion>`_
+            (see :attr:`FusedAttention` for more details on FusedAttention backends), and 3)
+            UnfusedDotProductAttention which is the native PyTorch implementation
+            with fused scaled masked softmax.
+
+        .. note::
+
+            Users can use environment variables :attr:`NVTE_FLASH_ATTN`, :attr:`NVTE_FUSED_ATTN`,
+            and :attr:`NVTE_FUSED_ATTN_BACKEND` to control which DotProductAttention backend,
+            and FusedAttention backend if applicable, to use. Transformer Engine prioritizes
+            FlashAttention over FusedAttention and over UnfusedDotProductAttention.
+            If FusedAttention is being used, users can also choose to switch to flash-attn's
+            implementation for backward by setting :attr:`NVTE_FUSED_ATTN_USE_FAv2_BWD=1`
+            (default: 0), because of the performance differences between various versions of
+            flash-attn and FusedAttention. Further, :attr:`NVTE_FUSED_ATTN_FORCE_WORKSPACE_OPT`
+            can be used to enable (:attr:`1`) or disable (:attr:`0`) the workspace related
+            optimizations in FusedAttention. When unset, Transformer Engine determines the code path
+            based on its internal logic. These optimizations trade memory for performance
+            and should be used with care.
+
+        .. note::
+            .. _cu_seqlens note:
+
+            When training data has variable sequence lengths, users have two options.
+
+            1. Manipulate the data and pad all sequences to the same length. Use
+               :attr:`qkv_format` = {"bshd", "sbhd"} and
+               :attr:`attn_mask_type` = {"padding", "padding_causal", "padding_causal_bottom_right"}.
+               Pass in :attr:`cu_seqlens_q` and :attr:`cu_seqlens_kv`, or :attr:`attention_mask`
+               (which will be converted to :attr:`cu_seqlens_q` and :attr:`cu_seqlens_kv`), to provide
+               the real sequence length information. For example, a batch of 3 sequences
+               [a a a b b c c c c] can be padded to [a a a PAD b b PAD PAD c c c c], and the cumulative
+               sequence length tensors would be
+               :attr:`cu_seqlens_q` = :attr:`cu_seqlens_kv` = [0, 3, 5, 9] for self-attention.
+
+            2. Do not perform padding on training data. Use :attr:`qkv_format` = "thd" and
+               :attr:`attn_mask_type` = {"padding", "padding_causal", "padding_causal_bottom_right"}.
+               Pass in :attr:`cu_seqlens_q` and :attr:`cu_seqlens_kv`, or :attr:`attention_mask`,
+               as in option 1. For example, a batch of 3 sequences [a a a b b c c c c] can be processed
+               without any padding, and the sequence length tensors would be
+               :attr:`cu_seqlens_q` = :attr:`cu_seqlens_kv` = [0, 3, 5, 9] for self-attention.
+
+               In certain use cases, a varying number of identifier tokens are inserted between
+               sequences. These tokens do not participate in the attention calculation.
+               :attr:`cu_seqlens_q_padded` and :attr:`cu_seqlens_kv_padded` must be specified
+               in such cases to correctly identify the start and end of each sequence in a batch.
+               For example, a batch of 3 sequences [a a a 1 b b 2 2 c c c c 3] would have
+               :attr:`cu_seqlens_q` = :attr:`cu_seqlens_kv` = [0, 3, 5, 9], and
+               :attr:`cu_seqlens_q_padded` = :attr:`cu_seqlens_kv_padded` = [0, 4, 8, 13]
+               for self-attention.
+
+        .. note::
+            .. _max_seqlen note:
+
+            When :attr:`qkv_format` = {"bshd", "sbhd"}, sequences are of equal length in a batch.
+            :attr:`max_seqlen_q` and :attr:`max_seqlen_kv` should be the same as the "s" dimension of
+            :attr:`query_layer` and :attr:`key_layer` tensors. When unset, Transformer Engine will
+            infer them as such.
+
+            When :attr:`qkv_format` = "thd", sequences have varying lengths. :attr:`max_seqlen_q` and
+            :attr:`max_seqlen_kv` should be the maximum query and key/value sequence length in a batch.
+            When unset, Transformer Engine deduces them from :attr:`cu_seqlens_q` and :attr:`cu_seqlens_kv`.
+            This deduction costs a small kernel and some CPU-GPU synchronization, and to avoid this
+            overhead, users are recommended to obtain the maximum sequence lengths from the data loaders
+            and pass them in.
+
+            - As the maximum sequence lengths, batch size, and number of tokens change from batch to batch,
+              dynamic shapes need to be supported for tensor construction. FlashAttention and
+              UnfusedDotProductAttention naturally do so, while FusedAttention requires parameters to be static
+              to create graphs before performance heuristics analysis. To reduce the number of graphs created
+              per run, Transformer Engine 1.13+ quantizes relevant parameters: for cuDNN < 9.6, {batch size,
+              :attr:`max_seqlen_q`, :attr:`max_seqlen_kv`}, and for cuDNN >= 9.6, {"t" dimension of
+              :attr:`query_layer`, "t" dimension of :attr:`key_layer`}.
+
+        Parameters
+        ----------
+        query_layer : torch.Tensor
+                     Query tensor.
+        key_layer : torch.Tensor
+                   Key tensor.
+        value_layer : torch.Tensor
+                     Value tensor.
+        attention_mask: Optional[Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]],
+             default = `None`. Boolean tensor(s) used to mask out attention softmax input.
+             It should be `None` for causal masks and "`no_mask`". For padding masks, it should be
+             a single tensor of [batch_size, 1, 1, seqlen_q] for self-attention, and a tuple of
+             two tensors in shapes [batch_size, 1, 1, seqlen_q] and [batch_size, 1, 1, seqlen_kv]
+             for cross-attention. For "`arbitrary`" mask, it should be in a shape broadcastable
+             to [batch_size, num_heads, max_seqlen_q, max_seqlen_kv]. A `True` value means
+             the corresponding position is masked out and a `False` means that position
+             is allowed to participate in attention.
+        qkv_format: str, default = `None`
+                   If provided, overrides :attr:`qkv_format` from initialization.
+        cu_seqlens_q: Optional[torch.Tensor], default = `None`
+                   Cumulative sum of sequence lengths (without offset) in a batch for `query_layer`,
+                   with shape [batch_size + 1] and dtype torch.int32.
+                   See :ref:`note<cu_seqlens note>` for more details.
+        cu_seqlens_kv: Optional[torch.Tensor], default = `None`
+                   Cumulative sum of sequence lengths (without offset) in a batch for `key_layer`
+                   and `value_layer`, with shape [batch_size + 1] and dtype torch.int32.
+                   See :ref:`note<cu_seqlens note>` for more details.
+        cu_seqlens_q_padded: Optional[torch.Tensor], default = `None`
+                   Cumulative sum of sequence lengths (with offset) in a batch for
+                   `query_layer`, with shape [batch_size + 1] and dtype torch.int32.
+                   When there is no padding between sequences in a batch,
+                   `cu_seqlens_q_padded = cu_seqlens_q`.
+                   See :ref:`note<cu_seqlens note>` for more details.
+        cu_seqlens_kv_padded: Optional[torch.Tensor], default = `None`
+                   Cumulative sum of sequence lengths (with offset) in a batch for `key_layer`
+                   and `value_layer`, with shape [batch_size + 1] and dtype torch.int32.
+                   When there is no padding between sequences in a batch,
+                   `cu_seqlens_kv_padded = cu_seqlens_kv`.
+                   See :ref:`note<cu_seqlens note>` for more details.
+        max_seqlen_q: Optional[int], default = `None`
+                      Maximum sequence length in `query_layer`.
+                      See :ref:`note<max_seqlen note>` for more details.
+        max_seqlen_kv: Optional[int], default = `None`
+                       Maximum sequence length in `key_layer` and `value_layer`.
+                       See :ref:`note<max_seqlen note>` for more details.
+        attn_mask_type: {'no_mask', 'padding', 'causal', 'padding,causal', 'causal,padding',
+                       'padding_causal', 'causal_bottom_right', 'padding_causal_bottom_right',
+                       'arbitrary'}, default = `None`. Type of attention mask passed into
+                       softmax operation. 'padding,causal', 'causal,padding' and 'padding_causal'
+                       are equivalent. By default, causal masks are aligned to the top left corner
+                       of the softmax matrix. When "`bottom_right`" is specified in the mask type,
+                       causal masks are aligned to the bottom right corner.
+        window_size: Optional[Tuple[int, int]], default = `None`
+                    Sliding window size for local attention.
+        checkpoint_core_attention : bool, default = `False`
+                                   If true, forward activations for attention are recomputed
+                                   during the backward pass in order to save memory that would
+                                   otherwise be occupied to store the forward activations until
+                                   backprop.
+        core_attention_bias_type: str, default = `no_bias`
+                    Bias type, {`no_bias`, `pre_scale_bias`, `post_scale_bias`, `alibi`}
+        core_attention_bias: Optional[torch.Tensor], default = `None`
+                    Bias tensor for Q * K.T, shape [1, num_head, max_seqlen_q, max_seqlen_kv].
+                    It should be 'None' for 'no_bias' and 'alibi' bias types.
+        alibi_slopes: Optional[torch.Tensor], default = `None`
+                     ALiBi slopes in FP32 and shape [nheads] or [batch_size, nheads].
+                     It adds a bias of (-alibi_slope * (i + seqlen_k - seqlen_q - j))
+                     to the attention score of query i and key j.
+        fast_zero_fill: bool, default = `True`
+                    Whether to use the fast path to set output tensors to 0 or not.
+        inference_params: Optional[InferenceParams], default = `None`
+            Optimizes execution performance during inference by caching Keys and Values of the
+            current decoding iteration. These cached values are appended to the K and V values
+            computed in previous iterations, eliminating the need to recalculate them for the
+            entire sequence.
+            Initialization of `inference_params` is required prior to use to ensure sufficient
+            memory allocation.
+            Adjustments of the sequence_len_offset should be done after a complete forward pass.
+            If rotary positional embeddings (RoPE) are utilized, they must be prepared beforehand.
+            Supports "sbhd" and "bshd" layouts, with the "sbhd" layout being more efficient.
+        pad_between_seqs: Optional[bool], default = `None`
+            If None, inferred from qkv_format, cu_seqlens and cu_seqlens_padded.
+            If true, there are padding tokens between individual sequences in a packed batch.
+        fp8_output: Optional[bool], default = `False`
+            Whether to enforce output to be in FP8 or not.
+        num_splits: Optional[int], default = `1`
+            Number of splits for FlashAttention-3 forward kernel.
+            Setting to 1 ensures deterministic backward across TP ranks.
+            Only supported with FlashAttention-3 backend.
+        """
+
+        with torch.cuda.device(query_layer.device), self.prepare_forward(
+            query_layer,
+            num_gemms=3,
+            allow_non_contiguous=True,
+            allow_different_data_and_param_types=self.softmax_type != "vanilla",
+        ) as query_layer:
+            # checks for RNG
+            if self.rng_states_tracker is not None and is_graph_capturing():
+                assert isinstance(
+                    self.rng_states_tracker, CudaRNGStatesTracker
+                ), "Unsupported RNG states tracker."
+                assert (
+                    graph_safe_rng_available()
+                ), "Upgrade PyTorch version to get RNG manipulation support for cuda graph capture."
+
+            # checks for FP8
+            if self.fp8:
+                if self.fp8_meta["recipe"].fp8_mha:
+                    if not self.fp8_meta["recipe"].fp8_dpa:
+                        self.fp8_meta["recipe"].fp8_dpa = True
+                        self.logger.warning(
+                            """Forcing fp8_meta["recipe"].fp8_dpa=True due to """
+                            """fp8_meta["recipe"].fp8_mha=True"""
+                        )
+            if self.fp8 and self.fp8_meta["recipe"].fp8_dpa:
+                forward_dtype = get_fp8_te_dtype(self.fp8_meta["recipe"], fprop_tensor=True)
+                backward_dtype = get_fp8_te_dtype(self.fp8_meta["recipe"], fprop_tensor=False)
+                assert forward_dtype in [
+                    tex.DType.kFloat8E4M3,
+                    tex.DType.kFloat8E5M2,
+                ] and backward_dtype in [
+                    tex.DType.kFloat8E4M3,
+                    tex.DType.kFloat8E5M2,
+                ], """DotProductAttention only supports "E4M3" and "E5M2" FP8 data types."""
+            else:
+                fp8_output = False
+
+            # checks for q/k/v shapes
+            assert (
+                query_layer.is_cuda and key_layer.is_cuda and value_layer.is_cuda
+            ), "DotProductAttention only supports CUDA tensors."
+            assert (
+                query_layer.dtype == key_layer.dtype and query_layer.dtype == value_layer.dtype
+            ), "Queries, keys and values must have the same data type!"
+            assert (
+                key_layer.shape[:-1] == value_layer.shape[:-1]
+            ), "Keys and values must have the same batch size, sequence length and number of heads!"
+            num_attention_heads = query_layer.shape[-2]
+            num_gqa_groups = key_layer.shape[-2]
+            assert (
+                query_layer.shape[-1] == key_layer.shape[-1]
+            ), "Queries and keys must have the same head dimension!"
+            head_dim_qk, head_dim_v = query_layer.shape[-1], value_layer.shape[-1]
+            assert (
+                head_dim_qk == self.hidden_size_per_attention_head_k
+            ), f"Keys have head_dim = {head_dim_qk}, "
+            "but expected head_dim = {self.hidden_size_per_attention_head_k}!"
+            assert (
+                head_dim_v == self.hidden_size_per_attention_head_v
+            ), f"Values have head_dim = {head_dim_v}, "
+            "but expected head_dim = {self.hidden_size_per_attention_head_v}!"
+            assert num_gqa_groups == self.num_gqa_groups_per_partition, (
+                "Keys and values must have num_gqa_group ="
+                f" {self.num_gqa_groups_per_partition} heads! Found {num_gqa_groups}."
+            )
+
+            # checks for attention mask
+            if attn_mask_type is None:
+                attn_mask_type = self.attn_mask_type
+            else:
+                attn_mask_type = attn_mask_type.replace(",", "_")
+                if attn_mask_type == "causal_padding":
+                    attn_mask_type = "padding_causal"
+            assert (
+                attn_mask_type in AttnMaskTypes
+            ), f"Attention mask type {attn_mask_type} is not supported!"
+
+            # checks for sliding window
+            if window_size is None:
+                window_size = self.window_size
+            window_size = dpa_utils.check_set_window_size(attn_mask_type, window_size)
+
+            # checks for qkv_format
+            if qkv_format is None:
+                qkv_format = self.qkv_format
+            assert qkv_format in [
+                "sbhd",
+                "bshd",
+                "thd",
+            ], "DotProductAttention only supports qkv_format = {'sbhd', 'bshd', 'thd'}!"
+            batch_size = None
+            if qkv_format in ["sbhd", "bshd"]:
+                assert all(
+                    len(x.shape) == 4 for x in (query_layer, key_layer, value_layer)
+                ), f"Queries, keys and values must be 4D tensors when {qkv_format=}!"
+                if qkv_format == "sbhd":
+                    batch_size = query_layer.shape[1]
+                    max_seqlen_q = query_layer.shape[0] if max_seqlen_q is None else max_seqlen_q
+                    max_seqlen_kv = key_layer.shape[0] if max_seqlen_kv is None else max_seqlen_kv
+                else:
+                    batch_size = query_layer.shape[0]
+                    max_seqlen_q = query_layer.shape[1] if max_seqlen_q is None else max_seqlen_q
+                    max_seqlen_kv = key_layer.shape[1] if max_seqlen_kv is None else max_seqlen_kv
+            if qkv_format == "thd":
+                assert all(
+                    len(x.shape) == 3 for x in (query_layer, key_layer, value_layer)
+                ), "Queries, keys and values must be 3D tensors when qkv_format = thd!"
+                assert (
+                    "padding" in attn_mask_type
+                ), "Attention mask type must be padding or padding_causal for qkv_format=thd!"
+                assert (
+                    cu_seqlens_q is not None and cu_seqlens_kv is not None
+                ), "cu_seqlens_q and cu_seqlens_kv can not be None when qkv_format = thd!"
+                assert (
+                    cu_seqlens_q.shape == cu_seqlens_kv.shape
+                    and len(cu_seqlens_q.shape) == 1
+                    and len(cu_seqlens_kv.shape) == 1
+                ), "cu_seqlens_q and cu_seqlens_q must both have shape [batch_size + 1]!"
+                assert (
+                    cu_seqlens_q.dtype == torch.int32 and cu_seqlens_kv.dtype == torch.int32
+                ), "cu_seqlens_q and cu_seqlens_q must both be in dtype torch.int32!"
+                batch_size = len(cu_seqlens_q) - 1
+                if max_seqlen_q is None:
+                    if cu_seqlens_q_padded is not None:
+                        seqlens_q = cu_seqlens_q_padded[1:] - cu_seqlens_q_padded[:-1]
+                    else:
+                        seqlens_q = cu_seqlens_q[1:] - cu_seqlens_q[:-1]
+                    max_seqlen_q = int((seqlens_q.max().item() + 63) // 64 * 64)
+                if max_seqlen_kv is None:
+                    if cu_seqlens_kv_padded is not None:
+                        seqlens_kv = cu_seqlens_kv_padded[1:] - cu_seqlens_kv_padded[:-1]
+                    else:
+                        seqlens_kv = cu_seqlens_kv[1:] - cu_seqlens_kv[:-1]
+                    max_seqlen_kv = int((seqlens_kv.max().item() + 63) // 64 * 64)
+
+            # update KV cache and retrieve saved tokens from cache for inference
+            if inference_params is not None:
+                assert self.layer_number is not None, "Layer number must be set!"
+
+                # convert top-left causal to bottom-right causal due to KV caching
+                # users can still use the same attention mask for inference as for training
+                assert "padding" in attn_mask_type, "KV caching requires padding mask!"
+                if attn_mask_type == "padding_causal":
+                    attn_mask_type = attn_mask_type + "_bottom_right"
+
+                self.attention_type = "cross"
+                self.flash_attention.attention_type = self.attention_type
+                self.fused_attention.attention_type = self.attention_type
+                self.unfused_attention.attention_type = self.attention_type
+
+                query_layer, key_layer, value_layer = [
+                    x.contiguous() if not x.is_contiguous() else x
+                    for x in [query_layer, key_layer, value_layer]
+                ]
+
+                # get full K/V tensors from cache and adjust cu_seqlens, qkv_format based on the cache
+                (
+                    key_layer,
+                    value_layer,
+                    cu_seqlens_q,
+                    cu_seqlens_kv,
+                    max_seqlen_kv,
+                    qkv_format,
+                ) = inference_params.step(
+                    self.layer_number,
+                    key_layer,
+                    value_layer,
+                    qkv_format,
+                )
+                cu_seqlens_q_padded = None
+                cu_seqlens_kv_padded = None
+
+            # get qkv's memory layout
+            if all(isinstance(x, Float8Tensor) for x in [query_layer, key_layer, value_layer]):
+                (
+                    qkv_layout,
+                    query_layer._data,
+                    key_layer._data,
+                    value_layer._data,
+                    q_format,
+                    kv_format,
+                ) = dpa_utils.get_qkv_layout(
+                    query_layer._data,
+                    key_layer._data,
+                    value_layer._data,
+                    qkv_format=qkv_format,
+                    inference_params=inference_params,
+                )
+            else:
+                (
+                    qkv_layout,
+                    query_layer,
+                    key_layer,
+                    value_layer,
+                    q_format,
+                    kv_format,
+                ) = dpa_utils.get_qkv_layout(
+                    query_layer,
+                    key_layer,
+                    value_layer,
+                    qkv_format=qkv_format,
+                    inference_params=inference_params,
+                )
+
+            # adjust max_seqlen and cu_seqlens for CP
+            cp_size = 1
+            if isinstance(self.cp_group, dist_group_type):
+                cp_size = get_distributed_world_size(self.cp_group)
+            elif isinstance(self.cp_group, list):
+                for group in self.cp_group:
+                    cp_size *= get_distributed_world_size(group)
+            context_parallel = cp_size > 1
+            if q_format in ["sbhd", "bshd"]:
+                max_seqlen_q *= cp_size
+                if cu_seqlens_q is None:
+                    if "padding" in attn_mask_type:
+                        assert (
+                            attention_mask is not None
+                        ), "Please provide attention_mask for padding!"
+                        if self.attention_type == "self":
+                            cu_seqlens_q = dpa_utils.get_cu_seqlens(attention_mask)
+                        else:
+                            cu_seqlens_q = dpa_utils.get_cu_seqlens(attention_mask[0])
+                    else:
+                        cu_seqlens_q = dpa_utils.get_full_cu_seqlens(
+                            batch_size,
+                            max_seqlen_q,
+                            query_layer.device,
+                        )
+            if kv_format in ["sbhd", "bshd"]:
+                max_seqlen_kv *= cp_size
+                if cu_seqlens_kv is None:
+                    if "padding" in attn_mask_type:
+                        assert (
+                            attention_mask is not None
+                        ), "Please provide attention_mask for padding!"
+                        if self.attention_type == "self":
+                            cu_seqlens_kv = dpa_utils.get_cu_seqlens(attention_mask)
+                        else:
+                            cu_seqlens_kv = dpa_utils.get_cu_seqlens(attention_mask[1])
+                    else:
+                        cu_seqlens_kv = dpa_utils.get_full_cu_seqlens(
+                            batch_size,
+                            max_seqlen_kv,
+                            key_layer.device,
+                        )
+
+            # set ALiBi attributes
+            global _alibi_cache
+            if alibi_slopes is not None:
+                assert (
+                    core_attention_bias_type == "alibi"
+                ), "core_attention_bias_type must be alibi in order to use alibi_slopes!"
+                if self.layer_number == 1:
+                    _alibi_cache["_alibi_slopes_require_update"] = True
+                    _alibi_cache["_alibi_bias_require_update"] = True
+            bottom_right_alignment = (attn_mask_type not in ["causal", "padding_causal"],)
+            if core_attention_bias_type == "alibi":
+                assert (
+                    core_attention_bias is None
+                ), "core_attention_bias must be None when core_attention_bias_type is alibi!"
+                if (
+                    _alibi_cache["_num_heads"] != query_layer.shape[-2]
+                    or _alibi_cache["_max_seqlen_q"] != max_seqlen_q
+                    or _alibi_cache["_max_seqlen_kv"] != max_seqlen_kv
+                    or _alibi_cache["_bottom_right_alignment"] != bottom_right_alignment
+                    or _alibi_cache["_alibi_slopes"] is None
+                ):
+                    _alibi_cache["_alibi_slopes_require_update"] = True
+                    _alibi_cache["_alibi_bias_require_update"] = True
+
+            # detect bias shape
+            core_attention_bias_shape = None
+            if core_attention_bias is not None:
+                if (
+                    core_attention_bias.shape[0] == batch_size
+                    and core_attention_bias.shape[1] == query_layer.shape[-2]
+                ):
+                    core_attention_bias_shape = "bhss"
+                elif (
+                    core_attention_bias.shape[0] == 1
+                    and core_attention_bias.shape[1] == query_layer.shape[-2]
+                ):
+                    core_attention_bias_shape = "1hss"
+                elif (
+                    core_attention_bias.shape[0] == batch_size and core_attention_bias.shape[1] == 1
+                ):
+                    core_attention_bias_shape = "b1ss"
+                elif core_attention_bias.shape[0] == 1 and core_attention_bias.shape[1] == 1:
+                    core_attention_bias_shape = "11ss"
+                else:
+                    assert (
+                        False
+                    ), "core_attention_bias must be in one of {bhss, 1hss, b1ss, 11ss} shapes"
+
+            # check if there is padding between sequences when qkv_format='thd'
+            if pad_between_seqs is None:
+                if qkv_format == "thd":
+                    pad_between_seqs = (
+                        cu_seqlens_q_padded is not None
+                        and not torch.equal(cu_seqlens_q_padded[:-1], cu_seqlens_q[:-1])
+                    ) or (
+                        cu_seqlens_kv_padded is not None
+                        and not torch.equal(cu_seqlens_kv_padded[:-1], cu_seqlens_kv[:-1])
+                    )
+                else:
+                    pad_between_seqs = False
+
+            # gather attention params for get_attention_backend
+            attention_params = dpa_utils.AttentionParams(
+                qkv_type=type(query_layer),
+                qkv_dtype=query_layer.dtype,
+                qkv_layout=qkv_layout,
+                batch_size=batch_size,
+                num_heads=num_attention_heads,
+                num_gqa_groups=num_gqa_groups,
+                max_seqlen_q=max_seqlen_q,
+                max_seqlen_kv=max_seqlen_kv,
+                head_dim_qk=head_dim_qk,
+                head_dim_v=head_dim_v,
+                attn_mask_type=attn_mask_type,
+                window_size=window_size,
+                alibi_slopes_shape=alibi_slopes.shape if alibi_slopes is not None else None,
+                core_attention_bias_type=core_attention_bias_type,
+                core_attention_bias_shape=core_attention_bias_shape,
+                core_attention_bias_requires_grad=(
+                    core_attention_bias.requires_grad if core_attention_bias is not None else False
+                ),
+                pad_between_seqs=pad_between_seqs,
+                attention_dropout=self.attention_dropout,
+                context_parallel=context_parallel,
+                cp_comm_type=self.cp_comm_type,
+                deterministic=self.deterministic,
+                is_training=self.training,
+                fp8=self.fp8,
+                fp8_meta=self.fp8_meta,
+                inference_params=inference_params,
+                softmax_type=self.softmax_type,
+                return_max_logit=self.return_max_logit,
+            )
+            global _attention_backends
+            if is_in_onnx_export_mode():
+                # We do not want to call get_attention_backend() in ONNX mode
+                # and we want to avoid using any global variables like _attention_backends.
+                use_flash_attention = False
+                use_fused_attention = False
+                use_unfused_attention = True
+            else:
+                if (
+                    _attention_backends["attention_params"] is None
+                    or attention_params != _attention_backends["attention_params"]
+                ):
+                    _attention_backends["attention_params"] = attention_params
+                    _attention_backends["backend_selection_requires_update"] = True
+                if _attention_backends["backend_selection_requires_update"]:
+                    (
+                        use_flash_attention,
+                        flash_attention_backend,
+                        use_fused_attention,
+                        fused_attention_backend,
+                        use_unfused_attention,
+                        _,
+                    ) = dpa_utils.get_attention_backend(attention_params)
+                    # Set global _attention_backends var using return value
+                    # from get_attention_backend()
+                    _attention_backends["use_flash_attention"] = use_flash_attention
+                    _attention_backends["flash_attention_backend"] = flash_attention_backend
+                    _attention_backends["use_fused_attention"] = use_fused_attention
+                    _attention_backends["fused_attention_backend"] = fused_attention_backend
+                    _attention_backends["use_unfused_attention"] = use_unfused_attention
+                    _attention_backends["backend_selection_requires_update"] = False
+                    if use_flash_attention:
+                        self.logger.info(
+                            "Running with FlashAttention backend (version %s)",
+                            flash_attention_backend,
+                        )
+                    elif use_fused_attention:
+                        self.logger.info(
+                            "Running with FusedAttention backend (sub-backend %s)",
+                            int(fused_attention_backend),
+                        )
+                    elif use_unfused_attention:
+                        self.logger.info("Running with UnfusedDotProductAttention backend")
+                else:
+                    use_flash_attention = _attention_backends["use_flash_attention"]
+                    flash_attention_backend = _attention_backends["flash_attention_backend"]
+                    use_fused_attention = _attention_backends["use_fused_attention"]
+                    fused_attention_backend = _attention_backends["fused_attention_backend"]
+                    use_unfused_attention = _attention_backends["use_unfused_attention"]
+
+            # raise exception if no backend is available
+            if sum([use_flash_attention, use_fused_attention, use_unfused_attention]) == 0:
+                raise ValueError(
+                    "No dot product attention backend is available for the provided inputs. Please"
+                    " run with NVTE_DEBUG=1 NVTE_DEBUG_LEVEL=2 to find out the reasons for"
+                    " disabling all backends."
+                )
+
+            # run attention
+            softmax_offset = (
+                self.softmax_offset.reshape(1, -1, 1, 1).to(torch.float32)
+                if self.softmax_offset is not None
+                else None
+            )
+
+            if use_flash_attention:
+                if core_attention_bias_type == "alibi":
+                    alibi_slopes, _ = dpa_utils.get_alibi(
+                        _alibi_cache,
+                        query_layer.shape[-2],
+                        max_seqlen_q,
+                        max_seqlen_kv,
+                        alibi_slopes=alibi_slopes,
+                    )
+                return self.flash_attention(
+                    query_layer,
+                    key_layer,
+                    value_layer,
+                    attention_mask=attention_mask,
+                    qkv_layout=qkv_layout,
+                    cu_seqlens_q=cu_seqlens_q,
+                    cu_seqlens_kv=cu_seqlens_kv,
+                    attn_mask_type=attn_mask_type,
+                    window_size=window_size,
+                    alibi_slopes=alibi_slopes,
+                    cp_group=self.cp_group,
+                    cp_global_ranks=self.cp_global_ranks,
+                    cp_stream=self.cp_stream,
+                    cp_comm_type=self.cp_comm_type,
+                    max_seqlen_q=max_seqlen_q,
+                    max_seqlen_kv=max_seqlen_kv,
+                    fp8=self.fp8 and self.fp8_meta["recipe"].fp8_dpa,
+                    fp8_meta=self.fp8_meta,
+                    quantizers=self.quantizers,
+                    inference_params=inference_params,
+                    flash_attention_backend=flash_attention_backend,
+                    fp8_output=fp8_output,
+                    num_splits=num_splits,
+                )
+
+            if use_fused_attention:
+                fu_core_attention_bias_type = core_attention_bias_type
+                fu_core_attention_bias = core_attention_bias
+                if core_attention_bias_type == "alibi" and (
+                    alibi_slopes is not None or max_seqlen_q != max_seqlen_kv
+                ):
+                    fu_core_attention_bias_type = "post_scale_bias"
+                    _, fu_core_attention_bias = dpa_utils.get_alibi(
+                        _alibi_cache,
+                        query_layer.shape[-2],
+                        max_seqlen_q,
+                        max_seqlen_kv,
+                        alibi_slopes=alibi_slopes,
+                        bias_dtype=query_layer.dtype,
+                        bottom_right_alignment=attn_mask_type not in ["causal", "padding_causal"],
+                    )
+                if checkpoint_core_attention:
+                    return self._checkpointed_attention_forward(
+                        self.fused_attention,
+                        query_layer,
+                        key_layer,
+                        value_layer,
+                        qkv_layout=qkv_layout,
+                        cu_seqlens_q=cu_seqlens_q,
+                        cu_seqlens_kv=cu_seqlens_kv,
+                        cu_seqlens_q_padded=cu_seqlens_q_padded,
+                        cu_seqlens_kv_padded=cu_seqlens_kv_padded,
+                        max_seqlen_q=max_seqlen_q,
+                        max_seqlen_kv=max_seqlen_kv,
+                        attn_mask_type=attn_mask_type,
+                        attention_mask=attention_mask,
+                        window_size=window_size,
+                        fused_attention_backend=fused_attention_backend,
+                        core_attention_bias_type=fu_core_attention_bias_type,
+                        core_attention_bias=fu_core_attention_bias,
+                        fast_zero_fill=fast_zero_fill,
+                        cp_group=self.cp_group,
+                        cp_global_ranks=self.cp_global_ranks,
+                        cp_stream=self.cp_stream,
+                        cp_comm_type=self.cp_comm_type,
+                        fp8=self.fp8 and self.fp8_meta["recipe"].fp8_dpa,
+                        fp8_meta=self.fp8_meta,
+                        quantizers=self.quantizers,
+                        pad_between_seqs=pad_between_seqs,
+                        inference_params=inference_params,
+                        softmax_offset=softmax_offset,
+                        fp8_output=fp8_output,
+                    )
+                return self.fused_attention(
+                    query_layer,
+                    key_layer,
+                    value_layer,
+                    qkv_layout=qkv_layout,
+                    cu_seqlens_q=cu_seqlens_q,
+                    cu_seqlens_kv=cu_seqlens_kv,
+                    cu_seqlens_q_padded=cu_seqlens_q_padded,
+                    cu_seqlens_kv_padded=cu_seqlens_kv_padded,
+                    max_seqlen_q=max_seqlen_q,
+                    max_seqlen_kv=max_seqlen_kv,
+                    attn_mask_type=attn_mask_type,
+                    attention_mask=attention_mask,
+                    window_size=window_size,
+                    fused_attention_backend=fused_attention_backend,
+                    core_attention_bias_type=fu_core_attention_bias_type,
+                    core_attention_bias=fu_core_attention_bias,
+                    fast_zero_fill=fast_zero_fill,
+                    cp_group=self.cp_group,
+                    cp_global_ranks=self.cp_global_ranks,
+                    cp_stream=self.cp_stream,
+                    cp_comm_type=self.cp_comm_type,
+                    fp8=self.fp8 and self.fp8_meta["recipe"].fp8_dpa,
+                    fp8_meta=self.fp8_meta,
+                    quantizers=self.quantizers,
+                    pad_between_seqs=pad_between_seqs,
+                    inference_params=inference_params,
+                    softmax_offset=softmax_offset,
+                    fp8_output=fp8_output,
+                )
+
+            from transformer_engine.pytorch.cpu_offload import CPUOffloadEnabled
+
+            if CPUOffloadEnabled:
+                warnings.warn(
+                    "Attention activation Offloading is only implemented"
+                    "with Flash Attention and Fused Attention!"
+                )
+
+            if use_unfused_attention:
+                allow_emulation = os.getenv("NVTE_UnfusedDPA_Emulate_FP8", "0") == "1"
+                if checkpoint_core_attention:
+                    return self._checkpointed_attention_forward(
+                        self.unfused_attention,
+                        _alibi_cache,
+                        query_layer,
+                        key_layer,
+                        value_layer,
+                        qkv_layout=qkv_layout,
+                        cu_seqlens_q=cu_seqlens_q,
+                        cu_seqlens_kv=cu_seqlens_kv,
+                        max_seqlen_q=max_seqlen_q,
+                        max_seqlen_kv=max_seqlen_kv,
+                        attn_mask_type=attn_mask_type,
+                        attention_mask=attention_mask,
+                        window_size=window_size,
+                        core_attention_bias_type=core_attention_bias_type,
+                        core_attention_bias=core_attention_bias,
+                        alibi_slopes=alibi_slopes,
+                        inference_params=inference_params,
+                        softmax_offset=softmax_offset,
+                        fp8=self.fp8 and self.fp8_meta["recipe"].fp8_dpa and allow_emulation,
+                        fp8_meta=self.fp8_meta,
+                        quantizers=self.quantizers,
+                        fp8_output=fp8_output,
+                    )
+                return self.unfused_attention(
+                    _alibi_cache,
+                    query_layer,
+                    key_layer,
+                    value_layer,
+                    qkv_layout=qkv_layout,
+                    cu_seqlens_q=cu_seqlens_q,
+                    cu_seqlens_kv=cu_seqlens_kv,
+                    max_seqlen_q=max_seqlen_q,
+                    max_seqlen_kv=max_seqlen_kv,
+                    attn_mask_type=attn_mask_type,
+                    attention_mask=attention_mask,
+                    window_size=window_size,
+                    core_attention_bias_type=core_attention_bias_type,
+                    core_attention_bias=core_attention_bias,
+                    alibi_slopes=alibi_slopes,
+                    inference_params=inference_params,
+                    softmax_offset=softmax_offset,
+                    fp8=self.fp8 and self.fp8_meta["recipe"].fp8_dpa and allow_emulation,
+                    fp8_meta=self.fp8_meta,
+                    quantizers=self.quantizers,
+                    fp8_output=fp8_output,
+                )
+            return None
diff --git a/examples/tp-numerics/patches/v2.9.0/transformer_engine/pytorch/module/layernorm_linear.py b/examples/tp-numerics/patches/v2.9.0/transformer_engine/pytorch/module/layernorm_linear.py
new file mode 100644
index 00000000000..8693b9adfeb
--- /dev/null
+++ b/examples/tp-numerics/patches/v2.9.0/transformer_engine/pytorch/module/layernorm_linear.py
@@ -0,0 +1,1913 @@
+# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+#
+# See LICENSE for license information.
+
+"""LayerNormLinear API"""
+import os
+import warnings
+from typing import Callable, Dict, Optional, Tuple, Union, List
+from functools import reduce
+from operator import mul as multiply_op
+
+import torch
+from torch.nn import init
+
+import transformer_engine_torch as tex
+
+from transformer_engine.common.recipe import Recipe
+from transformer_engine.pytorch import torch_version
+from transformer_engine.pytorch.tensor.utils import is_experimental
+from .base import (
+    fill_userbuffers_buffer_for_all_gather,
+    get_workspace,
+    get_ub,
+    TransformerEngineBaseModule,
+    get_dummy_wgrad,
+    _2X_ACC_FPROP,
+    _2X_ACC_DGRAD,
+    _2X_ACC_WGRAD,
+)
+from ..quantization import FP8GlobalStateManager
+from ..utils import (
+    assert_dim_for_fp8_exec,
+    assert_dim_for_all_gather,
+    cast_if_needed,
+    clear_tensor_data,
+    divide,
+    get_default_init_method,
+    init_method_constant,
+    nvtx_range_pop,
+    nvtx_range_push,
+    requires_grad,
+    needs_quantized_gemm,
+)
+from ..distributed import (
+    set_tensor_model_parallel_attributes,
+    get_distributed_world_size,
+    allreduce,
+    symmetric_all_reduce,
+    reduce_scatter_along_first_dim,
+    gather_along_first_dim,
+    in_fp8_activation_recompute_phase,
+    _fsdp_scatter_tensors,
+    _fsdp_gather_tensors,
+)
+from ..constants import GemmParallelModes, dist_group_type
+from ..jit import no_torch_dynamo
+from ..graph import is_graph_capturing
+from ._common import apply_normalization, noop_cat, WeightGradStore
+from ..tensor.quantized_tensor import (
+    QuantizedTensor,
+    QuantizedTensorStorage,
+    Quantizer,
+    prepare_for_saving,
+    restore_from_saved,
+)
+from ...debug.pytorch.debug_state import TEDebugState
+from ..tensor.float8_blockwise_tensor import Float8BlockQuantizer
+from ..tensor.mxfp8_tensor import MXFP8Quantizer
+from ..tensor.storage.float8_blockwise_tensor_storage import Float8BlockwiseQTensorStorage
+from ..tensor.storage.mxfp8_tensor_storage import MXFP8TensorStorage
+from ..export import is_in_onnx_export_mode, assert_warmed_up
+from ..cpu_offload import is_cpu_offload_enabled, mark_activation_offload
+
+from ..cpp_extensions import (
+    general_gemm,
+)
+
+__all__ = ["LayerNormLinear"]
+
+
+class _LayerNormLinear(torch.autograd.Function):
+    """LayerNormLinear semi-top level module
+    Calls custom cuda extensions.
+    """
+
+    @staticmethod
+    def forward(
+        ctx,
+        inp: torch.Tensor,
+        ln_weight: torch.Tensor,
+        ln_bias: Union[torch.Tensor, None],
+        weight: torch.Tensor,
+        bias: torch.Tensor,
+        eps: float,
+        is_first_microbatch: Union[bool, None],
+        fp8: bool,
+        fp8_calibration: bool,
+        wgrad_store: WeightGradStore,
+        fuse_wgrad_accumulation: bool,
+        input_quantizer: Optional[Quantizer],
+        weight_quantizer: Optional[Quantizer],
+        output_quantizer: Optional[Quantizer],
+        grad_input_quantizer: Optional[Quantizer],
+        grad_weight_quantizer: Optional[Quantizer],
+        grad_output_quantizer: Optional[Quantizer],
+        cpu_offloading: bool,
+        tp_group: Union[dist_group_type, None],
+        tp_size: int,
+        sequence_parallel: bool,
+        tensor_parallel: bool,
+        activation_dtype: torch.dtype,
+        parallel_mode: Union[str, None],
+        return_layernorm_output: bool,
+        return_layernorm_output_gathered: bool,
+        is_grad_enabled: bool,
+        fwd_ln_sm_margin: int,
+        bwd_ln_sm_margin: int,
+        zero_centered_gamma: bool,
+        normalization: str,
+        ub_overlap_ag_fprop: bool,
+        ub_overlap_rs_fprop: bool,
+        ub_overlap_ag_dgrad: bool,
+        ub_overlap_rs_dgrad: bool,
+        ub_bulk_wgrad: bool,
+        ub_bulk_dgrad: bool,
+        ub_name: str,
+        fsdp_group: Union[dist_group_type, None],
+        module: torch.nn.Module,
+        skip_fp8_weight_update: bool,
+        symmetric_ar_type: str,
+        debug: Optional[bool] = False,
+    ) -> Union[Tuple[torch.Tensor, ...], torch.Tensor]:
+        # pylint: disable=missing-function-docstring
+
+        # NVTX label for profiling
+        nvtx_label = "transformer_engine._LayerNormLinear.forward"
+        if ub_name is not None:
+            nvtx_label = f"{nvtx_label}.{ub_name}"
+
+        with_input_all_gather = parallel_mode == "column" and sequence_parallel
+
+        # Make sure input dimensions are compatible
+        out_features, in_features = weight.shape
+        inp_shape = inp.shape
+        inp_requires_grad = inp.requires_grad
+        assert inp_shape[-1] == in_features, "GEMM not possible"
+        inp = inp.view((-1, in_features))
+        inputmat = inp
+        if fp8:
+            assert_dim_for_fp8_exec(inputmat, weight)
+            assert_dim_for_all_gather(inputmat, with_input_all_gather, input_quantizer)
+
+        # Cast for native AMP
+        nvtx_range_push(f"{nvtx_label}.norm_input_cast")
+        inputmat = cast_if_needed(inputmat, activation_dtype)
+        ln_weight = cast_if_needed(ln_weight, activation_dtype)
+        if ln_bias is not None:
+            ln_bias = cast_if_needed(ln_bias, activation_dtype)
+        nvtx_range_pop(f"{nvtx_label}.norm_input_cast")
+
+        tp_world_size = get_distributed_world_size(tp_group)
+
+        weight_requires_grad = weight.requires_grad
+        backward_needs_input = is_grad_enabled and weight_requires_grad
+
+        # Configure Userbuffers communication (comm+GEMM overlap)
+        if debug:  # turn off userbuffers in debug mode
+            ub_overlap_ag_fprop = False
+            ub_overlap_rs_fprop = False
+            ub_overlap_ag_dgrad = False
+            ub_overlap_rs_dgrad = False
+            ub_bulk_wgrad = False
+            ub_bulk_dgrad = False
+        ub_obj = None
+        ub_type = None
+        ub_overlap_ag_fprop = (
+            ub_overlap_ag_fprop and is_grad_enabled and not return_layernorm_output
+        )
+        if ub_overlap_rs_fprop:
+            ub_obj = get_ub(ub_name + "_fprop", fp8)
+            ub_type = tex.CommOverlapType.RS
+        elif ub_overlap_ag_fprop:
+            ub_obj = get_ub(ub_name + "_fprop", fp8)
+            ub_type = tex.CommOverlapType.AG
+
+        # Configure quantizer for norm output
+        if fp8:
+            if input_quantizer is None:
+                raise ValueError("Missing quantizer for input tensor")
+            input_quantizer.set_usage(rowwise=True, columnwise=backward_needs_input)
+            if with_input_all_gather and input_quantizer.supports_only_rowwise_all_gather():
+                # All-gather is not supported with FP8 column-wise data
+                input_quantizer.set_usage(columnwise=False)
+
+        # Avoid quantized norm kernel if norm output will be returned
+        # or if a gather of ln_out must be in high precision.
+        experimental = is_experimental(input_quantizer)
+        with_quantized_norm = (
+            fp8
+            and not debug
+            and not return_layernorm_output
+            and not return_layernorm_output_gathered
+            and not experimental  # TODO(negvet): and not FP8GlobalStateManager.get_fp8_recipe().custom()
+        )
+
+        # Apply normalization
+        nvtx_range_push(f"{nvtx_label}.norm")
+        ln_out, mu, rsigma = apply_normalization(
+            inputmat,
+            None,  # ln_out
+            ln_weight,
+            ln_bias,
+            eps,
+            input_quantizer if with_quantized_norm else None,
+            inputmat.dtype,
+            normalization,
+            fwd_ln_sm_margin,
+            zero_centered_gamma,
+        )
+        nvtx_range_pop(f"{nvtx_label}.norm")
+
+        # Store unquantized layer norm output if we need to return it
+        ln_out_return = None
+        if return_layernorm_output or return_layernorm_output_gathered:
+            ln_out_return = ln_out
+
+        # ------------------------------------------------------
+        # Prepare GEMM input tensor
+        # Note: Cast to expected dtype and perform tensor-parallel communication
+        # ------------------------------------------------------
+        nvtx_range_push(f"{nvtx_label}.gemm_input_cast_comm")
+        ln_out_total = None
+        if with_input_all_gather:
+            if return_layernorm_output_gathered:
+                # Perform all-gather in high precision if gathered
+                # norm output will be returned
+                ln_out_total, _ = gather_along_first_dim(ln_out, tp_group)
+                ln_out_return = ln_out_total
+                if fp8 or debug:
+                    ln_out = input_quantizer(ln_out)
+                    input_quantizer.set_usage(rowwise=True, columnwise=False)
+                    if isinstance(input_quantizer, Float8BlockQuantizer):
+                        input_quantizer.all_gather_usage = False
+                    ln_out_total = input_quantizer(ln_out_total)
+            else:
+                quantizer = None
+                if fp8 or debug:
+                    quantizer = input_quantizer
+                    # experimental recipe doesn't need to support quantized AG
+                    if not with_quantized_norm and not experimental:
+                        ln_out = quantizer(ln_out)
+                    quantizer.set_usage(rowwise=True, columnwise=False)
+                if ub_overlap_ag_fprop:  # Initialize Userbuffers all-gather
+                    ln_out_total, _ = fill_userbuffers_buffer_for_all_gather(
+                        ub_obj,
+                        ln_out,
+                        quantizer,
+                        tp_group,
+                    )
+                else:  # Perform NCCL all-gather
+                    ln_out_total, _ = gather_along_first_dim(
+                        ln_out,
+                        tp_group,
+                        quantizer=quantizer,
+                    )
+        else:
+            if (fp8 or debug) and not with_quantized_norm:
+                ln_out = input_quantizer(ln_out)
+            ln_out_total = ln_out
+        nvtx_range_pop(f"{nvtx_label}.gemm_input_cast_comm")
+        # ------------------------------------------------------
+        # GEMM input tensor is ready...
+        # ------------------------------------------------------
+
+        # ------------------------------------------------------
+        # Prepare weight tensor
+        # ------------------------------------------------------
+        weightmat = weight
+        quantized_weight = False
+        if fp8 or debug:
+            quantized_weight = not isinstance(weight, QuantizedTensorStorage)
+
+            # Configure quantizer
+            if weight_quantizer is not None:
+                weight_quantizer.set_usage(rowwise=True, columnwise=is_grad_enabled)
+
+            # Get quantized weight
+            update_workspace = is_first_microbatch is None or is_first_microbatch
+            weightmat = module.get_weight_workspace(
+                tensor=weight,
+                quantizer=weight_quantizer,
+                cache_name=(None if is_first_microbatch is None else "weight"),
+                update_workspace=update_workspace,
+                skip_update_flag=skip_fp8_weight_update,
+                fsdp_group=fsdp_group,
+                workspace_dtype=activation_dtype,
+            )
+            weightmat.update_usage(rowwise_usage=True)
+
+        else:
+            weightmat = cast_if_needed(weightmat, activation_dtype)  # Cast for AMP
+        # ------------------------------------------------------
+        # Weight tensor is ready for GEMM...
+        # ------------------------------------------------------
+
+        # Cast bias to expected dtype
+        bias_dtype = activation_dtype
+        if needs_quantized_gemm(ln_out_total) and activation_dtype == torch.float32:
+            # cuBLAS does not support FP8 GEMM with FP32 bias, so we cast to BF16
+            bias_dtype = torch.bfloat16
+        bias = cast_if_needed(bias, bias_dtype) if bias is not None else bias
+
+        # Calibrate quantizers if needed
+        if not fp8 and fp8_calibration:
+            if input_quantizer is not None:
+                input_quantizer.calibrate(ln_out_total)
+            if weight_quantizer is not None:
+                weight_quantizer.calibrate(weight)
+
+        # Choose whether to use GEMM kernel with split accumulator
+        use_split_accumulator = _2X_ACC_FPROP
+        if fp8:
+            recipe = FP8GlobalStateManager.get_fp8_recipe()
+            if hasattr(recipe, "fp8_gemm_fprop"):
+                use_split_accumulator = recipe.fp8_gemm_fprop.use_split_accumulator
+
+        # Configure output quantizer
+        if output_quantizer is not None:
+            output_quantizer.set_usage(rowwise=True, columnwise=False)
+
+        # Output buffer for Userbuffers reduce-scatter
+        reduce_scatter_out = None
+        if ub_overlap_rs_fprop:
+            out_shape = list(inp_shape)
+            out_shape[0] //= tp_world_size
+            out_shape[-1] = out_features
+            reduce_scatter_out = torch.empty(out_shape, dtype=activation_dtype, device=inp.device)
+
+        # ------------------------------------------------------
+        # Forward GEMM
+        # Note: y = x * w^T
+        # No _tp_invariant_fwd needed here: LayerNormLinear is column-parallel,
+        # so forward GEMM K=hidden_size is constant across TP (invariant by construction).
+        # Only row-parallel forward (linear.py) needs the invariant path (K=in/TP varies).
+        # ------------------------------------------------------
+        if os.environ.get("NVTE_TP_INVARIANT_MODE", "0") == "1" and parallel_mode == "row":
+            assert False, (
+                "NVTE_TP_INVARIANT_MODE row-parallel forward is not implemented in "
+                "layernorm_linear.py. Use linear.py for row-parallel layers."
+            )
+        nvtx_range_push(f"{nvtx_label}.gemm")
+        gemm_out, *_, reduce_scatter_out = general_gemm(
+            weightmat,
+            ln_out_total,
+            get_workspace(),
+            quantization_params=output_quantizer,
+            out_dtype=activation_dtype,
+            bias=bias,
+            use_split_accumulator=use_split_accumulator,
+            ub=ub_obj,
+            ub_type=ub_type,
+            extra_output=reduce_scatter_out,
+        )
+        nvtx_range_pop(f"{nvtx_label}.gemm")
+        # ------------------------------------------------------
+        # Finished forward GEMM...
+        # ------------------------------------------------------
+
+        # Deallocate GEMM input tensor if no longer needed
+        if not weight.requires_grad and not return_layernorm_output:
+            clear_tensor_data(ln_out, ln_out_total)
+            ln_out = ln_out_total = None
+        elif with_input_all_gather and not return_layernorm_output_gathered:
+            clear_tensor_data(ln_out_total)
+            ln_out_total = None
+
+        # ------------------------------------------------------
+        # Prepare output tensor
+        # Note: Perform tensor-parallel communication
+        # ------------------------------------------------------
+        out = None
+        if ub_overlap_rs_fprop:
+            out = reduce_scatter_out
+        elif parallel_mode == "row" and tp_size > 1:
+            nvtx_range_push(f"{nvtx_label}.row_parallel_comm")
+            out = gemm_out
+            if sequence_parallel:
+                out, _ = reduce_scatter_along_first_dim(out, tp_group)
+            elif tensor_parallel:
+                if symmetric_ar_type is not None:
+                    out, _ = symmetric_all_reduce(out, tp_group, all_reduce_type=symmetric_ar_type)
+                else:
+                    out, _ = allreduce(out, tp_group)
+            nvtx_range_pop(f"{nvtx_label}.row_parallel_comm")
+        else:
+            out = gemm_out
+        out = out.view(-1, *inp_shape[1:-1], out_features)
+        # ------------------------------------------------------
+        # Output tensor is ready to return...
+        # ------------------------------------------------------
+
+        # ------------------------------------------------------
+        # Cache state for backward pass
+        # ------------------------------------------------------
+
+        if is_grad_enabled:
+            ctx.weight_quantizer = weight_quantizer
+            ctx.ln_out_needs_gather = (
+                weight.requires_grad and parallel_mode == "column" and sequence_parallel
+            )
+
+            # Input with column-wise usage is needed for wgrad GEMM.
+            if backward_needs_input:
+                if isinstance(ln_out, QuantizedTensorStorage):
+                    # For sequence parallel in vanilla FP8, rowwise data is
+                    # to gather the input. For MXFP8, columnwise only data
+                    # can be allgathered.
+                    if (
+                        isinstance(ln_out, (MXFP8TensorStorage, Float8BlockwiseQTensorStorage))
+                        or not ctx.ln_out_needs_gather
+                    ):
+                        ln_out.update_usage(rowwise_usage=False)
+
+            # Weight with column-wise usage is needed for dgrad GEMM.
+            if isinstance(weightmat, QuantizedTensorStorage):
+                weightmat.update_usage(columnwise_usage=True)
+
+            if cpu_offloading:
+                mark_activation_offload(inputmat, mu, rsigma, ln_out)
+
+            # Scatter intermediate/activation tensors saved for the backward pass
+            # NOTE: weight_fp8 = weight when ctx.fp8 == False and torch.disttributed.FSDP already
+            #       shards/unshards the base weights so we don't do it ourselves
+            nvtx_range_push(f"{nvtx_label}.fsdp_scatter")
+            ctx.fsdp_group = fsdp_group
+            ctx.fsdp_shapes = _fsdp_scatter_tensors(
+                fsdp_group,
+                mu,
+                rsigma,
+                weightmat if quantized_weight else None,
+                ln_out if weight.requires_grad else None,
+            )
+            nvtx_range_pop(f"{nvtx_label}.fsdp_scatter")
+
+            if cpu_offloading:
+                ctx.grad_added_to_main_grad = hasattr(weight, "grad_added_to_main_grad")
+
+                if ctx.grad_added_to_main_grad:
+                    # If you are passing torch.nn.Parameter through the Torch hooks, you will
+                    # get back torch.Tensor. Torch rips off the Parameter wrapper.
+                    # You need to preserve the weight object to have all the attributes user
+                    # sets for the weights. Because of this, it is not recommended to offload
+                    # weights if weights are externally touched outside this module
+                    ctx.weight_object = weight
+
+            tensors_to_save, tensor_objects = prepare_for_saving(
+                inputmat,
+                weightmat,
+                weight,
+                bias,
+                ln_weight,
+                ln_out,
+                mu,
+                rsigma,
+            )
+            ctx.save_for_backward(*tensors_to_save)
+            ctx.tensor_objects = tensor_objects
+            ctx.requires_dgrad = inp_requires_grad
+            ctx.requires_wgrad = weight.requires_grad
+            ctx.quantized_weight = quantized_weight
+            if fuse_wgrad_accumulation and weight.requires_grad:
+                # This check is needed to ensure that main_grad is not created
+                # during the forward pass when using MCore FSDP as it creates
+                # the main_grad buffer lazily before backprop
+                if hasattr(weight, "__fsdp_param__"):
+                    # MCore FSDP creates main_grad lazily before backward
+                    ctx.main_grad_func = weight.get_main_grad
+                else:
+                    ctx.main_grad_func = lambda: weight.main_grad
+            ctx.grad_input_quantizer = grad_input_quantizer
+            ctx.grad_weight_quantizer = grad_weight_quantizer
+            ctx.grad_output_quantizer = grad_output_quantizer
+            ctx.input_quantizer = input_quantizer
+            ctx.owns_input = inputmat is not inp
+            ctx.weight = weight
+            ctx.partition_stride = getattr(weight, 'partition_stride', 1)
+            ctx.activation_dtype = activation_dtype
+            ctx.fp8 = fp8
+            ctx.fp8_recipe = FP8GlobalStateManager.get_fp8_recipe() if fp8 else None
+            ctx.fuse_wgrad_accumulation = fuse_wgrad_accumulation
+            ctx.cpu_offloading = cpu_offloading
+            ctx.is_first_microbatch = is_first_microbatch
+            ctx.use_bias = bias is not None
+            ctx.sequence_parallel = sequence_parallel
+            ctx.tensor_parallel = tensor_parallel
+            ctx.inp_shape = inp_shape
+            ctx.parallel_mode = parallel_mode
+            ctx.tp_group = tp_group
+            ctx.tp_size = tp_size
+            ctx.return_layernorm_output = return_layernorm_output
+            ctx.return_layernorm_output_gathered = return_layernorm_output_gathered
+            ctx.bwd_ln_sm_margin = bwd_ln_sm_margin
+            ctx.zero_centered_gamma = zero_centered_gamma
+            ctx.ub_overlap_ag = ub_overlap_ag_dgrad
+            ctx.ub_overlap_rs_dgrad = ub_overlap_rs_dgrad
+            ctx.ub_bulk_wgrad = ub_bulk_wgrad
+            ctx.ub_bulk_dgrad = ub_bulk_dgrad
+            ctx.ub_name = ub_name
+            ctx.requires_dgrad = inp_requires_grad
+            ctx.normalization = normalization
+            ctx.reduce_and_update_bwd_fp8_tensors = False
+            if ctx.fp8 and requires_grad(inp, ln_weight, ln_bias, weight, bias):
+                _first_fp8_module = FP8GlobalStateManager.IS_FIRST_FP8_MODULE
+                ctx.reduce_and_update_bwd_fp8_tensors = FP8GlobalStateManager.is_first_fp8_module()
+                if in_fp8_activation_recompute_phase():
+                    FP8GlobalStateManager.IS_FIRST_FP8_MODULE = _first_fp8_module
+            ctx.wgrad_store = wgrad_store
+            ctx.debug = debug
+
+        # ------------------------------------------------------
+        # Cached state for backward pass is ready...
+        # ------------------------------------------------------
+
+        if return_layernorm_output:
+            if return_layernorm_output_gathered:
+                shape = list(inp_shape)
+                shape[0] *= tp_size if with_input_all_gather else 1
+                return out, ln_out_return.view(shape)
+            return out, ln_out_return.view(inp_shape)
+        return out
+
+    @staticmethod
+    def backward(
+        ctx, *grad_outputs: Tuple[torch.Tensor, ...]
+    ) -> Tuple[Union[torch.Tensor, None], ...]:
+        # pylint: disable=missing-function-docstring
+
+        # NVTX label for profiling
+        nvtx_label = "transformer_engine._LayerNormLinear.backward"
+        if ctx.ub_name is not None:
+            nvtx_label = f"{nvtx_label}.{ctx.ub_name}"
+
+        with torch.cuda.nvtx.range("_LayerNormLinear_backward"):
+            saved_tensors = ctx.saved_tensors
+            (  # pylint: disable=unbalanced-tuple-unpacking
+                inputmat,
+                weight,
+                origin_weight,
+                bias,
+                ln_weight,
+                ln_out,
+                mu,
+                rsigma,
+            ) = restore_from_saved(ctx.tensor_objects, saved_tensors)
+            # Delete the references to tensor objects once they've been consumed
+            # by the `restore_from_saved` method to construct back the actual tensors.
+            ctx.tensor_objects = None
+
+            # Since main_grad can be modified inplace, it should not be a part of saved_tensors
+            main_grad = (
+                ctx.main_grad_func()
+                if weight is not None and ctx.fuse_wgrad_accumulation and ctx.requires_wgrad
+                else None
+            )
+
+            # Gather intermediate/activation tensors if needed
+            # NOTE: weight_fp8 = weight when ctx.fp8 == False and torch.disttributed.FSDP already
+            #       shards/unshards the base weights so we don't do it ourselves
+            nvtx_range_push(f"{nvtx_label}.fsdp_gather")
+            _fsdp_gather_tensors(
+                ctx.fsdp_group,
+                ctx.fsdp_shapes,
+                mu,
+                rsigma,
+                weight if ctx.fp8 and ctx.quantized_weight else None,
+                ln_out,
+            )
+            nvtx_range_pop(f"{nvtx_label}.fsdp_gather")
+
+            # For CPU offloading, we offloaded weight and weight.main_grad to different tensors,
+            # we need to connect them into one.
+            if ctx.cpu_offloading:
+                if ctx.grad_added_to_main_grad:
+                    origin_weight = ctx.weight_object
+                if ctx.requires_wgrad and ctx.fuse_wgrad_accumulation:
+                    origin_weight.main_grad = main_grad
+
+            # Configure Userbuffers communication (comm+GEMM overlap)
+            ctx.ub_obj_gradout = None
+            ub_obj_dgrad = None
+            ub_obj_wgrad = None
+            ub_type_dgrad = None
+            ub_type_wgrad = None
+            dgrad_shape = [reduce(multiply_op, ctx.inp_shape[:-1]), ctx.inp_shape[-1]]
+            if ctx.ub_overlap_ag:
+                # Overlap grad_output all-gather with dgrad compute
+                ctx.ub_obj_gradout = get_ub(ctx.ub_name + "_dgrad", ctx.fp8)
+                ub_obj_dgrad = ctx.ub_obj_gradout
+                ub_type_dgrad = tex.CommOverlapType.AG
+            elif ctx.ub_overlap_rs_dgrad:
+                # Overlap dgrad reduce-scatter with dgrad compute
+                ctx.ub_obj_gradout = get_ub(ctx.ub_name + "_dgrad", ctx.fp8)
+                ub_obj_dgrad = ctx.ub_obj_gradout
+                ub_type_dgrad = tex.CommOverlapType.RS
+            else:
+                if ctx.ub_bulk_dgrad:
+                    # Overlap inputmat all-gather with dgrad compute
+                    ctx.ub_obj_gradout = get_ub(ctx.ub_name + "_dgrad", ctx.fp8)
+                    ub_obj_dgrad = ctx.ub_obj_gradout
+                    ub_type_dgrad = tex.CommOverlapType.AG
+                if ctx.ub_bulk_wgrad:
+                    # Overlap dgrad reduce-scatter with wgrad compute
+                    ub_obj_wgrad = get_ub(ctx.ub_name + "_wgrad", ctx.fp8)
+                    ub_type_wgrad = tex.CommOverlapType.RS
+
+            # --------------------------------------------------
+            # Prepare grad output tensor
+            # Note: Cast to expected dtype and perform tensor-parallel communication
+            # --------------------------------------------------
+
+            # Configure quantizer for grad output tensor
+            # Note: dgrad GEMM requires row-wise usage, wgrad GEMM
+            # requires column-wise usage
+            if ctx.grad_output_quantizer is not None:
+                quantizer = ctx.grad_output_quantizer
+                quantizer.set_usage(rowwise=True, columnwise=True)
+                if ctx.ub_overlap_ag:
+                    # Userbuffers only supports communication for one
+                    # tensor usage at a time. Configure quantizer with
+                    # usage for only dgrad GEMM.
+                    quantizer.set_usage(columnwise=False)
+
+            # Prepare grad output tensor
+            # Note: Cast to expected dtype and perform tensor-parallel communication
+            nvtx_range_push(f"{nvtx_label}.grad_output_preprocess")
+            (
+                grad_output,
+                grad_bias,
+            ) = TransformerEngineBaseModule.grad_output_preprocess(
+                ctx,
+                grad_outputs[0],
+                ctx.parallel_mode == "row",
+                ctx.grad_output_quantizer,
+            )
+            nvtx_range_pop(f"{nvtx_label}.grad_output_preprocess")
+
+            # --------------------------------------------------
+            # Grad output tensor is ready for computing grad input...
+            # --------------------------------------------------
+
+            # --------------------------------------------------
+            # Prepare GEMM input tensor
+            # Note: Input tensor is needed for wgrad GEMM.
+            # Tensor-parallel communication is overlapped with dgrad
+            # GEMM.
+            # --------------------------------------------------
+            ln_out_total = None
+            ln_out_total_work = None
+            if ctx.ln_out_needs_gather:
+                quantizer = None
+                if ctx.input_quantizer is not None:
+                    quantizer = ctx.input_quantizer
+                    if quantizer.supports_only_rowwise_all_gather():
+                        # If data is in FP8, we compute FP8 transposes manually
+                        quantizer.set_usage(rowwise=True, columnwise=False)
+                    else:
+                        # wgrad GEMM requires input with column-wise usage
+                        quantizer.set_usage(rowwise=False, columnwise=True)
+                if ctx.ub_bulk_dgrad:
+                    ln_out_total, _ = fill_userbuffers_buffer_for_all_gather(
+                        ub_obj_dgrad,
+                        ln_out,
+                        quantizer,
+                        ctx.tp_group,
+                    )
+                else:
+                    nvtx_range_push(f"{nvtx_label}.column_parallel_comm_input")
+                    ln_out_total, ln_out_total_work = gather_along_first_dim(
+                        ln_out,
+                        ctx.tp_group,
+                        async_op=True,
+                        quantizer=quantizer,
+                    )
+                    nvtx_range_pop(f"{nvtx_label}.column_parallel_comm_input")
+            else:
+                ln_out_total = ln_out
+            # --------------------------------------------------
+            # Input tensor is ready for computing grad weight...
+            # --------------------------------------------------
+
+            # --------------------------------------------------
+            # Compute grad input tensor
+            # Note: Gradient w.r.t. GEMM input (i.e. norm output).
+            # --------------------------------------------------
+
+            # Make sure required data is available
+            if isinstance(grad_output, QuantizedTensorStorage):
+                grad_output.update_usage(rowwise_usage=True)
+            if ctx.weight_quantizer is not None and isinstance(weight, QuantizedTensorStorage):
+                weight.update_usage(columnwise_usage=True)
+
+            # Choose whether to use GEMM kernel with split accumulator
+            use_split_accumulator = _2X_ACC_DGRAD
+            if ctx.fp8:
+                recipe = ctx.fp8_recipe
+                if hasattr(recipe, "fp8_gemm_dgrad"):
+                    use_split_accumulator = recipe.fp8_gemm_dgrad.use_split_accumulator
+
+            # Update grad input quantizer
+            if ctx.grad_input_quantizer is not None:
+                ctx.grad_input_quantizer.set_usage(rowwise=True, columnwise=False)
+
+            # Output buffers for Userbuffers reduce-scatter
+            gemm_out = None
+            reduce_scatter_out = None
+            if ctx.ub_overlap_rs_dgrad:
+                reduce_scatter_out = torch.empty(
+                    dgrad_shape, dtype=ctx.activation_dtype, device=grad_outputs[0].device
+                )
+            elif ctx.ub_bulk_wgrad:
+                gemm_out = ub_obj_wgrad.get_buffer(local_chunk=False)
+
+            # dgrad GEMM
+            # Note: dx = dy * w
+            dgrad = None
+            dgrad_work = None
+
+            _tp_invariant_bwd = (
+                os.environ.get("NVTE_TP_INVARIANT_MODE", "0") == "1"
+                and ctx.parallel_mode == "column"
+                and ctx.tp_size > 1
+            )
+
+            if _tp_invariant_bwd:
+                # TP-invariant diagnostic: full dgrad GEMM matching TP=1 accumulation.
+                assert not ctx.fp8, "NVTE_TP_INVARIANT_MODE does not support FP8"
+                nvtx_range_push(f"{nvtx_label}.tp_invariant_dgrad")
+
+                def allgather_along_dim(tensor, group, world_size, dim):
+                    chunks = [torch.empty_like(tensor) for _ in range(world_size)]
+                    torch.distributed.all_gather(
+                        chunks, tensor.contiguous(), group=group,
+                    )
+                    return torch.cat(chunks, dim=dim)
+
+                grad_output_gathered = allgather_along_dim(
+                    grad_output, ctx.tp_group, ctx.tp_size, dim=-1,
+                )  # [tokens, out/TP] -> [tokens, out]
+                weight_gathered = allgather_along_dim(
+                    weight, ctx.tp_group, ctx.tp_size, dim=0,
+                )  # [out/TP, in] -> [out, in]
+
+                # Deinterleave gathered tensors to match TP=1 K-dimension ordering.
+                # Only needed for gated linear units (partition_stride > 1, e.g. SwiGLU FC1)
+                # where each rank stores interleaved [gate_i | value_i].
+                # After all-gather: [gate_0|val_0 | gate_1|val_1 | ...].
+                # TP=1 layout is [gate_all | val_all]. Reorder to match.
+                #
+                # For non-gated layers (partition_stride == 1, e.g. QKV), each rank
+                # stores contiguous GQA groups, and the naive all-gather already
+                # produces the correct TP=1 ordering. Deinterleaving would corrupt it.
+                if ctx.partition_stride > 1:
+                    chunk_sz = weight.shape[0]  # out_features per rank
+                    half = chunk_sz // 2
+                    first_w = [weight_gathered[i * chunk_sz : i * chunk_sz + half] for i in range(ctx.tp_size)]
+                    second_w = [weight_gathered[i * chunk_sz + half : (i + 1) * chunk_sz] for i in range(ctx.tp_size)]
+                    weight_gathered = torch.cat(first_w + second_w, dim=0)
+
+                    g_dim = grad_output_gathered.shape[-1] // ctx.tp_size
+                    g_half = g_dim // 2
+                    first_g = [grad_output_gathered[..., i * g_dim : i * g_dim + g_half] for i in range(ctx.tp_size)]
+                    second_g = [grad_output_gathered[..., i * g_dim + g_half : (i + 1) * g_dim] for i in range(ctx.tp_size)]
+                    grad_output_gathered = torch.cat(first_g + second_g, dim=-1)
+
+                grad_output_2d = grad_output_gathered.reshape(
+                    -1, grad_output_gathered.shape[-1],
+                )
+                dgrad = general_gemm(
+                    weight_gathered, grad_output_2d, get_workspace(),
+                    layout="NN", grad=True,
+                    out_dtype=ctx.activation_dtype,
+                )
+                if isinstance(dgrad, tuple):
+                    dgrad = dgrad[0]
+
+                # SP: scatter to per-rank chunk along sequence dim.
+                if ctx.sequence_parallel:
+                    rank = torch.distributed.get_rank(ctx.tp_group)
+                    dgrad = dgrad.chunk(ctx.tp_size, dim=0)[rank].contiguous()
+
+                del grad_output_gathered, weight_gathered, grad_output_2d
+                nvtx_range_pop(f"{nvtx_label}.tp_invariant_dgrad")
+            else:
+                nvtx_range_push(f"{nvtx_label}.dgrad_gemm")
+                gemm_out, *_, reduce_scatter_out = general_gemm(
+                    weight,
+                    grad_output,
+                    get_workspace(),
+                    layout="NN",
+                    grad=True,
+                    quantization_params=ctx.grad_input_quantizer,
+                    out=gemm_out,
+                    out_dtype=ctx.activation_dtype,
+                    use_split_accumulator=use_split_accumulator,
+                    ub=ub_obj_dgrad,
+                    ub_type=ub_type_dgrad,
+                    extra_output=reduce_scatter_out,
+                    bulk_overlap=ctx.ub_bulk_dgrad,
+                )
+                nvtx_range_pop(f"{nvtx_label}.dgrad_gemm")
+
+                # Prepare grad input tensor
+                # Note: Perform tensor-parallel communication
+                if ctx.ub_overlap_rs_dgrad:
+                    dgrad = reduce_scatter_out
+                elif ctx.ub_bulk_wgrad:
+                    dgrad = ub_obj_wgrad.get_buffer(local_chunk=True)
+                elif ctx.parallel_mode == "column" and ctx.tp_size > 1:
+                    nvtx_range_push(f"{nvtx_label}.column_parallel_comm_dgrad")
+                    dgrad = gemm_out
+                    if ctx.sequence_parallel:
+                        dgrad, dgrad_work = reduce_scatter_along_first_dim(
+                            dgrad,
+                            ctx.tp_group,
+                            async_op=True,
+                        )
+                    else:
+                        dgrad, dgrad_work = allreduce(dgrad, ctx.tp_group, async_op=True)
+                    nvtx_range_pop(f"{nvtx_label}.column_parallel_comm_dgrad")
+                else:
+                    dgrad = gemm_out
+
+            # --------------------------------------------------
+            # Grad input tensor has been computed...
+            # --------------------------------------------------
+
+            # --------------------------------------------------
+            # Compute grad weight
+            # --------------------------------------------------
+
+            wgrad = None
+            if ctx.requires_wgrad:
+                # Prepare grad output tensor
+                # Note: Synchronize tensor-parallel communication and
+                # make sure required data is available
+                if ctx.ub_overlap_ag and isinstance(ctx.grad_output_quantizer, MXFP8Quantizer):
+                    # UB does not support pipelined overlapping grad output
+                    # all-gather with wgrad GEMM. Also, we can't
+                    # convert row-scaled MXFP8 to column-scaled, so we
+                    # can't reuse the grad output that was gathered
+                    # for the dgrad GEMM. We work around by explicitly
+                    # overlapping the AG operation with the dgrad GEMM.
+
+                    # Get the communication stream from the dgrad GEMM to use for the AG
+                    dgrad_send_stream, dgrad_recv_stream = ub_obj_dgrad.get_communication_stream()
+
+                    # This object is separate from the ub_obj_wgrad object which is passed to the GEMM
+                    ub_obj_overlap_wgrad = get_ub(ctx.ub_name + "_wgrad", ctx.fp8)
+
+                    ctx.grad_output_quantizer.set_usage(rowwise=False, columnwise=True)
+
+                    # We use the send stream to copy into the userbuffers.
+                    # This is the same stream that we will use to access the data in the AG,
+                    # so we dont need to add any syncs yet.
+                    with torch.cuda.stream(dgrad_send_stream):
+                        grad_output, _ = fill_userbuffers_buffer_for_all_gather(
+                            ub_obj_overlap_wgrad,
+                            grad_outputs[0],
+                            ctx.grad_output_quantizer,
+                            ctx.tp_group,
+                        )
+
+                    # Allgather grad_outputs[0] using the dgrad streams so we can overlap with the fc2_dgrad gemm
+                    tex.bulk_overlap_ag_with_external_gemm(
+                        ub_obj_overlap_wgrad, dgrad_send_stream, dgrad_recv_stream
+                    )
+
+                # Prepare input tensor
+                # Note: Synchronize tensor-parallel communication and
+                # make sure required data is available
+                if ln_out_total_work is not None:
+                    ln_out_total_work.wait()
+                    ln_out_total_work = None
+                if ctx.fp8 or ctx.debug:
+                    if isinstance(ln_out_total, QuantizedTensorStorage):
+                        ln_out_total.update_usage(columnwise_usage=True)
+                    else:
+                        ctx.input_quantizer.set_usage(rowwise=False, columnwise=True)
+                        ln_out_total = ctx.input_quantizer(ln_out_total)
+
+                if ctx.fp8 or ctx.debug:
+                    if isinstance(grad_output, QuantizedTensorStorage):
+                        grad_output.update_usage(columnwise_usage=True)
+                    else:
+                        ctx.grad_output_quantizer.set_usage(rowwise=False, columnwise=True)
+                        grad_output = ctx.grad_output_quantizer(grad_output)
+
+                # Figure out whether to use split accumulator
+                use_split_accumulator = _2X_ACC_WGRAD
+                if ctx.fp8:
+                    recipe = ctx.fp8_recipe
+                    if hasattr(recipe, "fp8_gemm_wgrad"):
+                        use_split_accumulator = recipe.fp8_gemm_wgrad.use_split_accumulator
+
+                # Figure out whether to output wgrad GEMM directly into main grad
+                if ctx.is_first_microbatch is not None:
+                    accumulate_wgrad_into_param_main_grad = (
+                        ctx.fuse_wgrad_accumulation and not ctx.is_first_microbatch
+                    )
+                else:
+                    accumulate_wgrad_into_param_main_grad = ctx.fuse_wgrad_accumulation
+
+                # Output buffer for overlapping FP8 grad input
+                # reduce-scatter with wgrad GEMM
+                reduce_scatter_out = None
+                if ctx.ub_bulk_wgrad and ub_obj_wgrad.is_fp8_ubuf():
+                    reduce_scatter_out = torch.empty(
+                        dgrad_shape, dtype=ctx.activation_dtype, device=grad_outputs[0].device
+                    )
+
+                # Arguments to include in wgrad GEMM closure
+                wgrad_gemm_kwargs = {
+                    "workspace": get_workspace(),
+                    "out_dtype": (
+                        main_grad.dtype if ctx.fuse_wgrad_accumulation else ctx.activation_dtype
+                    ),
+                    "quantization_params": ctx.grad_weight_quantizer,
+                    "accumulate": (
+                        accumulate_wgrad_into_param_main_grad
+                        if not getattr(weight, "overwrite_main_grad", False)
+                        else False
+                    ),
+                    "layout": "NT",
+                    "out": main_grad if ctx.fuse_wgrad_accumulation else None,
+                    "bias": (bias if (grad_bias is None and not ctx.fp8) else None),
+                    "use_split_accumulator": use_split_accumulator,
+                    "grad": True,
+                    "ub": ub_obj_wgrad,
+                    "ub_type": ub_type_wgrad,
+                    "extra_output": reduce_scatter_out,
+                    "bulk_overlap": ctx.ub_bulk_wgrad,
+                }
+
+                def wgrad_gemm(
+                    x: torch.Tensor,
+                    dy: torch.Tensor,
+                ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+                    """Perform wgrad GEMM: dw = dy^T * x
+
+                    May be fused with bgrad computation.
+
+                    May be called outside of this function to enable
+                    some advanced communication/compute overlapping.
+
+                    """
+                    nvtx_range_push(f"{nvtx_label}.wgrad_gemm")
+                    dw, db, *_ = general_gemm(x, dy, **wgrad_gemm_kwargs)
+                    nvtx_range_pop(f"{nvtx_label}.wgrad_gemm")
+                    return dw, db
+
+                # Choose whether to call wgrad GEMM now or delay
+                if ctx.wgrad_store is not None and ctx.wgrad_store.delay_wgrad_compute():
+                    if (
+                        wgrad_gemm_kwargs["ub"] is not None
+                        or wgrad_gemm_kwargs["ub_type"] is not None
+                        or wgrad_gemm_kwargs["extra_output"] is not None
+                        or wgrad_gemm_kwargs["bulk_overlap"]
+                    ):
+                        raise NotImplementedError(
+                            "Delayed weight grad computation is not supported "
+                            "with Userbuffers (tensor-parallel communication overlapping)"
+                        )
+                    ctx.wgrad_store.put([ln_out_total, grad_output], wgrad_gemm)
+                else:
+
+                    # Call wgrad GEMM now
+                    wgrad, grad_bias_ = wgrad_gemm(ln_out_total, grad_output)
+
+                    # Update grad bias if needed
+                    if grad_bias is None:
+                        grad_bias = grad_bias_
+                    del grad_bias_
+
+                    # Deallocate input tensors if permitted
+                    if not ctx.return_layernorm_output and not ctx.return_layernorm_output_gathered:
+                        # Input tensors have not been exposed externally
+                        clear_tensor_data(ln_out)
+                    elif ctx.ln_out_needs_gather and ctx.return_layernorm_output_gathered:
+                        # Non-gathered input has not been exposed externally
+                        clear_tensor_data(ln_out)
+                    if ctx.ln_out_needs_gather:
+                        # Gathered input is internal
+                        clear_tensor_data(ln_out_total)
+                    if ctx.parallel_mode == "row" and ctx.sequence_parallel:
+                        # Gathered grad output tensor is internal
+                        clear_tensor_data(grad_output)
+
+                # Update grad input if overlapping reduce-scatter with wgrad GEMM
+                if ctx.ub_bulk_wgrad:
+                    if ub_obj_wgrad.is_fp8_ubuf():
+                        dgrad = reduce_scatter_out
+                    else:
+                        dgrad = ub_obj_wgrad.get_buffer(local_chunk=True).clone()
+
+            # --------------------------------------------------
+            # Grad weight has been computed...
+            # --------------------------------------------------
+
+            # Don't return grad bias if not needed
+            if not ctx.use_bias:
+                grad_bias = None
+
+            # Synchronize tensor parallel communication
+            if ln_out_total_work is not None:
+                ln_out_total_work.wait()
+                ln_out_total_work = None
+            if dgrad_work is not None:
+                dgrad_work.wait()
+                dgrad_work = None
+
+            # Residual gradient
+            dgrad = dgrad.view(inputmat.shape)
+            if ctx.return_layernorm_output and not ctx.return_layernorm_output_gathered:
+                dgrad = dgrad + grad_outputs[1].view_as(dgrad)
+
+            # Norm gradient
+            dgamma = None
+            dbeta = None
+            nvtx_range_push(f"{nvtx_label}.norm")
+            if ctx.normalization == "LayerNorm":
+                dgrad, dgamma, dbeta = tex.layernorm_bwd(
+                    dgrad,
+                    inputmat,
+                    mu,
+                    rsigma,
+                    ln_weight,
+                    ctx.bwd_ln_sm_margin,
+                    ctx.zero_centered_gamma,
+                )
+                dgrad = dgrad.reshape(inputmat.size())
+            elif ctx.normalization == "RMSNorm":
+                dgrad, dgamma = tex.rmsnorm_bwd(
+                    dgrad,
+                    inputmat,
+                    rsigma,
+                    ln_weight,
+                    ctx.bwd_ln_sm_margin,
+                    ctx.zero_centered_gamma,
+                )
+                dgrad = dgrad.reshape(inputmat.size())
+                dbeta = None
+            nvtx_range_pop(f"{nvtx_label}.norm")
+            clear_tensor_data(mu)
+            clear_tensor_data(rsigma)
+
+        if ctx.requires_wgrad:
+            # Handle custom DDP from mcore.
+            if ctx.fuse_wgrad_accumulation and hasattr(origin_weight, "grad_added_to_main_grad"):
+                origin_weight.grad_added_to_main_grad = True
+                if getattr(origin_weight, "zero_out_wgrad", False):
+                    wgrad = get_dummy_wgrad(
+                        list(origin_weight.main_grad.shape),
+                        origin_weight.dtype,
+                        zero=True,
+                    )
+                else:
+                    wgrad = get_dummy_wgrad(
+                        list(origin_weight.main_grad.shape),
+                        origin_weight.dtype,
+                    )
+            elif ctx.fuse_wgrad_accumulation:
+                wgrad = None
+        else:
+            wgrad = None
+
+        if ctx.reduce_and_update_bwd_fp8_tensors and not is_graph_capturing():
+            nvtx_range_push(f"{nvtx_label}.reduce_and_update_fp8_tensors")
+            FP8GlobalStateManager.reduce_and_update_fp8_tensors(forward=False)
+            nvtx_range_pop(f"{nvtx_label}.reduce_and_update_fp8_tensors")
+
+        # Scatter fp8 weight buffers
+        # if ctx.fp8 and not isinstance(weight, QuantizedTensorStorage):
+        #    _fsdp_scatter_tensors(ctx.fsdp_group, weight_fp8)
+
+        return (
+            dgrad.view(ctx.inp_shape) if ctx.requires_dgrad else None,
+            dgamma,
+            dbeta,
+            wgrad,
+            grad_bias,
+            None,  # eps
+            None,  # is_first_microbatch
+            None,  # fp8
+            None,  # fp8_calibration
+            None,  # wgrad_store
+            None,  # fuse_wgrad_accumulation
+            None,  # input_quantizer
+            None,  # weight_quantizer
+            None,  # output_quantizer
+            None,  # grad_input_quantizer
+            None,  # grad_weight_quantizer
+            None,  # grad_output_quantizer
+            None,  # cpu_offloading
+            None,  # tp_group
+            None,  # tp_size
+            None,  # sequence_parallel
+            None,  # tensor_parallel
+            None,  # activation_dtype
+            None,  # parallel_mode
+            None,  # return_layernorm_output
+            None,  # return_layernorm_output_gathered
+            None,  # is_grad_enabled
+            None,  # fwd_ln_sm_margin
+            None,  # bwd_ln_sm_margin
+            None,  # zero_centered_gamma
+            None,  # normalization
+            None,  # ub_overlap_ag_fprop
+            None,  # ub_overlap_rs_fprop
+            None,  # ub_overlap_ag_dgrad
+            None,  # ub_overlap_rs_dgrad
+            None,  # ub_bulk_dgrad
+            None,  # ub_bulk_wgrad
+            None,  # ub_name
+            None,  # fsdp_group
+            None,  # debug
+            None,  # module
+            None,  # skip_fp8_weight_update
+            None,  # symmetric_ar_type
+        )
+
+
+class LayerNormLinear(TransformerEngineBaseModule):
+    r"""
+    Applies layer normalization followed by linear transformation to the incoming data.
+
+    Parameters
+    ----------
+    in_features : int
+                 size of each input sample.
+    out_features : int
+                  size of each output sample.
+    eps : float, default = 1e-5
+         a value added to the denominator of layer normalization for numerical stability.
+    bias : bool, default = `True`
+          if set to `False`, the layer will not learn an additive bias.
+    normalization : { 'LayerNorm', 'RMSNorm' }, default = 'LayerNorm'
+                   type of normalization applied.
+    init_method : Callable, default = `None`
+                 used for initializing weights in the following way: `init_method(weight)`.
+                 When set to `None`, defaults to `torch.nn.init.normal_(mean=0.0, std=0.023)`.
+    return_layernorm_output : bool, default = `False`
+                             if set to `True`, output of layernorm is returned from the forward
+                             together with the output of the linear transformation.
+                             Example use case: residual connection for transformer module is
+                             taken post layernorm.
+    return_layernorm_output_gathered : bool, default = `False`
+                             if set to `True`, output of layernorm is returned after the all
+                             gather operation. Ignored if return_layernorm_output is False.
+                             Example use case: with sequence parallel, input to residual connection
+                             for transformer module (e.g. LoRA) will need to be gathered.
+                             Returning layernorm output gathered will prevent a redundant gather.
+    parameters_split : Optional[Union[Tuple[str, ...], Dict[str, int]]], default = None
+                      Configuration for splitting the weight and bias tensors along dim 0 into
+                      multiple PyTorch parameters. If a list or tuple of strings is provided,
+                      they are used to make the names of equally-sized parameters. If a dict
+                      (preferably an OrderedDict) is provided, the keys are used as names and
+                      values as split sizes along dim 0. The resulting parameters will have
+                      names that end in `_weight` or `_bias`, so trailing underscores are
+                      stripped from any provided names.
+    zero_centered_gamma : bool, default = 'False'
+                         if set to 'True', gamma parameter in LayerNorm is initialized to 0 and
+                         the LayerNorm formula changes to
+
+                         .. math::
+                            y = \frac{x - \mathrm{E}[x]}{ \sqrt{\mathrm{Var}[x] + \varepsilon}} *
+                            (1 + \gamma) + \beta
+    device : Union[torch.device, str], default = "cuda"
+          The device on which the parameters of the model will be allocated. It is the user's
+          responsibility to ensure all parameters are moved to the GPU before running the
+          forward pass.
+    name: str, default = `None`
+        name of the module, currently used for debugging purposes.
+
+    Parallelism parameters
+    ----------------------
+    sequence_parallel : bool, default = `False`
+                       if set to `True`, uses sequence parallelism.
+    tp_group : ProcessGroup, default = `None`
+              tensor parallel process group.
+    tp_size : int, default = 1
+             used as TP (tensor parallel) world size when TP groups are not formed during
+             initialization. In this case, users must call the
+             `set_tensor_parallel_group(tp_group)` method on the initialized module before the
+             forward pass to supply the tensor parallel group needed for tensor and sequence
+             parallel collectives.
+    parallel_mode : {None, 'column', 'row'}, default = `None`
+                   used to decide whether this Linear layer is Column Parallel Linear or Row
+                   Parallel Linear as described `here <https://arxiv.org/pdf/1909.08053.pdf>`_.
+                   When set to `None`, no communication is performed.
+
+    Optimization parameters
+    -----------------------
+    fuse_wgrad_accumulation : bool, default = 'False'
+                             if set to `True`, enables fusing of creation and accumulation of
+                             the weight gradient. When enabled, it is assumed that the weights
+                             have an additional `main_grad` attribute (used instead of the
+                             regular `grad`) which is a pre-allocated buffer of the correct
+                             size to accumulate gradients in. This argument along with
+                             weight tensor having attribute 'overwrite_main_grad' set to True
+                             will overwrite `main_grad` instead of accumulating.
+    return_bias : bool, default = `False`
+                 when set to `True`, this module will not apply the additive bias itself, but
+                 instead return the bias value during the forward pass together with the
+                 output of the linear transformation :math:`y = xA^T`. This is useful when
+                 the bias addition can be fused to subsequent operations.
+    params_dtype : torch.dtype, default = `torch.get_default_dtype()`
+                  it controls the type used to allocate the initial parameters. Useful when
+                  the model is trained with lower precision and the original FP32 parameters
+                  would not fit in GPU memory.
+    delay_wgrad_compute : bool, default = `False`
+                         Whether or not to delay weight gradient computation. If set to `True`,
+                         it's the user's responsibility to call `module.backward_dw` to compute
+                         weight gradients.
+    symmetric_ar_type : {None, 'multimem_all_reduce', 'two_shot', 'one_shot'}, default = None
+                   Type of symmetric memory all-reduce to use during the forward pass.
+                   This can help in latency bound communication situations.
+                   Requires PyTorch version 2.7.0 or higher. When set to None, standard all-reduce
+                   is used.
+    """
+
+    def __init__(
+        self,
+        in_features: int,
+        out_features: int,
+        eps: float = 1e-5,
+        sequence_parallel: bool = False,
+        fuse_wgrad_accumulation: bool = False,
+        tp_group: Optional[dist_group_type] = None,
+        tp_size: int = 1,
+        get_rng_state_tracker: Optional[Callable] = None,
+        init_method: Optional[Callable] = None,
+        bias: bool = True,
+        normalization: str = "LayerNorm",
+        return_bias: bool = False,
+        params_dtype: Optional[torch.dtype] = None,
+        parallel_mode: Optional[str] = None,
+        return_layernorm_output: bool = False,
+        return_layernorm_output_gathered: bool = False,
+        parameters_split: Optional[Union[Tuple[str, ...], Dict[str, int]]] = None,
+        zero_centered_gamma: bool = False,
+        device: Union[torch.device, str] = "cuda",
+        ub_overlap_ag: bool = False,
+        ub_overlap_rs: bool = False,
+        ub_overlap_rs_dgrad: bool = False,
+        ub_bulk_wgrad: bool = False,
+        ub_bulk_dgrad: bool = False,
+        ub_name: Optional[str] = None,
+        delay_wgrad_compute: bool = False,
+        symmetric_ar_type: Optional[str] = None,
+        name: str = None,
+    ) -> None:
+        super().__init__()
+
+        params_dtype = torch.get_default_dtype() if params_dtype is None else params_dtype
+        self.in_features = in_features
+        self.out_features = out_features
+        self.fuse_wgrad_accumulation = fuse_wgrad_accumulation
+        self.normalization = normalization
+        assert normalization in ["LayerNorm", "RMSNorm"], "Unsupported normalization type!"
+        self.use_bias = bias
+        self.return_bias = return_bias
+        self.apply_bias = self.use_bias and not return_bias
+        self.return_layernorm_output = return_layernorm_output
+        self.return_layernorm_output_gathered = (
+            return_layernorm_output_gathered if return_layernorm_output else False
+        )
+        self.zero_centered_gamma = zero_centered_gamma
+        self.symmetric_ar_type = symmetric_ar_type
+
+        self.wgrad_store = WeightGradStore(delay_wgrad_compute, ub_bulk_wgrad)
+        self.name = name
+
+        if tp_group is None:
+            self.tp_size = tp_size
+            if tp_size == 1:
+                self.set_tensor_parallel_group(tp_group)
+        else:
+            self.tp_size = get_distributed_world_size(tp_group)
+            self.set_tensor_parallel_group(tp_group)
+        self.set_nccl_overlap_warning_if_tp()
+
+        self.parallel_mode = parallel_mode
+        assert (
+            self.parallel_mode in GemmParallelModes
+        ), f"parallel_mode {parallel_mode} not supported"
+
+        if self.parallel_mode == "column":
+            self.out_features = divide(self.out_features, self.tp_size)
+        elif self.parallel_mode == "row":
+            self.in_features = divide(self.in_features, self.tp_size)
+
+        if init_method is None:
+            init_method = get_default_init_method()
+
+        self.sequence_parallel = (self.tp_size > 1) and sequence_parallel
+
+        # Column-parallel overlaps
+        self.ub_overlap_ag_fprop = (
+            ub_overlap_ag and self.sequence_parallel and self.parallel_mode == "column"
+        )
+        self.ub_overlap_rs_dgrad = (
+            ub_overlap_rs_dgrad and self.sequence_parallel and self.parallel_mode == "column"
+        )
+        self.ub_bulk_wgrad = (
+            ub_bulk_wgrad
+            and self.sequence_parallel
+            and self.parallel_mode == "column"
+            and not self.ub_overlap_rs_dgrad
+        )
+        self.ub_bulk_dgrad = (
+            ub_bulk_dgrad
+            and self.sequence_parallel
+            and self.parallel_mode == "column"
+            and not self.ub_overlap_rs_dgrad
+        )
+
+        # Row-parallel overlaps
+        self.ub_overlap_rs_fprop = (
+            ub_overlap_rs and self.sequence_parallel and self.parallel_mode == "row"
+        )
+        self.ub_overlap_ag_dgrad = (
+            ub_overlap_ag and self.sequence_parallel and self.parallel_mode == "row"
+        )
+        if any(
+            [
+                self.ub_overlap_ag_fprop,
+                self.ub_overlap_rs_dgrad,
+                self.ub_bulk_dgrad,
+                self.ub_bulk_wgrad,
+                self.ub_overlap_rs_fprop,
+                self.ub_overlap_ag_dgrad,
+            ]
+        ):
+            assert ub_name is not None, "Userbuffer name [string] is not set."
+        self.ub_name = ub_name
+
+        if self.symmetric_ar_type is not None:
+            assert torch_version() >= (
+                2,
+                7,
+                0,
+            ), "Torch version must be at least 2.7 to use symmetric memory"
+
+        self.eps = eps
+        layer_norm_weight = torch.nn.Parameter(
+            torch.empty(self.in_features, device=device, dtype=params_dtype)
+        )
+        self.register_parameter(
+            "layer_norm_weight",
+            layer_norm_weight,
+            init_fn=init_method_constant(float(not self.zero_centered_gamma)),
+        )
+        if self.normalization != "RMSNorm":
+            layer_norm_bias = torch.nn.Parameter(
+                torch.empty(self.in_features, device=device, dtype=params_dtype)
+            )
+            self.register_parameter(
+                "layer_norm_bias", layer_norm_bias, init_fn=init_method_constant(0.0)
+            )
+        else:
+            self.layer_norm_bias = None
+
+        # Initialize params in FP8
+        with_fp8_params = FP8GlobalStateManager.with_fp8_parameters()
+
+        # Contiguous buffers for params
+        weight_tensor = torch.empty(
+            self.out_features,
+            self.in_features,
+            device=device,
+            dtype=params_dtype,
+        )
+        bias_tensor = None
+        if self.use_bias:
+            bias_tensor = torch.empty(
+                self.out_features,
+                device=device,
+                dtype=params_dtype,
+            )
+
+        # Configure parameter splits
+        self.weight_names = []
+        self.bias_names = []
+        self.parameter_split_sizes = []
+        if parameters_split is None:
+            # Split into a single parameter by default
+            self.weight_names = ["weight"]
+            self.bias_names = ["bias"]
+            self.parameter_split_sizes = [out_features]
+        elif not parameters_split:
+            raise ValueError("Cannot split weight buffer into 0 parameters")
+        elif isinstance(parameters_split, dict):
+            # Split parameters with provided sizes
+            for name, split_size in parameters_split.items():
+                self.weight_names.append(f"{name.rstrip('_')}_weight")
+                self.bias_names.append(f"{name.rstrip('_')}_bias")
+                self.parameter_split_sizes.append(split_size)
+        elif all(isinstance(name, str) for name in parameters_split):
+            # Split parameters evenly
+            split_size = out_features // len(parameters_split)
+            for name in parameters_split:
+                self.weight_names.append(f"{name.rstrip('_')}_weight")
+                self.bias_names.append(f"{name.rstrip('_')}_bias")
+                self.parameter_split_sizes.append(split_size)
+        else:
+            raise TypeError("Invalid configuration for parameters split")
+
+        # Make sure parameter splits are valid
+        if sum(self.parameter_split_sizes) != out_features:
+            raise ValueError(
+                f"Trying to split weight buffer ({out_features=}) "
+                f"with split sizes {self.parameter_split_sizes}"
+            )
+
+        # Adjust parameter splits for tensor-parallel distribution
+        if self.parallel_mode == "column":
+            for i, size in enumerate(self.parameter_split_sizes):
+                if size % self.tp_size != 0:
+                    raise RuntimeError(
+                        f"Attempting to distribute a parameter with out_features={size} "
+                        f"between {self.tp_size} tensor-parallel processes"
+                    )
+                self.parameter_split_sizes[i] = size // self.tp_size
+
+        # Construct weight parameters
+        # Note: Register weights together so that they are adjacent to
+        # each other in LayerNormLinear.parameters(). This makes it
+        # more likely that they will stay contiguous if the weights
+        # are manipulated externally, e.g. by FSDP.
+        offset = 0
+        for i, split_size in enumerate(self.parameter_split_sizes):
+            split_start = offset
+            offset += split_size
+            split_end = offset
+
+            # Check if parameters are subviews of buffers
+            is_subview = (split_start, split_end) != (0, self.out_features)
+            if is_subview and with_fp8_params:
+                raise RuntimeError(
+                    "Splitting QuantizedTensor into multiple params is not supported"
+                )
+
+            # Construct weight parameter
+            self.register_parameter(
+                self.weight_names[i],
+                torch.nn.Parameter(weight_tensor[split_start:split_end]),
+                init_fn=init_method,
+                get_rng_state_tracker=get_rng_state_tracker,
+                fp8_meta_index=tex.FP8FwdTensors.GEMM1_WEIGHT,
+            )
+
+        # Construct bias parameters if needed
+        if self.use_bias:
+            offset = 0
+            for i, split_size in enumerate(self.parameter_split_sizes):
+                split_start = offset
+                offset += split_size
+                split_end = offset
+                self.register_parameter(
+                    self.bias_names[i],
+                    torch.nn.Parameter(bias_tensor[split_start:split_end]),
+                    init_fn=init_method_constant(0.0),
+                )
+        else:
+            for name in self.bias_names:
+                bias = torch.Tensor().to(dtype=params_dtype, device=device)
+                setattr(self, name, bias)
+
+        if with_fp8_params:
+            self.init_fp8_metadata()
+
+        self.reset_parameters(defer_init=device == "meta")
+
+        # For RPL, bias has to be added after TP collectives
+        # So it cannot be fused with the GEMM
+        if self.parallel_mode == "row" and self.apply_bias:
+            self.gemm_bias_unfused_add = True
+        else:
+            self.gemm_bias_unfused_add = False
+
+        # These many SMs are subtracted from the total SM count when calling forward
+        # and backward LayerNorm C APIs. These envvars can be used to prevent the LN
+        # kernels from using all SMs in the device. This is useful for cases such as
+        # communication overlap with LN.
+        self.fwd_ln_sm_margin = int(os.getenv("NVTE_FWD_LAYERNORM_SM_MARGIN", "0"))
+        self.bwd_ln_sm_margin = int(os.getenv("NVTE_BWD_LAYERNORM_SM_MARGIN", "0"))
+        self.inf_ln_sm_margin = int(os.getenv("NVTE_INF_LAYERNORM_SM_MARGIN", "0"))
+
+        if self.wgrad_store.delay_wgrad_compute():
+            for name, param in self.named_parameters():
+                if name in self.weight_names or name in self.bias_names:
+                    param.skip_backward_post_hook = True
+
+    def set_meta_tensor(self, fwd: bool, recipe: Recipe) -> None:
+        """Init scales and amaxes for fwd | bwd."""
+        super().set_meta_tensor(fwd, recipe)
+
+        # customize quantizers based on each recipe & layer configs
+        recipe = FP8GlobalStateManager.get_fp8_recipe()
+        if recipe.float8_current_scaling():
+            self._customize_quantizers_float8_current_scaling(fwd, recipe)
+        elif recipe.float8_block_scaling():
+            self._customize_quantizers_float8_blockwise_scaling(fwd, recipe)
+        elif recipe.nvfp4():
+            self._customize_quantizers_nvfp4(fwd, recipe)
+        # elif other recipes (mxfp8, etc)
+
+    def reset_layer_norm_parameters(self) -> None:
+        """Init LN params"""
+        warnings.warn(
+            "This method will be deprecated in an upcoming release. "
+            "Update your code to use LayerNormLinear.reset_parameters() instead.",
+            DeprecationWarning,
+            stacklevel=2,
+        )
+        if not self.zero_centered_gamma:
+            init.ones_(self.layer_norm_weight)
+        else:
+            init.zeros_(self.layer_norm_weight)
+        if self.layer_norm_bias is not None:
+            init.zeros_(self.layer_norm_bias)
+
+    def reset_parameters(self, defer_init=False):
+        super().reset_parameters(defer_init=defer_init)
+
+        if not defer_init:
+            # Set parallelism attributes for layer norm parameters
+            setattr(self.layer_norm_weight, "sequence_parallel", self.sequence_parallel)
+            if self.normalization != "RMSNorm":
+                setattr(self.layer_norm_bias, "sequence_parallel", self.sequence_parallel)
+
+            # Set parallelism attributes for linear weights
+            for weight in self.weight_names:
+                set_tensor_model_parallel_attributes(
+                    tensor=getattr(self, weight),
+                    is_parallel=True,
+                    dim=1 if self.parallel_mode == "row" else 0,
+                    stride=1,
+                )
+
+            # Set parallelism attributes for linear biases
+            if self.use_bias:
+                for bias in self.bias_names:
+                    if self.parallel_mode == "row":
+                        setattr(getattr(self, bias), "sequence_parallel", self.sequence_parallel)
+                    elif self.parallel_mode == "column":
+                        set_tensor_model_parallel_attributes(getattr(self, bias), True, 0, 1)
+
+    @no_torch_dynamo()
+    def forward(
+        self,
+        inp: torch.Tensor,
+        is_first_microbatch: Optional[bool] = None,
+        fp8_output: Optional[bool] = False,
+        fp8_grad: Optional[bool] = False,
+    ) -> Union[torch.Tensor, Tuple[torch.Tensor, ...]]:
+        """
+        Apply layer normalization to the input followed by a linear transformation.
+
+        Parameters
+        ----------
+        inp : torch.Tensor
+             Input tensor.
+        is_first_microbatch : {True, False, None}, default = None
+                             During training using either gradient accumulation or
+                             pipeline parallelism a minibatch of data is further split
+                             into microbatches. Between the microbatches of the same minibatch
+                             the model weights are not updated. Setting this parameter indicates
+                             whether the current microbatch is the first in a minibatch or not.
+                             When set, this parameter enables additional optimizations:
+
+                             * during FP8 training, it allows caching of the FP8 versions of
+                               the weights
+                             * it also allows skipping gradient accumulation during the
+                               first microbatch (since it is the first gradient being
+                               produced)
+        """
+        if is_in_onnx_export_mode():
+            return self.onnx_forward(inp, fp8_output)
+
+        debug = self.is_debug_iter()
+
+        if FP8GlobalStateManager.fp8_graph_capturing():
+            skip_fp8_weight_update = FP8GlobalStateManager.get_skip_fp8_weight_update_tensor()
+        else:
+            skip_fp8_weight_update = None
+        if skip_fp8_weight_update is not None:
+            is_first_microbatch = False
+
+        if self.ub_overlap_rs_fprop:
+            if get_ub(
+                self.ub_name + "_fprop", FP8GlobalStateManager.is_fp8_enabled()
+            ).is_fp8_ubuf():
+                fp8_output = True
+        if self.ub_overlap_rs_dgrad:
+            if get_ub(
+                self.ub_name + "_dgrad", FP8GlobalStateManager.is_fp8_enabled()
+            ).is_fp8_ubuf():
+                fp8_grad = True
+
+        with torch.cuda.device(
+            getattr(self, list(self.named_parameters())[0][0]).device
+        ), self.prepare_forward(
+            inp, allow_non_contiguous=False  # removed .contiguous from inside the layer
+        ) as inp:
+
+            # Get concatenated weight and bias tensors
+            weight_tensor, bias_tensor = self._get_weight_and_bias_tensors()
+
+            quantizers = (
+                self._get_quantizers(fp8_output, fp8_grad)
+                if not debug
+                else self._get_debug_quantizers(fp8_output, fp8_grad)
+            )
+            if debug:
+                if self.no_debug_features_active(quantizers):
+                    debug = False
+                    quantizers = self._get_quantizers(fp8_output, fp8_grad)
+
+            (
+                input_quantizer,
+                weight_quantizer,
+                output_quantizer,
+                grad_input_quantizer,
+                grad_weight_quantizer,
+                grad_output_quantizer,
+            ) = quantizers
+
+            if torch.is_grad_enabled():
+                fwd_fn = _LayerNormLinear.apply
+                args = []
+            else:
+                fwd_fn = _LayerNormLinear.forward
+                args = [None]
+            args += (
+                inp,
+                self.layer_norm_weight,
+                self.layer_norm_bias,
+                weight_tensor,
+                bias_tensor if self.apply_bias and not self.gemm_bias_unfused_add else None,
+                self.eps,
+                is_first_microbatch,
+                self.fp8,
+                self.fp8_calibration,
+                self.wgrad_store,
+                self.fuse_wgrad_accumulation,
+                input_quantizer,
+                weight_quantizer,
+                output_quantizer,
+                grad_input_quantizer,
+                grad_weight_quantizer,
+                grad_output_quantizer,
+                is_cpu_offload_enabled(),
+                self.tp_group,
+                self.tp_size,
+                self.sequence_parallel,
+                self.tp_size > 1,
+                self.activation_dtype,
+                self.parallel_mode,
+                self.return_layernorm_output,
+                self.return_layernorm_output_gathered,
+                torch.is_grad_enabled(),
+                self.fwd_ln_sm_margin if torch.is_grad_enabled() else self.inf_ln_sm_margin,
+                self.bwd_ln_sm_margin,
+                self.zero_centered_gamma,
+                self.normalization,
+                self.ub_overlap_ag_fprop,
+                self.ub_overlap_rs_fprop,
+                self.ub_overlap_ag_dgrad,
+                self.ub_overlap_rs_dgrad,
+                self.ub_bulk_wgrad,
+                self.ub_bulk_dgrad,
+                self.ub_name,
+                self.fsdp_group,
+                self,
+                skip_fp8_weight_update,
+                self.symmetric_ar_type,
+                debug,
+            )
+            out = fwd_fn(*args)
+
+        if self.return_layernorm_output:
+            out, ln_out = out
+
+        if self.gemm_bias_unfused_add:
+            out = out + cast_if_needed(bias_tensor, self.activation_dtype)
+
+        if self.return_bias:
+            if self.return_layernorm_output:
+                return out, cast_if_needed(bias_tensor, self.activation_dtype), ln_out
+            return out, cast_if_needed(bias_tensor, self.activation_dtype)
+        if self.return_layernorm_output:
+            return out, ln_out
+        return out
+
+    def _get_quantizers(self, fp8_output, fp8_grad):
+        if not self.fp8:
+            return [None] * 6
+        grad_input_quantizer = None
+        grad_weight_quantizer = None
+        grad_output_quantizer = None
+        output_quantizer = None
+        input_quantizer = self.quantizers["scaling_fwd"][tex.FP8FwdTensors.GEMM1_INPUT]
+        input_quantizer.internal = True
+        (weight_quantizer,) = self._get_weight_quantizers()
+        if fp8_output:
+            output_quantizer = self.quantizers["scaling_fwd"][tex.FP8FwdTensors.GEMM1_OUTPUT]
+        if torch.is_grad_enabled():
+            grad_output_quantizer = self.quantizers["scaling_bwd"][tex.FP8BwdTensors.GRAD_OUTPUT1]
+            grad_output_quantizer.internal = True
+            if fp8_grad:
+                grad_input_quantizer = self.quantizers["scaling_bwd"][tex.FP8BwdTensors.GRAD_INPUT1]
+
+        return (
+            input_quantizer,
+            weight_quantizer,
+            output_quantizer,
+            grad_input_quantizer,
+            grad_weight_quantizer,
+            grad_output_quantizer,
+        )
+
+    def _get_debug_quantizers(self, fp8_output, fp8_grad):
+        original_quantizers = self._get_quantizers(fp8_output, fp8_grad)
+        assert TEDebugState.debug_enabled
+        from ...debug.pytorch.debug_quantization import DebugQuantizer
+
+        names = ["activation", "weight", "output", "dgrad", "wgrad", "gradient"]
+        return tuple(
+            DebugQuantizer(self.name, name, q, self.tp_group)
+            for name, q in zip(names, original_quantizers)
+        )
+
+    def _get_weight_and_bias_tensors(self):
+        # Get concatenated weight and bias tensors
+        unfused_weights = self._get_weight_tensors()
+
+        weight_tensor = noop_cat(unfused_weights)
+        if self.use_bias:
+            bias_tensor = noop_cat([getattr(self, name) for name in self.bias_names])
+        else:
+            bias_tensor = getattr(self, self.bias_names[0])  # Unused
+        return weight_tensor, bias_tensor
+
+    def onnx_forward(
+        self,
+        inp: torch.Tensor,
+        fp8_output: bool,
+    ) -> torch.Tensor:
+        """
+        ONNX-compatible version of the forward function that provides numerical equivalence
+        while only using operations that have defined ONNX symbolic translations.
+        This simplified implementation is designed specifically for inference scenarios.
+        """
+        from ..export import onnx_layernorm, onnx_gemm
+
+        assert not TEDebugState.debug_enabled, "Debug mode is not supported in ONNX export"
+        assert_warmed_up(self)
+        (
+            input_quantizer,
+            weight_quantizer,
+            output_quantizer,
+            *_,
+        ) = self._get_quantizers(fp8_output, fp8_grad=False)
+        inp_dtype = inp.dtype
+
+        weight_tensor, bias_tensor = self._get_weight_and_bias_tensors()
+        ln_out, ln_out_return = onnx_layernorm(
+            inp,
+            self.layer_norm_weight,
+            self.layer_norm_bias,
+            self.eps,
+            self.normalization,
+            self.zero_centered_gamma,
+            inp_dtype,
+            self.return_layernorm_output,
+            input_quantizer,
+        )
+
+        if weight_quantizer is not None:
+            weight_tensor_quantized = weight_quantizer.onnx_quantize(weight_tensor)
+            weight_tensor = weight_quantizer.onnx_dequantize(weight_tensor_quantized)
+        weight_tensor = weight_tensor.to(inp_dtype)
+
+        if bias_tensor is not None:
+            bias_tensor = bias_tensor.to(inp_dtype)
+
+        output = onnx_gemm(weight_tensor, ln_out, bias_tensor if self.apply_bias else None)
+
+        if output_quantizer is not None:
+            raise NotImplementedError("ONNX export of quantized output is not supported")
+        if self.return_layernorm_output and self.return_bias:
+            return output, bias_tensor.to(inp_dtype), ln_out_return
+        if self.return_layernorm_output:
+            return output, ln_out_return
+        if self.return_bias:
+            return output, bias_tensor.to(inp_dtype)
+        return output
+
+    def _customize_quantizers_float8_current_scaling(self, fwd: bool, recipe: Recipe) -> None:
+        """Customize quantizers based on current scaling recipe + layernorm_linear."""
+        assert (
+            recipe.float8_current_scaling()
+        ), "current scaling recipe quantizer customization here"
+        if fwd:
+            # set configs about amax epsilon and power_2_scale
+            self.quantizers["scaling_fwd"][
+                tex.FP8FwdTensors.GEMM1_INPUT
+            ].force_pow_2_scales = recipe.fp8_quant_fwd_inp.power_2_scale
+            self.quantizers["scaling_fwd"][
+                tex.FP8FwdTensors.GEMM1_INPUT
+            ].amax_epsilon = recipe.fp8_quant_fwd_inp.amax_epsilon
+            # also set weight quantizer with same amax_epsilon & power_2_scale
+            self.quantizers["scaling_fwd"][
+                tex.FP8FwdTensors.GEMM1_WEIGHT
+            ].force_pow_2_scales = recipe.fp8_quant_fwd_weight.power_2_scale
+            self.quantizers["scaling_fwd"][
+                tex.FP8FwdTensors.GEMM1_WEIGHT
+            ].amax_epsilon = recipe.fp8_quant_fwd_weight.amax_epsilon
+            # parallel related
+            if self.sequence_parallel and self.parallel_mode == "column":
+                # set input_quantizer with amax reduction TP group
+                self.quantizers["scaling_fwd"][
+                    tex.FP8FwdTensors.GEMM1_INPUT
+                ].with_amax_reduction = True
+                self.quantizers["scaling_fwd"][
+                    tex.FP8FwdTensors.GEMM1_INPUT
+                ].amax_reduction_group = self.tp_group
+        else:
+            # set grad_output_quantizer with amax epsilon and power_2_scale (no amax reduction here)
+            self.quantizers["scaling_bwd"][
+                tex.FP8BwdTensors.GRAD_OUTPUT1
+            ].force_pow_2_scales = recipe.fp8_quant_bwd_grad.power_2_scale
+            self.quantizers["scaling_bwd"][
+                tex.FP8BwdTensors.GRAD_OUTPUT1
+            ].amax_epsilon = recipe.fp8_quant_bwd_grad.amax_epsilon
+            # parallel related
+            if self.sequence_parallel and self.parallel_mode == "row":
+                # customize grad_output_quantizer with amax reduction TP group
+                self.quantizers["scaling_bwd"][
+                    tex.FP8BwdTensors.GRAD_OUTPUT1
+                ].with_amax_reduction = True
+                self.quantizers["scaling_bwd"][
+                    tex.FP8BwdTensors.GRAD_OUTPUT1
+                ].amax_reduction_group = self.tp_group
+
+    def _customize_quantizers_nvfp4(self, fwd: bool, recipe: Recipe) -> None:
+        """Customize quantizers based on current scaling recipe + layernorm_linear."""
+        assert recipe.nvfp4(), "Incorrect recipe."
+        if fwd:
+            if self.sequence_parallel and self.parallel_mode == "column":
+                # set input_quantizer with amax reduction TP group
+                self.quantizers["scaling_fwd"][
+                    tex.FP8FwdTensors.GEMM1_INPUT
+                ].with_amax_reduction = True
+                self.quantizers["scaling_fwd"][
+                    tex.FP8FwdTensors.GEMM1_INPUT
+                ].amax_reduction_group = self.tp_group
+        else:
+            if self.sequence_parallel and self.parallel_mode == "row":
+                # customize grad_output_quantizer with amax reduction TP group
+                self.quantizers["scaling_bwd"][
+                    tex.FP8BwdTensors.GRAD_OUTPUT1
+                ].with_amax_reduction = True
+                self.quantizers["scaling_bwd"][
+                    tex.FP8BwdTensors.GRAD_OUTPUT1
+                ].amax_reduction_group = self.tp_group
+
+    def _get_weight_tensors(self) -> List[Union[torch.Tensor, QuantizedTensorStorage]]:
+        """Get the weight tensors of the module."""
+        unfused_weights = [getattr(self, name) for name in self.weight_names]
+        if any(isinstance(w, QuantizedTensor) for w in unfused_weights):
+            if self.fp8:
+                if len(unfused_weights) != 1:
+                    raise RuntimeError(
+                        "Splitting QuantizedTensor into multiple params is not supported"
+                    )
+            else:
+                warnings.warn(
+                    "You are using quantized weights without quantized compute. "
+                    "Please make sure this is intentional."
+                )
+                unfused_weights = [w.dequantize() for w in unfused_weights]
+        return unfused_weights
+
+    def _get_weight_quantizers(self) -> List[Quantizer]:
+        """Get the weight quantizers of the module."""
+        if not self.fp8 and not self.fp8_calibration:
+            return [None]
+        weight_quantizer = self.quantizers["scaling_fwd"][tex.FP8FwdTensors.GEMM1_WEIGHT]
+        weight_quantizer.internal = True
+        return [weight_quantizer]
+
+    def _customize_quantizers_float8_blockwise_scaling(self, fwd: bool, recipe: Recipe) -> None:
+        """Customize quantizers based on blockwise scaling recipe + layernorm_linear."""
+        assert (
+            recipe.float8_block_scaling()
+        ), "blockwise scaling recipe quantizer customization here"
+        if fwd:
+            if self.sequence_parallel and self.parallel_mode == "column":
+                self.quantizers["scaling_fwd"][
+                    tex.FP8FwdTensors.GEMM1_INPUT
+                ].all_gather_usage = True
diff --git a/examples/tp-numerics/patches/v2.9.0/transformer_engine/pytorch/module/linear.py b/examples/tp-numerics/patches/v2.9.0/transformer_engine/pytorch/module/linear.py
new file mode 100644
index 00000000000..7d5374e7b2b
--- /dev/null
+++ b/examples/tp-numerics/patches/v2.9.0/transformer_engine/pytorch/module/linear.py
@@ -0,0 +1,1827 @@
+# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+#
+# See LICENSE for license information.
+
+"""Linear API"""
+import os
+from typing import Callable, Dict, Optional, Tuple, Union, List
+from functools import reduce
+from operator import mul as multiply_op
+import warnings
+
+import torch
+
+import transformer_engine_torch as tex
+
+from transformer_engine.common.recipe import Recipe
+from transformer_engine.pytorch import torch_version
+
+from .base import (
+    fill_userbuffers_buffer_for_all_gather,
+    get_dummy_wgrad,
+    get_ub,
+    get_workspace,
+    TransformerEngineBaseModule,
+    _2X_ACC_FPROP,
+    _2X_ACC_DGRAD,
+    _2X_ACC_WGRAD,
+)
+from ._common import noop_cat, WeightGradStore
+from ..quantization import FP8GlobalStateManager
+from ..utils import (
+    cast_if_needed,
+    clear_tensor_data,
+    divide,
+    init_method_constant,
+    requires_grad,
+    needs_quantized_gemm,
+    assert_dim_for_fp8_exec,
+    assert_dim_for_all_gather,
+    nvtx_range_pop,
+    nvtx_range_push,
+)
+from ..distributed import (
+    set_tensor_model_parallel_attributes,
+    get_distributed_world_size,
+    allreduce,
+    symmetric_all_reduce,
+    reduce_scatter_along_first_dim,
+    gather_along_first_dim,
+    is_fp8_activation_recompute_enabled,
+    in_fp8_activation_recompute_phase,
+    _fsdp_scatter_tensors,
+    _fsdp_gather_tensors,
+)
+from ..cpp_extensions import (
+    general_gemm,
+)
+from ..constants import GemmParallelModes, dist_group_type
+from ..jit import no_torch_dynamo
+from ..graph import is_graph_capturing
+from ..tensor.quantized_tensor import (
+    QuantizedTensor,
+    QuantizedTensorStorage,
+    Quantizer,
+    prepare_for_saving,
+    restore_from_saved,
+)
+from ..tensor.float8_tensor import Float8CurrentScalingQuantizer, Float8Quantizer
+from ..tensor.mxfp8_tensor import MXFP8Quantizer
+from ..tensor.utils import is_experimental
+from ..export import is_in_onnx_export_mode, assert_warmed_up
+from ..cpu_offload import is_cpu_offload_enabled, mark_activation_offload
+from ...debug.pytorch.debug_state import TEDebugState
+
+__all__ = ["Linear"]
+
+
+class _Linear(torch.autograd.Function):
+    """Linear semi-top level module
+    Calls custom cuda extensions.
+    """
+
+    @staticmethod
+    def forward(
+        ctx,
+        weight: torch.Tensor,
+        inp: torch.Tensor,
+        bias: Optional[torch.Tensor],
+        is_first_microbatch: Union[bool, None],
+        fp8: bool,
+        fp8_calibration: bool,
+        wgrad_store: WeightGradStore,
+        input_quantizer: Optional[Quantizer],
+        weight_quantizer: Optional[Quantizer],
+        output_quantizer: Optional[Quantizer],
+        grad_input_quantizer: Optional[Quantizer],
+        grad_weight_quantizer: Optional[Quantizer],
+        grad_output_quantizer: Optional[Quantizer],
+        fuse_wgrad_accumulation: bool,
+        cpu_offloading: bool,
+        tp_group: Union[dist_group_type, None],
+        tp_size: int,
+        sequence_parallel: bool,
+        tensor_parallel: bool,
+        activation_dtype: torch.dtype,
+        parallel_mode: Union[str, None],
+        is_grad_enabled: bool,
+        ub_overlap_rs_fprop: bool,
+        ub_overlap_ag_dgrad: bool,
+        ub_overlap_ag_fprop: bool,
+        ub_overlap_rs_dgrad: bool,
+        ub_bulk_dgrad: bool,
+        ub_bulk_wgrad: bool,
+        ub_name: str,
+        fp8_output: bool,  # pylint: disable=unused-argument
+        fsdp_group: Union[dist_group_type, None],
+        module: torch.nn.Module,
+        skip_fp8_weight_update: bool,
+        symmetric_ar_type: str,
+        save_original_input: bool = False,
+        debug: Optional[bool] = False,
+    ) -> torch.Tensor:
+        # pylint: disable=missing-function-docstring
+
+        # NVTX label for profiling
+        nvtx_label = "transformer_engine._Linear.forward"
+        if ub_name is not None:
+            nvtx_label = f"{nvtx_label}.{ub_name}"
+
+        # Make sure input dimensions are compatible
+        out_features, in_features = weight.shape
+        assert inp.shape[-1] == in_features, "GEMM not possible"
+
+        # Configure tensor-parallel communication
+        tp_world_size = get_distributed_world_size(tp_group)
+        backward_needs_input = is_grad_enabled and weight.requires_grad
+        with_input_all_gather_nccl = (
+            parallel_mode == "column" and sequence_parallel and not ub_overlap_ag_fprop
+        )
+
+        # Configure Userbuffers communication (comm+GEMM overlap)
+        if debug:  # turn off userbuffers in debug mode
+            ub_overlap_rs_fprop = False
+            ub_overlap_ag_fprop = False
+            ub_overlap_rs_dgrad = False
+            ub_bulk_wgrad = False
+            ub_bulk_dgrad = False
+        ub_obj = None
+        ub_type = None
+        if ub_overlap_rs_fprop:
+            ub_obj = get_ub(ub_name + "_fprop", fp8)
+            ub_type = tex.CommOverlapType.RS
+        elif ub_overlap_ag_fprop:
+            ub_obj = get_ub(ub_name + "_fprop", fp8)
+            ub_type = tex.CommOverlapType.AG
+
+        # experimental recipe check
+        experimental = is_experimental(input_quantizer) or is_experimental(weight_quantizer)
+
+        # ------------------------------------------------------
+        # Prepare input tensor
+        # Note: Cast to expected dtype and perform tensor-parallel communication
+        # ------------------------------------------------------
+        nvtx_range_push(f"{nvtx_label}.input_cast_comm")
+        inputmat = inp  # Input tensor to save for backward (maybe sharded)
+        inputmat_total = None  # Input tensor to pass to GEMM (gathered)
+        own_quantized_input = False
+        if fp8:
+            assert_dim_for_fp8_exec(inputmat, weight)
+            assert_dim_for_all_gather(inputmat, with_input_all_gather_nccl, input_quantizer)
+            if save_original_input:
+                assert not isinstance(
+                    input_quantizer, Float8Quantizer
+                ), "DelayedScaling recipe is not supported with save_original_input"
+
+        if with_input_all_gather_nccl or ub_overlap_ag_fprop:  # All-gather input tensor
+
+            # Cast local input tensor if needed
+            if fp8 or debug:
+                if input_quantizer is None:
+                    raise ValueError("Missing quantizer for input tensor")
+                if not isinstance(inputmat, QuantizedTensorStorage) and not experimental:
+                    own_quantized_input = True
+                    input_quantizer.set_usage(rowwise=True, columnwise=backward_needs_input)
+                    if isinstance(
+                        input_quantizer, (Float8Quantizer, Float8CurrentScalingQuantizer)
+                    ):
+                        # All-gather is not supported with FP8 column-wise data
+                        input_quantizer.set_usage(columnwise=False)
+                    if save_original_input:
+                        # No need for column-wise data since this
+                        # tensor will not be cached for backward pass
+                        input_quantizer.set_usage(columnwise=False)
+                        own_quantized_input = False
+                    inputmat = input_quantizer(inputmat)
+            else:
+                inputmat = cast_if_needed(inp, activation_dtype)  # Cast for AMP
+
+            # Initialize gathered input tensor
+            quantizer = None
+            if fp8 or debug:
+                quantizer = input_quantizer
+                quantizer.set_usage(rowwise=True, columnwise=False)
+            if with_input_all_gather_nccl:  # Perform NCCL all-gather
+                inputmat_total, _ = gather_along_first_dim(
+                    inputmat,
+                    tp_group,
+                    quantizer=quantizer,
+                )
+            elif ub_overlap_ag_fprop:  # Initialize Userbuffers all-gather
+                inputmat_total, _ = fill_userbuffers_buffer_for_all_gather(
+                    ub_obj,
+                    inputmat,
+                    quantizer,
+                    tp_group,
+                )
+
+        else:  # Do not all-gather input tensor
+            if fp8 or debug:
+                if isinstance(inputmat, QuantizedTensorStorage):
+                    inputmat.update_usage(rowwise_usage=True)
+                else:
+                    if input_quantizer is None:
+                        raise ValueError("Missing quantizer for input tensor")
+                    input_quantizer.set_usage(
+                        rowwise=True, columnwise=backward_needs_input and not save_original_input
+                    )
+                    inputmat = input_quantizer(inputmat)
+                    own_quantized_input = True
+            else:
+                inputmat = cast_if_needed(inp, activation_dtype)  # Cast for AMP
+            inputmat_total = inputmat
+        nvtx_range_pop(f"{nvtx_label}.input_cast_comm")
+        # ------------------------------------------------------
+        # Input tensor is ready for GEMM...
+        # ------------------------------------------------------
+
+        # ------------------------------------------------------
+        # Prepare weight tensor
+        # ------------------------------------------------------
+        weightmat = weight
+        if fp8 or debug:
+            # Configure quantizer
+            if weight_quantizer is not None:
+                columnwise_usage = is_grad_enabled and inp.requires_grad
+                if not columnwise_usage:
+                    columnwise_usage = (
+                        is_fp8_activation_recompute_enabled()
+                        and not in_fp8_activation_recompute_phase()
+                    )
+                weight_quantizer.set_usage(rowwise=True, columnwise=columnwise_usage)
+
+            # Get quantized weight
+            update_workspace = is_first_microbatch is None or is_first_microbatch
+            weightmat = module.get_weight_workspace(
+                tensor=weight,
+                quantizer=weight_quantizer,
+                cache_name=(None if is_first_microbatch is None else "weight"),
+                update_workspace=update_workspace,
+                skip_update_flag=skip_fp8_weight_update,
+                fsdp_group=fsdp_group,
+                workspace_dtype=activation_dtype,
+            )
+            weightmat.update_usage(rowwise_usage=True)
+
+        else:
+            weightmat = cast_if_needed(weightmat, activation_dtype)  # Cast for AMP
+        # ------------------------------------------------------
+        # Weight tensor is ready for GEMM...
+        # ------------------------------------------------------
+
+        # Cast bias to expected dtype
+        bias_dtype = activation_dtype
+        if needs_quantized_gemm(inputmat_total) and activation_dtype == torch.float32:
+            # cuBLAS does not support FP8 GEMM with FP32 bias, so we cast to BF16
+            bias_dtype = torch.bfloat16
+        bias = cast_if_needed(bias, bias_dtype) if bias is not None else bias
+
+        # Calibrate quantizers if needed
+        if not fp8 and fp8_calibration:
+            if input_quantizer is not None:
+                input_quantizer.calibrate(inputmat_total)
+            if weight_quantizer is not None:
+                weight_quantizer.calibrate(weight)
+
+        # Choose whether to use GEMM kernel with split accumulator
+        use_split_accumulator = _2X_ACC_FPROP
+        if fp8:
+            recipe = FP8GlobalStateManager.get_fp8_recipe()
+            if hasattr(recipe, "fp8_gemm_fprop"):
+                use_split_accumulator = recipe.fp8_gemm_fprop.use_split_accumulator
+
+        # Configure output quantizer
+        if output_quantizer is not None:
+            output_quantizer.set_usage(rowwise=True, columnwise=False)
+
+        # Output buffer for Userbuffers reduce-scatter
+        reduce_scatter_out = None
+        if ub_overlap_rs_fprop:
+            out_shape = list(inp.shape)
+            out_shape[0] //= tp_world_size
+            out_shape[-1] = out_features
+            reduce_scatter_out = torch.empty(out_shape, dtype=activation_dtype, device=inp.device)
+
+        # ------------------------------------------------------
+        # Forward GEMM
+        # Note: y = x * w^T
+        # ------------------------------------------------------
+        _fp32_tp_reduce = (
+            os.environ.get("NVTE_FP32_TP_REDUCE", "0") == "1"
+            and parallel_mode == "row"
+            and tp_size > 1
+        )
+        _gemm_out_dtype = torch.float32 if _fp32_tp_reduce else activation_dtype
+
+        _tp_invariant = (
+            os.environ.get("NVTE_TP_INVARIANT_MODE", "0") == "1"
+            and parallel_mode == "row"
+            and tp_size > 1
+        )
+
+        if _tp_invariant:
+            # TP-invariant diagnostic: full GEMM matching TP=1 accumulation order.
+            assert not fp8, "NVTE_TP_INVARIANT_MODE does not support FP8"
+            nvtx_range_push(f"{nvtx_label}.tp_invariant_gemm")
+
+            def allgather_along_dim(tensor, group, world_size, dim):
+                chunks = [torch.empty_like(tensor) for _ in range(world_size)]
+                torch.distributed.all_gather(chunks, tensor.contiguous(), group=group)
+                return torch.cat(chunks, dim=dim)
+
+            inputmat_gathered = allgather_along_dim(
+                inputmat_total, tp_group, tp_world_size, dim=-1,
+            )  # [tokens, hidden/TP] -> [tokens, hidden]
+            weight_gathered = allgather_along_dim(
+                weightmat, tp_group, tp_world_size, dim=-1,
+            )  # [out, hidden/TP] -> [out, hidden]
+
+            input_2d = inputmat_gathered.reshape(-1, inputmat_gathered.shape[-1])
+            out = general_gemm(
+                weight_gathered, input_2d, get_workspace(),
+                out_dtype=activation_dtype,
+                bias=bias,
+            )
+            if isinstance(out, tuple):
+                out = out[0]
+            out = out.reshape(inputmat_gathered.shape[:-1] + (weight_gathered.shape[0],))
+
+            # SP: scatter to per-rank chunk along sequence dim.
+            if sequence_parallel:
+                rank = torch.distributed.get_rank(tp_group)
+                out = out.chunk(tp_world_size, dim=0)[rank].contiguous()
+
+            del inputmat_gathered, weight_gathered, input_2d
+            nvtx_range_pop(f"{nvtx_label}.tp_invariant_gemm")
+        else:
+            nvtx_range_push(f"{nvtx_label}.gemm")
+            gemm_out, *_, reduce_scatter_out = general_gemm(
+                weightmat,
+                inputmat_total,
+                get_workspace(),
+                quantization_params=output_quantizer,
+                out_dtype=_gemm_out_dtype,
+                bias=bias,
+                use_split_accumulator=use_split_accumulator,
+                ub=ub_obj,
+                ub_type=ub_type,
+                extra_output=reduce_scatter_out,
+            )
+            nvtx_range_pop(f"{nvtx_label}.gemm")
+
+        # ------------------------------------------------------
+        # Finished forward GEMM...
+        # ------------------------------------------------------
+
+        # Deallocate GEMM input tensor if no longer needed
+        # TODO(yuzhongw, tmoon): Figure out why inputmat_total is not automatically
+        # deallocated by GC. Manually deallocating is a temporary hack.
+        if with_input_all_gather_nccl:
+            clear_tensor_data(inputmat_total)
+            inputmat_total = None
+
+        # ------------------------------------------------------
+        # Prepare output tensor
+        # Note: Perform tensor-parallel communication
+        # ------------------------------------------------------
+        if not _tp_invariant:
+            out = None
+            if ub_overlap_rs_fprop:
+                out = reduce_scatter_out
+            elif parallel_mode == "row" and tp_size > 1:
+                nvtx_range_push(f"{nvtx_label}.row_parallel_comm")
+                out = gemm_out
+                if sequence_parallel:
+                    out, _ = reduce_scatter_along_first_dim(out, tp_group)
+                elif tensor_parallel:
+                    if symmetric_ar_type is not None:
+                        out, _ = symmetric_all_reduce(out, tp_group, all_reduce_type=symmetric_ar_type)
+                    else:
+                        out, _ = allreduce(out, tp_group)
+                if _fp32_tp_reduce and out.dtype != activation_dtype:
+                    out = out.to(activation_dtype)
+                nvtx_range_pop(f"{nvtx_label}.row_parallel_comm")
+            else:
+                out = gemm_out
+        # ------------------------------------------------------
+        # Output tensor is ready to return...
+        # ------------------------------------------------------
+
+        # ------------------------------------------------------
+        # Cache state for backward pass
+        # ------------------------------------------------------
+
+        if is_grad_enabled:
+            if save_original_input:
+                inputmat = inp
+
+            ctx.weight_quantizer = weight_quantizer
+
+            ctx.backward_input_needs_gather = (
+                weight.requires_grad and parallel_mode == "column" and sequence_parallel
+            )
+
+            # Discard unneeded data in input tensor
+            if (
+                backward_needs_input
+                and own_quantized_input
+                and isinstance(inputmat, QuantizedTensorStorage)
+            ):
+                if (
+                    ctx.backward_input_needs_gather
+                    and weight_quantizer.supports_only_rowwise_all_gather()
+                ):
+                    # All-gather is not supported with FP8 column-wise data
+                    inputmat.update_usage(rowwise_usage=True, columnwise_usage=False)
+                else:
+                    # Discard row-wise data since it is not needed in backward pass
+                    inputmat.update_usage(rowwise_usage=False, columnwise_usage=True)
+
+            # Cached input tensor
+            saved_inputmat = None
+            if backward_needs_input:
+                saved_inputmat = inputmat
+
+            # Weight with column-wise usage is needed for dgrad GEMM.
+            if inp.requires_grad:
+                if isinstance(weightmat, QuantizedTensorStorage):
+                    weightmat.update_usage(columnwise_usage=True)
+
+            if cpu_offloading and saved_inputmat is not None:
+                mark_activation_offload(saved_inputmat)
+
+            # Scatter intermediate/activation tensors saved for the backward pass
+            # NOTE: FSDP sharding is not valid for models initialized with primary Fp8 weights
+            nvtx_range_push(f"{nvtx_label}.fsdp_scatter")
+            ctx.fsdp_group = fsdp_group
+            ctx.fsdp_shapes = _fsdp_scatter_tensors(
+                fsdp_group,
+                saved_inputmat,
+                weightmat if fp8 and not isinstance(weight, QuantizedTensorStorage) else None,
+            )
+            nvtx_range_pop(f"{nvtx_label}.fsdp_scatter")
+
+            if cpu_offloading:
+                ctx.grad_added_to_main_grad = hasattr(weight, "grad_added_to_main_grad")
+
+                if ctx.grad_added_to_main_grad:
+                    # If you are passing torch.nn.Parameter through the Torch hooks, you will
+                    # get back torch.Tensor. Torch rips off the Parameter wrapper.
+                    # You need to preserve the weight object to have all the attributes user
+                    # sets for the weights. Because of this, it is not recommended to offload
+                    # weights if weights are externally touched outside this module
+                    ctx.weight_object = weight
+
+            # TODO(ksivamani): Check memory usage
+            tensors_to_save, tensor_objects = prepare_for_saving(
+                saved_inputmat,
+                weightmat,
+                weight,
+                bias,
+            )
+            ctx.save_for_backward(*tensors_to_save)
+            ctx.tensor_objects = tensor_objects
+
+            ctx.activation_dtype = activation_dtype
+            ctx.fp8 = fp8
+            ctx.fp8_recipe = FP8GlobalStateManager.get_fp8_recipe() if fp8 else None
+            ctx.input_quantizer = input_quantizer
+            ctx.grad_input_quantizer = grad_input_quantizer
+            ctx.grad_weight_quantizer = grad_weight_quantizer
+            ctx.grad_output_quantizer = grad_output_quantizer
+            ctx.fuse_wgrad_accumulation = fuse_wgrad_accumulation
+            if fuse_wgrad_accumulation and weight.requires_grad:
+                # This check is needed to ensure that main_grad is not created
+                # during the forward pass when using MCore FSDP as it creates
+                # the main_grad buffer lazily before backprop
+                if hasattr(weight, "__fsdp_param__"):
+                    # MCore FSDP creates main_grad lazily before backward
+                    ctx.main_grad_func = weight.get_main_grad
+                else:
+                    ctx.main_grad_func = lambda: weight.main_grad
+
+            ctx.debug = debug
+            ctx.experimental = experimental
+            ctx.cpu_offloading = cpu_offloading
+            ctx.is_first_microbatch = is_first_microbatch
+            ctx.use_bias = bias is not None
+            ctx.sequence_parallel = sequence_parallel
+            ctx.tensor_parallel = tensor_parallel
+            ctx.inp_shape = inp.shape
+            ctx.parallel_mode = parallel_mode
+            ctx.tp_group = tp_group
+            ctx.ub_overlap_ag = ub_overlap_ag_dgrad
+            ctx.ub_overlap_rs_dgrad = ub_overlap_rs_dgrad
+            ctx.ub_bulk_dgrad = ub_bulk_dgrad
+            ctx.ub_bulk_wgrad = ub_bulk_wgrad
+            ctx.ub_name = ub_name
+            ctx.tp_size = tp_size
+            ctx.requires_dgrad = inp.requires_grad
+            ctx.requires_wgrad = weight.requires_grad
+            ctx.reduce_and_update_bwd_fp8_tensors = False
+
+            ctx.owns_input = saved_inputmat is not inp
+            if ctx.fp8 and requires_grad(inp, weight, bias):
+                _first_fp8_module = FP8GlobalStateManager.IS_FIRST_FP8_MODULE
+                ctx.reduce_and_update_bwd_fp8_tensors = FP8GlobalStateManager.is_first_fp8_module()
+                if in_fp8_activation_recompute_phase():
+                    FP8GlobalStateManager.IS_FIRST_FP8_MODULE = _first_fp8_module
+            ctx.wgrad_store = wgrad_store
+
+        # ------------------------------------------------------
+        # Cached state for backward pass is ready...
+        # ------------------------------------------------------
+
+        return out
+
+    @staticmethod
+    def backward(ctx, grad_output: torch.Tensor) -> Tuple[Union[torch.Tensor, None], ...]:
+        # pylint: disable=missing-function-docstring
+
+        # NVTX label for profiling
+        nvtx_label = "transformer_engine._Linear.backward"
+        if ctx.ub_name is not None:
+            nvtx_label = f"{nvtx_label}.{ctx.ub_name}"
+
+        with torch.cuda.nvtx.range("_Linear_backward"):
+            saved_tensors = ctx.saved_tensors
+            inputmat, weight_fp8, weight, bias = (  # pylint: disable=unbalanced-tuple-unpacking
+                restore_from_saved(ctx.tensor_objects, saved_tensors)
+            )
+
+            # Delete the references to tensor objects once they've been consumed
+            # by the `restore_from_saved` method to construct back the actual tensors.
+            ctx.tensor_objects = None
+
+            # Since main_grad can be modified inplace, it should not be a part of saved_tensors
+            main_grad = (
+                ctx.main_grad_func()
+                if weight is not None and ctx.fuse_wgrad_accumulation and ctx.requires_wgrad
+                else None
+            )
+
+            if ctx.cpu_offloading:
+                if ctx.grad_added_to_main_grad:
+                    weight = ctx.weight_object
+                if ctx.requires_wgrad and ctx.fuse_wgrad_accumulation:
+                    weight.main_grad = main_grad
+
+            # Gather intermediate/activation tensors if needed
+            # NOTE: weight_fp8 = weight when ctx.fp8 == False and torch.disttributed.FSDP already
+            #       shards/unshards the base weights so we don't do it ourselves
+            nvtx_range_push(f"{nvtx_label}.fsdp_gather")
+            _fsdp_gather_tensors(
+                ctx.fsdp_group,
+                ctx.fsdp_shapes,
+                inputmat,
+                weight_fp8,
+            )
+            nvtx_range_pop(f"{nvtx_label}.fsdp_gather")
+
+            # Configure Userbuffers communication (comm+GEMM overlap)
+            ctx.ub_obj_gradout = None
+            ub_obj_dgrad = None
+            ub_obj_wgrad = None
+            ub_type_dgrad = None
+            ub_type_wgrad = None
+            dgrad_shape = [reduce(multiply_op, ctx.inp_shape[:-1]), ctx.inp_shape[-1]]
+            if ctx.ub_overlap_ag:
+                # Overlap grad_output all-gather with dgrad compute
+                ctx.ub_obj_gradout = get_ub(ctx.ub_name + "_dgrad", ctx.fp8)
+                ub_obj_dgrad = ctx.ub_obj_gradout
+                ub_type_dgrad = tex.CommOverlapType.AG
+            elif ctx.ub_overlap_rs_dgrad:
+                # Overlap dgrad reduce-scatter with dgrad compute
+                ctx.ub_obj_gradout = get_ub(ctx.ub_name + "_dgrad", ctx.fp8)
+                ub_obj_dgrad = ctx.ub_obj_gradout
+                ub_type_dgrad = tex.CommOverlapType.RS
+            else:
+                if ctx.ub_bulk_dgrad:
+                    # Overlap inputmat all-gather with dgrad compute
+                    ctx.ub_obj_gradout = get_ub(ctx.ub_name + "_dgrad", ctx.fp8)
+                    ub_obj_dgrad = ctx.ub_obj_gradout
+                    ub_type_dgrad = tex.CommOverlapType.AG
+                if ctx.ub_bulk_wgrad:
+                    # Overlap dgrad reduce-scatter with wgrad compute
+                    ub_obj_wgrad = get_ub(ctx.ub_name + "_wgrad", ctx.fp8)
+                    ub_type_wgrad = tex.CommOverlapType.RS
+
+            # --------------------------------------------------
+            # Prepare grad output tensor
+            # Note: Cast to expected dtype and perform tensor-parallel communication
+            # --------------------------------------------------
+
+            # Unmodified grad output tensor
+            grad_output_arg = grad_output
+
+            # Configure quantizer for grad output tensor
+            # Note: dgrad GEMM requires row-wise usage, wgrad GEMM
+            # requires column-wise usage
+            if ctx.grad_output_quantizer is not None:
+                quantizer = ctx.grad_output_quantizer
+                quantizer.set_usage(rowwise=True, columnwise=True)
+                if ctx.ub_overlap_ag:
+                    # Userbuffers only supports communication for one
+                    # tensor usage at a time. Configure quantizer with
+                    # usage for only dgrad GEMM.
+                    quantizer.set_usage(columnwise=False)
+
+            # Adjust the quantization direction approach depending
+            # on whether wgrad calculations will be performed.
+            # NOTE: If requires_dgrad is False, disabling `rowwise` quantization and keeping `columnwise` quantization
+            #       results in `Assertion failed: output_tensor->has_data(). Quantizing in only the columnwise direction not supported yet!`
+            # NOTE: For `ctx.bias is True`, selected quantize kernel errors with
+            #       `cast_kernels.cuh:1322 in function fp8_quantize_arch_l_100: Not implemented scaling mode or fusion: NVTE_DELAYED_TENSOR_SCALING or IS_DBIAS=true on GPU with compute capability < 10.0.`
+            if (
+                not ctx.use_bias
+                and not ctx.requires_wgrad
+                and ctx.grad_output_quantizer is not None
+            ):
+                ctx.grad_output_quantizer.set_usage(columnwise=False)
+
+            # Prepare grad output tensor
+            nvtx_range_push(f"{nvtx_label}.grad_output_preprocess")
+            (
+                grad_output,
+                grad_bias,
+            ) = TransformerEngineBaseModule.grad_output_preprocess(
+                ctx,
+                grad_output,
+                ctx.parallel_mode == "row",
+                ctx.grad_output_quantizer,
+            )
+            nvtx_range_pop(f"{nvtx_label}.grad_output_preprocess")
+
+            # --------------------------------------------------
+            # Grad output tensor is ready for computing grad input...
+            # --------------------------------------------------
+
+            # --------------------------------------------------
+            # Prepare input tensor
+            # Note: Input tensor is needed for wgrad GEMM.
+            # Tensor-parallel communication is overlapped with dgrad
+            # GEMM.
+            # --------------------------------------------------
+            inputmat_total = None
+            inputmat_total_work = None
+            if ctx.requires_wgrad:
+                if ctx.fp8 or ctx.debug:
+                    if isinstance(inputmat, QuantizedTensorStorage):
+                        # Input tensor is already quantized
+                        pass
+                    elif ctx.debug or ctx.experimental:
+                        # Debug quantizer will be applied immediately before wgrad GEMM
+                        pass
+                    else:
+                        # Quantize input tensor
+                        quantizer = ctx.input_quantizer
+                        if quantizer.supports_only_rowwise_all_gather():
+                            # All-gather is not supported with FP8 column-wise data
+                            quantizer.set_usage(
+                                rowwise=True,
+                                columnwise=not ctx.backward_input_needs_gather,
+                            )
+                        else:
+                            quantizer.set_usage(rowwise=False, columnwise=True)
+                        inputmat = quantizer(inputmat)
+                else:
+                    if isinstance(inputmat, QuantizedTensorStorage):
+                        inputmat = inputmat.dequantize(dtype=ctx.activation_dtype)
+                    else:
+                        inputmat = cast_if_needed(inputmat, ctx.activation_dtype)
+            if ctx.backward_input_needs_gather:
+                quantizer = None
+                if ctx.fp8 or ctx.debug:
+                    quantizer = ctx.input_quantizer
+                    if quantizer.supports_only_rowwise_all_gather():
+                        # If data is in FP8, we compute FP8 transposes manually
+                        quantizer.set_usage(rowwise=True, columnwise=False)
+                    else:
+                        # wgrad GEMM requires input with column-wise usage
+                        quantizer.set_usage(rowwise=False, columnwise=True)
+                if ctx.ub_bulk_dgrad:
+                    inputmat_total, _ = fill_userbuffers_buffer_for_all_gather(
+                        ub_obj_dgrad,
+                        inputmat,
+                        quantizer,
+                        ctx.tp_group,
+                    )
+                else:
+                    nvtx_range_push(f"{nvtx_label}.column_parallel_comm_input")
+                    inputmat_total, inputmat_total_work = gather_along_first_dim(
+                        inputmat,
+                        ctx.tp_group,
+                        async_op=True,
+                        quantizer=quantizer,
+                    )
+                    nvtx_range_pop(f"{nvtx_label}.column_parallel_comm_input")
+            else:
+                inputmat_total = inputmat
+            # --------------------------------------------------
+            # Input tensor is ready for computing grad weight...
+            # --------------------------------------------------
+
+            # --------------------------------------------------
+            # Compute grad input tensor
+            # --------------------------------------------------
+
+            dgrad = None
+            dgrad_work = None
+            if ctx.requires_dgrad:
+
+                # Make sure required data is available
+                if isinstance(grad_output, QuantizedTensorStorage):
+                    grad_output.update_usage(rowwise_usage=True)
+                if ctx.weight_quantizer is not None and isinstance(
+                    weight_fp8, QuantizedTensorStorage
+                ):
+                    weight_fp8.update_usage(columnwise_usage=True)
+
+                # Choose whether to use GEMM kernel with split accumulator
+                use_split_accumulator = _2X_ACC_DGRAD
+                if ctx.fp8:
+                    recipe = ctx.fp8_recipe
+                    if hasattr(recipe, "fp8_gemm_dgrad"):
+                        use_split_accumulator = recipe.fp8_gemm_dgrad.use_split_accumulator
+
+                # Update grad input quantizer
+                if ctx.grad_input_quantizer is not None:
+                    ctx.grad_input_quantizer.set_usage(rowwise=True, columnwise=False)
+
+                # Output buffers for Userbuffers reduce-scatter
+                gemm_out = None
+                reduce_scatter_out = None
+                if ctx.ub_overlap_rs_dgrad:
+                    reduce_scatter_out = torch.empty(
+                        dgrad_shape, dtype=ctx.activation_dtype, device=grad_output_arg.device
+                    )
+                elif ctx.ub_bulk_wgrad:
+                    gemm_out = ub_obj_wgrad.get_buffer(local_chunk=False)
+
+                # dgrad GEMM
+                # Note: dx = dy * w
+                _fp32_tp_reduce_bwd = (
+                    os.environ.get("NVTE_FP32_TP_REDUCE", "0") == "1"
+                    and ctx.parallel_mode == "column"
+                    and ctx.tp_size > 1
+                )
+                _dgrad_out_dtype = torch.float32 if _fp32_tp_reduce_bwd else ctx.activation_dtype
+
+                _tp_invariant_bwd = (
+                    os.environ.get("NVTE_TP_INVARIANT_MODE", "0") == "1"
+                    and ctx.parallel_mode == "column"
+                    and ctx.tp_size > 1
+                )
+
+                if _tp_invariant_bwd:
+                    # TP-invariant diagnostic: full dgrad GEMM matching TP=1 accumulation.
+                    assert not ctx.fp8, "NVTE_TP_INVARIANT_MODE does not support FP8"
+                    nvtx_range_push(f"{nvtx_label}.tp_invariant_dgrad")
+
+                    def allgather_along_dim(tensor, group, world_size, dim):
+                        chunks = [torch.empty_like(tensor) for _ in range(world_size)]
+                        torch.distributed.all_gather(
+                            chunks, tensor.contiguous(), group=group,
+                        )
+                        return torch.cat(chunks, dim=dim)
+
+                    grad_output_gathered = allgather_along_dim(
+                        grad_output, ctx.tp_group, ctx.tp_size, dim=-1,
+                    )  # [tokens, out/TP] -> [tokens, out]
+                    weight_gathered = allgather_along_dim(
+                        weight_fp8, ctx.tp_group, ctx.tp_size, dim=0,
+                    )  # [out/TP, in] -> [out, in]
+
+                    grad_output_2d = grad_output_gathered.reshape(
+                        -1, grad_output_gathered.shape[-1],
+                    )
+                    dgrad = general_gemm(
+                        weight_gathered, grad_output_2d, get_workspace(),
+                        layout="NN", grad=True,
+                        out_dtype=ctx.activation_dtype,
+                    )
+                    if isinstance(dgrad, tuple):
+                        dgrad = dgrad[0]
+
+                    # SP: scatter to per-rank chunk along sequence dim.
+                    if ctx.sequence_parallel:
+                        rank = torch.distributed.get_rank(ctx.tp_group)
+                        dgrad = dgrad.chunk(ctx.tp_size, dim=0)[rank].contiguous()
+
+                    del grad_output_gathered, weight_gathered, grad_output_2d
+                    nvtx_range_pop(f"{nvtx_label}.tp_invariant_dgrad")
+                else:
+                    nvtx_range_push(f"{nvtx_label}.dgrad_gemm")
+                    gemm_out, *_, reduce_scatter_out = general_gemm(
+                        weight_fp8,
+                        grad_output,
+                        get_workspace(),
+                        layout="NN",
+                        grad=True,
+                        quantization_params=ctx.grad_input_quantizer,
+                        out=gemm_out,
+                        out_dtype=_dgrad_out_dtype,
+                        use_split_accumulator=use_split_accumulator,
+                        ub=ub_obj_dgrad,
+                        ub_type=ub_type_dgrad,
+                        extra_output=reduce_scatter_out,
+                        bulk_overlap=ctx.ub_bulk_dgrad,
+                    )
+                    nvtx_range_pop(f"{nvtx_label}.dgrad_gemm")
+
+                    # Prepare grad input tensor
+                    # Note: Perform tensor-parallel communication
+                    if ctx.ub_overlap_rs_dgrad:
+                        dgrad = reduce_scatter_out
+                    elif ctx.ub_bulk_wgrad:
+                        dgrad = ub_obj_wgrad.get_buffer(local_chunk=True)
+                    elif ctx.parallel_mode == "column" and ctx.tp_size > 1:
+                        nvtx_range_push(f"{nvtx_label}.column_parallel_comm_dgrad")
+                        dgrad = gemm_out
+                        if ctx.sequence_parallel:
+                            dgrad, dgrad_work = reduce_scatter_along_first_dim(
+                                dgrad,
+                                ctx.tp_group,
+                                async_op=True,
+                            )
+                        else:
+                            dgrad, dgrad_work = allreduce(dgrad, ctx.tp_group, async_op=True)
+                        if _fp32_tp_reduce_bwd and dgrad.dtype != ctx.activation_dtype:
+                            dgrad = dgrad.to(ctx.activation_dtype)
+                        nvtx_range_pop(f"{nvtx_label}.column_parallel_comm_dgrad")
+                    else:
+                        dgrad = gemm_out
+
+            # --------------------------------------------------
+            # Grad input tensor has been computed...
+            # --------------------------------------------------
+
+            # --------------------------------------------------
+            # Compute grad weight
+            # --------------------------------------------------
+
+            wgrad = None
+            if ctx.requires_wgrad:
+
+                # Prepare input tensor
+                # Note: Synchronize tensor-parallel communication and
+                # make sure required data is available
+                if inputmat_total_work is not None:
+                    inputmat_total_work.wait()
+                    inputmat_total_work = None
+                if ctx.fp8 or ctx.debug:
+                    if isinstance(inputmat_total, QuantizedTensorStorage):
+                        inputmat_total.update_usage(columnwise_usage=True)
+                    else:
+                        ctx.input_quantizer.set_usage(rowwise=False, columnwise=True)
+                        inputmat_total = ctx.input_quantizer(inputmat_total)
+
+                # Prepare grad output tensor
+                # Note: Synchronize tensor-parallel communication and
+                # make sure required data is available
+                if ctx.ub_overlap_ag and isinstance(ctx.grad_output_quantizer, MXFP8Quantizer):
+                    # UB does not support pipelined overlapping grad output
+                    # all-gather with wgrad GEMM. Also, we can't
+                    # convert row-scaled MXFP8 to column-scaled, so we
+                    # can't reuse the grad output that was gathered
+                    # for the dgrad GEMM. We work around by explicitly
+                    # overlapping the AG operation with the dgrad GEMM.
+
+                    # Get the communication stream from the dgrad GEMM to use for the AG
+                    dgrad_send_stream, dgrad_recv_stream = ub_obj_dgrad.get_communication_stream()
+
+                    # This object is separate from the ub_obj_wgrad object which is passed to the GEMM
+                    ub_obj_overlap_wgrad = get_ub(ctx.ub_name + "_wgrad", ctx.fp8)
+
+                    ctx.grad_output_quantizer.set_usage(rowwise=False, columnwise=True)
+
+                    # We use the send stream to copy into the userbuffers.
+                    # This is the same stream that we will use to access the data in the AG,
+                    # so we dont need to add any syncs yet.
+                    with torch.cuda.stream(dgrad_send_stream):
+                        grad_output, _ = fill_userbuffers_buffer_for_all_gather(
+                            ub_obj_overlap_wgrad,
+                            grad_output_arg,
+                            ctx.grad_output_quantizer,
+                            ctx.tp_group,
+                        )
+
+                    # Allgather grad_outputs[0] using the dgrad streams so we can overlap with the fc2_dgrad gemm
+                    tex.bulk_overlap_ag_with_external_gemm(
+                        ub_obj_overlap_wgrad, dgrad_send_stream, dgrad_recv_stream
+                    )
+
+                if ctx.fp8 or ctx.debug:
+                    if isinstance(grad_output, QuantizedTensorStorage):
+                        grad_output.update_usage(columnwise_usage=True)
+                    else:
+                        ctx.grad_output_quantizer.set_usage(rowwise=False, columnwise=True)
+                        grad_output = ctx.grad_output_quantizer(grad_output)
+
+                # Figure out whether to use split accumulator
+                use_split_accumulator = _2X_ACC_WGRAD
+                if ctx.fp8:
+                    recipe = ctx.fp8_recipe
+                    if hasattr(recipe, "fp8_gemm_wgrad"):
+                        use_split_accumulator = recipe.fp8_gemm_wgrad.use_split_accumulator
+
+                # Figure out whether to output wgrad GEMM directly into main grad
+                if ctx.is_first_microbatch is not None:
+                    accumulate_wgrad_into_param_main_grad = (
+                        ctx.fuse_wgrad_accumulation and not ctx.is_first_microbatch
+                    )
+                else:
+                    accumulate_wgrad_into_param_main_grad = ctx.fuse_wgrad_accumulation
+
+                # Output buffer for overlapping FP8 grad input
+                # reduce-scatter with wgrad GEMM
+                reduce_scatter_out = None
+                if ctx.ub_bulk_wgrad and ub_obj_wgrad.is_fp8_ubuf():
+                    reduce_scatter_out = torch.empty(
+                        dgrad_shape, dtype=ctx.activation_dtype, device=grad_output_arg.device
+                    )
+
+                # Arguments to include in wgrad GEMM closure
+                wgrad_gemm_kwargs = {
+                    "workspace": get_workspace(),
+                    "out_dtype": (
+                        main_grad.dtype if ctx.fuse_wgrad_accumulation else ctx.activation_dtype
+                    ),
+                    "quantization_params": ctx.grad_weight_quantizer,
+                    "accumulate": (
+                        accumulate_wgrad_into_param_main_grad
+                        if not getattr(weight, "overwrite_main_grad", False)
+                        else False
+                    ),
+                    "layout": "NT",
+                    "out": main_grad if ctx.fuse_wgrad_accumulation else None,
+                    "bias": (bias if (grad_bias is None and not ctx.fp8) else None),
+                    "use_split_accumulator": use_split_accumulator,
+                    "grad": True,
+                    "ub": ub_obj_wgrad,
+                    "ub_type": ub_type_wgrad,
+                    "extra_output": reduce_scatter_out,
+                    "bulk_overlap": ctx.ub_bulk_wgrad,
+                }
+
+                def wgrad_gemm(
+                    x: torch.Tensor,
+                    dy: torch.Tensor,
+                ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+                    """Perform wgrad GEMM: dw = dy^T * x
+
+                    May be fused with bgrad computation.
+
+                    May be called outside of this function to enable
+                    some advanced communication/compute overlapping.
+
+                    """
+                    nvtx_range_push(f"{nvtx_label}.wgrad_gemm")
+                    dw, db, *_ = general_gemm(x, dy, **wgrad_gemm_kwargs)
+                    nvtx_range_pop(f"{nvtx_label}.wgrad_gemm")
+                    return dw, db
+
+                # Choose whether to call wgrad GEMM now or delay
+                if ctx.wgrad_store is not None and ctx.wgrad_store.delay_wgrad_compute():
+                    if (
+                        wgrad_gemm_kwargs["ub"] is not None
+                        or wgrad_gemm_kwargs["ub_type"] is not None
+                        or wgrad_gemm_kwargs["extra_output"] is not None
+                        or wgrad_gemm_kwargs["bulk_overlap"]
+                    ):
+                        raise NotImplementedError(
+                            "Delayed weight grad computation is not supported "
+                            "with Userbuffers (tensor-parallel communication overlapping)"
+                        )
+                    ctx.wgrad_store.put([inputmat_total, grad_output], wgrad_gemm)
+                else:
+
+                    # Call wgrad GEMM now
+                    wgrad, grad_bias_ = wgrad_gemm(inputmat_total, grad_output)
+
+                    # Update grad bias if needed
+                    if grad_bias is None:
+                        grad_bias = grad_bias_
+                    del grad_bias_
+
+                    # Deallocate tensors if permitted
+                    if ctx.owns_input:
+                        # Input tensor is internal
+                        clear_tensor_data(inputmat_total)
+                    elif ctx.backward_input_needs_gather:
+                        # Gathered input tensor is internal
+                        clear_tensor_data(inputmat_total)
+                    if ctx.parallel_mode == "row" and ctx.sequence_parallel:
+                        # Gathered grad output tensor is internal
+                        clear_tensor_data(grad_output)
+
+                # Update grad input if overlapping reduce-scatter with wgrad GEMM
+                if ctx.ub_bulk_wgrad:
+                    if ub_obj_wgrad.is_fp8_ubuf():
+                        dgrad = reduce_scatter_out
+                    else:
+                        dgrad = ub_obj_wgrad.get_buffer(local_chunk=True).clone()
+
+            # --------------------------------------------------
+            # Grad weight has been computed...
+            # --------------------------------------------------
+
+            # Don't return grad bias if not needed
+            if not ctx.use_bias:
+                grad_bias = None
+
+            # Make sure all tensor-parallel communication is finished
+            if inputmat_total_work is not None:
+                inputmat_total_work.wait()
+                inputmat_total_work = None
+            if dgrad_work is not None:
+                dgrad_work.wait()
+                dgrad_work = None
+
+        if ctx.requires_wgrad:
+            # Handle custom DDP from mcore.
+            if (
+                ctx.fuse_wgrad_accumulation
+                and weight is not None
+                and hasattr(weight, "grad_added_to_main_grad")
+            ):
+                weight.grad_added_to_main_grad = True
+                if getattr(weight, "zero_out_wgrad", False):
+                    wgrad = get_dummy_wgrad(
+                        list(weight.main_grad.shape),
+                        weight.dtype,
+                        zero=True,
+                    )
+                else:
+                    wgrad = get_dummy_wgrad(
+                        list(weight.main_grad.shape),
+                        weight.dtype,
+                    )
+            elif ctx.fuse_wgrad_accumulation:
+                wgrad = None
+        else:
+            wgrad = None
+
+        # Update FP8 scaling factors if needed
+        if ctx.reduce_and_update_bwd_fp8_tensors and not is_graph_capturing():
+            nvtx_range_push(f"{nvtx_label}.reduce_and_update_fp8_tensors")
+            FP8GlobalStateManager.reduce_and_update_fp8_tensors(forward=False)
+            nvtx_range_pop(f"{nvtx_label}.reduce_and_update_fp8_tensors")
+
+        # Scatter fp8 weight buffers
+        if ctx.fp8 and not isinstance(weight, QuantizedTensorStorage):
+            _fsdp_scatter_tensors(ctx.fsdp_group, weight_fp8)
+        return (
+            wgrad,
+            dgrad.view(ctx.inp_shape) if ctx.requires_dgrad else None,
+            grad_bias,
+            None,  # is_first_microbatch
+            None,  # fp8
+            None,  # fp8_calibration
+            None,  # wgrad_store
+            None,  # input_quantizer
+            None,  # weight_quantizer
+            None,  # output_quantizer
+            None,  # grad_input_quantizer
+            None,  # grad_weight_quantizer
+            None,  # grad_output_quantizer
+            None,  # fuse_wgrad_accumulation
+            None,  # cpu_offloading
+            None,  # tp_group
+            None,  # tp_size
+            None,  # sequence_parallel
+            None,  # tensor_parallel
+            None,  # activation_dtype
+            None,  # parallel_mode
+            None,  # is_grad_enabled
+            None,  # ub_overlap_rs_fprop
+            None,  # ub_overlap_ag_dgrad
+            None,  # ub_overlap_ag_fprop
+            None,  # ub_overlap_rs_dgrad
+            None,  # ub_bulk_dgrad
+            None,  # ub_bulk_wgrad
+            None,  # ub_name
+            None,  # fp8_output
+            None,  # fsdp_group
+            None,  # module
+            None,  # skip_fp8_weight_update
+            None,  # symmetric_ar_type
+            None,  # save_original_input
+            None,  # debug
+        )
+
+
+class Linear(TransformerEngineBaseModule):
+    """Applies a linear transformation to the incoming data :math:`y = xA^T + b`
+
+    On NVIDIA GPUs it is a drop-in replacement for `torch.nn.Linear`.
+
+    Parameters
+    ----------
+    in_features : int
+                 size of each input sample.
+    out_features : int
+                  size of each output sample.
+    bias : bool, default = `True`
+          if set to `False`, the layer will not learn an additive bias.
+    init_method : Callable, default = `None`
+                 used for initializing weights in the following way: `init_method(weight)`.
+                 When set to `None`, defaults to `torch.nn.init.normal_(mean=0.0, std=0.023)`.
+    get_rng_state_tracker : Callable, default = `None`
+                 used to get the random number generator state tracker for initializing weights.
+    rng_tracker_name : str, default = `None`
+                 the param passed to get_rng_state_tracker to get the specific rng tracker.
+    parameters_split : Optional[Union[Tuple[str, ...], Dict[str, int]]], default = None
+                      Configuration for splitting the weight and bias tensors along dim 0 into
+                      multiple PyTorch parameters. If a list or tuple of strings is provided,
+                      they are used to make the names of equally-sized parameters. If a dict
+                      (preferably an OrderedDict) is provided, the keys are used as names and
+                      values as split sizes along dim 0. The resulting parameters will have
+                      names that end in `_weight` or `_bias`, so trailing underscores are
+                      stripped from any provided names.
+    device : Union[torch.device, str], default = "cuda"
+          The device on which the parameters of the model will be allocated. It is the user's
+          responsibility to ensure all parameters are moved to the GPU before running the
+          forward pass.
+    name: str, default = `None`
+        name of the module, currently used for debugging purposes.
+
+    Parallelism parameters
+    ----------------------
+    sequence_parallel : bool, default = `False`
+                       if set to `True`, uses sequence parallelism.
+    tp_group : ProcessGroup, default = `None`
+              tensor parallel process group.
+    tp_size : int, default = 1
+             used as TP (tensor parallel) world size when TP groups are not formed during
+             initialization. In this case, users must call the
+             `set_tensor_parallel_group(tp_group)` method on the initialized module before the
+             forward pass to supply the tensor parallel group needed for tensor and sequence
+             parallel collectives.
+    parallel_mode : {None, 'column', 'row'}, default = `None`
+                   used to decide whether this Linear layer is Column Parallel Linear or Row
+                   Parallel Linear as described `here <https://arxiv.org/pdf/1909.08053.pdf>`_.
+                   When set to `None`, no communication is performed.
+
+    Optimization parameters
+    -----------------------
+    fuse_wgrad_accumulation : bool, default = 'False'
+                             if set to `True`, enables fusing of creation and accumulation of
+                             the weight gradient. When enabled, it is assumed that the weights
+                             have an additional `main_grad` attribute (used instead of the
+                             regular `grad`) which is a pre-allocated buffer of the correct
+                             size to accumulate gradients in. This argument along with
+                             weight tensor having attribute 'overwrite_main_grad' set to True
+                             will overwrite `main_grad` instead of accumulating.
+    return_bias : bool, default = `False`
+                 when set to `True`, this module will not apply the additive bias itself, but
+                 instead return the bias value during the forward pass together with the
+                 output of the linear transformation :math:`y = xA^T`. This is useful when
+                 the bias addition can be fused to subsequent operations.
+    params_dtype : torch.dtype, default = `torch.get_default_dtype()`
+                  it controls the type used to allocate the initial parameters. Useful when
+                  the model is trained with lower precision and the original FP32 parameters
+                  would not fit in GPU memory.
+    delay_wgrad_compute : bool, default = `False`
+                         Whether or not to delay weight gradient computation. If set to `True`,
+                         it's the user's responsibility to call `module.backward_dw` to compute
+                         weight gradients.
+    symmetric_ar_type : {None, 'multimem_all_reduce', 'two_shot', 'one_shot'}, default = None
+                   Type of symmetric memory all-reduce to use during the forward pass.
+                   This can help in latency bound communication situations.
+                   Requires PyTorch version 2.7.0 or higher. When set to None, standard all-reduce
+                   is used.
+    save_original_input : bool, default = `False`
+                       If set to `True`, always saves the original input tensor rather than the
+                       cast tensor. In some scenarios, the input tensor is used by multiple modules,
+                       and saving the original input tensor may reduce the memory usage.
+                       Cannot work with FP8 DelayedScaling recipe.
+    """
+
+    def __init__(
+        self,
+        in_features: int,
+        out_features: int,
+        sequence_parallel: bool = False,
+        fuse_wgrad_accumulation: bool = False,
+        tp_group: Optional[dist_group_type] = None,
+        tp_size: int = 1,
+        get_rng_state_tracker: Optional[Callable] = None,
+        rng_tracker_name: Optional[str] = None,
+        init_method: Optional[Callable] = None,
+        bias: bool = True,
+        return_bias: bool = False,
+        params_dtype: Optional[torch.dtype] = None,
+        parallel_mode: Optional[str] = None,
+        parameters_split: Optional[Union[Tuple[str, ...], Dict[str, int]]] = None,
+        device: Union[torch.device, str] = "cuda",
+        ub_overlap_ag: bool = False,
+        ub_overlap_rs: bool = False,
+        ub_overlap_rs_dgrad: bool = False,
+        ub_bulk_dgrad: bool = False,
+        ub_bulk_wgrad: bool = False,
+        ub_name: Optional[str] = None,
+        delay_wgrad_compute: bool = False,
+        symmetric_ar_type: Optional[str] = None,
+        save_original_input: bool = False,
+        name: Optional[str] = None,
+    ) -> None:
+        super().__init__()
+
+        params_dtype = torch.get_default_dtype() if params_dtype is None else params_dtype
+        self.in_features = in_features
+        self.out_features = out_features
+        self.fuse_wgrad_accumulation = fuse_wgrad_accumulation
+        self.use_bias = bias
+        self.return_bias = return_bias
+        self.apply_bias = bias and not return_bias
+        self.get_rng_state_tracker = get_rng_state_tracker
+        self.rng_tracker_name = rng_tracker_name
+        self.symmetric_ar_type = symmetric_ar_type
+        self.save_original_input = save_original_input
+        self.name = name
+
+        self.wgrad_store = WeightGradStore(delay_wgrad_compute, ub_bulk_wgrad)
+
+        if device == "meta":
+            assert parameters_split is None, "Cannot split module parameters on 'meta' device."
+        if tp_group is None:
+            self.tp_size = tp_size
+            if tp_size == 1:
+                self.set_tensor_parallel_group(tp_group)
+        else:
+            self.tp_size = get_distributed_world_size(tp_group)
+            self.set_tensor_parallel_group(tp_group)
+        self.set_nccl_overlap_warning_if_tp()
+
+        self.parallel_mode = parallel_mode
+        assert (
+            self.parallel_mode in GemmParallelModes
+        ), f"parallel_mode {parallel_mode} not supported"
+
+        if self.parallel_mode == "column":
+            self.out_features = divide(self.out_features, self.tp_size)
+        elif self.parallel_mode == "row":
+            self.in_features = divide(self.in_features, self.tp_size)
+
+        self.sequence_parallel = (self.tp_size > 1) and sequence_parallel
+
+        # Column parallel TP overlap options
+        self.ub_overlap_ag_fprop = (
+            self.parallel_mode == "column" and self.sequence_parallel and ub_overlap_ag
+        )
+        self.ub_overlap_rs_dgrad = (
+            self.parallel_mode == "column" and self.sequence_parallel and ub_overlap_rs_dgrad
+        )
+        self.ub_bulk_dgrad = (
+            self.parallel_mode == "column"
+            and self.sequence_parallel
+            and ub_bulk_dgrad
+            and not self.ub_overlap_rs_dgrad
+        )
+        self.ub_bulk_wgrad = (
+            self.parallel_mode == "column"
+            and self.sequence_parallel
+            and ub_bulk_wgrad
+            and not self.ub_overlap_rs_dgrad
+        )
+
+        # Row parallel TP overlap options
+        self.ub_overlap_rs_fprop = (
+            self.parallel_mode == "row" and self.sequence_parallel and ub_overlap_rs
+        )
+        self.ub_overlap_ag_dgrad = (
+            self.parallel_mode == "row" and self.sequence_parallel and ub_overlap_ag
+        )
+
+        if any(
+            [
+                self.ub_overlap_rs_fprop,
+                self.ub_overlap_ag_dgrad,
+                self.ub_overlap_ag_fprop,
+                self.ub_overlap_rs_dgrad,
+                self.ub_bulk_dgrad,
+                self.ub_bulk_wgrad,
+            ]
+        ):
+            assert ub_name is not None, f"Comm+GEMM overlap layer '{ub_name}' is not initialized."
+        self.ub_name = ub_name
+
+        if self.symmetric_ar_type is not None:
+            assert torch_version() >= (
+                2,
+                7,
+                0,
+            ), "Torch version must be at least 2.7 to use symmetric memory"
+
+        # Initialize params in FP8
+        with_fp8_params = FP8GlobalStateManager.with_fp8_parameters()
+
+        # Contiguous buffers for params
+        weight_tensor = torch.empty(
+            self.out_features,
+            self.in_features,
+            device=device,
+            dtype=params_dtype,
+        )
+        bias_tensor = None
+        if self.use_bias:
+            bias_tensor = torch.empty(
+                self.out_features,
+                device=device,
+                dtype=params_dtype,
+            )
+
+        # Configure parameter splits
+        self.weight_names = []
+        self.bias_names = []
+        self.parameter_split_sizes = []
+        if parameters_split is None:
+            # Split into a single parameter by default
+            self.weight_names = ["weight"]
+            self.bias_names = ["bias"]
+            self.parameter_split_sizes = [out_features]
+        elif not parameters_split:
+            raise ValueError("Cannot split weight buffer into 0 parameters")
+        elif isinstance(parameters_split, dict):
+            # Split parameters with provided sizes
+            for name, split_size in parameters_split.items():
+                self.weight_names.append(f"{name.rstrip('_')}_weight")
+                self.bias_names.append(f"{name.rstrip('_')}_bias")
+                self.parameter_split_sizes.append(split_size)
+        elif all(isinstance(name, str) for name in parameters_split):
+            # Split parameters evenly
+            split_size = out_features // len(parameters_split)
+            for name in parameters_split:
+                self.weight_names.append(f"{name.rstrip('_')}_weight")
+                self.bias_names.append(f"{name.rstrip('_')}_bias")
+                self.parameter_split_sizes.append(split_size)
+        else:
+            raise TypeError("Invalid configuration for parameters split")
+
+        # Make sure parameter splits are valid
+        if sum(self.parameter_split_sizes) != out_features:
+            raise ValueError(
+                f"Trying to split weight buffer ({out_features=}) "
+                f"with split sizes {self.parameter_split_sizes}"
+            )
+
+        # Adjust parameter splits for tensor-parallel distribution
+        if self.parallel_mode == "column":
+            for i, size in enumerate(self.parameter_split_sizes):
+                if size % self.tp_size != 0:
+                    raise RuntimeError(
+                        f"Attempting to distribute a parameter with out_features={size} "
+                        f"between {self.tp_size} tensor-parallel processes"
+                    )
+                self.parameter_split_sizes[i] = size // self.tp_size
+
+        # Construct weight parameters
+        # Note: Register weights together so that they are adjacent to
+        # each other in Linear.parameters(). This makes it more likely
+        # that they will stay contiguous if the weights are
+        # manipulated externally, e.g. by FSDP.
+        offset = 0
+        for i, split_size in enumerate(self.parameter_split_sizes):
+            split_start = offset
+            offset += split_size
+            split_end = offset
+
+            # Check if parameters are subviews of buffers
+            is_subview = (split_start, split_end) != (0, self.out_features)
+            if is_subview and with_fp8_params:
+                raise RuntimeError(
+                    "Splitting QuantizedTensor into multiple params is not supported"
+                )
+
+            # Construct weight parameter
+            self.register_parameter(
+                self.weight_names[i],
+                torch.nn.Parameter(weight_tensor[split_start:split_end]),
+                init_fn=init_method,
+                get_rng_state_tracker=get_rng_state_tracker,
+                fp8_meta_index=tex.FP8FwdTensors.GEMM1_WEIGHT,
+            )
+
+        # Construct bias parameters if needed
+        if self.use_bias:
+            offset = 0
+            for i, split_size in enumerate(self.parameter_split_sizes):
+                split_start = offset
+                offset += split_size
+                split_end = offset
+                self.register_parameter(
+                    self.bias_names[i],
+                    torch.nn.Parameter(bias_tensor[split_start:split_end]),
+                    init_fn=init_method_constant(0.0),
+                )
+        else:
+            for name in self.bias_names:
+                bias = torch.Tensor().to(dtype=params_dtype, device=device)
+                setattr(self, name, bias)
+
+        if with_fp8_params:
+            self.init_fp8_metadata()
+
+        self.reset_parameters(defer_init=device == "meta")
+
+        # For RPL, bias has to be added after TP collectives
+        # So it cannot be fused with the GEMM
+        if self.parallel_mode == "row" and self.apply_bias:
+            self.gemm_bias_unfused_add = True
+        else:
+            self.gemm_bias_unfused_add = False
+
+        if self.wgrad_store.delay_wgrad_compute():
+            for name, param in self.named_parameters():
+                if name in self.weight_names or name in self.bias_names:
+                    param.skip_backward_post_hook = True
+
+    def set_meta_tensor(self, fwd: bool, recipe: Recipe) -> None:
+        """Init scales and amaxes for fwd | bwd."""
+        super().set_meta_tensor(fwd, recipe)
+
+        # customize quantizers based on each recipe & layer configs
+        recipe = FP8GlobalStateManager.get_fp8_recipe()
+        if recipe.float8_current_scaling():
+            self._customize_quantizers_float8_current_scaling(fwd, recipe)
+        elif recipe.float8_block_scaling():
+            self._customize_quantizers_float8_blockwise_scaling(fwd, recipe)
+        elif recipe.nvfp4():
+            self._customize_quantizers_nvfp4(fwd, recipe)
+        # elif for other recipes (mxfp8, etc.)
+
+    def reset_parameters(self, defer_init=False):
+        super().reset_parameters(defer_init=defer_init)
+
+        if not defer_init:
+            # Set parallelism attributes for linear weights
+            for weight in self.weight_names:
+                set_tensor_model_parallel_attributes(
+                    tensor=getattr(self, weight),
+                    is_parallel=True,
+                    dim=1 if self.parallel_mode == "row" else 0,
+                    stride=1,
+                )
+
+            # Set parallelism attributes for linear biases
+            if self.use_bias:
+                for bias in self.bias_names:
+                    if self.parallel_mode == "row":
+                        setattr(getattr(self, bias), "sequence_parallel", self.sequence_parallel)
+                    elif self.parallel_mode == "column":
+                        set_tensor_model_parallel_attributes(getattr(self, bias), True, 0, 1)
+
+    @no_torch_dynamo()
+    def forward(
+        self,
+        inp: torch.Tensor,
+        is_first_microbatch: Optional[bool] = None,
+        fp8_output: Optional[bool] = False,
+        fp8_grad: Optional[bool] = False,
+    ) -> Union[torch.Tensor, Tuple[torch.Tensor, ...]]:
+        """
+        Apply the linear transformation to the input.
+
+        Parameters
+        ----------
+        inp : torch.Tensor
+             Input tensor.
+        is_first_microbatch : {True, False, None}, default = None
+                             During training using either gradient accumulation or
+                             pipeline parallelism a minibatch of data is further split
+                             into microbatches. Between the microbatches of the same minibatch
+                             the model weights are not updated. Setting this parameter indicates
+                             whether the current microbatch is the first in a minibatch or not.
+                             When set, this parameter enables additional optimizations:
+
+                             * during FP8 training, it allows caching of the FP8 versions of
+                               the weights
+                             * it also allows skipping gradient accumulation during the
+                               first microbatch (since it is the first gradient being
+                               produced)
+        """
+        if is_in_onnx_export_mode():
+            return self.onnx_forward(inp, fp8_output)
+
+        debug = self.is_debug_iter()
+
+        if FP8GlobalStateManager.fp8_graph_capturing():
+            skip_fp8_weight_update = FP8GlobalStateManager.get_skip_fp8_weight_update_tensor()
+        else:
+            skip_fp8_weight_update = None
+        if skip_fp8_weight_update is not None:
+            is_first_microbatch = False
+
+        if self.ub_overlap_rs_fprop:
+            if get_ub(
+                self.ub_name + "_fprop", FP8GlobalStateManager.is_fp8_enabled()
+            ).is_fp8_ubuf():
+                fp8_output = True
+        if self.ub_overlap_rs_dgrad:
+            if get_ub(
+                self.ub_name + "_dgrad", FP8GlobalStateManager.is_fp8_enabled()
+            ).is_fp8_ubuf():
+                fp8_grad = True
+
+        with torch.cuda.device(
+            getattr(self, list(self.named_parameters())[0][0]).device
+        ), self.prepare_forward(
+            inp,
+            allow_non_contiguous=isinstance(inp, QuantizedTensor),
+        ) as inp:
+
+            weight_tensor, bias_tensor = self._get_weight_and_bias_tensors()
+
+            quantizers = (
+                self._get_quantizers(fp8_output, fp8_grad)
+                if not debug
+                else self._get_debug_quantizers(fp8_output, fp8_grad)
+            )
+            if debug:
+                if self.no_debug_features_active(quantizers):
+                    debug = False
+                    quantizers = self._get_quantizers(fp8_output, fp8_grad)
+
+            (
+                input_quantizer,
+                weight_quantizer,
+                output_quantizer,
+                grad_input_quantizer,
+                grad_weight_quantizer,
+                grad_output_quantizer,
+            ) = quantizers
+
+            if torch.is_grad_enabled():
+                linear_fn = _Linear.apply
+                args = []
+            else:
+                linear_fn = _Linear.forward
+                args = [None]
+            args += (
+                weight_tensor,
+                inp,
+                bias_tensor if (self.apply_bias and not self.gemm_bias_unfused_add) else None,
+                is_first_microbatch,
+                self.fp8,
+                self.fp8_calibration,
+                self.wgrad_store,
+                input_quantizer,
+                weight_quantizer,
+                output_quantizer,
+                grad_input_quantizer,
+                grad_weight_quantizer,
+                grad_output_quantizer,
+                self.fuse_wgrad_accumulation,
+                is_cpu_offload_enabled(),
+                self.tp_group,
+                self.tp_size,
+                self.sequence_parallel,
+                self.tp_size > 1,
+                self.activation_dtype,
+                self.parallel_mode,
+                torch.is_grad_enabled(),
+                self.ub_overlap_rs_fprop,
+                self.ub_overlap_ag_dgrad,
+                self.ub_overlap_ag_fprop,
+                self.ub_overlap_rs_dgrad,
+                self.ub_bulk_dgrad,
+                self.ub_bulk_wgrad,
+                self.ub_name,
+                fp8_output,
+                self.fsdp_group,
+                self,
+                skip_fp8_weight_update,
+                self.symmetric_ar_type,
+                self.save_original_input,
+                debug,
+            )
+            out = linear_fn(*args)
+        if self.gemm_bias_unfused_add:
+            out = out + cast_if_needed(bias_tensor, self.activation_dtype)
+
+        if self.return_bias:
+            return out, cast_if_needed(bias_tensor, self.activation_dtype)
+        return out
+
+    def _get_quantizers(self, fp8_output, fp8_grad):
+        if not self.fp8:
+            return [None] * 6
+        grad_input_quantizer = None
+        grad_weight_quantizer = None
+        grad_output_quantizer = None
+        output_quantizer = None
+        input_quantizer = self.quantizers["scaling_fwd"][tex.FP8FwdTensors.GEMM1_INPUT]
+        input_quantizer.internal = True
+        (weight_quantizer,) = self._get_weight_quantizers()
+        if fp8_output:
+            output_quantizer = self.quantizers["scaling_fwd"][tex.FP8FwdTensors.GEMM1_OUTPUT]
+        if torch.is_grad_enabled():
+            grad_output_quantizer = self.quantizers["scaling_bwd"][tex.FP8BwdTensors.GRAD_OUTPUT1]
+            grad_output_quantizer.internal = True
+            if fp8_grad:
+                grad_input_quantizer = self.quantizers["scaling_bwd"][tex.FP8BwdTensors.GRAD_INPUT1]
+        return (
+            input_quantizer,
+            weight_quantizer,
+            output_quantizer,
+            grad_input_quantizer,
+            grad_weight_quantizer,
+            grad_output_quantizer,
+        )
+
+    def _get_debug_quantizers(self, fp8_output, fp8_grad):
+        original_quantizers = self._get_quantizers(fp8_output, fp8_grad)
+        assert TEDebugState.debug_enabled
+        from ...debug.pytorch.debug_quantization import DebugQuantizer
+
+        names = ["activation", "weight", "output", "dgrad", "wgrad", "gradient"]
+        return tuple(
+            DebugQuantizer(self.name, name, q, self.tp_group)
+            for name, q in zip(names, original_quantizers)
+        )
+
+    def _get_weight_tensors(self) -> List[Union[torch.Tensor, QuantizedTensorStorage]]:
+        """Get the weight tensors of the module."""
+        unfused_weights = [getattr(self, name) for name in self.weight_names]
+        if any(isinstance(w, QuantizedTensor) for w in unfused_weights):
+            if self.fp8:
+                if len(unfused_weights) != 1:
+                    raise RuntimeError(
+                        "Splitting QuantizedTensor into multiple params is not supported"
+                    )
+            else:
+                warnings.warn(
+                    "You are using quantized weights without quantized compute. "
+                    "Please make sure this is intentional."
+                )
+                unfused_weights = [w.dequantize() for w in unfused_weights]
+        return unfused_weights
+
+    def _get_weight_and_bias_tensors(self) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
+        # Get concatenated weight and bias tensors
+        unfused_weights = self._get_weight_tensors()
+        if any(isinstance(w, QuantizedTensor) for w in unfused_weights):
+            if self.fp8:
+                if len(unfused_weights) != 1:
+                    raise RuntimeError(
+                        "Splitting QuantizedTensor into multiple params is not supported"
+                    )
+            else:
+                warnings.warn(
+                    "You are using quantized weights without quantized compute. "
+                    "Please make sure this is intentional."
+                )
+                unfused_weights = [w.dequantize() for w in unfused_weights]
+
+        weight_tensor = noop_cat(unfused_weights)
+        if self.use_bias:
+            bias_tensor = noop_cat([getattr(self, name) for name in self.bias_names])
+        else:
+            bias_tensor = None
+
+        return weight_tensor, bias_tensor
+
+    def onnx_forward(
+        self,
+        inp: torch.Tensor,
+        fp8_output: bool,
+    ) -> torch.Tensor:
+        """
+        ONNX-compatible version of the forward function that provides numerical equivalence
+        while only using operations that have defined ONNX symbolic translations.
+        This simplified implementation is designed specifically for inference scenarios.
+        """
+        from ..export import onnx_gemm
+
+        assert_warmed_up(self)
+        assert not TEDebugState.debug_enabled, "Debug mode is not supported in ONNX export."
+        weight_tensor, bias_tensor = self._get_weight_and_bias_tensors()
+        (
+            input_quantizer,
+            weight_quantizer,
+            output_quantizer,
+            *_,
+        ) = self._get_quantizers(fp8_output, False)
+        inp_dtype = inp.dtype
+
+        if input_quantizer is not None:
+            inp_q = input_quantizer.onnx_quantize(inp)
+            inp = input_quantizer.onnx_dequantize(inp_q)
+            inp = inp.to(inp_dtype)
+
+        if weight_quantizer is not None:
+            weight_q = weight_quantizer.onnx_quantize(weight_tensor)
+            weight_tensor = weight_quantizer.onnx_dequantize(weight_q)
+        if bias_tensor is not None:
+            bias_tensor = bias_tensor.to(inp_dtype)
+        weight_tensor = weight_tensor.to(inp_dtype)
+
+        if self.apply_bias:
+            output = onnx_gemm(weight_tensor, inp, bias_tensor)
+        else:
+            output = onnx_gemm(weight_tensor, inp, None)
+
+        if output_quantizer is not None:
+            raise NotImplementedError("ONNX export of quantized output is not supported")
+
+        if self.return_bias:
+            return output, bias_tensor
+
+        return output
+
+    def _customize_quantizers_float8_current_scaling(self, fwd: bool, recipe: Recipe) -> None:
+        """Customize quantizers based on current scaling recipe + linear."""
+        assert (
+            recipe.float8_current_scaling()
+        ), "current scaling recipe quantizer customization here"
+        if fwd:
+            # set configs about amax epsilon and power_2_scale
+            self.quantizers["scaling_fwd"][
+                tex.FP8FwdTensors.GEMM1_INPUT
+            ].force_pow_2_scales = recipe.fp8_quant_fwd_inp.power_2_scale
+            self.quantizers["scaling_fwd"][
+                tex.FP8FwdTensors.GEMM1_INPUT
+            ].amax_epsilon = recipe.fp8_quant_fwd_inp.amax_epsilon
+            # also set weight quantizer with same amax_epsilon & power_2_scale
+            self.quantizers["scaling_fwd"][
+                tex.FP8FwdTensors.GEMM1_WEIGHT
+            ].force_pow_2_scales = recipe.fp8_quant_fwd_weight.power_2_scale
+            self.quantizers["scaling_fwd"][
+                tex.FP8FwdTensors.GEMM1_WEIGHT
+            ].amax_epsilon = recipe.fp8_quant_fwd_weight.amax_epsilon
+            # paralle related
+            if self.sequence_parallel and self.parallel_mode == "column":
+                # customize input_quantizer with amax reduction TP group
+                self.quantizers["scaling_fwd"][
+                    tex.FP8FwdTensors.GEMM1_INPUT
+                ].with_amax_reduction = True
+                self.quantizers["scaling_fwd"][
+                    tex.FP8FwdTensors.GEMM1_INPUT
+                ].amax_reduction_group = self.tp_group
+        else:
+            # set grad_output_quantizer with amax epsilon and power_2_scale
+            self.quantizers["scaling_bwd"][
+                tex.FP8BwdTensors.GRAD_OUTPUT1
+            ].force_pow_2_scales = recipe.fp8_quant_bwd_grad.power_2_scale
+            self.quantizers["scaling_bwd"][
+                tex.FP8BwdTensors.GRAD_OUTPUT1
+            ].amax_epsilon = recipe.fp8_quant_bwd_grad.amax_epsilon
+            # parallel related
+            if self.sequence_parallel and self.parallel_mode == "row":
+                # customize grad_output_quantizer with amax reduction TP group
+                self.quantizers["scaling_bwd"][
+                    tex.FP8BwdTensors.GRAD_OUTPUT1
+                ].with_amax_reduction = True
+                self.quantizers["scaling_bwd"][
+                    tex.FP8BwdTensors.GRAD_OUTPUT1
+                ].amax_reduction_group = self.tp_group
+
+    def _customize_quantizers_nvfp4(self, fwd: bool, recipe: Recipe) -> None:
+        """Customize quantizers based on current scaling recipe + linear."""
+        assert recipe.nvfp4(), "Incorrect recipe."
+        if fwd:
+            if self.sequence_parallel and self.parallel_mode == "column":
+                # customize input_quantizer with amax reduction TP group
+                self.quantizers["scaling_fwd"][
+                    tex.FP8FwdTensors.GEMM1_INPUT
+                ].with_amax_reduction = True
+                self.quantizers["scaling_fwd"][
+                    tex.FP8FwdTensors.GEMM1_INPUT
+                ].amax_reduction_group = self.tp_group
+        else:
+            if self.sequence_parallel and self.parallel_mode == "row":
+                # customize grad_output_quantizer with amax reduction TP group
+                self.quantizers["scaling_bwd"][
+                    tex.FP8BwdTensors.GRAD_OUTPUT1
+                ].with_amax_reduction = True
+                self.quantizers["scaling_bwd"][
+                    tex.FP8BwdTensors.GRAD_OUTPUT1
+                ].amax_reduction_group = self.tp_group
+
+    def _get_weight_quantizers(self) -> List[Quantizer]:
+        """Get the weight quantizers of the module."""
+        if not self.fp8 and not self.fp8_calibration:
+            return [None]
+        weight_quantizer = self.quantizers["scaling_fwd"][tex.FP8FwdTensors.GEMM1_WEIGHT]
+        weight_quantizer.internal = True
+        return [weight_quantizer]
+
+    def _customize_quantizers_float8_blockwise_scaling(self, fwd: bool, recipe: Recipe) -> None:
+        """Customize quantizers based on blockwise scaling recipe + linear."""
+        assert (
+            recipe.float8_block_scaling()
+        ), "blockwise scaling recipe quantizer customization here"
+
+        if fwd:
+            if self.sequence_parallel and self.parallel_mode == "column":
+                # set compact for inp tensor X
+                self.quantizers["scaling_fwd"][
+                    tex.FP8FwdTensors.GEMM1_INPUT
+                ].all_gather_usage = True
+        else:
+            if self.sequence_parallel and self.parallel_mode == "row":
+                # set compact for grad_output tensor dY
+                self.quantizers["scaling_bwd"][
+                    tex.FP8BwdTensors.GRAD_OUTPUT1
+                ].all_gather_usage = True
diff --git a/examples/tp-numerics/plot_loss.py b/examples/tp-numerics/plot_loss.py
new file mode 100644
index 00000000000..dbd077a889f
--- /dev/null
+++ b/examples/tp-numerics/plot_loss.py
@@ -0,0 +1,68 @@
+"""Plot lm-loss curves from two Megatron training logs (overlay).
+
+Usage:
+    python plot_loss.py <log_a> <label_a> <log_b> <label_b> <out.png>
+
+Example (cross-arch B300 vs H100):
+    python plot_loss.py \\
+        b300_run.log "B300 TP=1 (tp-invariant)" \\
+        h100_run.log "H100 TP=1 (tp-invariant)" \\
+        b300_vs_h100_tp1_invariant.png
+"""
+import re
+import sys
+from pathlib import Path
+
+import matplotlib.pyplot as plt
+
+ITER_LOSS_RE = re.compile(r"iteration\s+(\d+)/\s+\d+.*?lm loss:\s+([\d.E+-]+)")
+
+
+def parse_log(path):
+    iters, losses = [], []
+    for line in Path(path).read_text().splitlines():
+        m = ITER_LOSS_RE.search(line)
+        if m:
+            iters.append(int(m.group(1)))
+            losses.append(float(m.group(2)))
+    return iters, losses
+
+
+def main():
+    if len(sys.argv) != 6:
+        print(__doc__, file=sys.stderr)
+        sys.exit(1)
+    log_a, label_a, log_b, label_b, out_png = sys.argv[1:]
+
+    a_it, a_loss = parse_log(log_a)
+    b_it, b_loss = parse_log(log_b)
+    assert a_it == b_it, f"iter index mismatch: {len(a_it)} vs {len(b_it)}"
+    print(f"{len(a_it)} iters; bitwise identical: {a_loss == b_loss}")
+
+    fig, ax = plt.subplots(figsize=(8, 6.5))
+    # B drawn first thicker (background), A on top thinner — both visible when overlapping
+    ax.plot(b_it, b_loss, color="#43a047", linewidth=2.8, label=label_b, alpha=0.9)
+    ax.plot(a_it, a_loss, color="#26c6da", linewidth=1.4, label=label_a)
+
+    ax.set_title("lm loss", fontsize=14, pad=14)
+    ax.set_xlim(0, max(a_it))
+    ax.tick_params(axis="both", labelsize=10)
+    ax.grid(True, color="#e0e0e0", linewidth=0.8)
+    ax.set_axisbelow(True)
+    for spine in ax.spines.values():
+        spine.set_edgecolor("#cfcfcf")
+    ax.text(max(a_it) - 0.5, ax.get_ylim()[0] + 0.05, "Step", color="#888", fontsize=10, ha="right")
+
+    leg = ax.legend(loc="upper center", bbox_to_anchor=(0.5, 1.0), ncol=2, frameon=False, fontsize=11,
+                    handlelength=1.5, columnspacing=2.2)
+    for text, color in zip(leg.get_texts(), ["#43a047", "#26c6da"]):
+        text.set_color(color)
+
+    Path(out_png).parent.mkdir(parents=True, exist_ok=True)
+    plt.tight_layout()
+    plt.savefig(out_png, dpi=150, bbox_inches="tight", facecolor="white")
+    print(f"saved {out_png}")
+
+
+if __name__ == "__main__":
+    main()
diff --git a/examples/tp-numerics/submit_qwen3_0.6b_tp_invariant.sh b/examples/tp-numerics/submit_qwen3_0.6b_tp_invariant.sh
new file mode 100644
index 00000000000..df26a30ddc4
--- /dev/null
+++ b/examples/tp-numerics/submit_qwen3_0.6b_tp_invariant.sh
@@ -0,0 +1,185 @@
+#!/bin/bash
+#
+# Qwen3-0.6B TP-invariant numerics validation. Raw Megatron-LM.
+#
+# Reproduces bitwise-identical training across hardware (B300 ≡ H100) with the
+# TP-invariant patch stack. Default config (TP=1, 100 iters): iter 1 lm loss
+# 1.213320E+01 / grad norm 18.259; iter 100 lm loss 8.277067E+00.
+#
+# Runs both ways:
+#   - Interactive (active salloc):   bash submit_qwen3_0.6b_tp_invariant.sh
+#   - SLURM batch:                   sbatch submit_qwen3_0.6b_tp_invariant.sh
+#
+# Tunable env (override at invocation):
+#   TP_SIZE       (default 1)        — tensor-parallel degree
+#   TRAIN_ITERS   (default 100)
+#   OUTPUT_DIR    (default ./output/qwen3_0.6b_tp${TP_SIZE})
+#   WORKSPACE     (default detected)
+
+#SBATCH --nodes=1
+#SBATCH --ntasks-per-node=1
+#SBATCH --gpus-per-node=8
+#SBATCH --time=00:30:00
+#SBATCH --job-name=qwen3_0.6b_tp_invariant
+#SBATCH --output=slurm_%j.log
+#SBATCH --exclusive
+
+set -euo pipefail
+
+# ---------------------------------------------------------------------------
+# Paths
+# ---------------------------------------------------------------------------
+WORKSPACE="${WORKSPACE:-/home/scratch.jinzex_sw/jinzex-ai-workspace/engineering/reflection-engineering}"
+MEGATRON_PATH="${MEGATRON_PATH:-${WORKSPACE}/third-party/Megatron-LM}"
+CONTAINER_IMAGE="${CONTAINER_IMAGE:-${WORKSPACE}/containers/nemo-25.11.01.sqsh}"
+CONTAINER_MOUNTS="${CONTAINER_MOUNTS:-/home/scratch.jinzex_sw:/home/scratch.jinzex_sw}"
+
+TP_SIZE=${TP_SIZE:-1}
+TRAIN_ITERS=${TRAIN_ITERS:-100}
+OUTPUT_DIR="${OUTPUT_DIR:-${WORKSPACE}/projects/Numerics/tp-numerics/results-b300/qwen3_0.6b_tp${TP_SIZE}_iter${TRAIN_ITERS}}"
+LOG_FILE="${OUTPUT_DIR}/training.log"
+
+mkdir -p "${OUTPUT_DIR}"
+
+# ---------------------------------------------------------------------------
+# Environment (TP-invariant numerics + determinism)
+# ---------------------------------------------------------------------------
+export NVTE_TP_INVARIANT_MODE=1
+export NVTE_ALLOW_NONDETERMINISTIC_ALGO=0
+export CUBLAS_WORKSPACE_CONFIG=:4096:8
+export NCCL_ALGO=^NVLS
+export NCCL_NVLS_ENABLE=0
+export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True
+# Force unfused attention (Thread A fix — FA2 leaks bits on SM_100).
+export NVTE_FLASH_ATTN=0
+export NVTE_FUSED_ATTN=0
+export CUDA_DEVICE_MAX_CONNECTIONS=1
+
+# Sequence parallel only when TP>1 (no-op at TP=1).
+SP_FLAG=""
+if [[ "${TP_SIZE}" -gt 1 ]]; then
+    SP_FLAG="--sequence-parallel"
+fi
+
+# TE TP-invariant patches: applied inside the same srun as training because
+# squashfs containers are read-only per-srun. Detects TE version and picks
+# the matching subdir under patches/; fails fast if unsupported.
+SETUP_CMD="\
+TE_DST=\$(python3 -c 'import transformer_engine.pytorch as p; print(p.__path__[0])') && \
+if ! grep -q NVTE_TP_INVARIANT_MODE \$TE_DST/module/linear.py; then \
+    P=${MEGATRON_PATH}/examples/tp-numerics/patches && \
+    TE_VERSION=\$(python3 -c 'import transformer_engine as te; print(te.__version__.split(\"+\")[0])') && \
+    PATCH_SRC=\$P/v\${TE_VERSION}/transformer_engine/pytorch && \
+    if [ ! -d \$PATCH_SRC ]; then echo \"ERROR: no TP-invariant patches for TE \$TE_VERSION (have: \$(ls \$P | grep ^v | tr '\n' ' '))\"; exit 1; fi && \
+    cp -r \$PATCH_SRC/* \$TE_DST/; \
+fi"
+
+# ---------------------------------------------------------------------------
+# Training command
+# Args mirror Bridge's qwen3_600m_pretrain_config + the validate_e2e
+# overrides (cpu_init=True, deterministic_mode=True, attn=unfused, BIK on,
+# cross-entropy fusion off, mock-data, seed 1234).
+# ---------------------------------------------------------------------------
+TRAINING_CMD="python ${MEGATRON_PATH}/pretrain_gpt.py \
+    --distributed-timeout-minutes 60 \
+    --tensor-model-parallel-size ${TP_SIZE} \
+    --pipeline-model-parallel-size 1 \
+    --use-distributed-optimizer \
+    ${SP_FLAG} \
+    --transformer-impl transformer_engine \
+    --attention-backend unfused \
+    --micro-batch-size 1 \
+    --global-batch-size 1 \
+    --train-iters ${TRAIN_ITERS} \
+    --eval-iters 32 \
+    --eval-interval $((TRAIN_ITERS + 1)) \
+    --num-layers 28 \
+    --hidden-size 1024 \
+    --ffn-hidden-size 3072 \
+    --num-attention-heads 16 \
+    --group-query-attention \
+    --num-query-groups 8 \
+    --kv-channels 128 \
+    --normalization RMSNorm \
+    --norm-epsilon 1e-06 \
+    --swiglu \
+    --qk-layernorm \
+    --disable-bias-linear \
+    --hidden-dropout 0.0 \
+    --attention-dropout 0.0 \
+    --no-bias-dropout-fusion \
+    --no-bias-swiglu-fusion \
+    --no-persist-layer-norm \
+    --no-rope-fusion \
+    --position-embedding-type rope \
+    --rotary-base 1000000 \
+    --rotary-percent 1.0 \
+    --max-position-embeddings 40960 \
+    --seq-length 2048 \
+    --vocab-size 32000 \
+    --null-tokenizer-vocab-includes-eod \
+    --padded-vocab-size 151936 \
+    --make-vocab-size-divisible-by 128 \
+    --init-method-std 0.02 \
+    --lr 3e-4 \
+    --min-lr 3e-5 \
+    --lr-decay-style cosine \
+    --lr-warmup-iters 0 \
+    --lr-decay-iters ${TRAIN_ITERS} \
+    --weight-decay 0.033 \
+    --adam-beta1 0.9 \
+    --adam-beta2 0.95 \
+    --adam-eps 1e-5 \
+    --clip-grad 1.0 \
+    --optimizer adam \
+    --bf16 \
+    --use-cpu-initialization \
+    --deterministic-mode \
+    --batch-invariant-mode \
+    --mock-data \
+    --tokenizer-type NullTokenizer \
+    --seed 1234 \
+    --log-interval 1 \
+    --tensorboard-dir ${OUTPUT_DIR}/tensorboard"
+
+# ---------------------------------------------------------------------------
+# Launch
+# ---------------------------------------------------------------------------
+# srun handles both interactive (in active salloc, --overlap to share with the
+# salloc shell's default step) and sbatch contexts.
+NUM_GPUS=${TP_SIZE}
+srun ${SRUN_EXTRA_ARGS:-} --mpi=pmix --overlap --cpu-bind=none \
+     --ntasks=${NUM_GPUS} --ntasks-per-node=${NUM_GPUS} \
+     --gpus-per-node=${NUM_GPUS} \
+     --no-container-mount-home \
+     --container-image="${CONTAINER_IMAGE}" \
+     --container-mounts="${CONTAINER_MOUNTS}" \
+     --container-workdir="${MEGATRON_PATH}" \
+     bash -c "${SETUP_CMD} && ${TRAINING_CMD}" 2>&1 | tee "${LOG_FILE}"
+
+EXPECTED_ITER1="lm loss: 1.213320E+01 grad norm: 18.259"
+EXPECTED_LAST="lm loss: 8.277067E+00 grad norm: 0.542"
+
+extract_iter() {
+    grep "iteration *$1/" "${LOG_FILE}" | head -1 \
+        | grep -oP '(lm loss|grad norm): \S+' | tr '\n' ' ' | sed 's/ $//'
+}
+ITER1=$(extract_iter 1)
+ITER_LAST=$(extract_iter "${TRAIN_ITERS}")
+
+echo
+echo "Done. Loss curve in: ${LOG_FILE}"
+echo "Iter 1   (this run):  ${ITER1}"
+echo "Iter ${TRAIN_ITERS} (this run):  ${ITER_LAST}"
+
+if [[ "${TP_SIZE}" == "1" && "${TRAIN_ITERS}" == "100" ]]; then
+    echo "Expected iter 1:      ${EXPECTED_ITER1}"
+    echo "Expected iter 100:    ${EXPECTED_LAST}"
+    if [[ "${ITER1}" != "${EXPECTED_ITER1}" || "${ITER_LAST}" != "${EXPECTED_LAST}" ]]; then
+        echo "FAIL: numerics mismatch (regression?)"
+        exit 1
+    fi
+    echo "PASS"
+else
+    echo "(No assertion: non-default TP_SIZE/TRAIN_ITERS — expected values only valid for default config.)"
+fi
diff --git a/examples/tp-numerics/submit_qwen3_8b_tp_invariant.sh b/examples/tp-numerics/submit_qwen3_8b_tp_invariant.sh
new file mode 100644
index 00000000000..e7af73bc937
--- /dev/null
+++ b/examples/tp-numerics/submit_qwen3_8b_tp_invariant.sh
@@ -0,0 +1,184 @@
+#!/bin/bash
+#
+# Qwen3-8B TP-invariant numerics validation. Raw Megatron-LM.
+#
+# Reproduces bitwise-identical training across hardware (B300 ≡ H100) with the
+# TP-invariant patch stack. Default config (TP=4, 100 iters): iter 1 lm loss
+# 1.275062E+01 / grad norm 90.104; iter 100 lm loss 7.337554E+00.
+#
+# Runs both ways:
+#   - Interactive (active salloc):   bash submit_qwen3_8b_tp_invariant.sh
+#   - SLURM batch:                   sbatch submit_qwen3_8b_tp_invariant.sh
+#
+# Tunable env (override at invocation):
+#   TP_SIZE       (default 4)        — tensor-parallel degree (8B at TP≤2 OOMs)
+#   TRAIN_ITERS   (default 100)
+#   OUTPUT_DIR    (default ./output/qwen3_8b_tp${TP_SIZE})
+#   WORKSPACE     (default detected)
+
+#SBATCH --nodes=1
+#SBATCH --ntasks-per-node=4
+#SBATCH --gpus-per-node=8
+#SBATCH --time=00:60:00
+#SBATCH --job-name=qwen3_8b_tp_invariant
+#SBATCH --output=slurm_%j.log
+#SBATCH --exclusive
+
+set -euo pipefail
+
+# ---------------------------------------------------------------------------
+# Paths
+# ---------------------------------------------------------------------------
+WORKSPACE="${WORKSPACE:-/home/scratch.jinzex_sw/jinzex-ai-workspace/engineering/reflection-engineering}"
+MEGATRON_PATH="${MEGATRON_PATH:-${WORKSPACE}/third-party/Megatron-LM}"
+CONTAINER_IMAGE="${CONTAINER_IMAGE:-${WORKSPACE}/containers/nemo-25.11.01.sqsh}"
+CONTAINER_MOUNTS="${CONTAINER_MOUNTS:-/home/scratch.jinzex_sw:/home/scratch.jinzex_sw}"
+
+TP_SIZE=${TP_SIZE:-4}
+TRAIN_ITERS=${TRAIN_ITERS:-100}
+OUTPUT_DIR="${OUTPUT_DIR:-${WORKSPACE}/projects/Numerics/tp-numerics/results-b300/qwen3_8b_tp${TP_SIZE}_iter${TRAIN_ITERS}}"
+LOG_FILE="${OUTPUT_DIR}/training.log"
+
+mkdir -p "${OUTPUT_DIR}"
+
+# ---------------------------------------------------------------------------
+# Environment (TP-invariant numerics + determinism)
+# ---------------------------------------------------------------------------
+export NVTE_TP_INVARIANT_MODE=1
+export NVTE_ALLOW_NONDETERMINISTIC_ALGO=0
+export CUBLAS_WORKSPACE_CONFIG=:4096:8
+export NCCL_ALGO=^NVLS
+export NCCL_NVLS_ENABLE=0
+export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True
+# Force unfused attention (Thread A fix — FA2 leaks bits on SM_100).
+export NVTE_FLASH_ATTN=0
+export NVTE_FUSED_ATTN=0
+export CUDA_DEVICE_MAX_CONNECTIONS=1
+
+# Sequence parallel only when TP>1 (default TP=4 → SP on).
+SP_FLAG=""
+if [[ "${TP_SIZE}" -gt 1 ]]; then
+    SP_FLAG="--sequence-parallel"
+fi
+
+# TE TP-invariant patches: applied inside the same srun as training because
+# squashfs containers are read-only per-srun. Detects TE version and picks
+# the matching subdir under patches/; fails fast if unsupported.
+SETUP_CMD="\
+TE_DST=\$(python3 -c 'import transformer_engine.pytorch as p; print(p.__path__[0])') && \
+if ! grep -q NVTE_TP_INVARIANT_MODE \$TE_DST/module/linear.py; then \
+    P=${MEGATRON_PATH}/examples/tp-numerics/patches && \
+    TE_VERSION=\$(python3 -c 'import transformer_engine as te; print(te.__version__.split(\"+\")[0])') && \
+    PATCH_SRC=\$P/v\${TE_VERSION}/transformer_engine/pytorch && \
+    if [ ! -d \$PATCH_SRC ]; then echo \"ERROR: no TP-invariant patches for TE \$TE_VERSION (have: \$(ls \$P | grep ^v | tr '\n' ' '))\"; exit 1; fi && \
+    cp -r \$PATCH_SRC/* \$TE_DST/; \
+fi"
+
+# ---------------------------------------------------------------------------
+# Training command
+# Args mirror Bridge's qwen3_8b_pretrain_config + the validate_e2e
+# overrides. Differences vs 0.6B: 36 layers, 4096 hidden, 32 attn heads,
+# 12288 FFN, untied embeddings (Qwen3-8B doesn't tie).
+# ---------------------------------------------------------------------------
+TRAINING_CMD="python ${MEGATRON_PATH}/pretrain_gpt.py \
+    --distributed-timeout-minutes 60 \
+    --tensor-model-parallel-size ${TP_SIZE} \
+    --pipeline-model-parallel-size 1 \
+    --use-distributed-optimizer \
+    ${SP_FLAG} \
+    --transformer-impl transformer_engine \
+    --attention-backend unfused \
+    --micro-batch-size 1 \
+    --global-batch-size 1 \
+    --train-iters ${TRAIN_ITERS} \
+    --eval-iters 32 \
+    --eval-interval $((TRAIN_ITERS + 1)) \
+    --num-layers 36 \
+    --hidden-size 4096 \
+    --ffn-hidden-size 12288 \
+    --num-attention-heads 32 \
+    --group-query-attention \
+    --num-query-groups 8 \
+    --kv-channels 128 \
+    --normalization RMSNorm \
+    --norm-epsilon 1e-06 \
+    --swiglu \
+    --qk-layernorm \
+    --disable-bias-linear \
+    --untie-embeddings-and-output-weights \
+    --hidden-dropout 0.0 \
+    --attention-dropout 0.0 \
+    --no-bias-dropout-fusion \
+    --no-bias-swiglu-fusion \
+    --no-persist-layer-norm \
+    --no-rope-fusion \
+    --position-embedding-type rope \
+    --rotary-base 1000000 \
+    --rotary-percent 1.0 \
+    --max-position-embeddings 40960 \
+    --seq-length 2048 \
+    --vocab-size 32000 \
+    --null-tokenizer-vocab-includes-eod \
+    --padded-vocab-size 151936 \
+    --make-vocab-size-divisible-by 128 \
+    --init-method-std 0.02 \
+    --lr 3e-4 \
+    --min-lr 3e-5 \
+    --lr-decay-style cosine \
+    --lr-warmup-iters 0 \
+    --lr-decay-iters ${TRAIN_ITERS} \
+    --weight-decay 0.033 \
+    --adam-beta1 0.9 \
+    --adam-beta2 0.95 \
+    --adam-eps 1e-5 \
+    --clip-grad 1.0 \
+    --optimizer adam \
+    --bf16 \
+    --use-cpu-initialization \
+    --deterministic-mode \
+    --batch-invariant-mode \
+    --mock-data \
+    --tokenizer-type NullTokenizer \
+    --seed 1234 \
+    --log-interval 1 \
+    --tensorboard-dir ${OUTPUT_DIR}/tensorboard"
+
+# ---------------------------------------------------------------------------
+# Launch
+# ---------------------------------------------------------------------------
+NUM_GPUS=${TP_SIZE}
+srun ${SRUN_EXTRA_ARGS:-} --mpi=pmix --overlap --cpu-bind=none \
+     --ntasks=${NUM_GPUS} --ntasks-per-node=${NUM_GPUS} \
+     --gpus-per-node=${NUM_GPUS} \
+     --no-container-mount-home \
+     --container-image="${CONTAINER_IMAGE}" \
+     --container-mounts="${CONTAINER_MOUNTS}" \
+     --container-workdir="${MEGATRON_PATH}" \
+     bash -c "${SETUP_CMD} && ${TRAINING_CMD}" 2>&1 | tee "${LOG_FILE}"
+
+EXPECTED_ITER1="lm loss: 1.275062E+01 grad norm: 90.104"
+EXPECTED_LAST="lm loss: 7.337554E+00 grad norm: 1.981"
+
+extract_iter() {
+    grep "iteration *$1/" "${LOG_FILE}" | head -1 \
+        | grep -oP '(lm loss|grad norm): \S+' | tr '\n' ' ' | sed 's/ $//'
+}
+ITER1=$(extract_iter 1)
+ITER_LAST=$(extract_iter "${TRAIN_ITERS}")
+
+echo
+echo "Done. Loss curve in: ${LOG_FILE}"
+echo "Iter 1   (this run):  ${ITER1}"
+echo "Iter ${TRAIN_ITERS} (this run):  ${ITER_LAST}"
+
+if [[ "${TP_SIZE}" == "4" && "${TRAIN_ITERS}" == "100" ]]; then
+    echo "Expected iter 1:      ${EXPECTED_ITER1}"
+    echo "Expected iter 100:    ${EXPECTED_LAST}"
+    if [[ "${ITER1}" != "${EXPECTED_ITER1}" || "${ITER_LAST}" != "${EXPECTED_LAST}" ]]; then
+        echo "FAIL: numerics mismatch (regression?)"
+        exit 1
+    fi
+    echo "PASS"
+else
+    echo "(No assertion: non-default TP_SIZE/TRAIN_ITERS — expected values only valid for default config.)"
+fi
diff --git a/examples/tp-numerics/submit_qwen3_moe_toy_tp_invariant.sh b/examples/tp-numerics/submit_qwen3_moe_toy_tp_invariant.sh
new file mode 100644
index 00000000000..9f84ee95ec3
--- /dev/null
+++ b/examples/tp-numerics/submit_qwen3_moe_toy_tp_invariant.sh
@@ -0,0 +1,195 @@
+#!/bin/bash
+#
+# Toy MoE TP-invariant numerics smoke test. Raw Megatron-LM.
+#
+# Toy MoE: Qwen3-30B-A3B base architecture + overrides:
+#   - 4 layers (vs 48 in full A3B)
+#   - 8 experts (vs 128)
+#   - top-2 (vs top-8)
+#   - moe_grouped_gemm=False (CRITICAL: bypasses BIK if True)
+#   - moe_token_dispatcher_type=alltoall
+#   - EP=1, GBS=1, MBS=1
+#   ⇒ fits on 1 GPU at TP=1
+#
+# Runs both ways:
+#   - Interactive (active salloc):   bash submit_qwen3_moe_toy_tp_invariant.sh
+#   - SLURM batch:                   sbatch submit_qwen3_moe_toy_tp_invariant.sh
+#
+# Tunable env (override at invocation):
+#   TP_SIZE       (default 1)        — tensor-parallel degree
+#   TRAIN_ITERS   (default 10)
+#   OUTPUT_DIR    (default ./output/qwen3_moe_toy_tp${TP_SIZE}_iter${TRAIN_ITERS})
+
+#SBATCH --nodes=1
+#SBATCH --ntasks-per-node=1
+#SBATCH --gpus-per-node=1
+#SBATCH --time=00:30:00
+#SBATCH --job-name=qwen3_moe_toy_tp_invariant
+#SBATCH --output=slurm_%j.log
+#SBATCH --exclusive
+
+set -euo pipefail
+
+# ---------------------------------------------------------------------------
+# Paths
+# ---------------------------------------------------------------------------
+WORKSPACE="${WORKSPACE:-/home/scratch.jinzex_sw/jinzex-ai-workspace/engineering/reflection-engineering}"
+MEGATRON_PATH="${MEGATRON_PATH:-${WORKSPACE}/third-party/Megatron-LM}"
+CONTAINER_IMAGE="${CONTAINER_IMAGE:-${WORKSPACE}/containers/nemo-25.11.01.sqsh}"
+CONTAINER_MOUNTS="${CONTAINER_MOUNTS:-/home/scratch.jinzex_sw:/home/scratch.jinzex_sw}"
+
+TP_SIZE=${TP_SIZE:-1}
+TRAIN_ITERS=${TRAIN_ITERS:-10}
+OUTPUT_DIR="${OUTPUT_DIR:-${WORKSPACE}/projects/Numerics/tp-numerics/results-b300/qwen3_moe_toy_tp${TP_SIZE}_iter${TRAIN_ITERS}}"
+LOG_FILE="${OUTPUT_DIR}/training.log"
+
+mkdir -p "${OUTPUT_DIR}"
+
+# ---------------------------------------------------------------------------
+# Environment (TP-invariant numerics + determinism)
+# ---------------------------------------------------------------------------
+export NVTE_TP_INVARIANT_MODE=1
+export NVTE_ALLOW_NONDETERMINISTIC_ALGO=0
+export CUBLAS_WORKSPACE_CONFIG=:4096:8
+export NCCL_ALGO=^NVLS
+export NCCL_NVLS_ENABLE=0
+export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True
+# Force unfused attention.
+export NVTE_FLASH_ATTN=0
+export NVTE_FUSED_ATTN=0
+export CUDA_DEVICE_MAX_CONNECTIONS=1
+
+# Sequence parallel only when TP>1.
+SP_FLAG=""
+if [[ "${TP_SIZE}" -gt 1 ]]; then
+    SP_FLAG="--sequence-parallel"
+fi
+
+# TE TP-invariant patches: applied inside the same srun as training because
+# squashfs containers are read-only per-srun. Detects TE version and picks
+# the matching subdir under patches/; fails fast if unsupported.
+SETUP_CMD="\
+TE_DST=\$(python3 -c 'import transformer_engine.pytorch as p; print(p.__path__[0])') && \
+if ! grep -q NVTE_TP_INVARIANT_MODE \$TE_DST/module/linear.py; then \
+    P=${MEGATRON_PATH}/examples/tp-numerics/patches && \
+    TE_VERSION=\$(python3 -c 'import transformer_engine as te; print(te.__version__.split(\"+\")[0])') && \
+    PATCH_SRC=\$P/v\${TE_VERSION}/transformer_engine/pytorch && \
+    if [ ! -d \$PATCH_SRC ]; then echo \"ERROR: no TP-invariant patches for TE \$TE_VERSION (have: \$(ls \$P | grep ^v | tr '\n' ' '))\"; exit 1; fi && \
+    cp -r \$PATCH_SRC/* \$TE_DST/; \
+fi"
+
+# ---------------------------------------------------------------------------
+# Training command — toy MoE config
+# ---------------------------------------------------------------------------
+TRAINING_CMD="python ${MEGATRON_PATH}/pretrain_gpt.py \
+    --distributed-timeout-minutes 60 \
+    --tensor-model-parallel-size ${TP_SIZE} \
+    --pipeline-model-parallel-size 1 \
+    --expert-model-parallel-size 1 \
+    --expert-tensor-parallel-size 1 \
+    --context-parallel-size 1 \
+    --use-distributed-optimizer \
+    ${SP_FLAG} \
+    --transformer-impl transformer_engine \
+    --attention-backend unfused \
+    --micro-batch-size 1 \
+    --global-batch-size 1 \
+    --train-iters ${TRAIN_ITERS} \
+    --eval-iters 0 \
+    --eval-interval $((TRAIN_ITERS + 1)) \
+    --num-layers 4 \
+    --hidden-size 2048 \
+    --ffn-hidden-size 6144 \
+    --num-attention-heads 32 \
+    --group-query-attention \
+    --num-query-groups 4 \
+    --kv-channels 128 \
+    --normalization RMSNorm \
+    --norm-epsilon 1e-06 \
+    --swiglu \
+    --qk-layernorm \
+    --disable-bias-linear \
+    --untie-embeddings-and-output-weights \
+    --hidden-dropout 0.0 \
+    --attention-dropout 0.0 \
+    --no-bias-dropout-fusion \
+    --no-bias-swiglu-fusion \
+    --no-persist-layer-norm \
+    --no-rope-fusion \
+    --position-embedding-type rope \
+    --rotary-base 1000000 \
+    --rotary-percent 1.0 \
+    --max-position-embeddings 40960 \
+    --seq-length 2048 \
+    --vocab-size 32000 \
+    --null-tokenizer-vocab-includes-eod \
+    --padded-vocab-size 151936 \
+    --make-vocab-size-divisible-by 128 \
+    --init-method-std 0.02 \
+    --lr 3e-4 \
+    --min-lr 3e-5 \
+    --lr-decay-style cosine \
+    --lr-warmup-iters 0 \
+    --lr-decay-iters ${TRAIN_ITERS} \
+    --weight-decay 0.033 \
+    --adam-beta1 0.9 \
+    --adam-beta2 0.95 \
+    --adam-eps 1e-5 \
+    --clip-grad 1.0 \
+    --optimizer adam \
+    --bf16 \
+    --use-cpu-initialization \
+    --deterministic-mode \
+    --batch-invariant-mode \
+    --mock-data \
+    --tokenizer-type NullTokenizer \
+    --seed 1234 \
+    --num-experts 8 \
+    --moe-ffn-hidden-size 768 \
+    --moe-router-load-balancing-type aux_loss \
+    --moe-router-topk 2 \
+    --moe-aux-loss-coeff 1e-3 \
+    --moe-permute-fusion \
+    --moe-token-dispatcher-type alltoall \
+    --log-interval 1 \
+    --tensorboard-dir ${OUTPUT_DIR}/tensorboard"
+
+# ---------------------------------------------------------------------------
+# Launch (1 GPU at TP=1, EP=1)
+# ---------------------------------------------------------------------------
+NUM_GPUS=${TP_SIZE}
+srun ${SRUN_EXTRA_ARGS:-} --mpi=pmix --overlap --cpu-bind=none -l \
+     --ntasks=${NUM_GPUS} --ntasks-per-node=${NUM_GPUS} \
+     --gpus-per-node=${NUM_GPUS} \
+     --no-container-mount-home \
+     --container-image="${CONTAINER_IMAGE}" \
+     --container-mounts="${CONTAINER_MOUNTS}" \
+     --container-workdir="${MEGATRON_PATH}" \
+     bash -c "${SETUP_CMD} && ${TRAINING_CMD}" 2>&1 | tee "${LOG_FILE}"
+
+EXPECTED_ITER1="lm loss: 1.232552E+01 grad norm: 6.592"
+EXPECTED_LAST="lm loss: 1.449088E+01 grad norm: 1.412"
+
+extract_iter() {
+    grep "iteration *$1/" "${LOG_FILE}" | head -1 \
+        | grep -oP '(lm loss|grad norm): \S+' | tr '\n' ' ' | sed 's/ $//'
+}
+ITER1=$(extract_iter 1)
+ITER_LAST=$(extract_iter "${TRAIN_ITERS}")
+
+echo
+echo "Done. Loss curve in: ${LOG_FILE}"
+echo "Iter 1  (this run):  ${ITER1}"
+echo "Iter ${TRAIN_ITERS} (this run):  ${ITER_LAST}"
+
+if [[ "${TP_SIZE}" == "1" && "${TRAIN_ITERS}" == "10" ]]; then
+    echo "Expected iter 1:     ${EXPECTED_ITER1}"
+    echo "Expected iter 10:    ${EXPECTED_LAST}"
+    if [[ "${ITER1}" != "${EXPECTED_ITER1}" || "${ITER_LAST}" != "${EXPECTED_LAST}" ]]; then
+        echo "FAIL: numerics mismatch (regression?)"
+        exit 1
+    fi
+    echo "PASS"
+else
+    echo "(No assertion: non-default TP_SIZE/TRAIN_ITERS — expected values only valid for default config.)"
+fi
diff --git a/megatron/core/extensions/transformer_engine.py b/megatron/core/extensions/transformer_engine.py
index 2a82a1e1cf2..9d68832a60a 100644
--- a/megatron/core/extensions/transformer_engine.py
+++ b/megatron/core/extensions/transformer_engine.py
@@ -1581,9 +1581,14 @@ def forward(
         # Default to constructor-provided num_splits unless explicitly overridden
         if num_splits is None:
             num_splits = self.num_splits
-        if num_splits is not None:
-            assert is_te_min_version("2.10.0"), (
-                f"Transformer-Engine v{get_te_version()} must be >= 2.10.0 to support" "num_splits."
+        if num_splits is not None and not is_te_min_version("2.10.0"):
+            import warnings
+            warnings.warn(
+                f"Transformer-Engine v{get_te_version()} < 2.10.0: num_splits support "
+                "requires cherry-picked DotProductAttention patches. "
+                "If patched backends.py/dot_product_attention.py are not installed, "
+                "num_splits will be silently ignored or cause a TypeError.",
+                stacklevel=2,
             )
 
         packed_seq_kwargs = (
diff --git a/megatron/core/optimizer/clip_grads.py b/megatron/core/optimizer/clip_grads.py
index 01095a50154..d2f701bda46 100644
--- a/megatron/core/optimizer/clip_grads.py
+++ b/megatron/core/optimizer/clip_grads.py
@@ -2,6 +2,8 @@
 
 """Gradient clipping."""
 
+import math
+import os
 from typing import List, Optional, Union
 
 import torch
@@ -90,6 +92,54 @@ def get_grad_norm_fp32(
     norm_type = float(norm_type)
     total_norm = 0.0
 
+    # TP-invariant grad norm computation.
+    #
+    # Root cause of TP-dependent norms: multi_tensor_l2norm (Apex CUDA kernel)
+    # reduces over different-shaped TP shards. Float32 CUDA reduction for 3M
+    # elements gives a different result than 1.5M elements for the same
+    # mathematical total → different clip_factors → ALL gradients scaled
+    # differently → training divergence across TP degrees.
+    #
+    # Fix: compute grad_norm in float64 (accurate for logging). The clip_coeff
+    # is rounded to BF16 in clip_grad_by_total_norm_fp32() for TP-invariance.
+    #
+    # Mismatch analysis (inter-TP float64 norm diff ≈ 6e-7 at norm≈18):
+    #   Float32 clip_coeff (no rounding): 40.6% mismatch rate — NOT viable
+    #   BF16-rounded clip_coeff: 0% mismatch, norm=[0.01, 200] — SAFE
+    #   Even with 10× larger models: 0% mismatch with BF16 rounding
+    #
+    # FIXED_GRAD_NORM=<value>: bypass all computation with a hardcoded norm (for debugging).
+    import os as _os
+
+    _fixed_norm = _os.environ.get("FIXED_GRAD_NORM", "")
+    if _fixed_norm:
+        return float(_fixed_norm)
+
+    _tp_invariant_norm = _os.environ.get("NVTE_TP_INVARIANT_MODE", "0") == "1"
+
+    if _tp_invariant_norm and grads_for_norm:
+        # Compute per-grad squared norm in float64 on GPU, accumulate in Python float64.
+        # Float64 reduction error is ~1e-10 relative (vs ~1e-4 in float32).
+        total_norm_sq = 0.0  # Python float = float64
+        for grad in grads_for_norm:
+            total_norm_sq += (grad.double() ** 2).sum().item()
+
+        # All-reduce in float64 across all parallel groups
+        total_norm_cuda = torch.tensor([total_norm_sq], dtype=torch.float64, device='cuda')
+        if data_parallel_group:
+            torch.distributed.all_reduce(
+                total_norm_cuda, op=torch.distributed.ReduceOp.SUM, group=data_parallel_group
+            )
+        torch.distributed.all_reduce(
+            total_norm_cuda, op=torch.distributed.ReduceOp.SUM, group=grad_stats_parallel_group
+        )
+        total_norm = total_norm_cuda.item() ** (1.0 / norm_type)
+
+        # Return accurate float64 norm (for logging). TP-invariant clip_coeff
+        # rounding is handled in clip_grad_by_total_norm_fp32().
+        return total_norm
+
+    # Original float32 norm computation (non-TP-invariant path).
     # Calculate norm.
     if norm_type == inf:
         total_norm = max(grad.abs().max() for grad in grads_for_norm)
@@ -190,10 +240,29 @@ def clip_grad_by_total_norm_fp32(
         multi_tensor_applier(
             multi_tensor_scale_tensor_impl, dummy_overflow_buf, [grads, grads], clip_coeff
         )
-    elif clip_coeff < 1.0:
-        multi_tensor_applier(
-            multi_tensor_scale_impl, dummy_overflow_buf, [grads, grads], clip_coeff
-        )
+    else:
+        # TP-invariant clip_coeff: round scalar clip_coeff to nearest power of 2
+        # (zero mantissa bits) so the tiny float64 inter-TP norm difference (~6e-7)
+        # can't produce different clip_coeffs. The exact clip value doesn't affect
+        # training — only TP-invariance.
+        #
+        # TP-mismatch rate (1M test points, norm=[0.01, 316], inter-TP diff=6e-7):
+        #   Approach           | Mismatches/1M |   Rate  | At 1000× model scale
+        #   Float32 (no round) |     403,003   | 40.300% |     100.000%
+        #   BF16 rounding      |           8   |  0.001% |       0.589%
+        #   Integer rounding   |           1   |  0.000% |       0.033%
+        #   Pow2 (no mantissa) |           0   |  0.000% |       0.001%
+        #
+        # Pow2 chosen: strongest TP-invariance guarantee (~0% even at 1000× scale).
+        # See plot_rounding_analysis.py for full analysis and visualization.
+        if os.environ.get("NVTE_TP_INVARIANT_MODE", "0") == "1":
+            if clip_coeff > 0:
+                clip_coeff = 2.0 ** round(math.log2(clip_coeff))
+
+        if clip_coeff < 1.0:
+            multi_tensor_applier(
+                multi_tensor_scale_impl, dummy_overflow_buf, [grads, grads], clip_coeff
+            )
 
 
 def count_zeros_fp32(
diff --git a/megatron/core/tensor_parallel/cross_entropy.py b/megatron/core/tensor_parallel/cross_entropy.py
index 27c8f063440..ddfeb8dce03 100644
--- a/megatron/core/tensor_parallel/cross_entropy.py
+++ b/megatron/core/tensor_parallel/cross_entropy.py
@@ -1,5 +1,6 @@
 # Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.
 
+import os
 from typing import Tuple
 
 import torch
@@ -12,6 +13,8 @@
 
 from .utils import VocabUtility
 
+_TP_INVARIANT_MODE = os.environ.get("NVTE_TP_INVARIANT_MODE", "0") == "1"
+
 
 class VocabParallelCrossEntropy:
     """
@@ -151,11 +154,21 @@ def forward(ctx, vocab_parallel_logits, target, label_smoothing=0.0):
             group=get_tensor_model_parallel_group(),
         )
 
-        torch.distributed.all_reduce(
-            sum_exp_logits,
-            op=torch.distributed.ReduceOp.SUM,
-            group=get_tensor_model_parallel_group(),
-        )
+        if _TP_INVARIANT_MODE and world_size > 1:
+            # All-gather exp_logits and recompute sum locally so the
+            # reduction order matches TP=1 (single contiguous sum).
+            gathered = [torch.empty_like(exp_logits) for _ in range(world_size)]
+            torch.distributed.all_gather(
+                gathered, exp_logits.contiguous(), group=get_tensor_model_parallel_group()
+            )
+            sum_exp_logits = torch.cat(gathered, dim=-1).sum(dim=-1)
+            del gathered
+        else:
+            torch.distributed.all_reduce(
+                sum_exp_logits,
+                op=torch.distributed.ReduceOp.SUM,
+                group=get_tensor_model_parallel_group(),
+            )
 
         exp_logits, loss = VocabParallelCrossEntropy.calculate_cross_entropy_loss(
             exp_logits, predicted_logits, sum_exp_logits
diff --git a/megatron/core/tensor_parallel/layers.py b/megatron/core/tensor_parallel/layers.py
index 610700f0a95..c5b2106bd70 100644
--- a/megatron/core/tensor_parallel/layers.py
+++ b/megatron/core/tensor_parallel/layers.py
@@ -7,6 +7,8 @@
 import os
 import warnings
 from functools import partial
+
+_TP_INVARIANT_MODE = os.environ.get("NVTE_TP_INVARIANT_MODE", "0") == "1"
 from typing import Any, Callable, List, Optional, Tuple
 
 import torch
@@ -526,7 +528,18 @@ def backward(ctx, grad_output):
                 total_input = all_gather_buffer
             else:
                 total_input = input
-        grad_input = grad_output.matmul(weight)
+        tp_size = tp_group.size()
+        _tp_inv = _TP_INVARIANT_MODE and tp_size > 1
+
+        if _tp_inv:
+            full_w = [torch.empty_like(weight) for _ in range(tp_size)]
+            torch.distributed.all_gather(full_w, weight.contiguous(), group=tp_group)
+            full_go = [torch.empty_like(grad_output) for _ in range(tp_size)]
+            torch.distributed.all_gather(full_go, grad_output.contiguous(), group=tp_group)
+            grad_input = torch.cat(full_go, dim=-1).matmul(torch.cat(full_w, dim=0))
+            del full_w, full_go
+        else:
+            grad_input = grad_output.matmul(weight)
 
         if ctx.sequence_parallel and wgrad_compute:
             # pylint: disable=possibly-used-before-assignment
@@ -537,13 +550,20 @@ def backward(ctx, grad_output):
                 grad_output, total_input
             )
 
-        if ctx.allreduce_dgrad:
+        if _tp_inv:
+            if ctx.sequence_parallel:
+                tp_rank = torch.distributed.get_rank(group=tp_group)
+                seq_chunk = grad_input.size(0) // tp_size
+                sub_grad_input = grad_input[tp_rank * seq_chunk : (tp_rank + 1) * seq_chunk].contiguous()
+                handle = None
+            else:
+                handle = None
+        elif ctx.allreduce_dgrad:
             # Asynchronous all-reduce
             handle = torch.distributed.all_reduce(grad_input, group=tp_group, async_op=True)
             # Here we rely on CUDA_DEVICE_MAX_CONNECTIONS=1 to ensure that the
             # all-reduce is scheduled before the weight gradient computation
-
-        if ctx.sequence_parallel:
+        elif ctx.sequence_parallel:
             assert not ctx.allreduce_dgrad
             dim_size = list(input.size())
             sub_grad_input = torch.empty(
@@ -555,6 +575,8 @@ def backward(ctx, grad_output):
             )
             # Here we rely on CUDA_DEVICE_MAX_CONNECTIONS=1 to ensure that the
             # reduce scatter is scheduled before the weight gradient computation
+        else:
+            handle = None
 
         if ctx.gradient_accumulation_fusion:
             if wgrad_compute:
@@ -629,13 +651,15 @@ def backward(ctx, grad_output):
         grad_bias = grad_output.sum(dim=0) if use_bias else None
 
         if ctx.sequence_parallel:
-            handle.wait()
+            if handle is not None:
+                handle.wait()
             # Need to return None's as gradient has to flow for all the input arguments
             # provided during forward
             return (sub_grad_input, grad_weight, grad_bias, None, None, None, None, None, None)
 
         if ctx.allreduce_dgrad:
-            handle.wait()
+            if handle is not None:
+                handle.wait()
 
         return grad_input, grad_weight, grad_bias, None, None, None, None, None, None
 
diff --git a/megatron/core/tokenizers/utils/build_tokenizer.py b/megatron/core/tokenizers/utils/build_tokenizer.py
index bf02451ae6c..2be563dc714 100644
--- a/megatron/core/tokenizers/utils/build_tokenizer.py
+++ b/megatron/core/tokenizers/utils/build_tokenizer.py
@@ -73,7 +73,10 @@ def build_tokenizer(args, **kwargs):
         )
         metadata = {'library': tokenizer_library}
         if args.vocab_size:
-            kwargs['vocab_size'] = args.vocab_size
+            if getattr(args, 'null_tokenizer_vocab_includes_eod', False):
+                kwargs['vocab_size'] = args.vocab_size - 1
+            else:
+                kwargs['vocab_size'] = args.vocab_size
         tokenizer = MegatronTokenizer.from_pretrained(metadata_path=metadata, **kwargs)
 
         # Add vocab size (if not already set from a checkpoint).
diff --git a/megatron/core/transformer/custom_layers/batch_invariant_kernels.py b/megatron/core/transformer/custom_layers/batch_invariant_kernels.py
index 6b4311fe540..bba037ac8c2 100644
--- a/megatron/core/transformer/custom_layers/batch_invariant_kernels.py
+++ b/megatron/core/transformer/custom_layers/batch_invariant_kernels.py
@@ -7,6 +7,7 @@
 import importlib
 import importlib.util
 import logging
+import os
 from collections import namedtuple
 from collections.abc import Callable
 from typing import Any, Dict, List, Optional
@@ -494,20 +495,48 @@ def _log_softmax_batch_invariant(input, dim, _half_to_float):
 
 
 def mean_batch_invariant(input, dim, keepdim=False, dtype: torch.dtype | None = None):
-    """Batch-invariant replacement for `aten::mean.dim` over one or more dimensions."""
+    """Batch-invariant replacement for `aten::mean.dim` over one or more dimensions.
+
+    Determinism across SM architectures:
+    - Full reduction (dim=[]): two-stage Triton via `mean_dim`. Stage 1 reduces
+      M=4096 chunks in parallel; stage 2 reduces the M partial means serially.
+      Both stages use fixed launch shape → cross-SM deterministic.
+    - Single-dim: `mean_dim` (Triton, fixed reduction order).
+    - Multi-dim partial: TODO — currently uses non-deterministic torch.sum.
+    """
     assert dtype is None or dtype == torch.float32, f"unsupported dtype: {dtype}"
+    if not dim:
+        # Full reduction: two-stage parallel deterministic via Triton mean_dim.
+        flat = input.reshape(-1)
+        n = flat.numel()
+        M = 4096  # fixed parallelism, SM-independent
+        if n % M == 0:
+            chunk = n // M
+            partials = mean_dim(flat.reshape(M, chunk), 1, keepdim=False)
+            return mean_dim(partials.reshape(1, M, 1), 1, keepdim=False).reshape(())
+        # Fallback: single-block (slow but deterministic).
+        return mean_dim(flat, 0, keepdim=False)
     if len(dim) == 1:
         return mean_dim(input, dim[0], keepdim=keepdim)
-    else:
-        assert input.dtype in {
-            torch.float16,
-            torch.bfloat16,
-            torch.float32,
-        }, "only float types supported for now"
-        n_elems = 1
-        for d in dim:
-            n_elems *= input.shape[d]
-        return torch.sum(input, dim=dim, keepdim=keepdim, dtype=torch.float32) / n_elems
+    # TODO: multi-dim partial reduction is non-deterministic across SMs.
+    # Use Triton (e.g., permute+reshape+mean_dim) to close this gap.
+    assert input.dtype in {
+        torch.float16,
+        torch.bfloat16,
+        torch.float32,
+    }, "only float types supported for now"
+    n_elems = 1
+    for d in dim:
+        n_elems *= input.shape[d]
+    return torch.sum(input, dim=dim, keepdim=keepdim, dtype=torch.float32) / n_elems
+
+
+# TODO: BIK does not patch aten::sum, aten::var, aten::std, aten::norm.
+# These reductions go through native PyTorch CUDA kernels and are not
+# cross-SM deterministic — meaning E2E loss reductions (e.g., Megatron's
+# `torch.sum(losses*mask) / mask.sum()`) will diverge slightly across
+# SM_90/SM_100 even with BIK on. Add Triton fixed-order kernels to close
+# this gap.
 
 
 AttentionBlockSize = namedtuple("AttentionBlockSize", ["block_m", "block_n"])
@@ -525,6 +554,7 @@ def get_batch_invariant_attention_block_size() -> AttentionBlockSize:
 _MEG_TE_GENERAL_GEMM_ORIG = None
 _TE_RMSNORM_FUNC_ORIGS: Dict[str, Any] = {}
 _TE_GEMM_FUNC_ORIGS: Dict[str, Any] = {}
+_TE_RMSNORM_BWD_ORIG = None
 
 
 def _import_module_if_available(name: str):
@@ -576,6 +606,69 @@ def _te_patch_for_batch_invariant():
         _MEG_TE_GENERAL_GEMM_ORIG = meg_te.general_gemm
         meg_te.general_gemm = _te_general_gemm_patched
 
+    # Patch tex.rmsnorm_bwd for TP-invariant dgamma in fused layers.
+    # Standalone RMSNorm is handled by BatchInvariantRMSNormFn.backward (no
+    # call to tex.rmsnorm_bwd), so there is no double-patching concern.
+    global _TE_RMSNORM_BWD_ORIG
+    import transformer_engine_torch as _tex
+
+    if _TE_RMSNORM_BWD_ORIG is None and hasattr(_tex, "rmsnorm_bwd"):
+        _TE_RMSNORM_BWD_ORIG = _tex.rmsnorm_bwd
+
+        def _te_rmsnorm_bwd_tp_invariant(dy, x, rsigma, gamma, sm_margin, zero_centered_gamma):
+            """Wrapper around tex.rmsnorm_bwd that recomputes dgamma in
+            Python FP32 with consistent accumulation order across all TP
+            degrees (including TP=1).
+
+            dx is left to the original CUDA kernel — its per-token BF16
+            computation is already TP-invariant (same result for each token
+            regardless of batch size).  Only dgamma needs fixing because
+            it's a reduction over tokens whose count differs under SP.
+            """
+            d_norm_x = dy.clone()
+            dx, dgamma_cuda = _TE_RMSNORM_BWD_ORIG(
+                dy, x, rsigma, gamma, sm_margin, zero_centered_gamma,
+            )
+
+            if os.environ.get("NVTE_TP_INVARIANT_MODE", "0") != "1":
+                return dx, dgamma_cuda
+            if not torch.distributed.is_initialized():
+                return dx, dgamma_cuda
+
+            from megatron.core.parallel_state import (
+                get_tensor_model_parallel_group,
+                get_tensor_model_parallel_world_size,
+                get_tensor_model_parallel_rank,
+            )
+
+            tp_size = get_tensor_model_parallel_world_size()
+
+            # Recompute dgamma in Python FP32 for consistent accumulation.
+            rsigma_f = rsigma.float()
+            if rsigma_f.ndim < d_norm_x.ndim:
+                rsigma_f = rsigma_f.unsqueeze(-1)
+            per_token = d_norm_x.float() * x.float() * rsigma_f
+            red_dims = tuple(range(0, per_token.ndim - 1))
+
+            if tp_size > 1:
+                tp_group = get_tensor_model_parallel_group()
+                gather_dim = 0 if per_token.ndim == 2 else per_token.ndim - 2
+                chunks = [torch.empty_like(per_token) for _ in range(tp_size)]
+                torch.distributed.all_gather(
+                    chunks, per_token.contiguous(), group=tp_group,
+                )
+                per_token = torch.cat(chunks, dim=gather_dim)
+
+            dgamma_python = per_token.sum(dim=red_dims).to(gamma.dtype)
+
+            if tp_size > 1:
+                if get_tensor_model_parallel_rank() != 0:
+                    dgamma_python.zero_()
+
+            return dx, dgamma_python
+
+        _tex.rmsnorm_bwd = _te_rmsnorm_bwd_tp_invariant
+
     # Patch RMSNorm.forward once (class may be on te or te.pytorch)
     rms_cls = getattr(te, "RMSNorm", None)
     if rms_cls is None:
@@ -838,6 +931,11 @@ def _te_general_gemm_patched(*args, **kwargs) -> List[torch.Tensor]:
         raise RuntimeError("TE general_gemm original not captured; patching order issue")
 
     A, B, out_dtype, layout, out, bias, grad = _extract_te_gemm_args(args, kwargs)
+
+    # Fall back to original for empty tensors (e.g., MoE experts with 0 tokens)
+    if A is not None and B is not None and (A.numel() == 0 or B.numel() == 0):
+        return _TE_GENERAL_GEMM_ORIG(*args, **kwargs)
+
     extra_output = kwargs.get("extra_output", None)
     ub = kwargs.get("ub", None)
     ub_type = kwargs.get("ub_type", None)
@@ -914,6 +1012,13 @@ def backward(ctx, grad_output: torch.Tensor):
 
         Computes gradients w.r.t. input and weight while matching TE's fp32
         accumulation and reduction behavior for numerical stability.
+
+        When NVTE_TP_INVARIANT_MODE=1, dgamma is computed in Python FP32
+        with consistent accumulation order across all TP degrees.  For
+        TP>1, per-token contributions are all-gathered before summation
+        and the result is kept on rank 0 only (others zeroed) so that the
+        subsequent all-reduce in finalize_model_grads recovers the correct
+        total without FP32 precision loss from divide-then-sum.
         """
         x, weight, rsigma = ctx.saved_tensors
         w_eff = (weight + 1.0) if ctx.zero_centered_gamma else weight
@@ -926,7 +1031,36 @@ def backward(ctx, grad_output: torch.Tensor):
         D = x.shape[-1]
 
         red_dims = tuple(range(0, go_fp32.ndim - 1))
-        g_w = (go_fp32 * x_fp32 * r).sum(dim=red_dims).to(weight.dtype)
+        per_token = go_fp32 * x_fp32 * r
+
+        _tp_invariant = os.environ.get("NVTE_TP_INVARIANT_MODE", "0") == "1"
+        if _tp_invariant and torch.distributed.is_initialized():
+            from megatron.core.parallel_state import (
+                get_tensor_model_parallel_group,
+                get_tensor_model_parallel_rank,
+                get_tensor_model_parallel_world_size,
+            )
+            tp_size = get_tensor_model_parallel_world_size()
+            if tp_size > 1:
+                tp_group = get_tensor_model_parallel_group()
+                # 2D [seq/TP, H]: SP-sharded RMSNorm → gather along seq (dim 0)
+                # 3D+ [seq, heads/TP, H]: TP-sharded qk_layernorm → gather along
+                #   the penultimate dim (heads)
+                gather_dim = 0 if per_token.ndim == 2 else per_token.ndim - 2
+                chunks = [torch.empty_like(per_token) for _ in range(tp_size)]
+                torch.distributed.all_gather(
+                    chunks, per_token.contiguous(), group=tp_group,
+                )
+                per_token = torch.cat(chunks, dim=gather_dim)
+                del chunks
+
+            g_w = per_token.sum(dim=red_dims).to(weight.dtype)
+
+            if tp_size > 1 and get_tensor_model_parallel_rank() != 0:
+                g_w.zero_()
+        else:
+            g_w = per_token.sum(dim=red_dims).to(weight.dtype)
+        del per_token
 
         s = (go_fp32 * x_fp32 * w_fp32).sum(dim=-1, keepdim=True)
         dx = go_fp32 * (w_fp32 * r) - (w_fp32 * r3) * (s * x_fp32) / D
diff --git a/megatron/core/transformer/transformer_config.py b/megatron/core/transformer/transformer_config.py
index 4224b0cfe4c..6ec91f08ab2 100644
--- a/megatron/core/transformer/transformer_config.py
+++ b/megatron/core/transformer/transformer_config.py
@@ -2703,9 +2703,9 @@ def __post_init__(self):
             )
 
         if self.batch_invariant_mode:
-            assert (
-                self.attention_backend == AttnBackend.flash
-            ), "Batch invariant mode only supports FlashAttention"
+            assert self.attention_backend in (
+                AttnBackend.flash, AttnBackend.auto, AttnBackend.unfused,
+            ), "Batch invariant mode only supports FlashAttention, auto, or unfused"
 
         if self.sequence_packing_scheduler is not None:
             # Check TE version.
diff --git a/megatron/training/arguments.py b/megatron/training/arguments.py
index 6a108a0d6d0..4e9dcbbb7e6 100644
--- a/megatron/training/arguments.py
+++ b/megatron/training/arguments.py
@@ -3974,6 +3974,14 @@ def _add_tokenizer_args(parser):
     group.add_argument(
         '--tokenizer-model', type=str, default=None, help='Sentencepiece tokenizer model.'
     )
+    group.add_argument(
+        '--null-tokenizer-vocab-includes-eod',
+        action='store_true',
+        default=False,
+        help='Treat --vocab-size as the total vocab including eod. '
+        'When set, NullTokenizer eod_id=N-1 (instead of default N). '
+        'For compatibility with Megatron-Bridge convention.',
+    )
     group.add_argument(
         '--tokenizer-metadata',
         type=str,
diff --git a/megatron/training/training.py b/megatron/training/training.py
index 8b4247be3b0..2d8173b3926 100644
--- a/megatron/training/training.py
+++ b/megatron/training/training.py
@@ -1065,7 +1065,11 @@ def pretrain(
 
         set_ideal_affinity_for_current_gpu()
     if args.batch_invariant_mode:
-        print_rank_0("Enabling batch invariant mode globally", flush=True)
+        from megatron.core.transformer.custom_layers.batch_invariant_kernels import (
+            enable_batch_invariant_mode,
+        )
+
+        print_rank_0("Enabling batch invariant mode globally")
         enable_batch_invariant_mode()
 
     if args.log_progress:
diff --git a/tests/unit_tests/transformer/test_te_layers_batch_invariant.py b/tests/unit_tests/transformer/test_te_layers_batch_invariant.py
index e2d52727925..2b6c7224dfd 100644
--- a/tests/unit_tests/transformer/test_te_layers_batch_invariant.py
+++ b/tests/unit_tests/transformer/test_te_layers_batch_invariant.py
@@ -735,3 +735,158 @@ def test_bik_te_general_gemm_numerical_parity(dtype):
         C_bik = _te_general_gemm(A, B, out_dtype=dtype, layout="TN")[0]
 
     torch.testing.assert_close(C_bik, C_ref, **_tols(dtype))
+
+
+# ============================================================================
+# Batch-Invariant bmm and softmax tests
+# ============================================================================
+
+try:
+    from megatron.core.transformer.custom_layers.batch_invariant_kernels import (
+        bmm_batch_invariant,
+        _softmax_batch_invariant,
+    )
+    HAVE_BMM_SOFTMAX_BIK = True
+except ImportError:
+    # Older BIK versions don't ship bmm/softmax helpers; skip dependent tests.
+    bmm_batch_invariant = None
+    _softmax_batch_invariant = None
+    HAVE_BMM_SOFTMAX_BIK = False
+
+_skip_no_bmm_softmax = pytest.mark.skipif(
+    not HAVE_BMM_SOFTMAX_BIK,
+    reason="bmm_batch_invariant / _softmax_batch_invariant not in this BIK build",
+)
+
+
+@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float32])
+def test_bmm_batch_invariant_chunking(dtype):
+    """Verify bmm gives identical per-element results regardless of batch size.
+
+    Simulates attention Q@K^T where batch=nheads varies with TP degree.
+    """
+    torch.manual_seed(42)
+    B1, B2, M, K, N = 5, 11, 64, 64, 64
+    a_full = torch.randn(B1 + B2, M, K, **_device(dtype))
+    b_full = torch.randn(B1 + B2, K, N, **_device(dtype))
+
+    c_full = bmm_batch_invariant(a_full, b_full)
+    c1 = bmm_batch_invariant(a_full[:B1].contiguous(), b_full[:B1].contiguous())
+    c2 = bmm_batch_invariant(a_full[B1:].contiguous(), b_full[B1:].contiguous())
+    c_cat = torch.cat([c1, c2], dim=0)
+
+    assert c_full.shape == c_cat.shape
+    assert torch.equal(c_full, c_cat), (
+        f"bmm results differ: max_diff={(c_full.float()-c_cat.float()).abs().max().item():.6e}"
+    )
+
+
+@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float32])
+def test_bmm_batch_invariant_attention_dims(dtype):
+    """Test with realistic attention dimensions (Qwen3-0.6B: seq=2048, d=128)."""
+    torch.manual_seed(42)
+    max_heads = 16
+    seq, head_dim = 2048, 128
+
+    # Q@K^T: [nheads, seq, head_dim] @ [nheads, head_dim, seq]
+    a_full = torch.randn(max_heads, seq, head_dim, **_device(dtype))
+    b_full = torch.randn(max_heads, head_dim, seq, **_device(dtype))
+
+    for nheads in [16, 8, 4]:
+        c = bmm_batch_invariant(a_full[:nheads].contiguous(), b_full[:nheads].contiguous())
+        c_ref = bmm_batch_invariant(a_full[:1].contiguous(), b_full[:1].contiguous())
+        assert torch.equal(c[0], c_ref[0]), (
+            f"nheads={nheads}: head 0 differs from nheads=1 by "
+            f"{(c[0].float()-c_ref[0].float()).abs().max().item():.6e}"
+        )
+
+
+@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float32])
+def test_bmm_batch_invariant_numerical_parity(dtype):
+    """Verify bmm_batch_invariant is numerically close to torch.bmm."""
+    torch.manual_seed(42)
+    B, M, K, N = 8, 64, 64, 64
+    a = torch.randn(B, M, K, **_device(dtype))
+    b = torch.randn(B, K, N, **_device(dtype))
+
+    c_bik = bmm_batch_invariant(a, b)
+    c_ref = torch.bmm(a.float(), b.float()).to(dtype)
+    torch.testing.assert_close(c_bik, c_ref, **_tols(dtype))
+
+
+@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float32])
+def test_softmax_batch_invariant_chunking(dtype):
+    """Verify softmax gives identical per-element results regardless of batch size."""
+    torch.manual_seed(42)
+    B1, B2, M, N = 5, 11, 64, 64
+    x_full = torch.randn(B1 + B2, M, N, **_device(dtype))
+
+    y_full = _softmax_batch_invariant(x_full, dim=-1, half_to_float=False)
+    y1 = _softmax_batch_invariant(x_full[:B1].contiguous(), dim=-1, half_to_float=False)
+    y2 = _softmax_batch_invariant(x_full[B1:].contiguous(), dim=-1, half_to_float=False)
+    y_cat = torch.cat([y1, y2], dim=0)
+
+    assert y_full.shape == y_cat.shape
+    assert torch.equal(y_full, y_cat), (
+        f"softmax results differ: max_diff={(y_full.float()-y_cat.float()).abs().max().item():.6e}"
+    )
+
+
+@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float32])
+def test_softmax_batch_invariant_attention_dims(dtype):
+    """Test softmax with realistic attention score dimensions."""
+    torch.manual_seed(42)
+    max_heads = 16
+    seq = 2048
+
+    scores = torch.randn(max_heads, seq, seq, **_device(dtype))
+
+    for nheads in [16, 8, 4]:
+        y = _softmax_batch_invariant(scores[:nheads].contiguous(), dim=-1, half_to_float=False)
+        y_ref = _softmax_batch_invariant(scores[:1].contiguous(), dim=-1, half_to_float=False)
+        assert torch.equal(y[0], y_ref[0]), (
+            f"nheads={nheads}: head 0 differs from nheads=1 by "
+            f"{(y[0].float()-y_ref[0].float()).abs().max().item():.6e}"
+        )
+
+
+@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
+def test_softmax_batch_invariant_half_to_float(dtype):
+    """Verify half_to_float mode returns fp32 output."""
+    torch.manual_seed(42)
+    x = torch.randn(4, 32, 32, **_device(dtype))
+
+    y = _softmax_batch_invariant(x, dim=-1, half_to_float=True)
+    assert y.dtype == torch.float32
+
+    # Row sums should be ~1.0
+    row_sums = y.sum(dim=-1)
+    assert (row_sums - 1.0).abs().max().item() < 1e-5
+
+
+@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float32])
+def test_softmax_batch_invariant_numerical_parity(dtype):
+    """Verify softmax_batch_invariant is numerically close to torch.softmax."""
+    torch.manual_seed(42)
+    x = torch.randn(8, 64, 64, **_device(dtype))
+
+    y_bik = _softmax_batch_invariant(x, dim=-1, half_to_float=False)
+    y_ref = torch.softmax(x.float(), dim=-1).to(dtype)
+    torch.testing.assert_close(y_bik, y_ref, **_tols(dtype))
+
+
+@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float32])
+def test_mean_batch_invariant_full_reduction_chunking_invariant(dtype):
+    """Full-reduction mean must be invariant to how the input is chunked.
+
+    Mean of equal-size halves' means should equal mean of the whole. Exercises
+    the two-stage Triton reduction path.
+    """
+    torch.manual_seed(42)
+    x = torch.randn(4096, 2, 4096, **_device(dtype))
+    flat = x.reshape(-1)
+    half = flat.numel() // 2
+    with set_batch_invariant_mode(True):
+        m_full = x.mean()
+        m_halves = (flat[:half].mean() + flat[half:].mean()) / 2
+    torch.testing.assert_close(m_full, m_halves, **_tols(dtype))
diff --git a/tests/unit_tests/transformer/test_tp_invariant.py b/tests/unit_tests/transformer/test_tp_invariant.py
new file mode 100644
index 00000000000..c40fce418a3
--- /dev/null
+++ b/tests/unit_tests/transformer/test_tp_invariant.py
@@ -0,0 +1,210 @@
+# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+"""TP-invariant numerics: TP=1 ≡ TP=2 bitwise on a small TransformerBlock.
+
+Validates the TP-invariant patch stack (TE column/row parallel linear with
+all-gather, LayerNormColumnParallelLinear, RowParallelLinear, etc.) produces
+forward + backward outputs that are bitwise-identical across tensor-parallel
+degrees, given the canonical deterministic config: CPU init, deterministic
+mode, unfused attention, and the patched MLM cross_entropy / clip_grads /
+batch_invariant_kernels.
+
+Requires:
+    - ≥2 GPUs (TP=2/4 cases skipped if not enough)
+    - TE TP-invariant patches applied (test skipped otherwise;
+      see examples/tp-numerics/README.md "Install TE patches")
+
+All required env vars are set by the test itself — no shell setup needed.
+
+Run:
+    torchrun --nproc_per_node=4 -m pytest \\
+        tests/unit_tests/transformer/test_tp_invariant.py -v -s
+
+Pass criterion: torch.equal(out_tp1, out_tpN) and torch.equal(grad_tp1, grad_tpN)
+for N in {2, 4}, dtype in {fp32, bf16}.
+"""
+import inspect
+import os
+
+# Required env vars for TP-invariant numerics; set BEFORE any TE/Megatron import
+# so the values take effect at module init. Self-contained: no shell-level setup
+# required to run this test.
+os.environ.setdefault("NVTE_TP_INVARIANT_MODE", "1")
+os.environ.setdefault("NVTE_ALLOW_NONDETERMINISTIC_ALGO", "0")
+# Force unfused attention path. MLM's `attention_backend=AttnBackend.unfused`
+# config field does NOT propagate to TE — TE picks backend from these env vars
+# (Thread A finding: FA2 on SM_100 leaks 2^-5 in backward; unfused is bitwise).
+os.environ.setdefault("NVTE_FLASH_ATTN", "0")
+os.environ.setdefault("NVTE_FUSED_ATTN", "0")
+
+import pytest
+import torch
+
+from megatron.core.models.gpt.gpt_layer_specs import get_gpt_layer_with_transformer_engine_spec
+from megatron.core.tensor_parallel.random import model_parallel_cuda_manual_seed
+from megatron.core.transformer.custom_layers.batch_invariant_kernels import (
+    enable_batch_invariant_mode,
+)
+from megatron.core.transformer.enums import AttnBackend
+from megatron.core.transformer.transformer_block import TransformerBlock
+from megatron.core.transformer.transformer_config import TransformerConfig
+
+from tests.unit_tests.test_utilities import Utils
+
+# Enable BIK at module-load time so all GEMMs use batch-invariant Triton kernels.
+# Required for bitwise TP-invariance: without BIK, GEMMs at different M-sizes
+# (full vs sharded) produce ULP-level differences that compound through bwd.
+enable_batch_invariant_mode()
+
+
+# --- Skip-conditions ---------------------------------------------------------
+
+def _te_tp_invariant_patches_applied() -> bool:
+    """Return True iff the TE TP-invariant patches have been cp'd into the
+    transformer_engine site-packages copy."""
+    try:
+        import transformer_engine.pytorch.module.linear as _te_linear
+        return "NVTE_TP_INVARIANT_MODE" in inspect.getsource(_te_linear)
+    except Exception:
+        return False
+
+
+_HAVE_TWO_GPUS = torch.cuda.is_available() and torch.cuda.device_count() >= 2
+_TE_PATCHED = _te_tp_invariant_patches_applied()
+
+
+# --- Test config -------------------------------------------------------------
+
+SEQ_LEN = 64
+BATCH = 2
+HIDDEN = 128
+NUM_HEADS = 4
+NUM_LAYERS = 1
+SEED = 42
+
+
+def _build_config(tp_size: int, dtype: torch.dtype) -> TransformerConfig:
+    bf16 = dtype == torch.bfloat16
+    fp16 = dtype == torch.float16
+    return TransformerConfig(
+        # arch
+        num_layers=NUM_LAYERS,
+        hidden_size=HIDDEN,
+        ffn_hidden_size=HIDDEN * 4,
+        num_attention_heads=NUM_HEADS,
+        num_query_groups=NUM_HEADS,  # plain MHA (no GQA); SelfAttention requires this set
+        kv_channels=HIDDEN // NUM_HEADS,  # SelfAttention requires this set; default None
+        normalization="RMSNorm",
+        attention_backend=AttnBackend.unfused,
+        # parallelism
+        tensor_model_parallel_size=tp_size,
+        pipeline_model_parallel_size=1,
+        sequence_parallel=False,  # SP=False keeps input replicated; simpler comparison
+        # determinism
+        use_cpu_initialization=True,
+        deterministic_mode=True,
+        hidden_dropout=0.0,
+        attention_dropout=0.0,
+        # precision
+        params_dtype=dtype,
+        bf16=bf16,
+        fp16=fp16,
+    )
+
+
+def _run_block(
+    tp_size: int,
+    dtype: torch.dtype,
+    input_full: torch.Tensor,
+    grad_full: torch.Tensor,
+) -> tuple[torch.Tensor | None, torch.Tensor | None]:
+    """Build a TransformerBlock at the given TP size, run forward + backward
+    with the supplied input and grad, return (output, input.grad) gathered to
+    rank 0 (None on other ranks).
+    """
+    Utils.initialize_model_parallel(
+        tensor_model_parallel_size=tp_size,
+        pipeline_model_parallel_size=1,
+    )
+    # Reset both global torch RNG and Megatron MP-cuda seed so weight init is
+    # deterministic and identical across TP=1 and TP=2 calls (between the two
+    # calls, the prior call's forward/backward has advanced global torch RNG).
+    torch.manual_seed(SEED)
+    torch.cuda.manual_seed_all(SEED)
+    model_parallel_cuda_manual_seed(SEED)
+
+    config = _build_config(tp_size=tp_size, dtype=dtype)
+    spec = get_gpt_layer_with_transformer_engine_spec()
+    block = TransformerBlock(config=config, spec=spec).cuda()
+    block.train()
+
+    inp = input_full.clone().detach().to(dtype=dtype, device="cuda").requires_grad_(True)
+    grad = grad_full.clone().detach().to(dtype=dtype, device="cuda")
+
+    out = block(hidden_states=inp, attention_mask=None)
+    # TransformerBlock.forward returns hidden_states (Tensor) or (Tensor, None)
+    # depending on whether context is used. Normalize.
+    if isinstance(out, tuple):
+        out = out[0]
+    out.backward(grad)
+
+    rank0 = torch.distributed.is_initialized() and torch.distributed.get_rank() == 0
+    if rank0:
+        out_cpu = out.detach().to(device="cpu", dtype=torch.float32)
+        grad_cpu = inp.grad.detach().to(device="cpu", dtype=torch.float32)
+        return out_cpu, grad_cpu
+    return None, None
+
+
+# --- Test --------------------------------------------------------------------
+
+@pytest.mark.skipif(not _HAVE_TWO_GPUS, reason="≥2 GPUs required for TP>1 path")
+@pytest.mark.skipif(
+    not _TE_PATCHED,
+    reason=(
+        "TE TP-invariant patches not applied. See "
+        "examples/tp-numerics/README.md 'Install TE patches' "
+        "to apply patches from the jinzex/tp-invariant-numerics TE branch."
+    ),
+)
+class TestTPInvariantTransformerBlock:
+    """Bitwise-identical fwd+bwd across TP=1 and TP=N (N in {2, 4})."""
+
+    def teardown_method(self, method):
+        Utils.destroy_model_parallel()
+
+    @pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16])
+    @pytest.mark.parametrize("compare_tp", [2, 4])
+    def test_tp1_equals_tpN_bitwise(self, dtype, compare_tp):
+        if torch.cuda.device_count() < compare_tp:
+            pytest.skip(f"{compare_tp} GPUs required for TP={compare_tp}; have {torch.cuda.device_count()}")
+        # Also skip if WORLD_SIZE doesn't match what we need for TP=compare_tp.
+        # `Utils.initialize_model_parallel` requires world_size to be a multiple of TP.
+        world_size = int(os.environ.get("WORLD_SIZE", "1"))
+        if world_size < compare_tp:
+            pytest.skip(f"WORLD_SIZE={world_size} < compare_tp={compare_tp}; launch torchrun --nproc_per_node={compare_tp}")
+
+        # Deterministic input on CPU (architecture-independent), replicated across ranks
+        gen = torch.Generator(device="cpu").manual_seed(SEED)
+        input_full = torch.randn(SEQ_LEN, BATCH, HIDDEN, generator=gen, dtype=torch.float32)
+        gen2 = torch.Generator(device="cpu").manual_seed(SEED + 1)
+        grad_full = torch.randn(SEQ_LEN, BATCH, HIDDEN, generator=gen2, dtype=torch.float32)
+
+        out_tp1, grad_tp1 = _run_block(
+            tp_size=1, dtype=dtype, input_full=input_full, grad_full=grad_full
+        )
+        Utils.destroy_model_parallel()
+        out_tpN, grad_tpN = _run_block(
+            tp_size=compare_tp, dtype=dtype, input_full=input_full, grad_full=grad_full
+        )
+
+        if torch.distributed.get_rank() == 0:
+            assert torch.equal(out_tp1, out_tpN), (
+                f"[{dtype}, TP=1 vs TP={compare_tp}] forward output differs: "
+                f"max_abs={(out_tp1 - out_tpN).abs().max().item():.3e}, "
+                f"mean_abs={(out_tp1 - out_tpN).abs().mean().item():.3e}"
+            )
+            assert torch.equal(grad_tp1, grad_tpN), (
+                f"[{dtype}, TP=1 vs TP={compare_tp}] input grad differs: "
+                f"max_abs={(grad_tp1 - grad_tpN).abs().max().item():.3e}, "
+                f"mean_abs={(grad_tp1 - grad_tpN).abs().mean().item():.3e}"
+            )