Skip to content

NVRTC NVFP4 quantization kernels #3301

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

Open
CarlosGomes98 wants to merge 6 commits into
base: main
Choose a base branch
from
Open

NVRTC NVFP4 quantization kernels #3301

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
6 commits
Select commit Hold shift + click to select a range
f6f2c54
[Common] Add architecture-aware NVRTC compilation
CarlosGomes98 Aug 3, 2026
e4eeeb9
[Common] Compile NVFP4 4over6 quantize with NVRTC
CarlosGomes98 Aug 3, 2026
124d874
[pre-commit.ci] auto fixes from pre-commit.com hooks
pre-commit-ci[bot] Aug 3, 2026
deb6591
[Common] Address NVFP4 NVRTC review feedback
CarlosGomes98 Aug 4, 2026
04df751
guard compile_quantize_4over6_rtc with FP4_TYPE_SUPPORTED
CarlosGomes98 Aug 4, 2026
b080615
add to CMake
CarlosGomes98 Aug 4, 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
25 changes: 25 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 @@ -181,6 +181,7 @@ list(APPEND transformer_engine_cpp_sources
cudnn_utils.cpp
transformer_engine.cpp
fused_attn/fused_attn.cpp
cast/nvfp4/rtc_dispatch.cpp
gemm/config.cpp
normalization/common.cpp
normalization/rtc_dispatch.cpp
Expand Down Expand Up @@ -565,6 +566,20 @@ make_string_header_from_file(utils.cuh
make_string_header_from_file(util/math.h
string_code_util_math_h)

# NVFP4 4over6 quantize: NVRTC bundled headers + RTC source file
make_string_header_from_file(include/transformer_engine/nvfp4/4over6.h
string_code_transformer_engine_nvfp4_4over6_h)
make_string_header_from_file(util/type_extrema.h
string_code_util_type_extrema_h)
make_string_header_from_file(util/ptx.cuh
string_code_util_ptx_cuh)
make_string_header_from_file(cast/nvfp4/core_nvfp4.cuh
string_code_cast_nvfp4_core_nvfp4_cuh)
make_string_header_from_file(cast/nvfp4/quantize_4over6_kernel.cuh
string_code_cast_nvfp4_quantize_4over6_kernel_cuh)
make_string_header_from_file(cast/nvfp4/rtc/quantize_4over6.cu
string_code_cast_nvfp4_rtc_quantize_4over6_cu)

# Norm NVRTC bundled headers + RTC source files
make_string_header_from_file(normalization/kernel_params.h
string_code_normalization_kernel_params_h)
Expand All @@ -589,6 +604,16 @@ make_string_header_from_file(normalization/rmsnorm/rtc/rmsnorm_bwd_kernel.cu
target_include_directories(transformer_engine PRIVATE
"${CMAKE_CURRENT_BINARY_DIR}/string_headers")

# Default OFF: the NVFP4 4over6 quantize kernel (384-way fanout) is compiled via
# NVRTC at runtime. Set ON to additionally statically instantiate it as the
# fallback selected when NVTE_DISABLE_NVRTC=1.
option(NVTE_BUILD_LEGACY_STATIC_NVFP4
"Also compile the static NVFP4 4over6 quantize kernel for NVTE_DISABLE_NVRTC fallback"
OFF)
target_compile_definitions(transformer_engine
PRIVATE
NVTE_BUILD_LEGACY_STATIC_NVFP4=$<BOOL:${NVTE_BUILD_LEGACY_STATIC_NVFP4}>)

option(NVTE_BUILD_LEGACY_STATIC_FUSED_SOFTMAX
"Also compile legacy static fused softmax kernels for NVTE_DISABLE_NVRTC fallback"
OFF)
Expand Down
21 changes: 21 additions & 0 deletions transformer_engine/common/cast/nvfp4/core_nvfp4.cuh
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -11,6 +11,7 @@
#ifndef TRANSFORMER_ENGINE_CORE_NVFP4_CUH_
#define TRANSFORMER_ENGINE_CORE_NVFP4_CUH_

#if !defined(__CUDACC_RTC__)
#include <cuda.h>
#include <cudaTypedefs.h>
#include <cuda_runtime.h>
Expand All @@ -22,11 +23,28 @@
#include "../../util/math.h"
#include "../../util/ptx.cuh"
#include "../../utils.cuh"
#else
// NVRTC build: common.h (host-only: cuDNN/cutlass) cannot be parsed by NVRTC.
// utils.cuh, util/math.h and ptx.cuh are injected as in-memory headers by the
// RTC dispatch and already provide the integer typedefs, detail::is_same, and
// the fp8 element types. util/type_extrema.h (also injected) provides the
// transformer_engine-namespace fp4 aliases and detail::TypeExtrema
// specializations that would otherwise come from common.h.
#include "ptx.cuh"
#include "util/math.h"
#include "utils.cuh"
#endif // __CUDACC_RTC__

#if FP4_TYPE_SUPPORTED
#include <cuda_fp4.h>
#endif // FP4_TYPE_SUPPORTED

#if defined(__CUDACC_RTC__)
namespace transformer_engine {
#include "util/type_extrema.h"
} // namespace transformer_engine
#endif // __CUDACC_RTC__

namespace transformer_engine {
namespace dispatch {
namespace nvfp4 {
Expand Down Expand Up @@ -94,6 +112,8 @@ __device__ __forceinline__ float compute_global_encode_scaling_factor_FP4(const
return global_encode_scale;
}

#if !defined(__CUDACC_RTC__)

@ptrendx ptrendx Aug 3, 2026

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

If you are not including this in the NVRTC build then how do you handle the stochastic rounding?

@CarlosGomes98 CarlosGomes98 Aug 4, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This particular kernel rejects stochastic rounding

  NVTE_CHECK(!quant_config->stochastic_rounding,
             "NVFP4 4over6 quantization does not support stochastic rounding.");

We can do a guard like this, or move this get_rbits to a separate header which we then include in the other kernels which use it.

// The RTC 4over6 path rejects stochastic rounding; static transpose kernels use this helper.
__device__ __forceinline__ uint32_t get_rbits(
transformer_engine::curanddx::detail::philox4x32_native_state<NVTE_BUILD_NUM_PHILOX_ROUNDS>
&rng,
Expand All @@ -108,6 +128,7 @@ __device__ __forceinline__ uint32_t get_rbits(
const uint32_t rbits = rbits_arr[rnd_idx++];
return rbits;
}
#endif // !__CUDACC_RTC__

#endif // FP4_TYPE_SUPPORTED

Expand Down
Loading
Loading