Skip to content

nvrtc MXFP8 kernels #3302

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
CarlosGomes98 wants to merge 4 commits into
base: main
Choose a base branch
from
Draft

nvrtc MXFP8 kernels #3302

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
f6f2c54
[Common] Add architecture-aware NVRTC compilation
CarlosGomes98 Aug 3, 2026
696c449
[Common] Compile specialized MXFP8 quantize kernels with NVRTC
CarlosGomes98 Aug 3, 2026
04801f4
cut down comment verbosity
CarlosGomes98 Aug 3, 2026
dc0d7f3
[pre-commit.ci] auto fixes from pre-commit.com hooks
pre-commit-ci[bot] Aug 3, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
27 changes: 27 additions & 0 deletions transformer_engine/common/CMakeLists.txt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -184,6 +184,7 @@ list(APPEND transformer_engine_cpp_sources
gemm/config.cpp
normalization/common.cpp
normalization/rtc_dispatch.cpp
cast/mxfp8/specialized/rtc_dispatch.cpp
normalization/layernorm/ln_api.cpp
normalization/rmsnorm/rmsnorm_api.cpp
util/cuda_driver.cpp
Expand Down Expand Up @@ -554,6 +555,21 @@ make_string_header_from_file(transpose/rtc/transpose.cu
string_code_transpose_rtc_transpose_cu)
make_string_header_from_file(transpose/rtc/swap_first_dims.cu
string_code_transpose_rtc_swap_first_dims_cu)

# MXFP8 specialized quantize: NVRTC bundled headers + RTC source file
make_string_header_from_file(util/ptx.cuh
string_code_util_ptx_cuh)
make_string_header_from_file(cast/mxfp8/specialized/swizzle.cuh
string_code_cast_mxfp8_specialized_swizzle_cuh)
make_string_header_from_file(cast/mxfp8/specialized/state_counter.cuh
string_code_cast_mxfp8_specialized_state_counter_cuh)
make_string_header_from_file(cast/mxfp8/specialized/quantize_mxfp8.cuh
string_code_cast_mxfp8_specialized_quantize_mxfp8_cuh)
make_string_header_from_file(cast/mxfp8/specialized/rtc/quantize_mxfp8_rowwise.cu
string_code_cast_mxfp8_specialized_rtc_quantize_mxfp8_rowwise_cu)
make_string_header_from_file(cast/mxfp8/specialized/rtc/quantize_mxfp8_bidimensional.cu
string_code_cast_mxfp8_specialized_rtc_quantize_mxfp8_bidimensional_cu)

make_string_header_from_file(fused_softmax/scaled_masked_softmax.cu
string_code_fused_softmax_scaled_masked_softmax_cu)
make_string_header_from_file(fused_softmax/scaled_upper_triang_masked_softmax.cu
Expand Down Expand Up @@ -586,6 +602,7 @@ make_string_header_from_file(normalization/rmsnorm/rtc/rmsnorm_fwd_kernel.cu
string_code_normalization_rmsnorm_rtc_rmsnorm_fwd_kernel_cu)
make_string_header_from_file(normalization/rmsnorm/rtc/rmsnorm_bwd_kernel.cu
string_code_normalization_rmsnorm_rtc_rmsnorm_bwd_kernel_cu)

target_include_directories(transformer_engine PRIVATE
"${CMAKE_CURRENT_BINARY_DIR}/string_headers")

Expand All @@ -607,6 +624,16 @@ target_compile_definitions(transformer_engine
PRIVATE
NVTE_BUILD_LEGACY_STATIC_NORM=$<BOOL:${NVTE_BUILD_LEGACY_STATIC_NORM}>)

# Default OFF: the 1x32 rowwise and 32x32 bidimensional MXFP8 specialized
# cast-only kernels are compiled via NVRTC at runtime. Set ON to additionally
# instantiate their static fallbacks for NVTE_DISABLE_NVRTC=1.
option(NVTE_BUILD_LEGACY_STATIC_MXFP8
"Also compile static specialized MXFP8 cast-only kernels for NVTE_DISABLE_NVRTC fallback"
OFF)
target_compile_definitions(transformer_engine
PRIVATE
NVTE_BUILD_LEGACY_STATIC_MXFP8=$<BOOL:${NVTE_BUILD_LEGACY_STATIC_MXFP8}>)

# Compiler options
set(nvte_sources_with_fast_math)
list(APPEND nvte_sources_with_fast_math fused_softmax/scaled_masked_softmax.cu
Expand Down
97 changes: 72 additions & 25 deletions transformer_engine/common/cast/mxfp8/quantize_mxfp8.cuh
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -19,9 +19,11 @@
#include "../../common.h"
#include "../../util/math.h"
#include "../../util/ptx.cuh"
#include "../../util/rtc.h"
#include "../../utils.cuh"
#include "../core/common.cuh"
#include "specialized/quantize_mxfp8.cuh"
#include "specialized/rtc_dispatch.cuh"
#include "swizzle.cuh"

namespace transformer_engine {
Expand Down Expand Up @@ -680,30 +682,50 @@ void quantize(const Tensor &input, const Tensor *act_input, const Tensor *noop,
switch (scaling_type) {
case ScalingType::ROWWISE: {
using traits = specialized::CastTraits<IType, OType, true, false>;
auto kernel = specialized::quantize_mxfp8_kernel_cast_only<traits>;

NVTE_CHECK_CUDA(cudaFuncSetAttribute(
kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, traits::smem));

dim3 block(traits::threadLayout::num, traits::warpLayout::N,
traits::warpLayout::M);
dim3 grid((cols + traits::blockDimN - 1) / traits::blockDimN,
(rows + traits::blockDimM - 1) / traits::blockDimM);
kernel<<<grid, block, traits::smem, stream>>>(
reinterpret_cast<typename traits::IType *>(input.data.dptr),
reinterpret_cast<typename traits::OType *>(output->data.dptr),
scales_rowwise_ptr, noop_ptr, rows, cols, scale_stride_rowwise,
scale_stride_colwise);
auto *rowwise_input =
reinterpret_cast<typename traits::IType *>(input.data.dptr);
auto *rowwise_output =
reinterpret_cast<typename traits::OType *>(output->data.dptr);

// Prefer the NVRTC-compiled kernel; fall back to the static
// instantiation only when it was compiled in and NVRTC is disabled.
#if NVTE_BUILD_LEGACY_STATIC_MXFP8
const bool use_rtc = rtc::is_enabled();
#else
constexpr bool use_rtc = true;
#endif
if (use_rtc) {
specialized::launch_rowwise_cast_only_rtc<typename traits::IType,
typename traits::OType>(
rowwise_input, rowwise_output, scales_rowwise_ptr, noop_ptr,
static_cast<int32_t>(rows), static_cast<int32_t>(cols),
static_cast<int32_t>(scale_stride_rowwise),
static_cast<int32_t>(scale_stride_colwise), stream);
} else {
#if NVTE_BUILD_LEGACY_STATIC_MXFP8
auto kernel = specialized::quantize_mxfp8_kernel_cast_only<traits>;
NVTE_CHECK_CUDA(cudaFuncSetAttribute(
kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, traits::smem));
dim3 block(traits::threadLayout::num, traits::warpLayout::N,
traits::warpLayout::M);
dim3 grid((cols + traits::blockDimN - 1) / traits::blockDimN,
(rows + traits::blockDimM - 1) / traits::blockDimM);
kernel<<<grid, block, traits::smem, stream>>>(
rowwise_input, rowwise_output, scales_rowwise_ptr, noop_ptr, rows, cols,
scale_stride_rowwise, scale_stride_colwise);
#else
NVTE_ERROR(
"MXFP8 rowwise specialized cast-only kernel requires NVRTC. Unset "
"NVTE_DISABLE_NVRTC, or rebuild with NVTE_BUILD_LEGACY_STATIC_MXFP8=ON "
"for the static fallback.");
#endif
}

break;
}
case ScalingType::BIDIMENSIONAL: {
using traits = specialized::CastTraits<IType, OType, true, true>;
auto kernel = specialized::quantize_mxfp8_kernel_cast_only<traits>;

NVTE_CHECK_CUDA(cudaFuncSetAttribute(
kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, traits::smem));
// TMA for loading, so that we don't need STS for transposing
alignas(64) CUtensorMap tensor_map_input{};
constexpr size_t input_type_bit_size = TypeInfo<IType>::size;
create_2D_tensor_map(tensor_map_input, input.data, rows, cols,
Expand All @@ -725,13 +747,38 @@ void quantize(const Tensor &input, const Tensor *act_input, const Tensor *noop,
cols, 0, output_type_bit_size,
traits::output_swizzle_pattern);

dim3 block(traits::rowThreadLayout::num, traits::numWarps);
dim3 grid((cols + traits::blockDIM::N - 1) / traits::blockDIM::N,
(rows + traits::blockDIM::M - 1) / traits::blockDIM::M);
kernel<<<grid, block, traits::smem, stream>>>(
tensor_map_input, tensor_map_rowwise_output, tensor_map_colwise_output,
scales_rowwise_ptr, scales_colwise_ptr, noop_ptr, rows, cols,
scale_stride_rowwise, scale_stride_colwise);
#if NVTE_BUILD_LEGACY_STATIC_MXFP8
const bool use_rtc = rtc::is_enabled();
#else
constexpr bool use_rtc = true;
#endif
if (use_rtc) {
specialized::launch_bidimensional_cast_only_rtc<typename traits::IType,
typename traits::OType>(
tensor_map_input, tensor_map_rowwise_output, tensor_map_colwise_output,
scales_rowwise_ptr, scales_colwise_ptr, noop_ptr,
static_cast<int32_t>(rows), static_cast<int32_t>(cols),
static_cast<int32_t>(scale_stride_rowwise),
static_cast<int32_t>(scale_stride_colwise), stream);
} else {
#if NVTE_BUILD_LEGACY_STATIC_MXFP8
auto kernel = specialized::quantize_mxfp8_kernel_cast_only<traits>;
NVTE_CHECK_CUDA(cudaFuncSetAttribute(
kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, traits::smem));
dim3 block(traits::rowThreadLayout::num, traits::numWarps);
dim3 grid((cols + traits::blockDIM::N - 1) / traits::blockDIM::N,
(rows + traits::blockDIM::M - 1) / traits::blockDIM::M);
kernel<<<grid, block, traits::smem, stream>>>(
tensor_map_input, tensor_map_rowwise_output, tensor_map_colwise_output,
scales_rowwise_ptr, scales_colwise_ptr, noop_ptr, rows, cols,
scale_stride_rowwise, scale_stride_colwise);
#else
NVTE_ERROR(
"MXFP8 bidimensional specialized cast-only kernel requires NVRTC. Unset "
"NVTE_DISABLE_NVRTC, or rebuild with NVTE_BUILD_LEGACY_STATIC_MXFP8=ON "
"for the static fallback.");
#endif
}

break;
}
Expand Down
Loading
Loading