[Common] Fix pointer arithmatic to generate correct LDS/STS instructions

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;

transformer_engine/common/cast/mxfp8/gated_mxfp8.cuh

@@ -124,11 +124,12 @@ __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));
+  // Keep the result derived from `dynamic_shmem` by pointer arithmetic: casting an
+  // integer back to a pointer loses the shared-memory provenance, so ptxas emits
+  // generic LD/ST instead of LDS/STS.
+  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;

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

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,11 @@ __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));
+  // Derive the aligned base from `smem` by pointer arithmetic. Round-tripping it
+  // through an integer and casting back loses the link to the `extern __shared__`
+  // object, so ptxas can no longer prove the address lives in the shared window
+  // and falls back to generic LD/ST instead of LDS/STS.
+  char *smemAligned = align_up(smem, CastTraits::smem_alignment);
 
   IType *sInput = reinterpret_cast<IType *>(smemAligned);
   inputUnitType *sInputUnit = reinterpret_cast<inputUnitType *>(sInput);
@@ -1179,8 +1178,11 @@ __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));
+  // Derive the aligned base from `smem` by pointer arithmetic. Round-tripping it
+  // through an integer and casting back loses the link to the `extern __shared__`
+  // object, so ptxas can no longer prove the address lives in the shared window
+  // and falls back to generic LD/ST instead of LDS/STS.
+  char *smemAligned = align_up(smem, CastTraits::smem_alignment);
   IType *sInput = reinterpret_cast<IType *>(smemAligned);
   inputUnitType *sInputUnit = reinterpret_cast<inputUnitType *>(sInput);
 

transformer_engine/common/cast/nvfp4/group_quantize_transpose_nvfp4.cuh

@@ -262,11 +262,12 @@ __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));
+  // Keep the result derived from `dynamic_shmem` by pointer arithmetic: casting an
+  // integer back to a pointer loses the shared-memory provenance, so ptxas emits
+  // generic LD/ST instead of LDS/STS.
+  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);

transformer_engine/common/cast/nvfp4/quantize_transpose_nvfp4.cuh

@@ -410,11 +410,12 @@ __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));
+  // Keep the result derived from `dynamic_shmem` by pointer arithmetic: casting an
+  // integer back to a pointer loses the shared-memory provenance, so ptxas emits
+  // generic LD/ST instead of LDS/STS.
+  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 +953,12 @@ __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));
+  // Keep the result derived from `dynamic_shmem` by pointer arithmetic: casting an
+  // integer back to a pointer loses the shared-memory provenance, so ptxas emits
+  // generic LD/ST instead of LDS/STS.
+  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);

transformer_engine/common/utils.cuh

@@ -32,6 +32,30 @@ static_assert(sizeof(uint64_t) == 8);
 
 // Minimal subset of <type_traits> used by RTC kernel headers. Keep these in a
 // project-owned namespace because adding primary templates to std is undefined.
+namespace transformer_engine {
+
+/*! \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));
+}
+
+}  // namespace transformer_engine
+
 namespace transformer_engine {
 namespace detail {