diff --git a/transformer_engine/common/cast/fp8/gated_fp8.cuh b/transformer_engine/common/cast/fp8/gated_fp8.cuh
index 522a9add8f..52ab6f49e9 100644
--- a/transformer_engine/common/cast/fp8/gated_fp8.cuh
+++ b/transformer_engine/common/cast/fp8/gated_fp8.cuh
@@ -68,11 +68,9 @@ __global__ void __launch_bounds__(THREADS_PER_CHUNK)
   const float scale = (scale_ptr != nullptr) ? *scale_ptr : 1;
 
   extern __shared__ char dynamic_shmem[];
-  uintptr_t base_shmem_ptr = reinterpret_cast<uintptr_t>(dynamic_shmem);
   // Manually align dynamic SHMEM per TMA requirements using padding
   // __align__(128) Does not guarantee the pointer to be aligned!
-  uintptr_t dshmem = (base_shmem_ptr + TMA_SHMEM_ALIGNMENT - 1) &
-                     ~(static_cast<uintptr_t>(TMA_SHMEM_ALIGNMENT - 1));
+  char *dshmem = align_up(dynamic_shmem, TMA_SHMEM_ALIGNMENT);
 
   constexpr size_t buff_elems = SHMEM_DIM_Y * SHMEM_DIM_X;
   constexpr size_t buff_elems_total = BUFFERS_NUM * buff_elems;
diff --git a/transformer_engine/common/cast/mxfp8/gated_mxfp8.cuh b/transformer_engine/common/cast/mxfp8/gated_mxfp8.cuh
index 83b5a49cae..879c6befc3 100644
--- a/transformer_engine/common/cast/mxfp8/gated_mxfp8.cuh
+++ b/transformer_engine/common/cast/mxfp8/gated_mxfp8.cuh
@@ -124,11 +124,9 @@ __global__ void __launch_bounds__(THREADS_PER_CHUNK)
   __shared__ float subamax_colwise_buff[SUBAMAX_BUFF_DIM_Y][CHUNK_DIM_X];
 
   extern __shared__ char dynamic_shmem[];
-  uintptr_t base_shmem_ptr = reinterpret_cast<uintptr_t>(dynamic_shmem);
   // Manually align dynamic SHMEM per TMA requirements using padding
   // __align__(128) Does not guarantee the pointer to be aligned!
-  uintptr_t dshmem = (base_shmem_ptr + TMA_SHMEM_ALIGNMENT - 1) &
-                     ~(static_cast<uintptr_t>(TMA_SHMEM_ALIGNMENT - 1));
+  char *dshmem = align_up(dynamic_shmem, TMA_SHMEM_ALIGNMENT);
 
   constexpr size_t buff_elems = BUFF_DIM_Y * BUFF_DIM_X;
   constexpr size_t buff_elems_total = BUFFS_NUM * buff_elems;
diff --git a/transformer_engine/common/cast/mxfp8/quantize_mxfp8.cuh b/transformer_engine/common/cast/mxfp8/quantize_mxfp8.cuh
index 8c57f112d2..9f84195a19 100644
--- a/transformer_engine/common/cast/mxfp8/quantize_mxfp8.cuh
+++ b/transformer_engine/common/cast/mxfp8/quantize_mxfp8.cuh
@@ -130,11 +130,9 @@ __global__ void __launch_bounds__(THREADS_PER_CHUNK)
   constexpr size_t out_mem_rowwise = (ROWWISE_SCALING ? buff_size_aligned_out : 0);
 
   extern __shared__ char dynamic_shmem[];
-  uintptr_t base_shmem_ptr = reinterpret_cast<uintptr_t>(dynamic_shmem);
   // Manually align dynamic SHMEM per TMA requirements using padding
   // __align__(128) Does not guarantee the pointer to be aligned!
-  uintptr_t dshmem = (base_shmem_ptr + TMA_SHMEM_ALIGNMENT - 1) &
-                     ~(static_cast<uintptr_t>(TMA_SHMEM_ALIGNMENT - 1));
+  char *dshmem = align_up(dynamic_shmem, TMA_SHMEM_ALIGNMENT);
 
   // The destination shared memory buffer of a bulk tensor operation should be 16-byte aligned
   IType *in_sh = reinterpret_cast<IType *>(dshmem);
diff --git a/transformer_engine/common/cast/mxfp8/specialized/quantize_mxfp8.cuh b/transformer_engine/common/cast/mxfp8/specialized/quantize_mxfp8.cuh
index 81e11c40b8..0c59733d31 100644
--- a/transformer_engine/common/cast/mxfp8/specialized/quantize_mxfp8.cuh
+++ b/transformer_engine/common/cast/mxfp8/specialized/quantize_mxfp8.cuh
@@ -670,10 +670,6 @@ struct CastTraits<_IType, _OType, /*rowwise=*/true, /*colwise=*/true> {
                                         smem_alignment + smem_rowwise_scale + smem_colwise_reduce);
 };
 
-__device__ __forceinline__ intptr_t align_to(intptr_t x, intptr_t align) {
-  return (x + align - 1) & ~((align)-1);
-}
-
 // 32x32
 template <typename CastTraits,
           std::enable_if_t<CastTraits::isRowwise && CastTraits::isColwise, int> = 0,
@@ -706,8 +702,7 @@ __global__ void quantize_mxfp8_kernel_cast_only(
   block_coords.x = blockIdx.x * CastTraits::blockDIM::N;
 
   extern __shared__ char smem[];
-  char *smemAligned = reinterpret_cast<char *>(
-      align_to(reinterpret_cast<intptr_t>(smem), CastTraits::smem_alignment));
+  char *smemAligned = align_up(smem, CastTraits::smem_alignment);
 
   IType *sInput = reinterpret_cast<IType *>(smemAligned);
   inputUnitType *sInputUnit = reinterpret_cast<inputUnitType *>(sInput);
@@ -1179,8 +1174,7 @@ __global__ void quantize_mxfp8_kernel_cast_only(
   block_coords.x = blockIdx.x * CastTraits::blockDIM::N;
 
   extern __shared__ char smem[];
-  char *smemAligned = reinterpret_cast<char *>(
-      align_to(reinterpret_cast<intptr_t>(smem), CastTraits::smem_alignment));
+  char *smemAligned = align_up(smem, CastTraits::smem_alignment);
   IType *sInput = reinterpret_cast<IType *>(smemAligned);
   inputUnitType *sInputUnit = reinterpret_cast<inputUnitType *>(sInput);
 
diff --git a/transformer_engine/common/cast/nvfp4/group_quantize_transpose_nvfp4.cuh b/transformer_engine/common/cast/nvfp4/group_quantize_transpose_nvfp4.cuh
index 91c6af26b5..0b980c5086 100644
--- a/transformer_engine/common/cast/nvfp4/group_quantize_transpose_nvfp4.cuh
+++ b/transformer_engine/common/cast/nvfp4/group_quantize_transpose_nvfp4.cuh
@@ -262,11 +262,9 @@ __global__ void __launch_bounds__(THREADS_NUM)
   constexpr size_t out_mem_rowwise_scales = 0;
 
   extern __shared__ char dynamic_shmem[];
-  uintptr_t base_shmem_ptr = reinterpret_cast<uintptr_t>(dynamic_shmem);
   // Manually align dynamic SHMEM per TMA requirements using padding
   // __align__(128) Does not guarantee the pointer to be aligned!
-  uintptr_t dshmem = (base_shmem_ptr + TMA_SHMEM_ALIGNMENT - 1) &
-                     ~(static_cast<uintptr_t>(TMA_SHMEM_ALIGNMENT - 1));
+  char *dshmem = align_up(dynamic_shmem, TMA_SHMEM_ALIGNMENT);
 
   // The destination shared memory buffer of a bulk tensor operation should be 16-byte aligned
   IType *in_sh = reinterpret_cast<IType *>(dshmem);
diff --git a/transformer_engine/common/cast/nvfp4/quantize_transpose_nvfp4.cuh b/transformer_engine/common/cast/nvfp4/quantize_transpose_nvfp4.cuh
index a38a620ebe..8fc4cc7c46 100644
--- a/transformer_engine/common/cast/nvfp4/quantize_transpose_nvfp4.cuh
+++ b/transformer_engine/common/cast/nvfp4/quantize_transpose_nvfp4.cuh
@@ -410,11 +410,9 @@ __global__ void __launch_bounds__(THREADS_NUM)
   constexpr size_t out_mem_rowwise_scales = 0;
 
   extern __shared__ char dynamic_shmem[];
-  uintptr_t base_shmem_ptr = reinterpret_cast<uintptr_t>(dynamic_shmem);
   // Manually align dynamic SHMEM per TMA requirements using padding
   // __align__(128) Does not guarantee the pointer to be aligned!
-  uintptr_t dshmem = (base_shmem_ptr + TMA_SHMEM_ALIGNMENT - 1) &
-                     ~(static_cast<uintptr_t>(TMA_SHMEM_ALIGNMENT - 1));
+  char *dshmem = align_up(dynamic_shmem, TMA_SHMEM_ALIGNMENT);
 
   // The destination shared memory buffer of a bulk tensor operation should be 16-byte aligned
   IType *in_sh = reinterpret_cast<IType *>(dshmem);
@@ -952,11 +950,9 @@ __global__ void __launch_bounds__(THREADS_NUM)
   constexpr size_t out_mem_rowwise_scales = 0;
 
   extern __shared__ char dynamic_shmem[];
-  uintptr_t base_shmem_ptr = reinterpret_cast<uintptr_t>(dynamic_shmem);
   // Manually align dynamic SHMEM per TMA requirements using padding
   // __align__(128) Does not guarantee the pointer to be aligned!
-  uintptr_t dshmem = (base_shmem_ptr + TMA_SHMEM_ALIGNMENT - 1) &
-                     ~(static_cast<uintptr_t>(TMA_SHMEM_ALIGNMENT - 1));
+  char *dshmem = align_up(dynamic_shmem, TMA_SHMEM_ALIGNMENT);
 
   // The destination shared memory buffer of a bulk tensor operation should be 16-byte aligned
   IType *in_sh = reinterpret_cast<IType *>(dshmem);
diff --git a/transformer_engine/common/common.h b/transformer_engine/common/common.h
index eb4dcc055c..8fd7d7480c 100644
--- a/transformer_engine/common/common.h
+++ b/transformer_engine/common/common.h
@@ -1153,6 +1153,26 @@ inline bool is_aligned_ptr(const void *ptr, size_t alignment) {
   return reinterpret_cast<uintptr_t>(ptr) % alignment == 0;
 }
 
+/*! \brief Align a shared-memory base pointer up to `align` bytes.
+ *
+ * The result is derived from `p` by pointer arithmetic on purpose, without losing its
+ * identity as a pointer in between, so the address is never rounded through an integer
+ * -- in which case the compiler would lose the link back to the `extern __shared__`
+ * object, and ptxas could no longer prove the address lives in the shared window and
+ * would fall back to generic address-space accesses (`LD.E`/`ST.E`) instead of
+ * `LDS`/`STS`.
+ *
+ * `align` must be a power of two.
+ */
+__device__ __forceinline__ char *align_up(char *p, uintptr_t align) {
+  const uintptr_t misalign = reinterpret_cast<uintptr_t>(p) & (align - 1);
+  // If p is not aligned, (align - misalign) & (align - 1) is the number of bytes to fill the gap between p and
+  // the next aligned address.
+  // If p is aligned, misalign is 0 and (align - misalign) is align itself, so we use & (align - 1)
+  // to make it 0 and return p itself.
+  return p + ((align - misalign) & (align - 1));
+}
+
 inline bool is_aligned_tensor_data(const Tensor &t, size_t alignment) {
   return is_aligned_ptr(static_cast<const void *>(t.data.dptr), alignment);
 }