pggPL · pggPL · Jul 7, 2026 · Jul 8, 2026 · Jul 9, 2026 · Jul 9, 2026
diff --git a/.clang-format b/.clang-format
@@ -261,7 +261,7 @@ SpacesInParensOptions:
   InEmptyParentheses: false
   Other:           false
 SpacesInSquareBrackets: false
-Standard:        Auto
+Standard:        c++17
 StatementAttributeLikeMacros:
   - Q_EMIT
 StatementMacros:
@@ -277,4 +277,3 @@ WhitespaceSensitiveMacros:
   - PP_STRINGIZE
   - STRINGIZE
 ...
-
diff --git a/.github/workflows/trigger-ci.yml b/.github/workflows/trigger-ci.yml
@@ -60,6 +60,9 @@ jobs:
            || github.actor == 'kainzhong'
            || github.actor == 'cspades'
            || github.actor == 'jomitchellnv'
+           || github.actor == 'fheinecke'
+           || github.actor == 'janekb04'
+           || github.actor == 'YangFei1990'
          )
     steps:
       - name: Check if comment is issued by authorized person

diff --git a/.gitmodules b/.gitmodules
@@ -4,6 +4,6 @@
 [submodule "3rdparty/cutlass"]
 	path = 3rdparty/cutlass
 	url = https://github.com/NVIDIA/cutlass.git
-[submodule "3rdparty/nccl"]
-	path = 3rdparty/nccl
-	url = https://github.com/NVIDIA/nccl.git
+[submodule "3rdparty/nccl-extensions"]
+	path = 3rdparty/nccl-extensions
+	url = https://github.com/NVIDIA/nccl-extensions.git
diff --git a/3rdparty/nccl b/3rdparty/nccl
diff --git a/3rdparty/nccl-extensions b/3rdparty/nccl-extensions
diff --git a/build_tools/VERSION.txt b/build_tools/VERSION.txt
@@ -1 +1 @@
-2.18.0.dev0
+2.19.0.dev0
diff --git a/docs/envvars.rst b/docs/envvars.rst
@@ -141,6 +141,24 @@ backend-selection overview.
    :Default: ``1``
    :Description: Enable or disable FlashAttention backend for DotProductAttention. When set to ``0``, FlashAttention will not be used.
 
+.. envvar:: NVTE_FLASH_ATTN_V2
+
+   :Type: ``int`` (0 or 1)
+   :Default: ``1``
+   :Description: Enable or disable FlashAttention 2 (the ``flash-attn`` package) for DotProductAttention, without affecting FlashAttention 3 or 4. When set to ``0``, FlashAttention 2 will not be used even if it is installed. Useful for pinning the FlashAttention version, e.g. so training-side attention runs the same kernel generation as an inference engine.
+
+.. envvar:: NVTE_FLASH_ATTN_V3
+
+   :Type: ``int`` (0 or 1)
+   :Default: ``1``
+   :Description: Enable or disable FlashAttention 3 (the ``flash-attn-3`` package) for DotProductAttention, without affecting FlashAttention 2 or 4. When set to ``0``, FlashAttention 3 will not be used even if it is installed.
+
+.. envvar:: NVTE_FLASH_ATTN_V4
+
+   :Type: ``int`` (0 or 1)
+   :Default: ``1``
+   :Description: Enable or disable FlashAttention 4 (the ``flash-attn-4`` package) for DotProductAttention, without affecting FlashAttention 2 or 3. When set to ``0``, FlashAttention 4 will not be used even if it is installed.
+
 .. envvar:: NVTE_FUSED_ATTN
 
    :Type: ``int`` (0 or 1)
@@ -274,7 +292,7 @@ Kernel Configuration
 
    :Type: ``int`` (0 or 1)
    :Default: ``0``
-   :Description: Disable NVRTC (CUDA Runtime Compilation) support. When set to ``1``, runtime kernel compilation is disabled. This can be useful in environments where NVRTC is not available or not desired.
+   :Description: Disable NVRTC (CUDA Runtime Compilation) support. When set to ``1``, runtime kernel compilation is disabled. Existing transpose operations select their static fallback automatically. Fused softmax and normalization paths require their corresponding ``NVTE_BUILD_LEGACY_STATIC_FUSED_SOFTMAX`` or ``NVTE_BUILD_LEGACY_STATIC_NORM`` CMake option to have been enabled when the library was built; otherwise no static fallback is available.
 
 .. envvar:: NVTE_USE_CUTLASS_GROUPED_GEMM
 

diff --git a/examples/jax/ep/run_test_ep.sh b/examples/jax/ep/run_test_ep.sh
@@ -32,8 +32,8 @@ export PYTHONPATH="${TE_PATH}${PYTHONPATH:+:${PYTHONPATH}}"
 COORD="${COORD:-127.0.0.1:12345}"
 TEST_TIMEOUT_S="${TEST_TIMEOUT_S:-300}"
 
-# Editable installs don't embed rpath; libtransformer_engine.so needs
-# libnccl_ep.so.0 from the TE editable location at dlopen time.
+# Editable installs don't embed rpath; the TE JAX extension needs
+# libtransformer_engine.so from the TE editable location at dlopen time.
 TE_LIB_PATH=$(pip3 show transformer-engine 2>/dev/null \
     | grep -E "Location:|Editable project location:" \
     | tail -n 1 | awk '{print $NF}')

diff --git a/examples/pytorch/ep/bench/run_ep_bench.sh b/examples/pytorch/ep/bench/run_ep_bench.sh
@@ -26,10 +26,8 @@ if [ "${NSYS}" -eq 1 ] && [ "${KINETO}" -eq 1 ]; then
 fi
 
 SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
-TE_REPO_ROOT="$(cd "${SCRIPT_DIR}/../../../.." && pwd)"
 RESULTS="${SCRIPT_DIR}/results"
 mkdir -p "${RESULTS}"
-export PYTHONPATH="${TE_REPO_ROOT}${PYTHONPATH:+:${PYTHONPATH}}"
 
 DETECTED_GPUS=$(nvidia-smi -L 2>/dev/null | wc -l)
 NUM_GPUS="${NUM_GPUS:-${DETECTED_GPUS}}"

diff --git a/examples/pytorch/ep/bench/run_nccl_ep_bench.sh b/examples/pytorch/ep/bench/run_nccl_ep_bench.sh
@@ -23,8 +23,8 @@ TE_REPO_ROOT="$(cd "${SCRIPT_DIR}/../../../.." && pwd)"
 RESULTS="${SCRIPT_DIR}/results"
 mkdir -p "${RESULTS}"
 
-BIN="${TE_REPO_ROOT}/3rdparty/nccl/build/test/nccl_ep/ep_bench"
-LIB="${TE_REPO_ROOT}/3rdparty/nccl/build/lib"
+BIN="${TE_REPO_ROOT}/3rdparty/nccl-extensions/build/test/nccl_ep/ep_bench"
+LIB="${TE_REPO_ROOT}/3rdparty/nccl-extensions/build/lib"
 [ -x "${BIN}" ] || { echo "ep_bench not built at ${BIN}" >&2; exit 2; }
 
 NUM_GPUS=$(nvidia-smi -L 2>/dev/null | wc -l)

diff --git a/examples/pytorch/ep/run_test_ep.sh b/examples/pytorch/ep/run_test_ep.sh
@@ -17,7 +17,6 @@ if [ "${NUM_GPUS}" -gt 8 ]; then NUM_GPUS=8; fi
 : ${TEST_TIMEOUT_S:=120}
 
 SCRIPT="${TE_PATH}/examples/pytorch/ep/ep_moe.py"
-export PYTHONPATH="${TE_PATH}${PYTHONPATH:+:${PYTHONPATH}}"
 
 # Stage JIT cubins on tmpfs for fast iteration.
 : ${NCCL_EP_JIT_CACHE_DIR:="${TMPDIR:-/tmp}/nccl_ep_jit_cache_$(id -u)"}

diff --git a/qa/L0_jax_distributed_unittest/test.sh b/qa/L0_jax_distributed_unittest/test.sh
@@ -41,12 +41,12 @@ wait
 TE_PATH=$TE_PATH bash $TE_PATH/examples/jax/ep/run_test_ep.sh || test_fail "run_test_ep.sh"
 wait
 
-# MoE custom_vjp distributed suite. Runs one Python process per GPU
-# via tests/jax/run_multiprocess_moe_vjp.sh (mirrors the pattern in
+# TE-EP MoE custom_vjp distributed suite. Runs one Python process per
+# GPU via tests/jax/run_te_ep_moe.sh (mirrors the pattern in
 # examples/jax/encoder/run_test_multiprocessing_encoder.sh). Requires
 # >=4 visible GPUs.
-TE_PATH=$TE_PATH bash $TE_PATH/tests/jax/run_multiprocess_moe_vjp.sh \
-    || test_fail "test_multiprocess_moe_vjp.py"
+TE_PATH=$TE_PATH bash $TE_PATH/tests/jax/run_te_ep_moe.sh \
+    || test_fail "test_te_ep_moe.py"
 # Exercise the multi-GPU tutorial in docs/examples/jax (needs >= 4 GPUs;
 # auto-skips otherwise).
 CUDA_VISIBLE_DEVICES=0,1,2,3 python3 -m pytest -c $TE_PATH/tests/jax/pytest.ini -v --junitxml=$XML_LOG_DIR/pytest_docs_examples_jax_distributed.xml -k multi_gpu $TE_PATH/docs/examples/jax/ || test_fail "docs/examples/jax (multi-GPU)"

diff --git a/qa/L0_pytorch_unittest/test.sh b/qa/L0_pytorch_unittest/test.sh
@@ -45,6 +45,7 @@ python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_gqa.xml $TE_PATH
 python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_fused_optimizer.xml $TE_PATH/tests/pytorch/test_fused_optimizer.py || test_fail "test_fused_optimizer.py"
 python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_multi_tensor.xml $TE_PATH/tests/pytorch/test_multi_tensor.py || test_fail "test_multi_tensor.py"
 python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_fusible_ops.xml $TE_PATH/tests/pytorch/test_fusible_ops.py || test_fail "test_fusible_ops.py"
+python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_distributed_weight.xml $TE_PATH/tests/pytorch/test_distributed_weight.py || test_fail "test_distributed_weight.py"
 python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_backward_override.xml $TE_PATH/tests/pytorch/test_backward_override.py || test_fail "test_backward_override.py"
 python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_permutation.xml $TE_PATH/tests/pytorch/test_permutation.py || test_fail "test_permutation.py"
 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"
@@ -66,6 +67,7 @@ python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_partial_cast.xml
 # Disable autotuning to make unittests faster. In addition, disable TF32 path to fully align with the pytorch reference implementation's precision
 NVTE_DISABLE_TRITON_AUTOTUNING=1 NVIDIA_TF32_OVERRIDE=0 python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_mhc.xml $TE_PATH/tests/pytorch/test_mhc.py || test_fail "test_mhc.py"
 PYTORCH_JIT=0 NVTE_TORCH_COMPILE=0 python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_grouped_linear.xml $TE_PATH/tests/pytorch/test_grouped_linear.py || test_fail "test_grouped_linear.py"
+PYTORCH_JIT=0 NVTE_TORCH_COMPILE=0 python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_ops_grouped_linear_distributed_weight.xml $TE_PATH/tests/pytorch/test_ops_grouped_linear_distributed_weight.py || test_fail "test_ops_grouped_linear_distributed_weight.py"
 NVTE_GROUPED_LINEAR_SINGLE_PARAM=1 NVTE_CUTEDSL_FUSED_GROUPED_MLP=1 python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_grouped_mlp.xml $TE_PATH/tests/pytorch/test_grouped_mlp.py || test_fail "test_grouped_mlp.py"
 
 if [ "$RET" -ne 0 ]; then

diff --git a/setup.py b/setup.py
@@ -25,6 +25,7 @@
     remove_dups,
     min_python_version_str,
     nccl_ep_enabled,
+    get_max_jobs_for_parallel_build,
 )
 
 frameworks = get_frameworks()
@@ -207,11 +208,11 @@ def _discover_nccl_home() -> str:
 
 
 def build_nccl_ep_submodule() -> str:
-    """Build libnccl_ep.a from the 3rdparty/nccl submodule and return NCCL_HOME."""
-    nccl_root = current_file_path / "3rdparty" / "nccl"
-    if not (nccl_root / "Makefile").exists():
+    """Build libnccl_ep.a from the 3rdparty/nccl-extensions submodule and return NCCL_HOME."""
+    nccl_root = current_file_path / "3rdparty" / "nccl-extensions"
+    if not (nccl_root / "nccl_ep" / "Makefile").exists():
         raise RuntimeError(
-            f"NCCL submodule not found at {nccl_root}. "
+            f"NCCL EP submodule not found at {nccl_root}. "
             "Run `git submodule update --init --recursive`."
         )
 
@@ -250,7 +251,7 @@ def build_nccl_ep_submodule() -> str:
         )
     gencode = " ".join(f"-gencode=arch=compute_{a},code=sm_{a}" for a in arch_list)
 
-    nproc = os.cpu_count() or 8
+    nproc = get_max_jobs_for_parallel_build()
     env = os.environ.copy()
     env["NVCC_GENCODE"] = gencode
     # NCCL EP needs the core NCCL headers + libnccl.so; write NCCL EP build
@@ -267,13 +268,13 @@ def build_nccl_ep_submodule() -> str:
                 "rebuilding libnccl_ep.a"
             )
             subprocess.check_call(
-                ["make", "-C", "contrib/nccl_ep", "clean"],
+                ["make", "-C", "nccl_ep", "clean"],
                 cwd=str(nccl_root),
                 env=env,
             )
         print(f"[NCCL EP] Building libnccl_ep.a (gencode='{gencode}')")
         subprocess.check_call(
-            ["make", "-j", str(nproc), "-C", "contrib/nccl_ep", "lib"],
+            ["make", "-j", str(nproc), "-C", "nccl_ep", "lib"],
             cwd=str(nccl_root),
             env=env,
         )

diff --git a/tests/cpp/operator/CMakeLists.txt b/tests/cpp/operator/CMakeLists.txt
@@ -27,13 +27,15 @@ add_executable(test_operator
                test_cast_transpose_dbias_dgelu.cu
                test_cast_transpose_dgeglu.cu
                test_act.cu
+               test_scaled_activation.cu
                test_normalization.cu
                test_normalization_mxfp8.cu
                test_memset.cu
                test_splits_to_offsets.cu
                test_multi_cast_transpose.cu
                test_multi_padding.cu
                test_multi_unpadding.cu
+               test_softmax.cu
                test_causal_softmax.cu
                test_swizzle.cu
                test_multi_swizzle.cu

diff --git a/tests/cpp/operator/test_cast_float8blockwise_grouped.cu b/tests/cpp/operator/test_cast_float8blockwise_grouped.cu
@@ -368,14 +368,15 @@ struct TestConfig {
   ScalingDir dir;
   std::vector<size_t> first_dims;
   size_t K;
+  bool force_pow_2_scales;
 };
 
 class GroupedFP8BlockwiseTestSuite : public ::testing::TestWithParam<TestConfig> {};
 
 TEST_P(GroupedFP8BlockwiseTestSuite, Test) {
   const TestConfig& cfg = GetParam();
   perform_test<bf16, fp8e4m3>(cfg.shape_rep, cfg.block_dim, cfg.dir, cfg.first_dims, cfg.K,
-                              /*force_pow_2_scales=*/false, /*epsilon=*/0.0f);
+                              cfg.force_pow_2_scales, /*epsilon=*/0.0f);
 }
 
 std::vector<TestConfig> make_configs() {
@@ -387,11 +388,13 @@ std::vector<TestConfig> make_configs() {
   for (auto bd : {BlockDim::ONE_D, BlockDim::TWO_D}) {
     for (auto dir : {ScalingDir::ROWWISE, ScalingDir::COLWISE, ScalingDir::BOTH}) {
       for (size_t K : Ks) {
-        for (const auto& v : uniform) {
-          configs.push_back({ShapeRep::SAME_BOTH_DIMS, bd, dir, v, K});
-        }
-        for (const auto& v : jagged) {
-          configs.push_back({ShapeRep::VARYING_FIRST_DIM, bd, dir, v, K});
+        for (bool pow2 : {false, true}) {
+          for (const auto& v : uniform) {
+            configs.push_back({ShapeRep::SAME_BOTH_DIMS, bd, dir, v, K, pow2});
+          }
+          for (const auto& v : jagged) {
+            configs.push_back({ShapeRep::VARYING_FIRST_DIM, bd, dir, v, K, pow2});
+          }
         }
       }
     }
@@ -408,6 +411,7 @@ std::string make_name(const ::testing::TestParamInfo<TestConfig>& info) {
   s += "_K" + std::to_string(c.K) + "_N" + std::to_string(c.first_dims.size());
   s += "_M";
   for (size_t m : c.first_dims) s += "_" + std::to_string(m);
+  s += (c.force_pow_2_scales ? "_POW2" : "_FP32SC");
   return s;
 }
 

diff --git a/tests/cpp/operator/test_cast_mxfp8_grouped.cu b/tests/cpp/operator/test_cast_mxfp8_grouped.cu
@@ -509,9 +509,9 @@ void performTest(const ProcessingMethod processing_method,
             break;
         }
         case ProcessingMethod::CAST_DBIAS: {
-            nvte_group_quantize_dbias(grad_group_tensor, out_group_tensor, output_dbias_tensor, workspace.data(), 0);
+            nvte_group_quantize_dbias(grad_group_tensor, out_group_tensor, output_dbias_tensor, workspace.data(), nullptr, 0);
             workspace = Tensor("workspace", workspace.rowwise_shape(), workspace.dtype());
-            nvte_group_quantize_dbias(grad_group_tensor, out_group_tensor, output_dbias_tensor, workspace.data(), 0);
+            nvte_group_quantize_dbias(grad_group_tensor, out_group_tensor, output_dbias_tensor, workspace.data(), nullptr, 0);
             break;
         }
         case ProcessingMethod::CAST_DBIAS_DACT: {