pggPL · pggPL · Jun 16, 2026 · Jun 22, 2026 · Jun 22, 2026 · Jun 29, 2026
diff --git a/.github/workflows/build.yml b/.github/workflows/build.yml
@@ -98,7 +98,7 @@ jobs:
     name: 'JAX'
     runs-on: ubuntu-latest
     container:
-      image: ghcr.io/nvidia/jax:jax
+      image: ghcr.io/nvidia/jax:jax-2026-07-21
       options: --user root
     steps:
       - name: 'Dependencies'
@@ -149,7 +149,7 @@ jobs:
 
       - name: Start named container
         run: |
-          docker run -v $(pwd):$(pwd) -w $(pwd) --name builder -d ghcr.io/nvidia/jax:jax sleep infinity
+          docker run -v $(pwd):$(pwd) -w $(pwd) --name builder -d ghcr.io/nvidia/jax:jax-2026-07-21 sleep infinity
 
       - name: 'Dependencies'
         run: |

diff --git a/.github/workflows/trigger-ci.yml b/.github/workflows/trigger-ci.yml
@@ -63,6 +63,7 @@ jobs:
            || github.actor == 'fheinecke'
            || github.actor == 'janekb04'
            || github.actor == 'YangFei1990'
+           || github.actor == 'sraman-rgb'
          )
     steps:
       - name: Check if comment is issued by authorized person

diff --git a/CONTRIBUTING.rst b/CONTRIBUTING.rst
@@ -13,7 +13,7 @@ Coding Guidelines
 -----------------
 
 * We follow `Google C++ Style Guide <https://google.github.io/styleguide/cppguide.html>`_. When no
-  rules can be found, follow the already occuring conventions. If there is no precedence in our
+  rules can be found, follow the already occurring conventions. If there is no precedence in our
   codebase we are open to discussion.
 * Prior to your contribution, please make sure that the code passes the linter check. We do both C++
   and Python linting. To invoke the check, please use

diff --git a/docs/envvars.rst b/docs/envvars.rst
@@ -379,6 +379,16 @@ Torch Compilation and Fusion
 LayerNorm/RMSNorm SM Margins
 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
 
+.. envvar:: NVTE_CUDNN_MXFP8_NORM_OUTPUT_IN_INPUT_DTYPE
+
+   :Type: ``int`` (0 or 1)
+   :Default: ``0``
+   :Description: With cuDNN 9.25.0 or later, use the normalization input datatype for the virtual
+                 LayerNorm/RMSNorm output consumed by cuDNN MXFP8 block-scale quantization. This
+                 enables cuDNN's fused MXFP8 normalization engine, which requires matching FP16 or
+                 BF16 input and normalization-output datatypes. When set to ``0``, or with an
+                 earlier cuDNN version, the virtual normalization output uses FP32.
+
 .. envvar:: NVTE_FWD_LAYERNORM_SM_MARGIN
 
    :Type: ``int``

diff --git a/examples/jax/encoder/requirements.txt b/examples/jax/encoder/requirements.txt
@@ -1,4 +1,4 @@
 datasets<4.0.0
 flax>=0.7.1
-nltk>=3.8.2
+nltk>=3.8.2,<3.10.1
 optax
diff --git a/examples/jax/ep/bench/ep_bench.py b/examples/jax/ep/bench/ep_bench.py
@@ -148,7 +148,7 @@ def main():
 
         @jax.jit
         def run_prepare(idx):
-            tc, hm = tex_ep.ep_prepare(cfg, idx)
+            tc, _trt, hm = tex_ep.ep_prepare(cfg, idx)
             return tc, hm
 
         @jax.jit
@@ -160,7 +160,7 @@ def run_dispatch(hm, idx, toks, w):
 
         @jax.jit
         def run_dispatch_vjp(idx, toks, w):
-            recv_t, recv_w, _hm, _tc = ep_dispatch(cfg, idx, toks, w, recv_capacity_per_rank)
+            recv_t, recv_w, _hm, _tc, _trt = ep_dispatch(cfg, idx, toks, w, recv_capacity_per_rank)
             recv_t = jax.lax.with_sharding_constraint(recv_t, NamedSharding(mesh, ep_spec_3d))
             recv_w = jax.lax.with_sharding_constraint(recv_w, NamedSharding(mesh, ep_spec_2d))
             return recv_t, recv_w

diff --git a/examples/jax/ep/ep_moe.py b/examples/jax/ep/ep_moe.py
@@ -231,7 +231,7 @@ def _moe_layer(args, cfg, mesh, topk_idx, tokens, topk_w, local_kernels):
     local_kernels = jax.lax.with_sharding_constraint(
         local_kernels, NamedSharding(mesh, kernel_spec)
     )
-    recv_tokens, recv_topk_w, handle_mem, _tc = ep_dispatch(
+    recv_tokens, recv_topk_w, handle_mem, _tc, _trt = ep_dispatch(
         cfg, topk_idx, tokens, topk_w, args.recv_capacity_per_rank
     )
     recv_tokens = jax.lax.with_sharding_constraint(recv_tokens, NamedSharding(mesh, ep3))

diff --git a/qa/L0_jax_unittest/test.sh b/qa/L0_jax_unittest/test.sh
@@ -21,7 +21,7 @@ FAILED_CASES=""
 
 export NVTE_JAX_TEST_TIMING=1
 
-pip3 install "nltk>=3.8.2" || error_exit "Failed to install nltk"
+pip3 install "nltk>=3.8.2,<3.10.1" || error_exit "Failed to install nltk"
 pip3 install pytest==8.2.1 || error_exit "Failed to install pytest"
 
 : ${TE_PATH:=/opt/transformerengine}

diff --git a/qa/L0_pytorch_unittest/test.sh b/qa/L0_pytorch_unittest/test.sh
@@ -54,6 +54,8 @@ python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_permutation.xml
 python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_parallel_cross_entropy.xml $TE_PATH/tests/pytorch/test_parallel_cross_entropy.py || test_fail "test_parallel_cross_entropy.py"
 python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_cpu_offloading.xml $TE_PATH/tests/pytorch/test_cpu_offloading.py || test_fail "test_cpu_offloading.py"
 NVTE_FLASH_ATTN=0 NVTE_CPU_OFFLOAD_V1=1 python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_cpu_offloading_v1.xml $TE_PATH/tests/pytorch/test_cpu_offloading_v1.py || test_fail "test_cpu_offloading_v1.py"
+python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_hybrid_quantization.xml $TE_PATH/tests/pytorch/test_hybrid_quantization.py || test_fail "test_hybrid_quantization.py"
+python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_identity_quantizer.xml $TE_PATH/tests/pytorch/test_identity_quantizer.py || test_fail "test_identity_quantizer.py"
 NVTE_ALLOW_UNSAFE_PICKLE_EXTRA_STATE=1 python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_attention.xml $TE_PATH/tests/pytorch/attention/test_attention.py || test_fail "test_attention.py"
 python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_flex_attention.xml $TE_PATH/tests/pytorch/attention/test_flex_attention.py || test_fail "test_flex_attention.py"
 NVTE_ALLOW_UNSAFE_PICKLE_EXTRA_STATE=1 NVTE_ALLOW_NONDETERMINISTIC_ALGO=0 python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_attention_deterministic.xml $TE_PATH/tests/pytorch/attention/test_attention.py || test_fail "NVTE_ALLOW_NONDETERMINISTIC_ALGO=0 test_attention.py"

diff --git a/qa/L1_jax_distributed_unittest/test.sh b/qa/L1_jax_distributed_unittest/test.sh
@@ -31,12 +31,12 @@ python3 -m pytest -c $TE_PATH/tests/jax/pytest.ini -v --junitxml=$XML_LOG_DIR/py
 
 python3 -m pytest -c $TE_PATH/tests/jax/pytest.ini -v --junitxml=$XML_LOG_DIR/pytest_dist_mlp.xml $TE_PATH/tests/jax/test_distributed_layernorm_mlp.py || test_fail "test_distributed_layernorm_mlp.py"
 
+python3 -m pytest -c $TE_PATH/tests/jax/pytest.ini -v --junitxml=$XML_LOG_DIR/pytest_dist_fused_attn.xml $TE_PATH/tests/jax/test_distributed_fused_attn.py || test_fail "test_distributed_fused_attn.py"
+
 # XLA_FLAGS to WAR for test_distributed_softmax issue with NCCL
-# TODO(Kshitij): remove when NCCL issue is fixed
+# TODO(KshitijLakhani): remove when NCCL issue is fixed
 XLA_FLAGS="$XLA_FLAGS --xla_gpu_enable_nccl_comm_splitting=false" python3 -m pytest -c $TE_PATH/tests/jax/pytest.ini -v --junitxml=$XML_LOG_DIR/pytest_dist_softmax.xml $TE_PATH/tests/jax/test_distributed_softmax.py || test_fail "test_distributed_softmax.py"
 
-python3 -m pytest -c $TE_PATH/tests/jax/pytest.ini -v --junitxml=$XML_LOG_DIR/pytest_dist_fused_attn.xml $TE_PATH/tests/jax/test_distributed_fused_attn.py || test_fail "test_distributed_fused_attn.py"
-
 # NCCL EP multi-process suite. Self-skips on <4 GPUs.
 TE_PATH=$TE_PATH bash $TE_PATH/tests/jax/multi_process_launch_ep.sh || test_fail "test_multi_process_ep.py"
 

diff --git a/qa/L1_pytorch_distributed_unittest/test.sh b/qa/L1_pytorch_distributed_unittest/test.sh
@@ -44,7 +44,7 @@ python3 -m pytest -v -s --junitxml=$XML_LOG_DIR/pytest_test_sanity.xml $TE_PATH/
 python3 -m pytest -v -s --junitxml=$XML_LOG_DIR/pytest_test_numerics.xml $TE_PATH/tests/pytorch/distributed/test_numerics.py || test_fail "test_numerics.py"
 python3 -m pytest -v -s --junitxml=$XML_LOG_DIR/pytest_test_numerics_exact.xml $TE_PATH/tests/pytorch/distributed/test_numerics_exact.py || test_fail "test_numerics_exact.py"
 python3 -m pytest -v -s --junitxml=$XML_LOG_DIR/pytest_test_fusible_ops.xml $TE_PATH/tests/pytorch/distributed/test_fusible_ops.py || test_fail "test_fusible_ops.py"
-python3 -m pytest -v -s --junitxml=$XML_LOG_DIR/pytest_test_torch_fsdp2.xml $TE_PATH/tests/pytorch/distributed/test_torch_fsdp2.py || test_fail "test_torch_fsdp2.py"
+python3 -m pytest -v -s --junitxml=$XML_LOG_DIR/pytest_test_torch_fsdp2.xml $TE_PATH/tests/pytorch/distributed/test_torch_fsdp2.py -k "not hybrid" || test_fail "test_torch_fsdp2.py"
 python3 -m pytest -v -s --junitxml=$XML_LOG_DIR/pytest_test_comm_gemm_overlap.xml $TE_PATH/tests/pytorch/distributed/test_comm_gemm_overlap.py || test_fail "test_comm_gemm_overlap.py"
 python3 -m pytest -v -s --junitxml=$XML_LOG_DIR/pytest_test_fusible_ops_with_userbuffers.xml $TE_PATH/tests/pytorch/distributed/test_fusible_ops_with_userbuffers.py || test_fail "test_fusible_ops_with_userbuffers.py"
 python3 -m pytest -v -s --junitxml=$XML_LOG_DIR/pytest_test_cp_utils.xml $TE_PATH/tests/pytorch/attention/test_cp_utils.py || test_fail "test_cp_utils.py"

diff --git a/qa/L1_pytorch_hybrid_distributed_unittest/test.sh b/qa/L1_pytorch_hybrid_distributed_unittest/test.sh
@@ -0,0 +1,33 @@
+# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+#
+# See LICENSE for license information.
+
+function error_exit() {
+    echo "Error: $1"
+    exit 1
+}
+
+function test_fail() {
+    RET=1
+    FAILED_CASES="$FAILED_CASES $1"
+    echo "Error: sub-test failed: $1"
+}
+
+RET=0
+FAILED_CASES=""
+
+: ${TE_PATH:=/opt/transformerengine}
+: ${XML_LOG_DIR:=/logs}
+mkdir -p "$XML_LOG_DIR"
+
+pip3 install pytest==8.2.1 || error_exit "Failed to install pytest"
+
+python3 -m pytest -v -s --junitxml=$XML_LOG_DIR/pytest_test_torch_fsdp2_hybrid.xml $TE_PATH/tests/pytorch/distributed/test_torch_fsdp2.py -k "hybrid" || test_fail "hybrid test_torch_fsdp2.py"
+python3 -m pytest -v -s --junitxml=$XML_LOG_DIR/pytest_test_hybrid_tp_sp.xml $TE_PATH/tests/pytorch/distributed/test_hybrid_tp_sp.py || test_fail "test_hybrid_tp_sp.py"
+
+if [ "$RET" -ne 0 ]; then
+    echo "Error in the following test cases:$FAILED_CASES"
+    exit 1
+fi
+echo "All tests passed"
+exit 0
diff --git a/qa/L2_jax_unittest/test.sh b/qa/L2_jax_unittest/test.sh
@@ -21,7 +21,7 @@ FAILED_CASES=""
 
 export NVTE_JAX_TEST_TIMING=1
 
-pip3 install "nltk>=3.8.2" || error_exit "Failed to install nltk"
+pip3 install "nltk>=3.8.2,<3.10.1" || error_exit "Failed to install nltk"
 pip3 install pytest==8.2.1 || error_exit "Failed to install pytest"
 
 : ${TE_PATH:=/opt/transformerengine}

diff --git a/tests/jax/test_distributed_fused_attn.py b/tests/jax/test_distributed_fused_attn.py
@@ -22,6 +22,7 @@
     _has_cudnn_frontend_python,
 )
 from utils import pytest_parametrize_wrapper
+from transformer_engine_jax import get_cudnn_version, get_device_compute_capability
 from transformer_engine.jax.attention import (
     is_fused_attn_kernel_available,
     AttnBiasType,
@@ -345,6 +346,59 @@ def test_softcap_score_mod_with_aux_params_backward(
     pytest.param([4, 256, 16, 64], id="4-256xCPx2-16-64"),
 ]
 
+DISTRIBUTED_CONTEXT_SELF_ATTN_D256_DATA_SHAPES = {
+    "L0": [],
+    "L1": [[2, 128, 16, 256]],
+    "L2": [],
+}
+
+# Keep these as explicit tuples instead of independent layout/mask/window as:
+# BSHD CP uses CAUSAL_MASK, THD CP uses PADDING_CAUSAL_MASK, SWA is
+# only valid for THD, and stripe_size behavior is different for
+# BSHD vs THD in these tests. Splitting the axes would mostly add
+# invalid BSHD+SWA and THD+CAUSAL combinations that fail or skip.
+DISTRIBUTED_CONTEXT_SELF_ATTN_D256_LAYOUTS_MASKS_WINDOWS = [
+    # BSHD with different layouts, but same causal mask and no sliding window
+    pytest.param(
+        QKVLayout.BSHD_BS2HD,
+        AttnMaskType.CAUSAL_MASK,
+        (-1, -1),
+        id="BSHD_KVPACKED-CAUSAL-NO_SWA",
+    ),
+    pytest.param(
+        QKVLayout.BSHD_BSHD_BSHD,
+        AttnMaskType.CAUSAL_MASK,
+        (-1, -1),
+        id="BSHD_SEPARATE-CAUSAL-NO_SWA",
+    ),
+    # THD with different sliding window sizes, but same packed layout and padding causal mask
+    pytest.param(
+        QKVLayout.THD_T2HD,
+        AttnMaskType.PADDING_CAUSAL_MASK,
+        (-1, -1),
+        id="THD_KVPACKED-PADDING_CAUSAL-NO_SWA",
+    ),
+    pytest.param(
+        QKVLayout.THD_T2HD,
+        AttnMaskType.PADDING_CAUSAL_MASK,
+        (20, 0),
+        id="THD_KVPACKED-PADDING_CAUSAL-SWA",
+    ),
+    # THD with different sliding window sizes, but same separate layout and padding causal mask
+    pytest.param(
+        QKVLayout.THD_THD_THD,
+        AttnMaskType.PADDING_CAUSAL_MASK,
+        (-1, -1),
+        id="THD_SEPARATE-PADDING_CAUSAL-NO_SWA",
+    ),
+    pytest.param(
+        QKVLayout.THD_THD_THD,
+        AttnMaskType.PADDING_CAUSAL_MASK,
+        (20, 0),
+        id="THD_SEPARATE-PADDING_CAUSAL-SWA",
+    ),
+]
+
 
 class TestDistributedContextParallelSelfAttn:
     # TODO(KshitijLakhani): parametrize num_segments_per_seq for all CP tests
@@ -638,6 +692,119 @@ def test_context_parallel_ring_attn(
             stripe_size=stripe_size,
         )
 
+    # CP ring and all-gather tests for D=256
+    # TODO(KshitijLakhani): Replace this with common-provided fused-attn disable reasons once
+    # they can be surfaced to framework tests.
+    @staticmethod
+    def skip_if_d256_cp_unsupported(qkv_layout):
+        compute_capability = get_device_compute_capability(0)
+        if not 100 <= compute_capability < 110:
+            pytest.skip("D=256 CP fused attention is only enabled on Blackwell server GPUs.")
+
+        required_cudnn_version = 92500 if qkv_layout.is_thd() else 92300
+        required_cudnn_version_label = "9.25" if qkv_layout.is_thd() else "9.23"
+        if get_cudnn_version() < required_cudnn_version:
+            pytest.skip(
+                f"D=256 CP fused attention with {qkv_layout} requires cuDNN"
+                f" {required_cudnn_version_label} or newer."
+            )
+
+    @pytest_parametrize_wrapper(
+        "device_count,mesh_shape,mesh_axes,mesh_resource",
+        generate_context_parallel_configs_for_attn(),
+    )
+    @pytest_parametrize_wrapper(
+        "data_shape",
+        DISTRIBUTED_CONTEXT_SELF_ATTN_D256_DATA_SHAPES,
+    )
+    @pytest.mark.parametrize(
+        "dtype",
+        [pytest.param(jnp.float16, id="FP16"), pytest.param(jnp.bfloat16, id="BF16")],
+    )
+    @pytest.mark.parametrize(
+        "qkv_layout, attn_mask_type, window_size",
+        DISTRIBUTED_CONTEXT_SELF_ATTN_D256_LAYOUTS_MASKS_WINDOWS,
+    )
+    def test_context_parallel_ring_attn_d256(
+        self,
+        device_count,
+        mesh_shape,
+        mesh_axes,
+        mesh_resource,
+        data_shape,
+        dtype,
+        qkv_layout,
+        attn_mask_type,
+        window_size,
+    ):
+        """D=256 CP ring coverage."""
+        self.skip_if_d256_cp_unsupported(qkv_layout)
+
+        self.impl_test_context_parallel_attn(
+            device_count,
+            mesh_shape,
+            mesh_axes,
+            mesh_resource,
+            data_shape,
+            1,
+            attn_mask_type,
+            dtype,
+            qkv_layout,
+            True,
+            CPStrategy.RING,
+            use_scan_ring=False,
+            window_size=window_size,
+            stripe_size=1 if qkv_layout.is_thd() else None,
+        )
+
+    @pytest_parametrize_wrapper(
+        "device_count,mesh_shape,mesh_axes,mesh_resource",
+        generate_context_parallel_configs_for_attn(),
+    )
+    @pytest_parametrize_wrapper(
+        "data_shape",
+        DISTRIBUTED_CONTEXT_SELF_ATTN_D256_DATA_SHAPES,
+    )
+    @pytest.mark.parametrize(
+        "dtype",
+        [pytest.param(jnp.float16, id="FP16"), pytest.param(jnp.bfloat16, id="BF16")],
+    )
+    @pytest.mark.parametrize(
+        "qkv_layout, attn_mask_type, window_size",
+        DISTRIBUTED_CONTEXT_SELF_ATTN_D256_LAYOUTS_MASKS_WINDOWS,
+    )
+    def test_context_parallel_allgather_attn_d256(
+        self,
+        device_count,
+        mesh_shape,
+        mesh_axes,
+        mesh_resource,
+        data_shape,
+        dtype,
+        qkv_layout,
+        attn_mask_type,
+        window_size,
+    ):
+        """D=256 CP all-gather coverage."""
+        self.skip_if_d256_cp_unsupported(qkv_layout)
+
+        self.impl_test_context_parallel_attn(
+            device_count,
+            mesh_shape,
+            mesh_axes,
+            mesh_resource,
+            data_shape,
+            1,
+            attn_mask_type,
+            dtype,
+            qkv_layout,
+            True,
+            CPStrategy.ALL_GATHER,
+            window_size=window_size,
+            stripe_size=128 if qkv_layout.is_thd() else None,
+            num_segments_per_seq=5 if qkv_layout.is_thd() else None,
+        )
+
 
 REORDER_CAUSAL_LOAD_BALANCING_DATA_SHAPES = {
     "L0": [[]],