NVFP4 cast/transpose without TMA

build_tools/hipify/custom_map.json

@@ -2,13 +2,20 @@
  "custom_map" : {
         "<cuda_bf16.h>" : "<hip/hip_bfloat16.h>",
         "<cuda_fp8.h>" : "\"common/amd_detail/hip_float8.h\"",
+        "<cuda_fp4.h>" : "<hip/hip_fp4.h>",
         "cuda_runtime.h\"" : "hip_runtime.h\"",
         "ATen/cudnn/Handle.h" : "ATen/miopen/Handle.h",
         "CUfunc_cache" : "hipFuncCache_t", 
         "<nvtx3/nvToolsExt.h>" : "<roctracer/roctx.h>",
         "cudaFuncSetAttribute(" : "hipFuncSetAttribute((const void*)",
         "__nv_bfloat162":"__hip_bfloat162",
-        "cuda::getCurrentCUDAStream" : "hip::getCurrentHIPStreamMasqueradingAsCUDA"
+        "cuda::getCurrentCUDAStream" : "hip::getCurrentHIPStreamMasqueradingAsCUDA",
+        "__nv_fp4_e2m1" : "__hip_fp4_e2m1",
+        "__nv_fp4x2_e2m1" : "__hip_fp4x2_e2m1",
+        "__nv_fp4x4_e2m1" : "__hip_fp4x4_e2m1",
+        "__nv_fp4x2_storage_t" : "__hip_fp4x2_storage_t",
+        "#include <cudaTypedefs.h>" : "",
+        "#include <cuda/barrier>" : ""
  }
 }
 

tests/cpp/operator/CMakeLists.txt

@@ -14,6 +14,7 @@ list(APPEND test_cuda_sources
             test_qdq.cu
             test_cast_mxfp8.cu
             test_dequantize_mxfp8.cu
+            test_cast_nvfp4_transpose.cu
             test_transpose.cu
             test_cast_transpose.cu
             test_cast_transpose_current_scaling.cu
@@ -32,7 +33,6 @@ list(APPEND test_cuda_sources
 	          ../test_common.cu)
 if(USE_CUDA)
   list(APPEND test_cuda_sources
-              test_cast_nvfp4_transpose.cu
               test_cast_float8blockwise.cu
               test_swizzle.cu)
 else()

tests/cpp/operator/test_cast_nvfp4_transpose.cu

@@ -1,4 +1,6 @@
 /*************************************************************************
+ * This file was modified for portability to AMDGPU
+ * Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
  * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
  *
  * See LICENSE for license information.
@@ -30,14 +32,29 @@ enum ActivationType {
     SReLU
 };
 
+#ifdef __HIP_PLATFORM_AMD__
+static constexpr float E2M1_LUT[16] = {
+     0.0f,  0.5f,  1.0f,  1.5f,  2.0f,  3.0f,  4.0f,  6.0f,
+    -0.0f, -0.5f, -1.0f, -1.5f, -2.0f, -3.0f, -4.0f, -6.0f,
+};
+#endif
+
 double2 cvt_fp4x2_to_double2(fp4e2m1x2 fp4_pair) {
+#ifdef __HIP_PLATFORM_AMD__
+    uint8_t raw = *reinterpret_cast<uint8_t*>(&fp4_pair);
+    // Decode manually
+    float lo = E2M1_LUT[raw & 0xF];
+    float hi = E2M1_LUT[(raw >> 4) & 0xF];
+    return {static_cast<double>(lo), static_cast<double>(hi)};
+#else
     const __half2_raw raw_truncated_to_fp4e2m1_pair =
         __nv_cvt_fp4x2_to_halfraw2(*reinterpret_cast<__nv_fp4x2_storage_t*>(&fp4_pair), __NV_E2M1);
 
     const __half2 truncated_to_fp4e2m1_pair(raw_truncated_to_fp4e2m1_pair);
     const double truncated_to_fp4e2m1_x = static_cast<double>(truncated_to_fp4e2m1_pair.x);
     const double truncated_to_fp4e2m1_y = static_cast<double>(truncated_to_fp4e2m1_pair.y);
     return {truncated_to_fp4e2m1_x, truncated_to_fp4e2m1_y};
+#endif
 }
 
 template <typename InputType>
@@ -567,7 +584,18 @@ void performTest(float (*OP)(const float),
     // Set 2nd stage NVFP4 scaling factor
     output.set_scale(amax);
 
+#ifndef __HIP_PLATFORM_AMD__
     bool use_2d_quantization = false;
+#else
+    // Test both 1D and 2D quantization paths on AMDGPU,
+    // as well as stochastic rounding.
+    hipDeviceProp_t prop;
+    hipGetDeviceProperties(&prop, 0);
+    const bool is_gfx950 = std::string(prop.gcnArchName).find("gfx950") != std::string::npos;
+    for (bool use_stochastic_rounding : (is_gfx950 ? std::vector<bool>{false, true}
+                                                   : std::vector<bool>{false})) {
+    for (bool use_2d_quantization : {false, true}) {
+#endif
 
     compute_ref<InputType>(OP,
                            input.rowwise_cpu_dptr<InputType>(),
@@ -589,7 +617,11 @@ void performTest(float (*OP)(const float),
     rng_state.rowwise_cpu_dptr<int64_t>()[0] = 123;  // rng_seed
     rng_state.rowwise_cpu_dptr<int64_t>()[1] = 321;  // rng_sequence
     rng_state.from_cpu();
+#ifdef __HIP_PLATFORM_AMD__
+    quant_config.set_stochastic_rounding(use_stochastic_rounding);
+#else
     quant_config.set_stochastic_rounding(false);
+#endif
     quant_config.set_rng_state(rng_state.data());
 
     // Set 2D quantization based on compile-time flag
@@ -631,15 +663,29 @@ void performTest(float (*OP)(const float),
     const fp8e4m3* ref_scales_t_ptr = ref_scales_t.get();
 
     size_t scale_mismatches_num = 0;
+#ifdef __HIP_PLATFORM_AMD__
+    std::vector<size_t> mismatches_scales_indices;
+#endif
+
     compare_scaling_factors<fp8e4m3>("scales", output.rowwise_cpu_scale_inv_ptr<fp8e4m3>(),
                                       ref_scales.get(),
                                       unpadded_blocks_Y, unpadded_blocks_X, scales_stride,
+#ifdef __HIP_PLATFORM_AMD__
+                                      mismatches_scales_indices,
+#endif
                                       scale_mismatches_num);
 
     compare_scaling_factors<fp8e4m3>("scales_t", output.columnwise_cpu_scale_inv_ptr<fp8e4m3>(),
                                       ref_scales_t.get(),
                                       unpadded_blocks_Y_t, unpadded_blocks_X_t, scales_stride_t,
+#ifdef __HIP_PLATFORM_AMD__
+                                      mismatches_scales_indices,
+#endif
                                       scale_mismatches_num);
+#ifdef __HIP_PLATFORM_AMD__
+    }
+    }
+#endif
 }
 
 std::vector<std::vector<size_t>> tensor_dims = {
@@ -674,10 +720,12 @@ class FusedCastTransposeNVFP4TestSuite : public ::testing::TestWithParam
                 transformer_engine::DType>> {};
 
 TEST_P(FusedCastTransposeNVFP4TestSuite, TestFusedCastTransposeNVFP4) {
+#ifndef __HIP_PLATFORM_AMD__
     // Skip tests for pre-Blackwell architectures
     if (getDeviceComputeCapability() < blackwellComputeCapability) {
         GTEST_SKIP();
     }
+#endif
 
     using namespace transformer_engine;
     using namespace test;

tests/cpp/test_common.cu

@@ -181,12 +181,21 @@ std::pair<scale_inv_meta, scale_inv_meta> get_scales(const NVTEShape& shape,
 
     scale_inv_meta ret_rowwise, ret_colwise;
 
-    size_t scale_dim_Y = DIVUP_TO_MULTIPLE(first_dim, scale_tensor_alignment_Y_rowwise);
-    size_t scale_dim_X = DIVUP_TO_MULTIPLE(DIVUP(last_dim, 16lu), scale_tensor_alignment_X_rowwise);
+#ifdef __HIP_PLATFORM_AMD__
+    // NVFP4 requires [128,4] padding on AMD regardless of MXFP8 alignment constants
+    constexpr size_t nvfp4_align_Y = 128;
+    constexpr size_t nvfp4_align_X = 4;
+#else
+    constexpr size_t nvfp4_align_Y = scale_tensor_alignment_Y_rowwise;
+    constexpr size_t nvfp4_align_X = scale_tensor_alignment_X_rowwise;
+#endif
+
+    size_t scale_dim_Y = DIVUP_TO_MULTIPLE(first_dim, nvfp4_align_Y);
+    size_t scale_dim_X = DIVUP_TO_MULTIPLE(DIVUP(last_dim, 16lu), nvfp4_align_X);
     ret_rowwise.shape = {scale_dim_Y, scale_dim_X};
 
-    size_t scale_dim_Y_t = DIVUP_TO_MULTIPLE(last_dim, scale_tensor_alignment_Y_rowwise);
-    size_t scale_dim_X_t = DIVUP_TO_MULTIPLE(DIVUP(first_dim, 16lu), scale_tensor_alignment_X_rowwise);
+    size_t scale_dim_Y_t = DIVUP_TO_MULTIPLE(last_dim, nvfp4_align_Y);
+    size_t scale_dim_X_t = DIVUP_TO_MULTIPLE(DIVUP(first_dim, 16lu), nvfp4_align_X);
     ret_colwise.shape = {scale_dim_Y_t, scale_dim_X_t};
 
     ret_rowwise.type = DType::kFloat8E4M3;
@@ -1219,12 +1228,21 @@ std::array<size_t, 4> get_scale_tensor_dims(const size_t rows,
     const bool is_rowwise = (block_size_rows == 1)
                             && ((block_size_cols == 32) || (block_size_cols == 16));
 
+#ifdef __HIP_PLATFORM_AMD__
+    // On AMD, MXFP8 scales (block_size=32) are allocated unpadded to match
+    // TE's internal allocation (which avoids padding for hipBLASlt compatibility).
+    // NVFP4 scales (block_size=16) still require [128,4] padding for kernel indexing.
+    const bool needs_padding = (block_size_cols == 16 || block_size_rows == 16);
+    const size_t alignment_Y = needs_padding ? (is_rowwise ? nvfp4_scale_tensor_alignment_Y_rowwise : nvfp4_scale_tensor_alignment_Y_colwise) : 1;
+    const size_t alignment_X = needs_padding ? (is_rowwise ? nvfp4_scale_tensor_alignment_X_rowwise : nvfp4_scale_tensor_alignment_X_colwise) : 1;
+#else
     const size_t alignment_Y = is_rowwise
                                ? scale_tensor_alignment_Y_rowwise
                                : scale_tensor_alignment_Y_colwise;
     const size_t alignment_X = is_rowwise
                                ? scale_tensor_alignment_X_rowwise
                                : scale_tensor_alignment_X_colwise;
+#endif
 
     const size_t unpadded_blocks_Y = divide_round_up(rows, block_size_rows);
     const size_t unpadded_blocks_X = divide_round_up(cols, block_size_cols);

tests/cpp/test_common.h

@@ -12,19 +12,19 @@
 #include <vector>
 #include <array>
 #include <random>
-
+#include <cudaTypedefs.h>
 #ifndef USE_ROCM
 #define FP4_TYPE_SUPPORTED (CUDA_VERSION >= 12080)
+#else
+#define FP4_TYPE_SUPPORTED (true)
+#endif
+
 #include <cuda_bf16.h>
+#include <cuda_fp16.h>
 #include <cuda_fp8.h>
 #if FP4_TYPE_SUPPORTED
 #include <cuda_fp4.h>
 #endif
-#else
-#define FP4_TYPE_SUPPORTED (false)
-#include <hip/hip_bfloat16.h>
-#include "amd_detail/hip_float8.h"
-#endif
 #include <cuda_runtime_api.h>
 
 #include <transformer_engine/transformer_engine.h>
@@ -339,10 +339,17 @@ constexpr uint32_t FP32_MANTISSA_BITS = 23;
 
 // [128,4] rowwise and [4,128] colwise alignment requirement
 #ifdef __HIP_PLATFORM_AMD__
+// For mxfp8:
 constexpr size_t scale_tensor_alignment_X_rowwise = 1;
 constexpr size_t scale_tensor_alignment_Y_rowwise = 1;
 constexpr size_t scale_tensor_alignment_X_colwise = 1;
 constexpr size_t scale_tensor_alignment_Y_colwise = 1;
+
+// For nvfp4:
+constexpr size_t nvfp4_scale_tensor_alignment_Y_rowwise = 128;
+constexpr size_t nvfp4_scale_tensor_alignment_X_rowwise = 4;
+constexpr size_t nvfp4_scale_tensor_alignment_Y_colwise = 4;
+constexpr size_t nvfp4_scale_tensor_alignment_X_colwise = 128;
 #else
 constexpr size_t scale_tensor_alignment_Y_rowwise = 128;
 constexpr size_t scale_tensor_alignment_X_rowwise = 4;

transformer_engine/common/CMakeLists.txt

@@ -222,7 +222,8 @@ list(APPEND transformer_engine_cuda_arch_specific_sources
      cast/cast.cu
      activation/gelu.cu
      activation/relu.cu
-     activation/swiglu.cu)
+     activation/swiglu.cu
+     transpose/quantize_transpose_vector_blockwise_fp4.cu)
 
 if(USE_CUDA)
 #NV specific source codes
@@ -246,7 +247,6 @@ if(USE_CUDA)
     list(APPEND transformer_engine_cuda_arch_specific_sources
         gemm/cutlass_grouped_gemm.cu
         transpose/quantize_transpose_square_blockwise.cu
-        transpose/quantize_transpose_vector_blockwise_fp4.cu
         hadamard_transform/hadamard_transform.cu
         hadamard_transform/hadamard_transform_cast_fusion.cu)
 else()

transformer_engine/common/cast/dispatch/quantize.cuh

@@ -91,7 +91,6 @@ void quantize_fwd_helper(const NVTETensor input, NVTETensor output,
           dummy_workspace_tensor, stream);
       break;
     }
-#ifndef __HIP_PLATFORM_AMD__
     case NVTE_NVFP4_1D_SCALING: {
       NVTE_CHECK(!IS_ACT, "IS_ACT is not supported by FWD NVTE_NVFP4_1D_SCALING");
 
@@ -108,6 +107,7 @@ void quantize_fwd_helper(const NVTETensor input, NVTETensor output,
                                   (cols % 32 == 0) && output_tensor->has_data();
 
       // Launch NVFP4 quantize kernel
+#ifndef __HIP_PLATFORM_AMD__
       if (use_optimized_kernel) {
         if (quant_config_cpp.nvfp4_2d_quantization) {
           nvfp4::quantize_transpose</*use_2d_quantization=*/true>(
@@ -117,10 +117,22 @@ void quantize_fwd_helper(const NVTETensor input, NVTETensor output,
               *input_tensor, noop_tensor, output_tensor, &quant_config_cpp, stream);
         }
       } else {
+#endif
         auto &global_amax = (output_tensor->amax.dptr != nullptr) ? output_tensor->amax
                                                                   : output_tensor->columnwise_amax;
+#ifdef __HIP_PLATFORM_AMD__
+        // Fix for upstream bug: if amax was not explicitly set, fall back to the
+        // scale field which holds the same value when set via set_scale().
+        NVTE_CHECK(global_amax.dptr != nullptr || output_tensor->scale.dptr != nullptr,
+                  "NVFP4 quantization requires global_amax (output_tensor->amax) "
+                  "or scale to be set. Call output.set_scale(amax_value) before quantizing.");
+        const SimpleTensor& effective_amax =
+            (global_amax.dptr != nullptr) ? global_amax : output_tensor->scale;
         quantize_transpose_vector_blockwise_fp4(
+            /*input=*/input_tensor->data, /*global_amax=*/effective_amax,
+#else
             /*input=*/input_tensor->data, /*global_amax=*/global_amax,
+#endif
             /*scale_inv=*/output_tensor->scale_inv,
             /*scale_inv_t=*/output_tensor->columnwise_scale_inv,
             /*output=*/output_tensor->data, /*output_t=*/output_tensor->columnwise_data,
@@ -131,9 +143,12 @@ void quantize_fwd_helper(const NVTETensor input, NVTETensor output,
             /*rng_state=*/quant_config_cpp.rng_state,
             /*use_2d_quantization=*/quant_config_cpp.nvfp4_2d_quantization,
             /*noop_tensor=*/noop_tensor->data, /*stream=*/stream);
+#ifndef __HIP_PLATFORM_AMD__
       }
+#endif
       break;
     }
+#ifndef __HIP_PLATFORM_AMD__
     case NVTE_BLOCK_SCALING_2D: {
       // TODO(kwyss): IS_ACT, ParamOP, OP parameters support.
       NVTE_CHECK(!IS_ACT, "IS_ACT is not implemented for FWD NVTE_BLOCK_SCALING_2D");
@@ -238,7 +253,6 @@ void quantize_bwd_helper(const NVTETensor grad, const NVTETensor input, NVTETens
           stream);
       break;
     }
-#ifndef __HIP_PLATFORM_AMD__
     case NVTE_NVFP4_1D_SCALING: {
       NVTE_CHECK((!IS_DBIAS && !IS_DACT),
                  "IS_DBIAS and IS_DACT are not supported by BWD NVTE_NVFP4_1D_SCALING");
@@ -256,6 +270,7 @@ void quantize_bwd_helper(const NVTETensor grad, const NVTETensor input, NVTETens
                                   (cols % 32 == 0) && output_tensor->has_data();
 
       // Launch NVFP4 quantize kernel
+#ifndef __HIP_PLATFORM_AMD__
       if (use_optimized_kernel) {
         if (quant_config_cpp.nvfp4_2d_quantization) {
           nvfp4::quantize_transpose</*use_2d_quantization=*/true>(
@@ -265,10 +280,22 @@ void quantize_bwd_helper(const NVTETensor grad, const NVTETensor input, NVTETens
               *grad_tensor, noop_tensor, output_tensor, &quant_config_cpp, stream);
         }
       } else {
+#endif
         auto &global_amax = (output_tensor->amax.dptr != nullptr) ? output_tensor->amax
                                                                   : output_tensor->columnwise_amax;
+#ifdef __HIP_PLATFORM_AMD__
+        // Fix for upstream bug: if amax was not explicitly set, fall back to the
+        // scale field which holds the same value when set via set_scale().
+        NVTE_CHECK(global_amax.dptr != nullptr || output_tensor->scale.dptr != nullptr,
+                  "NVFP4 quantization requires global_amax (output_tensor->amax) "
+                  "or scale to be set. Call output.set_scale(amax_value) before quantizing.");
+        const SimpleTensor& effective_amax =
+            (global_amax.dptr != nullptr) ? global_amax : output_tensor->scale;
         quantize_transpose_vector_blockwise_fp4(
-            /*input=*/grad_tensor->data, /*global_amax=*/global_amax,
+            /*input=*/input_tensor->data, /*global_amax=*/effective_amax,
+#else
+            /*input=*/input_tensor->data, /*global_amax=*/global_amax,
+#endif
             /*scale_inv=*/output_tensor->scale_inv,
             /*scale_inv_t=*/output_tensor->columnwise_scale_inv,
             /*output=*/output_tensor->data, /*output_t=*/output_tensor->columnwise_data,
@@ -279,9 +306,12 @@ void quantize_bwd_helper(const NVTETensor grad, const NVTETensor input, NVTETens
             /*rng_state=*/quant_config_cpp.rng_state,
             /*use_2d_quantization=*/quant_config_cpp.nvfp4_2d_quantization,
             /*noop_tensor=*/noop_tensor->data, /*stream=*/stream);
+#ifndef __HIP_PLATFORM_AMD__
       }
+#endif
       break;
     }
+#ifndef __HIP_PLATFORM_AMD__
     case NVTE_BLOCK_SCALING_2D: {
       // TODO(kwyss): IS_BIAS, IS_DACT, ParamOP, OP parameters support.
       NVTE_CHECK((!IS_DBIAS && !IS_DACT),

transformer_engine/common/common.h

@@ -13,7 +13,11 @@
 #include <cudaTypedefs.h>
 #define FP4_TYPE_SUPPORTED (CUDA_VERSION >= 12080)
 #else
+#ifdef __HIPCC__
+#define FP4_TYPE_SUPPORTED true
+#else
 #define FP4_TYPE_SUPPORTED false
+#endif
 #endif //#ifndef __HIP_PLATFORM_AMD__
 
 #include <cuda_bf16.h>
@@ -361,6 +365,11 @@ using fp4e2m1x4 = __nv_fp4x4_e2m1;
 using bf16 = hip_bfloat16;
 using fp8e4m3 = te_hip_fp8_e4m3;
 using fp8e5m2 = te_hip_fp8_e5m2;
+#if FP4_TYPE_SUPPORTED
+using fp4e2m1 = __hip_fp4_e2m1;
+using fp4e2m1x2 = __hip_fp4x2_e2m1;
+using fp4e2m1x4 = __hip_fp4x4_e2m1;
+#endif //FP4_TYPE_SUPPORTED
 #endif //__HIP_PLATFORM_AMD__
 
 using e8m0_t = uint8_t;
@@ -384,6 +393,9 @@ TRANSFORMER_ENGINE_TYPE_NAME(half)
 TRANSFORMER_ENGINE_TYPE_NAME(hip_bfloat16)
 TRANSFORMER_ENGINE_TYPE_NAME(te_hip_fp8_e4m3)
 TRANSFORMER_ENGINE_TYPE_NAME(te_hip_fp8_e5m2)
+#if FP4_TYPE_SUPPORTED
+TRANSFORMER_ENGINE_TYPE_NAME(__hip_fp4_e2m1)
+#endif
 #else
 TRANSFORMER_ENGINE_TYPE_NAME(nv_bfloat16)
 TRANSFORMER_ENGINE_TYPE_NAME(__nv_fp8_e4m3)