NVRTC NVFP4 quantization kernels by CarlosGomes98 · Pull Request #3301 · NVIDIA/TransformerEngine

CarlosGomes98 · 2026-08-03T15:28:54Z

Description

Migrate NVFP4 quantization kernels to NVRTC. #3054

These kernels represent a large chunk of TE's build time for sm100 + in the cast.cu TU.

Migrating them to NVFP4 I record a 70% reduction in build time, due to the quantize_4over6 template fan-out being removed.

Fixes # (issue)

Type of change

Documentation change (change only to the documentation, either a fix or a new content)
Bug fix (non-breaking change which fixes an issue)
New feature (non-breaking change which adds functionality)
Breaking change (fix or feature that would cause existing functionality to not work as expected)
Infra/Build change
Code refactoring

Changes

Introduce nvfp4 nvrtc support. Keeps support for non-nvrtc path by extracting out the kernel body.
Support for Architecture family specification for NVRTC
A few additional util files due to the usual pain point of not being able to include some host-side libs in nvrtc.

Checklist:

I have read and followed the contributing guidelines
The functionality is complete
I have commented my code, particularly in hard-to-understand areas
I have made corresponding changes to the documentation
My changes generate no new warnings
I have added tests that prove my fix is effective or that my feature works
New and existing unit tests pass locally with my changes

Signed-off-by: CarlosGomes98 <carlosmiguel.gomes@live.com.pt>

for more information, see https://pre-commit.ci

greptile-apps · 2026-08-03T15:32:23Z

Greptile Summary

The PR moves NVFP4 4over6 quantization kernels to runtime NVRTC compilation while retaining an optional statically compiled fallback.

Extracts a shared device kernel body used by NVRTC and legacy static entry points.
Adds NVFP4 RTC source embedding, specialization dispatch, and kernel caching.
Extends RTC compilation with minimum-architecture and architecture-specific target requirements.
Extracts device-safe FP4 types, extrema, and public 4over6 mode definitions into reusable headers.
Adds a CMake option for building the legacy static NVFP4 fallback.

Confidence Score: 5/5

The PR appears safe to merge because no blocking failure remains within the eligible follow-up-review scope.

No blocking failure remains.

Important Files Changed

Filename	Overview
transformer_engine/common/cast/nvfp4/quantize_4over6_kernel.cuh	Extracts the NVFP4 4over6 device implementation into a shared kernel body with validated group-aligned output stores.
transformer_engine/common/cast/nvfp4/quantize_4over6_nvfp4.cuh	Replaces unconditional static dispatch with NVRTC dispatch and an optional legacy static fallback.
transformer_engine/common/cast/nvfp4/rtc_dispatch.cpp	Embeds device sources and compiles architecture-specific NVFP4 specializations through KernelManager.
transformer_engine/common/cast/nvfp4/rtc_dispatch.h	Constructs specialization-complete cache keys and launches compiled NVFP4 kernels.
transformer_engine/common/util/rtc.cpp	Adds supported-architecture discovery and exact architecture-specific compilation targets.
transformer_engine/common/CMakeLists.txt	Registers RTC sources and generated headers while making static NVFP4 instantiation opt-in.
transformer_engine/common/util/type_extrema.h	Extracts device-safe FP4 aliases and numeric extrema for both static CUDA and NVRTC compilation.

Flowchart

%%{init: {'theme': 'neutral'}}%%
flowchart TD
  A[NVFP4 4over6 cast request] --> B{Legacy static kernel built?}
  B -- No --> C[NVRTC dispatch]
  B -- Yes --> D{NVRTC enabled?}
  D -- Yes --> C
  D -- No --> E[Static kernel dispatch]
  C --> F[Build specialization cache key]
  F --> G{Kernel cached for device?}
  G -- No --> H[Compile architecture-specific kernel]
  G -- Yes --> I[Launch cached kernel]
  H --> I
  E --> J[Launch statically instantiated kernel]
  I --> K[Shared quantize_4over6 body]
  J --> K

_{Reviews (2): Last reviewed commit: "add to CMake" | Re-trigger Greptile}

fheinecke · 2026-08-03T15:48:40Z

/te-ci

ptrendx · 2026-08-03T23:58:02Z

  return global_encode_scale;
 }

+#if !defined(__CUDACC_RTC__)


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

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.

ptrendx · 2026-08-04T00:00:44Z

@@ -0,0 +1,26 @@
+/*************************************************************************


Could we have that in some directory rather than in the top-level include directory?

ptrendx · 2026-08-04T00:03:43Z

+                                             scale_stride_t, noop);
+}
+
+#if !defined(__CUDACC_RTC__)


Could we instead split this file so that the kernel is in its own header included in this file?

Agreed. I split this out into its own header to minimize the code that we pass to rtc. This naturally resolved a lot of the other comments you had regarding the amount of RTC guards I added.

ptrendx · 2026-08-04T00:06:12Z

+#ifndef TRANSFORMER_ENGINE_CAST_NVFP4_RTC_DISPATCH_CUH_
+#define TRANSFORMER_ENGINE_CAST_NVFP4_RTC_DISPATCH_CUH_
+
+#if !defined(__CUDACC_RTC__)


Considering this entire file is excluded, why do we even try to pass it to nvrtc? Also why is it .cuh extension? It does not seem to actually need to be CUDA header, but rather regular c++ header, no?

True, this was not actually being passed in the header list for nvrtc, so there was no reason for the .cuh

ptrendx · 2026-08-04T00:07:39Z

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


Snippets like these should just live outside the RTC macro rather than being duplicated in 2 code paths.

Done, with a few other examples I found of this

ptrendx · 2026-08-04T00:13:15Z

+
+  // Keep this matrix synchronized with NVTE_SPECIFIC_ARCHS in CMakeLists.txt.
+  // The runtime NVRTC version is the JIT equivalent of CUDAToolkit_VERSION.
+  switch (sm_arch) {


We don't need to actually repeat the logic of the specific archs from cmake here - since we are building those kernels at runtime on the target machine, we can just always compile for the "a" variant of the specific chip that we run on (for 90+, since that's where this became a thing). The cmake compiles to "a" or "f" in order to limit the compilation targets, here we don't have this problem.

That way we will also have less maintenance since we will not need this switch at all.

That makes a lot of sense - I was too focused on trying to exactly replicate the compile-time rules. Now should compile always to "a"

ptrendx · 2026-08-04T00:15:56Z

-  static constexpr float max = 0x1.FFCp15;
-};
-
+// Host-only fallback for types without an explicit specialization above.


Please put that together with the rest of things in type_extrema.h (just guard it), as it belongs there.

ptrendx · 2026-08-04T00:16:17Z

+using int8_t = signed char;
+using int16_t = short int;                 // NOLINT(*)
+using int32_t = int;                       // NOLINT(*)
+using int64_t = long long int;             // NOLINT(*)
+using intptr_t = long long int;            // NOLINT(*)
+using uintptr_t = unsigned long long int;  // NOLINT(*)


Why do we need those?

int16_t, int32_t, and int64_t are needed because NVRTC now has to parse ptx.cuh, which uses them, while avoiding . The rest were leftovers from the mxfp8 work, I've removed them.

Signed-off-by: CarlosGomes98 <carlosmiguel.gomes@live.com.pt>

CarlosGomes98 added 2 commits August 3, 2026 15:23

[Common] Add architecture-aware NVRTC compilation

f6f2c54

Signed-off-by: CarlosGomes98 <carlosmiguel.gomes@live.com.pt>

[Common] Compile NVFP4 4over6 quantize with NVRTC

e4eeeb9

Signed-off-by: CarlosGomes98 <carlosmiguel.gomes@live.com.pt>

github-actions Bot added the community-contribution PRs from external contributor outside the core maintainers, representing community-driven work. label Aug 3, 2026

CarlosGomes98 marked this pull request as ready for review August 3, 2026 15:29

CarlosGomes98 requested review from Oleg-Goncharov and ptrendx as code owners August 3, 2026 15:29

[pre-commit.ci] auto fixes from pre-commit.com hooks

124d874

for more information, see https://pre-commit.ci

ptrendx reviewed Aug 3, 2026

View reviewed changes

ptrendx reviewed Aug 4, 2026

View reviewed changes

CarlosGomes98 added 3 commits August 4, 2026 16:09

[Common] Address NVFP4 NVRTC review feedback

deb6591

Signed-off-by: CarlosGomes98 <carlosmiguel.gomes@live.com.pt>

guard compile_quantize_4over6_rtc with FP4_TYPE_SUPPORTED

04df751

Signed-off-by: CarlosGomes98 <carlosmiguel.gomes@live.com.pt>

add to CMake

b080615

Signed-off-by: CarlosGomes98 <carlosmiguel.gomes@live.com.pt>

CarlosGomes98 force-pushed the cgomes/nvrtc-quantize-nvfp4 branch from 2748041 to b080615 Compare August 4, 2026 15:10

		@@ -0,0 +1,26 @@
		/*************************************************************************

Uh oh!

Conversation

CarlosGomes98 commented Aug 3, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Description

Type of change

Changes

Checklist:

Uh oh!

greptile-apps Bot commented Aug 3, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Greptile Summary

Confidence Score: 5/5

Important Files Changed

Flowchart

Uh oh!

fheinecke commented Aug 3, 2026

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

CarlosGomes98 commented Aug 3, 2026 •

edited

Loading

greptile-apps Bot commented Aug 3, 2026 •

edited

Loading