From fc9a88a09e1b0774076ac37594eb8e623aa4ed31 Mon Sep 17 00:00:00 2001 From: zjin-lcf Date: Wed, 19 Aug 2026 21:10:55 +0000 Subject: [PATCH] [ROCm] Route dense CP softmax-LSE correction through a native kernel The context-parallel softmax-LSE merge (flash_attn_fwd_softmax_lse_correction and its second-half variant) ran through @jit_fuser / torch.compile on ROCm. Dynamo can hold more than one compiled variant of these functions, and on ROCm the variants disagree by ~1 ULP because libdevice.log1p contracts differently under different Triton launch configurations (issue #693). A reference model and an actor model built from the same checkpoint then diverge. Add a fixed-configuration native kernel for the dense layout, mirroring the existing THD path: * nvte_cp_lse_correction C API + context_parallel::lse_correction kernel * tex.lse_correction PyTorch binding * route flash_attn_fwd_softmax_lse_correction and the second-half variant through the native kernel on ROCm (IS_HIP_EXTENSION), leaving the CUDA path unchanged The native kernel returns identical bits for every launch configuration, so the merge is bitwise-reproducible. Add regression tests asserting bitwise stability of both correction functions across intervening shapes. Fixes: https://github.com/ROCm/TransformerEngine/issues/693 --- tests/pytorch/attention/test_cp_utils.py | 61 +++++++++++++++++ .../common/fused_attn/context_parallel.cu | 66 ++++++++++++++++++- .../include/transformer_engine/fused_attn.h | 12 ++++ .../dot_product_attention/context_parallel.py | 13 +++- transformer_engine/pytorch/csrc/extensions.h | 4 +- .../pytorch/csrc/extensions/attention.cpp | 27 +++++++- .../pytorch/csrc/extensions/pybind.cpp | 5 +- 7 files changed, 182 insertions(+), 6 deletions(-) diff --git a/tests/pytorch/attention/test_cp_utils.py b/tests/pytorch/attention/test_cp_utils.py index e5051aab36..f73f886abd 100644 --- a/tests/pytorch/attention/test_cp_utils.py +++ b/tests/pytorch/attention/test_cp_utils.py @@ -1,3 +1,4 @@ +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. # Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. @@ -5,10 +6,14 @@ """Unit tests for context parallel utils.""" import torch import unittest +from torch.utils.cpp_extension import IS_HIP_EXTENSION + from transformer_engine.pytorch.attention.dot_product_attention.context_parallel import ( get_batch_on_this_cp_rank, pad_thd_sequences_for_cp, generate_positional_ids_for_cp, + flash_attn_fwd_softmax_lse_correction, + flash_attn_fwd_second_half_softmax_lse_correction, ) @@ -710,5 +715,61 @@ def test_integration_with_padding_and_cp_slicing(self): self.assertTrue(torch.equal(input_ids_r0, expected_input_ids_r0)) +@unittest.skipUnless( + torch.cuda.is_available() and IS_HIP_EXTENSION, + "Requires ROCm", +) +class TestSoftmaxLseCorrectionReproducibility(unittest.TestCase): + """Regression tests for https://github.com/ROCm/TransformerEngine/issues/693. + + The CP softmax-LSE merge has to return the same bits regardless of how many times it + has been called or with which shapes. When these functions are compiled, Dynamo can + hold more than one compiled variant, and on ROCm the variants may disagree by about + 1 ULP because log1p contracts differently under different Triton launch configurations. + A reference model and an actor model built from the same checkpoint then diverge. + """ + + @staticmethod + def _merged_lse(softmax_lse, softmax_lse_per_step): + max_scale = torch.max(softmax_lse, softmax_lse_per_step) + min_scale = torch.min(softmax_lse, softmax_lse_per_step) + return max_scale + torch.log1p(torch.exp(min_scale - max_scale)) + + @staticmethod + def _rand(*shape): + return torch.rand(*shape, device="cuda", dtype=torch.float32) * 10 + + def test_softmax_lse_correction_is_bitwise_stable(self): + """The full LSE correction matches eager after calls with other shapes.""" + b, h, s = 2, 4, 2053 + softmax_lse = self._rand(b, h, s) + softmax_lse_per_step = self._rand(b, h, s) + expected = self._merged_lse(softmax_lse, softmax_lse_per_step) + + # This second shape would make Dynamo build another variant if the function regressed + # to using jit_fuser. + flash_attn_fwd_softmax_lse_correction(self._rand(b, h, s // 2), self._rand(b, h, s // 2)) + + merged = softmax_lse.clone() + flash_attn_fwd_softmax_lse_correction(merged, softmax_lse_per_step) + self.assertTrue(torch.equal(merged, expected)) + + def test_second_half_softmax_lse_correction_is_bitwise_stable(self): + """The second-half LSE correction matches eager after calls with other shapes.""" + b, h, s = 2, 4, 2053 + softmax_lse = self._rand(b, h, 2, s) + softmax_lse_per_step = self._rand(b, h, s) + expected = softmax_lse.clone() + expected[..., 1, :] = self._merged_lse(softmax_lse[..., 1, :], softmax_lse_per_step) + + flash_attn_fwd_second_half_softmax_lse_correction( + self._rand(b, h, 2, s // 2), self._rand(b, h, s // 2) + ) + + merged = softmax_lse.clone() + flash_attn_fwd_second_half_softmax_lse_correction(merged, softmax_lse_per_step) + self.assertTrue(torch.equal(merged, expected)) + + if __name__ == "__main__": unittest.main() diff --git a/transformer_engine/common/fused_attn/context_parallel.cu b/transformer_engine/common/fused_attn/context_parallel.cu index cf1fffd94f..ce8921977c 100644 --- a/transformer_engine/common/fused_attn/context_parallel.cu +++ b/transformer_engine/common/fused_attn/context_parallel.cu @@ -1,4 +1,5 @@ /************************************************************************* + * Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * * See LICENSE for license information. @@ -15,7 +16,7 @@ namespace transformer_engine { namespace context_parallel { struct LseCorrectionFunctor { - __forceinline__ __device__ static void run(float *lse, float *half_lse, size_t idx, + __forceinline__ __device__ static void run(float *lse, const float *half_lse, size_t idx, size_t half_idx) { float val = lse[idx]; float val_per_step = half_lse[half_idx]; @@ -25,6 +26,60 @@ struct LseCorrectionFunctor { } }; +/*************************************************************************************************** + * Correct softmax LSE for dense Context Parallel layouts + **************************************************************************************************/ + +template +__global__ void lse_correction_kernel(float *lse, const float *lse_per_step, size_t rows, + size_t cols) { + size_t col = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + size_t row = blockIdx.y; + if (row >= rows || col >= cols) { + return; + } + + size_t half_idx = row * cols + col; + size_t idx = only_second_half ? (row * 2 + 1) * cols + col : half_idx; + LseCorrectionFunctor::run(lse, lse_per_step, idx, half_idx); +} + +void lse_correction(Tensor lse, const Tensor &lse_per_step, bool only_second_half, + cudaStream_t stream) { + using namespace transformer_engine; + NVTE_CHECK(lse.dtype() == DType::kFloat32); + NVTE_CHECK(lse_per_step.dtype() == DType::kFloat32); + NVTE_CHECK(lse_per_step.dim() >= 1); + + const auto cols = lse_per_step.shape().back(); + const auto step_numel = lse_per_step.numel(); + NVTE_CHECK(cols > 0 && step_numel > 0); + if (only_second_half) { + NVTE_CHECK(lse.dim() == lse_per_step.dim() + 1); + NVTE_CHECK(lse.shape()[lse.dim() - 2] == 2); + NVTE_CHECK(lse.shape().back() == cols); + for (size_t i = 0; i + 1 < lse_per_step.dim(); ++i) { + NVTE_CHECK(lse.shape()[i] == lse_per_step.shape()[i]); + } + } else { + NVTE_CHECK(lse.shape() == lse_per_step.shape()); + } + + const auto rows = step_numel / cols; + constexpr unsigned int block = 256; + dim3 grid((cols + block - 1) / block, rows); + if (only_second_half) { + lse_correction_kernel<<>>( + reinterpret_cast(lse.data.dptr), + reinterpret_cast(lse_per_step.data.dptr), rows, cols); + } else { + lse_correction_kernel<<>>( + reinterpret_cast(lse.data.dptr), + reinterpret_cast(lse_per_step.data.dptr), rows, cols); + } + NVTE_CHECK_CUDA(cudaGetLastError()); +} + struct ReadLseFunctor { __forceinline__ __device__ static void run(float *lse, float *half_lse, size_t idx, size_t half_idx) { @@ -681,6 +736,15 @@ void thd_get_partitioned_indices(const Tensor &cu_seqlens, Tensor output, int to } // namespace context_parallel } // namespace transformer_engine +void nvte_cp_lse_correction(NVTETensor lse, const NVTETensor &lse_per_step, int only_second_half, + cudaStream_t stream) { + NVTE_API_CALL(nvte_cp_lse_correction); + using namespace transformer_engine; + + context_parallel::lse_correction(*convertNVTETensorCheck(lse), + *convertNVTETensorCheck(lse_per_step), only_second_half, stream); +} + void nvte_cp_thd_read_half_tensor(const NVTETensor &tensor, const NVTETensor &cu_seqlens, NVTETensor half, int half_idx, cudaStream_t stream) { NVTE_API_CALL(nvte_thd_read_half_tensor); diff --git a/transformer_engine/common/include/transformer_engine/fused_attn.h b/transformer_engine/common/include/transformer_engine/fused_attn.h index 5221b0255f..0f4599cd12 100644 --- a/transformer_engine/common/include/transformer_engine/fused_attn.h +++ b/transformer_engine/common/include/transformer_engine/fused_attn.h @@ -476,6 +476,18 @@ void nvte_copy_to_kv_cache(NVTETensor new_k, NVTETensor new_v, NVTETensor k_cach int max_ctx_len, int max_seq_len, int max_pages_per_seq, int is_non_paged, cudaStream_t stream); +/*! \brief Correct softmax LSE (LogSumExp) for dense context parallel layouts. + * + * \warning This API is **experimental** and subject to change. + * + * \param[out] lse Output tensor. + * \param[in] lse_per_step Input tensor. + * \param[in] only_second_half Whether to correct only the second half of lse. + * \param[in] stream CUDA stream used for this operation. + */ +void nvte_cp_lse_correction(NVTETensor lse, const NVTETensor &lse_per_step, int only_second_half, + cudaStream_t stream); + /*! \brief Extract the first half (half_idx=0) or second half (half_idx=1) of a THD tensor. * * \warning This API is **experimental** and subject to change. diff --git a/transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py b/transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py index cca013461d..e88194c08c 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py @@ -184,24 +184,33 @@ def flash_attn_fwd_second_half_out_correction( out_.add_(out_corrected) -@jit_fuser +_lse_correction_fuser = (lambda func: func) if IS_HIP_EXTENSION else jit_fuser + + +@_lse_correction_fuser def flash_attn_fwd_softmax_lse_correction( softmax_lse: torch.Tensor, softmax_lse_per_step: torch.Tensor, ): """Merge softmax stats of each step in Attention with context parallelism""" + if IS_HIP_EXTENSION: + tex.lse_correction(softmax_lse, softmax_lse_per_step, False) + return max_scale = torch.max(softmax_lse, softmax_lse_per_step) min_scale = torch.min(softmax_lse, softmax_lse_per_step) new_scale = max_scale + torch.log1p(torch.exp(min_scale - max_scale)) softmax_lse.copy_(new_scale) -@jit_fuser +@_lse_correction_fuser def flash_attn_fwd_second_half_softmax_lse_correction( softmax_lse: torch.Tensor, softmax_lse_per_step: torch.Tensor, ): """Merge second half of softmax stats of each step in Attention with context parallelism""" + if IS_HIP_EXTENSION: + tex.lse_correction(softmax_lse, softmax_lse_per_step, True) + return softmax_lse_ = softmax_lse[..., 1, :] max_scale = torch.max(softmax_lse_, softmax_lse_per_step) min_scale = torch.min(softmax_lse_, softmax_lse_per_step) diff --git a/transformer_engine/pytorch/csrc/extensions.h b/transformer_engine/pytorch/csrc/extensions.h index 1ee3b087f0..8a6934b475 100644 --- a/transformer_engine/pytorch/csrc/extensions.h +++ b/transformer_engine/pytorch/csrc/extensions.h @@ -544,9 +544,11 @@ std::tuple swizzle_scales_and_pack_ptrs_for_ } // namespace grouped_mlp_experimental /*************************************************************************************************** - * Support THD format for Context Parallel + * Support Context Parallel **************************************************************************************************/ +void lse_correction(at::Tensor lse, const at::Tensor &lse_per_step, bool only_second_half); + at::Tensor thd_read_half_tensor(const at::Tensor &tensor, const at::Tensor &cu_seqlens, int half_idx); diff --git a/transformer_engine/pytorch/csrc/extensions/attention.cpp b/transformer_engine/pytorch/csrc/extensions/attention.cpp index b74c886909..d92ecfb92d 100644 --- a/transformer_engine/pytorch/csrc/extensions/attention.cpp +++ b/transformer_engine/pytorch/csrc/extensions/attention.cpp @@ -835,9 +835,34 @@ at::Tensor thd_read_half_tensor(const at::Tensor &tensor, const at::Tensor &cu_s } /*************************************************************************************************** - * Support THD format for Context Parallel: softmax_lse related operations + * Support Context Parallel: softmax_lse related operations **************************************************************************************************/ +void lse_correction(at::Tensor lse, const at::Tensor &lse_per_step, bool only_second_half) { + NVTE_CHECK(lse.is_cuda(), "lse must be a CUDA tensor"); + NVTE_CHECK(lse_per_step.is_cuda(), "lse_per_step must be a CUDA tensor"); + NVTE_CHECK(lse.scalar_type() == at::ScalarType::Float); + NVTE_CHECK(lse_per_step.scalar_type() == at::ScalarType::Float); + NVTE_CHECK(lse.is_contiguous(), "lse must be contiguous"); + NVTE_CHECK(lse_per_step.is_contiguous(), "lse_per_step must be contiguous"); + NVTE_CHECK(lse_per_step.dim() >= 1); + if (only_second_half) { + NVTE_CHECK(lse.dim() == lse_per_step.dim() + 1); + NVTE_CHECK(lse.size(-2) == 2); + NVTE_CHECK(lse.size(-1) == lse_per_step.size(-1)); + for (int64_t i = 0; i + 1 < lse_per_step.dim(); ++i) { + NVTE_CHECK(lse.size(i) == lse_per_step.size(i)); + } + } else { + NVTE_CHECK(lse.sizes() == lse_per_step.sizes()); + } + + auto te_lse = makeTransformerEngineTensor(lse); + auto te_lse_per_step = makeTransformerEngineTensor(lse_per_step); + nvte_cp_lse_correction(te_lse.data(), te_lse_per_step.data(), only_second_half, + at::cuda::getCurrentCUDAStream()); +} + void thd_second_half_lse_correction(at::Tensor lse, const at::Tensor &lse_per_step, const at::Tensor &cu_seqlens, bool lse_packed) { NVTE_CHECK(lse.scalar_type() == at::ScalarType::Float); diff --git a/transformer_engine/pytorch/csrc/extensions/pybind.cpp b/transformer_engine/pytorch/csrc/extensions/pybind.cpp index dd77a2ef32..8470f1eb43 100644 --- a/transformer_engine/pytorch/csrc/extensions/pybind.cpp +++ b/transformer_engine/pytorch/csrc/extensions/pybind.cpp @@ -582,7 +582,10 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { m.def("get_num_cublas_streams", &nvte_get_num_compute_streams, "Get number of compute streams", py::call_guard()); - // Support THD format for Context Parallel + // Support Context Parallel + m.def("lse_correction", &transformer_engine::pytorch::lse_correction, + "Correct the softmax_lse for dense context parallel layouts", + py::call_guard()); m.def("thd_read_half_tensor", &transformer_engine::pytorch::thd_read_half_tensor, "Read the first half(half_idx=0) or the second half(half_idx=1) of each sequence in a THD " "tensor",