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(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(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..c7c4aa0fa9 100644 --- a/transformer_engine/common/cast/mxfp8/gated_mxfp8.cuh +++ b/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(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(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; 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(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(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(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..95f16e384c 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 = 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( - align_to(reinterpret_cast(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(smemAligned); inputUnitType *sInputUnit = reinterpret_cast(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( - align_to(reinterpret_cast(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(smemAligned); inputUnitType *sInputUnit = reinterpret_cast(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..d25a2c25f8 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,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(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(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(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..0578b6fc9f 100644 --- a/transformer_engine/common/cast/nvfp4/quantize_transpose_nvfp4.cuh +++ b/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(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(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(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(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(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(dshmem); diff --git a/transformer_engine/common/utils.cuh b/transformer_engine/common/utils.cuh index 0b75da622c..75aa8701ce 100644 --- a/transformer_engine/common/utils.cuh +++ b/transformer_engine/common/utils.cuh @@ -32,6 +32,30 @@ static_assert(sizeof(uint64_t) == 8); // 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 { + +/*! \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(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 {