enable blockwise FP8 quantization on rocm by asdfvg123 · Pull Request #609 · ROCm/TransformerEngine

asdfvg123 · 2026-06-03T17:53:36Z

Description

Please include a brief summary of the changes, relevant motivation and context.

Enable blockwise FP8 quantization on rocm

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

Please list the changes introduced in this PR:

remove HIP guard in quantization.py
guard kernels using TMA in quantization.
add branch to handle rocm for different threads per wave

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

matthiasdiener · 2026-06-03T18:28:00Z

 # TODO replace with call to fp8.py when recipe added.
-recipe_available = not IS_HIP_EXTENSION and (get_device_compute_capability() >= (9, 0) and float(torch.version.cuda) >= 12.8)
+if IS_HIP_EXTENSION:
+    recipe_available = get_device_compute_capability() >= (9, 0)


Wouldn't this be always True on ROCm TE?

This test targets MI300 and MI350 so I set to (9,0)

matthiasdiener · 2026-06-03T18:29:11Z

@@ -1 +1 @@
 /*************************************************************************


Needs AMD copyright

matthiasdiener · 2026-06-03T18:30:03Z

+#ifndef __HIP_PLATFORM_AMD__
 #include <cudaTypedefs.h>
+#endif
 #include <cuda_bf16.h>
 #include <cuda_runtime.h>

 #include <cfloat>
+#ifndef __HIP_PLATFORM_AMD__
 #include <cuda/barrier>
+#endif

 #include "common/common.h"
 #include "common/recipe/recipe_common.cuh"
 #include "common/util/cuda_runtime.h"
+#ifndef __HIP_PLATFORM_AMD__
 #include "common/util/ptx.cuh"
+#endif


These #includes should be already disabled via hipify, so probably no need for the #ifndefs here.

matthiasdiener · 2026-06-03T18:43:26Z

+  static constexpr float max = 448.0f;
+  static constexpr float max_inverse = 1.0 / max;


Is this change necessary? fp8e4m3 max depends on the device type on AMD.

quantize_transpose_square_blockwise.cu and quantize_transpose_vector_blockwise.cu use
compute_scale_from_types<IType, fp8e4m3> for the first time, which exposed a latent bug in common.h

The #else branch of TypeExtrema<fp8e4m3> declared max as a static float,
This caused the constexpr static float max_finite_value initializer in TypeInfo in the same file to fail when the template was instantiated on the host.

The fix uses HIP_FP8_TYPE_FNUZ, used in hip_float8.h for selecting FNUZ at compile time, to make the host-pass branch constexpr as well.

If the value really used host size, it should be runtime detected. If it is only for host translation of GPU code (i.e.. results are discarded), you can keep 448, no extra ifdefs is needed

reverted to the original upstream. I instead changed the recipe_common.cuh following the other convention in quantize_transpose_vector_blockwise_fp4.cu L230

alextmagro · 2026-06-03T18:48:39Z

Could you give a description of what you want to achieve with this PR? My understanding is that block fp8 quantization relies on some upstream kernels that will need to be adapted for AMD.

If you're just trying to enable the interface, I would argue that we should do this last, after we have a working quantization and GEMM path (and enabled and passing C++/Python tests).

asdfvg123 · 2026-06-03T18:57:07Z

@alextmagro
This PR is to enable only the quantization in the AMD gpus, not the GEMM. There are two kernels in the upstream which uses TMA for the quantization and does not uses TMA for the quantization. I guarded the kernels which uses TMA and used the non-TMA kernels to quantize for AMD.

I tested with
tests/pytorch/test_float8blockwisetensor.py
and it passes [175 passed / 32 xpassed / 5 warnings]

alextmagro · 2026-06-03T19:10:43Z

@alextmagro This PR is to enable only the quantization in the AMD gpus, not the GEMM. There are two kernels in the upstream which uses TMA for the quantization and does not uses TMA for the quantization. I guarded the kernels which uses TMA and used the non-TMA kernels to quantize for AMD.

I tested with tests/pytorch/test_float8blockwisetensor.py and it passes [175 passed / 32 xpassed / 5 warnings]

OK, in that case we need to add the cpp blockwise tests to the CMake file, and the pytorch test file to ci/pytorch.sh.

…dant HIP guards, revert unnecessary common.h change

alextmagro · 2026-06-04T17:44:59Z

    run_default_fa 1 test_deferred_init.py
    run_default_fa 1 test_quantized_tensor.py
    run_default_fa 1 test_float8_current_scaling_exact.py
+    run_default_fa 1 test_float8blockwisetensor.py


nit: This test and test_float8_current_scaling_exact should be ahead of test_quantized_tensor.

alextmagro · 2026-06-04T17:45:25Z

  }

-  float qscale_inv = 1.0 / qscale;
+  float qscale_inv = 1.0f / qscale;


We had discussed this as a potential solution offline, but if the tolerances are needed then we should remove this

alextmagro · 2026-06-04T17:46:33Z

 void compare_scaling_factors(const std::string& name, const float* test, const float* ref,
                             const size_t row_blocks, const size_t col_blocks,
-                             const size_t test_stride, const size_t ref_stride) {
+                             const size_t test_stride, const size_t ref_stride,


Both function signature and the function change should be guarded. We might also just want to put the atol into the function body itself for now, since it is only used for blockwise scaling.

alextmagro · 2026-06-04T17:48:14Z

 # TODO replace with call to fp8.py when recipe added.
-recipe_available = not IS_HIP_EXTENSION and (get_device_compute_capability() >= (9, 0) and float(torch.version.cuda) >= 12.8)
+if IS_HIP_EXTENSION:
+    recipe_available = get_device_compute_capability() >= (9, 5)


I think we should be able to support MI300s for this, right? If so, > (9, 0) should be correct.

fixed, keeping (9,0)

alextmagro · 2026-06-04T17:50:40Z

+#ifdef __HIP_PLATFORM_AMD__
+  warp_tile_amax = blockwise_warp_reduce_max(amax);
+#else
  warp_tile_amax = warp_reduce_max<kThreadsPerWarp>(amax);


We should be able to just use warp_reduce_max here, and remove the kThreadsPerWarp=64 logic too. For the most part, the compiler will double up and we will be okay here.

warp_reduce_max in transformer_engine/common/utils.cuh uses THREADS_PER_WARP = 32 in the file which creates bug. Let me know if there is a better way

alextmagro · 2026-06-04T17:52:03Z

  }
 }

+#ifdef TMA_HW_SUPPORTED


I would prefer if we continued to use the AMD guard rather than TMA_HW_SUPPORTED here. Additionally, I am not sure we need to guard the def block around TMA_HW_SUPPORTED as all the cuda arch macros should be undefined.

Fixed and now using AMD guard. Keep the upstream's TMA_HW_SUPPORTED guard, to guard non-TMA kernel

alextmagro · 2026-06-04T17:53:37Z

                        scale_stride_x, scale_stride_y, scale_t_stride_x, scale_t_stride_y, epsilon,
                        tensor_map_output_trans, pow_2_scale, noop_ptr);
-              } else {
+              } else


Let's avoid splitting up the } else { line here. We can add another macro guard instead if needed.

I could not think of a better way than the current code. I guarded and called the kernel.

alextmagro · 2026-06-04T17:54:23Z

 constexpr int kNVecSMem = 2;   // The number of elements each LDS/STS touches
+
+#ifdef __HIP_PLATFORM_AMD__
+constexpr int kThreadsPerBlock = 512;  // Thread block size, 8 warps (wave64) in total


Are there actual performance improvements for increasing the # of threads and the threads per warp? If not, we should use the already present values for now.

The kernel expects 8 waves / block , so I increased the number of threads

alextmagro · 2026-06-04T17:55:18Z

 def check_fp8_block_scaling_support() -> Tuple[bool, str]:
    """Return if fp8 block scaling support is available"""
    if IS_HIP_EXTENSION:
-        return False, "FP8 block scaled gemm not yet supported for ROCm"


We need the same arch guard here as throughout the rest of the PR

alextmagro · 2026-06-04T18:12:20Z

By the way, to run CI you need to add a CI level label. L3 is required before merging, L1 is for lighter testing, mostly sGPU tests, if you are midway through the ticket and expect to make more changes

ipanfilo · 2026-06-05T02:52:57Z

ipanfilo · 2026-06-05T02:58:50Z


+#ifdef __HIP_PLATFORM_AMD__
+using WarpSyncMask = uint64_t;
+constexpr WarpSyncMask kFullWarpMask = 0xFFFFFFFFFFFFFFFFULL;


ROCm should not use it. See how *_sync calls are guarded in other places

removed the mask and use ROCm __shfl instead of __shfl_sync

ipanfilo · 2026-06-05T02:59:39Z

    }
  }
  // Reduce amax in the warp (32x32 tile)
+#ifdef __HIP_PLATFORM_AMD__


The whole this code is under #ifndef HIP_PLATFORM_AMD

removed the dead branch

ipanfilo · 2026-06-05T03:00:07Z

 // const values configuration

+#ifdef __HIP_PLATFORM_AMD__
+constexpr size_t kThreadsPerWarp = 64;


It is platform dependent.

fixed now guarded with gfx1250 for 32 threads

ipanfilo · 2026-06-05T03:03:22Z

     transpose/multi_cast_transpose.cu
-     transpose/quantize_transpose_vector_blockwise.cu #CUDA-only
+     transpose/quantize_transpose_vector_blockwise.cu
+     transpose/quantize_transpose_square_blockwise.cu


It should stay in transformer_engine_cuda_arch_specific_sources

ipanfilo · 2026-06-05T03:06:08Z

+  static constexpr float max = 448.0f;
+  static constexpr float max_inverse = 1.0 / max;


If the value really used host size, it should be runtime detected. If it is only for host translation of GPU code (i.e.. results are discarded), you can keep 448, no extra ifdefs is needed

asdfvg123 · 2026-06-10T20:07:06Z

MI300 has 64KB of LDS which makes overflow when loading 128 * 128 FP32 data into LDS. I created a helper and branched the kernel. When loading FP32 data, the kernel loads 128 * 64 chunk of data and iterate to quantize. From the host's view, the kernel quantizes 128 * 128 elements.

enable blockwise FP8 quantization on rocm

8335488

asdfvg123 requested review from alextmagro, wangye805 and wenchenvincent June 3, 2026 17:53

asdfvg123 requested a review from ipanfilo as a code owner June 3, 2026 17:53

matthiasdiener reviewed Jun 3, 2026

View reviewed changes

enable blockwise FP8 C++ tests on ROCm, fix wave64 bugs, remove redun…

6226301

…dant HIP guards, revert unnecessary common.h change

alextmagro requested changes Jun 4, 2026

View reviewed changes

fix constexpr chain, arch guard

bdf905e

asdfvg123 added the ci-level 1 CI test level 1 label Jun 4, 2026

ipanfilo requested changes Jun 5, 2026

View reviewed changes

asdfvg123 added 4 commits June 9, 2026 22:58

guard kThreadsPerWarp for wave32, derive fp8 max on AMD host

676d1f0

revert unnecessary guard, add license

f8a0fc5

use non-sync __shfl on ROCm, drop kFullWarpMask

231e381

fix fp32 LDS overflow in vector blockwise on gfx942

e158d3e

asdfvg123 requested review from alextmagro, ipanfilo and matthiasdiener June 10, 2026 20:07

		@@ -1 +1 @@
		/*************************************************************************

		static constexpr float max = 448.0f;
		static constexpr float max_inverse = 1.0 / max;

Conversation

asdfvg123 commented Jun 3, 2026

Description

Type of change

Changes

Checklist:

Uh oh!

Choose a reason for hiding this comment

Uh oh!

asdfvg123 Jun 4, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

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!

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

Uh oh!

Uh oh!

asdfvg123 commented Jun 3, 2026

Uh oh!

alextmagro commented Jun 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!

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!

alextmagro commented Jun 4, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

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

asdfvg123 Jun 4, 2026 •

edited

Loading

alextmagro commented Jun 3, 2026 •

edited

Loading

alextmagro commented Jun 4, 2026 •

edited

Loading