[ROCm] Route dense CP softmax-LSE correction through a native kernel

tests/pytorch/attention/test_cp_utils.py

@@ -1,14 +1,19 @@
+# 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.
 
 """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()

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 <bool only_second_half>
+__global__ void lse_correction_kernel(float *lse, const float *lse_per_step, size_t rows,
+                                      size_t cols) {
+  size_t col = static_cast<size_t>(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<true><<<grid, block, 0, stream>>>(
+        reinterpret_cast<float *>(lse.data.dptr),
+        reinterpret_cast<const float *>(lse_per_step.data.dptr), rows, cols);
+  } else {
+    lse_correction_kernel<false><<<grid, block, 0, stream>>>(
+        reinterpret_cast<float *>(lse.data.dptr),
+        reinterpret_cast<const float *>(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);

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.

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)

transformer_engine/pytorch/csrc/extensions.h

@@ -544,9 +544,11 @@ std::tuple<at::Tensor, at::Tensor, at::Tensor> 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);
 

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);

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<py::gil_scoped_release>());
 
-  // 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<py::gil_scoped_release>());
   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",