diff --git a/.clang-format b/.clang-format index aec13e3762..7860641ede 100644 --- 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 index de26531a98..c6c91906cd 100644 --- 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 index 07647e915e..c350f75227 100644 --- 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 deleted file mode 160000 index b87848fbc5..0000000000 --- a/3rdparty/nccl +++ /dev/null @@ -1 +0,0 @@ -Subproject commit b87848fbc52da65b5a898b4ac6633fcf51cec4ed diff --git a/3rdparty/nccl-extensions b/3rdparty/nccl-extensions new file mode 160000 index 0000000000..2c6135a721 --- /dev/null +++ b/3rdparty/nccl-extensions @@ -0,0 +1 @@ +Subproject commit 2c6135a721824ff792af7b72900b0ab758fa1f98 diff --git a/build_tools/VERSION.txt b/build_tools/VERSION.txt index 830a65a39c..e3daa3c6a7 100644 --- 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 index aa55621477..b3765a06bd 100644 --- 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 index 1305ca6fd1..86aa6ca087 100755 --- 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 index fefecd7fa9..3b0977e4c3 100755 --- 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 index 8f6da04a00..ac5fcacc25 100755 --- 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 index 13b41f4cb2..d8e6b50556 100755 --- 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 index f86cea284e..a563e6908d 100644 --- 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 index 91d3be63d0..5d767ba4d1 100644 --- 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 index 150d92969c..45e079c601 100644 --- 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 index 832177c637..06ed56cc5e 100644 --- a/tests/cpp/operator/CMakeLists.txt +++ b/tests/cpp/operator/CMakeLists.txt @@ -27,6 +27,7 @@ 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 @@ -34,6 +35,7 @@ add_executable(test_operator 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 index bc9f104e17..90418d6287 100644 --- a/tests/cpp/operator/test_cast_float8blockwise_grouped.cu +++ b/tests/cpp/operator/test_cast_float8blockwise_grouped.cu @@ -368,6 +368,7 @@ struct TestConfig { ScalingDir dir; std::vector first_dims; size_t K; + bool force_pow_2_scales; }; class GroupedFP8BlockwiseTestSuite : public ::testing::TestWithParam {}; @@ -375,7 +376,7 @@ class GroupedFP8BlockwiseTestSuite : public ::testing::TestWithParam TEST_P(GroupedFP8BlockwiseTestSuite, Test) { const TestConfig& cfg = GetParam(); perform_test(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 make_configs() { @@ -387,11 +388,13 @@ std::vector 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& 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 index de72299be1..80b80da0d8 100644 --- 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: { diff --git a/tests/cpp/operator/test_scaled_activation.cu b/tests/cpp/operator/test_scaled_activation.cu new file mode 100644 index 0000000000..cf12ca8633 --- /dev/null +++ b/tests/cpp/operator/test_scaled_activation.cu @@ -0,0 +1,291 @@ +/************************************************************************* + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * + * See LICENSE for license information. + ************************************************************************/ + +#include +#include +#include +#include +#include + +#include +#include + +#include + +#include "../test_common.h" + +using namespace transformer_engine; + +namespace { + +enum class ScaledActivationCase { + kSwiGLU, + kClampedSwiGLU, + kSReLU, +}; + +constexpr float kClampedLimit = 0.5f; +constexpr float kClampedAlpha = 1.702f; +constexpr float kClampedLinearOffset = 0.5f; + +const char *activation_name(ScaledActivationCase activation) { + switch (activation) { + case ScaledActivationCase::kSwiGLU: + return "scaled_swiglu"; + case ScaledActivationCase::kClampedSwiGLU: + return "scaled_clamped_swiglu"; + case ScaledActivationCase::kSReLU: + return "scaled_srelu"; + } + return "unknown"; +} + +inline void glu_indices(const size_t row, const size_t col, const size_t hidden, + const int64_t interleave, size_t *act_idx, size_t *linear_idx) { + if (interleave > 0) { + const size_t block = col / static_cast(interleave); + const size_t lane = col % static_cast(interleave); + const size_t base = row * hidden * 2 + block * static_cast(interleave) * 2 + lane; + *act_idx = base; + *linear_idx = base + static_cast(interleave); + } else { + const size_t base = row * hidden * 2; + *act_idx = base + col; + *linear_idx = base + hidden + col; + } +} + +inline void gated_grads(const ScaledActivationCase activation, const float act_in, + const float linear_in, float *dact, float *dlinear, float *unscaled) { + switch (activation) { + case ScaledActivationCase::kSwiGLU: { + const float act = test::silu(act_in); + *unscaled = act * linear_in; + *dact = test::dsilu(act_in) * linear_in; + *dlinear = act; + return; + } + case ScaledActivationCase::kClampedSwiGLU: { + const bool dlinear_mask = linear_in <= kClampedLimit && linear_in >= -kClampedLimit; + const float act = test::qgelu(fminf(kClampedLimit, act_in)); + const float dact_base = + act_in <= kClampedLimit ? test::dqgelu(fminf(kClampedLimit, act_in)) : 0.0f; + const float linear = + fminf(fmaxf(-kClampedLimit, linear_in), kClampedLimit) + kClampedLinearOffset; + *unscaled = act * linear; + *dact = dact_base * linear; + *dlinear = dlinear_mask ? act : 0.0f; + return; + } + case ScaledActivationCase::kSReLU: + *unscaled = test::srelu(act_in); + *dact = test::dsrelu(act_in); + *dlinear = 0.0f; + return; + } +} + +template +void compute_reference(ScaledActivationCase activation, const DataT *input, const ScaleT *scales, + const DataT *grad_output, DataT *output, DataT *grad_input, + DataT *grad_scales, const size_t rows, const size_t hidden, + const int64_t interleave, const bool compute_grad_scales) { + const bool is_gated = activation != ScaledActivationCase::kSReLU; + const size_t input_cols = is_gated ? hidden * 2 : hidden; + std::fill(grad_input, grad_input + rows * input_cols, static_cast(0.0f)); + + for (size_t row = 0; row < rows; ++row) { + const float scale = static_cast(scales[row]); + float scale_grad = 0.0f; + for (size_t col = 0; col < hidden; ++col) { + const size_t out_idx = row * hidden + col; + float unscaled = 0.0f; + float dact = 0.0f; + float dlinear = 0.0f; + if (is_gated) { + size_t act_idx = 0; + size_t linear_idx = 0; + glu_indices(row, col, hidden, interleave, &act_idx, &linear_idx); + const float act_in = static_cast(input[act_idx]); + const float linear_in = static_cast(input[linear_idx]); + gated_grads(activation, act_in, linear_in, &dact, &dlinear, &unscaled); + + const float scaled_grad = static_cast(grad_output[out_idx]) * scale; + grad_input[act_idx] = static_cast(scaled_grad * dact); + grad_input[linear_idx] = static_cast(scaled_grad * dlinear); + } else { + const float x = static_cast(input[out_idx]); + unscaled = test::srelu(x); + const float scaled_grad = static_cast(grad_output[out_idx]) * scale; + grad_input[out_idx] = static_cast(scaled_grad * test::dsrelu(x)); + } + + output[out_idx] = static_cast(unscaled * scale); + scale_grad += static_cast(grad_output[out_idx]) * unscaled; + } + if (compute_grad_scales) { + grad_scales[row] = static_cast(scale_grad); + } + } +} + +template +void run_scaled_activation_test(ScaledActivationCase activation, const size_t rows, + const size_t hidden, const int64_t interleave, + const bool compute_grad_scales) { + using namespace test; + const DType data_type = TypeInfo::dtype; + const DType scale_type = TypeInfo::dtype; + const bool is_gated = activation != ScaledActivationCase::kSReLU; + const size_t input_cols = is_gated ? hidden * 2 : hidden; + + Tensor input("input", std::vector{rows, input_cols}, data_type); + Tensor scales("act_scales", std::vector{rows}, scale_type); + Tensor output("output", std::vector{rows, hidden}, data_type); + Tensor grad_output("grad_output", std::vector{rows, hidden}, data_type); + Tensor grad_input("grad_input", std::vector{rows, input_cols}, data_type); + Tensor grad_scales("grad_scales", std::vector{rows}, data_type); + + fillUniform(&input); + fillUniform(&scales); + fillUniform(&grad_output); + + std::unique_ptr ref_output = std::make_unique(rows * hidden); + std::unique_ptr ref_grad_input = std::make_unique(rows * input_cols); + std::unique_ptr ref_grad_scales = std::make_unique(rows); + + compute_reference(activation, input.rowwise_cpu_dptr(), scales.rowwise_cpu_dptr(), + grad_output.rowwise_cpu_dptr(), ref_output.get(), + ref_grad_input.get(), ref_grad_scales.get(), rows, hidden, interleave, + compute_grad_scales); + + switch (activation) { + case ScaledActivationCase::kSwiGLU: + nvte_scaled_swiglu(input.data(), scales.data(), output.data(), interleave, 0); + nvte_scaled_dswiglu(grad_output.data(), input.data(), scales.data(), grad_input.data(), + compute_grad_scales ? grad_scales.data() : nullptr, interleave, 0); + break; + case ScaledActivationCase::kClampedSwiGLU: + nvte_scaled_clamped_swiglu(input.data(), scales.data(), output.data(), kClampedLimit, + kClampedAlpha, kClampedLinearOffset, interleave, 0); + nvte_scaled_clamped_dswiglu( + grad_output.data(), input.data(), scales.data(), grad_input.data(), + compute_grad_scales ? grad_scales.data() : nullptr, kClampedLimit, kClampedAlpha, + kClampedLinearOffset, interleave, 0); + break; + case ScaledActivationCase::kSReLU: + nvte_scaled_srelu(input.data(), scales.data(), output.data(), 0); + nvte_scaled_dsrelu(grad_output.data(), input.data(), scales.data(), grad_input.data(), + compute_grad_scales ? grad_scales.data() : nullptr, 0); + break; + } + + NVTE_CHECK_CUDA(cudaDeviceSynchronize()); + auto err = cudaGetLastError(); + ASSERT_EQ(err, cudaSuccess) << cudaGetErrorString(err); + + auto [atol, rtol] = getTolerances(data_type); + if (data_type == DType::kFloat32) { + atol = 5e-5; + rtol = 5e-5; + } + compareResults("scaled_activation_output", output, ref_output.get(), true, atol, rtol); + compareResults("scaled_activation_grad_input", grad_input, ref_grad_input.get(), true, atol, + rtol); + if (compute_grad_scales) { + compareResults("scaled_activation_grad_scales", grad_scales, ref_grad_scales.get(), true, atol, + rtol); + } +} + +class ScaledActivationTest + : public ::testing::TestWithParam< + std::tuple, int64_t, + bool>> { +}; + +std::string test_name_generator( + const testing::TestParamInfo &info) { + const auto activation = std::get<0>(info.param); + const auto data_type = std::get<1>(info.param); + const auto scale_type = std::get<2>(info.param); + const auto shape = std::get<3>(info.param); + const auto interleave = std::get<4>(info.param); + const auto compute_grad_scales = std::get<5>(info.param); + return std::string(activation_name(activation)) + "_data_" + test::typeName(data_type) + + "_scale_" + test::typeName(scale_type) + "_m_" + std::to_string(shape.first) + "_h_" + + std::to_string(shape.second) + "_interleave_" + std::to_string(interleave) + + (compute_grad_scales ? "_with_scale_grad" : "_no_scale_grad"); +} + +} // namespace + +TEST_P(ScaledActivationTest, ForwardBackward) { + const auto activation = std::get<0>(GetParam()); + const auto data_type = std::get<1>(GetParam()); + const auto scale_type = std::get<2>(GetParam()); + const auto shape = std::get<3>(GetParam()); + const auto interleave = std::get<4>(GetParam()); + const auto compute_grad_scales = std::get<5>(GetParam()); + + if (activation == ScaledActivationCase::kSReLU && interleave != 0) { + GTEST_SKIP() << "Interleave has no meaning for SReLU."; + } + if (activation != ScaledActivationCase::kSReLU && interleave > 0 && + shape.second % static_cast(interleave) != 0) { + GTEST_SKIP() << "Hidden size must be divisible by GLU interleave."; + } + + using namespace test; + TRANSFORMER_ENGINE_TYPE_SWITCH_ALL(data_type, DataT, { + TRANSFORMER_ENGINE_TYPE_SWITCH_ALL(scale_type, ScaleT, { + run_scaled_activation_test(activation, shape.first, shape.second, interleave, + compute_grad_scales); + }); + }); +} + +// Test axes (the six tuple elements consumed by ScaledActivationTest): +// 1. Activation : SwiGLU and ClampedSwiGLU are gated (input is [M, 2H]); +// SReLU is unary (input is [M, H], no gate split). +// 2. Data dtype : dtype of the activation input/output tensors. +// 3. Scale dtype : dtype of act_scales / grad_act_scales. +// 4. Shape {rows, hidden}: rows = M (tokens), hidden = H (output width; gated input is 2H). +// 5. GLU interleave : 0 = contiguous [a | b]; 32 = interleaved a/b blocks. Only valid +// for gated activations with hidden % 32 == 0; SReLU skips != 0. +// 6. compute_grad_scales : whether the backward also reduces grad_act_scales. + +// Interleave is swept over {0, 32}; invalid combinations -- SReLU with any nonzero interleave, or +// a gated activation whose hidden is not divisible by the interleave -- are skipped at runtime by +// the GTEST_SKIP guards in the test body. +INSTANTIATE_TEST_SUITE_P( + OperatorTest_ScaledActivation, ScaledActivationTest, + ::testing::Combine( + ::testing::Values(ScaledActivationCase::kSwiGLU, ScaledActivationCase::kClampedSwiGLU, + ScaledActivationCase::kSReLU), + ::testing::Values(DType::kFloat32, DType::kBFloat16), // data dtype + ::testing::Values(DType::kFloat32, DType::kBFloat16), // scale dtype + ::testing::Values(std::pair{17, 64}, // aligned + interleaved + std::pair{13, 100}, // scalar fallback + std::pair{1024, 2048}), // large FFN-ish width + ::testing::Values(0, 32), // contiguous + interleaved + ::testing::Values(false, true)), // grad_act_scales off / on + test_name_generator); + +// Keep FP16 coverage focused on representative aligned and scalar-fallback shapes instead of +// multiplying it across the full shape matrix above. +INSTANTIATE_TEST_SUITE_P( + OperatorTest_ScaledActivation_FP16, ScaledActivationTest, + ::testing::Combine( + ::testing::Values(ScaledActivationCase::kSwiGLU, ScaledActivationCase::kClampedSwiGLU, + ScaledActivationCase::kSReLU), + ::testing::Values(DType::kFloat16), // data dtype + ::testing::Values(DType::kFloat32, DType::kFloat16), // scale dtype + ::testing::Values(std::pair{17, 64}, // aligned/interleaved + std::pair{13, 100}), // scalar fallback + ::testing::Values(0, 32), // contiguous + interleaved + ::testing::Values(false, true)), // grad_act_scales off / on + test_name_generator); diff --git a/tests/cpp/operator/test_softmax.cu b/tests/cpp/operator/test_softmax.cu new file mode 100644 index 0000000000..b535c5c6cd --- /dev/null +++ b/tests/cpp/operator/test_softmax.cu @@ -0,0 +1,243 @@ +/************************************************************************* + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * + * See LICENSE for license information. + ************************************************************************/ + +#include +#include +#include + +#include +#include +#include +#include +#include +#include + +#include +#include "../test_common.h" + +using namespace transformer_engine; + +namespace { + +template +void ref_softmax_row(Type *out, const Type *in, const uint8_t *mask, int cols, float scale) { + float max_value = -10000.0f; + bool has_unmasked = false; + for (int j = 0; j < cols; ++j) { + if (mask != nullptr && mask[j] == 1) continue; + max_value = std::max(max_value, static_cast(in[j]) * scale); + has_unmasked = true; + } + float sum = 0.0f; + for (int j = 0; j < cols; ++j) { + if (mask != nullptr && mask[j] == 1) { + out[j] = static_cast(0.0f); + continue; + } + const float val = has_unmasked ? std::exp(static_cast(in[j]) * scale - max_value) : 0.0f; + sum += val; + out[j] = static_cast(val); + } + for (int j = 0; j < cols; ++j) { + out[j] = static_cast(static_cast(out[j]) / sum); + } +} + +template +void ref_softmax_bwd(Type *grad_in, const Type *grad, const Type *softmax, int cols, float scale) { + float sum = 0.0f; + for (int j = 0; j < cols; ++j) { + sum += static_cast(grad[j]) * static_cast(softmax[j]); + } + for (int j = 0; j < cols; ++j) { + grad_in[j] = + static_cast(scale * (static_cast(grad[j]) - sum) * + static_cast(softmax[j])); + } +} + +template +void ref_upper_row(Type *out, const Type *in, int row, int cols, float scale) { + float max_value = -10000.0f; + for (int j = 0; j <= row; ++j) { + max_value = std::max(max_value, static_cast(in[j]) * scale); + } + float sum = 0.0f; + for (int j = 0; j < cols; ++j) { + if (j <= row) { + const float val = std::exp(static_cast(in[j]) * scale - max_value); + sum += val; + out[j] = static_cast(val); + } else { + out[j] = static_cast(0.0f); + } + } + for (int j = 0; j <= row; ++j) { + out[j] = static_cast(static_cast(out[j]) / sum); + } +} + +template +void test_scaled_softmax(DType dtype, size_t cols) { + using namespace test; + constexpr size_t batches = 2; + constexpr size_t heads = 2; + constexpr size_t rows = 8; + constexpr float scale = 0.7f; + const size_t elements_total = batches * heads * rows * cols; + Tensor input("input", std::vector{batches, heads, rows, cols}, dtype); + Tensor softmax("softmax", std::vector{batches, heads, rows, cols}, dtype); + Tensor grad("grad", std::vector{batches, heads, rows, cols}, dtype); + Tensor grad_out("grad_out", std::vector{batches, heads, rows, cols}, dtype); + fillUniform(&input); + fillUniform(&grad); + nvte_scaled_softmax_forward(input.data(), softmax.data(), scale, 0); + nvte_scaled_softmax_backward(grad.data(), softmax.data(), grad_out.data(), scale, 0); + cudaDeviceSynchronize(); + ASSERT_EQ(cudaGetLastError(), cudaSuccess); + std::unique_ptr ref = std::make_unique(elements_total); + std::unique_ptr ref_grad = std::make_unique(elements_total); + const Type *input_cpu = input.rowwise_cpu_dptr(); + const Type *grad_cpu = grad.rowwise_cpu_dptr(); + for (size_t row = 0; row < elements_total / cols; ++row) { + ref_softmax_row(ref.get() + row * cols, input_cpu + row * cols, nullptr, cols, scale); + ref_softmax_bwd(ref_grad.get() + row * cols, grad_cpu + row * cols, ref.get() + row * cols, + cols, scale); + } + auto [atol, rtol] = getTolerances(dtype); + // Fused fp16/bf16 softmax backward differs from the fp32 reference by a few ULP. + atol = (dtype == DType::kBFloat16) ? 2e-3 : 1e-4; + compareResults("scaled_softmax_fwd", softmax, ref.get(), true, atol, rtol); + compareResults("scaled_softmax_bwd", grad_out, ref_grad.get(), true, atol, rtol); +} + +template +void test_masked_softmax(DType dtype, size_t cols) { + using namespace test; + constexpr size_t batches = 2; + constexpr size_t heads = 2; + constexpr size_t rows = 8; + constexpr float scale = -0.3f; + const size_t elements_total = batches * heads * rows * cols; + Tensor input("input", std::vector{batches, heads, rows, cols}, dtype); + Tensor mask("mask", std::vector{1, 1, rows, cols}, DType::kByte); + Tensor softmax("softmax", std::vector{batches, heads, rows, cols}, dtype); + Tensor grad("grad", std::vector{batches, heads, rows, cols}, dtype); + Tensor grad_out("grad_out", std::vector{batches, heads, rows, cols}, dtype); + fillUniform(&input); + fillUniform(&grad); + uint8_t *mask_cpu = mask.rowwise_cpu_dptr(); + for (size_t i = 0; i < rows * cols; ++i) { + mask_cpu[i] = (i % 7 == 0) ? 1 : 0; + } + mask.from_cpu(); + nvte_scaled_masked_softmax_forward(input.data(), mask.data(), softmax.data(), scale, 0); + nvte_scaled_masked_softmax_backward(grad.data(), softmax.data(), grad_out.data(), scale, 0); + cudaDeviceSynchronize(); + ASSERT_EQ(cudaGetLastError(), cudaSuccess); + std::unique_ptr ref = std::make_unique(elements_total); + std::unique_ptr ref_grad = std::make_unique(elements_total); + const Type *input_cpu = input.rowwise_cpu_dptr(); + const Type *grad_cpu = grad.rowwise_cpu_dptr(); + for (int row = 0; row < batches * heads * rows; ++row) { + const int mask_row = row % rows; + ref_softmax_row(ref.get() + row * cols, input_cpu + row * cols, mask_cpu + mask_row * cols, + cols, scale); + ref_softmax_bwd(ref_grad.get() + row * cols, grad_cpu + row * cols, ref.get() + row * cols, + cols, scale); + } + auto [atol, rtol] = getTolerances(dtype); + // Fused fp16/bf16 softmax backward differs from the fp32 reference by a few ULP. + atol = (dtype == DType::kBFloat16) ? 2e-3 : 1e-4; + compareResults("masked_softmax_fwd", softmax, ref.get(), true, atol, rtol); + compareResults("masked_softmax_bwd", grad_out, ref_grad.get(), true, atol, rtol); +} + +template +void test_upper_softmax(DType dtype, size_t seq) { + using namespace test; + // attn_batches must be a multiple of the kernel's batches_per_block (up to 8). + constexpr size_t attn_batches = 8; + constexpr float scale = 1.2f; + const size_t elements_total = attn_batches * seq * seq; + Tensor input("input", std::vector{attn_batches, seq, seq}, dtype); + Tensor softmax("softmax", std::vector{attn_batches, seq, seq}, dtype); + Tensor grad("grad", std::vector{attn_batches, seq, seq}, dtype); + Tensor grad_out("grad_out", std::vector{attn_batches, seq, seq}, dtype); + fillUniform(&input); + fillUniform(&grad); + nvte_scaled_upper_triang_masked_softmax_forward(input.data(), softmax.data(), scale, 0); + nvte_scaled_upper_triang_masked_softmax_backward(grad.data(), softmax.data(), grad_out.data(), + scale, 0); + cudaDeviceSynchronize(); + ASSERT_EQ(cudaGetLastError(), cudaSuccess); + std::unique_ptr ref = std::make_unique(elements_total); + std::unique_ptr ref_grad = std::make_unique(elements_total); + const Type *input_cpu = input.rowwise_cpu_dptr(); + const Type *grad_cpu = grad.rowwise_cpu_dptr(); + for (int batch = 0; batch < attn_batches; ++batch) { + for (int row = 0; row < seq; ++row) { + const size_t offset = (batch * seq + row) * seq; + ref_upper_row(ref.get() + offset, input_cpu + offset, row, seq, scale); + ref_softmax_bwd(ref_grad.get() + offset, grad_cpu + offset, ref.get() + offset, seq, scale); + for (int col = row + 1; col < seq; ++col) { + ref_grad[offset + col] = static_cast(0.0f); + } + } + } + auto [atol, rtol] = getTolerances(dtype); + // Upper-triangular backward diverges more near the causal diagonal; looser atol. + atol = (dtype == DType::kBFloat16) ? 1e-2 : 2e-3; + compareResults("upper_softmax_fwd", softmax, ref.get(), true, atol, rtol); + compareResults("upper_softmax_bwd", grad_out, ref_grad.get(), true, atol, rtol); +} + +} // namespace + +// Dispatch a 16-bit float dtype to a templated test body. Mirrors +// TRANSFORMER_ENGINE_TYPE_SWITCH_16BIT but uses the test harness's own fp16/bf16 +// aliases so we don't have to include common.h here -- doing so would make the +// test's Tensor type ambiguous with transformer_engine::Tensor. +#define SOFTMAX_TEST_DISPATCH_16BIT(dtype, fn, cols) \ + switch (dtype) { \ + case DType::kFloat16: \ + fn(dtype, cols); \ + break; \ + case DType::kBFloat16: \ + fn(dtype, cols); \ + break; \ + default: \ + GTEST_FAIL() << "Unsupported 16-bit dtype for test"; \ + } + +using SoftmaxTestParams = std::tuple; + +class SoftmaxApiTestSuite : public ::testing::TestWithParam {}; + +TEST_P(SoftmaxApiTestSuite, ScaledSoftmax) { + const auto [dtype, cols] = GetParam(); + SOFTMAX_TEST_DISPATCH_16BIT(dtype, test_scaled_softmax, cols); +} + +TEST_P(SoftmaxApiTestSuite, MaskedSoftmax) { + const auto [dtype, cols] = GetParam(); + SOFTMAX_TEST_DISPATCH_16BIT(dtype, test_masked_softmax, cols); +} + +TEST_P(SoftmaxApiTestSuite, UpperTriangularSoftmax) { + const auto [dtype, cols] = GetParam(); + SOFTMAX_TEST_DISPATCH_16BIT(dtype, test_upper_softmax, cols); +} + +INSTANTIATE_TEST_SUITE_P(OperatorTest, SoftmaxApiTestSuite, + ::testing::Combine( + ::testing::Values(DType::kFloat16, DType::kBFloat16), + ::testing::Values(32, 112, 1024)), + [](const testing::TestParamInfo &info) { + const DType dtype = std::get<0>(info.param); + const int cols = std::get<1>(info.param); + return test::typeName(dtype) + "_Cols" + std::to_string(cols); + }); diff --git a/tests/cpp/util/CMakeLists.txt b/tests/cpp/util/CMakeLists.txt index 6d70b7b84f..1dfd2fed4e 100644 --- a/tests/cpp/util/CMakeLists.txt +++ b/tests/cpp/util/CMakeLists.txt @@ -9,7 +9,15 @@ add_executable(test_util find_package(OpenMP REQUIRED) -target_link_libraries(test_util PUBLIC CUDA::cudart GTest::gtest_main ${TE_LIB} CUDA::nvrtc CUDNN::cudnn OpenMP::OpenMP_CXX) +find_package(Threads REQUIRED) +target_link_libraries(test_util PUBLIC + CUDA::cudart + GTest::gtest_main + ${TE_LIB} + CUDA::nvrtc + CUDNN::cudnn + OpenMP::OpenMP_CXX + Threads::Threads) target_compile_options(test_util PRIVATE -O2 -fopenmp) include(GoogleTest) diff --git a/tests/cpp/util/test_nvrtc.cpp b/tests/cpp/util/test_nvrtc.cpp index d41084449e..5228c60972 100644 --- a/tests/cpp/util/test_nvrtc.cpp +++ b/tests/cpp/util/test_nvrtc.cpp @@ -4,11 +4,14 @@ * See LICENSE for license information. ************************************************************************/ +#include + +#include +#include #include +#include #include -#include - #include "util/rtc.h" using namespace transformer_engine; @@ -19,10 +22,10 @@ TEST(UtilTest, NVRTC) { } // GPU data buffer - int *device_buffer; + int* device_buffer; std::vector host_buffer(2); - cudaMalloc((void**)&device_buffer, 2*sizeof(int)); // NOLINT(*) - cudaMemset(device_buffer, 0, 2*sizeof(int)); + cudaMalloc((void**)&device_buffer, 2 * sizeof(int)); // NOLINT(*) + cudaMemset(device_buffer, 0, 2 * sizeof(int)); // CUDA kernel implementations const char code1[] = R"code( @@ -38,45 +41,151 @@ __global__ void my_kernel(uint32_t *data) { data[0] = 789; data[1] = 12; } +)code"; + const char code3[] = R"code( +#ifndef NVTE_GTEST_RTC_VALUE +#error "NVTE_GTEST_RTC_VALUE must be provided" +#endif +__global__ void my_kernel(int *data) { + data[0] = NVTE_GTEST_RTC_VALUE; + data[1] = 34; +} +)code"; + const char header4[] = R"code( +#define NVTE_GTEST_RTC_HEADER_VALUE 78 +)code"; + const char code4[] = R"code( +#include "test_nvrtc_header.h" +__global__ void my_kernel(int *data) { + data[0] = NVTE_GTEST_RTC_HEADER_VALUE; + data[1] = 90; +} )code"; // Make sure kernels are not available auto& nvrtc_manager = rtc::KernelManager::instance(); EXPECT_FALSE(nvrtc_manager.is_compiled("my gtest kernel1")); EXPECT_FALSE(nvrtc_manager.is_compiled("my gtest kernel2")); - EXPECT_THROW(nvrtc_manager.launch("my gtest kernel1", 1, 1, 0, 0, - device_buffer), + EXPECT_THROW(nvrtc_manager.launch("my gtest kernel1", 1, 1, 0, 0, device_buffer), std::runtime_error); - EXPECT_THROW(nvrtc_manager.launch("my gtest kernel2", 1, 1, 0, 0, - device_buffer), + EXPECT_THROW(nvrtc_manager.launch("my gtest kernel2", 1, 1, 0, 0, device_buffer), std::runtime_error); // Compile and run first kernel - EXPECT_NO_THROW(nvrtc_manager.compile("my gtest kernel1", - "my_kernel", - code1, - "test_nvrtc_kernel1.cu")); + EXPECT_NO_THROW( + nvrtc_manager.compile("my gtest kernel1", "my_kernel", code1, "test_nvrtc_kernel1.cu")); EXPECT_TRUE(nvrtc_manager.is_compiled("my gtest kernel1")); EXPECT_FALSE(nvrtc_manager.is_compiled("my gtest kernel2")); - EXPECT_NO_THROW(nvrtc_manager.launch("my gtest kernel1", 1, 1, 0, 0, - device_buffer)); - EXPECT_EQ(cudaMemcpy(host_buffer.data(), device_buffer, 2*sizeof(int), - cudaMemcpyDeviceToHost), + EXPECT_NO_THROW(nvrtc_manager.launch("my gtest kernel1", 1, 1, 0, 0, device_buffer)); + EXPECT_EQ(cudaMemcpy(host_buffer.data(), device_buffer, 2 * sizeof(int), cudaMemcpyDeviceToHost), cudaSuccess); EXPECT_EQ(host_buffer[0], 123); EXPECT_EQ(host_buffer[1], -456); // Compile and run second kernel - EXPECT_NO_THROW(nvrtc_manager.compile("my gtest kernel2", - "my_kernel", - code2, - "test_nvrtc_kernel2.cu")); + EXPECT_NO_THROW( + nvrtc_manager.compile("my gtest kernel2", "my_kernel", code2, "test_nvrtc_kernel2.cu")); EXPECT_TRUE(nvrtc_manager.is_compiled("my gtest kernel1")); EXPECT_TRUE(nvrtc_manager.is_compiled("my gtest kernel2")); EXPECT_NO_THROW(nvrtc_manager.launch("my gtest kernel2", 1, 1, 0, 0, device_buffer)); - EXPECT_EQ(cudaMemcpy(host_buffer.data(), device_buffer, 2*sizeof(int), - cudaMemcpyDeviceToHost), + EXPECT_EQ(cudaMemcpy(host_buffer.data(), device_buffer, 2 * sizeof(int), cudaMemcpyDeviceToHost), cudaSuccess); EXPECT_EQ(host_buffer[0], 789); EXPECT_EQ(host_buffer[1], 12); + + // Compile and run kernel with extra compile options + EXPECT_NO_THROW(nvrtc_manager.compile("my gtest kernel3", "my_kernel", code3, + "test_nvrtc_kernel3.cu", {"-DNVTE_GTEST_RTC_VALUE=56"})); + EXPECT_TRUE(nvrtc_manager.is_compiled("my gtest kernel3")); + EXPECT_NO_THROW(nvrtc_manager.launch("my gtest kernel3", 1, 1, 0, 0, device_buffer)); + EXPECT_EQ(cudaMemcpy(host_buffer.data(), device_buffer, 2 * sizeof(int), cudaMemcpyDeviceToHost), + cudaSuccess); + EXPECT_EQ(host_buffer[0], 56); + EXPECT_EQ(host_buffer[1], 34); + + // Compile and run kernel with an extra in-memory header + EXPECT_NO_THROW(nvrtc_manager.compile("my gtest kernel4", "my_kernel", code4, + "test_nvrtc_kernel4.cu", {}, + {{header4, "test_nvrtc_header.h"}})); + EXPECT_TRUE(nvrtc_manager.is_compiled("my gtest kernel4")); + EXPECT_NO_THROW(nvrtc_manager.launch("my gtest kernel4", 1, 1, 0, 0, device_buffer)); + EXPECT_EQ(cudaMemcpy(host_buffer.data(), device_buffer, 2 * sizeof(int), cudaMemcpyDeviceToHost), + cudaSuccess); + EXPECT_EQ(host_buffer[0], 78); + EXPECT_EQ(host_buffer[1], 90); + + EXPECT_EQ(cudaFree(device_buffer), cudaSuccess); +} + +TEST(UtilTest, NVRTCConcurrentCompile) { + if (!rtc::is_enabled()) { + GTEST_SKIP() << "NVRTC not enabled, skipping tests"; + } + + constexpr int num_threads = 8; + constexpr char kernel_label[] = "my concurrent gtest kernel"; + const char code[] = R"code( +__global__ void my_kernel(int *data) { + data[0] = 314; + data[1] = 159; +} +)code"; + + int device_id = 0; + ASSERT_EQ(cudaGetDevice(&device_id), cudaSuccess); + + auto& nvrtc_manager = rtc::KernelManager::instance(); + ASSERT_FALSE(nvrtc_manager.is_compiled(kernel_label)); + + std::atomic ready{0}; + std::atomic start{false}; + std::vector errors(num_threads); + std::vector threads; + threads.reserve(num_threads); + for (int thread_id = 0; thread_id < num_threads; ++thread_id) { + threads.emplace_back([&, thread_id] { + ready.fetch_add(1, std::memory_order_release); + while (!start.load(std::memory_order_acquire)) { + std::this_thread::yield(); + } + try { + const cudaError_t status = cudaSetDevice(device_id); + if (status != cudaSuccess) { + throw std::runtime_error(cudaGetErrorString(status)); + } + (void)nvrtc_manager.is_compiled(kernel_label); + nvrtc_manager.compile(kernel_label, "my_kernel", code, "test_nvrtc_concurrent_kernel.cu"); + } catch (...) { + errors[thread_id] = std::current_exception(); + } + }); + } + + while (ready.load(std::memory_order_acquire) != num_threads) { + std::this_thread::yield(); + } + start.store(true, std::memory_order_release); + for (auto& thread : threads) { + thread.join(); + } + for (const auto& error : errors) { + if (error != nullptr) { + try { + std::rethrow_exception(error); + } catch (const std::exception& e) { + ADD_FAILURE() << e.what(); + } + } + } + + ASSERT_TRUE(nvrtc_manager.is_compiled(kernel_label)); + int* device_buffer = nullptr; + ASSERT_EQ(cudaMalloc(reinterpret_cast(&device_buffer), 2 * sizeof(int)), cudaSuccess); + ASSERT_NO_THROW(nvrtc_manager.launch(kernel_label, 1, 1, 0, 0, device_buffer)); + std::vector host_buffer(2); + ASSERT_EQ(cudaMemcpy(host_buffer.data(), device_buffer, 2 * sizeof(int), cudaMemcpyDeviceToHost), + cudaSuccess); + EXPECT_EQ(host_buffer[0], 314); + EXPECT_EQ(host_buffer[1], 159); + EXPECT_EQ(cudaFree(device_buffer), cudaSuccess); } diff --git a/tests/cpp_distributed/test_ep.cu b/tests/cpp_distributed/test_ep.cu index 7dbbcdce9d..47732196ed 100644 --- a/tests/cpp_distributed/test_ep.cu +++ b/tests/cpp_distributed/test_ep.cu @@ -675,7 +675,7 @@ class EPPipelineTest : public EpOpTestBase, public ::testing::WithParamInterface TEST_P(EPPipelineTest, FullForwardBackward) { const DType dtype = GetParam(); // NCCL EP backend currently asserts ncclBfloat16 in ncclEpDispatch - // (contrib/nccl_ep/nccl_ep.cc); skip FP16/FP32 until the backend supports them. + // (nccl_ep/nccl_ep.cc); skip FP16/FP32 until the backend supports them. if (dtype != DType::kBFloat16) { GTEST_SKIP() << test::typeName(dtype) << " not yet supported by NCCL EP backend"; } @@ -750,8 +750,9 @@ class EPZeroCopyTest : public EpOpTestBase { }; TYPED_TEST_SUITE(EPZeroCopyTest, EPBf16Only); -// Identity round-trip with symm-mem on dispatch i/o + combine input. Bit-exact -// vs HBM reference (same routing, same input). +// Identity round-trip with symm-mem on dispatch i/o + combine input. The combined +// result is bit-exact vs the HBM reference; the intermediate recv buffer is not, +// since zero-copy and HBM dispatch use different per-expert layouts. TYPED_TEST(EPZeroCopyTest, IdentityAllSymm) { using Tok = TypeParam; EP_PULL_FIXTURE(); @@ -776,10 +777,7 @@ TYPED_TEST(EPZeroCopyTest, IdentityAllSymm) { ref_t.result.data(), stream)); NVTE_CHECK_CUDA(cudaStreamSynchronize(stream)); - std::vector ref_recv(ref_buf.recv_capacity * hidden_dim_); std::vector ref_result(num_tokens_ * hidden_dim_); - NVTE_CHECK_CUDA(cudaMemcpy(ref_recv.data(), ref_buf.recv_tokens.get(), - ref_recv.size() * sizeof(Tok), cudaMemcpyDeviceToHost)); NVTE_CHECK_CUDA(cudaMemcpy(ref_result.data(), ref_buf.result.get(), ref_result.size() * sizeof(Tok), cudaMemcpyDeviceToHost)); @@ -818,24 +816,17 @@ TYPED_TEST(EPZeroCopyTest, IdentityAllSymm) { symm_window(sym_recv), sym_t.result.data(), stream)); NVTE_CHECK_CUDA(cudaStreamSynchronize(stream)); - std::vector sym_recv_host(sym_buf.recv_capacity * hidden_dim_); std::vector sym_result(num_tokens_ * hidden_dim_); - NVTE_CHECK_CUDA(cudaMemcpy(sym_recv_host.data(), sym_recv.ptr, - sym_recv_host.size() * sizeof(Tok), cudaMemcpyDeviceToHost)); NVTE_CHECK_CUDA(cudaMemcpy(sym_result.data(), sym_buf.result.get(), sym_result.size() * sizeof(Tok), cudaMemcpyDeviceToHost)); - // Compare per filled recv slot (HBM ref vs symm) and full result. - int total_recv = this->template read_total_recv(sym_buf); - for (int i = 0; i < total_recv * hidden_dim_; ++i) - ASSERT_EQ(tok_to_float(sym_recv_host[i]), tok_to_float(ref_recv[i])) - << "recv mismatch at " << i; + // Combined result is the cross-mode invariant (see note above). for (size_t i = 0; i < sym_result.size(); ++i) ASSERT_EQ(tok_to_float(sym_result[i]), tok_to_float(ref_result[i])) << "result mismatch at " << i; if (g_process_id == 0) - printf(" IdentityAllSymm: passed (recv_slots=%d, bit-exact vs HBM)\n", total_recv); + printf(" IdentityAllSymm: passed (result bit-exact vs HBM)\n"); NVTE_CHECK_CUDA(cudaStreamDestroy(stream)); } diff --git a/tests/jax/conftest.py b/tests/jax/conftest.py index 74cb91202c..d729bfd1c7 100644 --- a/tests/jax/conftest.py +++ b/tests/jax/conftest.py @@ -90,8 +90,8 @@ def pytest_addoption(parser): """CLI options used by multiprocess JAX tests. ``--num-process`` and ``--process-id`` let a multiprocess launcher - (see ``tests/jax/run_multiprocess_moe_vjp.sh``) fork one pytest - process per GPU and tell each child its rank, so the test module + (see ``tests/jax/run_te_ep_moe.sh``) fork one pytest process per + GPU and tell each child its rank, so the test module can call ``jax.distributed.initialize(...)`` with the right ``local_device_ids``. Both default to 0; non-multiprocess tests ignore them. diff --git a/tests/jax/multi_process_launch_ep.sh b/tests/jax/multi_process_launch_ep.sh index ff89f712eb..8547d77f2b 100755 --- a/tests/jax/multi_process_launch_ep.sh +++ b/tests/jax/multi_process_launch_ep.sh @@ -17,8 +17,8 @@ SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" TE_REPO_ROOT="$(cd "${SCRIPT_DIR}/../.." && pwd)" export PYTHONPATH="${TE_REPO_ROOT}${PYTHONPATH:+:${PYTHONPATH}}" -# 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/tests/jax/run_multiprocess_moe_vjp.sh b/tests/jax/run_te_ep_moe.sh similarity index 61% rename from tests/jax/run_multiprocess_moe_vjp.sh rename to tests/jax/run_te_ep_moe.sh index 8dc1d2eb04..32d5f21956 100755 --- a/tests/jax/run_multiprocess_moe_vjp.sh +++ b/tests/jax/run_te_ep_moe.sh @@ -3,46 +3,43 @@ # # See LICENSE for license information. # -# Multiprocess (one-GPU-per-process) launcher for the unified MoE VJP +# Multiprocess (one-GPU-per-process) launcher for the TE-EP MoE custom_vjp # test suite. Forks one pytest invocation per visible GPU, passing each -# its own --num-process=N --process-id=i, and waits for all of them. -# Each child calls jax.distributed.initialize(..., local_device_ids= -# process_id) so each Python process only sees its one GPU as a local -# device and the participating processes form a global mesh. +# its own --num-process=N --process-id=i, and waits for all of them. Each +# child calls jax.distributed.initialize(..., local_device_ids=process_id) +# so each Python process only sees its one GPU as a local device and the +# participating processes form a global (ep, fsdp) mesh. set -euo pipefail SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" TE_ROOT="$(cd "$SCRIPT_DIR/../.." && pwd)" -TEST_FILE="$TE_ROOT/tests/jax/test_multiprocess_moe_vjp.py" +TEST_FILE="$TE_ROOT/tests/jax/test_te_ep_moe.py" PYTEST_INI="$TE_ROOT/tests/jax/pytest.ini" NUM_GPUS="${NUM_GPUS:-$(nvidia-smi -L | wc -l)}" if [ "$NUM_GPUS" -lt 4 ]; then - echo "[run_multiprocess_moe_vjp.sh] need >=4 GPUs (got $NUM_GPUS); aborting" >&2 + echo "[run_te_ep_moe.sh] need >=4 GPUs (got $NUM_GPUS); aborting" >&2 exit 1 fi export XLA_PYTHON_CLIENT_PREALLOCATE="${XLA_PYTHON_CLIENT_PREALLOCATE:-false}" export XLA_PYTHON_CLIENT_MEM_FRACTION="${XLA_PYTHON_CLIENT_MEM_FRACTION:-0.5}" -export MOE_VJP_COORDINATOR_ADDRESS="${MOE_VJP_COORDINATOR_ADDRESS:-127.0.0.1:13456}" +export TE_EP_MOE_COORDINATOR_ADDRESS="${TE_EP_MOE_COORDINATOR_ADDRESS:-127.0.0.1:13457}" echo "============================================================" -echo "MoE VJP MULTIPROCESS test (one process per GPU, ${NUM_GPUS} GPUs)" +echo "TE-EP MoE MULTIPROCESS test (one process per GPU, ${NUM_GPUS} GPUs)" echo " test file : $TEST_FILE" -echo " coordinator : $MOE_VJP_COORDINATOR_ADDRESS" +echo " coordinator : $TE_EP_MOE_COORDINATOR_ADDRESS" echo " XLA_PYTHON_CLIENT_PREALLOCATE: $XLA_PYTHON_CLIENT_PREALLOCATE" echo " XLA_PYTHON_CLIENT_MEM_FRACTION: $XLA_PYTHON_CLIENT_MEM_FRACTION" echo "============================================================" -# Per-process logs. MOE_VJP_MP_LOG_DIR can be set to a host-mounted dir -# (e.g. when running inside a container that throws away /tmp on exit) -# so logs survive for postmortem inspection. Defaults to a fresh /tmp. -if [ -n "${MOE_VJP_MP_LOG_DIR:-}" ]; then - LOG_DIR="$MOE_VJP_MP_LOG_DIR" +if [ -n "${TE_EP_MOE_MP_LOG_DIR:-}" ]; then + LOG_DIR="$TE_EP_MOE_MP_LOG_DIR" mkdir -p "$LOG_DIR" else - LOG_DIR=$(mktemp -d -t moe_vjp_mp_XXXXXX) + LOG_DIR=$(mktemp -d -t te_ep_moe_mp_XXXXXX) fi echo "Per-process logs: $LOG_DIR" @@ -63,8 +60,6 @@ cleanup() { } trap cleanup EXIT INT TERM -# Launch one pytest per GPU. Process 0 streams to stdout; others log -# only to file so the live output isn't a mosaic. for i in $(seq 0 $((NUM_GPUS - 1))); do LOG_FILE="$LOG_DIR/proc_${i}.log" PYTEST_CMD=( @@ -84,7 +79,6 @@ for i in $(seq 0 $((NUM_GPUS - 1))); do PIDS+=("$!") done -# Wait for all and collect exit codes. EXITS=() for pid in "${PIDS[@]}"; do if wait "$pid"; then @@ -94,7 +88,6 @@ for pid in "${PIDS[@]}"; do fi done -# Summary. echo echo "============================================================" echo "Per-process exit codes:" @@ -102,12 +95,9 @@ for i in "${!EXITS[@]}"; do echo " proc $i -> ${EXITS[$i]}" done -# Final pass/fail. Any non-zero in any process fails the suite, but -# we tolerate non-zero on the non-zero processes only if proc 0 -# reports PASS (this matches the encoder launcher's logic). Simplest -# Treat exit 0 (pass) and exit 5 (pytest "no tests collected", which -# the file emits via ``pytest.skip(allow_module_level=True)`` on -# pre-Blackwell GPUs) as success. Anything else is a failure. +# Treat exit 0 (pass) and exit 5 (pytest "no tests collected", which the +# file emits via pytest.skip(allow_module_level=True) on pre-Blackwell +# GPUs) as success. FAILED=0 for e in "${EXITS[@]}"; do if [ "$e" != "0" ] && [ "$e" != "5" ]; then @@ -118,14 +108,14 @@ done echo if [ "$FAILED" -eq 0 ]; then - echo "[run_multiprocess_moe_vjp.sh] all processes PASSED" - if [ -z "${MOE_VJP_MP_LOG_DIR:-}" ]; then + echo "[run_te_ep_moe.sh] all processes PASSED" + if [ -z "${TE_EP_MOE_MP_LOG_DIR:-}" ]; then rm -rf "$LOG_DIR" fi exit 0 fi -echo "[run_multiprocess_moe_vjp.sh] at least one process FAILED" +echo "[run_te_ep_moe.sh] at least one process FAILED" echo " retaining logs at $LOG_DIR for diagnosis" echo " process 0 tail:" tail -20 "$LOG_DIR/proc_0.log" 2>/dev/null || true diff --git a/tests/jax/test_moe_vjp.py b/tests/jax/test_moe_vjp.py deleted file mode 100644 index cc458d039e..0000000000 --- a/tests/jax/test_moe_vjp.py +++ /dev/null @@ -1,443 +0,0 @@ -# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# -# See LICENSE for license information. - -"""Single-device tests for the unified MoE custom_vjp at -``transformer_engine.jax.moe.moe`` (and its Flax wrapper -``transformer_engine.jax.flax._MoEBlock``). - -Strategy --------- - -Rather than reproducing every internal kernel residual, we rely on a -single end-to-end pure-JAX *reference* implementation of the whole -MoE block (``_pure_jax_moe_reference`` below) and compare the TE -``moe(...)`` forward output AND parameter gradients against it. This -gives us coverage of: - -* the gate GEMM, -* the fused top-k routing primitive (and its bwd), -* the dispatch / per-expert FFN / combine pipeline (and their bwds - threaded through the absorbed primitives), -* the optional aux-loss path (and its bwd). - -The reference uses only ``jnp`` ops + ``jax.vjp``, so we get a -"definitive" pullback to compare against without needing the TE -primitive bwd kernels. - -Distributed (EP + FSDP) testing is intentionally NOT in this file -- -that needs a multi-device setup and lives in -``tests/jax/test_distributed_moe_vjp.py`` (follow-up). -""" - -from functools import partial -from typing import Optional, Tuple - -import jax -import jax.numpy as jnp -import numpy as np -import pytest - -from transformer_engine_jax import get_device_compute_capability -from transformer_engine.jax.flax import _MoEBlock as MoEBlock -from transformer_engine.jax.moe import PermutationBackend, moe - -# The MoE custom_vjp uses grouped GEMM, which is currently -# Blackwell-only (sm_100+). Skip the whole file on older arches. -if get_device_compute_capability(0) < 100: - pytest.skip( - "MoE custom_vjp tests require Blackwell (sm_100+) for grouped GEMM", - allow_module_level=True, - ) - -# Parametrize values for the dispatch / combine backend. Only the -# ``triton`` variant is gated by the ``triton`` marker (so the -# ``pure_jax`` variant still runs on environments without Triton). -BACKEND_PARAMS = [ - pytest.param("pure_jax", id="pure_jax"), - pytest.param("triton", id="triton", marks=pytest.mark.triton), -] - - -# ----------------------------------------------------------------------------- -# Test config -# ----------------------------------------------------------------------------- - -DTYPE = jnp.float32 # use fp32 for tighter parity assertions -BATCH_SIZE = 2 -SEQUENCE_LENGTH = 16 -HIDDEN_SIZE = 32 -INTERMEDIATE_SIZE = 64 -NUM_EXPERTS = 8 -NUM_EXPERTS_PER_TOK = 2 - - -def _make_inputs(key: jax.Array, *, batch=BATCH_SIZE, seq=SEQUENCE_LENGTH) -> jax.Array: - return jax.random.normal(key, (batch, seq, HIDDEN_SIZE), dtype=DTYPE) - - -# ----------------------------------------------------------------------------- -# Pure-JAX reference MoE -# ----------------------------------------------------------------------------- -# -# Implements EXACTLY the same math as ``moe(...)`` for the no-EP, -# softmax-routing, no-bias, silu activation, no-quantization path. -# Returns ``(output, aux_loss_or_zero)``. Used as ground truth for both -# fwd and bwd parity. - - -@partial( - jax.jit, - static_argnames=("num_experts", "num_experts_per_tok", "aux_loss_coeff"), -) -def _pure_jax_moe_reference( - x: jnp.ndarray, - gate_kernel: jnp.ndarray, - wi_0: jnp.ndarray, - wi_1: jnp.ndarray, - wo: jnp.ndarray, - *, - num_experts: int, - num_experts_per_tok: int, - aux_loss_coeff: float = 0.0, -) -> Tuple[jnp.ndarray, jnp.ndarray]: - """Reference no-EP MoE forward (pure JAX, no TE primitives). - - Mirrors :func:`transformer_engine.jax.moe._body_fwd` for the - PURE_JAX backend, no biases, softmax routing, silu activation, - no quantization. Linear ops only -- ``jax.vjp`` over this gives - the canonical bwd to compare against. - """ - B, S, H = x.shape - T = B * S - x_2d = x.reshape(T, H) - - # Gate - logits = x_2d @ gate_kernel # [T, E] - - # Softmax + topk (no expert_bias, no grouping, scale=1.0) - probs_full = jax.nn.softmax(logits, axis=-1) # [T, E] - # top-k by probability: - sorted_idx = jnp.argsort(probs_full, axis=-1) # ascending - selected = sorted_idx[:, -num_experts_per_tok:] # [T, K] - weights = jnp.take_along_axis(probs_full, selected, axis=-1) # [T, K] - # Normalize topk weights to sum to 1 (matches softmax->topk semantics - # of fused_topk_with_score_function with use_pre_softmax=False): - weights = weights / jnp.sum(weights, axis=-1, keepdims=True) - - # Build a sparse routing_map [T, E] with weights at selected positions - routing_weights_full = jnp.zeros_like(probs_full) - routing_weights_full = routing_weights_full.at[jnp.arange(T)[:, None], selected].set(weights) - - # Per-expert FFN: replicate each token K times, gather by expert, - # run through wi_0 / wi_1 / wo, gather back, weighted-sum. - # - # Vectorize the gather without sorting: for each (token, slot k), - # multiply the corresponding expert's FFN by routing_weights[t, k] - # and sum over experts. - # x_2d: [T, H], wi_0: [E, H, M], wi_1: [E, H, M], wo: [E, M, H] - # For each expert e: layer_w0_e = x_2d @ wi_0[e]; layer_w1_e = x_2d @ wi_1[e] - # intermediate_e = silu(layer_w0_e) * layer_w1_e - # expert_out_e = intermediate_e @ wo[e] - # output[t, h] = sum_e routing_weights_full[t, e] * expert_out_e[t, h] - layer_w0 = jnp.einsum("th,ehm->tem", x_2d, wi_0) # [T, E, M] - layer_w1 = jnp.einsum("th,ehm->tem", x_2d, wi_1) # [T, E, M] - intermediate = jax.nn.silu(layer_w0) * layer_w1 # [T, E, M] - expert_out = jnp.einsum("tem,emh->teh", intermediate, wo) # [T, E, H] - output_2d = jnp.einsum("te,teh->th", routing_weights_full, expert_out) # [T, H] - output = output_2d.reshape(B, S, H) - - if aux_loss_coeff > 0.0: - # aux scores: clean per-expert softmax (compute_aux_scores=True - # kernel uses a clean softmax, no bias, scale=1, no grouping). - aux_probs = jax.nn.softmax(logits.astype(jnp.float32), axis=-1) - # tokens_per_expert from REAL routing_map (post-grouping); here - # there's no grouping so == count of non-zero positions per expert. - routing_map = (routing_weights_full > 0).astype(jnp.int32) - tokens_per_expert = jnp.sum(routing_map, axis=0) # [E] - # aux_loss formula: (E * coeff / (k * T^2)) * sum_e - # (sum_t aux_probs[t, e]) * tokens_per_expert[e] - sum_probs_per_expert = jnp.sum(aux_probs, axis=0) # [E] - aux_loss = (num_experts * aux_loss_coeff / (num_experts_per_tok * (T**2))) * jnp.sum( - sum_probs_per_expert * tokens_per_expert.astype(jnp.float32) - ) - else: - aux_loss = jnp.zeros((), dtype=DTYPE) - - return output, aux_loss - - -# ----------------------------------------------------------------------------- -# Helpers -# ----------------------------------------------------------------------------- - - -def _init_params(key: jax.Array) -> dict: - k_g, k_w0, k_w1, k_wo = jax.random.split(key, 4) - init = jax.nn.initializers.variance_scaling(1.0, "fan_in", "truncated_normal") - return dict( - gate_kernel=init(k_g, (HIDDEN_SIZE, NUM_EXPERTS), DTYPE), - wi_0=init(k_w0, (NUM_EXPERTS, HIDDEN_SIZE, INTERMEDIATE_SIZE), DTYPE), - wi_1=init(k_w1, (NUM_EXPERTS, HIDDEN_SIZE, INTERMEDIATE_SIZE), DTYPE), - wo=init(k_wo, (NUM_EXPERTS, INTERMEDIATE_SIZE, HIDDEN_SIZE), DTYPE), - ) - - -@partial(jax.jit, static_argnames=("permutation_backend", "aux_loss_coeff")) -def _run_te_moe( - x: jnp.ndarray, - params: dict, - *, - permutation_backend, - aux_loss_coeff: float = 0.0, -) -> Tuple[jnp.ndarray, Optional[jnp.ndarray]]: - return moe( - x, - params["gate_kernel"], - params["wi_0"], - params["wi_1"], - params["wo"], - num_experts=NUM_EXPERTS, - num_experts_per_tok=NUM_EXPERTS_PER_TOK, - activation_type="silu", - score_function="softmax", - use_pre_softmax=False, - scaling_factor=1.0, - aux_loss_coeff=aux_loss_coeff, - permutation_backend=permutation_backend, - align_size=0, - dtype=DTYPE, - ) - - -@partial(jax.jit, static_argnames=("permutation_backend", "aux_loss_coeff")) -def _grads_te_main_loss(params, x, *, permutation_backend, aux_loss_coeff: float = 0.0): - """jit'd grad of ``mean(out**2)`` w.r.t. params (no aux contribution).""" - - def loss(params, x): - out, _ = _run_te_moe( - x, params, permutation_backend=permutation_backend, aux_loss_coeff=aux_loss_coeff - ) - return jnp.mean(out**2) - - return jax.grad(loss)(params, x) - - -@partial(jax.jit, static_argnames=("num_experts", "num_experts_per_tok", "aux_loss_coeff")) -def _grads_ref_main_loss(params, x, *, num_experts, num_experts_per_tok, aux_loss_coeff=0.0): - """jit'd grad of ``mean(out**2)`` w.r.t. params on the pure-JAX ref.""" - - def loss(params, x): - out, _ = _pure_jax_moe_reference( - x, - **params, - num_experts=num_experts, - num_experts_per_tok=num_experts_per_tok, - aux_loss_coeff=aux_loss_coeff, - ) - return jnp.mean(out**2) - - return jax.grad(loss)(params, x) - - -@partial(jax.jit, static_argnames=("permutation_backend",)) -def _grad_te_aux_only(params, x, *, permutation_backend): - """jit'd grad of just the aux loss scalar (no main contribution).""" - - def aux_only(params, x): - _, aux = _run_te_moe( - x, params, permutation_backend=permutation_backend, aux_loss_coeff=1e-2 - ) - return aux.astype(jnp.float32) - - return jax.grad(aux_only)(params, x) - - -# ----------------------------------------------------------------------------- -# Tests -# ----------------------------------------------------------------------------- - - -class TestMoeVjpForward: - """Forward shape / finiteness / parity vs pure-JAX reference.""" - - @pytest.mark.parametrize("backend_name", BACKEND_PARAMS) - def test_forward_shape_and_finite(self, backend_name): - backend = PermutationBackend(backend_name) - key = jax.random.PRNGKey(0) - kp, kx = jax.random.split(key) - params = _init_params(kp) - x = _make_inputs(kx) - out, aux = _run_te_moe(x, params, permutation_backend=backend) - assert out.shape == x.shape - assert out.dtype == x.dtype - assert jnp.all(jnp.isfinite(out)) - assert aux is None - - @pytest.mark.parametrize("backend_name", BACKEND_PARAMS) - def test_forward_parity_vs_pure_jax_reference(self, backend_name): - backend = PermutationBackend(backend_name) - key = jax.random.PRNGKey(1) - kp, kx = jax.random.split(key) - params = _init_params(kp) - x = _make_inputs(kx) - out_te, _ = _run_te_moe(x, params, permutation_backend=backend) - out_ref, _ = _pure_jax_moe_reference( - x, - **params, - num_experts=NUM_EXPERTS, - num_experts_per_tok=NUM_EXPERTS_PER_TOK, - ) - # FP32, small shapes -> tight tolerance - np.testing.assert_allclose(np.array(out_te), np.array(out_ref), atol=2e-5, rtol=2e-5) - - def test_pure_jax_triton_equivalence(self): - key = jax.random.PRNGKey(2) - kp, kx = jax.random.split(key) - params = _init_params(kp) - x = _make_inputs(kx) - out_pj, _ = _run_te_moe(x, params, permutation_backend=PermutationBackend.PURE_JAX) - out_tr, _ = _run_te_moe(x, params, permutation_backend=PermutationBackend.TRITON) - np.testing.assert_allclose(np.array(out_pj), np.array(out_tr), atol=2e-5, rtol=2e-5) - - -class TestMoeVjpBackward: - """Backward parity vs pure-JAX reference (which uses ``jax.vjp`` over - plain JAX ops, giving us the canonical pullback).""" - - @pytest.mark.parametrize("backend_name", BACKEND_PARAMS) - def test_grads_finite_and_nonzero(self, backend_name): - backend = PermutationBackend(backend_name) - key = jax.random.PRNGKey(3) - kp, kx = jax.random.split(key) - params = _init_params(kp) - x = _make_inputs(kx) - grads = _grads_te_main_loss(params, x, permutation_backend=backend) - for name in ("gate_kernel", "wi_0", "wi_1", "wo"): - g = grads[name] - assert jnp.all(jnp.isfinite(g)), f"{name} grad has NaN/Inf" - assert jnp.any(g != 0.0), f"{name} grad is identically zero" - - @pytest.mark.parametrize("backend_name", BACKEND_PARAMS) - def test_grads_match_pure_jax_reference(self, backend_name): - backend = PermutationBackend(backend_name) - key = jax.random.PRNGKey(4) - kp, kx = jax.random.split(key) - params = _init_params(kp) - x = _make_inputs(kx) - grads_te = _grads_te_main_loss(params, x, permutation_backend=backend) - grads_ref = _grads_ref_main_loss( - params, - x, - num_experts=NUM_EXPERTS, - num_experts_per_tok=NUM_EXPERTS_PER_TOK, - ) - # Loose-ish tol on grads: routing path has discrete topk so the - # softmax cotangent paths through the non-topk experts diverge - # slightly between TE (which uses the fused topk bwd) and the - # reference (which uses argsort-based take_along_axis). - # Tighter than the bf16 tests. - for name in ("wi_0", "wi_1", "wo"): - np.testing.assert_allclose( - np.array(grads_te[name]), - np.array(grads_ref[name]), - atol=5e-5, - rtol=5e-5, - err_msg=f"grad mismatch on {name}", - ) - # Gate grad has more error budget because it propagates through - # the topk derivative kernel (which differs in zero-pattern - # treatment from a plain take_along_axis). - np.testing.assert_allclose( - np.array(grads_te["gate_kernel"]), - np.array(grads_ref["gate_kernel"]), - atol=5e-4, - rtol=5e-4, - err_msg="grad mismatch on gate_kernel", - ) - - -class TestMoeVjpAuxLoss: - """Aux-loss path: forward + grad parity.""" - - @pytest.mark.parametrize("backend_name", BACKEND_PARAMS) - def test_aux_loss_returned_and_finite(self, backend_name): - backend = PermutationBackend(backend_name) - key = jax.random.PRNGKey(5) - kp, kx = jax.random.split(key) - params = _init_params(kp) - x = _make_inputs(kx) - _, aux = _run_te_moe(x, params, permutation_backend=backend, aux_loss_coeff=1e-2) - assert aux is not None - assert aux.shape == () - assert jnp.isfinite(aux) - assert jnp.abs(aux) < 1e2 - - @pytest.mark.parametrize("backend_name", BACKEND_PARAMS) - def test_aux_loss_parity_vs_reference(self, backend_name): - backend = PermutationBackend(backend_name) - key = jax.random.PRNGKey(6) - kp, kx = jax.random.split(key) - params = _init_params(kp) - x = _make_inputs(kx) - _, aux_te = _run_te_moe(x, params, permutation_backend=backend, aux_loss_coeff=1e-2) - _, aux_ref = _pure_jax_moe_reference( - x, - **params, - num_experts=NUM_EXPERTS, - num_experts_per_tok=NUM_EXPERTS_PER_TOK, - aux_loss_coeff=1e-2, - ) - np.testing.assert_allclose(float(aux_te), float(aux_ref), atol=1e-5, rtol=1e-5) - - @pytest.mark.parametrize("backend_name", BACKEND_PARAMS) - def test_aux_loss_grads_propagate_to_logits(self, backend_name): - """The aux-loss bwd path must produce non-zero gate-kernel grads - when only the aux-loss scalar is differentiated (no main-output - contribution).""" - backend = PermutationBackend(backend_name) - key = jax.random.PRNGKey(7) - kp, kx = jax.random.split(key) - params = _init_params(kp) - x = _make_inputs(kx) - g_gate = _grad_te_aux_only(params, x, permutation_backend=backend)["gate_kernel"] - assert jnp.all(jnp.isfinite(g_gate)) - assert jnp.any( - g_gate != 0.0 - ), "aux_loss bwd should propagate to gate_kernel via fused_topk bwd" - - -# ----------------------------------------------------------------------------- -# Flax wrapper smoke test -# ----------------------------------------------------------------------------- - - -class TestMoEBlockFlaxWrapper: - """Sanity-check the thin Flax wrapper: forward + grad on init.""" - - def test_init_and_apply(self): - block = MoEBlock( - num_experts=NUM_EXPERTS, - num_experts_per_tok=NUM_EXPERTS_PER_TOK, - intermediate_size=INTERMEDIATE_SIZE, - permutation_backend=PermutationBackend.PURE_JAX, - dtype=DTYPE, - ) - key = jax.random.PRNGKey(8) - ki, kx = jax.random.split(key) - x = _make_inputs(kx) - variables = jax.jit(block.init)(ki, x) - out, aux = jax.jit(block.apply)(variables, x) - assert out.shape == x.shape - assert aux is None - - @jax.jit - def grad_fn(variables, x): - return jax.grad(lambda v, x: jnp.mean(block.apply(v, x)[0] ** 2))(variables, x) - - grads = grad_fn(variables, x) - for name in ("gate_kernel", "wi_0", "wi_1", "wo"): - g = grads["params"][name] - g = g.value if hasattr(g, "value") else g - assert jnp.all(jnp.isfinite(g)), f"{name} grad NaN/Inf" - assert jnp.any(g != 0.0), f"{name} grad zero" diff --git a/tests/jax/test_multi_process_ep.py b/tests/jax/test_multi_process_ep.py index 0b8bb25f3f..d3effcdc07 100644 --- a/tests/jax/test_multi_process_ep.py +++ b/tests/jax/test_multi_process_ep.py @@ -36,6 +36,7 @@ ep_combine_fwd, get_ep_config, ) +from transformer_engine.jax.version_utils import is_collective_stream_supported # ── Test config ───────────────────────────────────────────────────────────── @@ -660,6 +661,55 @@ def run(idx, toks, w): expected = (("dp", "ep"),) if self.dp > 1 else ("ep",) self.assertEqual(tuple(compiled.output_shardings.spec), expected) + @unittest.skipUnless( + is_collective_stream_supported(), + "JAX/XLA lacks the gpu_stream:collective annotation (openxla/xla#39604)", + ) + def test_z_dispatch_combine_on_collective_stream(self): + """Every EP FFI custom call must carry the collective-stream annotation + so XLA schedules them on the collective stream instead of overlapping + them with other collectives.""" + T_dp, tokens, topk_idx, topk_w = self._make_random_inputs() + dp_spec = PartitionSpec(("dp", "ep"), None) + ep_spec_3d = PartitionSpec(("dp", "ep"), None, None) + ep_spec_2d = PartitionSpec(("dp", "ep"), None) + + with self.mesh, global_shard_guard(self.mr): + + @jax.jit + def run(idx, toks, w): + idx = jax.lax.with_sharding_constraint(idx, NamedSharding(self.mesh, dp_spec)) + toks = jax.lax.with_sharding_constraint(toks, NamedSharding(self.mesh, dp_spec)) + w = jax.lax.with_sharding_constraint(w, NamedSharding(self.mesh, dp_spec)) + recv_t, recv_w, hm, tc = ep_dispatch( + self.hk, idx, toks, w, self.recv_capacity_per_rank + ) + recv_t = jax.lax.with_sharding_constraint( + recv_t, NamedSharding(self.mesh, ep_spec_3d) + ) + recv_w = jax.lax.with_sharding_constraint( + recv_w, NamedSharding(self.mesh, ep_spec_2d) + ) + weighted = self._preweight_expert_out(recv_t, recv_w) + out = ep_combine(self.hk, hm, tc, weighted, T_dp, out_sharding=(("dp", "ep"), None)) + return jax.lax.with_sharding_constraint(out, NamedSharding(self.mesh, dp_spec)) + + hlo = run.lower(topk_idx, tokens, topk_w).compile().as_text() + + # Every te_ep_* FFI custom call must carry the collective-stream + # annotation so XLA places it on the collective stream. + ep_lines = [l for l in hlo.splitlines() if 'custom_call_target="te_ep_' in l] + self.assertTrue(ep_lines, f"no te_ep_* custom calls in compiled HLO:\n{hlo}") + missing = [ + l.strip()[:200] + for l in ep_lines + if '_xla_stream_annotation="collective"' not in l.replace(" ", "") + ] + self.assertFalse( + missing, + "te_ep_* custom calls missing collective-stream annotation:\n" + "\n".join(missing), + ) + def test_z_no_unexpected_reshard_in_hlo_bwd(self): """Compiled bwd HLO must not insert XLA collectives outside the EP FFI.""" T_dp, tokens, topk_idx, topk_w = self._make_random_inputs() diff --git a/tests/jax/test_multiprocess_moe_vjp.py b/tests/jax/test_multiprocess_moe_vjp.py deleted file mode 100644 index 97044780f0..0000000000 --- a/tests/jax/test_multiprocess_moe_vjp.py +++ /dev/null @@ -1,406 +0,0 @@ -# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# -# See LICENSE for license information. - -"""Multi-process (one-GPU-per-process) tests for the unified MoE custom_vjp. - -The launcher ``tests/jax/run_multiprocess_moe_vjp.sh`` forks one pytest -process per visible GPU (mirroring -``examples/jax/encoder/run_test_multiprocessing_encoder.sh``). Each -process binds to exactly one device via -``jax.distributed.initialize(..., local_device_ids=process_id)``; the -participating processes form a global mesh through JAX's distributed -runtime. - -How to run ----------- - -You typically do NOT invoke pytest on this file directly -- use the -launcher, which passes ``--num-process=N --process-id=i`` to each -forked process. Driving it directly with only one process will skip -every test because :func:`jax.distributed.initialize` requires -multiple participants. - - bash tests/jax/run_multiprocess_moe_vjp.sh - -CI invocation lives in ``qa/L0_jax_distributed_unittest/test.sh``. -""" - -import os - -# NCCL needs HBM headroom that JAX's default 90% preallocation does -# not leave. Set before any jax import below. -os.environ.setdefault("XLA_PYTHON_CLIENT_PREALLOCATE", "false") -os.environ.setdefault("XLA_PYTHON_CLIENT_MEM_FRACTION", "0.5") - -import sys - -import jax -import jax.numpy as jnp -import numpy as np -import pytest - -from jax.experimental import mesh_utils -from jax.sharding import Mesh, NamedSharding, PartitionSpec as P -from flax.linen import partitioning as nn_partitioning - - -# Per-process distributed bootstrap. Each pytest invocation initializes -# JAX with exactly one local device (its assigned GPU). Once -# initialized, the four processes form one global mesh of 4 devices. -def _init_distributed(num_process: int, process_id: int) -> bool: - """Initialize jax.distributed for this pytest process. - - Returns True if initialization succeeded (i.e. this is a real - multi-process launch), False if num_process == 0 / 1 meaning the - file is being collected without a launcher and tests should be - skipped at module level. - """ - if num_process <= 1: - return False - coord = os.environ.get("MOE_VJP_COORDINATOR_ADDRESS", "127.0.0.1:1234") - jax.distributed.initialize( - coordinator_address=coord, - num_processes=num_process, - process_id=process_id, - local_device_ids=process_id, - ) - assert jax.local_device_count() == 1, "one GPU per process is the whole point" - assert ( - jax.device_count() == num_process - ), f"global device_count {jax.device_count()} != num_process {num_process}" - return True - - -# Read --num-process / --process-id BEFORE pytest collects any tests so -# we can fast-skip the whole module when not in a multiprocess launch. -def _read_mp_options(): - # Use pytest's option lookup via the request fixture isn't available - # at module top-level; parse argv ourselves the same way encoder - # test does. CLI form is e.g. "pytest ... --num-process=4 --process-id=0". - num = int(os.environ.get("MP_NUM_PROCESS", "0") or "0") - pid = int(os.environ.get("MP_PROCESS_ID", "0") or "0") - for i, a in enumerate(sys.argv): - if a.startswith("--num-process="): - num = int(a.split("=", 1)[1]) - elif a == "--num-process" and i + 1 < len(sys.argv): - num = int(sys.argv[i + 1]) - elif a.startswith("--process-id="): - pid = int(a.split("=", 1)[1]) - elif a == "--process-id" and i + 1 < len(sys.argv): - pid = int(sys.argv[i + 1]) - return num, pid - - -_MP_NUM_PROCESS, _MP_PROCESS_ID = _read_mp_options() -_MP_ACTIVE = _init_distributed(_MP_NUM_PROCESS, _MP_PROCESS_ID) - -if not _MP_ACTIVE: - # Skip the entire module if not launched via the multiprocess - # runner. Lets `pytest tests/jax/` collect this file harmlessly. - pytest.skip( - "test_multiprocess_moe_vjp.py requires the multiprocess launcher " - "(run_multiprocess_moe_vjp.sh). Skipping.", - allow_module_level=True, - ) - -from transformer_engine_jax import get_device_compute_capability - -# Grouped GEMM in the MoE custom_vjp currently requires Blackwell -# (sm_100+). Skip the whole file on older arches. -if get_device_compute_capability(0) < 100: - pytest.skip( - "MoE custom_vjp tests require Blackwell (sm_100+) for grouped GEMM", - allow_module_level=True, - ) - -import transformer_engine.jax as te -from transformer_engine.common import recipe as te_recipe -from transformer_engine.jax.flax import _MoEBlock as MoEBlock -from transformer_engine.jax.moe import PermutationBackend -from transformer_engine.jax.sharding import MeshResource, global_shard_guard - -# Parametrize values for the dispatch / combine backend. Only the -# ``triton`` variant carries the ``triton`` marker, so the -# ``pure_jax`` variant still runs on environments without Triton. -BACKEND_PARAMS = [ - pytest.param("pure_jax", id="pure_jax"), - pytest.param("triton", id="triton", marks=pytest.mark.triton), -] - - -EP_AXIS = "ep" -FSDP_AXIS = "fsdp" -EP_SIZE = 2 -# FSDP_SIZE adapts to whatever the launcher gave us: dlcluster GB200 -# gives 4 GPUs (FSDP=2), CI B200 gives 8 GPUs (FSDP=4). Both stay -# 128-aligned for MXFP8 and divide num_experts/topk cleanly. -assert ( - jax.device_count() % EP_SIZE == 0 -), f"device_count {jax.device_count()} must be divisible by EP_SIZE={EP_SIZE}" -FSDP_SIZE = jax.device_count() // EP_SIZE -NUM_DEVICES_REQUIRED = EP_SIZE * FSDP_SIZE - -LOGICAL_AXIS_RULES = ( - ("exp", EP_AXIS), - ("embed", FSDP_AXIS), - ("mlp", None), - ("batch", (EP_AXIS, FSDP_AXIS)), -) - - -@pytest.fixture(scope="module") -def mesh(): - if jax.device_count() < NUM_DEVICES_REQUIRED: - pytest.skip( - f"Need >={NUM_DEVICES_REQUIRED} devices for ep={EP_SIZE} x fsdp={FSDP_SIZE};" - f" have {jax.device_count()}" - ) - devices = mesh_utils.create_device_mesh((EP_SIZE, FSDP_SIZE)) - return Mesh(devices, axis_names=(EP_AXIS, FSDP_AXIS)) - - -# ``recipe`` parametrize values used across all tests below. ``None`` -# = plain bf16; the named recipes route through TE's autocast and -# exercise the FP8/MXFP8 quantization paths in _body_fwd/_body_bwd. -# Only recipes that work on TE Blackwell are included; older GPUs -# skip via the ``hardware_supports`` guard below. -RECIPE_NAMES = ("bf16", "MXFP8BlockScaling") - - -def _resolve_recipe(name): - """Return ``(use_fp8, recipe_instance)`` for the parametrize id.""" - if name == "bf16": - return False, None - if name == "MXFP8BlockScaling": - return True, te_recipe.MXFP8BlockScaling() - raise ValueError(f"unknown recipe name: {name!r}") - - -def _hardware_supports(recipe_name): - """Skip an FP8 recipe on GPUs that don't have the hw for it.""" - if recipe_name == "bf16": - return True - from transformer_engine_jax import get_device_compute_capability - - arch = get_device_compute_capability(0) - if recipe_name == "MXFP8BlockScaling": - return arch >= 100 - return False - - -def _autocast_ctx(recipe_name): - """Context manager that turns FP8 on for non-bf16 recipes.""" - use_fp8, recipe_inst = _resolve_recipe(recipe_name) - return te.autocast(enabled=use_fp8, recipe=recipe_inst) - - -def _tol_finite_grad(recipe_name): - """Per-recipe absolute tolerance for parity grad comparison.""" - if recipe_name == "bf16": - return 5e-2 - # MXFP8 grads carry block-scale quantization noise; loosen accordingly. - return 3e-1 - - -# ----------------------------------------------------------------------------- -# Helpers -# ----------------------------------------------------------------------------- - - -def _make_block( - *, - num_experts, - num_experts_per_tok, - intermediate_size, - permutation_backend, - aux_loss_coeff=0.0, - dtype=jnp.bfloat16, - align_size=0, -): - return MoEBlock( - num_experts=num_experts, - num_experts_per_tok=num_experts_per_tok, - intermediate_size=intermediate_size, - permutation_backend=permutation_backend, - data_parallelism_axes=(FSDP_AXIS,), - aux_loss_coeff=aux_loss_coeff, - dtype=dtype, - _align_size=align_size, - ) - - -def _shard_inputs(x, mesh): - return jax.lax.with_sharding_constraint( - x, NamedSharding(mesh, P((EP_AXIS, FSDP_AXIS), None, None)) - ) - - -def _init_apply(block, mesh, x, key): - with mesh, global_shard_guard( - MeshResource(ep_resource=EP_AXIS, fsdp_resource=FSDP_AXIS) - ), nn_partitioning.axis_rules(LOGICAL_AXIS_RULES): - x = _shard_inputs(x, mesh) - variables = jax.jit(block.init)(key, x) - jax.block_until_ready(jax.tree_util.tree_leaves(variables)[0]) - output, aux = jax.jit(block.apply)(variables, x) - jax.block_until_ready(output) - return variables, output, aux - - -def _grad_step(block, variables, mesh, x): - with mesh, global_shard_guard( - MeshResource(ep_resource=EP_AXIS, fsdp_resource=FSDP_AXIS) - ), nn_partitioning.axis_rules(LOGICAL_AXIS_RULES): - x = _shard_inputs(x, mesh) - - def loss_fn(variables, x): - output, aux = block.apply(variables, x) - main = jnp.mean(output.astype(jnp.float32) ** 2) - return main + (aux.astype(jnp.float32) if aux is not None else 0.0) - - grads = jax.jit(jax.grad(loss_fn))(variables, x) - jax.block_until_ready(jax.tree_util.tree_leaves(grads)[0]) - return grads - - -def _unwrap(x): - return x.value if hasattr(x, "value") else x - - -def _local_shard(x): - """Return the local (this-process) shard of a global JAX Array as numpy. - - Every assertion in this file is structural (finite-ness, non-zero, - parity within tolerance). For all of these, checking the local - shard on each process is sufficient and avoids any cross-process - collective in the test machinery. ``arr.addressable_data(0)`` - returns the local-device view of the sharded array -- with one - GPU per process there is exactly one addressable shard. - """ - return np.asarray(jax.device_get(x.addressable_data(0))) - - -# ----------------------------------------------------------------------------- -# Mixtral-style shapes, sized to fit on a single 4-GPU bf16 box (a -# 4-way data-parallel shard of a Mixtral-8 block). -# ----------------------------------------------------------------------------- - -BATCH = EP_SIZE * FSDP_SIZE * 4 # 16 on 4-GPU, 32 on 8-GPU -SEQ = 2048 -HIDDEN = 1024 -INTER = 4096 -NUM_EXPERTS = 8 -TOPK = 2 - - -class TestMoeVjpMultiprocess: - """Multiprocess (one-GPU-per-process) correctness checks for the - unified MoE custom_vjp. - """ - - @pytest.mark.parametrize("backend_name", BACKEND_PARAMS) - @pytest.mark.parametrize("recipe_name", RECIPE_NAMES) - def test_fwd_and_bwd(self, mesh, backend_name, recipe_name): - if not _hardware_supports(recipe_name): - pytest.skip(f"recipe {recipe_name} not supported on this GPU") - backend = PermutationBackend(backend_name) - block = _make_block( - num_experts=NUM_EXPERTS, - num_experts_per_tok=TOPK, - intermediate_size=INTER, - permutation_backend=backend, - ) - x = jax.random.normal( - jax.random.PRNGKey(0), - (BATCH, SEQ, HIDDEN), - dtype=jnp.bfloat16, - ) - with _autocast_ctx(recipe_name): - variables, output, aux = _init_apply(block, mesh, x, jax.random.PRNGKey(1)) - # Local-shard checks (see _local_shard docstring for why). - out_local = _local_shard(output) - assert output.dtype == x.dtype - assert np.all(np.isfinite(out_local)), "output has NaN/Inf" - assert aux is None - with _autocast_ctx(recipe_name): - grads = _grad_step(block, variables, mesh, x) - for name in ("gate_kernel", "wi_0", "wi_1", "wo"): - g_local = _local_shard(_unwrap(grads["params"][name])) - assert np.all(np.isfinite(g_local)), f"{name} grad has NaN/Inf" - assert np.any(g_local != 0.0), f"{name} grad is identically zero" - - @pytest.mark.parametrize("backend_name", BACKEND_PARAMS) - @pytest.mark.parametrize("recipe_name", RECIPE_NAMES) - def test_aux_loss(self, mesh, backend_name, recipe_name): - if not _hardware_supports(recipe_name): - pytest.skip(f"recipe {recipe_name} not supported on this GPU") - backend = PermutationBackend(backend_name) - block = _make_block( - num_experts=NUM_EXPERTS, - num_experts_per_tok=TOPK, - intermediate_size=INTER, - permutation_backend=backend, - aux_loss_coeff=1e-2, - ) - x = jax.random.normal( - jax.random.PRNGKey(4), - (BATCH, SEQ, HIDDEN), - dtype=jnp.bfloat16, - ) - with _autocast_ctx(recipe_name): - variables, output, aux = _init_apply(block, mesh, x, jax.random.PRNGKey(5)) - out_local = _local_shard(output) - assert np.all(np.isfinite(out_local)), "output has NaN/Inf under aux" - assert aux is not None - assert aux.shape == () - aux_local = _local_shard(aux) - assert np.isfinite(aux_local), "aux is NaN/Inf" - with _autocast_ctx(recipe_name): - grads = _grad_step(block, variables, mesh, x) - g_gate_local = _local_shard(_unwrap(grads["params"]["gate_kernel"])) - assert np.all(np.isfinite(g_gate_local)), "gate grad NaN/Inf under aux" - - @pytest.mark.parametrize("recipe_name", RECIPE_NAMES) - def test_pure_jax_triton_parity(self, mesh, recipe_name): - if not _hardware_supports(recipe_name): - pytest.skip(f"recipe {recipe_name} not supported on this GPU") - block_pj = _make_block( - num_experts=NUM_EXPERTS, - num_experts_per_tok=TOPK, - intermediate_size=INTER, - permutation_backend=PermutationBackend.PURE_JAX, - ) - block_tr = _make_block( - num_experts=NUM_EXPERTS, - num_experts_per_tok=TOPK, - intermediate_size=INTER, - permutation_backend=PermutationBackend.TRITON, - ) - x = jax.random.normal( - jax.random.PRNGKey(6), - (BATCH, SEQ, HIDDEN), - dtype=jnp.bfloat16, - ) - tol = _tol_finite_grad(recipe_name) - with _autocast_ctx(recipe_name): - variables, out_pj, _ = _init_apply(block_pj, mesh, x, jax.random.PRNGKey(7)) - with mesh, global_shard_guard( - MeshResource(ep_resource=EP_AXIS, fsdp_resource=FSDP_AXIS) - ), nn_partitioning.axis_rules(LOGICAL_AXIS_RULES): - x_sh = _shard_inputs(x, mesh) - out_tr, _ = jax.jit(block_tr.apply)(variables, x_sh) - - out_pj_local = _local_shard(out_pj) - out_tr_local = _local_shard(out_tr) - diff = float(np.max(np.abs(out_pj_local - out_tr_local))) - assert diff < tol, f"forward parity breach: max_abs_diff={diff} (tol={tol})" - - with _autocast_ctx(recipe_name): - grads_pj = _grad_step(block_pj, variables, mesh, x) - grads_tr = _grad_step(block_tr, variables, mesh, x) - for name in ("gate_kernel", "wi_0", "wi_1", "wo"): - g_pj = _local_shard(_unwrap(grads_pj["params"][name])) - g_tr = _local_shard(_unwrap(grads_tr["params"][name])) - d = float(np.max(np.abs(g_pj - g_tr))) - assert d < tol, f"grad parity breach on {name}: max_abs_diff={d} (tol={tol})" diff --git a/tests/jax/test_te_ep_moe.py b/tests/jax/test_te_ep_moe.py new file mode 100644 index 0000000000..d08765e184 --- /dev/null +++ b/tests/jax/test_te_ep_moe.py @@ -0,0 +1,745 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""Multi-process (one-GPU-per-process) tests for the TE-EP MoE custom_vjp. + +The launcher ``tests/jax/run_te_ep_moe.sh`` forks one pytest process per +visible GPU. Each process binds to exactly one device via +``jax.distributed.initialize(..., local_device_ids=process_id)``; the +participating processes form a global ``(ep, fsdp)`` mesh through JAX's +distributed runtime. + +How to run +---------- + +You typically do NOT invoke pytest on this file directly -- use the +launcher, which passes ``--num-process=N --process-id=i`` to each +forked process. Driving it directly with only one process will skip +every test because :func:`jax.distributed.initialize` requires +multiple participants, and the TE EP NCCL primitives require at +least four ranks. + + bash tests/jax/run_te_ep_moe.sh + +What this suite covers +---------------------- + +Each test exercises one MoE-block run and bundles every check that +single run supports — shape, dtype, +finiteness AND numerical parity vs a pure-JAX reference. Variations +on the block are pytest parametrize values rather than separate test +classes: + +* ``test_forward`` covers the forward across a curated set of + configurations (softmax/sigmoid scoring, optional non-zero + expert_bias). Each config asserts shape, dtype, finiteness and + numerical parity vs the reference in one run. +* ``test_backward`` mirrors that for gradients. +* ``TestTeEpMoeAuxLoss`` covers the second return value end-to-end + (returned + parity + aux-only grad propagates to gate + combined + main+aux grads stay finite) in two consolidated tests. +""" + +import os + +os.environ.setdefault("XLA_PYTHON_CLIENT_PREALLOCATE", "false") +os.environ.setdefault("XLA_PYTHON_CLIENT_MEM_FRACTION", "0.5") + +import sys +from functools import partial + +import jax +import jax.numpy as jnp +import numpy as np +import pytest + +from jax.experimental import mesh_utils +from jax.sharding import Mesh, NamedSharding, PartitionSpec as P +from flax.linen import partitioning as nn_partitioning + + +def _init_distributed(num_process: int, process_id: int) -> bool: + """Initialize jax.distributed for this pytest process. + + Returns True on a real multi-process launch, False otherwise so + the module can fast-skip when pytest collects it without the + launcher. + """ + if num_process <= 1: + return False + coord = os.environ.get("TE_EP_MOE_COORDINATOR_ADDRESS", "127.0.0.1:13457") + jax.distributed.initialize( + coordinator_address=coord, + num_processes=num_process, + process_id=process_id, + local_device_ids=process_id, + ) + assert jax.local_device_count() == 1, "one GPU per process is required for TE EP" + assert ( + jax.device_count() == num_process + ), f"global device_count {jax.device_count()} != num_process {num_process}" + return True + + +def _read_mp_options(): + num = int(os.environ.get("MP_NUM_PROCESS", "0") or "0") + pid = int(os.environ.get("MP_PROCESS_ID", "0") or "0") + for i, a in enumerate(sys.argv): + if a.startswith("--num-process="): + num = int(a.split("=", 1)[1]) + elif a == "--num-process" and i + 1 < len(sys.argv): + num = int(sys.argv[i + 1]) + elif a.startswith("--process-id="): + pid = int(a.split("=", 1)[1]) + elif a == "--process-id" and i + 1 < len(sys.argv): + pid = int(sys.argv[i + 1]) + return num, pid + + +_MP_NUM_PROCESS, _MP_PROCESS_ID = _read_mp_options() +_MP_ACTIVE = _init_distributed(_MP_NUM_PROCESS, _MP_PROCESS_ID) + +if not _MP_ACTIVE: + pytest.skip( + "test_te_ep_moe.py requires the multiprocess launcher (run_te_ep_moe.sh). Skipping.", + allow_module_level=True, + ) + +from transformer_engine_jax import get_device_compute_capability + +# Grouped GEMM in the MoE custom_vjp requires Blackwell (sm_100+). The +# TE EP NCCL primitives themselves need SM>=90, but the FFN body uses +# grouped_gemm, so the file as a whole gates on sm_100+. +if get_device_compute_capability(0) < 100: + pytest.skip( + "MoE TE EP tests require Blackwell (sm_100+) for grouped GEMM", + allow_module_level=True, + ) + +from transformer_engine.jax.flax import _MoEBlock as MoEBlock +from transformer_engine.jax.moe import _ALIGN_SIZE, moe, record_ep_bootstrap_signature_for_moe +from transformer_engine.jax.ep import ep_bootstrap +from transformer_engine.jax.sharding import MeshResource, global_shard_guard + + +# ----------------------------------------------------------------------------- +# Mesh / shape config +# ----------------------------------------------------------------------------- + +EP_AXIS = "ep" +FSDP_AXIS = "fsdp" +EP_SIZE = 2 +assert ( + jax.device_count() % EP_SIZE == 0 +), f"device_count {jax.device_count()} must be divisible by EP_SIZE={EP_SIZE}" +FSDP_SIZE = jax.device_count() // EP_SIZE +NUM_DEVICES_REQUIRED = EP_SIZE * FSDP_SIZE + +LOGICAL_AXIS_RULES = ( + ("exp", EP_AXIS), + ("embed", FSDP_AXIS), + ("mlp", None), + ("batch", (EP_AXIS, FSDP_AXIS)), +) + +# Small shapes so the parity tests stay tight on bf16. The block still +# has all four ranks participating in dispatch/combine. +DTYPE = jnp.bfloat16 +BATCH = EP_SIZE * FSDP_SIZE * 2 # 8 on 4-GPU, 16 on 8-GPU +SEQ = 32 +HIDDEN = 64 +INTER = 128 +NUM_EXPERTS = 8 +TOPK = 2 + +# bf16 grouped_gemm + softmax-topk + ep all-to-all stack drifts ~1e-1 vs a +# fp32 numpy reference. Keep these tight enough to catch real bugs but +# loose enough to absorb expected bf16 rounding. +FWD_ATOL = 5e-2 +FWD_RTOL = 5e-2 +GRAD_FFN_ATOL = 1e-1 +GRAD_FFN_RTOL = 1e-1 +GRAD_GATE_ATOL = 5e-1 +GRAD_GATE_RTOL = 5e-1 + +# Two TE EP runs that should be bitwise-equal modulo XLA fusion order +# (slot alignment rounding, etc.). +TE_TO_TE_ATOL = 5e-3 +TE_TO_TE_RTOL = 5e-3 + +# Aux loss is computed in float32 from the SAME logits as the routing +# path. Numerical drift between TE-EP and the reference is dominated by +# the bf16-rounded softmax inside the topk kernel. +AUX_ATOL = 1e-3 +AUX_RTOL = 1e-3 + + +# ----------------------------------------------------------------------------- +# Fixtures +# ----------------------------------------------------------------------------- + + +def _compute_worst_case_recv_pr(): + """Per-rank recv buffer the bootstrap must reserve. + + NCCL EP HT expert-major uses one flat recv buffer with variable + per-expert zones. Each non-empty expert zone is padded to + ``_ALIGN_SIZE`` slots, so the reserve must cover the worst-case + total assignments plus independent per-zone padding. + """ + num_procs = jax.device_count() + num_local_experts = NUM_EXPERTS // EP_SIZE + max_tokens_per_rank = (BATCH // num_procs) * SEQ + tokens_per_ep_group = EP_SIZE * max_tokens_per_rank + max_local_assignments = tokens_per_ep_group * min(TOPK, num_local_experts) + max_nonempty_experts = min(num_local_experts, max_local_assignments) + padded_total_bound = max_local_assignments + (_ALIGN_SIZE - 1) * max_nonempty_experts + aligned_total_bound = ((padded_total_bound + _ALIGN_SIZE - 1) // _ALIGN_SIZE) * _ALIGN_SIZE + per_expert_bound = ( + num_local_experts * ((tokens_per_ep_group + _ALIGN_SIZE - 1) // _ALIGN_SIZE) * _ALIGN_SIZE + ) + return min(per_expert_bound, aligned_total_bound) + + +@pytest.fixture(scope="module") +def mesh(): + if jax.device_count() < NUM_DEVICES_REQUIRED: + pytest.skip( + f"Need >={NUM_DEVICES_REQUIRED} devices for ep={EP_SIZE} x fsdp={FSDP_SIZE};" + f" have {jax.device_count()}" + ) + # ``ep`` must be the inner axis: ``ep_bootstrap`` forms NCCL EP groups + # from consecutive global ranks via ``dp_color = rank // ep_size``, so + # only an (outer_fsdp, inner_ep) device layout groups ranks correctly. + devices = mesh_utils.create_device_mesh((FSDP_SIZE, EP_SIZE)) + mesh_obj = Mesh(devices, axis_names=(FSDP_AXIS, EP_AXIS)) + + num_procs = jax.process_count() + max_tokens_per_rank = (BATCH // num_procs) * SEQ + recv_capacity_per_rank = _compute_worst_case_recv_pr() + + # Eager bootstrap: ep_bootstrap does a host-side NCCL UID allgather + # and cannot run from inside jax.jit. Sized to the worst-case recv_pr + # across _CONFIGS so every parametrized config is bootstrap-compatible. + with mesh_obj, global_shard_guard(MeshResource(ep_resource=EP_AXIS, fsdp_resource=FSDP_AXIS)): + ep_bootstrap( + world_size=num_procs, + rank=jax.process_index(), + num_experts=NUM_EXPERTS, + max_tokens_per_rank=max_tokens_per_rank, + recv_capacity_per_rank=recv_capacity_per_rank, + hidden_dim=HIDDEN, + max_token_dtype=DTYPE, + ) + record_ep_bootstrap_signature_for_moe( + num_experts=NUM_EXPERTS, + max_tokens_per_rank=max_tokens_per_rank, + recv_capacity_per_rank=recv_capacity_per_rank, + hidden_dim=HIDDEN, + ep_size=EP_SIZE, + ) + return mesh_obj + + +# ----------------------------------------------------------------------------- +# Pure-JAX reference MoE (no EP). Mirrors the exact math of TE's fused +# router primitive (see tests/jax/test_fused_router.py for the same +# reference applied to the standalone router kernel): +# +# softmax + post-softmax (use_pre_softmax=False, the default): +# 1. top_k by raw logits +# 2. softmax over just the K selected logits (so weights sum to 1) +# +# sigmoid + optional expert_bias: +# 1. scores = sigmoid(logits) +# 2. top_k by (scores + expert_bias) [bias only steers selection] +# 3. weights = scores at top_k positions, normalized when K > 1 +# +# Then for both: +# * weights *= scaling_factor (we leave scaling_factor=1.0 in this +# suite, matching _make_block's default). +# * per-expert FFN: silu(layer_w0) * layer_w1 → wo. +# ----------------------------------------------------------------------------- + + +@partial( + jax.jit, + static_argnames=( + "num_experts", + "num_experts_per_tok", + "aux_loss_coeff", + "score_function", + ), +) +def _pure_jax_moe_reference( + x, + gate_kernel, + wi_0, + wi_1, + wo, + expert_bias=None, + *, + num_experts, + num_experts_per_tok, + aux_loss_coeff: float = 0.0, + score_function: str = "softmax", +): + B, S, H = x.shape + T = B * S + K = num_experts_per_tok + x_2d = x.reshape(T, H) + + gate_kernel_cast = gate_kernel.astype(x.dtype) + logits = (x_2d @ gate_kernel_cast).astype(jnp.float32) # [T, E] + + if score_function == "softmax": + # use_pre_softmax=False: topk on raw logits, then softmax over K. + top_logits, top_indices = jax.lax.top_k(logits, k=K) + weights = jax.nn.softmax(top_logits, axis=-1) # [T, K], sums to 1 + elif score_function == "sigmoid": + scores = jax.nn.sigmoid(logits) # [T, E] + if expert_bias is not None and expert_bias.shape != (0,): + scores_for_routing = scores + expert_bias.astype(jnp.float32)[None, :] + _, top_indices = jax.lax.top_k(scores_for_routing, k=K) + weights = jnp.take_along_axis(scores, top_indices, axis=-1) + else: + weights, top_indices = jax.lax.top_k(scores, k=K) + # Sigmoid weights are normalized when K > 1 (matches the kernel). + if K > 1: + weights = weights / (weights.sum(axis=-1, keepdims=True) + 1e-20) + else: + raise ValueError(f"Unsupported score_function={score_function!r}") + + routing_weights_full = jnp.zeros((T, num_experts), dtype=jnp.float32) + routing_weights_full = routing_weights_full.at[jnp.arange(T)[:, None], top_indices].set(weights) + + # FFN. ``apply_topk_weights_early`` is a fusion knob that doesn't + # change the math (wo is linear), so the reference is identical for + # both placements. + layer_w0 = jnp.einsum("th,ehm->tem", x_2d, wi_0) + layer_w1 = jnp.einsum("th,ehm->tem", x_2d, wi_1) + # Activation runs in x.dtype (typically bf16) to mirror the impl -- + # the impl keeps silu+multiply in the wi GEMM output dtype because + # storing higher precision than the consumer (wo) GEMM buys nothing. + intermediate = jax.nn.silu(layer_w0) * layer_w1 + expert_out = jnp.einsum("tem,emh->teh", intermediate, wo) # [T, E, H] + output_2d = jnp.einsum("te,teh->th", routing_weights_full.astype(x.dtype), expert_out) + output = output_2d.reshape(B, S, H).astype(x.dtype) + + if aux_loss_coeff > 0.0: + # tex.fused_moe_aux_loss formula (matches the same + # reference_aux_loss helper from test_fused_router.py). The + # "aux scores" use the same score_function but always with + # K-normalised sigmoid (when sigmoid) / plain softmax (when + # softmax) — see tex.fused_topk_with_score_function_fwd with + # compute_aux_scores=True. + if score_function == "softmax": + aux_scores = jax.nn.softmax(logits, axis=-1) + else: # sigmoid + aux_scores = jax.nn.sigmoid(logits) + if K > 1: + aux_scores = aux_scores / (aux_scores.sum(axis=-1, keepdims=True) + 1e-20) + routing_map = (routing_weights_full > 0).astype(jnp.int32) + tokens_per_expert = jnp.sum(routing_map, axis=0) # [E] + sum_probs_per_expert = jnp.sum(aux_scores, axis=0) # [E] + aux_loss = (num_experts * aux_loss_coeff / (K * (T**2))) * jnp.sum( + sum_probs_per_expert * tokens_per_expert.astype(jnp.float32) + ) + aux_loss = aux_loss.astype(x.dtype) + else: + aux_loss = jnp.zeros((), dtype=x.dtype) + return output, aux_loss + + +# ----------------------------------------------------------------------------- +# Helpers +# ----------------------------------------------------------------------------- + + +def _make_block( + *, + apply_topk_weights_early=False, + aux_loss_coeff=0.0, + use_expert_routing_bias=False, + score_function="softmax", + expert_bias_init=None, +): + kwargs = dict( + num_experts=NUM_EXPERTS, + num_experts_per_tok=TOPK, + intermediate_size=INTER, + data_parallelism_axes=(FSDP_AXIS,), + apply_topk_weights_early=apply_topk_weights_early, + aux_loss_coeff=aux_loss_coeff, + use_expert_routing_bias=use_expert_routing_bias, + score_function=score_function, + dtype=DTYPE, + ) + # Custom expert_bias_init lets tests inject a non-zero expert_bias without + # poking variables['params'] post-init. + if expert_bias_init is not None: + kwargs["expert_bias_init"] = expert_bias_init + return MoEBlock(**kwargs) + + +def _strong_expert_bias_init(key, shape, dtype): + """Half +5, half -5 — large enough to force topk onto the +ve half.""" + del key + n = shape[0] + return jnp.concatenate( + [ + jnp.full((n // 2,), 5.0, dtype=dtype), + jnp.full((n - n // 2,), -5.0, dtype=dtype), + ] + ) + + +def _shard_inputs(x, mesh): + # Match the layout moe.py re-pins to: outer dp axes, then ep innermost. + return jax.lax.with_sharding_constraint( + x, NamedSharding(mesh, P((FSDP_AXIS, EP_AXIS), None, None)) + ) + + +def _ctx(mesh): + """Combined mesh + global_shard_guard + axis_rules context.""" + + class _Combo: + def __enter__(self_inner): + self_inner._m = mesh.__enter__() + self_inner._gs = global_shard_guard( + MeshResource(ep_resource=EP_AXIS, fsdp_resource=FSDP_AXIS) + ) + self_inner._gs.__enter__() + self_inner._ar = nn_partitioning.axis_rules(LOGICAL_AXIS_RULES) + self_inner._ar.__enter__() + return self_inner._m + + def __exit__(self_inner, *args): + self_inner._ar.__exit__(*args) + self_inner._gs.__exit__(*args) + mesh.__exit__(*args) + + return _Combo() + + +def _init_apply(block, mesh, x, key): + with _ctx(mesh): + x_sh = _shard_inputs(x, mesh) + variables = jax.jit(block.init)(key, x_sh) + jax.block_until_ready(jax.tree_util.tree_leaves(variables)[0]) + output, aux = jax.jit(block.apply)(variables, x_sh) + jax.block_until_ready(output) + return variables, output, aux + + +def _grad_step(block, variables, mesh, x, *, include_aux=False): + """Run jax.grad of mean(out^2) [+ aux if include_aux] vs (params, x). + + Returns ``(grads_variables, grad_x)`` so callers can check both the + weight gradients and the input-activation gradient that propagates + back to the previous layer. + """ + with _ctx(mesh): + x_sh = _shard_inputs(x, mesh) + + def loss_fn(variables, x): + output, aux = block.apply(variables, x) + loss = jnp.mean(output.astype(jnp.float32) ** 2) + if include_aux and aux is not None: + loss = loss + aux.astype(jnp.float32) + return loss + + grads_v, grad_x = jax.jit(jax.grad(loss_fn, argnums=(0, 1)))(variables, x_sh) + jax.block_until_ready(jax.tree_util.tree_leaves(grads_v)[0]) + jax.block_until_ready(grad_x) + return grads_v, grad_x + + +def _grad_aux_only(block, variables, mesh, x): + """Jit'd grad of just the aux loss scalar — proves it reaches the + gate even when no main-output contribution is present.""" + with _ctx(mesh): + x_sh = _shard_inputs(x, mesh) + + def aux_only(variables, x): + _, aux = block.apply(variables, x) + return aux.astype(jnp.float32) + + grads = jax.jit(jax.grad(aux_only))(variables, x_sh) + jax.block_until_ready(jax.tree_util.tree_leaves(grads)[0]) + return grads + + +def _unwrap(x): + return x.value if hasattr(x, "value") else x + + +def _to_global_numpy(arr, mesh): + """Replicate a sharded JAX array onto every rank and return as numpy. + + Triggers an all-gather inside JIT. The resulting addressable_data(0) + contains the full global array on every process, so we can run the + pure-JAX reference and compare against it from any process. + """ + rep = NamedSharding(mesh, P()) + with mesh: + full = jax.jit(lambda a: jax.lax.with_sharding_constraint(a, rep))(arr) + full.block_until_ready() + return np.asarray(jax.device_get(full.addressable_data(0))) + + +def _params_global_numpy(variables, mesh): + """Pull every entry of variables['params'] to a replicated numpy array.""" + params = variables["params"] + return {name: _to_global_numpy(_unwrap(p), mesh) for name, p in params.items()} + + +def _make_inputs(key): + """Generate a globally-identical input tensor on every process.""" + return jax.random.normal(key, (BATCH, SEQ, HIDDEN), dtype=DTYPE) + + +# ----------------------------------------------------------------------------- +# Tests +# ----------------------------------------------------------------------------- + + +# ----------------------------------------------------------------------------- +# Parametrize variants exercised by both the forward and the backward +# parity tests. Each config is one MoE-block configuration the suite +# wants covered; the test body checks shape, dtype, finiteness AND +# numerical parity vs the same pure-JAX reference (which understands +# the same set of knobs). +# ----------------------------------------------------------------------------- + +_CONFIGS = [ + pytest.param( + dict(score_function="softmax"), + id="softmax", + ), + pytest.param( + dict(score_function="softmax", apply_topk_weights_early=True), + id="softmax-early-weighting", + ), + pytest.param( + dict(score_function="sigmoid"), + id="sigmoid", + ), + # NOTE: a ``sigmoid-bias-zero`` config (use_expert_routing_bias=True + # with a zero-initialised bias buffer) was previously exercised + # here. It was dropped because the routing math collapses to the + # no-bias case when the buffer is zero -- ``sigmoid`` already + # covers that numerical path. The bias-aware codepath is still + # exercised by ``sigmoid-bias-strong`` below, which uses a + # non-zero bias. + pytest.param( + dict( + score_function="sigmoid", + use_expert_routing_bias=True, + expert_bias_init=_strong_expert_bias_init, + ), + id="sigmoid-bias-strong", + ), +] + + +def _reference_kwargs_from_config(config, params_np): + """Pick out the reference-relevant pieces of a parametrize config.""" + return dict( + score_function=config.get("score_function", "softmax"), + expert_bias=( + jnp.asarray(params_np["expert_bias"]) + if config.get("use_expert_routing_bias", False) + else None + ), + ) + + +class TestTeEpMoeForward: + """Per-config forward correctness in a single run: shape, dtype, + finiteness AND numerical parity vs the pure-JAX reference.""" + + @pytest.mark.parametrize("config", _CONFIGS) + def test_forward(self, mesh, config): + block = _make_block(**config) + x = _make_inputs(jax.random.PRNGKey(0)) + variables, output, aux = _init_apply(block, mesh, x, jax.random.PRNGKey(1)) + + # Shape / dtype / finiteness (cheap; on the local shard). + assert output.shape == x.shape + assert output.dtype == x.dtype + out_local = np.asarray(jax.device_get(output.addressable_data(0))) + assert np.all(np.isfinite(out_local)), "output has NaN/Inf" + assert aux is None, "aux_loss should be None when aux_loss_coeff == 0" + + # Numerical parity (replicated global view -> single rank's numpy). + params_np = _params_global_numpy(variables, mesh) + x_np = np.asarray(jax.device_get(x)) + out_te_np = _to_global_numpy(output, mesh) + + out_ref, _ = _pure_jax_moe_reference( + jnp.asarray(x_np), + jnp.asarray(params_np["gate_kernel"]), + jnp.asarray(params_np["wi_0"]), + jnp.asarray(params_np["wi_1"]), + jnp.asarray(params_np["wo"]), + num_experts=NUM_EXPERTS, + num_experts_per_tok=TOPK, + **_reference_kwargs_from_config(config, params_np), + ) + np.testing.assert_allclose( + out_te_np.astype(np.float32), + np.asarray(jax.device_get(out_ref)).astype(np.float32), + atol=FWD_ATOL, + rtol=FWD_RTOL, + err_msg=f"forward parity breach for config={config}", + ) + + +class TestTeEpMoeBackward: + """Per-config backward correctness in a single run: per-tensor + grads finite, non-zero AND parity vs the pure-JAX reference.""" + + @pytest.mark.parametrize("config", _CONFIGS) + def test_backward(self, mesh, config): + block = _make_block(**config) + x = _make_inputs(jax.random.PRNGKey(2)) + variables, _, _ = _init_apply(block, mesh, x, jax.random.PRNGKey(3)) + grads_te, grad_x_te = _grad_step(block, variables, mesh, x) + + # Reference grads via jax.grad over the pure-JAX MoE with the + # same config. argnums=(0, 1) so the reference also produces a + # d_x for the propagated-gradient parity check below. + params_np = _params_global_numpy(variables, mesh) + x_np = np.asarray(jax.device_get(x)) + ref_kwargs = _reference_kwargs_from_config(config, params_np) + ref_expert_bias = ref_kwargs.pop("expert_bias") + + def loss_fn(params, x): + out, _ = _pure_jax_moe_reference( + x, + params["gate_kernel"], + params["wi_0"], + params["wi_1"], + params["wo"], + ref_expert_bias, + num_experts=NUM_EXPERTS, + num_experts_per_tok=TOPK, + **ref_kwargs, + ) + return jnp.mean(out.astype(jnp.float32) ** 2) + + grads_ref, grad_x_ref = jax.jit(jax.grad(loss_fn, argnums=(0, 1)))( + {k: jnp.asarray(v) for k, v in params_np.items() if k != "expert_bias"}, + jnp.asarray(x_np), + ) + grads_ref_np = {k: np.asarray(jax.device_get(v)) for k, v in grads_ref.items()} + grad_x_ref_np = np.asarray(jax.device_get(grad_x_ref)) + + for name in ("gate_kernel", "wi_0", "wi_1", "wo"): + # Per-tensor: finite + non-zero + parity in one pass. + g_te = _to_global_numpy(_unwrap(grads_te["params"][name]), mesh) + assert np.all(np.isfinite(g_te)), f"{name} grad has NaN/Inf [config={config}]" + assert np.any(g_te != 0.0), f"{name} grad identically zero [config={config}]" + atol, rtol = ( + (GRAD_GATE_ATOL, GRAD_GATE_RTOL) + if name == "gate_kernel" + else (GRAD_FFN_ATOL, GRAD_FFN_RTOL) + ) + np.testing.assert_allclose( + g_te.astype(np.float32), + grads_ref_np[name].astype(np.float32), + atol=atol, + rtol=rtol, + err_msg=f"grad parity breach on {name} [config={config}]", + ) + + # d_x: the gradient propagated back to the previous layer. Checks + # shape, dtype (must match x.dtype — protects the + # _with_sharding_constraint_cast_bwd wrapper that casts the + # fp32-promoted gate path back to bf16), finiteness, non-zero + # AND numerical parity vs the pure-JAX reference d_x. + grad_x_te_np = _to_global_numpy(grad_x_te, mesh) + assert ( + grad_x_te.shape == x.shape + ), f"d_x shape {grad_x_te.shape} != x.shape {x.shape} [config={config}]" + assert ( + grad_x_te.dtype == x.dtype + ), f"d_x dtype {grad_x_te.dtype} != x.dtype {x.dtype} [config={config}]" + assert np.all(np.isfinite(grad_x_te_np)), f"d_x has NaN/Inf [config={config}]" + assert np.any(grad_x_te_np != 0.0), f"d_x identically zero [config={config}]" + np.testing.assert_allclose( + grad_x_te_np.astype(np.float32), + grad_x_ref_np.astype(np.float32), + atol=GRAD_FFN_ATOL, + rtol=GRAD_FFN_RTOL, + err_msg=f"d_x parity breach [config={config}]", + ) + + +class TestTeEpMoeAuxLoss: + """Aux-loss path. Consolidated into: + * ``test_aux_loss``: one run that checks the returned scalar's + shape / dtype / finiteness / magnitude AND numerical parity vs the + reference AND that the aux-only bwd propagates to gate_kernel. + * ``test_combined_loss_grads``: one run for joint main+aux bwd + finite + non-zero per tensor. + """ + + def test_aux_loss(self, mesh): + coeff = 1e-2 + block = _make_block(aux_loss_coeff=coeff) + x = _make_inputs(jax.random.PRNGKey(20)) + variables, _, aux = _init_apply(block, mesh, x, jax.random.PRNGKey(21)) + + # Shape / dtype / finiteness / magnitude. + assert aux is not None, "aux_loss should be returned when coeff > 0" + assert aux.shape == (), f"aux_loss must be 0-d scalar, got {aux.shape}" + assert aux.dtype == DTYPE, f"aux_loss dtype {aux.dtype} != {DTYPE}" + aux_np = _to_global_numpy(aux, mesh) + assert np.isfinite(aux_np), "aux_loss is NaN/Inf" + assert abs(float(aux_np)) < 1e2, f"aux_loss looks unreasonable: {aux_np}" + + # Numerical parity vs the reference. + params_np = _params_global_numpy(variables, mesh) + x_np = np.asarray(jax.device_get(x)) + _, aux_ref = _pure_jax_moe_reference( + jnp.asarray(x_np), + jnp.asarray(params_np["gate_kernel"]), + jnp.asarray(params_np["wi_0"]), + jnp.asarray(params_np["wi_1"]), + jnp.asarray(params_np["wo"]), + num_experts=NUM_EXPERTS, + num_experts_per_tok=TOPK, + aux_loss_coeff=coeff, + ) + np.testing.assert_allclose( + float(aux_np), + float(jax.device_get(aux_ref)), + atol=AUX_ATOL, + rtol=AUX_RTOL, + ) + + # Aux-only bwd must propagate to gate_kernel — proves the + # fused_moe_aux_loss_bwd → topk(compute_aux_scores)_bwd chain is + # wired. + aux_grads = _grad_aux_only(block, variables, mesh, x) + g_gate = np.asarray( + jax.device_get(_unwrap(aux_grads["params"]["gate_kernel"]).addressable_data(0)) + ) + assert np.all(np.isfinite(g_gate)), "gate grad NaN/Inf under aux-only loss" + assert np.any(g_gate != 0.0), "aux bwd should propagate to gate_kernel" + + def test_combined_loss_grads(self, mesh): + """Joint main + aux loss bwd: per-tensor finite + non-zero in + one pass.""" + block = _make_block(aux_loss_coeff=1e-2) + x = _make_inputs(jax.random.PRNGKey(22)) + variables, _, _ = _init_apply(block, mesh, x, jax.random.PRNGKey(23)) + grads, _ = _grad_step(block, variables, mesh, x, include_aux=True) + for name in ("gate_kernel", "wi_0", "wi_1", "wo"): + g_local = np.asarray(jax.device_get(_unwrap(grads["params"][name]).addressable_data(0))) + assert np.all(np.isfinite(g_local)), f"{name} grad NaN/Inf under main+aux" + assert np.any(g_local != 0.0), f"{name} grad zero under main+aux" diff --git a/tests/pytorch/attention/test_attention.py b/tests/pytorch/attention/test_attention.py index f2ab4a8495..7bcdd10daf 100644 --- a/tests/pytorch/attention/test_attention.py +++ b/tests/pytorch/attention/test_attention.py @@ -132,6 +132,7 @@ def test_dot_product_attention( qkv_layout, swa, pad_between_seqs, + declarative_packed=False, ): """Test DotProductAttention module""" @@ -149,6 +150,8 @@ def test_dot_product_attention( qkv_layout = "bshd_bs2hd" if not is_mla and not is_mqa_gqa else "bshd_bshd_bshd" if "3" in qkv_layout and config.attn_type == "cross": pytest.skip("No need to test this layout for cross attention") + if declarative_packed and not any(c.isdigit() for c in qkv_layout): + pytest.skip("Declarative packed inputs only apply to packed qkv layouts.") if config.window_size == (-1, -1) and swa: config.window_size = [2, 2] @@ -221,6 +224,7 @@ def test_dot_product_attention( qkv_layout, pad_between_seqs, is_training, + declarative_packed=declarative_packed, ) # FlashAttention backend @@ -233,6 +237,7 @@ def test_dot_product_attention( qkv_layout, pad_between_seqs, is_training, + declarative_packed=declarative_packed, ) # Compare results @@ -924,6 +929,26 @@ def test_dpa_qkv_layout(dtype, model_configs, model, qkv_layout): test_dot_product_attention(dtype, model_configs, model, False, qkv_layout, False, False) +qkv_layouts_packed = [l for l in qkv_layouts if any(c.isdigit() for c in l)] + + +@pytest.mark.skipif(get_cudnn_version() < (8, 9, 5), reason="cuDNN 8.9.5+ is required.") +@pytest.mark.parametrize("dtype", param_types_lean) +@pytest.mark.parametrize("model_configs", [model_configs_layout]) +@pytest.mark.parametrize("model", ["layout_1_1", "layout_1_2"]) +@pytest.mark.parametrize("qkv_layout", qkv_layouts_packed) +def test_dpa_qkv_layout_declarative(dtype, model_configs, model, qkv_layout): + """Declarative packed inputs: the packed buffer is passed to + DotProductAttention via qkv_layer/kv_layer (declared layout, gradients read + off the packed buffer) instead of q/k/v views + pointer-based detection. + Layout coverage is complete; the model-config dimension is trimmed to one + self-attention and one cross-attention config, since past the input + handling the backend code is identical to test_dpa_qkv_layout.""" + test_dot_product_attention( + dtype, model_configs, model, False, qkv_layout, False, False, declarative_packed=True + ) + + qkv_layouts_thd = ["t3hd", "th3d", "thd_t2hd", "thd_th2d", "thd_thd_thd"] model_configs_layout_thd = { # test: ModelConfig(b, sq, hq, dqk) @@ -988,7 +1013,7 @@ def test_dpa_qkv_layout(dtype, model_configs, model, qkv_layout): @pytest.mark.parametrize("model_configs", [model_configs_layout_thd]) @pytest.mark.parametrize("model", model_configs_layout_thd.keys()) @pytest.mark.parametrize("qkv_layout", qkv_layouts_thd) -def test_dpa_qkv_layout_thd(dtype, model_configs, model, qkv_layout): +def test_dpa_qkv_layout_thd(dtype, model_configs, model, qkv_layout, declarative_packed=False): """Test DotProductAttention module with different QKV layouts""" config = model_configs[model] if config.num_heads != config.num_gqa_groups and "3" in qkv_layout: @@ -996,17 +1021,47 @@ def test_dpa_qkv_layout_thd(dtype, model_configs, model, qkv_layout): logging.info("[test_dpa_qkv_layout_thd]: pad_between_seqs = True") pad_between_seqs = True test_dot_product_attention( - dtype, model_configs, model, False, qkv_layout, False, pad_between_seqs + dtype, + model_configs, + model, + False, + qkv_layout, + False, + pad_between_seqs, + declarative_packed=declarative_packed, ) if get_cudnn_version() >= (9, 3, 0): logging.info("[test_dpa_qkv_layout_thd]: pad_between_seqs = False") # cuDNN 9.3.0+ is required to run pad_between_seqs = False/True in the same run pad_between_seqs = False test_dot_product_attention( - dtype, model_configs, model, False, qkv_layout, False, pad_between_seqs + dtype, + model_configs, + model, + False, + qkv_layout, + False, + pad_between_seqs, + declarative_packed=declarative_packed, ) +qkv_layouts_thd_packed = [l for l in qkv_layouts_thd if any(c.isdigit() for c in l)] + + +@pytest.mark.skipif(get_cudnn_version() < (9, 0, 0), reason="cuDNN 9.0.0+ is required.") +@pytest.mark.skipif( + get_device_compute_capability() < (9, 0), reason="THD is only supported on Hopper+." +) +@pytest.mark.parametrize("dtype", param_types_lean) +@pytest.mark.parametrize("model_configs", [model_configs_layout_thd]) +@pytest.mark.parametrize("model", ["layout_0_0"]) +@pytest.mark.parametrize("qkv_layout", qkv_layouts_thd_packed) +def test_dpa_qkv_layout_thd_declarative(dtype, model_configs, model, qkv_layout): + """Declarative packed thd inputs, see test_dpa_qkv_layout_declarative.""" + test_dpa_qkv_layout_thd(dtype, model_configs, model, qkv_layout, declarative_packed=True) + + def _run_dot_product_attention( dtype: torch.dtype, config: ModelConfig, @@ -1015,8 +1070,14 @@ def _run_dot_product_attention( qkv_layout: str, pad_between_seqs: bool, is_training: bool, + declarative_packed: bool = False, ) -> Tuple[torch.Tensor, Tuple[torch.Tensor, torch.Tensor, torch.Tensor]]: - """Run DotProductAttention module with one forward pass and one backward pass""" + """Run DotProductAttention module with one forward pass and one backward pass. + + With declarative_packed=True (packed qkv_layout only), the packed buffer is + passed to DotProductAttention directly via qkv_layer/kv_layer instead of + slicing it into q/k/v views, and input gradients are read off the packed + buffer itself.""" # Set RNG and environment varables reset_rng_states() os.environ["NVTE_FLASH_ATTN"] = "0" @@ -1213,6 +1274,12 @@ def _run_dot_product_attention( tensor_count = int(l) split_dim = dim break + if declarative_packed and split_dim != 0: + # The packed buffer is the autograd leaf; q/k/v below are non-leaf + # views of it, and DPA receives the buffer via qkv_layer/kv_layer. + tensor.requires_grad_() + packed_tensor = tensor + packed_interleave_dim = split_dim - tensor.dim() tensors = torch.split(tensor, 1, dim=split_dim) if split_dim != 0 else [tensor] tensors_orig = ( torch.split(tensor_orig, 1, dim=split_dim) if split_dim != 0 else [tensor_orig] @@ -1225,8 +1292,10 @@ def _run_dot_product_attention( inp.append(tensors[j]) inp_orig.append(tensors_orig[j]) for i in range(3): - inp[i].requires_grad = True - inp_orig[i].requires_grad = True + if inp[i].is_leaf: + inp[i].requires_grad = True + if inp_orig[i].is_leaf: + inp_orig[i].requires_grad = True # Create output gradient qkv_format_kv = "_".join(qkv_format) @@ -1308,10 +1377,21 @@ def get_dummy_cuda_rng_tracker() -> CudaRNGStatesTracker: k = inp[1] v = inp[2] d_out = out_grad + packed_kwargs = {} + if declarative_packed: + assert backend in ["FusedAttention", "FlashAttention"] + packed_kwargs["qkv_interleave_dim"] = packed_interleave_dim + if len(qkv_layout.split("_")) == 1: + packed_kwargs["qkv_layer"] = packed_tensor + q, k, v = None, None, None + else: + packed_kwargs["kv_layer"] = packed_tensor + k, v = None, None out = block( q, k, v, + **packed_kwargs, window_size=config.window_size, attention_mask=attention_mask, qkv_format=qkv_format, @@ -1341,15 +1421,33 @@ def get_dummy_cuda_rng_tracker() -> CudaRNGStatesTracker: if is_training: out.backward(d_out) + q_grad, k_grad, v_grad = None, None, None + if is_training: + if declarative_packed: + # Input gradients live on the packed buffer; slice them back out so + # the cross-backend comparisons below stay uniform with the + # separate-q/k/v path. + assert ( + packed_tensor.grad is not None and packed_tensor.grad.shape == packed_tensor.shape + ) + packed_grads = [ + packed_tensor.grad.select(packed_interleave_dim, j) + for j in range(packed_tensor.shape[packed_interleave_dim]) + ] + if len(qkv_layout.split("_")) == 1: + q_grad, k_grad, v_grad = packed_grads + else: + q_grad = q.grad + k_grad, v_grad = packed_grads + else: + q_grad, k_grad, v_grad = q.grad, k.grad, v.grad + d_softmax_offset = None if is_training and config.softmax_type != "vanilla": d_softmax_offset = block.softmax_offset.grad if backend in ["UnfusedDotProductAttention"]: - if is_training: - return out, max_logit, (q.grad, k.grad, v.grad, d_softmax_offset) - else: - return out, max_logit, (None, None, None, d_softmax_offset) + return out, max_logit, (q_grad, k_grad, v_grad, d_softmax_offset) if backend in ["FusedAttention", "FlashAttention"]: if qkv_format == "thd" and pad_between_seqs: out_orig = torch.Tensor([]).to(device="cuda", dtype=dtype) @@ -1369,13 +1467,13 @@ def get_dummy_cuda_rng_tracker() -> CudaRNGStatesTracker: out_orig = torch.cat([out_orig, out[valid_range_q[0] : valid_range_q[1]]], dim=0) if is_training: q_grad_orig = torch.cat( - [q_grad_orig, q.grad[valid_range_q[0] : valid_range_q[1]]], dim=0 + [q_grad_orig, q_grad[valid_range_q[0] : valid_range_q[1]]], dim=0 ) k_grad_orig = torch.cat( - [k_grad_orig, k.grad[valid_range_kv[0] : valid_range_kv[1]]], dim=0 + [k_grad_orig, k_grad[valid_range_kv[0] : valid_range_kv[1]]], dim=0 ) v_grad_orig = torch.cat( - [v_grad_orig, v.grad[valid_range_kv[0] : valid_range_kv[1]]], dim=0 + [v_grad_orig, v_grad[valid_range_kv[0] : valid_range_kv[1]]], dim=0 ) if is_training: return ( @@ -1386,10 +1484,7 @@ def get_dummy_cuda_rng_tracker() -> CudaRNGStatesTracker: else: return out_orig, max_logit, (None, None, None, d_softmax_offset) else: - if is_training: - return out, max_logit, (q.grad, k.grad, v.grad, d_softmax_offset) - else: - return out, max_logit, (None, None, None, d_softmax_offset) + return out, max_logit, (q_grad, k_grad, v_grad, d_softmax_offset) model_configs_te_layer = { diff --git a/tests/pytorch/distributed/run_ep.py b/tests/pytorch/distributed/run_ep.py index ee6a97ffea..534fd9642a 100644 --- a/tests/pytorch/distributed/run_ep.py +++ b/tests/pytorch/distributed/run_ep.py @@ -19,6 +19,7 @@ ep_dispatch, ep_combine, symm_mem_alloc, + is_symm_backed, _ep_combine_raw, _ep_dispatch_raw, ) @@ -155,13 +156,7 @@ def setUp(self): ): self.skipTest("not exercised in zero-copy mode") - def _make_buffer( - self, - alignment=0, - top_k=TOP_K, - dispatch_recv_tokens=None, - combine_grad_expert_out=None, - ): + def _make_buffer(self, alignment=0, top_k=TOP_K): return EpBuffer( top_k=top_k, max_tokens_per_rank=TOKENS_PER_RANK, @@ -169,8 +164,6 @@ def _make_buffer( hidden_dim=HIDDEN_DIM, num_local_experts=NUM_LOCAL_EXPERTS, alignment=alignment, - dispatch_recv_tokens=dispatch_recv_tokens, - combine_grad_expert_out=combine_grad_expert_out, ) def _expert_out(self, expert_out): @@ -253,10 +246,10 @@ def test_dispatch_autograd(self): ] for label, recv_tokens in cases: with self.subTest(case=label): - buf = self._make_buffer(dispatch_recv_tokens=recv_tokens) + buf = self._make_buffer() topk_idx, tokens, w = _make_identity_inputs(self.cfg.rank, self.cfg.ep_size) tokens_p = tokens.detach().clone().requires_grad_(True) - rt, rw, _tc = ep_dispatch(buf, tokens_p, topk_idx, w) + rt, rw, _tc = ep_dispatch(buf, tokens_p, topk_idx, w, recv_tokens=recv_tokens) if recv_tokens is not None: # caller-supplied recv_tokens must be used in place self.assertEqual(rt.data_ptr(), recv_tokens.data_ptr()) rt = self._stage_grad_symm(rt) @@ -269,20 +262,17 @@ def test_dispatch_autograd(self): @_zero_copy_test_include def test_caller_provides_dispatch_recv_tokens(self): - """Caller-supplied recv_tokens: EpBuffer adopts it (recv_topk_weights stays - owned) and ep_dispatch returns a view of the caller's buffer.""" + """Caller-supplied recv_tokens (symm-mem-backed in zero-copy): ep_dispatch + writes into it and returns a view of the caller's buffer.""" if ZERO_COPY: rc = self.cfg.recv_capacity_per_rank rt_buf = symm_mem_alloc((rc, HIDDEN_DIM), torch.bfloat16, self.ep_group) else: rt_buf, _rw_buf, _ = self._make_raw_recv() - buf = self._make_buffer(dispatch_recv_tokens=rt_buf) - self.assertEqual(buf.recv_tokens_symm_buf.data_ptr(), rt_buf.data_ptr()) - if ZERO_COPY: # recv_topk_weights is always buffer-owned in zero-copy - self.assertIsNotNone(buf.recv_topk_weights_symm_buf) + buf = self._make_buffer() topk_idx, tokens, w = _make_identity_inputs(self.cfg.rank, self.cfg.ep_size) tokens_p = tokens.detach().clone().requires_grad_(True) - rt, rw, _ = ep_dispatch(buf, tokens_p, topk_idx, w) + rt, rw, _ = ep_dispatch(buf, tokens_p, topk_idx, w, recv_tokens=rt_buf) self.assertEqual(rt.data_ptr(), rt_buf.data_ptr()) rt = self._stage_grad_symm(rt) rw = self._stage_grad_symm(rw) @@ -294,22 +284,45 @@ def test_caller_provides_dispatch_recv_tokens(self): @_zero_copy_test_include def test_caller_provides_grad_expert_out(self): - """Caller-supplied grad_expert_out: EpBuffer adopts it as the combine - backward grad target (symm-mem under zero-copy).""" + """Caller-supplied grad_out (symm-mem-backed in zero-copy): ep_combine's + backward scatters the expert-out grad into it.""" rc = self.cfg.recv_capacity_per_rank if ZERO_COPY: gbuf = symm_mem_alloc((rc, HIDDEN_DIM), torch.bfloat16, self.ep_group) else: gbuf = torch.empty(rc, HIDDEN_DIM, dtype=torch.bfloat16, device=self.cfg.device) - buf = self._make_buffer(combine_grad_expert_out=gbuf) - self.assertEqual(buf.grad_expert_out_symm_buf.data_ptr(), gbuf.data_ptr()) + gbuf.zero_() + buf = self._make_buffer() topk_idx, tokens, w = _make_identity_inputs(self.cfg.rank, self.cfg.ep_size) tokens_p = tokens.detach().clone().requires_grad_(True) recv_t, recv_w, _ = ep_dispatch(buf, tokens_p, topk_idx, w) recv_t = self._stage_grad_symm(recv_t) recv_w = self._stage_grad_symm(recv_w) expert_out = self._expert_out(self._weighted(recv_t, recv_w)) - out = ep_combine(buf, expert_out) + out = ep_combine(buf, expert_out, grad_out=gbuf) + (0.5 * (out.float() ** 2).sum()).backward() + torch.cuda.synchronize() + torch.testing.assert_close(out.float(), tokens.float(), atol=5e-2, rtol=5e-2) + torch.testing.assert_close(tokens_p.grad.float(), tokens.float(), atol=5e-2, rtol=5e-2) + # the caller-owned buffer was used as the combine-bwd scatter target + self.assertGreater(gbuf.abs().sum().item(), 0.0) + + @_zero_copy_test_include + def test_zero_copy_pool_auto_alloc(self): + """Zero-copy with recv/grad left None: ep_dispatch/ep_combine allocate their IO + tensors from the symm-mem pool (is_symm_backed). This is the primary mcore + path — mcore hands no caller buffers, TE pools them on the fly.""" + if not ZERO_COPY: + self.skipTest("zero-copy pool auto-alloc only") + buf = self._make_buffer() + topk_idx, tokens, w = _make_identity_inputs(self.cfg.rank, self.cfg.ep_size) + tokens_p = tokens.detach().clone().requires_grad_(True) + recv_t, recv_w, _ = ep_dispatch(buf, tokens_p, topk_idx, w) # recv_tokens=None -> pool + self.assertTrue(is_symm_backed(recv_t)) # dispatch recv came from the symm-mem pool + recv_t = self._stage_grad_symm(recv_t) + recv_w = self._stage_grad_symm(recv_w) + expert_out = self._expert_out(self._weighted(recv_t, recv_w)) + out = ep_combine(buf, expert_out) # grad_out=None -> bwd allocs the grad from the pool (0.5 * (out.float() ** 2).sum()).backward() torch.cuda.synchronize() torch.testing.assert_close(out.float(), tokens.float(), atol=5e-2, rtol=5e-2) diff --git a/tests/pytorch/distributed/run_test_ep.sh b/tests/pytorch/distributed/run_test_ep.sh index 68b691f787..62c9a0207f 100755 --- a/tests/pytorch/distributed/run_test_ep.sh +++ b/tests/pytorch/distributed/run_test_ep.sh @@ -9,8 +9,6 @@ set -uo pipefail SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" -TE_REPO_ROOT="$(cd "${SCRIPT_DIR}/../../.." && pwd)" -export PYTHONPATH="${TE_REPO_ROOT}${PYTHONPATH:+:${PYTHONPATH}}" DETECTED_GPUS=$(nvidia-smi -L 2>/dev/null | wc -l) if [ "${DETECTED_GPUS}" -lt 4 ]; then diff --git a/tests/pytorch/test_distributed_weight.py b/tests/pytorch/test_distributed_weight.py new file mode 100644 index 0000000000..69b5917d90 --- /dev/null +++ b/tests/pytorch/test_distributed_weight.py @@ -0,0 +1,172 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""Unit tests for the GTP-agnostic DistributedWeight protocol and dispatchers. + +These tests exercise TE's weight-parallelism extension point in isolation, with +a tiny in-repo ``FakeDistributedWeight`` stub standing in for a real implementer +(e.g. Megatron's GTPShardedParam). No GPU, process group, or Megatron import is +required: the whole contract lives in TE and is verified against the fake. +""" + +import pytest +import torch + +from transformer_engine.pytorch.distributed_weight import ( + DistributedWeight, + is_distributed_weight, + materialize_weight_for_forward, + materialize_weight_for_backward, + finalize_weight_grads, +) + + +class FakeDistributedWeight(torch.Tensor): + """Minimal DistributedWeight implementer for dispatcher tests. + + Records how it was called and returns marker tensors so the dispatcher's + behavior (delegation, list normalization, no-op fallback) is observable. + """ + + is_distributed_weight = True + + def __new__(cls, group_size=1): + t = torch.zeros(1).as_subclass(cls) + t.group_size = group_size + t.calls = [] + return t + + def materialize_group_for_forward(self): + self.calls.append("fwd") + out = [torch.full((2, 2), float(i)) for i in range(self.group_size)] + # Match the real GTP contract: single weight returns a bare tensor. + return out if self.group_size > 1 else out[0] + + def materialize_group_for_backward(self, **kwargs): + self.calls.append(("bwd", kwargs)) + out = [torch.full((2, 2), float(10 + i)) for i in range(self.group_size)] + return out if self.group_size > 1 else out[0] + + def finalize_group_grads(self, wgrads, **kwargs): + self.calls.append(("finalize", wgrads)) + wl = wgrads if isinstance(wgrads, (list, tuple)) else [wgrads] + out = [w + 100 for w in wl] + return out if self.group_size > 1 else out[0] + + def grad_buffer(self): + return torch.full((2, 2), -1.0) + + +class FakeNonTensorWeight: + """DistributedWeight-shaped object that is NOT a torch.Tensor (contract violation).""" + + is_distributed_weight = True + + def materialize_group_for_forward(self): + return torch.zeros(2, 2) + + def materialize_group_for_backward(self, **kwargs): + return torch.zeros(2, 2) + + def finalize_group_grads(self, wgrads, **kwargs): + return wgrads + + def grad_buffer(self): + return torch.zeros(2, 2) + + +def test_protocol_runtime_checkable(): + """A conforming object passes isinstance; a plain tensor does not.""" + assert isinstance(FakeDistributedWeight(), DistributedWeight) + assert not isinstance(torch.zeros(2), DistributedWeight) + + +def test_is_distributed_weight(): + assert is_distributed_weight(FakeDistributedWeight()) + assert not is_distributed_weight(torch.zeros(2)) + assert not is_distributed_weight(torch.nn.Parameter(torch.zeros(2))) + + +def test_non_tensor_implementer_rejected(): + """Implementers must be torch.Tensor subclasses; a non-Tensor fails loudly.""" + with pytest.raises(TypeError, match="torch.Tensor subclass"): + is_distributed_weight(FakeNonTensorWeight()) + + +def test_forward_noop_on_plain_tensor(): + """Plain weights pass through unchanged — the critical non-regression.""" + w = torch.nn.Parameter(torch.randn(4, 4)) + out = materialize_weight_for_forward(w) + assert out == [w] + assert out[0] is w + + +@pytest.mark.parametrize("group_size", [1, 3]) +def test_forward_dispatches(group_size): + """Linear (N=1) and GroupedLinear (N=k): one coalesced call, full list returned.""" + w = FakeDistributedWeight(group_size=group_size) + out = materialize_weight_for_forward(w) + assert isinstance(out, list) and len(out) == group_size + # Leader is delegated to exactly once (coalesced), not once per weight. + assert w.calls == ["fwd"] + + +def test_forward_accepts_weight_list(): + """The dispatcher accepts the full per-expert list; the leader (index 0) coalesces it.""" + leader = FakeDistributedWeight(group_size=3) + followers = [torch.zeros(2, 2), torch.zeros(2, 2)] + out = materialize_weight_for_forward([leader, *followers]) + assert isinstance(out, list) and len(out) == 3 + # Leader delegated exactly once; the follower entries are not materialized separately. + assert leader.calls == ["fwd"] + + +def test_forward_noop_on_plain_weight_list(): + """A non-distributed weight list passes through unchanged (all N returned).""" + ws = [torch.nn.Parameter(torch.randn(2, 2)) for _ in range(3)] + out = materialize_weight_for_forward(ws) + assert out == ws + + +@pytest.mark.parametrize("group_size", [1, 2]) +def test_backward_dispatches(group_size): + w = FakeDistributedWeight(group_size=group_size) + out = materialize_weight_for_backward(w) + assert len(out) == group_size + assert torch.equal(out[0], torch.full((2, 2), 10.0)) + + +def test_backward_accepts_weight_list(): + """Backward dispatcher also accepts the full per-expert list; the leader coalesces it.""" + leader = FakeDistributedWeight(group_size=2) + out = materialize_weight_for_backward([leader, torch.zeros(2, 2)]) + assert isinstance(out, list) and len(out) == 2 + assert torch.equal(out[0], torch.full((2, 2), 10.0)) + + +def test_backward_noop_on_plain_tensor(): + plain = torch.zeros(2) + assert materialize_weight_for_backward(plain) == [plain] + + +def test_backward_noop_on_plain_weight_list(): + """A non-distributed weight list passes through unchanged (all N returned).""" + ws = [torch.nn.Parameter(torch.randn(2, 2)) for _ in range(3)] + assert materialize_weight_for_backward(ws) == ws + + +@pytest.mark.parametrize("group_size", [1, 2]) +def test_finalize_grads_dispatches(group_size): + w = FakeDistributedWeight(group_size=group_size) + wgrads = [torch.zeros(2, 2) for _ in range(group_size)] + out = finalize_weight_grads(w, wgrads) + assert len(out) == group_size + assert torch.equal(out[0], torch.full((2, 2), 100.0)) + + +def test_finalize_grads_noop_on_plain_tensor(): + """No-op path leaves the grads untouched.""" + plain_w = torch.nn.Parameter(torch.zeros(2)) + g = [torch.ones(2)] + assert finalize_weight_grads(plain_w, g) == g diff --git a/tests/pytorch/test_fused_optimizer.py b/tests/pytorch/test_fused_optimizer.py index a2863cba98..6832ef89dd 100644 --- a/tests/pytorch/test_fused_optimizer.py +++ b/tests/pytorch/test_fused_optimizer.py @@ -166,6 +166,19 @@ def test_frozen_model(self): torch.testing.assert_close(ref_param, tst_param) + def test_empty_param_at_end_of_group(self): + tensors = [ + torch.ones(4, dtype=torch.float, device="cuda"), + torch.empty(0, dtype=torch.float, device="cuda"), + ] + ref_param, tst_param, ref_optim, tst_optim = self.gen_param_optim(tensors, self.options) + + self.gen_grad(ref_param, tst_param) + ref_optim.step() + tst_optim.step() + + torch.testing.assert_close(ref_param, tst_param) + def gen_precision_aware_test( self, use_fp8_params, @@ -796,6 +809,19 @@ def test_float(self): def test_half(self): self.gen_single_type_test(param_type=torch.float16) + def test_empty_param_at_end_of_group(self): + tensors = [ + torch.ones(4, dtype=torch.float, device="cuda"), + torch.empty(0, dtype=torch.float, device="cuda"), + ] + ref_param, tst_param, ref_optim, tst_optim = self.gen_param_optim(tensors, self.options) + + self.gen_grad(ref_param, tst_param) + ref_optim.step() + tst_optim.step() + + torch.testing.assert_close(ref_param, tst_param) + class Model(torch.nn.Module): def __init__(self): diff --git a/tests/pytorch/test_fusible_ops.py b/tests/pytorch/test_fusible_ops.py index 3ad15b6a05..5ae657c5d2 100644 --- a/tests/pytorch/test_fusible_ops.py +++ b/tests/pytorch/test_fusible_ops.py @@ -19,6 +19,10 @@ import transformer_engine.common.recipe import transformer_engine.pytorch as te import transformer_engine.pytorch.ops as te_ops +from transformer_engine.pytorch.ops.basic.grouped_linear import ( + OUTPUT_BUFFER_KEY, + GRAD_INPUT_BUFFER_KEY, +) from transformer_engine.pytorch._extra_state import UNSAFE_PICKLE_EXTRA_STATE_ENV from transformer_engine.pytorch.ops.fused import ( @@ -2196,6 +2200,86 @@ def test_grouped_linear( if bias: assert_close_grads(getattr(op, f"bias{group_idx}"), bs_ref[group_idx], **tols) + def test_grouped_linear_caller_buffers( + self, + *, + dtype: torch.dtype = torch.bfloat16, + device: torch.device = "cuda", + ) -> None: + """Caller output/grad_input buffers routed to the last/first op of a Sequential. + + The chain has two GroupedLinears with distinct inner dims, so ``output`` + can only fit the last op's output and ``grad_input`` only the first op's + dgrad -- a mis-route would fail the buffer's shape check. + """ + group_size = 3 + in_features, hidden, out_features = 128, 256, 64 + split_sizes = torch.tensor([128, 256, 128], dtype=torch.int32, device=device) + num_tokens = int(split_sizes.sum()) + + torch.manual_seed(1234) + x = (0.1 * torch.randn(num_tokens, in_features, device=device)).to(dtype) + dy = (0.1 * torch.randn(num_tokens, out_features, device=device)).to(dtype) + + def build() -> te_ops.Sequential: + fc1 = te_ops.GroupedLinear( + group_size, in_features, hidden, bias=False, device=device, dtype=dtype + ) + fc2 = te_ops.GroupedLinear( + group_size, hidden, out_features, bias=False, device=device, dtype=dtype + ) + return te_ops.Sequential(fc1, fc2) + + # Reference: internal allocation. + model_ref = build() + x_ref = x.detach().clone().requires_grad_(True) + y_ref = model_ref(x_ref, split_sizes, split_sizes) + y_ref.backward(dy) + + # Caller-provided buffers, same weights as the reference. + model = build() + with torch.no_grad(): + for op_idx in range(2): + for i in range(group_size): + getattr(model[op_idx], f"weight{i}").copy_( + getattr(model_ref[op_idx], f"weight{i}") + ) + sentinel = 7.0 + out_buf = torch.full((num_tokens, out_features), sentinel, dtype=dtype, device=device) + dgrad_buf = torch.full((num_tokens, in_features), sentinel, dtype=dtype, device=device) + x_test = x.detach().clone().requires_grad_(True) + y = model( + x_test, + split_sizes, + split_sizes, + op_kwargs={ + model[0]: {GRAD_INPUT_BUFFER_KEY: dgrad_buf}, + model[1]: {OUTPUT_BUFFER_KEY: out_buf}, + }, + ) + + # Forward output aliases the last op's output buffer with no copy. + assert y.data_ptr() == out_buf.data_ptr() + torch.testing.assert_close(y, y_ref, rtol=0, atol=0) + + y.backward(dy) + + # grad_input written into the first op's dgrad buffer. + assert not torch.all(dgrad_buf == sentinel) + torch.testing.assert_close(dgrad_buf, x_ref.grad, rtol=0, atol=0) + torch.testing.assert_close(x_test.grad, x_ref.grad, rtol=0, atol=0) + + # A buffer whose shape does not match the output is rejected. + bad = torch.empty(num_tokens + 1, out_features, dtype=dtype, device=device) + with pytest.raises(ValueError): + model_bad = build() + model_bad( + x.detach(), + split_sizes, + split_sizes, + op_kwargs={model_bad[1]: {OUTPUT_BUFFER_KEY: bad}}, + ) + @pytest.mark.parametrize("in_shape", ((71, 192), (5, 7, 128))) @pytest.mark.parametrize("input_requires_grad", (False, True)) @pytest.mark.parametrize("scales_requires_grad", (False, True)) diff --git a/tests/pytorch/test_grouped_linear.py b/tests/pytorch/test_grouped_linear.py index b561506b31..10d0dbd227 100644 --- a/tests/pytorch/test_grouped_linear.py +++ b/tests/pytorch/test_grouped_linear.py @@ -1499,6 +1499,9 @@ def test_fp8_grouped_gemm(shape, accumulate): _fp8_available, _reason_for_no_fp8 = fp8_available, reason_for_no_fp8 _mxfp8_available, _reason_for_no_mxfp8 = mxfp8_available, reason_for_no_mxfp8 _nvfp4_available, _reason_for_no_nvfp4 = nvfp4_available, reason_for_no_nvfp4 +_fp8_block_scaling_available, _reason_for_no_fp8_block_scaling = te.is_fp8_block_scaling_available( + return_reason=True +) @pytest.fixture(autouse=True) @@ -1590,8 +1593,14 @@ def _run_grouped_linear_path( recipe.NVFP4BlockScaling(disable_stochastic_rounding=True), marks=pytest.mark.skipif(not _nvfp4_available, reason=_reason_for_no_nvfp4), ), + pytest.param( + recipe.Float8BlockScaling(), + marks=pytest.mark.skipif( + not _fp8_block_scaling_available, reason=_reason_for_no_fp8_block_scaling + ), + ), ], - ids=["bf16", "fp8_current_scaling", "mxfp8", "nvfp4"], + ids=["bf16", "fp8_current_scaling", "mxfp8", "nvfp4", "fp8_block_scaling"], ) @pytest.mark.parametrize("bias", _ALL_BOOLEAN) @pytest.mark.parametrize("fp8_model_params", _ALL_BOOLEAN) @@ -1605,10 +1614,13 @@ def test_grouped_linear_grouped_tensor_path_matches_legacy( pytest.skip( "GroupedTensor grouped GEMM path requires Hopper (SM90) or Blackwell (SM10x and SM110)." ) - # MXFP8/NVFP4 grouped quantization kernels require Blackwell, but FP8 per-tensor - # current scaling also runs on the Hopper grouped GEMM path. + # MXFP8/NVFP4 grouped quantization kernels require Blackwell; FP8 per-tensor + # current scaling runs on Hopper and Blackwell; FP8 block scaling is Hopper-only. is_current_scaling = use_fp8 and fp8_recipe.float8_current_scaling() - if use_fp8 and not is_current_scaling and device_capability < (10, 0): + is_block_scaling = use_fp8 and fp8_recipe.float8_block_scaling() + if is_block_scaling and not (9, 0) <= device_capability < (10, 0): + pytest.skip("Fused grouped FP8 block-scaling requires Hopper (SM90).") + if use_fp8 and not is_current_scaling and not is_block_scaling and device_capability < (10, 0): pytest.skip( "Quantized GroupedTensor grouped GEMM path (MXFP8/NVFP4) requires Blackwell (SM100+)." ) @@ -1710,6 +1722,102 @@ def test_grouped_linear_grouped_tensor_path_single_grouped_bias_delay_wgrad(monk grouped_linear.backward_dw() +@pytest.mark.parametrize("use_fused_path", [False, True], ids=["legacy", "grouped_tensor"]) +@pytest.mark.parametrize("supply", ["out", "dgrad_out", "both"]) +def test_grouped_linear_caller_output_buffers(use_fused_path, supply, monkeypatch): + """Caller-provided forward out and/or backward dgrad_out buffers. + + Checks that a supplied buffer is written in place (bit-for-bit vs internal allocation) + and, on the fused path with a padded input, that only the valid rows are touched. + """ + if use_fused_path: + device_capability = torch.cuda.get_device_capability() + if not (9, 0) <= device_capability <= (11, 0): + pytest.skip( + "GroupedTensor grouped GEMM path requires Hopper (SM90) or Blackwell" + " (SM10x and SM110)." + ) + cublaslt_version = tex.get_cublasLt_version() + if device_capability < (10, 0) and cublaslt_version < 130400: + pytest.skip("Grouped GEMM on Hopper requires cuBLAS 13.4+.") + if cublaslt_version < 130300: + pytest.skip("Grouped GEMM requires cuBLAS 13.3+.") + + monkeypatch.setenv(_FUSED_GROUPED_GEMM_ENV, "1" if use_fused_path else "0") + give_out = supply in ("out", "both") + give_dgrad = supply in ("dgrad_out", "both") + + dtype = torch.bfloat16 + num_gemms = 3 + in_features = 128 + out_features = 128 + m_splits_list = [64, 96, 80] + valid_tokens = sum(m_splits_list) # 240 + # The fused path supports a padded input; the legacy path requires tight packing. + num_rows = valid_tokens + (80 if use_fused_path else 0) + m_splits = ( + torch.tensor(m_splits_list, dtype=torch.int64, device="cuda") + if use_fused_path + else m_splits_list + ) + + torch.manual_seed(1234) + x_base = (0.1 * torch.randn(num_rows, in_features, device="cuda")).to(dtype) + dy = torch.zeros(num_rows, out_features, dtype=dtype, device="cuda") + dy[:valid_tokens] = (0.1 * torch.randn(valid_tokens, out_features, device="cuda")).to(dtype) + + grouped_linear = GroupedLinear( + num_gemms, + in_features, + out_features, + bias=False, + params_dtype=dtype, + device="cuda", + ) + + # Reference: internal allocation. + x_ref = x_base.detach().clone().requires_grad_(True) + y_ref = grouped_linear(x_ref, m_splits) + y_ref.backward(dy) + + # Caller-provided buffers with a sentinel-filled tail (only the requested ones). + sentinel = 7.0 + out_buf = ( + torch.full((num_rows, out_features), sentinel, dtype=dtype, device="cuda") + if give_out + else None + ) + dgrad_buf = ( + torch.full((num_rows, in_features), sentinel, dtype=dtype, device="cuda") + if give_dgrad + else None + ) + x = x_base.detach().clone().requires_grad_(True) + y = grouped_linear(x, m_splits, out=out_buf, dgrad_out=dgrad_buf) + + if give_out: + # Forward output is the caller buffer itself (no copy); padded tail untouched. + assert y.data_ptr() == out_buf.data_ptr() + assert tuple(y.shape) == (num_rows, out_features) + assert torch.all(out_buf[valid_tokens:] == sentinel) + torch.testing.assert_close(y[:valid_tokens], y_ref[:valid_tokens], rtol=0, atol=0) + + y.backward(dy) + + if give_dgrad: + # dgrad written into the caller buffer; padded tail untouched. + assert torch.all(dgrad_buf[valid_tokens:] == sentinel) + torch.testing.assert_close( + dgrad_buf[:valid_tokens], x_ref.grad[:valid_tokens], rtol=0, atol=0 + ) + torch.testing.assert_close(x.grad[:valid_tokens], x_ref.grad[:valid_tokens], rtol=0, atol=0) + + # A buffer whose row count does not match the input rows is rejected. + bad_out = torch.empty(num_rows + 1, out_features, dtype=dtype, device="cuda") + with pytest.raises(ValueError): + grouped_linear(x, m_splits, out=bad_out) + + @pytest.mark.skipif(not _nvfp4_available, reason=_reason_for_no_nvfp4) def test_grouped_linear_grouped_tensor_path_skips_non_rht_nvfp4(monkeypatch): """Non-RHT NVFP4 falls back to the legacy path; check it stays numerically correct. @@ -1810,8 +1918,14 @@ def test_grouped_linear_grouped_tensor_path_skips_non_rht_nvfp4(monkeypatch): recipe.NVFP4BlockScaling(disable_stochastic_rounding=True), marks=pytest.mark.skipif(not _nvfp4_available, reason=_reason_for_no_nvfp4), ), + pytest.param( + recipe.Float8BlockScaling(), + marks=pytest.mark.skipif( + not _fp8_block_scaling_available, reason=_reason_for_no_fp8_block_scaling + ), + ), ], - ids=["bf16", "fp8_current_scaling", "mxfp8", "nvfp4"], + ids=["bf16", "fp8_current_scaling", "mxfp8", "nvfp4", "fp8_block_scaling"], ) @pytest.mark.parametrize("bias", _ALL_BOOLEAN) def test_grouped_linear_fused_path_cuda_graph_safe(fp8_recipe, bias, monkeypatch): @@ -1822,10 +1936,13 @@ def test_grouped_linear_fused_path_cuda_graph_safe(fp8_recipe, bias, monkeypatch pytest.skip( "GroupedTensor grouped GEMM path requires Hopper (SM90) or Blackwell (SM10x and SM110)." ) - # MXFP8/NVFP4 grouped quantization kernels require Blackwell, but FP8 per-tensor - # current scaling also runs on the Hopper grouped GEMM path. + # MXFP8/NVFP4 grouped quantization kernels require Blackwell; FP8 per-tensor + # current scaling runs on Hopper and Blackwell; FP8 block scaling is Hopper-only. is_current_scaling = use_fp8 and fp8_recipe.float8_current_scaling() - if use_fp8 and not is_current_scaling and device_capability < (10, 0): + is_block_scaling = use_fp8 and fp8_recipe.float8_block_scaling() + if is_block_scaling and not (9, 0) <= device_capability < (10, 0): + pytest.skip("Fused grouped FP8 block-scaling requires Hopper (SM90).") + if use_fp8 and not is_current_scaling and not is_block_scaling and device_capability < (10, 0): pytest.skip( "Quantized GroupedTensor grouped GEMM path (MXFP8/NVFP4) requires Blackwell (SM100+)." ) @@ -1857,6 +1974,15 @@ def test_grouped_linear_fused_path_cuda_graph_safe(fp8_recipe, bias, monkeypatch params_dtype=dtype, device=device, ) + reference_grouped_linear = GroupedLinear( + num_gemms, + in_features, + out_features, + bias=bias, + params_dtype=dtype, + device=device, + ) + reference_grouped_linear.load_state_dict(grouped_linear.state_dict()) static_x = torch.randn(total_tokens, in_features, dtype=dtype, device=device) static_x.requires_grad_(True) @@ -1919,7 +2045,7 @@ def _train_step(x, dy, out_buf, *, use_graphed): expected_x = fresh_x.detach().clone().requires_grad_(True) expected_dy = fresh_dy.detach().clone() with autocast(enabled=use_fp8, recipe=fp8_recipe): - expected_out = grouped_linear(expected_x, static_m_splits) + expected_out = reference_grouped_linear(expected_x, static_m_splits) expected_out.backward(expected_dy) tols = dict(rtol=1e-2, atol=5e-3) @@ -1927,11 +2053,29 @@ def _train_step(x, dy, out_buf, *, use_graphed): tols = dict(rtol=0.05, atol=0.05) torch.testing.assert_close(graph_out.float(), expected_out.float(), **tols) torch.testing.assert_close(graph_dx.float(), expected_x.grad.float(), **tols) - for graph_grad, param in zip(graph_param_grads, grouped_linear.parameters()): + for graph_grad, param in zip(graph_param_grads, reference_grouped_linear.parameters()): assert param.grad is not None torch.testing.assert_close(graph_grad.float(), param.grad.float(), **tols) +@pytest.mark.skipif(not _fp8_block_scaling_available, reason=_reason_for_no_fp8_block_scaling) +@pytest.mark.skipif( + not (10, 0) <= torch.cuda.get_device_capability() <= (11, 0), + reason="Error path only triggers on Blackwell (SM100/SM110).", +) +def test_grouped_linear_fused_path_fp8_block_scaling_blackwell_error(monkeypatch): + """FP8BS + fused env var on Blackwell must raise, not silently fall back.""" + monkeypatch.setenv(_FUSED_GROUPED_GEMM_ENV, "1") + FP8GlobalStateManager.reset() + dtype = torch.bfloat16 + grouped_linear = GroupedLinear(2, 128, 128, bias=False, params_dtype=dtype, device="cuda") + x = torch.randn(256, 128, device="cuda", dtype=dtype, requires_grad=True) + m_splits = torch.tensor([128, 128], dtype=torch.int64, device="cuda") + with pytest.raises(RuntimeError, match="Hopper-only"): + with autocast(enabled=True, recipe=recipe.Float8BlockScaling()): + grouped_linear(x, m_splits) + + @pytest.mark.parametrize("swizzle_type", ["mxfp8_rowwise", "mxfp8_columnwise", "nvfp4"]) def test_swizzle_scales_and_pack_ptrs_for_discrete_weights( swizzle_type: str, diff --git a/tests/pytorch/test_grouped_mlp.py b/tests/pytorch/test_grouped_mlp.py index e24fff9049..76550db5f8 100644 --- a/tests/pytorch/test_grouped_mlp.py +++ b/tests/pytorch/test_grouped_mlp.py @@ -19,6 +19,10 @@ _cudnn_frontend_supports_grouped_gemm_srelu, _cudnn_frontend_version_supported, ) +from transformer_engine.pytorch.ops.basic.grouped_linear import ( + OUTPUT_BUFFER_KEY, + GRAD_INPUT_BUFFER_KEY, +) from transformer_engine.pytorch import ( QuantizedTensor, Float8CurrentScalingQuantizer, @@ -515,7 +519,8 @@ def test_grouped_linear_cuda_graph_safe( split_sizes = torch.tensor(split_sizes, dtype=torch.int, device=device) # Pad input tokens to validate the sync-free flow - in_shape = (split_sizes.sum().item() + token_padding, in_features) + num_active_tokens = split_sizes.sum().item() + in_shape = (num_active_tokens + token_padding, in_features) out_shape = (in_shape[0], out_features) recipe = make_recipe(quantization) @@ -531,33 +536,45 @@ def test_grouped_linear_cuda_graph_safe( single_grouped_weight=single_grouped_weight, single_grouped_bias=single_grouped_bias, ) + reference_op = te.ops.GroupedLinear( + group_size, + in_features, + out_features, + bias=bias, + device=device, + dtype=dtype, + accumulate_into_main_grad=accumulate_into_main_grad, + single_grouped_weight=single_grouped_weight, + single_grouped_bias=single_grouped_bias, + ) + reference_op.load_state_dict(op.state_dict()) - def _weight_params() -> list[torch.nn.Parameter]: + def _weight_params(module: torch.nn.Module) -> list[torch.nn.Parameter]: if single_grouped_weight: - return [op.weight] - return [getattr(op, f"weight{i}") for i in range(group_size)] + return [module.weight] + return [getattr(module, f"weight{i}") for i in range(group_size)] - def _bias_params() -> list[torch.nn.Parameter]: + def _bias_params(module: torch.nn.Module) -> list[torch.nn.Parameter]: if not bias: return [] if single_grouped_bias: - return [op.bias] - return [getattr(op, f"bias{i}") for i in range(group_size)] + return [module.bias] + return [getattr(module, f"bias{i}") for i in range(group_size)] - def _init_main_grads(value: float = 0.0) -> None: + def _init_main_grads(module: torch.nn.Module, value: float = 0.0) -> None: if not accumulate_into_main_grad: return with torch.no_grad(): - for w in _weight_params(): + for w in _weight_params(module): if getattr(w, "main_grad", None) is None: w.main_grad = torch.empty(w.size(), device=device, dtype=torch.float32) w.main_grad.fill_(value) - def _collect_main_grads() -> list[torch.Tensor]: - return [w.main_grad.detach().clone() for w in _weight_params()] + def _collect_main_grads(module: torch.nn.Module) -> list[torch.Tensor]: + return [w.main_grad.detach().clone() for w in _weight_params(module)] - def _zero_param_grads() -> None: - for param in op.parameters(): + def _zero_param_grads(module: torch.nn.Module) -> None: + for param in module.parameters(): if param.grad is None: param.grad = torch.zeros_like(param) else: @@ -582,7 +599,7 @@ def train_step( out_buf.copy_(out) return out_buf - _init_main_grads(0.0) + _init_main_grads(op, 0.0) static_x = torch.randn(in_shape, device=device, dtype=dtype, requires_grad=True) static_dy = torch.randn(out_shape, device=device, dtype=dtype) @@ -605,8 +622,8 @@ def train_step( static_dy.copy_(fresh_dy) # Reset grads & main_grads so the captured iteration starts fresh. - _zero_param_grads() - _init_main_grads(0.5) + _zero_param_grads(op) + _init_main_grads(op, 0.5) if static_x.grad is not None: static_x.grad.zero_() @@ -617,35 +634,36 @@ def train_step( torch.cuda.synchronize() graph_dx = static_x.grad.detach().clone() if accumulate_into_main_grad: - graph_main_grads = _collect_main_grads() + graph_main_grads = _collect_main_grads(op) graph_param_grads: list[torch.Tensor] = [] else: graph_main_grads = [] graph_param_grads = [param.grad.detach().clone() for param in op.parameters()] - # Reference: same op invoked eagerly with the same fresh inputs and - # the same starting grad/main_grad state. - _zero_param_grads() - _init_main_grads(0.5) - static_x.grad.zero_() + # Reference: an independent op invoked eagerly with the same fresh + # inputs and starting grad/main_grad state. + _zero_param_grads(reference_op) + _init_main_grads(reference_op, 0.5) expected_x = fresh_x.detach().clone().requires_grad_(True) expected_dy = fresh_dy.detach().clone() with te.autocast(enabled=quantization is not None, recipe=recipe): - expected_out = op(expected_x, static_split_sizes) + expected_out = reference_op(expected_x, static_split_sizes) expected_out.backward(expected_dy) tols = dtype_tols(dtype) if quantization is not None: tols = quantization_tols(quantization) - assert_close(graph_out, expected_out, **tols) - assert_close(graph_dx, expected_x.grad, **tols) + # Grouped GEMM only defines rows covered by split_sizes. The padded tail + # is intentionally outside every group and remains uninitialized. + assert_close(graph_out[:num_active_tokens], expected_out[:num_active_tokens], **tols) + assert_close(graph_dx[:num_active_tokens], expected_x.grad[:num_active_tokens], **tols) if accumulate_into_main_grad: - for g, w in zip(graph_main_grads, _weight_params()): + for g, w in zip(graph_main_grads, _weight_params(reference_op)): assert_close(g, w.main_grad, **tols) else: - for g, param in zip(graph_param_grads, op.parameters()): + for g, param in zip(graph_param_grads, reference_op.parameters()): assert_close(g, param.grad, **tols) @@ -1317,6 +1335,103 @@ def _run_case(single_grouped_weight: bool) -> tuple[torch.Tensor, ...]: torch.testing.assert_close(fc1_db_false, fc1_db_true, **bias_tols) torch.testing.assert_close(fc2_db_false, fc2_db_true, **bias_tols) + @pytest.mark.parametrize("quantization", ("mxfp8", "nvfp4_rht")) + def test_grouped_mlp_caller_buffers( + self, + quantization: str, + *, + dtype: torch.dtype = torch.bfloat16, + device: torch.device = "cuda", + group_size: int = 4, + hidden_size: int = 256, + split_alignment: int = 256, + glu_interleave_size: int = 32, + ) -> None: + """Caller-provided output/grad_input buffers on the fused MXFP8/NVFP4 grouped MLP.""" + if quantization == "mxfp8" and not mxfp8_available: + pytest.skip(reason_for_no_mxfp8) + if quantization == "nvfp4_rht" and not nvfp4_available: + pytest.skip(reason_for_no_nvfp4) + if not te.ops.fused.GroupedMLP_CuTeGEMMGLU.is_supported(): + pytest.skip("Fused grouped MLP is not supported on this system") + + split_sizes = torch.tensor( + [split_alignment * (i + 1) for i in range(group_size)], + dtype=torch.int64, + device=device, + ) + num_tokens = int(split_sizes.sum()) + recipe = make_recipe(quantization) + + x_base = torch.empty((num_tokens, hidden_size), device=device, dtype=dtype).uniform_( + -0.25, 0.25 + ) + probs_base = torch.empty((num_tokens,), device=device, dtype=dtype).uniform_(-0.25, 0.25) + dy_base = torch.empty((num_tokens, hidden_size), device=device, dtype=dtype).uniform_( + -0.25, 0.25 + ) + fc1_ws_base = [ + torch.empty((2 * hidden_size, hidden_size), device=device, dtype=dtype).uniform_( + -0.25, 0.25 + ) + for _ in range(group_size) + ] + fc2_ws_base = [ + torch.empty((hidden_size, hidden_size), device=device, dtype=dtype).uniform_( + -0.25, 0.25 + ) + for _ in range(group_size) + ] + + def build() -> te.ops.Sequential: + with te.quantized_model_init(enabled=True, recipe=recipe): + fc1 = te.ops.GroupedLinear( + group_size, hidden_size, 2 * hidden_size, bias=False, device=device, dtype=dtype + ) + scaled_act = te.ops.ScaledSwiGLU(glu_interleave_size=glu_interleave_size) + fc2 = te.ops.GroupedLinear( + group_size, hidden_size, hidden_size, bias=False, device=device, dtype=dtype + ) + module = te.ops.Sequential(fc1, scaled_act, fc2) + with torch.no_grad(): + for i in range(group_size): + getattr(fc1, f"weight{i}").copy_(fc1_ws_base[i]) + getattr(fc2, f"weight{i}").copy_(fc2_ws_base[i]) + return module + + def run(module, *, output=None, grad_input=None): + x = x_base.detach().clone().requires_grad_(True) + probs = probs_base.detach().clone().requires_grad_(True) + op_kwargs = {} + if grad_input is not None: + op_kwargs[module[0]] = {GRAD_INPUT_BUFFER_KEY: grad_input} + if output is not None: + op_kwargs[module[2]] = {OUTPUT_BUFFER_KEY: output} + with te.autocast(enabled=True, recipe=recipe): + y = module(x, split_sizes, probs, split_sizes, op_kwargs=op_kwargs or None) + y.backward(dy_base) + return y, x.grad + + # Reference: internal allocation. + y_ref, dx_ref = run(build()) + + # Caller-provided buffers. + module = build() + sentinel = 7.0 + out_buf = torch.full((num_tokens, hidden_size), sentinel, dtype=dtype, device=device) + dgrad_buf = torch.full((num_tokens, hidden_size), sentinel, dtype=dtype, device=device) + y, _ = run(module, output=out_buf, grad_input=dgrad_buf) + + # The fused op ran, the returned output aliases the buffer, and both buffers were written. + assert isinstance( + module._module_groups[0]._forward_ops[0][0], + te.ops.fused.GroupedMLP_CuTeGEMMGLU, + ) + assert y.data_ptr() == out_buf.data_ptr() + assert not torch.all(dgrad_buf == sentinel) + torch.testing.assert_close(y, y_ref, rtol=0, atol=0) + torch.testing.assert_close(dgrad_buf, dx_ref, rtol=0, atol=0) + @pytest.mark.parametrize("single_grouped_weight", (False, True)) @pytest.mark.parametrize("delay_wgrad_compute", (False, True)) @pytest.mark.parametrize("zero_out_wgrad", (False, True)) @@ -1491,7 +1606,8 @@ def test_grouped_mlp_cuda_graph_safe_mxfp8( random.shuffle(split_sizes) split_sizes = torch.tensor(split_sizes, dtype=torch.int64, device=device) # Pad the input tokens to validate the sync-free MOE - in_shape = (split_sizes.sum().item() + token_padding, hidden_size) + num_active_tokens = split_sizes.sum().item() + in_shape = (num_active_tokens + token_padding, hidden_size) recipe = make_recipe("mxfp8") with te.quantized_model_init(enabled=True, recipe=recipe): fc1 = te.ops.GroupedLinear( @@ -1524,8 +1640,43 @@ def test_grouped_mlp_cuda_graph_safe_mxfp8( scaled_act, fc2, ) + reference_fc1 = te.ops.GroupedLinear( + group_size, + hidden_size, + 2 * hidden_size, + bias=False, + device=device, + dtype=dtype, + single_grouped_weight=single_grouped_weight, + accumulate_into_main_grad=accumulate_into_main_grad, + ) + reference_fc2 = te.ops.GroupedLinear( + group_size, + hidden_size, + hidden_size, + bias=False, + device=device, + dtype=dtype, + single_grouped_weight=single_grouped_weight, + accumulate_into_main_grad=accumulate_into_main_grad, + ) + reference_scaled_act = ( + te.ops.ScaledSwiGLU(glu_interleave_size=glu_interleave_size) + if activation == "scaled_swiglu" + else te.ops.ScaledClampedQGeGLU(glu_interleave_size=glu_interleave_size) + ) + reference_module = te.ops.Sequential( + reference_fc1, + reference_scaled_act, + reference_fc2, + ) + reference_module.load_state_dict(module.state_dict()) - def _init_main_grads(value: float = 0.0) -> None: + def _init_main_grads( + fc1: torch.nn.Module, + fc2: torch.nn.Module, + value: float = 0.0, + ) -> None: if not accumulate_into_main_grad: return with torch.no_grad(): @@ -1563,7 +1714,10 @@ def _init_main_grads(value: float = 0.0) -> None: fc1_weight.main_grad.fill_(value) fc2_weight.main_grad.fill_(value) - def _collect_main_grads() -> tuple[torch.Tensor, torch.Tensor]: + def _collect_main_grads( + fc1: torch.nn.Module, + fc2: torch.nn.Module, + ) -> tuple[torch.Tensor, torch.Tensor]: if single_grouped_weight: fc1_main_grad = fc1.weight.main_grad.detach().clone() fc2_main_grad = fc2.weight.main_grad.detach().clone() @@ -1604,7 +1758,7 @@ def train_step( out_buf.copy_(out) return out_buf - _init_main_grads(0.0) + _init_main_grads(fc1, fc2, 0.0) static_x = torch.randn(in_shape, device=device, dtype=dtype, requires_grad=True) static_probs = torch.randn((in_shape[0],), device=device, dtype=dtype, requires_grad=True) @@ -1643,7 +1797,7 @@ def train_step( for param in module.parameters(): param.grad = torch.zeros_like(param) - _init_main_grads(0.5) + _init_main_grads(fc1, fc2, 0.5) if static_x.grad is not None: static_x.grad.zero_() if static_probs.grad is not None: @@ -1658,21 +1812,19 @@ def train_step( graph_dx = static_x.grad.detach().clone() graph_dprobs = static_probs.grad.detach().clone() if accumulate_into_main_grad: - graph_fc1_main_grad, graph_fc2_main_grad = _collect_main_grads() + graph_fc1_main_grad, graph_fc2_main_grad = _collect_main_grads(fc1, fc2) else: graph_param_grads = [param.grad.detach().clone() for param in module.parameters()] - for param in module.parameters(): - param.grad.zero_() - _init_main_grads(0.5) - static_x.grad.zero_() - static_probs.grad.zero_() + for param in reference_module.parameters(): + param.grad = torch.zeros_like(param) + _init_main_grads(reference_fc1, reference_fc2, 0.5) expected_x = fresh_x.detach().clone().requires_grad_(True) expected_probs = fresh_probs.detach().clone().requires_grad_(True) expected_dy = fresh_dy.detach().clone() with te.autocast(enabled=True, recipe=recipe): - expected_out = module( + expected_out = reference_module( expected_x, static_split_sizes, expected_probs, @@ -1681,15 +1833,24 @@ def train_step( expected_out.backward(expected_dy) tols = dtype_tols(dtype) - assert_close(graph_out, expected_out, **tols) - assert_close(graph_dx, expected_x.grad, **tols) - assert_close(graph_dprobs, expected_probs.grad, **tols) + # The padded tail is outside every expert and its outputs/gradients are + # intentionally left uninitialized. + assert_close(graph_out[:num_active_tokens], expected_out[:num_active_tokens], **tols) + assert_close(graph_dx[:num_active_tokens], expected_x.grad[:num_active_tokens], **tols) + assert_close( + graph_dprobs[:num_active_tokens], + expected_probs.grad[:num_active_tokens], + **tols, + ) if accumulate_into_main_grad: - expected_fc1_main_grad, expected_fc2_main_grad = _collect_main_grads() + expected_fc1_main_grad, expected_fc2_main_grad = _collect_main_grads( + reference_fc1, + reference_fc2, + ) assert_close(graph_fc1_main_grad, expected_fc1_main_grad, **tols) assert_close(graph_fc2_main_grad, expected_fc2_main_grad, **tols) else: - for graph_grad, param in zip(graph_param_grads, module.parameters()): + for graph_grad, param in zip(graph_param_grads, reference_module.parameters()): assert_close(graph_grad, param.grad, **tols) diff --git a/tests/pytorch/test_grouped_tensor.py b/tests/pytorch/test_grouped_tensor.py index eeb6e7a394..4dd52ee2bd 100644 --- a/tests/pytorch/test_grouped_tensor.py +++ b/tests/pytorch/test_grouped_tensor.py @@ -1008,11 +1008,19 @@ def _assert_fp8_cs_group_quantize_matches_reference( @pytest.mark.parametrize("shape_case", ["uniform", "varying_first"]) @pytest.mark.parametrize("direction", ["rowwise", "columnwise", "both"]) @pytest.mark.parametrize("output_dbias", [False, True]) + @pytest.mark.parametrize( + "force_pow_2_scales", [False, True], ids=["fp32_scales", "pow2_scales"] + ) @pytest.mark.skipif( not fp8_block_scaling_grouped_available, reason=reason_for_no_fp8_block_scaling_grouped ) def test_quantize_grouped_fp8_blockwise( - self, block_scaling_dim: int, shape_case: str, direction: str, output_dbias: bool + self, + block_scaling_dim: int, + shape_case: str, + direction: str, + output_dbias: bool, + force_pow_2_scales: bool, ) -> None: """Test grouped FP8 block-scaling quantization against per-tensor quantization. @@ -1061,7 +1069,7 @@ def test_quantize_grouped_fp8_blockwise( fp8_dtype=tex.DType.kFloat8E4M3, rowwise=rowwise, columnwise=columnwise, - force_pow_2_scales=False, + force_pow_2_scales=force_pow_2_scales, amax_epsilon=0.0, block_scaling_dim=block_scaling_dim, ) @@ -1081,7 +1089,7 @@ def test_quantize_grouped_fp8_blockwise( fp8_dtype=tex.DType.kFloat8E4M3, rowwise=True, columnwise=True, - force_pow_2_scales=False, + force_pow_2_scales=force_pow_2_scales, amax_epsilon=0.0, block_scaling_dim=block_scaling_dim, ) diff --git a/tests/pytorch/test_ops_grouped_linear_distributed_weight.py b/tests/pytorch/test_ops_grouped_linear_distributed_weight.py new file mode 100644 index 0000000000..700be4138e --- /dev/null +++ b/tests/pytorch/test_ops_grouped_linear_distributed_weight.py @@ -0,0 +1,131 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""DistributedWeight dispatch in the fusible ``ops.GroupedLinear`` (unfused fallback path). + +TE ships no DistributedWeight implementer (GTP etc. live in the caller), so this validates the +*dispatch wiring* with an in-repo fake. The fake applies a DISTINCT, observable scale in each +materialize hook so every plumbing point is independently checked -- a wiring bug fails a specific +assertion: + + * ``materialize_group_for_forward`` scales the weights by ``FWD_SCALE`` -> the fwd GEMM must use + the materialized weights, so ``out == FWD_SCALE * plain_out``. + * ``materialize_group_for_backward`` scales by ``BWD_SCALE`` -> dgrad must use the + re-materialized weights, so ``dgrad == BWD_SCALE * plain_dgrad``. + * ``finalize_group_grads`` reduce-scatters the wgrads into ``main_grad`` in-place and returns a + dummy -> the real wgrad lands in ``main_grad`` and the ops path returns a throwaway ``.grad`` + (it discards finalize's return; see DistributedWeight.finalize_group_grads). + +Real all-gather / reduce-scatter math is exercised by the caller's distributed tests; here the fake +is single-process and only proves the ops.GroupedLinear integration routes weights/grads through +the hooks (and does not bypass them). +""" + +import pytest +import torch + +import transformer_engine.pytorch as te + +# Distinct powers of two so each scale commutes exactly through the (possibly TF32) GEMM rounding, +# making the linear scale relations bit-exact under a tight tolerance. +FWD_SCALE = 2.0 +BWD_SCALE = 4.0 + + +class _FakeDistWeight(torch.nn.Parameter): + """Single-process fake DistributedWeight leader for dispatch testing. + + Each materialize hook applies its own scale (see module docstring) so fwd/bwd routing is + observable; ``finalize_group_grads`` models the real main-grad contract -- reduce-scatter (here + an identity accumulate) the wgrads into each shard's ``main_grad`` in-place, flag + ``grad_added_to_main_grad``, and return a dummy that the ops path discards. + """ + + is_distributed_weight = True + + def materialize_group_for_forward(self): + self.calls["fwd"] += 1 + return [w * FWD_SCALE for w in self._group] + + def materialize_group_for_backward(self, **kwargs): + self.calls["bwd"] += 1 + return [w * BWD_SCALE for w in self._group] + + def finalize_group_grads(self, wgrads, **kwargs): + self.calls["finalize"] += 1 + wl = list(wgrads) if isinstance(wgrads, (list, tuple)) else [wgrads] + for w, g in zip(self._group, wl): + w.main_grad.add_(g.to(w.main_grad.dtype)) # in-place accumulate into main_grad + w.grad_added_to_main_grad = True + return [torch.zeros_like(g) for g in wl] # dummy grads (real value is now in main_grad) + + def grad_buffer(self): + return self.data + + +def _make_fake_dist_leader(op, num_gemms): + """Replace ``op.weight0`` with a fake distributed leader referencing the whole group.""" + w0 = op.weight0 + leader = _FakeDistWeight(w0.data) + leader.calls = {"fwd": 0, "bwd": 0, "finalize": 0} + leader._group = [leader] + [getattr(op, f"weight{i}") for i in range(1, num_gemms)] + op.weight0 = leader + return leader + + +@pytest.mark.parametrize("num_gemms", [2, 4]) +def test_ops_grouped_linear_distributed_weight_dispatch(num_gemms): + """Every DistributedWeight hook must be routed through the GEMM flow (and not bypassed). + + fwd/bwd use the scaled materialized weights; finalize reduce-scatters the wgrad into + ``main_grad`` in-place, and the ops path returns a throwaway dummy ``.grad``. + """ + if not torch.cuda.is_available(): + pytest.skip("requires CUDA") + torch.manual_seed(0) + # fp32 with power-of-two scales: each scale commutes exactly through the GEMM rounding (even + # TF32), so the linear scale relations below are bit-exact under a tight tolerance. + in_f, out_f, total_tokens = 32, 64, num_gemms * 8 + dtype, device = torch.float32, "cuda" + + op = te.ops.GroupedLinear(num_gemms, in_f, out_f, bias=False, device=device, dtype=dtype) + reference = te.ops.GroupedLinear(num_gemms, in_f, out_f, bias=False, device=device, dtype=dtype) + reference.load_state_dict(op.state_dict()) + + leader = _make_fake_dist_leader(op, num_gemms) + for i in range(num_gemms): + w = getattr(op, f"weight{i}") + w.main_grad = torch.zeros((out_f, in_f), dtype=torch.float32, device=device) + w.grad_added_to_main_grad = False # DDP initializes this on every param + + m_splits = [total_tokens // num_gemms] * num_gemms + m_splits[-1] += total_tokens - sum(m_splits) + split_sizes = torch.tensor(m_splits, dtype=torch.int64, device=device) + + x = torch.randn(total_tokens, in_f, dtype=dtype, device=device, requires_grad=True) + ref_x = x.detach().clone().requires_grad_(True) + + out = op(x, split_sizes) + out.sum().backward() + ref_out = reference(ref_x, split_sizes) + ref_out.sum().backward() + + # All three dispatch hooks actually fired. + assert leader.calls["fwd"] > 0 and leader.calls["bwd"] > 0 and leader.calls["finalize"] > 0 + + tols = dict(rtol=1e-5, atol=1e-5) + # fwd used the materialized (FWD_SCALE) weights. + torch.testing.assert_close(out, FWD_SCALE * ref_out, **tols) + # dgrad used the re-materialized (BWD_SCALE) weights. + torch.testing.assert_close(x.grad, BWD_SCALE * ref_x.grad, **tols) + for i in range(num_gemms): + w = getattr(op, f"weight{i}") + ref_w = getattr(reference, f"weight{i}") + # The distributed op's grad of record is main_grad (finalize reduce-scattered the identity + # wgrad there); compare it to the plain reference module's ordinary autograd .grad. + torch.testing.assert_close(w.main_grad.to(dtype), ref_w.grad, **tols) + # main_grad is flagged, and op's own .grad is a discarded dummy (not the real grad, which + # lives in main_grad). Nobody writes the real grad into op.weight.grad by design. + assert w.grad_added_to_main_grad is True + assert w.grad is not None, f"weight{i} should receive a dummy .grad" diff --git a/transformer_engine/common/CMakeLists.txt b/transformer_engine/common/CMakeLists.txt index be64fcb2be..8eba515e5e 100644 --- a/transformer_engine/common/CMakeLists.txt +++ b/transformer_engine/common/CMakeLists.txt @@ -183,6 +183,7 @@ list(APPEND transformer_engine_cpp_sources fused_attn/fused_attn.cpp gemm/config.cpp normalization/common.cpp + normalization/rtc_dispatch.cpp normalization/layernorm/ln_api.cpp normalization/rmsnorm/rmsnorm_api.cpp util/cuda_driver.cpp @@ -252,6 +253,9 @@ list(APPEND transformer_engine_cuda_arch_specific_sources activation/relu_dbias.cu activation/relu_grouped.cu activation/relu_grouped_dbias.cu + activation/scaled_activation.cu + activation/scaled_srelu.cu + activation/scaled_swiglu.cu activation/swiglu.cu activation/swiglu_dbias.cu activation/swiglu_grouped.cu @@ -442,10 +446,10 @@ option(NVTE_WITH_NCCL_EP "Build NCCL EP into libtransformer_engine.so" ON) if(NVTE_WITH_NCCL_EP) # SM>=90 and NCCL>=2.30.4 are gated at runtime in EPBackend::initialize. # -- NCCL EP headers -------------------------------------------------------- -# Headers + libs are produced by the in-tree 3rdparty/nccl submodule build +# Headers + libs are produced by the in-tree 3rdparty/nccl-extensions submodule build # (auto-built by setup.py via build_nccl_ep_submodule). set(NCCL_EP_SUBMODULE_ROOT - "${CMAKE_CURRENT_SOURCE_DIR}/../../3rdparty/nccl") + "${CMAKE_CURRENT_SOURCE_DIR}/../../3rdparty/nccl-extensions") set(NCCL_EP_INCLUDE_DIR "${NCCL_EP_SUBMODULE_ROOT}/build/include") if(NOT EXISTS "${NCCL_EP_INCLUDE_DIR}/nccl_ep.h") message(FATAL_ERROR @@ -551,13 +555,59 @@ make_string_header_from_file(transpose/rtc/transpose.cu string_code_transpose_rtc_transpose_cu) make_string_header_from_file(transpose/rtc/swap_first_dims.cu string_code_transpose_rtc_swap_first_dims_cu) +make_string_header_from_file(fused_softmax/scaled_masked_softmax.cu + string_code_fused_softmax_scaled_masked_softmax_cu) +make_string_header_from_file(fused_softmax/scaled_upper_triang_masked_softmax.cu + string_code_fused_softmax_scaled_upper_triang_masked_softmax_cu) +make_string_header_from_file(fused_softmax/scaled_aligned_causal_masked_softmax.cu + string_code_fused_softmax_scaled_aligned_causal_masked_softmax_cu) make_string_header_from_file(utils.cuh string_code_utils_cuh) make_string_header_from_file(util/math.h string_code_util_math_h) + +# Norm NVRTC bundled headers + RTC source files +make_string_header_from_file(normalization/kernel_params.h + string_code_normalization_kernel_params_h) +make_string_header_from_file(normalization/kernel_traits.h + string_code_normalization_kernel_traits_h) +make_string_header_from_file(normalization/layernorm/ln_fwd_kernels.cuh + string_code_normalization_layernorm_ln_fwd_kernels_cuh) +make_string_header_from_file(normalization/layernorm/ln_bwd_kernels.cuh + string_code_normalization_layernorm_ln_bwd_kernels_cuh) +make_string_header_from_file(normalization/rmsnorm/rmsnorm_fwd_kernels.cuh + string_code_normalization_rmsnorm_rmsnorm_fwd_kernels_cuh) +make_string_header_from_file(normalization/rmsnorm/rmsnorm_bwd_kernels.cuh + string_code_normalization_rmsnorm_rmsnorm_bwd_kernels_cuh) +make_string_header_from_file(normalization/layernorm/rtc/ln_fwd_kernel.cu + string_code_normalization_layernorm_rtc_ln_fwd_kernel_cu) +make_string_header_from_file(normalization/layernorm/rtc/ln_bwd_kernel.cu + string_code_normalization_layernorm_rtc_ln_bwd_kernel_cu) +make_string_header_from_file(normalization/rmsnorm/rtc/rmsnorm_fwd_kernel.cu + string_code_normalization_rmsnorm_rtc_rmsnorm_fwd_kernel_cu) +make_string_header_from_file(normalization/rmsnorm/rtc/rmsnorm_bwd_kernel.cu + string_code_normalization_rmsnorm_rtc_rmsnorm_bwd_kernel_cu) target_include_directories(transformer_engine PRIVATE "${CMAKE_CURRENT_BINARY_DIR}/string_headers") +option(NVTE_BUILD_LEGACY_STATIC_FUSED_SOFTMAX + "Also compile legacy static fused softmax kernels for NVTE_DISABLE_NVRTC fallback" + OFF) +target_compile_definitions(transformer_engine + PRIVATE + NVTE_BUILD_LEGACY_STATIC_FUSED_SOFTMAX=$) + +# Default OFF: LayerNorm/RMSNorm require NVRTC. Set ON to additionally compile +# the legacy static template instantiations, which are selected only when +# NVTE_DISABLE_NVRTC=1 at runtime. This compiles all 556 +# REGISTER_NORM_LAUNCHER variants up front. +option(NVTE_BUILD_LEGACY_STATIC_NORM + "Also compile static norm kernels for NVTE_DISABLE_NVRTC fallback" + OFF) +target_compile_definitions(transformer_engine + PRIVATE + NVTE_BUILD_LEGACY_STATIC_NORM=$) + # Compiler options set(nvte_sources_with_fast_math) list(APPEND nvte_sources_with_fast_math fused_softmax/scaled_masked_softmax.cu @@ -583,6 +633,9 @@ if (NVTE_BUILD_ACTIVATION_WITH_FAST_MATH) activation/relu_dbias.cu activation/relu_grouped.cu activation/relu_grouped_dbias.cu + activation/scaled_activation.cu + activation/scaled_srelu.cu + activation/scaled_swiglu.cu activation/swiglu.cu activation/swiglu_dbias.cu activation/swiglu_grouped.cu diff --git a/transformer_engine/common/activation/scaled_activation.cu b/transformer_engine/common/activation/scaled_activation.cu new file mode 100644 index 0000000000..4dce3185e0 --- /dev/null +++ b/transformer_engine/common/activation/scaled_activation.cu @@ -0,0 +1,642 @@ +/************************************************************************* + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * + * See LICENSE for license information. + ************************************************************************/ + +/* Vectorization model (gated): + * + * With no GLU interleave, the row is laid out as: + * [ act[0:H] | gate[0:H] ] + * With GLU interleave (e.g. 32): + * [ act[0:32] | gate[0:32] | act[32:64] | gate[32:64] | ... ] + * + * Backward uses fused SiLU / ClampedSiLU specializations (one sigmoid) matching + * gated_fp8.cuh; other ActOP/DActOP pairs call the ops directly. + */ + +#include +#include + +#include "../common.h" +#include "../util/math.h" +#include "./scaled_activation.h" + +namespace transformer_engine { +namespace { + +using namespace detail::scaled_activation; + +using WarpReducer = Reducer; + +// blockDim.x must be a multiple of warp size and <= kReductionThreads. +__device__ __forceinline__ float block_reduce_sum(float value, float *smem) { + const int lane = threadIdx.x % THREADS_PER_WARP; + const int warp = threadIdx.x / THREADS_PER_WARP; + const int num_warps = blockDim.x / THREADS_PER_WARP; + Empty params = {}; + WarpReducer reducer(params, /*bidm=*/0, /*bidn=*/0, /*warp_m=*/0, /*warp_n=*/0, lane, + /*smem=*/nullptr); + Sum sum; + + value = reducer.reduce(value, sum); + if (lane == 0) { + smem[warp] = value; + } + __syncthreads(); + + value = threadIdx.x < num_warps ? smem[lane] : 0.0f; + return warp == 0 ? reducer.reduce(value, sum) : value; +} + +// --------------------------------------------------------------------------- +// Device helpers: fused Act / DAct (IEEE sigmoid / expf) +// --------------------------------------------------------------------------- + +template +__device__ __forceinline__ float gated_forward_value(const float act_in, const float gate_in, + const ParamOP ¶m) { + if constexpr (std::is_same::value) { + const float gate = fminf(fmaxf(-param.limit, gate_in), param.limit) + param.glu_linear_offset; + return ActOP(act_in, param) * gate; + } else { + return ActOP(act_in, param) * gate_in; + } +} + +template +__device__ __forceinline__ void gated_backward_values(const float act_in, const float gate_in, + const ParamOP ¶m, float *dact, + float *dgate, float *unscaled) { + Empty empty = {}; + float act_x = 0.0f; + float dact_x = 0.0f; + float gate = gate_in; + bool dgate_mask = true; + + if constexpr (std::is_same::value) { + dgate_mask = gate_in <= param.limit && gate_in >= -param.limit; + gate = fminf(fmaxf(-param.limit, gate_in), param.limit) + param.glu_linear_offset; + const bool dact_mask = act_in <= param.limit; + const float clamped_act_in = fminf(act_in, param.limit); + const float s = sigmoid(param.alpha * clamped_act_in, empty); + act_x = clamped_act_in * s; + dact_x = dact_mask ? s + param.alpha * clamped_act_in * s * (1.0f - s) : 0.0f; + } else if constexpr ((ActOP == &silu) && (DActOP == &dsilu)) { + const float s = sigmoid(act_in, empty); + act_x = act_in * s; + dact_x = s + act_in * s * (1.0f - s); + } else { + act_x = ActOP(act_in, param); + dact_x = DActOP(act_in, param); + } + + *unscaled = act_x * gate; + *dact = dact_x * gate; + *dgate = dgate_mask ? act_x : 0.0f; +} + +// --------------------------------------------------------------------------- +// Gated kernels +// --------------------------------------------------------------------------- + +template +__global__ void __launch_bounds__(kThreads, 4) + scaled_gated_forward_kernel(const InputT *__restrict__ input, + const ScaleT *__restrict__ act_scales, OutputT *__restrict__ output, + const size_t rows, const size_t hidden, const size_t segment_size, + const size_t num_segments, const size_t num_vectors_per_segment, + const ParamOP param) { + const size_t total_vectors = rows * num_segments * num_vectors_per_segment; + for (size_t tid = blockIdx.x * blockDim.x + threadIdx.x; tid < total_vectors; + tid += gridDim.x * blockDim.x) { + const size_t vector_idx = tid % num_vectors_per_segment; + const size_t segment = (tid / num_vectors_per_segment) % num_segments; + const size_t row = tid / (num_vectors_per_segment * num_segments); + const size_t input_segment_offset = row * hidden * 2 + segment * segment_size * 2; + const size_t output_segment_offset = row * hidden + segment * segment_size; + + VectorizedLoader act_loader(input + input_segment_offset, segment_size); + VectorizedLoader gate_loader(input + input_segment_offset + segment_size, + segment_size); + VectorizedStorer output_storer(output + output_segment_offset, + segment_size); + act_loader.load(vector_idx, segment_size); + gate_loader.load(vector_idx, segment_size); + const float scale = static_cast(act_scales[row]); +#pragma unroll + for (int lane = 0; lane < nvec; ++lane) { + const float unscaled = gated_forward_value( + static_cast(act_loader.separate()[lane]), + static_cast(gate_loader.separate()[lane]), param); + output_storer.separate()[lane] = static_cast(unscaled * scale); + } + output_storer.store(vector_idx, segment_size); + } +} + +template +__global__ void __launch_bounds__(kReductionThreads, 4) + scaled_gated_backward_kernel(const GradT *__restrict__ grad_output, + const InputT *__restrict__ input, + const ScaleT *__restrict__ act_scales, + OutputT *__restrict__ grad_input, + GradScaleT *__restrict__ grad_act_scales, const size_t rows, + const size_t hidden, const size_t segment_size, + const size_t num_segments, const size_t num_vectors_per_segment, + const ParamOP param) { + __shared__ float smem[kReductionWarps]; + const size_t row = blockIdx.x; + (void)rows; + float scale_grad = 0.0f; + const float scale = static_cast(act_scales[row]); + + const size_t row_vectors = num_segments * num_vectors_per_segment; + for (size_t row_vector_idx = threadIdx.x; row_vector_idx < row_vectors; + row_vector_idx += blockDim.x) { + const size_t segment = row_vector_idx / num_vectors_per_segment; + const size_t vector_idx = row_vector_idx % num_vectors_per_segment; + const size_t input_segment_offset = row * hidden * 2 + segment * segment_size * 2; + const size_t output_segment_offset = row * hidden + segment * segment_size; + VectorizedLoader grad_loader(grad_output + output_segment_offset, + segment_size); + VectorizedLoader act_loader(input + input_segment_offset, segment_size); + VectorizedLoader gate_loader(input + input_segment_offset + segment_size, + segment_size); + VectorizedStorer act_storer(grad_input + input_segment_offset, + segment_size); + VectorizedStorer gate_storer( + grad_input + input_segment_offset + segment_size, segment_size); + + grad_loader.load(vector_idx, segment_size); + act_loader.load(vector_idx, segment_size); + gate_loader.load(vector_idx, segment_size); +#pragma unroll + for (int lane = 0; lane < nvec; ++lane) { + float dact = 0.0f; + float dgate = 0.0f; + float unscaled = 0.0f; + gated_backward_values( + static_cast(act_loader.separate()[lane]), + static_cast(gate_loader.separate()[lane]), param, &dact, &dgate, &unscaled); + const float grad = static_cast(grad_loader.separate()[lane]); + if constexpr (ComputeScaleGrad) { + scale_grad += grad * unscaled; + } + + const float scaled_grad = grad * scale; + act_storer.separate()[lane] = static_cast(scaled_grad * dact); + gate_storer.separate()[lane] = static_cast(scaled_grad * dgate); + } + act_storer.store(vector_idx, segment_size); + gate_storer.store(vector_idx, segment_size); + } + + if constexpr (ComputeScaleGrad) { + scale_grad = block_reduce_sum(scale_grad, smem); + if (threadIdx.x == 0) { + grad_act_scales[row] = static_cast(scale_grad); + } + } +} + +// --------------------------------------------------------------------------- +// Unary kernels +// --------------------------------------------------------------------------- + +template +__global__ void __launch_bounds__(kThreads, 4) + scaled_unary_forward_kernel(const InputT *__restrict__ input, + const ScaleT *__restrict__ act_scales, OutputT *__restrict__ output, + const size_t rows, const size_t hidden, + const size_t num_vectors_per_row, const ParamOP param) { + const size_t total_vectors = rows * num_vectors_per_row; + for (size_t tid = blockIdx.x * blockDim.x + threadIdx.x; tid < total_vectors; + tid += gridDim.x * blockDim.x) { + const size_t vector_idx = tid % num_vectors_per_row; + const size_t row = tid / num_vectors_per_row; + VectorizedLoader input_loader(input + row * hidden, hidden); + VectorizedStorer output_storer(output + row * hidden, hidden); + input_loader.load(vector_idx, hidden); + const float scale = static_cast(act_scales[row]); +#pragma unroll + for (int lane = 0; lane < nvec; ++lane) { + const float unscaled = ActOP(static_cast(input_loader.separate()[lane]), param); + output_storer.separate()[lane] = static_cast(unscaled * scale); + } + output_storer.store(vector_idx, hidden); + } +} + +template +__global__ void __launch_bounds__(kReductionThreads, 4) + scaled_unary_backward_kernel(const GradT *__restrict__ grad_output, + const InputT *__restrict__ input, + const ScaleT *__restrict__ act_scales, + OutputT *__restrict__ grad_input, + GradScaleT *__restrict__ grad_act_scales, const size_t rows, + const size_t hidden, const size_t num_vectors_per_row, + const ParamOP param) { + __shared__ float smem[kReductionWarps]; + const size_t row = blockIdx.x; + (void)rows; + float scale_grad = 0.0f; + const float scale = static_cast(act_scales[row]); + + VectorizedLoader grad_loader(grad_output + row * hidden, hidden); + VectorizedLoader input_loader(input + row * hidden, hidden); + VectorizedStorer grad_input_storer(grad_input + row * hidden, hidden); + for (size_t vector_idx = threadIdx.x; vector_idx < num_vectors_per_row; + vector_idx += blockDim.x) { + grad_loader.load(vector_idx, hidden); + input_loader.load(vector_idx, hidden); +#pragma unroll + for (int lane = 0; lane < nvec; ++lane) { + const float x = static_cast(input_loader.separate()[lane]); + const float grad = static_cast(grad_loader.separate()[lane]); + if constexpr (ComputeScaleGrad) { + scale_grad += grad * ActOP(x, param); + } + grad_input_storer.separate()[lane] = static_cast(grad * scale * DActOP(x, param)); + } + grad_input_storer.store(vector_idx, hidden); + } + + if constexpr (ComputeScaleGrad) { + scale_grad = block_reduce_sum(scale_grad, smem); + if (threadIdx.x == 0) { + grad_act_scales[row] = static_cast(scale_grad); + } + } +} + +// --------------------------------------------------------------------------- +// Tensor checks +// --------------------------------------------------------------------------- + +void check_gated_forward_tensors(const Tensor *input, const Tensor *act_scales, + const Tensor *output, const int64_t glu_interleave_size, + const char *api_name, size_t *rows, size_t *hidden) { + const auto input_dims = input->flat_2d_dims(); + const auto output_dims = output->flat_2d_dims(); + NVTE_CHECK(input_dims[0] == output_dims[0], api_name, ": input/output row mismatch."); + NVTE_CHECK(input_dims[1] == output_dims[1] * 2, api_name, + ": gated input last dimension must be twice output last dimension."); + NVTE_CHECK(glu_interleave_size >= 0, api_name, ": glu_interleave_size must be non-negative."); + if (glu_interleave_size > 0) { + NVTE_CHECK(glu_interleave_size % 32 == 0, api_name, + ": nonzero glu_interleave_size must be a multiple of 32."); + NVTE_CHECK(output_dims[1] % static_cast(glu_interleave_size) == 0, api_name, + ": output last dimension must be divisible by glu_interleave_size."); + } + NVTE_CHECK(act_scales->numel() == input_dims[0], api_name, + ": act_scales must have one value per row."); + *rows = input_dims[0]; + *hidden = output_dims[1]; +} + +void check_gated_backward_tensors(const Tensor *grad_output, const Tensor *input, + const Tensor *act_scales, const Tensor *grad_input, + const Tensor *grad_act_scales, const int64_t glu_interleave_size, + const char *api_name, size_t *rows, size_t *hidden) { + const auto grad_dims = grad_output->flat_2d_dims(); + const auto input_dims = input->flat_2d_dims(); + const auto grad_input_dims = grad_input->flat_2d_dims(); + NVTE_CHECK(grad_dims[0] == input_dims[0] && input_dims[0] == grad_input_dims[0], api_name, + ": input/grad row mismatch."); + NVTE_CHECK(input_dims[1] == grad_dims[1] * 2 && grad_input_dims[1] == input_dims[1], api_name, + ": gated backward dimensions are inconsistent."); + NVTE_CHECK(glu_interleave_size >= 0, api_name, ": glu_interleave_size must be non-negative."); + if (glu_interleave_size > 0) { + NVTE_CHECK(glu_interleave_size % 32 == 0, api_name, + ": nonzero glu_interleave_size must be a multiple of 32."); + NVTE_CHECK(grad_dims[1] % static_cast(glu_interleave_size) == 0, api_name, + ": grad last dimension must be divisible by glu_interleave_size."); + } + NVTE_CHECK(act_scales->numel() == input_dims[0], api_name, + ": act_scales must have one value per row."); + if (grad_act_scales != nullptr) { + NVTE_CHECK(grad_act_scales->numel() == input_dims[0], api_name, + ": grad_act_scales must have one value per row."); + } + *rows = input_dims[0]; + *hidden = grad_dims[1]; +} + +void check_unary_forward_tensors(const Tensor *input, const Tensor *act_scales, + const Tensor *output, const char *api_name, size_t *rows, + size_t *hidden) { + const auto input_dims = input->flat_2d_dims(); + const auto output_dims = output->flat_2d_dims(); + NVTE_CHECK(input_dims[0] == output_dims[0] && input_dims[1] == output_dims[1], api_name, + ": input/output shapes must match."); + NVTE_CHECK(act_scales->numel() == input_dims[0], api_name, + ": act_scales must have one value per row."); + *rows = input_dims[0]; + *hidden = output_dims[1]; +} + +void check_unary_backward_tensors(const Tensor *grad_output, const Tensor *input, + const Tensor *act_scales, const Tensor *grad_input, + const Tensor *grad_act_scales, const char *api_name, size_t *rows, + size_t *hidden) { + const auto grad_dims = grad_output->flat_2d_dims(); + const auto input_dims = input->flat_2d_dims(); + const auto grad_input_dims = grad_input->flat_2d_dims(); + NVTE_CHECK(grad_dims[0] == input_dims[0] && input_dims[0] == grad_input_dims[0], api_name, + ": input/grad row mismatch."); + NVTE_CHECK(grad_dims[1] == input_dims[1] && input_dims[1] == grad_input_dims[1], api_name, + ": unary backward dimensions are inconsistent."); + NVTE_CHECK(act_scales->numel() == input_dims[0], api_name, + ": act_scales must have one value per row."); + if (grad_act_scales != nullptr) { + NVTE_CHECK(grad_act_scales->numel() == input_dims[0], api_name, + ": grad_act_scales must have one value per row."); + } + *rows = input_dims[0]; + *hidden = grad_dims[1]; +} + +} // namespace + +// --------------------------------------------------------------------------- +// Launch implementations +// --------------------------------------------------------------------------- + +using namespace detail::scaled_activation; + +template +void launch_scaled_gated_forward(const NVTETensor nvte_input, const NVTETensor nvte_act_scales, + NVTETensor nvte_output, ParamOP param, int64_t glu_interleave_size, + cudaStream_t stream, const char *api_name) { + const Tensor *input = convertNVTETensorCheck(nvte_input); + const Tensor *act_scales = convertNVTETensorCheck(nvte_act_scales); + Tensor *output = convertNVTETensorCheck(nvte_output); + size_t rows = 0; + size_t hidden = 0; + check_gated_forward_tensors(input, act_scales, output, glu_interleave_size, api_name, &rows, + &hidden); + if (rows == 0 || hidden == 0) return; + + TRANSFORMER_ENGINE_TYPE_SWITCH_NON_FP8ONLY(input->data.dtype, InputT, { + TRANSFORMER_ENGINE_TYPE_SWITCH_NON_FP8ONLY(act_scales->data.dtype, ScaleT, { + TRANSFORMER_ENGINE_TYPE_SWITCH_NON_FP8ONLY(output->data.dtype, OutputT, { + constexpr int nvec = 32 / static_cast(std::max(sizeof(InputT), sizeof(OutputT))); + const auto input_ptr = reinterpret_cast(input->data.dptr); + const auto scale_ptr = reinterpret_cast(act_scales->data.dptr); + auto output_ptr = reinterpret_cast(output->data.dptr); + const size_t segment_size = + glu_interleave_size > 0 ? static_cast(glu_interleave_size) : hidden; + const size_t num_segments = glu_interleave_size > 0 ? hidden / segment_size : 1; + const auto align = row_vector_alignment(segment_size, nvec, input_ptr, + input_ptr + segment_size, output_ptr); + const bool use_vector = align == Alignment::SAME_ALIGNED; + const size_t num_vectors = + use_vector ? get_num_aligned_elements(input_ptr, segment_size, nvec, sizeof(InputT)) + : segment_size; + const int blocks = static_cast(std::min( + DIVUP(rows * num_segments * num_vectors, static_cast(kThreads)), 65535)); + if (use_vector) { + scaled_gated_forward_kernel + <<>>(input_ptr, scale_ptr, output_ptr, rows, hidden, + segment_size, num_segments, num_vectors, param); + } else { + scaled_gated_forward_kernel<1, InputT, ScaleT, OutputT, ParamOP, ActOP> + <<>>(input_ptr, scale_ptr, output_ptr, rows, hidden, + segment_size, num_segments, segment_size, param); + } + }); + }); + }); + NVTE_CHECK_CUDA(cudaGetLastError()); +} + +template +void launch_scaled_gated_backward(const NVTETensor nvte_grad_output, const NVTETensor nvte_input, + const NVTETensor nvte_act_scales, NVTETensor nvte_grad_input, + NVTETensor nvte_grad_act_scales, ParamOP param, + int64_t glu_interleave_size, cudaStream_t stream, + const char *api_name) { + const Tensor *grad_output = convertNVTETensorCheck(nvte_grad_output); + const Tensor *input = convertNVTETensorCheck(nvte_input); + const Tensor *act_scales = convertNVTETensorCheck(nvte_act_scales); + Tensor *grad_input = convertNVTETensorCheck(nvte_grad_input); + Tensor *grad_act_scales = + nvte_grad_act_scales == nullptr ? nullptr : convertNVTETensorCheck(nvte_grad_act_scales); + size_t rows = 0; + size_t hidden = 0; + check_gated_backward_tensors(grad_output, input, act_scales, grad_input, grad_act_scales, + glu_interleave_size, api_name, &rows, &hidden); + if (rows == 0 || hidden == 0) return; + + TRANSFORMER_ENGINE_TYPE_SWITCH_NON_FP8ONLY(grad_output->data.dtype, GradT, { + TRANSFORMER_ENGINE_TYPE_SWITCH_NON_FP8ONLY(input->data.dtype, InputT, { + TRANSFORMER_ENGINE_TYPE_SWITCH_NON_FP8ONLY(act_scales->data.dtype, ScaleT, { + TRANSFORMER_ENGINE_TYPE_SWITCH_NON_FP8ONLY(grad_input->data.dtype, OutputT, { + constexpr int nvec = + 32 / static_cast(std::max({sizeof(GradT), sizeof(InputT), sizeof(OutputT)})); + const auto grad_ptr = reinterpret_cast(grad_output->data.dptr); + const auto input_ptr = reinterpret_cast(input->data.dptr); + const auto scale_ptr = reinterpret_cast(act_scales->data.dptr); + auto grad_input_ptr = reinterpret_cast(grad_input->data.dptr); + const size_t segment_size = + glu_interleave_size > 0 ? static_cast(glu_interleave_size) : hidden; + const size_t num_segments = glu_interleave_size > 0 ? hidden / segment_size : 1; + const auto align = row_vector_alignment(segment_size, nvec, grad_ptr, input_ptr, + input_ptr + segment_size, grad_input_ptr, + grad_input_ptr + segment_size); + const bool use_vector = align == Alignment::SAME_ALIGNED; + const size_t num_vectors = + use_vector ? get_num_aligned_elements(input_ptr, segment_size, nvec, sizeof(InputT)) + : segment_size; + const int reduction_threads = choose_reduction_threads(num_segments * num_vectors); + if (grad_act_scales == nullptr) { + if (use_vector) { + scaled_gated_backward_kernel + <<(rows), reduction_threads, 0, stream>>>( + grad_ptr, input_ptr, scale_ptr, grad_input_ptr, nullptr, rows, hidden, + segment_size, num_segments, num_vectors, param); + } else { + scaled_gated_backward_kernel<1, false, GradT, InputT, ScaleT, OutputT, OutputT, + ParamOP, ActOP, DActOP> + <<(rows), reduction_threads, 0, stream>>>( + grad_ptr, input_ptr, scale_ptr, grad_input_ptr, nullptr, rows, hidden, + segment_size, num_segments, segment_size, param); + } + } else { + TRANSFORMER_ENGINE_TYPE_SWITCH_NON_FP8ONLY(grad_act_scales->data.dtype, GradScaleT, { + auto grad_act_scales_ptr = reinterpret_cast(grad_act_scales->data.dptr); + if (use_vector) { + scaled_gated_backward_kernel + <<(rows), reduction_threads, 0, stream>>>( + grad_ptr, input_ptr, scale_ptr, grad_input_ptr, grad_act_scales_ptr, rows, + hidden, segment_size, num_segments, num_vectors, param); + } else { + scaled_gated_backward_kernel<1, true, GradT, InputT, ScaleT, OutputT, GradScaleT, + ParamOP, ActOP, DActOP> + <<(rows), reduction_threads, 0, stream>>>( + grad_ptr, input_ptr, scale_ptr, grad_input_ptr, grad_act_scales_ptr, rows, + hidden, segment_size, num_segments, segment_size, param); + } + }); + } + }); + }); + }); + }); + NVTE_CHECK_CUDA(cudaGetLastError()); +} + +template +void launch_scaled_unary_forward(const NVTETensor nvte_input, const NVTETensor nvte_act_scales, + NVTETensor nvte_output, ParamOP param, cudaStream_t stream, + const char *api_name) { + const Tensor *input = convertNVTETensorCheck(nvte_input); + const Tensor *act_scales = convertNVTETensorCheck(nvte_act_scales); + Tensor *output = convertNVTETensorCheck(nvte_output); + size_t rows = 0; + size_t hidden = 0; + check_unary_forward_tensors(input, act_scales, output, api_name, &rows, &hidden); + if (rows == 0 || hidden == 0) return; + + TRANSFORMER_ENGINE_TYPE_SWITCH_NON_FP8ONLY(input->data.dtype, InputT, { + TRANSFORMER_ENGINE_TYPE_SWITCH_NON_FP8ONLY(act_scales->data.dtype, ScaleT, { + TRANSFORMER_ENGINE_TYPE_SWITCH_NON_FP8ONLY(output->data.dtype, OutputT, { + constexpr int nvec = 32 / static_cast(std::max(sizeof(InputT), sizeof(OutputT))); + const auto input_ptr = reinterpret_cast(input->data.dptr); + const auto scale_ptr = reinterpret_cast(act_scales->data.dptr); + auto output_ptr = reinterpret_cast(output->data.dptr); + const auto align = row_vector_alignment(hidden, nvec, input_ptr, output_ptr); + const bool use_vector = align == Alignment::SAME_ALIGNED; + const size_t num_vectors = + use_vector ? get_num_aligned_elements(input_ptr, hidden, nvec, sizeof(InputT)) : hidden; + const int blocks = static_cast( + std::min(DIVUP(rows * num_vectors, static_cast(kThreads)), 65535)); + if (use_vector) { + scaled_unary_forward_kernel + <<>>(input_ptr, scale_ptr, output_ptr, rows, hidden, + num_vectors, param); + } else { + scaled_unary_forward_kernel<1, InputT, ScaleT, OutputT, ParamOP, ActOP> + <<>>(input_ptr, scale_ptr, output_ptr, rows, hidden, + hidden, param); + } + }); + }); + }); + NVTE_CHECK_CUDA(cudaGetLastError()); +} + +template +void launch_scaled_unary_backward(const NVTETensor nvte_grad_output, const NVTETensor nvte_input, + const NVTETensor nvte_act_scales, NVTETensor nvte_grad_input, + NVTETensor nvte_grad_act_scales, ParamOP param, + cudaStream_t stream, const char *api_name) { + const Tensor *grad_output = convertNVTETensorCheck(nvte_grad_output); + const Tensor *input = convertNVTETensorCheck(nvte_input); + const Tensor *act_scales = convertNVTETensorCheck(nvte_act_scales); + Tensor *grad_input = convertNVTETensorCheck(nvte_grad_input); + Tensor *grad_act_scales = + nvte_grad_act_scales == nullptr ? nullptr : convertNVTETensorCheck(nvte_grad_act_scales); + size_t rows = 0; + size_t hidden = 0; + check_unary_backward_tensors(grad_output, input, act_scales, grad_input, grad_act_scales, + api_name, &rows, &hidden); + if (rows == 0 || hidden == 0) return; + + TRANSFORMER_ENGINE_TYPE_SWITCH_NON_FP8ONLY(grad_output->data.dtype, GradT, { + TRANSFORMER_ENGINE_TYPE_SWITCH_NON_FP8ONLY(input->data.dtype, InputT, { + TRANSFORMER_ENGINE_TYPE_SWITCH_NON_FP8ONLY(act_scales->data.dtype, ScaleT, { + TRANSFORMER_ENGINE_TYPE_SWITCH_NON_FP8ONLY(grad_input->data.dtype, OutputT, { + constexpr int nvec = + 32 / static_cast(std::max({sizeof(GradT), sizeof(InputT), sizeof(OutputT)})); + const auto grad_ptr = reinterpret_cast(grad_output->data.dptr); + const auto input_ptr = reinterpret_cast(input->data.dptr); + const auto scale_ptr = reinterpret_cast(act_scales->data.dptr); + auto grad_input_ptr = reinterpret_cast(grad_input->data.dptr); + const auto align = + row_vector_alignment(hidden, nvec, grad_ptr, input_ptr, grad_input_ptr); + const bool use_vector = align == Alignment::SAME_ALIGNED; + const size_t num_vectors = + use_vector ? get_num_aligned_elements(input_ptr, hidden, nvec, sizeof(InputT)) + : hidden; + const int reduction_threads = choose_reduction_threads(num_vectors); + if (grad_act_scales == nullptr) { + if (use_vector) { + scaled_unary_backward_kernel + <<(rows), reduction_threads, 0, stream>>>( + grad_ptr, input_ptr, scale_ptr, grad_input_ptr, nullptr, rows, hidden, + num_vectors, param); + } else { + scaled_unary_backward_kernel<1, false, GradT, InputT, ScaleT, OutputT, OutputT, + ParamOP, ActOP, DActOP> + <<(rows), reduction_threads, 0, stream>>>( + grad_ptr, input_ptr, scale_ptr, grad_input_ptr, nullptr, rows, hidden, hidden, + param); + } + } else { + TRANSFORMER_ENGINE_TYPE_SWITCH_NON_FP8ONLY(grad_act_scales->data.dtype, GradScaleT, { + auto grad_act_scales_ptr = reinterpret_cast(grad_act_scales->data.dptr); + if (use_vector) { + scaled_unary_backward_kernel + <<(rows), reduction_threads, 0, stream>>>( + grad_ptr, input_ptr, scale_ptr, grad_input_ptr, grad_act_scales_ptr, rows, + hidden, num_vectors, param); + } else { + scaled_unary_backward_kernel<1, true, GradT, InputT, ScaleT, OutputT, GradScaleT, + ParamOP, ActOP, DActOP> + <<(rows), reduction_threads, 0, stream>>>( + grad_ptr, input_ptr, scale_ptr, grad_input_ptr, grad_act_scales_ptr, rows, + hidden, hidden, param); + } + }); + } + }); + }); + }); + }); + NVTE_CHECK_CUDA(cudaGetLastError()); +} + +// --------------------------------------------------------------------------- +// Explicit instantiations +// --------------------------------------------------------------------------- + +template void launch_scaled_gated_forward>(const NVTETensor, + const NVTETensor, NVTETensor, + Empty, int64_t, cudaStream_t, + const char *); +template void launch_scaled_gated_backward, dsilu>( + const NVTETensor, const NVTETensor, const NVTETensor, NVTETensor, NVTETensor, Empty, int64_t, + cudaStream_t, const char *); + +template void launch_scaled_gated_forward>( + const NVTETensor, const NVTETensor, NVTETensor, ClampedSwiGLUParam, int64_t, cudaStream_t, + const char *); +template void launch_scaled_gated_backward, + clamped_dsilu>( + const NVTETensor, const NVTETensor, const NVTETensor, NVTETensor, NVTETensor, + ClampedSwiGLUParam, int64_t, cudaStream_t, const char *); + +template void launch_scaled_unary_forward>(const NVTETensor, + const NVTETensor, NVTETensor, + Empty, cudaStream_t, + const char *); +template void launch_scaled_unary_backward, dsrelu>( + const NVTETensor, const NVTETensor, const NVTETensor, NVTETensor, NVTETensor, Empty, + cudaStream_t, const char *); + +} // namespace transformer_engine diff --git a/transformer_engine/common/activation/scaled_activation.h b/transformer_engine/common/activation/scaled_activation.h new file mode 100644 index 0000000000..391bff6841 --- /dev/null +++ b/transformer_engine/common/activation/scaled_activation.h @@ -0,0 +1,90 @@ +/************************************************************************* + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * + * See LICENSE for license information. + ************************************************************************/ + +/* Scaled activations apply an activation, multiply by a per-row scale + * (act_scales[row]), do all math in fp32, and cast once at the store. The + * backward path optionally also reduces the gradient of the per-row scale. + * + * Public launch APIs are templated on ParamOP / ActOP / DActOP (same shape as + * gated_act_fn). Kernel definitions and explicit instantiations live in + * scaled_activation.cu. + */ + +#ifndef TRANSFORMER_ENGINE_COMMON_ACTIVATION_SCALED_ACTIVATION_H_ +#define TRANSFORMER_ENGINE_COMMON_ACTIVATION_SCALED_ACTIVATION_H_ + +#include + +#include "../common.h" +#include "../util/math.h" +#include "../util/vectorized_pointwise.h" +#include "../utils.cuh" + +namespace transformer_engine { +namespace detail { +namespace scaled_activation { + +constexpr int kThreads = unary_kernel_threads; +constexpr int kReductionThreads = 256; +constexpr int kReductionWarps = kReductionThreads / THREADS_PER_WARP; + +// Pick a CTA size for one-block-per-row scale-grad: enough threads to cover +// row_vectors, rounded up to a warp multiple, capped at kReductionThreads. +inline int choose_reduction_threads(const size_t row_vectors) { + if (row_vectors >= static_cast(kReductionThreads)) { + return kReductionThreads; + } + const int needed = static_cast(row_vectors); + int rounded = (needed + THREADS_PER_WARP - 1) / THREADS_PER_WARP * THREADS_PER_WARP; + if (rounded < THREADS_PER_WARP) { + rounded = THREADS_PER_WARP; + } + return rounded; +} + +template +Alignment row_vector_alignment(const size_t lead_dim, const int nvec, const Ptrs... ptrs) { + if (nvec == 1) { + return Alignment::SAME_ALIGNED; + } + if (lead_dim % static_cast(nvec) != 0) { + return Alignment::DIFFERENT; + } + const auto align = CheckAlignment(lead_dim, nvec, ptrs...); + return align == Alignment::SAME_ALIGNED ? Alignment::SAME_ALIGNED : Alignment::DIFFERENT; +} + +} // namespace scaled_activation +} // namespace detail + +template +void launch_scaled_gated_forward(const NVTETensor input, const NVTETensor act_scales, + NVTETensor output, ParamOP param, int64_t glu_interleave_size, + cudaStream_t stream, const char *api_name); + +template +void launch_scaled_gated_backward(const NVTETensor grad, const NVTETensor input, + const NVTETensor act_scales, NVTETensor grad_input, + NVTETensor grad_act_scales, ParamOP param, + int64_t glu_interleave_size, cudaStream_t stream, + const char *api_name); + +template +void launch_scaled_unary_forward(const NVTETensor input, const NVTETensor act_scales, + NVTETensor output, ParamOP param, cudaStream_t stream, + const char *api_name); + +template +void launch_scaled_unary_backward(const NVTETensor grad, const NVTETensor input, + const NVTETensor act_scales, NVTETensor grad_input, + NVTETensor grad_act_scales, ParamOP param, cudaStream_t stream, + const char *api_name); + +} // namespace transformer_engine + +#endif // TRANSFORMER_ENGINE_COMMON_ACTIVATION_SCALED_ACTIVATION_H_ diff --git a/transformer_engine/common/activation/scaled_srelu.cu b/transformer_engine/common/activation/scaled_srelu.cu new file mode 100644 index 0000000000..2e81bab669 --- /dev/null +++ b/transformer_engine/common/activation/scaled_srelu.cu @@ -0,0 +1,29 @@ +/************************************************************************* + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * + * See LICENSE for license information. + ************************************************************************/ + +#include + +#include "../common.h" +#include "../util/math.h" +#include "./scaled_activation.h" + +void nvte_scaled_srelu(const NVTETensor input, const NVTETensor act_scales, NVTETensor output, + cudaStream_t stream) { + NVTE_API_CALL(nvte_scaled_srelu); + using namespace transformer_engine; + Empty param = {}; + launch_scaled_unary_forward>(input, act_scales, output, param, stream, + "nvte_scaled_srelu"); +} + +void nvte_scaled_dsrelu(const NVTETensor grad, const NVTETensor input, const NVTETensor act_scales, + NVTETensor grad_input, NVTETensor grad_act_scales, cudaStream_t stream) { + NVTE_API_CALL(nvte_scaled_dsrelu); + using namespace transformer_engine; + Empty param = {}; + launch_scaled_unary_backward, dsrelu>( + grad, input, act_scales, grad_input, grad_act_scales, param, stream, "nvte_scaled_dsrelu"); +} diff --git a/transformer_engine/common/activation/scaled_swiglu.cu b/transformer_engine/common/activation/scaled_swiglu.cu new file mode 100644 index 0000000000..ee0692b1de --- /dev/null +++ b/transformer_engine/common/activation/scaled_swiglu.cu @@ -0,0 +1,56 @@ +/************************************************************************* + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * + * See LICENSE for license information. + ************************************************************************/ + +#include + +#include "../common.h" +#include "../util/math.h" +#include "./scaled_activation.h" + +void nvte_scaled_swiglu(const NVTETensor input, const NVTETensor act_scales, NVTETensor output, + int64_t glu_interleave_size, cudaStream_t stream) { + NVTE_API_CALL(nvte_scaled_swiglu); + using namespace transformer_engine; + Empty param = {}; + launch_scaled_gated_forward>( + input, act_scales, output, param, glu_interleave_size, stream, "nvte_scaled_swiglu"); +} + +void nvte_scaled_dswiglu(const NVTETensor grad, const NVTETensor input, const NVTETensor act_scales, + NVTETensor grad_input, NVTETensor grad_act_scales, + int64_t glu_interleave_size, cudaStream_t stream) { + NVTE_API_CALL(nvte_scaled_dswiglu); + using namespace transformer_engine; + Empty param = {}; + launch_scaled_gated_backward, dsilu>( + grad, input, act_scales, grad_input, grad_act_scales, param, glu_interleave_size, stream, + "nvte_scaled_dswiglu"); +} + +void nvte_scaled_clamped_swiglu(const NVTETensor input, const NVTETensor act_scales, + NVTETensor output, float limit, float alpha, + float glu_linear_offset, int64_t glu_interleave_size, + cudaStream_t stream) { + NVTE_API_CALL(nvte_scaled_clamped_swiglu); + using namespace transformer_engine; + ClampedSwiGLUParam param = {limit, alpha, glu_linear_offset}; + launch_scaled_gated_forward>( + input, act_scales, output, param, glu_interleave_size, stream, "nvte_scaled_clamped_swiglu"); +} + +void nvte_scaled_clamped_dswiglu(const NVTETensor grad, const NVTETensor input, + const NVTETensor act_scales, NVTETensor grad_input, + NVTETensor grad_act_scales, float limit, float alpha, + float glu_linear_offset, int64_t glu_interleave_size, + cudaStream_t stream) { + NVTE_API_CALL(nvte_scaled_clamped_dswiglu); + using namespace transformer_engine; + ClampedSwiGLUParam param = {limit, alpha, glu_linear_offset}; + launch_scaled_gated_backward, + clamped_dsilu>( + grad, input, act_scales, grad_input, grad_act_scales, param, glu_interleave_size, stream, + "nvte_scaled_clamped_dswiglu"); +} diff --git a/transformer_engine/common/cast/cast_grouped_dbias.cu b/transformer_engine/common/cast/cast_grouped_dbias.cu index 5290255a00..b7ced30b11 100644 --- a/transformer_engine/common/cast/cast_grouped_dbias.cu +++ b/transformer_engine/common/cast/cast_grouped_dbias.cu @@ -11,7 +11,8 @@ #include "dispatch/quantize.cuh" void nvte_group_quantize_dbias(const NVTEGroupedTensor input, NVTEGroupedTensor output, - NVTEGroupedTensor dbias, NVTETensor workspace, cudaStream_t stream) { + NVTEGroupedTensor dbias, NVTETensor workspace, + const NVTEQuantizationConfig quant_config, cudaStream_t stream) { NVTE_API_CALL(nvte_group_quantize_dbias); using namespace transformer_engine; @@ -20,5 +21,5 @@ void nvte_group_quantize_dbias(const NVTEGroupedTensor input, NVTEGroupedTensor constexpr const NVTEGroupedTensor activation_input = nullptr; dispatch::group_quantize_bwd_helper( - input, activation_input, output, dbias, workspace, nullptr, stream); + input, activation_input, output, dbias, workspace, quant_config, stream); } diff --git a/transformer_engine/common/cast/dispatch/quantize.cuh b/transformer_engine/common/cast/dispatch/quantize.cuh index 033d464bcf..54f374b551 100644 --- a/transformer_engine/common/cast/dispatch/quantize.cuh +++ b/transformer_engine/common/cast/dispatch/quantize.cuh @@ -475,22 +475,16 @@ void group_quantize_fwd_helper(const NVTEGroupedTensor input, NVTEGroupedTensor } case NVTE_BLOCK_SCALING_1D: { NVTE_CHECK(!IS_ACT, "IS_ACT is not implemented for grouped NVTE_BLOCK_SCALING_1D."); - NVTE_CHECK(!quant_config_cpp.force_pow_2_scales, - "Fused grouped FP8 block-scaling quantize does not support " - "force_pow_2_scales=True. Set force_pow_2_scales=False, or use the unfused " - "split-quantize path (NVTE_GROUPED_LINEAR_USE_FUSED_GROUPED_GEMM=0)."); fp8_blockwise::group_quantize_blockwise_1d(input_tensor, output_tensor, noop_tensor, - quant_config_cpp.amax_epsilon, stream); + quant_config_cpp.amax_epsilon, + quant_config_cpp.force_pow_2_scales, stream); break; } case NVTE_BLOCK_SCALING_2D: { NVTE_CHECK(!IS_ACT, "IS_ACT is not implemented for grouped NVTE_BLOCK_SCALING_2D."); - NVTE_CHECK(!quant_config_cpp.force_pow_2_scales, - "Fused grouped FP8 block-scaling quantize does not support " - "force_pow_2_scales=True. Set force_pow_2_scales=False, or use the unfused " - "split-quantize path (NVTE_GROUPED_LINEAR_USE_FUSED_GROUPED_GEMM=0)."); fp8_blockwise::group_quantize_blockwise_2d(input_tensor, output_tensor, noop_tensor, - quant_config_cpp.amax_epsilon, stream); + quant_config_cpp.amax_epsilon, + quant_config_cpp.force_pow_2_scales, stream); break; } default: @@ -537,22 +531,18 @@ void group_quantize_bwd_helper(const NVTEGroupedTensor grad, const NVTEGroupedTe case NVTE_BLOCK_SCALING_1D: case NVTE_BLOCK_SCALING_2D: { NVTE_CHECK(!IS_DACT, "IS_DACT is not implemented for grouped FP8 block scaling."); - NVTE_CHECK(!quant_config_cpp.force_pow_2_scales, - "Fused grouped FP8 block-scaling quantize does not support " - "force_pow_2_scales=True. Set force_pow_2_scales=False, or use the unfused " - "split-quantize path (NVTE_GROUPED_LINEAR_USE_FUSED_GROUPED_GEMM=0)."); // dbias is computed in-kernel and reduced per-expert inside group_quantize_blockwise_{1d,2d} // (mirrors MXFP8); those also handle the two-call workspace sizing protocol. GroupedTensor *dbias_arg = IS_DBIAS ? dbias_tensor : nullptr; Tensor *workspace_arg = IS_DBIAS ? workspace_tensor : nullptr; if (scaling_mode == NVTE_BLOCK_SCALING_1D) { - fp8_blockwise::group_quantize_blockwise_1d(grad_tensor, output_tensor, noop_tensor, - quant_config_cpp.amax_epsilon, stream, dbias_arg, - workspace_arg); + fp8_blockwise::group_quantize_blockwise_1d( + grad_tensor, output_tensor, noop_tensor, quant_config_cpp.amax_epsilon, + quant_config_cpp.force_pow_2_scales, stream, dbias_arg, workspace_arg); } else { - fp8_blockwise::group_quantize_blockwise_2d(grad_tensor, output_tensor, noop_tensor, - quant_config_cpp.amax_epsilon, stream, dbias_arg, - workspace_arg); + fp8_blockwise::group_quantize_blockwise_2d( + grad_tensor, output_tensor, noop_tensor, quant_config_cpp.amax_epsilon, + quant_config_cpp.force_pow_2_scales, stream, dbias_arg, workspace_arg); } break; } diff --git a/transformer_engine/common/cast/fp8_blockwise/group_quantize_fp8_blockwise.cuh b/transformer_engine/common/cast/fp8_blockwise/group_quantize_fp8_blockwise.cuh index 1fd1738f93..33aadb5bfa 100644 --- a/transformer_engine/common/cast/fp8_blockwise/group_quantize_fp8_blockwise.cuh +++ b/transformer_engine/common/cast/fp8_blockwise/group_quantize_fp8_blockwise.cuh @@ -284,7 +284,8 @@ __global__ void __launch_bounds__(kThreadsPerBlock, 4) group_block_scaled_2d_tma CType* __restrict__ scale_inv_t_base, const int64_t* __restrict__ tensor_offsets_ptr, const size_t num_tensors, const size_t common_first_dim_blocks, const size_t K, const size_t total_row_blocks, const size_t blocks_X, const size_t scale_stride_y, - const float epsilon, const float* __restrict__ noop_ptr, float* __restrict__ dbias_workspace) { + const float epsilon, const bool pow_2_scales, const float* __restrict__ noop_ptr, + float* __restrict__ dbias_workspace) { #if __CUDA_ARCH__ >= 900 && __CUDA_ARCH__ < 1000 if (noop_ptr != nullptr && noop_ptr[0] == 1.0f) return; @@ -378,8 +379,7 @@ __global__ void __launch_bounds__(kThreadsPerBlock, 4) group_block_scaled_2d_tma for (int w = 1; w < kNumWarps; ++w) { block_amax = fmaxf(block_amax, warp_amaxes[w]); } - const CType scale = - compute_scale_from_types(block_amax, epsilon, /*pow_2_scaling=*/false); + const CType scale = compute_scale_from_types(block_amax, epsilon, pow_2_scales); // The 2D colwise per-expert scale offset requires a CTA-cooperative prefix // sum in the VARYING_FIRST_DIM case, so compute it across all threads before @@ -475,7 +475,7 @@ __global__ void __launch_bounds__(kThreadsPerBlock) const size_t num_tensors, const size_t common_first_dim_blocks, const size_t K, const size_t total_row_blocks, const size_t R_total, const float epsilon, - const float* __restrict__ noop_ptr) { + const bool pow_2_scales, const float* __restrict__ noop_ptr) { #if __CUDA_ARCH__ >= 900 && __CUDA_ARCH__ < 1000 if (noop_ptr != nullptr && noop_ptr[0] == 1.0f) return; @@ -530,8 +530,7 @@ __global__ void __launch_bounds__(kThreadsPerBlock) CType amax = compute_row_amax(in_vec[it]); amax = subwarp_reduce_max_broadcast(amax); - const CType scale = - compute_scale_from_types(amax, epsilon, /*pow_2_scaling=*/false); + const CType scale = compute_scale_from_types(amax, epsilon, pow_2_scales); const CType scale_inv = 1.f / scale; if (thr_col == 0 && r_global < R_total) { // Per-expert layout: (blocks_X, roundup(M_t, 4)). Compute expert base @@ -572,8 +571,8 @@ __global__ void __launch_bounds__(kThreadsPerBlock) group_block_scaled_1d_tma_ke CType* __restrict__ scale_inv_t_base, const int64_t* __restrict__ tensor_offsets_ptr, const size_t num_tensors, const size_t common_first_dim_blocks, const size_t K, const size_t total_row_blocks, const size_t blocks_X, const size_t scale_t_stride_aligned_K, - const size_t R_total, const float epsilon, const float* __restrict__ noop_ptr, - float* __restrict__ dbias_workspace) { + const size_t R_total, const float epsilon, const bool pow_2_scales, + const float* __restrict__ noop_ptr, float* __restrict__ dbias_workspace) { #if __CUDA_ARCH__ >= 900 && __CUDA_ARCH__ < 1000 if (noop_ptr != nullptr && noop_ptr[0] == 1.0f) return; @@ -656,8 +655,7 @@ __global__ void __launch_bounds__(kThreadsPerBlock) group_block_scaled_1d_tma_ke CType amax = compute_row_amax(in_vec); amax = subwarp_reduce_max_broadcast(amax); - const CType scale = - compute_scale_from_types(amax, epsilon, /*pow_2_scaling=*/false); + const CType scale = compute_scale_from_types(amax, epsilon, pow_2_scales); const CType scale_inv = 1.f / scale; const size_t r_global = global_row_base + row_local; @@ -733,8 +731,7 @@ __global__ void __launch_bounds__(kThreadsPerBlock) group_block_scaled_1d_tma_ke } amax = subwarp_reduce_max_broadcast(amax); - const CType scale = - compute_scale_from_types(amax, epsilon, /*pow_2_scaling=*/false); + const CType scale = compute_scale_from_types(amax, epsilon, pow_2_scales); const CType scale_inv = 1.f / scale; const size_t c_global = global_col_base + col_local; @@ -814,7 +811,8 @@ inline GroupedBlockwiseLaunchInfo prepare_grouped_blockwise_launch(const Grouped // reports the [total_row_blocks, K] fp32 shape and returns without launching. inline void group_quantize_blockwise_2d(const GroupedTensor* input, GroupedTensor* output, const Tensor* noop, const float epsilon, - cudaStream_t stream, GroupedTensor* dbias = nullptr, + const bool pow_2_scales, cudaStream_t stream, + GroupedTensor* dbias = nullptr, Tensor* workspace = nullptr) { const int sm = transformer_engine::cuda::sm_arch(); NVTE_CHECK(sm >= 90 && sm < 100, @@ -883,7 +881,7 @@ inline void group_quantize_blockwise_2d(const GroupedTensor* input, GroupedTenso : nullptr, info.tensor_offsets_d, info.num_tensors, info.common_first_dim_blocks, info.K, info.total_row_blocks, info.blocks_X, scale_stride_y, epsilon, - noop_ptr, dbias_workspace); + pow_2_scales, noop_ptr, dbias_workspace); if (dbias_workspace != nullptr) { const ShapeRepresentation shape_rep = info.same_both_dims ? ShapeRepresentation::SAME_BOTH_DIMS @@ -907,7 +905,8 @@ inline void group_quantize_blockwise_2d(const GroupedTensor* input, GroupedTenso // per-tile column partial can be computed. inline void group_quantize_blockwise_1d(const GroupedTensor* input, GroupedTensor* output, const Tensor* noop, const float epsilon, - cudaStream_t stream, GroupedTensor* dbias = nullptr, + const bool pow_2_scales, cudaStream_t stream, + GroupedTensor* dbias = nullptr, Tensor* workspace = nullptr) { const int sm = transformer_engine::cuda::sm_arch(); NVTE_CHECK(sm >= 90 && sm < 100, @@ -967,7 +966,7 @@ inline void group_quantize_blockwise_1d(const GroupedTensor* input, GroupedTenso reinterpret_cast(output->scale_inv.dptr), info.tensor_offsets_d, info.num_tensors, info.common_first_dim_blocks, info.K, info.total_row_blocks, - info.R_total, epsilon, noop_ptr); + info.R_total, epsilon, pow_2_scales, noop_ptr); } } else if constexpr (kRowwise || kColwise) { // CW-only, BOTH, or RW-only WITH dbias: smem-cached TMA kernel. @@ -998,7 +997,7 @@ inline void group_quantize_blockwise_1d(const GroupedTensor* input, GroupedTenso : nullptr, info.tensor_offsets_d, info.num_tensors, info.common_first_dim_blocks, info.K, info.total_row_blocks, info.blocks_X, scale_t_stride_aligned_K, - info.R_total, epsilon, noop_ptr, dbias_workspace); + info.R_total, epsilon, pow_2_scales, noop_ptr, dbias_workspace); if (dbias_workspace != nullptr) { const ShapeRepresentation shape_rep = info.same_both_dims ? ShapeRepresentation::SAME_BOTH_DIMS diff --git a/transformer_engine/common/fused_attn/fused_attn_f16_arbitrary_seqlen.cu b/transformer_engine/common/fused_attn/fused_attn_f16_arbitrary_seqlen.cu index 6df7ad35c8..bf34758a35 100644 --- a/transformer_engine/common/fused_attn/fused_attn_f16_arbitrary_seqlen.cu +++ b/transformer_engine/common/fused_attn/fused_attn_f16_arbitrary_seqlen.cu @@ -94,6 +94,20 @@ void fused_attn_arbitrary_seqlen_fwd_impl( NVTE_CHECK(is_padding, "Paged attention requires padding mask!"); } + // Newer versions of cuDNN SDPA can accept sequence lengths directly as a cumulative + // tensor, and can accept ragged offsets in arbitrary units (such as tokens) instead + // of elements. Take advantage of this if possible to avoid 2 extra kernel calls. + const bool use_cu_seqlens_directly = + CUDNN_FRONTEND_VERSION >= 12500 && + // The frontend gates cu_seq_len support on min(compile-time, runtime) cuDNN + // version, so we'll do the same. + (CUDNN_VERSION >= 92400 && cudnn_runtime_version >= 92400) && + // This extra restriction is needed because cuDNN frontend doesn't yet allow + // the combination of dropout and stats generation for the fprop unified engine, + // so any such request would always get routed to the old composite SDPA engine + // (which doesn't support cu_seqlens). Remove this restriction when possible. + !is_dropout; + // keep original batch size because cu_seqlens are created with [b+1] shape int64_t actual_b = b; if ((is_ragged_q || is_ragged_kv) && cudnn_runtime_version >= 90600) { @@ -103,14 +117,27 @@ void fused_attn_arbitrary_seqlen_fwd_impl( // so the check passes; ragged offset still provides variable-length boundaries. if (sm_arch_ != 120) { // replace batch size and maximum sequence lengths with maximum token counts - // for query and key/value so the graph is static within each quantization bucket - b = max_b; + // for query and key/value so the graph is static within each quantization bucket. + // When passing cu_seqlens* directly to cuDNN SDPA, keep the true batch size: + // cuDNN reads the user's [actual_b+1] cu_seqlens buffers, so a quantized batch + // would read out of bounds. + if (!use_cu_seqlens_directly) { + b = max_b; + } s_q = is_ragged_q ? max_t_q : s_q; s_kv = is_ragged_kv ? max_t_kv : s_kv; } } - const DType ragged_offset_type = cudnn_runtime_version >= 90500 ? DType::kInt64 : DType::kInt32; + const DType ragged_offset_type = + use_cu_seqlens_directly + ? DType::kInt32 // cu_seqlens* are given to us as int32; keep it that way. + : (cudnn_runtime_version >= 90500 ? DType::kInt64 : DType::kInt32); + + // Ragged offset multipliers (elements per token); shared with the legacy conversion + // kernel (cu_seqlens_padded_to_offsets) so the two paths cannot drift apart. + const RaggedOffsetMultipliers offset_mults(layout_group, h, hg, d_qk, d_v); + bool generate_stats = true; // Always return stats try { FADescriptor_v1 descriptor{ @@ -166,8 +193,8 @@ void fused_attn_arbitrary_seqlen_fwd_impl( std::shared_ptr, // S2 std::shared_ptr, // bias std::shared_ptr, // softmax_offset - std::shared_ptr, // seq_q - std::shared_ptr, // seq_kv + std::shared_ptr, // seq_q / cu_seq_len_q + std::shared_ptr, // seq_kv / cu_seq_len_kv std::shared_ptr, // page_table_k std::shared_ptr, // page_table_v std::shared_ptr, // offset_q @@ -231,6 +258,9 @@ void fused_attn_arbitrary_seqlen_fwd_impl( .set_stride({1, 1, 1, 1}) .set_data_type(get_cudnn_fe_dtype(ragged_offset_type))); Q->set_ragged_offset(offset_q); + if (use_cu_seqlens_directly) { + Q->set_ragged_offset_multiplier(offset_mults.q); + } } K = mha_graph->tensor(fe::graph::Tensor_attributes().set_name("K").set_stride(k_stride)); V = mha_graph->tensor(fe::graph::Tensor_attributes().set_name("V").set_stride(v_stride)); @@ -250,6 +280,10 @@ void fused_attn_arbitrary_seqlen_fwd_impl( .set_data_type(get_cudnn_fe_dtype(ragged_offset_type))); K->set_dim({b, hg, s_kv, d_qk}).set_ragged_offset(offset_k); V->set_dim({b, hg, s_kv, d_v}).set_ragged_offset(offset_v); + if (use_cu_seqlens_directly) { + K->set_ragged_offset_multiplier(offset_mults.k); + V->set_ragged_offset_multiplier(offset_mults.v); + } } else { K->set_dim({b, hg, s_kv, d_qk}); V->set_dim({b, hg, s_kv, d_v}); @@ -293,17 +327,38 @@ void fused_attn_arbitrary_seqlen_fwd_impl( } if (is_padding) { - seq_q = mha_graph->tensor(fe::graph::Tensor_attributes() - .set_name("seq_q") - .set_dim({b, 1, 1, 1}) - .set_stride({1, 1, 1, 1}) - .set_data_type(fe::DataType_t::INT32)); - seq_kv = mha_graph->tensor(fe::graph::Tensor_attributes() - .set_name("seq_kv") - .set_dim({b, 1, 1, 1}) - .set_stride({1, 1, 1, 1}) - .set_data_type(fe::DataType_t::INT32)); - sdpa_options.set_padding_mask(is_padding).set_seq_len_q(seq_q).set_seq_len_kv(seq_kv); + if (use_cu_seqlens_directly) { + // seq_q/seq_kv keep their tuple slots but hold (b+1)-shaped cu_seqlen tensors. + seq_q = mha_graph->tensor(fe::graph::Tensor_attributes() + .set_name("cu_seq_len_q") + .set_dim({b + 1, 1, 1, 1}) + .set_stride({1, 1, 1, 1}) + .set_data_type(fe::DataType_t::INT32)); + seq_kv = mha_graph->tensor(fe::graph::Tensor_attributes() + .set_name("cu_seq_len_kv") + .set_dim({b + 1, 1, 1, 1}) + .set_stride({1, 1, 1, 1}) + .set_data_type(fe::DataType_t::INT32)); + sdpa_options.set_padding_mask(is_padding) + .set_cu_seq_len_q(seq_q) + .set_cu_seq_len_kv(seq_kv); + // cu_seq_len (and the ragged offset multiplier) are unified-engine-only. + // Pin the implementation so an unsupported config fails with the unified + // engine's specific error instead of auto-selection's generic failure. + sdpa_options.set_implementation(fe::AttentionImplementation_t::UNIFIED); + } else { + seq_q = mha_graph->tensor(fe::graph::Tensor_attributes() + .set_name("seq_q") + .set_dim({b, 1, 1, 1}) + .set_stride({1, 1, 1, 1}) + .set_data_type(fe::DataType_t::INT32)); + seq_kv = mha_graph->tensor(fe::graph::Tensor_attributes() + .set_name("seq_kv") + .set_dim({b, 1, 1, 1}) + .set_stride({1, 1, 1, 1}) + .set_data_type(fe::DataType_t::INT32)); + sdpa_options.set_padding_mask(is_padding).set_seq_len_q(seq_q).set_seq_len_kv(seq_kv); + } } if (is_paged_kv) { @@ -363,6 +418,9 @@ void fused_attn_arbitrary_seqlen_fwd_impl( .set_data_type(fe::DataType_t::FLOAT)); if (use_ragged_stats) { Max->set_stride({h * s_q, 1, h, 1}).set_ragged_offset(offset_stats); + if (use_cu_seqlens_directly) { + Max->set_ragged_offset_multiplier(offset_mults.stats); + } } else { Max->set_stride({h * s_q, s_q, 1, 1}); } @@ -382,11 +440,17 @@ void fused_attn_arbitrary_seqlen_fwd_impl( .set_stride({1, 1, 1, 1}) .set_data_type(get_cudnn_fe_dtype(ragged_offset_type))); O->set_ragged_offset(offset_o); + if (use_cu_seqlens_directly) { + O->set_ragged_offset_multiplier(offset_mults.o); + } } Stats->set_output(true).set_data_type(fe::DataType_t::FLOAT).set_dim({b, h, s_q, 1}); - if (is_ragged_q && cudnn_runtime_version >= 90600) { + if (use_ragged_stats) { Stats->set_stride({h * s_q, 1, h, 1}).set_ragged_offset(offset_stats); + if (use_cu_seqlens_directly) { + Stats->set_ragged_offset_multiplier(offset_mults.stats); + } } else { Stats->set_stride({h * s_q, s_q, 1, 1}); } @@ -437,18 +501,23 @@ void fused_attn_arbitrary_seqlen_fwd_impl( // Exit to request upper level API to allocate memory if needed // n.b. Care should be taken to align each of the added worksapce tensors to their type. // We do this by adding padding at the end of each separate allocation. + // When passing cu_seqlens* directly to cuDNN SDPA, no conversion workspace is + // needed: cuDNN consumes the user's cu_seqlens buffers as-is. auto plan_workspace_size = alignTo<16>(mha_graph->get_workspace_size()); const size_t num_bytes_per_seqlen = alignTo<16>(b * sizeof(int32_t)); - const size_t actual_seqlen_workspace_size = is_padding ? 2 * num_bytes_per_seqlen : 0; const size_t num_bytes_per_ragged_offset = alignTo<16>(((b + 1) * typeToNumBits(ragged_offset_type)) / 8); + size_t actual_seqlen_workspace_size = 0; size_t seqlen_offsets_workspace_size = 0; - if (is_ragged_q || is_ragged_kv) { - size_t count = 2 * (static_cast(is_ragged_q) + static_cast(is_ragged_kv)); - if (use_ragged_stats) { - seqlen_offsets_workspace_size = (count + 1) * num_bytes_per_ragged_offset; - } else { - seqlen_offsets_workspace_size = count * num_bytes_per_ragged_offset; + if (!use_cu_seqlens_directly) { + if (is_padding) { + actual_seqlen_workspace_size = 2 * num_bytes_per_seqlen; + } + if (is_ragged_q || is_ragged_kv) { + const size_t count = + 2 * (static_cast(is_ragged_q) + static_cast(is_ragged_kv)); + seqlen_offsets_workspace_size = + (use_ragged_stats ? count + 1 : count) * num_bytes_per_ragged_offset; } } if (workspace == nullptr) { @@ -475,17 +544,22 @@ void fused_attn_arbitrary_seqlen_fwd_impl( } if (is_padding) { - constexpr size_t nthreads_per_block = 128; - const size_t grid = (b + nthreads_per_block - 1) / nthreads_per_block; - void *devActualSeqlenQ = static_cast(workspace) + plan_workspace_size; - void *devActualSeqlenKV = static_cast(devActualSeqlenQ) + num_bytes_per_seqlen; - cu_seqlens_to_actual_seqlens<<>>( - actual_b, b, static_cast(devPtrCuSeqlensQ), - static_cast(devPtrCuSeqlensKV), static_cast(devActualSeqlenQ), - static_cast(devActualSeqlenKV)); - NVTE_CHECK_CUDA(cudaGetLastError()); - variant_pack[seq_q] = devActualSeqlenQ; - variant_pack[seq_kv] = devActualSeqlenKV; + if (use_cu_seqlens_directly) { + variant_pack[seq_q] = devPtrCuSeqlensQ; + variant_pack[seq_kv] = devPtrCuSeqlensKV; + } else { + constexpr size_t nthreads_per_block = 128; + const size_t grid = (b + nthreads_per_block - 1) / nthreads_per_block; + void *devActualSeqlenQ = static_cast(workspace) + plan_workspace_size; + void *devActualSeqlenKV = static_cast(devActualSeqlenQ) + num_bytes_per_seqlen; + cu_seqlens_to_actual_seqlens<<>>( + actual_b, b, static_cast(devPtrCuSeqlensQ), + static_cast(devPtrCuSeqlensKV), + static_cast(devActualSeqlenQ), static_cast(devActualSeqlenKV)); + NVTE_CHECK_CUDA(cudaGetLastError()); + variant_pack[seq_q] = devActualSeqlenQ; + variant_pack[seq_kv] = devActualSeqlenKV; + } } if (is_paged_kv) { @@ -493,7 +567,22 @@ void fused_attn_arbitrary_seqlen_fwd_impl( variant_pack[page_table_v] = devPtrPageTableV; } - if (is_ragged_q || is_ragged_kv) { + if (use_cu_seqlens_directly) { + // The token-unit cu_seqlens_padded buffers serve as the ragged offsets; the engine + // applies the per-tensor multipliers set at graph build time. + if (is_ragged_q) { + variant_pack[offset_q] = devPtrSeqOffsetsQ; + variant_pack[offset_o] = devPtrSeqOffsetsQ; + } + if (is_ragged_kv) { + void *devOffsetsKV = offset_mults.kv_from_q ? devPtrSeqOffsetsQ : devPtrSeqOffsetsKV; + variant_pack[offset_k] = devOffsetsKV; + variant_pack[offset_v] = devOffsetsKV; + } + if (use_ragged_stats) { + variant_pack[offset_stats] = devPtrSeqOffsetsQ; + } + } else if (is_ragged_q || is_ragged_kv) { constexpr size_t nthreads_per_block = 128; const size_t grid = (b + nthreads_per_block) / nthreads_per_block; void *devOffsets = @@ -517,9 +606,8 @@ void fused_attn_arbitrary_seqlen_fwd_impl( (static_cast(is_ragged_q) + static_cast(is_ragged_kv)) * 2 * num_bytes_per_ragged_offset; } - const NVTE_QKV_Layout_Group layout_group = nvte_get_qkv_layout_group(qkv_layout); cu_seqlens_padded_to_offsets<<>>( - layout_group, actual_b, b, h, hg, d_qk, d_v, static_cast(devPtrSeqOffsetsQ), + offset_mults, actual_b, b, static_cast(devPtrSeqOffsetsQ), static_cast(devPtrSeqOffsetsKV), ragged_offset_type, devOffsetsQ, devOffsetsK, devOffsetsV, devOffsetsO, devOffsetsS); NVTE_CHECK_CUDA(cudaGetLastError()); @@ -549,7 +637,7 @@ void fused_attn_arbitrary_seqlen_fwd_impl( } catch (cudnn_frontend::cudnnException &e) { NVTE_ERROR(e.what()); } -} +} // NOLINT(readability/fn_size) void fused_attn_arbitrary_seqlen_bwd_impl( int64_t b, int64_t h, int64_t hg, int64_t s_q, int64_t s_kv, int64_t d_qk, int64_t d_v, @@ -1034,9 +1122,10 @@ void fused_attn_arbitrary_seqlen_bwd_impl( (static_cast(is_ragged_q) + static_cast(is_ragged_kv)) * 2 * num_bytes_per_ragged_offset; } - const NVTE_QKV_Layout_Group layout_group = nvte_get_qkv_layout_group(qkv_layout); + const RaggedOffsetMultipliers offset_mults(nvte_get_qkv_layout_group(qkv_layout), h, hg, d_qk, + d_v); cu_seqlens_padded_to_offsets<<>>( - layout_group, actual_b, b, h, hg, d_qk, d_v, static_cast(devPtrSeqOffsetsQ), + offset_mults, actual_b, b, static_cast(devPtrSeqOffsetsQ), static_cast(devPtrSeqOffsetsKV), ragged_offset_type, devOffsetsQ, devOffsetsK, devOffsetsV, devOffsetsO, devOffsetsS); NVTE_CHECK_CUDA(cudaGetLastError()); @@ -1142,7 +1231,10 @@ void fused_attn_arbitrary_seqlen_fwd( Tensor *output_S = convertNVTETensorCheck(Aux_CTX_Tensors->tensors[i++]); output_S->data.dptr = nullptr; - if (q_format == NVTE_QKV_Format::NVTE_THD && cudnn_runtime_version >= 90600) { + // sm120 does not use ragged stats: the graph declares a dense + // [b, h, s_q, 1] stats tensor, so allocate to match (same as Max below). + if ((q_format == NVTE_QKV_Format::NVTE_THD && cudnn_runtime_version >= 90600) && + (sm_arch_ != 120)) { output_S->data.shape = {num_tokens_q, num_attn_heads, 1}; } else { output_S->data.shape = {batch, num_attn_heads, max_seqlen_q, 1}; diff --git a/transformer_engine/common/fused_attn/fused_attn_fp8.cu b/transformer_engine/common/fused_attn/fused_attn_fp8.cu index eab1ae02e6..000af41aee 100644 --- a/transformer_engine/common/fused_attn/fused_attn_fp8.cu +++ b/transformer_engine/common/fused_attn/fused_attn_fp8.cu @@ -60,6 +60,24 @@ void fused_attn_fp8_fwd_impl( NVTE_CHECK(!is_mxfp8 || cudnn_runtime_version >= 92100, "MXFP8 fused attention requires cuDNN 9.21.0 or later!"); + // Newer versions of cuDNN SDPA can accept sequence lengths directly as a cumulative + // tensor. Take advantage of this if possible to avoid 1 extra kernel call. (Unlike + // the F16 path, the FP8 path has no THD/ragged-offset support, so only the + // cu_seqlens_to_actual_seqlens conversion applies here. Also note that the + // needed versions of cuDNN backend and frontend are higher than for F16.) + const bool use_cu_seqlens_directly = + // Frontend 1.26 supports fp8+cu_seqlens (for the C++ API). + // Note: For the Python API, 1.27 is required. + CUDNN_FRONTEND_VERSION >= 12600 && + // The frontend gates cu_seq_len support on min(compile-time, runtime) cuDNN + // version, so we'll do the same. + (CUDNN_VERSION >= 92500 && cudnn_runtime_version >= 92500) && + // This extra restriction is needed because cuDNN frontend doesn't yet allow + // the combination of dropout and stats generation for the fprop unified engine, + // so any such request would always get routed to the old composite SDPA engine + // (which doesn't support cu_seqlens). Remove this restriction when possible. + !is_dropout; + try { FADescriptor_v1 descriptor{b, h, @@ -262,17 +280,38 @@ void fused_attn_fp8_fwd_impl( // } if (is_padding) { - seq_q = mha_graph->tensor(fe::graph::Tensor_attributes() - .set_name("seq_q") - .set_dim({b, 1, 1, 1}) - .set_stride({1, 1, 1, 1}) - .set_data_type(fe::DataType_t::INT32)); - seq_kv = mha_graph->tensor(fe::graph::Tensor_attributes() - .set_name("seq_kv") - .set_dim({b, 1, 1, 1}) - .set_stride({1, 1, 1, 1}) - .set_data_type(fe::DataType_t::INT32)); - sdpa_options.set_padding_mask(is_padding).set_seq_len_q(seq_q).set_seq_len_kv(seq_kv); + if (use_cu_seqlens_directly) { + // seq_q/seq_kv keep their tuple slots but hold (b+1)-shaped cu_seqlen tensors. + seq_q = mha_graph->tensor(fe::graph::Tensor_attributes() + .set_name("cu_seq_len_q") + .set_dim({b + 1, 1, 1, 1}) + .set_stride({1, 1, 1, 1}) + .set_data_type(fe::DataType_t::INT32)); + seq_kv = mha_graph->tensor(fe::graph::Tensor_attributes() + .set_name("cu_seq_len_kv") + .set_dim({b + 1, 1, 1, 1}) + .set_stride({1, 1, 1, 1}) + .set_data_type(fe::DataType_t::INT32)); + sdpa_options.set_padding_mask(is_padding) + .set_cu_seq_len_q(seq_q) + .set_cu_seq_len_kv(seq_kv); + // cu_seq_len (and the ragged offset multiplier) are unified-engine-only. + // Pin the implementation so an unsupported config fails with the unified + // engine's specific error instead of auto-selection's generic failure. + sdpa_options.set_implementation(fe::AttentionImplementation_t::UNIFIED); + } else { + seq_q = mha_graph->tensor(fe::graph::Tensor_attributes() + .set_name("seq_q") + .set_dim({b, 1, 1, 1}) + .set_stride({1, 1, 1, 1}) + .set_data_type(fe::DataType_t::INT32)); + seq_kv = mha_graph->tensor(fe::graph::Tensor_attributes() + .set_name("seq_kv") + .set_dim({b, 1, 1, 1}) + .set_stride({1, 1, 1, 1}) + .set_data_type(fe::DataType_t::INT32)); + sdpa_options.set_padding_mask(is_padding).set_seq_len_q(seq_q).set_seq_len_kv(seq_kv); + } } if (is_dropout) { @@ -379,8 +418,10 @@ void fused_attn_fp8_fwd_impl( auto plan_workspace_size = mha_graph->get_workspace_size(); - // Exit to request upper level API to allocate memory if needed - size_t actual_seqlen_workspace_size = 2 * b * sizeof(int32_t); + // Exit to request upper level API to allocate memory if needed. + // When passing cu_seqlens* directly to cuDNN SDPA, no conversion workspace is + // needed: cuDNN consumes the user's cu_seqlens buffers as-is. + size_t actual_seqlen_workspace_size = use_cu_seqlens_directly ? 0 : 2 * b * sizeof(int32_t); if (workspace == nullptr) { *workspace_size = plan_workspace_size + actual_seqlen_workspace_size; return; @@ -417,17 +458,22 @@ void fused_attn_fp8_fwd_impl( } */ if (is_padding) { - constexpr size_t nthreads_per_block = 128; - const size_t grid = (b + nthreads_per_block - 1) / nthreads_per_block; - void* devActualSeqlenQ = static_cast(workspace) + plan_workspace_size; - void* devActualSeqlenKV = static_cast(devActualSeqlenQ) + b * sizeof(int32_t); - cu_seqlens_to_actual_seqlens<<>>( - b, b, static_cast(devPtrcuSeqlensQ), // TODO(pass max_b) - static_cast(devPtrcuSeqlensKV), static_cast(devActualSeqlenQ), - static_cast(devActualSeqlenKV)); - NVTE_CHECK_CUDA(cudaGetLastError()); - variant_pack[seq_q] = devActualSeqlenQ; - variant_pack[seq_kv] = devActualSeqlenKV; + if (use_cu_seqlens_directly) { + variant_pack[seq_q] = devPtrcuSeqlensQ; + variant_pack[seq_kv] = devPtrcuSeqlensKV; + } else { + constexpr size_t nthreads_per_block = 128; + const size_t grid = (b + nthreads_per_block - 1) / nthreads_per_block; + void* devActualSeqlenQ = static_cast(workspace) + plan_workspace_size; + void* devActualSeqlenKV = static_cast(devActualSeqlenQ) + b * sizeof(int32_t); + cu_seqlens_to_actual_seqlens<<>>( + b, b, static_cast(devPtrcuSeqlensQ), // TODO(pass max_b) + static_cast(devPtrcuSeqlensKV), static_cast(devActualSeqlenQ), + static_cast(devActualSeqlenKV)); + NVTE_CHECK_CUDA(cudaGetLastError()); + variant_pack[seq_q] = devActualSeqlenQ; + variant_pack[seq_kv] = devActualSeqlenKV; + } } if (is_dropout) { diff --git a/transformer_engine/common/fused_attn/utils.cu b/transformer_engine/common/fused_attn/utils.cu index 3e628b6581..9b54a64cbe 100644 --- a/transformer_engine/common/fused_attn/utils.cu +++ b/transformer_engine/common/fused_attn/utils.cu @@ -429,71 +429,43 @@ __global__ void cu_seqlens_to_actual_seqlens(int64_t actual_b, int64_t max_b, // convert cu_seqlens_padded to offsets template __device__ void cu_seqlens_padded_to_offsets_impl( - NVTE_QKV_Layout_Group layout_group, int64_t actual_b, int64_t max_b, int64_t h, int64_t hg, - int64_t d_qk, int64_t d_v, const int32_t *cu_seqlens_q_padded, - const int32_t *cu_seqlens_kv_padded, OFFSETS_T *offsets_q, OFFSETS_T *offsets_k, - OFFSETS_T *offsets_v, OFFSETS_T *offsets_o, OFFSETS_T *offsets_s) { + const RaggedOffsetMultipliers &mults, int64_t actual_b, int64_t max_b, + const int32_t *cu_seqlens_q_padded, const int32_t *cu_seqlens_kv_padded, OFFSETS_T *offsets_q, + OFFSETS_T *offsets_k, OFFSETS_T *offsets_v, OFFSETS_T *offsets_o, OFFSETS_T *offsets_s) { size_t tid = blockIdx.x * blockDim.x + threadIdx.x; auto cu_seqlens_id = min(tid, actual_b); if (tid <= max_b) { if (offsets_s != nullptr) { - offsets_s[tid] = h * cu_seqlens_q_padded[cu_seqlens_id]; + offsets_s[tid] = mults.stats * cu_seqlens_q_padded[cu_seqlens_id]; } if (offsets_q != nullptr && offsets_o != nullptr) { - offsets_o[tid] = h * d_v * cu_seqlens_q_padded[cu_seqlens_id]; - switch (layout_group) { - case NVTE_QKV_Layout_Group::NVTE_HD_HD_HD: - case NVTE_QKV_Layout_Group::NVTE_Paged_KV_HD_HD_HD: - offsets_q[tid] = h * d_qk * cu_seqlens_q_padded[cu_seqlens_id]; - break; - case NVTE_QKV_Layout_Group::NVTE_3HD: - case NVTE_QKV_Layout_Group::NVTE_H3D: - offsets_q[tid] = 3 * h * d_qk * cu_seqlens_q_padded[cu_seqlens_id]; - break; - case NVTE_QKV_Layout_Group::NVTE_HD_2HD: - case NVTE_QKV_Layout_Group::NVTE_HD_H2D: - offsets_q[tid] = h * d_qk * cu_seqlens_q_padded[cu_seqlens_id]; - break; - } + offsets_q[tid] = mults.q * cu_seqlens_q_padded[cu_seqlens_id]; + offsets_o[tid] = mults.o * cu_seqlens_q_padded[cu_seqlens_id]; } if (offsets_k != nullptr && offsets_v != nullptr) { - switch (layout_group) { - case NVTE_QKV_Layout_Group::NVTE_HD_HD_HD: - case NVTE_QKV_Layout_Group::NVTE_Paged_KV_HD_HD_HD: - offsets_k[tid] = hg * d_qk * cu_seqlens_kv_padded[cu_seqlens_id]; - offsets_v[tid] = hg * d_v * cu_seqlens_kv_padded[cu_seqlens_id]; - break; - case NVTE_QKV_Layout_Group::NVTE_3HD: - case NVTE_QKV_Layout_Group::NVTE_H3D: - offsets_k[tid] = 3 * h * d_qk * cu_seqlens_q_padded[cu_seqlens_id]; - offsets_v[tid] = offsets_k[cu_seqlens_id]; - break; - case NVTE_QKV_Layout_Group::NVTE_HD_2HD: - case NVTE_QKV_Layout_Group::NVTE_HD_H2D: - offsets_k[tid] = 2 * hg * d_qk * cu_seqlens_kv_padded[cu_seqlens_id]; - offsets_v[tid] = offsets_k[cu_seqlens_id]; - break; - } + const int32_t *cu_seqlens_kv_src = + mults.kv_from_q ? cu_seqlens_q_padded : cu_seqlens_kv_padded; + offsets_k[tid] = mults.k * cu_seqlens_kv_src[cu_seqlens_id]; + offsets_v[tid] = mults.v * cu_seqlens_kv_src[cu_seqlens_id]; } } } -__global__ void cu_seqlens_padded_to_offsets(NVTE_QKV_Layout_Group layout_group, int64_t actual_b, - int64_t max_b, int64_t h, int64_t hg, int64_t d_qk, - int64_t d_v, const int32_t *cu_seqlens_q_padded, +__global__ void cu_seqlens_padded_to_offsets(RaggedOffsetMultipliers mults, int64_t actual_b, + int64_t max_b, const int32_t *cu_seqlens_q_padded, const int32_t *cu_seqlens_kv_padded, DType offset_dtype, void *offsets_q, void *offsets_k, void *offsets_v, void *offsets_o, void *offsets_s) { if (offset_dtype == DType::kInt32) { cu_seqlens_padded_to_offsets_impl( - layout_group, actual_b, max_b, h, hg, d_qk, d_v, cu_seqlens_q_padded, cu_seqlens_kv_padded, + mults, actual_b, max_b, cu_seqlens_q_padded, cu_seqlens_kv_padded, reinterpret_cast(offsets_q), reinterpret_cast(offsets_k), reinterpret_cast(offsets_v), reinterpret_cast(offsets_o), reinterpret_cast(offsets_s)); } else { assert(offset_dtype == DType::kInt64 && "expect int64"); cu_seqlens_padded_to_offsets_impl( - layout_group, actual_b, max_b, h, hg, d_qk, d_v, cu_seqlens_q_padded, cu_seqlens_kv_padded, + mults, actual_b, max_b, cu_seqlens_q_padded, cu_seqlens_kv_padded, reinterpret_cast(offsets_q), reinterpret_cast(offsets_k), reinterpret_cast(offsets_v), reinterpret_cast(offsets_o), reinterpret_cast(offsets_s)); diff --git a/transformer_engine/common/fused_attn/utils.h b/transformer_engine/common/fused_attn/utils.h index b2bbd31f68..1864f9417d 100644 --- a/transformer_engine/common/fused_attn/utils.h +++ b/transformer_engine/common/fused_attn/utils.h @@ -309,14 +309,45 @@ struct FADescriptor_v1 { } }; +// Per-tensor scale factors relating cu_seqlens_padded (token units) to tensor-element +// ragged offsets, as a function of the QKV layout group. Single source of truth shared +// by the cu_seqlens_padded_to_offsets conversion kernel and the direct-seqlens path +// (which passes them to cuDNN as ragged offset multipliers). +struct RaggedOffsetMultipliers { + RaggedOffsetMultipliers(NVTE_QKV_Layout_Group layout_group, int64_t h, int64_t hg, int64_t d_qk, + int64_t d_v) + : q(h * d_qk), k(hg * d_qk), v(hg * d_v), o(h * d_v), stats(h), kv_from_q(false) { + switch (layout_group) { + case NVTE_QKV_Layout_Group::NVTE_3HD: + case NVTE_QKV_Layout_Group::NVTE_H3D: + q = k = v = 3 * h * d_qk; + kv_from_q = true; + break; + case NVTE_QKV_Layout_Group::NVTE_HD_2HD: + case NVTE_QKV_Layout_Group::NVTE_HD_H2D: + k = v = 2 * hg * d_qk; + break; + default: + break; + } + } + + int64_t q; + int64_t k; + int64_t v; + int64_t o; + int64_t stats; + // K/V offsets scale the Q-side cu_seqlens_padded (interleaved QKV layouts) + bool kv_from_q; +}; + __global__ void cu_seqlens_to_actual_seqlens(int64_t actual_b, int64_t max_b, int32_t const *const q_cu_seqlens, int32_t const *const kv_cu_seqlens, int32_t *q_seqlens, int32_t *kv_seqlens); -__global__ void cu_seqlens_padded_to_offsets(NVTE_QKV_Layout_Group layout_group, int64_t actual_b, - int64_t max_b, int64_t h, int64_t hg, int64_t d_qk, - int64_t d_v, const int32_t *cu_seqlens_q_padded, +__global__ void cu_seqlens_padded_to_offsets(RaggedOffsetMultipliers mults, int64_t actual_b, + int64_t max_b, const int32_t *cu_seqlens_q_padded, const int32_t *cu_seqlens_kv_padded, DType offset_dtype, void *offsets_q, void *offsets_k, void *offsets_v, void *offsets_o, void *offsets_s); diff --git a/transformer_engine/common/fused_softmax/scaled_aligned_causal_masked_softmax.cu b/transformer_engine/common/fused_softmax/scaled_aligned_causal_masked_softmax.cu index 6ea8017e07..1f8e748479 100644 --- a/transformer_engine/common/fused_softmax/scaled_aligned_causal_masked_softmax.cu +++ b/transformer_engine/common/fused_softmax/scaled_aligned_causal_masked_softmax.cu @@ -4,9 +4,11 @@ * See LICENSE for license information. ************************************************************************/ +#ifdef __CUDACC_RTC__ +#include "utils.cuh" +#else #include #include -#include #include #include #include @@ -19,10 +21,23 @@ #include "../common.h" #include "../util/logging.h" +#include "../util/rtc.h" +#include "../util/string.h" #include "../utils.cuh" +#include "string_code_fused_softmax_scaled_aligned_causal_masked_softmax_cu.h" +#endif namespace transformer_engine { +#ifdef __CUDACC_RTC__ +using fp16 = half; +using bf16 = nv_bfloat16; +#endif + +#ifndef NVTE_BUILD_LEGACY_STATIC_FUSED_SOFTMAX +#define NVTE_BUILD_LEGACY_STATIC_FUSED_SOFTMAX 0 +#endif + template __device__ __inline__ void copy_vector(Datatype *dst, const Datatype *src); @@ -92,7 +107,7 @@ struct Max { template __device__ __forceinline__ T WARP_SHFL_XOR_NATIVE(T value, int laneMask, int width = warpSize, unsigned int mask = 0xffffffff) { -#if CUDA_VERSION >= 9000 +#if defined(__CUDACC_RTC__) || CUDA_VERSION >= 9000 return __shfl_xor_sync(mask, value, laneMask, width); #else return __shfl_xor(value, laneMask, width); @@ -345,6 +360,42 @@ __global__ void scaled_aligned_causal_masked_softmax_warp_backward( } } +#ifdef __CUDACC_RTC__ + +#else + +namespace { + +constexpr const char *kRtcSourceFile = + "transformer_engine/common/fused_softmax/scaled_aligned_causal_masked_softmax.cu"; + +template +std::string make_softmax_rtc_code(int log2_elements) { + (void)log2_elements; + std::string code = string_code_fused_softmax_scaled_aligned_causal_masked_softmax_cu; + return code; +} + +template +std::string make_softmax_rtc_label(const char *direction, int log2_elements) { + return concat_strings("fused_softmax,variant=aligned_causal,direction=", direction, + ",type=", TypeInfo::name, ",log2=", log2_elements, ",fast_math=1"); +} + +template +std::string make_softmax_rtc_kernel_name(const char *kernel_name, int log2_elements) { + return concat_strings("&::transformer_engine::", kernel_name, "<", TypeInfo::name, ",", + TypeInfo::name, ",float,", log2_elements, ">"); +} + +void throw_nvrtc_required(const char *direction) { + NVTE_ERROR("Fused aligned causal softmax RTC path is disabled for ", direction, + ". Set NVTE_DISABLE_NVRTC=0 or rebuild with " + "NVTE_BUILD_LEGACY_STATIC_FUSED_SOFTMAX=ON."); +} + +} // namespace + template void call_kernel_scaled_aligned_causal_masked_softmax_forward( dim3 grid_size, dim3 block_size, const int shmem_size, cudaStream_t stream, output_t *dst, @@ -454,6 +505,24 @@ void dispatch_scaled_aligned_causal_masked_softmax_forward(output_t *dst, const dim3 block_size(warp_width, warps_per_block); dim3 grid_size(blocks); + if (rtc::is_enabled()) { + auto &rtc_manager = rtc::KernelManager::instance(); + const std::string kernel_label = make_softmax_rtc_label("forward", log2_elements); + if (!rtc_manager.is_compiled(kernel_label)) { + rtc_manager.compile(kernel_label, + make_softmax_rtc_kernel_name( + "scaled_aligned_causal_masked_softmax_warp_forward", log2_elements), + make_softmax_rtc_code(log2_elements), kRtcSourceFile, + {"--use_fast_math"}); + } + const acc_t rtc_scale = static_cast(scale); + rtc_manager.launch(kernel_label, grid_size, block_size, 0, stream, dst, src, rtc_scale, + microbatches, query_seq_len, key_seq_len); + NVTE_CHECK_CUDA(cudaGetLastError()); + return; + } + +#if NVTE_BUILD_LEGACY_STATIC_FUSED_SOFTMAX // create an array of pointers to functions using ForwardFuncType = typename FunctionWrapper::ForwardType; static std::array forwardFunctionArray; @@ -466,6 +535,9 @@ void dispatch_scaled_aligned_causal_masked_softmax_forward(output_t *dst, const // Call the corresponding kernel forwardFunctionArray[log2_elements](grid_size, block_size, 0, stream, dst, src, scale, microbatches, query_seq_len, key_seq_len); +#else + throw_nvrtc_required("forward"); +#endif } template @@ -499,6 +571,23 @@ void dispatch_scaled_aligned_causal_masked_softmax_backward( dim3 block_size(warp_width, warps_per_block); dim3 grid_size(blocks); + if (rtc::is_enabled()) { + auto &rtc_manager = rtc::KernelManager::instance(); + const std::string kernel_label = make_softmax_rtc_label("backward", log2_elements); + if (!rtc_manager.is_compiled(kernel_label)) { + rtc_manager.compile(kernel_label, + make_softmax_rtc_kernel_name( + "scaled_aligned_causal_masked_softmax_warp_backward", log2_elements), + make_softmax_rtc_code(log2_elements), kRtcSourceFile, + {"--use_fast_math"}); + } + rtc_manager.launch(kernel_label, grid_size, block_size, 0, stream, grad_input, grad, output, + scale, microbatches, query_seq_len, key_seq_len); + NVTE_CHECK_CUDA(cudaGetLastError()); + return; + } + +#if NVTE_BUILD_LEGACY_STATIC_FUSED_SOFTMAX // create an array of pointers to functions using BackwardFuncType = typename FunctionWrapper::BackwardType; static std::array backwardFunctionArray; @@ -511,6 +600,9 @@ void dispatch_scaled_aligned_causal_masked_softmax_backward( // Call the corresponding kernel backwardFunctionArray[log2_elements](grid_size, block_size, 0, stream, grad_input, grad, output, scale, microbatches, query_seq_len, key_seq_len); +#else + throw_nvrtc_required("backward"); +#endif } void scaled_aligned_causal_masked_softmax_forward(const Tensor &input, Tensor *softmax_results, @@ -546,8 +638,13 @@ void scaled_aligned_causal_masked_softmax_backward(Tensor output_grads, const Te reinterpret_cast(softmax_results.data.dptr), scale_factor, query_seq_len, key_seq_len, batches, attn_heads, stream);); } + +#endif // __CUDACC_RTC__ + } // end namespace transformer_engine +#ifndef __CUDACC_RTC__ + void nvte_scaled_aligned_causal_masked_softmax_forward(const NVTETensor input, NVTETensor softmax_results, float scale_factor, cudaStream_t stream) { @@ -568,3 +665,5 @@ void nvte_scaled_aligned_causal_masked_softmax_backward(const NVTETensor incomin *convertNVTETensorCheck(output_grads), *convertNVTETensorCheck(incoming_grads), *convertNVTETensorCheck(softmax_results), scale_factor, stream); } + +#endif // __CUDACC_RTC__ diff --git a/transformer_engine/common/fused_softmax/scaled_masked_softmax.cu b/transformer_engine/common/fused_softmax/scaled_masked_softmax.cu index 27f86673c5..6f295ad00a 100644 --- a/transformer_engine/common/fused_softmax/scaled_masked_softmax.cu +++ b/transformer_engine/common/fused_softmax/scaled_masked_softmax.cu @@ -4,9 +4,11 @@ * See LICENSE for license information. ************************************************************************/ +#ifdef __CUDACC_RTC__ +#include "utils.cuh" +#else #include #include -#include #include #include #include @@ -17,10 +19,27 @@ #include "../common.h" #include "../util/logging.h" +#include "../util/rtc.h" +#include "../util/string.h" #include "../utils.cuh" +#include "string_code_fused_softmax_scaled_masked_softmax_cu.h" +#endif namespace transformer_engine { +#ifdef __CUDACC_RTC__ +using bf16 = nv_bfloat16; +#endif + +#ifndef NVTE_BUILD_LEGACY_STATIC_FUSED_SOFTMAX +#define NVTE_BUILD_LEGACY_STATIC_FUSED_SOFTMAX 0 +#endif + +template +__device__ __forceinline__ T neg_infinity() { + return -static_cast(__int_as_float(0x7f800000)); +} + template __device__ __inline__ void copy_vector(Datatype *dst, const Datatype *src); @@ -67,7 +86,7 @@ struct Max { template __device__ __forceinline__ T WARP_SHFL_XOR_NATIVE(T value, int laneMask, int width = warpSize, unsigned int mask = 0xffffffff) { -#if CUDA_VERSION >= 9000 +#if defined(__CUDACC_RTC__) || CUDA_VERSION >= 9000 return __shfl_xor_sync(mask, value, laneMask, width); #else return __shfl_xor(value, laneMask, width); @@ -144,7 +163,7 @@ __global__ void scaled_softmax_warp_forward(output_t *dst, const input_t *src, c } else { #pragma unroll for (int element = 0; element < ELEMENTS_PER_LDG_STG; ++element) { - elements[i][it + element] = -std::numeric_limits::infinity(); + elements[i][it + element] = neg_infinity(); } } } @@ -167,7 +186,7 @@ __global__ void scaled_softmax_warp_forward(output_t *dst, const input_t *src, c for (int i = 0; i < WARP_BATCH; ++i) { #pragma unroll for (int it = 0; it < WARP_ITERATIONS; ++it) { - elements[i][it] = std::exp((elements[i][it] - max_value[i])); + elements[i][it] = expf((elements[i][it] - max_value[i])); sum[i] += elements[i][it]; } } @@ -269,7 +288,7 @@ __global__ void scaled_masked_softmax_warp_forward(output_t *dst, const input_t } else { #pragma unroll for (int element = 0; element < ELEMENTS_PER_LDG_STG; ++element) { - elements[i][it + element] = -std::numeric_limits::infinity(); + elements[i][it + element] = neg_infinity(); } } } @@ -299,7 +318,7 @@ __global__ void scaled_masked_softmax_warp_forward(output_t *dst, const input_t for (int i = 0; i < WARP_BATCH; ++i) { #pragma unroll for (int it = 0; it < WARP_ITERATIONS; ++it) { - elements[i][it] = std::exp((elements[i][it] - max_value[i])); + elements[i][it] = expf((elements[i][it] - max_value[i])); sum[i] += elements[i][it]; } } @@ -420,6 +439,42 @@ __global__ void scaled_masked_softmax_warp_backward(output_t *gradInput, const i } } +#ifdef __CUDACC_RTC__ + +#else + +namespace { + +constexpr const char *kRtcSourceFile = + "transformer_engine/common/fused_softmax/scaled_masked_softmax.cu"; + +template +std::string make_softmax_rtc_code(int log2_elements) { + (void)log2_elements; + std::string code = string_code_fused_softmax_scaled_masked_softmax_cu; + return code; +} + +template +std::string make_softmax_rtc_label(const char *variant, const char *direction, int log2_elements) { + return concat_strings("fused_softmax,variant=", variant, ",direction=", direction, + ",type=", TypeInfo::name, ",log2=", log2_elements, ",fast_math=1"); +} + +template +std::string make_softmax_rtc_kernel_name(const char *kernel_name, int log2_elements) { + return concat_strings("&::transformer_engine::", kernel_name, "<", TypeInfo::name, ",", + TypeInfo::name, ",float,", log2_elements, ">"); +} + +void throw_nvrtc_required(const char *variant) { + NVTE_ERROR("Fused softmax RTC path is disabled for ", variant, + ". Set NVTE_DISABLE_NVRTC=0 or rebuild with " + "NVTE_BUILD_LEGACY_STATIC_FUSED_SOFTMAX=ON."); +} + +} // namespace + template void dispatch_scaled_softmax_forward(output_t *dst, const input_t *src, const input_t scale, int query_seq_len, int key_seq_len, int batches, @@ -448,70 +503,89 @@ void dispatch_scaled_softmax_forward(output_t *dst, const input_t *src, const in NVTE_CHECK(query_seq_len % batches_per_block == 0, "Unsupported shape."); dim3 blocks(query_seq_len / batches_per_block, attn_heads, batches); dim3 threads(warp_size, warps_per_block, 1); - // Launch code would be more elegant if C++ supported FOR CONSTEXPR - switch (log2_elements) { - case 0: // 1 - scaled_softmax_warp_forward - <<>>(dst, src, scale, batch_count, key_seq_len); - break; - case 1: // 2 - scaled_softmax_warp_forward - <<>>(dst, src, scale, batch_count, key_seq_len); - break; - case 2: // 4 - scaled_softmax_warp_forward - <<>>(dst, src, scale, batch_count, key_seq_len); - break; - case 3: // 8 - scaled_softmax_warp_forward - <<>>(dst, src, scale, batch_count, key_seq_len); - break; - case 4: // 16 - scaled_softmax_warp_forward - <<>>(dst, src, scale, batch_count, key_seq_len); - break; - case 5: // 32 - scaled_softmax_warp_forward - <<>>(dst, src, scale, batch_count, key_seq_len); - break; - case 6: // 64 - scaled_softmax_warp_forward - <<>>(dst, src, scale, batch_count, key_seq_len); - break; - case 7: // 128 - scaled_softmax_warp_forward - <<>>(dst, src, scale, batch_count, key_seq_len); - break; - case 8: // 256 - scaled_softmax_warp_forward - <<>>(dst, src, scale, batch_count, key_seq_len); - break; - case 9: // 512 - scaled_softmax_warp_forward - <<>>(dst, src, scale, batch_count, key_seq_len); - break; - case 10: // 1024 - scaled_softmax_warp_forward - <<>>(dst, src, scale, batch_count, key_seq_len); - break; - case 11: // 2048 - scaled_softmax_warp_forward - <<>>(dst, src, scale, batch_count, key_seq_len); - break; - case 12: // 4096 - scaled_softmax_warp_forward - <<>>(dst, src, scale, batch_count, key_seq_len); - break; - case 13: // 8192 - scaled_softmax_warp_forward - <<>>(dst, src, scale, batch_count, key_seq_len); - break; - case 14: // 16384 - scaled_softmax_warp_forward - <<>>(dst, src, scale, batch_count, key_seq_len); - break; - default: - break; + if (rtc::is_enabled()) { + auto &rtc_manager = rtc::KernelManager::instance(); + const std::string kernel_label = + make_softmax_rtc_label("scaled", "forward", log2_elements); + if (!rtc_manager.is_compiled(kernel_label)) { + rtc_manager.compile( + kernel_label, + make_softmax_rtc_kernel_name("scaled_softmax_warp_forward", log2_elements), + make_softmax_rtc_code(log2_elements), kRtcSourceFile, {"--use_fast_math"}); + } + const acc_t rtc_scale = static_cast(scale); + rtc_manager.launch(kernel_label, blocks, threads, 0, stream, dst, src, rtc_scale, batch_count, + key_seq_len); + } else { +#if NVTE_BUILD_LEGACY_STATIC_FUSED_SOFTMAX + // Launch code would be more elegant if C++ supported FOR CONSTEXPR + switch (log2_elements) { + case 0: // 1 + scaled_softmax_warp_forward + <<>>(dst, src, scale, batch_count, key_seq_len); + break; + case 1: // 2 + scaled_softmax_warp_forward + <<>>(dst, src, scale, batch_count, key_seq_len); + break; + case 2: // 4 + scaled_softmax_warp_forward + <<>>(dst, src, scale, batch_count, key_seq_len); + break; + case 3: // 8 + scaled_softmax_warp_forward + <<>>(dst, src, scale, batch_count, key_seq_len); + break; + case 4: // 16 + scaled_softmax_warp_forward + <<>>(dst, src, scale, batch_count, key_seq_len); + break; + case 5: // 32 + scaled_softmax_warp_forward + <<>>(dst, src, scale, batch_count, key_seq_len); + break; + case 6: // 64 + scaled_softmax_warp_forward + <<>>(dst, src, scale, batch_count, key_seq_len); + break; + case 7: // 128 + scaled_softmax_warp_forward + <<>>(dst, src, scale, batch_count, key_seq_len); + break; + case 8: // 256 + scaled_softmax_warp_forward + <<>>(dst, src, scale, batch_count, key_seq_len); + break; + case 9: // 512 + scaled_softmax_warp_forward + <<>>(dst, src, scale, batch_count, key_seq_len); + break; + case 10: // 1024 + scaled_softmax_warp_forward + <<>>(dst, src, scale, batch_count, key_seq_len); + break; + case 11: // 2048 + scaled_softmax_warp_forward + <<>>(dst, src, scale, batch_count, key_seq_len); + break; + case 12: // 4096 + scaled_softmax_warp_forward + <<>>(dst, src, scale, batch_count, key_seq_len); + break; + case 13: // 8192 + scaled_softmax_warp_forward + <<>>(dst, src, scale, batch_count, key_seq_len); + break; + case 14: // 16384 + scaled_softmax_warp_forward + <<>>(dst, src, scale, batch_count, key_seq_len); + break; + default: + break; + } +#else + throw_nvrtc_required("scaled softmax forward"); +#endif } NVTE_CHECK_CUDA(cudaGetLastError()); } @@ -546,85 +620,105 @@ void dispatch_scaled_masked_softmax_forward(output_t *dst, const input_t *src, c NVTE_CHECK(query_seq_len % batches_per_block == 0, "Unsupported shape."); dim3 blocks(query_seq_len / batches_per_block, attn_heads, batches); dim3 threads(warp_size, warps_per_block, 1); - // Launch code would be more elegant if C++ supported FOR CONSTEXPR - switch (log2_elements) { - case 0: // 1 - scaled_masked_softmax_warp_forward - <<>>(dst, src, mask, scale, batch_count, key_seq_len, - pad_batches); - break; - case 1: // 2 - scaled_masked_softmax_warp_forward - <<>>(dst, src, mask, scale, batch_count, key_seq_len, - pad_batches); - break; - case 2: // 4 - scaled_masked_softmax_warp_forward - <<>>(dst, src, mask, scale, batch_count, key_seq_len, - pad_batches); - break; - case 3: // 8 - scaled_masked_softmax_warp_forward - <<>>(dst, src, mask, scale, batch_count, key_seq_len, - pad_batches); - break; - case 4: // 16 - scaled_masked_softmax_warp_forward - <<>>(dst, src, mask, scale, batch_count, key_seq_len, - pad_batches); - break; - case 5: // 32 - scaled_masked_softmax_warp_forward - <<>>(dst, src, mask, scale, batch_count, key_seq_len, - pad_batches); - break; - case 6: // 64 - scaled_masked_softmax_warp_forward - <<>>(dst, src, mask, scale, batch_count, key_seq_len, - pad_batches); - break; - case 7: // 128 - scaled_masked_softmax_warp_forward - <<>>(dst, src, mask, scale, batch_count, key_seq_len, - pad_batches); - break; - case 8: // 256 - scaled_masked_softmax_warp_forward - <<>>(dst, src, mask, scale, batch_count, key_seq_len, - pad_batches); - break; - case 9: // 512 - scaled_masked_softmax_warp_forward - <<>>(dst, src, mask, scale, batch_count, key_seq_len, - pad_batches); - break; - case 10: // 1024 - scaled_masked_softmax_warp_forward - <<>>(dst, src, mask, scale, batch_count, key_seq_len, - pad_batches); - break; - case 11: // 2048 - scaled_masked_softmax_warp_forward - <<>>(dst, src, mask, scale, batch_count, key_seq_len, - pad_batches); - break; - case 12: // 4096 - scaled_masked_softmax_warp_forward - <<>>(dst, src, mask, scale, batch_count, key_seq_len, - pad_batches); - break; - case 13: // 8192 - scaled_masked_softmax_warp_forward - <<>>(dst, src, mask, scale, batch_count, key_seq_len, - pad_batches); - break; - case 14: // 16384 - scaled_masked_softmax_warp_forward - <<>>(dst, src, mask, scale, batch_count, key_seq_len, - pad_batches); - break; - default: - break; + if (rtc::is_enabled()) { + auto &rtc_manager = rtc::KernelManager::instance(); + const std::string kernel_label = + make_softmax_rtc_label("masked", "forward", log2_elements); + if (!rtc_manager.is_compiled(kernel_label)) { + rtc_manager.compile(kernel_label, + make_softmax_rtc_kernel_name( + "scaled_masked_softmax_warp_forward", log2_elements), + make_softmax_rtc_code(log2_elements), kRtcSourceFile, + {"--use_fast_math"}); + } + const acc_t rtc_scale = static_cast(scale); + rtc_manager.launch(kernel_label, blocks, threads, 0, stream, dst, src, mask, rtc_scale, + batch_count, key_seq_len, pad_batches); + } else { +#if NVTE_BUILD_LEGACY_STATIC_FUSED_SOFTMAX + // Launch code would be more elegant if C++ supported FOR CONSTEXPR + switch (log2_elements) { + case 0: // 1 + scaled_masked_softmax_warp_forward + <<>>(dst, src, mask, scale, batch_count, key_seq_len, + pad_batches); + break; + case 1: // 2 + scaled_masked_softmax_warp_forward + <<>>(dst, src, mask, scale, batch_count, key_seq_len, + pad_batches); + break; + case 2: // 4 + scaled_masked_softmax_warp_forward + <<>>(dst, src, mask, scale, batch_count, key_seq_len, + pad_batches); + break; + case 3: // 8 + scaled_masked_softmax_warp_forward + <<>>(dst, src, mask, scale, batch_count, key_seq_len, + pad_batches); + break; + case 4: // 16 + scaled_masked_softmax_warp_forward + <<>>(dst, src, mask, scale, batch_count, key_seq_len, + pad_batches); + break; + case 5: // 32 + scaled_masked_softmax_warp_forward + <<>>(dst, src, mask, scale, batch_count, key_seq_len, + pad_batches); + break; + case 6: // 64 + scaled_masked_softmax_warp_forward + <<>>(dst, src, mask, scale, batch_count, key_seq_len, + pad_batches); + break; + case 7: // 128 + scaled_masked_softmax_warp_forward + <<>>(dst, src, mask, scale, batch_count, key_seq_len, + pad_batches); + break; + case 8: // 256 + scaled_masked_softmax_warp_forward + <<>>(dst, src, mask, scale, batch_count, key_seq_len, + pad_batches); + break; + case 9: // 512 + scaled_masked_softmax_warp_forward + <<>>(dst, src, mask, scale, batch_count, key_seq_len, + pad_batches); + break; + case 10: // 1024 + scaled_masked_softmax_warp_forward + <<>>(dst, src, mask, scale, batch_count, key_seq_len, + pad_batches); + break; + case 11: // 2048 + scaled_masked_softmax_warp_forward + <<>>(dst, src, mask, scale, batch_count, key_seq_len, + pad_batches); + break; + case 12: // 4096 + scaled_masked_softmax_warp_forward + <<>>(dst, src, mask, scale, batch_count, key_seq_len, + pad_batches); + break; + case 13: // 8192 + scaled_masked_softmax_warp_forward + <<>>(dst, src, mask, scale, batch_count, key_seq_len, + pad_batches); + break; + case 14: // 16384 + scaled_masked_softmax_warp_forward + <<>>(dst, src, mask, scale, batch_count, key_seq_len, + pad_batches); + break; + default: + break; + } +#else + throw_nvrtc_required("scaled masked softmax forward"); +#endif } NVTE_CHECK_CUDA(cudaGetLastError()); } @@ -658,85 +752,104 @@ void dispatch_scaled_masked_softmax_backward(output_t *grad_input, const input_t int batches_per_block = warps_per_block * batches_per_warp; int blocks = batch_count / batches_per_block; dim3 threads(warp_size, warps_per_block, 1); - // Launch code would be more elegant if C++ supported FOR CONSTEXPR - switch (log2_elements) { - case 0: // 1 - scaled_masked_softmax_warp_backward - <<>>(grad_input, grad, output, scale, batch_count, - key_seq_len); - break; - case 1: // 2 - scaled_masked_softmax_warp_backward - <<>>(grad_input, grad, output, scale, batch_count, - key_seq_len); - break; - case 2: // 4 - scaled_masked_softmax_warp_backward - <<>>(grad_input, grad, output, scale, batch_count, - key_seq_len); - break; - case 3: // 8 - scaled_masked_softmax_warp_backward - <<>>(grad_input, grad, output, scale, batch_count, - key_seq_len); - break; - case 4: // 16 - scaled_masked_softmax_warp_backward - <<>>(grad_input, grad, output, scale, batch_count, - key_seq_len); - break; - case 5: // 32 - scaled_masked_softmax_warp_backward - <<>>(grad_input, grad, output, scale, batch_count, - key_seq_len); - break; - case 6: // 64 - scaled_masked_softmax_warp_backward - <<>>(grad_input, grad, output, scale, batch_count, - key_seq_len); - break; - case 7: // 128 - scaled_masked_softmax_warp_backward - <<>>(grad_input, grad, output, scale, batch_count, - key_seq_len); - break; - case 8: // 256 - scaled_masked_softmax_warp_backward - <<>>(grad_input, grad, output, scale, batch_count, - key_seq_len); - break; - case 9: // 512 - scaled_masked_softmax_warp_backward - <<>>(grad_input, grad, output, scale, batch_count, - key_seq_len); - break; - case 10: // 1024 - scaled_masked_softmax_warp_backward - <<>>(grad_input, grad, output, scale, batch_count, - key_seq_len); - break; - case 11: // 2048 - scaled_masked_softmax_warp_backward - <<>>(grad_input, grad, output, scale, batch_count, - key_seq_len); - break; - case 12: // 4096 - scaled_masked_softmax_warp_backward - <<>>(grad_input, grad, output, scale, batch_count, - key_seq_len); - break; - case 13: // 8192 - scaled_masked_softmax_warp_backward - <<>>(grad_input, grad, output, scale, batch_count, - key_seq_len); - break; - case 14: // 16384 - scaled_masked_softmax_warp_backward - <<>>(grad_input, grad, output, scale, batch_count, - key_seq_len); - break; - default: - break; + if (rtc::is_enabled()) { + auto &rtc_manager = rtc::KernelManager::instance(); + const std::string kernel_label = + make_softmax_rtc_label("masked", "backward", log2_elements); + if (!rtc_manager.is_compiled(kernel_label)) { + rtc_manager.compile(kernel_label, + make_softmax_rtc_kernel_name( + "scaled_masked_softmax_warp_backward", log2_elements), + make_softmax_rtc_code(log2_elements), kRtcSourceFile, + {"--use_fast_math"}); + } + rtc_manager.launch(kernel_label, blocks, threads, 0, stream, grad_input, grad, output, scale, + batch_count, key_seq_len); + } else { +#if NVTE_BUILD_LEGACY_STATIC_FUSED_SOFTMAX + // Launch code would be more elegant if C++ supported FOR CONSTEXPR + switch (log2_elements) { + case 0: // 1 + scaled_masked_softmax_warp_backward + <<>>(grad_input, grad, output, scale, batch_count, + key_seq_len); + break; + case 1: // 2 + scaled_masked_softmax_warp_backward + <<>>(grad_input, grad, output, scale, batch_count, + key_seq_len); + break; + case 2: // 4 + scaled_masked_softmax_warp_backward + <<>>(grad_input, grad, output, scale, batch_count, + key_seq_len); + break; + case 3: // 8 + scaled_masked_softmax_warp_backward + <<>>(grad_input, grad, output, scale, batch_count, + key_seq_len); + break; + case 4: // 16 + scaled_masked_softmax_warp_backward + <<>>(grad_input, grad, output, scale, batch_count, + key_seq_len); + break; + case 5: // 32 + scaled_masked_softmax_warp_backward + <<>>(grad_input, grad, output, scale, batch_count, + key_seq_len); + break; + case 6: // 64 + scaled_masked_softmax_warp_backward + <<>>(grad_input, grad, output, scale, batch_count, + key_seq_len); + break; + case 7: // 128 + scaled_masked_softmax_warp_backward + <<>>(grad_input, grad, output, scale, batch_count, + key_seq_len); + break; + case 8: // 256 + scaled_masked_softmax_warp_backward + <<>>(grad_input, grad, output, scale, batch_count, + key_seq_len); + break; + case 9: // 512 + scaled_masked_softmax_warp_backward + <<>>(grad_input, grad, output, scale, batch_count, + key_seq_len); + break; + case 10: // 1024 + scaled_masked_softmax_warp_backward + <<>>(grad_input, grad, output, scale, batch_count, + key_seq_len); + break; + case 11: // 2048 + scaled_masked_softmax_warp_backward + <<>>(grad_input, grad, output, scale, batch_count, + key_seq_len); + break; + case 12: // 4096 + scaled_masked_softmax_warp_backward + <<>>(grad_input, grad, output, scale, batch_count, + key_seq_len); + break; + case 13: // 8192 + scaled_masked_softmax_warp_backward + <<>>(grad_input, grad, output, scale, batch_count, + key_seq_len); + break; + case 14: // 16384 + scaled_masked_softmax_warp_backward + <<>>(grad_input, grad, output, scale, batch_count, + key_seq_len); + break; + default: + break; + } +#else + throw_nvrtc_required("scaled softmax backward"); +#endif } NVTE_CHECK_CUDA(cudaGetLastError()); } @@ -812,8 +925,12 @@ void scaled_masked_softmax_backward(Tensor output_grads, const Tensor incoming_g query_seq_len, key_seq_len, batches, attn_heads, stream);); } +#endif // __CUDACC_RTC__ + } // end namespace transformer_engine +#ifndef __CUDACC_RTC__ + void nvte_scaled_softmax_forward(const NVTETensor input, NVTETensor softmax_results, float scale_factor, cudaStream_t stream) { NVTE_API_CALL(nvte_scaled_softmax_forward); @@ -850,3 +967,5 @@ void nvte_scaled_masked_softmax_backward(const NVTETensor incoming_grads, *convertNVTETensorCheck(incoming_grads), *convertNVTETensorCheck(softmax_results), scale_factor, stream); } + +#endif // __CUDACC_RTC__ diff --git a/transformer_engine/common/fused_softmax/scaled_upper_triang_masked_softmax.cu b/transformer_engine/common/fused_softmax/scaled_upper_triang_masked_softmax.cu index 431148cd1d..c18444f646 100644 --- a/transformer_engine/common/fused_softmax/scaled_upper_triang_masked_softmax.cu +++ b/transformer_engine/common/fused_softmax/scaled_upper_triang_masked_softmax.cu @@ -4,9 +4,11 @@ * See LICENSE for license information. ************************************************************************/ +#ifdef __CUDACC_RTC__ +#include "utils.cuh" +#else #include #include -#include #include #include #include @@ -17,10 +19,28 @@ #include "../common.h" #include "../util/logging.h" +#include "../util/rtc.h" +#include "../util/string.h" #include "../utils.cuh" +#include "string_code_fused_softmax_scaled_upper_triang_masked_softmax_cu.h" +#endif namespace transformer_engine { +#ifdef __CUDACC_RTC__ +using fp16 = half; +using bf16 = nv_bfloat16; +#endif + +#ifndef NVTE_BUILD_LEGACY_STATIC_FUSED_SOFTMAX +#define NVTE_BUILD_LEGACY_STATIC_FUSED_SOFTMAX 0 +#endif + +template +__device__ __forceinline__ T neg_infinity() { + return -static_cast(__int_as_float(0x7f800000)); +} + template __device__ __inline__ void copy_vector(Datatype *dst, const Datatype *src); @@ -90,7 +110,7 @@ struct Max { template __device__ __forceinline__ T WARP_SHFL_XOR_NATIVE(T value, int laneMask, int width = warpSize, unsigned int mask = 0xffffffff) { -#if CUDA_VERSION >= 9000 +#if defined(__CUDACC_RTC__) || CUDA_VERSION >= 9000 return __shfl_xor_sync(mask, value, laneMask, width); #else return __shfl_xor(value, laneMask, width); @@ -166,13 +186,13 @@ __global__ void scaled_upper_triang_masked_softmax_warp_forward(output_t *dst, c if ((element_index + element) < batch_element_count) { elements[i][it + element] = (acc_t)temp_data[element] * scale; } else { - elements[i][it + element] = -std::numeric_limits::infinity(); + elements[i][it + element] = neg_infinity(); } } } else { #pragma unroll for (int element = 0; element < ELEMENTS_PER_LDG_STG; ++element) { - elements[i][it + element] = -std::numeric_limits::infinity(); + elements[i][it + element] = neg_infinity(); } } } @@ -196,7 +216,7 @@ __global__ void scaled_upper_triang_masked_softmax_warp_forward(output_t *dst, c #pragma unroll for (int it = 0; it < WARP_ITERATIONS; ++it) { if (it < warp_iteration_limit) { - elements[i][it] = std::exp((elements[i][it] - max_value[i])); + elements[i][it] = expf((elements[i][it] - max_value[i])); sum[i] += elements[i][it]; } } @@ -333,6 +353,42 @@ __global__ void scaled_upper_triang_masked_softmax_warp_backward(output_t *gradI } } +#ifdef __CUDACC_RTC__ + +#else + +namespace { + +constexpr const char *kRtcSourceFile = + "transformer_engine/common/fused_softmax/scaled_upper_triang_masked_softmax.cu"; + +template +std::string make_softmax_rtc_code(int log2_elements) { + (void)log2_elements; + std::string code = string_code_fused_softmax_scaled_upper_triang_masked_softmax_cu; + return code; +} + +template +std::string make_softmax_rtc_label(const char *direction, int log2_elements) { + return concat_strings("fused_softmax,variant=upper_triang,direction=", direction, + ",type=", TypeInfo::name, ",log2=", log2_elements, ",fast_math=1"); +} + +template +std::string make_softmax_rtc_kernel_name(const char *kernel_name, int log2_elements) { + return concat_strings("&::transformer_engine::", kernel_name, "<", TypeInfo::name, ",", + TypeInfo::name, ",float,", log2_elements, ">"); +} + +void throw_nvrtc_required(const char *direction) { + NVTE_ERROR("Fused upper-triangular softmax RTC path is disabled for ", direction, + ". Set NVTE_DISABLE_NVRTC=0 or rebuild with " + "NVTE_BUILD_LEGACY_STATIC_FUSED_SOFTMAX=ON."); +} + +} // namespace + template void dispatch_scaled_upper_triang_masked_softmax_forward(output_t *dst, const input_t *src, const input_t scale, int softmax_elements, @@ -365,85 +421,104 @@ void dispatch_scaled_upper_triang_masked_softmax_forward(output_t *dst, const in int blocks_per_seq = attn_batches / batches_per_block; dim3 blocks(seq_len, blocks_per_seq, 1); dim3 threads(warp_size, warps_per_block, 1); - // Launch code would be more elegant if C++ supported FOR CONSTEXPR - switch (log2_elements) { - case 0: // 1 - scaled_upper_triang_masked_softmax_warp_forward - <<>>(dst, src, scale, batch_count, softmax_elements_stride, - softmax_elements); - break; - case 1: // 2 - scaled_upper_triang_masked_softmax_warp_forward - <<>>(dst, src, scale, batch_count, softmax_elements_stride, - softmax_elements); - break; - case 2: // 4 - scaled_upper_triang_masked_softmax_warp_forward - <<>>(dst, src, scale, batch_count, softmax_elements_stride, - softmax_elements); - break; - case 3: // 8 - scaled_upper_triang_masked_softmax_warp_forward - <<>>(dst, src, scale, batch_count, softmax_elements_stride, - softmax_elements); - break; - case 4: // 16 - scaled_upper_triang_masked_softmax_warp_forward - <<>>(dst, src, scale, batch_count, softmax_elements_stride, - softmax_elements); - break; - case 5: // 32 - scaled_upper_triang_masked_softmax_warp_forward - <<>>(dst, src, scale, batch_count, softmax_elements_stride, - softmax_elements); - break; - case 6: // 64 - scaled_upper_triang_masked_softmax_warp_forward - <<>>(dst, src, scale, batch_count, softmax_elements_stride, - softmax_elements); - break; - case 7: // 128 - scaled_upper_triang_masked_softmax_warp_forward - <<>>(dst, src, scale, batch_count, softmax_elements_stride, - softmax_elements); - break; - case 8: // 256 - scaled_upper_triang_masked_softmax_warp_forward - <<>>(dst, src, scale, batch_count, softmax_elements_stride, - softmax_elements); - break; - case 9: // 512 - scaled_upper_triang_masked_softmax_warp_forward - <<>>(dst, src, scale, batch_count, softmax_elements_stride, - softmax_elements); - break; - case 10: // 1024 - scaled_upper_triang_masked_softmax_warp_forward - <<>>(dst, src, scale, batch_count, softmax_elements_stride, - softmax_elements); - break; - case 11: // 2048 - scaled_upper_triang_masked_softmax_warp_forward - <<>>(dst, src, scale, batch_count, softmax_elements_stride, - softmax_elements); - break; - case 12: // 4096 - scaled_upper_triang_masked_softmax_warp_forward - <<>>(dst, src, scale, batch_count, softmax_elements_stride, - softmax_elements); - break; - case 13: // 8192 - scaled_upper_triang_masked_softmax_warp_forward - <<>>(dst, src, scale, batch_count, softmax_elements_stride, - softmax_elements); - break; - case 14: // 16384 - scaled_upper_triang_masked_softmax_warp_forward - <<>>(dst, src, scale, batch_count, softmax_elements_stride, - softmax_elements); - break; - default: - break; + if (rtc::is_enabled()) { + auto &rtc_manager = rtc::KernelManager::instance(); + const std::string kernel_label = make_softmax_rtc_label("forward", log2_elements); + if (!rtc_manager.is_compiled(kernel_label)) { + rtc_manager.compile(kernel_label, + make_softmax_rtc_kernel_name( + "scaled_upper_triang_masked_softmax_warp_forward", log2_elements), + make_softmax_rtc_code(log2_elements), kRtcSourceFile, + {"--use_fast_math"}); + } + const acc_t rtc_scale = static_cast(scale); + rtc_manager.launch(kernel_label, blocks, threads, 0, stream, dst, src, rtc_scale, batch_count, + softmax_elements_stride, softmax_elements); + } else { +#if NVTE_BUILD_LEGACY_STATIC_FUSED_SOFTMAX + // Launch code would be more elegant if C++ supported FOR CONSTEXPR + switch (log2_elements) { + case 0: // 1 + scaled_upper_triang_masked_softmax_warp_forward + <<>>(dst, src, scale, batch_count, + softmax_elements_stride, softmax_elements); + break; + case 1: // 2 + scaled_upper_triang_masked_softmax_warp_forward + <<>>(dst, src, scale, batch_count, + softmax_elements_stride, softmax_elements); + break; + case 2: // 4 + scaled_upper_triang_masked_softmax_warp_forward + <<>>(dst, src, scale, batch_count, + softmax_elements_stride, softmax_elements); + break; + case 3: // 8 + scaled_upper_triang_masked_softmax_warp_forward + <<>>(dst, src, scale, batch_count, + softmax_elements_stride, softmax_elements); + break; + case 4: // 16 + scaled_upper_triang_masked_softmax_warp_forward + <<>>(dst, src, scale, batch_count, + softmax_elements_stride, softmax_elements); + break; + case 5: // 32 + scaled_upper_triang_masked_softmax_warp_forward + <<>>(dst, src, scale, batch_count, + softmax_elements_stride, softmax_elements); + break; + case 6: // 64 + scaled_upper_triang_masked_softmax_warp_forward + <<>>(dst, src, scale, batch_count, + softmax_elements_stride, softmax_elements); + break; + case 7: // 128 + scaled_upper_triang_masked_softmax_warp_forward + <<>>(dst, src, scale, batch_count, + softmax_elements_stride, softmax_elements); + break; + case 8: // 256 + scaled_upper_triang_masked_softmax_warp_forward + <<>>(dst, src, scale, batch_count, + softmax_elements_stride, softmax_elements); + break; + case 9: // 512 + scaled_upper_triang_masked_softmax_warp_forward + <<>>(dst, src, scale, batch_count, + softmax_elements_stride, softmax_elements); + break; + case 10: // 1024 + scaled_upper_triang_masked_softmax_warp_forward + <<>>(dst, src, scale, batch_count, + softmax_elements_stride, softmax_elements); + break; + case 11: // 2048 + scaled_upper_triang_masked_softmax_warp_forward + <<>>(dst, src, scale, batch_count, + softmax_elements_stride, softmax_elements); + break; + case 12: // 4096 + scaled_upper_triang_masked_softmax_warp_forward + <<>>(dst, src, scale, batch_count, + softmax_elements_stride, softmax_elements); + break; + case 13: // 8192 + scaled_upper_triang_masked_softmax_warp_forward + <<>>(dst, src, scale, batch_count, + softmax_elements_stride, softmax_elements); + break; + case 14: // 16384 + scaled_upper_triang_masked_softmax_warp_forward + <<>>(dst, src, scale, batch_count, + softmax_elements_stride, softmax_elements); + break; + default: + break; + } +#else + throw_nvrtc_required("forward"); +#endif } NVTE_CHECK_CUDA(cudaGetLastError()); } @@ -482,85 +557,103 @@ void dispatch_scaled_upper_triang_masked_softmax_backward(output_t *grad_input, int blocks_per_seq = attn_batches / batches_per_block; dim3 blocks(seq_len, blocks_per_seq, 1); dim3 threads(warp_size, warps_per_block, 1); - // Launch code would be more elegant if C++ supported FOR CONSTEXPR - switch (log2_elements) { - case 0: // 1 - scaled_upper_triang_masked_softmax_warp_backward - <<>>(grad_input, grad, output, scale, batch_count, - softmax_elements_stride, softmax_elements); - break; - case 1: // 2 - scaled_upper_triang_masked_softmax_warp_backward - <<>>(grad_input, grad, output, scale, batch_count, - softmax_elements_stride, softmax_elements); - break; - case 2: // 4 - scaled_upper_triang_masked_softmax_warp_backward - <<>>(grad_input, grad, output, scale, batch_count, - softmax_elements_stride, softmax_elements); - break; - case 3: // 8 - scaled_upper_triang_masked_softmax_warp_backward - <<>>(grad_input, grad, output, scale, batch_count, - softmax_elements_stride, softmax_elements); - break; - case 4: // 16 - scaled_upper_triang_masked_softmax_warp_backward - <<>>(grad_input, grad, output, scale, batch_count, - softmax_elements_stride, softmax_elements); - break; - case 5: // 32 - scaled_upper_triang_masked_softmax_warp_backward - <<>>(grad_input, grad, output, scale, batch_count, - softmax_elements_stride, softmax_elements); - break; - case 6: // 64 - scaled_upper_triang_masked_softmax_warp_backward - <<>>(grad_input, grad, output, scale, batch_count, - softmax_elements_stride, softmax_elements); - break; - case 7: // 128 - scaled_upper_triang_masked_softmax_warp_backward - <<>>(grad_input, grad, output, scale, batch_count, - softmax_elements_stride, softmax_elements); - break; - case 8: // 256 - scaled_upper_triang_masked_softmax_warp_backward - <<>>(grad_input, grad, output, scale, batch_count, - softmax_elements_stride, softmax_elements); - break; - case 9: // 512 - scaled_upper_triang_masked_softmax_warp_backward - <<>>(grad_input, grad, output, scale, batch_count, - softmax_elements_stride, softmax_elements); - break; - case 10: // 1024 - scaled_upper_triang_masked_softmax_warp_backward - <<>>(grad_input, grad, output, scale, batch_count, - softmax_elements_stride, softmax_elements); - break; - case 11: // 2048 - scaled_upper_triang_masked_softmax_warp_backward - <<>>(grad_input, grad, output, scale, batch_count, - softmax_elements_stride, softmax_elements); - break; - case 12: // 4096 - scaled_upper_triang_masked_softmax_warp_backward - <<>>(grad_input, grad, output, scale, batch_count, - softmax_elements_stride, softmax_elements); - break; - case 13: // 8192 - scaled_upper_triang_masked_softmax_warp_backward - <<>>(grad_input, grad, output, scale, batch_count, - softmax_elements_stride, softmax_elements); - break; - case 14: // 16384 - scaled_upper_triang_masked_softmax_warp_backward - <<>>(grad_input, grad, output, scale, batch_count, - softmax_elements_stride, softmax_elements); - break; - default: - break; + if (rtc::is_enabled()) { + auto &rtc_manager = rtc::KernelManager::instance(); + const std::string kernel_label = make_softmax_rtc_label("backward", log2_elements); + if (!rtc_manager.is_compiled(kernel_label)) { + rtc_manager.compile(kernel_label, + make_softmax_rtc_kernel_name( + "scaled_upper_triang_masked_softmax_warp_backward", log2_elements), + make_softmax_rtc_code(log2_elements), kRtcSourceFile, + {"--use_fast_math"}); + } + rtc_manager.launch(kernel_label, blocks, threads, 0, stream, grad_input, grad, output, scale, + batch_count, softmax_elements_stride, softmax_elements); + } else { +#if NVTE_BUILD_LEGACY_STATIC_FUSED_SOFTMAX + // Launch code would be more elegant if C++ supported FOR CONSTEXPR + switch (log2_elements) { + case 0: // 1 + scaled_upper_triang_masked_softmax_warp_backward + <<>>(grad_input, grad, output, scale, batch_count, + softmax_elements_stride, softmax_elements); + break; + case 1: // 2 + scaled_upper_triang_masked_softmax_warp_backward + <<>>(grad_input, grad, output, scale, batch_count, + softmax_elements_stride, softmax_elements); + break; + case 2: // 4 + scaled_upper_triang_masked_softmax_warp_backward + <<>>(grad_input, grad, output, scale, batch_count, + softmax_elements_stride, softmax_elements); + break; + case 3: // 8 + scaled_upper_triang_masked_softmax_warp_backward + <<>>(grad_input, grad, output, scale, batch_count, + softmax_elements_stride, softmax_elements); + break; + case 4: // 16 + scaled_upper_triang_masked_softmax_warp_backward + <<>>(grad_input, grad, output, scale, batch_count, + softmax_elements_stride, softmax_elements); + break; + case 5: // 32 + scaled_upper_triang_masked_softmax_warp_backward + <<>>(grad_input, grad, output, scale, batch_count, + softmax_elements_stride, softmax_elements); + break; + case 6: // 64 + scaled_upper_triang_masked_softmax_warp_backward + <<>>(grad_input, grad, output, scale, batch_count, + softmax_elements_stride, softmax_elements); + break; + case 7: // 128 + scaled_upper_triang_masked_softmax_warp_backward + <<>>(grad_input, grad, output, scale, batch_count, + softmax_elements_stride, softmax_elements); + break; + case 8: // 256 + scaled_upper_triang_masked_softmax_warp_backward + <<>>(grad_input, grad, output, scale, batch_count, + softmax_elements_stride, softmax_elements); + break; + case 9: // 512 + scaled_upper_triang_masked_softmax_warp_backward + <<>>(grad_input, grad, output, scale, batch_count, + softmax_elements_stride, softmax_elements); + break; + case 10: // 1024 + scaled_upper_triang_masked_softmax_warp_backward + <<>>(grad_input, grad, output, scale, batch_count, + softmax_elements_stride, softmax_elements); + break; + case 11: // 2048 + scaled_upper_triang_masked_softmax_warp_backward + <<>>(grad_input, grad, output, scale, batch_count, + softmax_elements_stride, softmax_elements); + break; + case 12: // 4096 + scaled_upper_triang_masked_softmax_warp_backward + <<>>(grad_input, grad, output, scale, batch_count, + softmax_elements_stride, softmax_elements); + break; + case 13: // 8192 + scaled_upper_triang_masked_softmax_warp_backward + <<>>(grad_input, grad, output, scale, batch_count, + softmax_elements_stride, softmax_elements); + break; + case 14: // 16384 + scaled_upper_triang_masked_softmax_warp_backward + <<>>(grad_input, grad, output, scale, batch_count, + softmax_elements_stride, softmax_elements); + break; + default: + break; + } +#else + throw_nvrtc_required("backward"); +#endif } NVTE_CHECK_CUDA(cudaGetLastError()); } @@ -595,8 +688,12 @@ void scaled_upper_triang_masked_softmax_backward(Tensor output_grads, const Tens seq_len, attn_batches, stream);); } +#endif // __CUDACC_RTC__ + } // end namespace transformer_engine +#ifndef __CUDACC_RTC__ + void nvte_scaled_upper_triang_masked_softmax_forward(const NVTETensor input, NVTETensor softmax_results, float scale_factor, cudaStream_t stream) { @@ -615,3 +712,5 @@ void nvte_scaled_upper_triang_masked_softmax_backward(const NVTETensor incoming_ *convertNVTETensorCheck(output_grads), *convertNVTETensorCheck(incoming_grads), *convertNVTETensorCheck(softmax_results), scale_factor, stream); } + +#endif // __CUDACC_RTC__ diff --git a/transformer_engine/common/hadamard_transform/customized_pipeline.cuh b/transformer_engine/common/hadamard_transform/customized_pipeline.cuh index bc46341e88..11754fbb50 100644 --- a/transformer_engine/common/hadamard_transform/customized_pipeline.cuh +++ b/transformer_engine/common/hadamard_transform/customized_pipeline.cuh @@ -18,7 +18,7 @@ namespace detail { // by producer_commit. Use case, accumulator generation as // the result of MMA instructions. template , - class AtomThrShape_MNK_ = Shape<_1, _1, _1> > + class AtomThrShape_MNK_ = Shape<_1, _1, _1>> class CustomizedPipelineTmaUmmaAsync { public: static constexpr uint32_t Stages = Stages_; diff --git a/transformer_engine/common/include/transformer_engine/activation.h b/transformer_engine/common/include/transformer_engine/activation.h index 4ed083740d..3df84954bf 100644 --- a/transformer_engine/common/include/transformer_engine/activation.h +++ b/transformer_engine/common/include/transformer_engine/activation.h @@ -368,6 +368,43 @@ void nvte_clamped_swiglu(const NVTETensor input, NVTETensor output, float limit, void nvte_clamped_swiglu_v2(const NVTETensor input, NVTETensor output, float limit, float alpha, float glu_linear_offset, cudaStream_t stream); +/*! \brief Computes ScaledSwiGLU without materializing GLU deinterleave. + * + * Computes output = SwiGLU(input) * act_scales[:, None]. + * If glu_interleave_size > 0, input is interpreted as interleaved + * [activation_block, linear_block] chunks of that size. + * + * \param[in] input Input tensor of shape [N, H * 2]. + * \param[in] act_scales Row-wise activation scales of shape [N]. + * \param[in,out] output Output tensor of shape [N, H]. + * \param[in] glu_interleave_size 0 for non-interleaved layout; otherwise a positive + * multiple of 32 that divides H. + * \param[in] stream CUDA stream used for the operation. + */ +void nvte_scaled_swiglu(const NVTETensor input, const NVTETensor act_scales, NVTETensor output, + int64_t glu_interleave_size, cudaStream_t stream); + +/*! \brief Computes ScaledClampedSwiGLU without materializing GLU deinterleave. + * + * Computes output = ClampedSwiGLU(input) * act_scales[:, None]. + * This uses the same clamping, alpha, and linear-offset semantics as + * nvte_clamped_swiglu_v2. + * + * \param[in] input Input tensor of shape [N, H * 2]. + * \param[in] act_scales Row-wise activation scales of shape [N]. + * \param[in,out] output Output tensor of shape [N, H]. + * \param[in] limit Clipping limit. + * \param[in] alpha Activation sigmoid alpha. + * \param[in] glu_linear_offset Offset added to linear component after clamping. + * \param[in] glu_interleave_size 0 for non-interleaved layout; otherwise a positive + * multiple of 32 that divides H. + * \param[in] stream CUDA stream used for the operation. + */ +void nvte_scaled_clamped_swiglu(const NVTETensor input, const NVTETensor act_scales, + NVTETensor output, float limit, float alpha, + float glu_linear_offset, int64_t glu_interleave_size, + cudaStream_t stream); + /*! \brief Computes the gated ReLU activation of the input. * If the scaling mode of the output tensor is set to NVTE_MXFP8_1D_SCALING, * the block quantization (MXFP8) of the specified shape of the block will be used. @@ -473,6 +510,47 @@ void nvte_clamped_dswiglu_v2(const NVTETensor grad, const NVTETensor input, NVTE float limit, float alpha, float glu_linear_offset, cudaStream_t stream); +/*! \brief Computes ScaledSwiGLU backward without materializing GLU deinterleave. + * + * The optional grad_act_scales tensor may be null. When present, it receives + * sum(dY * SwiGLU(input), dim=-1). + * + * \param[in] grad Incoming gradient of shape [N, H]. + * \param[in] input Forward input tensor of shape [N, H * 2]. + * \param[in] act_scales Row-wise activation scales of shape [N]. + * \param[in,out] grad_input Outgoing gradient of shape [N, H * 2]. + * \param[in,out] grad_act_scales Optional row-wise scale gradient of shape [N], or null. + * \param[in] glu_interleave_size 0 for non-interleaved layout; otherwise a positive + * multiple of 32 that divides H. + * \param[in] stream CUDA stream used for the operation. + */ +void nvte_scaled_dswiglu(const NVTETensor grad, const NVTETensor input, const NVTETensor act_scales, + NVTETensor grad_input, NVTETensor grad_act_scales, + int64_t glu_interleave_size, cudaStream_t stream); + +/*! \brief Computes ScaledClampedSwiGLU backward without materializing GLU deinterleave. + * + * The optional grad_act_scales tensor may be null. When present, it receives + * sum(dY * ClampedSwiGLU(input), dim=-1). + * + * \param[in] grad Incoming gradient of shape [N, H]. + * \param[in] input Forward input tensor of shape [N, H * 2]. + * \param[in] act_scales Row-wise activation scales of shape [N]. + * \param[in,out] grad_input Outgoing gradient of shape [N, H * 2]. + * \param[in,out] grad_act_scales Optional row-wise scale gradient of shape [N], or null. + * \param[in] limit Clipping limit. + * \param[in] alpha Activation sigmoid alpha. + * \param[in] glu_linear_offset Offset added to linear component after clamping. + * \param[in] glu_interleave_size 0 for non-interleaved layout; otherwise a positive + * multiple of 32 that divides H. + * \param[in] stream CUDA stream used for the operation. + */ +void nvte_scaled_clamped_dswiglu(const NVTETensor grad, const NVTETensor input, + const NVTETensor act_scales, NVTETensor grad_input, + NVTETensor grad_act_scales, float limit, float alpha, + float glu_linear_offset, int64_t glu_interleave_size, + cudaStream_t stream); + /*! \brief Computes the gated ReLU activation gradient. * If the scaling mode of the output tensor is set to NVTE_MXFP8_1D_SCALING, * the block quantization (MXFP8) of the specified shape of the block will be used. @@ -509,6 +587,33 @@ void nvte_dqgeglu(const NVTETensor grad, const NVTETensor input, NVTETensor outp void nvte_dsreglu(const NVTETensor grad, const NVTETensor input, NVTETensor output, cudaStream_t stream); +/*! \brief Computes ScaledSReLU. + * + * Computes output = SReLU(input) * act_scales[:, None]. + * + * \param[in] input Input tensor for activation. + * \param[in] act_scales Row-wise activation scales of shape [N]. + * \param[in,out] output Output tensor. + * \param[in] stream CUDA stream used for the operation. + */ +void nvte_scaled_srelu(const NVTETensor input, const NVTETensor act_scales, NVTETensor output, + cudaStream_t stream); + +/*! \brief Computes ScaledSReLU backward. + * + * The optional grad_act_scales tensor may be null. When present, it receives + * sum(dY * SReLU(input), dim=-1). + * + * \param[in] grad Incoming gradient. + * \param[in] input Forward input tensor. + * \param[in] act_scales Row-wise activation scales of shape [N]. + * \param[in,out] grad_input Outgoing input gradient. + * \param[in,out] grad_act_scales Optional row-wise scale gradient of shape [N], or null. + * \param[in] stream CUDA stream used for the operation. + */ +void nvte_scaled_dsrelu(const NVTETensor grad, const NVTETensor input, const NVTETensor act_scales, + NVTETensor grad_input, NVTETensor grad_act_scales, cudaStream_t stream); + #ifdef __cplusplus } // extern "C" #endif diff --git a/transformer_engine/common/include/transformer_engine/cast.h b/transformer_engine/common/include/transformer_engine/cast.h index 554d8c1ac9..4d6d24ba65 100644 --- a/transformer_engine/common/include/transformer_engine/cast.h +++ b/transformer_engine/common/include/transformer_engine/cast.h @@ -161,10 +161,12 @@ void nvte_quantize_dbias(const NVTETensor input, NVTETensor output, NVTETensor d * \param[in,out] output Output grouped FP8/MXFP8 tensor. * \param[out] dbias Result of the reduction of the input along columns. * \param[out] workspace Workspace tensor. + * \param[in] quant_config Quantization configuration. * \param[in] stream CUDA stream used for the operation. */ void nvte_group_quantize_dbias(const NVTEGroupedTensor input, NVTEGroupedTensor output, - NVTEGroupedTensor dbias, NVTETensor workspace, cudaStream_t stream); + NVTEGroupedTensor dbias, NVTETensor workspace, + const NVTEQuantizationConfig quant_config, cudaStream_t stream); /*! \brief Computes backward of GeLU operation on the input, then casts to FP8/MXFP8. * Additionally, reduces the result of the GeLU backward along columns. diff --git a/transformer_engine/common/multi_tensor/multi_tensor_apply.cuh b/transformer_engine/common/multi_tensor/multi_tensor_apply.cuh index 6710a161b3..d6ac23d4af 100644 --- a/transformer_engine/common/multi_tensor/multi_tensor_apply.cuh +++ b/transformer_engine/common/multi_tensor/multi_tensor_apply.cuh @@ -75,6 +75,9 @@ void multi_tensor_apply(int64_t block_size, int64_t chunk_size, loc_tensor_info++; auto chunks_this_tensor = (tensor_lists[0][t]->numel() + chunk_size - 1) / chunk_size; + NVTE_CHECK(chunks_this_tensor > 0, + "multi_tensor_apply expects tensors with at least one chunk; zero-sized tensors " + "must be filtered before launch because they skip the chunk loop"); for (auto chunk = 0; chunk < chunks_this_tensor; chunk++) { tl.block_to_tensor[loc_block_info] = loc_tensor_info - 1; diff --git a/transformer_engine/common/normalization/common.h b/transformer_engine/common/normalization/common.h index f5dce64193..31a547f62c 100644 --- a/transformer_engine/common/normalization/common.h +++ b/transformer_engine/common/normalization/common.h @@ -24,6 +24,7 @@ #include "../common.h" #include "../cudnn_utils.h" #include "../util/system.h" +#include "kernel_params.h" namespace transformer_engine { @@ -51,97 +52,6 @@ struct LaunchParams { } }; -struct KernelParamsBase { - KernelParamsBase() - : ctas_per_col(0), - rows(0), - cols(0), - x(nullptr), - mu(nullptr), - rs(nullptr), - gamma(nullptr), - workspace(nullptr), - barrier(nullptr), - zero_centered_gamma(false) {} - - // For Multi-CTA, number of different CTA groups. Otherwise same as gridDim.x. - int ctas_per_col; - // Size of CTA group. - int ctas_per_row; - - // Input is interpreted as matrix. We normalize across columns. - int rows; - int cols; - - // Common data pointers. - void* x; - void* mu; - void* rs; - void* gamma; - - // Multi-CTA workspace in gmem. - void* workspace; - - // Multi-CTA sync barriers in gmem. - int* barrier; - - // Whether gamma is centered around 0 - bool zero_centered_gamma; -}; - -struct ForwardKernelParams : public KernelParamsBase { - ForwardKernelParams() - : KernelParamsBase(), z(nullptr), beta(nullptr), epsilon(0.f), fp8_out(false) {} - - // Output of LN FWD. - void* z; - void* beta; - float epsilon; - - // Scaling factor - void* scale; - int scale_byte_size; - - // Inverse of scaling factor - void* scale_inv; - - // AMax output - void* amax; - int amax_byte_size; - - // Whether to compute scale and amax - bool fp8_out; -}; - -struct BackwardKernelParams : public KernelParamsBase { - BackwardKernelParams() - : KernelParamsBase(), - dz(nullptr), - dbeta_part(nullptr), - dgamma_part(nullptr), - dx(nullptr), - dbeta(nullptr), - dgamma(nullptr) {} - - // Input: gradient wrt. LN FWD output. - void* dz; - - // Input: extra tensor to add for fused backward+add - void* add; - - // Workspace for Wgrad pre-reduction. - void* dbeta_part; - void* dgamma_part; - - // Output: Dgrad. - void* dx; - // Output: Wgrad. - void* dbeta; - void* dgamma; -}; - -using BackwardAddKernelParams = BackwardKernelParams; - enum class NVTE_Norm_Backend { Te, Cudnn }; enum class NVTE_Norm_Stage { Forward, Backward, BackwardAdd }; @@ -191,6 +101,16 @@ class TeNormalizationRegistry { return 0; } + // Overload for capturing-callable dispatchers (e.g. NVRTC closures). + static int registerFunction(TupleKeyType key, Function func) { + auto [general_key, batch_size, hidden_size, is_tuned] = key; + if (is_tuned) + getInstance().tuned_function_map.emplace(key, std::move(func)); + else + getInstance().general_function_map[general_key].emplace(hidden_size, std::move(func)); + return 0; + } + static Function getKernel(TupleKeyType key) { auto& instance = getInstance(); auto [general_key, batch_size, hidden_size, is_tuned] = key; diff --git a/transformer_engine/common/normalization/kernel_params.h b/transformer_engine/common/normalization/kernel_params.h new file mode 100644 index 0000000000..42dc423dd0 --- /dev/null +++ b/transformer_engine/common/normalization/kernel_params.h @@ -0,0 +1,88 @@ +/************************************************************************* + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * + * See LICENSE for license information. + ************************************************************************/ + +#ifndef TRANSFORMER_ENGINE_COMMON_NORM_KERNEL_PARAMS_H_ +#define TRANSFORMER_ENGINE_COMMON_NORM_KERNEL_PARAMS_H_ + +// POD kernel-parameter structs shared between the host norm dispatchers and +// the NVRTC kernel sources. This header is intentionally free of host-only +// includes (no , no cudnn) so it can be compiled by +// NVRTC as well. + +namespace transformer_engine { +namespace normalization { + +struct KernelParamsBase { + // For Multi-CTA, number of different CTA groups. Otherwise same as gridDim.x. + int ctas_per_col = 0; + // Size of CTA group. + int ctas_per_row = 0; + + // Input is interpreted as matrix. We normalize across columns. + int rows = 0; + int cols = 0; + + // Common data pointers. + void* x = nullptr; + void* mu = nullptr; + void* rs = nullptr; + void* gamma = nullptr; + + // Multi-CTA workspace in gmem. + void* workspace = nullptr; + + // Multi-CTA sync barriers in gmem. + int* barrier = nullptr; + + // Whether gamma is centered around 0 + bool zero_centered_gamma = false; +}; + +struct ForwardKernelParams : public KernelParamsBase { + // Output of LN FWD. + void* z = nullptr; + void* beta = nullptr; + float epsilon = 0.f; + + // Scaling factor + void* scale = nullptr; + int scale_byte_size = 0; + + // Inverse of scaling factor + void* scale_inv = nullptr; + + // AMax output + void* amax = nullptr; + int amax_byte_size = 0; + + // Whether to compute scale and amax + bool fp8_out = false; +}; + +struct BackwardKernelParams : public KernelParamsBase { + // Input: gradient wrt. LN FWD output. + void* dz = nullptr; + + // Input: extra tensor to add for fused backward+add + void* add = nullptr; + + // Workspace for Wgrad pre-reduction. + void* dbeta_part = nullptr; + void* dgamma_part = nullptr; + + // Output: Dgrad. + void* dx = nullptr; + // Output: Wgrad. + void* dbeta = nullptr; + void* dgamma = nullptr; +}; + +using BackwardAddKernelParams = BackwardKernelParams; + +} // namespace normalization +} // namespace transformer_engine + +#endif // TRANSFORMER_ENGINE_COMMON_NORM_KERNEL_PARAMS_H_ diff --git a/transformer_engine/common/normalization/kernel_traits.h b/transformer_engine/common/normalization/kernel_traits.h index 12fc095c38..41be82e742 100644 --- a/transformer_engine/common/normalization/kernel_traits.h +++ b/transformer_engine/common/normalization/kernel_traits.h @@ -7,12 +7,29 @@ #ifndef TRANSFORMER_ENGINE_COMMON_NORM_KERNEL_TRAITS_H_ #define TRANSFORMER_ENGINE_COMMON_NORM_KERNEL_TRAITS_H_ +#ifdef __CUDACC_RTC__ +#include "utils.cuh" +#else #include "../common.h" #include "../utils.cuh" +#endif namespace transformer_engine { namespace normalization { +#ifdef __CUDACC_RTC__ +// Under NVRTC the kernel sources do not include common.h (it pulls in cuDNN and +// other host-only headers), so the dtype aliases that the RTC name expressions +// reference are nototherwise visible (e.g. ::transformer_engine::normalization::fp16). +// Mirror the definitions from common.h here for the RTC. +// The underlying CUDA types come from utils.cuh (cuda_fp16/bf16/fp8 headers). +using fp32 = float; +using fp16 = half; +using bf16 = nv_bfloat16; +using fp8e4m3 = __nv_fp8_e4m3; +using fp8e5m2 = __nv_fp8_e5m2; +#endif // __CUDACC_RTC__ + template struct Kernel_traits_base { @@ -31,7 +48,7 @@ template > + compute_t_, index_t_, THREADS_PER_CTA_>> struct Kernel_traits_finalize : public Base { enum { ROWS_PER_CTA = Base::THREADS_PER_CTA / Base::THREADS_PER_WARP }; static_assert(static_cast(ROWS_PER_CTA) <= static_cast(Base::THREADS_PER_WARP)); @@ -69,7 +86,7 @@ template > + WARPS_M_ * WARPS_N_ * THREADS_PER_WARP>> struct Kernel_traits : public Base { using input_t = typename Base::input_t; using weight_t = typename Base::weight_t; diff --git a/transformer_engine/common/normalization/layernorm/ln_bwd_kernels.cuh b/transformer_engine/common/normalization/layernorm/ln_bwd_kernels.cuh index c4b00b87c3..0480b588b9 100644 --- a/transformer_engine/common/normalization/layernorm/ln_bwd_kernels.cuh +++ b/transformer_engine/common/normalization/layernorm/ln_bwd_kernels.cuh @@ -7,8 +7,13 @@ #ifndef TRANSFORMER_ENGINE_COMMON_LAYER_NORM_LN_BWD_KERNELS_CUH_ #define TRANSFORMER_ENGINE_COMMON_LAYER_NORM_LN_BWD_KERNELS_CUH_ +#ifdef __CUDACC_RTC__ +#include "kernel_params.h" +#include "utils.cuh" +#else #include "../../utils.cuh" #include "../common.h" +#endif namespace transformer_engine { namespace normalization { diff --git a/transformer_engine/common/normalization/layernorm/ln_bwd_semi_cuda_kernel.cu b/transformer_engine/common/normalization/layernorm/ln_bwd_semi_cuda_kernel.cu index 68aa0942c1..2e3d5fd622 100644 --- a/transformer_engine/common/normalization/layernorm/ln_bwd_semi_cuda_kernel.cu +++ b/transformer_engine/common/normalization/layernorm/ln_bwd_semi_cuda_kernel.cu @@ -7,6 +7,7 @@ #include "../../common.h" #include "../common.h" #include "../kernel_traits.h" +#include "../rtc_dispatch.h" #include "ln_bwd_kernels.cuh" using namespace transformer_engine::normalization; @@ -132,6 +133,30 @@ void launch_ln_bwd_general_(LaunchParams &launch_params, NVTE_CHECK_CUDA(cudaGetLastError()); } +#define REGISTER_NORM_LAUNCHER_LN_BWD_tuned(HIDDEN_SIZE, WTYPE, ITYPE, OTYPE, CTYPE, CTAS_PER_ROW, \ + WARPS_M, WARPS_N, BL_MAIN, BL_FINAL, STATIC_FALLBACK) \ + [[maybe_unused]] static const int \ + _ln_bwd_tuned_##HIDDEN_SIZE##_##WTYPE##_##ITYPE##_##OTYPE##_##CTYPE##_##CTAS_PER_ROW##_##WARPS_M##_##WARPS_N##_##BL_MAIN##_##BL_FINAL = \ + ([] { \ + ::transformer_engine::normalization::rtc_norm::register_ln_bwd_tuned( \ + TypeToDType::value, TypeToDType::value, TypeToDType::value, \ + TypeToDType::value, HIDDEN_SIZE, CTAS_PER_ROW, WARPS_M, WARPS_N, BL_MAIN, \ + BL_FINAL, STATIC_FALLBACK); \ + return 0; \ + })() +#define REGISTER_NORM_LAUNCHER_LN_BWD_general(HIDDEN_SIZE, WTYPE, ITYPE, OTYPE, CTYPE, WARPS_M, \ + WARPS_N, BL_MAIN, BL_FINAL, STATIC_FALLBACK) \ + [[maybe_unused]] static const int \ + _ln_bwd_general_##HIDDEN_SIZE##_##WTYPE##_##ITYPE##_##OTYPE##_##CTYPE##_##WARPS_M##_##WARPS_N##_##BL_MAIN##_##BL_FINAL = \ + ([] { \ + ::transformer_engine::normalization::rtc_norm::register_ln_bwd_general( \ + TypeToDType::value, TypeToDType::value, TypeToDType::value, \ + TypeToDType::value, HIDDEN_SIZE, WARPS_M, WARPS_N, BL_MAIN, BL_FINAL, \ + STATIC_FALLBACK); \ + return 0; \ + })() + +#if NVTE_BUILD_LEGACY_STATIC_NORM #define REGISTER_NORM_LAUNCHER(NORM_TYPE, NORM_STAGE, LAUNCH_TYPE, HIDDEN_SIZE, WTYPE, ITYPE, \ OTYPE, CTYPE, ...) \ namespace { \ @@ -141,10 +166,16 @@ void launch_ln_bwd_general_(LaunchParams &launch_params, launch_ln_bwd_##LAUNCH_TYPE##_(launch_params, configure_params); \ } \ - REGISTER_NORM_BASE( \ - NORM_TYPE, NORM_STAGE, LAUNCH_TYPE, HIDDEN_SIZE, WTYPE, ITYPE, OTYPE, CTYPE, \ + REGISTER_NORM_LAUNCHER_LN_BWD_##LAUNCH_TYPE( \ + HIDDEN_SIZE, WTYPE, ITYPE, OTYPE, CTYPE, __VA_ARGS__, \ norm_##NORM_TYPE##_##NORM_STAGE##_##LAUNCH_TYPE##_##HIDDEN_SIZE##_##WTYPE##_##ITYPE##_##OTYPE##_##CTYPE); \ - } // namespace + } // namespace +#else +#define REGISTER_NORM_LAUNCHER(NORM_TYPE, NORM_STAGE, LAUNCH_TYPE, HIDDEN_SIZE, WTYPE, ITYPE, \ + OTYPE, CTYPE, ...) \ + REGISTER_NORM_LAUNCHER_LN_BWD_##LAUNCH_TYPE(HIDDEN_SIZE, WTYPE, ITYPE, OTYPE, CTYPE, \ + __VA_ARGS__, nullptr) +#endif // NVTE_BUILD_LEGACY_STATIC_NORM // Create tuned launch function and register. Macro signature: // HIDDEN_SIZE, WTYPE, ITYPE, OTYPE, CTYPE, CTAS_PER_ROW, ... diff --git a/transformer_engine/common/normalization/layernorm/ln_fwd_cuda_kernel.cu b/transformer_engine/common/normalization/layernorm/ln_fwd_cuda_kernel.cu index 464df8d276..5be13357b1 100644 --- a/transformer_engine/common/normalization/layernorm/ln_fwd_cuda_kernel.cu +++ b/transformer_engine/common/normalization/layernorm/ln_fwd_cuda_kernel.cu @@ -6,6 +6,7 @@ #include "../common.h" #include "../kernel_traits.h" +#include "../rtc_dispatch.h" #include "ln_fwd_kernels.cuh" using namespace transformer_engine::normalization; @@ -101,6 +102,33 @@ void launch_ln_fwd_general_(LaunchParams &launch_params, } } +// Register a single RTC-first dispatcher. When the static fallback is enabled, +// its function pointer is passed to the dispatcher and selected only when +// NVTE_DISABLE_NVRTC=1. +#define REGISTER_NORM_LAUNCHER_LN_FWD_tuned(HIDDEN_SIZE, WTYPE, ITYPE, OTYPE, CTYPE, CTAS_PER_ROW, \ + WARPS_M, WARPS_N, BYTES_PER_LDG, STATIC_FALLBACK) \ + [[maybe_unused]] static const int \ + _ln_fwd_tuned_##HIDDEN_SIZE##_##WTYPE##_##ITYPE##_##OTYPE##_##CTYPE##_##CTAS_PER_ROW##_##WARPS_M##_##WARPS_N##_##BYTES_PER_LDG = \ + ([] { \ + ::transformer_engine::normalization::rtc_norm::register_ln_fwd_tuned( \ + TypeToDType::value, TypeToDType::value, TypeToDType::value, \ + TypeToDType::value, HIDDEN_SIZE, CTAS_PER_ROW, WARPS_M, WARPS_N, \ + BYTES_PER_LDG, STATIC_FALLBACK); \ + return 0; \ + })() +#define REGISTER_NORM_LAUNCHER_LN_FWD_general(HIDDEN_SIZE, WTYPE, ITYPE, OTYPE, CTYPE, WARPS_M, \ + WARPS_N, BYTES_PER_LDG, STATIC_FALLBACK) \ + [[maybe_unused]] static const int \ + _ln_fwd_general_##HIDDEN_SIZE##_##WTYPE##_##ITYPE##_##OTYPE##_##CTYPE##_##WARPS_M##_##WARPS_N##_##BYTES_PER_LDG = \ + ([] { \ + ::transformer_engine::normalization::rtc_norm::register_ln_fwd_general( \ + TypeToDType::value, TypeToDType::value, TypeToDType::value, \ + TypeToDType::value, HIDDEN_SIZE, WARPS_M, WARPS_N, BYTES_PER_LDG, \ + STATIC_FALLBACK); \ + return 0; \ + })() + +#if NVTE_BUILD_LEGACY_STATIC_NORM #define REGISTER_NORM_LAUNCHER(NORM_TYPE, NORM_STAGE, LAUNCH_TYPE, HIDDEN_SIZE, WTYPE, ITYPE, \ OTYPE, CTYPE, ...) \ namespace { \ @@ -110,10 +138,16 @@ void launch_ln_fwd_general_(LaunchParams &launch_params, launch_ln_fwd_##LAUNCH_TYPE##_(launch_params, configure_params); \ } \ - REGISTER_NORM_BASE( \ - NORM_TYPE, NORM_STAGE, LAUNCH_TYPE, HIDDEN_SIZE, WTYPE, ITYPE, OTYPE, CTYPE, \ + REGISTER_NORM_LAUNCHER_LN_FWD_##LAUNCH_TYPE( \ + HIDDEN_SIZE, WTYPE, ITYPE, OTYPE, CTYPE, __VA_ARGS__, \ norm_##NORM_TYPE##_##NORM_STAGE##_##LAUNCH_TYPE##_##HIDDEN_SIZE##_##WTYPE##_##ITYPE##_##OTYPE##_##CTYPE); \ } // namespace +#else +#define REGISTER_NORM_LAUNCHER(NORM_TYPE, NORM_STAGE, LAUNCH_TYPE, HIDDEN_SIZE, WTYPE, ITYPE, \ + OTYPE, CTYPE, ...) \ + REGISTER_NORM_LAUNCHER_LN_FWD_##LAUNCH_TYPE(HIDDEN_SIZE, WTYPE, ITYPE, OTYPE, CTYPE, \ + __VA_ARGS__, nullptr) +#endif // NVTE_BUILD_LEGACY_STATIC_NORM // Create tuned launch function and register. Macro signature: // HIDDEN_SIZE, WTYPE, ITYPE, OTYPE, CTYPE, CTAS_PER_ROW, WARPS_M, WARPS_N, BYTES_PER_LDG diff --git a/transformer_engine/common/normalization/layernorm/ln_fwd_kernels.cuh b/transformer_engine/common/normalization/layernorm/ln_fwd_kernels.cuh index 5a37cf46da..0e350146ea 100644 --- a/transformer_engine/common/normalization/layernorm/ln_fwd_kernels.cuh +++ b/transformer_engine/common/normalization/layernorm/ln_fwd_kernels.cuh @@ -7,11 +7,16 @@ #ifndef TRANSFORMER_ENGINE_COMMON_LAYER_NORM_LN_FWD_KERNELS_CUH_ #define TRANSFORMER_ENGINE_COMMON_LAYER_NORM_LN_FWD_KERNELS_CUH_ +#ifdef __CUDACC_RTC__ +#include "kernel_params.h" +#include "utils.cuh" +#else #include #include #include "../../utils.cuh" #include "../common.h" +#endif namespace transformer_engine { namespace normalization { @@ -147,7 +152,7 @@ __global__ __launch_bounds__(Ktraits::THREADS_PER_CTA) void ln_fwd_tuned_kernel( if (requires_amax) { amax = reduce_max(amax, warp); if (threadIdx.x == 0) { - static_assert(std::is_same::value); + static_assert(transformer_engine::detail::is_same::value); atomicMaxFloat(reinterpret_cast(params.amax), amax); } } @@ -323,7 +328,7 @@ __global__ __launch_bounds__(Ktraits::THREADS_PER_CTA) void ln_fwd_general_kerne if (requires_amax) { amax = reduce_max(amax, warp); if (threadIdx.x == 0) { - static_assert(std::is_same::value); + static_assert(transformer_engine::detail::is_same::value); atomicMaxFloat(reinterpret_cast(params.amax), amax); } } diff --git a/transformer_engine/common/normalization/layernorm/rtc/ln_bwd_kernel.cu b/transformer_engine/common/normalization/layernorm/rtc/ln_bwd_kernel.cu new file mode 100644 index 0000000000..f93f9e6edb --- /dev/null +++ b/transformer_engine/common/normalization/layernorm/rtc/ln_bwd_kernel.cu @@ -0,0 +1,8 @@ +/************************************************************************* + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * + * See LICENSE for license information. + ************************************************************************/ + +#include "kernel_traits.h" +#include "ln_bwd_kernels.cuh" diff --git a/transformer_engine/common/normalization/layernorm/rtc/ln_fwd_kernel.cu b/transformer_engine/common/normalization/layernorm/rtc/ln_fwd_kernel.cu new file mode 100644 index 0000000000..5c34f05e8d --- /dev/null +++ b/transformer_engine/common/normalization/layernorm/rtc/ln_fwd_kernel.cu @@ -0,0 +1,12 @@ +/************************************************************************* + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * + * See LICENSE for license information. + ************************************************************************/ + +// NVRTC source file for LayerNorm forward kernels. The host bundles this as a +// string header; specific template instantiations are requested via +// nvrtcAddNameExpression at runtime. + +#include "kernel_traits.h" +#include "ln_fwd_kernels.cuh" diff --git a/transformer_engine/common/normalization/rmsnorm/rmsnorm_bwd_kernels.cuh b/transformer_engine/common/normalization/rmsnorm/rmsnorm_bwd_kernels.cuh index d620ee5260..492edd65f6 100644 --- a/transformer_engine/common/normalization/rmsnorm/rmsnorm_bwd_kernels.cuh +++ b/transformer_engine/common/normalization/rmsnorm/rmsnorm_bwd_kernels.cuh @@ -7,10 +7,13 @@ #ifndef TRANSFORMER_ENGINE_COMMON_RMSNORM_RMSNORM_BWD_KERNELS_CUH_ #define TRANSFORMER_ENGINE_COMMON_RMSNORM_RMSNORM_BWD_KERNELS_CUH_ -#include - +#ifdef __CUDACC_RTC__ +#include "kernel_params.h" +#include "utils.cuh" +#else #include "../../utils.cuh" #include "../common.h" +#endif namespace transformer_engine { namespace normalization { @@ -18,19 +21,34 @@ namespace normalization { struct maybe_not_t {}; template -using maybe_t = std::conditional_t; - -template +using maybe_t = transformer_engine::detail::conditional_t; + +// dx and add share storage; `add` is positioned at the tail of the `dx` +// storage via leading padding. NeedsPadding is false when dx_t and add_t are +// the same size (or add_t is larger), in which case the padding array would be +// zero-length -- legal as a GNU extension under nvcc but rejected by NVRTC. The +// no-padding specialization below covers that case so both compilers are happy +// while keeping an identical layout. +template sizeof(maybe_t))> union dx_add_t { using add_t = maybe_t; using dx_t = Ivec; struct { - char _padding[sizeof(dx_t) > sizeof(add_t) ? sizeof(dx_t) - sizeof(add_t) : 0]; + char _padding[sizeof(dx_t) - sizeof(add_t)]; add_t add; }; dx_t dx; }; +template +union dx_add_t { + using add_t = maybe_t; + using dx_t = Ivec; + add_t add; + dx_t dx; +}; + template __global__ __launch_bounds__(Ktraits::THREADS_PER_CTA) void rmsnorm_bwd_tuned_kernel( BackwardKernelParams params) { diff --git a/transformer_engine/common/normalization/rmsnorm/rmsnorm_bwd_semi_cuda_kernel.cu b/transformer_engine/common/normalization/rmsnorm/rmsnorm_bwd_semi_cuda_kernel.cu index 60238f256d..7e8a1cd624 100644 --- a/transformer_engine/common/normalization/rmsnorm/rmsnorm_bwd_semi_cuda_kernel.cu +++ b/transformer_engine/common/normalization/rmsnorm/rmsnorm_bwd_semi_cuda_kernel.cu @@ -6,6 +6,7 @@ #include "../common.h" #include "../kernel_traits.h" +#include "../rtc_dispatch.h" #include "rmsnorm_bwd_kernels.cuh" using namespace transformer_engine::normalization; @@ -133,6 +134,60 @@ void launch_rmsnorm_bwd_general_(LaunchParams &launch_para NVTE_CHECK_CUDA(cudaGetLastError()); } +// Two-level dispatch: stage (Backward / BackwardAdd) and launch type +// (tuned / general). The optional static fallback is selected only when +// NVTE_DISABLE_NVRTC=1. +#define REGISTER_NORM_LAUNCHER_RMSN_BWD_tuned_Backward(HIDDEN_SIZE, WTYPE, ITYPE, OTYPE, CTYPE, \ + CTAS_PER_ROW, WARPS_M, WARPS_N, BL_MAIN, \ + BL_FINAL, STATIC_FALLBACK) \ + [[maybe_unused]] static const int \ + _rmsn_bwd_tuned_##HIDDEN_SIZE##_##WTYPE##_##ITYPE##_##OTYPE##_##CTYPE##_##CTAS_PER_ROW##_##WARPS_M##_##WARPS_N##_##BL_MAIN##_##BL_FINAL = \ + ([] { \ + ::transformer_engine::normalization::rtc_norm::register_rmsnorm_bwd_tuned( \ + TypeToDType::value, TypeToDType::value, TypeToDType::value, \ + TypeToDType::value, HIDDEN_SIZE, CTAS_PER_ROW, WARPS_M, WARPS_N, BL_MAIN, \ + BL_FINAL, false, STATIC_FALLBACK); \ + return 0; \ + })() +#define REGISTER_NORM_LAUNCHER_RMSN_BWD_tuned_BackwardAdd(HIDDEN_SIZE, WTYPE, ITYPE, OTYPE, CTYPE, \ + CTAS_PER_ROW, WARPS_M, WARPS_N, BL_MAIN, \ + BL_FINAL, ADD_FLAG, STATIC_FALLBACK) \ + static_assert(ADD_FLAG, "RMSNorm BackwardAdd registrations require ADD_FLAG=true"); \ + [[maybe_unused]] static const int \ + _rmsn_bwd_tuned_add_##HIDDEN_SIZE##_##WTYPE##_##ITYPE##_##OTYPE##_##CTYPE##_##CTAS_PER_ROW##_##WARPS_M##_##WARPS_N##_##BL_MAIN##_##BL_FINAL = \ + ([] { \ + ::transformer_engine::normalization::rtc_norm::register_rmsnorm_bwd_tuned( \ + TypeToDType::value, TypeToDType::value, TypeToDType::value, \ + TypeToDType::value, HIDDEN_SIZE, CTAS_PER_ROW, WARPS_M, WARPS_N, BL_MAIN, \ + BL_FINAL, ADD_FLAG, STATIC_FALLBACK); \ + return 0; \ + })() +#define REGISTER_NORM_LAUNCHER_RMSN_BWD_general_Backward( \ + HIDDEN_SIZE, WTYPE, ITYPE, OTYPE, CTYPE, WARPS_M, WARPS_N, BL_MAIN, BL_FINAL, STATIC_FALLBACK) \ + [[maybe_unused]] static const int \ + _rmsn_bwd_general_##HIDDEN_SIZE##_##WTYPE##_##ITYPE##_##OTYPE##_##CTYPE##_##WARPS_M##_##WARPS_N##_##BL_MAIN##_##BL_FINAL = \ + ([] { \ + ::transformer_engine::normalization::rtc_norm::register_rmsnorm_bwd_general( \ + TypeToDType::value, TypeToDType::value, TypeToDType::value, \ + TypeToDType::value, HIDDEN_SIZE, WARPS_M, WARPS_N, BL_MAIN, BL_FINAL, \ + false, STATIC_FALLBACK); \ + return 0; \ + })() +#define REGISTER_NORM_LAUNCHER_RMSN_BWD_general_BackwardAdd(HIDDEN_SIZE, WTYPE, ITYPE, OTYPE, \ + CTYPE, WARPS_M, WARPS_N, BL_MAIN, \ + BL_FINAL, ADD_FLAG, STATIC_FALLBACK) \ + static_assert(ADD_FLAG, "RMSNorm BackwardAdd registrations require ADD_FLAG=true"); \ + [[maybe_unused]] static const int \ + _rmsn_bwd_general_add_##HIDDEN_SIZE##_##WTYPE##_##ITYPE##_##OTYPE##_##CTYPE##_##WARPS_M##_##WARPS_N##_##BL_MAIN##_##BL_FINAL = \ + ([] { \ + ::transformer_engine::normalization::rtc_norm::register_rmsnorm_bwd_general( \ + TypeToDType::value, TypeToDType::value, TypeToDType::value, \ + TypeToDType::value, HIDDEN_SIZE, WARPS_M, WARPS_N, BL_MAIN, BL_FINAL, \ + ADD_FLAG, STATIC_FALLBACK); \ + return 0; \ + })() + +#if NVTE_BUILD_LEGACY_STATIC_NORM #define REGISTER_NORM_LAUNCHER(NORM_TYPE, NORM_STAGE, LAUNCH_TYPE, HIDDEN_SIZE, WTYPE, ITYPE, \ OTYPE, CTYPE, ...) \ namespace { \ @@ -142,10 +197,16 @@ void launch_rmsnorm_bwd_general_(LaunchParams &launch_para launch_rmsnorm_bwd_##LAUNCH_TYPE##_(launch_params, configure_params); \ } \ - REGISTER_NORM_BASE( \ - NORM_TYPE, NORM_STAGE, LAUNCH_TYPE, HIDDEN_SIZE, WTYPE, ITYPE, OTYPE, CTYPE, \ + REGISTER_NORM_LAUNCHER_RMSN_BWD_##LAUNCH_TYPE##_##NORM_STAGE( \ + HIDDEN_SIZE, WTYPE, ITYPE, OTYPE, CTYPE, __VA_ARGS__, \ norm_##NORM_TYPE##_##NORM_STAGE##_##LAUNCH_TYPE##_##HIDDEN_SIZE##_##WTYPE##_##ITYPE##_##OTYPE##_##CTYPE); \ } // namespace +#else +#define REGISTER_NORM_LAUNCHER(NORM_TYPE, NORM_STAGE, LAUNCH_TYPE, HIDDEN_SIZE, WTYPE, ITYPE, \ + OTYPE, CTYPE, ...) \ + REGISTER_NORM_LAUNCHER_RMSN_BWD_##LAUNCH_TYPE##_##NORM_STAGE(HIDDEN_SIZE, WTYPE, ITYPE, OTYPE, \ + CTYPE, __VA_ARGS__, nullptr) +#endif // NVTE_BUILD_LEGACY_STATIC_NORM // Create rmsnorm bwd tuned launch function and register. Macro signature: // HIDDEN_SIZE, WTYPE, ITYPE, OTYPE, CTYPE, CTAS_PER_ROW, ... diff --git a/transformer_engine/common/normalization/rmsnorm/rmsnorm_fwd_cuda_kernel.cu b/transformer_engine/common/normalization/rmsnorm/rmsnorm_fwd_cuda_kernel.cu index 5522fd5c6b..c3308cbce0 100644 --- a/transformer_engine/common/normalization/rmsnorm/rmsnorm_fwd_cuda_kernel.cu +++ b/transformer_engine/common/normalization/rmsnorm/rmsnorm_fwd_cuda_kernel.cu @@ -6,6 +6,7 @@ #include "../common.h" #include "../kernel_traits.h" +#include "../rtc_dispatch.h" #include "rmsnorm_fwd_kernels.cuh" using namespace transformer_engine::normalization; @@ -102,6 +103,31 @@ void launch_rmsnorm_fwd_general_(LaunchParams &launch_param } } +#define REGISTER_NORM_LAUNCHER_RMSN_FWD_tuned(HIDDEN_SIZE, WTYPE, ITYPE, OTYPE, CTYPE, \ + CTAS_PER_ROW, WARPS_M, WARPS_N, BYTES_PER_LDG, \ + STATIC_FALLBACK) \ + [[maybe_unused]] static const int \ + _rmsn_fwd_tuned_##HIDDEN_SIZE##_##WTYPE##_##ITYPE##_##OTYPE##_##CTYPE##_##CTAS_PER_ROW##_##WARPS_M##_##WARPS_N##_##BYTES_PER_LDG = \ + ([] { \ + ::transformer_engine::normalization::rtc_norm::register_rmsnorm_fwd_tuned( \ + TypeToDType::value, TypeToDType::value, TypeToDType::value, \ + TypeToDType::value, HIDDEN_SIZE, CTAS_PER_ROW, WARPS_M, WARPS_N, \ + BYTES_PER_LDG, STATIC_FALLBACK); \ + return 0; \ + })() +#define REGISTER_NORM_LAUNCHER_RMSN_FWD_general(HIDDEN_SIZE, WTYPE, ITYPE, OTYPE, CTYPE, WARPS_M, \ + WARPS_N, BYTES_PER_LDG, STATIC_FALLBACK) \ + [[maybe_unused]] static const int \ + _rmsn_fwd_general_##HIDDEN_SIZE##_##WTYPE##_##ITYPE##_##OTYPE##_##CTYPE##_##WARPS_M##_##WARPS_N##_##BYTES_PER_LDG = \ + ([] { \ + ::transformer_engine::normalization::rtc_norm::register_rmsnorm_fwd_general( \ + TypeToDType::value, TypeToDType::value, TypeToDType::value, \ + TypeToDType::value, HIDDEN_SIZE, WARPS_M, WARPS_N, BYTES_PER_LDG, \ + STATIC_FALLBACK); \ + return 0; \ + })() + +#if NVTE_BUILD_LEGACY_STATIC_NORM #define REGISTER_NORM_LAUNCHER(NORM_TYPE, NORM_STAGE, LAUNCH_TYPE, HIDDEN_SIZE, WTYPE, ITYPE, \ OTYPE, CTYPE, ...) \ namespace { \ @@ -111,10 +137,16 @@ void launch_rmsnorm_fwd_general_(LaunchParams &launch_param launch_rmsnorm_fwd_##LAUNCH_TYPE##_(launch_params, configure_params); \ } \ - REGISTER_NORM_BASE( \ - NORM_TYPE, NORM_STAGE, LAUNCH_TYPE, HIDDEN_SIZE, WTYPE, ITYPE, OTYPE, CTYPE, \ + REGISTER_NORM_LAUNCHER_RMSN_FWD_##LAUNCH_TYPE( \ + HIDDEN_SIZE, WTYPE, ITYPE, OTYPE, CTYPE, __VA_ARGS__, \ norm_##NORM_TYPE##_##NORM_STAGE##_##LAUNCH_TYPE##_##HIDDEN_SIZE##_##WTYPE##_##ITYPE##_##OTYPE##_##CTYPE); \ } // namespace +#else +#define REGISTER_NORM_LAUNCHER(NORM_TYPE, NORM_STAGE, LAUNCH_TYPE, HIDDEN_SIZE, WTYPE, ITYPE, \ + OTYPE, CTYPE, ...) \ + REGISTER_NORM_LAUNCHER_RMSN_FWD_##LAUNCH_TYPE(HIDDEN_SIZE, WTYPE, ITYPE, OTYPE, CTYPE, \ + __VA_ARGS__, nullptr) +#endif // NVTE_BUILD_LEGACY_STATIC_NORM // Create rmsnorm tuned launch function and register. Macro signature: // HIDDEN_SIZE, WTYPE, ITYPE, OTYPE, CTYPE, CTAS_PER_ROW, WARPS_M, WARPS_N, BYTES_PER_LDG diff --git a/transformer_engine/common/normalization/rmsnorm/rmsnorm_fwd_kernels.cuh b/transformer_engine/common/normalization/rmsnorm/rmsnorm_fwd_kernels.cuh index 900fb58be2..bb284a9ab1 100644 --- a/transformer_engine/common/normalization/rmsnorm/rmsnorm_fwd_kernels.cuh +++ b/transformer_engine/common/normalization/rmsnorm/rmsnorm_fwd_kernels.cuh @@ -7,11 +7,16 @@ #ifndef TRANSFORMER_ENGINE_COMMON_RMSNORM_RMSNORM_FWD_KERNELS_CUH_ #define TRANSFORMER_ENGINE_COMMON_RMSNORM_RMSNORM_FWD_KERNELS_CUH_ +#ifdef __CUDACC_RTC__ +#include "kernel_params.h" +#include "utils.cuh" +#else #include #include #include "../../utils.cuh" #include "../common.h" +#endif namespace transformer_engine { namespace normalization { @@ -139,7 +144,7 @@ __global__ __launch_bounds__(Ktraits::THREADS_PER_CTA) void rmsnorm_fwd_tuned_ke if (requires_amax) { amax = reduce_max(amax, warp); if (threadIdx.x == 0) { - static_assert(std::is_same::value); + static_assert(transformer_engine::detail::is_same::value); atomicMaxFloat(reinterpret_cast(params.amax), amax); } } @@ -298,7 +303,7 @@ __global__ __launch_bounds__(Ktraits::THREADS_PER_CTA) void rmsnorm_fwd_general_ if (requires_amax) { amax = reduce_max(amax, warp); if (threadIdx.x == 0) { - static_assert(std::is_same::value); + static_assert(transformer_engine::detail::is_same::value); atomicMaxFloat(reinterpret_cast(params.amax), amax); } } diff --git a/transformer_engine/common/normalization/rmsnorm/rtc/rmsnorm_bwd_kernel.cu b/transformer_engine/common/normalization/rmsnorm/rtc/rmsnorm_bwd_kernel.cu new file mode 100644 index 0000000000..706e79870b --- /dev/null +++ b/transformer_engine/common/normalization/rmsnorm/rtc/rmsnorm_bwd_kernel.cu @@ -0,0 +1,8 @@ +/************************************************************************* + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * + * See LICENSE for license information. + ************************************************************************/ + +#include "kernel_traits.h" +#include "rmsnorm_bwd_kernels.cuh" diff --git a/transformer_engine/common/normalization/rmsnorm/rtc/rmsnorm_fwd_kernel.cu b/transformer_engine/common/normalization/rmsnorm/rtc/rmsnorm_fwd_kernel.cu new file mode 100644 index 0000000000..4d1bd10a7b --- /dev/null +++ b/transformer_engine/common/normalization/rmsnorm/rtc/rmsnorm_fwd_kernel.cu @@ -0,0 +1,8 @@ +/************************************************************************* + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * + * See LICENSE for license information. + ************************************************************************/ + +#include "kernel_traits.h" +#include "rmsnorm_fwd_kernels.cuh" diff --git a/transformer_engine/common/normalization/rtc_dispatch.cpp b/transformer_engine/common/normalization/rtc_dispatch.cpp new file mode 100644 index 0000000000..5ae81f4233 --- /dev/null +++ b/transformer_engine/common/normalization/rtc_dispatch.cpp @@ -0,0 +1,742 @@ +/************************************************************************* + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * + * See LICENSE for license information. + ************************************************************************/ + +// NVRTC-backed registry registration for LayerNorm/RMSNorm forward + backward +// launchers. Each per-config REGISTER_NORM_LAUNCHER macro expands to a call +// into one of the register_*_tuned/general functions defined here. When +// NVTE_BUILD_LEGACY_STATIC_NORM=ON, that call also supplies a static fallback. +// The registered closure prefers NVRTC and selects the fallback only when +// NVTE_DISABLE_NVRTC=1. + +#include "rtc_dispatch.h" + +#include +#include +#include + +#include "../util/cuda_driver.h" +#include "../util/rtc.h" +#include "../util/string.h" +#include "common.h" + +// NVRTC source strings for the four kernel families. These are tiny stub +// files (each is a couple of #include lines that pull in kernel_traits + the +// matching kernel header); NVRTC then instantiates a specific +// (Kernel_traits) on demand via nvrtcAddNameExpression. +#include "string_code_normalization_kernel_params_h.h" +#include "string_code_normalization_kernel_traits_h.h" +#include "string_code_normalization_layernorm_ln_bwd_kernels_cuh.h" +#include "string_code_normalization_layernorm_ln_fwd_kernels_cuh.h" +#include "string_code_normalization_layernorm_rtc_ln_bwd_kernel_cu.h" +#include "string_code_normalization_layernorm_rtc_ln_fwd_kernel_cu.h" +#include "string_code_normalization_rmsnorm_rmsnorm_bwd_kernels_cuh.h" +#include "string_code_normalization_rmsnorm_rmsnorm_fwd_kernels_cuh.h" +#include "string_code_normalization_rmsnorm_rtc_rmsnorm_bwd_kernel_cu.h" +#include "string_code_normalization_rmsnorm_rtc_rmsnorm_fwd_kernel_cu.h" + +namespace transformer_engine { +namespace normalization { +namespace rtc_norm { + +namespace { + +const std::vector& norm_headers() { + static const std::vector headers = { + {string_code_normalization_kernel_params_h, "kernel_params.h"}, + {string_code_normalization_kernel_traits_h, "kernel_traits.h"}, + {string_code_normalization_layernorm_ln_fwd_kernels_cuh, "ln_fwd_kernels.cuh"}, + {string_code_normalization_layernorm_ln_bwd_kernels_cuh, "ln_bwd_kernels.cuh"}, + {string_code_normalization_rmsnorm_rmsnorm_fwd_kernels_cuh, "rmsnorm_fwd_kernels.cuh"}, + {string_code_normalization_rmsnorm_rmsnorm_bwd_kernels_cuh, "rmsnorm_bwd_kernels.cuh"}, + }; + return headers; +} + +void compile_norm_kernel(rtc::KernelManager& manager, const std::string& label, + const std::string& kernel_expr, const char* rtc_source, + const char* filename) { + manager.compile(label, kernel_expr, rtc_source, filename, {}, norm_headers()); +} + +// Map our DType enum onto the C++ type names used inside the norm RTC sources. +// Aliases come from normalization/common.h (`using bf16 = nv_bfloat16;` etc.). +const char* cpp_name_for(DType dt) { + switch (dt) { + case DType::kFloat32: + return "::transformer_engine::normalization::fp32"; + case DType::kFloat16: + return "::transformer_engine::normalization::fp16"; + case DType::kBFloat16: + return "::transformer_engine::normalization::bf16"; + case DType::kFloat8E4M3: + return "::transformer_engine::normalization::fp8e4m3"; + case DType::kFloat8E5M2: + return "::transformer_engine::normalization::fp8e5m2"; + default: + NVTE_ERROR("Unsupported DType for norm RTC dispatch"); + } +} + +int byte_size_of(DType dt) { + switch (dt) { + case DType::kFloat32: + return 4; + case DType::kFloat16: + case DType::kBFloat16: + return 2; + case DType::kFloat8E4M3: + case DType::kFloat8E5M2: + return 1; + default: + NVTE_ERROR("Unsupported DType for norm RTC dispatch"); + } +} + +// Build a Kernel_traits template-argument list as a string. The C++ argument +// list matches: +// Kernel_traits +std::string kernel_traits_expr(DType wt, DType it, DType ot, DType ct, int hidden_size, + int ctas_per_row, int warps_m, int warps_n, int bytes_per_ldg) { + return concat_strings("::transformer_engine::normalization::Kernel_traits<", cpp_name_for(wt), + ", ", cpp_name_for(it), ", ", cpp_name_for(ot), ", ", cpp_name_for(ct), + ", uint32_t, ", hidden_size, ", ", ctas_per_row, ", ", warps_m, ", ", + warps_n, ", ", bytes_per_ldg, ">"); +} + +// Stats::SMEM_BYTES expressed in host code. +// stats_t = TypeToVec2::Type — for compute_t == fp32 that's float2 +// (8 bytes); for fp16 it's half2 (4 bytes); for bf16 it's nv_bfloat162 (4 +// bytes). Matches the formulas in utils.cuh: +// WARPS_N == 1 → 0 +// else → WARPS_M * WARPS_N * sizeof(stats_t) * 2 +int stats_smem_bytes(DType ctype, int warps_m, int warps_n) { + if (warps_n == 1) return 0; + int sizeof_stats_t; + switch (ctype) { + case DType::kFloat32: + sizeof_stats_t = 8; + break; + case DType::kFloat16: + case DType::kBFloat16: + sizeof_stats_t = 4; + break; + default: + NVTE_ERROR("Unsupported compute dtype for norm smem calc"); + } + return warps_m * warps_n * sizeof_stats_t * 2; +} + +// Reducer::SMEM_BYTES. +// reduce_t = TypeToVec2::Type — same sizes as stats_t. +int reducer_smem_bytes(DType ctype, int warps_m, int warps_n) { + return stats_smem_bytes(ctype, warps_m, warps_n); +} + +// Kernel_traits::SMEM_BYTES (used by backward launchers): +// SMEM_BYTES_DGRAD = Reducer::SMEM_BYTES +// SMEM_BYTES_WGRAD = (CTAS_PER_ROW > 1) ? 0 : WARPS_M * HIDDEN * sizeof(compute_t) +// SMEM_BYTES = DGRAD + WGRAD +int bwd_smem_bytes(DType ctype, int hidden_size, int ctas_per_row, int warps_m, int warps_n) { + const int dgrad = reducer_smem_bytes(ctype, warps_m, warps_n); + const int wgrad = (ctas_per_row > 1) ? 0 : warps_m * hidden_size * byte_size_of(ctype); + return dgrad + wgrad; +} + +template +bool try_static_fallback(StaticFallback static_fallback, + LaunchParams& launch_params, const bool configure_params) { + if (rtc::is_enabled()) { + return false; + } + NVTE_CHECK(static_fallback != nullptr, + "NVRTC is disabled and this normalization build has no static fallback. Rebuild with " + "NVTE_BUILD_LEGACY_STATIC_NORM=ON."); + static_fallback(launch_params, configure_params); + return true; +} + +// Common per-launch configure/launch helper. `kernel_expr` is the full +// templated kernel symbol (used as the nvrtcAddNameExpression argument); the +// closure captures everything needed. +template +void register_launcher(const std::string& label, const std::string& kernel_expr, + const char* rtc_source, const char* filename, TupleKeyType key, + int threads_per_cta, int dynamic_smem_bytes, int ctas_per_row, + bool needs_cooperative, int barrier_bytes_per_col, + int workspace_bytes_per_col, int dgamma_part_bytes_per_col, + StaticFallback static_fallback) { + auto closure = [label, kernel_expr, rtc_source, filename, threads_per_cta, dynamic_smem_bytes, + ctas_per_row, needs_cooperative, barrier_bytes_per_col, workspace_bytes_per_col, + dgamma_part_bytes_per_col, static_fallback](LaunchParams& launch_params, + const bool configure_params) { + if (try_static_fallback(static_fallback, launch_params, configure_params)) { + return; + } + auto& mgr = rtc::KernelManager::instance(); + if (!mgr.is_compiled(label)) { + compile_norm_kernel(mgr, label, kernel_expr, rtc_source, filename); + } + + if (configure_params) { + const int ctas_per_sm = + mgr.occupancy_max_active_blocks_per_sm(label, threads_per_cta, dynamic_smem_bytes); + launch_params.params.ctas_per_row = ctas_per_row; + launch_params.params.ctas_per_col = + launch_params.multiprocessorCount * ctas_per_sm / ctas_per_row; + if (ctas_per_row > 1) { + launch_params.barrier_bytes = barrier_bytes_per_col * launch_params.params.ctas_per_col; + launch_params.workspace_bytes = workspace_bytes_per_col * launch_params.params.ctas_per_col; + } + if (dgamma_part_bytes_per_col > 0) { + launch_params.dgamma_part_bytes = + dgamma_part_bytes_per_col * launch_params.params.ctas_per_col; + } + return; + } + + // Real launch. + if (dynamic_smem_bytes >= 48 * 1024) { + mgr.set_function_attribute(label, CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES, + dynamic_smem_bytes); + } + const auto stream = launch_params.stream; + const auto ctas_per_col = launch_params.params.ctas_per_col; + if (ctas_per_row == 1) { + mgr.launch(label, dim3(ctas_per_col), dim3(threads_per_cta), dynamic_smem_bytes, stream, + launch_params.params); + } else { + mgr.launch_cooperative(label, dim3(ctas_per_row * ctas_per_col), dim3(threads_per_cta), + dynamic_smem_bytes, stream, launch_params.params); + } + (void)needs_cooperative; + }; + TeNormalizationRegistry::registerFunction(key, std::move(closure)); +} + +} // namespace + +// ============================================================================ +// LayerNorm Forward +// ============================================================================ + +void register_ln_fwd_tuned(DType wt, DType it, DType ot, DType ct, int hidden, int cr, int wm, + int wn, int bl, StaticFallback static_fallback) { + const auto key = get_key(NVTE_Norm_Backend::Te, NVTE_Norm_Type::LayerNorm, + NVTE_Norm_Stage::Forward, wt, it, ot, ct, 0, hidden, false, true); + const std::string label = + concat_strings("ln_fwd_tuned,w=", static_cast(wt), ",i=", static_cast(it), + ",o=", static_cast(ot), ",c=", static_cast(ct), ",h=", hidden, + ",cr=", cr, ",wm=", wm, ",wn=", wn, ",bl=", bl); + const std::string traits_expr = kernel_traits_expr(wt, it, ot, ct, hidden, cr, wm, wn, bl); + const std::string kernel_expr = concat_strings( + "&::transformer_engine::normalization::ln_fwd_tuned_kernel<", traits_expr, ">"); + const int threads_per_cta = wm * wn * 32; + const int smem_bytes = stats_smem_bytes(ct, wm, wn); + // tuned path multi-CTA workspace formula: + // barrier_bytes = 2 * ctas_per_col * sizeof(index_t == uint32_t == 4 bytes) + // workspace_bytes = ctas_per_col * WARPS_M * CTAS_PER_ROW * sizeof(stats_t) * 2 + const int sizeof_stats_t = (ct == DType::kFloat32) ? 8 : 4; // float2 vs half2/bf162 + const int barrier_per_col = 2 * 4; + const int workspace_per_col = wm * cr * sizeof_stats_t * 2; + register_launcher( + label, kernel_expr, string_code_normalization_layernorm_rtc_ln_fwd_kernel_cu, + "ln_fwd_kernel.cu", key, threads_per_cta, smem_bytes, cr, /*needs_cooperative=*/cr > 1, + barrier_per_col, workspace_per_col, /*dgamma_part_bytes_per_col=*/0, static_fallback); +} + +void register_ln_fwd_general(DType wt, DType it, DType ot, DType ct, int hidden, int wm, int wn, + int bl, StaticFallback static_fallback) { + const auto key = get_key(NVTE_Norm_Backend::Te, NVTE_Norm_Type::LayerNorm, + NVTE_Norm_Stage::Forward, wt, it, ot, ct, 0, hidden, false, false); + const std::string label = + concat_strings("ln_fwd_general,w=", static_cast(wt), ",i=", static_cast(it), + ",o=", static_cast(ot), ",c=", static_cast(ct), ",h=", hidden, + ",wm=", wm, ",wn=", wn, ",bl=", bl); + // "general" path always uses CTAS_PER_ROW=1 in the Kernel_traits. + const std::string traits_expr = kernel_traits_expr(wt, it, ot, ct, hidden, 1, wm, wn, bl); + const std::string kernel_expr = concat_strings( + "&::transformer_engine::normalization::ln_fwd_general_kernel<", traits_expr, ">"); + const int threads_per_cta = wm * wn * 32; + const auto closure = [label, kernel_expr, threads_per_cta, hidden, wm, ct, static_fallback]( + LaunchParams& launch_params, + const bool configure_params) { + if (try_static_fallback(static_fallback, launch_params, configure_params)) { + return; + } + auto& mgr = rtc::KernelManager::instance(); + if (!mgr.is_compiled(label)) { + compile_norm_kernel(mgr, label, kernel_expr, + string_code_normalization_layernorm_rtc_ln_fwd_kernel_cu, + "ln_fwd_kernel.cu"); + } + auto ceil_div = [](int x, int y) { return (x + y - 1) / y; }; + const int rows = launch_params.params.rows; + const int cols = launch_params.params.cols; + int ctas_per_col = launch_params.params.ctas_per_col; + int ctas_per_row = launch_params.params.ctas_per_row; + if (configure_params) { + const int ctas_per_sm = + mgr.occupancy_max_active_blocks_per_sm(label, threads_per_cta, /*smem=*/0); + const int max_ctas = launch_params.multiprocessorCount * ctas_per_sm; + ctas_per_row = ceil_div(cols, hidden); + ctas_per_col = std::min(ceil_div(rows, wm), max_ctas / std::max(ctas_per_row, 1)); + launch_params.params.ctas_per_row = ctas_per_row; + launch_params.params.ctas_per_col = ctas_per_col; + if (ctas_per_row > 1) { + launch_params.barrier_bytes = 2 * ctas_per_col * sizeof(int); + // compute_t bytes + const int ctype_bytes = byte_size_of(ct); + launch_params.workspace_bytes = ctas_per_col * wm * ctas_per_row * ctype_bytes * 2; + } + return; + } + const auto stream = launch_params.stream; + if (ctas_per_row == 1) { + mgr.launch(label, dim3(ctas_per_row * ctas_per_col), dim3(threads_per_cta), 0, stream, + launch_params.params); + } else { + mgr.launch_cooperative(label, dim3(ctas_per_row * ctas_per_col), dim3(threads_per_cta), 0, + stream, launch_params.params); + } + }; + TeNormalizationRegistry::registerFunction(key, std::move(closure)); +} + +// ============================================================================ +// RMSNorm Forward (same shape as LayerNorm Forward) +// ============================================================================ + +void register_rmsnorm_fwd_tuned(DType wt, DType it, DType ot, DType ct, int hidden, int cr, int wm, + int wn, int bl, + StaticFallback static_fallback) { + const auto key = get_key(NVTE_Norm_Backend::Te, NVTE_Norm_Type::RMSNorm, NVTE_Norm_Stage::Forward, + wt, it, ot, ct, 0, hidden, false, true); + const std::string label = + concat_strings("rmsnorm_fwd_tuned,w=", static_cast(wt), ",i=", static_cast(it), + ",o=", static_cast(ot), ",c=", static_cast(ct), ",h=", hidden, + ",cr=", cr, ",wm=", wm, ",wn=", wn, ",bl=", bl); + const std::string traits_expr = kernel_traits_expr(wt, it, ot, ct, hidden, cr, wm, wn, bl); + const std::string kernel_expr = concat_strings( + "&::transformer_engine::normalization::rmsnorm_fwd_tuned_kernel<", traits_expr, ">"); + const int threads_per_cta = wm * wn * 32; + const int smem_bytes = stats_smem_bytes(ct, wm, wn); + const int sizeof_stats_t = (ct == DType::kFloat32) ? 8 : 4; + const int barrier_per_col = 2 * 4; + const int workspace_per_col = wm * cr * sizeof_stats_t * 2; + register_launcher( + label, kernel_expr, string_code_normalization_rmsnorm_rtc_rmsnorm_fwd_kernel_cu, + "rmsnorm_fwd_kernel.cu", key, threads_per_cta, smem_bytes, cr, /*needs_cooperative=*/cr > 1, + barrier_per_col, workspace_per_col, 0, static_fallback); +} + +void register_rmsnorm_fwd_general(DType wt, DType it, DType ot, DType ct, int hidden, int wm, + int wn, int bl, + StaticFallback static_fallback) { + const auto key = get_key(NVTE_Norm_Backend::Te, NVTE_Norm_Type::RMSNorm, NVTE_Norm_Stage::Forward, + wt, it, ot, ct, 0, hidden, false, false); + const std::string label = + concat_strings("rmsnorm_fwd_general,w=", static_cast(wt), ",i=", static_cast(it), + ",o=", static_cast(ot), ",c=", static_cast(ct), ",h=", hidden, + ",wm=", wm, ",wn=", wn, ",bl=", bl); + const std::string traits_expr = kernel_traits_expr(wt, it, ot, ct, hidden, 1, wm, wn, bl); + const std::string kernel_expr = concat_strings( + "&::transformer_engine::normalization::rmsnorm_fwd_general_kernel<", traits_expr, ">"); + const int threads_per_cta = wm * wn * 32; + const auto closure = [label, kernel_expr, threads_per_cta, hidden, wm, ct, static_fallback]( + LaunchParams& launch_params, + const bool configure_params) { + if (try_static_fallback(static_fallback, launch_params, configure_params)) { + return; + } + auto& mgr = rtc::KernelManager::instance(); + if (!mgr.is_compiled(label)) { + compile_norm_kernel(mgr, label, kernel_expr, + string_code_normalization_rmsnorm_rtc_rmsnorm_fwd_kernel_cu, + "rmsnorm_fwd_kernel.cu"); + } + auto ceil_div = [](int x, int y) { return (x + y - 1) / y; }; + const int rows = launch_params.params.rows; + const int cols = launch_params.params.cols; + int ctas_per_col = launch_params.params.ctas_per_col; + int ctas_per_row = launch_params.params.ctas_per_row; + if (configure_params) { + const int ctas_per_sm = mgr.occupancy_max_active_blocks_per_sm(label, threads_per_cta, 0); + const int max_ctas = launch_params.multiprocessorCount * ctas_per_sm; + ctas_per_row = ceil_div(cols, hidden); + ctas_per_col = std::min(ceil_div(rows, wm), max_ctas / std::max(ctas_per_row, 1)); + launch_params.params.ctas_per_row = ctas_per_row; + launch_params.params.ctas_per_col = ctas_per_col; + if (ctas_per_row > 1) { + launch_params.barrier_bytes = 2 * ctas_per_col * sizeof(int); + const int ctype_bytes = byte_size_of(ct); + launch_params.workspace_bytes = ctas_per_col * wm * ctas_per_row * ctype_bytes * 2; + } + return; + } + const auto stream = launch_params.stream; + if (ctas_per_row == 1) { + mgr.launch(label, dim3(ctas_per_row * ctas_per_col), dim3(threads_per_cta), 0, stream, + launch_params.params); + } else { + mgr.launch_cooperative(label, dim3(ctas_per_row * ctas_per_col), dim3(threads_per_cta), 0, + stream, launch_params.params); + } + }; + TeNormalizationRegistry::registerFunction(key, std::move(closure)); +} + +// ============================================================================ +// LayerNorm Backward (main kernel + finalize kernel) +// ============================================================================ + +void register_ln_bwd_tuned(DType wt, DType it, DType ot, DType ct, int hidden, int cr, int wm, + int wn, int bl_main, int bl_final, + StaticFallback static_fallback) { + const auto key = get_key(NVTE_Norm_Backend::Te, NVTE_Norm_Type::LayerNorm, + NVTE_Norm_Stage::Backward, wt, it, ot, ct, 0, hidden, false, true); + const std::string label = + concat_strings("ln_bwd_tuned,w=", static_cast(wt), ",i=", static_cast(it), + ",o=", static_cast(ot), ",c=", static_cast(ct), ",h=", hidden, + ",cr=", cr, ",wm=", wm, ",wn=", wn, ",bl=", bl_main, ",blf=", bl_final); + const std::string main_label = concat_strings(label, ",main"); + const std::string finalize_label = concat_strings(label, ",finalize"); + const std::string main_traits = kernel_traits_expr(wt, it, ot, ct, hidden, cr, wm, wn, bl_main); + const std::string main_kexpr = concat_strings( + "&::transformer_engine::normalization::ln_bwd_tuned_kernel<", main_traits, ">"); + // Kernel_traits_finalize + const std::string finalize_traits = + concat_strings("::transformer_engine::normalization::Kernel_traits_finalize<", hidden, ", ", + cpp_name_for(wt), ", ", cpp_name_for(it), ", ", cpp_name_for(ot), ", ", + cpp_name_for(ct), ", uint32_t, 1024, ", bl_final, ">"); + const std::string finalize_kexpr = concat_strings( + "&::transformer_engine::normalization::ln_bwd_finalize_tuned_kernel<", finalize_traits, ">"); + const int threads_per_cta = wm * wn * 32; + const int smem_bytes = bwd_smem_bytes(ct, hidden, cr, wm, wn); + const int sizeof_reduce_t = (ct == DType::kFloat32) ? 8 : 4; + // tuned backward: workspace = ctas_per_col * WARPS_M * CTAS_PER_ROW * sizeof(reduce_t) * 2 + // dgamma_part_bytes = ctas_per_col * cols * sizeof(compute_t) + // barrier_bytes = 2 * ctas_per_col * sizeof(index_t) + const int barrier_per_col = 2 * 4; + const int workspace_per_col = wm * cr * sizeof_reduce_t * 2; + const int dgamma_part_per_col = hidden * byte_size_of(ct); + + // Finalize kernel dims: Kernel_traits_finalize::THREADS_PER_CTA == 1024 + // Kernel_traits_finalize::CTAS = HIDDEN_SIZE / 32 (since COLS%32==0) + // COLS = HIDDEN_SIZE * sizeof(compute_t) / BYTES_PER_LDG_FINAL + // CTAS = COLS / 32 + const int colspass = hidden * byte_size_of(ct) / bl_final; + const int finalize_ctas = colspass / 32; + const int finalize_threads_per_cta = 1024; + + auto closure = [main_label, main_kexpr, finalize_label, finalize_kexpr, threads_per_cta, + smem_bytes, cr, barrier_per_col, workspace_per_col, dgamma_part_per_col, + finalize_ctas, finalize_threads_per_cta, + static_fallback](LaunchParams& launch_params, + const bool configure_params) { + if (try_static_fallback(static_fallback, launch_params, configure_params)) { + return; + } + auto& mgr = rtc::KernelManager::instance(); + if (!mgr.is_compiled(main_label)) { + compile_norm_kernel(mgr, main_label, main_kexpr, + string_code_normalization_layernorm_rtc_ln_bwd_kernel_cu, + "ln_bwd_kernel.cu"); + } + if (!mgr.is_compiled(finalize_label)) { + compile_norm_kernel(mgr, finalize_label, finalize_kexpr, + string_code_normalization_layernorm_rtc_ln_bwd_kernel_cu, + "ln_bwd_kernel.cu"); + } + if (configure_params) { + const int ctas_per_sm = + mgr.occupancy_max_active_blocks_per_sm(main_label, threads_per_cta, smem_bytes); + launch_params.params.ctas_per_row = cr; + launch_params.params.ctas_per_col = launch_params.multiprocessorCount * ctas_per_sm / cr; + if (cr > 1) { + launch_params.barrier_bytes = barrier_per_col * launch_params.params.ctas_per_col; + launch_params.workspace_bytes = workspace_per_col * launch_params.params.ctas_per_col; + } + launch_params.dgamma_part_bytes = dgamma_part_per_col * launch_params.params.ctas_per_col; + return; + } + if (smem_bytes >= 48 * 1024) { + mgr.set_function_attribute(main_label, CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES, + smem_bytes); + } + const auto stream = launch_params.stream; + const auto ctas_per_col = launch_params.params.ctas_per_col; + if (cr == 1) { + mgr.launch(main_label, dim3(ctas_per_col), dim3(threads_per_cta), smem_bytes, stream, + launch_params.params); + } else { + mgr.launch_cooperative(main_label, dim3(cr * ctas_per_col), dim3(threads_per_cta), smem_bytes, + stream, launch_params.params); + } + mgr.launch(finalize_label, dim3(finalize_ctas), dim3(finalize_threads_per_cta), 0, stream, + launch_params.params); + }; + TeNormalizationRegistry::registerFunction(key, std::move(closure)); +} + +void register_ln_bwd_general(DType wt, DType it, DType ot, DType ct, int hidden, int wm, int wn, + int bl_main, int bl_final, + StaticFallback static_fallback) { + const auto key = get_key(NVTE_Norm_Backend::Te, NVTE_Norm_Type::LayerNorm, + NVTE_Norm_Stage::Backward, wt, it, ot, ct, 0, hidden, false, false); + const std::string label = + concat_strings("ln_bwd_general,w=", static_cast(wt), ",i=", static_cast(it), + ",o=", static_cast(ot), ",c=", static_cast(ct), ",h=", hidden, + ",wm=", wm, ",wn=", wn, ",bl=", bl_main, ",blf=", bl_final); + const std::string main_label = concat_strings(label, ",main"); + const std::string finalize_label = concat_strings(label, ",finalize"); + const std::string traits = kernel_traits_expr(wt, it, ot, ct, hidden, 1, wm, wn, bl_main); + const std::string main_kexpr = + concat_strings("&::transformer_engine::normalization::ln_bwd_general_kernel<", traits, ">"); + // ln_bwd_finalize_general_kernel + const std::string finalize_kexpr = + concat_strings("&::transformer_engine::normalization::ln_bwd_finalize_general_kernel<", + cpp_name_for(wt), ", ", cpp_name_for(ct), ", 4, 1, ", bl_final, ", 32>"); + const int threads_per_cta = wm * wn * 32; + // general bwd uses ctas_per_row = ceil_div(cols, HIDDEN_SIZE); smem=0 for main kernel call. + const int ctype_bytes = byte_size_of(ct); + const int finalize_threads_per_warp = 32; + const int finalize_warps_n = 1; + const int finalize_warps_m = 4; + const int finalize_elts_n_per_cta = + finalize_threads_per_warp * finalize_warps_n * bl_final / ctype_bytes; + + auto closure = [main_label, main_kexpr, finalize_label, finalize_kexpr, threads_per_cta, hidden, + wm, ctype_bytes, finalize_elts_n_per_cta, finalize_warps_n, finalize_warps_m, + finalize_threads_per_warp, + static_fallback](LaunchParams& launch_params, + const bool configure_params) { + if (try_static_fallback(static_fallback, launch_params, configure_params)) { + return; + } + auto& mgr = rtc::KernelManager::instance(); + if (!mgr.is_compiled(main_label)) { + compile_norm_kernel(mgr, main_label, main_kexpr, + string_code_normalization_layernorm_rtc_ln_bwd_kernel_cu, + "ln_bwd_kernel.cu"); + } + if (!mgr.is_compiled(finalize_label)) { + compile_norm_kernel(mgr, finalize_label, finalize_kexpr, + string_code_normalization_layernorm_rtc_ln_bwd_kernel_cu, + "ln_bwd_kernel.cu"); + } + auto ceil_div = [](int x, int y) { return (x + y - 1) / y; }; + const int rows = launch_params.params.rows; + const int cols = launch_params.params.cols; + int ctas_per_col = launch_params.params.ctas_per_col; + int ctas_per_row = launch_params.params.ctas_per_row; + if (configure_params) { + const int ctas_per_sm = + mgr.occupancy_max_active_blocks_per_sm(main_label, threads_per_cta, 0); + const int max_ctas = launch_params.multiprocessorCount * ctas_per_sm; + ctas_per_row = ceil_div(cols, hidden); + ctas_per_col = std::min(ceil_div(rows, wm), max_ctas / std::max(ctas_per_row, 1)); + launch_params.params.ctas_per_row = ctas_per_row; + launch_params.params.ctas_per_col = ctas_per_col; + if (ctas_per_row > 1) { + launch_params.barrier_bytes = 2 * ctas_per_col * sizeof(int); + launch_params.workspace_bytes = ctas_per_col * wm * ctas_per_row * (ctype_bytes * 2) * 2; + } + launch_params.dgamma_part_bytes = ctas_per_col * cols * ctype_bytes; + return; + } + const auto stream = launch_params.stream; + if (ctas_per_row == 1) { + mgr.launch(main_label, dim3(ctas_per_row * ctas_per_col), dim3(threads_per_cta), 0, stream, + launch_params.params); + } else { + mgr.launch_cooperative(main_label, dim3(ctas_per_row * ctas_per_col), dim3(threads_per_cta), + 0, stream, launch_params.params); + } + const dim3 fin_block(finalize_threads_per_warp * finalize_warps_n, finalize_warps_m); + const dim3 fin_grid(ceil_div(cols, finalize_elts_n_per_cta), 1); + mgr.launch(finalize_label, fin_grid, fin_block, 0, stream, launch_params.params); + }; + TeNormalizationRegistry::registerFunction(key, std::move(closure)); +} + +// ============================================================================ +// RMSNorm Backward (main kernel + finalize kernel; same shape as LayerNorm bwd) +// ============================================================================ + +void register_rmsnorm_bwd_tuned(DType wt, DType it, DType ot, DType ct, int hidden, int cr, int wm, + int wn, int bl_main, int bl_final, bool with_add, + StaticFallback static_fallback) { + const auto stage = with_add ? NVTE_Norm_Stage::BackwardAdd : NVTE_Norm_Stage::Backward; + const auto key = get_key(NVTE_Norm_Backend::Te, NVTE_Norm_Type::RMSNorm, stage, wt, it, ot, ct, 0, + hidden, false, true); + const std::string add_tag = with_add ? "_add" : ""; + const std::string label = concat_strings( + "rmsnorm_bwd_tuned", add_tag, ",w=", static_cast(wt), ",i=", static_cast(it), + ",o=", static_cast(ot), ",c=", static_cast(ct), ",h=", hidden, ",cr=", cr, + ",wm=", wm, ",wn=", wn, ",bl=", bl_main, ",blf=", bl_final); + const std::string main_label = concat_strings(label, ",main"); + const std::string finalize_label = concat_strings(label, ",finalize"); + const std::string traits = kernel_traits_expr(wt, it, ot, ct, hidden, cr, wm, wn, bl_main); + const char* add_flag = with_add ? "true" : "false"; + const std::string main_kexpr = + concat_strings("&::transformer_engine::normalization::rmsnorm_bwd_tuned_kernel<", traits, + ", ", add_flag, ">"); + const std::string finalize_traits = + concat_strings("::transformer_engine::normalization::Kernel_traits_finalize<", hidden, ", ", + cpp_name_for(wt), ", ", cpp_name_for(it), ", ", cpp_name_for(ot), ", ", + cpp_name_for(ct), ", uint32_t, 1024, ", bl_final, ">"); + const std::string finalize_kexpr = + concat_strings("&::transformer_engine::normalization::rmsnorm_bwd_finalize_tuned_kernel<", + finalize_traits, ">"); + const int threads_per_cta = wm * wn * 32; + const int smem_bytes = bwd_smem_bytes(ct, hidden, cr, wm, wn); + const int sizeof_reduce_t = (ct == DType::kFloat32) ? 8 : 4; + const int barrier_per_col = 2 * 4; + const int workspace_per_col = wm * cr * sizeof_reduce_t * 2; + const int dgamma_part_per_col = hidden * byte_size_of(ct); + const int colspass = hidden * byte_size_of(ct) / bl_final; + const int finalize_ctas = colspass / 32; + const int finalize_threads_per_cta = 1024; + + auto closure = [main_label, main_kexpr, finalize_label, finalize_kexpr, threads_per_cta, + smem_bytes, cr, barrier_per_col, workspace_per_col, dgamma_part_per_col, + finalize_ctas, finalize_threads_per_cta, + static_fallback](LaunchParams& launch_params, + const bool configure_params) { + if (try_static_fallback(static_fallback, launch_params, configure_params)) { + return; + } + auto& mgr = rtc::KernelManager::instance(); + if (!mgr.is_compiled(main_label)) { + compile_norm_kernel(mgr, main_label, main_kexpr, + string_code_normalization_rmsnorm_rtc_rmsnorm_bwd_kernel_cu, + "rmsnorm_bwd_kernel.cu"); + } + if (!mgr.is_compiled(finalize_label)) { + compile_norm_kernel(mgr, finalize_label, finalize_kexpr, + string_code_normalization_rmsnorm_rtc_rmsnorm_bwd_kernel_cu, + "rmsnorm_bwd_kernel.cu"); + } + if (configure_params) { + const int ctas_per_sm = + mgr.occupancy_max_active_blocks_per_sm(main_label, threads_per_cta, smem_bytes); + launch_params.params.ctas_per_row = cr; + launch_params.params.ctas_per_col = launch_params.multiprocessorCount * ctas_per_sm / cr; + if (cr > 1) { + launch_params.barrier_bytes = barrier_per_col * launch_params.params.ctas_per_col; + launch_params.workspace_bytes = workspace_per_col * launch_params.params.ctas_per_col; + } + launch_params.dgamma_part_bytes = dgamma_part_per_col * launch_params.params.ctas_per_col; + return; + } + if (smem_bytes >= 48 * 1024) { + mgr.set_function_attribute(main_label, CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES, + smem_bytes); + } + const auto stream = launch_params.stream; + const auto ctas_per_col = launch_params.params.ctas_per_col; + if (cr == 1) { + mgr.launch(main_label, dim3(ctas_per_col), dim3(threads_per_cta), smem_bytes, stream, + launch_params.params); + } else { + mgr.launch_cooperative(main_label, dim3(cr * ctas_per_col), dim3(threads_per_cta), smem_bytes, + stream, launch_params.params); + } + mgr.launch(finalize_label, dim3(finalize_ctas), dim3(finalize_threads_per_cta), 0, stream, + launch_params.params); + }; + TeNormalizationRegistry::registerFunction(key, std::move(closure)); +} + +void register_rmsnorm_bwd_general(DType wt, DType it, DType ot, DType ct, int hidden, int wm, + int wn, int bl_main, int bl_final, bool with_add, + StaticFallback static_fallback) { + const auto stage = with_add ? NVTE_Norm_Stage::BackwardAdd : NVTE_Norm_Stage::Backward; + const auto key = get_key(NVTE_Norm_Backend::Te, NVTE_Norm_Type::RMSNorm, stage, wt, it, ot, ct, 0, + hidden, false, false); + const std::string add_tag = with_add ? "_add" : ""; + const std::string label = concat_strings( + "rmsnorm_bwd_general", add_tag, ",w=", static_cast(wt), ",i=", static_cast(it), + ",o=", static_cast(ot), ",c=", static_cast(ct), ",h=", hidden, ",wm=", wm, + ",wn=", wn, ",bl=", bl_main, ",blf=", bl_final); + const std::string main_label = concat_strings(label, ",main"); + const std::string finalize_label = concat_strings(label, ",finalize"); + const std::string traits = kernel_traits_expr(wt, it, ot, ct, hidden, 1, wm, wn, bl_main); + const char* add_flag = with_add ? "true" : "false"; + const std::string main_kexpr = + concat_strings("&::transformer_engine::normalization::rmsnorm_bwd_general_kernel<", traits, + ", ", add_flag, ">"); + const std::string finalize_kexpr = + concat_strings("&::transformer_engine::normalization::rmsnorm_bwd_finalize_general_kernel<", + cpp_name_for(wt), ", ", cpp_name_for(ct), ", 4, 1, ", bl_final, ", 32>"); + const int threads_per_cta = wm * wn * 32; + const int ctype_bytes = byte_size_of(ct); + const int finalize_warps_m = 4; + const int finalize_warps_n = 1; + const int finalize_threads_per_warp = 32; + const int finalize_elts_n_per_cta = + finalize_threads_per_warp * finalize_warps_n * bl_final / ctype_bytes; + + auto closure = [main_label, main_kexpr, finalize_label, finalize_kexpr, threads_per_cta, hidden, + wm, ctype_bytes, finalize_elts_n_per_cta, finalize_warps_n, finalize_warps_m, + finalize_threads_per_warp, + static_fallback](LaunchParams& launch_params, + const bool configure_params) { + if (try_static_fallback(static_fallback, launch_params, configure_params)) { + return; + } + auto& mgr = rtc::KernelManager::instance(); + if (!mgr.is_compiled(main_label)) { + compile_norm_kernel(mgr, main_label, main_kexpr, + string_code_normalization_rmsnorm_rtc_rmsnorm_bwd_kernel_cu, + "rmsnorm_bwd_kernel.cu"); + } + if (!mgr.is_compiled(finalize_label)) { + compile_norm_kernel(mgr, finalize_label, finalize_kexpr, + string_code_normalization_rmsnorm_rtc_rmsnorm_bwd_kernel_cu, + "rmsnorm_bwd_kernel.cu"); + } + auto ceil_div = [](int x, int y) { return (x + y - 1) / y; }; + const int rows = launch_params.params.rows; + const int cols = launch_params.params.cols; + int ctas_per_col = launch_params.params.ctas_per_col; + int ctas_per_row = launch_params.params.ctas_per_row; + if (configure_params) { + const int ctas_per_sm = + mgr.occupancy_max_active_blocks_per_sm(main_label, threads_per_cta, 0); + const int max_ctas = launch_params.multiprocessorCount * ctas_per_sm; + ctas_per_row = ceil_div(cols, hidden); + ctas_per_col = std::min(ceil_div(rows, wm), max_ctas / std::max(ctas_per_row, 1)); + launch_params.params.ctas_per_row = ctas_per_row; + launch_params.params.ctas_per_col = ctas_per_col; + if (ctas_per_row > 1) { + launch_params.barrier_bytes = 2 * ctas_per_col * sizeof(int); + launch_params.workspace_bytes = ctas_per_col * wm * ctas_per_row * (ctype_bytes * 2) * 2; + } + launch_params.dgamma_part_bytes = ctas_per_col * cols * ctype_bytes; + return; + } + const auto stream = launch_params.stream; + if (ctas_per_row == 1) { + mgr.launch(main_label, dim3(ctas_per_row * ctas_per_col), dim3(threads_per_cta), 0, stream, + launch_params.params); + } else { + mgr.launch_cooperative(main_label, dim3(ctas_per_row * ctas_per_col), dim3(threads_per_cta), + 0, stream, launch_params.params); + } + const dim3 fin_block(finalize_threads_per_warp * finalize_warps_n, finalize_warps_m); + const dim3 fin_grid(ceil_div(cols, finalize_elts_n_per_cta), 1); + mgr.launch(finalize_label, fin_grid, fin_block, 0, stream, launch_params.params); + }; + TeNormalizationRegistry::registerFunction(key, std::move(closure)); +} + +} // namespace rtc_norm +} // namespace normalization +} // namespace transformer_engine diff --git a/transformer_engine/common/normalization/rtc_dispatch.h b/transformer_engine/common/normalization/rtc_dispatch.h new file mode 100644 index 0000000000..314f150f70 --- /dev/null +++ b/transformer_engine/common/normalization/rtc_dispatch.h @@ -0,0 +1,65 @@ +/************************************************************************* + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * + * See LICENSE for license information. + ************************************************************************/ + +#ifndef TRANSFORMER_ENGINE_COMMON_NORM_RTC_DISPATCH_H_ +#define TRANSFORMER_ENGINE_COMMON_NORM_RTC_DISPATCH_H_ + +#include + +#include "common.h" + +namespace transformer_engine { +namespace normalization { +namespace rtc_norm { + +template +using StaticFallback = void (*)(LaunchParams&, const bool); + +// Register an RTC-backed launcher for a single LayerNorm Forward "tuned" +// (multi-CTA-capable) config. Compiles and launches via NVRTC on first use. +void register_ln_fwd_tuned(DType wtype, DType itype, DType otype, DType ctype, int hidden_size, + int ctas_per_row, int warps_m, int warps_n, int bytes_per_ldg, + StaticFallback static_fallback = nullptr); + +// Register an RTC-backed launcher for a single LayerNorm Forward "general" +// (no multi-CTA) config. +void register_ln_fwd_general(DType wtype, DType itype, DType otype, DType ctype, int hidden_size, + int warps_m, int warps_n, int bytes_per_ldg, + StaticFallback static_fallback = nullptr); + +// Register an RTC-backed launcher for a single LayerNorm Backward "tuned" config. +void register_ln_bwd_tuned(DType wtype, DType itype, DType otype, DType ctype, int hidden_size, + int ctas_per_row, int warps_m, int warps_n, int bytes_per_ldg_main, + int bytes_per_ldg_final, + StaticFallback static_fallback = nullptr); + +// Register an RTC-backed launcher for a single LayerNorm Backward "general" config. +void register_ln_bwd_general(DType wtype, DType itype, DType otype, DType ctype, int hidden_size, + int warps_m, int warps_n, int bytes_per_ldg_main, + int bytes_per_ldg_final, + StaticFallback static_fallback = nullptr); + +// Same set for RMSNorm. +void register_rmsnorm_fwd_tuned(DType wtype, DType itype, DType otype, DType ctype, int hidden_size, + int ctas_per_row, int warps_m, int warps_n, int bytes_per_ldg, + StaticFallback static_fallback = nullptr); +void register_rmsnorm_fwd_general(DType wtype, DType itype, DType otype, DType ctype, + int hidden_size, int warps_m, int warps_n, int bytes_per_ldg, + StaticFallback static_fallback = nullptr); +void register_rmsnorm_bwd_tuned(DType wtype, DType itype, DType otype, DType ctype, int hidden_size, + int ctas_per_row, int warps_m, int warps_n, int bytes_per_ldg_main, + int bytes_per_ldg_final, bool with_add, + StaticFallback static_fallback = nullptr); +void register_rmsnorm_bwd_general(DType wtype, DType itype, DType otype, DType ctype, + int hidden_size, int warps_m, int warps_n, int bytes_per_ldg_main, + int bytes_per_ldg_final, bool with_add, + StaticFallback static_fallback = nullptr); + +} // namespace rtc_norm +} // namespace normalization +} // namespace transformer_engine + +#endif // TRANSFORMER_ENGINE_COMMON_NORM_RTC_DISPATCH_H_ diff --git a/transformer_engine/common/recipe/__init__.py b/transformer_engine/common/recipe/__init__.py index 8a03f2f51a..8c209ace15 100644 --- a/transformer_engine/common/recipe/__init__.py +++ b/transformer_engine/common/recipe/__init__.py @@ -404,6 +404,12 @@ class Float8BlockScaling(Recipe): NOTE: To relax the default constraint that scales be powers of 2, set env variable NVTE_FP8_BLOCK_SCALING_FP32_SCALES=1 to override it for the recipe defaults. + NOTE: FP8 block scaling requires split accumulation for numerical accuracy, so + ``fp8_gemm_fprop``/``fp8_gemm_dgrad``/``fp8_gemm_wgrad`` all fix + ``use_split_accumulator=True`` (enforced in ``__post_init__``). The fused grouped + GEMM path (GroupedLinear) always uses split accumulation for FP8 block scaling and + ignores any caller- or recipe-supplied ``use_split_accumulator`` value. + Parameters ---------- fp8_format : {Format.E4M3, Format.HYBRID}, default = Format.E4M3 diff --git a/transformer_engine/common/util/rtc.cpp b/transformer_engine/common/util/rtc.cpp index 70024a202c..20616f8cb6 100644 --- a/transformer_engine/common/util/rtc.cpp +++ b/transformer_engine/common/util/rtc.cpp @@ -47,12 +47,7 @@ inline int max_supported_sm_arch() { } // namespace bool is_enabled() { - static bool is_enabled_ = false; - static bool need_to_check_env = true; - if (need_to_check_env) { - is_enabled_ = !getenv("NVTE_DISABLE_NVRTC"); - need_to_check_env = false; - } + static const bool is_enabled_ = !getenv("NVTE_DISABLE_NVRTC"); return is_enabled_; } @@ -132,6 +127,18 @@ void Kernel::set_function_cache_config(int device_id, CUfunc_cache cache_config) NVTE_CALL_CHECK_CUDA_DRIVER(cuFuncSetCacheConfig, get_function(device_id), cache_config); } +void Kernel::set_function_attribute(int device_id, CUfunction_attribute attr, int value) { + NVTE_CALL_CHECK_CUDA_DRIVER(cuFuncSetAttribute, get_function(device_id), attr, value); +} + +int Kernel::occupancy_max_active_blocks_per_sm(int device_id, int block_size, + std::size_t dynamic_smem_bytes) { + int num_blocks = 0; + NVTE_CALL_CHECK_CUDA_DRIVER(cuOccupancyMaxActiveBlocksPerMultiprocessor, &num_blocks, + get_function(device_id), block_size, dynamic_smem_bytes); + return num_blocks; +} + KernelManager& KernelManager::instance() { NVTE_CHECK(is_enabled(), "NVRTC support is not enabled"); static KernelManager instance_; @@ -139,11 +146,17 @@ KernelManager& KernelManager::instance() { } void KernelManager::compile(const std::string& kernel_label, const std::string& kernel_name, - const std::string& code, const std::string& filename) { - std::lock_guard lock_guard_(lock_); + const std::string& code, const std::string& filename, + const std::vector& extra_options, + const std::vector
& extra_headers) { + const int device_id = cuda::current_device(); + const auto key = get_kernel_cache_key(kernel_label, device_id); + std::unique_lock lock_guard_(lock_); + if (kernel_cache_.count(key) > 0) { + return; + } // Choose whether to compile to PTX or cubin - const int device_id = cuda::current_device(); const int sm_arch_ = cuda::sm_arch(device_id); const int compile_sm_arch = std::min(sm_arch_, max_supported_sm_arch()); const bool compile_ptx = sm_arch_ != compile_sm_arch; @@ -160,6 +173,7 @@ void KernelManager::compile(const std::string& kernel_label, const std::string& opts.push_back(concat_strings("--gpu-architecture=sm_", compile_sm_arch)); } opts.push_back(concat_strings("-I", cuda::include_directory(true))); + opts.insert(opts.end(), extra_options.begin(), extra_options.end()); std::vector opts_ptrs; for (const auto& opt : opts) { opts_ptrs.push_back(opt.c_str()); @@ -167,11 +181,19 @@ void KernelManager::compile(const std::string& kernel_label, const std::string& // Compile source nvrtcProgram program; - constexpr int num_headers = 2; - constexpr const char* headers[num_headers] = {string_code_utils_cuh, string_code_util_math_h}; - constexpr const char* include_names[num_headers] = {"utils.cuh", "util/math.h"}; - NVTE_CHECK_NVRTC(nvrtcCreateProgram(&program, code.c_str(), filename.c_str(), num_headers, - headers, include_names)); + std::vector headers = {string_code_utils_cuh, string_code_util_math_h}; + std::vector include_names = {"utils.cuh", "util/math.h"}; + headers.reserve(headers.size() + extra_headers.size()); + include_names.reserve(include_names.size() + extra_headers.size()); + for (const auto& header : extra_headers) { + NVTE_CHECK(header.content != nullptr && header.include_name != nullptr, + "NVRTC header content and include name must not be null"); + headers.push_back(header.content); + include_names.push_back(header.include_name); + } + NVTE_CHECK_NVRTC(nvrtcCreateProgram(&program, code.c_str(), filename.c_str(), + static_cast(headers.size()), headers.data(), + include_names.data())); NVTE_CHECK_NVRTC(nvrtcAddNameExpression(program, kernel_name.c_str())); const nvrtcResult compile_result = nvrtcCompileProgram(program, opts_ptrs.size(), opts_ptrs.data()); @@ -239,7 +261,6 @@ void KernelManager::compile(const std::string& kernel_label, const std::string& } // Cache compiled code - const auto key = get_kernel_cache_key(kernel_label, device_id); kernel_cache_.insert({key, Kernel(mangled_name, std::move(compiled_code))}); kernel_cache_.at(key).get_function(device_id); // Make sure kernel is available on device @@ -250,12 +271,34 @@ void KernelManager::compile(const std::string& kernel_label, const std::string& void KernelManager::set_cache_config(const std::string& kernel_label, CUfunc_cache cache_config) { const int device_id = cuda::current_device(); const auto key = get_kernel_cache_key(kernel_label, device_id); + std::shared_lock lock_guard_(lock_); NVTE_CHECK(kernel_cache_.count(key) > 0, "Attempted to configure RTC kernel before compilation"); kernel_cache_.at(key).set_function_cache_config(device_id, cache_config); } +void KernelManager::set_function_attribute(const std::string& kernel_label, + CUfunction_attribute attr, int value) { + const int device_id = cuda::current_device(); + const auto key = get_kernel_cache_key(kernel_label, device_id); + std::shared_lock lock_guard_(lock_); + NVTE_CHECK(kernel_cache_.count(key) > 0, "Attempted to configure RTC kernel before compilation"); + kernel_cache_.at(key).set_function_attribute(device_id, attr, value); +} + +int KernelManager::occupancy_max_active_blocks_per_sm(const std::string& kernel_label, + int block_size, + std::size_t dynamic_smem_bytes) { + const int device_id = cuda::current_device(); + const auto key = get_kernel_cache_key(kernel_label, device_id); + std::shared_lock lock_guard_(lock_); + NVTE_CHECK(kernel_cache_.count(key) > 0, "Attempted to query occupancy before compilation"); + return kernel_cache_.at(key).occupancy_max_active_blocks_per_sm(device_id, block_size, + dynamic_smem_bytes); +} + bool KernelManager::is_compiled(const std::string& kernel_label, int device_id) const { const auto key = get_kernel_cache_key(kernel_label, device_id); + std::shared_lock lock_guard_(lock_); return kernel_cache_.count(key) > 0; } diff --git a/transformer_engine/common/util/rtc.h b/transformer_engine/common/util/rtc.h index 65faf7bcc2..745566ee60 100644 --- a/transformer_engine/common/util/rtc.h +++ b/transformer_engine/common/util/rtc.h @@ -13,6 +13,7 @@ #include #include +#include #include #include #include @@ -33,6 +34,12 @@ namespace rtc { */ bool is_enabled(); +/*! \brief Header made available to an NVRTC program */ +struct Header { + const char *content; + const char *include_name; +}; + /*! \brief Wrapper class for a runtime-compiled CUDA kernel */ class Kernel { public: @@ -79,6 +86,29 @@ class Kernel { */ void set_function_cache_config(int device_id, CUfunc_cache cache_config); + /*! \brief Set a kernel function attribute (driver-API wrapper of + * cuFuncSetAttribute, e.g. for dynamic shared memory size). + */ + void set_function_attribute(int device_id, CUfunction_attribute attr, int value); + + /*! \brief Wrapper of cuOccupancyMaxActiveBlocksPerMultiprocessor for a + * runtime-compiled function. + */ + int occupancy_max_active_blocks_per_sm(int device_id, int block_size, + std::size_t dynamic_smem_bytes); + + /*! \brief Cooperative launch of an RTC kernel via cuLaunchCooperativeKernel. + */ + template + void launch_cooperative(int device_id, const dim3 grid_dim, const dim3 block_dim, + unsigned int shared_mem_bytes, cudaStream_t stream, ArgTs &&...args) { + cuda_driver::ensure_context_exists(); + void *arg_ptrs[] = {const_cast(static_cast(&args))...}; + NVTE_CALL_CHECK_CUDA_DRIVER(cuLaunchCooperativeKernel, get_function(device_id), grid_dim.x, + grid_dim.y, grid_dim.z, block_dim.x, block_dim.y, block_dim.z, + shared_mem_bytes, static_cast(stream), arg_ptrs); + } + private: /*! \brief Mangled function name */ std::string mangled_name_; @@ -113,9 +143,13 @@ class KernelManager { * \param[in] code Kernel source code * \param[in] filename Path to associate with source code, * primarily for debugging + * \param[in] extra_options Additional NVRTC compiler options + * \param[in] extra_headers Additional in-memory headers available to the program */ void compile(const std::string &kernel_label, const std::string &kernel_name, - const std::string &code, const std::string &filename); + const std::string &code, const std::string &filename, + const std::vector &extra_options = {}, + const std::vector
&extra_headers = {}); /*! \brief Whether CUDA kernel has been compiled for CUDA device * @@ -143,6 +177,7 @@ class KernelManager { unsigned int shared_mem_bytes, cudaStream_t stream, ArgTs &&...args) { const int device_id = cuda::current_device(); const auto key = get_kernel_cache_key(kernel_label, device_id); + std::shared_lock lock_guard_(lock_); NVTE_CHECK(kernel_cache_.count(key) > 0, "Attempted to launch RTC kernel before compilation"); kernel_cache_.at(key).launch(device_id, grid_dim, block_dim, shared_mem_bytes, stream, std::forward(args)...); @@ -157,11 +192,32 @@ class KernelManager { */ void set_cache_config(const std::string &kernel_label, CUfunc_cache cache_config); + /*! \brief Set a function attribute (e.g. cuFuncAttributeMaxDynamicSharedMemorySize). */ + void set_function_attribute(const std::string &kernel_label, CUfunction_attribute attr, + int value); + + /*! \brief Query cuOccupancyMaxActiveBlocksPerMultiprocessor for a compiled kernel. */ + int occupancy_max_active_blocks_per_sm(const std::string &kernel_label, int block_size, + std::size_t dynamic_smem_bytes); + + /*! \brief Cooperative launch wrapper (cuLaunchCooperativeKernel). */ + template + void launch_cooperative(const std::string &kernel_label, const dim3 grid_dim, + const dim3 block_dim, unsigned int shared_mem_bytes, cudaStream_t stream, + ArgTs &&...args) { + const int device_id = cuda::current_device(); + const auto key = get_kernel_cache_key(kernel_label, device_id); + std::shared_lock lock_guard_(lock_); + NVTE_CHECK(kernel_cache_.count(key) > 0, "Attempted to launch RTC kernel before compilation"); + kernel_cache_.at(key).launch_cooperative(device_id, grid_dim, block_dim, shared_mem_bytes, + stream, std::forward(args)...); + } + private: /*! \brief Compiled kernels */ std::unordered_map kernel_cache_; - /*! \brief Mutex for thread-safe compilation */ - std::mutex lock_; + /*! \brief Mutex for thread-safe cache access */ + mutable std::shared_mutex lock_; KernelManager() = default; ~KernelManager() = default; diff --git a/transformer_engine/common/utils.cuh b/transformer_engine/common/utils.cuh index 635c7a36d2..0b75da622c 100644 --- a/transformer_engine/common/utils.cuh +++ b/transformer_engine/common/utils.cuh @@ -30,6 +30,37 @@ static_assert(sizeof(uint32_t) == 4); static_assert(sizeof(uint64_t) == 8); #endif +// Minimal subset of used by RTC kernel headers. Keep these in a +// project-owned namespace because adding primary templates to std is undefined. +namespace transformer_engine { +namespace detail { + +template +struct is_same { + static constexpr bool value = false; +}; +template +struct is_same { + static constexpr bool value = true; +}; + +template +inline constexpr bool is_same_v = is_same::value; + +template +struct conditional { + using type = T; +}; +template +struct conditional { + using type = F; +}; +template +using conditional_t = typename conditional::type; + +} // namespace detail +} // namespace transformer_engine + //////////////////////////////////////////////////////////////////////////////////////////////////// constexpr uint32_t THREADS_PER_WARP = 32; @@ -940,13 +971,13 @@ __device__ __forceinline__ float ordered_uint_to_float(unsigned int u) { template __device__ __forceinline__ T abs_val(T val) { - if constexpr (std::is_same_v) { + if constexpr (detail::is_same_v) { #if __CUDA_ARCH__ >= 800 return __habs(val); #else return static_cast<__nv_bfloat16>(fabsf(static_cast(val))); #endif - } else if constexpr (std::is_same_v) { + } else if constexpr (detail::is_same_v) { return __habs(val); } else { return fabsf(val); @@ -955,13 +986,13 @@ __device__ __forceinline__ T abs_val(T val) { template __device__ __forceinline__ T max_val(T a, T b) { - if constexpr (std::is_same_v) { + if constexpr (detail::is_same_v) { #if __CUDA_ARCH__ >= 800 return __hmax(a, b); #else return static_cast<__nv_bfloat16>(fmaxf(static_cast(a), static_cast(b))); #endif - } else if constexpr (std::is_same_v) { + } else if constexpr (detail::is_same_v) { return __hmax(a, b); } else { return fmaxf(a, b); diff --git a/transformer_engine/jax/cpp_extensions/ep.py b/transformer_engine/jax/cpp_extensions/ep.py index 34ec0a33b6..806e7ae480 100644 --- a/transformer_engine/jax/cpp_extensions/ep.py +++ b/transformer_engine/jax/cpp_extensions/ep.py @@ -23,7 +23,19 @@ import transformer_engine_jax from .base import BasePrimitive, register_primitive -from ..sharding import global_mesh_resource +from ..sharding import global_mesh_resource, get_mesh_axis_size +from ..version_utils import is_collective_stream_supported + + +def _on_collective_stream(func): + """Pin ``func``'s ops to XLA's collective stream so the scheduler serializes + them with native collectives. No-op on JAX that lacks the annotation.""" + if not is_collective_stream_supported(): + return func + from jax.experimental.compute_on import compute_on + + return compute_on("gpu_stream:collective")(func) # pylint: disable=not-callable + __all__ = [ "EpConfig", @@ -125,8 +137,15 @@ def _ep_outer_axis(): When set, EP-output globals carry an extra leading ``dp_size`` dim so SPMD sees each DP color's slab as distinct (rather than replicated across DP). + + A dp/fsdp axis that is sized 1 in the active mesh is treated as absent so + we don't pin EP-output specs to a degenerate axis that JAX may collapse. """ gsr = global_mesh_resource() + if gsr.dp_resource is not None and get_mesh_axis_size(gsr.dp_resource) > 1: + return gsr.dp_resource + if gsr.fsdp_resource is not None and get_mesh_axis_size(gsr.fsdp_resource) > 1: + return gsr.fsdp_resource return gsr.dp_resource or gsr.fsdp_resource @@ -887,6 +906,7 @@ def shardy_sharding_rule(*args): # ── Public-ish helpers (used by jax/ep.py) ────────────────────────────────── +@_on_collective_stream def ep_prepare(cfg: EpLayerConfig, topk_idx): """Exchange routing metadata for ``cfg``; return ``(token_counts, handle_mem)``.""" return EpPreparePrimitive.outer_primitive.bind( @@ -897,6 +917,7 @@ def ep_prepare(cfg: EpLayerConfig, topk_idx): ) +@_on_collective_stream def ep_dispatch_fwd( cfg: EpLayerConfig, handle_mem, topk_idx, tokens, topk_weights, recv_capacity_per_rank ): @@ -913,6 +934,7 @@ def ep_dispatch_fwd( ) +@_on_collective_stream def ep_combine_fwd( cfg: EpLayerConfig, handle_mem, expert_out, num_local_tokens, out_partition_spec=None ): @@ -928,6 +950,7 @@ def ep_combine_fwd( ) +@_on_collective_stream def ep_dispatch_bwd( cfg: EpLayerConfig, handle_mem, @@ -949,6 +972,7 @@ def ep_dispatch_bwd( ) +@_on_collective_stream def ep_combine_bwd(cfg: EpLayerConfig, handle_mem, grad, recv_capacity_per_rank): """Backward of combine; returns grad_expert_out [num_procs, recv_capacity_per_rank, H].""" return EpCombineBwdPrimitive.outer_primitive.bind( diff --git a/transformer_engine/jax/cpp_extensions/router.py b/transformer_engine/jax/cpp_extensions/router.py index 8cc94fcaaf..46f51c9d33 100644 --- a/transformer_engine/jax/cpp_extensions/router.py +++ b/transformer_engine/jax/cpp_extensions/router.py @@ -412,6 +412,11 @@ def partition( arg_infos, result_infos, ): + # NOTE: do NOT include ``routing_map_format`` in this ``del``: the + # ``sharded_impl`` closure below resolves it by name at call time + # (when XLA invokes the partitioned impl), so deleting it here + # raises ``NameError: cannot access free variable 'routing_map_format'`` + # at execution time of the bwd custom_partitioning. del result_infos grad_spec = get_padded_spec(arg_infos[2]) out_sharding = NamedSharding(mesh, PartitionSpec(*grad_spec)) @@ -645,7 +650,14 @@ def shardy_sharding_rule(*args): # backward reconstructs the full [num_tokens, num_experts] grad_probs from # scalar inputs. Shardy will leave num_tokens unsharded, which matches the # replicated PartitionSpec(None, None) in partition(). - return "const_buf_one, num_experts, grad_one -> i num_experts" + # + # grad_aux_loss is the cotangent of a scalar loss and is therefore + # rank-0; the third operand entry is empty (no factor labels). Declaring + # it with the spurious "grad_one" factor gave it rank-1 and tripped + # JAX's custom_partitioning_sharding_rule check once the MoE block + # lifted its aux-loss path out of shard_map (the rule is skipped under + # shard_map, which is why this surfaces only at global view). + return "const_buf_one, num_experts, -> i num_experts" register_primitive(FusedMoEAuxLossBwdPrimitive) diff --git a/transformer_engine/jax/flax/moe.py b/transformer_engine/jax/flax/moe.py index 91346a7a48..3629346e33 100644 --- a/transformer_engine/jax/flax/moe.py +++ b/transformer_engine/jax/flax/moe.py @@ -37,8 +37,7 @@ # import P`` without a second jax.sharding import. from jax.sharding import PartitionSpec as P # noqa: F401 # pylint: disable=unused-import -from ..moe import PermutationBackend, moe -from ..quantize import noop_quantizer_set +from ..moe import moe from ..router import ScoreFunction from ..sharding import get_active_resource_axis from .module import TransformerEngineBase @@ -50,7 +49,7 @@ Initializer = Callable[[PRNGKey, Shape, DType], Array] -__all__ = ["PermutationBackend", "_MoEBlock"] +__all__ = ["_MoEBlock"] class _MoEBlock(TransformerEngineBase): @@ -82,10 +81,11 @@ class _MoEBlock(TransformerEngineBase): Grouped top-k knobs (DeepSeek-style). ``None`` disables grouping. scaling_factor : float Multiplier on the routing weights. - use_expert_bias : bool - If ``True``, registers a per-expert routing bias (shape ``[E]``). - Only meaningful with ``score_function="sigmoid"``; the underlying - primitive validates the pairing. + use_expert_routing_bias : bool + If ``True``, registers a per-expert routing bias (shape ``[E]``) + used by the topk selection. Only meaningful with + ``score_function="sigmoid"``; the underlying primitive validates + the pairing. aux_loss_coeff : float If ``> 0``, return the MoE auxiliary load-balancing loss scalar in addition to the main output. @@ -100,23 +100,27 @@ class _MoEBlock(TransformerEngineBase): replicated across non-EP axes within an EP group; set e.g. ``("fsdp",)`` for true FSDP-of-batch where each device owns a unique slice of the batch. - permutation_backend : PermutationBackend - ``PURE_JAX`` (default) or ``TRITON``. - _align_size : int - Per-expert group-size alignment (``0`` disables; required > 0 - for quantized grouped GEMM). Internal knob; will be inferred - from the active quantization recipe in a follow-up PR. + apply_topk_weights_early : bool + If ``True``, multiply expert outputs by their top-k weights + *inside* each shard before ``ep_combine`` (saves one global + reduction at the cost of an extra broadcast). Default ``False``. + + The per-expert dispatch-slot alignment is fixed internally at 128 + tokens (see ``moe._ALIGN_SIZE``) -- the value required by NCCL EP + HT and satisfied by every current TE grouped-GEMM recipe -- and is + therefore not exposed as a per-instance knob. dtype : jnp.dtype Compute / parameter dtype. kernel_init, bias_init, expert_bias_init : Initializers. - use_bias : bool - Register per-expert FFN biases. + use_ffn_bias : bool + Register per-expert FFN biases (``wi_0_bias``, ``wi_1_bias``, + ``wo_bias``). Quantization is currently configured via the standard TE autocast - context (``fp8_autocast``/``with_quantizer_set``); per-call - quantizer sets can also be passed through ``__call__``'s - ``quantizer_sets`` keyword once we stabilise the recipe pipeline. + context (``fp8_autocast``/``with_quantizer_set``) and threaded + through ``moe()`` internally; this wrapper does not expose a + per-call ``quantizer_sets`` knob yet. """ # Architecture @@ -131,7 +135,7 @@ class _MoEBlock(TransformerEngineBase): num_groups: Optional[int] = None group_topk: Optional[int] = None scaling_factor: float = 1.0 - use_expert_bias: bool = False + use_expert_routing_bias: bool = False aux_loss_coeff: float = 0.0 # Sharding (logical axes) @@ -143,16 +147,15 @@ class _MoEBlock(TransformerEngineBase): # Parallelism data_parallelism_axes: Tuple[str, ...] = () - # Permutation - permutation_backend: PermutationBackend = PermutationBackend.PURE_JAX - _align_size: int = 0 + # MoE knobs forwarded to ``moe()`` + apply_topk_weights_early: bool = False # Dtypes / init / misc dtype: DType = jnp.float32 kernel_init: Optional[Initializer] = None bias_init: Initializer = nn.initializers.zeros expert_bias_init: Initializer = nn.initializers.zeros - use_bias: bool = False + use_ffn_bias: bool = False def __post_init__(self): if self.kernel_init is None: @@ -163,11 +166,6 @@ def __post_init__(self): 1.0, "fan_in", "truncated_normal", dtype=self.dtype ), ) - if not isinstance(self.permutation_backend, PermutationBackend): - raise TypeError( - "permutation_backend must be a PermutationBackend, got" - f" {self.permutation_backend!r}" - ) super().__post_init__() @nn.compact @@ -221,7 +219,7 @@ def __call__(self, inputs: Array) -> Tuple[Array, Optional[Array]]: self.dtype, ) wi_0_bias = wi_1_bias = wo_bias = None - if self.use_bias: + if self.use_ffn_bias: wi_0_bias = self.param( "wi_0_bias", nn.with_logical_partitioning(self.bias_init, ("exp", "mlp")), @@ -241,12 +239,14 @@ def __call__(self, inputs: Array) -> Tuple[Array, Optional[Array]]: self.dtype, ) expert_bias = None - if self.use_expert_bias: + if self.use_expert_routing_bias: + # The router logits are promoted to fp32 before fused top-k; keep + # the routing bias in the same dtype so it only affects selection. expert_bias = self.param( "expert_bias", nn.with_logical_partitioning(self.expert_bias_init, ("exp",)), (self.num_experts,), - self.dtype, + jnp.float32, ) ep_axis = get_active_resource_axis("ep_resource") @@ -270,15 +270,12 @@ def __call__(self, inputs: Array) -> Tuple[Array, Optional[Array]]: group_topk=self.group_topk, scaling_factor=self.scaling_factor, aux_loss_coeff=self.aux_loss_coeff, - permutation_backend=self.permutation_backend, - align_size=self._align_size, - gate_inside_vjp=True, + apply_topk_weights_early=self.apply_topk_weights_early, ep_axis=ep_axis, data_parallelism_axes=self.data_parallelism_axes, input_axes=self.input_axes, gate_kernel_axes=self.gate_kernel_axes, wi_kernel_axes=self.wi_kernel_axes, wo_kernel_axes=self.wo_kernel_axes, - quantizer_sets=(noop_quantizer_set, noop_quantizer_set, noop_quantizer_set), dtype=self.dtype, ) diff --git a/transformer_engine/jax/moe.py b/transformer_engine/jax/moe.py index 2a1c818cb3..887e005de6 100644 --- a/transformer_engine/jax/moe.py +++ b/transformer_engine/jax/moe.py @@ -1,76 +1,51 @@ # Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. -"""Functional Mixture-of-Experts (MoE) entry point with a single fused VJP. - -This module exposes :func:`moe`, the framework-agnostic flat function that -implements an entire MoE block (gate -> top-k routing -> token dispatch -> -per-expert FFN -> token combine, plus optional expert parallelism via a -shard_map / ragged_all_to_all collective) under a *single* -``jax.custom_vjp``. It is the moral analog of -:func:`transformer_engine.jax.layernorm_mlp.layernorm_mlp` for MoE: one -custom_vjp boundary covers the whole block so future fusions (FP8 over the -EP wire, fused ``ragged_all_to_all + grouped_gemm``, gate+route+dispatch -fusion) can land without re-architecting the call site. - -Design rationale ----------------- - -The earlier MoE block (:class:`transformer_engine.jax.flax.moe._MoEBlock`) -composed many narrower custom_vjps -- one per :func:`grouped_dense`, one -per :func:`token_dispatch`, etc. Every nested custom_vjp is a place where -a quantized :class:`ScaledTensor` cannot survive (JAX requires custom_vjp -inputs / outputs to be plain ``jnp.ndarray`` ish pytrees). To enable -end-to-end FP8 flow -- in particular FP8 carried over the EP -ragged_all_to_all -- the dispatch's quantize, the a2a, the per-expert -FFN, the inverse a2a, and the combine all have to live inside the same -VJP. This file collapses them into one. - -Implementation conventions --------------------------- - -* No nested ``custom_vjp``. Every primitive's ``_fwd`` and ``_bwd`` is - called directly (e.g. :func:`tex.fused_topk_with_score_function_fwd` / - ``_bwd``, :func:`unpermute_with_mask_map`, - :func:`unpermute_bwd_with_merging_probs`, - :func:`sort_chunks_by_map(is_forward=False)`, - forward + reverse :func:`jax.lax.ragged_all_to_all`) so the outer - ``_moe_bwd_rule`` controls the bwd graph end-to-end without invoking - ``jax.vjp`` for re-linearization. -* The fwd/bwd context (``ctx``) is a plain ``dict`` whose keys depend on - the static configuration (permutation backend, EP active or not, - presence of biases, aux loss enabled). The ``_moe_fwd_rule`` builds a - matching ``ctx_specs`` dict in lockstep when opening the EP shard_map - so ``out_specs`` structurally matches the body's return. -* :func:`_dispatch` is the helper that wraps - ``permute -> a2a -> local_permute`` (forward); :func:`_combine` is its - inverse. Their ``_bwd`` siblings drive the inverse collectives in the - bwd rule. None of these helpers form a custom_vjp boundary. +"""Mixture-of-Experts (MoE) layer for TransformerEngine JAX. + +This module exposes :func:`moe`, a single fused MoE forward pass + bwd +built on top of TE's NCCL-backed Expert Parallelism primitives +(``tex.ep_dispatch`` / ``tex.ep_combine``). The block runs:: + + gate -> topk -> ep_dispatch -> per-expert FFN (grouped GEMMs) + -> ep_combine -> output + +under a single ``jax.custom_vjp`` so the routing, dispatch, FFN and +combine steps fuse cleanly under XLA without leaking intermediate +residuals into the user-facing autograd graph. + +Sharding model +-------------- +* Inbound activations are 3D ``[B, S, H]`` sharded + ``((*data_parallelism_axes, ep_axis), None, None)``. The public + :func:`moe` soft-repins this on entry and warns when a reshard is + inserted. +* The EP primitives operate at global view (their custom_partitioning + rules handle per-shard execution). The FFN GEMMs run per-shard inside + a small ``shard_map`` whose ``in_specs`` and ``out_specs`` mirror the + same ``((dp, ep), ...)`` layout. + +Out-of-scope (for now) +---------------------- +FP8 / MXFP8 quantizer sets are not yet wired on this path; turning +them on requires recipe-aware residual specs and ``ScaledTensor`` +leaves across the ``shard_map`` boundary. ``aux_loss_coeff`` and +``expert_bias`` are supported (the former forces a per-step +all-gather over the routing-side logits, which lives off the critical +path and overlaps with the dispatch collective). """ -import math -from dataclasses import dataclass -from enum import Enum from functools import partial -from typing import Any, NewType, Optional, Tuple, Union +from typing import Any, Optional, Tuple, Union +import warnings +import flax.struct import jax import jax.numpy as jnp -from flax import struct as flax_struct -from jax.sharding import PartitionSpec as P +from jax.sharding import NamedSharding, PartitionSpec as P from . import cpp_extensions as tex -from .permutation import ( - PureJaxPermState, - compute_ragged_all_to_all_params, - compute_reverse_ragged_all_to_all_params, - pure_jax_token_combine, - pure_jax_token_dispatch, - routing_map_to_selected_experts, -) from .quantize import ( - QuantizerSet, - ScaledTensor, TensorUsage, noop_quantizer_set, with_sharding_constraint_by_logical_axes, @@ -79,1070 +54,251 @@ from .router import ScoreFunction, _validate_score_function from .sharding import _get_mesh -# Triton-backed primitives are imported lazily: callers on the PURE_JAX -# permutation backend should not need ``triton`` installed. The TRITON -# branches in this module call ``_require_triton()`` first to raise a -# clear error if the import failed. -try: - from .triton_extensions.permutation import ( - make_chunk_sort_map, - make_row_id_map, - permute_with_mask_map, - permute_with_mask_map_and_pad, - sort_chunks_by_map, - unpermute_bwd_with_merging_probs, - unpermute_bwd_with_merging_probs_and_unpad, - unpermute_with_mask_map, - unpermute_with_mask_map_and_unpad, - ) - - _TRITON_AVAILABLE = True -except ImportError: - _TRITON_AVAILABLE = False - make_chunk_sort_map = None - make_row_id_map = None - permute_with_mask_map = None - permute_with_mask_map_and_pad = None - sort_chunks_by_map = None - unpermute_bwd_with_merging_probs = None - unpermute_bwd_with_merging_probs_and_unpad = None - unpermute_with_mask_map = None - unpermute_with_mask_map_and_unpad = None - - -def _require_triton(): - """Raise a clear error if Triton permutation kernels are unavailable.""" - if not _TRITON_AVAILABLE: - raise ImportError( - "PermutationBackend.TRITON requires" - " ``transformer_engine.jax.triton_extensions`` (and ``triton``)." - " Install Triton or pass PermutationBackend.PURE_JAX." - ) - - -PRNGKey = Any -Shape = Tuple[int, ...] -DType = NewType("DType", jnp.dtype) -Array = NewType("Array", jnp.ndarray) - +__all__ = ["moe"] -__all__ = ["moe", "PermutationBackend"] +# Per-expert dispatch-slot alignment fed to ``tex.ep_prepare`` as +# ``dispatch_output_per_expert_alignment``. NCCL EP HT requires the +# per-expert recv block to be at least 128-token aligned, and all current +# TE grouped-GEMM recipes (bf16/fp16/fp8/mxfp8) are satisfied by the +# same 128-token tile, so a single constant covers every supported path. +_ALIGN_SIZE = 128 -# ============================================================================= -# Enums -# ============================================================================= +def _with_sharding_constraint_cast_bwd(x: jnp.ndarray, sharding) -> jnp.ndarray: + """Sharding constraint that keeps bwd cotangents in the primal dtype. -class PermutationBackend(Enum): - """Token-dispatch / combine backend used by :func:`moe`. - - * ``TRITON``: TE's fused Triton kernels. Faster than ``PURE_JAX`` - on current hardware and the recommended default. - * ``PURE_JAX``: ``jnp.argsort`` + gather paths compiled as plain - XLA; useful as a numerical reference and on builds without - Triton available. - """ + Plain ``jax.lax.with_sharding_constraint`` is identity on the fwd + but does not constrain the dtype of the cotangent that flows back + through it. In this MoE bwd, ``d_x`` is built from two paths: - PURE_JAX = "pure_jax" - TRITON = "triton" + * ``d_x_from_dispatch`` from ``ep_dispatch_bwd`` -- primal dtype + (bf16 in mixed precision). + * ``d_x_from_gate = d_logits_2d @ gate_kernel.T`` where + ``d_logits_2d`` is produced by + ``fused_topk_with_score_function_bwd``. That primitive runs at + fp32 because the fwd promoted ``logits_2d`` to fp32 (the fused + topk/softmax/sigmoid kernels are only validated at fp32). - -# ============================================================================= -# Dispatch-state records (carried _dispatch -> _combine / *_bwd) -# ============================================================================= -# -# Two NamedTuples (one per permutation backend) so we get type -# discrimination at the consumer side via ``isinstance``. The backend- -# specific residuals are required fields; the EP-only residuals are -# Optional and are populated only when the run is EP-active. Each field -# is either an ``ndarray`` or ``None`` -- nothing static, since these -# values cross the shard_map pytree boundary and would otherwise be -# coerced into JitTracers. - - -@flax_struct.dataclass -class _PureJaxDispatchState: - """Residuals saved by :func:`_dispatch` on the PURE_JAX path. - - Registered as a JAX pytree via ``flax.struct.dataclass``: each - annotated field is a leaf, ``None`` is a non-leaf sentinel. The - matching spec built by :func:`_build_dispatch_specs` mirrors this - layout so shard_map's value and spec trees line up. + JAX's type promotion then makes ``d_x_from_gate + d_x_from_dispatch`` + fp32, so the user-visible ``d_x`` ends up wider than ``x``. That + doubles activation-grad bandwidth and breaks any downstream kernel + that pins a bf16 input layout. This wrapper inserts an explicit + cast back to the primal dtype on the bwd side and re-asserts the + same sharding there as well. """ - group_sizes: jnp.ndarray - sorted_indices: jnp.ndarray - routing_weights: jnp.ndarray - # EP-only: - all_shards_tokens_per_expert: Optional[jnp.ndarray] = None - local_perm_row_id_map: Optional[jnp.ndarray] = None - - -@flax_struct.dataclass -class _TritonDispatchState: - """Residuals saved by :func:`_dispatch` on the TRITON path.""" - - group_sizes: jnp.ndarray - row_id_map: jnp.ndarray - pad_offsets: Optional[jnp.ndarray] # populated only when align_size > 0 - merging_probs: jnp.ndarray - # EP-only: - all_shards_tokens_per_expert: Optional[jnp.ndarray] = None - local_perm_row_id_map: Optional[jnp.ndarray] = None - - -_DispatchState = Union[_PureJaxDispatchState, _TritonDispatchState] - + @jax.custom_vjp + def _constraint(y): + return jax.lax.with_sharding_constraint(y, sharding) -@flax_struct.dataclass -class _BodyCtx: - """Residuals carried fwd_rule -> bwd_rule by :func:`_body_fwd`. + def _constraint_fwd(y): + return jax.lax.with_sharding_constraint(y, sharding), jnp.zeros((), dtype=y.dtype) - Optional fields (``expert_bias``, ``aux_*``) are ``None`` when the - matching feature is disabled. :func:`_build_ctx_specs` mirrors that - layout so the shard_map spec and value trees match leaf-for-leaf. - """ + def _constraint_bwd(dtype_ref, grad): + return (jax.lax.with_sharding_constraint(grad.astype(dtype_ref.dtype), sharding),) - # Always present. - x: Any - gate_kernel: Any - logits_2d: Any - saved_scores: Any - routing_map: Any - dispatch: Any # _DispatchState - casted_sorted_x_lhs_trans: Any - casted_wi_rhs_trans: Any # combined [E, H, 2M] residual for fused wi_0|wi_1 bwd - gate_proj_out: Any - up_proj_out: Any - casted_intermediate_lhs_trans: Any - casted_wo_rhs_trans: Any - expert_outputs: Any - local_group_sizes: Any - # Feature-gated. - expert_bias: Any = None - aux_const_buf: Any = None - aux_tokens_per_expert: Any = None - aux_logits_for_score: Any = None - aux_saved_scores: Any = None + _constraint.defvjp(_constraint_fwd, _constraint_bwd) + return _constraint(x) # ============================================================================= -# ctx / dispatch-state key conventions +# Process-level NCCL EP bootstrap (must run eagerly, outside jax.jit) # ============================================================================= # -# Both ``ctx`` (carried fwd_rule -> bwd_rule) and the dispatch state -# (carried _dispatch -> _combine / _dispatch_bwd / _combine_bwd) are plain -# python dicts. Using a dict (rather than a flax_struct.dataclass) lets us -# vary the populated keys with the static config without breaking -# ``shard_map``'s ``out_specs`` structural match: the spec dict and the -# value dict are built with the SAME keys via :func:`_build_ctx_specs`. -# -# Below is the key glossary so the rest of the file reads cleanly. -# -# DispatchState (dict): values are jnp.ndarray unless noted -# Always present: -# "group_sizes" [n_groups] per-expert token counts -# (n_groups = E for no-EP, -# E_local for EP) -# "ep_active" bool (carried as a Python flag, -# not in the dict; passed -# alongside) -# PURE_JAX backend: -# "sorted_indices" [num_real + padding] argsort indices -# "routing_weights" [num_tokens, topk] per-token-per-expert weights -# TRITON backend: -# "row_id_map" [num_tokens, 2*E + 1] -# "pad_offsets" [E] or None -# "merging_probs" [num_tokens, E] -# EP-only: -# "all_shards_tokens_per_expert" [num_ep, E] -# "local_perm_row_id_map" [recv_buffer_rows] -# "local_perm_inv_row_id_map" [recv_buffer_rows] -# -# NOTE: per-shard compile-time-constant shapes (num_real_tokens, -# padding_size, pre/post_a2a_buffer_shape) are NOT stored in this -# dict; they are recomputed in _body_fwd/_body_bwd via -# _compute_static_shape_info and passed as Python ints / int tuples to -# the dispatch/combine helpers. Storing them in the dict would cause -# JAX's pytree-flatten across the shard_map boundary to coerce them -# into JitTracer 0-d arrays, which breaks Python-level control flow -# (e.g. ``if padding > 0``) and ``jnp.zeros(shape)`` in the bwd. -# -# See :class:`_BodyCtx` (NamedTuple) for the ctx layout and field -# documentation. :func:`_build_ctx_specs` returns a matching ``_BodyCtx`` -# of ``P(...)`` specs so shard_map's value/spec trees line up -# leaf-for-leaf. - - -# ============================================================================= -# Static shape helper -# ============================================================================= -# -# A set of per-shard shape/size values that the dispatch and combine -# helpers (both fwd and bwd) need. They're all derivable from existing -# static args, so we recompute them in both ``_body_fwd`` and -# ``_body_bwd`` and pass them as Python ints / int-tuples through -# explicit kwargs. We MUST NOT stash them inside the dynamic -# ``state`` / ``ctx`` dict: when the dict crosses the EP shard_map's -# out_specs/in_specs boundary, JAX's pytree-flatten coerces any Python -# int leaves into traced 0-d arrays, which then breaks dependent Python -# code in the bwd (e.g. ``if padding > 0`` and ``jnp.zeros(shape)``). - - -@dataclass(frozen=True) -class _StaticShapeInfo: - """Per-shard compile-time-constant shape info used by dispatch / - combine fwd and bwd. Fields are Python ints / int tuples (NOT jnp - arrays) so they can be passed as ordinary static keyword args. - - Attributes - ---------- - num_real_tokens : int - Per-shard count of real (non-padding) permuted tokens, - i.e. ``per_shard_num_tokens * num_experts_per_tok``. - padding_size : int - Per-shard number of alignment-padding tokens appended to the - sort buffer (``num_experts * (align_size - 1)`` when - ``align_size > 0``, else ``0``). - pre_a2a_buffer_shape : tuple[int, int] - ``(num_real_tokens + padding_size, hidden)`` -- the per-shard - shape of the sorted-inputs buffer sent over the EP - ragged_all_to_all in the fwd direction. - post_a2a_buffer_shape : Optional[tuple[int, int]] - ``(recv_buffer_rows, hidden)`` when EP is active, ``None`` - otherwise. - """ +# ``tex.ep_bootstrap`` does a NCCL UID allgather over the JAX runtime, which +# cannot run from inside a jit-traced function. The caller must bootstrap +# eagerly once per process before any jitted MoE call, then record the +# bootstrap signature via ``record_ep_bootstrap_signature_for_moe``. The +# per-call check below verifies the recorded signature is wide enough for +# the current MoE invocation (smaller per-call usage is fine since the C++ +# backend reserves worst-case buffers at bootstrap time). - num_real_tokens: int - padding_size: int - pre_a2a_buffer_shape: Tuple[int, int] - post_a2a_buffer_shape: Optional[Tuple[int, int]] +_te_ep_bootstrap_signature: Optional[Tuple[int, int, int, int, int]] = None -def _compute_static_shape_info( - *, - batch_size: int, - sequence_length: int, - hidden: int, +def record_ep_bootstrap_signature_for_moe( num_experts: int, - num_experts_per_tok: int, - align_size: int, - ep_active: bool, - num_ep: int = 1, - fsdp_sizes: Tuple[int, ...] = (), - recv_buffer_rows: int = 0, - batch_is_per_shard: bool = True, -) -> _StaticShapeInfo: - """Build a :class:`_StaticShapeInfo` for the current rank. - - ``batch_is_per_shard`` controls whether ``batch_size`` is already - sharded (True -- e.g. when this is called from inside a shard_map - body, where ``x.shape[0]`` reports the per-shard batch size) or - global (False -- e.g. when computing from x.shape outside the - shard_map body). + max_tokens_per_rank: int, + recv_capacity_per_rank: int, + hidden_dim: int, + ep_size: int, +) -> None: + """Record the params passed to ``ep_bootstrap`` so the per-call check + in ``_moe_fwd_rule`` can verify compatibility. Call this once per + process immediately after ``ep_bootstrap``. """ - if ep_active and not batch_is_per_shard: - dp_size = math.prod(fsdp_sizes) if fsdp_sizes else 1 - per_shard_batch = batch_size // (num_ep * dp_size) - else: - per_shard_batch = batch_size - per_shard_num_tokens = per_shard_batch * sequence_length - num_real_tokens = per_shard_num_tokens * num_experts_per_tok - padding_size = num_experts * (align_size - 1) if align_size > 0 else 0 - pre_a2a_buffer_shape = (num_real_tokens + padding_size, hidden) - post_a2a_buffer_shape = (recv_buffer_rows, hidden) if ep_active else None - return _StaticShapeInfo( - num_real_tokens=num_real_tokens, - padding_size=padding_size, - pre_a2a_buffer_shape=pre_a2a_buffer_shape, - post_a2a_buffer_shape=post_a2a_buffer_shape, + global _te_ep_bootstrap_signature + _te_ep_bootstrap_signature = ( + num_experts, + max_tokens_per_rank, + recv_capacity_per_rank, + hidden_dim, + ep_size, ) -# ============================================================================= -# Dispatch / combine helpers (no VJP boundary -- pure Python) -# ============================================================================= - - -def _dispatch( - inputs_2d: jnp.ndarray, - sparse_probs: jnp.ndarray, - routing_map: jnp.ndarray, - *, - backend: PermutationBackend, +def _te_ep_assert_compatible_bootstrap( num_experts: int, - num_experts_per_tok: int, - align_size: int, - # EP-only: - ep_active: bool, - ep_axis: Optional[str], - num_ep: int, - recv_buffer_rows: int, - shard_id: Optional[jnp.ndarray] = None, -) -> Tuple[jnp.ndarray, dict]: - """``permute -> (a2a -> local_permute) iff ep_active``. - - Returns ``(sorted_x, state)`` where ``sorted_x`` has shape - ``[buffer_rows, hidden]`` -- ``E`` groups (no-EP) or ``E_local`` groups - (EP) -- and ``state`` is a dict carrying everything :func:`_combine` - and the bwd helpers need to reverse the operation. - - Bypasses the ``custom_vjp``-wrapped public ``token_dispatch`` / - ``pure_jax_token_dispatch`` wrappers (well, mostly: PURE_JAX still - composes through ``pure_jax_token_dispatch`` because that helper has - no ``custom_vjp`` itself -- only its inner ``_sort_activations`` does, - which is fine since we never auto-diff through it from this layer). - For TRITON we call the underlying ``permute_with_mask_map`` / - ``permute_with_mask_map_and_pad`` primitives directly. - """ - num_tokens, hidden = inputs_2d.shape - topk = num_experts_per_tok - - # Backend-specific residuals collected here, then packaged into the - # appropriate _*DispatchState below. - sorted_indices = None - routing_weights_kept = None - row_id_map = None - pad_offsets = None - merging_probs = None - - # ------------------------------------------------------------------ - # Step 1: global permute (every shard routes its own tokens over the - # full expert axis). Backend-specific. - # ------------------------------------------------------------------ - if backend is PermutationBackend.PURE_JAX: - selected_experts, routing_weights = routing_map_to_selected_experts( - sparse_probs, routing_map, topk - ) - sorted_inputs, perm_state, group_sizes = pure_jax_token_dispatch( - inputs_2d, - selected_experts, - num_experts=num_experts, - num_experts_per_tok=topk, - align_size=align_size, - ) - # NOTE: ``perm_state.num_real_tokens`` and ``perm_state.padding_size`` - # are compile-time Python ints; intentionally NOT stored in the - # returned state (would be coerced to JitTracer 0-d arrays under - # the EP shard_map's pytree flatten). Recompute via - # ``_compute_static_shape_info`` in the bwd / EP-combine - # call sites that need them. - sorted_indices = perm_state.sorted_indices - routing_weights_kept = routing_weights - else: - # TRITON backend -- inline the underlying primitive sequence - # (mirrors ``_token_dispatch_fwd_rule`` but exposes the residuals - # to our ctx instead of saving them inside another custom_vjp). - num_out_tokens = num_tokens * topk - row_id_map = make_row_id_map(routing_map, num_tokens, num_experts) - tokens_per_expert = jnp.sum(routing_map, axis=0).astype(jnp.int32) - if align_size > 0: - target_tokens_per_expert = ( - jnp.ceil(tokens_per_expert / align_size) * align_size - ).astype(jnp.int32) - pad_lengths = target_tokens_per_expert - tokens_per_expert - cum_pad = jnp.cumsum(pad_lengths) - pad_offsets = jnp.concatenate([jnp.array([0], dtype=cum_pad.dtype), cum_pad[:-1]]) - worst_case_out_tokens = ( - (num_out_tokens + num_experts * (align_size - 1)) // align_size - ) * align_size - sorted_inputs, _ = permute_with_mask_map_and_pad( - inputs_2d, - row_id_map, - None, - pad_offsets, - num_tokens, - num_experts, - worst_case_out_tokens, - hidden, - align_size=align_size, - ) - group_sizes = target_tokens_per_expert - else: - sorted_inputs, _ = permute_with_mask_map( - inputs_2d, - row_id_map, - None, - num_tokens, - num_experts, - num_out_tokens, - hidden, - ) - pad_offsets = None - group_sizes = tokens_per_expert - merging_probs = sparse_probs - - def _build_state(group_sizes_val, ep_all=None, ep_local=None): - if backend is PermutationBackend.PURE_JAX: - return _PureJaxDispatchState( - group_sizes=group_sizes_val, - sorted_indices=sorted_indices, - routing_weights=routing_weights_kept, - all_shards_tokens_per_expert=ep_all, - local_perm_row_id_map=ep_local, - ) - return _TritonDispatchState( - group_sizes=group_sizes_val, - row_id_map=row_id_map, - pad_offsets=pad_offsets, - merging_probs=merging_probs, - all_shards_tokens_per_expert=ep_all, - local_perm_row_id_map=ep_local, - ) - - if not ep_active: - return sorted_inputs, _build_state(group_sizes) - - # ------------------------------------------------------------------ - # Step 2 (EP only): all_gather per-expert counts so every shard knows - # the [num_ep, num_experts] token-count matrix. - # ------------------------------------------------------------------ - all_shards_tokens_per_expert = jax.lax.all_gather( - group_sizes[None, :], - axis_name=ep_axis, - axis=0, - tiled=True, - ) - - # ------------------------------------------------------------------ - # Step 3 (EP only): forward ragged_all_to_all over the EP axis. - # ------------------------------------------------------------------ - in_off, send_sz, out_off, recv_sz = compute_ragged_all_to_all_params( - all_shards_tokens_per_expert, shard_id, num_ep - ) - post_a2a_buffer_shape = (recv_buffer_rows, hidden) - recv_buf = jnp.zeros(post_a2a_buffer_shape, dtype=sorted_inputs.dtype) - x_recv = jax.lax.ragged_all_to_all( - sorted_inputs, recv_buf, in_off, send_sz, out_off, recv_sz, axis_name=ep_axis - ) - - # ------------------------------------------------------------------ - # Step 4 (EP only): local permute -- (source_shard, expert) -> - # (expert, shard). Inlined ``local_permute_after_a2a`` so we control - # both the row_id_map and its inverse for the bwd. - # ------------------------------------------------------------------ - num_experts_local = num_experts // num_ep - local_expert_start = shard_id * num_experts_local - local_expert_columns = jax.lax.dynamic_slice( - all_shards_tokens_per_expert, - start_indices=(0, local_expert_start), - slice_sizes=(num_ep, num_experts_local), - ) - split_sizes = local_expert_columns.reshape(-1) # source-major - indices_matrix = jnp.arange(num_ep * num_experts_local, dtype=jnp.int32).reshape( - num_ep, num_experts_local - ) - sorted_chunk_indices = indices_matrix.T.reshape(-1) # source-major -> expert-major - num_chunks = num_ep * num_experts_local - # Build a SINGLE row_id_map. ``is_forward=True`` permutes - # source-major -> expert-major; ``is_forward=False`` is the exact - # inverse (this is exactly what ``_sort_chunks_by_index_bwd_rule`` - # uses on the saved residual). _MoEBlock builds two row_id_maps - # only because it calls ``sort_chunks_by_index`` twice -- once in - # ``local_permute_after_a2a`` and again in ``local_unpermute_before_a2a``; - # each of those wrappers calls ``make_chunk_sort_map`` internally. - # Here we share one map across (fwd permute, fwd inverse-permute, - # bwd permute, bwd inverse-permute). - local_perm_row_id_map = make_chunk_sort_map( - split_sizes, sorted_chunk_indices, recv_buffer_rows, num_chunks - ) - sorted_x, _ = sort_chunks_by_map( - x_recv, local_perm_row_id_map, None, recv_buffer_rows, hidden, is_forward=True - ) - local_group_sizes = jnp.sum(local_expert_columns, axis=0) - - # NOTE: pre_a2a_buffer_shape and post_a2a_buffer_shape are compile- - # time int tuples; intentionally NOT stored in the returned state - # (would be coerced to JitTracer 0-d arrays under the EP shard_map's - # pytree flatten). Recompute via ``_compute_static_shape_info`` in - # the bwd call sites that need them. For EP, ``group_sizes`` here is - # the per-local-expert count (the FFN runs over E_local groups, not - # E). The global ``group_sizes`` lives inside - # ``all_shards_tokens_per_expert`` if anyone needs it for - # diagnostics. - return sorted_x, _build_state( - local_group_sizes, - ep_all=all_shards_tokens_per_expert, - ep_local=local_perm_row_id_map, - ) - - -def _combine( - expert_outputs: jnp.ndarray, - state: _DispatchState, - *, - backend: PermutationBackend, - ep_active: bool, - batch_size: int, - sequence_length: int, - dtype: jnp.dtype, - num_experts_per_tok: int, - # Per-shard compile-time-constant shape info (Python ints / int tuples). - # Computed by _compute_static_shape_info in the caller, passed here - # rather than stored in ``state`` to survive shard_map crossings. - num_real_tokens: int, - padding_size: int, - pre_a2a_buffer_shape: Tuple[int, int], - # EP-only: - ep_axis: Optional[str], - shard_id: Optional[jnp.ndarray] = None, - num_ep: int = 1, -) -> Tuple[jnp.ndarray, jnp.ndarray]: - """Inverse of :func:`_dispatch`. - - Returns ``(output, expert_outputs_post_ep)``. ``output`` is the - ``[B, S, H]`` combined activations. ``expert_outputs_post_ep`` is - the FFN-output tensor in the shape that Step 3 of the combine - actually consumed (i.e. after the reverse ragged_all_to_all on EP - runs, or the original input on non-EP). The caller stashes this as - the bwd residual so that ``_combine_bwd``'s Step-3 inverse sees - the same tensor the forward Step 3 used. - """ - if ep_active: - # Step 1 (EP): inverse local permute. Reuse the SAME row_id_map - # built in _dispatch by setting is_forward=False (this is the - # exact inverse, identical to what - # ``_sort_chunks_by_index_bwd_rule`` does with the saved residual). - recv_buffer_rows, hidden = expert_outputs.shape - x_send_back, _ = sort_chunks_by_map( - expert_outputs, - state.local_perm_row_id_map, - None, - recv_buffer_rows, - hidden, - is_forward=False, - ) - # Step 2 (EP): reverse ragged_all_to_all. - in_off_r, send_sz_r, out_off_r, recv_sz_r = compute_reverse_ragged_all_to_all_params( - state.all_shards_tokens_per_expert, shard_id, num_ep - ) - send_back_buf = jnp.zeros(pre_a2a_buffer_shape, dtype=expert_outputs.dtype) - expert_outputs = jax.lax.ragged_all_to_all( - x_send_back, - send_back_buf, - in_off_r, - send_sz_r, - out_off_r, - recv_sz_r, - axis_name=ep_axis, - ) - - # Step 3: global combine. ``expert_outputs`` here is the post-A2A - # tensor under EP, or the original input under non-EP -- whichever - # value Step 3 actually consumes. Returned as the second tuple - # element so the caller can stash it as the bwd residual. - if backend is PermutationBackend.PURE_JAX: - # Reuse the reference pure-jax implementation; it has no - # custom_vjp on its outer surface so we can call it freely. - perm_state = PureJaxPermState( - sorted_indices=state.sorted_indices, - num_real_tokens=num_real_tokens, - padding_size=padding_size, - ) - output = pure_jax_token_combine( - expert_outputs, - perm_state, - state.routing_weights, - num_experts_per_tok=num_experts_per_tok, - batch_size=batch_size, - sequence_length=sequence_length, + max_tokens_per_rank: int, + recv_capacity_per_rank: int, + hidden_dim: int, + ep_size: int, +) -> None: + """Verify a prior eager ``ep_bootstrap`` is wide enough for this call.""" + if _te_ep_bootstrap_signature is None: + raise RuntimeError( + "TE EP was not bootstrapped. Call" + " transformer_engine.jax.ep.ep_bootstrap(...) eagerly (outside" + " any jax.jit) once per process, then" + " transformer_engine.jax.moe.record_ep_bootstrap_signature_for_moe(...)" + " with the same params, before invoking moe()." ) - return output, expert_outputs - # TRITON - num_tokens = state.row_id_map.shape[0] - num_experts = (state.row_id_map.shape[1] - 1) // 2 - hidden = expert_outputs.shape[-1] - if state.pad_offsets is not None: - out_2d, _ = unpermute_with_mask_map_and_unpad( - expert_outputs, - state.row_id_map, - state.merging_probs, - None, - state.pad_offsets, - num_tokens, - num_experts, - hidden, - ) - else: - out_2d, _ = unpermute_with_mask_map( - expert_outputs, - state.row_id_map, - state.merging_probs, - None, - num_tokens, - num_experts, - hidden, + b_num_experts, b_max_tpr, b_recv_pr, b_hidden, b_ep_size = _te_ep_bootstrap_signature + if ( + num_experts != b_num_experts + or hidden_dim != b_hidden + or ep_size != b_ep_size + or max_tokens_per_rank > b_max_tpr + or recv_capacity_per_rank > b_recv_pr + ): + raise ValueError( + "TE EP was already bootstrapped with signature" + f" (num_experts={b_num_experts}, max_tokens_per_rank={b_max_tpr}," + f" recv_capacity_per_rank={b_recv_pr}, hidden_dim={b_hidden}," + f" ep_size={b_ep_size}); this moe() call needs" + f" (num_experts={num_experts}, max_tokens_per_rank={max_tokens_per_rank}," + f" recv_capacity_per_rank={recv_capacity_per_rank}, hidden_dim={hidden_dim}," + f" ep_size={ep_size}). Re-bootstrap with wider params (or matching exact" + " sizes) is required." ) - return out_2d.reshape(batch_size, sequence_length, hidden).astype(dtype), expert_outputs -def _combine_bwd( # pylint: disable=unused-argument - d_output: jnp.ndarray, - state: _DispatchState, - expert_outputs: jnp.ndarray, - *, - backend: PermutationBackend, - ep_active: bool, - batch_size: int, - sequence_length: int, - dtype: jnp.dtype, - num_experts: int, - num_experts_per_tok: int, - # Per-shard compile-time-constant shape info (Python ints / int tuples). - # See ``_compute_static_shape_info`` and the note in ``_dispatch`` - # for why these are kwargs rather than state-dict entries. - num_real_tokens: int, - padding_size: int, - post_a2a_buffer_shape: Optional[Tuple[int, int]], - # EP-only: - ep_axis: Optional[str], - shard_id: Optional[jnp.ndarray] = None, - num_ep: int = 1, -) -> Tuple[jnp.ndarray, Optional[jnp.ndarray]]: - """Inverse of :func:`_combine` on the cotangent. - - Returns ``(d_expert_outputs, d_routing_weights_or_merging_probs)``. +# ============================================================================= +# Residual container threaded fwd -> bwd +# ============================================================================= - ``expert_outputs`` is the *forward* output of the FFN (same value the - fwd handed to :func:`_combine`). It's required by the TRITON - combine_bwd kernel; for PURE_JAX we don't need it but accept it for - a symmetric signature. - """ - # Step 3 inverse: global combine bwd. - d_output_2d = d_output.reshape(-1, d_output.shape[-1]) - if backend is PermutationBackend.PURE_JAX: - # The pure-jax combine is: - # unsort = _sort_activations(expert_outputs, argsort(sorted_indices)) - # if pad: unsort = unsort[:num_real] - # reshape -> einsum BKE,BK -> BE -> reshape to BSE - # Hand-derive the bwd in plain JAX (no custom_vjp involved): - unsort_indices = jnp.argsort(state.sorted_indices) - topk = num_experts_per_tok - num_real = num_real_tokens - padding = padding_size - # Recover the unsorted intermediate that the fwd produced (we - # need it for the d_routing_weights pullback). Apply the same - # gather the fwd did. - unsort_intermediate = expert_outputs[unsort_indices] - if padding > 0: - unsort_intermediate = unsort_intermediate[:num_real] - # Bwd of einsum/reshape: - # output[B, E] = sum_K intermediate[B, K, E] * weights[B, K] - # d_intermediate[B, K, E] = d_output[B, E] * weights[B, K] - # d_weights[B, K] = sum_E d_output[B, E] * intermediate[B, K, E] - rw = state.routing_weights.reshape(-1, topk) - intermediate_3d = unsort_intermediate.reshape(rw.shape[0], topk, -1) - rw_cast = rw.astype(intermediate_3d.dtype) - d_intermediate_3d = jnp.einsum("BE,BK -> BKE", d_output_2d, rw_cast) - d_routing_weights = jnp.einsum("BE,BKE -> BK", d_output_2d, intermediate_3d).astype( - state.routing_weights.dtype - ) - d_routing_weights = d_routing_weights.reshape(state.routing_weights.shape) - d_unsort_intermediate = d_intermediate_3d.reshape(num_real, -1) - # Pad back with zeros if the fwd stripped padding. - if padding > 0: - d_unsort_intermediate = jnp.concatenate( - [ - d_unsort_intermediate, - jnp.zeros( - (padding, d_unsort_intermediate.shape[-1]), - dtype=d_unsort_intermediate.dtype, - ), - ], - axis=0, - ) - # Bwd of the gather is gather-by-original-indices: - # sorted = unsort[argsort(sorted_indices)] - # d_sorted = scatter d_unsort via argsort(sorted_indices) - # = d_unsort[sorted_indices] (gather by original sorted_indices, - # which is the inverse of argsort(sorted_indices)). - d_expert_outputs_global = d_unsort_intermediate[state.sorted_indices] - else: - # TRITON combine bwd: requires fwd_input (expert_outputs). - num_tokens = state.row_id_map.shape[0] - n_experts = (state.row_id_map.shape[1] - 1) // 2 - hidden = d_output_2d.shape[-1] - num_out_tokens = expert_outputs.shape[0] - if state.pad_offsets is not None: - d_expert_outputs_global, d_merging_probs = unpermute_bwd_with_merging_probs_and_unpad( - d_output_2d, - state.row_id_map, - expert_outputs, - state.merging_probs, - state.pad_offsets, - num_tokens, - n_experts, - num_out_tokens, - hidden, - ) - # The kernel only writes positions tokens map to; padded - # positions may contain NaN. Replace with zeros (matches - # ``_token_combine_bwd_rule``). - d_expert_outputs_global = jnp.where( - jnp.isnan(d_expert_outputs_global), 0.0, d_expert_outputs_global - ) - else: - d_expert_outputs_global, d_merging_probs = unpermute_bwd_with_merging_probs( - d_output_2d, - state.row_id_map, - expert_outputs, - state.merging_probs, - num_tokens, - n_experts, - num_out_tokens, - hidden, - ) - d_routing_weights = d_merging_probs - - if not ep_active: - return d_expert_outputs_global, d_routing_weights - - # Step 2 (EP) inverse: bwd of reverse ragged_all_to_all is a forward - # ragged_all_to_all using the SAME forward parameters (sender / - # receiver roles swap from the reverse direction back to forward). - in_off_f, send_sz_f, out_off_f, recv_sz_f = compute_ragged_all_to_all_params( - state.all_shards_tokens_per_expert, shard_id, num_ep - ) - recv_buf_for_bwd = jnp.zeros(post_a2a_buffer_shape, dtype=d_expert_outputs_global.dtype) - d_x_send_back = jax.lax.ragged_all_to_all( - d_expert_outputs_global, - recv_buf_for_bwd, - in_off_f, - send_sz_f, - out_off_f, - recv_sz_f, - axis_name=ep_axis, - ) - # Step 1 (EP) inverse: combine fwd applied is_forward=False; the - # bwd is is_forward=True with the SAME row_id_map. - recv_buffer_rows, hidden = d_x_send_back.shape - d_expert_outputs, _ = sort_chunks_by_map( - d_x_send_back, - state.local_perm_row_id_map, - None, - recv_buffer_rows, - hidden, - is_forward=True, - ) - return d_expert_outputs, d_routing_weights +@flax.struct.dataclass +class _Ctx: + """Residuals carried from the fwd rule into the bwd rule. -def _dispatch_bwd( - d_sorted_x: jnp.ndarray, - state: _DispatchState, - inputs_2d_shape: Tuple[int, ...], - *, - backend: PermutationBackend, - ep_active: bool, - num_experts: int, - num_experts_per_tok: int, - # Per-shard compile-time-constant shape info (Python ints / int tuples). - # See ``_compute_static_shape_info`` and the note in ``_dispatch`` - # for why these are kwargs rather than state-dict entries. - num_real_tokens: int, - padding_size: int, - pre_a2a_buffer_shape: Tuple[int, int], - # EP-only: - ep_axis: Optional[str], - shard_id: Optional[jnp.ndarray] = None, - num_ep: int = 1, -) -> jnp.ndarray: - """Inverse of :func:`_dispatch` on the cotangent. Returns ``d_inputs_2d``. - - The probs path through dispatch is always discarded (PURE_JAX never - threads probs through dispatch; TRITON technically does but the - caller drops ``permuted_probs``, so its cotangent is structurally - zero). The probs gradient instead flows back through - :func:`_combine_bwd`. + Flattened automatically by jax.custom_vjp; ``cfg`` is the only + static field (the rest are jnp.ndarray, GroupedNoScaleTensor, or + None when aux_loss_coeff == 0). """ - if ep_active: - # Step 4 inverse: dispatch fwd applied is_forward=True; bwd is - # is_forward=False with the SAME row_id_map. - recv_buffer_rows, hidden = d_sorted_x.shape - d_x_recv, _ = sort_chunks_by_map( - d_sorted_x, - state.local_perm_row_id_map, - None, - recv_buffer_rows, - hidden, - is_forward=False, - ) - # Step 3 inverse: bwd of forward ragged_a2a is the reverse-direction - # ragged_a2a using the SAME params with sender/receiver swapped. - in_off_r, send_sz_r, out_off_r, recv_sz_r = compute_reverse_ragged_all_to_all_params( - state.all_shards_tokens_per_expert, shard_id, num_ep - ) - recv_buf_pre = jnp.zeros(pre_a2a_buffer_shape, dtype=d_x_recv.dtype) - d_sorted_x = jax.lax.ragged_all_to_all( - d_x_recv, - recv_buf_pre, - in_off_r, - send_sz_r, - out_off_r, - recv_sz_r, - axis_name=ep_axis, - ) - # Step 1 inverse: global permute bwd. - if backend is PermutationBackend.PURE_JAX: - # Fwd was: replicated = repeat(inputs_2d, topk, axis=0) - # padded = pad(replicated, (0, padding_size)) - # sorted = padded[sorted_indices] - # Bwd: d_padded = scatter via sorted_indices - # = d_sorted[argsort(sorted_indices)] - # d_replicated = d_padded[:num_real] - # d_inputs_2d = d_replicated.reshape(T, topk, H).sum(axis=1) - sorted_indices = state.sorted_indices - num_real = num_real_tokens - padding = padding_size - topk = num_experts_per_tok - unsort_indices = jnp.argsort(sorted_indices) - d_padded = d_sorted_x[unsort_indices] - if padding > 0: - d_replicated = d_padded[:num_real] - else: - d_replicated = d_padded - num_tokens = inputs_2d_shape[0] - hidden = inputs_2d_shape[-1] - d_inputs_2d = d_replicated.reshape(num_tokens, topk, hidden).sum(axis=1) - return d_inputs_2d - - # TRITON: bwd is unpermute_with_mask_map[_and_unpad]. - num_tokens = inputs_2d_shape[0] - hidden = inputs_2d_shape[-1] - if state.pad_offsets is not None: - d_inputs_2d, _ = unpermute_with_mask_map_and_unpad( - d_sorted_x, - state.row_id_map, - None, - None, - state.pad_offsets, - num_tokens, - num_experts, - hidden, - ) - else: - d_inputs_2d, _ = unpermute_with_mask_map( - d_sorted_x, - state.row_id_map, - None, - None, - num_tokens, - num_experts, - hidden, - ) - return d_inputs_2d + x: jnp.ndarray + gate_kernel: jnp.ndarray + expert_bias: jnp.ndarray + logits_2d: jnp.ndarray + saved_scores: jnp.ndarray + routing_map: jnp.ndarray + cfg: Any = flax.struct.field(pytree_node=False) + handle_mem: jnp.ndarray + token_counts: jnp.ndarray + recv_topk_weights: jnp.ndarray + casted_sorted_x_lhs_trans: Any + casted_wi_rhs_trans: Any + gate_proj_out: jnp.ndarray + up_proj_out: jnp.ndarray + casted_intermediate_lhs_trans: Any + casted_wo_rhs_trans: Any + expert_outputs: jnp.ndarray + local_group_sizes: jnp.ndarray + aux_const_buf: Any = None + aux_tokens_per_expert: Any = None + aux_saved_scores: Any = None # ============================================================================= -# Per-shard body +# Per-shard FFN body (runs inside shard_map) # ============================================================================= -def _body_fwd( # pylint: disable=unused-argument - captured: dict, +def _ffn_fwd_per_shard( + recv_tokens_local: jnp.ndarray, + recv_topk_weights_local: jnp.ndarray, + token_counts_local: jnp.ndarray, + wi_0: jnp.ndarray, + wi_1: jnp.ndarray, + wo: jnp.ndarray, + wi_0_bias: Optional[jnp.ndarray], + wi_1_bias: Optional[jnp.ndarray], + wo_bias: Optional[jnp.ndarray], *, - # Statics - num_experts: int, - num_experts_per_tok: int, + num_local_experts: int, activation_type: str, - score_function: ScoreFunction, - use_pre_softmax: bool, - num_groups: Optional[int], - group_topk: Optional[int], - scaling_factor: float, - aux_loss_coeff: float, - permutation_backend: PermutationBackend, - align_size: int, - gate_inside_vjp: bool, - quantizer_sets: Tuple[QuantizerSet, QuantizerSet, QuantizerSet], - dtype: jnp.dtype, - # EP-only statics - ep_active: bool, - ep_axis: Optional[str], - data_parallelism_axes: Tuple[str, ...], - fsdp_sizes: Tuple[int, ...], - num_ep: int, - num_experts_local: int, - recv_buffer_rows: int, -) -> Tuple[jnp.ndarray, jnp.ndarray, dict]: - """Per-shard forward body. Returns ``(output, aux_loss, ctx_dict)``. - - ``aux_loss`` is always materialized (zeros scalar when disabled) so - the ``shard_map``'s ``out_specs`` has a static structure. + apply_topk_weights_early: bool, +): + """Per-shard FFN forward. + + Operates on the shard-local ``[1, recv_pr, H]`` slice that + ``tex.ep_dispatch`` produces. Returns the expert outputs (shaped + ``[1, recv_pr, H_out]`` so the surrounding ``shard_map`` reassembles + them as ``[num_procs, recv_pr, H_out]``) plus the residuals consumed + by the bwd. + + ``token_counts_local`` (``[1, num_local_experts]``, from + ``tex.ep_prepare``) is passed to ``grouped_gemm`` as ``group_sizes`` + so cuBLAS skips both 0-token-routed experts and the dispatch + overalloc tail. """ - if not gate_inside_vjp: - raise NotImplementedError( - "gate_inside_vjp=False is deferred to a follow-up PR; for now" - " the gate GEMM lives inside the MoE VJP." - ) - - x = captured["inputs"] - gate_kernel = captured["gate_kernel"] - wi_0 = captured["wi_0"] - wi_1 = captured["wi_1"] - wo = captured["wo"] - wi_0_bias = captured.get("wi_0_bias") - wi_1_bias = captured.get("wi_1_bias") - wo_bias = captured.get("wo_bias") - expert_bias = captured.get("expert_bias") - - batch_size, sequence_length, hidden = x.shape - - # ---------------- Stage 1: gate ---------------- - gate_kernel_cast = gate_kernel.astype(x.dtype) - gate_logits = jnp.einsum("bsh,he->bse", x, gate_kernel_cast) - logits_2d = gate_logits.reshape(-1, num_experts) - inputs_2d = x.reshape(-1, hidden) - - # ---------------- Stage 2: routing ---------------- - # Under EP, expert_bias is sharded P(ep_axis); the router needs the - # full E-dim view, so all_gather it. - if ep_active and expert_bias is not None: - full_expert_bias = jax.lax.all_gather(expert_bias, axis_name=ep_axis, tiled=True) - else: - full_expert_bias = expert_bias - # Pass an empty array sentinel when expert_bias is unused (the - # underlying primitive expects a real ndarray, not None). - eb_arg = ( - full_expert_bias if full_expert_bias is not None else jnp.zeros((0,), dtype=jnp.float32) - ) - sparse_probs, routing_map, saved_scores = tex.fused_topk_with_score_function_fwd( - logits_2d, - topk=num_experts_per_tok, - use_pre_softmax=use_pre_softmax, - num_groups=-1 if num_groups is None else num_groups, - group_topk=-1 if group_topk is None else group_topk, - scaling_factor=scaling_factor, - score_function=score_function, - expert_bias=eb_arg, - compute_aux_scores=False, - ) - sparse_probs = sparse_probs.astype(dtype) - - # ---------------- Stage 2b: aux loss ---------------- - if aux_loss_coeff > 0.0: - if ep_active: - collective_axes: Any = ( - ep_axis if not data_parallelism_axes else (ep_axis, *data_parallelism_axes) - ) - global_logits_2d = jax.lax.all_gather( - logits_2d, axis_name=collective_axes, axis=0, tiled=True - ) - _, global_routing_map, _ = tex.fused_topk_with_score_function_fwd( - global_logits_2d, - topk=num_experts_per_tok, - use_pre_softmax=use_pre_softmax, - num_groups=-1 if num_groups is None else num_groups, - group_topk=-1 if group_topk is None else group_topk, - scaling_factor=scaling_factor, - score_function=score_function, - expert_bias=eb_arg, - compute_aux_scores=False, - ) - aux_tokens_per_expert = jnp.sum(global_routing_map.astype(jnp.int32), axis=0) - aux_logits_for_score = global_logits_2d - else: - aux_tokens_per_expert = jnp.sum(routing_map.astype(jnp.int32), axis=0) - aux_logits_for_score = logits_2d - # Aux-side scores: clean per-expert scores (no grouped routing, - # no bias). compute_aux_scores=True takes a separate path that - # ignores the grouping knobs. - aux_probs, _aux_routing_map, aux_saved_scores = tex.fused_topk_with_score_function_fwd( - aux_logits_for_score.astype(jnp.float32), - topk=num_experts_per_tok, - use_pre_softmax=False, - num_groups=-1, - group_topk=-1, - scaling_factor=1.0, - score_function=score_function, - expert_bias=jnp.zeros((0,), dtype=jnp.float32), - compute_aux_scores=True, - ) - aux_loss, aux_const_buf = tex.fused_moe_aux_loss_fwd( - aux_probs.astype(jnp.float32), - aux_tokens_per_expert.astype(jnp.int32), - topk=num_experts_per_tok, - coeff=aux_loss_coeff, - ) - else: - aux_loss = jnp.zeros((), dtype=dtype) - aux_const_buf = None - aux_tokens_per_expert = None - aux_logits_for_score = None - aux_saved_scores = None - - # ---------------- Stage 3: dispatch ---------------- - shard_id = jax.lax.axis_index(ep_axis) if ep_active else None - sorted_x, dispatch_state = _dispatch( - inputs_2d, - sparse_probs, - routing_map, - backend=permutation_backend, - num_experts=num_experts, - num_experts_per_tok=num_experts_per_tok, - align_size=align_size, - ep_active=ep_active, - ep_axis=ep_axis, - num_ep=num_ep, - recv_buffer_rows=recv_buffer_rows, - shard_id=shard_id, - ) - local_group_sizes = dispatch_state.group_sizes - - # ---------------- Stage 4: per-expert FFN (inlined) ---------------- - q_set_w0, q_set_w1, q_set_wo = quantizer_sets - if q_set_w0 == noop_quantizer_set: - wi_0 = wi_0.astype(sorted_x.dtype) - if q_set_w1 == noop_quantizer_set: - wi_1 = wi_1.astype(sorted_x.dtype) - if q_set_wo == noop_quantizer_set: - wo = wo.astype(sorted_x.dtype) - - # GEMM 1+2 (fused): up_proj_combined = sorted_x @ wi where - # wi := concat([wi_0, wi_1], axis=-1) -> shape [E, H, 2M] - # combined_out := sorted_x @ wi -> shape [T, 2M] - # Splitting the output back into ``gate_proj_out`` / ``up_proj_out`` - # is free (it's a slicing reshape). This collapses two grouped - # GEMMs and two grouped quantizes of ``sorted_x`` (one per kernel) - # into one of each. Bias is concatenated the same way. - # - # FP8/MXFP8 caveat: per-expert amax is now computed over [H, 2M] - # rather than [H, M] for each of wi_0 / wi_1 separately, so the - # representable range for one of the two halves may shift slightly - # vs. the pre-fusion code. Numerics tests cover this. - inter_M = wi_0.shape[-1] + hidden = recv_tokens_local.shape[-1] + sorted_x = recv_tokens_local.reshape(-1, hidden) + recv_w_flat = recv_topk_weights_local.reshape(-1) + local_group_sizes = token_counts_local.reshape(-1).astype(jnp.int32) + + wi_0 = wi_0.astype(sorted_x.dtype) + wi_1 = wi_1.astype(sorted_x.dtype) + wo = wo.astype(sorted_x.dtype) + + # Concat wi_0/wi_1 along the trailing axis (NOT stack on a new + # axis). grouped_gemm requires the 3D (G, K, N) weight layout with + # contracting_dims=((1,), (1,)); a 4D stack variant walks off the + # end of the RHS and returns NaN. wi_combined = jnp.concatenate([wi_0, wi_1], axis=-1) wi_combined_bias = ( jnp.concatenate([wi_0_bias, wi_1_bias], axis=-1) if wi_0_bias is not None else None ) - casted_sorted_x = tex.grouped_quantize(sorted_x, q_set_w0.x, local_group_sizes, flatten_axis=-1) - casted_wi = tex.grouped_quantize(wi_combined, q_set_w0.kernel, flatten_axis=-1) + + q_set = noop_quantizer_set + casted_sorted_x = tex.grouped_quantize(sorted_x, q_set.x, local_group_sizes, flatten_axis=-1) + casted_wi = tex.grouped_quantize(wi_combined, q_set.kernel, flatten_axis=-1) combined_out = tex.grouped_gemm( casted_sorted_x.get_tensor(usage=TensorUsage.LHS), casted_wi.get_tensor(usage=TensorUsage.RHS), contracting_dims=((1,), (1,)), bias=wi_combined_bias, ) - gate_proj_out = combined_out[..., :inter_M] - up_proj_out = combined_out[..., inter_M:] + gate_proj_out, up_proj_out = jnp.split(combined_out, 2, axis=-1) casted_sorted_x_lhs_trans = casted_sorted_x.get_tensor(usage=TensorUsage.LHS_TRANS) casted_wi_rhs_trans = casted_wi.get_tensor(usage=TensorUsage.RHS_TRANS) - if isinstance(casted_sorted_x_lhs_trans, ScaledTensor): - casted_sorted_x_lhs_trans = casted_sorted_x_lhs_trans.checkpoint(q_set_w0.x) - if isinstance(casted_wi_rhs_trans, ScaledTensor): - casted_wi_rhs_trans = casted_wi_rhs_trans.checkpoint(q_set_w0.kernel) - # Activation: intermediate = act(gate_proj_out) * up_proj_out + # Activation inputs (gate_proj_out, up_proj_out) stay in the wi GEMM + # output dtype; the activation output (`intermediate`) stays in the + # dtype the wo GEMM / wo's quantized input consumes. For bf16 compute + # that's all bf16; for FP8/FP4 the downstream grouped_quantize is what + # transitions to the target precision. act_fn = _convert_to_activation_function(activation_type) intermediate = act_fn(gate_proj_out) * up_proj_out - # GEMM 3: expert_outputs = intermediate @ wo + if apply_topk_weights_early: + # Fold the per-token combine weights into the FFN intermediate; + # the downstream wo GEMM is linear so this is equivalent to the + # late-weighting path. Padded recv slots can contain uninitialized + # data, so overwrite inactive rows with literal zeros instead of + # relying on multiplication by a zero mask (IEEE NaN * 0 = NaN). + # ``w_b`` is cast to ``intermediate.dtype`` so the multiply doesn't + # promote expert_outputs above the EP buffer's element width. + w_b = recv_w_flat[:, None].astype(intermediate.dtype) + active = (recv_w_flat != 0)[:, None] + intermediate = jnp.where(active, intermediate * w_b, jnp.zeros_like(intermediate)) + casted_intermediate = tex.grouped_quantize( - intermediate, q_set_wo.x, local_group_sizes, flatten_axis=-1 + intermediate, q_set.x, local_group_sizes, flatten_axis=-1 ) - casted_wo = tex.grouped_quantize(wo, q_set_wo.kernel, flatten_axis=-1) + casted_wo = tex.grouped_quantize(wo, q_set.kernel, flatten_axis=-1) expert_outputs = tex.grouped_gemm( casted_intermediate.get_tensor(usage=TensorUsage.LHS), casted_wo.get_tensor(usage=TensorUsage.RHS), @@ -1151,524 +307,148 @@ def _body_fwd( # pylint: disable=unused-argument ) casted_intermediate_lhs_trans = casted_intermediate.get_tensor(usage=TensorUsage.LHS_TRANS) casted_wo_rhs_trans = casted_wo.get_tensor(usage=TensorUsage.RHS_TRANS) - if isinstance(casted_intermediate_lhs_trans, ScaledTensor): - casted_intermediate_lhs_trans = casted_intermediate_lhs_trans.checkpoint(q_set_wo.x) - if isinstance(casted_wo_rhs_trans, ScaledTensor): - casted_wo_rhs_trans = casted_wo_rhs_trans.checkpoint(q_set_wo.kernel) - - # ---------------- Stage 5: combine ---------------- - # Compute per-shard static shape info once and pass through both - # _combine and (later) the bwd helpers via kwargs -- never via the - # state dict, which gets pytree-flattened across shard_map and would - # coerce Python ints into JitTracer 0-d arrays. - _static_shape = _compute_static_shape_info( - batch_size=batch_size, - sequence_length=sequence_length, - hidden=hidden, - num_experts=num_experts, - num_experts_per_tok=num_experts_per_tok, - align_size=align_size, - ep_active=ep_active, - num_ep=num_ep, - fsdp_sizes=fsdp_sizes, - recv_buffer_rows=recv_buffer_rows, - ) - # ``expert_outputs_residual`` is the post-A2A FFN-output tensor that - # Step 3 of the combine actually consumed. Saving this (rather than - # the pre-A2A shard-local FFN output) is what makes - # ``_combine_bwd``'s Step-3 inverse see the same value the forward - # Step 3 saw -- otherwise EP + TRITON yields wrong d_expert_outputs. - output, expert_outputs_residual = _combine( - expert_outputs, - dispatch_state, - backend=permutation_backend, - ep_active=ep_active, - batch_size=batch_size, - sequence_length=sequence_length, - dtype=dtype, - num_experts_per_tok=num_experts_per_tok, - num_real_tokens=_static_shape.num_real_tokens, - padding_size=_static_shape.padding_size, - pre_a2a_buffer_shape=_static_shape.pre_a2a_buffer_shape, - ep_axis=ep_axis, - shard_id=shard_id, - num_ep=num_ep, - ) - # ---------------- Build ctx ---------------- - aux_enabled = aux_loss_coeff > 0.0 - ctx = _BodyCtx( - x=x, - gate_kernel=gate_kernel, - logits_2d=logits_2d, - saved_scores=saved_scores, - routing_map=routing_map, - dispatch=dispatch_state, - casted_sorted_x_lhs_trans=casted_sorted_x_lhs_trans, - casted_wi_rhs_trans=casted_wi_rhs_trans, - gate_proj_out=gate_proj_out, - up_proj_out=up_proj_out, - casted_intermediate_lhs_trans=casted_intermediate_lhs_trans, - casted_wo_rhs_trans=casted_wo_rhs_trans, - expert_outputs=expert_outputs_residual, - local_group_sizes=local_group_sizes, - expert_bias=expert_bias if expert_bias is not None else None, - aux_const_buf=aux_const_buf if aux_enabled else None, - aux_tokens_per_expert=aux_tokens_per_expert if aux_enabled else None, - aux_logits_for_score=aux_logits_for_score if aux_enabled else None, - aux_saved_scores=aux_saved_scores if aux_enabled else None, + expert_outputs_3d = expert_outputs.reshape(1, expert_outputs.shape[0], expert_outputs.shape[1]) + # Reshape local_group_sizes to (1, num_local_experts) so the + # surrounding shard_map can stitch per-shard counts back into the + # global (num_procs, num_local_experts) layout matching token_counts. + local_group_sizes_3d = local_group_sizes.reshape(1, num_local_experts) + residuals = ( + casted_sorted_x_lhs_trans, + casted_wi_rhs_trans, + gate_proj_out, + up_proj_out, + casted_intermediate_lhs_trans, + casted_wo_rhs_trans, + local_group_sizes_3d, ) - - return output, aux_loss, ctx - - -def _body_bwd( # pylint: disable=unused-argument - ctx: _BodyCtx, - dy_pair: Tuple[jnp.ndarray, jnp.ndarray], + return expert_outputs_3d, residuals + + +def _ffn_bwd_per_shard( + d_expert_outputs_local: jnp.ndarray, + casted_sorted_x_lhs_trans, + casted_wi_rhs_trans, + gate_proj_out: jnp.ndarray, + up_proj_out: jnp.ndarray, + casted_intermediate_lhs_trans, + casted_wo_rhs_trans, + local_group_sizes: jnp.ndarray, + recv_topk_weights_local: jnp.ndarray, *, - num_experts: int, - num_experts_per_tok: int, activation_type: str, - score_function: ScoreFunction, - use_pre_softmax: bool, - num_groups: Optional[int], - group_topk: Optional[int], - scaling_factor: float, - aux_loss_coeff: float, - permutation_backend: PermutationBackend, - align_size: int, - gate_inside_vjp: bool, - quantizer_sets: Tuple[QuantizerSet, QuantizerSet, QuantizerSet], - dtype: jnp.dtype, - ep_active: bool, - ep_axis: Optional[str], - data_parallelism_axes: Tuple[str, ...], - fsdp_sizes: Tuple[int, ...], - num_ep: int, - num_experts_local: int, - recv_buffer_rows: int, - # Static side info (kept here rather than inside ctx because they're - # python flags / shapes, not array leaves): - has_wi_bias: bool, - has_wo_bias: bool, - has_expert_bias: bool, - x_shape: Tuple[int, ...], -) -> dict: - """Per-shard backward body. Returns a dict of grads keyed identically - to the ``captured`` dict consumed by :func:`_body_fwd`.""" - if not gate_inside_vjp: - raise NotImplementedError("gate_inside_vjp=False is deferred to a follow-up PR.") - - d_output, d_aux_loss = dy_pair - # The fused FFN bwd quantizes via ``q_set_w0`` only (one quantize for - # the [E, H, 2M] fused wi tensor and one for the [T, 2M] fused dgrad), - # so ``q_set_w1`` is intentionally unused here. - q_set_w0, _q_set_w1, q_set_wo = quantizer_sets - batch_size, sequence_length, hidden = x_shape - shard_id = jax.lax.axis_index(ep_axis) if ep_active else None - - # Recompute per-shard static shape info from existing statics - # (Python ints / int tuples). Plumbed via kwargs to _combine_bwd - # and _dispatch_bwd -- NOT through the ctx dict, because the - # dict gets pytree-flattened across the bwd shard_map's in_specs - # and Python ints would be coerced into JitTracer 0-d arrays - # (breaking ``if padding > 0`` and ``jnp.zeros(shape)`` callsites). - # ``batch_size`` here is the GLOBAL batch size (captured in - # ``x_shape`` by the outer fwd rule), hence ``batch_is_per_shard=False``. - _static_shape = _compute_static_shape_info( - batch_size=batch_size, - sequence_length=sequence_length, - hidden=hidden, - num_experts=num_experts, - num_experts_per_tok=num_experts_per_tok, - align_size=align_size, - ep_active=ep_active, - num_ep=num_ep, - fsdp_sizes=fsdp_sizes, - recv_buffer_rows=recv_buffer_rows, - batch_is_per_shard=False, - ) - - # Compute per-shard input shape: under the EP shard_map body, the - # gradient tensors live at per-shard shape, so the dispatch_bwd - # reshape target and ``d_x_from_dispatch.reshape(x_shape)`` below - # must use the per-shard shape rather than the captured global - # ``x_shape``. - if ep_active: - dp_size = math.prod(fsdp_sizes) if fsdp_sizes else 1 - per_shard_batch = batch_size // (num_ep * dp_size) - per_shard_x_shape: Tuple[int, ...] = (per_shard_batch, sequence_length, hidden) - else: - per_shard_x_shape = x_shape - - # ---------------- Combine bwd ---------------- - d_expert_outputs, d_routing_weights = _combine_bwd( - d_output, - ctx.dispatch, - ctx.expert_outputs, - backend=permutation_backend, - ep_active=ep_active, - batch_size=batch_size, - sequence_length=sequence_length, - dtype=dtype, - num_experts=num_experts, - num_experts_per_tok=num_experts_per_tok, - num_real_tokens=_static_shape.num_real_tokens, - padding_size=_static_shape.padding_size, - post_a2a_buffer_shape=_static_shape.post_a2a_buffer_shape, - ep_axis=ep_axis, - shard_id=shard_id, - num_ep=num_ep, - ) + apply_topk_weights_early: bool, + has_bias: bool, +): + """Per-shard FFN backward. - # ---------------- FFN bwd: GEMM 3 (wo) ---------------- - casted_d_eo = tex.grouped_quantize( - d_expert_outputs, q_set_wo.dgrad, ctx.local_group_sizes, flatten_axis=-1 - ) + Mirrors :func:`_ffn_fwd_per_shard`. Returns + ``(d_sorted_x [1, recv_pr, H], d_recv_w [1, recv_pr], + d_wi_0, d_wi_1, d_wo, d_wi_0_bias, d_wi_1_bias, d_wo_bias)``. + """ + local_group_sizes = local_group_sizes.reshape(-1).astype(jnp.int32) + d_eo_2d = d_expert_outputs_local.reshape(-1, d_expert_outputs_local.shape[-1]) + recv_w_flat = recv_topk_weights_local.reshape(-1) + q_set = noop_quantizer_set + # cuBLAS grouped_gemm skips size_g == 0 groups without zero-filling + # the output slice; mask 0-token-expert wgrads to zero so the + # optimizer never sees uninit memory. + wgrad_group_active = (local_group_sizes > 0)[:, None, None] + + # wo bwd + casted_d_eo = tex.grouped_quantize(d_eo_2d, q_set.dgrad, local_group_sizes, flatten_axis=-1) + _casted_d_eo_lhs = casted_d_eo.get_tensor(usage=TensorUsage.LHS) + _casted_d_eo_rhs = casted_d_eo.get_tensor(usage=TensorUsage.RHS) d_intermediate = tex.grouped_gemm( - casted_d_eo.get_tensor(usage=TensorUsage.LHS), - ctx.casted_wo_rhs_trans, + _casted_d_eo_lhs, + casted_wo_rhs_trans, contracting_dims=((1,), (2,)), ) d_wo = tex.grouped_gemm( - ctx.casted_intermediate_lhs_trans, - casted_d_eo.get_tensor(usage=TensorUsage.RHS), + casted_intermediate_lhs_trans, + _casted_d_eo_rhs, contracting_dims=((0,), (0,)), ) - d_wo_bias = tex.grouped_dbias(d_expert_outputs, ctx.local_group_sizes) if has_wo_bias else None + d_wo = jnp.where(wgrad_group_active, d_wo, jnp.zeros_like(d_wo)) + d_wo_bias = tex.grouped_dbias(d_eo_2d, local_group_sizes) if has_bias else None - # ---------------- Activation bwd ---------------- - # intermediate = act(gate_proj_out) * up_proj_out - # d(gate_proj_out) = vjp(act, gate_proj_out)(d_intermediate * up_proj_out) - # d(up_proj_out) = d_intermediate * act(gate_proj_out) act_fn = _convert_to_activation_function(activation_type) - act_gate_proj_out, dact_gate_proj_pullback = jax.vjp(act_fn, ctx.gate_proj_out) - d_up_proj_out = d_intermediate * act_gate_proj_out - (d_gate_proj_out,) = dact_gate_proj_pullback(d_intermediate * ctx.up_proj_out) - - # ---------------- FFN bwd: GEMM 1+2 fused (wi_0 | wi_1) ---------------- - # Concat the two upstream grads along the output (M) axis, do one - # grouped quantize + one dgrad GEMM + one wgrad GEMM, then split. - # ``ctx.casted_wi_rhs_trans`` has shape [E, H, 2M] from the fwd - # fused quantize, so the dgrad math is: - # d_sorted_x = [d_gate | d_up] @ wi_rhs_trans - # = d_gate @ wi_0^T + d_up @ wi_1^T - inter_M = d_gate_proj_out.shape[-1] + if apply_topk_weights_early: + # intermediate' = intermediate * w * mask. Split the cotangent + # across both factors before the activation bwd consumes it. Padded + # recv slots may still be NaN in the saved activation residuals, so + # use zero-filled residuals on inactive rows before the activation VJP. + w_b = recv_w_flat[:, None].astype(d_intermediate.dtype) + active = (recv_w_flat != 0)[:, None] + gate_proj_for_bwd = jnp.where(active, gate_proj_out, jnp.zeros_like(gate_proj_out)) + up_proj_for_bwd = jnp.where(active, up_proj_out, jnp.zeros_like(up_proj_out)) + intermediate_unweighted = act_fn(gate_proj_for_bwd) * up_proj_for_bwd + d_recv_w_from_intermediate = jnp.sum( + d_intermediate * intermediate_unweighted, + axis=-1, + ).astype(recv_w_flat.dtype) + d_intermediate = jnp.where(active, d_intermediate * w_b, jnp.zeros_like(d_intermediate)) + else: + gate_proj_for_bwd = gate_proj_out + up_proj_for_bwd = up_proj_out + d_recv_w_from_intermediate = jnp.zeros_like(recv_w_flat) + + # Activation bwd, symmetric with the fwd: silu' and the two + # elementwise products run in the GEMM dtype (no fp32 island), so + # the chain rule composes through at the same precision the wi/wo + # GEMMs consume. + act_gp, dact_pullback = jax.vjp(act_fn, gate_proj_for_bwd) + d_up_proj_out = d_intermediate * act_gp + (d_gate_proj_out,) = dact_pullback(d_intermediate * up_proj_for_bwd) + + # wi bwd (fused gate/up via concat). Mirror the fused fwd: pack the + # gate/up cotangents along the trailing axis, run a single + # grouped_quantize + two grouped_gemm pair (one dgrad, one wgrad) + # against the fused casted_wi_rhs_trans residual, then split the + # wgrad result back into d_wi_0 / d_wi_1 halves with jnp.split. d_combined = jnp.concatenate([d_gate_proj_out, d_up_proj_out], axis=-1) casted_d_combined = tex.grouped_quantize( - d_combined, q_set_w0.dgrad, ctx.local_group_sizes, flatten_axis=-1 + d_combined, q_set.dgrad, local_group_sizes, flatten_axis=-1 ) d_sorted_x = tex.grouped_gemm( casted_d_combined.get_tensor(usage=TensorUsage.LHS), - ctx.casted_wi_rhs_trans, + casted_wi_rhs_trans, contracting_dims=((1,), (2,)), ) d_wi_combined = tex.grouped_gemm( - ctx.casted_sorted_x_lhs_trans, + casted_sorted_x_lhs_trans, casted_d_combined.get_tensor(usage=TensorUsage.RHS), contracting_dims=((0,), (0,)), ) - d_wi_0 = d_wi_combined[..., :inter_M] - d_wi_1 = d_wi_combined[..., inter_M:] - if has_wi_bias: - d_wi_combined_bias = tex.grouped_dbias(d_combined, ctx.local_group_sizes) - d_wi_0_bias = d_wi_combined_bias[..., :inter_M] - d_wi_1_bias = d_wi_combined_bias[..., inter_M:] + d_wi_combined = jnp.where(wgrad_group_active, d_wi_combined, jnp.zeros_like(d_wi_combined)) + d_wi_0, d_wi_1 = jnp.split(d_wi_combined, 2, axis=-1) + if has_bias: + d_wi_combined_bias = tex.grouped_dbias(d_combined, local_group_sizes) + d_wi_0_bias, d_wi_1_bias = jnp.split(d_wi_combined_bias, 2, axis=-1) else: d_wi_0_bias = None d_wi_1_bias = None - # ---------------- Dispatch bwd ---------------- - inputs_2d_shape = (per_shard_x_shape[0] * per_shard_x_shape[1], hidden) - d_inputs_2d = _dispatch_bwd( - d_sorted_x, - ctx.dispatch, - inputs_2d_shape=inputs_2d_shape, - backend=permutation_backend, - ep_active=ep_active, - num_experts=num_experts, - num_experts_per_tok=num_experts_per_tok, - num_real_tokens=_static_shape.num_real_tokens, - padding_size=_static_shape.padding_size, - pre_a2a_buffer_shape=_static_shape.pre_a2a_buffer_shape, - ep_axis=ep_axis, - shard_id=shard_id, - num_ep=num_ep, - ) - d_x_from_dispatch = d_inputs_2d.reshape(per_shard_x_shape) - - # ---------------- Routing bwd ---------------- - # The probs cotangent comes from _combine_bwd. For PURE_JAX it's the - # cotangent of routing_weights (post-routing_map_to_selected_experts); - # we need to bridge back to sparse_probs. For TRITON it's already the - # cotangent of merging_probs == sparse_probs. - if d_routing_weights is not None: - if permutation_backend is PermutationBackend.PURE_JAX: - # routing_map_to_selected_experts: - # selected_experts = argsort(routing_map)[..., -topk:] - # weights = take_along_axis(sparse_probs, selected_experts, axis=-1) - # routing_map is bool (non-diff); the gradient of weights - # w.r.t. sparse_probs is a scatter-into-zero along the - # selected_experts indices. - selected_experts = jnp.argsort(ctx.routing_map, axis=-1)[..., -num_experts_per_tok:] - d_sparse_probs = jnp.zeros_like(ctx.saved_scores).astype(d_routing_weights.dtype) - d_sparse_probs = jnp.take_along_axis(d_sparse_probs, selected_experts, axis=-1) - # Actually scatter: build via jnp.zeros + .at[].set - d_sparse_probs = jnp.zeros(ctx.routing_map.shape, dtype=d_routing_weights.dtype) - d_sparse_probs = d_sparse_probs.at[ - jnp.arange(ctx.routing_map.shape[0])[:, None], selected_experts - ].set(d_routing_weights) - else: - d_sparse_probs = d_routing_weights.astype(jnp.float32) - else: - d_sparse_probs = jnp.zeros(ctx.routing_map.shape, dtype=jnp.float32) - - # Topk bwd primitive: returns d_logits (no d_expert_bias). - d_logits_2d_main = tex.fused_topk_with_score_function_bwd( - ctx.routing_map, - ctx.saved_scores, - d_sparse_probs.astype(ctx.saved_scores.dtype), - topk=num_experts_per_tok, - use_pre_softmax=use_pre_softmax, - scaling_factor=scaling_factor, - score_function=score_function, - compute_aux_scores=False, - ) - - # ---------------- Aux loss bwd ---------------- - if aux_loss_coeff > 0.0: - # Step 1: aux_loss bwd -> d_aux_probs - aux_num_tokens = ctx.aux_logits_for_score.shape[0] - d_aux_probs = tex.fused_moe_aux_loss_bwd( - ctx.aux_const_buf, - ctx.aux_tokens_per_expert.astype(jnp.int32), - d_aux_loss.reshape(()), - num_tokens=aux_num_tokens, - ) - # Step 2: aux-side topk bwd (compute_aux_scores=True path). - # The routing_map argument is ignored in this branch (the kernel - # uses saved_scores); pass any shape-correct integer tensor. - d_aux_logits = tex.fused_topk_with_score_function_bwd( - jnp.zeros(ctx.aux_logits_for_score.shape, dtype=jnp.bool_), - ctx.aux_saved_scores, - d_aux_probs.astype(ctx.aux_saved_scores.dtype), - topk=num_experts_per_tok, - use_pre_softmax=False, - scaling_factor=1.0, - score_function=score_function, - compute_aux_scores=True, - ) - # Step 3: under EP the aux logits were all_gathered along - # ``(ep_axis, *data_parallelism_axes)`` (the latter being FSDP - # axes that shard the batch). The bwd is the inverse of that - # multi-axis tiled all_gather: ``dynamic_slice`` to pick out - # this shard's local rows from the global cotangent. - # - # JAX's convention for tiled ``all_gather(axis_name=(a, b, ...))`` - # is row-major over the tuple: the shard at mesh position - # ``(i_a, i_b, ...)`` writes to rows - # ``[(i_a * size_b * ... + i_b * ... + ...) * local_T : - # + local_T)``. We invert that by computing the same flat - # index here and slicing. - if ep_active: - local_T_aux = ctx.logits_2d.shape[0] - flat_shard = shard_id # ep is the outermost axis in the gather tuple - for ax, sz in zip(data_parallelism_axes, fsdp_sizes): - flat_shard = flat_shard * sz + jax.lax.axis_index(ax) - d_aux_logits_local = jax.lax.dynamic_slice( - d_aux_logits.astype(ctx.logits_2d.dtype), - start_indices=(flat_shard * local_T_aux, 0), - slice_sizes=(local_T_aux, num_experts), - ) - else: - d_aux_logits_local = d_aux_logits.astype(d_logits_2d_main.dtype) - d_logits_2d = d_logits_2d_main + d_aux_logits_local.astype(d_logits_2d_main.dtype) - else: - d_logits_2d = d_logits_2d_main - - # ---------------- Gate bwd ---------------- - d_gate_logits = d_logits_2d.reshape(per_shard_x_shape[0], per_shard_x_shape[1], num_experts) - gate_kernel_cast = ctx.gate_kernel.astype(ctx.x.dtype) - d_x_from_gate = jnp.einsum("bse,he->bsh", d_gate_logits, gate_kernel_cast) - d_gate_kernel = jnp.einsum("bsh,bse->he", ctx.x, d_gate_logits).astype(ctx.gate_kernel.dtype) - d_x = d_x_from_gate + d_x_from_dispatch - - # Reduce per-rank partial contributions to match the out_specs - # declared by _build_grads_specs: - # gate_kernel : P() -> psum across (ep, *fsdp) - # wi_0/wi_1/wo : P(ep_axis, ...) -> psum across (*fsdp) only - # inputs : P((ep, fsdp), ...) -> already shard-local, no reduction - if ep_active: - replicate_all = (ep_axis,) + tuple(data_parallelism_axes) - d_gate_kernel = jax.lax.psum(d_gate_kernel, axis_name=replicate_all) - if data_parallelism_axes: - replicate_fsdp = tuple(data_parallelism_axes) - d_wi_0 = jax.lax.psum(d_wi_0, axis_name=replicate_fsdp) - d_wi_1 = jax.lax.psum(d_wi_1, axis_name=replicate_fsdp) - d_wo = jax.lax.psum(d_wo, axis_name=replicate_fsdp) - if has_wi_bias: - d_wi_0_bias = jax.lax.psum(d_wi_0_bias, axis_name=replicate_fsdp) - d_wi_1_bias = jax.lax.psum(d_wi_1_bias, axis_name=replicate_fsdp) - if has_wo_bias: - d_wo_bias = jax.lax.psum(d_wo_bias, axis_name=replicate_fsdp) - - grads: dict = { - "inputs": d_x, - "gate_kernel": d_gate_kernel, - "wi_0": d_wi_0, - "wi_1": d_wi_1, - "wo": d_wo, - } - if has_wi_bias: - grads["wi_0_bias"] = d_wi_0_bias - grads["wi_1_bias"] = d_wi_1_bias - if has_wo_bias: - grads["wo_bias"] = d_wo_bias - if has_expert_bias: - # expert_bias has no gradient through topk (the topk bwd returns - # None for it). Emit a structural zero so the outer rule has - # something to package. - grads["expert_bias"] = jnp.zeros_like(ctx.expert_bias) - return grads - - -# ============================================================================= -# Spec builders for shard_map (lockstep with ctx_dict / captured_dict) -# ============================================================================= - - -def _build_in_specs( - ep_axis: str, - batch_pspec_axis: Any, - *, - has_bias: bool, - has_expert_bias: bool, -) -> dict: - """Build the ``in_specs`` dict for the EP fwd shard_map.""" - specs: dict = { - "inputs": P(batch_pspec_axis, None, None), - "gate_kernel": P(), - "wi_0": P(ep_axis, None, None), - "wi_1": P(ep_axis, None, None), - "wo": P(ep_axis, None, None), - } - if has_bias: - for name in ("wi_0_bias", "wi_1_bias", "wo_bias"): - specs[name] = P(ep_axis, None) - if has_expert_bias: - specs["expert_bias"] = P(ep_axis) - return specs - - -def _build_dispatch_specs( # pylint: disable=unused-argument - ep_axis: str, - *, - backend: PermutationBackend, - ep_active: bool, - align_size: int, -) -> _DispatchState: - """Build the shard_map ``out_specs`` for the dispatch state. - - Returns a :data:`_DispatchState` (either :class:`_PureJaxDispatchState` - or :class:`_TritonDispatchState`) whose fields are - :class:`PartitionSpec` placeholders. Optional fields are set to - ``P()`` when populated by :func:`_dispatch` and to ``None`` when - intentionally omitted, so the spec's pytree structure mirrors the - value's structure leaf-for-leaf. - """ - ep_all = P() if ep_active else None - ep_local = P() if ep_active else None - if backend is PermutationBackend.PURE_JAX: - return _PureJaxDispatchState( - group_sizes=P(), - sorted_indices=P(), - routing_weights=P(), - all_shards_tokens_per_expert=ep_all, - local_perm_row_id_map=ep_local, - ) - return _TritonDispatchState( - group_sizes=P(), - row_id_map=P(), - pad_offsets=P() if align_size > 0 else None, - merging_probs=P(), - all_shards_tokens_per_expert=ep_all, - local_perm_row_id_map=ep_local, - ) - - -def _build_ctx_specs( # pylint: disable=unused-argument - ep_axis: str, - batch_pspec_axis: Any, - *, - backend: PermutationBackend, - ep_active: bool, - has_bias: bool, - has_expert_bias: bool, - aux_loss_enabled: bool, - align_size: int, -) -> _BodyCtx: - """Build the spec :class:`_BodyCtx` mirroring :func:`_body_fwd`'s ctx. - - Fields gated off by the static config (``expert_bias``, ``aux_*``) - are ``None`` here so the spec pytree matches the value pytree - leaf-for-leaf. - """ - return _BodyCtx( - # Per-shard local activations along the batch axis. - x=P(batch_pspec_axis, None, None), - gate_kernel=P(), - logits_2d=P(batch_pspec_axis, None), - saved_scores=P(batch_pspec_axis, None), - routing_map=P(batch_pspec_axis, None), - dispatch=_build_dispatch_specs( - ep_axis, backend=backend, ep_active=ep_active, align_size=align_size - ), - # FFN residuals: the LHS_TRANS / RHS_TRANS variants of - # grouped_quantize have leading "rows"/"experts" dims that are - # already shard-local (post-dispatch). Use P(ep_axis,...) on - # leading dim; that works whether the leaf is a plain ndarray - # or a ScaledTensor (shard_map applies the spec leaf-wise to - # the registered ScaledTensor pytree). - casted_sorted_x_lhs_trans=P(), - casted_wi_rhs_trans=P(ep_axis, None, None), - gate_proj_out=P(), - up_proj_out=P(), - casted_intermediate_lhs_trans=P(), - casted_wo_rhs_trans=P(ep_axis, None, None), - expert_outputs=P(), - local_group_sizes=P(), - expert_bias=P(ep_axis) if has_expert_bias else None, - aux_const_buf=P() if aux_loss_enabled else None, - aux_tokens_per_expert=P() if aux_loss_enabled else None, - aux_logits_for_score=P() if aux_loss_enabled else None, - aux_saved_scores=P() if aux_loss_enabled else None, - ) - - -def _build_grads_specs( - ep_axis: str, - batch_pspec_axis: Any, - *, - has_bias: bool, - has_expert_bias: bool, -) -> dict: - """Spec dict for the grads dict returned by :func:`_body_bwd`.""" - return _build_in_specs( - ep_axis, - batch_pspec_axis, - has_bias=has_bias, - has_expert_bias=has_expert_bias, + d_sorted_x_3d = d_sorted_x.reshape(1, d_sorted_x.shape[0], d_sorted_x.shape[1]) + d_recv_w_3d = d_recv_w_from_intermediate.reshape(1, -1) + return ( + d_sorted_x_3d, + d_recv_w_3d, + d_wi_0, + d_wi_1, + d_wo, + d_wi_0_bias, + d_wi_1_bias, + d_wo_bias, ) # ============================================================================= -# Top-level VJP rules +# Full fwd / bwd rules (custom_vjp halves) # ============================================================================= -def _moe_fwd_rule( # pylint: disable=unused-argument - # Args MUST match the positional order of ``_moe`` (diff first, - # then nondiff). See ``_moe_bwd_rule`` for the opposite convention. +def _moe_fwd_rule( x, gate_kernel, wi_0, @@ -1687,170 +467,328 @@ def _moe_fwd_rule( # pylint: disable=unused-argument group_topk, scaling_factor, aux_loss_coeff, - permutation_backend, - align_size, - gate_inside_vjp, ep_axis, data_parallelism_axes, input_axes, gate_kernel_axes, wi_kernel_axes, wo_kernel_axes, - quantizer_sets, dtype, + apply_topk_weights_early, ): - x = with_sharding_constraint_by_logical_axes(x, input_axes) - ep_active = ep_axis is not None - body_kwargs = { - "num_experts": num_experts, - "num_experts_per_tok": num_experts_per_tok, - "activation_type": activation_type, - "score_function": score_function, - "use_pre_softmax": use_pre_softmax, - "num_groups": num_groups, - "group_topk": group_topk, - "scaling_factor": scaling_factor, - "aux_loss_coeff": aux_loss_coeff, - "permutation_backend": permutation_backend, - "align_size": align_size, - "gate_inside_vjp": gate_inside_vjp, - "quantizer_sets": quantizer_sets, - "dtype": dtype, - "ep_axis": ep_axis, - "data_parallelism_axes": data_parallelism_axes, - } - captured: dict = { - "inputs": x, - "gate_kernel": gate_kernel, - "wi_0": wi_0, - "wi_1": wi_1, - "wo": wo, - } - has_bias = wi_0_bias is not None - has_expert_bias = expert_bias is not None - if has_bias: - captured["wi_0_bias"] = wi_0_bias - captured["wi_1_bias"] = wi_1_bias - captured["wo_bias"] = wo_bias - if has_expert_bias: - captured["expert_bias"] = expert_bias - - if not ep_active: - output, aux_loss, ctx = _body_fwd( - captured, - **body_kwargs, - ep_active=False, - fsdp_sizes=(), - num_ep=1, - num_experts_local=num_experts, - recv_buffer_rows=0, - ) - # Carry static side info to the bwd rule alongside ctx. These - # are Python ints/bools/tuples (NOT pytree leaves), so we - # bundle them as a plain dict rather than putting them on the - # ``_BodyCtx`` NamedTuple where shard_map would try to flatten - # them into JitTracers. - static = { - "has_wi_bias": has_bias, - "has_wo_bias": has_bias, - "has_expert_bias": has_expert_bias, - "x_shape": x.shape, - "num_experts_local": num_experts, - "recv_buffer_rows": 0, - } - return (output, aux_loss), (ctx, static) - - # ---------------- EP path ---------------- + """Forward: gate -> topk -> ep_dispatch -> shard_map(FFN) -> ep_combine. + + Returns ``(output, aux_loss)``. ``aux_loss`` is a zero scalar when + ``aux_loss_coeff == 0``. + """ + del gate_kernel_axes, wi_kernel_axes, wo_kernel_axes # used in bwd only from jax.experimental.shard_map import shard_map + x = with_sharding_constraint_by_logical_axes(x, input_axes) + mesh = _get_mesh() if mesh is None or mesh.empty: - raise ValueError("moe(...) requires an active jax.sharding.Mesh when ep_axis is set.") + raise ValueError("moe(...) requires an active jax.sharding.Mesh.") + if ep_axis is None: + raise ValueError("moe(...) requires ep_axis to be set (TE EP backend).") num_ep = mesh.shape[ep_axis] if num_experts % num_ep != 0: raise ValueError(f"num_experts={num_experts} must be divisible by EP size={num_ep}") - num_experts_local = num_experts // num_ep + num_local_experts = num_experts // num_ep - # Reject overlapping EP / FSDP axes. Listing ep_axis in - # data_parallelism_axes would produce a duplicate-axis PartitionSpec - # ((ep, ep, ...)) which JAX rejects, and would also double-count - # num_ep in dp_size (under-sizing recv_buffer_rows by a factor of - # num_ep). Catch it up front with a clear error. + dp_size = 1 for ax in data_parallelism_axes: - if ax not in mesh.shape: - raise ValueError( - f"data_parallelism_axes contains {ax!r} but mesh has" - f" axes {tuple(mesh.shape.keys())}" - ) - if ax == ep_axis: - raise ValueError( - f"data_parallelism_axes={data_parallelism_axes!r} contains the EP" - f" axis {ep_axis!r}; EP is implicit in the batch sharding and must" - " not also be listed as a data-parallel axis." - ) + dp_size *= mesh.shape[ax] + num_procs = num_ep * dp_size + + B, S, H = x.shape + K = num_experts_per_tok + if B % num_procs != 0: + raise ValueError(f"batch={B} not divisible by ep*dp={num_procs}") + + # Per-rank send capacity: B/num_procs rows x S tokens per rank. + max_tokens_per_rank = (B // num_procs) * S + # Per-rank receive capacity. NCCL EP HT expert-major lays out variable + # per-expert zones in one flat recv buffer, with each non-empty zone padded + # to ``dispatch_output_per_expert_alignment``. + tokens_per_ep_group = num_ep * max_tokens_per_rank + max_local_assignments = tokens_per_ep_group * min(K, num_local_experts) + max_nonempty_experts = min(num_local_experts, max_local_assignments) + padded_total_bound = max_local_assignments + (_ALIGN_SIZE - 1) * max_nonempty_experts + aligned_total_bound = ((padded_total_bound + _ALIGN_SIZE - 1) // _ALIGN_SIZE) * _ALIGN_SIZE + per_expert_bound = ( + num_local_experts * ((tokens_per_ep_group + _ALIGN_SIZE - 1) // _ALIGN_SIZE) * _ALIGN_SIZE + ) + recv_pr = min(per_expert_bound, aligned_total_bound) + + _te_ep_assert_compatible_bootstrap( + num_experts=num_experts, + max_tokens_per_rank=max_tokens_per_rank, + recv_capacity_per_rank=recv_pr, + hidden_dim=H, + ep_size=num_ep, + ) if not data_parallelism_axes: batch_pspec_axis: Any = ep_axis else: - batch_pspec_axis = (ep_axis, *data_parallelism_axes) - dp_size = 1 - for ax in data_parallelism_axes: - dp_size *= mesh.shape[ax] + # ep must be innermost: ep_bootstrap forms NCCL EP comms from + # consecutive global ranks (dp_color = rank // ep_size), so the + # comm only stays within one model replica under (outer_dp, ep). + batch_pspec_axis = (*data_parallelism_axes, ep_axis) + ep3_spec = P(batch_pspec_axis, None, None) + ep2_spec = P(batch_pspec_axis, None) + x = jax.lax.with_sharding_constraint(x, NamedSharding(mesh, ep3_spec)) + + # ---------------- Gate (global view) ---------------- + # tex.fused_topk_with_score_function is only validated against its + # pytorch reference at fp32 (see tests/pytorch/test_fused_router.py: + # parametrize gates dtype on torch.float32 only; the tolerance helper + # raises NotImplementedError for any other dtype). Keeping logits in + # the activation dtype (e.g. bf16) lets sigmoid / softmax / topk + # accumulate at low precision and silently produce NaNs on tokens + # whose normalised weights underflow. Cast to fp32 here to stay in + # the validated regime. + gate_kernel_cast = gate_kernel.astype(x.dtype) + gate_logits = jnp.einsum("bsh,he->bse", x, gate_kernel_cast) + logits_2d = gate_logits.reshape(-1, num_experts).astype(jnp.float32) + + # ---------------- Routing (global view) ---------------- + # expert_bias is an empty (shape-(0,)) sentinel when the caller did + # not enable it; the primitive treats that as "no bias". + eb_arg = expert_bias if expert_bias.shape != (0,) else jnp.zeros((0,), dtype=jnp.float32) + sparse_probs, routing_map, saved_scores = tex.fused_topk_with_score_function_fwd( + logits_2d, + topk=K, + use_pre_softmax=use_pre_softmax, + num_groups=-1 if num_groups is None else num_groups, + group_topk=-1 if group_topk is None else group_topk, + scaling_factor=scaling_factor, + score_function=score_function, + expert_bias=eb_arg, + compute_aux_scores=False, + ) + sparse_probs = sparse_probs.astype(dtype) - global_batch_size, sequence_length, _hidden = x.shape - topk = num_experts_per_tok - if global_batch_size % (num_ep * dp_size) != 0: - raise ValueError(f"batch={global_batch_size} not divisible by ep*dp={num_ep * dp_size}") - recv_buffer_rows = (global_batch_size // dp_size) * sequence_length * topk - if align_size > 0: - recv_buffer_rows += num_experts * (align_size - 1) + # ---------------- Aux loss (global view, replicated) ---------------- + # ``fused_moe_aux_loss_fwd`` sums probs and tokens_per_expert across + # all tokens, which is wrong when T is sharded. Force-replicate the + # gate logits and recompute the routing map at global view so the + # kernel sees a complete [T_global, E] tensor. The replication is a + # single all-gather over (*dp, ep) and lives off the dispatch + # critical path. + if aux_loss_coeff > 0.0: + global_logits_2d = jax.lax.with_sharding_constraint(logits_2d, NamedSharding(mesh, P())) + _, global_routing_map, _ = tex.fused_topk_with_score_function_fwd( + global_logits_2d, + topk=K, + use_pre_softmax=use_pre_softmax, + num_groups=-1 if num_groups is None else num_groups, + group_topk=-1 if group_topk is None else group_topk, + scaling_factor=scaling_factor, + score_function=score_function, + expert_bias=eb_arg, + compute_aux_scores=False, + ) + aux_tokens_per_expert = jnp.sum(global_routing_map.astype(jnp.int32), axis=0) + # compute_aux_scores=True takes a separate kernel path: clean + # per-expert softmax, no grouping / bias / scaling. + aux_probs, _aux_rm, aux_saved_scores = tex.fused_topk_with_score_function_fwd( + global_logits_2d.astype(jnp.float32), + topk=K, + use_pre_softmax=False, + num_groups=-1, + group_topk=-1, + scaling_factor=1.0, + score_function=score_function, + expert_bias=jnp.zeros((0,), dtype=jnp.float32), + compute_aux_scores=True, + ) + aux_loss, aux_const_buf = tex.fused_moe_aux_loss_fwd( + aux_probs.astype(jnp.float32), + aux_tokens_per_expert.astype(jnp.int32), + topk=K, + coeff=aux_loss_coeff, + ) + aux_loss = aux_loss.astype(dtype) + else: + aux_loss = jnp.zeros((), dtype=dtype) + aux_const_buf = None + aux_tokens_per_expert = None + aux_saved_scores = None - in_specs = _build_in_specs( - ep_axis, - batch_pspec_axis, - has_bias=has_bias, - has_expert_bias=has_expert_bias, + # ---------------- Routing -> (topk_idx, topk_w) at 3D ---------------- + # argsort on a bool tensor places True last (False=0 < True=1), so the + # last K indices are the selected expert IDs. + selected_experts = jnp.argsort(routing_map, axis=-1)[..., -K:] + routing_weights = jnp.take_along_axis(sparse_probs, selected_experts, axis=-1) + topk_idx_3d = selected_experts.reshape(B, S, K).astype(jnp.int32) + topk_w_3d = routing_weights.reshape(B, S, K).astype(jnp.float32) + # tex.ep_prepare/dispatch's partition only folds ep_axis into a replicated + # leading dim, not the outer dp/fsdp axes, so a replicated topk_idx makes + # each rank see B/ep rows (not B/num_procs) and overrun the bootstrap-sized + # send buffer. Pin both routing tensors to the (outer, ep) leading sharding + # so per-rank token counts match max_tokens_per_rank. + topk_idx_3d = jax.lax.with_sharding_constraint(topk_idx_3d, NamedSharding(mesh, ep3_spec)) + topk_w_3d = jax.lax.with_sharding_constraint(topk_w_3d, NamedSharding(mesh, ep3_spec)) + + # ---------------- TE EP dispatch (global view) ---------------- + cfg = tex.EpLayerConfig( + top_k=K, + dispatch_output_per_expert_alignment=_ALIGN_SIZE, ) - output_spec = P(batch_pspec_axis, None, None) - aux_spec = P() - ctx_spec = _build_ctx_specs( - ep_axis, - batch_pspec_axis, - backend=permutation_backend, - ep_active=True, - has_bias=has_bias, - has_expert_bias=has_expert_bias, - aux_loss_enabled=(aux_loss_coeff > 0.0), - align_size=align_size, + token_counts, handle_mem = tex.ep_prepare(cfg, topk_idx_3d) + recv_tokens, recv_topk_weights = tex.ep_dispatch_fwd( + cfg, handle_mem, topk_idx_3d, x, topk_w_3d, recv_pr + ) + recv_tokens = jax.lax.with_sharding_constraint(recv_tokens, NamedSharding(mesh, ep3_spec)) + recv_topk_weights = jax.lax.with_sharding_constraint( + recv_topk_weights, NamedSharding(mesh, ep2_spec) + ) + + # ---------------- FFN (per-shard via shard_map) ---------------- + has_bias = wi_0_bias is not None + kernel_spec = P(ep_axis, None, None) + bias_spec = P(ep_axis, None) if has_bias else None + # token_counts is the per-shard (1, num_local_experts) padded + # per-expert count from ep_prepare; piped into _ffn_fwd_per_shard + # as the grouped_gemm group_sizes so cuBLAS skips both 0-token + # experts and the trailing overalloc tail. + ffn_in_specs = (ep3_spec, ep2_spec, ep2_spec, kernel_spec, kernel_spec, kernel_spec) + ffn_in_args = [recv_tokens, recv_topk_weights, token_counts, wi_0, wi_1, wo] + if has_bias: + ffn_in_specs = ffn_in_specs + (bias_spec, bias_spec, bias_spec) + ffn_in_args.extend([wi_0_bias, wi_1_bias, wo_bias]) + + # FFN residuals live entirely on the local ep rank, so the leading + # "experts" / "rows" dims map to P() (already shard-local). wi is + # fused via jnp.concatenate along the trailing (output) axis + # (see _ffn_fwd_per_shard for rationale), so the residual is a + # single 3D casted_wi_rhs_trans of shape + # (num_local_experts, hidden, 2*H_inter). local_group_sizes is + # now per-shard dynamic (= per-shard token_counts), so its + # residual spec mirrors ep2_spec (one row per ep rank). + residuals_spec = ( + P(), # casted_sorted_x_lhs_trans + P(ep_axis, None, None), # casted_wi_rhs_trans + P(), # gate_proj_out + P(), # up_proj_out + P(), # casted_intermediate_lhs_trans + P(ep_axis, None, None), # casted_wo_rhs_trans + ep2_spec, # local_group_sizes (1, num_local_experts) per shard ) + out_specs = (ep3_spec, residuals_spec) - _fsdp_sizes: Tuple[int, ...] = tuple(mesh.shape[ax] for ax in data_parallelism_axes) - - def _shardmap_body(captured_local): - return _body_fwd( - captured_local, - **body_kwargs, - ep_active=True, - fsdp_sizes=_fsdp_sizes, - num_ep=num_ep, - num_experts_local=num_experts_local, - recv_buffer_rows=recv_buffer_rows, + def _body(*args): + if has_bias: + (r_tok, r_w, tc, w0, w1, w_o, w0b, w1b, wob) = args + else: + (r_tok, r_w, tc, w0, w1, w_o) = args + w0b = w1b = wob = None + # NOTE: tex.ep_dispatch_fwd's NCCL EP HT path leaves the recv + # buffer uninitialised on fully-empty-receiver ranks (and at + # padded slots on partially-loaded ranks). We don't need a + # zero-init guard here anymore because: + # 1. ``tc`` (per-expert padded counts) is plumbed into + # grouped_gemm as group_sizes, so cuBLAS skips both + # 0-token experts and the trailing overalloc tail. + # 2. The per-group wgrad masks in _ffn_bwd_per_shard zero + # ``d_wo`` / ``d_wi_combined`` slices for 0-token-globally + # experts (cuBLAS skips size_g==0 groups without + # zero-filling, which would otherwise leak NaN into the + # user's optimizer). + # 3. All other downstream consumers (ep_combine, + # ep_dispatch_bwd) are handle_mem-aware and read only + # valid positions. + # If a future caller adds a non-group-aware reader of r_tok + # (e.g. an inspect probe over the full recv tile), re-add the + # ``jax.lax.cond(jnp.any(r_w != 0), identity, zeros_like)`` + # guard here. + return _ffn_fwd_per_shard( + r_tok, + r_w, + tc, + w0, + w1, + w_o, + w0b, + w1b, + wob, + num_local_experts=num_local_experts, + activation_type=activation_type, + apply_topk_weights_early=apply_topk_weights_early, ) - output, aux_loss, ctx = shard_map( - _shardmap_body, + expert_outputs, ffn_residuals = shard_map( + _body, mesh=mesh, - in_specs=(in_specs,), - out_specs=(output_spec, aux_spec, ctx_spec), + in_specs=ffn_in_specs, + out_specs=out_specs, check_rep=False, - )(captured) + )(*ffn_in_args) + expert_outputs = jax.lax.with_sharding_constraint(expert_outputs, NamedSharding(mesh, ep3_spec)) + + # ---------------- TE EP combine (global view) ---------------- + out_partition_spec = (batch_pspec_axis, None, None) + if apply_topk_weights_early: + # expert_outputs is already weighted upstream. + output = tex.ep_combine_fwd( + cfg, + handle_mem, + expert_outputs, + num_local_tokens=(B, S), + out_partition_spec=out_partition_spec, + ) + else: + # HT combine is unweighted; apply routing weights before calling it. + # Padded recv slots are ignored by combine via handle_mem metadata. + w = recv_topk_weights[..., None].astype(expert_outputs.dtype) + weighted = expert_outputs * w + output = tex.ep_combine_fwd( + cfg, + handle_mem, + weighted, + num_local_tokens=(B, S), + out_partition_spec=out_partition_spec, + ) + + ( + casted_sorted_x_lhs_trans, + casted_wi_rhs_trans, + gate_proj_out, + up_proj_out, + casted_intermediate_lhs_trans, + casted_wo_rhs_trans, + local_group_sizes, + ) = ffn_residuals + + ctx = _Ctx( + x=x, + gate_kernel=gate_kernel, + expert_bias=expert_bias, + logits_2d=logits_2d, + saved_scores=saved_scores, + routing_map=routing_map, + cfg=cfg, + handle_mem=handle_mem, + token_counts=token_counts, + recv_topk_weights=recv_topk_weights, + casted_sorted_x_lhs_trans=casted_sorted_x_lhs_trans, + casted_wi_rhs_trans=casted_wi_rhs_trans, + gate_proj_out=gate_proj_out, + up_proj_out=up_proj_out, + casted_intermediate_lhs_trans=casted_intermediate_lhs_trans, + casted_wo_rhs_trans=casted_wo_rhs_trans, + expert_outputs=expert_outputs, + local_group_sizes=local_group_sizes, + aux_const_buf=aux_const_buf, + aux_tokens_per_expert=aux_tokens_per_expert, + aux_saved_scores=aux_saved_scores, + ) static = { - "has_wi_bias": has_bias, - "has_wo_bias": has_bias, - "has_expert_bias": has_expert_bias, + "has_bias": has_bias, "x_shape": x.shape, - "num_experts_local": num_experts_local, - "recv_buffer_rows": recv_buffer_rows, + "recv_pr": recv_pr, } return (output, aux_loss), (ctx, static) @@ -1865,128 +803,259 @@ def _moe_bwd_rule( group_topk, scaling_factor, aux_loss_coeff, - permutation_backend, - align_size, - gate_inside_vjp, ep_axis, data_parallelism_axes, input_axes, gate_kernel_axes, wi_kernel_axes, wo_kernel_axes, - quantizer_sets, dtype, - ctx, - dy_pair, + apply_topk_weights_early, + residuals, + cotangents, ): - ctx, static = ctx # split tensor residuals from static side info - has_wi_bias = static["has_wi_bias"] - has_wo_bias = static["has_wo_bias"] - has_expert_bias = static["has_expert_bias"] - x_shape = static["x_shape"] - num_experts_local = static["num_experts_local"] - recv_buffer_rows = static["recv_buffer_rows"] + """Backward mirror of :func:`_moe_fwd_rule`.""" + del num_groups, group_topk, dtype # captured in residuals / unused in bwd + from jax.experimental.shard_map import shard_map - ep_active = ep_axis is not None - mesh = _get_mesh() if ep_active else None - fsdp_sizes: Tuple[int, ...] = ( - tuple(mesh.shape[ax] for ax in data_parallelism_axes) if ep_active else () - ) - body_kwargs = { - "num_experts": num_experts, - "num_experts_per_tok": num_experts_per_tok, - "activation_type": activation_type, - "score_function": score_function, - "use_pre_softmax": use_pre_softmax, - "num_groups": num_groups, - "group_topk": group_topk, - "scaling_factor": scaling_factor, - "aux_loss_coeff": aux_loss_coeff, - "permutation_backend": permutation_backend, - "align_size": align_size, - "gate_inside_vjp": gate_inside_vjp, - "quantizer_sets": quantizer_sets, - "dtype": dtype, - "ep_axis": ep_axis, - "data_parallelism_axes": data_parallelism_axes, - "fsdp_sizes": fsdp_sizes, - "num_ep": 1 if not ep_active else mesh.shape[ep_axis], - "num_experts_local": num_experts_local, - "recv_buffer_rows": recv_buffer_rows, - "has_wi_bias": has_wi_bias, - "has_wo_bias": has_wo_bias, - "has_expert_bias": has_expert_bias, - "x_shape": x_shape, - } + d_output, d_aux_loss = cotangents - if not ep_active: - grads = _body_bwd(ctx, dy_pair, ep_active=False, **body_kwargs) - # Apply sharding constraints on grads. - grads["gate_kernel"] = with_sharding_constraint_by_logical_axes( - grads["gate_kernel"], gate_kernel_axes - ) - grads["wi_0"] = with_sharding_constraint_by_logical_axes(grads["wi_0"], wi_kernel_axes) - grads["wi_1"] = with_sharding_constraint_by_logical_axes(grads["wi_1"], wi_kernel_axes) - grads["wo"] = with_sharding_constraint_by_logical_axes(grads["wo"], wo_kernel_axes) - grads["inputs"] = with_sharding_constraint_by_logical_axes(grads["inputs"], input_axes) - return _grads_dict_to_tuple(grads, has_wi_bias, has_wo_bias, has_expert_bias) + ctx, static = residuals + has_bias = static["has_bias"] + x_shape = static["x_shape"] + recv_pr = static["recv_pr"] - from jax.experimental.shard_map import shard_map + mesh = _get_mesh() + if mesh is None or mesh.empty: + raise ValueError("moe(...) requires an active jax.sharding.Mesh.") + dp_size = 1 + for ax in data_parallelism_axes: + dp_size *= mesh.shape[ax] + B, S, _ = x_shape + K = num_experts_per_tok if not data_parallelism_axes: batch_pspec_axis: Any = ep_axis else: - batch_pspec_axis = (ep_axis, *data_parallelism_axes) - ctx_spec = _build_ctx_specs( - ep_axis, - batch_pspec_axis, - backend=permutation_backend, - ep_active=True, - has_bias=has_wi_bias, - has_expert_bias=has_expert_bias, - aux_loss_enabled=(aux_loss_coeff > 0.0), - align_size=align_size, + batch_pspec_axis = (*data_parallelism_axes, ep_axis) + ep3_spec = P(batch_pspec_axis, None, None) + ep2_spec = P(batch_pspec_axis, None) + out_partition_spec = (batch_pspec_axis, None, None) + + # ---------------- Combine bwd (global view) ---------------- + d_output = jax.lax.with_sharding_constraint(d_output, NamedSharding(mesh, ep3_spec)) + grad_pre_combine = tex.ep_combine_bwd(ctx.cfg, ctx.handle_mem, d_output, recv_pr) + grad_pre_combine = jax.lax.with_sharding_constraint( + grad_pre_combine, NamedSharding(mesh, ep3_spec) ) - dy_specs = (P(batch_pspec_axis, None, None), P()) - grads_spec = _build_grads_specs( - ep_axis, batch_pspec_axis, has_bias=has_wi_bias, has_expert_bias=has_expert_bias + + if apply_topk_weights_early: + # combine_fwd consumed already-weighted expert_outputs; the recv_w + # cotangent flows through the early-weighting step inside the FFN bwd. + d_expert_outputs = grad_pre_combine + d_recv_w_from_combine = jnp.zeros_like(ctx.recv_topk_weights) + else: + # Reverse the late-weighting multiply. Padded expert-major rows are + # part of the physical grouped-GEMM ranges, so write literal zero + # cotangents for inactive rows instead of relying on NaN * 0. + w = ctx.recv_topk_weights[..., None].astype(grad_pre_combine.dtype) + mask_bool = (ctx.recv_topk_weights != 0)[..., None] + d_expert_outputs = jnp.where( + mask_bool, grad_pre_combine * w, jnp.zeros_like(grad_pre_combine) + ) + d_recv_w_from_combine = (grad_pre_combine * ctx.expert_outputs).sum(axis=-1) + d_recv_w_from_combine = d_recv_w_from_combine.astype(ctx.recv_topk_weights.dtype) + + # ---------------- FFN bwd (per-shard via shard_map) ---------------- + kernel_spec = P(ep_axis, None, None) + bias_spec = P(ep_axis, None) if has_bias else None + + bwd_in_specs = ( + ep3_spec, # d_expert_outputs + P(), # casted_sorted_x_lhs_trans + P(ep_axis, None, None), # casted_wi_rhs_trans + P(), # gate_proj_out + P(), # up_proj_out + P(), # casted_intermediate_lhs_trans + P(ep_axis, None, None), # casted_wo_rhs_trans + ep2_spec, # local_group_sizes (1, num_local_experts) per shard + ep2_spec, # recv_topk_weights + ) + bwd_in_args = [ + d_expert_outputs, + ctx.casted_sorted_x_lhs_trans, + ctx.casted_wi_rhs_trans, + ctx.gate_proj_out, + ctx.up_proj_out, + ctx.casted_intermediate_lhs_trans, + ctx.casted_wo_rhs_trans, + ctx.local_group_sizes, + ctx.recv_topk_weights, + ] + bwd_out_specs = ( + ep3_spec, # d_sorted_x + ep2_spec, # d_recv_w_from_intermediate + kernel_spec, # d_wi_0 + kernel_spec, # d_wi_1 + kernel_spec, # d_wo + bias_spec if has_bias else None, # d_wi_0_bias + bias_spec if has_bias else None, # d_wi_1_bias + bias_spec if has_bias else None, # d_wo_bias ) - def _bwd_body(ctx_local, dy_local): - return _body_bwd(ctx_local, dy_local, ep_active=True, **body_kwargs) + def _bwd_body(*args): + ( + d_sorted_x_3d, + d_recv_w_3d, + d_wi_0, + d_wi_1, + d_wo, + d_wi_0_bias, + d_wi_1_bias, + d_wo_bias, + ) = _ffn_bwd_per_shard( + *args, + activation_type=activation_type, + apply_topk_weights_early=apply_topk_weights_early, + has_bias=has_bias, + ) + # Weight grads accumulate per-DP-shard inside the body; psum across + # DP axes so each replica sees the full sum (matches out_specs + # P(ep_axis, ...) which is DP-replicated). + if data_parallelism_axes: + dp = tuple(data_parallelism_axes) + d_wi_0 = jax.lax.psum(d_wi_0, axis_name=dp) + d_wi_1 = jax.lax.psum(d_wi_1, axis_name=dp) + d_wo = jax.lax.psum(d_wo, axis_name=dp) + if has_bias: + d_wi_0_bias = jax.lax.psum(d_wi_0_bias, axis_name=dp) + d_wi_1_bias = jax.lax.psum(d_wi_1_bias, axis_name=dp) + d_wo_bias = jax.lax.psum(d_wo_bias, axis_name=dp) + return ( + d_sorted_x_3d, + d_recv_w_3d, + d_wi_0, + d_wi_1, + d_wo, + d_wi_0_bias, + d_wi_1_bias, + d_wo_bias, + ) - grads = shard_map( + ( + d_sorted_x, + d_recv_w_from_intermediate, + d_wi_0, + d_wi_1, + d_wo, + d_wi_0_bias, + d_wi_1_bias, + d_wo_bias, + ) = shard_map( _bwd_body, mesh=mesh, - in_specs=(ctx_spec, dy_specs), - out_specs=grads_spec, + in_specs=bwd_in_specs, + out_specs=bwd_out_specs, check_rep=False, - )(ctx, dy_pair) + )( + *bwd_in_args + ) - grads["gate_kernel"] = with_sharding_constraint_by_logical_axes( - grads["gate_kernel"], gate_kernel_axes + d_recv_w_total = d_recv_w_from_combine + d_recv_w_from_intermediate + + # ---------------- Dispatch bwd (global view) ---------------- + d_sorted_x = jax.lax.with_sharding_constraint(d_sorted_x, NamedSharding(mesh, ep3_spec)) + d_recv_w_total = jax.lax.with_sharding_constraint(d_recv_w_total, NamedSharding(mesh, ep2_spec)) + d_x_from_dispatch, d_topk_w = tex.ep_dispatch_bwd( + ctx.cfg, + ctx.handle_mem, + d_sorted_x, + d_recv_w_total, + num_local_tokens=(B, S), + out_partition_spec=out_partition_spec, ) - grads["wi_0"] = with_sharding_constraint_by_logical_axes(grads["wi_0"], wi_kernel_axes) - grads["wi_1"] = with_sharding_constraint_by_logical_axes(grads["wi_1"], wi_kernel_axes) - grads["wo"] = with_sharding_constraint_by_logical_axes(grads["wo"], wo_kernel_axes) - grads["inputs"] = with_sharding_constraint_by_logical_axes(grads["inputs"], input_axes) - return _grads_dict_to_tuple(grads, has_wi_bias, has_wo_bias, has_expert_bias) + # ---------------- Routing bwd (global view) ---------------- + # The cotangent on routing_weights is a sparse scatter into sparse_probs + # at the selected_experts indices. + selected_experts = jnp.argsort(ctx.routing_map, axis=-1)[..., -K:] + d_topk_w_flat = d_topk_w.reshape(-1, K) + d_sparse_probs = jnp.zeros(ctx.routing_map.shape, dtype=d_topk_w_flat.dtype) + d_sparse_probs = d_sparse_probs.at[ + jnp.arange(ctx.routing_map.shape[0])[:, None], selected_experts + ].set(d_topk_w_flat) + + d_logits_2d = tex.fused_topk_with_score_function_bwd( + ctx.routing_map, + ctx.saved_scores, + d_sparse_probs.astype(ctx.saved_scores.dtype), + topk=K, + use_pre_softmax=use_pre_softmax, + scaling_factor=scaling_factor, + score_function=score_function, + compute_aux_scores=False, + ) + + # ---------------- Aux loss bwd (global view, replicated) ---------------- + # Reverse the fwd's all-gather/aux pipeline: aux_loss_bwd produces + # d_aux_probs, then topk_bwd(compute_aux_scores=True) produces the + # extra d_logits contribution. The replicated tensor adds into the + # T-sharded routing-side d_logits via JAX's normal broadcast. + if aux_loss_coeff > 0.0: + T_global = ctx.logits_2d.shape[0] + d_aux_loss_scalar = d_aux_loss.reshape(()).astype(jnp.float32) + d_aux_probs = tex.fused_moe_aux_loss_bwd( + ctx.aux_const_buf, + ctx.aux_tokens_per_expert.astype(jnp.int32), + d_aux_loss_scalar, + num_tokens=int(T_global), + ) + # routing_map is ignored by the kernel when compute_aux_scores=True, + # so pass a zero placeholder of the right shape/dtype. + zero_routing_map = jnp.zeros(ctx.aux_saved_scores.shape, dtype=ctx.routing_map.dtype) + d_logits_aux = tex.fused_topk_with_score_function_bwd( + zero_routing_map, + ctx.aux_saved_scores, + d_aux_probs.astype(ctx.aux_saved_scores.dtype), + topk=K, + use_pre_softmax=False, + scaling_factor=1.0, + score_function=score_function, + compute_aux_scores=True, + ) + d_logits_2d = d_logits_2d + d_logits_aux.astype(d_logits_2d.dtype) + + # ---------------- Gate bwd (global view) ---------------- + d_gate_logits = d_logits_2d.reshape(B, S, num_experts) + gate_kernel_cast = ctx.gate_kernel.astype(ctx.x.dtype) + d_x_from_gate = jnp.einsum("bse,he->bsh", d_gate_logits, gate_kernel_cast) + d_gate_kernel = jnp.einsum("bsh,bse->he", ctx.x, d_gate_logits).astype(ctx.gate_kernel.dtype) + d_x = d_x_from_gate + d_x_from_dispatch + + # Pin output grads to the declared logical axes so downstream + # optimizers see consistent shardings. + d_x = with_sharding_constraint_by_logical_axes(d_x, input_axes) + d_gate_kernel = with_sharding_constraint_by_logical_axes(d_gate_kernel, gate_kernel_axes) + d_wi_0 = with_sharding_constraint_by_logical_axes(d_wi_0, wi_kernel_axes) + d_wi_1 = with_sharding_constraint_by_logical_axes(d_wi_1, wi_kernel_axes) + d_wo = with_sharding_constraint_by_logical_axes(d_wo, wo_kernel_axes) + + # expert_bias has no learnable bwd path through fused_topk: the + # primitive's bwd returns None for the bias slot. Match that with a + # zero cotangent of the right shape so custom_vjp's arity check + # passes. + d_expert_bias = jnp.zeros_like(ctx.expert_bias) -def _grads_dict_to_tuple( - grads: dict, has_wi_bias: bool, has_wo_bias: bool, has_expert_bias: bool -) -> Tuple: - """Pack the body_bwd's grads dict into the positional tuple JAX expects.""" return ( - grads["inputs"], - grads["gate_kernel"], - grads["wi_0"], - grads["wi_1"], - grads["wo"], - grads.get("wi_0_bias") if has_wi_bias else None, - grads.get("wi_1_bias") if has_wi_bias else None, - grads.get("wo_bias") if has_wo_bias else None, - grads.get("expert_bias") if has_expert_bias else None, + d_x, + d_gate_kernel, + d_wi_0, + d_wi_1, + d_wo, + d_wi_0_bias if has_bias else None, + d_wi_1_bias if has_bias else None, + d_wo_bias if has_bias else None, + d_expert_bias, ) @@ -1995,7 +1064,7 @@ def _grads_dict_to_tuple( # ============================================================================= -@partial(jax.custom_vjp, nondiff_argnums=tuple(range(9, 29))) +@partial(jax.custom_vjp, nondiff_argnums=tuple(range(9, 26))) def _moe( x, gate_kernel, @@ -2015,23 +1084,16 @@ def _moe( group_topk, scaling_factor, aux_loss_coeff, - permutation_backend, - align_size, - gate_inside_vjp, ep_axis, data_parallelism_axes, input_axes, gate_kernel_axes, wi_kernel_axes, wo_kernel_axes, - quantizer_sets, dtype, + apply_topk_weights_early, ): - # Call in `_moe`'s own signature order to match what JAX will pass - # the fwd rule via ``_argnums_partial``. See the comment block at - # the top of ``_moe_fwd_rule`` for why this differs from - # ``_moe_bwd_rule``'s convention. - output_pair, _ = _moe_fwd_rule( + primal, _ = _moe_fwd_rule( x, gate_kernel, wi_0, @@ -2050,19 +1112,16 @@ def _moe( group_topk, scaling_factor, aux_loss_coeff, - permutation_backend, - align_size, - gate_inside_vjp, ep_axis, data_parallelism_axes, input_axes, gate_kernel_axes, wi_kernel_axes, wo_kernel_axes, - quantizer_sets, dtype, + apply_topk_weights_early, ) - return output_pair + return primal _moe.defvjp(_moe_fwd_rule, _moe_bwd_rule) @@ -2079,56 +1138,106 @@ def moe( wo_bias: Optional[jnp.ndarray] = None, expert_bias: Optional[jnp.ndarray] = None, *, - # Architecture num_experts: int, num_experts_per_tok: int, activation_type: str = "silu", - # Routing score_function: Union[str, ScoreFunction] = "softmax", use_pre_softmax: bool = False, num_groups: Optional[int] = None, group_topk: Optional[int] = None, scaling_factor: float = 1.0, aux_loss_coeff: float = 0.0, - # Permutation - permutation_backend: PermutationBackend = PermutationBackend.PURE_JAX, - align_size: int = 0, - # Gate placement (Phuong: "perhaps as an option") - gate_inside_vjp: bool = True, - # Parallelism (resolved by caller from MeshResource) - ep_axis: Optional[str] = None, + apply_topk_weights_early: bool = False, + ep_axis: str, data_parallelism_axes: Tuple[str, ...] = (), - # Logical axes for sharding constraints input_axes: Tuple[Optional[str], ...] = (), gate_kernel_axes: Tuple[Optional[str], ...] = (), wi_kernel_axes: Tuple[Optional[str], ...] = ("exp", "embed", "mlp"), wo_kernel_axes: Tuple[Optional[str], ...] = ("exp", "mlp", "embed"), - # Quantization - quantizer_sets: Tuple[QuantizerSet, QuantizerSet, QuantizerSet] = ( - noop_quantizer_set, - noop_quantizer_set, - noop_quantizer_set, - ), dtype: jnp.dtype = jnp.float32, ) -> Tuple[jnp.ndarray, Optional[jnp.ndarray]]: - """Run a full MoE block under a single fused custom_vjp. + """Run a full MoE block under a single fused custom_vjp on the TE EP path. + + Returns ``(output, aux_loss)``. ``aux_loss`` is ``None`` when + ``aux_loss_coeff == 0`` and a 0-d scalar otherwise. - Parameters and return are documented at the call site of - ``_MoEBlock.__call__``. See module docstring for design rationale. + Parameters + ---------- + expert_bias : Optional[jnp.ndarray] + ``[num_experts]`` learnable router bias added before the top-k + when ``score_function='sigmoid'``. Pass ``None`` to disable. + The bias has no gradient through the top-k primitive itself (it + only steers expert selection); a zero cotangent is returned for + it. + aux_loss_coeff : float + Per-step expert-load-balance loss coefficient. ``0.0`` (default) + disables the aux loss entirely. When non-zero, an extra + all-gather over the routing-side logits is inserted so the + ``fused_moe_aux_loss`` kernel sees a global ``[T_global, E]`` + view; this lives off the dispatch critical path. + + Note that the per-expert dispatch-slot alignment is fixed internally + at 128 tokens (``_ALIGN_SIZE``); see that constant's docstring for + rationale and how to extend if a future recipe needs >128. + + Axis-name parameters: + + * ``ep_axis`` and ``data_parallelism_axes`` are *physical mesh + axis names* -- they index ``jax.sharding.Mesh.shape`` directly + (to compute ``num_ep`` / ``dp_size`` and to construct + ``P((dp..., ep), None, None)`` for the per-shard + ``jax.lax.with_sharding_constraint`` calls that JAX requires + to refer to real mesh axes). + * ``input_axes``, ``gate_kernel_axes``, ``wi_kernel_axes``, + ``wo_kernel_axes`` are *logical axis names* (e.g. + ``"batch"``, ``"embed"``, ``"mlp"``, ``"exp"``) -- they get + resolved via the active Flax logical-axis rules and consumed + by ``with_sharding_constraint_by_logical_axes``. They are + ``Optional[str]`` tuples so a rule of ``None`` means + "replicated on this axis". + + Logical-axis support for ``ep_axis`` / ``data_parallelism_axes`` + is intentionally out of scope: the EP comm-group construction + (``dp_color = rank // ep_size``) and the bootstrap signature + check both require concrete integer sizes, so a logical name + would have to be resolved to a physical one anyway before any + EP primitive is called. If a downstream pipeline needs to plumb + logical names all the way to ``moe()``, do the rule lookup at + the call site. + + See module docstring for the rest of the parameter semantics and the + surrounding design rationale. """ - if not isinstance(permutation_backend, PermutationBackend): - raise TypeError( - f"permutation_backend must be a PermutationBackend, got {permutation_backend!r}" - ) - if permutation_backend is PermutationBackend.TRITON: - _require_triton() - # Normalize string score_function ("softmax" / "sigmoid") to the - # ScoreFunction enum once here. The underlying primitive - # ``tex.fused_topk_with_score_function_fwd`` expects an int-coercible - # value (the enum has integer .value), and the public router wrapper - # we bypass also normalizes here. score_function = _validate_score_function(score_function) + # Enforce ((outer_dp..., ep), None, None) on inbound activations. The + # EP comm groups consecutive global ranks (dp_color = rank // ep_size), + # so ep MUST be innermost in the partition spec. Soft re-pin: free if + # upstream already matches, single reshard otherwise. + mesh = _get_mesh() + if mesh is None or mesh.empty: + raise ValueError("moe(...) requires an active jax.sharding.Mesh.") + expected_leading: Any = (*data_parallelism_axes, ep_axis) if data_parallelism_axes else ep_axis + expected_spec = P(expected_leading, None, None) + actual_spec = getattr(getattr(x, "sharding", None), "spec", None) + if actual_spec is not None and tuple(actual_spec) != tuple(expected_spec): + warnings.warn( + f"moe(...): inbound x sharding {actual_spec} does not match expected " + f"{expected_spec}; inserting a reshard. Apply " + "jax.lax.with_sharding_constraint upstream to avoid this overhead.", + UserWarning, + stacklevel=2, + ) + x = _with_sharding_constraint_cast_bwd(x, NamedSharding(mesh, expected_spec)) + + # custom_vjp can't trace through None args; lower expert_bias to an + # empty shape-(0,) tensor that fused_topk_with_score_function treats + # as "no bias". + if expert_bias is None: + expert_bias_arg = jnp.zeros((0,), dtype=jnp.float32) + else: + expert_bias_arg = expert_bias.astype(jnp.float32) + output, aux_loss = _moe( x, gate_kernel, @@ -2138,28 +1247,26 @@ def moe( wi_0_bias, wi_1_bias, wo_bias, - expert_bias, - num_experts=num_experts, - num_experts_per_tok=num_experts_per_tok, - activation_type=activation_type, - score_function=score_function, - use_pre_softmax=use_pre_softmax, - num_groups=num_groups, - group_topk=group_topk, - scaling_factor=scaling_factor, - aux_loss_coeff=aux_loss_coeff, - permutation_backend=permutation_backend, - align_size=align_size, - gate_inside_vjp=gate_inside_vjp, - ep_axis=ep_axis, - data_parallelism_axes=data_parallelism_axes, - input_axes=input_axes, - gate_kernel_axes=gate_kernel_axes, - wi_kernel_axes=wi_kernel_axes, - wo_kernel_axes=wo_kernel_axes, - quantizer_sets=quantizer_sets, - dtype=dtype, + expert_bias_arg, + num_experts, + num_experts_per_tok, + activation_type, + score_function, + use_pre_softmax, + num_groups, + group_topk, + scaling_factor, + float(aux_loss_coeff), + ep_axis, + data_parallelism_axes, + input_axes, + gate_kernel_axes, + wi_kernel_axes, + wo_kernel_axes, + dtype, + apply_topk_weights_early, ) if aux_loss_coeff <= 0.0: aux_loss = None + assert output.dtype == x.dtype, f"moe() output dtype {output.dtype} != input dtype {x.dtype}" return output, aux_loss diff --git a/transformer_engine/jax/version_utils.py b/transformer_engine/jax/version_utils.py index e4619d8670..8765a4fce0 100644 --- a/transformer_engine/jax/version_utils.py +++ b/transformer_engine/jax/version_utils.py @@ -64,6 +64,16 @@ def is_triton_autotuned_alias_safe() -> bool: return v >= PkgVersion(_TRITON_AUTOTUNED_ALIAS_STABLE_FLOOR) +# XLA gained the ``gpu_stream:collective`` stream annotation in openxla/xla#39604, +# which ships in the JAX 0.10.0 release. Older XLA fatally fails on it. +_COLLECTIVE_STREAM_MIN_JAX_VERSION = "0.10.0" + + +def is_collective_stream_supported() -> bool: + """Return True if the installed JAX supports the gpu_stream:collective annotation.""" + return jax_version_meet_requirement(_COLLECTIVE_STREAM_MIN_JAX_VERSION) + + def is_triton_extension_supported() -> bool: """Return True if the current JAX version supports Triton kernel dispatch. @@ -77,6 +87,7 @@ def is_triton_extension_supported() -> bool: __all__ = [ "jax_version_meet_requirement", "is_triton_autotuned_alias_safe", + "is_collective_stream_supported", "is_triton_extension_supported", "TRITON_EXTENSION_MIN_JAX_VERSION", "TRITON_EXTENSION_CUDA_GRAPH_MIN_JAX_VERSION", diff --git a/transformer_engine/pytorch/attention/dot_product_attention/backends.py b/transformer_engine/pytorch/attention/dot_product_attention/backends.py index 785e438cda..d6acc1aaf7 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/backends.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/backends.py @@ -1328,6 +1328,8 @@ def forward( fp8_output, layer_number, return_max_logit, + packed_qkv=None, + packed_kv=None, ): # pylint: disable=missing-function-docstring @@ -1390,7 +1392,14 @@ def forward( q_fp8, k_fp8, v_fp8 = q, k, v else: q_fp8, k_fp8, v_fp8, qkv_layout, qkv_scale_inv_format = combine_and_quantize( - qkv_layout, q, k, v, QKV_quantizer, used_in_backward=is_training + qkv_layout, + q, + k, + v, + QKV_quantizer, + used_in_backward=is_training, + combined_qkv=packed_qkv, + combined_kv=packed_kv, ) # print quantizers @@ -1920,6 +1929,8 @@ def backward(ctx, d_out, *_args): None, None, None, + None, # packed_qkv + None, # packed_kv ) @@ -2014,6 +2025,8 @@ def forward( score_mod_bprop: Optional[Callable] = None, score_mod_tensors: Optional[Dict[str, torch.Tensor]] = None, score_mod_bprop_tensors: Optional[Dict[str, torch.Tensor]] = None, + packed_qkv: Optional[torch.Tensor] = None, + packed_kv: Optional[torch.Tensor] = None, ) -> torch.Tensor: """fused attention fprop""" assert ( @@ -2233,6 +2246,8 @@ def forward( fp8_output, self.layer_number, self.return_max_logit, + packed_qkv, + packed_kv, ) if self.return_max_logit: diff --git a/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py b/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py index dd812f15e9..4cc4cab1b8 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py @@ -196,6 +196,111 @@ def _trim_output(attn_out, num_attention_heads, padded_head_dim_v, orig_head_dim return attn_out[..., :orig_head_dim_v].reshape(*out_shape, -1) +def _unpack_packed_qkv( + qkv_layer: Optional[torch.Tensor], + kv_layer: Optional[torch.Tensor], + query_layer: Optional[torch.Tensor], + key_layer: Optional[torch.Tensor], + value_layer: Optional[torch.Tensor], + qkv_format: str, + qkv_interleave_dim: int, + inference_params: Optional[InferenceParams], +) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, Optional[str]]: + """Resolve declarative packed inputs into q/k/v. + + Derives q/k/v as zero-copy views of the packed buffer (``qkv_layer`` or + ``kv_layer``) and constructs the exact layout string from the declaration. + The layout enum is truthful by construction, so no pointer-based detection + is needed downstream (this also covers thd and FP8 DPA). + + Returns ``(query_layer, key_layer, value_layer, declared_qkv_layout)``; + the layout is ``None`` when no packed input is given. + """ + if qkv_layer is None and kv_layer is None: + if query_layer is None or key_layer is None or value_layer is None: + raise ValueError( + "query_layer, key_layer and value_layer are required unless packed" + " inputs (qkv_layer or query_layer + kv_layer) are provided." + ) + return query_layer, key_layer, value_layer, None + + if qkv_layer is not None and kv_layer is not None: + raise ValueError("qkv_layer and kv_layer are mutually exclusive.") + if inference_params is not None: + raise ValueError( + "Packed inputs (qkv_layer/kv_layer) are not supported with KV caching" + " (inference_params); pass separate query/key/value tensors instead." + ) + if qkv_interleave_dim not in (-3, -2): + raise ValueError( + "qkv_interleave_dim must be -3 (e.g. bs3hd) or -2 (e.g. bsh3d), got" + f" {qkv_interleave_dim}." + ) + packed = qkv_layer if qkv_layer is not None else kv_layer + # The declared layout describes the packed buffer's memory, so it must have + # stride 1 in its last dimension (the check get_qkv_layout would otherwise + # perform on the derived q/k/v views). + if packed.stride(-1) != 1: + raise ValueError( + "The packed tensor (qkv_layer/kv_layer) must have stride 1 in its last" + f" dimension, got strides {tuple(packed.stride())}." + ) + + def _packed_layout(fmt: str, num: int) -> str: + # bshd + 3 @ -3 -> bs3hd; bshd + 3 @ -2 -> bsh3d; thd + 2 @ -2 -> th2d + pos = len(fmt) + qkv_interleave_dim + 1 + return fmt[:pos] + str(num) + fmt[pos:] + + if qkv_layer is not None: + if any(x is not None for x in (query_layer, key_layer, value_layer)): + raise ValueError( + "qkv_layer already packs Q, K and V: query_layer, key_layer and" + " value_layer must be None when qkv_layer is provided." + ) + expected_rank = 4 if qkv_format == "thd" else 5 + if qkv_layer.dim() != expected_rank: + raise ValueError( + f"qkv_layer must be a {expected_rank}D tensor for" + f" qkv_format={qkv_format!r}, got {qkv_layer.dim()}D." + ) + if qkv_layer.shape[qkv_interleave_dim] != 3: + raise ValueError( + f"qkv_layer must have size 3 at dim {qkv_interleave_dim}" + f" (qkv_interleave_dim), got shape {tuple(qkv_layer.shape)}." + ) + query_layer, key_layer, value_layer = ( + qkv_layer.select(qkv_interleave_dim, i) for i in range(3) + ) + return query_layer, key_layer, value_layer, _packed_layout(qkv_format, 3) + + if query_layer is None: + raise ValueError( + "kv_layer packs only K and V: query_layer is required when kv_layer is provided." + ) + if key_layer is not None or value_layer is not None: + raise ValueError( + "kv_layer already packs K and V: key_layer and value_layer must be" + " None when kv_layer is provided." + ) + if kv_layer.dim() != query_layer.dim() + 1: + raise ValueError( + "kv_layer must have one more dimension than query_layer, got" + f" {kv_layer.dim()}D kv_layer and {query_layer.dim()}D query_layer." + ) + if kv_layer.shape[qkv_interleave_dim] != 2: + raise ValueError( + f"kv_layer must have size 2 at dim {qkv_interleave_dim}" + f" (qkv_interleave_dim), got shape {tuple(kv_layer.shape)}." + ) + key_layer, value_layer = (kv_layer.select(qkv_interleave_dim, i) for i in range(2)) + return ( + query_layer, + key_layer, + value_layer, + f"{qkv_format}_{_packed_layout(qkv_format, 2)}", + ) + + class DotProductAttention(TransformerEngineBaseModule): r"""Allows the model to jointly attend to information from different representation subspaces as described in the paper: @@ -988,9 +1093,9 @@ def get_quantizer_roles( @no_torch_dynamo(recursive=False) def forward( self, - query_layer: torch.Tensor, - key_layer: torch.Tensor, - value_layer: torch.Tensor, + query_layer: Optional[torch.Tensor] = None, + key_layer: Optional[torch.Tensor] = None, + value_layer: Optional[torch.Tensor] = None, attention_mask: Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]] = None, qkv_format: str = None, cu_seqlens_q: torch.Tensor = None, @@ -1015,6 +1120,9 @@ def forward( score_mod_bprop: Optional[Callable] = None, score_mod_tensors: Optional[Dict[str, torch.Tensor]] = None, score_mod_bprop_tensors: Optional[Dict[str, torch.Tensor]] = None, + qkv_layer: Optional[torch.Tensor] = None, + kv_layer: Optional[torch.Tensor] = None, + qkv_interleave_dim: int = -3, ) -> torch.Tensor: r""" Dot Product Attention Layer. @@ -1113,12 +1221,15 @@ def forward( Parameters ---------- - query_layer : torch.Tensor - Query tensor. - key_layer : torch.Tensor - Key tensor. - value_layer : torch.Tensor - Value tensor. + query_layer : Optional[torch.Tensor], default = None + Query tensor. Required unless a packed input (``qkv_layer``, or + ``kv_layer`` together with ``query_layer``) is provided instead. + key_layer : Optional[torch.Tensor], default = None + Key tensor. Required unless a packed input (``qkv_layer`` or ``kv_layer``) + is provided instead. + value_layer : Optional[torch.Tensor], default = None + Value tensor. Required unless a packed input (``qkv_layer`` or ``kv_layer``) + is provided instead. 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 @@ -1226,8 +1337,44 @@ def forward( Runtime tensors exposed to score_mod_bprop as cuDNN graph tensors. Keys are user-defined string names consumed by the callback through ``tensors[name]``; there is no predefined set of accepted keys. + qkv_layer: Optional[torch.Tensor], default = None + Fully packed QKV tensor. When the QKV projection produces one packed buffer + (e.g. a fused QKV GEMM), it can be passed here directly instead of slicing + it into :attr:`query_layer`/:attr:`key_layer`/:attr:`value_layer` views. + For :attr:`qkv_format` = {"bshd", "sbhd"}, it must be a 5D tensor of shape + ``[b, s, 3, h, d]``/``[s, b, 3, h, d]`` (:attr:`qkv_interleave_dim` = -3) or + ``[b, s, h, 3, d]``/``[s, b, h, 3, d]`` (:attr:`qkv_interleave_dim` = -2); + for :attr:`qkv_format` = "thd", a 4D tensor of shape ``[t, 3, h, d]`` or + ``[t, h, 3, d]``. Q/K/V are derived as zero-copy views and the memory layout + (e.g. ``bs3hd``) is declared from the packing itself, so no pointer-based + layout detection runs on this path -- including for "thd" and FP8 attention. + Mutually exclusive with :attr:`query_layer`, :attr:`key_layer`, + :attr:`value_layer` and :attr:`kv_layer`. + kv_layer: Optional[torch.Tensor], default = None + Packed KV tensor, used together with :attr:`query_layer` + (e.g. ``[b, s, 2, hg, d]`` for :attr:`qkv_interleave_dim` = -3, or + ``[b, s, hg, 2, d]`` for :attr:`qkv_interleave_dim` = -2). K/V are derived + as zero-copy views and the layout (e.g. ``bshd_bs2hd``) is declared, not + detected. Mutually exclusive with :attr:`key_layer`, :attr:`value_layer` + and :attr:`qkv_layer`. + qkv_interleave_dim: int, default = -3 + Dimension of :attr:`qkv_layer`/:attr:`kv_layer` where the 3 (QKV) or 2 (KV) + interleave sits; must be -3 (e.g. ``bs3hd``) or -2 (e.g. ``bsh3d``, + Megatron-style). This is an explicit knob rather than shape inference, + since e.g. ``h == 3`` would make the shapes ambiguous. """ + query_layer, key_layer, value_layer, declared_qkv_layout = _unpack_packed_qkv( + qkv_layer, + kv_layer, + query_layer, + key_layer, + value_layer, + qkv_format if qkv_format is not None else self.qkv_format, + qkv_interleave_dim, + inference_params, + ) + with self.prepare_forward_ctx( query_layer, num_gemms=3, @@ -1413,7 +1560,15 @@ def forward( cu_seqlens_kv_padded = None # get qkv's memory layout - if all( + if declared_qkv_layout is not None: + # Packed inputs (qkv_layer/kv_layer) declare the layout: the enum is + # truthful by construction, so the pointer-based detection in + # get_qkv_layout is skipped entirely -- for dense, thd (t3hd/th3d) + # and FP8 DPA alike. + qkv_layout = declared_qkv_layout + q_format = qkv_format + kv_format = qkv_format + elif all( isinstance(x, Float8TensorStorage) for x in [query_layer, key_layer, value_layer] ): ( @@ -1801,6 +1956,8 @@ def forward( inference_params=inference_params, softmax_offset=softmax_offset, fp8_output=fp8_output, + packed_qkv=qkv_layer, + packed_kv=kv_layer, ) return self.fused_attention( query_layer, @@ -1836,6 +1993,8 @@ def forward( score_mod_bprop=score_mod_bprop, score_mod_tensors=score_mod_tensors, score_mod_bprop_tensors=score_mod_bprop_tensors, + packed_qkv=qkv_layer, + packed_kv=kv_layer, ) if use_unfused_attention: diff --git a/transformer_engine/pytorch/attention/dot_product_attention/utils.py b/transformer_engine/pytorch/attention/dot_product_attention/utils.py index 7be94a6fa1..1713a98129 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/utils.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/utils.py @@ -35,6 +35,7 @@ META_DP, ) from transformer_engine.pytorch.attention.inference import InferenceParams +from transformer_engine.pytorch.cpu_offload import is_cpu_offload_enabled from transformer_engine.pytorch.quantized_tensor import QuantizedTensorStorage from transformer_engine.pytorch.tensor.float8_tensor import ( Float8Tensor, @@ -455,18 +456,18 @@ def get_attention_backend( # Filter: Environment variables use_flash_attention = int(os.getenv("NVTE_FLASH_ATTN", "1")) - use_flash_attention_2 = use_flash_attention - use_flash_attention_3 = use_flash_attention - use_flash_attention_4 = use_flash_attention + use_flash_attention_2 = use_flash_attention and int(os.getenv("NVTE_FLASH_ATTN_V2", "1")) + use_flash_attention_3 = use_flash_attention and int(os.getenv("NVTE_FLASH_ATTN_V3", "1")) + use_flash_attention_4 = use_flash_attention and int(os.getenv("NVTE_FLASH_ATTN_V4", "1")) flash_attention_backend = None use_fused_attention = int(os.getenv("NVTE_FUSED_ATTN", "1")) use_unfused_attention = int(os.getenv("NVTE_UNFUSED_ATTN", "1")) if not use_flash_attention_2 and FlashAttentionUtils.is_installed: - logger.debug("Disabling FlashAttention 2 due to NVTE_FLASH_ATTN=0") + logger.debug("Disabling FlashAttention 2 due to NVTE_FLASH_ATTN=0 or NVTE_FLASH_ATTN_V2=0") if not use_flash_attention_3 and FlashAttentionUtils.v3_is_installed: - logger.debug("Disabling FlashAttention 3 due to NVTE_FLASH_ATTN=0") + logger.debug("Disabling FlashAttention 3 due to NVTE_FLASH_ATTN=0 or NVTE_FLASH_ATTN_V3=0") if not use_flash_attention_4 and FlashAttentionUtils.v4_is_installed: - logger.debug("Disabling FlashAttention 4 due to NVTE_FLASH_ATTN=0") + logger.debug("Disabling FlashAttention 4 due to NVTE_FLASH_ATTN=0 or NVTE_FLASH_ATTN_V4=0") if not use_fused_attention: logger.debug("Disabling FusedAttention due to NVTE_FUSED_ATTN=0") if not use_unfused_attention: @@ -2432,6 +2433,22 @@ def run_iteratively(q, k, v): if qkv_layout == "not_supported": raise RuntimeError("The provided qkv memory layout is not supported!") + if len(qkv_layout.split("_")) < 3: + # q/k/v were recognized as views of a packed buffer only by inspecting + # their data pointers, strides and storage offsets. Skip the nudge while + # CPU offloading is enabled: offloading forces MultiheadAttention onto + # its sliced-views fallback, so packed views reaching detection are + # expected there and the caller has no migration option. + if not is_cpu_offload_enabled(): + warnings.warn( + "Relying on pointer-based detection of packed q/k/v layouts" + f" (detected {qkv_layout!r}) is deprecated: pass the packed buffer" + " explicitly via qkv_layer/kv_layer (with qkv_interleave_dim) to" + " DotProductAttention instead.", + DeprecationWarning, + stacklevel=2, + ) + if inference_params is not None and inference_params.is_paged: qkv_layout = "paged_kv_" + qkv_layout @@ -2787,8 +2804,18 @@ def combine_and_quantize( used_in_forward=True, used_in_backward=False, keep_same_data_and_scale_inv_format=False, + combined_qkv: Optional[torch.Tensor] = None, + combined_kv: Optional[torch.Tensor] = None, ): - """Combine Q, K, V tensors based on qkv_layout and quantize them together.""" + """Combine Q, K, V tensors based on qkv_layout and quantize them together. + + When ``combined_qkv`` (for ``qkv_group=1`` layouts such as ``bs3hd``) or + ``combined_kv`` (for ``qkv_group=2`` layouts such as ``bshd_bs2hd``) is provided, it must be the + caller's original packed buffer that q/k/v are views of. It is then quantized + directly instead of re-deriving the packed buffer from the q/k/v views via + ``combine_tensors`` (which rebuilds it with a raw ``set_`` under a silent + adjacency/interleave assumption). Ignored for MXFP8 quantization. + """ if isinstance(qkv_quantizer, MXFP8Quantizer): qkv_format, q_format, kv_format = get_qkv_format(qkv_layout) assert qkv_format in ("bshd", "sbhd"), ( @@ -2888,12 +2915,26 @@ def combine_and_quantize( match qkv_group: case 1: dim = qkv_layout.find("3") - qkv = combine_tensors([q, k, v], dim) + if combined_qkv is not None: + assert combined_qkv.shape[dim] == 3, ( + f"combined_qkv does not match qkv_layout {qkv_layout}: expected" + f" size 3 at dim {dim}, got shape {tuple(combined_qkv.shape)}." + ) + qkv = combined_qkv + else: + qkv = combine_tensors([q, k, v], dim) qkv_fp8 = qkv_quantizer(qkv) q_data, k_data, v_data = SplitAlongDim.apply(qkv_fp8._data, dim, [1, 1, 1], True) case 2: dim = qkv_layout.split("_")[1].find("2") - kv = combine_tensors([k, v], dim) + if combined_kv is not None: + assert combined_kv.shape[dim] == 2, ( + f"combined_kv does not match qkv_layout {qkv_layout}: expected" + f" size 2 at dim {dim}, got shape {tuple(combined_kv.shape)}." + ) + kv = combined_kv + else: + kv = combine_tensors([k, v], dim) tensors = [q, kv] num_tensors = len(tensors) shapes = [x.shape for x in tensors] diff --git a/transformer_engine/pytorch/attention/multi_head_attention.py b/transformer_engine/pytorch/attention/multi_head_attention.py index 70ae9dfc21..0b73daf667 100644 --- a/transformer_engine/pytorch/attention/multi_head_attention.py +++ b/transformer_engine/pytorch/attention/multi_head_attention.py @@ -9,6 +9,7 @@ import torch from transformer_engine.pytorch.quantization import FP8GlobalStateManager, QuantizerRole +from transformer_engine.pytorch.quantized_tensor import QuantizedTensorStorage from transformer_engine.pytorch.tensor.float8_tensor import Float8Tensor from transformer_engine.pytorch.module.base import TransformerEngineBaseModule from transformer_engine.pytorch.module import LayerNormLinear, Linear, RMSNorm, LayerNorm @@ -903,6 +904,23 @@ def forward( self._update_output_quantizer_roles(qkv_fp8_output, proj_fp8_grad, dpa_fp8_output) + # Packed pass-through to DotProductAttention: the fused QKV/KV projection + # already produces one packed buffer, which DPA accepts directly via its + # declarative qkv_layer/kv_layer arguments (deriving q/k/v as zero-copy + # views and skipping pointer-based layout detection). Only possible when + # no per-tensor operation (RoPE, QK normalization, KV caching, CPU + # offloading) needs the individual q/k/v slices. + packed_dpa_eligible = ( + rotary_pos_emb is None + and self.q_norm is None + and self.k_norm is None + and inference_params is None + and not is_cpu_offload_enabled() + ) + packed_qkv_layer = None + packed_kv_layer = None + packed_interleave_dim = -3 + layernorm_output = None if self.attention_type == "self": # Attention heads [sq, b, h] --> [sq, b, ng * (np/ng + 2) * hn] @@ -947,28 +965,43 @@ def forward( mixed_x_layer = mixed_x_layer.view(*new_tensor_shape) - # qkv_weight_interleaved: - # [sq, b, ng, (np/ng + 2), hn] - # --> [sq, b, ng, np/ng, hn], [sq, b, ng, 1, hn], [sq, b, ng, 1, hn] - # not qkv_weight_interleaved: - # [sq, b, (np/ng + 2), ng, hn] - # --> [sq, b, np/ng, np, hn], [sq, b, 1, ng, hn], [sq, b, 1, ng, hn] - query_layer, key_layer, value_layer = SplitAlongDim.apply( - mixed_x_layer, split_dim, (num_queries_per_key_value, 1, 1) - ) - - if self.qkv_format == "thd": - query_layer, key_layer, value_layer = ( - x.reshape(x.size(0), -1, self.hidden_size_per_attention_head) - for x in (query_layer, key_layer, value_layer) - ) + if ( + num_queries_per_key_value == 1 + and packed_dpa_eligible + and not isinstance(mixed_x_layer, QuantizedTensorStorage) + ): + # np == ng: the projection output is a uniform 3-interleave + # ([.., h, 3, d] interleaved / [.., 3, h, d] otherwise), which + # DotProductAttention accepts directly as a declared packed + # qkv_layer -- no slicing here, no layout detection there. + packed_qkv_layer = mixed_x_layer + packed_interleave_dim = split_dim + query_layer = None + key_layer = None + value_layer = None else: - # query: -> [sq, b, np, hn] - # key, value: -> [sq, b, ng, hn] - query_layer, key_layer, value_layer = ( - x.reshape(x.size(0), x.size(1), -1, self.hidden_size_per_attention_head) - for x in (query_layer, key_layer, value_layer) + # qkv_weight_interleaved: + # [sq, b, ng, (np/ng + 2), hn] + # --> [sq, b, ng, np/ng, hn], [sq, b, ng, 1, hn], [sq, b, ng, 1, hn] + # not qkv_weight_interleaved: + # [sq, b, (np/ng + 2), ng, hn] + # --> [sq, b, np/ng, np, hn], [sq, b, 1, ng, hn], [sq, b, 1, ng, hn] + query_layer, key_layer, value_layer = SplitAlongDim.apply( + mixed_x_layer, split_dim, (num_queries_per_key_value, 1, 1) ) + + if self.qkv_format == "thd": + query_layer, key_layer, value_layer = ( + x.reshape(x.size(0), -1, self.hidden_size_per_attention_head) + for x in (query_layer, key_layer, value_layer) + ) + else: + # query: -> [sq, b, np, hn] + # key, value: -> [sq, b, ng, hn] + query_layer, key_layer, value_layer = ( + x.reshape(x.size(0), x.size(1), -1, self.hidden_size_per_attention_head) + for x in (query_layer, key_layer, value_layer) + ) elif self.attention_type == "cross": # Attention heads [sk, b, h] --> [sk, b, (ng * 2 * hn)] mixed_kv_layer = self.key_value( @@ -996,34 +1029,56 @@ def forward( mixed_kv_layer = mixed_kv_layer.view(*new_tensor_shape) - # mixed_kv_layer --> 2 [sk, b, ng, hn] - key_layer, value_layer = SplitAlongDim.apply( - mixed_kv_layer, - split_dim, - mixed_kv_layer.shape[split_dim] // 2, - ) - key_layer, value_layer = ( - x.reshape( - x.size(0), - x.size(1), - -1, - self.hidden_size_per_attention_head, - ) - for x in (key_layer, value_layer) - ) - - if self.qkv_format == "thd": - key_layer, value_layer = ( - x.reshape(x.size(0), -1, self.hidden_size_per_attention_head) - for x in (key_layer, value_layer) - ) + if packed_dpa_eligible and not isinstance(mixed_kv_layer, QuantizedTensorStorage): + # Declare the packed KV to DotProductAttention instead of + # slicing it: expose the 2-interleave as its own dimension. + if self.qkv_weight_interleaved: + # [.., ng, 2 * hn] --> [.., ng, 2, hn] + packed_kv_shape = mixed_kv_layer.size()[:-1] + ( + 2, + self.hidden_size_per_attention_head, + ) + packed_interleave_dim = -2 + else: + # [.., 2 * ng, hn] --> [.., 2, ng, hn] + packed_kv_shape = mixed_kv_layer.size()[:-2] + ( + 2, + self.num_gqa_groups_per_partition, + self.hidden_size_per_attention_head, + ) + packed_interleave_dim = -3 + packed_kv_layer = mixed_kv_layer.view(*packed_kv_shape) + key_layer = None + value_layer = None else: - # key, value: -> [sq, b, ng, hn] + # mixed_kv_layer --> 2 [sk, b, ng, hn] + key_layer, value_layer = SplitAlongDim.apply( + mixed_kv_layer, + split_dim, + mixed_kv_layer.shape[split_dim] // 2, + ) key_layer, value_layer = ( - x.reshape(x.size(0), x.size(1), -1, self.hidden_size_per_attention_head) + x.reshape( + x.size(0), + x.size(1), + -1, + self.hidden_size_per_attention_head, + ) for x in (key_layer, value_layer) ) + if self.qkv_format == "thd": + key_layer, value_layer = ( + x.reshape(x.size(0), -1, self.hidden_size_per_attention_head) + for x in (key_layer, value_layer) + ) + else: + # key, value: -> [sq, b, ng, hn] + key_layer, value_layer = ( + x.reshape(x.size(0), x.size(1), -1, self.hidden_size_per_attention_head) + for x in (key_layer, value_layer) + ) + # Attention head [sq, b, h] --> [sq, b, hp] if self.input_layernorm: layernorm_query_outputs = self.layernorm_query( @@ -1143,6 +1198,9 @@ def forward( inference_params=inference_params, pad_between_seqs=pad_between_seqs, fp8_output=dpa_fp8_output, + qkv_layer=packed_qkv_layer, + kv_layer=packed_kv_layer, + qkv_interleave_dim=packed_interleave_dim, ) # =================== diff --git a/transformer_engine/pytorch/cpp_extensions/gemm.py b/transformer_engine/pytorch/cpp_extensions/gemm.py index f3b066d50b..d5446e7b61 100644 --- a/transformer_engine/pytorch/cpp_extensions/gemm.py +++ b/transformer_engine/pytorch/cpp_extensions/gemm.py @@ -13,6 +13,7 @@ from ..utils import get_sm_count, _empty_tensor from ..quantized_tensor import Quantizer +from ..tensor.float8_blockwise_tensor import Float8BlockQuantizer from ..tensor.storage.float8_blockwise_tensor_storage import Float8BlockwiseQTensorStorage from ..tensor.storage.grouped_tensor_storage import GroupedTensorStorage from ..tensor.storage.nvfp4_tensor_storage import NVFP4TensorStorage @@ -435,6 +436,30 @@ def _get_fp32_zeros_tensor(num_tensors: int, device: torch.device) -> torch.Tens return torch.zeros(num_tensors, dtype=torch.float32, device=device) +@functools.lru_cache(maxsize=None) +def _get_grouped_gemm_setup_workspace(device: int, num_tensors: int) -> torch.Tensor: + """Persistent setup workspace (per-group pointer/dim arrays) for grouped-tensor GEMM.""" + return torch.empty( + get_grouped_gemm_setup_workspace_size(num_tensors), + dtype=torch.uint8, + device=device, + ) + + +@functools.lru_cache(maxsize=None) +def _get_grouped_cublas_workspace(device: int, layout: str) -> torch.Tensor: + """Persistent cuBLAS workspace for the grouped-tensor GEMM path, one per GEMM layout. + + Grouped cuBlasLt GEMM kernels in cuBLAS versions <= 13.7 leave behind stale descriptors in the + workspace that cause back-to-back GEMM kernels to crash/deadlock on 2nd CUDA-graph replay. As a + workaround, we allocate a different workspace for each GEMM layout (TN, NN, NT) to avoid + contamination between subsequent GEMM calls (when there is no other graph node between GEMM + kernels). + """ + assert layout in ("TN", "NN", "NT"), f"unexpected grouped GEMM layout {layout}" + return torch.empty(get_cublas_workspace_size_bytes(), dtype=torch.uint8, device=device) + + def general_grouped_gemm_for_grouped_tensor( A, B, @@ -474,6 +499,20 @@ def general_grouped_gemm_for_grouped_tensor( if isinstance(out, GroupedTensorStorage) and out.row_scaled_nvfp4: raise NotImplementedError("Row-scaled NVFP4 GroupedTensor GEMM is not supported yet.") + def _is_fp8_blockwise(operand) -> bool: + if isinstance(operand, (list, tuple)): + return any(isinstance(t, Float8BlockwiseQTensorStorage) for t in operand) + if isinstance(operand, GroupedTensorStorage): + return isinstance(operand.quantizer, Float8BlockQuantizer) + return False + + if _is_fp8_blockwise(A) or _is_fp8_blockwise(B): + # The fused grouped FP8 block-scaling GEMM only supports split accumulation, + # so force it on and intentionally override any caller-supplied value. This + # matches the Float8BlockScaling recipe, which fixes use_split_accumulator=True + # for all of fprop/dgrad/wgrad, so no user-configurable setting is discarded. + use_split_accumulator = True + if is_discrete_out: # wgrad case. grouped_gemm_impl = tex.te_general_grouped_gemm_for_discrete_out @@ -513,16 +552,12 @@ def general_grouped_gemm_for_grouped_tensor( if not alpha.is_cuda or not beta.is_cuda: raise ValueError("alpha and beta must be CUDA tensors.") - workspace_setup = torch.empty( - get_grouped_gemm_setup_workspace_size(num_tensors), - dtype=torch.uint8, - device=device, - ) - workspace_cublas = torch.empty( - get_cublas_workspace_size_bytes(), - dtype=torch.uint8, - device=device, - ) + workspace_setup = _get_grouped_gemm_setup_workspace(device.index, num_tensors) + # Each grouped-GEMM layout gets its own persistent cuBLAS workspace: two grouped + # GEMMs sharing one workspace can deadlock under CUDA-graph replay (see + # _get_grouped_cublas_workspace). wgrad (NT) is the case seen in TE; fprop (TN) and + # dgrad (NN) have also been reported to conflict, so all three layouts are isolated. + workspace_cublas = _get_grouped_cublas_workspace(device.index, layout) sm_count = get_sm_count() sm_count = sm_count - int(os.getenv("NVTE_EXT_MARGIN_SM", str(sm_count))) diff --git a/transformer_engine/pytorch/csrc/extensions/cast.cpp b/transformer_engine/pytorch/csrc/extensions/cast.cpp index 8d77a9e349..8b1cd384aa 100644 --- a/transformer_engine/pytorch/csrc/extensions/cast.cpp +++ b/transformer_engine/pytorch/csrc/extensions/cast.cpp @@ -500,6 +500,15 @@ py::object bgrad_group_quantize(const at::Tensor &tensor, py::handle quantizer, auto quantizer_cpp = convert_quantizer(quantizer); + // MXFP8 carries no quantization knobs; FP8 block-scaling reads scale constraints + // off the quantizer, matching the forward group_quantize dispatch. + QuantizationConfigWrapper quant_config_cpp; + if (detail::IsFloat8BlockwiseQuantizers(quantizer.ptr())) { + auto *fp8_block_quantizer_cpp = static_cast(quantizer_cpp.get()); + quant_config_cpp.set_force_pow_2_scales(fp8_block_quantizer_cpp->force_pow_2_scales); + quant_config_cpp.set_amax_epsilon(fp8_block_quantizer_cpp->amax_epsilon); + } + auto grouped_input_tensor = GroupedTensorWrapper(num_tensors, logical_shape); grouped_input_tensor.set_rowwise_data(tensor.data_ptr(), GetTransformerEngineDType(tensor.scalar_type()), @@ -530,7 +539,8 @@ py::object bgrad_group_quantize(const at::Tensor &tensor, py::handle quantizer, auto stream = at::cuda::getCurrentCUDAStream(); NVTE_SCOPED_GIL_RELEASE({ nvte_group_quantize_dbias(grouped_input_tensor.data(), grouped_output_tensor_cpp.data(), - grouped_dbias.data(), workspace_nvte.data(), stream); + grouped_dbias.data(), workspace_nvte.data(), quant_config_cpp, + stream); }); if (workspace_nvte.ndim() > 0 && workspace_nvte.numel() > 0) { at::Tensor workspace_torch = allocateSpace(workspace_nvte.shape(), workspace_nvte.dtype()); @@ -539,7 +549,8 @@ py::object bgrad_group_quantize(const at::Tensor &tensor, py::handle quantizer, } NVTE_SCOPED_GIL_RELEASE({ nvte_group_quantize_dbias(grouped_input_tensor.data(), grouped_output_tensor_cpp.data(), - grouped_dbias.data(), workspace_nvte.data(), stream); + grouped_dbias.data(), workspace_nvte.data(), quant_config_cpp, + stream); }); return py::make_tuple(py::reinterpret_borrow(grouped_output_py), py::cast(std::move(dbias_torch))); diff --git a/transformer_engine/pytorch/csrc/extensions/ep.cpp b/transformer_engine/pytorch/csrc/extensions/ep.cpp index 8173df947e..118f14a01f 100644 --- a/transformer_engine/pytorch/csrc/extensions/ep.cpp +++ b/transformer_engine/pytorch/csrc/extensions/ep.cpp @@ -145,7 +145,13 @@ void ep_initialize(uintptr_t comm_ptr, const std::string& group_name, int64_t nu .max_token_dtype = static_cast(GetTransformerEngineDType(torch_dtype)), .zero_copy = zero_copy ? 1 : 0, }; - nvte_ep_initialize(static_cast(ep_comm), &cfg); + // Release the GIL only around the native init. It must stay held while pybind11 casts + // the ``max_token_dtype`` object above and destroys the by-value ``pybind11::object`` + // parameter on return; releasing it across those trips pybind11's dec_ref GIL assertion. + { + pybind11::gil_scoped_release nogil; + nvte_ep_initialize(static_cast(ep_comm), &cfg); + } g_zero_copy_enabled.store(zero_copy, std::memory_order_relaxed); g_ep_initialized = true; g_ep_group_name = group_name; @@ -366,8 +372,7 @@ void register_ep_bindings(pybind11::module_& m) { "Initialize the EP backend; borrows torch's NCCL comm pointed to by ``comm_ptr``.", py::arg("comm_ptr"), py::arg("group_name"), py::arg("num_experts"), py::arg("max_tokens_per_rank"), py::arg("max_recv_tokens_per_rank"), py::arg("hidden_dim"), - py::arg("max_num_sms") = 0, py::arg("max_token_dtype"), py::arg("zero_copy") = false, - py::call_guard()); + py::arg("max_num_sms") = 0, py::arg("max_token_dtype"), py::arg("zero_copy") = false); m.def("ep_finalize", &ep_finalize, "Tear down the EP backend. Idempotent.", py::call_guard()); m.def("ep_get_zero_copy", &ep_get_zero_copy, "Return the current EP zero-copy toggle state."); diff --git a/transformer_engine/pytorch/csrc/extensions/pybind.cpp b/transformer_engine/pytorch/csrc/extensions/pybind.cpp index 9c9ec36138..7e9d114be8 100644 --- a/transformer_engine/pytorch/csrc/extensions/pybind.cpp +++ b/transformer_engine/pytorch/csrc/extensions/pybind.cpp @@ -716,6 +716,9 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { int num_max_streams, int comm_cga_size, int gemm_priority, int comm_priority, int num_comm_sm, bool set_sm_margin, bool atomic_gemm, bool rs_overlap_first_gemm) { + // Release the GIL only around the native construction (blocking collectives) to avoid + // tripping pybind11's inc_ref/dec_ref GIL assertions. + py::gil_scoped_release nogil; if (use_cublasmp) { return std::make_shared(helper, helper->mylocal, tp_size, comm_type, buffer_shape, buffer_dtype, num_comm_sm, @@ -726,8 +729,7 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { comm_cga_size, gemm_priority, comm_priority, num_comm_sm, set_sm_margin, atomic_gemm, rs_overlap_first_gemm); }), - py::call_guard(), py::arg("buffer_shape"), - py::arg("buffer_dtype"), py::arg("helper"), py::arg("tp_size"), + py::arg("buffer_shape"), py::arg("buffer_dtype"), py::arg("helper"), py::arg("tp_size"), py::arg("use_cublasmp") = false, py::arg("comm_type") = transformer_engine::CommOverlapType::RS, py::arg("num_splits") = 4, py::arg("num_max_streams") = NVTE_COMM_OVERLAP_MAX_STREAMS, @@ -751,6 +753,9 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { int comm_cga_size, int gemm_priority, int comm_priority, int num_comm_sm, bool set_sm_margin, bool atomic_gemm, bool use_ce, bool aggregate, bool use_cublasmp) { + // Release the GIL only around the native construction (blocking collectives) to avoid + // tripping pybind11's inc_ref/dec_ref GIL assertions. + py::gil_scoped_release nogil; if (use_cublasmp) { return std::make_shared(helper, helper->mylocal, tp_size, comm_type, buffer_shape, buffer_dtype, num_comm_sm, @@ -761,12 +766,12 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { gemm_priority, comm_priority, num_comm_sm, set_sm_margin, atomic_gemm, use_ce, aggregate); }), - py::call_guard(), py::arg("buffer_shape"), - py::arg("buffer_dtype"), py::arg("helper"), py::arg("tp_size"), py::arg("comm_type"), - py::arg("num_max_streams") = NVTE_COMM_OVERLAP_MAX_STREAMS, py::arg("comm_cga_size") = 1, - py::arg("gemm_priority") = 0, py::arg("comm_priority") = 0, py::arg("num_comm_sm") = 1, - py::arg("set_sm_margin") = false, py::arg("atomic_gemm") = false, - py::arg("use_ce") = true, py::arg("aggregate") = false, py::arg("use_cublasmp") = false) + py::arg("buffer_shape"), py::arg("buffer_dtype"), py::arg("helper"), py::arg("tp_size"), + py::arg("comm_type"), py::arg("num_max_streams") = NVTE_COMM_OVERLAP_MAX_STREAMS, + py::arg("comm_cga_size") = 1, py::arg("gemm_priority") = 0, py::arg("comm_priority") = 0, + py::arg("num_comm_sm") = 1, py::arg("set_sm_margin") = false, + py::arg("atomic_gemm") = false, py::arg("use_ce") = true, py::arg("aggregate") = false, + py::arg("use_cublasmp") = false) .def("copy_into_buffer", static_cast( &CommOverlapP2P::copy_into_buffer), diff --git a/transformer_engine/pytorch/csrc/quantizer.cpp b/transformer_engine/pytorch/csrc/quantizer.cpp index 4308464c52..17237fa9b8 100644 --- a/transformer_engine/pytorch/csrc/quantizer.cpp +++ b/transformer_engine/pytorch/csrc/quantizer.cpp @@ -1154,14 +1154,6 @@ std::pair Float8BlockQuantizer::create_grouped const size_t logical_last_dim) const { using namespace pybind11::literals; - // The fused grouped FP8 block-scaling path uses unconstrained FP32 scales and does not - // implement power-of-2 scaling. Reject force_pow_2_scales rather than silently ignoring it; - // the unfused per-tensor split-quantize path still honors it. - NVTE_CHECK(!force_pow_2_scales, - "Fused grouped FP8 block-scaling quantize does not support force_pow_2_scales=True. " - "Set force_pow_2_scales=False, or use the unfused split-quantize path " - "(NVTE_GROUPED_LINEAR_USE_FUSED_GROUPED_GEMM=0) which supports power-of-2 scales."); - const auto tensor_offsets = resolve_grouped_tensor_offsets(num_tensors, first_dims, last_dims, precomputed_tensor_offsets, logical_first_dim, logical_last_dim); diff --git a/transformer_engine/pytorch/distributed.py b/transformer_engine/pytorch/distributed.py index c050f26869..d1525b53f0 100644 --- a/transformer_engine/pytorch/distributed.py +++ b/transformer_engine/pytorch/distributed.py @@ -6,7 +6,7 @@ from __future__ import annotations from collections.abc import Iterable -from contextlib import contextmanager, AbstractContextManager, ContextDecorator +from contextlib import contextmanager, AbstractContextManager, ContextDecorator, nullcontext from functools import lru_cache from dataclasses import dataclass import math @@ -926,7 +926,10 @@ def fork(self, name: str = "model-parallel-rng"): def reduce_scatter_along_first_dim( - inp: torch.Tensor, tp_group: dist_group_type, async_op: bool = False + inp: torch.Tensor, + tp_group: dist_group_type, + async_op: bool = False, + output: torch.Tensor = None, ) -> Tuple[torch.Tensor, Optional[torch.distributed.Work]]: """Reduce-scatter the input tensor across model parallel group.""" world_size = get_distributed_world_size(tp_group) @@ -944,7 +947,8 @@ def reduce_scatter_along_first_dim( dim_size[0] = dim_size[0] // world_size - output = torch.empty(dim_size, dtype=inp.dtype, device=torch.cuda.current_device()) + if output is None: + output = torch.empty(dim_size, dtype=inp.dtype, device=torch.cuda.current_device()) handle = torch.distributed.reduce_scatter_tensor( output, inp.contiguous(), group=tp_group, async_op=async_op ) @@ -1311,7 +1315,8 @@ def wait(self) -> None: """Wait for the async operation to complete and post-process the tensor.""" if self._synchronized: return - self.async_handle.wait() + if self.async_handle is not None: + self.async_handle.wait() _post_process_nvfp4_gather( self.output, self.columnwise_data_interleaved, @@ -1328,6 +1333,8 @@ def _all_gather_nvfp4( async_op: bool = False, quantizer: NVFP4Quantizer, out_shape: Optional[list[int]] = None, + output_tensor=None, + external_coalescing=False, ) -> tuple[NVFP4TensorStorage, Optional[torch.distributed.Work]]: """All-gather NVFP4 tensor along first dimension.""" @@ -1404,15 +1411,23 @@ def _all_gather_nvfp4( inp = quantizer(inp.dequantize(dtype=dtype)) # Construct NVFP4 output tensor - out = quantizer.make_empty(out_shape, dtype=dtype, device=device) + if output_tensor is not None: + out = output_tensor + else: + out = quantizer.make_empty(out_shape, dtype=dtype, device=device) # Coalesce NCCL collectives for gathering data and scale inverses. - with torch.distributed._coalescing_manager( - group=process_group, - device=device, - async_ops=async_op, - ) as gather_coalescing_manager: + if not external_coalescing: + gather_coalescing_manager = torch.distributed._coalescing_manager( + group=process_group, + device=device, + async_ops=async_op, + ) + else: + # Caller owns an outer coalescing manager (managers cannot nest); step aside. + gather_coalescing_manager = nullcontext() + with gather_coalescing_manager as coalesced_handle: # Gather NVFP4 data for row-wise usage if quantizer.rowwise_usage: @@ -1493,10 +1508,10 @@ def _all_gather_nvfp4( # Transfer amax to output. out._amax_columnwise = inp._amax_columnwise - handle = gather_coalescing_manager if async_op else None + handle = coalesced_handle if async_op else None # Fixes interleaved data for transposed tensor/scale inv and pads scale inv if needed. - if async_op and quantizer.columnwise_usage: + if (async_op or external_coalescing) and quantizer.columnwise_usage: handle = _NVFP4AllGatherAsyncHandle( out, out_columnwise_data, out_scale_inv, world_size, handle ) @@ -1513,6 +1528,8 @@ def _all_gather_mxfp8( async_op: bool = False, quantizer: MXFP8Quantizer, out_shape: Optional[list[int]] = None, + output_tensor: torch.Tensor = None, + external_coalescing: bool = False, ) -> tuple[MXFP8TensorStorage, Optional[torch.distributed.Work]]: """All-gather MXFP8 tensor along first dimension.""" @@ -1578,15 +1595,23 @@ def _all_gather_mxfp8( inp = quantizer(inp.dequantize(dtype=dtype)) # Construct MXFP8 output tensor - out = quantizer.make_empty(out_shape, dtype=dtype, device=device) + if output_tensor is not None: + out = output_tensor + else: + out = quantizer.make_empty(out_shape, dtype=dtype, device=device) - # Coalesce NCCL collectives - with torch.distributed._coalescing_manager( - group=process_group, - device=device, - async_ops=async_op, - ) as coalescing_manager: + if not external_coalescing: + # Coalesce NCCL collectives for gathering data and scale inverses. + gather_coalescing_manager = torch.distributed._coalescing_manager( + group=process_group, + device=device, + async_ops=async_op, + ) + else: + # Caller owns an outer coalescing manager (managers cannot nest); step aside. + gather_coalescing_manager = nullcontext() + with gather_coalescing_manager as coalesced_handle: # Gather MXFP8 data for row-wise usage if quantizer.rowwise_usage: @@ -1633,7 +1658,7 @@ def _all_gather_mxfp8( group=process_group, ) - handle = coalescing_manager if async_op else None + handle = coalesced_handle if async_op else None return out, handle @@ -1642,9 +1667,17 @@ def gather_along_first_dim( process_group: dist_group_type, async_op: bool = False, quantizer: Optional[Quantizer] = None, + output_tensor: torch.Tensor = None, + external_coalescing: bool = False, ) -> tuple[torch.Tensor, Optional[torch.distributed.Work]]: """ All-gather tensors and concatenate along first dimension. + + ``external_coalescing``: composability flag for callers that batch several gathers into + one outer ``torch.distributed._coalescing_manager``. Coalescing managers cannot nest, so + when set this call skips opening its own manager and defers any post-gather fixup + (e.g. NVFP4 columnwise de-interleave) into the returned handle; ``handle.wait()`` completes + it once the outer manager has closed. Leave ``False`` for standalone gathers. """ # Return immediately if no communication is required @@ -1732,6 +1765,8 @@ def gather_along_first_dim( async_op=async_op, quantizer=quantizer, out_shape=out_shape, + output_tensor=output_tensor, + external_coalescing=external_coalescing, ) # NVFP4 case @@ -1746,6 +1781,8 @@ def gather_along_first_dim( async_op=async_op, quantizer=quantizer, out_shape=out_shape, + output_tensor=output_tensor, + external_coalescing=external_coalescing, ) # High-precision communication for quantized tensors @@ -1775,19 +1812,20 @@ def gather_along_first_dim( inp = inp.dequantize() # Communication for plain PyTorch tensors - out = torch.empty( - out_shape, - dtype=inp.dtype, - device=inp.device, - memory_format=torch.contiguous_format, - ) + if output_tensor is None: + output_tensor = torch.empty( + out_shape, + dtype=inp.dtype, + device=inp.device, + memory_format=torch.contiguous_format, + ) handle = torch.distributed.all_gather_into_tensor( - out, + output_tensor, inp.contiguous(), group=process_group, async_op=async_op, ) - return out, handle + return output_tensor, handle # Global cache to store symmetric memory tensors @@ -1842,13 +1880,54 @@ def get_symmetric_memory_tensor(tensor_numel, tensor_dtype, tensor_device, tp_gr return msg +_SYMM_MEM_POOL = None +_SYMM_MEM_POOL_BACKEND = None + + +def _get_symm_mem_pool(device: torch.device, backend: str = "NCCL"): + """Process-wide torch MemPool backed by the symmetric-memory allocator, created once (each rank + drives one device). The pool/allocator captures the backend at creation and there is no per-pool + backend arg, so the (process-global) backend is always set before the pool is created. The + collective rendezvous cost is amortized across allocations (paid per new segment, not per buffer). + """ + global _SYMM_MEM_POOL, _SYMM_MEM_POOL_BACKEND + if _SYMM_MEM_POOL is None: + symm_mem.set_backend(backend) + _SYMM_MEM_POOL_BACKEND = backend + if hasattr(symm_mem, "get_mem_pool"): + _SYMM_MEM_POOL = symm_mem.get_mem_pool(device) + elif hasattr(torch.cuda, "MemPool") and hasattr(symm_mem, "get_mempool_allocator"): + _SYMM_MEM_POOL = torch.cuda.MemPool(symm_mem.get_mempool_allocator(device)) + else: + raise RuntimeError( + "No symmetric-memory MemPool API available (need torch symm-mem get_mem_pool, or " + "torch.cuda.MemPool + get_mempool_allocator)." + ) + elif backend != _SYMM_MEM_POOL_BACKEND: + raise RuntimeError( + f"symm-mem pool already created with backend {_SYMM_MEM_POOL_BACKEND!r}; " + f"cannot switch to {backend!r}" + ) + return _SYMM_MEM_POOL + + def symm_mem_alloc( shape, dtype: torch.dtype, ep_group: dist_group_type, device: Optional[torch.device] = None, + use_pool: bool = False, + backend: str = "NCCL", ) -> torch.Tensor: - """Allocate and rendezvous a symm-mem buffer on ep_group. Collective on ep_group.""" + """Allocate a symm-mem buffer on ep_group. + + ``use_pool=False`` (default): freshly allocate and do one explicit collective ``rendezvous`` per + buffer, for caller-owned static buffers (fp8 zero-copy). ``use_pool=True``: allocate from a + process-wide symm-mem MemPool whose segments are auto-registered (implicit mempool), so no explicit + rendezvous is needed and torch manages the tensor lifecycle (freed back to the pool) — for + lifecycle-managed zero-copy, e.g. bf16, where the recv buffer is saved for backward and so cannot + be a shared static buffer. ``backend`` selects the symm-mem backend (default NCCL; for the pool it + is captured at pool creation).""" if device is None: device = torch.device("cuda", torch.cuda.current_device()) if not HAS_TORCH_SYMMETRIC: @@ -1856,10 +1935,15 @@ def symm_mem_alloc( "torch.distributed._symmetric_memory is unavailable; symm_mem_alloc " "requires PyTorch built with NCCL symm-mem support." ) - if symm_mem.get_backend(device) != "NCCL": - symm_mem.set_backend("NCCL") - t = symm_mem.empty(*shape, dtype=dtype, device=device) - symm_mem.rendezvous(t, group=ep_group) + if use_pool: + pool = _get_symm_mem_pool(device, backend) + with torch.cuda.use_mem_pool(pool): + t = torch.empty(*shape, dtype=dtype, device=device) + else: + if symm_mem.get_backend(device) != backend: + symm_mem.set_backend(backend) + t = symm_mem.empty(*shape, dtype=dtype, device=device) + symm_mem.rendezvous(t, group=ep_group) return t diff --git a/transformer_engine/pytorch/distributed_weight.py b/transformer_engine/pytorch/distributed_weight.py new file mode 100644 index 0000000000..0d491430e5 --- /dev/null +++ b/transformer_engine/pytorch/distributed_weight.py @@ -0,0 +1,117 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""GTP-agnostic weight-parallelism extension point. + +TE owns this contract but ships no implementation; the caller (e.g. Megatron GTP) implements the +protocol on the weight and injects it at construction. Dispatchers are list-shaped (Linear -> 1, +GroupedLinear -> N; leader is ``weights[0]``) and no-op on plain tensors. +""" + +from typing import Any, List, Protocol, runtime_checkable + +import torch + +__all__ = [ + "DistributedWeight", + "is_distributed_weight", + "materialize_weight_for_forward", + "materialize_weight_for_backward", + "finalize_weight_grads", +] + + +@runtime_checkable +class DistributedWeight(Protocol): + """Structural interface for a custom-weight-parallel weight (AG for the GEMM, reduce/RS the + grad, re-materialize in backward). Duck-typed ``typing.Protocol``: implementers need not + subclass it, and all state (shards, group, async handles) lives outside TE on the implementer. + + Implementers MUST be ``torch.Tensor`` subclasses (needed by ``ctx.save_for_backward``, DDP + backward hooks, and ``torch.compile``); enforced at runtime by :func:`is_distributed_weight`. + """ + + # Capability marker: True on an implementer, absent on plain tensors; TE's fwd/bwd gate on it. + is_distributed_weight: bool + + def materialize_group_for_forward(self) -> Any: + """Return the tensor(s) to feed the forward GEMM (may all-gather shards).""" + + def materialize_group_for_backward(self) -> Any: + """Re-materialize the full weight(s) for the backward GEMMs.""" + + def finalize_group_grads(self, wgrads: Any) -> Any: + """Post-process freshly computed weight grad(s) (e.g. reduce-scatter). + + May consume ``wgrads`` in-place -- reduce-scatter into ``main_grad`` and set + ``grad_added_to_main_grad`` -- returning a dummy grad (or ``None`` for an async collective) + that callers use as the parameter grad(s) or discard. + """ + + def grad_buffer(self) -> torch.Tensor: + """The gradient accumulation buffer for this weight.""" + + +def is_distributed_weight(weight: Any) -> bool: + """True if ``weight`` participates in custom weight parallelism (False on plain tensors). + + Enforces the :class:`DistributedWeight` requirement that an implementer be a ``torch.Tensor`` + subclass, failing loudly here rather than silently breaking autograd downstream. + """ + flag = bool(getattr(weight, "is_distributed_weight", False)) + if flag and not isinstance(weight, torch.Tensor): + raise TypeError( + "DistributedWeight implementers must be torch.Tensor subclasses; got " + f"{type(weight).__name__}." + ) + return flag + + +def materialize_weight_for_forward(weights: Any) -> List[Any]: + """Prepare the weight(s) fed to the forward GEMM, always returned as a list. + + Args: + weights: the module's weight(s) -- a single weight (Linear) or the full per-expert list + (GroupedLinear). A bare weight is treated as a one-element list. + + Returns: + - Distributed group: the leader ``weights[0]`` all-gathers/coalesces the whole group and + returns all N materialized weights; the follower entries ``weights[1:]`` are ignored + here (the leader already holds references to its group). + - Otherwise: the input weights, unchanged. + """ + if not isinstance(weights, (list, tuple)): + weights = [weights] + leader = weights[0] + if is_distributed_weight(leader): + out = leader.materialize_group_for_forward() + return list(out) if isinstance(out, (list, tuple)) else [out] + return list(weights) + + +def materialize_weight_for_backward(weights: Any) -> List[Any]: + """Backward-GEMM mirror of :func:`materialize_weight_for_forward` (same contract).""" + if not isinstance(weights, (list, tuple)): + weights = [weights] + leader = weights[0] + if is_distributed_weight(leader): + out = leader.materialize_group_for_backward() + return list(out) if isinstance(out, (list, tuple)) else [out] + return list(weights) + + +def finalize_weight_grads(weights: Any, wgrads: List[Any]) -> List[Any]: + """Finalize a weight group's grad(s), mirroring :func:`materialize_weight_for_backward`. + + Delegates to the leader's :meth:`DistributedWeight.finalize_group_grads` (which defines the + in-place / dummy / async-``None`` return contract); returns ``wgrads`` unchanged when not + distributed. + """ + if not isinstance(weights, (list, tuple)): + weights = [weights] + leader = weights[0] + if is_distributed_weight(leader): + out = leader.finalize_group_grads(wgrads if len(wgrads) > 1 else wgrads[0]) + return list(out) if isinstance(out, (list, tuple)) else [out] + return list(wgrads) diff --git a/transformer_engine/pytorch/ep.py b/transformer_engine/pytorch/ep.py index e57caa1f42..0ae4805caa 100644 --- a/transformer_engine/pytorch/ep.py +++ b/transformer_engine/pytorch/ep.py @@ -25,6 +25,7 @@ "ep_dispatch", "ep_combine", "symm_mem_alloc", + "is_symm_backed", ] @@ -158,18 +159,32 @@ def ep_finalize() -> None: _EP_GROUP = None +def is_symm_backed(t: torch.Tensor) -> bool: + """Whether ``t`` is symm-mem-backed on the EP group. Prefer torch's local ``is_symm_mem_tensor`` + when the build provides it (no collective, no exception); otherwise fall back to the rendezvous + probe the C++ ep kernel uses (``maybe_make_window``): cached for an already-registered tensor, + raises for a plain one.""" + from torch.distributed import _symmetric_memory as _symm + + if hasattr(_symm, "is_symm_mem_tensor"): + return bool(_symm.is_symm_mem_tensor(t)) + if _EP_GROUP is None: + raise RuntimeError( + "is_symm_backed called before ensure_nccl_ep_bootstrapped(); no EP group registered." + ) + try: + _symm.rendezvous(t, _EP_GROUP.group_name) + return True + except Exception: # pylint: disable=broad-exception-caught + return False + + # Buffer class EpBuffer: - """Per-microbatch EP layer state holding handle_mem and token_counts. + """Per-microbatch EP layer state: handle_mem, token_counts, and shape/dtype config. Use one EpBuffer per concurrently-in-flight call (e.g. per PP-1F1B microbatch). - - In zero-copy mode the buffer owns the symm-mem buffers the one-sided path - requires: the dispatch recv outputs (recv_tokens, recv_topk_weights) and the - combine backward grad target. One set per buffer, so each layer/microbatch is - isolated. In normal mode these are None and allocated in-flight instead (recv - outputs in the dispatch forward, the combine grad in the backward). """ __slots__ = ( @@ -184,39 +199,8 @@ class EpBuffer: "device", "token_counts", "zero_copy", - "recv_tokens_symm_buf", - "recv_topk_weights_symm_buf", - "grad_expert_out_symm_buf", ) - def _alloc_symm_buffers(self) -> None: - """Fill in buffer-owned symm-mem buffers the caller did not supply. - recv_topk_weights is always owned. In normal mode caller-supplied - tensors are kept as-is and the rest stay None (allocated in-flight).""" - if not self.zero_copy: - self.recv_topk_weights_symm_buf = None - return - if _EP_GROUP is None: - raise RuntimeError( - "ep_bootstrap must be called before constructing a zero-copy EpBuffer" - ) - rc, h = self.recv_capacity_per_rank, self.hidden_dim - # Persistent across microbatches; keep resident under CPU offloading. - self.recv_topk_weights_symm_buf = symm_mem_alloc( - (rc,), torch.float32, _EP_GROUP, device=self.device - ) - mark_not_offload(self.recv_topk_weights_symm_buf) - if self.recv_tokens_symm_buf is None: - self.recv_tokens_symm_buf = symm_mem_alloc( - (rc, h), self.payload_dtype, _EP_GROUP, device=self.device - ) - mark_not_offload(self.recv_tokens_symm_buf) - if self.grad_expert_out_symm_buf is None: - self.grad_expert_out_symm_buf = symm_mem_alloc( - (rc, h), self.payload_dtype, _EP_GROUP, device=self.device - ) - mark_not_offload(self.grad_expert_out_symm_buf) - def __init__( self, top_k: int, @@ -227,14 +211,7 @@ def __init__( alignment: int = 0, payload_dtype: torch.dtype = torch.bfloat16, device: Optional[torch.device] = None, - dispatch_recv_tokens: Optional[torch.Tensor] = None, - combine_grad_expert_out: Optional[torch.Tensor] = None, ) -> None: - """Pass ``dispatch_recv_tokens`` (dispatch recv output) and/or - ``combine_grad_expert_out`` (combine backward grad target) to use caller-owned - buffers; the buffer then skips allocating them. Both must be symm-mem-backed - under zero-copy. Whatever is left None is buffer-owned (zero-copy) or allocated - in-flight (normal mode). recv_topk_weights is always owned by the buffer.""" if device is None: device = torch.device("cuda", torch.cuda.current_device()) alignment = int(alignment) @@ -249,15 +226,12 @@ def __init__( self.payload_dtype = payload_dtype self.device = device self.zero_copy = bool(tex.ep_get_zero_copy()) - self.recv_tokens_symm_buf = dispatch_recv_tokens - self.grad_expert_out_symm_buf = combine_grad_expert_out size_bytes = tex.ep_handle_mem_size(self.top_k, self.alignment) self.handle_mem = torch.empty(int(size_bytes), dtype=torch.uint8, device=device) self.token_counts = torch.empty(self.num_local_experts, dtype=torch.int32, device=device) # Persistent tensor; keep resident if activation CPU offloading is on. mark_not_offload(self.handle_mem) - self._alloc_symm_buffers() # torch.library custom ops (so they don't graph-break under torch.compile) @@ -476,14 +450,15 @@ def backward(ctx, g_recv_tokens, g_recv_topk_weights, _g_token_counts): # type: class _EpCombine(torch.autograd.Function): """Autograd combine. - bwd scatters the expert_out grad into ``grad_symm_buf`` (EpBuffer-owned - symm-mem, one-sided) in zero-copy mode, or into a plain tensor allocated - in-flight here otherwise. The latter keeps allocation torch.compile / - CUDA-graph safe and lets autograd own the grad's lifetime. + bwd scatters the expert_out grad into ``grad_out``. When the caller supplies it + (mcore-managed mode) that buffer is used as-is; otherwise it is allocated on the + fly here — from the symm-mem pool in zero-copy mode (one-sided target), or a plain + tensor in normal mode (keeps allocation torch.compile / CUDA-graph safe and lets + autograd own the grad's lifetime). - ``grad_symm_buf`` is the backward's scatter target (an output it writes, never - reads), so it is stashed as a plain ctx attribute rather than via - save_for_backward, which would version-track a tensor we mutate. + ``grad_out`` is the backward's scatter target (an output it writes, never reads), + so it is stashed as a plain ctx attribute rather than via save_for_backward, which + would version-track a tensor we mutate. """ @staticmethod @@ -492,7 +467,7 @@ def forward( # type: ignore[override] handle_mem: torch.Tensor, num_local_tokens: int, hidden_dim: int, - grad_symm_buf: Optional[torch.Tensor], + grad_out: Optional[torch.Tensor], expert_out: torch.Tensor, ): """Combine fwd; stashes the bwd grad target or expert_out shape to size it.""" @@ -500,8 +475,8 @@ def forward( # type: ignore[override] result = torch.empty(num_local_tokens, hidden_dim, dtype=expert_out.dtype, device=device) torch.ops.transformer_engine_ep.combine(handle_mem, expert_out, result) ctx.save_for_backward(handle_mem) - ctx.grad_symm_buf = grad_symm_buf - if grad_symm_buf is None: + ctx.grad_out = grad_out + if grad_out is None: ctx.expert_out_shape = expert_out.shape ctx.expert_out_dtype = expert_out.dtype ctx.device = device @@ -513,17 +488,17 @@ def backward(ctx, g_result): # type: ignore[override] if not g_result.is_contiguous(): g_result = g_result.contiguous() (handle_mem,) = ctx.saved_tensors - grad_expert_out = ctx.grad_symm_buf + grad_expert_out = ctx.grad_out if grad_expert_out is None: - grad_expert_out = torch.empty( - ctx.expert_out_shape, dtype=ctx.expert_out_dtype, device=ctx.device + grad_expert_out = _alloc_io( + ctx.expert_out_shape, ctx.expert_out_dtype, ctx.device, tex.ep_get_zero_copy() ) torch.ops.transformer_engine_ep.combine_bwd(handle_mem, g_result, grad_expert_out) return ( None, # handle_mem None, # num_local_tokens None, # hidden_dim - None, # grad_symm_buf + None, # grad_out grad_expert_out, ) @@ -539,18 +514,33 @@ def _require_bf16(name: str, t: torch.Tensor) -> None: ) +def _alloc_io(shape, dtype: torch.dtype, device, zero_copy: bool) -> torch.Tensor: + """Allocate a dispatch/combine IO tensor the caller did not supply: from the symm-mem pool in + zero-copy mode (auto-registered segment, lifecycle managed by torch refcount), else plain.""" + if zero_copy: + t = symm_mem_alloc(shape, dtype, _EP_GROUP, device=device, use_pool=True) + # symm-mem storage is non-resizable; exempt it from CPU activation offloading (which + # releases via storage.resize_(0)). Matters for bf16 recv_tokens (the saved activation). + mark_not_offload(t) + return t + return torch.empty(*shape, dtype=dtype, device=device) + + def ep_dispatch( buffer: EpBuffer, tokens: torch.Tensor, topk_idx: torch.Tensor, topk_weights: torch.Tensor, + *, + recv_tokens: Optional[torch.Tensor] = None, + recv_topk_weights: Optional[torch.Tensor] = None, ): """Prepare + dispatch with autograd. topk_idx must be int32 or int64. - recv_tokens comes from the EpBuffer (caller-supplied or buffer-owned under - zero-copy) or is allocated in-flight (normal mode). recv_topk_weights is always - owned by the buffer. Returns (recv_tokens, recv_topk_weights, token_counts); - token_counts is non-diff. + ``recv_tokens`` / ``recv_topk_weights`` are the dispatch recv outputs: pass caller-owned buffers + (mcore-managed mode; in zero-copy they must be symm-mem-backed) or leave them None to allocate on + the fly (zero-copy: symm-mem pool; normal: plain). Returns (recv_tokens, recv_topk_weights, + token_counts); token_counts is non-diff. """ _require_bf16("tokens", tokens) if topk_weights.dtype is not torch.float32: @@ -558,19 +548,17 @@ def ep_dispatch( f"topk_weights must be float32; got dtype={topk_weights.dtype}. " "Cast with topk_weights.float() before calling." ) - recv_tokens = buffer.recv_tokens_symm_buf if recv_tokens is None: - recv_tokens = torch.empty( - buffer.recv_capacity_per_rank, - buffer.hidden_dim, - dtype=buffer.payload_dtype, - device=buffer.device, + recv_tokens = _alloc_io( + (buffer.recv_capacity_per_rank, buffer.hidden_dim), + buffer.payload_dtype, + buffer.device, + buffer.zero_copy, + ) + if recv_topk_weights is None: + recv_topk_weights = _alloc_io( + (buffer.recv_capacity_per_rank,), torch.float32, buffer.device, buffer.zero_copy ) - recv_topk_weights = ( - buffer.recv_topk_weights_symm_buf - if buffer.zero_copy - else torch.empty(buffer.recv_capacity_per_rank, dtype=torch.float32, device=buffer.device) - ) return _EpDispatch.apply( buffer.handle_mem, buffer.top_k, @@ -589,22 +577,23 @@ def ep_combine( expert_out: torch.Tensor, *, num_local_tokens: Optional[int] = None, + grad_out: Optional[torch.Tensor] = None, ): """Combine with autograd; caller pre-applies topk weighting. - The backward scatters the expert_out grad into the EpBuffer grad target - (caller-supplied or buffer-owned under zero-copy), or a tensor allocated - in-flight (normal mode). Result shape is (num_local_tokens, buffer.hidden_dim); - defaults to buffer.max_tokens_per_rank rows. + ``expert_out`` is the combine input (always caller-supplied; in zero-copy it must be symm-mem- + backed). ``grad_out`` is the backward's grad target: pass a caller-owned buffer (mcore-managed + mode) or leave it None to allocate on the fly in the backward (zero-copy: symm-mem pool; normal: + plain). Result shape is (num_local_tokens, buffer.hidden_dim); defaults to + buffer.max_tokens_per_rank rows. """ _require_bf16("expert_out", expert_out) if num_local_tokens is None: num_local_tokens = buffer.max_tokens_per_rank - grad_expert_out = buffer.grad_expert_out_symm_buf return _EpCombine.apply( buffer.handle_mem, num_local_tokens, buffer.hidden_dim, - grad_expert_out, + grad_out, expert_out, ) diff --git a/transformer_engine/pytorch/module/grouped_linear.py b/transformer_engine/pytorch/module/grouped_linear.py index 82ee7953c1..32963c07ca 100644 --- a/transformer_engine/pytorch/module/grouped_linear.py +++ b/transformer_engine/pytorch/module/grouped_linear.py @@ -46,6 +46,12 @@ is_fp8_activation_recompute_enabled, in_fp8_activation_recompute_phase, ) +from ..distributed_weight import ( + is_distributed_weight, + materialize_weight_for_forward, + materialize_weight_for_backward, + finalize_weight_grads, +) from ..cpp_extensions import ( general_grouped_gemm, general_grouped_gemm_for_grouped_tensor, @@ -55,7 +61,13 @@ from ..cpu_offload import is_cpu_offload_enabled, mark_not_offload, start_offload from ..triton.grouped_dbias_dscales import compute_grouped_dbias -from ..tensor import Float8CurrentScalingQuantizer, Float8Quantizer, MXFP8Quantizer, NVFP4Quantizer +from ..tensor import ( + Float8BlockQuantizer, + Float8CurrentScalingQuantizer, + Float8Quantizer, + MXFP8Quantizer, + NVFP4Quantizer, +) from ..quantized_tensor import ( QuantizedTensorStorage, Quantizer, @@ -103,11 +115,15 @@ def _is_grouped_tensor_path_supported( and be incompatible with CUDA Graphs. Supported Compute Capability (CC) and precisions: - * Hopper (CC 9.0): BF16/FP16 and FP8 per-tensor current scaling. + * Hopper (CC 9.0): BF16/FP16, FP8 per-tensor current scaling, and FP8 + block scaling (1D/2D, including power-of-2 scales). * Blackwell (CC 10.x and 11.0): BF16/FP16/MXFP8/NVFP4 with RHT and FP8 per-tensor current scaling. - FP8 delayed scaling and FP8 block scaling are not supported because the - corresponding grouped quantization kernels are missing. + FP8 delayed scaling is not supported because the corresponding grouped + quantization kernels are missing. FP8 block scaling on Blackwell (SM100 and + SM110) raises instead of falling back: the fused path is Hopper-only and has + no MXFP8-broadcast emulation. Architectures outside the fused-path window + (e.g. SM120) fall back to the legacy path like every other recipe. Grouped GEMM requires cuBLAS 13.3+ (13.4+ on Hopper, 13.5+ for FP8 per-tensor current scaling on Hopper); otherwise the legacy path is used. Non-RHT NVFP4 falls back to the legacy path because graph-safe grouped quantization @@ -147,6 +163,21 @@ def _is_grouped_tensor_path_supported( if device_capability < (10, 0) and cublaslt_version < 130500: return False return True + if all(isinstance(q, Float8BlockQuantizer) for q in input_quantizers): + # Grouped FP8 block-scaling quantize kernels and cuBLASLt grouped GEMM + # scale modes are Hopper-only, and the fused path has no MXFP8-broadcast + # emulation. On Blackwell (SM100/SM110, the only other arch that reaches + # this branch) fail loudly rather than silently falling back to the + # unfused path the user explicitly opted out of. + if get_device_compute_capability() >= (10, 0): + raise RuntimeError( + "NVTE_GROUPED_LINEAR_USE_FUSED_GROUPED_GEMM=1 does not support the" + " FP8 block-scaling recipe on Blackwell GPUs: the fused grouped" + " FP8 block-scaling path is Hopper-only. Unset" + " NVTE_GROUPED_LINEAR_USE_FUSED_GROUPED_GEMM to use the unfused" + " path (emulated via MXFP8 GEMM on Blackwell)." + ) + return True # MXFP8 and NVFP4 require Blackwell+. if not (10, 0) <= device_capability <= (11, 0): return False @@ -236,6 +267,41 @@ def _prepare_weights_for_grouped_tensor_gemm( weights_for_gemm.append(weight_fp8) return weights_for_gemm, new_workspaces + @staticmethod + def _validate_or_alloc_output( + buffer: Optional[torch.Tensor], + rows: int, + cols: int, + dtype: torch.dtype, + device: torch.device, + ) -> torch.Tensor: + """Validate and return the caller's output buffer, or allocate one if it is None. + + The buffer must be a 2D, contiguous, non-grad tensor matching the required shape, + dtype, and device. Validation reads host-side metadata only, with no device sync. + """ + if buffer is None: + return torch.empty((rows, cols), dtype=dtype, device=device) + if buffer.dim() != 2: + raise ValueError(f"Output buffer must be 2D, got {buffer.dim()}D.") + if buffer.size(0) != rows: + raise ValueError(f"Output buffer rows {buffer.size(0)} must match input rows {rows}.") + if buffer.size(1) != cols: + raise ValueError( + f"Output buffer last dim {buffer.size(1)} does not match required {cols}." + ) + if buffer.dtype != dtype: + raise ValueError(f"Output buffer dtype {buffer.dtype} does not match required {dtype}.") + if buffer.device != device: + raise ValueError( + f"Output buffer device {buffer.device} does not match required {device}." + ) + if not buffer.is_contiguous(): + raise ValueError("Output buffer must be contiguous.") + if buffer.requires_grad: + raise ValueError("Output buffer must not require gradient.") + return buffer + @staticmethod def _forward_grouped_tensor( ctx, @@ -259,6 +325,8 @@ def _forward_grouped_tensor( skip_fp8_weight_update: Optional[torch.Tensor], weights: Tuple[torch.Tensor, ...], biases: Tuple[torch.Tensor, ...], + out: Optional[torch.Tensor] = None, + dgrad_out: Optional[torch.Tensor] = None, ) -> Tuple[torch.Tensor, list]: """Forward path backed by GroupedTensor + cuBLASLt grouped GEMM.""" num_gemms = len(m_splits) @@ -303,10 +371,12 @@ def _forward_grouped_tensor( cache_weight=cache_weight, ) - out = torch.empty( - [x.size(0), out_features], - dtype=activation_dtype, - device=device, + out = _GroupedLinear._validate_or_alloc_output( + out, + x.size(0), + out_features, + activation_dtype, + device, ) grouped_out = _GroupedLinear._make_grouped_tensor( out, @@ -381,6 +451,7 @@ def _forward_grouped_tensor( lambda j=i: weights[j].main_grad for i in range(num_gemms) ] ctx.device = device + ctx.dgrad_out = dgrad_out ctx.m_splits = None ctx.num_gemms = num_gemms ctx.activation_dtype = activation_dtype @@ -413,6 +484,8 @@ def forward( inp: torch.Tensor, m_splits: torch.Tensor, non_tensor_args: Tuple, + out: Optional[torch.Tensor], + dgrad_out: Optional[torch.Tensor], *weights_and_biases, ) -> Tuple[torch.Tensor, list]: # pylint: disable=missing-function-docstring @@ -457,6 +530,11 @@ def forward( device = inp.device weight_requires_grad = weights[0].requires_grad + origin_weights = weights + is_dist_weight = is_distributed_weight(weights[0]) + if is_dist_weight: + weights = materialize_weight_for_forward(weights) + # Configure quantizers if save_original_input and isinstance(input_quantizers[0], Float8Quantizer): if FP8GlobalStateManager.get_fp8_recipe().custom(): @@ -544,6 +622,8 @@ def forward( skip_fp8_weight_update=skip_fp8_weight_update, weights=weights, biases=biases, + out=out, + dgrad_out=dgrad_out, ) # Convert splits to list of ints for compatibility with split functions @@ -598,10 +678,12 @@ def forward( bias_dtype = torch.bfloat16 # FP8 GEMM only supports BF16/FP16 bias biases = [cast_if_needed(bias, bias_dtype) for bias in biases] if use_bias else biases # Initialize output tensor - out = torch.empty( - [sum(m_splits), weights_fp8[0].size(0)], - dtype=activation_dtype, - device=device, + out = _GroupedLinear._validate_or_alloc_output( + out, + sum(m_splits), + weights_fp8[0].size(0), + activation_dtype, + device, ) # Choose whether to use split accumulator @@ -663,6 +745,10 @@ def forward( if backward_override == "high_precision" and inp.requires_grad else [None] * num_gemms ) + if is_dist_weight: + # GTP: gathered workspace is transient (re-gathered in backward), don't save it. + weights_fp8 = [None] * num_gemms + saved_weights = origin_weights tensors_to_save, tensor_objects = prepare_for_saving( *inputmats, *weights_fp8, @@ -690,6 +776,8 @@ def forward( if hasattr(weights[0], "__fsdp_param__"): # MCore FSDP creates main_grad lazily before backward ctx.main_grad_funcs = [weights[i].get_main_grad for i in range(num_gemms)] + elif is_dist_weight: + ctx.main_grad_funcs = [origin_weights[i].grad_buffer for i in range(num_gemms)] else: ctx.main_grad_funcs = [ lambda j=i: weights[j].main_grad for i in range(num_gemms) @@ -719,6 +807,7 @@ def forward( ctx.debug = debug ctx.save_original_input = save_original_input ctx.input_quantizers = input_quantizers + ctx.dgrad_out = dgrad_out # backward overrides if backward_override is not None: @@ -773,7 +862,12 @@ def _backward_grouped_tensor( columnwise=ctx.weights_requires_grad, ) grad_output_quantizer.optimize_for_gemm = True - if ctx.use_bias and isinstance(grad_output_quantizer, MXFP8Quantizer): + # The grouped FP8 block-scaling bgrad kernel computes dbias in the rowwise + # pass, so the fusion needs rowwise output (i.e. dgrad required). + fuse_bgrad = isinstance(grad_output_quantizer, MXFP8Quantizer) or ( + isinstance(grad_output_quantizer, Float8BlockQuantizer) and ctx.requires_dgrad + ) + if ctx.use_bias and fuse_bgrad: grouped_dy, dbias_packed = tex.bgrad_group_quantize( dy_2d, grad_output_quantizer, @@ -813,10 +907,12 @@ def _backward_grouped_tensor( for weight in weights: if isinstance(weight, QuantizedTensorStorage): weight.update_usage(columnwise_usage=True) - dgrad = torch.empty( - (dy_2d.size(0), ctx.weights_shape_1), - dtype=ctx.activation_dtype, - device=ctx.device, + dgrad = _GroupedLinear._validate_or_alloc_output( + ctx.dgrad_out, + dy_2d.size(0), + ctx.weights_shape_1, + ctx.activation_dtype, + ctx.device, ) grouped_dgrad = _GroupedLinear._make_grouped_tensor( dgrad, @@ -918,6 +1014,8 @@ def handle_custom_ddp_from_mcore(weight, main_grad, wgrad): dgrad.view(ctx.inp_shape) if ctx.requires_dgrad else None, None, # m_splits None, # non_tensor_args + None, # out + None, # dgrad_out *wgrad_list, *grad_biases, ) @@ -943,7 +1041,12 @@ def backward( # Only needed when fuse_wgrad_accumulation is enabled. origin_weights = [None] * N main_grads = [None] * N - if ctx.fuse_wgrad_accumulation and ctx.weights_requires_grad: + is_dist_weight = is_distributed_weight(saved_weights[0]) + if is_dist_weight: + origin_weights = saved_weights + if ctx.fuse_wgrad_accumulation and ctx.weights_requires_grad: + main_grads = [main_grad_func() for main_grad_func in ctx.main_grad_funcs] + elif ctx.fuse_wgrad_accumulation and ctx.weights_requires_grad: origin_weight_refs = ctx.origin_weight_refs ctx.origin_weight_refs = None origin_weights = [ref() if ref is not None else None for ref in origin_weight_refs] @@ -1004,13 +1107,18 @@ def backward( ctx.m_splits, ) - if ctx.is_first_microbatch is not None: + if is_dist_weight: + accumulate_wgrad_into_param_main_grad = False + elif 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 + if is_dist_weight: + weights = materialize_weight_for_backward(origin_weights) + if ctx.requires_dgrad: dgrad_gemm_use_split_accumulator = _2X_ACC_DGRAD if ctx.fp8 or ctx.debug: @@ -1019,10 +1127,12 @@ def backward( dgrad_gemm_use_split_accumulator = ( recipe.fp8_gemm_dgrad.use_split_accumulator ) - dgrad = torch.empty( - (sum(ctx.m_splits), ctx.weights_shape_1), - dtype=ctx.activation_dtype, - device=ctx.device, + dgrad = _GroupedLinear._validate_or_alloc_output( + ctx.dgrad_out, + sum(ctx.m_splits), + ctx.weights_shape_1, + ctx.activation_dtype, + ctx.device, ) weights_for_dgrad = weights if ctx.backward_override == "dequantized": @@ -1079,6 +1189,9 @@ def backward( device=ctx.device, ) wgrad_list = [wgrad_packed[i] for i in range(ctx.num_gemms)] + if is_dist_weight: + # Gathered weights are no longer needed after dgrad GEMM. + del weights if ctx.save_original_input: inp = inputmats[0] @@ -1129,7 +1242,8 @@ def backward( use_split_accumulator=wgrad_gemm_use_split_accumulator, accumulate=( accumulate_wgrad_into_param_main_grad - if not getattr(ctx, "origin_weights_overwrite_main_grad", False) + if not is_dist_weight + and not getattr(ctx, "origin_weights_overwrite_main_grad", False) else False ), ) @@ -1171,10 +1285,13 @@ def handle_custom_ddp_from_mcore(weight, main_grad, wgrad): wgrad = None return wgrad - wgrad_list = [ - handle_custom_ddp_from_mcore(weight, main_grad, wgrad) - for weight, main_grad, wgrad in zip(origin_weights, main_grads, wgrad_list) - ] + if is_dist_weight: + wgrad_list = finalize_weight_grads(origin_weights, wgrad_list) + else: + wgrad_list = [ + handle_custom_ddp_from_mcore(weight, main_grad, wgrad) + for weight, main_grad, wgrad in zip(origin_weights, main_grads, wgrad_list) + ] else: wgrad_list = [None] * ctx.num_gemms @@ -1191,6 +1308,8 @@ def handle_custom_ddp_from_mcore(weight, main_grad, wgrad): dgrad.view(ctx.inp_shape) if ctx.requires_dgrad else None, None, # m_splits None, # non_tensor_args + None, # out + None, # dgrad_out *wgrad_list, *grad_biases, ) @@ -1684,6 +1803,8 @@ def forward( inp: torch.Tensor, m_splits: torch.Tensor, is_first_microbatch: Optional[bool] = None, + out: Optional[torch.Tensor] = None, + dgrad_out: Optional[torch.Tensor] = None, ) -> Union[torch.Tensor, Tuple[torch.Tensor, ...]]: """ Apply the linear transformation to the input. @@ -1707,6 +1828,19 @@ def forward( * it also allows skipping gradient accumulation during the first microbatch (since it is the first gradient being produced) + out : torch.Tensor, default = None + Optional preallocated buffer for the forward output; the returned tensor + aliases it with no copy. Must be a 2D, contiguous, non-grad tensor of shape + [num_tokens, out_features] in the activation dtype. Only the first + sum(m_splits) rows are written; any padded trailing rows are left unchanged. + Can be given independently of dgrad_out. If the buffer is reused across + iterations, pass ``buffer.detach()`` so autograd does not set its + ``requires_grad`` (which would trip the non-grad check on the next call). + dgrad_out : torch.Tensor, default = None + Optional preallocated buffer for the backward input gradient, of shape + [num_tokens, in_features] with the same constraints as out. Receives the + final gradient only when inp has a single consumer in the autograd graph; + otherwise autograd accumulates into a new tensor. """ debug = self.is_debug_iter() is_grad_enabled = torch.is_grad_enabled() @@ -1805,7 +1939,14 @@ def forward( debug, ) out, new_workspaces = linear_fn( - *autograd_ctx, inp, m_splits, non_tensor_args, *weight_tensors, *bias_tensors + *autograd_ctx, + inp, + m_splits, + non_tensor_args, + out, + dgrad_out, + *weight_tensors, + *bias_tensors, ) if cache_weight: diff --git a/transformer_engine/pytorch/module/layernorm_linear.py b/transformer_engine/pytorch/module/layernorm_linear.py index 43799b003c..cffc5ba1e6 100644 --- a/transformer_engine/pytorch/module/layernorm_linear.py +++ b/transformer_engine/pytorch/module/layernorm_linear.py @@ -56,6 +56,12 @@ _fsdp_scatter_tensors, _fsdp_gather_tensors, ) +from ..distributed_weight import ( + is_distributed_weight, + materialize_weight_for_forward, + materialize_weight_for_backward, + finalize_weight_grads, +) from ..constants import FP8BwdTensorIdx, FP8FwdTensorIdx, GemmParallelModes, dist_group_type from ..jit import no_torch_dynamo from ..graph import is_graph_capturing @@ -310,6 +316,11 @@ def forward( # ------------------------------------------------------ # Prepare weight tensor # ------------------------------------------------------ + origin_weight = weight + is_dist_weight = is_distributed_weight(origin_weight) + if is_dist_weight: + weight = materialize_weight_for_forward(weight)[0] + out_features = weight.shape[0] new_weight_workspace = None weightmat = weight is_weight_param_quantized = False @@ -500,10 +511,15 @@ def forward( wt_save = weightmat if is_fsdp2 and weightmat is not weight: wt_save = None + # Distributed weight (e.g. GTP): don't save the gathered quantized workspace; + # backward re-gathers from the saved (sharded) weight and re-quantizes. + if is_dist_weight: + wt_save = None + tensors_to_save, tensor_objects = prepare_for_saving( inputmat, wt_save, - weight, + origin_weight, bias, ln_weight, ln_out_to_save, @@ -530,6 +546,8 @@ def forward( if hasattr(weight, "__fsdp_param__"): # MCore FSDP creates main_grad lazily before backward ctx.main_grad_func = weight.get_main_grad + elif is_dist_weight: + ctx.main_grad_func = origin_weight.grad_buffer else: ctx.main_grad_func = lambda: weight.main_grad ctx.grad_input_quantizer = grad_input_quantizer @@ -623,6 +641,9 @@ def backward( rsigma, ) = restore_from_func_ctx(ctx) + is_dist_weight = is_distributed_weight(saved_weight) + if is_dist_weight: + weight = materialize_weight_for_backward(saved_weight)[0] # Restore from weakref to get original weight python object # (preserves attributes like main_grad, grad_added_to_main_grad, etc.) # Only needed when fuse_wgrad_accumulation is enabled. @@ -640,7 +661,7 @@ def backward( ), "weight was removed while fuse_wgrad_accumulation=True" # 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 else None - if main_grad is not None: + if main_grad is not None and not is_dist_weight: origin_weight.main_grad = main_grad # Gather intermediate/activation tensors if needed @@ -984,7 +1005,10 @@ def backward( 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: + if is_dist_weight: + # Distributed weight (e.g. GTP): accumulation happens downstream in finalize. + accumulate_wgrad_into_param_main_grad = False + elif ctx.is_first_microbatch is not None: accumulate_wgrad_into_param_main_grad = ( ctx.fuse_wgrad_accumulation and not ctx.is_first_microbatch ) @@ -1056,6 +1080,9 @@ def wgrad_gemm( # Call wgrad GEMM now wgrad, grad_bias_ = wgrad_gemm(ln_out_total, grad_output) + if is_dist_weight: + wgrad = finalize_weight_grads(saved_weight, [wgrad])[0] + # Update grad bias if needed if grad_bias is None: grad_bias = grad_bias_ diff --git a/transformer_engine/pytorch/module/linear.py b/transformer_engine/pytorch/module/linear.py index 78a4d31852..cef6d3d533 100644 --- a/transformer_engine/pytorch/module/linear.py +++ b/transformer_engine/pytorch/module/linear.py @@ -55,6 +55,12 @@ _fsdp_scatter_tensors, _fsdp_gather_tensors, ) +from ..distributed_weight import ( + is_distributed_weight, + materialize_weight_for_forward, + materialize_weight_for_backward, + finalize_weight_grads, +) from ..cpp_extensions import ( general_gemm, ) @@ -265,6 +271,7 @@ def _linear_forward_impl( """ weight = args.weight + is_dist_weight = is_distributed_weight(args.weight) inp = args.inp bias = args.bias input_quantizer = args.input_quantizer @@ -413,6 +420,14 @@ def _linear_forward_impl( # ------------------------------------------------------ # Prepare weight tensor # ------------------------------------------------------ + # Distributed weight (e.g. GTP): rebind `weight` to the all-gathered tensor; + # `args.weight` keeps the sharded-param reference for backward re-gather / grad + # finalize. No-op for a plain weight. + if is_dist_weight: + weight = materialize_weight_for_forward(args.weight)[0] + # Refresh out_features from the gathered weight (captured sharded above, pre-gather). + out_features = weight.shape[0] + new_weight_workspace = None weightmat = weight if fp8 or debug: @@ -599,6 +614,9 @@ def _linear_forward_impl( wt_save = weightmat if is_fsdp2 and weightmat is not weight: wt_save = None + # Distributed weight (e.g. GTP): don't save the workspace; backward re-gathers it. + if is_dist_weight: + wt_save = None # Dedup save slots that alias forward inputs; ``_linear_setup_ctx`` # rebuilds the refs from ``inp`` / ``weight`` / ``bias``. @@ -704,6 +722,8 @@ def _linear_setup_ctx( bwd_args.origin_weight_overwrites_main_grad = getattr(weight, "overwrite_main_grad", False) if hasattr(weight, "__fsdp_param__"): bwd_args.main_grad_func = weight.get_main_grad + elif is_distributed_weight(weight): + bwd_args.main_grad_func = weight.grad_buffer else: bwd_args.main_grad_func = lambda: weight.main_grad @@ -749,6 +769,7 @@ def _linear_backward(args: LinearBwdArgs) -> Tuple[Union[torch.Tensor, None], .. inputmat = args.inputmat weight_fp8 = args.weight_fp8 saved_weight = args.saved_weight + is_dist_weight = is_distributed_weight(saved_weight) bias = args.bias input_quantizer = args.input_quantizer weight_quantizer = args.weight_quantizer @@ -793,7 +814,8 @@ def _linear_backward(args: LinearBwdArgs) -> Tuple[Union[torch.Tensor, None], .. origin_weight_python_object is not None ), "weight was removed while fuse_wgrad_accumulation=True" main_grad = bwd_args.main_grad_func() - origin_weight_python_object.main_grad = main_grad + if not is_dist_weight: + origin_weight_python_object.main_grad = main_grad # Gather intermediate/activation tensors if needed # NOTE: weight_fp8 = weight when bwd_args.fp8 == False and torch.disttributed.FSDP already @@ -963,6 +985,12 @@ def _linear_backward(args: LinearBwdArgs) -> Tuple[Union[torch.Tensor, None], .. dgrad = None dgrad_work = None + + # Distributed weight (e.g. GTP): re-gather the sharded weight; runs even when + # requires_dgrad=False so the prev_w prefetch is issued for the next layer's bwd. + if is_dist_weight: + weight_fp8 = materialize_weight_for_backward(saved_weight)[0] + if bwd_args.requires_dgrad: # FSDP2: Re-create workspace from all-gathered weight when @@ -977,6 +1005,16 @@ def _linear_backward(args: LinearBwdArgs) -> Tuple[Union[torch.Tensor, None], .. elif bwd_args.weight_quantizer is not None: bwd_args.weight_quantizer.set_usage(rowwise=True, columnwise=True) weight_fp8 = bwd_args.weight_quantizer(saved_weight) + elif ( + is_dist_weight + and bwd_args.fp8 + and bwd_args.weight_quantizer is not None + and not isinstance(weight_fp8, QuantizedTensorStorage) + ): + # Distributed weight re-gathered a BF16 weight: quantize with the layer quantizer + # so the dgrad operand isn't cast by the delayed recipe. + bwd_args.weight_quantizer.set_usage(rowwise=True, columnwise=True) + weight_fp8 = bwd_args.weight_quantizer(weight_fp8) # Make sure required data is available if isinstance(grad_output, QuantizedTensorStorage): @@ -1163,7 +1201,10 @@ def _linear_backward(args: LinearBwdArgs) -> Tuple[Union[torch.Tensor, None], .. use_split_accumulator = bwd_args.wgrad_use_split_accumulator # Figure out whether to output wgrad GEMM directly into main grad - if bwd_args.is_first_microbatch is not None: + if is_dist_weight: + # Distributed weight (e.g. GTP): accumulation happens downstream in finalize. + accumulate_wgrad_into_param_main_grad = False + elif bwd_args.is_first_microbatch is not None: accumulate_wgrad_into_param_main_grad = ( bwd_args.fuse_wgrad_accumulation and not bwd_args.is_first_microbatch ) @@ -1239,6 +1280,11 @@ def wgrad_gemm( # Call wgrad GEMM now wgrad, grad_bias_ = wgrad_gemm(inputmat_total, grad_output) + # Distributed weight (e.g. GTP): reduce-scatter the freshly computed wgrad + # (async; overlap with the next layer's bwd via the cascade). + if is_dist_weight: + wgrad = finalize_weight_grads(saved_weight, [wgrad])[0] + # Update grad bias if needed if grad_bias is None: grad_bias = grad_bias_ @@ -1287,15 +1333,19 @@ def wgrad_gemm( origin_weight_python_object, "grad_added_to_main_grad" ): origin_weight_python_object.grad_added_to_main_grad = True + # Use the param's local shape (sharded under GTP) so the dummy wgrad + # matches the saved weight shape; main_grad_func() under GTP returns + # an unsharded scratch and would otherwise mismatch. + wgrad_shape = list(origin_weight_python_object.shape) if getattr(origin_weight_python_object, "zero_out_wgrad", False): wgrad = get_dummy_wgrad( - list(main_grad.shape), + wgrad_shape, origin_weight_python_object.dtype, zero=True, ) else: wgrad = get_dummy_wgrad( - list(main_grad.shape), + wgrad_shape, origin_weight_python_object.dtype, ) elif bwd_args.fuse_wgrad_accumulation: diff --git a/transformer_engine/pytorch/ops/_common.py b/transformer_engine/pytorch/ops/_common.py index 675a102240..607346ce30 100644 --- a/transformer_engine/pytorch/ops/_common.py +++ b/transformer_engine/pytorch/ops/_common.py @@ -18,6 +18,36 @@ from ..utils import canonicalize_dtype +def validate_or_alloc_output( + buffer: Optional[torch.Tensor], + shape: tuple[int, ...] | list[int], + dtype: torch.dtype, + device: torch.device, +) -> torch.Tensor: + """Return the caller's output buffer, or allocate one if it is None. + + The buffer must be a contiguous, non-grad tensor matching the required + shape, dtype, and device. Validation reads host-side metadata only. If the + buffer is reused across iterations, pass ``buffer.detach()`` so autograd does + not set its ``requires_grad`` (which would trip the non-grad check here on the + next call). + """ + shape = tuple(shape) + if buffer is None: + return torch.empty(shape, dtype=dtype, device=device) + if tuple(buffer.shape) != shape: + raise ValueError(f"Output buffer shape {tuple(buffer.shape)} does not match {shape}.") + if buffer.dtype != dtype: + raise ValueError(f"Output buffer dtype {buffer.dtype} does not match {dtype}.") + if buffer.device != device: + raise ValueError(f"Output buffer device {buffer.device} does not match {device}.") + if not buffer.is_contiguous(): + raise ValueError("Output buffer must be contiguous.") + if buffer.requires_grad: + raise ValueError("Output buffer must not require gradient.") + return buffer + + def is_quantized_tensor(tensor: torch.Tensor | QuantizedTensorStorage) -> bool: """Check if tensor is a quantized tensor""" return isinstance(tensor, QuantizedTensorStorage) diff --git a/transformer_engine/pytorch/ops/basic/grouped_linear.py b/transformer_engine/pytorch/ops/basic/grouped_linear.py index 5c96d4658e..af3f8b5930 100644 --- a/transformer_engine/pytorch/ops/basic/grouped_linear.py +++ b/transformer_engine/pytorch/ops/basic/grouped_linear.py @@ -48,9 +48,16 @@ get_main_grad_from_param, is_quantized_tensor, maybe_dequantize, + validate_or_alloc_output, view_main_grad_as_grouped_buffer, ) from ..op import BasicOperation, OperationContext +from ...distributed_weight import ( + finalize_weight_grads, + is_distributed_weight, + materialize_weight_for_backward, + materialize_weight_for_forward, +) from ...tensor import GroupedTensor, GroupedTensorStorage from ...triton.grouped_dbias_dscales import ( compute_grouped_dbias, @@ -58,6 +65,12 @@ ) +# Keys for passing caller-provided output and grad-input buffers to a grouped +# linear (or fused grouped MLP) through Sequential's ``op_kwargs``. +OUTPUT_BUFFER_KEY = "output" +GRAD_INPUT_BUFFER_KEY = "grad_input" + + class GroupedLinear(BasicOperation): r"""Apply multiple linear transformations: :math:``y_i = x_i W_i^T + b_i`` @@ -898,6 +911,26 @@ def _get_weight_tensors(self) -> list[torch.nn.Parameter]: return [self.weight] return [getattr(self, f"weight{idx}") for idx in range(self.num_groups)] + def _forward_weight_list(self) -> list[torch.Tensor]: + """Per-expert forward weights, materialized (all-gathered) when distributed.""" + weights = [getattr(self, f"weight{idx}") for idx in range(self.num_groups)] + if is_distributed_weight(weights[0]): + weights = materialize_weight_for_forward(weights) + return weights + + def _backward_weight_setup(self): + """Return ``(origin_weights, is_dist_weight, dgrad_weights)``; dgrad weights are the + re-materialized (all-gathered) weights when distributed, else ``None``.""" + origin_weights = self._get_weight_tensors() + is_dist_weight = is_distributed_weight(origin_weights[0]) + dgrad_weights = materialize_weight_for_backward(origin_weights) if is_dist_weight else None + return origin_weights, is_dist_weight, dgrad_weights + + def _is_distributed_weight(self) -> bool: + """Whether this op's weights are distributed (materialized per fwd/bwd, not saved).""" + leader = self.weight if self.single_grouped_weight else self.weight0 + return is_distributed_weight(leader) + def _get_grouped_bias_for_gemm( self, dtype: torch.dtype, @@ -986,6 +1019,9 @@ def fuser_forward( if self._scale_bias: scales = basic_op_extra_inputs[0][1] + # Caller-provided output buffer (backward grad-input buffer is read in save_ctx). + out_buffer = basic_op_kwargs[0].get(OUTPUT_BUFFER_KEY) + # Dispatch: graph-safe GroupedTensor flow whenever it can be used. # See ``_is_graph_safe_path_supported`` for the gating rationale -- # in short it requires Hopper (SM90+) plus a supported dtype / @@ -1010,6 +1046,7 @@ def fuser_forward( input_requires_grad=input_requires_grad, weight_requires_grad=weight_requires_grad, device=device, + out_buffer=out_buffer, ) else: out, tensors_to_save = self._fuser_forward_split_quantize( @@ -1023,6 +1060,7 @@ def fuser_forward( input_requires_grad=input_requires_grad, weight_requires_grad=weight_requires_grad, device=device, + out_buffer=out_buffer, ) # Save tensors and autograd metadata on the basic-op context. @@ -1111,6 +1149,8 @@ def fuser_forward_save_ctx( ctx.dtype = weight_param.dtype ctx.input_requires_grad = requires_grad[0] ctx.weight_requires_grad = requires_grad[0] and weight_param.requires_grad + # Caller-provided backward grad-input buffer. + ctx.dgrad_out = basic_op_kwargs[0].get(GRAD_INPUT_BUFFER_KEY) # ================================================================== # Legacy `tex.split_quantize` + `general_grouped_gemm` flow. @@ -1129,6 +1169,7 @@ def _fuser_forward_split_quantize( input_requires_grad: bool, weight_requires_grad: bool, device: torch.device, + out_buffer: Optional[torch.Tensor] = None, ) -> tuple[torch.Tensor, tuple[Optional[torch.Tensor], ...]]: """Legacy ``tex.split_quantize`` + ``general_grouped_gemm`` flow.""" num_groups = self.num_groups @@ -1143,7 +1184,7 @@ def _fuser_forward_split_quantize( if weights is None: weights = self.weight.split_into_quantized_tensors() else: - weights = [getattr(self, f"weight{idx}") for idx in range(num_groups)] + weights = self._forward_weight_list() # materialized when distributed bs = None if has_bias: bs = self._get_bias_tensors(dtype) @@ -1173,7 +1214,7 @@ def _fuser_forward_split_quantize( # Allocate output tensor in_shape = list(input_.size()) out_shape = in_shape[:-1] + [self.out_features] - out = torch.empty(out_shape, dtype=dtype, device=device) + out = validate_or_alloc_output(out_buffer, out_shape, dtype, device) # Perform GEMMs use_gemm_bias = has_bias and not self._scale_bias @@ -1198,7 +1239,8 @@ def _fuser_forward_split_quantize( out_splits[i].add_(bs[i].unsqueeze(0) * scales_splits[i].unsqueeze(-1)) # Prepare weight tensors for backward pass - if not input_requires_grad: + # Distributed weights are re-materialized in backward, so we never save the gathered weight + if not input_requires_grad or self._is_distributed_weight(): ws = [None] * num_groups elif with_quantized_compute: for w, weight_param in zip(ws, weights): @@ -1239,6 +1281,7 @@ def _fuser_forward_grouped_tensor( input_requires_grad: bool, weight_requires_grad: bool, device: torch.device, + out_buffer: Optional[torch.Tensor] = None, ) -> tuple[torch.Tensor, tuple[Optional[torch.Tensor], ...]]: """Graph-safe GroupedTensor forward path (pure compute). Returns ``(output, tensors_to_save)``. ``split_sizes``, @@ -1290,7 +1333,7 @@ def _fuser_forward_grouped_tensor( else: # Discrete weights grouped_weights = self._get_discrete_weights_for_gemm( - [getattr(self, f"weight{idx}") for idx in range(num_groups)], + self._forward_weight_list(), weight_quantizers, columnwise_usage=input_requires_grad, with_quantized_compute=with_quantized_compute, @@ -1299,7 +1342,7 @@ def _fuser_forward_grouped_tensor( # Allocate output buffer and wrap as a GroupedTensor view. out_shape = original_shape[:-1] + [self.out_features] - out = torch.empty(out_shape, dtype=dtype, device=device) + out = validate_or_alloc_output(out_buffer, out_shape, dtype, device) grouped_out = GroupedTensorStorage( shape=(total_tokens, self.out_features), dtype=dtype, @@ -1333,7 +1376,8 @@ def _fuser_forward_grouped_tensor( bias_scale=bias_scale, ) - if not input_requires_grad: + # Distributed weights are re-materialized in backward, so never save the gathered weight. + if not input_requires_grad or self._is_distributed_weight(): grouped_weights = None if self.single_grouped_weight else [None] * num_groups if not weight_requires_grad: @@ -1393,7 +1437,7 @@ def _fuser_backward_split_quantize( ]: num_groups = self.num_groups has_bias = self.has_bias - weights = self._get_weight_tensors() + weights, is_dist_weight, dist_dgrad_weights = self._backward_weight_setup() device = weights[0].device # Saved tensors from forward pass. Layout: @@ -1445,7 +1489,8 @@ def _fuser_backward_split_quantize( grad_biases = [dbias_packed[idx].to(dtype=ctx.dtype) for idx in range(num_groups)] # Initialize grad weight buffers. - accumulate_into_main_grad = self._accumulate_into_main_grad + # Distributed weights reduce their own grads (finalize); never accumulate into main_grad. + accumulate_into_main_grad = self._accumulate_into_main_grad and not is_dist_weight grad_weights = [None] * num_groups final_weight_grads: list[Optional[torch.Tensor]] = ( [None] if self.single_grouped_weight else [None] * num_groups @@ -1489,13 +1534,11 @@ def _fuser_backward_split_quantize( if ctx.input_requires_grad: out_shape = list(grad_output.size()) in_shape = out_shape[:-1] + [self.in_features] - grad_input = torch.empty( - in_shape, - dtype=ctx.dtype, - device=device, + grad_input = validate_or_alloc_output( + getattr(ctx, "dgrad_out", None), in_shape, ctx.dtype, device ) general_grouped_gemm( - ws, + dist_dgrad_weights if is_dist_weight else ws, dys, [grad_input], [None] * num_groups, # quantization_params @@ -1540,11 +1583,16 @@ def _fuser_backward_split_quantize( if not delay_wgrad: clear_tensor_data(*xs) + # Distributed weights: finalize (e.g. reduce-scatter) the freshly computed wgrads per shard. + # Return discarded (see finalize_weight_grads); the dummy is returned below instead. + if ctx.weight_requires_grad and is_dist_weight: + assert not delay_wgrad, "delayed wgrad unsupported with distributed weights." + finalize_weight_grads(weights, grad_weights) # Megatron-LM wgrad fusion: regardless of overwrite vs. accumulate, # signal that ``main_grad`` already carries the wgrad and replace # ``.grad`` with a dummy so DDP/FSDP hooks won't add ``.grad`` into # ``main_grad`` again. - if ctx.weight_requires_grad and self._accumulate_into_main_grad: + if ctx.weight_requires_grad and (is_dist_weight or self._accumulate_into_main_grad): final_weight_grads = get_dummy_wgrads_for_params(weights) elif ctx.weight_requires_grad and delay_wgrad: final_weight_grads = [None] if self.single_grouped_weight else [None] * num_groups @@ -1578,7 +1626,7 @@ def _fuser_backward_grouped_tensor( ]: num_groups = self.num_groups has_bias = self.has_bias - weights = self._get_weight_tensors() + weights, is_dist_weight, dist_dgrad_weights = self._backward_weight_setup() device = weights[0].device dtype = ctx.dtype @@ -1668,7 +1716,9 @@ def _fuser_backward_grouped_tensor( grad_input = None if ctx.input_requires_grad: grad_input_shape = list(grad_output.size())[:-1] + [self.in_features] - grad_input = torch.empty(grad_input_shape, dtype=dtype, device=device) + grad_input = validate_or_alloc_output( + getattr(ctx, "dgrad_out", None), grad_input_shape, dtype, device + ) grouped_grad_input = GroupedTensorStorage( shape=(total_tokens, self.in_features), dtype=dtype, @@ -1679,7 +1729,7 @@ def _fuser_backward_grouped_tensor( tensor_offsets=base_split_offsets * self.in_features, ) general_grouped_gemm_for_grouped_tensor( - ws, + dist_dgrad_weights if is_dist_weight else ws, grouped_dy, grouped_grad_input, layout="NN", @@ -1725,7 +1775,8 @@ def _fuser_backward_grouped_tensor( final_weight_grads[0] = grouped_wgrad.rowwise_data.view(num_groups, *weight_shape) wgrad_output = grouped_wgrad else: - if self._accumulate_into_main_grad: + # Distributed weights finalize wgrads (below); never accumulate into main_grad. + if self._accumulate_into_main_grad and not is_dist_weight: final_weight_grads = [ get_main_grad_from_param(w, op_label="GroupedLinear") for w in weights ] @@ -1755,11 +1806,16 @@ def _fuser_backward_grouped_tensor( else: wgrad_gemm(grouped_x, grouped_dy, wgrad_output) + # Distributed weights: finalize (e.g. reduce-scatter) the freshly computed wgrads per shard. + # Return discarded (see finalize_weight_grads); the dummy is returned below instead. + if ctx.weight_requires_grad and is_dist_weight: + assert not delay_wgrad, "delayed wgrad unsupported with distributed weights." + finalize_weight_grads(weights, final_weight_grads) # Megatron-LM wgrad fusion: regardless of overwrite vs. accumulate, # signal that ``main_grad`` already carries the wgrad and replace # ``.grad`` with a dummy so DDP/FSDP hooks won't add ``.grad`` into # ``main_grad`` again. - if ctx.weight_requires_grad and self._accumulate_into_main_grad: + if ctx.weight_requires_grad and (is_dist_weight or self._accumulate_into_main_grad): final_weight_grads = get_dummy_wgrads_for_params(weights) elif ctx.weight_requires_grad and delay_wgrad: final_weight_grads = [None] if self.single_grouped_weight else [None] * num_groups diff --git a/transformer_engine/pytorch/ops/fused/grouped_mlp.py b/transformer_engine/pytorch/ops/fused/grouped_mlp.py index e0961044ad..83954a9b3d 100644 --- a/transformer_engine/pytorch/ops/fused/grouped_mlp.py +++ b/transformer_engine/pytorch/ops/fused/grouped_mlp.py @@ -19,6 +19,12 @@ from ...constants import MXFP8_BLOCK_SCALING_SIZE, NVFP4_BLOCK_SCALING_SIZE from ...cpu_offload import is_cpu_offload_enabled, mark_activation_offload, start_offload from ...cpp_extensions import general_gemm, general_grouped_gemm_for_grouped_tensor +from ...distributed_weight import ( + is_distributed_weight, + materialize_weight_for_forward, + materialize_weight_for_backward, + finalize_weight_grads, +) from ...module.base import _2X_ACC_WGRAD from ...quantization import Recipe from ...tensor import NVFP4Quantizer, NVFP4Tensor, NVFP4TensorStorage, Quantizer @@ -48,8 +54,10 @@ get_main_grad_from_param, is_quantized_tensor, maybe_dequantize, + validate_or_alloc_output, view_main_grad_as_grouped_buffer, ) +from ..basic.grouped_linear import GRAD_INPUT_BUFFER_KEY, OUTPUT_BUFFER_KEY @functools.lru_cache(maxsize=None) @@ -539,6 +547,7 @@ def _compute_grad_params( wgrad_output = None op_label = f"Grouped MLP fused backward ({label})" if label else "Grouped MLP fused backward" weights = fc_op._get_weight_tensors() + is_dist_weight = is_distributed_weight(weights[0]) if fc_op.single_grouped_weight: w_list = [None] if ctx.weight_requires_grad: @@ -571,7 +580,9 @@ def _compute_grad_params( else: w_list = [None] * num_groups if ctx.weight_requires_grad: - if fc_op._accumulate_into_main_grad: + # Distributed weight: the GEMM produces full-sized wgrads but main_grad is sharded, + # so use a full-sized scratch buffer (the reduce-scatter below lands it in main_grad). + if fc_op._accumulate_into_main_grad and not is_dist_weight: w_list = [get_main_grad_from_param(w, op_label=op_label) for w in weights] accumulate_into_main_grad = get_accumulate_flag_in_param(weights[0]) else: @@ -587,6 +598,10 @@ def _compute_grad_params( if ctx.weight_requires_grad: # Launch or defer the GEMM delay_wgrad = fc_op.wgrad_store is not None and fc_op.wgrad_store.delay_wgrad_compute() + if is_dist_weight and delay_wgrad: + raise RuntimeError( + "distributed-weight fused grouped-MLP requires delay_wgrad_compute=False." + ) if cudnn_wgrad_kernel_fn is not None: offsets = offsets if offsets.dtype == torch.int32 else offsets.to(dtype=torch.int32) gemm_fn = functools.partial( @@ -625,9 +640,14 @@ def _compute_grad_params( fc_op.wgrad_store.put([grouped_x, grouped_dy, wgrad_output], gemm_fn) else: gemm_fn(grouped_x, grouped_dy, wgrad_output) + # Distributed weight: reduce-scatter the wgrads into main_grad. + # Return discarded (see finalize_weight_grads); dummy wgrads returned below. + if is_dist_weight: + finalize_weight_grads(weights, w_list) # Need to return dummy wgrads for Megatron-LM wgrad fusion if grad is already added - if fc_op._accumulate_into_main_grad: + # (wgrad fusion, or the distributed-weight reduce-scatter above) so it doesn't double-add. + if fc_op._accumulate_into_main_grad or is_dist_weight: w_list = get_dummy_wgrads_for_params(weights) elif delay_wgrad: w_list = [None] if fc_op.single_grouped_weight else [None] * num_groups @@ -768,7 +788,9 @@ def fuse_grouped_mlp_ops( class _GroupedMLP_CuTeGEMMBase(FusedOperation): """Joint fused op for block-scaled GroupedLinear + activation + GroupedLinear. - Uses experimental CuTe DSL kernels from cuDNN front-end. + MXFP8 uses CuTe DSL grouped-GEMM kernels via the cuDNN front-end. When + caller output/grad_input buffers are provided, the GEMMs write into them + directly. """ @@ -879,6 +901,21 @@ def fuser_forward( fc1_op, activation_op, fc2_op = self.basic_ops fc1_ctx, _activation_ctx, fc2_ctx = basic_op_ctxs + # Caller-provided buffers: FC2 forward output and FC1 backward grad-input. Guard against + # a buffer routed to the wrong op (output belongs on FC2, grad-input on FC1). + if OUTPUT_BUFFER_KEY in basic_op_kwargs[0]: + raise ValueError( + f"'{OUTPUT_BUFFER_KEY}' buffer can only be provided to FC2 (the last op) of " + "the fused grouped MLP, not FC1." + ) + if GRAD_INPUT_BUFFER_KEY in basic_op_kwargs[-1]: + raise ValueError( + f"'{GRAD_INPUT_BUFFER_KEY}' buffer can only be provided to FC1 (the first op) of " + "the fused grouped MLP, not FC2." + ) + output_buffer = basic_op_kwargs[-1].get(OUTPUT_BUFFER_KEY) + fc1_ctx.dgrad_out = basic_op_kwargs[0].get(GRAD_INPUT_BUFFER_KEY) + # Tensor properties fc1_weight_shape = (fc1_op.out_features, fc1_op.in_features) fc2_weight_shape = (fc2_op.out_features, fc2_op.in_features) @@ -890,6 +927,19 @@ def fuser_forward( num_groups = fc1_op.num_groups fc1_weight_param = fc1_op.weight if fc1_op.single_grouped_weight else fc1_op.weight0 fc2_weight_param = fc2_op.weight if fc2_op.single_grouped_weight else fc2_op.weight0 + + # Distributed weight: expert weights are sharded 1/N along out_features; the fused kernels + # read the full shape, so all-gather the full weight first. A plain weight is a no-op. + fc1_is_dist = is_distributed_weight(fc1_weight_param) + fc2_is_dist = is_distributed_weight(fc2_weight_param) + assert fc1_is_dist == fc2_is_dist, "FC1/FC2 must share one distributed-weight group." + if fc1_is_dist: + assert ( + not fc1_op.single_grouped_weight and not fc2_op.single_grouped_weight + ), "distributed-weight fused grouped-MLP requires single_grouped_weight=False." + assert fc1_op.weight0.is_routed_expert and fc1_op.weight0.weight_list is not None + assert fc2_op.weight0.is_routed_expert and fc2_op.weight0.weight_list is not None + device = fc1_weight_param.device if torch.is_autocast_enabled(): dtype = torch.get_autocast_dtype("cuda") @@ -978,6 +1028,8 @@ def fuser_forward( else: quantized_fc1_weights.append(weight) grouped_fc1_weight = quantized_fc1_weights + if fc1_is_dist: + grouped_fc1_weight = materialize_weight_for_forward(grouped_fc1_weight) # Prepare FC2 grouped weight tensor for fused kernels. if fc2_op.single_grouped_weight: @@ -1016,10 +1068,6 @@ def fuser_forward( grouped_fc1_weight, "_with_gemm_swizzled_scales" ): grouped_fc1_weight._with_gemm_swizzled_scales = False - if isinstance(grouped_fc2_weight, GroupedTensor) and not hasattr( - grouped_fc2_weight, "_with_gemm_swizzled_scales" - ): - grouped_fc2_weight._with_gemm_swizzled_scales = False # Group-quantize input tensor and convert dtypes if needed fc1_input_quantizer.set_usage(rowwise=True, columnwise=weight_requires_grad) @@ -1252,6 +1300,13 @@ def fuser_forward( else: fc1_kernel_out = self.grouped_gemm_activation_kernel()(**fc1_activation_kwargs) + if fc2_is_dist: + grouped_fc2_weight = materialize_weight_for_forward(grouped_fc2_weight) + if isinstance(grouped_fc2_weight, GroupedTensor) and not hasattr( + grouped_fc2_weight, "_with_gemm_swizzled_scales" + ): + grouped_fc2_weight._with_gemm_swizzled_scales = False + # Unpack kernel outputs # Note: Fused kernel outputs tensors with non-contiguous # logical dims. @@ -1301,7 +1356,7 @@ def fuser_forward( tensor_offsets=fc2_x_tensor_offsets, ) - fc2_out_buf = torch.empty(fc2_out_shape, dtype=dtype, device=device) + fc2_out_buf = validate_or_alloc_output(output_buffer, fc2_out_shape, dtype, device) if ( num_groups == 1 and grouped_fc2_x.columnwise_data is not None @@ -1329,9 +1384,9 @@ def fuser_forward( fc2_bias_packed.transpose(0, 1).contiguous().expand(in_shape[0], -1) ) if fc2_scales is not None: - fc2_out_buf = fc2_out_buf + token_bias * fc2_scales.view(-1, 1) + fc2_out_buf += token_bias * fc2_scales.view(-1, 1) else: - fc2_out_buf = fc2_out_buf + token_bias + fc2_out_buf += token_bias else: fc2_out_grouped = GroupedTensorStorage( shape=(in_shape[0], fc2_weight_shape[0]), @@ -1434,8 +1489,15 @@ def fuser_forward( fc2_quant_kwargs["b_dtype"] = torch.float8_e4m3fn fc2_quant_kwargs["b_major"] = "k" - fc2_kernel_out = self.grouped_gemm_quant_kernel()(**fc2_quant_kwargs) - fc2_out = fc2_kernel_out["d_tensor"].permute(2, 0, 1).view(fc2_out_shape).contiguous() + # Always allocate the output (the caller's buffer if provided, else a fresh one) and + # pass it as the kernel's d_tensor, so the kernel writes in place and the call is uniform. + output_buffer = validate_or_alloc_output(output_buffer, fc2_out_shape, dtype, device) + fc2_quant_kwargs["d_tensor"] = output_buffer.as_strided( + (in_shape[0], fc2_weight_shape[0], 1), + (fc2_weight_shape[0], 1, in_shape[0] * fc2_weight_shape[0]), + ) + self.grouped_gemm_quant_kernel()(**fc2_quant_kwargs) + fc2_out = output_buffer # Save state for backward pass if requires_grad: @@ -1478,6 +1540,11 @@ def fuser_forward( fc2_weight_tensors = ( [grouped_fc2_weight] if fc2_op.single_grouped_weight else grouped_fc2_weight ) + # Save the joint op's internal layout; distributed weights save the small shards. + if fc1_is_dist: + fc1_weight_tensors = fc1_weights + if fc2_is_dist: + fc2_weight_tensors = fc2_weights fc1_ctx.save_for_backward( split_sizes, base_split_offsets, @@ -1741,6 +1808,10 @@ def fuser_backward( glu_clamp_min=self._cudnn_glu_clamp_min, ) + fc2_leader = fc2_op.weight if fc2_op.single_grouped_weight else fc2_op.weight0 + if is_distributed_weight(fc2_leader): + grouped_fc2_weight = materialize_weight_for_backward(fc2_leader) + if fc2_op.single_grouped_weight: # Clone and swizzle scales for GEMM fc2_weight_for_gemm = grouped_fc2_weight.copy() @@ -1962,11 +2033,16 @@ def fuser_backward( # FC1 dgrad GEMM grad_input = None + grad_input_buffer = getattr(fc1_ctx, "dgrad_out", None) if fc1_ctx.input_requires_grad: in_shape = out_shape[:-1] + [fc1_weight_shape[1]] + fc1_leader = fc1_op.weight if fc1_op.single_grouped_weight else fc1_op.weight0 + if is_distributed_weight(fc1_leader): + grouped_fc1_weight = materialize_weight_for_backward(fc1_leader) + if use_nvfp4: - grad_input = torch.empty(in_shape, dtype=dtype, device=device) + grad_input = validate_or_alloc_output(grad_input_buffer, in_shape, dtype, device) if num_groups == 1: if fc1_op.single_grouped_weight: fc1_w_single = grouped_fc1_weight.split_into_quantized_tensors()[0] @@ -2063,8 +2139,17 @@ def fuser_backward( fc1_dgrad_kwargs["b_dtype"] = torch.float8_e4m3fn fc1_dgrad_kwargs["b_major"] = "n" - fc1_dgrad_kernel_out = self.grouped_gemm_quant_kernel()(**fc1_dgrad_kwargs) - grad_input = fc1_dgrad_kernel_out["d_tensor"].view(in_shape) + # Always allocate the grad-input (the caller's buffer if provided, else a fresh + # one) and pass it as the kernel's d_tensor, decoupling it from the kernel call. + grad_input_buffer = validate_or_alloc_output( + grad_input_buffer, in_shape, dtype, device + ) + fc1_dgrad_kwargs["d_tensor"] = grad_input_buffer.as_strided( + (out_shape[0], fc1_weight_shape[1], 1), + (fc1_weight_shape[1], 1, out_shape[0] * fc1_weight_shape[1]), + ) + self.grouped_gemm_quant_kernel()(**fc1_dgrad_kwargs) + grad_input = grad_input_buffer # FC1 wgrad GEMM fc1_grad_params = _compute_grad_params( diff --git a/transformer_engine/pytorch/ops/sequential.py b/transformer_engine/pytorch/ops/sequential.py index 592ddae23a..cb5dfecb9f 100644 --- a/transformer_engine/pytorch/ops/sequential.py +++ b/transformer_engine/pytorch/ops/sequential.py @@ -6,11 +6,11 @@ from __future__ import annotations from collections.abc import Iterable, Iterator -from typing import Optional +from typing import Any, Optional import torch -from transformer_engine.pytorch.ops.op import FusibleOperation +from transformer_engine.pytorch.ops.op import BasicOperation, FusibleOperation from transformer_engine.pytorch.ops.fuser import OperationFuser @@ -168,23 +168,34 @@ def forward( self, input: torch.Tensor, # pylint: disable=redefined-builtin *extra_inputs: torch.Tensor, + op_kwargs: Optional[dict[torch.nn.Module | int, dict[str, Any]]] = None, ) -> torch.Tensor | tuple[torch.Tensor, ...]: - """Forward pass""" + """Forward pass. + + op_kwargs : optional mapping from a contained op, keyed by module + or index, to extra keyword arguments forwarded to that op. + Only fusible operations can be targeted, for example to pass + preallocated output or grad input buffers to a grouped linear + or grouped MLP. + """ # Create module groups if needed if self._module_groups is None: self._module_groups = self._make_module_groups(self._modules.values()) + # Route op kwargs to each module group's basic ops + group_op_kwargs = self._resolve_op_kwargs(op_kwargs) + # Forward pass for each module group x = input extra_outputs: list[torch.Tensor] = [] - for module_group in self._module_groups: + for group_idx, module_group in enumerate(self._module_groups): if isinstance(module_group, OperationFuser): xs, extra_inputs = ( (x,) + extra_inputs[: module_group.num_extra_inputs], extra_inputs[module_group.num_extra_inputs :], ) - xs = module_group(*xs) + xs = module_group(*xs, basic_op_kwargs=group_op_kwargs[group_idx]) if isinstance(xs, tuple): x, ys = xs[0], xs[1:] extra_outputs.extend(ys) @@ -196,3 +207,39 @@ def forward( if extra_outputs: return (x,) + tuple(extra_outputs) return x + + def _resolve_op_kwargs( + self, + op_kwargs: Optional[dict[torch.nn.Module | int, dict[str, Any]]], + ) -> list[Optional[list[dict[str, Any]]]]: + """Map per-op kwargs onto each module group's basic-op kwargs list.""" + group_kwargs: list[Optional[list[dict[str, Any]]]] = [None] * len(self._module_groups) + if not op_kwargs: + return group_kwargs + + # Construct map from basic-op id to its kwarg dict + resolved: dict[int, dict[str, Any]] = {} + for key, kwargs in op_kwargs.items(): + module = self[key] if isinstance(key, int) else key + if not isinstance(module, BasicOperation): + raise ValueError( + f"Attempted to provide forward kwargs to {type(module).__name__}, but " + "Sequential only allows providing forward kwargs to a BasicOperation." + ) + resolved[id(module)] = kwargs + + # Slot each op's kwargs into its module group by matching basic-op identity + for group_idx, module_group in enumerate(self._module_groups): + if not isinstance(module_group, OperationFuser): + continue + # pylint: disable-next=protected-access + group_kwargs[group_idx] = [{} for _ in module_group._basic_ops] + # pylint: disable-next=protected-access + for idx, op in enumerate(module_group._basic_ops): + if id(op) in resolved: + group_kwargs[group_idx][idx] = resolved.pop(id(op)) + + # Any keys left over target an op that is not in this Sequential + if resolved: + raise ValueError("op_kwargs contains keys that are not in this Sequential") + return group_kwargs diff --git a/transformer_engine/pytorch/optimizers/fused_sgd.py b/transformer_engine/pytorch/optimizers/fused_sgd.py index d7ab3fe9fe..10151e1406 100644 --- a/transformer_engine/pytorch/optimizers/fused_sgd.py +++ b/transformer_engine/pytorch/optimizers/fused_sgd.py @@ -295,20 +295,19 @@ def step(self, closure=None): for _, (launch_set, first_run) in enumerate(zip(launch_sets, first_runs)): assert len(launch_set[0]) == len(launch_set[1]) assert len(launch_set[0]) == len(launch_set[2]) - if len(launch_set[0]) > 0: - multi_tensor_applier( - self.multi_tensor_sgd, - self._dummy_overflow_buf, - launch_set, - weight_decay, - momentum, - dampening, - group["lr"], - nesterov, - first_run, - self.wd_after_momentum, - 1.0 / self.most_recent_scale, - ) + multi_tensor_applier( + self.multi_tensor_sgd, + self._dummy_overflow_buf, + launch_set, + weight_decay, + momentum, + dampening, + group["lr"], + nesterov, + first_run, + self.wd_after_momentum, + 1.0 / self.most_recent_scale, + ) self.most_recent_scale = 1.0 self.scale_set_by_backward = False diff --git a/transformer_engine/pytorch/optimizers/multi_tensor_apply.py b/transformer_engine/pytorch/optimizers/multi_tensor_apply.py index a5cbd27337..e7d4fa0db3 100644 --- a/transformer_engine/pytorch/optimizers/multi_tensor_apply.py +++ b/transformer_engine/pytorch/optimizers/multi_tensor_apply.py @@ -18,6 +18,16 @@ def __call__(self, op, noop_flag_buffer, tensor_lists, *args): if isinstance(t, DTensor): tensor_lists[i][j] = t._local_tensor + if any(len(tensors) != len(tensor_lists[0]) for tensors in tensor_lists): + raise RuntimeError("Expected aligned multi-tensor lists.") + + keep_slot = [tensor.numel() > 0 for tensor in tensor_lists[0]] + for i, tensors in enumerate(tensor_lists): + tensor_lists[i] = [tensor for tensor, keep in zip(tensors, keep_slot) if keep] + + if not tensor_lists[0]: + return None + return op(self.chunk_size, noop_flag_buffer, tensor_lists, *args)